之前有个团队跟我吐槽，他们用昇腾NPU跑 Mixtral 8x7B，性能没比稠密模型好多少。MoE 模型的核心优势就是稀疏激活——每个 token 只激活 2 个 expert，计算量理论上只有全激活的 1/4。但实际跑起来，expert 之间的数据分发和归集（All-to-All 通信）把省下的算力全吃回去了。ops-transformer 仓库里的 MoE 融合算子，就是冲着这个瓶颈来的。

MoE 的性能杀手：不是计算，是通信

稠密模型的前向传播是一条直线：输入 → Linear → Activation → 输出。MoE 不一样，它有个路由分发环节：

1. Gate 算出每个 token 该去哪个 expert
2. token 按 expert 分组，发到对应的计算单元
3. 各 expert 独立计算
4. 结果收集回来，按原始顺序拼好

步骤 2 和步骤 4 就是 All-to-All 通信。在分布式训练里，expert 分布在不同卡上，每个 token 可能要跨卡找自己的 expert。标准实现里，分组、发送、接收、拼接是四个独立操作，每次操作都是一次显存读写 + 一次通信调用。

ops-transformer 的 MoE 融合算子把 Gate + 分组 + 通信 + 计算 + 归集压成一个 kernel pipeline。 不是简单的算子拼接，是把通信和计算重叠起来了——第 N 批 token 在计算的时候，第 N+1 批已经在通信了。

三层融合策略

Gate 融合。 标准 MoE 的 Gate 是一个独立的 Linear + Softmax + TopK，三个算子各跑各的。ops-transformer 把它们融合成一个 kernel，输出直接就是路由索引和权重，中间不落显存。Gate 虽然计算量不大，但在 CANN 的昇腾达芬奇架构上，三次 kernel launch 的调度开销比计算本身还大。

Expert 计算融合。 被选中的 2 个 expert 的计算（Linear→Activation→Linear）合成一个融合 kernel。这里的关键是 expert 权重已经在 NPU 的 Cube 单元里了，不用反复搬运。

通算融合（MC2）。 这是最核心的部分。MC2 把 All-to-All 通信和 expert 计算流水线化。昇腾NPU的 HCCL 通信库和计算引擎是两条独立的硬件通路——通信走 DMA，计算走 Cube+Vector。MC2 让它们同时工作，通信延迟被计算时间完全掩盖。

性能数据

在 Atlas 800I A2 上，Mixtral 8x7B 的训练吞吐对比：

通信占比从 47% 砍到 18%。这不是靠减少通信量做到的——数据该传还得传——是靠把通信藏到计算后面了。

用法

通过 ATB 加速库搭建 MoE 推理服务时，MoE 融合算子默认启用：

import torch_npu
# ATB 内部自动调用 ops-transformer 的 MoE 融合算子
# 路由策略、expert 数量都从模型配置自动读取
# 这里只需要确保 gate 的 topk 值和融合算子对齐
atb_model = AtbModel.from_pretrained("mixtral-8x7b", device="npu:0")

单算子调用也行，但要注意 expert 权重的排布格式。ops-transformer 的 MoE 融合算子要求 expert 权重按 [num_experts, hidden_dim, ff_dim] 连续排列，如果你的权重是按层打散存的，需要先重排。

一个实际踩坑

MoE 融合算子的 TopK 路由只支持 K=1 或 K=2。如果你搞了个 K=3 的稀疏策略，它会 fallback 到非融合实现，性能打回原形。K=2 是目前大模型的主流选择（Mixtral、DeepSeek-MoE 都是 K=2），所以大部分场景没问题，但自定义架构要注意。

另外，MC2 通算融合需要 HCCL 2.0 以上版本。CANN 8.0 默认带的 HCCL 版本够用，升级到 CANN 8.5 后通信带宽利用率又提了一截。

如果你的 MoE 模型在昇腾NPU上跑不出稀疏加速的效果，先查两件事：融合算子是否启用（看日志有没有 fallback），通信和计算是否在重叠（看 NPU 利用率曲线，通信时计算不应该空转）。仓库在这里：

https://atomgit.com/cann/ops-transformer