aclnnAddRmsNormDynamicQuant

Applicable Products

Product	Supported
Ascend 950PR/Ascend 950DT	√
Atlas A3 training products/Atlas A3 inference products	√
Atlas A2 training products/Atlas A2 inference products	√
Atlas 200I/500 A2 inference products	×
Atlas inference products	×
Atlas training products	×

Function

Description: The RmsNorm operator is a commonly used normalization operation in large models. Compared with the LayerNorm operator, it removes the part of subtracting the mean. The DynamicQuant operator performs symmetric dynamic quantization on the input tensor. The AddRmsNormDynamicQuant operator fuses the Add operator before RmsNorm and one or two DynamicQuant operators that receive the RmsNorm normalized output, reducing data movement in and out.
Formula:
$x=x_{1}+x_{2}$ $y = \operatorname{RmsNorm}(x)=\frac{x}{\operatorname{Rms}(\mathbf{x})}\cdot gamma, \quad \text { where } \operatorname{Rms}(\mathbf{x})=\sqrt{\frac{1}{n} \sum_{i=1}^n x_i^2+epsilon}$
- If neither smoothScale1Optional nor smoothScale2Optional is input, the outputs y2Out and scale2Out are meaningless. The computation process is as follows:
$scale1Out=row\_max(abs(y))/127$ $y1Out=round(y/scale1Out)$
- If only smoothScale1Optional is input, the outputs y2Out and scale2Out are meaningless. The computation process is as follows:
$input = y\cdot smoothScale1Optional$ $scale1Out=row\_max(abs(input))/127$ $y1Out=round(input/scale1Out)$
- If both smoothScale1Optional and smoothScale2Optional are input, all five outputs of the operator are valid outputs. The computation process is as follows:
$input1 = y\cdot smoothScale1Optional$ $input2 = y\cdot smoothScale2Optional$ $scale1Out=row\_max(abs(input1))/127$ $scale2Out=row\_max(abs(input2))/127$ $y1Out=round(input1/scale1Out)$ $y2Out=round(input2/scale2Out)$
where row_max represents the maximum value of each row.

Function Prototype

Each operator consists of a two-phase API. You must first call the aclnnAddRmsNormDynamicQuantGetWorkspaceSize API to obtain the input parameters and the required workspace size based on the computation process, and then call the aclnnAddRmsNormDynamicQuant API to perform the computation.

aclnnStatus aclnnAddRmsNormDynamicQuantGetWorkspaceSize(
  const aclTensor *x1,
  const aclTensor *x2,
  const aclTensor *gamma,
  const aclTensor *smoothScale1Optional,
  const aclTensor *smoothScale2Optional,
  double           epsilon,
  aclTensor       *y1Out,
  aclTensor       *y2Out,
  aclTensor       *xOut,
  aclTensor       *scale1Out,
  aclTensor       *scale2Out,
  uint64_t        *workspaceSize,
  aclOpExecutor  **executor)

aclnnStatus aclnnAddRmsNormDynamicQuant(
  void          *workspace,
  uint64_t       workspaceSize,
  aclOpExecutor *executor,
  aclrtStream    stream)

aclnnAddRmsNormDynamicQuantGetWorkspaceSize

Parameters

Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
x1 (aclTensor*)	Input	Represents the source data tensor in the normalization process. Corresponds to `x1` in the formula.	Supports empty tensors.	FLOAT16, BFLOAT16	ND	2-8	√
x2 (aclTensor*)	Input	Represents the source data tensor in the normalization process. Corresponds to `x2` in the formula.	Supports null Tensor. The shape and data type must be consistent with `x1`.	FLOAT16, BFLOAT16	ND	2-8	√
gamma (aclTensor*)	Input	Represents the weight tensor in the normalization process. Corresponds to `gamma` in the formula.	Supports null tensor. The data type must be consistent with `x1`. The shape must be consistent with the last dimension of `x1`.	FLOAT16, BFLOAT16	ND	1	√
smoothScale1Optional（aclTensor*）	Input	The smoothScale tensor used to obtain `y1Out` during quantization. Corresponds to `smoothScale1Optional` in the formula.	Supports empty tensor. Optional parameter; supports passing a null pointer. The shape and data type must be consistent with `gamma`.	FLOAT16, BFLOAT16	ND	1	√
smoothScale2Optional (aclTensor*)	Input	Represents the smoothScale tensor used to obtain `y2Out` during the quantization process. Corresponds to `smoothScale2Optional` in the formula.	Supports null Tensor. Optional parameter. A null pointer can be passed. Must be used together with smoothScale1Optional. The shape and data type must be consistent with `gamma`.	FLOAT16, BFLOAT16	ND	1	√
epsilon (double)	Input	Indicates the value used to prevent division-by-zero errors, corresponding to `epsilon` in the formula.	It is recommended to pass a small positive number, such as 1e-6.	-	-	-	-
y1Out (aclTensor*)	Output	Represents the quantized output tensor, corresponding to `y1Out` in the formula.	Null Tensor is supported. The shape must be consistent with the input `x1`/`x2`, or be two-dimensional with the first dimension equal to the product of all dimensions of `x1` except the last dimension, and the second dimension equal to the last dimension of `x1`.	INT8, HIFLOAT8, FLOAT8_E5M2, FLOAT8_E4M3FN	ND	2-8	√
y2Out (aclTensor*)	Output	Represents the quantized output tensor, corresponding to `y2Out` in the formula.	Supports null tensors. When smoothScale2Optional does not exist, this output is meaningless. If the `y2Out` output is meaningless, the shape is [1]; in other scenarios, the shape must be consistent with `y1Out`.	INT8, HIFLOAT8, FLOAT8_E5M2, FLOAT8_E4M3FN	ND	2-8	√
xOut (aclTensor*)	Output	Represents the sum of x1 and x2, corresponding to `x` in the formula.	Supports empty tensors. The shape and data type must be consistent with the input `x1`/`x2`.	FLOAT16, BFLOAT16	ND	2-8	√
scale1Out (aclTensor*)	Output	Represents the output of the first quantization path, corresponding to `scale1Out` in the formula.	Supports null tensors. The shape must be consistent with the shape of the input `x1` after removing the last dimension, or consistent with the product of `x1` excluding the last dimension.	FLOAT32	ND	1-8	√
scale2Out (aclTensor*)	Output	Represents the output of the second quantization path, corresponding to `scale2Out` in the formula.	Supports null tensor. When smoothScale2Optional does not exist, this output is meaningless. The shape must be consistent with `scale1Out`.	FLOAT32	ND	1-8	√
workspaceSize (uint64_t*)	Output	Returns the workspace size to be requested on the Device side.	-	-	-	-	-
executor (aclOpExecutor**)	Output	Returns the op executor, which contains the operator computation process.	-	-	-	-	-

Atlas A3 training products/Atlas A3 inference products, Atlas A2 training products/Atlas A2 inference products:
- The output parameters y1Out and y2Out only support INT8.

Return Value

aclnnStatus: return code. For details, see aclnn Return Code.

The first-phase API performs input parameter validation and returns an error in the following scenarios:

Return Code	Error Code	Description
ACLNN_ERR_PARAM_NULLPTR	161001	Returns 161001 if the passed parameter is a required input, output, or required attribute and is a null pointer.
ACLNN_ERR_PARAM_INVALID	161002	The data type of the input or output is not within the supported range.

aclnnAddRmsNormDynamicQuant

Parameters

Parameter	Input/Output	Description
workspace	Input	Address of the workspace memory applied for on the Device side.
workspaceSize	Input	Size of the workspace applied for on the Device side, obtained through the first-phase API aclnnAddRmsNormDynamicQuantGetWorkspaceSize.
executor	Input	Op executor, which contains the operator computation process.
stream	Input	Stream that specifies the task execution.

Return Value

aclnnStatus: return code. For details, see aclnn Return Code.

Constraints

Supported Type Description

Whether empty tensors are supported: empty input yields empty output.

Data Format Description

It is recommended to use the ND format for all input and output tensors. Other data formats are converted to the ND format by the framework by default.

Supported Data Types by Product Model

Atlas A3 training products/Atlas A3 inference products, Atlas A2 training products/Atlas A2 inference products:

`x1` Data Type	`x2` Data Type	`gamma` Data Type	`smoothScale1Optional` Data Type	`smoothScale2Optional` Data Type	`y1Out` Data Type	`y2Out` Data Type	`xOut` Data Type	`scale1Out` Data Type	`scale2Out` Data Type
FLOAT16	FLOAT16	FLOAT16	FLOAT16	FLOAT16	INT8	INT8	FLOAT16	FLOAT32	FLOAT32
BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	INT8	INT8	BFLOAT16	FLOAT32	FLOAT32

Ascend 950PR/Ascend 950DT:

`x1` Data Type	`x2` Data Type	`gamma` Data Type	`smoothScale1Optional` Data Type	`smoothScale2Optional` Data Type	`y1Out` Data Type	`y2Out` Data Type	`xOut` Data Type	`scale1Out` Data Type	`scale2Out` Data Type
FLOAT16	FLOAT16	FLOAT16	FLOAT16	FLOAT16	INT8	INT8	FLOAT16	FLOAT32	FLOAT32
BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	INT8	INT8	BFLOAT16	FLOAT32	FLOAT32
FLOAT16	FLOAT16	FLOAT16	FLOAT16	FLOAT16	HIFLOAT8	HIFLOAT8	FLOAT16	FLOAT32	FLOAT32
BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	HIFLOAT8	HIFLOAT8	BFLOAT16	FLOAT32	FLOAT32
FLOAT16	FLOAT16	FLOAT16	FLOAT16	FLOAT16	FLOAT8_E5M2	FLOAT8_E5M2	FLOAT16	FLOAT32	FLOAT32
BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	FLOAT8_E5M2	FLOAT8_E5M2	BFLOAT16	FLOAT32	FLOAT32
FLOAT16	FLOAT16	FLOAT16	FLOAT16	FLOAT16	FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT16	FLOAT32	FLOAT32
BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	FLOAT8_E4M3FN	FLOAT8_E4M3FN	BFLOAT16	FLOAT32	FLOAT32

Deterministic computation:
- aclnnAddRmsNormDynamicQuant defaults to a deterministic implementation.

Example

The following provides example code for reference only. For details about compilation and running, see Compile and Run Samples.

#include <iostream>
#include <vector>
#include "acl/acl.h"
#include "aclnnop/aclnn_add_rms_norm_dynamic_quant.h"

#define CHECK_RET(cond, return_expr) \
    do {                             \
        if (!(cond)) {               \
            return_expr;             \
        }                            \
    } while (0)

#define LOG_PRINT(message, ...)         \
    do {                                \
        printf(message, ##__VA_ARGS__); \
    } while (0)

int64_t GetShapeSize(const std::vector<int64_t>& shape)
{
    int64_t shape_size = 1;
    for (auto i : shape) {
        shape_size *= i;
    }
    return shape_size;
}

int Init(int32_t deviceId, aclrtStream* stream)
{
    // Boilerplate: resource initialization.
    auto ret = aclInit(nullptr);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
    ret = aclrtSetDevice(deviceId);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
    ret = aclrtCreateStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);

    return 0;
}

template <typename T>
int CreateAclTensor(
    const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr, aclDataType dataType,
    aclTensor** tensor)
{
    auto size = GetShapeSize(shape) * sizeof(T);
    // Call aclrtMalloc to allocate device-side memory.
    auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);

    // Call aclrtMemcpy to copy data from the host to the device memory.
    ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);

    // Compute strides for a contiguous tensor.
    std::vector<int64_t> strides(shape.size(), 1);
    for (int64_t i = shape.size() - 2; i >= 0; i--) {
        strides[i] = shape[i + 1] * strides[i + 1];
    }

    // Call the aclCreateTensor API to create an aclTensor.
    *tensor = aclCreateTensor(
        shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(),
        *deviceAddr);
    return 0;
}

int main()
{
    // 1. (Boilerplate) Initialize the device/stream. See the ACL API manual.
    // Set deviceId based on the actual device.
    int32_t deviceId = 0;
    aclrtStream stream;
    auto ret = Init(deviceId, &stream);
    // Handle the check as needed.
    CHECK_RET(ret == 0, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
    // 2. Construct the input and output. Customize the construction based on the API.
    std::vector<int64_t> xShape = {2, 8};
    std::vector<int64_t> gammaShape = {8};
    std::vector<int64_t> reduceShape = {2, 1};

    void* x1DeviceAddr = nullptr;
    void* x2DeviceAddr = nullptr;
    void* gammaDeviceAddr = nullptr;
    void* smooth1DeviceAddr = nullptr;
    void* smooth2DeviceAddr = nullptr;

    void* y1DeviceAddr = nullptr;
    void* y2DeviceAddr = nullptr;
    void* xDeviceAddr = nullptr;
    void* scale1DeviceAddr = nullptr;
    void* scale2DeviceAddr = nullptr;

    aclTensor* x1 = nullptr;
    aclTensor* x2 = nullptr;
    aclTensor* gamma = nullptr;
    aclTensor* smooth1 = nullptr;
    aclTensor* smooth2 = nullptr;
    aclTensor* y1 = nullptr;
    aclTensor* y2 = nullptr;
    aclTensor* x = nullptr;
    aclTensor* scale1 = nullptr;
    aclTensor* scale2 = nullptr;

    int64_t xShapeSize = GetShapeSize(xShape);
    int64_t gammaShapeSize = GetShapeSize(gammaShape);
    int64_t reduceShapeSize = GetShapeSize(reduceShape);

    std::vector<short> x1HostData(xShapeSize, 0x3800);
    std::vector<short> x2HostData(xShapeSize, 0x3800);
    std::vector<short> gammaHostData(gammaShapeSize, 0x3e00);
    std::vector<short> smooth1HostData(gammaShapeSize, 0x3e00);
    std::vector<short> smooth2HostData(gammaShapeSize, 0x3e00);

    std::vector<short> y1HostData(xShapeSize, 0);
    std::vector<short> y2HostData(xShapeSize, 0);
    std::vector<short> xHostData(xShapeSize, 0);
    std::vector<short> scale1HostData(reduceShapeSize, 0);
    std::vector<short> scale2HostData(reduceShapeSize, 0);

    float epsilon = 1e-6;

    // Create the x1 aclTensor.
    ret = CreateAclTensor(x1HostData, xShape, &x1DeviceAddr, aclDataType::ACL_FLOAT16, &x1);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the x2 aclTensor.
    ret = CreateAclTensor(x2HostData, xShape, &x2DeviceAddr, aclDataType::ACL_FLOAT16, &x2);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the gamma aclTensor.
    ret = CreateAclTensor(gammaHostData, gammaShape, &gammaDeviceAddr, aclDataType::ACL_FLOAT16, &gamma);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the smooth1 aclTensor.
    ret = CreateAclTensor(smooth1HostData, gammaShape, &smooth1DeviceAddr, aclDataType::ACL_FLOAT16, &smooth1);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the smooth2 aclTensor.
    ret = CreateAclTensor(smooth2HostData, gammaShape, &smooth2DeviceAddr, aclDataType::ACL_FLOAT16, &smooth2);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    // Create the y1 aclTensor.
    ret = CreateAclTensor(y1HostData, xShape, &y1DeviceAddr, aclDataType::ACL_INT8, &y1);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the y2 aclTensor.
    ret = CreateAclTensor(y2HostData, xShape, &y2DeviceAddr, aclDataType::ACL_INT8, &y2);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the x aclTensor.
    ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_FLOAT16, &x);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the outScale1 aclTensor.
    ret = CreateAclTensor(scale1HostData, reduceShape, &scale1DeviceAddr, aclDataType::ACL_FLOAT, &scale1);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the outScale1 aclTensor.
    ret = CreateAclTensor(scale2HostData, reduceShape, &scale2DeviceAddr, aclDataType::ACL_FLOAT, &scale2);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    // 3. Call the CANN operator library API. Replace it with the specific API.
    uint64_t workspaceSize = 0;
    aclOpExecutor* executor;
    // Call the first-phase API of aclnnAddRmsNormDynamicQuant.
    ret = aclnnAddRmsNormDynamicQuantGetWorkspaceSize(
        x1, x2, gamma, smooth1, smooth2, epsilon, y1, y2, x, scale1, scale2, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnAddRmsNormDynamicQuantGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
    // Apply for device memory based on the workspaceSize calculated by the first-phase API.
    void* workspaceAddr = nullptr;
    if (workspaceSize > 0) {
        ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret;);
    }
    // Call the second-phase API of aclnnAddRmsNormDynamicQuant.
    ret = aclnnAddRmsNormDynamicQuant(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnAddRmsNormDynamicQuant failed. ERROR: %d\n", ret); return ret);
    // 4. (Boilerplate) Synchronously wait for task execution to complete.
    ret = aclrtSynchronizeStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);

    // 5. Obtain the output values and copy the results from the device memory to the host memory. Modify based on the specific API definition.
    auto size = GetShapeSize(xShape);
    std::vector<int8_t> y1Ret(size, 0);
    ret = aclrtMemcpy(
        y1Ret.data(), y1Ret.size() * sizeof(y1Ret[0]), y1DeviceAddr, size * sizeof(int8_t), ACL_MEMCPY_DEVICE_TO_HOST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);
    for (int64_t i = 0; i < size; i++) {
        LOG_PRINT("result[%ld] is: %d\n", i, y1Ret[i]);
    }

    // 6. Release aclTensor and aclScalar. Modify based on the specific API definition.
    aclDestroyTensor(x1);
    aclDestroyTensor(x2);
    aclDestroyTensor(gamma);
    aclDestroyTensor(smooth1);
    aclDestroyTensor(smooth2);
    aclDestroyTensor(y1);
    aclDestroyTensor(y2);
    aclDestroyTensor(x);
    aclDestroyTensor(scale1);
    aclDestroyTensor(scale2);

    // 7. Release device resources. Modify based on the specific API definition.
    aclrtFree(x1DeviceAddr);
    aclrtFree(x2DeviceAddr);
    aclrtFree(gammaDeviceAddr);
    aclrtFree(smooth1DeviceAddr);
    aclrtFree(smooth2DeviceAddr);
    aclrtFree(y1DeviceAddr);
    aclrtFree(y2DeviceAddr);
    aclrtFree(xDeviceAddr);
    aclrtFree(scale1DeviceAddr);
    aclrtFree(scale2DeviceAddr);

    if (workspaceSize > 0) {
        aclrtFree(workspaceAddr);
    }
    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    aclFinalize();
    return 0;
}