人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载AddRmsNormDynamicQuantV2 是 CANN ops-nn 神经网络算子库中面向大模型推理/训练场景的融合算子它把 Add → RmsNorm → 1 路或 2 路 DynamicQuant 对称动态量化 三段计算合并为一次 NPU 算子执行减少中间张量的搬入搬出。本文以 norm/add_rms_norm_dynamic_quant_v2/README.md 为主线结合仓库内算子 IR 定义、Host 侧 InferShape/Tiling 实现、Kernel 侧实现与 GE 图融合 Pass 源码完整讲解其功能、计算公式、全部参数与约束并给出可直接运行的图模式调用示例帮助开发者在 CANN 环境中正确构造与使用该算子。一、算子定位为什么需要融合 Add、RmsNorm 与 DynamicQuant在 LLM大语言模型的 Transformer 结构中RmsNormRoot Mean Square Layer Normalization是最常用的归一化算子相比 LayerNorm它去掉了减去均值的步骤只做均方根归一化节省了一次规约与减法开销。而在量化推理管线中归一化输出通常还要紧接着送入对称动态量化算子DynamicQuant转成 INT8 等低比特数据以匹配后续 MatMul 的量化输入。如果 Add、RmsNorm、DynamicQuant 各自独立执行中间结果x1x2、归一化输出y都要在 GM全局内存与计算单元之间多次搬移。AddRmsNormDynamicQuantV2 的设计目标正是把这些算子融合为一个 Kernel在片上完成加法、RmsNorm 归一化并将归一化输出分别送入 1 个或 2 个 DynamicQuant 量化支路减少搬入搬出操作、降低访存开销。从 op_graph/fusion_pass/add_rms_norm_dynamic_quant_v2_fusion_pass.cpp 的文件头注释可以直观看到该算子的来源形态x1 x2 gamma smooth1 \ | / | AddRmsNorm | / | \ | x y \ | | \ | Cast DynamicQuant | / \ y3 y1 scale1 x1, x2, gamma, smooth1 -- AddRmsNormDynamicQuantV2 outputs: y1, y3(Cast), y4(AddRmsNorm.y), x, scale1即图编译阶段由名为AddRmsNormDynamicQuantV2FusionPass的融合 Pass 识别 AddRmsNorm Cast DynamicQuant含 smooth1/smooth2 两路 子图替换为单个 AddRmsNormDynamicQuantV2 算子节点。该 Pass 通过IsTargetPlatform()检查目标 SOC目前仅在Ascend910B与Ascend950上启用见 add_rms_norm_dynamic_quant_v2_fusion_pass.cpp 中isPlatform910B/isPlatform950判断并有对应的图模式单测 test_add_rms_norm_dynamic_quant_v2_fusion_pass.cpp 覆盖。二、产品支持情况当前算子支持的产品以 README 为准产品是否支持Ascend 950PR 950DT 系列产品√Atlas A3 系列产品×Atlas A2 系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×Kirin X90 处理器系列产品√Kirin 9030 处理器系列产品√这一支持矩阵与 op_host/add_rms_norm_dynamic_quant_v2_def.cpp 中OP_ADD注册的 AICore 配置一致代码里显式AddConfig(ascend910b)、AddConfig(kirinx90)、AddConfig(kirin9030)以及AddConfig(ascend950)同时 op_host/config 目录下提供了ascend910b/ascend950/kirin9030/kirinx90四个平台的 binary 配置文件如 ascend950/add_rms_norm_dynamic_quant_v2_binary.json用于编译生成对应平台的算子二进制。Host 侧 Tiling 也按 arch22910B 系列与 arch35950 系列分别实现见 op_host/arch22/add_rms_norm_dynamic_quant_v2_tiling.cpp 与 op_host/arch35/add_rms_norm_dynamic_quant_v2_tiling_arch35.cpp。三、功能说明与计算公式3.1 计算流程算子按如下顺序执行计算两路输入之和x x1 x2对x做 RmsNorm 归一化乘上gamma可选加beta得到y将y分别送入 1 路或 2 路对称动态量化每路可配置可选的 smoothScale输出量化结果y1/y2与量化尺度scale1/scale2可选地输出 FP32 版本y3、原始输入类型版本y4以及求和结果x供下游算子使用。3.2 数学公式加法与 RmsNorm 归一化$$ xx_{1}x_{2} $$$$ y \operatorname{RmsNorm}(x)\frac{x}{\operatorname{Rms}(\mathbf{x})}\cdot gamma, \quad \text { where } \operatorname{Rms}(\mathbf{x})\sqrt{\frac{1}{n} \sum_{i1}^n x_i^2epsilon} $$FP32 输出对应y3$$ yFP32\begin{cases} cast(y) outputMask[2]True\ ||\ outputMask\ null \ 无效输出 outputMask[2]False \end{cases} $$加入偏置项对应公式中的beta即y_input y beta后两路量化前的输入分别为$$ y_inputybeta $$$$ input1 \begin{cases} y_input \cdot smoothScale1Optional \ \ smoothScale1Optional\ ! null \ y_input \ \ smoothScale1Optional\ null \end{cases} $$$$ input2 \begin{cases} y_input \cdot smoothScale2Optional \ \ smoothScale2Optional\ ! null \ y_input \ \ smoothScale2Optional\ null \end{cases} $$对称动态量化INT8 为例缩放因子取每行最大绝对值除以 127量化结果为round(input/scale)$$ scale1Out\begin{cases} row_max(abs(input1))/127 outputMask[0]True\ ||\ outputMask\ null \ 无效输出 outputMask[0]False \end{cases} $$$$ y1Out\begin{cases} round(input1/scale1Out) outputMask[0]True\ ||\ outputMask\ null \ 无效输出 outputMask[0]False \end{cases} $$$$ scale2Out\begin{cases} row_max(abs(input2))/127 outputMask[1]True\ ||\ (outputMask\ null\ \ smoothScale1Optional\ ! null\ \ smoothScale2Optional\ ! null) \ 无效输出 outputMask[1]False\ ||\ (outputMask\ null\ \ (smoothScale1Optional\ null\ ||\ smoothScale2Optional\ null)) \end{cases} $$$$ y2Out\begin{cases} round(input2/scale2Out) outputMask[1]True\ ||\ (outputMask\ null\ \ smoothScale1Optional\ ! null\ \ smoothScale2Optional\ ! null)\ 无效输出 outputMask[1]False\ ||\ (outputMask\ null\ \ (smoothScale1Optional\ null\ ||\ smoothScale2Optional\ null)) \end{cases} $$其中row_max表示按最后一维行求最大值当outputMask[3]False时不输出y即y4无效。四、参数说明算子完整的输入/输出/属性定义位于 op_graph/add_rms_norm_dynamic_quant_v2_proto.hREG_OP(AddRmsNormDynamicQuantV2)以及 Host 侧注册文件 op_host/add_rms_norm_dynamic_quant_v2_def.cpp。下表为 README 给出的完整参数清单参数名输入/输出/属性描述数据类型数据格式x1输入标准化过程中的源数据张量对应公式x1。支持空 Tensor当输出y1或y2类型为 INT4 时x1的尾轴必须能被 2 整除FLOAT16、BFLOAT16NDx2输入标准化过程中的源数据张量对应公式x2shape 和数据类型与x1一致。支持空 TensorFLOAT16、BFLOAT16NDgamma输入标准化权重张量对应公式gamma数据类型与x1一致shape 需与x1最后一维一致。支持空 TensorFLOAT16、BFLOAT16NDsmooth_scale1可选输入量化得到 y1 使用的 smoothScale 张量对应smoothScale1Optionalshape 与数据类型需与gamma一致。支持空 TensorFLOAT16、BFLOAT16NDsmooth_scale2可选输入量化得到 y2 使用的 smoothScale 张量对应smoothScale2Optionalshape 与数据类型需与gamma一致。支持空 TensorFLOAT16、BFLOAT16NDbeta可选输入标准化过程中的偏置项对应公式betashape 与数据类型需与gamma一致。支持空 TensorFLOAT16、BFLOAT16NDepsilon可选属性防止除 0 错误对应公式epsilon。默认值 1e-6FLOAT-output_mask可选属性输出掩码对应outputMask。只支持长度为 0 或 4 的数组。默认值{}LISTBOOL-dst_type可选属性指定y1和y2的输出数据类型。取值范围{2, 29, 34, 35, 36}分别对应{INT8, INT4, HIFLOAT8, FLOAT8_E5M2, FLOAT8_E4M3FN}。默认值 2INT8INT-y1输出第一路量化输出对应y1Out。有效输出时 shape 和数据类型需与输入x1保持一致。支持空 TensorINT8、INT4、HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FNNDy2输出第二路量化输出对应y2Out。有效输出时 shape 和数据类型需与输入x1保持一致。支持空 TensorINT8、INT4、HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FNNDy3输出RmsNorm 的 FLOAT32 类型输出对应yFP32。有效输出时 shape 需与输入x1保持一致。支持空 TensorFLOAT32NDy4输出RmsNorm 的原始输入类型输出对应y。有效输出时 shape 和数据类型需与输入x1保持一致。支持空 TensorFLOAT16、BFLOAT16NDx输出x1与x2的和对应公式x。shape 和数据类型需与输入x1保持一致。支持空 TensorFLOAT16、BFLOAT16NDscale1输出第一路量化输出尺度对应scale1Out。有效输出时 shape 为x1去掉最后一维后的 shapeFLOAT32NDscale2输出第二路量化输出尺度对应scale2Out。有效输出时 shape 为x1去掉最后一维后的 shapeFLOAT32ND几点源码印证默认值REG_OP中.ATTR(epsilon, Float, 1e-6)、.ATTR(output_mask, ListBool, {})、.ATTR(dst_type, Int, DT_INT8)与 README 表格完全一致输出数据类型推断add_rms_norm_dynamic_quant_v2_infershape.cpp 的InferDataType4AddRmsNormDynamicQuantV2会把dst_type属性直接写入y1/y2的输出数据类型y3固定为DT_FLOATy4/x与x1同类型scale1/scale2固定为DT_FLOATscale 的 shape 推导InferReduceShape用xDimNum - gammaDimNum得到 reduce 维度数即 scale 输出比x1少最后一维INT4 尾轴对齐README 中 x1尾轴必须能被 2 整除 的约束与 Kernel 侧量化按 32 元素对齐numLastDimAligned见 add_rms_norm_dynamic_quant_v2_base.h 中注释 Quantize better be aligned to 32 elements的实现取向一致。4.1 Atlas A2 与 Kirin 系列的平台差异化限制README 特别指出在Atlas A2 系列产品、Kirin X90 处理器系列产品、Kirin 9030 处理器系列产品上算子行为有额外限制x1、x2、gamma、smooth_scale1、smooth_scale2、y4和x的数据类型不支持 BFLOAT16仅 FLOAT16y1和y2的数据类型仅支持 INT8beta、output_mask和dst_type的配置无效即配置了也不生效y1和y2的输出情况仅与smooth_scale1和smooth_scale2的输入情况有关且仅y2可不输出。这一限制同样能在 add_rms_norm_dynamic_quant_v2_def.cpp 的GetKirinCoreConfig()中得到印证Kirin 平台配置里所有输入输出数据类型只注册了DT_FLOAT16y1/y2只注册了DT_INT8且smooth_scale1/smooth_scale2/beta均为OPTIONAL。因此在 Kirin 平台上构建算子时应使用 FLOAT16 INT8 的组合不要依赖output_mask与dst_type来控制输出。五、约束说明README 对输出有效性的约束如下构造图时必须严格遵守当output_mask不为空时长度为 4参数smooth_scale1有值时output_mask[0]必须为 True参数smooth_scale2有值时output_mask[1]必须为 Trueoutput_mask[0]和output_mask[1]不能同时为 False各输出有效性由output_mask统一控制对应位置为 True 时yy1/y2/y3/y4与scalescale1/scale2为有效输出为 False 时为无效输出。当output_mask为空时长度为 0参数smooth_scale2有值时参数smooth_scale1不能为空即不允许只有第二路 smoothScale的配置y1、y3、y4和scale1始终为有效输出y2和scale2只有在smooth_scale1与smooth_scale2均有效时才为有效输出否则为无效输出。上述规则在 InferShape 源码中有直接对应实现add_rms_norm_dynamic_quant_v2_infershape.cpp 中当output_mask非空但长度不为 4 时直接返回GRAPH_FAILEDFillKnownRankShapesV2中的OP_CHECK_IF(outputMaskLen ! NUM_FOUR, ...)当output_mask为空且(!smooth1Exist) smooth2Exist时同样报错错误信息为 When output_mask is NULL, AddRmsNormDynamicQuantV2 Not support only have scale2.未知 shaperank 不确定场景由HandleUnknownRankShapesV2单独处理output_mask对应位为 False 时 scale 输出置为Shape({1})。另外gamma与smooth_scale1/smooth_scale2/beta的 shape 一致性也在 InferShape 中强制校验GammaShape is not same to smooth1Shape.等报错分支。六、调用方式图模式示例README 给出的调用方式为图模式即通过算子 IR 构图add_rms_norm_dynamic_quant_v2_proto.h方式在计算图中创建算子节点参考样例为 examples/test_geir_add_rms_norm_dynamic_quant_v2.cpp。调用方式样例代码说明图模式test_geir_add_rms_norm_dynamic_quant_v2.cpp通过算子 IR 构图方式调用 AddRmsNormDynamicQuantV2 算子6.1 核心构图步骤拆解该示例展示了用ge::op::AddRmsNormDynamicQuantV2在 Graph 中构图并运行的核心流程构造算子节点auto add1 op::AddRmsNormDynamicQuantV2(add1);声明 shape示例中使用x1Shape {4, 1, 8}、x2Shape {4, 1, 8}、gammaShape {8}、scale1Shape {8}、scale2Shape {8}即最后一维 D8gamma/smoothScale 的 shape 与尾轴一致符合第四章参数约束绑定输入通过ADD_INPUT宏依次为x1、x2、gamma、smooth_scale1、smooth_scale2创建op::Data占位节点并set_input_*连接到算子同时构造全 1 的 Host 侧 Tensor 数据GenOnesData/GenOnesDataFloat32设置输出outputs.push_back(add1)将算子节点整体作为图输出初始化 GE 会话ge::GEInitialize(global_options)其中global_options配置了{ge.exec.deviceId, 0}与{ge.graphRunMode, 1}建图并运行session-AddGraph(graph_id, graph, graph_options)后调用session-RunGraph(graph_id, input, output)结果导出运行结束后把输入/输出数据按tc_ge_irrun_test_0008_npu_input_i.bin/tc_ge_irrun_test_0008_npu_output_i.bin的命名写入当前目录便于离线比对同时可用aclgrphDumpGraph(graph, ./dump, ...)导出图文件用于调试。示例默认使用DT_BF16作为输入类型DataType inDtype DT_BF16;在 910B/950 平台可直接运行若目标平台为 Atlas A2/Kirin 系列需按 4.1 节限制改为DT_FLOAT16。6.2 属性设置与输出掩码示例在测试用例 tests/ut/op_host/test_AddRmsNormDynamicQuantV2_infershape.cpp 中可以找到属性的标准设置写法op.SetAttr(epsilon, static_castfloat(1e-6)); std::vectorbool out_shape {true, true, true, true}; // output_mask 长度为 4 op.SetAttr(output_mask, out_shape); op.SetAttr(dst_type, 2); // INT8对应输入x1/x2为{8, 64}、gamma/smooth_scale1/smooth_scale2为{64}时InferShape 期望的验证结果为y1/y2/y3/y4/x均为{8, 64}scale1/scale2均为{8}即x1去掉最后一维。测试还覆盖了未知 rank{-2}场景以及output_mask为空{}时smooth_scale2依赖smooth_scale1的逻辑分支。七、从源码看实现纵深Tiling 策略与 Kernel 分派为了帮助读者理解算子在 NPU 上如何高效执行这里补充说明 Host 侧 Tiling 与 Kernel 侧的配合关系属源码级补充不影响 README 给出的使用方式。7.1 Tiling 三种策略op_host/arch22/add_rms_norm_dynamic_quant_v2_tiling.cpp 中定义了三种 UB Tiling 策略并通过context_-SetTilingKey(tilingKey)写入 tiling keyUB_TILING_POLICY_NORMALkey1常规分块按firstDimPerCore将行维均分到多核UB_TILING_POLICY_SINGLE_ROWkey2单行处理适用于尾轴较长、单行即可占满 UB 的场景UB_TILING_POLICY_SLICE_Dkey3沿 D尾轴切片SLICE_COL_LEN 8864为单次切片的列长度用于尾轴超长时的分段处理并需要额外申请 workspaceuseCore * numLastDim * sizeof(float) * workspaceRowsNum字节。Tiling 数据字段useCore、numFirstDim、numLastDim、numLastDimAligned、firstDimPerLoop、lastDimSliceLen、smoothNum、epsilon、avgFactor等通过AddRmsNormDynamicQuantV2TilingData序列化后传给 Kernel其中smoothNum0/1/2直接决定一路还是两路量化生效。7.2 Kernel 入口与分派Kernel 入口 op_kernel/add_rms_norm_dynamic_quant_v2.cpp 根据 Tiling key 分派到三种模板实现if (TILING_KEY_IS(0)) { // 0 Tiling, Do Nothing. } else if (TILING_KEY_IS(1)) { KernelAddRmsNormDynamicQuantV2NormalDTYPE_X1, 1 op(pipe); INIT_AND_PROCESS; } else if (TILING_KEY_IS(2)) { KernelAddRmsNormDynamicQuantV2SingleRowDTYPE_X1, 2 op(pipe); INIT_AND_PROCESS; } else if (TILING_KEY_IS(3)) { KernelAddRmsNormDynamicQuantV2SliceDDTYPE_X1, 3 op(pipe); INIT_AND_PROCESS; }三种 Kernel 实现分别位于 add_rms_norm_dynamic_quant_v2_normal_kernel.h、add_rms_norm_dynamic_quant_v2_single_row_kernel.h 与 add_rms_norm_dynamic_quant_v2_cut_d_kernel.h公共逻辑抽在 add_rms_norm_dynamic_quant_v2_base.h 中多核行划分、smooth1Exist/smooth2Exist判断、双 scale 缓冲等。ascend950 平台另有独立的 arch35 Kernel 实现 op_kernel/arch35/add_rms_norm_dynamic_quant_v2.cpp。7.3 测试覆盖仓库为该算子提供了完整的单测矩阵tests/utop_hostInferShape 单测 test_AddRmsNormDynamicQuantV2_infershape.cpp、arch22 与 arch35 的 Tiling 单测op_graph融合 Pass 单测 test_add_rms_norm_dynamic_quant_v2_fusion_pass.cppop_kernelKernel 行为单测 test_add_rms_norm_dynamic_quant_v2.cpp。这些测试文件可作为理解算子边界行为与自测算子移植的第一手参考。八、总结AddRmsNormDynamicQuantV2 是 CANN ops-nn 中面向大模型推理量化场景的典型融合算子它将 Add、RmsNorm 与至多两路对称动态量化合并为单次 Kernel 执行通过smooth_scale1/smooth_scale2、beta、output_mask、dst_type等参数灵活控制两路量化的启用与各输出的有效性。使用时需重点注意三点一是output_mask为空时不允许仅第二路 smoothScale的配置二是 Atlas A2/Kirin 平台仅支持 FLOAT16 INT8 且beta/output_mask/dst_type配置无效三是 INT4 输出要求x1尾轴能被 2 整除。开发者可参照 图模式示例 与 算子 IR 定义 快速完成构图接入并通过仓库内 Tiling/Kernel/InferShape 源码与单测深入了解其 NPU 上的执行细节。赞分享人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载相关推荐CANN ops-nn AdaLayerNormV2 算子深度解析自适应 LayerNorm 融合算子原理、参数与 aclnn 调用实战CANN ops nn AdaLayerNormV2 算子深度解析自适应 LayerNorm 融合算子原理、参数与 aclnn 调用实战 导读 AdaLaye人工智能算子库深度学习CANNAscendCANN ops-nn IndexFill 算子深度解析原理、参数与 aclnn 调用实战CANN ops nn IndexFill 算子深度解析原理、参数与 aclnn 调用实战 导读 IndexFill 是 CANN ops nn 神经网络算子人工智能算子库深度学习CANNAscendCANN ops-nn PReluGradUpdate 算子深度解析原理、aclnn 调用与图融合实现CANN ops nn PReluGradUpdate 算子深度解析原理、aclnn 调用与图融合实现 PReluGradUpdate 是 CANN ops人工智能算子库深度学习CANNAscend上一篇Security-101 IAM 核心概念精讲身份与访问管理、最小权限与认证授权的实战指南下一篇StarRocks ds_theta_union 标量函数详解Theta Sketch 两两并集与去重统计实战创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考