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

布丁爱吃鱼ya

12人浏览 · 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)
}