当前大模型发展演进中， DSA、GDN、mHC [1-3] 等新型网络架构如雨后春笋不断涌现，如何在短时间内将创新架构高效落地于昇腾训练集群，成为关键挑战。昇腾训练团队通过mHC高性能算子设计、编译预热的整网适配、细粒度重计算技术，在万卡级昇腾Atlas 900 A3 SuperPoD超节点上实现典型shape下单算子性能提升5.5 倍，1500卡环境下吞吐提升近100B tokens/day。本文旨在复盘这一过程中的关键技术决策、工程实践与系统思考，为后续昇腾训练集群基础设施建设提供可复用的经验。

DeepSeek通过引入双随机矩阵的流形约束mHC技术解决HC训练不稳定问题，mHC将HC过程分为了三部分：负责多条流混合为多条流的H^res、多条流缩小为一条流的H^pre、一条流扩展为多条流的H^post。为保证各阶段的数学性质，在具体实现上通过引入一个满足元素非负且行、列和均为1的双随机矩阵（具有范数保持、组合封闭性和凸组合几何解释三大特性）对多流混合成多条流时进行混合加权，在多流到单流的压缩以及单流到多流的扩展则分别引入取值范围为 (0, 1) 及 (0, 2) 的sigmoid函数作为权重系数。 

mHC 示意图

Triton是面向深度学习与高性能计算的领域特定语言（DSL）及编译器系统，旨在兼顾开发效率与接近硬件极限的性能。不同于线程级编程模型，Triton将抽象粒度提升至“分块”（Tile/Block）级别：用户仅需描述张量分块方式与分块上的计算逻辑，内存管理、数据搬运、调度优化等底层细节均由编译器自动处理。凭借“快速收敛到高性能”的特性，Triton已成为构建高效自定义算子的首选工具。

在mHC Triton算子中，我们并未将Triton仅视为一个孤立的计算 kernel 编写工具，而是将其定位为算子编排的“粘合层”，灵活协同PyTorch原生算子、昇腾NPU自定义算子与Triton kernel，以实现整体性能最优。具体而言，

对于NPU已高度优化的基础算子：如MatMul，直接复用PyTorch的高效实现；

对于NPU提供硬件加速支持的特定操作：如RMSNorm，在实测对比Triton实现与昇腾自定义算子的性能后，选择调用后者以获得更高性能；

对于mHC架构中特有的计算模式：包括Sigmoid激活与 Sinkhorn迭代归一化 , 通过Triton编写定制化融合kernel实现。

这种“按需选型、优势互补”的策略，既避免了重复造轮子，又充分释放了硬件潜力，最终在端到端训练中实现了显著的性能提升，并成功支撑万卡规模昇腾超节点上生产训练的实践工作。

mHC 融合计算示意图 

在mHC算子涉及的BMM计算中，其矩阵shape呈现出小N（e.g., 4）、大D（e.g., 5120）的特点。若实现方式不合理极易出现 Memory Bound 明显、Scalar指令占比偏高以及 MTE 传输压力集中的问题。我们选择在Vector单元上实现定制化的BMM kernel，以减少数据搬运成本并提升整体计算效率：一方面，Vector kernel所需的中间数据规模可以控制在UB（192KB）容量范围内；另一方面，也能够避免Cube与Vector之间频繁的GM中转。

在具体实现中，我们进一步利用昇腾AIV架构的48个Vector核（A3 设备）进行并行计算，基于负载均衡的分核并行策略，通过入参形状与硬件核数动态获取Group划分粒度，将batch维度按照 Group的形式拆分到多个Vector核上执行。每个AIV负责处理一组batch数据，并在D维上进行向量化计算，使计算任务能够在 Vector 核之间均衡分布，该方式不仅提升了并行度，还降低了单核寄存器压力与UB占用，从而获得较高计算性能。

基于整网数据的形状特征，我们在两种典型配置下（Batch Size=1, Seq Len=1024, N=4, Dim=5120 与Batch Size=1, Seq Len=2048, N=4, Dim=2560）对比了Triton算子与原生 Torch小算子的执行耗时。实验结果表明，相较于原生实现，Triton算子在上述场景中均实现了约5.5倍的加速比。具体数据如下表所示：

将模型训练部署扩展至上千卡的大规模分布式集群时，我们发现首 loss 的计算耗时明显高于未使用该算子的配置。根本原因在于在大规模多机分布式训练的高并发场景下，大量训练进程几乎同时启动并首次调用Triton算子，从而并发触发其即时编译（JIT）过程，不仅造成训练初始化时间大幅延长，还因瞬时对编译资源（如CPU、内存或文件系统）的激烈竞争，被训练框架误判为异常，进而激活“快恢复”逻辑，严重影响训练的稳定性与集群资源的利用效率。

因此，我们在模型正式训练前引入显式的Triton算子预热（warmup）机制：即在主训练循环启动前，主动调用模型中所有将被使用的Triton算子，并传入与实际训练尺寸兼容的张量进行一次算子调用。该操作可提前完成JIT编译、内核缓存及上下文初始化，确保后续多机并发训练阶段直接加载已编译的内核，避免运行时编译开销。

mHC算法的引入使得每层输入内存扩大了N倍，在计算过程中采用fp32计算，额外增加了2N倍的激活显存需要保留用于反向梯度计算。由于分别在Transformer的attn和mlp/moe模块引入了mHC算法，所以每层激活增加了2*(7N-3)*B*S*D的内存。

mHC的激活内存表现

在常规配置下(B=1, S=4096, D=5120, N=4),  仅激活内存增加就接近1G，并随着层数线性增长，在HBM资源(64GB)本就吃紧的情况下无疑雪上加霜，因此我们设计了两种细粒度重计算节省了85%以上的新增显存占用，以适配不同的应用场景：

mHC层内重计算：通过对hc_pre和hc_post Triton算子进行精细化重计算控制，并结合Transformer框架attn以及mlp/moe前后的Norm重计算进行协同优化，实现单层内存的极致优化；具体的，在层内重计算中，仅保留主干上attn hc_pre的输入，attn的输出以及mlp/moe的输出。

mHC跨层重计算：为了进一步压榨内存，在PP/VPP场景下，我们在一个chunk中，仅保留第一个Transformer层h_pre的输入，其他层的h_pre输入通过前一层mHC的重计算来还原，并与PP异步通信进行通算掩盖。

在Triton算子应用在正式生产之前，我们特别做了精度对比实验，小算子和Triton算子分出在1500卡的规模中基于同一CKPT采用相同配置进行“陪跑”训练。如图所示，在5000步的“陪跑”对比实验中，算子精度符合预期，吞吐提升明显，训练单步耗时由14.73s降低至10.48s，提升29%。

“陪跑”对比实验中单步耗时对比结果

本文介绍了mHC的基本计算流程，具体阐述了基于Triton的负载均衡的分核计算优化、内存与计算均衡优化性能优化方法，实现单算子5.5倍性能提升。此外，昇腾团队还将基于Triton实现的 mHC计算优化应用于万卡规模的昇腾超节点中，通过算子预热的策略实现在大规模昇腾训练集群下Triton算子可用性与系统稳定性。通过将mHC结构拆解并与细粒度重计算框架协同设计，有效缓解了其内存开销。得益于Triton算子的灵活性及其与框架重计算机制的无缝兼容，该方案得以快速落地验证，最终实现了整网性能提升，支撑多个规格的模型在总计万卡规模的昇腾超节点中生产训练2个月。

mHC-Triton训练关键技术：
https://gitcode.com/Ascend/MindSpeed-Ops/pull/35

MindSpeed-Ops仓库提供昇腾优化的训练业务自定义算子实现，欢迎关注和共建：
https://gitcode.com/Ascend/MindSpeed-Ops

参考文献

[1] Liu, Aixin, et al. Deepseek-v3. 2: Pushing the frontier of open large language models. arXiv preprint arXiv:2512.02556 (2025).

[2] Yang, Songlin, et al. Gated delta networks: Improving mamba2 with delta rule. arXiv preprint arXiv:2412.06464 (2024).

[3] Xie, Zhenda, et al. mhc: Manifold-constrained hyper-connections. arXiv preprint arXiv:2512.24880 (2025).