MoE 推理完整机制:从 Router 到 Expert Parallelism 的一切
MoE 推理完整机制:从 Router 到 Expert Parallelism 的一切
前情说明:技术深度 D3。当代最大的开源模型基本都是 MoE——DeepSeek V4 1.6T、Kimi K3、Qwen3-30B-A3B、Mixtral、Grok。参数量爆炸但激活参数少是 MoE 的魅力,也是它推理侧的头号难题。本篇挖 MoE 推理的完整机制:从 Router 到 Expert 分发、到 fused kernel、再到 Expert Parallelism 的编排。
素材来自 vLLM Fused MoE 官方文档 + DeepGEMM 公开资料 + Elastic Expert Parallelism blog。
起手:MoE 到底为什么难
普通 dense LLM 推理你可以想象成"一根流水线"——每一层都是稠密矩阵乘,很规整。
MoE 层完全不一样——它长这样:
输入 batch of tokens (N × D)
↓
Router (小 gate 网络) 给每个 token 打分
↓
每个 token 选 Top-K 个 experts (通常 K=2 或 8)
↓
按 expert 重新组织 tokens:
Expert 0 收到: [token_3, token_17, token_42, ...]
Expert 1 收到: [token_1, token_9, token_88, ...]
Expert 2 收到: [token_5, token_11, ...] ← 负载可能不均衡
...
Expert 255 收到: [...]
↓
每个 expert 独立跑一个 FFN (两个 GEMM)
↓
把结果按原始 token 顺序汇合
↓
乘以 router 权重加权求和
↓
输出 (N × D)
三个致命难点:
-
动态、稀疏、不规则——每个 token 走哪些 expert 是运行时决定的,不能提前编译;expert 负载不均衡(有的 expert 特别热门)
-
通信开销——如果 experts 分布在多个 GPU 上(Expert Parallelism),需要 all-to-all 通信把 token 送到对应 GPU
-
kernel 效率——256 个 expert 每个都单独跑一个小 GEMM,每个都太小 GPU 吃不饱
MoE 推理引擎的所有优化都在攻这三个点。
第一层:Router 与 Top-K 的算法细节
Router 是一个很小的 gate 网络——通常只是一个 Linear(hidden_dim → num_experts)。给定输入 token embedding x,输出 logits:
logits = W_gate @ x # 形状 (N, num_experts)
然后对每一行做 Top-K 选择 + softmax:
top_k_logits, top_k_indices = torch.topk(logits, k=K, dim=-1) # 选前 K 个
top_k_weights = softmax(top_k_logits) # 归一化
关键 detail:softmax 是只对 top-K 做,不是对所有 experts。这样归一化后权重和为 1,可以直接加权求和。
DeepSeek 的 sigmoid + norm 变体
DeepSeek V2/V3/V4 用了另一种归一化方案——sigmoid + 手工归一化:
# 每个 expert 的 gate 独立 sigmoid
scores = torch.sigmoid(logits)
top_k_scores, top_k_indices = torch.topk(scores, k=K, dim=-1)
top_k_weights = top_k_scores / top_k_scores.sum(dim=-1, keepdim=True) # 归一化
为什么这么做? 因为 softmax 有个副作用——它是竞争式的(一个 expert 分高,其他就分低)。Sigmoid 则是独立的——多个 expert 都可以有高分。这在训练时能让 gate 输出更"稳定",梯度回传更健康。
推理时:算法不同,但 top-K + 加权求和的骨架是一样的——kernel 层可以复用大部分逻辑。
Auxiliary Loss 负载均衡(训练时)
MoE 有个致命的训练-推理不一致问题:如果不干预,训练时 gate 会退化到只用一两个 expert(因为 gradient 会集中优化那些),其他 expert 变成 dead weight。
训练时用 auxiliary load balancing loss 强制 gate 均匀分配 token。推理时这些 loss 全消失,只剩 top-K 选择——如果训练不够均衡,推理时依然会有个别 expert 特别热门。
这就是 hot expert 问题——推理侧引擎必须解决。
第二层:Fused MoE Kernel 的两种范式
Router 之后,需要把 token 送到对应 expert。这一步是 MoE 推理最热的战场。
范式 A:Contiguous / Non-Batched(标准 MoE)
思想:先按 expert 重排 token,再对每个 expert 跑一次 GEMM:
Step 1: Permutation(token 重排)
原 tokens: [t0, t1, t2, t3, t4, t5, t6, t7]
Router 分配: t0→e0, t1→e2, t2→e0, t3→e1, t4→e2, t5→e0, t6→e1, t7→e2
重排后:
Expert 0 段: [t0, t2, t5] ┐
Expert 1 段: [t3, t6] ├─ 一个连续 buffer
Expert 2 段: [t1, t4, t7] ┘
额外存: expert_offsets = [0, 3, 5, 8]
Step 2: 每个 expert 跑一次 batched GEMM
Expert 0: shape (3, D) → (3, D_ffn) → (3, D)
Expert 1: shape (2, D) → (2, D_ffn) → (2, D)
Expert 2: shape (3, D) → (3, D_ffn) → (3, D)
可以用 **Grouped GEMM** 或 **Batched GEMM** 一次调用
Step 3: Un-permutation(还原顺序)
按原始 token 顺序把结果放回,同时乘上 router 权重
关键 kernel:Grouped GEMM
CUTLASS 提供的 Grouped GEMM 允许一次 kernel launch 完成多个不同大小的 GEMM。避免了 256 次单独 kernel launch 的开销。
优化亮点:
-
Permutation 用 scatter/gather kernel 完成(一次 HBM 读+写)
-
Expert offsets 存 metadata 用于后续 kernel
-
Grouped GEMM 内部动态分配 SM 到不同 expert
主要问题:hot expert 会拖慢整个 grouped GEMM——因为 grouped GEMM 要等最慢的那个完成。
范式 B:Batched(Padded)
思想:给每个 expert 分配固定 capacity,超过就 drop token(或者其他策略):
capacity_per_expert = ceil(1.5 * N * K / num_experts) # 每个 expert 的固定容量
Expert 0 buffer: [t0, t2, t5, _, _, _, _, _] ← padding to capacity
Expert 1 buffer: [t3, t6, _, _, _, _, _, _]
...
优点:
-
所有 expert 是同样大小的 tensor → 可以用普通 batched matmul
-
简单粗暴,kernel 效率极高
缺点:
-
padding 浪费算力(很多 slot 是空的)
-
溢出的 token 被 drop 或送到 shared expert(质量损失)
这种范式主要用在训练时(容忍质量小损失换稳定性),推理侧几乎都用范式 A。
vLLM 的 Fused MoE Modular Kernel
vLLM 把 MoE 抽象成一个可插拔的模块化 kernel(官方文档):
FusedMoEModularKernel:
- prepare() # 量化 + dispatch(把 token 送到对应 expert 的位置)
- fused_experts() # 跑 grouped GEMM
- finalize() # 合并结果 + 加权
三个方法都可以独立替换——比如 dispatch 可以用 DeepEP(专门的 all-to-all 库),expert kernel 可以用 DeepGEMM(DeepSeek 出的 grouped GEMM 库)。
这个抽象的强大之处:新硬件(如 Ascend)或新算法(如 FP4 experts)只需要实现自己的 modular kernel,不用碰 vLLM 主逻辑。
第三层:Expert Parallelism(EP)—— 多卡怎么分 expert
单卡装不下所有 expert(1.6T 模型 256 个 expert,每个几 GB)。必须多卡拆 expert——这就是 Expert Parallelism (EP)。
三种并行策略的对比
| 并行策略 | 拆分方式 | 通信 | 显存占用 |
|---|---|---|---|
| TP (Tensor Parallelism) | 每个 expert 拆到多卡(按 hidden dim 切) | AllReduce | 每卡都有全部 expert 的一部分 |
| EP (Expert Parallelism) | 不同 expert 放不同卡 | All-to-All | 每卡只有部分 expert 完整副本 |
| EP + TP 混合 | Expert 分组,组内 TP | 两种都有 | 灵活 |
MoE 场景 EP 通常最优——因为 MoE 的参数量集中在 expert(router 很小),拆 expert 天然节省显存。
EP 的核心问题:All-to-All 通信
考虑 8 卡 EP,每卡有 32 个 expert(共 256)。Router 给一个 token 分了 top-2 expert,这两个 expert 可能在不同卡上。
流程:
Step 1: 每张 GPU 计算 router 决定,输出每个 token 要去的 expert
GPU 0: token_0 → expert_5 (本地) + expert_100 (在 GPU 3)
GPU 1: token_5 → expert_45 (在 GPU 1) + expert_200 (在 GPU 6)
...
Step 2: All-to-All Dispatch
把每个 token 的 hidden state 发送到对应 expert 所在的 GPU
这是全网状通信——每张 GPU 都要跟所有其他 GPU 通信
Step 3: 每张 GPU 用本地 expert 处理收到的 token
Step 4: All-to-All Combine
把处理结果送回原始 token 所在的 GPU
两次 all-to-all 是 MoE EP 的最大开销——延迟高、带宽敏感。
DeepEP:DeepSeek 的 all-to-all 库
DeepSeek 开源了 DeepEP,专门优化 MoE 的 all-to-all:
-
NVLink 优化:单节点内用 NVLink,带宽 900 GB/s
-
RDMA 优化:跨节点用 IBGDA (InfiniBand GPU Direct Async)
-
Prefetch + overlap:把 dispatch/combine 跟 expert 计算 overlap
-
低精度传输:token hidden state 可以 FP8 传输,节省一半带宽
收益:比 NCCL 的通用 all-to-all 快 2-3x。
DeepEP 之所以能这么快是因为它专为 MoE 通信模式设计——不是通用 all-to-all,是**"每 token 只发到 top-K 个 target"**这种特殊模式。
第四层:DeepGEMM & MegaMoE Kernel
即使有了 DeepEP,Expert 计算侧的 GEMM 效率也是瓶颈。DeepSeek 开源了 DeepGEMM 攻这个。
DeepGEMM 的 MegaMoE Kernel
核心思想:把所有 expert 的 GEMM 融合成一个大 kernel:
传统 Grouped GEMM:
launch kernel → 256 个 GEMM 各自跑 → 等最慢的完成
DeepGEMM MegaMoE:
launch 1 个巨大 kernel → 内部动态分配 SM 到不同 expert
→ 同一个 SM 可能处理多个 expert
→ 完全消除 kernel launch overhead
关键技术:
-
用 CUTLASS 底层原语,绕过 CuBLAS 的通用性开销
-
FP8 MoE 深度优化(DeepSeek V4 原生 FP8 权重)
-
支持任意 M 维度(每个 expert 收到的 token 数)
-
Persistent kernel——kernel 常驻在 SM 上,处理完一个 expert 立刻处理下一个
收益:比原始 grouped GEMM 快 30-50%,是 DeepSeek 模型能在低成本硬件上跑起来的关键。
vLLM 已经集成 DeepGEMM MegaMoE kernel(DeepSeek V4 部署 recipe 里可以看到)。
第五层:Elastic Expert Parallelism(vLLM 2026-05)
MoE 生产环境有个头疼问题:hot expert 会导致负载严重不均衡——某几个 expert 的所在 GPU 打爆,其他 GPU 闲着。
传统方案:只能重启服务,重新分配 expert。生产上这不能接受。
vLLM Elastic EP 的做法
Elastic EP blog(2026-05-14)介绍了运行时 expert 迁移:
监控 → 发现 hot expert
↓
决定 expert 迁移方案(比如把 expert_47 从 GPU 3 迁到 GPU 5)
↓
在线迁移 expert weights(用 CUDA IPC 或 RDMA)
↓
更新 router table(所有 GPU 同步新映射)
↓
新请求走新路径,旧请求继续用旧路径直到完成
关键技术:
-
双缓冲——迁移中的 expert 在源 GPU 和目标 GPU 都保留副本,直到确认
-
原子路由表切换——用一次 broadcast 更新所有 GPU 的路由表,避免中间态
-
DP 副本变化——如果不只是重分配,还可以增/减 DP worker 数量
收益:
-
Hot spot 缓解:热点 expert 可以复制到多张卡分散负载
-
Elastic scaling:workload 上升自动加 DP worker,下降自动减
-
零停机——生产环境的关键
这是 MoE 推理走向"云原生"的关键一步——就像 K8s 的 HPA 自动扩缩容。
一张表看清 MoE 推理的六个优化点
| 层级 | 优化 | 代表实现 | 收益 |
|---|---|---|---|
| 算法层 | Sigmoid + norm 归一化 | DeepSeek V2/V3/V4 | 训练稳定性↑ |
| Kernel 层 | Grouped GEMM | CUTLASS | 消除 256 次 launch |
| Kernel 层 | MegaMoE (persistent + fused) | DeepGEMM | +30-50% |
| 通信层 | MoE 专用 all-to-all | DeepEP | +2-3x |
| 通信层 | FP8 token 传输 | DeepEP + vLLM | 带宽↓ 50% |
| 调度层 | Elastic EP | vLLM 2026-05 | 消除 hot spot + 弹性伸缩 |
MoE 推理的每一层都需要专门优化——这就是为什么 vLLM 花了两年才把 MoE 优化到当前水平。
三个通用启示
启示 1:稀疏是免费的算力,但代价是"不规则"
MoE 的核心承诺——参数量 10x,激活参数不变——听起来很美。但这是用规则性换来的:
-
Dense GEMM: 每个 token 的每个参数都参与计算 → 完美规则 → tensor core 打满
-
Sparse MoE: 每个 token 只用一部分 expert → 不规则 → 需要 permutation / all-to-all / dynamic dispatch
没有免费的午餐——每次算力"稀疏了",都要在其他地方(内存/通信/调度)付出代价。
这也是为什么 MoE 训练比 dense 训练难得多——负载均衡、专家路由、通信开销都是新问题。
启示 2:MoE 是"通信 bound",不是"计算 bound"
Dense LLM:主要是 memory bandwidth bound(decode)或 compute bound(prefill)。
MoE:这两个都不是——通信 bound。All-to-All dispatch/combine 占了 MoE forward 的 30-50% 时间。
这就是为什么 DeepEP / DeepGEMM 是 DeepSeek 能开源出 SOTA MoE 的关键——它们本质上都在降低通信成本。
Nvidia 也看到了这个趋势——Blackwell 的 NVLink 5 直接把带宽翻倍,就是给 MoE 用的。
启示 3:MoE 的商业逻辑:更大模型,更便宜服务
站远一点看——MoE 的核心价值是训练成本和推理成本解耦:
-
训练时:全部 1.6T 参数都在优化(成本高,但一次性)
-
推理时:只激活 30-50B 参数(成本低,每个请求都省)
结果:用大模型的智能,付小模型的价格。
这就是为什么 2026 年头部开源模型全是 MoE——DeepSeek V4、Kimi K3、Qwen3-30B-A3B、Grok。Dense 大模型(除了 Llama)几乎绝迹——从商业逻辑上根本竞争不过 MoE。
MoE 是"训练时代黄昏 + 推理时代崛起"的完美产物——它把大部分成本放在推理侧,而推理侧的每一分优化都直接变成利润。
附:素材来源
主线
-
vLLM Docs · Fused MoE Modular Kernel
-
vLLM Blog · Elastic Expert Parallelism(2026-05-14)
-
DeepWiki · FusedMoE Layer Architecture
代码 & 库
-
vllm-project/vllm - MoE modular kernel 实现
-
deepseek-ai/DeepEP - MoE all-to-all 通信库
-
deepseek-ai/DeepGEMM - MegaMoE grouped GEMM 库
-
CUTLASS · NVIDIA/cutlass
论文
-
Switch Transformer · arxiv 2101.03961
-
DeepSeek V3 Technical Report · arxiv 2412.19437
姊妹篇
整理于 2026-07-31 深夜。原文技术细节归属 vLLM 团队与 DeepSeek 团队。中文表达可能有解读偏差,请以原文为准。

Written by
Zoe
AI Infra Engineer · LLM Serving · GPU/RDMA · 造工具的偏执狂