aclnnAddRmsNormQuant

支持的产品型号

• Atlas 推理系列产品
• Atlas A2 训练系列产品/Atlas 800I A2 推理产品

接口原型

每个算子分为两段式接口，必须先调用aclnnAddRmsNormQuantGetWorkspaceSize接口获取入参并根据计算流程所需workspace大小，再调用aclnnAddRmsNormQuant接口执行计算。

• aclnnStatus aclnnAddRmsNormQuantGetWorkspaceSize(const aclTensor *x1, const aclTensor *x2, const aclTensor *gamma, const aclTensor *scales1, const aclTensor *scales2Optional, const aclTensor *zeroPoints1Optional, const aclTensor *zeroPoints2Optional, int64_t axis, double epsilon, bool divMode, aclTensor *y1Out, aclTensor *y2Out, aclTensor *xOut, uint64_t *workspaceSize, aclOpExecutor **executor)
• aclnnStatus aclnnAddRmsNormQuant(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

功能描述

• 算子功能：RmsNorm算子是大模型常用的标准化操作，相比LayerNorm算子，其去掉了减去均值的部分。AddRmsNormQuant算子将RmsNorm前的Add算子以及RmsNorm后的Quantize算子融合起来，减少搬入搬出操作。
• 计算公式：$$x_i={x1}_{i1}+{x2}_{i2}$$ $$y_i=\frac{x_i}{\operatorname{Rms}(\mathbf{x})} g_i, \quad \text { where } \operatorname{Rms}(\mathbf{x})=\sqrt{\frac{1}{n} \sum_{i=1}^n x_i^2+eps}$$

divMod为True时：

$$y1=round((y/scales1)+zero\_points1)$$ $$y2=round((y/scales2)+zero\_points2)$$

divMod为False时：

$$y1=round((y*scales1)+zero\_points1)$$ $$y2=round((y*scales2)+zero\_points2)$$

aclnnAddRmsNormQuantGetWorkspaceSize

• 参数说明：

  • x1(aclTensor*，计算输入)：表示标准化过程中的源数据张量，公式中的x1，Device侧的aclTensor。shape支持1-8维。数据格式支持ND，支持非连续的Tensor，不支持空Tensor。
    • Atlas 推理系列产品：数据类型支持FLOAT16。
    • Atlas A2 训练系列产品/Atlas 800I A2 推理产品：数据类型支持FLOAT16、BFLOAT16。
  • x2(aclTensor*，计算输入)：表示标准化过程中的源数据张量，公式中的x2，Device侧的aclTensor。shape支持1-8维，shape和数据类型需要与x1保持一致。数据格式支持ND，支持非连续的Tensor，不支持空Tensor。
    • Atlas 推理系列产品：数据类型支持FLOAT16。
    • Atlas A2 训练系列产品/Atlas 800I A2 推理产品：数据类型支持FLOAT16、BFLOAT16。
  • gamma(aclTensor*，计算输入)：表示标准化过程中的权重张量，公式中的g，Device侧的aclTensor。shape支持1-8维，shape需要与x1需要Norm的维度保持一致，数据类型需要与x1保持一致。数据格式支持ND，支持非连续的Tensor，不支持空Tensor。
    • Atlas 推理系列产品：数据类型支持FLOAT16。
    • Atlas A2 训练系列产品/Atlas 800I A2 推理产品：数据类型支持FLOAT16、BFLOAT16。
  • scales1（aclTensor*，计算输入）：表示量化过程中得到y1进行的scales张量，公式中的scales1，Device侧的aclTensor。shape需要与gamma保持一致。数据格式支持ND，支持非连续的Tensor，不支持空Tensor。
    • Atlas 推理系列产品：数据类型支持FLOAT32。
    • Atlas A2 训练系列产品/Atlas 800I A2 推理产品：数据类型支持FLOAT32、BFLOAT16。
  • scales2Optional（aclTensor*，计算输入）：表示量化过程中得到y2进行的scales张量，公式中的scales2，Device侧的aclTensor。可选参数，支持传入空指针。shape需要与gamma保持一致，数据类型需要与scales1保持一致。数据格式支持ND，支持非连续的Tensor，不支持空Tensor。
    • Atlas 推理系列产品：数据类型支持FLOAT32。
    • Atlas A2 训练系列产品/Atlas 800I A2 推理产品：数据类型支持FLOAT32、BFLOAT16。
  • zeroPoints1Optional（aclTensor*，计算输入）：表示量化过程中得到y1进行的offset张量，公式中的zero_points1，Device侧的aclTensor。可选参数，支持传入空指针。shape需要与gamma保持一致。数据格式支持ND，支持非连续的Tensor，不支持空Tensor。
    • Atlas 推理系列产品：数据类型支持INT32。
    • Atlas A2 训练系列产品/Atlas 800I A2 推理产品：数据类型支持INT32、BFLOAT16。
  • zeroPoints2Optional（aclTensor*，计算输入）：表示量化过程中得到y2进行的offset张量，公式中的zero_points2，Device侧的aclTensor。可选参数，支持传入空指针。shape需要与gamma保持一致，数据类型需要与zeroPoints1Optional保持一致。数据格式支持ND，支持非连续的Tensor，不支持空Tensor。
    • Atlas 推理系列产品：数据类型支持INT32。
    • Atlas A2 训练系列产品/Atlas 800I A2 推理产品：数据类型支持INT32、BFLOAT16。
  • axis（int64_t，计算输入）：Host侧的整型，表示需要进行量化的elewise轴，其他的轴做broadcast，指定的轴不能超过输入x的维度数。当前仅支持-1，传其他值均不生效。
  • epsilon(double，计算输入)： 公式中的输入eps，用于防止除0错误，数据类型为double。建议传较小的正数。
  • divMode(bool，计算输入)： 公式中决定量化公式是否使用除法的参数，数据类型为bool，当前仅支持True，传其他值均不生效。
  • y1Out（aclTensor*，计算输出）：表示量化输出Tensor，公式中的y1，Device侧的aclTensor。shape支持1-8维度，shape需要与输入x1/x2一致，数据类型支持INT8，数据格式支持ND，支持非连续的Tensor。
  • y2Out（aclTensor*，计算输出）：表示量化输出Tensor，公式中的y2，Device侧的aclTensor。shape支持1-8维度，shape需要与输入x1/x2一致，数据类型支持INT8，数据格式支持ND，支持非连续的Tensor。
  • xOut(aclTensor*，计算输出)：表示x1和x2的和，公式中的x，Device侧的aclTensor。shape支持1-8维度，shape和数据类型需要与输入x1/x2一致，数据格式支持ND，支持非连续的Tensor。
    • Atlas 推理系列产品：数据类型支持FLOAT16。
    • Atlas A2 训练系列产品/Atlas 800I A2 推理产品：数据类型支持FLOAT16、BFLOAT16。
  • workspaceSize(uint64_t*，出参)：返回需要在Device侧申请的workspace大小。
  • executor(aclOpExecutor**，出参)：返回op执行器，包含了算子计算流程。

• 返回值：

  aclnnStatus：返回状态码。（参见aclnn返回码）

  第一段接口完成入参校验，出现以下场景时报错：
  
  - 161001 (ACLNN_ERR_PARAM_NULLPTR)：如果传入参数是必选输入，输出或者必选属性，且是空指针，则返回161001。
  - 161002 (ACLNN_ERR_PARAM_INVALID)：输入和输出的数据类型不在支持的范围之内。

aclnnAddRmsNormQuant

• 参数说明：

  • workspace(void*，入参)：在Device侧申请的workspace内存地址。
  • workspaceSize(uint64_t，入参)：在Device侧申请的workspace大小，由第一段接口aclnnAddRmsNormQuantGetWorkspaceSize获取。
  • executor(aclOpExecutor*，入参)：op执行器，包含了算子计算流程。
  • stream(aclrtStream，入参)：指定执行任务的AscendCL stream流。

• 返回值：

  aclnnStatus：返回状态码。（具体参见aclnn返回码）

约束与限制

• 各产品型号支持数据类型说明

  • Atlas A2 训练系列产品/Atlas 800I A2 推理产品：

  x1 数据类型	x2 数据类型	gamma 数据类型	scales1 数据类型	scales2Optional 数据类型	zeroPoints1Optional 数据类型	zeroPoints2Optional 数据类型	y1Out 数据类型	y2Out 数据类型	xOut 数据类型
  float16	float16	float16	float32	float32	int32	int32	int8	int8	float16
  bfloat16	bfloat16	bfloat16	bfloat16	bfloat16	bfloat16	bfloat16	int8	int8	bfloat16

  • Atlas 推理系列产品：

  x1 数据类型	x2 数据类型	gamma 数据类型	scales1 数据类型	scales2Optional 数据类型	zeroPoints1Optional 数据类型	zeroPoints2Optional 数据类型	y1Out 数据类型	y2Out 数据类型	xOut 数据类型
  float16	float16	float16	float32	float32	int32	int32	int8	int8	float16

调用示例

示例编译和执行请参考编译与运行样例。

#include <iostream>
#include <vector>
#include "acl/acl.h"
#include "aclnnop/aclnn_add_rms_norm_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) {
  // 固定写法，AscendCL初始化
  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);
  // 调用aclrtMalloc申请device侧内存
  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);

  // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上
  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);

  // 计算连续tensor的strides
  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];
  }

  // 调用aclCreateTensor接口创建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. （固定写法）device/stream初始化, 参考acl对外接口列表
  // 根据自己的实际device填写deviceId
  int32_t deviceId = 0;
  aclrtStream stream;
  auto ret = Init(deviceId, &stream);
  // check根据自己的需要处理
  CHECK_RET(ret == 0, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
  // 2. 构造输入与输出，需要根据API的接口自定义构造
  std::vector<int64_t> xShape = {2, 16};
  std::vector<int64_t> gammaShape = {16};
  std::vector<int64_t> yShape = {2, 16};
  long long xShapeSize = GetShapeSize(xShape);
  long long gammaShapeSize = GetShapeSize(gammaShape);
  void* x1DeviceAddr = nullptr;
  void* x2DeviceAddr = nullptr;
  void* gammaDeviceAddr = nullptr;
  void* scales1DeviceAddr = nullptr;
  void* zeroPoints1DeviceAddr = nullptr;
  void* y1DeviceAddr = nullptr;
  void* y2DeviceAddr = nullptr;
  void* xDeviceAddr = nullptr;
  aclTensor* x1 = nullptr;
  aclTensor* x2 = nullptr;
  aclTensor* gamma = nullptr;
  aclTensor* scales1 = nullptr;
  aclTensor* zeroPoints1 = nullptr;
  aclTensor* y1 = nullptr;
  aclTensor* y2 = nullptr;
  aclTensor* x = nullptr;
  std::vector<int16_t> x1HostData(xShapeSize, 0);
  std::vector<int16_t> x2HostData(xShapeSize, 0);
  std::vector<int16_t> gammaHostData(gammaShapeSize, 0);
  std::vector<float> scales1HostData(gammaShapeSize, 1);
  std::vector<int32_t> zeroPoints1HostData(gammaShapeSize, 100);
  std::vector<int8_t> y1HostData(xShapeSize, 0);
  std::vector<int8_t> y2HostData(xShapeSize, 0);
  std::vector<int16_t> xHostData(xShapeSize, 0);
  float epsilon = 1e-6;
  int64_t axis = -1;
  bool divMode = true;
  // 创建x1 aclTensor
  ret = CreateAclTensor(x1HostData, xShape, &x1DeviceAddr, aclDataType::ACL_FLOAT16, &x1);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // 创建x2 aclTensor
  ret = CreateAclTensor(x2HostData, xShape, &x2DeviceAddr, aclDataType::ACL_FLOAT16, &x2);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // 创建gamma aclTensor
  ret = CreateAclTensor(gammaHostData, gammaShape, &gammaDeviceAddr, aclDataType::ACL_FLOAT16, &gamma);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // 创建scales1 aclTensor
  ret = CreateAclTensor(scales1HostData, gammaShape, &scales1DeviceAddr, aclDataType::ACL_FLOAT, &scales1);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // 创建zeroPoints1 aclTensor
  ret = CreateAclTensor(zeroPoints1HostData, gammaShape, &zeroPoints1DeviceAddr, aclDataType::ACL_INT32, &zeroPoints1);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // 创建y1 aclTensor
  ret = CreateAclTensor(y1HostData, yShape, &y1DeviceAddr, aclDataType::ACL_INT8, &y1);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // 创建y2 aclTensor
  ret = CreateAclTensor(y2HostData, yShape, &y2DeviceAddr, aclDataType::ACL_INT8, &y2);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // 创建x aclTensor
   ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_FLOAT16, &x);
   CHECK_RET(ret == ACL_SUCCESS, return ret);
  // 3. 调用CANN算子库API，需要修改为具体的API
  uint64_t workspaceSize = 0;
  aclOpExecutor* executor;
  // 调用aclnnAddRmsNormQuant第一段接口
  ret = aclnnAddRmsNormQuantGetWorkspaceSize(x1, x2, gamma, scales1, nullptr, zeroPoints1, nullptr, axis, epsilon, divMode, y1, y2, x, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnRmsNormQuantGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
  // 根据第一段接口计算出的workspaceSize申请device内存
  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;);
  }
  // 调用aclnnAddRmsNormQuant第二段接口
  ret = aclnnAddRmsNormQuant(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnAddRmsNormQuant failed. ERROR: %d\n", ret); return ret);
  // 4. （固定写法）同步等待任务执行结束
  ret = aclrtSynchronizeStream(stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
  // 5. 获取输出的值，将device侧内存上的结果拷贝至host侧，需要根据具体API的接口定义修改
  auto size = GetShapeSize(yShape);
  std::vector<int8_t> resultData(size, 0);
  ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[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, resultData[i]);
  }

  // 6. 释放aclTensor和aclScalar，需要根据具体API的接口定义修改
  aclDestroyTensor(x1);
  aclDestroyTensor(x2);
  aclDestroyTensor(gamma);
  aclDestroyTensor(scales1);
  aclDestroyTensor(zeroPoints1);
  aclDestroyTensor(y1);
  aclDestroyTensor(y2);
  aclDestroyTensor(x);

  // 7. 释放device资源，需要根据具体API的接口定义修改
  aclrtFree(x1DeviceAddr);
  aclrtFree(x2DeviceAddr);
  aclrtFree(gammaDeviceAddr);
  aclrtFree(scales1DeviceAddr);
  aclrtFree(zeroPoints1DeviceAddr);
  aclrtFree(y1DeviceAddr);
  aclrtFree(y2DeviceAddr);
  aclrtFree(xDeviceAddr);

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