# MoE 推理优化：Mixtral 8×7B 在昇腾上的吞吐提升 4 倍-平芜编程栈

摘要

Mixtral 8×7B 部署到 910B 4 卡，吞吐 180 TPS。客户说 GPU 上能到 1200 TPS，差了 6 倍。跑 msprof 一看：MoE Router + AllReduce 占 42% 时间。

MoE（Mixture of Experts）模型的推理瓶颈不在计算，在通信和调度。Router 动态路由 + Expert 通信开销吃掉 40% 时间。优化后吞吐从 180 提到 720 TPS，涨了 4 倍。

MoE 推理瓶颈分析

MoE 的核心结构：

输入 x ↓ Router（选择 top-k expert） ↓ Expert 1, Expert 2, ..., Expert n（只激活 top-k 个） ↓ 加权求和 ↓ 输出 y

Mixtral 8×7B：8 个 expert，每次激活 top-2。

瓶颈拆解：

阶段	耗时占比	问题
Router（Gating + TopK）	18%	动态路由开销
Expert 计算	35%	8 个 expert 并行度不足
AllReduce 通信	24%	Expert 间通信带宽利用率低
其他	23%	-

Router + AllReduce 占 42% 时间，这才是 MoE 推理慢的根本原因。

工程经验：很多人以为 MoE 慢是因为 expert 多、计算量大。其实 MoE 每次只激活 top-k 个 expert，计算量比同等参数的 Dense 模型还小。真正的瓶颈是 Router 动态路由和 Expert 间通信。

GatingTopK 算子融合

Router 的计算流程：

1. Gating：x @ W_gate → gate_score（每个 token 对 8 个 expert 的得分） 2. TopK：从 8 个得分中选 top-2 3. Softmax：对 top-2 得分做 softmax，得到权重

不融合的实现：

Gating 输出 → 写 HBM → TopK 读 HBM → 写 HBM → Softmax 读 HBM

3 次 HBM 读写，开销 36μs（每次 12μs）。

融合后的实现：

Gating 输出 → L1 → TopK → L1 → Softmax

中间结果走 L1，不落 HBM。省掉 2 次 HBM 读写，开销降到 12μs。

ops-transformer 的 GatingTopK 算子：

fromops_transformerimportGatingTopK gating_topk=GatingTopK(num_experts=8,top_k=2,fusion_mode="full"# Gating + TopK + Softmax 全融合)# 融合计算expert_ids,expert_weights=gating_topk(x,W_gate)

性能收益：

模型	融合前 Router 耗时	融合后 Router 耗时	节省
Mixtral 8×7B	18μs/层	6μs/层	12μs/层
32 层总计	576μs	192μs	384μs

工程经验：GatingTopK 融合前，Router 输出要写 HBM 再给 TopK 读。融合后直接 L1 传数据，省一次 HBM 读写。单层省 12μs，32 层省 384μs。decode 每个 token 快 0.38ms，1000 个 token 差 0.38 秒。

Expert 并行策略

MoE 模型有两种并行策略：Tensor Parallel（TP）和 Expert Parallel（EP）。

Tensor Parallel（TP）：

把每个 expert 的权重切分到多卡，每卡算一部分。

Expert 1: 卡1 算 1/2，卡2 算 1/2，AllReduce 合并 Expert 2: 卡1 算 1/2，卡2 算 1/2，AllReduce 合并 ...

Expert Parallel（EP）：

每个卡放完整的 expert，不同卡算不同 expert。

卡1: Expert 1, Expert 2 卡2: Expert 3, Expert 4 卡3: Expert 5, Expert 6 卡4: Expert 7, Expert 8

对比：

策略	通信模式	适用场景
TP	AllReduce（每个 expert 都要）	Expert 数量少、单 expert 大
EP	All-to-All（激活的 expert 通信）	Expert 数量多、单 expert 小

Mixtral 8×7B：8 个 expert，每个 expert 7B 参数。EP 更优。

实测数据（Mixtral 8×7B，910B 4 卡）：

并行策略	吞吐 (TPS)	通信开销
TP	420	28%
EP	570	18%
混合（TP=2, EP=2）	610	15%

混合策略最优：TP=2（每个 expert 切 2 份），EP=2（每卡放 4 个 expert）。

工程经验：MoE 场景下 EP 比 TP 快 35%。原因：TP 每个 expert 都要 AllReduce，EP 只在激活的 expert 之间通信。Mixtral 每次只激活 2 个 expert，通信量比 TP 小。

HCCL 通信优化

HCCL 是昇腾的集合通信库，对标 NCCL。

AllReduce 带宽利用率对比：

库	带宽利用率	吞吐 (TPS)
HCCL 默认配置	45%	420
HCCL 优化后	89%	720

优化点：

1. 算法选择

HCCL 支持多种 AllReduce 算法：Ring、Mesh、NHR（Non-hierarchical Ring）。

小数据量（< 1MB）：Mesh 最快 中数据量（1MB-100MB）：NHR 最快 大数据量（> 100MB）：Ring 最快

MoE 的 AllReduce 数据量通常 10-50MB，选 NHR。

配置方式：

importhccl hccl.set_allreduce_algorithm("NHR")

2. 通信与计算重叠

Expert 计算和 AllReduce 通信并行：

Expert 1 计算 → Expert 2 计算 → AllReduce（Expert 1） Expert 3 计算 → AllReduce（Expert 2）

Expert 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 分组，每组单独处理：

Batch=32, seq=2048 → 65536 tokens ↓ 按激活的 expert 分组 Expert 1: 29500 tokens Expert 2: 22900 tokens Expert 3: 9800 tokens Expert 4: 3300 tokens ... ↓ 每个 expert 独立处理

实现方式：

fromops_transformerimportMoEDynamicBatch moe_batch=MoEDynamicBatch(num_experts=8,top_k=2,capacity_factor=1.25# 允许 25% 的容量冗余)# 动态 batch 处理output=moe_batch(x,experts,router_output)