人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载AddRmsNorm 是 CANN ops-nn 算子库中面向大模型LLM场景提供的一个融合算子它将残差连接Add与 RMS 归一化RmsNorm合并为一次内核调度从而减少张量在全局内存与片上缓存之间的搬入搬出操作。本文以 norm/add_rms_norm 目录下的 README.md 与 aclnnAddRmsNorm.md 为主体结合算子定义、tiling 与 kernel 源码完整讲解其产品支持情况、数学模型、参数语义、两段式 aclnn 调用流程以及 Device 侧执行细节读完即可在 NPU 上正确配置并调用该算子。一、算子背景与设计动机在大模型 Transformer 类网络中残差连接加归一化是每个 Attention 与 MLP 子层几乎必然出现的计算模式。传统的实现方式是先用独立的 Add 算子完成x x1 x2再把结果整体读入内存交给归一化算子处理同一份数据需要经历写出-再读入两个过程带来额外的访存开销。RmsNormRoot Mean Square Normalization本身是大模型常用的归一化操作相比 LayerNorm 去掉了减去均值的部分仅基于均方根对向量做缩放。AddRmsNorm 算子将 RmsNorm 之前的 Add 算子直接融合进来把加法结果留在片上完成后续归一化从而减少搬入搬出数据搬运操作。从源码结构看该算子内部存在 ADD_RMS_NORM、PRE_RMS_NORM、POST_RMS_NORM 等多种计算模式见 op_host/add_rms_norm_tiling.cpp分别对应完整输出、仅输出归一化结果等不同场景灵活支撑前向计算与后续反向如 InplaceAddRmsNorm的复用。二、产品支持情况根据 README.md 的产品支持矩阵AddRmsNorm 的硬件适配情况如下产品是否支持Ascend 950PR 950DT 系列产品√Atlas A3 系列产品√Atlas A2 系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品√Atlas 训练系列产品×Kirin X90 处理器系列产品√Kirin 9030 处理器系列产品√这一支持矩阵与算子注册代码一一对应。在 op_host/add_rms_norm_def.cpp 中可以看到算子通过AICore().AddConfig()注册了ascend910b、ascend910_93、ascend310p、kirinx90、kirin9030、ascend950、ascend350等多套 AICore 配置其中ascend310pAtlas 推理系列及其同配置的 Kirin 平台对输入输出数据类型做了裁剪仅支持 FLOAT16 与 FLOAT32不支持 BFLOAT16。注意README 的表格中Atlas 训练系列产品910 系列标注为不支持但源码同时注册了ascend910b与ascend910_93两套配置二者对应关系以实际发布版本的适配清单为准本文以仓库源码与文档共同表述的事实为准。三、数学模型与功能说明AddRmsNorm 的完整计算分为两步残差相加将两个输入逐元素相加。$$ x_i x1_i x2_i $$均方根归一化对相加结果按最后一组需要归一的维度计算均方根并用其倒数缩放最后乘上缩放因子权重。$$ \operatorname{RmsNorm}(x_i)\frac{x_i}{\operatorname{Rms}(\mathbf{x})} g_i, \quad \text { where } \operatorname{Rms}(\mathbf{x})\sqrt{\frac{1}{n} \sum_{i1}^n x_i^2\varepsilon} $$其中 $\varepsilon$ 为epsilong为gamma缩放因子n为参与归一化的最后一维或后几维的元素个数。从算子 IR 定义 op_graph/add_rms_norm_proto.h 中可看到更简洁的等价描述x x1 x2rstd rsqrt(mean(x^2, reduce_axis, keepdimsTrue) epsilon)y gamma * (x * rstd)。这清晰地揭示了三个输出y、rstd、x各自的语义y归一化后的最终结果rstd归一化后标准差的倒数即Rms(x)的倒数对应反向传播中常用的中间量xAdd 相加的结果同样常被反向计算复用。融合算子一次性产出三个输出避免反向阶段重新计算 Add 与统计量是其在 LLM 训练场景下价值所在。四、参数说明算子图模式下表完整列出 AddRmsNorm 在图模式算子 IR下的输入、输出与属性参数与 README.md 的参数表一致参数名输入/输出/属性描述数据类型数据格式x1输入用于 Add 计算的第一个输入对应公式中的x1。FLOAT32、FLOAT16、BFLOAT16NDx2输入用于 Add 计算的第二个输入对应公式中的x2。FLOAT32、FLOAT16、BFLOAT16NDgamma输入RmsNorm 的缩放因子权重对应公式中的g。shape 需要与x1后几维保持一致后几维为x1需要 norm 的维度。FLOAT32、FLOAT16、BFLOAT16NDepsilon可选属性添加到分母中的值用于数值稳定、防止除 0 错误需大于等于零对应公式中的eps。默认值为 1e-6。FLOAT-y输出最终的归一化输出对应公式中的RmsNorm(x)。FLOAT32、FLOAT16、BFLOAT16NDrstd输出归一化后标准差的倒数对应公式中Rms(x)的倒数。FLOAT32NDx输出Add 计算的结果对应公式中的x。FLOAT32、FLOAT16、BFLOAT16ND平台差异约束需在配置参数时注意Atlas 推理系列产品x1、x2、gamma、y、x的数据类型不支持 BFLOAT16。Kirin X90 处理器系列产品、Kirin 9030 处理器系列产品x1、x2、gamma、y、x的数据类型同样不支持 BFLOAT16。上述约束与 add_rms_norm_def.cpp 中ascend310p配置仅注册{ge::DT_FLOAT16, ge::DT_FLOAT}的实现完全吻合。关于gamma的 shape 约定可以这样理解gamma的维度数决定了被归一化的后几维其余维度视为 batch/行方向逐行独立做归一化。例如x1shape 为(2, 3, 4, 8)时gamma可取(8)或(4, 8)分别表示只对最后一维、或对后两维联合做均方根归一化。五、aclnn 接口调用实战5.1 两段式接口原型与 CANN 其他算子一致aclnnAddRmsNorm 采用两段式接口设计仓库内对应说明见 docs/zh/context/two_phase_api.md必须先调用aclnnAddRmsNormGetWorkspaceSize获取计算所需的 workspace 大小与执行器再调用aclnnAddRmsNorm真正执行计算。接口声明位于 op_host/op_api/aclnn_add_rms_norm.h。aclnnStatus aclnnAddRmsNormGetWorkspaceSize( const aclTensor *x1, const aclTensor *x2, const aclTensor *gamma, double epsilon, aclTensor *yOut, aclTensor *rstdOut, aclTensor *xOut, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnAddRmsNorm( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)5.2 第一段接口参数详解第一段接口aclnnAddRmsNormGetWorkspaceSize的参数与图模式参数一一对应但以aclTensor形式承载并补充了维度与连续性约束。完整参数语义如下参数输入/输出描述数据类型数据格式维度(shape)非连续 Tensorx1输入Add 的第一个输入对应公式x1。FLOAT32、FLOAT16、BFLOAT16ND1-8√x2输入Add 的第二个输入shape 与数据类型需与x1一致。FLOAT32、FLOAT16、BFLOAT16ND1-8√gamma输入RmsNorm 缩放因子数据类型需与x1一致。FLOAT32、FLOAT16、BFLOAT16ND1-8√epsilon输入分母附加值需大于等于零建议值为 1e-6。----yOut输出归一化输出shape、数据类型与x1一致。FLOAT32、FLOAT16、BFLOAT16ND1-8√rstdOut输出标准差的倒数维度数与x1一致需要 norm 的维度为 1见下例。FLOAT32ND1-8√xOut输出Add 的结果shape、数据类型与x1一致。FLOAT32、FLOAT16、BFLOAT16ND1-8√workspaceSize输出返回需要在 Device 侧申请的 workspace 大小。----executor输出返回 op 执行器包含算子计算流程。----rstdOut 的 shape 推导规则维度数与x1保持一致其中不需要 norm 的维度与x1对应维度相同需要 norm 的维度与gamma维度数相同的后几维全部置为 1。官方示例x1shape(2, 3, 4, 8)、gammashape(8)→rstdOutshape(2, 3, 4, 1)x1shape(2, 3, 4, 8)、gammashape(4, 8)→rstdOutshape(2, 3, 1, 1)。这一规则与 add_rms_norm_infershape.cpp 中的 InferShape 实现一致rstd的每个维度当rmsIdx xDimNum - gammaDimNum时继承x1对应维度否则置 1。特殊语义rstdOut与xOut均可传入nullptr。当rstdOut传入nullptr时该输出无效xOut传入nullptr时同理。yOut、x1、x2、gamma支持空 Tensorshape 中有 0 维度。5.3 返回码与校验第一段接口完成入参校验出现以下场景时报错返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的x1、x2、gamma、yOut为空指针或当rstdOut传入的预置值不为nullptr时xOut传入的预置值为nullptrACLNN_ERR_PARAM_INVALID161002输入或输出的数据类型不在支持范围之内或输入输出参数不满足参数说明中的约束aclnn 返回码的完整含义可参考仓库文档 docs/zh/context/aclnn_return_code.md。5.4 第二段接口参数详解参数输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。5.5 完整调用示例仓库提供了可直接编译运行的示例程序 examples/test_aclnn_add_rms_norm.cppUT 测试 tests/ut/op_host/op_api/test_aclnn_add_rms_norm.cpp 也复用了同样的调用骨架。下面给出带注释的完整流程以 FLOAT32、shape(2, 16)为例#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_add_rms_norm.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::vectorint64_t shape) { int64_t shape_size 1; for (auto i : shape) { shape_size * i; } return shape_size; } int Init(int32_t deviceId, aclrtStream* stream) { // 固定写法资源初始化 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); aclFinalize(); return ret); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); aclrtResetDevice(deviceId); aclFinalize(); 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); // 申请 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); // 将 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::vectorint64_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 API 手册 int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret 0, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出 std::vectorint64_t xShape {2, 16}; std::vectorint64_t gammaShape {16}; std::vectorint64_t yShape {2, 16}; std::vectorint64_t rstdShape {2, 1}; // 注意需要 norm 的维度为 1 void* x1DeviceAddr nullptr; void* x2DeviceAddr nullptr; void* gammaDeviceAddr nullptr; void* yDeviceAddr nullptr; void* rstdDeviceAddr nullptr; void* xDeviceAddr 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 {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 rstdHostData {1, 2}; std::vectorfloat xHostData {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}; float epsilon 1e-6; ret CreateAclTensor(x1HostData, xShape, x1DeviceAddr, aclDataType::ACL_FLOAT, x1); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(x2HostData, xShape, x2DeviceAddr, aclDataType::ACL_FLOAT, x2); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(gammaHostData, gammaShape, gammaDeviceAddr, aclDataType::ACL_FLOAT, gamma); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(yHostData, yShape, yDeviceAddr, aclDataType::ACL_FLOAT, y); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(rstdHostData, rstdShape, rstdDeviceAddr, aclDataType::ACL_FLOAT, rstd); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(xHostData, xShape, xDeviceAddr, aclDataType::ACL_FLOAT, x); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用 CANN 算子库 API两段式 uint64_t workspaceSize 0; aclOpExecutor* executor; // 第一段接口计算 workspace 大小并获取执行器 ret aclnnAddRmsNormGetWorkspaceSize(x1, x2, gamma, epsilon, y, rstd, x, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnAddRmsNormGetWorkspaceSize 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); } // 第二段接口执行计算 ret aclnnAddRmsNorm(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnAddRmsNorm 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 侧 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); 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(y result[%ld] is: %f\n, i, resultData[i]); } // 6. 释放 aclTensor aclDestroyTensor(x1); aclDestroyTensor(x2); aclDestroyTensor(gamma); aclDestroyTensor(y); aclDestroyTensor(rstd); aclDestroyTensor(x); // 7. 释放 device 资源 aclrtFree(x1DeviceAddr); aclrtFree(x2DeviceAddr); aclrtFree(xDeviceAddr); aclrtFree(gammaDeviceAddr); aclrtFree(yDeviceAddr); aclrtFree(rstdDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }该示例的编译与运行方式遵循 CANN 通用样例流程可参考仓库文档 docs/zh/context/compile_and_run_sample.md。工程侧还存在对应的 ST 用例tests/st/aclnnAddRmsNorm/executor_aclnnAddRmsNorm.py与 golden 脚本tests/assets/golden.py可对照验证结果。六、数据类型与输出组合约束6.1 各产品的输入输出组合aclnn 接口下x1、x2、gamma、yOut、rstdOut、xOut的支持组合因产品而异Atlas A2 系列产品、Atlas A3 系列产品x1x2gammayOutrstdOutxOutFLOAT32FLOAT32shape 与x1后几维一致FLOAT32FLOAT32必选必选FLOAT32FLOAT16FLOAT16shape 与x1后几维一致FLOAT16FLOAT16必选必选FLOAT16BFLOAT16BFLOAT16shape 与x1后几维一致BFLOAT16BFLOAT16必选必选BFLOAT16FLOAT16FLOAT16shape 为[1, x1 的最后一维]FLOAT16FLOAT16空指针空指针FLOAT16BFLOAT16BFLOAT16shape 为[1, x1 的最后一维]BFLOAT16BFLOAT16空指针空指针BFLOAT16FLOAT16FLOAT16shape 为[1, x1 的最后一维]FLOAT16FLOAT16空指针必选FLOAT16BFLOAT16BFLOAT16shape 为[1, x1 的最后一维]BFLOAT16BFLOAT16空指针必选BFLOAT16Ascend 950PR 950DT 系列产品x1x2gammayOutrstdOutxOutFLOAT32FLOAT32shape 与x1后几维一致FLOAT32FLOAT32必选必选FLOAT32FLOAT16FLOAT16shape 与x1后几维一致FLOAT16FLOAT16必选必选FLOAT16BFLOAT16BFLOAT16shape 与x1后几维一致BFLOAT16BFLOAT16必选必选BFLOAT16Atlas 推理系列产品x1x2gammayOutrstdOutxOutFLOAT32FLOAT32shape 与x1后几维一致FLOAT32FLOAT32必选必选FLOAT32FLOAT16FLOAT16shape 与x1后几维一致FLOAT16FLOAT16必选必选FLOAT16FLOAT16FLOAT16shape 为[1, x1 的最后一维]FLOAT16FLOAT16空指针空指针FLOAT16FLOAT16FLOAT16shape 为[1, x1 的最后一维]FLOAT16FLOAT16空指针必选FLOAT166.2 边界值行为当输入是 Inf 时输出为 Inf。当输入是 NaN 时输出为 NaN。6.3 确定性aclnnAddRmsNorm 默认采用确定性实现多次运行同一输入可得到一致的输出结果。七、Device 侧执行原理源码级解析7.1 从算子注册到图模式构图除了 aclnn 接口调用方式AddRmsNorm 还支持图模式调用通过算子 IR op_graph/add_rms_norm_proto.h 中REG_OP(AddRmsNorm)注册的AddRmsNorm算子节点构图图中可直接设置epsilon属性。对应地框架侧还提供了 ONNX 插件 framework/npu_add_rms_norm_onnx_plugin.cpp用于将 ONNX 模型中的 Add RmsNorm 子图识别并融合为 AddRmsNorm 算子节点。7.2 Shape 推导op_host/add_rms_norm_infershape.cpp 实现了InferShape4AddRmsNorm与InferDataType4AddRmsNorm输出y与x的 shape 直接继承x1的 shape输出rstd的维度数与x1一致前xDimNum - gammaDimNum维继承x1对应维度后gammaDimNum维全部置 1输出y与x的数据类型继承x1输出rstd固定为DT_FLOATFLOAT32。同一实现同时注册给了AddRmsNorm与InplaceAddRmsNorm原地版本用于支持反向传播场景下的复用说明该算子与同目录 norm 家族的其他融合算子共享了部分归一化基础设施。7.3 Tiling 调度策略tiling 在 Host 侧完成负责把全局 shape 拆分为适配各 AI Core 的切块并写入 tiling data。op_host/add_rms_norm_tiling.h 定义了多套 tiling 数据结构每一套对应一种内核策略AddRMSNormTilingData通用策略字段包括num_row、num_col、block_factor每核行数、row_factor每次循环处理行数、ub_factor单次进 UB 的列数、epsilon、avg_factor以及 fp16/fp32 分离的乘法循环参数等AddRMSNormRegbaseRFullLoadTilingDatakey1000R行维度全部载入AddRMSNormRegbaseTilingDatakey2000基础 regbase 策略AddRMSNormRegbaseSplitARTilingDatakey3000按 R 分核每核遍历全部归一化行AddRMSNormRegbaseTransTilingDatakey4000物理转置后沿归一化维度向量化AddRMSNormRegbaseReduceEmptyTilingDatakey5000R 为空时仅切分并写出rstd。从 add_rms_norm_tiling.cpp 的常量可以推断tiling 会根据数据类型FP16/FP32/BF16 分别对应 dtypeKey 1/2/3、归一化列数num_col、UB 容量与 16/8 元素对齐粒度BLOCK_ALIGN_NUM16、FLOAT_BLOCK_ALIGN_NUM8等条件在多种模式MODE_NORMAL、MODE_SPLIT_D、MODE_MERGE_N、MODE_SINGLE_N、MODE_MULTI_N中选择合适的切分方案并通过getPerformanceFlag在 910B 上针对 2/3 维、低行数、FP16/BF16 的典型 LLM shape 开启性能模式。7.4 Kernel 执行流程Device 侧内核模板 op_kernel/add_rms_norm.h 中KernelAddRmsNormT, MODE是核心实现Init 阶段根据 tiling data 读取numRow、numCol、blockFactor、rowFactor、ubFactor、epsilon等参数按GetBlockIdx()计算当前核负责的行区间blockIdx GetBlockNum()-1时处理blockFactor行最后一个核处理剩余尾行并按行偏移建立x1Gm、x2Gm、gammaGm、yGm、rstdGm、xGm的全局内存视图UB 缓冲区为输入队列、gamma 队列、输出 y 队列、rstd 队列申请片上缓冲区FP16/BF16 输入额外申请 FP32 转换缓冲xFp32Buf与平方和缓冲sqxBuf供归一化统计计算使用Process 阶段先CopyInGamma载入权重再按rowFactor分块循环调用SubProcess每个子块内逐行完成CopyIn → Compute → CopyOutY在 ADD_RMS_NORM 模式下最后统一CopyOutRstd写出标准差倒数多模式支持通过MODE ADD_RMS_NORM_MODE / PRE_RMS_NORM_MODE等编译期分支决定是否写出rstd与x与 tiling 阶段根据输出 shape 是否为 0 计算的norm_key0 表示完整模式、100 表示 PRE、1000 表示 POST相配合。kernel 目录下还包含多种架构专用实现arch35/下按 1000/2000/3000/4000/5000 key 拆分的add_rms_norm_regbase*.h系列以及面向 950 平台的 APTAoT/预编译实现 add_rms_norm_apt.cpp配合 op_host/config 下各产品的 binary 配置文件add_rms_norm_binary.json覆盖 ascend910b、ascend910_93、ascend350、ascend950、ascend310p、kirinx90、kirin9030完成算子的二进制形态分发。八、测试与验证仓库为 AddRmsNorm 提供了完整的测试覆盖可用于验证本文所述行为UTHost 侧tests/ut/op_host/test_AddRmsNorm_infershape.cpp 验证 shape 推导包括rstd置 1 规则tests/ut/op_host/test_add_rms_norm_tiling.cpp 验证 tiling 参数计算tests/ut/op_host/op_api/test_aclnn_add_rms_norm.cpp 验证两段式接口调用UTKernel 侧tests/ut/op_kernel/test_add_rms_norm.cpp 与 test_add_rms_norm_regbase.cpp 验证 Device 内核计算结果ST系统测试tests/st/aclnnAddRmsNorm/executor_aclnnAddRmsNorm.py 配合 atk_aclnnAddRmsNorm.json 执行端到端用例arch35 目录下还有ttk_kernel_add_rms_norm_st.csv与ttk_kernel_add_rms_norm_perf.csv记录 35 架构的内核 ST 与性能测试项。九、使用建议与注意事项优先使用融合算子凡是在大模型中出现的x1 x2 → RmsNorm模式直接用 AddRmsNorm 替代 Add RmsNorm 两个算子可减少一次全局内存往返正确设置 gamma 与 rstd 的 shapegamma的维度数决定归一化范围rstdOut需要 norm 的维度必须为 1否则会触发参数校验错误161002留意产品差异Atlas 推理系列与 Kirin 平台不支持 BFLOAT16Atlas 200I/500 A2 推理产品与 Atlas 训练系列910 系列在 README 中标注不支持合理利用可选输出不需要rstd或x时传nullptr让 tiling 走 PRE/POST 裁剪模式减少写回开销epsilon 取值需大于等于零默认 1e-6建议保持默认或使用与训练一致的数值避免数值稳定性问题确定性该算子默认确定性实现适合对结果可复现性有要求的训练场景。十、小结AddRmsNorm 通过Add RmsNorm的算子融合将大模型最频繁的残差归一化模式压缩为单次内核执行并额外产出rstd与x两个中间量供反向复用。本文从产品支持矩阵、数学模型、图模式参数、两段式 aclnn 接口、Device 侧 tiling 与 kernel 执行链路四个层面做了完整拆解并给出了可直接运行的调用示例与测试入口。读者可以进一步阅读 README.md、aclnnAddRmsNorm.md或直接基于 examples/test_aclnn_add_rms_norm.cpp 展开实验。赞分享人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载相关推荐GBrain 正则纪律约定如何判断该用正则还是让模型做判断GBrain 正则纪律约定如何判断该用正则还是让模型做判断 正则表达式regex是每一个 AI Agent 工程里最常见的工具之一也是被误用得最多的工具人工智能算子库深度学习CANNAscendCANN ops-nn 算子指南aclnnAddRmsNorm 融合算子Add RmsNorm的接口原理与实战调用CANN ops nn 算子指南aclnnAddRmsNorm 融合算子Add RmsNorm的接口原理与实战调用 AddRmsNorm 是 CANN人工智能算子库深度学习CANNAscendCube Sandbox AgentHub数字助手部署与配置实战基于 OpenClaw 的 AI Agent 托管指南Cube Sandbox AgentHub数字助手部署与配置实战基于 OpenClaw 的 AI Agent 托管指南 AgentHub数字助手是 C人工智能算子库深度学习CANNAscend上一篇大麦自动抢票工具 ticket-purchase 实战Selenium Appium 双端毫秒级下单下一篇如何永久保存微信聊天记录WeChatMsg完整指南帮你轻松实现数据自由创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载AddRmsNorm 是 CANN ops-nn 算子库中面向大模型LLM场景提供的一个融合算子它将残差连接Add与 RMS 归一化RmsNorm合并为一次内核调度从而减少张量在全局内存与片上缓存之间的搬入搬出操作。本文以 norm/add_rms_norm 目录下的 README.md 与 aclnnAddRmsNorm.md 为主体结合算子定义、tiling 与 kernel 源码完整讲解其产品支持情况、数学模型、参数语义、两段式 aclnn 调用流程以及 Device 侧执行细节读完即可在 NPU 上正确配置并调用该算子。一、算子背景与设计动机在大模型 Transformer 类网络中残差连接加归一化是每个 Attention 与 MLP 子层几乎必然出现的计算模式。传统的实现方式是先用独立的 Add 算子完成x x1 x2再把结果整体读入内存交给归一化算子处理同一份数据需要经历写出-再读入两个过程带来额外的访存开销。RmsNormRoot Mean Square Normalization本身是大模型常用的归一化操作相比 LayerNorm 去掉了减去均值的部分仅基于均方根对向量做缩放。AddRmsNorm 算子将 RmsNorm 之前的 Add 算子直接融合进来把加法结果留在片上完成后续归一化从而减少搬入搬出数据搬运操作。从源码结构看该算子内部存在 ADD_RMS_NORM、PRE_RMS_NORM、POST_RMS_NORM 等多种计算模式见 op_host/add_rms_norm_tiling.cpp分别对应完整输出、仅输出归一化结果等不同场景灵活支撑前向计算与后续反向如 InplaceAddRmsNorm的复用。二、产品支持情况根据 README.md 的产品支持矩阵AddRmsNorm 的硬件适配情况如下产品是否支持Ascend 950PR 950DT 系列产品√Atlas A3 系列产品√Atlas A2 系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品√Atlas 训练系列产品×Kirin X90 处理器系列产品√Kirin 9030 处理器系列产品√这一支持矩阵与算子注册代码一一对应。在 op_host/add_rms_norm_def.cpp 中可以看到算子通过AICore().AddConfig()注册了ascend910b、ascend910_93、ascend310p、kirinx90、kirin9030、ascend950、ascend350等多套 AICore 配置其中ascend310pAtlas 推理系列及其同配置的 Kirin 平台对输入输出数据类型做了裁剪仅支持 FLOAT16 与 FLOAT32不支持 BFLOAT16。注意README 的表格中Atlas 训练系列产品910 系列标注为不支持但源码同时注册了ascend910b与ascend910_93两套配置二者对应关系以实际发布版本的适配清单为准本文以仓库源码与文档共同表述的事实为准。三、数学模型与功能说明AddRmsNorm 的完整计算分为两步残差相加将两个输入逐元素相加。$$ x_i x1_i x2_i $$均方根归一化对相加结果按最后一组需要归一的维度计算均方根并用其倒数缩放最后乘上缩放因子权重。$$ \operatorname{RmsNorm}(x_i)\frac{x_i}{\operatorname{Rms}(\mathbf{x})} g_i, \quad \text { where } \operatorname{Rms}(\mathbf{x})\sqrt{\frac{1}{n} \sum_{i1}^n x_i^2\varepsilon} $$其中 $\varepsilon$ 为epsilong为gamma缩放因子n为参与归一化的最后一维或后几维的元素个数。从算子 IR 定义 op_graph/add_rms_norm_proto.h 中可看到更简洁的等价描述x x1 x2rstd rsqrt(mean(x^2, reduce_axis, keepdimsTrue) epsilon)y gamma * (x * rstd)。这清晰地揭示了三个输出y、rstd、x各自的语义y归一化后的最终结果rstd归一化后标准差的倒数即Rms(x)的倒数对应反向传播中常用的中间量xAdd 相加的结果同样常被反向计算复用。融合算子一次性产出三个输出避免反向阶段重新计算 Add 与统计量是其在 LLM 训练场景下价值所在。四、参数说明算子图模式下表完整列出 AddRmsNorm 在图模式算子 IR下的输入、输出与属性参数与 README.md 的参数表一致参数名输入/输出/属性描述数据类型数据格式x1输入用于 Add 计算的第一个输入对应公式中的x1。FLOAT32、FLOAT16、BFLOAT16NDx2输入用于 Add 计算的第二个输入对应公式中的x2。FLOAT32、FLOAT16、BFLOAT16NDgamma输入RmsNorm 的缩放因子权重对应公式中的g。shape 需要与x1后几维保持一致后几维为x1需要 norm 的维度。FLOAT32、FLOAT16、BFLOAT16NDepsilon可选属性添加到分母中的值用于数值稳定、防止除 0 错误需大于等于零对应公式中的eps。默认值为 1e-6。FLOAT-y输出最终的归一化输出对应公式中的RmsNorm(x)。FLOAT32、FLOAT16、BFLOAT16NDrstd输出归一化后标准差的倒数对应公式中Rms(x)的倒数。FLOAT32NDx输出Add 计算的结果对应公式中的x。FLOAT32、FLOAT16、BFLOAT16ND平台差异约束需在配置参数时注意Atlas 推理系列产品x1、x2、gamma、y、x的数据类型不支持 BFLOAT16。Kirin X90 处理器系列产品、Kirin 9030 处理器系列产品x1、x2、gamma、y、x的数据类型同样不支持 BFLOAT16。上述约束与 add_rms_norm_def.cpp 中ascend310p配置仅注册{ge::DT_FLOAT16, ge::DT_FLOAT}的实现完全吻合。关于gamma的 shape 约定可以这样理解gamma的维度数决定了被归一化的后几维其余维度视为 batch/行方向逐行独立做归一化。例如x1shape 为(2, 3, 4, 8)时gamma可取(8)或(4, 8)分别表示只对最后一维、或对后两维联合做均方根归一化。五、aclnn 接口调用实战5.1 两段式接口原型与 CANN 其他算子一致aclnnAddRmsNorm 采用两段式接口设计仓库内对应说明见 docs/zh/context/two_phase_api.md必须先调用aclnnAddRmsNormGetWorkspaceSize获取计算所需的 workspace 大小与执行器再调用aclnnAddRmsNorm真正执行计算。接口声明位于 op_host/op_api/aclnn_add_rms_norm.h。aclnnStatus aclnnAddRmsNormGetWorkspaceSize( const aclTensor *x1, const aclTensor *x2, const aclTensor *gamma, double epsilon, aclTensor *yOut, aclTensor *rstdOut, aclTensor *xOut, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnAddRmsNorm( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)5.2 第一段接口参数详解第一段接口aclnnAddRmsNormGetWorkspaceSize的参数与图模式参数一一对应但以aclTensor形式承载并补充了维度与连续性约束。完整参数语义如下参数输入/输出描述数据类型数据格式维度(shape)非连续 Tensorx1输入Add 的第一个输入对应公式x1。FLOAT32、FLOAT16、BFLOAT16ND1-8√x2输入Add 的第二个输入shape 与数据类型需与x1一致。FLOAT32、FLOAT16、BFLOAT16ND1-8√gamma输入RmsNorm 缩放因子数据类型需与x1一致。FLOAT32、FLOAT16、BFLOAT16ND1-8√epsilon输入分母附加值需大于等于零建议值为 1e-6。----yOut输出归一化输出shape、数据类型与x1一致。FLOAT32、FLOAT16、BFLOAT16ND1-8√rstdOut输出标准差的倒数维度数与x1一致需要 norm 的维度为 1见下例。FLOAT32ND1-8√xOut输出Add 的结果shape、数据类型与x1一致。FLOAT32、FLOAT16、BFLOAT16ND1-8√workspaceSize输出返回需要在 Device 侧申请的 workspace 大小。----executor输出返回 op 执行器包含算子计算流程。----rstdOut 的 shape 推导规则维度数与x1保持一致其中不需要 norm 的维度与x1对应维度相同需要 norm 的维度与gamma维度数相同的后几维全部置为 1。官方示例x1shape(2, 3, 4, 8)、gammashape(8)→rstdOutshape(2, 3, 4, 1)x1shape(2, 3, 4, 8)、gammashape(4, 8)→rstdOutshape(2, 3, 1, 1)。这一规则与 add_rms_norm_infershape.cpp 中的 InferShape 实现一致rstd的每个维度当rmsIdx xDimNum - gammaDimNum时继承x1对应维度否则置 1。特殊语义rstdOut与xOut均可传入nullptr。当rstdOut传入nullptr时该输出无效xOut传入nullptr时同理。yOut、x1、x2、gamma支持空 Tensorshape 中有 0 维度。5.3 返回码与校验第一段接口完成入参校验出现以下场景时报错返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的x1、x2、gamma、yOut为空指针或当rstdOut传入的预置值不为nullptr时xOut传入的预置值为nullptrACLNN_ERR_PARAM_INVALID161002输入或输出的数据类型不在支持范围之内或输入输出参数不满足参数说明中的约束aclnn 返回码的完整含义可参考仓库文档 docs/zh/context/aclnn_return_code.md。5.4 第二段接口参数详解参数输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。5.5 完整调用示例仓库提供了可直接编译运行的示例程序 examples/test_aclnn_add_rms_norm.cppUT 测试 tests/ut/op_host/op_api/test_aclnn_add_rms_norm.cpp 也复用了同样的调用骨架。下面给出带注释的完整流程以 FLOAT32、shape(2, 16)为例#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_add_rms_norm.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::vectorint64_t shape) { int64_t shape_size 1; for (auto i : shape) { shape_size * i; } return shape_size; } int Init(int32_t deviceId, aclrtStream* stream) { // 固定写法资源初始化 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); aclFinalize(); return ret); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); aclrtResetDevice(deviceId); aclFinalize(); 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); // 申请 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); // 将 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::vectorint64_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 API 手册 int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret 0, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出 std::vectorint64_t xShape {2, 16}; std::vectorint64_t gammaShape {16}; std::vectorint64_t yShape {2, 16}; std::vectorint64_t rstdShape {2, 1}; // 注意需要 norm 的维度为 1 void* x1DeviceAddr nullptr; void* x2DeviceAddr nullptr; void* gammaDeviceAddr nullptr; void* yDeviceAddr nullptr; void* rstdDeviceAddr nullptr; void* xDeviceAddr 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 {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 rstdHostData {1, 2}; std::vectorfloat xHostData {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}; float epsilon 1e-6; ret CreateAclTensor(x1HostData, xShape, x1DeviceAddr, aclDataType::ACL_FLOAT, x1); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(x2HostData, xShape, x2DeviceAddr, aclDataType::ACL_FLOAT, x2); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(gammaHostData, gammaShape, gammaDeviceAddr, aclDataType::ACL_FLOAT, gamma); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(yHostData, yShape, yDeviceAddr, aclDataType::ACL_FLOAT, y); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(rstdHostData, rstdShape, rstdDeviceAddr, aclDataType::ACL_FLOAT, rstd); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(xHostData, xShape, xDeviceAddr, aclDataType::ACL_FLOAT, x); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用 CANN 算子库 API两段式 uint64_t workspaceSize 0; aclOpExecutor* executor; // 第一段接口计算 workspace 大小并获取执行器 ret aclnnAddRmsNormGetWorkspaceSize(x1, x2, gamma, epsilon, y, rstd, x, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnAddRmsNormGetWorkspaceSize 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); } // 第二段接口执行计算 ret aclnnAddRmsNorm(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnAddRmsNorm 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 侧 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); 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(y result[%ld] is: %f\n, i, resultData[i]); } // 6. 释放 aclTensor aclDestroyTensor(x1); aclDestroyTensor(x2); aclDestroyTensor(gamma); aclDestroyTensor(y); aclDestroyTensor(rstd); aclDestroyTensor(x); // 7. 释放 device 资源 aclrtFree(x1DeviceAddr); aclrtFree(x2DeviceAddr); aclrtFree(xDeviceAddr); aclrtFree(gammaDeviceAddr); aclrtFree(yDeviceAddr); aclrtFree(rstdDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }该示例的编译与运行方式遵循 CANN 通用样例流程可参考仓库文档 docs/zh/context/compile_and_run_sample.md。工程侧还存在对应的 ST 用例tests/st/aclnnAddRmsNorm/executor_aclnnAddRmsNorm.py与 golden 脚本tests/assets/golden.py可对照验证结果。六、数据类型与输出组合约束6.1 各产品的输入输出组合aclnn 接口下x1、x2、gamma、yOut、rstdOut、xOut的支持组合因产品而异Atlas A2 系列产品、Atlas A3 系列产品x1x2gammayOutrstdOutxOutFLOAT32FLOAT32shape 与x1后几维一致FLOAT32FLOAT32必选必选FLOAT32FLOAT16FLOAT16shape 与x1后几维一致FLOAT16FLOAT16必选必选FLOAT16BFLOAT16BFLOAT16shape 与x1后几维一致BFLOAT16BFLOAT16必选必选BFLOAT16FLOAT16FLOAT16shape 为[1, x1 的最后一维]FLOAT16FLOAT16空指针空指针FLOAT16BFLOAT16BFLOAT16shape 为[1, x1 的最后一维]BFLOAT16BFLOAT16空指针空指针BFLOAT16FLOAT16FLOAT16shape 为[1, x1 的最后一维]FLOAT16FLOAT16空指针必选FLOAT16BFLOAT16BFLOAT16shape 为[1, x1 的最后一维]BFLOAT16BFLOAT16空指针必选BFLOAT16Ascend 950PR 950DT 系列产品x1x2gammayOutrstdOutxOutFLOAT32FLOAT32shape 与x1后几维一致FLOAT32FLOAT32必选必选FLOAT32FLOAT16FLOAT16shape 与x1后几维一致FLOAT16FLOAT16必选必选FLOAT16BFLOAT16BFLOAT16shape 与x1后几维一致BFLOAT16BFLOAT16必选必选BFLOAT16Atlas 推理系列产品x1x2gammayOutrstdOutxOutFLOAT32FLOAT32shape 与x1后几维一致FLOAT32FLOAT32必选必选FLOAT32FLOAT16FLOAT16shape 与x1后几维一致FLOAT16FLOAT16必选必选FLOAT16FLOAT16FLOAT16shape 为[1, x1 的最后一维]FLOAT16FLOAT16空指针空指针FLOAT16FLOAT16FLOAT16shape 为[1, x1 的最后一维]FLOAT16FLOAT16空指针必选FLOAT166.2 边界值行为当输入是 Inf 时输出为 Inf。当输入是 NaN 时输出为 NaN。6.3 确定性aclnnAddRmsNorm 默认采用确定性实现多次运行同一输入可得到一致的输出结果。七、Device 侧执行原理源码级解析7.1 从算子注册到图模式构图除了 aclnn 接口调用方式AddRmsNorm 还支持图模式调用通过算子 IR op_graph/add_rms_norm_proto.h 中REG_OP(AddRmsNorm)注册的AddRmsNorm算子节点构图图中可直接设置epsilon属性。对应地框架侧还提供了 ONNX 插件 framework/npu_add_rms_norm_onnx_plugin.cpp用于将 ONNX 模型中的 Add RmsNorm 子图识别并融合为 AddRmsNorm 算子节点。7.2 Shape 推导op_host/add_rms_norm_infershape.cpp 实现了InferShape4AddRmsNorm与InferDataType4AddRmsNorm输出y与x的 shape 直接继承x1的 shape输出rstd的维度数与x1一致前xDimNum - gammaDimNum维继承x1对应维度后gammaDimNum维全部置 1输出y与x的数据类型继承x1输出rstd固定为DT_FLOATFLOAT32。同一实现同时注册给了AddRmsNorm与InplaceAddRmsNorm原地版本用于支持反向传播场景下的复用说明该算子与同目录 norm 家族的其他融合算子共享了部分归一化基础设施。7.3 Tiling 调度策略tiling 在 Host 侧完成负责把全局 shape 拆分为适配各 AI Core 的切块并写入 tiling data。op_host/add_rms_norm_tiling.h 定义了多套 tiling 数据结构每一套对应一种内核策略AddRMSNormTilingData通用策略字段包括num_row、num_col、block_factor每核行数、row_factor每次循环处理行数、ub_factor单次进 UB 的列数、epsilon、avg_factor以及 fp16/fp32 分离的乘法循环参数等AddRMSNormRegbaseRFullLoadTilingDatakey1000R行维度全部载入AddRMSNormRegbaseTilingDatakey2000基础 regbase 策略AddRMSNormRegbaseSplitARTilingDatakey3000按 R 分核每核遍历全部归一化行AddRMSNormRegbaseTransTilingDatakey4000物理转置后沿归一化维度向量化AddRMSNormRegbaseReduceEmptyTilingDatakey5000R 为空时仅切分并写出rstd。从 add_rms_norm_tiling.cpp 的常量可以推断tiling 会根据数据类型FP16/FP32/BF16 分别对应 dtypeKey 1/2/3、归一化列数num_col、UB 容量与 16/8 元素对齐粒度BLOCK_ALIGN_NUM16、FLOAT_BLOCK_ALIGN_NUM8等条件在多种模式MODE_NORMAL、MODE_SPLIT_D、MODE_MERGE_N、MODE_SINGLE_N、MODE_MULTI_N中选择合适的切分方案并通过getPerformanceFlag在 910B 上针对 2/3 维、低行数、FP16/BF16 的典型 LLM shape 开启性能模式。7.4 Kernel 执行流程Device 侧内核模板 op_kernel/add_rms_norm.h 中KernelAddRmsNormT, MODE是核心实现Init 阶段根据 tiling data 读取numRow、numCol、blockFactor、rowFactor、ubFactor、epsilon等参数按GetBlockIdx()计算当前核负责的行区间blockIdx GetBlockNum()-1时处理blockFactor行最后一个核处理剩余尾行并按行偏移建立x1Gm、x2Gm、gammaGm、yGm、rstdGm、xGm的全局内存视图UB 缓冲区为输入队列、gamma 队列、输出 y 队列、rstd 队列申请片上缓冲区FP16/BF16 输入额外申请 FP32 转换缓冲xFp32Buf与平方和缓冲sqxBuf供归一化统计计算使用Process 阶段先CopyInGamma载入权重再按rowFactor分块循环调用SubProcess每个子块内逐行完成CopyIn → Compute → CopyOutY在 ADD_RMS_NORM 模式下最后统一CopyOutRstd写出标准差倒数多模式支持通过MODE ADD_RMS_NORM_MODE / PRE_RMS_NORM_MODE等编译期分支决定是否写出rstd与x与 tiling 阶段根据输出 shape 是否为 0 计算的norm_key0 表示完整模式、100 表示 PRE、1000 表示 POST相配合。kernel 目录下还包含多种架构专用实现arch35/下按 1000/2000/3000/4000/5000 key 拆分的add_rms_norm_regbase*.h系列以及面向 950 平台的 APTAoT/预编译实现 add_rms_norm_apt.cpp配合 op_host/config 下各产品的 binary 配置文件add_rms_norm_binary.json覆盖 ascend910b、ascend910_93、ascend350、ascend950、ascend310p、kirinx90、kirin9030完成算子的二进制形态分发。八、测试与验证仓库为 AddRmsNorm 提供了完整的测试覆盖可用于验证本文所述行为UTHost 侧tests/ut/op_host/test_AddRmsNorm_infershape.cpp 验证 shape 推导包括rstd置 1 规则tests/ut/op_host/test_add_rms_norm_tiling.cpp 验证 tiling 参数计算tests/ut/op_host/op_api/test_aclnn_add_rms_norm.cpp 验证两段式接口调用UTKernel 侧tests/ut/op_kernel/test_add_rms_norm.cpp 与 test_add_rms_norm_regbase.cpp 验证 Device 内核计算结果ST系统测试tests/st/aclnnAddRmsNorm/executor_aclnnAddRmsNorm.py 配合 atk_aclnnAddRmsNorm.json 执行端到端用例arch35 目录下还有ttk_kernel_add_rms_norm_st.csv与ttk_kernel_add_rms_norm_perf.csv记录 35 架构的内核 ST 与性能测试项。九、使用建议与注意事项优先使用融合算子凡是在大模型中出现的x1 x2 → RmsNorm模式直接用 AddRmsNorm 替代 Add RmsNorm 两个算子可减少一次全局内存往返正确设置 gamma 与 rstd 的 shapegamma的维度数决定归一化范围rstdOut需要 norm 的维度必须为 1否则会触发参数校验错误161002留意产品差异Atlas 推理系列与 Kirin 平台不支持 BFLOAT16Atlas 200I/500 A2 推理产品与 Atlas 训练系列910 系列在 README 中标注不支持合理利用可选输出不需要rstd或x时传nullptr让 tiling 走 PRE/POST 裁剪模式减少写回开销epsilon 取值需大于等于零默认 1e-6建议保持默认或使用与训练一致的数值避免数值稳定性问题确定性该算子默认确定性实现适合对结果可复现性有要求的训练场景。十、小结AddRmsNorm 通过Add RmsNorm的算子融合将大模型最频繁的残差归一化模式压缩为单次内核执行并额外产出rstd与x两个中间量供反向复用。本文从产品支持矩阵、数学模型、图模式参数、两段式 aclnn 接口、Device 侧 tiling 与 kernel 执行链路四个层面做了完整拆解并给出了可直接运行的调用示例与测试入口。读者可以进一步阅读 README.md、aclnnAddRmsNorm.md或直接基于 examples/test_aclnn_add_rms_norm.cpp 展开实验。赞分享人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载相关推荐GBrain 正则纪律约定如何判断该用正则还是让模型做判断GBrain 正则纪律约定如何判断该用正则还是让模型做判断 正则表达式regex是每一个 AI Agent 工程里最常见的工具之一也是被误用得最多的工具人工智能算子库深度学习CANNAscendCANN ops-nn 算子指南aclnnAddRmsNorm 融合算子Add RmsNorm的接口原理与实战调用CANN ops nn 算子指南aclnnAddRmsNorm 融合算子Add RmsNorm的接口原理与实战调用 AddRmsNorm 是 CANN人工智能算子库深度学习CANNAscendCube Sandbox AgentHub数字助手部署与配置实战基于 OpenClaw 的 AI Agent 托管指南Cube Sandbox AgentHub数字助手部署与配置实战基于 OpenClaw 的 AI Agent 托管指南 AgentHub数字助手是 C人工智能算子库深度学习CANNAscend上一篇大麦自动抢票工具 ticket-purchase 实战Selenium Appium 双端毫秒级下单下一篇如何永久保存微信聊天记录WeChatMsg完整指南帮你轻松实现数据自由创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
