CANN ops-nn 融合算子 FusedAddRmsNorm 全解析：ScaledAdd 与 RmsNorm 融合原理、aclnn 两段式接口调用与 NPU 源码实现

CANN ops-nn 融合算子 FusedAddRmsNorm 全解析：ScaledAdd 与 RmsNorm 融合原理、aclnn 两段式接口调用与 NPU 源码实现

【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn

本文围绕 CANN 神经网络算子库 ops-nn 中 experimental/norm 目录下的 FusedAddRmsNorm 融合算子展开：它面向大模型推理/训练中高频出现的“先加后归一”计算模式，将 RmsNorm 前的 ScaledAdd（x = x1 * scale + x2）与 RmsNorm 本身融合为单个算子，从而减少 NPU 上的数据搬入搬出。读完本文，你将掌握该算子的计算公式与全部参数语义、产品支持情况、aclnn 两段式接口（GetWorkspaceSize + 执行）的完整调用范式、图模式 IR 构图方式，以及从 op_def / infershape / tiling 到 AICore kernel 的源码级实现细节。

一、产品支持情况

FusedAddRmsNorm 算子在当前仓库中声明支持以下产品，均以 √ 表示支持：

从算子定义源码 fused_add_rms_norm_def.cpp 可以看到与上述产品对应的 AICore 配置：基础配置（BaseConfig）注册了ascend910b与ascend910_93两类内核；ascend310p/kirinx90共享一套仅支持 FLOAT16、FLOAT 的配置；ascend910_95则额外开启了动态编译、动态 rank 与动态 shape 支持。这从实现侧印证了表格中的支持范围，同时也说明不同硬件内核的数据类型支持存在差异（310P 系列不支持 BFLOAT16，详见下文参数说明）。

二、功能说明与计算公式

2.1 融合动机

RmsNorm（Root Mean Square Layer Normalization）是大模型常用的归一化操作，相比 LayerNorm，它去掉了“减去均值”的步骤，只基于均方根对输入做缩放。在实际模型中，RmsNorm 之前往往紧跟一个 ScaledAdd（典型如 MiniCPM 类模型的残差缩放结构）：先计算x = x1 * scale + x2，再做 RmsNorm。如果分开执行，中间结果x需要从 Device 侧内存搬出再搬入，产生额外的访存开销。

FusedAddRmsNorm 算子正是把这两步融合进一个算子：内部完成 ScaledAdd 计算后直接就地做 RmsNorm，减少搬入搬出操作，同时对外暴露中间结果x与逆标准差rstd（后者可被反向算子复用，避免反向传播时重复计算）。

2.2 计算公式

第一步，ScaledAdd 融合加法：

$$ x_i = x1_{i} \times scale + x2_{i} $$

第二步，对结果做 RmsNorm：

$$ \operatorname{RmsNorm}(x_i) = \frac{x_i}{\operatorname{Rms}(\mathbf{x})} g_i, \quad \text{其中 } \operatorname{Rms}(\mathbf{x}) = \sqrt{\frac{1}{n} \sum_{i=1}^{n} x_i^2 + \epsilon} $$

其中g为可学习的缩放权重（gamma），n为需要归一化的维度大小，epsilon用于数值稳定、防止除零。算子图 IR 原型 fused_add_rms_norm_proto.h 中以 NumPy 形式给出了等价描述：rstd = np.rsqrt(np.mean(x^2, reduce_axis, keepdims=True) + epsilon)，y = gamma * (x * rstd)。

三、参数说明（算子级）

下表完整列出了算子层级的输入、输出与属性参数，来自 README.md 的参数说明章节：

几个需要特别留意的要点：

• gamma 与归一化维度：gamma 的 shape 必须与x1的后几维一致，例如x1shape 为 (2, 3, 4, 8) 时，gamma 可取 (8) 或 (4, 8)，表示对最后 1 维或最后 2 维做归一化。
• rstd 的 shape 推导：rstd 的数据类型固定为 FLOAT32（与输入类型无关），其 shape 等于x1去掉 gamma 覆盖的后几维、并把这些维度压成 1。以 aclnn 接口文档 aclnnFusedAddRmsNorm.md 中的示例：x1shape 为 (2, 3, 4, 8)、gamma shape 为 (8) 时，rstdOut shape 为 (2, 3, 4, 1)；gamma shape 为 (4, 8) 时，rstdOut shape 为 (2, 3, 1, 1)。
• 属性默认值：epsilon默认 1e-6、scale默认 1.0，在 fused_add_rms_norm_def.cpp（op.Attr("epsilon").AttrType(OPTIONAL).Float(1e-6)、op.Attr("scale").AttrType(OPTIONAL).Float(1.0)）与 fused_add_rms_norm_proto.h（.ATTR(epsilon, Float, 1e-6f)、.ATTR(scale, Float, 1.0f)）两处均有登记，完全一致。
• 310P 差异：算子定义中ARCH310P_DTYPE仅包含 FLOAT16、FLOAT，即ascend310p/kirinx90内核不支持 BFLOAT16 输入，使用时需结合目标硬件选择数据类型。

四、aclnn 接口调用指南（两段式接口）

4.1 两段式调用模型

在 CANN 的 aclnn 单算子调用体系中，FusedAddRmsNorm 遵循标准的[两段式接口]模式：第一段aclnnFusedAddRmsNormGetWorkspaceSize负责入参校验并计算算子执行所需的 workspace 大小、创建执行器；第二段aclnnFusedAddRmsNorm使用第一段返回的 workspace 与 executor 在指定 Stream 上真正下发执行。

4.2 函数原型

// 第一段：获取 workspace 大小并创建执行器 aclnnStatus aclnnFusedAddRmsNormGetWorkspaceSize( const aclTensor *x1, const aclTensor *x2, const aclTensor *gamma, double epsilon, double scale, const aclTensor *yOut, const aclTensor *rstdOut, const aclTensor *xOut, uint64_t *workspaceSize, aclOpExecutor **executor) // 第二段：执行计算 aclnnStatus aclnnFusedAddRmsNorm( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

4.3 aclnnFusedAddRmsNormGetWorkspaceSize 参数说明

需要注意：接口层epsilon、scale以DOUBLE（C 语言 double）传入，而算子属性层为 FLOAT，实际值一致；接口层允许传入非连续 Tensor（表中“非连续Tensor”列为 √）。

4.4 第一段接口返回码（入参校验）

第一段接口会完成入参校验，出现以下场景时报错：

aclnnFusedAddRmsNorm第二段接口参数（workspace、workspaceSize、executor、stream）分别表示 Device 侧 workspace 内存地址、由第一段接口获取的 workspace 大小、算子执行器以及执行任务所在的 Stream；其返回值同为 aclnnStatus 状态码。

4.5 约束说明

• 边界值：输入为 Inf 时输出为 Inf；输入为 NaN 时输出为 NaN。
• 确定性：aclnnFusedAddRmsNorm默认为确定性实现，多次运行结果可复现。
• 除此之外算子本身无其他约束（README 约束说明为“无”）。

五、完整调用示例

以下示例来自 examples/test_aclnn_fused_add_rms_norm.cpp（aclnn 接口文档中的调用示例与此一致），演示了从环境初始化、Tensor 构造、两段式调用到结果回拷与资源释放的完整流程。编译与执行的整体过程可参考仓库 docs 中的编译运行样例说明。

#include <cstdio> #include <iostream> #include <vector> #include "acl/acl.h" #include "aclnnop/aclnn_fused_add_rms_norm.h" #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) int64_t GetShapeSize(const std::vector<int64_t>& shape) { int64_t shapeSize = 1; for (auto dim : shape) { shapeSize *= dim; } return shapeSize; } int Init(int32_t deviceId, aclrtStream* stream) { auto ret = aclInit(nullptr); if (ret != ACL_SUCCESS) { LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret; } ret = aclrtSetDevice(deviceId); if (ret != ACL_SUCCESS) { LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret; } ret = aclrtCreateStream(stream); if (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); auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); if (ret != ACL_SUCCESS) { LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret; } ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); if (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 = static_cast<int64_t>(shape.size()) - 2; i >= 0; i--) { strides[i] = shape[i + 1] * strides[i + 1]; } *tensor = aclCreateTensor( shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { int32_t deviceId = 0; aclrtStream stream = nullptr; auto ret = Init(deviceId, &stream); if (ret != 0) { LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret; } // 构造输入输出 shape：x1/x2/y/x 为 (2, 16)，gamma 为 (16)，rstd 为 (2, 1) std::vector<int64_t> xShape = {2, 16}; std::vector<int64_t> gammaShape = {16}; std::vector<int64_t> yShape = {2, 16}; std::vector<int64_t> rstdShape = {2, 1}; void* x1DeviceAddr = nullptr; void* x2DeviceAddr = nullptr; void* gammaDeviceAddr = nullptr; void* yDeviceAddr = nullptr; void* rstdDeviceAddr = nullptr; void* xDeviceAddr = nullptr; void* workspaceAddr = nullptr; aclTensor* x1 = nullptr; aclTensor* x2 = nullptr; aclTensor* gamma = nullptr; aclTensor* y = nullptr; aclTensor* rstd = nullptr; aclTensor* x = nullptr; std::vector<float> x1HostData = {0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7}; std::vector<float> x2HostData = {0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7}; std::vector<float> gammaHostData = {0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7}; std::vector<float> yHostData(GetShapeSize(yShape), 0); std::vector<float> rstdHostData(GetShapeSize(rstdShape), 0); std::vector<float> xHostData(GetShapeSize(xShape), 0); // 创建 6 个 aclTensor（x1/x2/gamma 输入，y/rstd/x 输出） ret = CreateAclTensor(x1HostData, xShape, &x1DeviceAddr, aclDataType::ACL_FLOAT, &x1); if (ret != ACL_SUCCESS) { return ret; } ret = CreateAclTensor(x2HostData, xShape, &x2DeviceAddr, aclDataType::ACL_FLOAT, &x2); if (ret != ACL_SUCCESS) { return ret; } ret = CreateAclTensor(gammaHostData, gammaShape, &gammaDeviceAddr, aclDataType::ACL_FLOAT, &gamma); if (ret != ACL_SUCCESS) { return ret; } ret = CreateAclTensor(yHostData, yShape, &yDeviceAddr, aclDataType::ACL_FLOAT, &y); if (ret != ACL_SUCCESS) { return ret; } ret = CreateAclTensor(rstdHostData, rstdShape, &rstdDeviceAddr, aclDataType::ACL_FLOAT, &rstd); if (ret != ACL_SUCCESS) { return ret; } ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_FLOAT, &x); if (ret != ACL_SUCCESS) { return ret; } float epsilon = 1e-6F; // 建议值 1e-6 float scale = 0.5F; // 示例中 x1 的缩放系数取 0.5 uint64_t workspaceSize = 0; aclOpExecutor* executor = nullptr; // 第一段接口：校验入参并获取 workspace 大小、创建执行器 ret = aclnnFusedAddRmsNormGetWorkspaceSize( x1, x2, gamma, epsilon, scale, y, rstd, x, &workspaceSize, &executor); if (ret != ACL_SUCCESS) { LOG_PRINT("aclnnFusedAddRmsNormGetWorkspaceSize failed. ERROR: %d\n", ret); return ret; } // 按需申请 Device 侧 workspace if (workspaceSize > 0) { ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); if (ret != ACL_SUCCESS) { LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret; } } // 第二段接口：在指定 Stream 上执行计算 ret = aclnnFusedAddRmsNorm(workspaceAddr, workspaceSize, executor, stream); if (ret != ACL_SUCCESS) { LOG_PRINT("aclnnFusedAddRmsNorm failed. ERROR: %d\n", ret); return ret; } // 同步等待任务执行结束 ret = aclrtSynchronizeStream(stream); if (ret != ACL_SUCCESS) { LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret; } // 将 y 结果从 Device 侧回拷到 Host 侧并打印 auto size = GetShapeSize(yShape); std::vector<float> resultData(size, 0); ret = aclrtMemcpy( resultData.data(), resultData.size() * sizeof(resultData[0]), yDeviceAddr, size * sizeof(float), ACL_MEMCPY_DEVICE_TO_HOST); if (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("y result[%ld] is: %f\n", i, resultData[i]); } // 释放 aclTensor 与 Device 资源 aclDestroyTensor(x1); aclDestroyTensor(x2); aclDestroyTensor(gamma); aclDestroyTensor(y); aclDestroyTensor(rstd); aclDestroyTensor(x); aclrtFree(x1DeviceAddr); aclrtFree(x2DeviceAddr); aclrtFree(xDeviceAddr); aclrtFree(gammaDeviceAddr); aclrtFree(yDeviceAddr); aclrtFree(rstdDeviceAddr); if (workspaceSize > 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }

该示例的关键点在于：使用aclCreateTensor构造 ND 格式 aclTensor 时必须正确计算 strides（连续布局下从最后一维向前累乘）；两段式接口中 workspace 仅在第一段返回非 0 时才需要申请；执行结束后必须aclrtSynchronizeStream同步后再回拷结果。

六、图模式调用（算子 IR 构图）

除 aclnn 单算子调用外，FusedAddRmsNorm 还支持图模式调用：通过算子 IR 原型 fused_add_rms_norm_proto.h 构图（即 GE 图上的REG_OP(FusedAddRmsNorm)注册），该头文件同时就是算子原型定义，声明了：

• 3 个输入x1、x2、gamma，类型均为{DT_FLOAT, DT_FLOAT16, DT_BF16}；
• 3 个输出y、rstd、x，其中rstd固定为{DT_FLOAT, DT_FLOAT, DT_FLOAT}，其余与输入同类型；
• 2 个属性epsilon（默认 1e-6f）与scale（默认 1.0f）。

README 的调用说明汇总如下：

七、源码级实现纵深：host 侧到 kernel 侧

7.1 算子定义与数据约束（op_host）

fused_add_rms_norm_def.cpp 中通过OP_ADD(FusedAddRmsNorm)完成算子注册，核心信息包括：

• 输入统一使用.ParamType(REQUIRED)、.AutoContiguous()，并指定 DataType / Format / UnknownShapeFormat，保证动态 shape 场景下输入自动连续化；
• 基础 dtype 集合为{DT_FLOAT16, DT_FLOAT, DT_BF16}，rstd输出恒为DT_FLOAT；
• 属性epsilon默认 1e-6、scale默认 1.0，类型 FLOAT，均为 OPTIONAL；
• 按架构差异化注册 AICore 配置（ascend910b/ascend910_93、ascend310p/kirinx90、ascend910_95），310P 系列不支持 BF16，910_95 支持动态 rank/shape。

7.2 shape 与数据类型推导（op_host）

fused_add_rms_norm_infershape.cpp 实现了InferShape与InferDataType两个推导逻辑：

• y、x的 shape 直接拷贝x1的 shape；
• rstd保持与x1相同的维度数，但 gamma 覆盖的后gammaDimNum维全部置为 1（即i < xDimNum - gammaDimNum的维度取x1原值，其余维度取 1），与接口文档中的示例完全吻合；
• 输出数据类型：y/x与x1一致，rstd固定为DT_FLOAT。

仓库配套的单测 test_FusedAddRmsNorm_infershape.cpp 覆盖了上述 shape 推导规则。

7.3 融合计算的 AICore kernel 实现（op_kernel）

kernel 侧以 fused_add_rms_norm.h 为入口，复用rms_norm算子的公共基础（包含 rms_norm_base.h），并按照归一化维度的不同拆分为多个实现文件：fused_add_rms_norm_single_n.h、fused_add_rms_norm_multi_n.h、fused_add_rms_norm_merge_n.h、fused_add_rms_norm_split_d.h，分别应对不同的切分策略。

融合实现的关键路径（以通用KernelFusedAddRmsNorm为例）：

1. 数据搬入：按行（row）切分任务，通过CopyIn将x1、x2行数据搬入 UB；gamma 整行搬入并复用。
2. ScaledAdd 就地融合：FLOAT32 直接Muls(x1, x1, scale)再Add(x1, x1, x2)；FLOAT16/BFLOAT16 则先Cast到 FLOAT32 计算（保证精度），再Cast回原类型写出中间结果x（对 BF16 使用CAST_RINT舍入）。
3. RmsNorm 计算：公共函数BuildRstd（fused_add_rms_norm_common.h）依次执行Mul求平方、Muls乘1/n平均、ReduceSumCustom求和、Adds加epsilon、Sqrt开方、Div(1, ...)取倒数，得到 rstd 标量，写入rstdGm。
4. 输出 y：Muls(x, x, rstdValue)再Mul(y, gamma, y)，把结果经CopyOutY写回yGm；中间结果x也独立写出，供上层复用。

从FUSED_ADD_RMS_NORM_INIT_ROW_COMMON宏可以看到 kernel 的任务划分逻辑：按num_row（总行数）、num_col（归一化维度大小）、blockFactor（每核行数）、rowFactor（每轮处理行数）、ubFactor（UB 缓冲元素数）等 tiling 参数驱动，多核并行时每个核处理blockFactor行，最后一个核处理尾行，行与行之间通过rowFactor分轮流水。tiling 参数由 host 侧 fused_add_rms_norm_tiling.cpp 计算，并有对应的单测 test_fused_add_rms_norm_tiling.cpp 与 kernel 单测 test_fused_add_rms_norm.cpp 验证数值正确性。

八、使用建议总结

1. 参数选取：epsilon建议取 1e-6；scale默认为 1.0，MiniCPM 等带残差缩放权重的模型应按模型配置显式传入。
2. shape 规划：x2与x1同 shape；gamma对齐x1的后几维（归一化维度）；rstd无需手动设计，shape 由x1、gamma共同决定（前段维度保留、归一化维度压 1）。
3. 精度策略：FLOAT16/BFLOAT16 输入在 kernel 内部会经 FLOAT32 中间态计算，精度可控；BF16 输出使用CAST_RINT舍入。
4. 调试与验证：仓库在tests/ut下提供了 host 侧 infershape/tiling 与 kernel 侧数值单测，可作为新硬件适配或参数调整时的回归依据。
5. 硬件匹配：BFLOAT16 输入在 310P 系列上不受支持，请以目标产品的实际支持矩阵为准（README 与 def.cpp 中的架构配置可相互印证）。

【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn