# MoE 推理优化：Mixtral 8×7B 在昇腾上的吞吐提升 4 倍

摘要：Mixtral 8×7B模型在910B 4卡部署时，初始吞吐仅180 TPS，远低于客户预期的1200 TPS。分析发现MoE推理瓶颈主要来自通信和调度，其中Router动态路由和AllReduce通信耗时占比达42%。通过GatingTopK算子融合（节省384μs/32层）、Expert并行策略优化（EP比TP快35%）、HCCL通信优化（带宽利用率提升至89%）以及动态batch负载均

waitingforloveJJ

24人浏览 · 2026-05-22 13:36:45

waitingforloveJJ · 2026-05-22 13:36:45 发布

摘要

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 算子：

from ops_transformer import GatingTopK

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。

配置方式：

import hccl

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 独立处理

实现方式：

from ops_transformer import MoEDynamicBatch

moe_batch = MoEDynamicBatch(
    num_experts=8,
    top_k=2,
    capacity_factor=1.25  # 允许 25% 的容量冗余
)

# 动态 batch 处理
output = moe_batch(x, experts, router_output)