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

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__voidMoERouterTopK(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)intnum_tokens,intnum_experts,inttop_k,LocalTensorfloat16_tgateLogitsLocal,LocalTensorint32_tlocalTopKIdx,LocalTensorfloat16_tlocalTopKGate){// 阶段1: Cube核算gate logits// x W_r - gate_logits// 这里用 MatMul 算全量 logitsshape: (num_tokens, num_experts)MatMul(gateLogitsLocal,x,W_r);// 阶段2: Vector核做局部Top-K// 每个AI Core处理 num_tokens / num_cores 个token// 每个token取局部top-1intlocal_token_countnum_tokens/TILING_DIM;for(inti0;ilocal_token_count;i){// 找这个token在局部expert中的最大值float16_tmax_valgateLogitsLocal[i*num_experts];int32_tmax_idx0;for(inte1;enum_experts;e){float16_tcurgateLogitsLocal[i*num_expertse];if(curmax_val){max_valcur;max_idxe;}}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 ParallelEP把专家分布到多张卡上。Mixtral-8x7B 有 8 个专家如果每张卡放 1 个专家EP8。EP 推理时一个 token 被路由到专家 0但专家 0 在卡 1 上所以 token 要跨卡发送到卡 1算完再发回来。每个 token 都要做一次 All-to-All 通信。All-to-All 通信的开销随 EP 规模线性增长。如果路由的 top_k2理论上每个 token 要跨 2 次卡——实际可能更多因为两个专家可能在同一张卡上可以合并通信。MC2Merge-Communicate-Split融合是降低 All-to-All 开销的核心技术。思路是不等单个 token 的两个专家都算完再收数据而是在等第一个专家结果的时候就把第二个专家的请求发出去——通信和计算流水起来。ops-transformer 仓的moe_mc2_fusion.cpp实现了 MC2 融合算子CANN 8.0 之后默认开启。top_k 参数的选择top_k1、2、4 分别对应 MQA、Top-2、Top-4 MoE。top_k1等价 MQAKV Cache 最省但路由没有多样性模型质量掉得最明显。适合显存紧张、batch size 大的场景。top_k2这是 Mixtral 的默认配置。社区验证下来精度和 top_k1 差距明显但 top_k4 差距不大。是一个性价比配置。top_k4精度损失最小但每个 token 要激活 4 个专家Expert Parallel 的通信量翻倍。如果 EP 规模大8 卡以上通信时间占比会很高可能抵消精度收益。EP 规模也会影响选择。EP2 的时候top_k4 有两张卡可选通信量翻倍但利用率高。EP8 的时候top_k4 有 8 张卡可选通信量变成 4x but 每张卡参与的概率下降。CANN 8.0 之后moe_topk_sweep.py脚本可以帮助快速测试不同 top_k 配置的性能和精度 trade-off。脚本会跑指定的模型输出不同 top_k 下的延迟和 perplexity。ops-transformer 仓的 MoE 相关实现在ops/moe/目录moe_router_topk.cpp是 Top-K 路由的核心moe_mc2_fusion.cpp是 MC2 融合。https://atomgit.com/cann/ops-transformer