CANN HcPre 融合算子解析mHC 结构前处理 custom.npu_hc_pre 与 npu_hc_pre_v2 实战指南【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer导读HcPre 是 CANN 推理优化样例中面向 mHCManifold-Constrained Hyper-Connections架构的前处理融合算子它在单个算子内完成 InvRms 归一化、混合系数投影MatMul与 Sinkhorn 迭代三类计算并通过npu_hc_pre/npu_hc_pre_v2两个自定义算子接口暴露给 PyTorch 侧。本文以 custom-npu_hc_pre.md 为主体结合仓库内 test_npu_hc_pre.py 等示例源码与 mhc.py 接入实现完整说明该算子的函数原型、参数与返回值语义、shape 约束、融合路由策略以及精度验证方法帮助你在 DeepSeek-V4 类模型的 mHC 前处理链路上正确选用并验证该算子。mHC 架构与 HcPre 的定位mHCManifold-Constrained Hyper-Connections是传统残差连接的扩展将hidden_state从单一路径扩展为多路在 Attention/MoE 计算前通过Pre Mapping将多路融合回一路保持 Attention/MoE 计算过程不变对 Attention 和 MoE 的输出再通过Post Mapping扩展回多路并与多路残差经 Res Mapping 特征融合后的结果相加得到 mHC 输出参见 deepseek_v4_inference_guide.md。在 mhc.py 中mHC 模块的融合算子由hc_pre与hc_post两部分组成hc_pre承担前处理即本文主角hc_post承担后处理其计算式为out post.unsqueeze(-1) * x.unsqueeze(-2) torch.sum(comb.unsqueeze(-1) * residual.unsqueeze(-2), dim2)详见 custom-npu_hc_post.md。本仓库针对 mHC 架构提供了基于 AscendC 的 HCPre / HCPost 融合算子并同时开源了 AscendC、PyPTO 两种hc_pre实现版本。产品支持情况产品是否支持Atlas A3 推理系列产品√Ascend 950PR / Ascend 950DT√该算子面向推理场景设计并支持 aclgraph 入图。与 PyTorch 配合使用时需要保证 CANN 相关包与 PyTorch 相关包的版本匹配。功能说明与计算流程HcPre 是 mHC 结构前处理部分的融合算子。从 test_npu_hc_pre.py 中的 golden 实现_hc_pre可以看到整个前处理可以拆解为以下三个步骤InvRms 归一化对输入x按 hc 轴计算 RMS 的倒数即rsqrt torch.rsqrt(x.square().mean(-1, keepdimTrue) norm_eps)。其独立算子形态为custom.npu_hc_pre_inv_rms公式为InvRms(x) 1 / RMS(x)其中RMS(x) sqrt(1/N * Σ x_i² epsilon)详见 custom-npu_hc_pre_inv_rms.md。混合系数投影MatMul将归一化后的输入与投影权重hc_fn做线性变换得到mixesmixes F.linear(x, hc_fn) * rsqrt。其中hc_fn的 shape 为[hc_mix, hc_mult * d]。Sinkhorn 迭代将mixes按最后一维切分为pre前hc_mult个、post中间hc_mult个、comb末尾hc_mult * hc_mult个三部分分别做带hc_scale/hc_base的仿射变换其中pre经 sigmoid 后加hc_epspost为2 * sigmoidcomb重塑为[hc_mult, hc_mult]后做 softmax 并进行 20 轮行列归一化的 Sinkhorn 迭代。最后y Σ_hc pre_hc * x[hc]即用pre作为加权系数对多路输入在 hc 轴做加权求和。独立的小算子形态为custom.npu_hc_pre_sinkhorn其输入除mixes外还接收rsqrt见 custom-npu_hc_pre_sinkhorn.md。该算子通过 CubeMatMul单元与 Vector 单元协同完成上述计算并在 MatMul 阶段使用 HF32 模式SetHF32Mode(1) SetHF32TransMode(1)详见下文精度验证章节。函数原型custom.npu_hc_pre(Tensor x, Tensor hc_fn, Tensor hc_scale, Tensor hc_base, *, int hc_mult4, int hc_sinkhorn_iters20, float norm_eps1e-6, float hc_eps1e-6) - (Tensor, Tensor, Tensor) custom.npu_hc_pre_v2(Tensor x, Tensor hc_fn, Tensor hc_scale, Tensor hc_base, Tensor pre_mixNone, *, int hc_mult4, int hc_sinkhorn_iters20, float norm_eps1e-6, float hc_eps1e-6) - (Tensor, Tensor, Tensor, Tensor)两个接口的区别在于npu_hc_pre_v2增加可选输入pre_mix且总是额外返回第 4 个输出hc_pre传入pre_mix时输出y直接使用pre_mix作为加权系数典型用法填入上一轮hc_pre的输出不传入时y使用本轮内部计算的hc_pre进行加权求和行为与npu_hc_pre一致。参数说明说明bbatch size表示输入样本批量大小、ssequence length表示输入样本序列长度、hchead count表示注意力头数、dhead dimension表示注意力头的维度数、T 表示 bs 合轴后的大小。必选参数*之前位置相关参数类型说明数据类型shapexTensormHC 结构的输入数据不支持非连续bfloat16[T, hc_mult, d]或[b, s, hc_mult, d]hc_fnTensor混合系数投影权重不支持非连续float32[hc_mix, hc_mult * d]hc_scaleTensor仿射变换缩放系数不支持非连续float[3]hc_baseTensor仿射变换偏置不支持非连续float[hc_mix]npu_hc_pre_v2额外支持的可选输入参数类型说明数据类型shapepre_mixTensor, 可选仅npu_hc_pre_v2加权求和权重。传入时y 的加权求和直接使用 pre_mix 作为权重典型用法填入上一轮 hc_pre 的输出不传入时y 使用本轮内部计算的 hc_pre 进行加权求和float32[T, hc_mult]或[b, s, hc_mult]可选参数*之后键值对赋值与位置无关*之前的参数是位置相关的必须按照顺序输入属于必选参数*之后的参数是键值对赋值与位置无关属于可选参数不传入会使用默认值。参数类型默认值说明hc_multint4固定为 4hc_sinkhorn_itersint20取值固定为 20norm_epsfloat1e-06InvRms 计算过程中的 ε 参数Host 侧参数仅支持 double 类型hc_epsfloat1e-06Sinkhorn 计算过程中的 ε 参数Host 侧参数仅支持 double 类型返回值说明输出类型说明数据类型shapeyTensor加权求和后的主输出bfloat16[T, d]或[b, s, d]postTensor后处理系数2 * sigmoid结果float[T, hc_mult]或[b, s, hc_mult]comb_fragTensorSinkhorn 迭代后的组合矩阵float[T, hc_mult, hc_mult]或[b, s, hc_mult, hc_mult]hc_preTensor仅npu_hc_pre_v2返回本轮内部计算的 pre 结果sigmoid 后的加权系数float32[T, hc_mult]或[b, s, hc_mult]可作为下一轮 hc_pre 的 pre_mix 输入其中post与comb_frag会被下游hc_post算子继续消费hc_post按公式out post * x Σ(comb * residual)完成多路残差的扩展与融合参见 custom-npu_hc_post.md。约束说明shape 字段取值范围约束字段名取值规则与说明hc_mult取值固定为 4d取值固定为 4096hc_mix取值固定为 24注hc_mix 24 与(hc_mult 2) * hc_mult 6 * 4 24一致mixes按[hc_mult, hc_mult, hc_mult * hc_mult]切分为 pre/post/comb 三部分见 mhc.py。从测试代码看d 亦覆盖了 5120、7168 的扩展取值。使用限制该接口支持推理场景下使用。该接口支持 aclgraph 入图。该接口与 PyTorch 配合使用时需要保证 CANN 相关包与 PyTorch 相关包的版本匹配。输入 tensor 不支持非连续数据格式支持 ND。融合路由策略性能相关当T/bs小于等于 128 且能被 16 整除时会使能hc_pre 融合算子性能较高其他场景会使能hc_pre_inv_rms 与 hc_pre_sinkhorn 小算子拼接性能较低。需要说明的是从 test_npu_hc_pre.py 中的 Ascend950 专项回归用例来看实际路由还会覆盖更细的边界例如 bs ≤ 512 或 bs % 8192 0 时走 fused融合算子路径否则走 composite小算子拼接路径测试注释与route fused if (bs 512 or bs % 8192 0) else composite逻辑见 test_npu_hc_pre.py。具体以所使用 CANN 版本的算子路由实现为准建议在实际部署时结合当前文档约束与芯片型号验证。调用示例完整可运行示例见 test_npu_hc_pre.py。其核心调用方式如下。Eager 模式调用import torch import torch_npu import custom_ops # 注册自定义算子绑定 DEVICE_ID 0 torch_npu.npu.set_device(DEVICE_ID) # 输入构造示例 shape b, s, hc_mult, d, hc_mix 1, 192, 4, 4096, 24 x torch.randn(b, s, hc_mult, d, dtypetorch.bfloat16) hc_fn torch.randn(hc_mix, hc_mult * d, dtypetorch.float32) hc_scale torch.rand(3, dtypetorch.float32) hc_base torch.rand(hc_mix, dtypetorch.float32) x_npu x.to(npu:%d % DEVICE_ID) hc_fn_npu hc_fn.to(npu:%d % DEVICE_ID) hc_scale_npu hc_scale.to(npu:%d % DEVICE_ID) hc_base_npu hc_base.to(npu:%d % DEVICE_ID) # v1返回 (y, post, comb_frag) y, post, comb_frag torch.ops.custom.npu_hc_pre( x_npu, hc_fn_npu, hc_scale_npu, hc_base_npu, hc_multhc_mult, hc_sinkhorn_iters20, norm_eps1e-6, hc_eps1e-6) # v2传入 pre_mix 时返回 (y, post, comb_frag, hc_pre)pre_mix 可为 None pre_mix torch.rand(b, s, hc_mult, dtypetorch.float32).to(npu:%d % DEVICE_ID) y, post, comb_frag, hc_pre torch.ops.custom.npu_hc_pre_v2( x_npu, hc_fn_npu, hc_scale_npu, hc_base_npu, pre_mix, hc_multhc_mult, hc_sinkhorn_iters20, norm_eps1e-6, hc_eps1e-6)Graph 模式调用torch.compile示例同时覆盖了torch.compile成图路径将算子封装进nn.Module的forward并通过torchair.get_npu_backendCompilerConfig(modereduce-overhead)编译执行import torch.nn as nn from torchair.configs.compiler_config import CompilerConfig class Network(nn.Module): def forward(self, x, hc_fn, hc_scale, hc_base, hc_mult, hc_sinkhorn_iters, norm_eps, hc_eps): y, post, comb_frag torch.ops.custom.npu_hc_pre( x, hc_fn, hc_scale, hc_base, hc_multhc_mult, hc_sinkhorn_itershc_sinkhorn_iters, norm_epsnorm_eps, hc_epshc_eps) return y, post, comb_frag config CompilerConfig() config.mode reduce-overhead npu_backend torchair.get_npu_backend(compiler_configconfig) npu_mode torch.compile(Network().to(npu:%d % DEVICE_ID), fullgraphTrue, backendnpu_backend, dynamicFalse) y, post, comb_frag npu_mode(x_npu, hc_fn_npu, hc_scale_npu, hc_base_npu, hc_mult, 20, 1e-6, 1e-6)除默认后端外test_npu_hc_pre.py 还演示了 aclgraph 路径torch.compile(..., backendnpugraph_ex)的成图验证方式与文档中该接口支持 aclgraph 入图的约束相互印证。精度验证golden 实现与比对阈值test_npu_hc_pre.py 提供了完整的精度验证框架可作为算子正确性验收的参考Golden 算法_hc_pre/_hc_pre_v2在 CPU 侧以 fp32 复现 InvRms → MatMul → Sinkhorn → 加权求和的完整流程。HF32 模拟函数to_hf32模拟算子 Cube MatMul 的 HF32 模式——在乘法前将 L0A/L0B 中的 fp32 操作数按 round-toward-zero 截断为 1 符号位 8 指数位 10 尾数位即清空 fp32 位模式低 13 位尾数累加仍保持 fp32。该处理使比对反映真实算法误差而非 HF32 与 fp32 的系统差距见 test_npu_hc_pre.py。比对阈值y为 bfloat16 输出使用diff_thd0.004≈ bf16 ULP、pct_thd0.02post、comb_frag、pre为 float32 输出使用diff_thd0.0001、pct_thd0.005。覆盖矩阵测试覆盖 eager / graph / aclgraph 三种执行路径d 取 4096 / 5120 / 7168Ascend950 上对 fused / composite 两条路由做泛化 bs 扫描含 1、2、7、511、512、513、768、3079、8193、12288、16384、32768 等边界值v2 还覆盖 3D/4D 两种输入形态以及pre_mix传与不传两种组合。测试中还特别提到输入构造的数值范围考量hc_fn按1/(hc_mult*d)缩放以保持mixes量级为 O(1)避免 comb softmax logits 过大导致 Sinkhorn 迭代数值奇异x直接以 bf16 生成以避免 numpy float64 中间态引发大 bs 场景下的内存峰值见 test_npu_hc_pre.py。这些细节对自行构造精度用例同样有参考价值。在 mHC 模块中的接入方式在 mhc.py 中hc_pre通过register_op_impl注册了多个实现版本hc_pre_ascendc_a3func_keyhc_pre_ascendc_a3A3 芯片使用完整的融合自定义算子直接调用torch.ops.custom.npu_hc_pre见 mhc.pyhc_pre_ascendcfunc_keyhc_pre_ascendc通过torch.ops.cann_ops_transformer.mhc_pre_sinkhorn走 transformer 库算子路径见 mhc.pyhc_pre_pypto_a3func_keyhc_pre_pypto_a3PyPTO 实现调用ops.pypto_python.impl.hc_pre_pypto见 mhc.pyhc_pre_native默认纯 PyTorch 原生实现作为对照 golden见 mhc.py。可见custom.npu_hc_pre是 A3 路径上的首选融合实现其余版本作为备选与精度对照。接入时需注意A3 的 fused 实现要求hc_mult4、d4096hc_mix24等 shape 约束并确保 CANN 与 PyTorch 版本匹配。小结custom.npu_hc_pre/custom.npu_hc_pre_v2是 CANN 侧面向 mHC 前处理设计的融合算子将 InvRms、混合系数投影与 Sinkhorn 迭代合并为一次算子下发配合hc_post形成完整的前后处理闭环。使用时需重点关注四点shape 约束hc_mult4、d4096、hc_mix24、T/bs ≤ 128 且可被 16 整除的融合路由条件、npu_hc_pre_v2的pre_mix轮次复用语义以及 eager / graph / aclgraph 三种执行路径下的精度验收方法。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN HcPre 融合算子解析：mHC 结构前处理 custom.npu_hc_pre 与 npu_hc_pre_v2 实战指南
CANN HcPre 融合算子解析mHC 结构前处理 custom.npu_hc_pre 与 npu_hc_pre_v2 实战指南【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer导读HcPre 是 CANN 推理优化样例中面向 mHCManifold-Constrained Hyper-Connections架构的前处理融合算子它在单个算子内完成 InvRms 归一化、混合系数投影MatMul与 Sinkhorn 迭代三类计算并通过npu_hc_pre/npu_hc_pre_v2两个自定义算子接口暴露给 PyTorch 侧。本文以 custom-npu_hc_pre.md 为主体结合仓库内 test_npu_hc_pre.py 等示例源码与 mhc.py 接入实现完整说明该算子的函数原型、参数与返回值语义、shape 约束、融合路由策略以及精度验证方法帮助你在 DeepSeek-V4 类模型的 mHC 前处理链路上正确选用并验证该算子。mHC 架构与 HcPre 的定位mHCManifold-Constrained Hyper-Connections是传统残差连接的扩展将hidden_state从单一路径扩展为多路在 Attention/MoE 计算前通过Pre Mapping将多路融合回一路保持 Attention/MoE 计算过程不变对 Attention 和 MoE 的输出再通过Post Mapping扩展回多路并与多路残差经 Res Mapping 特征融合后的结果相加得到 mHC 输出参见 deepseek_v4_inference_guide.md。在 mhc.py 中mHC 模块的融合算子由hc_pre与hc_post两部分组成hc_pre承担前处理即本文主角hc_post承担后处理其计算式为out post.unsqueeze(-1) * x.unsqueeze(-2) torch.sum(comb.unsqueeze(-1) * residual.unsqueeze(-2), dim2)详见 custom-npu_hc_post.md。本仓库针对 mHC 架构提供了基于 AscendC 的 HCPre / HCPost 融合算子并同时开源了 AscendC、PyPTO 两种hc_pre实现版本。产品支持情况产品是否支持Atlas A3 推理系列产品√Ascend 950PR / Ascend 950DT√该算子面向推理场景设计并支持 aclgraph 入图。与 PyTorch 配合使用时需要保证 CANN 相关包与 PyTorch 相关包的版本匹配。功能说明与计算流程HcPre 是 mHC 结构前处理部分的融合算子。从 test_npu_hc_pre.py 中的 golden 实现_hc_pre可以看到整个前处理可以拆解为以下三个步骤InvRms 归一化对输入x按 hc 轴计算 RMS 的倒数即rsqrt torch.rsqrt(x.square().mean(-1, keepdimTrue) norm_eps)。其独立算子形态为custom.npu_hc_pre_inv_rms公式为InvRms(x) 1 / RMS(x)其中RMS(x) sqrt(1/N * Σ x_i² epsilon)详见 custom-npu_hc_pre_inv_rms.md。混合系数投影MatMul将归一化后的输入与投影权重hc_fn做线性变换得到mixesmixes F.linear(x, hc_fn) * rsqrt。其中hc_fn的 shape 为[hc_mix, hc_mult * d]。Sinkhorn 迭代将mixes按最后一维切分为pre前hc_mult个、post中间hc_mult个、comb末尾hc_mult * hc_mult个三部分分别做带hc_scale/hc_base的仿射变换其中pre经 sigmoid 后加hc_epspost为2 * sigmoidcomb重塑为[hc_mult, hc_mult]后做 softmax 并进行 20 轮行列归一化的 Sinkhorn 迭代。最后y Σ_hc pre_hc * x[hc]即用pre作为加权系数对多路输入在 hc 轴做加权求和。独立的小算子形态为custom.npu_hc_pre_sinkhorn其输入除mixes外还接收rsqrt见 custom-npu_hc_pre_sinkhorn.md。该算子通过 CubeMatMul单元与 Vector 单元协同完成上述计算并在 MatMul 阶段使用 HF32 模式SetHF32Mode(1) SetHF32TransMode(1)详见下文精度验证章节。函数原型custom.npu_hc_pre(Tensor x, Tensor hc_fn, Tensor hc_scale, Tensor hc_base, *, int hc_mult4, int hc_sinkhorn_iters20, float norm_eps1e-6, float hc_eps1e-6) - (Tensor, Tensor, Tensor) custom.npu_hc_pre_v2(Tensor x, Tensor hc_fn, Tensor hc_scale, Tensor hc_base, Tensor pre_mixNone, *, int hc_mult4, int hc_sinkhorn_iters20, float norm_eps1e-6, float hc_eps1e-6) - (Tensor, Tensor, Tensor, Tensor)两个接口的区别在于npu_hc_pre_v2增加可选输入pre_mix且总是额外返回第 4 个输出hc_pre传入pre_mix时输出y直接使用pre_mix作为加权系数典型用法填入上一轮hc_pre的输出不传入时y使用本轮内部计算的hc_pre进行加权求和行为与npu_hc_pre一致。参数说明说明bbatch size表示输入样本批量大小、ssequence length表示输入样本序列长度、hchead count表示注意力头数、dhead dimension表示注意力头的维度数、T 表示 bs 合轴后的大小。必选参数*之前位置相关参数类型说明数据类型shapexTensormHC 结构的输入数据不支持非连续bfloat16[T, hc_mult, d]或[b, s, hc_mult, d]hc_fnTensor混合系数投影权重不支持非连续float32[hc_mix, hc_mult * d]hc_scaleTensor仿射变换缩放系数不支持非连续float[3]hc_baseTensor仿射变换偏置不支持非连续float[hc_mix]npu_hc_pre_v2额外支持的可选输入参数类型说明数据类型shapepre_mixTensor, 可选仅npu_hc_pre_v2加权求和权重。传入时y 的加权求和直接使用 pre_mix 作为权重典型用法填入上一轮 hc_pre 的输出不传入时y 使用本轮内部计算的 hc_pre 进行加权求和float32[T, hc_mult]或[b, s, hc_mult]可选参数*之后键值对赋值与位置无关*之前的参数是位置相关的必须按照顺序输入属于必选参数*之后的参数是键值对赋值与位置无关属于可选参数不传入会使用默认值。参数类型默认值说明hc_multint4固定为 4hc_sinkhorn_itersint20取值固定为 20norm_epsfloat1e-06InvRms 计算过程中的 ε 参数Host 侧参数仅支持 double 类型hc_epsfloat1e-06Sinkhorn 计算过程中的 ε 参数Host 侧参数仅支持 double 类型返回值说明输出类型说明数据类型shapeyTensor加权求和后的主输出bfloat16[T, d]或[b, s, d]postTensor后处理系数2 * sigmoid结果float[T, hc_mult]或[b, s, hc_mult]comb_fragTensorSinkhorn 迭代后的组合矩阵float[T, hc_mult, hc_mult]或[b, s, hc_mult, hc_mult]hc_preTensor仅npu_hc_pre_v2返回本轮内部计算的 pre 结果sigmoid 后的加权系数float32[T, hc_mult]或[b, s, hc_mult]可作为下一轮 hc_pre 的 pre_mix 输入其中post与comb_frag会被下游hc_post算子继续消费hc_post按公式out post * x Σ(comb * residual)完成多路残差的扩展与融合参见 custom-npu_hc_post.md。约束说明shape 字段取值范围约束字段名取值规则与说明hc_mult取值固定为 4d取值固定为 4096hc_mix取值固定为 24注hc_mix 24 与(hc_mult 2) * hc_mult 6 * 4 24一致mixes按[hc_mult, hc_mult, hc_mult * hc_mult]切分为 pre/post/comb 三部分见 mhc.py。从测试代码看d 亦覆盖了 5120、7168 的扩展取值。使用限制该接口支持推理场景下使用。该接口支持 aclgraph 入图。该接口与 PyTorch 配合使用时需要保证 CANN 相关包与 PyTorch 相关包的版本匹配。输入 tensor 不支持非连续数据格式支持 ND。融合路由策略性能相关当T/bs小于等于 128 且能被 16 整除时会使能hc_pre 融合算子性能较高其他场景会使能hc_pre_inv_rms 与 hc_pre_sinkhorn 小算子拼接性能较低。需要说明的是从 test_npu_hc_pre.py 中的 Ascend950 专项回归用例来看实际路由还会覆盖更细的边界例如 bs ≤ 512 或 bs % 8192 0 时走 fused融合算子路径否则走 composite小算子拼接路径测试注释与route fused if (bs 512 or bs % 8192 0) else composite逻辑见 test_npu_hc_pre.py。具体以所使用 CANN 版本的算子路由实现为准建议在实际部署时结合当前文档约束与芯片型号验证。调用示例完整可运行示例见 test_npu_hc_pre.py。其核心调用方式如下。Eager 模式调用import torch import torch_npu import custom_ops # 注册自定义算子绑定 DEVICE_ID 0 torch_npu.npu.set_device(DEVICE_ID) # 输入构造示例 shape b, s, hc_mult, d, hc_mix 1, 192, 4, 4096, 24 x torch.randn(b, s, hc_mult, d, dtypetorch.bfloat16) hc_fn torch.randn(hc_mix, hc_mult * d, dtypetorch.float32) hc_scale torch.rand(3, dtypetorch.float32) hc_base torch.rand(hc_mix, dtypetorch.float32) x_npu x.to(npu:%d % DEVICE_ID) hc_fn_npu hc_fn.to(npu:%d % DEVICE_ID) hc_scale_npu hc_scale.to(npu:%d % DEVICE_ID) hc_base_npu hc_base.to(npu:%d % DEVICE_ID) # v1返回 (y, post, comb_frag) y, post, comb_frag torch.ops.custom.npu_hc_pre( x_npu, hc_fn_npu, hc_scale_npu, hc_base_npu, hc_multhc_mult, hc_sinkhorn_iters20, norm_eps1e-6, hc_eps1e-6) # v2传入 pre_mix 时返回 (y, post, comb_frag, hc_pre)pre_mix 可为 None pre_mix torch.rand(b, s, hc_mult, dtypetorch.float32).to(npu:%d % DEVICE_ID) y, post, comb_frag, hc_pre torch.ops.custom.npu_hc_pre_v2( x_npu, hc_fn_npu, hc_scale_npu, hc_base_npu, pre_mix, hc_multhc_mult, hc_sinkhorn_iters20, norm_eps1e-6, hc_eps1e-6)Graph 模式调用torch.compile示例同时覆盖了torch.compile成图路径将算子封装进nn.Module的forward并通过torchair.get_npu_backendCompilerConfig(modereduce-overhead)编译执行import torch.nn as nn from torchair.configs.compiler_config import CompilerConfig class Network(nn.Module): def forward(self, x, hc_fn, hc_scale, hc_base, hc_mult, hc_sinkhorn_iters, norm_eps, hc_eps): y, post, comb_frag torch.ops.custom.npu_hc_pre( x, hc_fn, hc_scale, hc_base, hc_multhc_mult, hc_sinkhorn_itershc_sinkhorn_iters, norm_epsnorm_eps, hc_epshc_eps) return y, post, comb_frag config CompilerConfig() config.mode reduce-overhead npu_backend torchair.get_npu_backend(compiler_configconfig) npu_mode torch.compile(Network().to(npu:%d % DEVICE_ID), fullgraphTrue, backendnpu_backend, dynamicFalse) y, post, comb_frag npu_mode(x_npu, hc_fn_npu, hc_scale_npu, hc_base_npu, hc_mult, 20, 1e-6, 1e-6)除默认后端外test_npu_hc_pre.py 还演示了 aclgraph 路径torch.compile(..., backendnpugraph_ex)的成图验证方式与文档中该接口支持 aclgraph 入图的约束相互印证。精度验证golden 实现与比对阈值test_npu_hc_pre.py 提供了完整的精度验证框架可作为算子正确性验收的参考Golden 算法_hc_pre/_hc_pre_v2在 CPU 侧以 fp32 复现 InvRms → MatMul → Sinkhorn → 加权求和的完整流程。HF32 模拟函数to_hf32模拟算子 Cube MatMul 的 HF32 模式——在乘法前将 L0A/L0B 中的 fp32 操作数按 round-toward-zero 截断为 1 符号位 8 指数位 10 尾数位即清空 fp32 位模式低 13 位尾数累加仍保持 fp32。该处理使比对反映真实算法误差而非 HF32 与 fp32 的系统差距见 test_npu_hc_pre.py。比对阈值y为 bfloat16 输出使用diff_thd0.004≈ bf16 ULP、pct_thd0.02post、comb_frag、pre为 float32 输出使用diff_thd0.0001、pct_thd0.005。覆盖矩阵测试覆盖 eager / graph / aclgraph 三种执行路径d 取 4096 / 5120 / 7168Ascend950 上对 fused / composite 两条路由做泛化 bs 扫描含 1、2、7、511、512、513、768、3079、8193、12288、16384、32768 等边界值v2 还覆盖 3D/4D 两种输入形态以及pre_mix传与不传两种组合。测试中还特别提到输入构造的数值范围考量hc_fn按1/(hc_mult*d)缩放以保持mixes量级为 O(1)避免 comb softmax logits 过大导致 Sinkhorn 迭代数值奇异x直接以 bf16 生成以避免 numpy float64 中间态引发大 bs 场景下的内存峰值见 test_npu_hc_pre.py。这些细节对自行构造精度用例同样有参考价值。在 mHC 模块中的接入方式在 mhc.py 中hc_pre通过register_op_impl注册了多个实现版本hc_pre_ascendc_a3func_keyhc_pre_ascendc_a3A3 芯片使用完整的融合自定义算子直接调用torch.ops.custom.npu_hc_pre见 mhc.pyhc_pre_ascendcfunc_keyhc_pre_ascendc通过torch.ops.cann_ops_transformer.mhc_pre_sinkhorn走 transformer 库算子路径见 mhc.pyhc_pre_pypto_a3func_keyhc_pre_pypto_a3PyPTO 实现调用ops.pypto_python.impl.hc_pre_pypto见 mhc.pyhc_pre_native默认纯 PyTorch 原生实现作为对照 golden见 mhc.py。可见custom.npu_hc_pre是 A3 路径上的首选融合实现其余版本作为备选与精度对照。接入时需注意A3 的 fused 实现要求hc_mult4、d4096hc_mix24等 shape 约束并确保 CANN 与 PyTorch 版本匹配。小结custom.npu_hc_pre/custom.npu_hc_pre_v2是 CANN 侧面向 mHC 前处理设计的融合算子将 InvRms、混合系数投影与 Sinkhorn 迭代合并为一次算子下发配合hc_post形成完整的前后处理闭环。使用时需重点关注四点shape 约束hc_mult4、d4096、hc_mix24、T/bs ≤ 128 且可被 16 整除的融合路由条件、npu_hc_pre_v2的pre_mix轮次复用语义以及 eager / graph / aclgraph 三种执行路径下的精度验收方法。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考