DeepSeek-V4 mHC 融合算子深度解析:HcPre 与 HcPost 在 CANN A3 架构上的实现与优化
2026/9/18 8:18:52 网站建设 项目流程

DeepSeek-V4 mHC 融合算子深度解析:HcPre 与 HcPost 在 CANN A3 架构上的实现与优化

【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法,提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer

mHC(manifold-constrained hyper-connections,流形约束超连接)是 DeepSeek-V4 中引入的注意力与 MoE 层间连接结构。为消除其在小算子拼接下的内存墙问题,本仓库发布两个全新的融合算子HcPreHcPost,分别覆盖 mHC 前处理与后处理的全流程计算。本文以 docs/models/deepseek_v4/deepseek_v4_mHC_guide.md 为核心,结合仓库中算子实现、模型调用链与算子文档,系统讲解这两个算子的计算语义、A3 分离架构下的 Tiling 与流水设计、性能表现及工程落地方式,读者可据此理解融合算子的设计方法论并直接复用仓库中的接口与配置。

mHC 结构背景与融合动机

mHC 结构出自 DeepSeek 于 2025 年底公开的论文(arXiv 编号 2512.24880)。在该结构下,每一层 Transformer(包括注意力子层与 MoE 子层)的输入都需要经过"前处理 → 主计算 → 后处理"的完整链路,其中前处理负责将输入混合投影到若干条超连接分支并计算加权系数,后处理则将这些分支按系数融合回残差流。

在小算子(kernel-by-kernel)逐段实现的方案下,mHC 的计算被拆分为大量细粒度算子,中间结果反复经由 HBM 读写,形成了典型的"内存墙"瓶颈。本仓库的解法是将其整体融合为两个大算子:

  • HcPre:将 mHC 前处理全部融合为一个算子,通过极致的性能优化将中间过程输出封闭在核内,脚本侧只感知一个简单接口;
  • HcPost:将 mHC 后处理融合为一个算子,消除小算子方案下的内存墙问题。

目标硬件:A3 分离架构

两个算子均面向 Atlas A3 推理系列产品实现。A3 为分离架构:Cube 核与 Vector 核数量比为 1:2,每 1 个 Cube 核与 2 个 Vector 核构成一个组核(AICore)。组核内部,L1、L0A、L0B、L0C 是 Cube 核的片上 Buffer,UB(Unified Buffer)是 Vector 核的片上 Buffer。

该架构对算子设计的关键影响在于:Cube 核与 Vector 核的数据交换必须借助 HBM(workspace)中转,无法像统一架构那样通过片上直接传递。因此 A3 上的融合算子设计需要显式规划 CV 流水与 HBM 访存节奏,这也是本文后续 Tiling 与流水设计的出发点。

HcPre 融合算子实现

HcPre 覆盖 mHC 前处理的全部计算:从输入 x 出发,经 RMS 归一化、与投影矩阵 hc_fn 的矩阵乘(mixes 计算)、pre/post 系数非线性变换、comb 矩阵的 Sinkhorn 迭代归一化,最终输出主分支 y 以及供后处理使用的 post、comb。

计算语义与算子接口

仓库中算子原型(见 ops/ascendc/docs/custom-npu_hc_pre.md)为:

custom.npu_hc_pre(x, hc_fn, hc_scale, hc_base, *, hc_mult=4, hc_sinkhorn_iters=20, norm_eps=1e-6, hc_eps=1e-6) -> (y, post, comb_frag) custom.npu_hc_pre_v2(x, hc_fn, hc_scale, hc_base, pre_mix=None, *, ...) -> (y, post, comb_frag, hc_pre)

关键参数与约束如下:

参数说明shape / 取值
xmHC 结构输入,bfloat16,ND 格式,要求连续[T, hc_mult, d] 或 [b, s, hc_mult, d]
hc_fn混合投影权重,float32[hc_mix, hc_mult * d]
hc_scale非线性变换缩放系数,float[3]
hc_base非线性变换偏置,float[hc_mix]
pre_mix(v2 可选)外部指定的 pre 权重,float32[T, hc_mult] 或 [b, s, hc_mult]
hc_mult超连接分支数,固定为 4固定值 4
hc_sinkhorn_itersSinkhorn 迭代次数固定为 20
norm_epsInvRms 计算的 epsilon默认 1e-6
hc_epsSinkhorn 计算的 epsilon默认 1e-6

其中维度关系为hc_mix = (2 + hc_mult) * hc_mult = 24d = 4096。返回值中 y 为 bfloat16 的主分支输出(shape [T, d]),post 与 comb_frag 为 float32 的系数(shape 分别为 [T, hc_mult] 与 [T, hc_mult, hc_mult])。npu_hc_pre_v2额外返回本轮内部计算出的 pre 系数,可作为下一轮hc_prepre_mix输入复用。

计算语义在仓库中有一份可读性极高的 PyTorch 参考实现(golden),见 models/deepseek_v4/models/modules/op_impls/mhc.py 与 ops/ascendc/examples/test_npu_hc_pre.py:先对 x 做 RMS 归一化得到 rsqrt,经F.linear(x, hc_fn)得到 mixes;随后将 mixes 沿最后一维切分为[hc_mult, hc_mult, hc_mult*hc_mult]三段,分别经 sigmoid 变换得到 pre、post,以及经 softmax 与多轮行/列归一化(Sinkhorn 迭代)得到 comb;最终主分支y = Σ_hc pre_hc * x[hc]

值得注意的实现细节:测试代码中对 Cube 矩阵乘使用了 HF32 模式模拟(SetHF32Mode(1)+SetHF32TransMode(1),即对 fp32 操作数截断低位 13 bit 尾数后再乘、fp32 累加),并在 golden 中镜像该行为,从而让精度对比反映真实算法误差而非 HF32 与 fp32 之间的固有差距。此外,算子文档明确:当 T/bs ≤ 128 且能被 16 整除时会启用完整的 hc_pre 融合算子(性能较高),其他场景退化为 hc_pre_inv_rms 与 hc_pre_sinkhorn 小算子拼接(性能较低)——这正是融合算子针对典型 Decode shape 的定向优化。

Decode 场景的 A3 实现设计

Decode 场景下 m(batch × seq 合轴)很小,若按 token 分核会导致大部分核空闲,因此采用多核切 K 模板:沿 K 轴(hidden 维度)将计算切分到多个核上以提升并行度与 Cube 利用率。

整体计算流与 Buffer 使用如图(A3 Decode 计算步骤):

流水设计分为两个阶段:

  • 阶段一流水:以"搬运 → 类型转换 → Cube 计算"的形式构成 CV 流水。由于中间内存较小,数据基本命中 L2,有效提升了流水效率,掩盖了搬运与计算之间的相互等待。
  • 阶段二流水:提前搬运 x(为后处理做准备),用搬运过程掩盖 pre、post 与 combine 的 Vector 计算开销,实现"算搬并行"。

A3 Decode 场景性能测试结果(规格 b*s / hc / d,单位 us):

规格大小性能(us)
b*s=16; hc=4; d=409632.8
b*s=32; hc=4; d=409637
b*s=64; hc=4; d=409644
b*s=128; hc=4; d=409660

可以看到在 m 从 16 增长到 128 的过程中,耗时仅从 32.8us 增加到 60us,呈近似亚线性增长,说明多核切 K 模板下 Cube 计算被充分并行化,性能主要受限于 K 轴切分带来的同步与访存开销而非计算量本身。

Prefill 场景的 A3 实现设计

Prefill 场景下 token 数(bs)很大,天然适合按 token 分核:将 bs 切分到多个核上,每个核独立完成自己份额内的 HcPre 计算,形成 CV 流水。

整体计算流与 Buffer 使用如图(A3 Prefill 计算步骤):

Tiling 与流水设计上,Prefill 模板的两个阶段计算当前为串行执行。文档明确指出后续可进一步优化:让两个 Vector 核分别承担阶段一的 Vector 计算与阶段二的计算,以流水方式进一步掩盖开销(A3 组核内恰好是 1 Cube + 2 Vector 的结构,为这一优化留出了空间)。

A3 Prefill 场景性能测试结果(规格 b*s / hc / d,单位 us):

规格大小性能(us)
b*s=4K; hc=4; d=40961254.916
b*s=16K; hc=4; d=40964784.204
b*s=64K; hc=4; d=409618545.39

耗时随 token 数近似线性增长(4K→16K 约 3.8 倍,16K→64K 约 3.9 倍),符合按 token 分核后每核工作量与总 token 数成正比的特征。

HcPost 融合算子实现

HcPost 覆盖 mHC 后处理:将主计算输出 x 与残差 residual 按 pre、post、comb 系数融合回残差流。小算子方案下该过程的瓶颈在于内存墙——中间张量反复经 HBM 读写。融合后,数据不出 UBuffer(核内 buffer),大幅减少 HBM 访存,从而显著提升性能。

计算语义与算子接口

HcPost 的 PyTorch 参考实现(见 models/deepseek_v4/models/modules/op_impls/mhc.py)为:

y = post.unsqueeze(-1) * x.unsqueeze(-2) + torch.sum(comb.unsqueeze(-1) * residual.unsqueeze(-2), dim=2)

即主分支加权与残差分支按 comb 融合后相加,对应昇腾侧算子torch.ops.cann_ops_transformer.mhc_post(residual, comb, x, post)。在 ops/ascendc/docs/custom-npu_hc_pre.md 之外,mhc 相关算子与 HcPre 的 inv_rms / sinkhorn 拆分版本同样在 ops/ascendc/README.md 中登记(hc_pre、hc_pre_inv_rms、hc_pre_sinkhorn),并可通过bash build.sh -n "...;hc_pre"这类方式参与算子编译。

A3 实现设计

在 A3 上,Vector 的中间计算结果位于 UBuffer,因此 HcPost 融合的关键收益是让整个后处理链路的数据流停留在 UBuffer 内,避免 HBM 往返。

整体计算流与 Buffer 设计(A3 HcPost 计算步骤):

Tiling 与流水设计:在 b、s 以及 d 轴上进行分核,b、s 合并为一个轴 bs,按 (m, n) 的 tile 块划分。极小 case 下多核切分需权衡启动开销与处理开销,因此设置了两个约束:每个核最低搬运数据量诉求为 4KB,n 方向数据量不小于 1KB 且 512B 对齐。

同时,依据 buffer 占用与流水掩盖的权衡,决定是否在 hc 轴上做循环搬运流水,分为两类模板:

  1. hc 较小的 case:可以一次性将 residual 的 tile 块 (m, hc, n) 全部搬入 UBuffer,以最大程度的流水掩盖换取简单直接的调度:

  2. hc 泛化 case:当 hc 较大导致 UBuffer 可能不够用时,在 residual 的 hc 轴上开循环做流水,每次只搬运 t 个 hc 分量:

    循环场景下的计算步骤示意图如下:

性能测试结果(规格 b*s / hc / d,单位 us):

规格大小性能(us)
b*s=8; hc=4; d=40965.09
b*s=16; hc=4; d=40965.37
b*s=32; hc=4; d=40966.99
b*s=64; hc=4; d=40969.16
b*s=128; hc=4; d=409613.85
b*s=4K; hc=4; d=4096319.86
b*s=16K; hc=4; d=40961264.68
b*s=64K; hc=4; d=40965044.53

与 HcPre 对比可见:HcPost 在同等规格下耗时约为 HcPre 的四分之一到三分之一(如 b*s=4K 时为 319.86us vs 1254.916us),这正是"数据不出 UBuffer、消除内存墙"带来的直接收益——后处理本身计算量小,融合后瓶颈从访存转移到了更可控的计算与调度上。

模型侧落地:调用链与多后端实现

在 DeepSeek-V4 推理脚本中,mHC 融合算子已接入完整的 Transformer 前向流程。以 models/deepseek_v4/models/modeling_deepseek.py 为例,每个 Transformer 层内:

  1. 保存残差后调用OpKernel.hc_pre(hidden_states, hc_attn_fn, hc_attn_scale, hc_attn_base, hc_mult, hc_sinkhorn_iters, norm_eps, hc_eps)得到主分支与 post、comb 系数;
  2. 主分支经过注意力子层计算;
  3. 调用OpKernel.hc_post(hidden_states, residual, post, comb)将结果融合回残差流;
  4. MoE 子层重复同样的"hc_pre → ffn → hc_post"模式(分别使用hc_ffn_fnhc_ffn_scalehc_ffn_base参数)。

对应的 mHC 参数在模型初始化时按mix_hc = (2 + hc_mult) * hc_mult分配hc_attn_fn/hc_ffn_fn(shape [mix_hc, hc_mult * hidden_size])、hc_attn_base/hc_ffn_base(shape [mix_hc])以及hc_attn_scale/hc_ffn_scale(shape [3]),并支持通过配置项hc_multhc_sinkhorn_itershc_eps控制算子行为。此外,DSpark 分支(models/deepseek_v4/models/modeling_dspark.py)在保留 proposal 轴的前提下复用同一套 hc_pre/hc_post 接口,并在 head 预测侧使用独立的hc_head_fn

算子后端采用注册表机制管理多种实现(见 models/deepseek_v4/models/modules/op_impls/mhc.py):

  • hc_pre_ascendc_a3:A3 上完整的自定义融合实现,通过torch.ops.custom.npu_hc_pre暴露,即本文所述的"单接口大算子";
  • hc_pre_ascendc:调用cann_ops_transformer.mhc_pre_sinkhorn的组合实现;
  • hc_pre_pypto_a3:基于 ops/pypto_python/impl/hc_pre_pypto.py 的 PyPTO 版本;
  • hc_pre_native/hc_post_native:PyTorch 参考实现,用于精度对照;
  • hc_post_ascendc:调用cann_ops_transformer.mhc_post的融合实现。

这种"同一语义、多后端注册"的设计,既保证了模型脚本侧的调用接口统一(脚本侧对 HcPre 的极致优化中间过程不感知),又为不同硬件形态与工具链(AscendC / PyPTO)保留了独立演进空间。

小结

HcPre 与 HcPost 是本仓库针对 DeepSeek-V4 mHC 结构给出的融合算子方案,其核心方法论可以归纳为三点:融合消除内存墙(HcPost 数据不出 UBuffer,HcPre 中间过程封闭在核内)、CV 流水并行(在 A3 分离架构下通过"搬运-转换-Cube 计算"与"提前搬运掩盖 Vector 开销"两个阶段的流水设计)、以及极致的 Tiling 与流水(Decode 多核切 K、Prefill 按 token 分核、hc 轴循环搬运三档模板)。实测数据表明,该方案在 Decode 与 Prefill 的典型规格下均获得了与计算量匹配的近线性扩展,并在同等规格下显著优于小算子拼接的内存墙瓶颈。对于希望复用该方案的开发者,可直接在 models/deepseek_v4/models/modeling_deepseek.py 中观察调用模式,在 ops/ascendc/docs/custom-npu_hc_pre.md 与 ops/ascendc/examples/test_npu_hc_pre.py 中获取完整的接口语义与精度对照实现。

【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法,提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询