NormRopeReshapeOperation

Applicable Products

Hardware Model	Supported or Not
Atlas 350 accelerator card	x
Atlas A3 inference products/Atlas A3 training products	√
Atlas A2 training products/Atlas A2 inference products	√
Atlas training products	x
Atlas inference products	x
Atlas 200I/500 A2 inference products	x

Description

Performs the RmsNorm operation on the input tensor x, performs the Rope operation on keyrope, combines the outputs of the two operations, and performs the ReshapeAndCache operation.

Figure 1 Computation process

Definition

struct NormRopeReshapeParam {
    uint32_t precisionMode = 0;
    uint32_t rotaryCoeff = 2;
    float epsilon = 1e-5;
    uint8_t rsv[16] = {0};
};

Parameters

Member	Type	Default Value	Value Range	Mandatory or Not	Description
precisionMode	uint32_t	0	-	No	Precision mode. Currently, only 0 is supported.
epsilon	float	1e-5	-	No	It is added to the denominator during normalization to prevent division by 0.
rotaryCoeff	uint32_t	2	-	No	Rotation coefficient calculated by the Rope part in the operator.
rsv	uint8_t	uint8_t	-	No	Reserved

Input

Parameter	Dimension	Data Type	Format	Description
X	[ntokens, 1, head_size_x]	float16	ND	Currently, only [ntokens, 1, 512] is supported.
gamma	[head_size_x]	float16	ND	Currently, only [512] is supported.
keyRope	[ntokens, hiddenSizeK]	float16	ND	[ntokens, hiddenSizeK] hiddenSizeK = headNum × head_size (headNum = 1). Currently, only [ntokens, 64] is supported.
cos	[ntokens, head_size]	float16	ND	[ntokens, head_size]. The size is the same as hiddenSizeK (headNum = 1). Currently, only [ntokens, 64] is supported.
sin	[ntokens, head_size]	float16	ND	Same as that of cos.
slotMapping	[dnslot]	int32	ND	The element value cannot exceed blockNum multiplied by blockSize. Currently, only [64] is supported.
keycachein	[blockNum, blockSize, 1,dnrac]	float16	ND	dnrac = head_size_x + hiddenSizeK. Currently, only [blockNum, blockSize, 1, 576] is supported. blockNum × blockSize ≥ ntokens.

Output

Parameter	Dimension	Data Type	Format	Description
keycacheout	[blockNum, blockSize, 1, dnrac]	Same as that of keycachein	ND	Same as that of keycachein.

Constraints

The size of the last dimension of x or gamma is the same and 32-byte aligned.
The data types of keyRope, cos, and sin must be the same, and the values of hiddenSizeK and head_size must be the same.
11 × dn × sizeof(float16) + dnrac × sizeof(float16) < maxUbSize(196352).
The value of ntokens must be less than or equal to 65536.

Parent topic: atb/infer_op_params.h