CANN ops-transformer：MoE 路由算子的负载均衡策略

布丁爱吃鱼ya

441人浏览 · 2026-05-22 12:20:02

布丁爱吃鱼ya · 2026-05-22 12:20:02 发布

在这里插入图片描述

MoE 模型的推理瓶颈不在计算，而在路由。每个 token 要选 2-4 个专家，路由决策本身的开销在 Mixtral-8x7B 上能占到 Prefill 阶段的 18%。更麻烦的是负载不均衡——80% 的 token 挤到同一个专家，其他专家闲着。ops-transformer 仓库的 MoE 路由算子针对昇腾 NPU 的 Cube/Vector 双核架构做了专门的流水优化，CANN 8.0 之后这个算子成了 MoE 推理的事实标准。

负载不均衡的根源

先说清楚负载不均衡是怎么来的。

MoE 的路由机制是 Top-K 选择：给每个 token 选打分最高的 K 个专家。打分用什么？用一个轻量路由器——把 token 的表示 x 过一个线性层 W_r 得到 logits，再取 top-K。

问题在于 softmax 的输出天然是偏的。如果 W_r 学出来的分布不够均匀，大多数 token 都会指向同一批专家。极端情况下，90% 的 token 都选了专家 0，其他 7 个专家只处理 10% 的 token——等于 7 个专家的算力白费了。

辅助损失函数（auxiliary loss）是解决这个问题的主流方法。在训练时加一个负载均衡损失，惩罚那些被过度使用的专家，鼓励均匀分配。但辅助损失和主损失之间存在 trade-off——辅助损失权重设太高，主任务的性能会下降。

还有一个问题是 token 丢弃（token dropping）。有些实现会主动丢弃负载过高的 expert 处理的 token，但这会引入训练-推理不一致，导致精度下降。

Top-K 路由的并行化

推理的时候，Top-K 路由的开销不容忽视。

标准的 Top-K 实现要做两件事：计算每个专家的 logits（一次矩阵乘法），然后做 Top-K 选择。矩阵乘法用 Cube 核，Top-K 选择本身是排序操作，在昇腾 NPU 上要用 Vector 核。

Top-K 的难点在于 K 的不确定性。Top-1 简单，只需要找最大值；Top-2 也不难；但 Top-8 就麻烦了，因为要找到第 8 大的值，排序的复杂度从 O(N) 变成 O(N log K)。

ops-transformer 的 MoE 路由算子采用了一个巧妙的策略：把 Top-K 选择分成两阶段。第一阶段用 Vector 核做局部 Top-K（每个 AI Core 独立选本地 token 的 top 专家），第二阶段用 AllReduce 收集全局信息，再在每张卡上独立完成最终选择。

// MoE Top-K 路由的 Ascend C 实现（两阶段示意）

__aicore__ void MoERouterTopK(
    GM_ADDR x,        // 输入: token embeddings (num_tokens, hidden_dim)
    GM_ADDR topk_idx, // 输出: top-k expert indices (num_tokens, top_k)
    GM_ADDR topk_gate,// 输出: top-k gate values (num_tokens, top_k)
    int num_tokens,
    int num_experts,
    int top_k,
    LocalTensor<float16_t> gateLogitsLocal,
    LocalTensor<int32_t> localTopKIdx,
    LocalTensor<float16_t> localTopKGate
) {
    // 阶段1: Cube 核算 gate logits
    // x @ W_r -> gate_logits
    // 这里用 MatMul 计算全量 logits，shape: (num_tokens, num_experts)
    MatMul(gateLogitsLocal, x, W_r);
    
    // 阶段2: Vector 核做局部 Top-K
    // 每个 AI Core 处理 num_tokens / num_cores 个 token
    // 每个 token 取局部 top-1
    int local_token_count = num_tokens / TILING_DIM;
    for (int i = 0; i < local_token_count; ++i) {
        // 找这个 token 在局部 expert 中的最大值
        float16_t max_val = gateLogitsLocal[i * num_experts];
        int32_t max_idx = 0;
        for (int e = 1; e < num_experts; ++e) {
            float16_t cur = gateLogitsLocal[i * num_experts + e];
            if (cur > max_val) {
                max_val = cur;
                max_idx = e;
            }
        }
        localTopKIdx[i] = max_idx;
        localTopKGate[i] = max_val;
    }
    
    // 阶段3: AllReduce 收集各卡结果
    // 每张卡得到所有卡上局部 top-1 的综合
    // 从中选出真正的全局 top-K
    // (这一步依赖 HCCL AllReduce，具体实现在 hccl_op.h)
}

Expert Parallel 下的 All-to-All 通信

MoE 的 Expert Parallel（EP）把专家分布到多张卡上。Mixtral-8x7B 有 8 个专家，如果每张卡放 1 个专家，EP=8。

在 EP 推理模式下，当一个 token 被路由到专家 0，而该专家恰好位于卡 1 上时，该 token 就需要先跨卡发送到卡 1，待计算完成后再将结果发送回来。这意味着每个 token 都需要经历一次 All-to-All 通信。

All-to-All 通信的开销会随着 EP 规模的扩大而线性增长。以 top_k=2 的路由为例，理论上每个 token 需要跨卡两次。不过，实际开销可能更高，也可能更低——如果两个专家位于同一张卡上，通信可以合并，从而减少开销。

为了降低 All-to-All 通信的开销，MC2（Merge-Communicate-Split）融合技术应运而生。其核心思路是：不再等待单个 token 的所有专家都计算完毕后再进行数据收集，而是在等待第一个专家结果的同时，就将第二个专家的计算请求发送出去，从而实现通信与计算的流水线化，提升整体效率。

MC2 融合的详细工作流程如下：

Merge（合并）：当一个 token 被路由到多个专家时，系统会先将这些专家的计算请求合并成一个批次。这减少了通信启动的次数，提高了带宽利用率。
Communicate（通信）：合并后的请求被一次性发送到目标计算卡。在等待第一个专家计算完成的同时，系统可以并行地将第二个专家的请求发送出去，或者接收其他 token 的计算结果，从而隐藏通信延迟。
Split（拆分）：目标卡完成计算后，将结果返回。系统再将合并的结果拆分，并路由回对应的 token 流中。

这种流水线化的设计，使得通信和计算可以部分重叠，有效降低了端到端的延迟。尤其是在专家分布在不同计算卡（EP>1）且 top_k>1 的场景下，收益更为显著。

MC2 融合带来的性能提升：

降低延迟：通过通信与计算的重叠，减少了 token 在专家间传输的等待时间。
提升吞吐：合并请求减少了通信次数，更充分地利用了网络带宽。
提高硬件利用率：计算卡在等待网络数据时，可以处理其他计算任务，减少了空闲时间。

ops-transformer 仓库中的 moe_mc2_fusion.cpp 文件实现了 MC2 融合算子。该实现针对昇腾 NPU 的架构进行了深度优化，例如利用 HCCL（Huawei Collective Communication Library）进行高效的多卡通信，并设计了专门的数据结构来管理合并与拆分操作。自 CANN 8.0 起，该功能已默认开启，成为昇腾平台上运行 MoE 模型的关键优化之一。