aclnnAddRmsNormDynamicQuant

支持的产品型号

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

接口原型

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

• 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)

功能描述

• 算子功能：RmsNorm算子是大模型常用的归一化操作，相比LayerNorm算子，其去掉了减去均值的部分。DynamicQuant算子则是为输入张量进行对称动态量化的算子。AddRmsNormDynamicQuant 算子将 RmsNorm 前的 Add 算子和 RmsNorm 归一化输出给到的 1 个或 2 个 DynamicQuant 算子融合起来，减少搬入搬出操作。

• 计算公式：

  $$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}$$
  • 若 smoothScale1Optional 和 smoothScale2Optional 均不输入，则 y2Out 和 scale2Out 输出无实际意义。计算过程如下所示：
  $$scale1Out=row\_max(abs(y))/127$$ $$y1Out=round(y/scalOut)$$
  • 若仅输入 smoothScale1Optional ，则 y2Out 和 scale2Out 输出无实际意义。计算过程如下所示：$$input = y\cdot smoothScale1Optional$$ $$scale1Out=row\_max(abs(input))/127$$
  $$y1Out=round(input/scale1Out)$$
  • 若 smoothScale1Optional 和 smoothScale2Optional 均输入 ，则算子的五个输出均为有效输出。计算过程如下所示：$$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)$$

  其中row_max代表每行求最大值。

aclnnAddRmsNormDynamicQuantGetWorkspaceSize

• 参数说明：

  • x1(aclTensor*，计算输入)：表示标准化过程中的源数据张量，Device侧的aclTensor。shape支持2-8维，数据类型支持FLOAT16、BFLOAT16。数据格式支持ND，支持非连续的Tensor。
  • x2(aclTensor*，计算输入)：表示标准化过程中的源数据张量，Device侧的aclTensor。shape和数据类型需要与x1保持一致。数据格式支持ND，支持非连续的Tensor。
  • gamma(aclTensor*，计算输入)：表示标准化过程中的权重张量，Device侧的aclTensor。shape支持1维，shape需要与x1最后一维一致，数据类型需要与x1保持一致。数据格式支持ND，支持非连续的Tensor。
  • smoothScale1Optional（aclTensor*，计算输入）：表示量化过程中得到y1使用的smoothScale张量，Device侧的aclTensor。可选参数，支持传入空指针。shape和数据类型需要与gamma保持一致。数据格式支持ND，支持非连续的Tensor。
  • smoothScale2Optional（aclTensor*，计算输入）：表示量化过程中得到y2使用的smoothScale张量，Device侧的aclTensor。可选参数，支持传入空指针。必须与smoothScale1Optional配套使用。shape和数据类型需要与gamma保持一致。数据格式支持ND，支持非连续的Tensor。
  • epsilon（double，计算输入）：用于防止除0错误，数据类型为double，建议传入较小正数，如1e-6。
  • y1Out（aclTensor*，计算输出）：表示量化输出Tensor，Device侧的aclTensor。shape需要与输入x1/x2一致，或者是二维并且第一维等于x1除最后一维的维度乘积，第二维等于x1的最后一维，数据类型支持INT8，数据格式支持ND，支持非连续的Tensor。
  • y2Out（aclTensor*，计算输出）：表示量化输出Tensor，Device侧的aclTensor。当smoothScale2Optional不存在时，此输出无意义。shape需要与y1Out一致，数据类型支持INT8，数据格式支持ND，支持非连续的Tensor。
  • xOut（aclTensor*，计算输出）：表示x1和x2的和，Device侧的aclTensor。shape和数据类型需要与输入x1/x2一致，数据格式支持ND，支持非连续的Tensor。
  • scale1Out(aclTensor*，计算输出)：第一路量化的输出，Device侧的aclTensor。shape需要与输入x1除最后一维后的shape一致，或者与x1除最后一维的乘积一致，数据类型支持FLOAT32。数据格式支持ND，支持非连续的Tensor。
  • scale2Out(aclTensor*，计算输出)：第二路量化的输出，Device侧的aclTensor。当smoothScale2Optional不存在时，此输出无意义。shape需要与scale1Out一致，数据类型支持FLOAT32。数据格式支持ND，支持非连续的Tensor。
  • workspaceSize(uint64_t*，出参)：返回需要在Device侧申请的workspace大小。
  • executor(aclOpExecutor**，出参)：返回op执行器，包含了算子计算流程。

• 返回值：

  aclnnStatus: 返回状态码，具体参见aclnn返回码。

  第一段接口完成入参校验，出现以下场景时报错：
  
  - 161001 (ACLNN_ERR_PARAM_NULLPTR)：如果传入参数是必选输入，输出或者必选属性，且是空指针，则返回161001。
  - 161002 (ACLNN_ERR_PARAM_INVALID)：输入和输出的数据类型不在支持的范围之内。
  - 561002 (ACLNN_ERR_INNER_TILING_ERROR):
    1. 输入smoothScale2Optional，而没有输入smoothScale1Optional。
    2. 输入/输出的shape关系不符合预期。

aclnnAddRmsNormDynamicQuant

• 参数说明：

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

• 返回值：

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

约束与限制

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

  x1 数据类型	x2 数据类型	gamma 数据类型	smoothScale1Optional 数据类型	smoothScale2Optional 数据类型	y1Out 数据类型	y2Out 数据类型	scale1Out 数据类型	scale2Out 数据类型
  float16	float16	float16	float16	float16	int8	int8	float32	float32
  bfloat16	bfloat16	bfloat16	bfloat16	bfloat16	int8	int8	float32	float32

调用示例

示例代码如下，仅供参考，具体编译和执行过程请参考编译与运行样例。

#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) {
  // 固定写法，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, 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;

  // 创建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);
  // 创建 smooth1 aclTensor
  ret = CreateAclTensor(smooth1HostData, gammaShape, &smooth1DeviceAddr, aclDataType::ACL_FLOAT16, &smooth1);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // 创建 smooth2 aclTensor
  ret = CreateAclTensor(smooth2HostData, gammaShape, &smooth2DeviceAddr, aclDataType::ACL_FLOAT16, &smooth2);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  // 创建y1 aclTensor
  ret = CreateAclTensor(y1HostData, xShape, &y1DeviceAddr, aclDataType::ACL_INT8, &y1);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // 创建y2 aclTensor
  ret = CreateAclTensor(y2HostData, xShape, &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);
  // 创建outScale1 aclTensor
  ret = CreateAclTensor(scale1HostData, reduceShape, &scale1DeviceAddr, aclDataType::ACL_FLOAT, &scale1);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // 创建outScale1 aclTensor
  ret = CreateAclTensor(scale2HostData, reduceShape, &scale2DeviceAddr, aclDataType::ACL_FLOAT, &scale2);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  // 3. 调用CANN算子库API，需要修改为具体的API
  uint64_t workspaceSize = 0;
  aclOpExecutor* executor;
  // 调用aclnnAddRmsNormDynamicQuant第一段接口
  ret = aclnnAddRmsNormDynamicQuantGetWorkspaceSize(x1, x2, gamma, smooth1, smooth2, epsilon, y1, y2, x, scale1, scale2, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnRmsNormGetWorkspaceSize 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;);
  }
  // 调用aclnnAddRmsNormDynamicQuant第二段接口
  ret = aclnnAddRmsNormDynamicQuant(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnAddRmsNormDynamicQuant 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(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. 释放aclTensor和aclScalar，需要根据具体API的接口定义修改
  aclDestroyTensor(x1);
  aclDestroyTensor(x2);
  aclDestroyTensor(gamma);
  aclDestroyTensor(smooth1);
  aclDestroyTensor(smooth2);
  aclDestroyTensor(y1);
  aclDestroyTensor(y2);
  aclDestroyTensor(x);
  aclDestroyTensor(scale1);
  aclDestroyTensor(scale2);

  // 7. 释放device资源，需要根据具体API的接口定义修改
  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;
}