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 前的 ScaledAddx x1 * scale x2与 RmsNorm 本身融合为单个算子从而减少 NPU 上的数据搬入搬出。读完本文你将掌握该算子的计算公式与全部参数语义、产品支持情况、aclnn 两段式接口GetWorkspaceSize 执行的完整调用范式、图模式 IR 构图方式以及从 op_def / infershape / tiling 到 AICore kernel 的源码级实现细节。一、产品支持情况FusedAddRmsNorm 算子在当前仓库中声明支持以下产品均以 √ 表示支持产品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 训练系列产品 / Atlas A3 推理系列产品√Atlas A2 训练系列产品 / Atlas A2 推理系列产品√从算子定义源码 fused_add_rms_norm_def.cpp 可以看到与上述产品对应的 AICore 配置基础配置BaseConfig注册了ascend910b与ascend910_93两类内核ascend310p/kirinx90共享一套仅支持 FLOAT16、FLOAT 的配置ascend910_95则额外开启了动态编译、动态 rank 与动态 shape 支持。这从实现侧印证了表格中的支持范围同时也说明不同硬件内核的数据类型支持存在差异310P 系列不支持 BFLOAT16详见下文参数说明。二、功能说明与计算公式2.1 融合动机RmsNormRoot 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_{i1}^{n} x_i^2 \epsilon} $$其中g为可学习的缩放权重gamman为需要归一化的维度大小epsilon用于数值稳定、防止除零。算子图 IR 原型 fused_add_rms_norm_proto.h 中以 NumPy 形式给出了等价描述rstd np.rsqrt(np.mean(x^2, reduce_axis, keepdimsTrue) epsilon)y gamma * (x * rstd)。三、参数说明算子级下表完整列出了算子层级的输入、输出与属性参数来自 README.md 的参数说明章节参数名输入/输出/属性描述数据类型数据格式x1输入用于 ScaledAdd 计算的第一个输入对应公式中的x1FLOAT32、FLOAT16、BFLOAT16NDx2输入用于 ScaledAdd 计算的第二个输入对应公式中的x2FLOAT32、FLOAT16、BFLOAT16NDgamma输入RmsNorm 的缩放因子权重对应公式中的g。shape 需要与x1后几维保持一致后几维为x1需要 norm 的维度FLOAT32、FLOAT16、BFLOAT16NDepsilon可选属性添加到分母中的值确保数值稳定防止除 0 错误对应公式中的eps默认值为 1e-6fFLOAT-scale可选属性ScaledAdd 阶段中x1的缩放系数对应公式中的scale默认值为 1.0fFLOAT-y输出最终输出Device 侧 aclTensor对应公式中的RmsNorm(x)FLOAT32、FLOAT16、BFLOAT16NDrstd输出归一化后的标准差的倒数对应公式中的Rms(x)FLOAT32NDx输出ScaledAdd 计算的结果对应公式中的xFLOAT32、FLOAT16、BFLOAT16ND几个需要特别留意的要点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.cppop.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 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续Tensorx1输入用于 ScaledAdd 计算的第一个输入对应公式中的x1不支持空 TensorFLOAT32、FLOAT16、BFLOAT16ND1-8√x2输入用于 ScaledAdd 计算的第二个输入对应公式中的x2不支持空 Tensorshape 和数据类型需与x1保持一致FLOAT32、FLOAT16、BFLOAT16ND1-8√gamma输入RmsNorm 的缩放因子权重对应公式中的gamma不支持空 Tensor数据类型与x1一致shape 与x1后几维一致FLOAT32、FLOAT16、BFLOAT16ND1-8√epsilon输入添加到分母中的值确保数值稳定对应公式中的epsilon建议值 1e-6DOUBLE---scale输入ScaledAdd 阶段中x1的缩放系数对应公式中的scale默认值 1.0MiniCPM 场景通常按模型配置传入DOUBLE---yOut输出最终输出对应公式中的RmsNorm(x)不支持空 Tensorshape、数据类型与x1一致FLOAT32、FLOAT16、BFLOAT16ND1-8√rstdOut输出归一化后标准差的倒数对应公式中Rms(x)的倒数不支持空 Tensorshape 与x1前几维一致前几维为不需要 norm 的维度后几维置 1FLOAT32ND1-8√xOut输出ScaledAdd 计算的结果对应公式中的x不支持空 Tensorshape、数据类型与x1一致FLOAT32、FLOAT16、BFLOAT16ND1-8√workspaceSize输出需要在 Device 侧申请的 workspace 大小-----executor输出op 执行器包含算子计算流程-----需要注意接口层epsilon、scale以DOUBLEC 语言 double传入而算子属性层为 FLOAT实际值一致接口层允许传入非连续 Tensor表中“非连续Tensor”列为 √。4.4 第一段接口返回码入参校验第一段接口会完成入参校验出现以下场景时报错返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001必选输入、输出或必选属性为空指针时返回 161001ACLNN_ERR_PARAM_INVALID161002输入或输出的数据类型不在支持范围内或输入输出参数不满足参数说明中的约束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.cppaclnn 接口文档中的调用示例与此一致演示了从环境初始化、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::vectorint64_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::vectorT hostData, const std::vectorint64_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::vectorint64_t strides(shape.size(), 1); for (int64_t i static_castint64_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; } // 构造输入输出 shapex1/x2/y/x 为 (2, 16)gamma 为 (16)rstd 为 (2, 1) std::vectorint64_t xShape {2, 16}; std::vectorint64_t gammaShape {16}; std::vectorint64_t yShape {2, 16}; std::vectorint64_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::vectorfloat 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::vectorfloat 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::vectorfloat gammaHostData {0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7}; std::vectorfloat yHostData(GetShapeSize(yShape), 0); std::vectorfloat rstdHostData(GetShapeSize(rstdShape), 0); std::vectorfloat xHostData(GetShapeSize(xShape), 0); // 创建 6 个 aclTensorx1/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::vectorfloat 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 的调用说明汇总如下调用方式样例代码说明aclnn 接口test_aclnn_fused_add_rms_norm通过 aclnnFusedAddRmsNorm 接口方式调用 FusedAddRmsNorm 算子图模式-通过算子IR构图方式调用 FusedAddRmsNorm 算子七、源码级实现纵深host 侧到 kernel 侧7.1 算子定义与数据约束op_hostfused_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_95310P 系列不支持 BF16910_95 支持动态 rank/shape。7.2 shape 与数据类型推导op_hostfused_add_rms_norm_infershape.cpp 实现了InferShape与InferDataType两个推导逻辑y、x的 shape 直接拷贝x1的 shaperstd保持与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_kernelkernel 侧以 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为例数据搬入按行row切分任务通过CopyIn将x1、x2行数据搬入 UBgamma 整行搬入并复用。ScaledAdd 就地融合FLOAT32 直接Muls(x1, x1, scale)再Add(x1, x1, x2)FLOAT16/BFLOAT16 则先Cast到 FLOAT32 计算保证精度再Cast回原类型写出中间结果x对 BF16 使用CAST_RINT舍入。RmsNorm 计算公共函数BuildRstdfused_add_rms_norm_common.h依次执行Mul求平方、Muls乘1/n平均、ReduceSumCustom求和、Adds加epsilon、Sqrt开方、Div(1, ...)取倒数得到 rstd 标量写入rstdGm。输出 yMuls(x, x, rstdValue)再Mul(y, gamma, y)把结果经CopyOutY写回yGm中间结果x也独立写出供上层复用。从FUSED_ADD_RMS_NORM_INIT_ROW_COMMON宏可以看到 kernel 的任务划分逻辑按num_row总行数、num_col归一化维度大小、blockFactor每核行数、rowFactor每轮处理行数、ubFactorUB 缓冲元素数等 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 验证数值正确性。八、使用建议总结参数选取epsilon建议取 1e-6scale默认为 1.0MiniCPM 等带残差缩放权重的模型应按模型配置显式传入。shape 规划x2与x1同 shapegamma对齐x1的后几维归一化维度rstd无需手动设计shape 由x1、gamma共同决定前段维度保留、归一化维度压 1。精度策略FLOAT16/BFLOAT16 输入在 kernel 内部会经 FLOAT32 中间态计算精度可控BF16 输出使用CAST_RINT舍入。调试与验证仓库在tests/ut下提供了 host 侧 infershape/tiling 与 kernel 侧数值单测可作为新硬件适配或参数调整时的回归依据。硬件匹配BFLOAT16 输入在 310P 系列上不受支持请以目标产品的实际支持矩阵为准README 与 def.cpp 中的架构配置可相互印证。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考