- 人工智能
- 算子库
- 深度学习
- CANN
- Ascend
【免费下载链接】ops-nn
本项目是CANN提供的神经网络类计算算子库,实现网络在NPU上加速计算。
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 结构中,RmsNorm(Root Mean Square Layer Normalization)是最常用的归一化算子:相比 LayerNorm,它去掉了"减去均值"的步骤,只做均方根归一化,节省了一次规约与减法开销。而在量化推理管线中,归一化输出通常还要紧接着送入对称动态量化算子(DynamicQuant)转成 INT8 等低比特数据,以匹配后续 MatMul 的量化输入。
如果 Add、RmsNorm、DynamicQuant 各自独立执行,中间结果x1+x2、归一化输出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 也按 arch22(910B 系列)与 arch35(950 系列)分别实现,见 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 归一化:
$$ x=x_{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_{i=1}^n x_i^2+epsilon} $$
FP32 输出(对应y3):
$$ yFP32=\begin{cases} cast(y) & outputMask[2]=True\ ||\ outputMask\ = null \ 无效输出 & outputMask[2]=False \end{cases} $$
加入偏置项(对应公式中的beta,即y_input = y + beta)后,两路量化前的输入分别为:
$$ y_input=y+beta $$
$$ 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.h(REG_OP(AddRmsNormDynamicQuantV2))以及 Host 侧注册文件 op_host/add_rms_norm_dynamic_quant_v2_def.cpp。下表为 README 给出的完整参数清单:
| 参数名 | 输入/输出/属性 | 描述 | 数据类型 | 数据格式 |
|---|---|---|---|---|
| x1 | 输入 | 标准化过程中的源数据张量,对应公式x1。支持空 Tensor;当输出y1或y2类型为 INT4 时,x1的尾轴必须能被 2 整除 | FLOAT16、BFLOAT16 | ND |
| x2 | 输入 | 标准化过程中的源数据张量,对应公式x2,shape 和数据类型与x1一致。支持空 Tensor | FLOAT16、BFLOAT16 | ND |
| gamma | 输入 | 标准化权重张量,对应公式gamma,数据类型与x1一致,shape 需与x1最后一维一致。支持空 Tensor | FLOAT16、BFLOAT16 | ND |
| smooth_scale1 | 可选输入 | 量化得到 y1 使用的 smoothScale 张量,对应smoothScale1Optional,shape 与数据类型需与gamma一致。支持空 Tensor | FLOAT16、BFLOAT16 | ND |
| smooth_scale2 | 可选输入 | 量化得到 y2 使用的 smoothScale 张量,对应smoothScale2Optional,shape 与数据类型需与gamma一致。支持空 Tensor | FLOAT16、BFLOAT16 | ND |
| beta | 可选输入 | 标准化过程中的偏置项,对应公式beta,shape 与数据类型需与gamma一致。支持空 Tensor | FLOAT16、BFLOAT16 | ND |
| epsilon | 可选属性 | 防止除 0 错误,对应公式epsilon。默认值 1e-6 | FLOAT | - |
| output_mask | 可选属性 | 输出掩码,对应outputMask。只支持长度为 0 或 4 的数组。默认值{} | LISTBOOL | - |
| dst_type | 可选属性 | 指定y1和y2的输出数据类型。取值范围{2, 29, 34, 35, 36},分别对应{INT8, INT4, HIFLOAT8, FLOAT8_E5M2, FLOAT8_E4M3FN}。默认值 2(INT8) | INT | - |
| y1 | 输出 | 第一路量化输出,对应y1Out。有效输出时 shape 和数据类型需与输入x1保持一致。支持空 Tensor | INT8、INT4、HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FN | ND |
| y2 | 输出 | 第二路量化输出,对应y2Out。有效输出时 shape 和数据类型需与输入x1保持一致。支持空 Tensor | INT8、INT4、HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FN | ND |
| y3 | 输出 | RmsNorm 的 FLOAT32 类型输出,对应yFP32。有效输出时 shape 需与输入x1保持一致。支持空 Tensor | FLOAT32 | ND |
| y4 | 输出 | RmsNorm 的原始输入类型输出,对应y。有效输出时 shape 和数据类型需与输入x1保持一致。支持空 Tensor | FLOAT16、BFLOAT16 | ND |
| x | 输出 | x1与x2的和,对应公式x。shape 和数据类型需与输入x1保持一致。支持空 Tensor | FLOAT16、BFLOAT16 | ND |
| scale1 | 输出 | 第一路量化输出尺度,对应scale1Out。有效输出时 shape 为x1去掉最后一维后的 shape | FLOAT32 | ND |
| scale2 | 输出 | 第二路量化输出尺度,对应scale2Out。有效输出时 shape 为x1去掉最后一维后的 shape | FLOAT32 | ND |
几点源码印证:
- 默认值:
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_FLOAT,y4/x与x1同类型,scale1/scale2固定为DT_FLOAT; - scale 的 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(仅 FLOAT16);y1和y2的数据类型仅支持 INT8;beta、output_mask和dst_type的配置无效(即配置了也不生效);y1和y2的输出情况仅与smooth_scale1和smooth_scale2的输入情况有关,且仅y2可不输出。
这一限制同样能在 add_rms_norm_dynamic_quant_v2_def.cpp 的GetKirinCoreConfig()中得到印证:Kirin 平台配置里所有输入输出数据类型只注册了DT_FLOAT16,y1/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]必须为 True; output_mask[0]和output_mask[1]不能同时为 False;- 各输出有效性由
output_mask统一控制:对应位置为 True 时y(y1/y2/y3/y4)与scale(scale1/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_FAILED(FillKnownRankShapesV2中的OP_CHECK_IF(outputMaskLen != NUM_FOUR, ...)); - 当
output_mask为空且(!smooth1Exist) && smooth2Exist时同样报错,错误信息为 "When output_mask is NULL, AddRmsNormDynamicQuantV2 Not support only have scale2."; - 未知 shape(rank 不确定)场景由
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},即最后一维 D=8,gamma/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_cast<float>(1e-6)); std::vector<bool> 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 key:
UB_TILING_POLICY_NORMAL(key=1):常规分块,按firstDimPerCore将行维均分到多核;UB_TILING_POLICY_SINGLE_ROW(key=2):单行处理,适用于尾轴较长、单行即可占满 UB 的场景;UB_TILING_POLICY_SLICE_D(key=3):沿 D(尾轴)切片,SLICE_COL_LEN = 8864为单次切片的列长度,用于尾轴超长时的分段处理,并需要额外申请 workspace(useCore * numLastDim * sizeof(float) * workspaceRowsNum字节)。
Tiling 数据字段(useCore、numFirstDim、numLastDim、numLastDimAligned、firstDimPerLoop、lastDimSliceLen、smoothNum、epsilon、avgFactor等)通过AddRmsNormDynamicQuantV2TilingData序列化后传给 Kernel,其中smoothNum(0/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)) { KernelAddRmsNormDynamicQuantV2Normal<DTYPE_X1, 1> op(&pipe); INIT_AND_PROCESS; } else if (TILING_KEY_IS(2)) { KernelAddRmsNormDynamicQuantV2SingleRow<DTYPE_X1, 2> op(&pipe); INIT_AND_PROCESS; } else if (TILING_KEY_IS(3)) { KernelAddRmsNormDynamicQuantV2SliceD<DTYPE_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/ut):
- op_host:InferShape 单测 test_AddRmsNormDynamicQuantV2_infershape.cpp、arch22 与 arch35 的 Tiling 单测;
- op_graph:融合 Pass 单测 test_add_rms_norm_dynamic_quant_v2_fusion_pass.cpp;
- op_kernel:Kernel 行为单测 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 上的执行细节。
- 人工智能
- 算子库
- 深度学习
- CANN
- Ascend
【免费下载链接】ops-nn
本项目是CANN提供的神经网络类计算算子库,实现网络在NPU上加速计算。
相关推荐
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
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考