RmsNormWithStrideOperation

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

To enable RmsNorm to support non-contiguous tensors, the x tensor is represented by (shape, strides, offset). strides indicates the interval between two adjacent elements in a dimension, and offset indicates the offset of the first element of the tensor relative to the address.

Example:

For a tensor with shape=(6, 5), strides=(10, 1), and offset=22, its memory layout is as follows.

The boldfaced values represent the memory layout of the tensor. The complete tensor is discontiguous in memory layout. strides indicates the interval between two adjacent elements in the tensor dimension. If the stride of dimension 1 is 1, the dimension is contiguous. If the stride of dimension 0 is 10, adjacent elements are separated by 10 elements, that is, they are discontiguous elements. offset indicates the offset of the first element of the tensor relative to the address.

Formula

$\text{[math]}$

gamma is a trainable parameter, and epsilon is a floating-point value whose value range is (0, 1], and is used to avoid division by zero.

Definition

struct RmsNormWithStrideParam {
    enum RmsNormType : int {
        RMS_NORM_UNDEFINED = 0, 
        RMS_NORM_NORM,  
        RMS_NORM_PRENORM,   
        RMS_NORM_POSTNORM,      
    };
    enum PrecisionMode : int {
        HIGH_PRECISION_MODE = 0, 
        HIGH_PERFORMANCE_MODE,  
    };
    enum ModelType : int {
        LLAMA_MODEL = 0, 
        GEMMA_MODEL,    
    };
    struct NormParam {
        QuantType quantType = QUANT_UNQUANT;
        float epsilon = 1e-5;
        double layerNormEps = 1e-5;
        bool rstd = false;
        PrecisionMode precisionMode = HIGH_PRECISION_MODE;
        ModelType modelType = LLAMA_MODEL;
        DynamicQuantType dynamicQuantType = DYNAMIC_QUANT_UNDEFINED;
        uint8_t rsv[32] = {0};
    };
    struct PreNormParam {
        QuantType quantType = QUANT_UNQUANT;
        float epsilon = 1e-5;
        bool hasBias = false;
        uint8_t rsv[23] = {0};
    };
    struct PostNormParam {
        QuantType quantType = QUANT_UNQUANT;
        float epsilon = 1e-5;
        bool hasBias = false;
        uint8_t rsv[23] = {0};
    };
    RmsNormType layerType = RMS_NORM_UNDEFINED;
    NormParam normParam;
    PreNormParam preNormParam;
    PostNormParam postNormParam;
    uint8_t rsv[8] = {0};
};

Parameters

Member	Description
layerType	Normalization type. RMS_NORM_UNDEFINED (default): undefined. RMS_NORM_NORM: NORM parameter. RMS_NORM_PRENORM: PRENORM parameter RMS_NORM_POSTNORM: POSTNORM parameter.
normParam	NORM parameter. For details, see Table 1.
preNormParam	PRENORM parameter. For details, see Table 2.
postNormParam	POSTNORM parameter. For details, see Table 3.
rsv[8]	Reserved

**Table 1** normParam members
Member	Type	Default Value	Value Range	Mandatory or Not	Description
quantType	QuantType	QUANT_UNQUANT	QUANT_UNQUANT	Yes	Quantization is performed after the operator computation is complete. Currently, only the type QUANT_UNQUANT is supported.
epsilon	float	1e-5	1e-5	Yes	Epsilon, which is added to the denominator during normalization to prevent division by 0.
layerNormEps	double	1e-5	1e-5	Yes	Epsilon. The default value is 1e-5. This parameter is not used currently.
rstd	bool	false	true/false	Yes	The default value is false. This parameter cannot be set together with precisionMode and modelType. rstd is not supported in quantization scenarios.
precisionMode	PrecisionMode	HIGH_PRECISION_MODE	HIGH_PRECISION_MODE HIGH_PERFORMANCE_MODE	Yes	HIGH_PRECISION_MODE (default): The float type is used for intermediate calculation. HIGH_PERFORMANCE_MODE: The float16 type is used for intermediate calculation. This parameter cannot be set together with rstd and modelType. The input type can only be float16. In a quantization scenario, precisionMode is not supported. If this parameter is set in this scenario, the error message "ERROR_INVALID_PARAM" is returned.
modelType	ModelType	LLAMA_MODEL	LLAMA_MODEL GEMMA_MODEL	Yes	When the value is set to GEMMA_MODEL, the rmsnorm calculation formula of the gemma model is used. LLAMA_MODEL (default ): rms norm calculation formula of Llama. GEMMA_MODEL: rms norm calculation formula of GEMMA. This parameter cannot be used together with rstd and precisionMode. In a quantization scenario, modelType is not supported. If this parameter is set in this scenario, the error message "ERROR_INVALID_PARAM" is returned.
dynamicQuantType	DynamicQuantType	DYNAMIC_QUANT_UNDEFINED	DYNAMIC_QUANT_UNDEFINED	No	Dynamic quantization type. The default value is DYNAMIC_QUANT_UNDEFINED. The current version does not support asymmetric dynamic quantization.
rsv[32]	uint8_t	{0}	[0]	No	Reserved

**Table 2** PreNormParam members
Member	Type	Default Value	Value Range	Mandatory or Not	Description
quantType	QuantType	QUANT_UNQUANT	QUANT_UNQUANT QUANT_INT8	Yes	Quantization type.
epsilon	float	1e-5	1e-5	Yes	Epsilon, which is added to the denominator during normalization to prevent division by 0.
hasBias	bool	false	true/false	No	Whether to add a bias. Set this parameter to true when beta needs to be input. In a quantization scenario, hasBias is not supported. If this parameter is set in this scenario, the error message "ERROR_INVALID_PARAM" is returned.
rsv[23]	uint8_t	{0}	[0]	No	Reserved

**Table 3** PostNormParam members
Member	Type	Default Value	Value Range	Mandatory or Not	Description
quantType	QuantType	QUANT_UNQUANT	QUANT_UNQUANT	Yes	Quantization type. Currently, only QUANT_UNQUANT is supported.
epsilon	float	1e-5	1e-5	Yes	Epsilon, which is added to the denominator during normalization to prevent division by 0.
hasBias	bool	false	true/false	No	Whether to add a bias. Set this parameter to true when beta needs to be input.
rsv[23]	uint8_t	{0}	[0]	No	Reserved

Input

Parameter	Dimension	Data Type	Format	Description
x	[dim_0, ..., dim_n-1, dim_n]	float16/bf16	ND	The size of the last dimension must be a multiple of 16.
gamma	[1, dim_n]/[dim_i, ..., dim_n](0<=i<=n)	float16/bf16	ND	The dimension is [1, dim_n]. The last dimension is the same as that of x.
xStrides	[x_shape]	int64	ND	Stride of x. x_shape indicates the number of dimensions of x. The information is passed as a host tensor, with all values greater than 0, and the last value being 1.
xOffset	[1]	int64	ND	Offset of x. The offset must be greater than 0. The default value is 0. The information is passed as a host tensor.

Output

Parameter	Dimension	Data Type	Format	Description
output	[dim_0, ..., dim_n-1, dim_n]	float16/bf16	ND	The dimension must be the same as that of x.

Constraints

Quantization and rstd are not supported in the NORM scenario.
Currently, the preNormParam and postNormParam members are not supported.

Parent topic: atb/infer_op_params.h