尧图网页设计 
人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载导读AddRmsNormCast 是 CANN ops-nn 神经网络算子库中面向大模型LLM场景的融合归一化算子它把Add逐元素相加、RmsNorm均方根归一化与Cast类型转换三个操作融合为一次 Kernel 执行通过减少数据在内存与计算单元之间的搬入搬出次数来降低算子调用开销。本文以 norm/add_rms_norm_cast/README.md 与配套的 aclnnAddRmsNormCast 接口文档 为主线结合该算子在仓库中的算子定义、shape 推导、tiling 策略、Kernel 分发与测试代码完整讲解其产品支持情况、数学原理、输入输出参数、aclnn 两段式调用方法与图模式构图方式。读完本文你将掌握如何在 Atlas A2/A3、Ascend 950 系列与 Kirin 系列产品上正确构造并调用该算子并能理解其底层实现机制。算子定位为什么要把 Add、RmsNorm、Cast 融合在一起RmsNormRoot Mean Square Normalization是大模型常用的归一化操作。在典型的 Transformer 网络结构中残差相加Add之后往往紧跟 RmsNorm而部分计算图还会在 RmsNorm 之后衔接类型转换Cast以满足后续算子对精度的要求。如果按照原始计算图逐个算子执行Add、RmsNorm、Cast 各自都需要独立申请内存、启动 Kernel、完成数据搬入搬出会引入可观的额外开销。AddRmsNormCast 算子的核心设计目标正是消除这些开销将 AddRmsNorm 之后的 Cast 算子融合进归一化计算中减少搬入搬出操作见 README.md 功能说明。一次 Kernel 调用即可完成求和 → 归一化 → 类型转换同时顺带输出 RmsNorm 中常用的中间量rstd标准差的倒数与x归一化前的数据和方便上层框架在反向传播中复用避免重复计算。从仓库目录结构看该算子是一个完整的 CANN 算子工程包含接口文档、调用示例、构图原型、Host 侧定义/推导/tiling、Kernel 侧实现与单元/系统测试接口文档aclnn 单算子调用接口说明调用示例完整的可编译示例算子 IR 原型图模式构图接口op_host、op_kernel、testsHost 侧与 Kernel 侧实现及测试。产品支持情况根据 README.mdAddRmsNormCast 算子在不同产品上的支持情况如下产品是否支持Ascend 950PR950DT系列产品√Atlas A3系列产品√Atlas A2系列产品√Atlas 200I/500 A2推理产品×Atlas推理系列产品×Atlas训练系列产品×Kirin X90处理器系列产品√Kirin 9030处理器系列产品√这一支持矩阵与算子定义文件 op_host/add_rms_norm_cast_def.cpp 中注册的 AICore 配置一一对应ascend910bAtlas A2 系列、ascend910_93Atlas A3 系列、ascend950Ascend 950 系列、kirinx90与kirin9030。同时 op_host/config 目录下也分别维护了ascend910b、ascend910_93、ascend950、kirinx90、kirin9030五份算子二进制配置add_rms_norm_cast_binary.json与上表的产品范围一致。需要注意的版本差异Kirin X90 处理器系列产品与 Kirin 9030 处理器系列产品上x1、x2、gamma、y2和x的数据类型不支持 BFLOAT16。这一点在算子定义中也有体现——Kirin 系列专属配置为x1、x2、gamma、y2、x只注册了DT_FLOAT16而其他产品则同时注册DT_FLOAT16与DT_BF16。功能说明与计算公式计算流程AddRmsNormCast 一次执行完成三步计算此处参数命名遵循 README.md与 aclnn 接口文档中的命名差异见下文参数说明对两个输入做逐元素求和得到x对x做 RmsNorm 归一化得到归一化结果y2将y2转换为更高精度的y1。计算公式求和$$ x_ix1_{i}x2_{i} $$RmsNorm 归一化$$ y_2\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^2eps} $$类型转换$$ y_1float(y_2) $$其中x1、x2需要归一化的原始数据输入g数据缩放因子即gamma逐元素乘到归一化结果上eps即epsilon添加到分母根号内的常数用于防止除零并保证数值稳定y2归一化后、未做类型转换的输出FLOAT16 / BFLOAT16y1归一化后经过类型转换的输出FLOAT32中间量Rms(x)的倒数即输出rstdx1x2的和即输出x。该计算流程与算子 IR 原型 op_graph/add_rms_norm_cast_proto.h 中注释描述的实现完全一致x float(x1) float(x2) rstd np.rsqrt(np.mean(np.power(x,2), reduce_axis, keepdimsTrue) epsilon) y1 gamma * (x * rstd) y2 cast(y1)此外tests/assets/golden.py 中的 golden 函数复用add_rms_norm算子的 golden 实现并通过_post_actioncast追加类型转换以纯 NumPy 复现Add → RmsNorm → Cast的完整计算链是验证算子数值正确性的参考实现。参数说明下表完整列出算子各参数的定义数据格式均为 ND参数名输入/输出/属性描述数据类型数据格式x1输入需要归一化的原始数据输入公式中的输入x1。FLOAT16、BFLOAT16NDx2输入需要归一化的原始数据输入公式中的输入x2。FLOAT16、BFLOAT16NDgamma可选输入数据缩放因子公式中的输入g。shape 需要与x1后几维保持一致后几维为x1需要 norm 的维度。FLOAT16、BFLOAT16NDepsilon可选属性添加到分母中的值以确保数值稳定用于防止除 0 错误对应公式中的eps。默认值为 1e-6。FLOAT32-y1输出归一化后经过类型转换的输出数据公式中的输出y1。FLOAT32NDy2输出归一化后的输出数据公式中的输出y2。FLOAT16、BFLOAT16NDrstd输出x 的标准差公式中的输出Rms(x)。FLOAT32NDx输出归一化的数据和公式中的输出x。FLOAT16、BFLOAT16ND参数背后的实现约束gamma 与 norm 维度gamma 的 shape 必须与x1的后几维保持一致其中后几维就是执行归一化reduce的维度。例如x1shape 为(2,3,4,8)时若gammashape 为(8)则只有最后一维参与归一化若gammashape 为(4,8)则后两维参与归一化。rstd 的 shape 推导rstd 需要与x1数据格式一致维度数与x1相同其中不需要 norm 的维度x1的维度数减去gamma的维度数后的前几维与x1对应维度保持一致需要 norm 的维度与gamma维度数相同的后几维均为 1。举例若x1shape 为(2,3,4,8)、gammashape 为(8)则rstdshape 为(2,3,4,1)若gammashape 为(4,8)则rstdshape 为(2,3,1,1)。y1 / y2 / x 的 shape与x1保持一致x2的 shape 与数据类型也需要与x1保持一致。数据类型联动y1固定为 FLOAT32y2与x的数据类型跟随x1rstd固定为 FLOAT32。在 op_host/add_rms_norm_cast_infershape.cpp 的InferDataType4AddRmsNormCast中可以看到这一规则的实现y1与rstd直接置为DT_FLOATy2与x取x1的输入类型。关于接口文档中输出命名的说明README 的参数表将 FLOAT32 输出命名为y1、FLOAT16/BFLOAT16 输出命名为y2而 aclnnAddRmsNormCast 接口文档 中 FLOAT32 输出命名为y1Out归一化输出、FLOAT16/BFLOAT16 输出命名为y2Outcast 输出。两者对归一化结果与类型转换结果的命名顺序相反但类型映射一致FLOAT32 一个、FLOAT16/BFLOAT16 一个实际调用时以你使用的接口文档为准、按数据类型对应即可。约束说明输出不支持非连续 Tensory1Out、y2Out、rstdOut、xOut均要求连续。维度边界x1、x2、gamma、y1、y2、rstd、x的 shape 中每一维大小都不大于 INT32 最大值 2147483647各张量维度数限制为 1~8tiling 校验代码 op_host/add_rms_norm_cast_tiling.cpp 中的MAX_DIM_NUM8与之对应。空 Tensor 边界不支持非 Norm 维度元素总数大于 0 且 Norm 维度元素总数为 0的空 Tensor 场景。特殊数值传递输入为 Inf 时输出为 Inf输入为 NaN 时输出为 NaN。确定性计算aclnnAddRmsNormCast默认确定性实现同一输入多次运行结果可复现。epsilon 取值建议值为 1e-6tiling 中还会校验 epsilon 不小于 0。Kirin 系列 dtype 限制Kirin X90 / Kirin 9030 上x1、x2、gamma、y2、x不支持 BFLOAT16。调用说明AddRmsNormCast 支持两种调用方式通过aclnnAddRmsNormCast接口单算子调用对应 示例代码与通过算子 IR 构图图模式对应 add_rms_norm_cast_proto.h。调用方式样例代码说明aclnn接口test_aclnn_add_rms_norm_cast通过 aclnnAddRmsNormCast 接口方式调用 AddRmsNormCast 算子。图模式-通过 算子IR 构图方式调用 AddRmsNormCast 算子。aclnn 两段式接口每个算子分为两段式接口必须先调用aclnnAddRmsNormCastGetWorkspaceSize获取计算所需 workspace 大小以及包含了算子计算流程的执行器executor再调用aclnnAddRmsNormCast执行计算。aclnnStatus aclnnAddRmsNormCastGetWorkspaceSize( const aclTensor *x1, const aclTensor *x2, const aclTensor *gamma, double epsilon, const aclTensor *y1Out, const aclTensor *y2Out, const aclTensor *rstdOut, const aclTensor *xOut, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnAddRmsNormCast( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)aclnnAddRmsNormCastGetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续Tensorx1aclTensor*输入表示用于 Add 计算的第一个输入。对应公式中的x1。支持空 Tensor。BFLOAT16、FLOAT16ND1-8√x2aclTensor*输入表示用于 Add 计算的第二个输入。对应公式中的x2。支持空 Tensorshape 和数据类型需要与x1的 shape 和数据类型保持一致。FLOAT16、BFLOAT16ND1-8√gammaaclTensor*输入表示 RmsNorm 的缩放因子权重。对应公式中的gamma。支持空 Tensor数据类型与x1保持一致shape 需要与x1后几维保持一致。FLOAT16、BFLOAT16ND1-8√epsilondouble输入表示添加到分母中的值以确保数值稳定。对应公式中的epsilon。建议值为 1e-6。----y1OutaclTensor*输出表示归一化后的输出数据。支持空 Tensorshape、数据格式与入参x1保持一致。FLOAT32ND1-8×y2OutaclTensor*输出表示归一化后经过类型转换的输出数据。支持空 Tensorshape、数据格式、数据类型均与入参x1保持一致。FLOAT16、BFLOAT16ND1-8×rstdOutaclTensor*输出表示归一化后的标准差的倒数。支持空 Tensor与入参x1数据格式一致不需要 norm 的维度与x1对应维度一致需要 norm 的维度均为 1。FLOAT32ND1-8×xOutaclTensor*输出表示 Add 计算的结果。支持空 Tensorshape、数据格式、数据类型均与入参x1保持一致。FLOAT16、BFLOAT16ND1-8×workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小。-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程。-----返回码aclnnStatus返回状态码。第一段接口完成入参校验出现以下场景时报错返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001如果传入参数是必选输入、输出或者必选属性且是空指针。ACLNN_ERR_PARAM_INVALID161002输入或输出的数据类型不在支持的范围之内。ACLNN_ERR_INNER_TILING_ERROR561002输入和输出不符合上述参数说明内的要求。aclnnAddRmsNormCast 参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口aclnnAddRmsNormCastGetWorkspaceSize获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。完整调用示例以下代码取自 examples/test_aclnn_add_rms_norm_cast.cpp展示了从环境初始化、Tensor 构造、两段式接口调用到结果回拷与资源释放的完整流程#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_add_rms_norm_cast.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; } // 固定写法acl 资源初始化 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; } // 申请 device 内存、拷贝数据并创建 ND 格式 aclTensor 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); CHECK_RET(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); 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]; } *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 初始化 int32_t deviceId 0; // 根据自己的实际 device 填写 aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret 0, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出示例 shape 为 x1/x2/y1/y2/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* y1DeviceAddr nullptr; void* y2DeviceAddr nullptr; void* rstdDeviceAddr nullptr; void* xDeviceAddr nullptr; aclTensor* x1 nullptr; aclTensor* x2 nullptr; aclTensor* gamma nullptr; aclTensor* y1 nullptr; aclTensor* y2 nullptr; aclTensor* rstd nullptr; aclTensor* x nullptr; std::vectorshort x1HostData {0x0000, 0x3C00, 0x4000, 0x4200, 0x4400, 0x4500, 0x4600, 0x4700, 0x0000, 0x3C00, 0x4000, 0x4200, 0x4400, 0x4500, 0x4600, 0x4700, 0x0000, 0x3C00, 0x4000, 0x4200, 0x4400, 0x4500, 0x4600, 0x4700, 0x0000, 0x3C00, 0x4000, 0x4200, 0x4400, 0x4500, 0x4600, 0x4700}; std::vectorshort x2HostData x1HostData; // 实际示例中与 x1HostData 相同 std::vectorshort gammaHostData {0x0000, 0x3C00, 0x4000, 0x4200, 0x4400, 0x4500, 0x4600, 0x4700, 0x0000, 0x3C00, 0x4000, 0x4200, 0x4400, 0x4500, 0x4600, 0x4700}; std::vectorfloat y1HostData {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::vectorshort y2HostData x1HostData; // 实际示例中与 x1HostData 相同 std::vectorfloat rstdHostData {1, 2}; std::vectorshort xHostData x1HostData; // 实际示例中与 x1HostData 相同 float epsilon 1e-6; // 创建各 aclTensorACL_FLOAT16 对应 FLOAT16 输入ACL_FLOAT 对应 FLOAT32 输出 ret CreateAclTensor(x1HostData, xShape, x1DeviceAddr, aclDataType::ACL_FLOAT16, x1); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(x2HostData, xShape, x2DeviceAddr, aclDataType::ACL_FLOAT16, x2); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(gammaHostData, gammaShape, gammaDeviceAddr, aclDataType::ACL_FLOAT16, gamma); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(y1HostData, yShape, y1DeviceAddr, aclDataType::ACL_FLOAT, y1); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(y2HostData, yShape, y2DeviceAddr, aclDataType::ACL_FLOAT16, y2); 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_FLOAT16, x); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用第一段接口获取 workspace 大小与执行器 uint64_t workspaceSize 0; aclOpExecutor* executor; ret aclnnAddRmsNormCastGetWorkspaceSize(x1, x2, gamma, epsilon, y1, y2, rstd, x, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnAddRmsNormCastGetWorkspaceSize 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 aclnnAddRmsNormCast(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnAddRmsNormCast 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]), y1DeviceAddr, size * sizeof(resultData[0]), 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(y1 result[%ld] is: %f\n, i, resultData[i]); } // 6. 释放 aclTensor aclDestroyTensor(x1); aclDestroyTensor(x2); aclDestroyTensor(gamma); aclDestroyTensor(y1); aclDestroyTensor(y2); aclDestroyTensor(rstd); aclDestroyTensor(x); // 7. 释放 device 资源 aclrtFree(x1DeviceAddr); aclrtFree(x2DeviceAddr); aclrtFree(xDeviceAddr); aclrtFree(gammaDeviceAddr); aclrtFree(y2DeviceAddr); aclrtFree(y1DeviceAddr); aclrtFree(rstdDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例中的关键点输出命名对应示例中y1对应接口文档的y1OutFLOAT32y2对应y2OutFLOAT16/BFLOAT16rstd 的 shape 构造x1为(2,16)、gamma为(16)时rstd为(2,1)——前 1 维与x1相同、norm 维压缩为 1与 shape 推导规则一致workspace 处理只有当workspaceSize 0时才需要申请内存执行完毕后同样需要释放完整的编译与运行流程可参考 CANN 文档中编译与运行样例的相关指引。图模式构图在图模式下通过算子 IR 在计算图中插入 AddRmsNormCast 节点完成调用。算子原型定义在 op_graph/add_rms_norm_cast_proto.h输入x1FLOAT16/BFLOAT16ND、x2FLOAT16/BFLOAT16ND、gammaFLOAT16/BFLOAT16ND输出y1FLOAT32、y2FLOAT16/BFLOAT16、rstdFLOAT32、xFLOAT16/BFLOAT16格式均为 ND属性epsilonFloat 类型默认值1e-6f。即使用REG_OP(AddRmsNormCast)注册的图算子节点框架侧可通过ge::Operator按上述输入输出与属性构造节点加入图。源码实现纵深Host 侧与 Kernel 侧的关键机制算子定义op_host/add_rms_norm_cast_def.cpp在 add_rms_norm_cast_def.cpp 中所有输入输出都声明为REQUIRED必选数据类型与格式与文档一致epsilon声明为OPTIONAL属性默认1e-6ascend950配置开启动态编译静态化DynamicCompileStaticFlag、动态 rank 与动态 shape 支持DynamicRankSupportFlag、DynamicShapeSupportFlag并指定扩展编译配置add_rms_norm_cast_apt对应 kernel 侧 add_rms_norm_cast_apt.cppKirin 系列kirinx90、kirin9030使用专属配置只支持 FLOAT16、开启动态格式/动态 rank/动态 shape、关闭支持检查NeedCheckSupportFlag(false)并开启精度降低PrecisionReduceFlag(true)。shape 推导op_host/add_rms_norm_cast_infershape.cppadd_rms_norm_cast_infershape.cpp 实现了 shape 与 dtype 推导y1、y2、x的 shape 直接继承x1rstd的 shape 按前xDimNum - gammaDimNum维与x1相同、norm 维为 1规则构造x1的维数小于gamma的维数时报错数据类型y1、rstd固定 FLOATy2、x跟随x1。tiling 策略op_host/add_rms_norm_cast_tiling.cpptiling 是算子性能的关键。从 add_rms_norm_cast_tiling.cpp 可以看到其核心思路数据视图把输入按矩阵方式看待numRow为x1前x1DimNum - gammaDimNum维的元素乘积即归一化行数numCol为gamma的元素个数即归一化列数avgFactor 1/numCol作为均值因子参与计算tiling key 编码tilingKey dtypeKey * 10 modeKey其中dtypeKey1FLOAT16、2FLOAT32、3BFLOAT16与modeKey0Normal、1SplitD、2MergeN、3SingleN、4MultiN组合出多种实现变体。例如 10/30 对应 Normal 模式、11/31 对应 SplitD、13/33 对应 SingleN、14 对应 MultiN模式选择numCol超过 UB 容量UB_FACTOR_B168704等阈值时切到 SplitD 模式对归一化维做切分blockFactor 为 1 且非特定 SoC 时走 SingleN 模式workspace 规划默认申请约 16MB 系统 workspace 加 256B 用户 workspace校验调用前对x1/x2/y1/y2/x的 shape 一致性、gamma与x1后几维的一致性、rstd与x1前几维的一致性、维度数范围1~8以及epsilon 0做完整检查。此外add_rms_norm_cast_tiling.h 中注册了多套 tiling 数据类AddRMSNormCastTilingData与AddRmsNormCastRegbaseTilingData其中AddRmsNormCast_100~_103、_199等 tiling key 对应 arch35 平台的 RegBase 实现说明该算子在 Atlas A3arch35等平台上有独立的 regbase tiling 路径其 kernel 实现在 op_kernel/arch35 目录下含add_rms_norm_cast_regbase.h、add_rms_norm_cast_regbase_high_performance.h、add_rms_norm_cast_regbase_single_n.h、add_rms_norm_cast_regbase_spilt_reduce.h等。Kernel 分发op_kernel/add_rms_norm_cast.cppadd_rms_norm_cast.cpp 是 kernel 入口按 tiling key 将计算分发给不同实现类TILING_KEY_IS(10/30)→KernelAddRmsNormCastNormalhalf / bfloat16_tTILING_KEY_IS(11/31)→KernelAddRmsNormCastSplitD归一化维切分TILING_KEY_IS(13/33)→KernelAddRmsNormCastSingleNTILING_KEY_IS(14)→KernelAddRmsNormCastMultiNBF16 多 N 场景。各实现类add_rms_norm_cast_single_n.h、add_rms_norm_cast_multi_n.h、add_rms_norm_cast_split_d.h都通过统一的Init(x1, x2, gamma, y1, y2, rstd, x, workspace, tilingData)Process()接口执行先 Add 求和、再按 tiling 数据计算rstd与归一化结果、最后完成 cast 输出。二进制配置与测试每个产品目录下的 add_rms_norm_cast_binary.jsonascend910b、ascend910_93、ascend950、kirinx90、kirin9030各一份以bin_filename如AddRmsNormCast_fp16、AddRmsNormCast_bf16绑定不同 dtype 的二进制声明输入输出 shape 为动态-2、epsilon默认值0.000001供编译框架生成对应的算子二进制。测试方面Host 侧单测tests/ut/op_host/test_AddRmsNormCast_infershape.cpp 覆盖 shape 推导含 rstd 各维度规则tests/ut/op_host/test_add_rms_norm_cast_tiling.cpp 覆盖 tiling 参数计算Kernel 侧单测tests/ut/op_kernel/test_add_rms_norm_cast.cpp 与 tests/ut/op_kernel/test_add_rms_norm_cast_regbase.cpp 覆盖常规与 arch35 regbase 实现系统测试arch35 平台 ST 用例在 tests/st/arch35/ttk_kernel_add_rms_norm_cast_st.csv数值 goldentests/assets/golden.py 复用 add_rms_norm 的 golden 并通过_post_actioncast生成参考结果。总结AddRmsNormCast 是 CANN ops-nn 中一个典型的融合 多输出归一化算子它将 Add、RmsNorm、Cast 三合一减少数据搬入搬出除归一化结果外还输出rstd与求和结果x为反向传播提供可复用中间量。使用时需重点把握三点一是gamma决定归一化维度、rstd的 shape 随之确定二是通过 aclnn 两段式接口GetWorkspaceSize 执行完成单算子调用输出不支持非连续 Tensor三是按产品选择支持的 dtypeKirin 系列不支持 BFLOAT16。从源码看其 tiling 层针对不同数据规模选择了 Normal / SplitD / SingleN / MultiN 多种实现模式并在 arch35 平台提供独立的 RegBase 高性能路径体现了面向大模型归一化场景的性能优化思路。若需在计算图中使用可直接基于 add_rms_norm_cast_proto.h 的算子 IR 构图。赞分享人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载相关推荐CANN ops-nn 算子 aclnnAddRmsNormCast 深度解析AddRmsNormCast 三合一融合算子的接口、原理与实战CANN ops nn 算子 aclnnAddRmsNormCast 深度解析AddRmsNormCast 三合一融合算子的接口、原理与实战 本技术指南以人工智能算子库深度学习CANNAscendCANN ops-nn 融合算子 add_rms_norm_dynamic_quant 完全指南Add RmsNorm MX 动态量化CANN ops nn 融合算子 add_rms_norm_dynamic_quant 完全指南Add RmsNorm MX 动态量化 本篇技术指南围人工智能算子库深度学习CANNAscendTiptap 2.x 到 3.x 升级5 步完成富文本编辑器迁移Tiptap 2.x 到 3.x 升级5 步完成富文本编辑器迁移 本文面向仍停留在 2.x 的 Tiptap 使用者讲清 Tiptap 3.x 迁移要动哪些人工智能算子库深度学习CANNAscend上一篇Hyperframes Cinematic 字幕模式实战一个引擎、十种 DNA 视觉语言的纯嵌入字幕管线下一篇Remix UI Accordion 组件全解从 remix/ui/accordion 样式组件到 primitives 无头原语创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN ops-nn AddRmsNormCast 算子深度解析：Add 融合 RmsNorm 与 Cast 的归一化算子上手指南
人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载导读AddRmsNormCast 是 CANN ops-nn 神经网络算子库中面向大模型LLM场景的融合归一化算子它把Add逐元素相加、RmsNorm均方根归一化与Cast类型转换三个操作融合为一次 Kernel 执行通过减少数据在内存与计算单元之间的搬入搬出次数来降低算子调用开销。本文以 norm/add_rms_norm_cast/README.md 与配套的 aclnnAddRmsNormCast 接口文档 为主线结合该算子在仓库中的算子定义、shape 推导、tiling 策略、Kernel 分发与测试代码完整讲解其产品支持情况、数学原理、输入输出参数、aclnn 两段式调用方法与图模式构图方式。读完本文你将掌握如何在 Atlas A2/A3、Ascend 950 系列与 Kirin 系列产品上正确构造并调用该算子并能理解其底层实现机制。算子定位为什么要把 Add、RmsNorm、Cast 融合在一起RmsNormRoot Mean Square Normalization是大模型常用的归一化操作。在典型的 Transformer 网络结构中残差相加Add之后往往紧跟 RmsNorm而部分计算图还会在 RmsNorm 之后衔接类型转换Cast以满足后续算子对精度的要求。如果按照原始计算图逐个算子执行Add、RmsNorm、Cast 各自都需要独立申请内存、启动 Kernel、完成数据搬入搬出会引入可观的额外开销。AddRmsNormCast 算子的核心设计目标正是消除这些开销将 AddRmsNorm 之后的 Cast 算子融合进归一化计算中减少搬入搬出操作见 README.md 功能说明。一次 Kernel 调用即可完成求和 → 归一化 → 类型转换同时顺带输出 RmsNorm 中常用的中间量rstd标准差的倒数与x归一化前的数据和方便上层框架在反向传播中复用避免重复计算。从仓库目录结构看该算子是一个完整的 CANN 算子工程包含接口文档、调用示例、构图原型、Host 侧定义/推导/tiling、Kernel 侧实现与单元/系统测试接口文档aclnn 单算子调用接口说明调用示例完整的可编译示例算子 IR 原型图模式构图接口op_host、op_kernel、testsHost 侧与 Kernel 侧实现及测试。产品支持情况根据 README.mdAddRmsNormCast 算子在不同产品上的支持情况如下产品是否支持Ascend 950PR950DT系列产品√Atlas A3系列产品√Atlas A2系列产品√Atlas 200I/500 A2推理产品×Atlas推理系列产品×Atlas训练系列产品×Kirin X90处理器系列产品√Kirin 9030处理器系列产品√这一支持矩阵与算子定义文件 op_host/add_rms_norm_cast_def.cpp 中注册的 AICore 配置一一对应ascend910bAtlas A2 系列、ascend910_93Atlas A3 系列、ascend950Ascend 950 系列、kirinx90与kirin9030。同时 op_host/config 目录下也分别维护了ascend910b、ascend910_93、ascend950、kirinx90、kirin9030五份算子二进制配置add_rms_norm_cast_binary.json与上表的产品范围一致。需要注意的版本差异Kirin X90 处理器系列产品与 Kirin 9030 处理器系列产品上x1、x2、gamma、y2和x的数据类型不支持 BFLOAT16。这一点在算子定义中也有体现——Kirin 系列专属配置为x1、x2、gamma、y2、x只注册了DT_FLOAT16而其他产品则同时注册DT_FLOAT16与DT_BF16。功能说明与计算公式计算流程AddRmsNormCast 一次执行完成三步计算此处参数命名遵循 README.md与 aclnn 接口文档中的命名差异见下文参数说明对两个输入做逐元素求和得到x对x做 RmsNorm 归一化得到归一化结果y2将y2转换为更高精度的y1。计算公式求和$$ x_ix1_{i}x2_{i} $$RmsNorm 归一化$$ y_2\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^2eps} $$类型转换$$ y_1float(y_2) $$其中x1、x2需要归一化的原始数据输入g数据缩放因子即gamma逐元素乘到归一化结果上eps即epsilon添加到分母根号内的常数用于防止除零并保证数值稳定y2归一化后、未做类型转换的输出FLOAT16 / BFLOAT16y1归一化后经过类型转换的输出FLOAT32中间量Rms(x)的倒数即输出rstdx1x2的和即输出x。该计算流程与算子 IR 原型 op_graph/add_rms_norm_cast_proto.h 中注释描述的实现完全一致x float(x1) float(x2) rstd np.rsqrt(np.mean(np.power(x,2), reduce_axis, keepdimsTrue) epsilon) y1 gamma * (x * rstd) y2 cast(y1)此外tests/assets/golden.py 中的 golden 函数复用add_rms_norm算子的 golden 实现并通过_post_actioncast追加类型转换以纯 NumPy 复现Add → RmsNorm → Cast的完整计算链是验证算子数值正确性的参考实现。参数说明下表完整列出算子各参数的定义数据格式均为 ND参数名输入/输出/属性描述数据类型数据格式x1输入需要归一化的原始数据输入公式中的输入x1。FLOAT16、BFLOAT16NDx2输入需要归一化的原始数据输入公式中的输入x2。FLOAT16、BFLOAT16NDgamma可选输入数据缩放因子公式中的输入g。shape 需要与x1后几维保持一致后几维为x1需要 norm 的维度。FLOAT16、BFLOAT16NDepsilon可选属性添加到分母中的值以确保数值稳定用于防止除 0 错误对应公式中的eps。默认值为 1e-6。FLOAT32-y1输出归一化后经过类型转换的输出数据公式中的输出y1。FLOAT32NDy2输出归一化后的输出数据公式中的输出y2。FLOAT16、BFLOAT16NDrstd输出x 的标准差公式中的输出Rms(x)。FLOAT32NDx输出归一化的数据和公式中的输出x。FLOAT16、BFLOAT16ND参数背后的实现约束gamma 与 norm 维度gamma 的 shape 必须与x1的后几维保持一致其中后几维就是执行归一化reduce的维度。例如x1shape 为(2,3,4,8)时若gammashape 为(8)则只有最后一维参与归一化若gammashape 为(4,8)则后两维参与归一化。rstd 的 shape 推导rstd 需要与x1数据格式一致维度数与x1相同其中不需要 norm 的维度x1的维度数减去gamma的维度数后的前几维与x1对应维度保持一致需要 norm 的维度与gamma维度数相同的后几维均为 1。举例若x1shape 为(2,3,4,8)、gammashape 为(8)则rstdshape 为(2,3,4,1)若gammashape 为(4,8)则rstdshape 为(2,3,1,1)。y1 / y2 / x 的 shape与x1保持一致x2的 shape 与数据类型也需要与x1保持一致。数据类型联动y1固定为 FLOAT32y2与x的数据类型跟随x1rstd固定为 FLOAT32。在 op_host/add_rms_norm_cast_infershape.cpp 的InferDataType4AddRmsNormCast中可以看到这一规则的实现y1与rstd直接置为DT_FLOATy2与x取x1的输入类型。关于接口文档中输出命名的说明README 的参数表将 FLOAT32 输出命名为y1、FLOAT16/BFLOAT16 输出命名为y2而 aclnnAddRmsNormCast 接口文档 中 FLOAT32 输出命名为y1Out归一化输出、FLOAT16/BFLOAT16 输出命名为y2Outcast 输出。两者对归一化结果与类型转换结果的命名顺序相反但类型映射一致FLOAT32 一个、FLOAT16/BFLOAT16 一个实际调用时以你使用的接口文档为准、按数据类型对应即可。约束说明输出不支持非连续 Tensory1Out、y2Out、rstdOut、xOut均要求连续。维度边界x1、x2、gamma、y1、y2、rstd、x的 shape 中每一维大小都不大于 INT32 最大值 2147483647各张量维度数限制为 1~8tiling 校验代码 op_host/add_rms_norm_cast_tiling.cpp 中的MAX_DIM_NUM8与之对应。空 Tensor 边界不支持非 Norm 维度元素总数大于 0 且 Norm 维度元素总数为 0的空 Tensor 场景。特殊数值传递输入为 Inf 时输出为 Inf输入为 NaN 时输出为 NaN。确定性计算aclnnAddRmsNormCast默认确定性实现同一输入多次运行结果可复现。epsilon 取值建议值为 1e-6tiling 中还会校验 epsilon 不小于 0。Kirin 系列 dtype 限制Kirin X90 / Kirin 9030 上x1、x2、gamma、y2、x不支持 BFLOAT16。调用说明AddRmsNormCast 支持两种调用方式通过aclnnAddRmsNormCast接口单算子调用对应 示例代码与通过算子 IR 构图图模式对应 add_rms_norm_cast_proto.h。调用方式样例代码说明aclnn接口test_aclnn_add_rms_norm_cast通过 aclnnAddRmsNormCast 接口方式调用 AddRmsNormCast 算子。图模式-通过 算子IR 构图方式调用 AddRmsNormCast 算子。aclnn 两段式接口每个算子分为两段式接口必须先调用aclnnAddRmsNormCastGetWorkspaceSize获取计算所需 workspace 大小以及包含了算子计算流程的执行器executor再调用aclnnAddRmsNormCast执行计算。aclnnStatus aclnnAddRmsNormCastGetWorkspaceSize( const aclTensor *x1, const aclTensor *x2, const aclTensor *gamma, double epsilon, const aclTensor *y1Out, const aclTensor *y2Out, const aclTensor *rstdOut, const aclTensor *xOut, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnAddRmsNormCast( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)aclnnAddRmsNormCastGetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续Tensorx1aclTensor*输入表示用于 Add 计算的第一个输入。对应公式中的x1。支持空 Tensor。BFLOAT16、FLOAT16ND1-8√x2aclTensor*输入表示用于 Add 计算的第二个输入。对应公式中的x2。支持空 Tensorshape 和数据类型需要与x1的 shape 和数据类型保持一致。FLOAT16、BFLOAT16ND1-8√gammaaclTensor*输入表示 RmsNorm 的缩放因子权重。对应公式中的gamma。支持空 Tensor数据类型与x1保持一致shape 需要与x1后几维保持一致。FLOAT16、BFLOAT16ND1-8√epsilondouble输入表示添加到分母中的值以确保数值稳定。对应公式中的epsilon。建议值为 1e-6。----y1OutaclTensor*输出表示归一化后的输出数据。支持空 Tensorshape、数据格式与入参x1保持一致。FLOAT32ND1-8×y2OutaclTensor*输出表示归一化后经过类型转换的输出数据。支持空 Tensorshape、数据格式、数据类型均与入参x1保持一致。FLOAT16、BFLOAT16ND1-8×rstdOutaclTensor*输出表示归一化后的标准差的倒数。支持空 Tensor与入参x1数据格式一致不需要 norm 的维度与x1对应维度一致需要 norm 的维度均为 1。FLOAT32ND1-8×xOutaclTensor*输出表示 Add 计算的结果。支持空 Tensorshape、数据格式、数据类型均与入参x1保持一致。FLOAT16、BFLOAT16ND1-8×workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小。-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程。-----返回码aclnnStatus返回状态码。第一段接口完成入参校验出现以下场景时报错返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001如果传入参数是必选输入、输出或者必选属性且是空指针。ACLNN_ERR_PARAM_INVALID161002输入或输出的数据类型不在支持的范围之内。ACLNN_ERR_INNER_TILING_ERROR561002输入和输出不符合上述参数说明内的要求。aclnnAddRmsNormCast 参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口aclnnAddRmsNormCastGetWorkspaceSize获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。完整调用示例以下代码取自 examples/test_aclnn_add_rms_norm_cast.cpp展示了从环境初始化、Tensor 构造、两段式接口调用到结果回拷与资源释放的完整流程#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_add_rms_norm_cast.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; } // 固定写法acl 资源初始化 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; } // 申请 device 内存、拷贝数据并创建 ND 格式 aclTensor 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); CHECK_RET(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); 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]; } *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 初始化 int32_t deviceId 0; // 根据自己的实际 device 填写 aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret 0, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出示例 shape 为 x1/x2/y1/y2/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* y1DeviceAddr nullptr; void* y2DeviceAddr nullptr; void* rstdDeviceAddr nullptr; void* xDeviceAddr nullptr; aclTensor* x1 nullptr; aclTensor* x2 nullptr; aclTensor* gamma nullptr; aclTensor* y1 nullptr; aclTensor* y2 nullptr; aclTensor* rstd nullptr; aclTensor* x nullptr; std::vectorshort x1HostData {0x0000, 0x3C00, 0x4000, 0x4200, 0x4400, 0x4500, 0x4600, 0x4700, 0x0000, 0x3C00, 0x4000, 0x4200, 0x4400, 0x4500, 0x4600, 0x4700, 0x0000, 0x3C00, 0x4000, 0x4200, 0x4400, 0x4500, 0x4600, 0x4700, 0x0000, 0x3C00, 0x4000, 0x4200, 0x4400, 0x4500, 0x4600, 0x4700}; std::vectorshort x2HostData x1HostData; // 实际示例中与 x1HostData 相同 std::vectorshort gammaHostData {0x0000, 0x3C00, 0x4000, 0x4200, 0x4400, 0x4500, 0x4600, 0x4700, 0x0000, 0x3C00, 0x4000, 0x4200, 0x4400, 0x4500, 0x4600, 0x4700}; std::vectorfloat y1HostData {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::vectorshort y2HostData x1HostData; // 实际示例中与 x1HostData 相同 std::vectorfloat rstdHostData {1, 2}; std::vectorshort xHostData x1HostData; // 实际示例中与 x1HostData 相同 float epsilon 1e-6; // 创建各 aclTensorACL_FLOAT16 对应 FLOAT16 输入ACL_FLOAT 对应 FLOAT32 输出 ret CreateAclTensor(x1HostData, xShape, x1DeviceAddr, aclDataType::ACL_FLOAT16, x1); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(x2HostData, xShape, x2DeviceAddr, aclDataType::ACL_FLOAT16, x2); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(gammaHostData, gammaShape, gammaDeviceAddr, aclDataType::ACL_FLOAT16, gamma); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(y1HostData, yShape, y1DeviceAddr, aclDataType::ACL_FLOAT, y1); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(y2HostData, yShape, y2DeviceAddr, aclDataType::ACL_FLOAT16, y2); 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_FLOAT16, x); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用第一段接口获取 workspace 大小与执行器 uint64_t workspaceSize 0; aclOpExecutor* executor; ret aclnnAddRmsNormCastGetWorkspaceSize(x1, x2, gamma, epsilon, y1, y2, rstd, x, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnAddRmsNormCastGetWorkspaceSize 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 aclnnAddRmsNormCast(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnAddRmsNormCast 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]), y1DeviceAddr, size * sizeof(resultData[0]), 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(y1 result[%ld] is: %f\n, i, resultData[i]); } // 6. 释放 aclTensor aclDestroyTensor(x1); aclDestroyTensor(x2); aclDestroyTensor(gamma); aclDestroyTensor(y1); aclDestroyTensor(y2); aclDestroyTensor(rstd); aclDestroyTensor(x); // 7. 释放 device 资源 aclrtFree(x1DeviceAddr); aclrtFree(x2DeviceAddr); aclrtFree(xDeviceAddr); aclrtFree(gammaDeviceAddr); aclrtFree(y2DeviceAddr); aclrtFree(y1DeviceAddr); aclrtFree(rstdDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例中的关键点输出命名对应示例中y1对应接口文档的y1OutFLOAT32y2对应y2OutFLOAT16/BFLOAT16rstd 的 shape 构造x1为(2,16)、gamma为(16)时rstd为(2,1)——前 1 维与x1相同、norm 维压缩为 1与 shape 推导规则一致workspace 处理只有当workspaceSize 0时才需要申请内存执行完毕后同样需要释放完整的编译与运行流程可参考 CANN 文档中编译与运行样例的相关指引。图模式构图在图模式下通过算子 IR 在计算图中插入 AddRmsNormCast 节点完成调用。算子原型定义在 op_graph/add_rms_norm_cast_proto.h输入x1FLOAT16/BFLOAT16ND、x2FLOAT16/BFLOAT16ND、gammaFLOAT16/BFLOAT16ND输出y1FLOAT32、y2FLOAT16/BFLOAT16、rstdFLOAT32、xFLOAT16/BFLOAT16格式均为 ND属性epsilonFloat 类型默认值1e-6f。即使用REG_OP(AddRmsNormCast)注册的图算子节点框架侧可通过ge::Operator按上述输入输出与属性构造节点加入图。源码实现纵深Host 侧与 Kernel 侧的关键机制算子定义op_host/add_rms_norm_cast_def.cpp在 add_rms_norm_cast_def.cpp 中所有输入输出都声明为REQUIRED必选数据类型与格式与文档一致epsilon声明为OPTIONAL属性默认1e-6ascend950配置开启动态编译静态化DynamicCompileStaticFlag、动态 rank 与动态 shape 支持DynamicRankSupportFlag、DynamicShapeSupportFlag并指定扩展编译配置add_rms_norm_cast_apt对应 kernel 侧 add_rms_norm_cast_apt.cppKirin 系列kirinx90、kirin9030使用专属配置只支持 FLOAT16、开启动态格式/动态 rank/动态 shape、关闭支持检查NeedCheckSupportFlag(false)并开启精度降低PrecisionReduceFlag(true)。shape 推导op_host/add_rms_norm_cast_infershape.cppadd_rms_norm_cast_infershape.cpp 实现了 shape 与 dtype 推导y1、y2、x的 shape 直接继承x1rstd的 shape 按前xDimNum - gammaDimNum维与x1相同、norm 维为 1规则构造x1的维数小于gamma的维数时报错数据类型y1、rstd固定 FLOATy2、x跟随x1。tiling 策略op_host/add_rms_norm_cast_tiling.cpptiling 是算子性能的关键。从 add_rms_norm_cast_tiling.cpp 可以看到其核心思路数据视图把输入按矩阵方式看待numRow为x1前x1DimNum - gammaDimNum维的元素乘积即归一化行数numCol为gamma的元素个数即归一化列数avgFactor 1/numCol作为均值因子参与计算tiling key 编码tilingKey dtypeKey * 10 modeKey其中dtypeKey1FLOAT16、2FLOAT32、3BFLOAT16与modeKey0Normal、1SplitD、2MergeN、3SingleN、4MultiN组合出多种实现变体。例如 10/30 对应 Normal 模式、11/31 对应 SplitD、13/33 对应 SingleN、14 对应 MultiN模式选择numCol超过 UB 容量UB_FACTOR_B168704等阈值时切到 SplitD 模式对归一化维做切分blockFactor 为 1 且非特定 SoC 时走 SingleN 模式workspace 规划默认申请约 16MB 系统 workspace 加 256B 用户 workspace校验调用前对x1/x2/y1/y2/x的 shape 一致性、gamma与x1后几维的一致性、rstd与x1前几维的一致性、维度数范围1~8以及epsilon 0做完整检查。此外add_rms_norm_cast_tiling.h 中注册了多套 tiling 数据类AddRMSNormCastTilingData与AddRmsNormCastRegbaseTilingData其中AddRmsNormCast_100~_103、_199等 tiling key 对应 arch35 平台的 RegBase 实现说明该算子在 Atlas A3arch35等平台上有独立的 regbase tiling 路径其 kernel 实现在 op_kernel/arch35 目录下含add_rms_norm_cast_regbase.h、add_rms_norm_cast_regbase_high_performance.h、add_rms_norm_cast_regbase_single_n.h、add_rms_norm_cast_regbase_spilt_reduce.h等。Kernel 分发op_kernel/add_rms_norm_cast.cppadd_rms_norm_cast.cpp 是 kernel 入口按 tiling key 将计算分发给不同实现类TILING_KEY_IS(10/30)→KernelAddRmsNormCastNormalhalf / bfloat16_tTILING_KEY_IS(11/31)→KernelAddRmsNormCastSplitD归一化维切分TILING_KEY_IS(13/33)→KernelAddRmsNormCastSingleNTILING_KEY_IS(14)→KernelAddRmsNormCastMultiNBF16 多 N 场景。各实现类add_rms_norm_cast_single_n.h、add_rms_norm_cast_multi_n.h、add_rms_norm_cast_split_d.h都通过统一的Init(x1, x2, gamma, y1, y2, rstd, x, workspace, tilingData)Process()接口执行先 Add 求和、再按 tiling 数据计算rstd与归一化结果、最后完成 cast 输出。二进制配置与测试每个产品目录下的 add_rms_norm_cast_binary.jsonascend910b、ascend910_93、ascend950、kirinx90、kirin9030各一份以bin_filename如AddRmsNormCast_fp16、AddRmsNormCast_bf16绑定不同 dtype 的二进制声明输入输出 shape 为动态-2、epsilon默认值0.000001供编译框架生成对应的算子二进制。测试方面Host 侧单测tests/ut/op_host/test_AddRmsNormCast_infershape.cpp 覆盖 shape 推导含 rstd 各维度规则tests/ut/op_host/test_add_rms_norm_cast_tiling.cpp 覆盖 tiling 参数计算Kernel 侧单测tests/ut/op_kernel/test_add_rms_norm_cast.cpp 与 tests/ut/op_kernel/test_add_rms_norm_cast_regbase.cpp 覆盖常规与 arch35 regbase 实现系统测试arch35 平台 ST 用例在 tests/st/arch35/ttk_kernel_add_rms_norm_cast_st.csv数值 goldentests/assets/golden.py 复用 add_rms_norm 的 golden 并通过_post_actioncast生成参考结果。总结AddRmsNormCast 是 CANN ops-nn 中一个典型的融合 多输出归一化算子它将 Add、RmsNorm、Cast 三合一减少数据搬入搬出除归一化结果外还输出rstd与求和结果x为反向传播提供可复用中间量。使用时需重点把握三点一是gamma决定归一化维度、rstd的 shape 随之确定二是通过 aclnn 两段式接口GetWorkspaceSize 执行完成单算子调用输出不支持非连续 Tensor三是按产品选择支持的 dtypeKirin 系列不支持 BFLOAT16。从源码看其 tiling 层针对不同数据规模选择了 Normal / SplitD / SingleN / MultiN 多种实现模式并在 arch35 平台提供独立的 RegBase 高性能路径体现了面向大模型归一化场景的性能优化思路。若需在计算图中使用可直接基于 add_rms_norm_cast_proto.h 的算子 IR 构图。赞分享人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载相关推荐CANN ops-nn 算子 aclnnAddRmsNormCast 深度解析AddRmsNormCast 三合一融合算子的接口、原理与实战CANN ops nn 算子 aclnnAddRmsNormCast 深度解析AddRmsNormCast 三合一融合算子的接口、原理与实战 本技术指南以人工智能算子库深度学习CANNAscendCANN ops-nn 融合算子 add_rms_norm_dynamic_quant 完全指南Add RmsNorm MX 动态量化CANN ops nn 融合算子 add_rms_norm_dynamic_quant 完全指南Add RmsNorm MX 动态量化 本篇技术指南围人工智能算子库深度学习CANNAscendTiptap 2.x 到 3.x 升级5 步完成富文本编辑器迁移Tiptap 2.x 到 3.x 升级5 步完成富文本编辑器迁移 本文面向仍停留在 2.x 的 Tiptap 使用者讲清 Tiptap 3.x 迁移要动哪些人工智能算子库深度学习CANNAscend上一篇Hyperframes Cinematic 字幕模式实战一个引擎、十种 DNA 视觉语言的纯嵌入字幕管线下一篇Remix UI Accordion 组件全解从 remix/ui/accordion 样式组件到 primitives 无头原语创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考