GatherPreRmsNormOperation

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

Performs the Gather index operation on ResIn, adds the result to X, and then performs the RmsNorm calculation.

Gather operation: Index the corresponding row in ResIn based on the value in indices to obtain ReslnAfterGather.
Add operation: Add X and ResInAfterGather to obtain ResOut.
Norm operation:
$\text{[math]}$

$\text{[math]}$

$\text{[math]}$

struct GatherPreRmsNormParam {
    float epsilon = 1e-5;
    uint8_t rsv[28] = {0};
};

Member	Type	Default Value	Value Range	Mandatory or Not	Description
epsilon	float	1e-5	(2e-38, 1e-5]	Yes	It is added to the denominator during normalization to prevent division by 0.
rsv[28]	uint8_t	{0}	[0]	No	Reserved

Parameter	Dimension	Data Type	Format	Description
x	[dim_m, dim_n]	float16/bf16	ND	2D. The last axis dimension must be 32-byte aligned (a multiple of 16) and cannot be greater than 7680.
resIn	[dim_res, dim_n]	float16/bf16	ND	2D. The data type is the same as that of x. The values of dim_m and dim_res may be different.
indices	[dim_m]	int32	ND	1D. The element value range is [0, dim_res), where dim_res is the first dimension value of resIn.
gamma	[1, dim_n] or [dim_n]	float16/bf16	ND	1D or 2D. The data type is the same as that of x.

Parameter	Dimension	Data Type	Format	Description
y	[dim_m, dim_n]	float	ND	The data type, number of dimensions, and dimension value are the same as those of x.
ResOut	[dim_m, dim_n]	float16/bf16	ND	The dimension is the same as that of x.

Parent topic: atb/infer_op_params.h