RmsNorm

Applicability

Product	Supported
Atlas 350 Accelerator Card	√
Atlas A3 training product/Atlas A3 inference product	√
Atlas A2 training product/Atlas A2 inference product	√
Atlas 200I/500 A2 inference product	x
Atlas inference product AI Core	√
Atlas inference product Vector Core	x
Atlas training product	x

Function Usage

Normalizes input data whose shape is [B, S, H] using RmsNorm. Below is the formula.

$\text{[math]}$

γ is the scale coefficient, and ε is the weight coefficient for preventing division by zero.

Prototype

Pass the temporary space through the sharedTmpBuffer input parameter.

template <typename T, bool isBasicBlock = false>
__aicore__ inline void RmsNorm(const LocalTensor<T>& dstLocal, const LocalTensor<T>& srcLocal, const LocalTensor<T>& gammaLocal, const LocalTensor<uint8_t>& sharedTmpBuffer, const T epsilon, const RmsNormTiling& tiling)

Allocate the temporary space through the API framework.

template <typename T, bool isBasicBlock = false>
__aicore__ inline void RmsNorm(const LocalTensor<T>& dstLocal, const LocalTensor<T>& srcLocal, const LocalTensor<T>& gammaLocal, const T epsilon, const RmsNormTiling& tiling)

Due to the complex computation involved in the internal implementation of this API, additional temporary space is required to store intermediate variables generated during computation. The temporary space can be allocated through the API framework or passed by developers through the sharedTmpBuffer input parameter.

When the API framework is used for temporary space allocation, developers do not need to allocate the space, but must reserve the required size for the temporary space.

When the sharedTmpBuffer input parameter is used for passing the temporary space, the tensor serves as the temporary space. In this case, the API framework is not required for temporary space allocation. This enables developers to manage the sharedTmpBuffer space and reuse the buffer after calling the API, so that the buffer is not repeatedly allocated or deallocated, improving the flexibility and buffer utilization.

If the API framework is used, developers must reserve the temporary space. If sharedTmpBuffer is used, developers must allocate space for the tensor. The method of obtaining the temporary space size (BufferSize) is as follows: Obtain the required maximum and minimum temporary space sizes using the GetRmsNormMaxMinTmpSize API provided in RmsNorm Tiling. The minimum space can ensure correct functionality, while the maximum space is used to improve performance.

Parameters

**Table 1** Template parameters
Parameter	Description
T	Data type of the operand. For the Atlas 350 Accelerator Card, the supported data types are half and float. For the Atlas A3 training product/Atlas A3 inference product, the supported data types are half and float. For the Atlas A2 training product/Atlas A2 inference product, the supported data types are half and float. For the Atlas inference product AI Core, the supported data types are half and float.
isBasicBlock	If the shape information and the tiling policy of both srcTensor and dstTensor meet the basic block requirements, this parameter can be enabled to improve performance. By default, this parameter is disabled. For basic blocks, the shapes of srcTensor and dstTensor must meet the following requirements: The length of the last axis (H axis) is a multiple of 64 but less than 2048. The length (*BS**) of a non-last axis is a multiple of 8.

**Table 2** API parameters
Parameter	Input/Output	Description
dstLocal	Output	Destination operand. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT. The shape of dstLocal must be the same as that of the source operand srcLocal.
srcLocal	Input	Source operand. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT. If the shape is [B, S, H], the length of the last axis (H axis) must be 32-byte aligned.
gammaLocal	Input	Scaling coefficient. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT. The length of its shape must be the same as that of the last axis (H axis) of srcLocal and dstLocal, that is, the shape is [H].
sharedTmpBuffer	Input	Temporary space. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT. This parameter is used to store intermediate variables during complex internal API computation and is provided by developers. For details about how to obtain the temporary space size (BufferSize), see RmsNorm Tiling.
epsilon	Input	Weight coefficient for preventing division by zero. The data type must be the same as that of srcLocal/dstLocal.
tiling	Input	Tiling information required for RmsNorm computation. For details about how to obtain the tiling information, see RmsNorm Tiling.

Returns

None

Constraints

The tensor space of dstLocal and gammaLocal cannot be reused.
Currently, only the ND format is supported.
For details about the operand address alignment requirements, see General Address Alignment Restrictions.
If the H axis in the original shape of srcLocal is not 32-byte aligned, you need to pad the original input on the H axis to ensure 32-byte alignment. The computation result of the API will overwrite the data in dstLocal corresponding to the padded position of srcLocal.

Examples

For a complete call example, see RmsNorm operator sample.

// dstLocal: tensor for storing the RmsNorm computation result
// srcLocal: input tensor involved in computation
// gammaLocal: input tensor, scaling coefficient γ of the normalized data
// epsilon: weight coefficient ε for preventing division by zero
// tiling: tiling data, obtained from the host

// // If the length of the last axis (H) does not exceed 2040 and is a multiple of 64, and the length of the non-last axes (B*S) is a multiple of 8, you can set isBasicBlock to true to improve performance.
AscendC::RmsNorm<dataType, isBasicBlock>(dstLocal, srcLocal, gammaLocal, epsilon, tiling);

The following is an example:

Input (srcLocal, shape: [1, 1, 16]):
[ 0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 ]
Input (gammaLocal, shape: [16]):
[ 0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 ]
Output (dstLocal):
[ 0.          0.11359233  0.4543693   1.022331    1.8174772   2.8398082   4.089324    5.566024    
  7.269909    9.200979   11.359233   13.744672   16.357296   19.197104   22.264095   25.558275 ]

Parent topic: Normalization