摘要Mixtral 8×7B 部署到 910B 4 卡吞吐 180 TPS。客户说 GPU 上能到 1200 TPS差了 6 倍。跑 msprof 一看MoE Router AllReduce 占 42% 时间。MoEMixture of Experts模型的推理瓶颈不在计算在通信和调度。Router 动态路由 Expert 通信开销吃掉 40% 时间。优化后吞吐从 180 提到 720 TPS涨了 4 倍。MoE 推理瓶颈分析MoE 的核心结构输入 x ↓ Router选择 top-k expert ↓ Expert 1, Expert 2, ..., Expert n只激活 top-k 个 ↓ 加权求和 ↓ 输出 yMixtral 8×7B8 个 expert每次激活 top-2。瓶颈拆解阶段耗时占比问题RouterGating TopK18%动态路由开销Expert 计算35%8 个 expert 并行度不足AllReduce 通信24%Expert 间通信带宽利用率低其他23%-Router AllReduce 占 42% 时间这才是 MoE 推理慢的根本原因。工程经验很多人以为 MoE 慢是因为 expert 多、计算量大。其实 MoE 每次只激活 top-k 个 expert计算量比同等参数的 Dense 模型还小。真正的瓶颈是 Router 动态路由和 Expert 间通信。GatingTopK 算子融合Router 的计算流程1. Gatingx W_gate → gate_score每个 token 对 8 个 expert 的得分 2. TopK从 8 个得分中选 top-2 3. Softmax对 top-2 得分做 softmax得到权重不融合的实现Gating 输出 → 写 HBM → TopK 读 HBM → 写 HBM → Softmax 读 HBM3 次 HBM 读写开销 36μs每次 12μs。融合后的实现Gating 输出 → L1 → TopK → L1 → Softmax中间结果走 L1不落 HBM。省掉 2 次 HBM 读写开销降到 12μs。ops-transformer 的 GatingTopK 算子fromops_transformerimportGatingTopK gating_topkGatingTopK(num_experts8,top_k2,fusion_modefull# Gating TopK Softmax 全融合)# 融合计算expert_ids,expert_weightsgating_topk(x,W_gate)性能收益模型融合前 Router 耗时融合后 Router 耗时节省Mixtral 8×7B18μs/层6μs/层12μs/层32 层总计576μs192μs384μs工程经验GatingTopK 融合前Router 输出要写 HBM 再给 TopK 读。融合后直接 L1 传数据省一次 HBM 读写。单层省 12μs32 层省 384μs。decode 每个 token 快 0.38ms1000 个 token 差 0.38 秒。Expert 并行策略MoE 模型有两种并行策略Tensor ParallelTP和 Expert ParallelEP。Tensor ParallelTP把每个 expert 的权重切分到多卡每卡算一部分。Expert 1: 卡1 算 1/2卡2 算 1/2AllReduce 合并 Expert 2: 卡1 算 1/2卡2 算 1/2AllReduce 合并 ...Expert ParallelEP每个卡放完整的 expert不同卡算不同 expert。卡1: Expert 1, Expert 2 卡2: Expert 3, Expert 4 卡3: Expert 5, Expert 6 卡4: Expert 7, Expert 8对比策略通信模式适用场景TPAllReduce每个 expert 都要Expert 数量少、单 expert 大EPAll-to-All激活的 expert 通信Expert 数量多、单 expert 小Mixtral 8×7B8 个 expert每个 expert 7B 参数。EP 更优。实测数据Mixtral 8×7B910B 4 卡并行策略吞吐 (TPS)通信开销TP42028%EP57018%混合TP2, EP261015%混合策略最优TP2每个 expert 切 2 份EP2每卡放 4 个 expert。工程经验MoE 场景下 EP 比 TP 快 35%。原因TP 每个 expert 都要 AllReduceEP 只在激活的 expert 之间通信。Mixtral 每次只激活 2 个 expert通信量比 TP 小。HCCL 通信优化HCCL 是昇腾的集合通信库对标 NCCL。AllReduce 带宽利用率对比库带宽利用率吞吐 (TPS)HCCL 默认配置45%420HCCL 优化后89%720优化点1. 算法选择HCCL 支持多种 AllReduce 算法Ring、Mesh、NHRNon-hierarchical Ring。小数据量 1MBMesh 最快 中数据量1MB-100MBNHR 最快 大数据量 100MBRing 最快MoE 的 AllReduce 数据量通常 10-50MB选 NHR。配置方式importhccl hccl.set_allreduce_algorithm(NHR)2. 通信与计算重叠Expert 计算和 AllReduce 通信并行Expert 1 计算 → Expert 2 计算 → AllReduceExpert 1 Expert 3 计算 → AllReduceExpert 2Expert 2 计算时Expert 1 的结果已经在 AllReduce。实现方式# 开启通信与计算重叠hccl.enable_overlap(True)3. 梯度压缩用于训练AllReduce 前梯度 FP16→FP8通信量减半。hccl.set_quantization(fp8)工程经验HCCL 默认配置下 AllReduce 带宽利用率只有 45%。改成 NHR 算法 通信计算重叠后利用率拉到 89%。吞吐从 420 TPS 涨到 720 TPS。动态路由的负载均衡MoE 的动态路由会导致负载不均某些 expert 激活频繁某些 expert 空闲。问题场景Expert 1: 激活 45% 的 token Expert 2: 激活 35% 的 token Expert 3: 激活 15% 的 token Expert 4: 激活 5% 的 token ...Expert 1 忙死Expert 4 闲死。最忙的 expert 决定整体吞吐。解决动态 batch把 batch 内的 token 按 expert 分组每组单独处理Batch32, seq2048 → 65536 tokens ↓ 按激活的 expert 分组 Expert 1: 29500 tokens Expert 2: 22900 tokens Expert 3: 9800 tokens Expert 4: 3300 tokens ... ↓ 每个 expert 独立处理实现方式fromops_transformerimportMoEDynamicBatch moe_batchMoEDynamicBatch(num_experts8,top_k2,capacity_factor1.25# 允许 25% 的容量冗余)# 动态 batch 处理outputmoe_batch(x,experts,router_output)性能收益策略吞吐 (TPS)Expert 空闲率静态 batch57023%动态 batch720 5%动态 batch 把 Expert 空闲率从 23% 压到 5%吞吐涨 26%。完整优化路径优化项吞吐 (TPS)累计提升基线180-GatingTopK 融合22022%EP 并行策略420133%HCCL 优化720300%动态 batch780333%最终吞吐180→780 TPS涨了 4.3 倍。踩坑实录坑 1TP 并行在 MoE 场景慢TP 每个 expert 都要 AllReduce通信开销大。MoE 场景用 EP 或混合并行。坑 2Router 输出不融合开销大Router 输出写 HBM 再给 TopK 读开销 12μs/层。融合后降到 6μs/层。坑 3Expert 激活不均有空闲动态路由导致某些 expert 忙死、某些闲死。用动态 batch 解决。坑 4HCCL 默认配置带宽利用率低默认配置下 AllReduce 带宽利用率 45%。改成 NHR 算法 通信计算重叠后拉到 89%。https://atomgit.com/cann/ops-transformerhttps://atomgit.com/cann/hcclhttps://atomgit.com/cann/cann-recipes-infer
