Instructions

Interface Definition

Initialize RMSNormalizationLayerFWD. During construction, the tensor information of the input matrix, scaling matrix, and output matrix needs to be passed. statsinfo describes the mean and variance.

RMSNormalizationLayerFWD(const TensorInfo &srcInfo, const TensorInfo &statsInfo, const TensorInfo &scaleInfo,const TensorInfo &dstInfo, NormalizationFlags flags)->void

**Table 1** Input parameters of RMSNormalizationLayerFWD
Parameter	Data Type	Description	Value Range
srcInfo	KuDNN::TensorInfo	Input matrix information.	{shape{A, ... , D}, type, layout}
statsInfo	KuDNN::TensorInfo	Mean and variance information.	{shape{A, ...}, type, layout}
scaleShiftInfo	KuDNN::TensorInfo	Scaling matrix information.	{shape{D}, type, layout}
dstInfo	KuDNN::TensorInfo	Output matrix information.	{shape{A, ... , D}, type, layout}
flags	KuDNN::NormalizationFlags	Enumeration of normalization modes.	The default value is NONE. The options are as follows: USE_GLOBAL_STATS: uses the input mean and var for calculation. USE_SCALE: scales the normalized result. USE_SHIFT: shifts the normalized result.

Assume that the input data shape is {A, B, C, D}. Normalization, scaling, and shifting operations are performed in the last dimension, that is, dimension D.

Run the operator. src and dst are the input and output pointers. scale and shift are the scaling factor and shift value. mean and variance are the pointers to the mean and variance. saveStats indicates whether to save the calculated mean and variance. eps is used to avoid the division-by-zero error.

Run(const void *src, void *dst, const void *scale, float *variance,bool saveStats, const float eps) ->void

**Table 2** Input parameters of Run
Parameter	Data Type	Description	Value Range
src	void*	Input pointer	-
dst	void*	Output pointer	-
scale	void*	Scaling pointer	-
shift	void*	Shift pointer	-
mean	float*	Mean pointer	-
variance	float*	Variance pointer	-
saveStats	bool	Whether to save the calculated variance	Boolean
eps	float	Used to avoid the division-by-zero error	Floating-point number

ValidateInput is used to validate the RMSNormalizationLayerFWD input parameters and is automatically triggered during operator construction.

ValidateInput(const TensorInfo &srcInfo, const TensorInfo &statsInfo,const TensorInfo &scaleInfo, const TensorInfo &dstInfo,NormalizationFlags flags) ->KuDNN::Status

**Table 3** Input parameters of ValidateInput
Parameter	Data Type	Description	Value Range
srcInfo	KuDNN::TensorInfo	Input matrix information.	{shape{A, ... , D}, type, layout}
statsInfo	KuDNN::TensorInfo	Mean and variance information.	{shape{A, ...}, type, layout}
scaleShiftInfo	KuDNN::TensorInfo	Scaling matrix information.	{shape{D}, type, layout}
dstInfo	KuDNN::TensorInfo	Output matrix information.	{shape{A, ... , D}, type, layout}
flags	KuDNN::NormalizationFlags	Enumeration of normalization modes.	The default value is NONE. The options are as follows: USE_GLOBAL_STATS: uses the input mean and var for calculation. USE_SCALE: scales the normalized result. USE_SHIFT: shifts the normalized result.

Supported Data Types

RMSNorm supports the FP16 data type. (The Shape, Type, and Layout parameters need to be passed during the initialization of the TensorInfo object. The following lists the data types supported by Type.)
Table 4 Type supported during TensorInfo object initialization
srcInfo

statInfo

scaleShiftInfo

KuDNN::Element::Type::F16(fp16)

KuDNN::Element::Type::F16(fp16)

KuDNN::Element::Type::F16(fp16)

**Table 4** Type supported during TensorInfo object initialization
srcInfo	statInfo	scaleShiftInfo
KuDNN::Element::Type::F16(fp16)	KuDNN::Element::Type::F16(fp16)	KuDNN::Element::Type::F16(fp16)

A maximum of 5 dimensions are supported. The supported sequential data layouts include a, ab, abc, abcd, and abcde.
They correspond to KuDNN::Layout::A, KuDNN::Layout::AB, KuDNN::Layout::ABC, KuDNN::Layout::ABCD, and KuDNN::Layout::ABCDE, respectively.

Table 5 Layout supported during TensorInfo object initialization
Dimension

srcInfo Data Layout

dstInfo Data Layout

2D

ab

ab

3D

abc

abc

4D

abcd

abcd

5D

abcde

abcde

**Table 5** Layout supported during TensorInfo object initialization
Dimension	srcInfo Data Layout	dstInfo Data Layout
2D	ab	ab
3D	abc	abc
4D	abcd	abcd
5D	abcde	abcde

Examples

Perform root mean square normalization using FP16. Data layout configurations are: KuDNN::Layout::AB for srcInfo data, KuDNN::Layout::A for statInfo data, KuDNN::Layout::A for scaleShiftInfo data, and KuDNN::Layout::AB for dstInfo data.

using SizeType = KuDNN::SizeType;
using Shape = KuDNN::Shape;
using Type KuDNN::Element::TypeT
Shape shape(100, 100);
// Define tensor information.
TensorInfo srcInfo = {shape, Type::F16, KuDNN::Layout::AB};
TensorInfo statInfo = {{shape[0]}, Type::F16, KuDNN::Layout::A};
TensorInfo scaleInfo = {{shape[1]}, Type::F16, KuDNN::Layout::A};
TensorInfo dstInfo = {shape, Type::F16, KuDNN::Layout::AB};
KuDNN::NormalizationFlags flags = KuDNN::NormalizationFlags::NONE;
// Construct the operator: KuDNN::RMSNormalizationLayerFWD rmsLayer1(srcInfo, statInfo, scaleInfo, dstInfo, flags);
// Initialize matrix data.
SizeType srcSize = srcInfo.GetTotalTensorSize();
SizeType dstSize = dstInfo.GetTotalTensorSize();
SizeType statSize = statInfo.GetTotalTensorSize();
SizeType innerSize = scaleInfo.GetTotalTensorSize();
__fp16 *src = (__fp16 *)malloc(srcSize * sizeof(__fp16));
__fp16 *dst = (__fp16 *)malloc(dstSize * sizeof(__fp16));
__fp16 *dstRef = (__fp16 *)malloc(dstSize * sizeof(__fp16));
float *variance = (float *)malloc(statSize * sizeof(float));
__fp16 *scale = (__fp16 *)malloc(innerSize * sizeof(__fp16));
float eps = 1e-5;
// Run the operator.
rmsLayer1.Run(src, dst, scale, variance, true, eps);

Parent topic: RMSNorm