CANN 自定义算子 npu_hc_pre_inv_rms 详解mHC 结构 HcPre 前处理中的 InvRms 计算与 PyTorch 调用实战【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer导读本文以 custom-npu_hc_pre_inv_rms.md 为技术主线深入讲解 CANN 昇腾 NPU 上自定义算子custom.npu_hc_pre_inv_rms的功能定位、数学原理、函数原型、参数约束与调用方式。该算子是 mHCmulti-Head Combination注意力结构前处理HcPre流程中被拆出的 InvRms 计算小算子用于求输入张量沿最后一个维度均方根的倒数即 RMSNorm 归一化系数的倒数。读完本文你将掌握该算子的完整规格shape/数据类型/epsilon 约束、它在 HcPre 融合/拆分策略中的角色以及如何通过 PyTorch eager 与torch.compile图模式两种方式调用并验证该算子。一、功能定位HcPre 前处理中的 InvRms 计算npu_hc_pre_inv_rms是 mHC 结构 HcPrepre-processing of multi-Head Combination流程中用于计算InvRmsInverse RMS均方根倒数的算子。关于 HcPre 的完整计算过程可参考同目录下的 custom-npu_hc_pre.mdHcPre 是 mHC 结构前处理部分的融合算子其内部涉及 InvRms 计算与 Sinkhorn 迭代等步骤。当输入规模满足特定条件时T/bs 小于等于 128 且能被 16 整除会使能hc_pre融合算子以获得较高性能其他场景下则会退化为hc_pre_inv_rms与hc_pre_sinkhorn两个小算子拼接执行。因此本文所述的npu_hc_pre_inv_rms正是该拆分路径中负责 InvRms 计算的独立算子是理解 HcPre 整体行为不可或缺的一环。说明本文涉及算子的形状参数中bbatch size表示输入样本批量大小ssequence length表示输入样本序列长度hchead count表示注意力头数dhead dimension表示注意力头维度数T 表示 b 与 s 合轴flatten后的大小。二、产品支持情况产品是否支持Atlas A3 推理系列产品√Ascend 950PR / Ascend 950DT√该算子定位于推理场景使用与 CANN 自定义算子工程ops/ascendc中的其他算子保持一致的硬件支持范围。三、数学原理InvRms 计算公式算子的计算逻辑即 InvRmsInverse RMS其定义如下$$ \text{InvRms}(x) \frac{1}{\text{RMS}(x)} $$其中 RMS 为均方根$$ \text{RMS}(x) \sqrt{\frac{1}{N} \sum_{i1}^{N} x_i^2 \epsilon} $$即对输入张量x沿最后一个维度求平方均值加上 epsilon 后再开方取倒数。从测试用例 test_npu_hc_pre_inv_rms.py 中提供的 CPU 参考实现_hc_pre_inv_rms可以直观看出其等价计算过程def _hc_pre_inv_rms(x, epsilon1e-20): if x.dim() 4: x x.flatten(2) # (b, s, hc, d) - (b, s, hc * d) elif x.dim() 3: x x.flatten(1) # (T, hc, d) - (T, hc * d) x x.float() y torch.rsqrt(x.square().mean(-1, keepdimTrue) epsilon) return y要点参考实现先将 hc 与 d 两维合轴再沿最后一个维度求square().mean()加上 epsilon 后使用torch.rsqrt求倒数。这与算子文档中shape 为 [T, hc_mult, d] 或 [b, s, hc_mult, d] 的输入、输出 shape 为 [T, 1] 或 [b, s, 1]的规格完全对应——即最终规约维度恒为最后一个维度输出在每个样本位置只保留一个标量。四、函数原型与参数说明4.1 函数原型custom.npu_hc_pre_inv_rms(Tensor x, *, float epsilon1e-20) - Tensor该原型在源码中有两处对应注册算子 schema 定义位于 ops_def_registration.cppm.def(npu_hc_pre_inv_rms(Tensor x, *, float epsilon1e-20) - Tensor);NPU 与 Meta 设备的前向实现注册位于 npu_hc_pre_inv_rms.cppTORCH_LIBRARY_IMPL(custom, PrivateUse1, ...)与TORCH_LIBRARY_IMPL(custom, Meta, ...)。4.2 参数说明xTensor必选参数输入张量即公式中参与 InvRms 计算的输入。数据格式ND数据类型bfloat16C 侧实现同时校验并接受float16、float32详见下文源码实现解析shape[T, hc_mult, d]或[b, s, hc_mult, d]不支持非连续non-contiguous张量*符号含义*之前的参数是位置相关参数必须按照顺序输入必选*之后的参数以键值对形式赋值与位置无关可选不传入时使用默认值。epsilonfloat可选参数RmsNorm 公式中的 ε 参数。Host 侧参数C 实现中按double类型接收npu_hc_pre_inv_rms_npu(const at::Tensor x, double epsilon1e-20)默认值为1e-20约束必须大于等于 0C 实现通过TORCH_CHECK(epsilon 0, epsilon should be greater than 0.)强制校验4.3 返回值说明yTensorInvRms 的计算输出张量。数据格式ND数据类型floatshape[T, 1]或[b, s, 1]输出 shape 的推导逻辑在 npu_hc_pre_inv_rms.cpp 的construct_hc_pre_inv_rms_output_tensor中实现保留输入前xDims - 2个维度最后一个维度固定为 1即x: (b, s, hc, d) -- y: (b, s, 1)或x: (b * s, hc, d) -- y: (b * s, 1)输出张量 dtype 通过options.dtype(at::kFloat)强制设为 float32。五、约束说明5.1 shape 取值范围约束字段名取值规则与说明hc_mult取值固定为4d取值固定为4096注意算子文档将 d 约束为 4096但测试用例中同时覆盖了d 4096与d 7168两组场景见 test_npu_hc_pre_inv_rms.py实际使用时应以目标 CANN 版本算子规格为准并按需通过测试验证扩展 shape 的可用性。5.2 其他约束该接口支持推理场景下使用。该接口支持aclgraph 入图即通过torch.compiletorchair后端进行图模式编译。该接口与 PyTorch 配合使用时需要保证 CANN 相关包与 PyTorch 相关包的版本匹配。六、源码实现解析从 PyTorch 算子到 AClNN 执行npu_hc_pre_inv_rms的完整调用链可从仓库源码中梳理出来分为三层6.1 C 封装层NPU 与 Meta 双实现npu_hc_pre_inv_rms.cpp 中实现了算子前向逻辑npu_hc_pre_inv_rms_npuNPU 设备实现。执行三类校验——输入非空x.numel() 0、epsilon 非负、输入 dtype 为kFloat / kHalf / kBFloat16三者之一随后构造输出张量并通过宏EXEC_NPU_CMD_V1(aclnnHcPreInvRms, x, epsilon, yOut)下发到底层 AClNN 接口aclnnHcPreInvRms执行。npu_hc_pre_inv_rms_metaMeta 设备实现仅做参数校验与输出 shape 推导供torch.compile的 meta 模式使用。两个实现分别通过TORCH_LIBRARY_IMPL(custom, PrivateUse1, ...)与TORCH_LIBRARY_IMPL(custom, Meta, ...)注册到custom命名空间最终对外暴露为torch.ops.custom.npu_hc_pre_inv_rms。6.2 Graph 模式转换层fx2ge converter在torch.compile图模式下算子通过 converter/npu_hc_pre_inv_rms.py 注册的register_fx_node_ge_converter(torch.ops.custom.npu_hc_pre_inv_rms.default)转换器将 fx 节点映射为 GE 自定义算子节点register_fx_node_ge_converter(torch.ops.custom.npu_hc_pre_inv_rms.default) def convert_npu_hc_pre_inv_rms(x: Tensor, *, epsilon: float): return torchair.ge.custom_op( HcPreInvRms, inputs{x: x}, attrs{epsilon: epsilon}, outputs[y] )由此实现 aclgraph 入图算子名HcPreInvRms、输入x、属性epsilon、输出y与文档规格一一对应。该 converter 位于 custom_ops/converter并在 custom_ops/__init__.py 中被导入导出。6.3 依赖与版本前提调用该算子需要环境具备PyTorch、torch_npu、torchair以及 CANN 工具链算子实际计算逻辑aclnnHcPreInvRms由 CANN 软件栈提供。测试脚本 test_npu_hc_pre_inv_rms.py 顶部展示了运行所需的依赖import torch import torch_npu import torchair import custom_ops import numpy as np import torch.nn as nn其中custom_ops为 torch_ops_extension/custom_ops 编译生成的算子扩展包。七、调用示例与精度验证官方调用示例位于 examples/test_npu_hc_pre_inv_rms.py覆盖了eager 模式与graph 模式两种调用路径并内置了 CPU 参考实现对比的精度校验。7.1 Eager 模式调用eager 模式直接通过torch.ops.custom.npu_hc_pre_inv_rms调用测试覆盖了 4 维输入(b, s, hc, d)与 3 维输入(T, hc, d)、d 为 4096 与 7168、以及 float32 / bfloat16 / float16 三种输入 dtype 的组合b, s, hc, d 4, 4, 4, 4096 eps 1e-6 # 构造输入并搬运到 NPU x torch.tensor(np.random.uniform(-1, 1, (b, s, hc, d))).to(torch.float32) x x.to(npu:0) # 调用自定义算子 npu_y torch.ops.custom.npu_hc_pre_inv_rms(x, epsiloneps)7.2 Graph 模式调用aclgraph 入图graph 模式通过torch.compile搭配torchair的 NPU 后端实现整图编译验证算子支持 aclgraph 入图的能力from torchair.configs.compiler_config import CompilerConfig class Network(nn.Module): def forward(self, x, epsilon): return torch.ops.custom.npu_hc_pre_inv_rms(x, epsilonepsilon) config CompilerConfig() config.mode reduce-overhead npu_backend torchair.get_npu_backend(compiler_configconfig) npu_mode torch.compile(Network().to(npu:0), fullgraphTrue, backendnpu_backend, dynamicFalse) npu_y npu_mode(x_npu, eps)7.3 精度验证方式测试脚本通过data_compare与cal_relative_diff_np实现 NPU 输出与 CPU 参考输出的相对误差比对默认diff_thd0.0001、pct_thd0.0005、max_diff_hd0.0001满足通过百分比不低于阈值且最大相对误差不超限即判定为Pass。每个用例最终通过assert(compare_y[0] Pass)断言可作为复现算子正确性的最小验证集。八、注意事项与常见问题epsilon 参数校验epsilon 必须为非负数C 层会强制TORCH_CHECK文档默认值为1e-20测试用例中常显式传入1e-6。输入 dtype 差异文档声明输入数据类型支持bfloat16但 C 实现校验允许float32、float16、bfloat16三种 dtype测试用例也覆盖了这三种类型实际可用范围以源码校验逻辑为准。连续性问题输入不支持非连续张量调用前应保证张量内存连续。版本匹配与 PyTorch 配合使用时须确保 CANN 相关包与 PyTorch 相关包的版本匹配否则可能出现算子注册或 AClNN 下发失败。与 HcPre 融合算子的关系当 T/bs 较小≤128 且能被 16 整除时框架会倾向使用hc_pre融合算子以获得更优性能npu_hc_pre_inv_rms主要用于拆分路径是理解 HcPre 内部 InvRms 步骤的独立算子详见 custom-npu_hc_pre.md。九、相关资源算子规格文档custom-npu_hc_pre_inv_rms.md调用示例与测试用例test_npu_hc_pre_inv_rms.pyC 封装实现npu_hc_pre_inv_rms.cpp算子 schema 注册ops_def_registration.cppGraph 模式转换器npu_hc_pre_inv_rms.py上游 HcPre 融合算子文档custom-npu_hc_pre.mdHcPre 完整参考实现test_npu_hc_pre.pyAscendC 自定义算子工程总览ops/ascendc/README.md【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN 自定义算子 npu_hc_pre_inv_rms 详解：mHC 结构 HcPre 前处理中的 InvRms 计算与 PyTorch 调用实战
CANN 自定义算子 npu_hc_pre_inv_rms 详解mHC 结构 HcPre 前处理中的 InvRms 计算与 PyTorch 调用实战【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer导读本文以 custom-npu_hc_pre_inv_rms.md 为技术主线深入讲解 CANN 昇腾 NPU 上自定义算子custom.npu_hc_pre_inv_rms的功能定位、数学原理、函数原型、参数约束与调用方式。该算子是 mHCmulti-Head Combination注意力结构前处理HcPre流程中被拆出的 InvRms 计算小算子用于求输入张量沿最后一个维度均方根的倒数即 RMSNorm 归一化系数的倒数。读完本文你将掌握该算子的完整规格shape/数据类型/epsilon 约束、它在 HcPre 融合/拆分策略中的角色以及如何通过 PyTorch eager 与torch.compile图模式两种方式调用并验证该算子。一、功能定位HcPre 前处理中的 InvRms 计算npu_hc_pre_inv_rms是 mHC 结构 HcPrepre-processing of multi-Head Combination流程中用于计算InvRmsInverse RMS均方根倒数的算子。关于 HcPre 的完整计算过程可参考同目录下的 custom-npu_hc_pre.mdHcPre 是 mHC 结构前处理部分的融合算子其内部涉及 InvRms 计算与 Sinkhorn 迭代等步骤。当输入规模满足特定条件时T/bs 小于等于 128 且能被 16 整除会使能hc_pre融合算子以获得较高性能其他场景下则会退化为hc_pre_inv_rms与hc_pre_sinkhorn两个小算子拼接执行。因此本文所述的npu_hc_pre_inv_rms正是该拆分路径中负责 InvRms 计算的独立算子是理解 HcPre 整体行为不可或缺的一环。说明本文涉及算子的形状参数中bbatch size表示输入样本批量大小ssequence length表示输入样本序列长度hchead count表示注意力头数dhead dimension表示注意力头维度数T 表示 b 与 s 合轴flatten后的大小。二、产品支持情况产品是否支持Atlas A3 推理系列产品√Ascend 950PR / Ascend 950DT√该算子定位于推理场景使用与 CANN 自定义算子工程ops/ascendc中的其他算子保持一致的硬件支持范围。三、数学原理InvRms 计算公式算子的计算逻辑即 InvRmsInverse RMS其定义如下$$ \text{InvRms}(x) \frac{1}{\text{RMS}(x)} $$其中 RMS 为均方根$$ \text{RMS}(x) \sqrt{\frac{1}{N} \sum_{i1}^{N} x_i^2 \epsilon} $$即对输入张量x沿最后一个维度求平方均值加上 epsilon 后再开方取倒数。从测试用例 test_npu_hc_pre_inv_rms.py 中提供的 CPU 参考实现_hc_pre_inv_rms可以直观看出其等价计算过程def _hc_pre_inv_rms(x, epsilon1e-20): if x.dim() 4: x x.flatten(2) # (b, s, hc, d) - (b, s, hc * d) elif x.dim() 3: x x.flatten(1) # (T, hc, d) - (T, hc * d) x x.float() y torch.rsqrt(x.square().mean(-1, keepdimTrue) epsilon) return y要点参考实现先将 hc 与 d 两维合轴再沿最后一个维度求square().mean()加上 epsilon 后使用torch.rsqrt求倒数。这与算子文档中shape 为 [T, hc_mult, d] 或 [b, s, hc_mult, d] 的输入、输出 shape 为 [T, 1] 或 [b, s, 1]的规格完全对应——即最终规约维度恒为最后一个维度输出在每个样本位置只保留一个标量。四、函数原型与参数说明4.1 函数原型custom.npu_hc_pre_inv_rms(Tensor x, *, float epsilon1e-20) - Tensor该原型在源码中有两处对应注册算子 schema 定义位于 ops_def_registration.cppm.def(npu_hc_pre_inv_rms(Tensor x, *, float epsilon1e-20) - Tensor);NPU 与 Meta 设备的前向实现注册位于 npu_hc_pre_inv_rms.cppTORCH_LIBRARY_IMPL(custom, PrivateUse1, ...)与TORCH_LIBRARY_IMPL(custom, Meta, ...)。4.2 参数说明xTensor必选参数输入张量即公式中参与 InvRms 计算的输入。数据格式ND数据类型bfloat16C 侧实现同时校验并接受float16、float32详见下文源码实现解析shape[T, hc_mult, d]或[b, s, hc_mult, d]不支持非连续non-contiguous张量*符号含义*之前的参数是位置相关参数必须按照顺序输入必选*之后的参数以键值对形式赋值与位置无关可选不传入时使用默认值。epsilonfloat可选参数RmsNorm 公式中的 ε 参数。Host 侧参数C 实现中按double类型接收npu_hc_pre_inv_rms_npu(const at::Tensor x, double epsilon1e-20)默认值为1e-20约束必须大于等于 0C 实现通过TORCH_CHECK(epsilon 0, epsilon should be greater than 0.)强制校验4.3 返回值说明yTensorInvRms 的计算输出张量。数据格式ND数据类型floatshape[T, 1]或[b, s, 1]输出 shape 的推导逻辑在 npu_hc_pre_inv_rms.cpp 的construct_hc_pre_inv_rms_output_tensor中实现保留输入前xDims - 2个维度最后一个维度固定为 1即x: (b, s, hc, d) -- y: (b, s, 1)或x: (b * s, hc, d) -- y: (b * s, 1)输出张量 dtype 通过options.dtype(at::kFloat)强制设为 float32。五、约束说明5.1 shape 取值范围约束字段名取值规则与说明hc_mult取值固定为4d取值固定为4096注意算子文档将 d 约束为 4096但测试用例中同时覆盖了d 4096与d 7168两组场景见 test_npu_hc_pre_inv_rms.py实际使用时应以目标 CANN 版本算子规格为准并按需通过测试验证扩展 shape 的可用性。5.2 其他约束该接口支持推理场景下使用。该接口支持aclgraph 入图即通过torch.compiletorchair后端进行图模式编译。该接口与 PyTorch 配合使用时需要保证 CANN 相关包与 PyTorch 相关包的版本匹配。六、源码实现解析从 PyTorch 算子到 AClNN 执行npu_hc_pre_inv_rms的完整调用链可从仓库源码中梳理出来分为三层6.1 C 封装层NPU 与 Meta 双实现npu_hc_pre_inv_rms.cpp 中实现了算子前向逻辑npu_hc_pre_inv_rms_npuNPU 设备实现。执行三类校验——输入非空x.numel() 0、epsilon 非负、输入 dtype 为kFloat / kHalf / kBFloat16三者之一随后构造输出张量并通过宏EXEC_NPU_CMD_V1(aclnnHcPreInvRms, x, epsilon, yOut)下发到底层 AClNN 接口aclnnHcPreInvRms执行。npu_hc_pre_inv_rms_metaMeta 设备实现仅做参数校验与输出 shape 推导供torch.compile的 meta 模式使用。两个实现分别通过TORCH_LIBRARY_IMPL(custom, PrivateUse1, ...)与TORCH_LIBRARY_IMPL(custom, Meta, ...)注册到custom命名空间最终对外暴露为torch.ops.custom.npu_hc_pre_inv_rms。6.2 Graph 模式转换层fx2ge converter在torch.compile图模式下算子通过 converter/npu_hc_pre_inv_rms.py 注册的register_fx_node_ge_converter(torch.ops.custom.npu_hc_pre_inv_rms.default)转换器将 fx 节点映射为 GE 自定义算子节点register_fx_node_ge_converter(torch.ops.custom.npu_hc_pre_inv_rms.default) def convert_npu_hc_pre_inv_rms(x: Tensor, *, epsilon: float): return torchair.ge.custom_op( HcPreInvRms, inputs{x: x}, attrs{epsilon: epsilon}, outputs[y] )由此实现 aclgraph 入图算子名HcPreInvRms、输入x、属性epsilon、输出y与文档规格一一对应。该 converter 位于 custom_ops/converter并在 custom_ops/__init__.py 中被导入导出。6.3 依赖与版本前提调用该算子需要环境具备PyTorch、torch_npu、torchair以及 CANN 工具链算子实际计算逻辑aclnnHcPreInvRms由 CANN 软件栈提供。测试脚本 test_npu_hc_pre_inv_rms.py 顶部展示了运行所需的依赖import torch import torch_npu import torchair import custom_ops import numpy as np import torch.nn as nn其中custom_ops为 torch_ops_extension/custom_ops 编译生成的算子扩展包。七、调用示例与精度验证官方调用示例位于 examples/test_npu_hc_pre_inv_rms.py覆盖了eager 模式与graph 模式两种调用路径并内置了 CPU 参考实现对比的精度校验。7.1 Eager 模式调用eager 模式直接通过torch.ops.custom.npu_hc_pre_inv_rms调用测试覆盖了 4 维输入(b, s, hc, d)与 3 维输入(T, hc, d)、d 为 4096 与 7168、以及 float32 / bfloat16 / float16 三种输入 dtype 的组合b, s, hc, d 4, 4, 4, 4096 eps 1e-6 # 构造输入并搬运到 NPU x torch.tensor(np.random.uniform(-1, 1, (b, s, hc, d))).to(torch.float32) x x.to(npu:0) # 调用自定义算子 npu_y torch.ops.custom.npu_hc_pre_inv_rms(x, epsiloneps)7.2 Graph 模式调用aclgraph 入图graph 模式通过torch.compile搭配torchair的 NPU 后端实现整图编译验证算子支持 aclgraph 入图的能力from torchair.configs.compiler_config import CompilerConfig class Network(nn.Module): def forward(self, x, epsilon): return torch.ops.custom.npu_hc_pre_inv_rms(x, epsilonepsilon) config CompilerConfig() config.mode reduce-overhead npu_backend torchair.get_npu_backend(compiler_configconfig) npu_mode torch.compile(Network().to(npu:0), fullgraphTrue, backendnpu_backend, dynamicFalse) npu_y npu_mode(x_npu, eps)7.3 精度验证方式测试脚本通过data_compare与cal_relative_diff_np实现 NPU 输出与 CPU 参考输出的相对误差比对默认diff_thd0.0001、pct_thd0.0005、max_diff_hd0.0001满足通过百分比不低于阈值且最大相对误差不超限即判定为Pass。每个用例最终通过assert(compare_y[0] Pass)断言可作为复现算子正确性的最小验证集。八、注意事项与常见问题epsilon 参数校验epsilon 必须为非负数C 层会强制TORCH_CHECK文档默认值为1e-20测试用例中常显式传入1e-6。输入 dtype 差异文档声明输入数据类型支持bfloat16但 C 实现校验允许float32、float16、bfloat16三种 dtype测试用例也覆盖了这三种类型实际可用范围以源码校验逻辑为准。连续性问题输入不支持非连续张量调用前应保证张量内存连续。版本匹配与 PyTorch 配合使用时须确保 CANN 相关包与 PyTorch 相关包的版本匹配否则可能出现算子注册或 AClNN 下发失败。与 HcPre 融合算子的关系当 T/bs 较小≤128 且能被 16 整除时框架会倾向使用hc_pre融合算子以获得更优性能npu_hc_pre_inv_rms主要用于拆分路径是理解 HcPre 内部 InvRms 步骤的独立算子详见 custom-npu_hc_pre.md。九、相关资源算子规格文档custom-npu_hc_pre_inv_rms.md调用示例与测试用例test_npu_hc_pre_inv_rms.pyC 封装实现npu_hc_pre_inv_rms.cpp算子 schema 注册ops_def_registration.cppGraph 模式转换器npu_hc_pre_inv_rms.py上游 HcPre 融合算子文档custom-npu_hc_pre.mdHcPre 完整参考实现test_npu_hc_pre.pyAscendC 自定义算子工程总览ops/ascendc/README.md【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考