MoE 推理完整机制:从 Router 到 Expert Parallelism 的一切

17 分钟阅读

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)

三个致命难点:

  1. 动态、稀疏、不规则——每个 token 走哪些 expert 是运行时决定的,不能提前编译;expert 负载不均衡(有的 expert 特别热门)

  2. 通信开销——如果 experts 分布在多个 GPU 上(Expert Parallelism),需要 all-to-all 通信把 token 送到对应 GPU

  3. 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 是"训练时代黄昏 + 推理时代崛起"的完美产物——它把大部分成本放在推理侧,而推理侧的每一分优化都直接变成利润。


附:素材来源

主线

代码 & 库

论文

  • Switch Transformer · arxiv 2101.03961

  • DeepSeek V3 Technical Report · arxiv 2412.19437

姊妹篇


整理于 2026-07-31 深夜。原文技术细节归属 vLLM 团队与 DeepSeek 团队。中文表达可能有解读偏差,请以原文为准。

Zoe

Written by

Zoe

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

评论