的接口原理与实战调用)
人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载AddRmsNorm 是 CANN ops-nn 神经网络算子库中面向大模型场景的归一化融合算子它把 RmsNorm 之前的元素级 Add 算子合并进归一化计算一次完成x x1 x2与RmsNorm(x)两步操作减少张量在内存中的搬入搬出。本文以 norm/add_rms_norm/docs/aclnnAddRmsNorm.md 为核心结合仓库中的源码与测试系统讲解该算子的产品支持情况、数学原理、两段式 aclnn 接口签名、全部参数约束、返回码语义、确定性保证并给出可直接编译运行的 C 调用示例帮助开发者掌握在 NPU 上正确使用 aclnnAddRmsNorm 的完整方法。算子概述为什么需要 AddRmsNormRmsNormRoot Mean Square Layer Normalization是大模型如 Transformer、LLM中常用的归一化操作。与 LayerNorm 不同RmsNorm去掉了减去均值的部分只基于均方根对输入做缩放因此在保证训练稳定性的同时省去了均值归约计算更轻量。其数学定义为$$ x_i x1_{i} x2_{i} $$$$ \operatorname{RmsNorm}(x_i) \frac{x_i}{\operatorname{Rms}(\mathbf{x})} gamma_i, \quad \text { where } \operatorname{Rms}(\mathbf{x})\sqrt{\frac{1}{n} \sum_{i1}^n x_i^2epsilon} $$其中x1、x2为两个输入张量gamma为逐元素缩放权重epsilon为分母中的数值稳定项。算子先完成加法得到x再对x沿归一化维度计算均方根倒数rstd最后逐元素乘上gamma得到输出y。在大模型推理/训练图例如 DeepNorm、Add Norm 结构中残差连接Residual Connection的Add与后续的RmsNorm天然相邻。若分成两个独立算子执行中间结果x需要整体写回 Global Memory 再被下一个算子重新搬入。AddRmsNorm 通过算子融合将两步合并中间结果直接留在片上UB/寄存器级完成归一化从而显著减少搬入搬出开销——这正是该算子被设计出来的核心动机从 op_host/op_api/aclnn_add_rms_norm.cpp 中Add RmsNorm 的融合算子的接口注释可以印证。产品支持情况根据文档aclnnAddRmsNorm 在不同硬件平台上的支持情况如下产品是否支持Ascend 950PR 950DT 系列产品支持Atlas A3 系列产品支持Atlas A2 系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品支持Atlas 训练系列产品不支持仓库根目录 README.md 的补充信息显示除上述产品外Kirin X90 处理器系列产品与 Kirin 9030 处理器系列产品同样支持本算子。文档同时给出了平台相关的类型差异Atlas 推理系列产品上x1、x2、gamma、yOut、xOut的数据类型不支持 BFLOAT16且rstdOut在当前产品使用场景下无效Kirin 系列产品同样不支持 BFLOAT16。两段式接口GetWorkspaceSize 与执行接口与 CANN 大多数 aclnn 算子一样aclnnAddRmsNorm 采用两段式接口设计参见 两段式接口说明必须先调用aclnnAddRmsNormGetWorkspaceSize完成入参校验、构图并返回 workspace 大小与执行器再调用aclnnAddRmsNorm真正下发计算。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)从源码实现看第一段接口内部完成了如下关键动作aclnn_add_rms_norm.cpp创建OpExecutor并依次执行空指针、数据类型、shape、epsilon 属性校验将三个输入通过l0op::Contiguous转为连续 Tensor调用l0op::AddRmsNorm构图把内部计算输出通过l0op::ViewCopy映射到用户传入的yOut、rstdOut、xOut通过executor-GetWorkspaceSize()返回所需 workspace 大小。aclnnAddRmsNormGetWorkspaceSize 参数详解参数说明下表完整列出了第一段接口的全部参数及其约束参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续Tensorx1aclTensor*输入用于 Add 计算的第一个输入对应公式中的x1支持空 TensorFLOAT32、FLOAT16、BFLOAT16ND1-8√x2aclTensor*输入用于 Add 计算的第二个输入对应公式中的x2支持空 Tensorshape 和数据类型需与x1一致FLOAT32、FLOAT16、BFLOAT16ND1-8√gammaaclTensor*输入RmsNorm 的缩放因子权重对应公式中的gamma支持空 Tensor数据类型与x1一致FLOAT32、FLOAT16、BFLOAT16ND1-8√epsilondouble输入添加到分母中的值确保数值稳定对应公式中的epsilon值需大于等于零建议值为 1e-6----yOutaclTensor*输出最后的输出对应公式中的RmsNorm(x)支持空 Tensorshape、数据类型与x1一致FLOAT32、FLOAT16、BFLOAT16ND1-8√rstdOutaclTensor*输出归一化后的标准差的倒数对应公式中Rms(x)的倒数支持空 Tensor维度数与x1保持一致无需 norm 的维度与x1对应维度一致需要 norm 的维度与gamma维度数相同的后几维均为 1FLOAT32ND1-8√xOutaclTensor*输出Add 计算的结果对应公式中的x支持空 Tensorshape、数据类型与x1一致FLOAT32、FLOAT16、BFLOAT16ND1-8√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含算子计算流程-----rstdOut 的 shape 推导规则是理解本算子的关键文档给出了两个具体示例若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)。即gamma的维度对应x1的后几维被归一化的维度rstdOut在这些维度上压缩为 1其余维度与x1保持一致。这一规则在源码中得到了精确印证——add_rms_norm_infershape.cpp 中rstdShape的前xDimNum - gammaDimNum维直接拷贝x1对应维度剩余维度全部置 1同时y与x的 shape 直接继承x1。接口层 aclnn_add_rms_norm.cpp 的CheckGammaAndRstdShape也按同样规则构造expectedGammaShape与expectedRstdShape并逐一比对。Atlas 推理系列产品的平台差异x1、x2、gamma、yOut、xOut的数据类型不支持 BFLOAT16rstdOut在当前产品使用场景下无效。返回值与错误码第一段接口返回aclnnStatus状态码完整语义可参考 aclnn 返回码。在入参校验阶段出现以下场景时分别报错返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 x1、x2、gamma、yOut 是空指针ACLNN_ERR_PARAM_NULLPTR161001当 rstdOut 传入的预置值不为 nullptr 时xOut 传入的预置值为 nullptrACLNN_ERR_PARAM_INVALID161002输入或输出的数据类型不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002输入和输出参数不满足参数说明中的约束从 aclnn_add_rms_norm.cpp 的CheckParams可以看到更细的校验逻辑CheckNotNull会先判断x1/x2/gamma/yOut是否为 nullptrCheckDtypeValid校验三个输入数据类型一致且在支持列表内、yOut与x1类型一致、rstdOut必须为 FLOAT32CheckShapeDim校验维度在 1~8 之间、x1与x2/yOut/xOutshape 一致、gamma维度数不大于x1维度数CheckAttr校验epsilon 0。值得注意的是空指针相关的返回码ACLNN_ERR_PARAM_NULLPTR对应数值 161001参数非法返回码ACLNN_ERR_PARAM_INVALID对应 161002。aclnnAddRmsNorm 执行接口参数说明第二段接口只需 4 个参数负责把第一段构造好的执行器在指定 Stream 上真正执行参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnAddRmsNormGetWorkspaceSize 获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream其实现非常简洁——aclnn_add_rms_norm.cpp 中第二段接口直接调用CommonOpExecutorRun(workspace, workspaceSize, executor, stream)完成计算下发返回值同样为aclnnStatus。约束说明边界值与数据类型组合边界值场景当输入是 Inf 时输出为 Inf当输入是 NaN 时输出为 NaN。该语义符合公式的数值传递规律即归一化不会修正非有限输入。确定性计算aclnnAddRmsNorm 默认确定性实现即相同输入在同一平台上的多次执行结果一致便于调试与精度比对。输入/输出支持组合随产品系列不同而不同Atlas A2 系列产品、Atlas A3 系列产品x1x2gammayOutrstdOutxOutFLOAT32FLOAT32shape 与 x1 后几维保持一致后几维为需要 norm 的维度FLOAT32FLOAT32必选必选FLOAT32FLOAT16FLOAT16同上FLOAT16FLOAT16必选必选FLOAT16BFLOAT16BFLOAT16同上BFLOAT16BFLOAT16必选必选BFLOAT16FLOAT16FLOAT16shape 为 [1, x1 的最后一维]FLOAT16FLOAT16空指针空指针FLOAT16BFLOAT16BFLOAT16shape 为 [1, x1 的最后一维]BFLOAT16BFLOAT16空指针空指针BFLOAT16FLOAT16FLOAT16shape 为 [1, x1 的最后一维]FLOAT16FLOAT16空指针必选FLOAT16BFLOAT16BFLOAT16shape 为 [1, x1 的最后一维]BFLOAT16BFLOAT16空指针必选BFLOAT16Ascend 950PR 950DT 系列产品x1x2gammayOutrstdOutxOutFLOAT32FLOAT32shape 与 x1 后几维保持一致FLOAT32FLOAT32必选必选FLOAT32FLOAT16FLOAT16同上FLOAT16FLOAT16必选必选FLOAT16BFLOAT16BFLOAT16同上BFLOAT16BFLOAT16必选必选BFLOAT16Atlas 推理系列产品x1x2gammayOutrstdOutxOutFLOAT32FLOAT32shape 与 x1 后几维保持一致FLOAT32FLOAT32必选必选FLOAT32FLOAT16FLOAT16同上FLOAT16FLOAT16必选必选FLOAT16FLOAT16FLOAT16shape 为 [1, x1 的最后一维]FLOAT16FLOAT16空指针空指针FLOAT16FLOAT16FLOAT16shape 为 [1, x1 的最后一维]FLOAT16FLOAT16空指针必选FLOAT16从源码看三种输出组合模式细心的读者会发现上表反复出现rstdOut 空指针 / xOut 空指针的组合。源码揭示了其背后的设计aclnn_add_rms_norm.cpp 中根据xOut与rstdOut是否为 nullptr 自动推导出三种模式AddRmsNorm 模式mode0默认xOut与rstdOut均非空输出 y、rstd、x 三个结果PreRmsNorm 模式mode1xOut非空、rstdOut为空即先 Add 再 RmsNorm仅输出 y 和 xPostRmsNorm 模式mode2xOut与rstdOut均为空即仅 RmsNorm只输出 y。该 mode 会作为属性透传给内核构图l0op::AddRmsNorm的第 5 个参数并在 tiling 侧add_rms_norm_tiling.cpp通过norm_keyRMS_NORM0、PRE_RMS_NORM100、POST_RMS_NORM1000进一步区分归一化 key 参与算子调度。在开启 Regbase寄存器底座路径时源码强制要求rstdOut与xOut均非空Regbase kernels always write both outputs即 Regbase 内核统一按 AddRmsNorm 模式执行。调用示例完整的 C 工程文档给出了完整可编译的示例仓库中的 examples/test_aclnn_add_rms_norm.cpp 与文档示例保持一致。下面按步骤拆解这个示例完整代码见上具体编译与执行过程参考 编译与运行样例。#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); // 调用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::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手册 // 根据自己的实际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::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; 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; // 创建x1 aclTensor ret CreateAclTensor(x1HostData, xShape, x1DeviceAddr, aclDataType::ACL_FLOAT, x1); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建x2 aclTensor ret CreateAclTensor(x2HostData, xShape, x2DeviceAddr, aclDataType::ACL_FLOAT, x2); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建gamma aclTensor ret CreateAclTensor(gammaHostData, gammaShape, gammaDeviceAddr, aclDataType::ACL_FLOAT, gamma); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建y aclTensor ret CreateAclTensor(yHostData, yShape, yDeviceAddr, aclDataType::ACL_FLOAT, y); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建rstd aclTensor ret CreateAclTensor(rstdHostData, rstdShape, rstdDeviceAddr, aclDataType::ACL_FLOAT, rstd); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建x aclTensor ret CreateAclTensor(xHostData, xShape, xDeviceAddr, aclDataType::ACL_FLOAT, x); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用CANN算子库API需要修改为具体的API uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnAddRmsNorm第一段接口 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); } // 调用aclnnAddRmsNorm第二段接口 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侧需要根据具体API的接口定义修改 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需要根据具体API的接口定义修改 aclDestroyTensor(x1); aclDestroyTensor(x2); aclDestroyTensor(gamma); aclDestroyTensor(y); aclDestroyTensor(rstd); aclDestroyTensor(x); // 7. 释放device资源需要根据具体API的接口定义修改 aclrtFree(x1DeviceAddr); aclrtFree(x2DeviceAddr); aclrtFree(xDeviceAddr); aclrtFree(gammaDeviceAddr); aclrtFree(yDeviceAddr); aclrtFree(rstdDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例关键点解读资源初始化步骤 1aclInit→aclrtSetDevice→aclrtCreateStream是 CANN 编程的固定起手式Tensor 构造步骤 2CreateAclTensor模板完成申请 Device 内存 → Host 数据拷贝到 Device → 计算连续 strides →aclCreateTensor创建 aclTensor四步。本例使用 shape 为(2, 16)的 FLOAT32 输入、(16)的 gamma对应rstdshape 为(2, 1)——正好落在前面介绍的 rstdOut shape 推导规则上两段式调用步骤 3先调第一段接口拿到workspaceSize与executor若workspaceSize 0则用aclrtMalloc申请 workspace再调第二段接口执行。若传入空 Tensor第一段接口会在构图前识别x1-IsEmpty() || gamma-IsEmpty()workspace 大小会被置 0 并直接返回成功无需申请内存同步与取数步骤 4、5aclrtSynchronizeStream等待任务完成再通过aclrtMemcpy将结果从 Device 拷回 Host 打印资源释放步骤 6、7依次aclDestroyTensor销毁 aclTensoraclrtFree释放 Device 内存含 workspace最后销毁 Stream、Reset Device 并aclFinalize。编译与运行示例的编译与执行遵循 CANN aclnn 算子样例的通用流程参考 编译与运行样例将上述代码与头文件路径、libascendcl等链接库配合编译生成可执行文件后在已安装 CANN 工具链的昇腾环境上运行运行前需确保deviceId示例中为 0对应的 NPU 设备可用。仓库的 examples/test_aclnn_add_rms_norm.cpp 即为可直接参照的工程文件。图模式调用与算子 IR 注册除 aclnn 接口外AddRmsNorm 还支持图模式调用。算子原型在 op_graph/add_rms_norm_proto.h 中通过REG_OP注册REG_OP(AddRmsNorm) .INPUT(x1, TensorType({DT_FLOAT, DT_FLOAT16, DT_BF16})) .INPUT(x2, TensorType({DT_FLOAT, DT_FLOAT16, DT_BF16})) .INPUT(gamma, TensorType({DT_FLOAT, DT_FLOAT16, DT_BF16})) .OUTPUT(y, TensorType({DT_FLOAT, DT_FLOAT16, DT_BF16})) .OUTPUT(rstd, TensorType({DT_FLOAT, DT_FLOAT, DT_FLOAT})) .OUTPUT(x, TensorType({DT_FLOAT, DT_FLOAT16, DT_BF16})) .ATTR(epsilon, Float, 1e-6f) .OP_END_FACTORY_REG(AddRmsNorm)该原型声明了与 aclnn 接口一致的三输入三输出并将epsilon声明为Float 类型的可选属性默认值 1e-6与文档建议值为 1e-6呼应。README 中列出的调用方式对照关系如下调用方式样例代码说明aclnn 接口test_aclnn_add_rms_norm通过 aclnnAddRmsNorm 接口方式调用 AddRmsNorm 算子图模式-通过 算子IR 构图方式调用 AddRmsNorm 算子底层实现从 tiling 到 kernel 的分发路径多套内核实现与 tiling 策略从 add_rms_norm_tiling.h 可以看到AddRmsNorm 针对不同场景注册了多套 tiling 数据与内核AddRmsNormkey0通用路径tiling 字段包括num_row、num_col、block_factor、row_factor、ub_factor、epsilon、avg_factor等按行分块、多核并行AddRmsNorm_1000/2000/3000/4000/5000Regbase 路径分别对应整行装载R 全载、按 UB 切块、按 R 分核、每核遍历全部 A 行、物理转置后沿 A 向量化、R 为空时仅切分写出 rstd 五种策略每种内核均同时注册了InplaceAddRmsNorm变体支持原地计算场景。这些 tiling 结构中的epsilon、avgFactor 1/numCol即公式中 1/n 的预计算值直接参与核函数计算add_rms_norm.cpp 中的KernelAddRmsNorm类在Init阶段即从 tiling 取出这些字段。tiling 阶段还会根据归一化维度与平台生成差异化策略例如 add_rms_norm_tiling.cpp 中的getPerformanceFlag仅在 Ascend 910B 平台、2~3 维输入、单维 gamma、FP16/BF16 且列数 ≤ 5120 时启用性能路径。平台二进制配置不同产品通过 op_host/config 下的 JSON 配置选择算子二进制例如ascend950目录中的 add_rms_norm_binary.json 为 fp16/fp32/bf16 三种类型分别声明了AddRmsNorm_fp16、AddRmsNorm_fp32、AddRmsNorm_bf16二进制每个条目列出三输入三输出的 dtype/formatND与epsilon 0.000001默认属性。这也是输入输出支持组合表在编译产物层面的落地体现。测试与验证仓库为 AddRmsNorm 提供了完整的测试体系可作为精度验证与二次开发的参考UT单算子测试tests/ut/op_host/op_api/test_aclnn_add_rms_norm.cpp 覆盖 aclnn 接口的参数校验与执行tests/ut/op_host/test_AddRmsNorm_infershape.cpp 与 tests/ut/op_host/test_add_rms_norm_tiling.cpp 分别验证 shape 推导与 tiling 计算kernel 侧还有 tests/ut/op_kernel/test_add_rms_norm.cpp 与 tests/ut/op_kernel/test_add_rms_norm_regbase.cppST系统测试tests/st/aclnnAddRmsNorm/executor_aclnnAddRmsNorm.py 配合atk_aclnnAddRmsNorm.json用例描述执行端到端验证tests/assets/golden.py 提供 golden 参考实现可按公式直接计算期望输出用于精度对比。总结aclnnAddRmsNorm 是 CANN ops-nn 中一个典型的融合归一化算子将大模型残差连接中的 Add 与 RmsNorm 合并为一次片上计算返回归一化结果y、均方根倒数rstd与加法结果x三个输出。使用时需严格遵循两段式接口流程注意三个平台维度的差异产品支持范围Atlas 200I/500 A2 与 Atlas 训练系列不支持、数据类型Atlas 推理系列与 Kirin 系列不支持 BFLOAT16以及rstdOut/xOut 的可选组合三种输出模式由是否传空指针自动推导。结合仓库的 infershape、tiling、kernel 与测试代码开发者可以快速定位问题并完成算子的集成、调试与精度验证。赞分享人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载相关推荐CANN ops-transformer 算子实战aclnnMoeDistributeCombineAddRmsNormV2 通信整合与 AddRmsNorm 融合算子详解CANN ops transformer 算子实战aclnnMoeDistributeCombineAddRmsNormV2 通信整合与 AddRmsNor算子库人工智能大模型深度学习CANNAscendCANN ops-nn AddLayerNormGrad 算子全解析Add 与 LayerNormGrad 融合反向计算的原理、接口与实现CANN ops nn AddLayerNormGrad 算子全解析Add 与 LayerNormGrad 融合反向计算的原理、接口与实现 AddLayerN人工智能算子库深度学习CANNAscendCANN ops-nn 融合算子 FusedAddRmsNorm 全解析ScaledAdd 与 RmsNorm 融合原理、aclnn 两段式接口调用与 NPU 源码实现CANN ops nn 融合算子 FusedAddRmsNorm 全解析ScaledAdd 与 RmsNorm 融合原理、aclnn 两段式接口调用与 NPU人工智能算子库深度学习CANNAscend上一篇10分钟掌握Godot PCK解包提取游戏资源的完整实战指南下一篇reth 可观测性实战基于 etc 目录搭建 Prometheus Grafana Loki 监控与日志栈创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考