Lecture 8 · 混合专家模型的推理:专家并行、MegaScale-Infer 与 Helix · Mixture of Experts

AI Inference Infrastructure · Lecture 08。个人学习讲义与配套实验,按主题整理并持续修订。

课程源码 ZIP ↓

个人整理的学习笔记,非官方讲义;课程来源与整理说明见课程介绍。配套示例未逐一运行验证。

本篇目录

对应 Maps to:Stanford CS349D AI Inference Infrastructure(Spring 2026),#12 Wed May 6 Mixture of Experts。必读:R. Zhu, Z. Jiang, C. Jin 等, MegaScale-Infer: Serving Mixture-of-Experts at Scale with Disaggregated Expert Parallelism(字节跳动 Seed / 北京大学,2025);N. Bhatia, A. More, R. Borkar 等, Helix Parallelism: Rethinking Sharding Strategies for Interactive Multi-Million-Token LLM Decoding(NVIDIA,2025)。选读:N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, J. Dean, Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer(ICLR 2017);O. J. Aimuyo, B. Oh, R. Singh, FlashMoE: Fast Distributed MoE in a Single Kernel(NeurIPS 2025)。 依据说明:原课没有公开 slides。本讲依据四篇论文的全文。§2”MoE 推理的代价”的分析、§3 里关于 DeepSeek-V3 部署的数字(来自 DeepSeek-V3 技术报告和 DeepSeek 2025 年 2 月公开的推理系统概述)、§7”之后的发展”是我补充的。讲义里的纸上数字全部来自配套代码的实际运行。 配套代码 Code:code/lecture08/(numpy;5 个脚本 + 1 个模块 moe_lab.py:真实的 MoE 层与专家并行、MoE 的算子级代价模型、大规模专家并行的 decode、MegaScale-Infer 的 ping-pong 模拟与部署搜索;Helix 的 KV 并行注意力在代码 05 里。代码 03 约 10 秒,其余几秒) 预计阅读时间:3–4 小时(含运行代码)


0. 本讲导航 · How to read this lecture

0.1 学习目标 Learning goals

学完这一讲,你应该能够:

  1. 说清 MoE 层的结构(门控、top-k、专家、共享专家)、总参数与激活参数的区别,以及 Shazeer 2017 为什么需要负载均衡损失;解释推理引擎为什么不像训练那样”丢 token”。
  2. 用”一批 token 会碰到几个专家”和”每个专家分到多少 token”两个量,解释 MoE 的 decode 为什么在批 1 时便宜、在中等批时比同尺寸的稠密模型还贵、要很大的批才能把算力用满。
  3. 说清专家并行的数据流(两次 all-to-all,中间是”汇聚后的批量计算”),估算 all-to-all 的数据量和时间;解释 DeepSeek 为什么用几十到几百块卡组成一个专家并行组跑 decode,以及负载不均为什么会吃掉它的收益、冗余专家怎样补救。
  4. 说清 MegaScale-Infer 的注意力 / 专家解耦:为什么要拆、ping-pong 流水线怎样把两边填满(三个条件和最少微批数)、怎样搜索部署方案、为什么异构硬件能省钱、M2N 通信库解决了 NCCL 的什么问题。
  5. 说清 Helix:长上下文 decode 中张量并行为什么在 KV 头数处”撞墙”,KV 并行怎样精确地切开注意力(部分结果 + LSE + 一次 all-to-all),以及为什么同一组 GPU 在 FFN 阶段要换一种切法。
  6. 知道 FlashMoE 把整个分布式 MoE 层融合进一个常驻内核要解决的两个问题:几百次内核启动、对称集合通信的补零。

0.2 先修与本讲范围 Prerequisites & scope

默认你已经会 Assumed本讲只给直觉,后面再展开 Previewed only
第 1 讲:roofline、算术强度、decode 是 memory-boundGPU 内核、SM、warp 的细节(前置讲 P1)
第 4 讲:注意力状态 (O, LSE) 与合并 ⊕、split-KV集合通信的硬件细节:NVSwitch、RDMA(前置讲 P5)
第 6 讲:张量并行、流水线并行、all-to-all、α–β 模型混合架构与稀疏注意力(第 10 讲 DeepSeek-V3.2)
第 7 讲:prefill / decode 解耦、goodput、异构部署MoE 的训练(负载均衡损失只讲到推理需要的程度)

0.3 本讲地图 Map

§1  MoE 基础:Shazeer 2017——门控、稀疏、负载均衡、"批越分越小"            ← 代码 01 演示 1–2、代码 02 演示 1
§2  MoE 推理的代价:碰到几个专家、算子级 roofline、每个专家的批             ← 代码 02 演示 2–4
§3  专家并行:两次 all-to-all、大规模 EP 的 decode、负载不均与冗余专家       ← 代码 01 演示 3、代码 03
§4  MegaScale-Infer:注意力 / 专家解耦、ping-pong、部署搜索、异构、M2N      ← 代码 04
§5  Helix:长上下文 decode——注意力按 KV 切、FFN 按权重切                   ← 代码 05
§6  FlashMoE:把分布式 MoE 层放进一个内核                                  ← 代码 01 演示 3(补零)
§7  之后的发展
§8–12  小结、常见误区、练习、延伸阅读、下一讲预告

0.4 记号 Notation

符号含义
EE、kk路由专家数、每个 token 选几个专家(top-k);E/kE/k 是”稀疏度”
G(x)G(x)门控输出:长度 EE 的稀疏权重向量,只有 kk 个非零
总参数 / 激活参数全部专家都算 / 一个 token 实际用到的(kk 个路由专家 + 共享专家 + 其余部分)
EP、DP、TP、PP专家并行、数据并行、张量并行、流水线并行
TaT_a、TeT_e、TcT_cMegaScale-Infer:一个微批在注意力节点、专家节点上一层的计算时间,单向通信时间;Tf=max⁡(Ta,Te)T_f = \max(T_a, T_e)
mm、LL微批数、层数
tpatp_a、tpetp_e、nan_a注意力 / 专家节点内的张量并行度、注意力节点数
bab_a、beb_e每个注意力节点、每个专家每个微批的 token 数;be=banak/Eb_e = b_a n_a k / E
KVP、TPA、TPFHelix:KV 并行宽度(沿序列切 KV)、注意力的张量并行宽度(≤ KV 头数)、FFN 的张量并行宽度
TTLHelix:token 间延迟(token-to-token latency),即 TBT

0.5 关于数字 · About the numbers

本讲的数字有三类:(1) numpy 上的真实计算:MoE 层、专家并行、KV 并行注意力的输出核对(代码 01、代码 05 演示 1);(2) 纸上 GPU:第 1 讲 hw_model 的 GPU 参数和效率系数,但改成算子级 roofline(§2.2),通信用第 6 讲的 α–β 模型,Helix 的 NVL72 通信另有假设(§5.5);它只代表数量级和趋势,而且比真实系统乐观;(3) 模拟:按热度抽样的路由、ping-pong 流水线的离散事件模拟。论文里的数字写成”论文报告”。


1. MoE 基础:稀疏门控的混合专家层 · Sparsely-gated MoE

1.1 动机:参数多一些,计算别跟着涨

一个模型能”装下”多少知识,大致取决于它有多少参数;可每多一个参数,每个 token 的计算也多一份(第 1 讲:每 token 约 2N 次运算)。有没有办法让参数涨、计算不涨?条件计算(conditional computation)的想法很早就有:每个输入只激活网络的一部分。Shazeer 等人 2017 年第一次让它在大规模上真正奏效。论文开头列了条件计算长期做不起来的五个原因,几乎每一条都是系统问题:

  1. GPU 擅长算术、不擅长分支:所以要”成块地”开关网络,而不是逐个神经元;
  2. 大批是效率的关键,而条件计算会把每个被激活部分的批变小(“批越分越小”,shrinking batch problem);
  3. 网络带宽可能成为瓶颈:GPU 集群的算力比设备间带宽高几千倍,算法的”计算 / 通信比”必须超过这个比例;
  4. 需要额外的损失项来控制稀疏程度和负载均衡;
  5. 模型容量只有在超大数据集上才有意义。

这五条里,第 2、3 条正是今天 MoE 推理系统要解决的问题,只是从训练搬到了推理,而且更严重(§2)。

1.2 MoE 层的结构

MoE 层由 EE 个”专家”(每个是一个前馈网络 FFN)和一个”门控网络”组成。输出是

y=∑i=1EG(x)i Ei(x)y = \sum_{i=1}^{E} G(x)_i \, E_i(x)

其中 G(x)G(x) 是稀疏的:只有 kk 个分量非零,G(x)i=0G(x)_i = 0 的专家根本不用算。在 Transformer 里,MoE 层替换的是每层(或每隔几层)的 FFN,注意力不变。

门控怎么选。Shazeer 的 Noisy Top-K Gating:

H(x)i=(xWg)i+StandardNormal()⋅Softplus((xWnoise)i)H(x)_i = (x W_g)_i + \text{StandardNormal}() \cdot \text{Softplus}((x W_{noise})_i)

G(x)=Softmax(KeepTopK(H(x),k))G(x) = \text{Softmax}(\text{KeepTopK}(H(x), k))

KeepTopK 把前 kk 名以外的分数设成 −∞-\infty(softmax 之后为 0)。噪声在训练时帮助负载均衡,推理时通常关掉。今天的模型有几种写法(代码 01 演示 1 对比了三种):

  • 先 top-k 再 softmax(Shazeer、Mixtral):kk 个权重之和为 1;
  • 先 softmax 再 top-k、不再归一化(DeepSeek-V2 等):权重之和小于 1;
  • sigmoid 亲和度、选出后再归一化(DeepSeek-V3):并且可以给”选择”加一个偏置,只影响选谁、不影响权重——DeepSeek-V3 用它代替辅助损失做负载均衡:某个专家太忙就把它的偏置调低。

前两种写法选中的专家完全相同(softmax 不改变大小顺序),区别只在权重的尺度。

专家的粒度。Mixtral-8x22B 是”少而大”的专家(8 个,选 2 个);DeepSeek-MoE 系列走向”多而小”(DeepSeek-V3:256 个路由专家选 8 个,再加 1 个所有 token 都经过的共享专家),Qwen3-235B 是 128 选 8。专家越细,组合越灵活,但稀疏度 E/kE/k 越高——后面会看到,这对推理系统是更大的挑战。

1.3 总参数与激活参数

模型总参数激活参数激活比例BF16 权重放下权重至少要几块 H100E/kE/k
Mixtral-8x7B46.7B12.9B28%93 GB24
Mixtral-8x22B140.6B39.2B28%281 GB44
DBRX131.6B36.5B28%263 GB44
Qwen3-235B-A22B235.1B22.2B9%470 GB716
DeepSeek-V3671.0B37.6B6%1342 GB1932
Llama-3-70B(稠密)70.6B70.6B100%141 GB21

(代码 02 演示 1,按公开结构逐项计算,与公开的总参数、激活参数一致。)

这张表是理解 MoE 推理的起点:计算按激活参数算,显存按总参数算。DeepSeek-V3 每个 token 只用 37.6B 参数做计算,却要把 671B 参数全部放进显存——BF16 要 19 块 H100,DeepSeek 用 FP8 存权重也要 10 块。所以 MoE 在计算上像一个小模型,在显存上像一个大模型。

1.4 负载均衡:训练时的问题,推理时的后果

Shazeer 观察到:门控会自我强化地偏向少数几个专家——被选得多的专家训练得快,于是被选得更多。论文定义了两个量(代码 01 演示 1):

  • importance:一批里所有 token 给某个专家的门控权重之和;
  • load:一批里有多少个 token 选了这个专家。

并用”变异系数的平方” CV2\text{CV}^2 做损失,鼓励所有专家一样重要、一样忙。为什么要两个?因为 importance 均衡不代表 load 均衡:一个专家可以收少量大权重的 token,另一个收大量小权重的 token;而分布式硬件上决定时间的是 load。论文附录 A 的实验(论文报告):不加任何均衡损失时测试困惑度 39.8,加了之后约 35.6;最忙的专家是平均的 17.80 倍,两个损失的权重都取 1.0 时降到 1.07 倍。

训练好的模型也只是”大致”均衡。代码 01 演示 1 里,随机输入下最忙的专家是平均的 1.27 倍;输入按”话题”聚集时(一个批里 70% 的 token 来自同一个话题),最忙的是平均的 2.8 倍。MegaScale-Infer 在真实流量里也看到明显的冷热不均(论文图 16)。

容量因子与丢 token。GShard、Switch Transformer 等训练系统给每个专家设一个容量 C=容量因子×token 数×kEC = \text{容量因子} \times \frac{\text{token 数} \times k}{E},超出的 token 被丢掉(这个 token 在这个专家上的贡献变成 0,残差连接还在)。这样每个专家的计算量固定,方便编译和通信。推理时能不能这样做?代码 01 演示 2 的答案是”不能”:

  • 容量因子 1.0–1.25 时,话题聚集的输入会丢掉三到四成的 (token, 专家) 对,输出的相对误差在 0.36–0.48;
  • 更严重的是,有容量限制时一个 token 的输出取决于同一批里还有谁:同一个请求的 16 个 token,换一批”邻居”一起处理,输出最多相差 2.15;不限容量时两次的输出逐位相同。

这违反了推理系统”不改变输出”的不变量(第 1 讲),还让结果不可复现。所以 vLLM、SGLang、TensorRT-LLM 等推理引擎都是 dropless 的:一个专家收多少就算多少。代价是:负载不均时最忙的专家拖慢整层(§3.3)。

1.5 Shazeer 的系统设计:怎样让每个专家看到大批

“批越分越小”:bb 个 token 的批,每个专家只分到约 kb/Ekb/E 个。Shazeer 的解法是混合数据并行和模型并行:注意力(论文里是 LSTM)和门控按数据并行复制到 dd 个设备上,每个设备处理不同的批;而每个专家只有一份,放在某个设备上;所有设备的 token 汇到一起再交给专家。于是每个专家收到约 kbd/Ekbd/E 个 token——专家批变大了 dd 倍。设备数越多,可以放的专家越多,每个专家的批却保持不变。这正是今天”注意力数据并行 + 专家并行”的原型(§3)。

网络带宽:专家固定不动,网络上传的是专家的输入和输出。一个隐藏层大小为 hh 的专家,“计算 / 输入输出”之比正好是 hh,所以只要专家的隐藏层够大,计算就能盖过通信。

论文的结果(论文报告):在 1000 亿词的语料上,137B 参数的 MoE 层(131,072 个专家);65,536 个专家时(层稀疏度 99.994%)困惑度比计算量相当的基线低 39%,计算效率仍有 0.72 TFLOPS/GPU(K40)。

项目内容
论文Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer,Shazeer 等,ICLR 2017(Google Brain)
一句话用可训练的稀疏门控从成千上万个专家里为每个 token 选 k 个,把模型容量提高 1000 倍以上而计算几乎不变
问题条件计算在理论上诱人,但分支、批变小、网络带宽、均衡损失和数据规模让它一直没有在大规模上奏效
关键洞察稀疏到”整块专家”级别;用数据并行 + 模型并行的混合把各设备的小批汇成专家的大批;专家的隐藏层大小就是计算 / 通信比
机制Noisy Top-K Gating;importance 与 load 两个 CV² 损失;分层 MoE;把 MoE 层”卷积式”地应用到所有时间步以增大批
评测语言建模和机器翻译上以更少的计算超过当时最好的结果;最大 137B 参数的 MoE 层(论文报告)
局限与假设训练时用容量 / 均衡损失控制负载;专家仍会冷热不均;MoE 层在 LSTM 之间,而不是今天的 Transformer FFN
与本课的关系今天 MoE 推理的每个问题(批变小、all-to-all、负载均衡)在这篇论文里都已出现;§3 的”注意力 DP + 专家 EP”就是它的混合并行

2. MoE 推理的代价 · What MoE costs at inference time

(这一节的分析是我补充的,§2.4 的论证来自 MegaScale-Infer §2.3。)

2.1 一批 token 会碰到几个专家

decode 一步时,每个被至少一个 token 选中的专家都要把权重从显存读一遍。nn 个 token、每个均匀地选 kk 个不同专家时,某个专家不被一个 token 选中的概率是 1−k/E1 - k/E,不被任何一个选中的概率是 (1−k/E)n(1-k/E)^n,所以

E[被碰到的专家数]=E⋅(1−(1−k/E)n)\mathbb{E}[\text{被碰到的专家数}] = E \cdot \left(1 - (1 - k/E)^n\right)

代码 02 演示 2 用按热度抽样的路由核对了这个公式:

(E,k)(E, k)批 1批 4批 16批 64批 256
(8, 2):Mixtral25.57.988
(128, 8):Qwen3829.182.4125.9128
(256, 8):DeepSeek-V3830.5102.0222.4255.9

批 1 时只读 kk 个专家;批稍微大一点,几乎所有专家都被碰到——而此时每个专家平均只分到很少的 token(DeepSeek-V3 批 256 时每个专家平均 8 个)。专家热度偏斜时碰到的少一些,但”中等批就要读几乎全部权重”的结论不变。

这就是 MoE decode 的尴尬:读的字节很快涨到总参数,做的运算却只按激活参数算。

2.2 为什么本讲换成”算子级 roofline”

第 1 讲的 hw_model 把一整步的运算量和字节数各自加总,再取 max(计算时间, 访存时间)。这等于假设”访存慢的算子”和”计算慢的算子”能完美重叠。可 GPU 是一个算子接一个算子地执行的:decode 时注意力在等显存读 KV,这段时间里 FFN 的矩阵乘法并不能同时进行。MegaScale-Infer 的出发点恰恰是”注意力 memory-bound、FFN 本可以 compute-bound,两者串行时各自浪费一部分”。所以本讲按算子分别取 max 再相加:

t层=max⁡(注意力的计算,注意力的访存)+max⁡(MoE 的计算,MoE 的访存)+通信+固定开销t_{\text{层}} = \max(\text{注意力的计算}, \text{注意力的访存}) + \max(\text{MoE 的计算}, \text{MoE 的访存}) + \text{通信} + \text{固定开销}

它对稠密模型给出的时间比 hw_model 略长(更真实),趋势不变。

2.3 Mixtral 与稠密的 Llama-3-70B:谁的 decode 更快?

代码 02 演示 3:8 块 H100、张量并行 8 路、上下文 1024。

批大小Mixtral-8x22B 一步Llama-3-70B 一步Mixtral / Llama
17.33 ms11.69 ms0.63
412.85 ms11.78 ms1.09
1616.93 ms12.12 ms1.40
6417.91 ms13.48 ms1.33
25621.29 ms18.93 ms1.12
51225.81 ms31.37 ms0.82
102434.85 ms56.33 ms0.62

三段走势,各有原因:

  • 批 1:Mixtral 只读 2 个专家(激活 39B),比读 70B 的 Llama 少 37% 的时间——这就是”大模型、小延迟”。
  • 批 4–256:Mixtral 很快碰到全部 8 个专家,每步要读 141B 参数,MoE 部分每层的访存(225 µs)是 Llama FFN(66 µs)的 3.4 倍,整体反而慢 9%–40%。在这个区间,“MoE 更省”不成立。
  • 批 ≥ 512:FFN 转向 compute-bound,时间按激活参数增长,Mixtral 重新变快;显存允许的最大批下,Mixtral 的吞吐是 Llama 的 1.6 倍(30,144 vs 18,624 token/s)。

所以 MoE 的 decode 曲线是”两头好、中间差”。真正省钱的前提是每步的批足够大,让每个专家分到足够多的 token。

2.4 每个专家要分到多少 token?

MegaScale-Infer §2.3 的论证:一次 b×hb \times h 乘 h×nh \times n 的矩阵乘法要做 2bhn2bhn 次运算、读 2hn2hn 字节的 BF16 权重。设 GPU 算力为 FF、带宽为 BB,要 compute-bound 需要 2bhn/F≥2hn/B2bhn/F \ge 2hn/B,即 b≥F/Bb \ge F/B。A100:312 TFLOPS ÷ 2 TB/s = 156 个 token。对稠密模型这是整批的大小;对 MoE,这是每个专家的批,整批要大 E/kE/k 倍:

FFN 的利用率=min⁡(kE⋅BF⋅b, 1)\text{FFN 的利用率} = \min\left(\frac{k}{E} \cdot \frac{B}{F} \cdot b, \ 1\right)

论文的例子:Mixtral 8x22B 在 A100 上批 156 时,每个专家平均只分到 39 个 token,FFN 的理论 MFU 只有 25%。

打上本课程的效率折扣后,H100 上每个专家要分到约 258 个 token。代码 02 演示 4 算了显存允许的最大批下(上下文 2048)每个专家实际分到多少:

模型部署最大批每专家 token整批要多大才够MoE 算力用了几成
Mixtral-8x22B8 块 H100(TP8)6271571,03361%
Mixtral-8x22B16 块(TP8×PP2)9262321,03390%
Qwen3-235B(FP8)8 块(TP8)432274,13310%
DeepSeek-V3(FP8)16 块(TP8×PP2)20868,2663%
DeepSeek-V3(FP8)32 块(TP8×PP4)354118,2664%

8 选 2 是”温和”的稀疏,Mixtral 的专家能分到一两百个 token;DeepSeek-V3 在这种”传统”部署下每个专家只分到个位数的 token,MoE 部分的算力几乎全部闲置。DeepSeek-V3 的最大批小得出奇还有一个原因:MLA 只缓存一份所有头共用的潜向量,张量并行无法按头切开它,8 个 TP 分片各存一份完整的 KV(§5.1 还会遇到这个问题)。

第 8 讲所有系统设计的共同目标,就是让每个专家看到更大的批。三类办法:

  1. 把更多 GPU 上的请求汇到同一组专家:大规模专家并行(§3,DeepSeek 的做法);
  2. 把注意力拿出去单独扩展,让注意力那边放更多 KV、凑出更大的批:注意力 / 专家解耦(§4,MegaScale-Infer);
  3. 用更便宜的硬件跑注定 memory-bound 的部分(§4.5)。

3. 专家并行 · Expert parallelism

3.1 两次 all-to-all

专家并行(EP)把 EE 个专家分到多个设备上,每个设备只存其中一部分。一个 MoE 层的数据流(代码 01 演示 3 在 4 个虚拟设备上真实地实现了它):

  1. 路由:每个设备用(复制的)门控为自己的 token 选专家;
  2. dispatch:一次 all-to-all,把每个 token 发到它选中的专家所在的设备;同一个 token 选了同一设备上的两个专家时只发一次,门控权重随数据一起发;
  3. 专家计算:每个设备把从所有设备收到的 token 拼成一批,每个本地专家做一次批量计算——各设备的小批在这里汇成大批;
  4. combine:再一次 all-to-all,把结果发回 token 原来的设备,按门控权重加起来。

注意力、门控、共享专家通常按数据并行复制,每个设备只处理自己的请求。输出与单卡完全相同(代码 01 演示 3:最大差 0)。

通信量。每个设备每层发出约”本设备 token 数 × kk × D−1D\frac{D-1}{D}“份 token 向量(DD 是专家并行组的大小),dispatch 一次、combine 一次。DeepSeek-V3 的做法是 dispatch 用 FP8、combine 用 BF16。例子(练习 3):每块卡 128 个 token、64 卡的专家并行组,dispatch 约 7.2 MB、combine 约 14.5 MB;走 400 Gb/s 的 InfiniBand(有效 40 GB/s)分别约 0.19 ms 和 0.37 ms,58 个 MoE 层合计约 32 ms——必须和计算重叠,否则通信就吃掉了 decode 的全部时间预算。

3.2 大规模专家并行的 decode

DeepSeek-V3 的部署把 Shazeer 的混合并行推到极致:注意力(MLA)按数据并行,每块卡服务自己的一批请求;路由专家按专家并行,分到几十到上百块卡上。代码 03 演示 2 在纸上模拟了它(FP8 权重、上下文 4096、H800 + 400 Gb/s IB):

EP 组每卡序列每专家 token每层:注意力 / 路由专家 / dispatch / combine每卡吞吐(串行)每卡吞吐(双微批交替)
164924152 / 257 / 71 / 137 µs1,2531,006
32120120275 / 128 / 172 / 338 µs2,1182,910
64155310335 / 79 / 224 / 442 µs2,3303,718
128173692366 / 88 / 251 / 497 µs2,3473,717
2561821,456381 / 93 / 265 / 525 µs2,3543,715

(8 块卡放不下:FP8 权重 671 GB 加上 KV 超过 8 × 72 GB。)

看三件事:

  1. 专家并行组越大,每卡越”轻”:16 卡时每卡要存约 41 GB 专家权重,128 卡时约 5 GB;省下的显存放 KV,每卡的序列数从 49 涨到 182;全组一起路由,每个专家分到的 token 从 24 涨到上千。
  2. 64 卡以后每卡吞吐不再涨:每卡序列数受 KV 限制基本饱和,注意力(读 KV)和两次 all-to-all(每层合计约 0.7 ms)成了大头。
  3. 双微批交替(DeepSeek 的 dual micro-batch overlap:一个微批算注意力和专家时,另一个做 all-to-all)在 32 卡以上把吞吐再提高 37%–60%;但在 16 卡时反而变差——那里路由专家是 memory-bound 的,切成两半要把专家权重读两遍。

DeepSeek 公开的部署:技术报告里 decode 的最小部署单元是 40 个节点、320 块卡(注意力 TP4 + DP80,MoE 部分 EP320,每块卡一个专家,另有 64 块卡放冗余专家和共享专家);2025 年 2 月的推理系统概述里,decode 是每个部署单元 18 个节点、144 块卡(路由专家 EP144,MLA 与共享专家 DP144,加 32 个冗余专家),prefill 是 4 个节点、32 块卡(EP32)。纸上的结论与”越大越好、但必须掩盖通信”一致。

3.3 负载不均:最忙的卡决定这一层的时间

专家并行里,每层都要等所有设备算完、all-to-all 完成,才能进入下一层。最忙的设备决定这一层的时间。不均衡有两个来源(代码 03 演示 1):

随机波动(“球放进盒子”)。即使每个专家的热度完全相同,随机分配也不均匀:每个专家平均分到 1 个 token 时,最忙的专家是平均的 4.9 倍;平均 1,024 个时也多 9%。均值越小,相对波动越大(泊松分布的标准差约为 均值\sqrt{\text{均值}})。按”每块卡 8 个专家”汇总后波动小得多(1.69 → 1.02)。

系统性的冷热不均。热度偏斜时(本讲取 Zipf 分布 s=0.6s = 0.6,最热的专家是平均的 12.2 倍),最忙的专家稳定地是平均的 10.5 倍左右,每卡汇总后约 2.1 倍,批再大也降不下来——这只能靠放置解决。

代码 03 演示 3 把不均衡代入 decode 的时间模型(最忙的卡的专家要多算、网卡要多收发):

EP 组不均衡(连续放置)不均衡(冗余专家)每卡吞吐:均衡 / 连续放置 / 冗余专家
161.581.091,006 / 1,006 / 1,006
322.081.062,910 / 1,881 / 2,910
643.401.073,718 / 1,166 / 3,496
1285.551.093,717 / 717 / 3,438
  • 连续放置(每块卡固定放一段编号连续的专家)在卡多时非常糟:每块卡上的专家越少,热门专家越没法和冷门专家”互相平均”。128 卡时每卡吞吐从约 3,700 掉到约 700,扩大专家并行组的收益全被吃掉。16 卡时几乎没影响:那里瓶颈是读专家权重,多来几个 token 是”免费”的。
  • 冗余专家(DeepSeek 的 EPLB、MegaScale-Infer §6 都是这个思路):每块卡多留一两个”专家位”,给热门专家放副本,token 平均分给副本;再用 LPT 贪心(负载大的副本先放到最闲的卡上)分配。不均衡压到 1.1 倍以内,吞吐恢复到均衡时的 92%–100%。MegaScale-Infer 把它写成优化问题:最小化 max⁡jCj\max_j C_j,Cj=∑ixi,jmax⁡(ai,K)C_j = \sum_i x_{i,j} \max(a_i, K),其中 aia_i 是专家 ii 的活跃 token 代价,KK 是冷门专家的最低代价(读权重的固定成本)——max⁡(ai,K)\max(a_i, K) 正好对应”memory-bound 时多几个 token 不要钱”。

3.4 decode 稳定,prefill 波动

MegaScale-Infer §8 的部署经验:decode 时各专家的负载在不同批之间相当稳定,prefill 时波动大。代码 03 演示 4 用”每个请求有自己的话题”的模型解释了原因:

单个请求的话题批的构成连续放置按历史热度放置按本批负载放置
很集中decode:4,096 个请求各 1 个 token2.091.061.00
很集中prefill:4 个请求各 1,024 个 token2.151.471.00
接近长期平均prefill2.071.111.01

decode 的批由几千个请求各贡献 1 个 token,请求之间的话题差异被平均掉,批的负载几乎就是长期平均热度,静态方案就够了;prefill 的批只有几个请求、每个贡献上千个 token,批的负载取决于这几个请求的话题。所以 MegaScale-Infer 对 decode 用静态或周期性的专家放置、对 prefill 更频繁地调整;DeepSeek-V3 技术报告也写到根据线上统计周期性地(例如每 10 分钟)调整冗余专家。调整不能太频繁:搬一个专家就是搬几十 MB 的权重。


4. MegaScale-Infer:把注意力和专家拆开 · MegaScale-Infer

4.1 动机:一层里有两种性格的算子

MegaScale-Infer 论文图 1 用三张示意图讲清了动机:

  • 稠密模型:decode 时注意力要读每个请求各自的 KV,批再大也是 memory-bound,GPU 利用率低;FFN 的权重被所有 token 共享,批大了就 compute-bound。显存和延迟允许的最大批下,FFN 能把算力用满。
  • MoE:同样的最大批,每个专家只分到 k/Ek/E,FFN 也变成 memory-bound——两部分都低效。
  • MegaScale-Infer:把注意力和 FFN 放到不同的 GPU 上。注意力部分按数据并行复制很多份,每份服务一批请求;所有副本的 token 汇到专家那边,专家的批变大、重新 compute-bound。

这就是注意力 / 专家解耦(论文叫 disaggregated expert parallelism)。它带来两个好处:两部分可以独立扩展、各用各的并行方式;还可以用不同的硬件(注意力用显存带宽便宜的卡,专家用算力便宜的卡)。它和第 7 讲的 prefill / decode 解耦是同一个思想在模型内部的延伸;MegaScale-Infer 本身也把 prefill 和 decode 放在不同的集群上,论文只讨论 decode。

4.2 架构

一个 MegaScale-Infer 实例(论文图 3):

  • 注意力节点(nan_a 个):每个节点 1–8 块卡,存注意力的权重(QKV 投影、输出投影)和 KV cache;节点内张量并行(tpatp_a),节点之间是数据并行的副本;门控也在这里算。
  • 专家节点(EE 个):每个节点 1–8 块卡,存一个专家在所有层的权重;节点内张量并行(tpetp_e)。所有专家节点组成一个专家并行组。
  • 两类节点之间每层两次通信:注意力 → 专家(发 token),专家 → 注意力(发结果)。因为两边的节点数和并行度可以任意选,原来的 all-to-all 变成了 M2N 和 N2M。

两个新问题随之而来:一个微批在注意力那边算的时候专家闲着,反过来也一样,通信时两边都闲着(§4.3);M2N 通信要足够快、足够稳(§4.6)。

4.3 ping-pong 流水线

把一批请求切成 mm 个微批,让它们在注意力节点和专家节点之间”你来我往”:微批 1 在专家那边算的时候,注意力节点在算微批 2。代码 04 演示 1 画出了 2 层、Ta=Te=1T_a = T_e = 1、Tc=0.5T_c = 0.5 时的时间线(第 1 层的微批记作 1 2 3 4,第 2 层记作 a b c d):

m = 1:总时间 6.00,两边利用率 33%
    注意力节点 |1111········aaaa········|
    A→E 链路  |····11··········aa······|
    专家节点   |······1111········aaaa··|
    E→A 链路  |··········11··········aa|

m = 4:总时间 10.00,两边利用率 80%
    注意力节点 |1111222233334444aaaabbbbccccdddd········|
    A→E 链路  |····11··22··33··44··aa··bb··cc··dd······|
    专家节点   |······1111222233334444aaaabbbbccccdddd··|
    E→A 链路  |··········11··22··33··44··aa··bb··cc··dd|

论文给出三个条件:

Ta≈Te(1),Tc<Tf(2),m⋅Tf≥2(Tf+Tc)(3)T_a \approx T_e \quad (1), \qquad T_c < T_f \quad (2), \qquad m \cdot T_f \ge 2(T_f + T_c) \quad (3)

(1) 让两边的计算时间接近,否则快的一边总在等;(2) 通信比计算快,才可能被计算盖住;(3) 一层里 mm 个微批的计算总时间,要盖住一个微批的一次往返(注意力 → 发送 → 专家 → 发回)。由 (3) 得到最少微批数

m≥2(1+TcTf)m \ge 2\left(1 + \frac{T_c}{T_f}\right)

通信快(Tc<Tf/2T_c < T_f/2)时 3 个微批就够,慢一些要 4 个。条件成立时,一次迭代的时间是

Ttotal=(Ta+Te+2Tc)+Tf⋅(mL−1)(5)T_{\text{total}} = (T_a + T_e + 2T_c) + T_f \cdot (mL - 1) \quad (5)

公式 (5) 什么时候不准? 条件不成立时它会严重低估:(Ta,Te,Tc)=(1,1,0.3)(T_a, T_e, T_c) = (1, 1, 0.3)、L=56L = 56、m=1m = 1 时公式给 57.6,模拟是 145.6(代码 04 演示 1)。我推导了一个处处成立的闭式公式:一个微批走完一层要 R=Ta+Te+2TcR = T_a + T_e + 2T_c;每种资源每层要为 mm 个微批各服务一次。稳定后每层的周期是

C=max⁡(mTa, mTe, mTc, R)C = \max(m T_a,\ m T_e,\ m T_c,\ R)

要么被最忙的资源卡住(流水线满),要么被单个微批的往返卡住(微批太少)。总时间

T=R+(L−1) C+(m−1)max⁡(Ta,Te,Tc)T = R + (L-1)\,C + (m-1)\max(T_a, T_e, T_c)

它在 2000 组随机参数上与离散事件模拟的最大差是 4×10−144 \times 10^{-14};三个条件成立时 C=mTfC = m T_f,就化简为公式 (5)。顺带看出:论文的条件 (3) 用 2(Tf+Tc)2(T_f + T_c) 代替 RR,在 Ta≠TeT_a \ne T_e 时偏保守(练习 4)。

微批数的影响(代码 04 演示 2,论文图 14):mm 从 1 到 2,每 GPU 吞吐约翻倍(论文实测 1.9 倍);从 2 到 3 再涨 37%–49%(论文:1.10–1.38 倍);m=4m = 4 是否还有用取决于 Tc/TfT_c / T_f——Mixtral 的 Tc/Tf≈0.45T_c/T_f \approx 0.45,3 个就够;DBRX 和 Scaled-MoE 的 TcT_c 接近 TfT_f,第 4 个微批又带来 20%–26%。代价是 TBT 变长(每层都要等 mm 个微批轮一遍),所以论文把 mm 的上限设为 4。

4.4 部署方案搜索

给定模型、硬件、负载和 TBT 的 SLO,要选:注意力和专家各自的张量并行度 tpatp_a、tpetp_e,注意力节点数 nan_a,微批数 mm,以及批大小。论文的算法 1 是一个小枚举:

for tpe in {1, …, Me}:
  for tpa in {1, …, Ma}:
    if 显存放得下注意力权重和一个专家:
      na ← balance(tpa, tpe)               # 让 Ta ≈ Te
      for m in {3, …, Nm}:
        二分查找满足 SLO 的最大批 B,模拟得到每单位价格的吞吐
        保留最好的方案

性能模型:TaT_a 和 TeT_e 都近似是批大小的线性函数(Ta=k1ba+k2T_a = k_1 b_a + k_2、Te=k3be+k4T_e = k_3 b_e + k_4,系数由 profiling 得到),于是 na=k1Ek3Kn_a = \frac{k_1 E}{k_3 K} 能让两者平衡;TcT_c 取发送端和接收端中较慢的一个(公式 6);约束还包括 TBT ≤ SLO 和注意力节点的显存(KV + 注意力权重 < tpa×tp_a \times 显存)。目标是每单位价格的吞吐:

tpuc=B/Ttotaltpana⋅Costa+tpeE⋅Coste\text{tpuc} = \frac{B / T_{\text{total}}}{tp_a n_a \cdot \text{Cost}_a + tp_e E \cdot \text{Cost}_e}

注意力节点数的取舍(代码 04 演示 3,论文图 15):nan_a 小时专家每个微批只分到几十到一百多个 token,专家节点在等注意力,TBT 几乎不变、每 GPU 吞吐快速上涨;nan_a 大到瓶颈转到专家后,TeT_e 随 token 数线性增长、TBT 变长,注意力节点开始闲着,每 GPU 吞吐下降。我们的模型里峰值在 na=4n_a = 4,比 Ta=TeT_a = T_e 的平衡点(nan_a 在 2 和 3 之间)稍大:每加一个注意力节点只多 2 块卡,16 块专家卡是固定成本。论文的峰值恰好在平衡点(na=8n_a = 8)。

同构集群上的提升(代码 04 演示 4:A800、TBT ≤ 150 ms、平均上下文 651):

模型共置(节点内 TP8,必要时跨节点 PP)每 GPU 吞吐搜出的解耦方案 (tpa,tpe,na,m,ba)(tp_a, tp_e, n_a, m, b_a)每 GPU 吞吐提升
Mixtral-8x22BTP8,每微批 1973 序列1,853(4, 1, 1, 3, 620),12 卡2,1061.14×
DBRXTP8,每微批 2658 序列2,215(2, 1, 3, 4, 321),22 卡2,4371.10×
Scaled-MoETP8×PP2,每微批 1998 序列1,665(2, 1, 7, 4, 252),46 卡1,8851.13×

纸上的提升只有 10%–14%,比论文小得多(论文报告:相对 TensorRT-LLM 1.28–1.90 倍,相对 vLLM 2.56–7.11 倍)。原因是这里的共置基线是理想化的:显存允许的最大批下,每个专家每个微批已经分到 250–665 个 token,MoE 部分本来就接近 compute-bound;真实引擎在小批 GEMM、注意力内核和跨节点通信上的低效不在模型里。解耦在纸上省下的是两样东西:共置时注意力(每层 400–560 µs)和 MoE 只能串行;共置时 TP8 的 all-reduce 每层要 590–770 µs,而解耦后 tpatp_a = 2–4、tpetp_e = 1,张量并行的通信少得多,换成了能被 ping-pong 掩盖的 M2N 通信。

4.5 异构部署

论文表 3 列出了几种 GPU 的价格(按 L20 归一化)和每单位价格能买到的资源:

GPU价格显存带宽算力(BF16)每单位价格:GB / GB/s / TFLOPS
H201.8596 GB4096 GB/s14852 / 2214 / 80
L40S1.0848 GB864 GB/s36244 / 800 / 335
A8002.2680 GB2039 GB/s31235 / 902 / 138
H8005.2880 GB3430 GB/s98915 / 650 / 187

H20 每单位价格的显存带宽是 L40S 的 2.8 倍,适合 memory-bound、要存 KV 的注意力;L40S 每单位价格的算力是 H20 的 4 倍多,适合能凑够大批的专家。共置部署只能选一种卡,必然有一部分在用不划算的资源;解耦后各取所长。代码 04 演示 5 的搜索结果:注意力 H20 + 专家 L40S 在三个模型上都是每单位价格吞吐最高的方案,是最好的共置部署的 1.98–2.63 倍(论文报告:相对 H20 上的 TensorRT-LLM 1.86 倍、vLLM 3.24 倍);反过来(注意力 L40S + 专家 H20)明显更差。

和 §4.4 对比可以看出:在纸上 GPU 里,解耦的大部分收益来自”能用不同的卡”,而不是同构集群里的流水线。论文还报告异构部署让每瓦吞吐提高 1.80 倍(decode)。

4.6 M2N 通信库

token 的去向是门控在运行时决定的,所以只能用点对点通信。论文先用 NCCL 和一个 CPU 端的 RDMA 基准(perftest)对比”一个发送者给 N 个接收者各发 128 KB”,发现 NCCL 的中位延迟明显更高、P99 延迟在接收者多时急剧恶化。原因:

  • NCCL 的网络传输要先把数据从 GPU 显存复制到 CPU 代理,再由 CPU 发出;
  • 点对点的分组操作每批最多 8 个,接收者一多就要分很多批;
  • 作为通用库,每次分组操作都有准备、启动、校验的开销;
  • GPU 同步和显存访问是常见的不稳定来源。

MegaScale-Infer 自己写了一个 M2N 库:发送端用 CUDA event 等前面的内核写好数据,用 CUDA driver 操作阻塞后续的流,CPU 端的通信库直接用 RDMA write with immediate 把预先注册好的 GPU 显存发出去,轮询完成队列,最后更新一个共享内存标志解除阻塞;接收端类似,用 GDRCopy 做一次 flush 保证 GPU 看到的数据一致。没有 GPU 到 CPU 的复制、没有分组初始化、没有 GPU 同步。还有两项针对流量的优化:把 ACK 包放进高优先级队列(双向的 ping-pong 流量里 ACK 排队会拖慢发送端)、调整拥塞控制(各接收者数据量不均时)。

与 DeepSeek 的 DeepEP 相比:DeepEP 用 GPU 直接发起通信(不经过 CPU 代理),并行度高、包很小时吞吐更好,但要占用 GPU 的 SM,还要用 PTX 指令小心地避免和计算内核争抢 L2 缓存;MegaScale-Infer 用 CPU 发起,不占 SM——每对收发之间几百 KB 的数据量下,单线程 CPU 足以打满带宽。论文的估算:Mixtral 8x22B、微批 128、tpa=2tp_a = 2 时,每块注意力 GPU 给每块专家 GPU 平均发 128 × 2/8 × 6144 × 2 / 2 = 196,608 字节。

4.7 评测与论文卡片

论文在 Mixtral-8x22B(141B)、DBRX(132B)和一个 317B 的 Scaled-MoE 上评测,负载来自生产(输入中位数 571、输出中位数 159 个 token),TBT 的 SLO 为 150 ms(论文报告):

  • 同构(Ampere 80GB,NVLink 400 GB/s,每卡 200 Gb/s 网卡):每 GPU decode 吞吐相对 vLLM 和 TensorRT-LLM 分别最高 2.56 倍和 1.28 倍(Mixtral / DBRX)、7.11 倍和 1.90 倍(Scaled-MoE);TBT 与基线相当;把 prefill 算进去的端到端吞吐最高 1.18 倍。
  • 异构(H20 + L40S):每单位价格的 decode 吞吐相对 H20 上的 vLLM / TensorRT-LLM 最高 3.24 倍 / 1.86 倍;端到端 1.66 倍;每瓦吞吐 1.80 倍(decode)、1.72 倍(端到端)。
  • M2N:256 KB 时中位延迟降低 68.2%、P99 降低 92.9%、吞吐 4.2 倍;增加收发数时尾延迟降低 54.7%–96.9%。
  • 消融:只做解耦、通信仍用 NCCL,就比共置高最多 4.66 倍;换成 M2N 库再高最多 1.53 倍。
  • 部署:已在近 1 万块 GPU 的生产集群上运行,异构部署下服务同样流量的成本降低 1.5–2.0 倍。
项目内容
论文MegaScale-Infer: Serving Mixture-of-Experts at Scale with Disaggregated Expert Parallelism,Zhu 等,2025(字节跳动 Seed、北京大学)
一句话把每层的注意力和专家拆到不同的 GPU 上,注意力数据并行、专家专家并行,用 ping-pong 流水线和专用的 M2N 通信让两边都忙起来
问题MoE 的稀疏让 FFN 在 decode 时也变成 memory-bound;共置部署下注意力和 FFN 都用不满 GPU
关键洞察注意力和专家的资源需求不同:前者要显存容量和带宽,后者要算力和大批;拆开后可以独立扩展、选不同的硬件
机制解耦的专家并行;ping-pong 流水线(三个条件、最少 3–4 个微批);基于性能模型的部署搜索;异构部署;CPU 发起、零复制的 M2N 通信库;基于热度的冗余专家
评测每 GPU decode 吞吐相对 TensorRT-LLM 最高 1.90×;异构下每单位价格最高 1.86×(相对 TensorRT-LLM)、3.24×(相对 vLLM)(论文报告)
局限与假设每层两次跨节点通信,单个请求的延迟不会降低(只能靠流水线掩盖);需要 Tc<TfT_c < T_f,依赖高速网络;方案依赖负载(序列长度、批大小),负载变化要重新规划;只优化 decode
与本课的关系第 7 讲”按阶段解耦”推进到”按模块解耦”;纸上 GPU 显示它的大部分收益来自异构硬件(§4.5)

5. Helix:长上下文 decode 的分片 · Helix Parallelism

5.1 问题:百万 token 的 KV 和张量并行的墙

当上下文长到几十万、上百万 token 时,decode 有两个瓶颈:

  1. 读 KV:KV cache 随上下文长度和批大小线性增长,每步都要全部读一遍;
  2. 读 FFN 权重:每生成一个 token 都要读一遍 FFN 权重,批小时这部分无法摊薄。

张量并行能同时缓解两者,但它在 KV 头数处撞墙。GQA 把 QQ 个 query 头合并成 KK 个 KV 头(KK 常常 ≤ 8),MLA 甚至只有一份所有头共用的潜向量。TP 宽度 ≤ KK 时,KV 被干净地切开;TP > KK 时,每个分片都得保存它的 query 头需要的那个完整 KV 头——再加卡,每块卡的 KV 也不再减少(第 6 讲代码 02 演示 2 里我们见过这种复制)。

代码 05 演示 2 复现了 Helix 论文图 1(稠密模型,批 8,Q=128Q = 128、K=8K = 8,FP4,8 TB/s,上下文 100 万):

TP 宽度1248163264
每层读 KV1.02 ms512 µs256 µs128 µs128 µs128 µs128 µs
每层读权重237 µs118 µs59 µs30 µs15 µs7.6 µs3.9 µs

两点:上下文 100 万时读 KV 是读权重的 4 倍多,长上下文 decode 的瓶颈是 KV;TP 超过 8 以后读 KV 的时间停在 128 µs。对 MLA(DeepSeek-R1)更极端:一份潜向量,TP 在 KV 上完全无效,批 8、上下文 100 万时每块卡每步读 KV 要 17.6 ms,用多少路 TP 都一样。

此前的 Medha 用 KV 并行(KVP:沿序列方向把一个请求的 KV 切到 NN 块卡上)解决注意力部分,但注意力之后把结果收拢到固定的一组 TP 卡(比如 8 块)上做 FFN——KVP 再宽,FFN 的权重读取也帮不上。

5.2 Helix 的注意力阶段

Helix 的核心想法:同一组 GPU,在一层之内的注意力和 FFN 两个阶段用不同的切法。

注意力阶段把 N=KVP×TPAN = \text{KVP} \times \text{TPA} 块卡排成二维网格(TPA ≤ KK):TPA 方向按头切,KVP 方向按序列切 KV。步骤(代码 05 演示 1 在 numpy 上真实实现):

  1. 为了避免先 all-gather query,每个 KVP rank 都对整批 token 计算自己那组头的完整 QKV 投影(代价:QKV 投影的权重在 KVP 方向上被重复读取);
  2. 每块卡对自己那段 KV 跑 FlashAttention,得到部分输出 OO 和每个 token 每个头一个 log-sum-exp 标量——这就是第 4 讲的注意力状态;
  3. 同一头组内的 KVP 块卡做一次 all-to-all(沿 query 头的维度交换碎片);
  4. 每块卡把收到的碎片按 LSE 重新缩放、相加(第 4 讲的 ⊕),得到它负责的那几个头精确的注意力输出。

代码 05 演示 1 的结果:切成 (2, 2)、(4, 2)、(2, 4)、(1, 8) 都与不切的结果只差 10−710^{-7} 量级(浮点舍入);每块卡 all-to-all 发送的字节数在上下文 256、1024、4096 时完全相同——交换的是 B×B \times 头数 ×\times 头维度的部分输出和每个头一个 LSE,和上下文长度无关(论文 §2.1.2)。上下文到 100 万时,这一点决定了 KV 并行能不能用。

all-to-all 结束时,每块卡持有 H/NH / N 个隐藏维度的、整批的、已归一化的注意力输出——正好是后面做 NN 路张量并行所需要的切分。

新 token 的 KV 追加到哪里(§2.3):如果总是追加到最后一个 rank,它会先被填满。Helix 交错追加:前 16 个 decode 步的 KV 写到 rank 0,接下来 16 个写到 rank 1……轮流来。代码 05 演示 1:上下文 1000 时 4 个 rank 分别存 256、256、248、240 个 token,最多差一个块。

5.3 Helix 的 FFN 阶段

注意力之后,同样的 NN 块卡立刻”重新编组”:

  • 输出投影:NN 路张量并行(每块卡持有 WOW_O 的 H/NH/N 行),一次 all-reduce 得到完整的 B×HB \times H 输出;
  • 稠密 FFN:TPF=N\text{TPF} = N 路张量并行,所有卡一起摊薄权重读取;
  • MoE FFN:把 NN 块卡重新排成 TPF×EP\text{TPF} \times \text{EP} 的网格,按专家的大小和个数选择两者的比例,token 路由到专家组,组内张量并行。

于是 FFN 的张量并行宽度可以远超 KV 头数,又不会在注意力阶段重新引入 KV 复制。论文 §5 指出,短上下文(< 4K)时 Helix 退化为”注意力数据并行 + FFN 张量并行”,这已经是推理框架里常见的做法(例如 SGLang 对 DeepSeek 模型的 DP attention)。

HOP-B(批维度上的通信 / 计算重叠):第一个请求的注意力一算完就开始它的 all-to-all,同时算第二个请求的注意力。论文图 3 的例子:8 个请求,每个注意力 2 个单位、通信 1.2 个单位,不重叠时 16 + 9.6 = 25.6;重叠后只有最后一个请求的通信露在外面,16 + 1.2 = 17.2(论文图里写 17)。

5.4 评测与论文卡片

论文用 NVIDIA 内部的高保真模拟器模拟 GB200 NVL72(FP4 的权重、KV 和运算),上下文 100 万 token,在 1–64 块卡上枚举了十万多种配置(TP、PP、EP、vanilla KVP 与批大小),画出”每卡吞吐 vs 每用户交互性”的帕累托前沿(论文报告,数字都相对基线归一化):

  • DeepSeek-R1(671B MoE,MLA):交互性最高提高 1.5 倍,同样的延迟预算下并发用户数(以及每卡吞吐)最高 32 倍;
  • Llama-405B(稠密,128 个 query 头、8 个 KV 头):交互性最高提高 1.13 倍,吞吐和批容量 4 倍;
  • 关掉 HOP-B:Llama-405B 的每用户 token/s 最多下降 12%,DeepSeek-R1 只下降约 1%(它的 all-to-all 只占 TTL 的约 1%)。
项目内容
论文Helix Parallelism: Rethinking Sharding Strategies for Interactive Multi-Million-Token LLM Decoding,Bhatia 等,2025(NVIDIA)
一句话一层之内,注意力用”KV 并行 × 按头 TP”切、FFN 用”N 路 TP(或 TP × EP)“切,同一组 GPU 两种切法,中间用一次与上下文长度无关的 all-to-all 衔接
问题百万 token 上下文的 decode 同时受限于读 KV 和读 FFN 权重;TP 超过 KV 头数后 KV 被复制,KV 并行(Medha)又没有把多出来的卡用于 FFN
关键洞察注意力的自然切法是”按序列切 KV”,FFN 的自然切法是”按权重切”;两者可以在同一组卡上交替使用,部分注意力结果用 LSE 精确合并
机制KVP × TPA 的注意力(每个 KVP rank 独立计算 QKV);一次 all-to-all + LSE 重缩放;TP = N 的输出投影和 FFN(MoE 为 TPF × EP);交错追加 KV;HOP-B 重叠
评测模拟的 GB200 NVL72 上,DeepSeek-R1 交互性 1.5×、批容量 32×;Llama-405B 1.13×、4×(论文报告)
局限与假设结果来自模拟器、相对值;依赖大 NVLink 域(NVL72)和低延迟的集合通信;每个 KVP rank 重复读 QKV 权重;没有讨论 prefill 和 KV 的写入 / 驱逐
与本课的关系第 4 讲 split-KV 的 ⊕ 从”线程块之间”扩展到”GPU 之间”;第 6 讲”不同模块可以用不同的并行方式”在一层之内的体现;和 MegaScale-Infer 一样,核心是让注意力和 FFN 各用各的最优方案

5.5 我们的复现:TTL 与吞吐的前沿

代码 05 演示 3 用论文附录 A 的读取公式加上本课程的 roofline 习惯,搜索 Llama-405B、上下文 100 万、1–64 块 GB200 上三种方式的最优配置:

TTL 上限纯 TPMedha(KVP,FFN 绑在 TP 上)Helix
4 ms4.1(批 1,64 卡)达不到27.4(批 7,8×8)
6 ms11.8(批 1,16 卡)8.1(批 3,8×8)33.9(批 13,8×8)
8 ms17.3(批 1,8 卡)18.1(批 9,8×8)36.6(批 18,8×8)
15 ms33.6(批 4,8 卡)33.6(批 4,8 卡)41.3(批 39,8×8)
30 ms39.9(批 8,8 卡)40.1(批 17,8×2)44.3(批 39,8×4)

(单元格:每卡每秒生成的 token 数;“8×8”是 TPA × KVP。)最低 TTL:Helix 2.29 ms,纯 TP 3.83 ms,Medha 5.40 ms;最大批:Helix 80,Medha 70,纯 TP 10。

  • 纯 TP 的批最多 10:每个序列 100 万 token 的 KV 在 FP4 下约 129 GB,TP 超过 8 后每块卡要存每个序列 16 GB 的 KV。KVP 把它切到 64 块卡上,批可以到 80。
  • TTL 上限紧时差别最大:6–8 ms 时 Helix 的每卡吞吐是纯 TP 的 2.1–2.9 倍;放宽到 20–30 ms,差距缩小到 11%–14%。
  • Medha 能扩大批,却降不下延迟:FFN 只用 8 块卡读权重(每层 1.3 GB 的 FP4 权重,126 层合计约 3.2 ms),比纯 TP 的最低 TTL 还慢。

这些数字依赖一个关键假设:NVL72 里一次集合通信的固定延迟取 3 µs(经过 NVSwitch 的 one-shot / 交换机内归约,延迟基本不随卡数增长)。我最初沿用第 6 讲的 ring 模型(每次 all-reduce 的固定延迟是 2(k−1)α2(k-1)\alpha),64 卡时每次 all-reduce 要 0.13 ms,所有大 NN 的配置都失去意义,三种方式的最优配置全部退回到 8 卡、结果完全相同。Helix 是为”大 NVLink 域 + 低延迟集合通信”的硬件设计的,换到普通集群上结论不成立。

我们的提升(最低 TTL 低 1.67 倍、最大批 8 倍)比论文报告的(1.13 倍、4 倍)大,可能因为本讲把激活按 BF16 传、没有计 TP 组内的计算 / 通信重叠、FP4 算力等参数与 NVIDIA 的模拟器不同。HOP-B 在我们的模型里带来 0%–8%(代码 05 演示 4):批 1 时 all-to-all 占 TTL 的 16.5%,但没有”下一个请求”可以重叠;批 8 时提升最大(7.7%);批再大,读 KV 的时间随批线性增长,通信占比降到 3% 左右。


6. FlashMoE:一个内核里的分布式 MoE · FlashMoE

分布式 MoE 层的典型实现是一串内核:门控 → 排列 token → all-to-all(dispatch)→ 专家 GEMM → all-to-all(combine)→ 恢复顺序、加权求和。FlashMoE 指出三类低效:

  1. 同步的集合通信:all-to-all 必须所有 GPU 都到齐才能完成,一块慢卡拖住所有卡(straggler);通信可占 MoE 运行时间的 68%。
  2. 内核启动开销:论文用 Nsight Systems 数了一个 MoE 层(2 块 A100,每块 32 个专家)的 GPU 操作数——COMET 33 个、Megatron-LM CUTLASS 85 个、Megatron-LM TE 261 个、Megatron-LM + DeepEP 432 个、DeepSpeedMoE 550 个。每次启动都带来不确定的开始时间、额外的同步点和全局显存的往返;CUDA Graph 可以缓解启动开销,但与 MoE 动态的路由(形状每次不同)不兼容。
  3. 补零的载荷:对称的集合通信 API 要求各设备的缓冲区大小相同,token 不均匀地路由时就要补零;空数据浪费带宽,有的实现还会在零矩阵上做无用的计算。代码 01 演示 3 算了一笔账:按”最坏情况”留缓冲区(永不溢出)时,4 个设备的有效载荷只有 54%–57%,上限是 k/(D−1)k / (D - 1)——DeepSeek-V3 在 64 块卡上做专家并行时只有 8/63 ≈ 13%;按容量因子 1.25 留缓冲区有效载荷 87%–91%,但话题聚集时会溢出。

FlashMoE 的设计:整个分布式 MoE 层是一个常驻的(persistent)GPU 内核。

  • actor 模型:每块 GPU 上 N 个线程块中,N − 1 个是 Processor(做 GEMM、逐元素运算和发送 tile),1 个是”操作系统”块:其中 1 个 warp 当 Scheduler(把就绪的任务分给空闲的 Processor),3 个 warp 当 Subscriber(解码从其他 GPU 收到的 tile 包,生成任务)。
  • tile 级的任务:输入按 (128, 64) 的 tile 切开,每个 tile 是一个任务(GEMM0、GEMM1、combine),依赖关系在 tile 粒度上满足后就执行,计算和通信自然地流水起来。
  • 设备端发起的单边通信:用 NVSHMEM 建立全局地址空间,Processor 直接把 tile 写到远端 GPU(DMA / RDMA),不经过 CPU、不需要所有 GPU 同步;只发真实的 token,不补零(payload efficiency)。

结果(论文报告;8 块 H100、最多 128 个专家、16K token 的前向):GPU 利用率最高 9 倍、延迟低 6 倍、吞吐 5.7 倍、重叠效率 4 倍——而且 FlashMoE 用 FP32,基线用 FP16。SM 平均利用率 93.17%,DeepEP 约 14%。

项目内容
论文FlashMoE: Fast Distributed MoE in a Single Kernel,Aimuyo、Oh、Singh,NeurIPS 2025(Cornell)
一句话把门控、dispatch、专家计算、combine 全部融合进一个常驻内核,用 tile 级的任务和设备发起的单边通信代替 CPU 调度的内核串和同步的 all-to-all
问题分布式 MoE 层要启动几十到几百个内核,同步的 all-to-all 受慢卡拖累,对称 API 迫使补零
关键洞察MoE 的依赖关系可以在 tile 粒度上满足;GPU 内部可以自己调度(actor 模型),通信可以由 GPU 直接发起
机制常驻内核;Processor / Scheduler / Subscriber 三种 actor;(128, 64) 的 tile;NVSHMEM 单边传输与对称张量布局;不补零
评测8×H100:GPU 利用率最高 9×、延迟 6× 更低、吞吐 5.7×(论文报告)
局限与假设评测是单节点、大批 token 的前向(偏训练场景),用 FP32;多节点、decode 小批的效果需要另外验证;巨型内核的开发和调试难度很高
与本课的关系第 4 讲”融合内核省掉显存往返”的思想扩展到”融合计算与通信”;推理时小批 decode 里启动开销和补零的占比更大

7. 之后的发展 · Since then

(这一节是我的补充,不是原课内容。)

  • 开源的专家并行基础设施:DeepSeek 在 2025 年 2 月开源了 DeepEP(专家并行的 all-to-all 通信库,有高吞吐和低延迟两种内核)、EPLB(冗余专家的负载均衡算法)和 DeepGEMM(FP8 的分组矩阵乘法),并公开了 V3 / R1 推理系统的部署方式(§3.2)。SGLang、vLLM、TensorRT-LLM 随后都支持了”prefill / decode 解耦 + 大规模专家并行”的 DeepSeek 式部署。
  • 注意力 / FFN 解耦被更多系统采用:例如阶跃星辰的 Step-3 把它和模型结构一起设计(论文称 attention-FFN disaggregation,AFD);它与 prefill / decode 解耦叠加,形成”阶段 × 模块”两个维度的解耦。
  • 更大的 NVLink 域:GB200 NVL72 让 72 块卡共享一个 NVLink 域,专家并行的 all-to-all 可以完全走 NVLink(“wide EP”),Helix 这类依赖低延迟集合通信的方案也建立在它之上。
  • 更细、更稀疏的专家:新一代开源 MoE(DeepSeek-V3、Qwen3、Kimi K2 等)普遍采用几百个细粒度专家加共享专家,稀疏度更高,§2.4 的”每个专家的批”问题更突出;低精度(FP8、FP4)的专家权重既省显存,也降低 memory-bound 部分的时间。
  • 长上下文与稀疏注意力:Helix 的未来工作提到原生稀疏注意力(NSA)这类方法减少 KV 的读取量但不减少容量;第 10 讲的 DeepSeek-V3.2 稀疏注意力就是这条路线。

8. 小结 · Summary

  1. MoE 用门控为每个 token 从 EE 个专家里选 kk 个,计算按激活参数算、显存按总参数算(DeepSeek-V3:激活 37.6B、总 671B)。训练时用均衡损失和容量因子控制负载;推理引擎是 dropless 的,因为丢 token 会让一个请求的输出取决于同批的邻居(代码 01)。
  2. MoE 的 decode 很尴尬:批稍大就碰到几乎所有专家(读总参数),每个专家却只分到整批的 k/Ek/E(算激活参数)。Mixtral-8x22B 在批 4–256 时比稠密的 Llama-3-70B 还慢 9%–40%;每个专家要分到 150–260 个 token 才 compute-bound,DeepSeek-V3 在传统部署下只分到个位数(代码 02)。所有系统设计都在想办法让每个专家看到更大的批。
  3. 专家并行:两次 all-to-all 中间夹一次汇聚后的批量计算(Shazeer 的混合并行)。DeepSeek 式的”注意力 DP + 专家 EP”在纸上把每卡吞吐从 1,253(16 卡)提高到 2,354(64 卡以上),再用双微批交替掩盖通信到约 3,700;但负载偏斜时最忙的卡决定一切(128 卡、连续放置时掉到 717),冗余专家把它拉回 3,438;decode 的负载稳定、prefill 的需要频繁调整(代码 03)。
  4. MegaScale-Infer:注意力和专家放到不同的 GPU 上,ping-pong 流水线让两边都忙(条件:Ta≈TeT_a \approx T_e、Tc<TfT_c < T_f、m≥2(1+Tc/Tf)m \ge 2(1 + T_c/T_f));公式 (5) 只在条件成立时准,本讲的闭式公式处处成立。纸上同构部署只提升 10%–14%,异构(H20 注意力 + L40S 专家)每单位价格提升 1.98–2.63 倍;论文报告相对 TensorRT-LLM 最高 1.90 倍(代码 04)。
  5. Helix:长上下文 decode 中张量并行在 KV 头数处撞墙;Helix 让注意力按 KV 切(部分结果 + LSE + 一次与上下文长度无关的 all-to-all,结果精确),同一组卡在 FFN 阶段换成 NN 路张量并行。在大 NVLink 域、低延迟集合通信的前提下,Llama-405B 百万上下文的最大批从 10 到 80、最低 TTL 从 3.83 到 2.29 ms(代码 05)。
  6. FlashMoE:把分布式 MoE 层融合进一个常驻内核,用 tile 级任务和设备发起的单边通信省掉几百次内核启动和对称通信的补零。

9. 常见误区 · Common pitfalls

  1. “MoE 的推理成本按激活参数算。” 只有计算是。显存按总参数算,而 decode 时只要批稍大,读的字节也接近总参数(代码 02 演示 2–3)。
  2. “MoE 的 decode 一定比同等质量的稠密模型快。” 批 1 时快;中等批时可能更慢(Mixtral vs Llama-3-70B 慢 9%–40%);很大的批时又快。要看你的批落在哪一段(代码 02 演示 3)。
  3. “推理时也可以像训练那样用容量因子。” 会丢 token、改变输出,而且输出取决于同批的其他请求,结果不可复现(代码 01 演示 2)。
  4. “专家热度均衡了,专家并行就均衡了。” 小批时随机波动本身就很大(每个专家平均 1 个 token 时最忙的是平均的 4.9 倍);真实流量还有系统性的冷热不均,而且卡越多越难平均(代码 03 演示 1、3)。
  5. “专家并行组越大越好。” 每卡吞吐在 64 卡以后就饱和了,all-to-all 越来越贵,必须靠微批交替掩盖;负载不均的代价也随卡数增加(代码 03)。
  6. “解耦(不管是 prefill / decode 还是注意力 / 专家)总能大幅提升吞吐。” 在纸上 GPU 的同构集群里,注意力 / 专家解耦只提升 10%–14%;它的大头收益来自能用不同的硬件(代码 04 演示 4、5)。
  7. “ping-pong 流水线的微批越多越好。” 微批多了 TBT 变长、每个微批的 GEMM 变小(memory-bound 时权重要多读几遍);够盖住通信即可(m≥2(1+Tc/Tf)m \ge 2(1 + T_c/T_f),代码 04 演示 2)。
  8. “张量并行加卡就能降低长上下文 decode 的延迟。” TP 超过 KV 头数后 KV 被复制,读 KV 的时间不再下降;MLA 的模型上 TP 对 KV 完全无效(代码 05 演示 2)。
  9. “KV 切到很多卡上,通信量会随上下文变大。” Helix 交换的是部分输出和 LSE,与上下文长度无关(代码 05 演示 1);真正的约束是集合通信的固定延迟(§5.5)。

10. 练习 · Exercises

练习 1、3、4、6、7 的参考答案里的数字是用 moe_lab.py 和代码 05 的函数实际算出来的。

练习 1(碰到几个专家)。(a) E=64E = 64、k=4k = 4 时,批 1、16、64 分别期望碰到几个专家?(b) Mixtral-8x7B(每个专家 1.76 亿参数,32 层,BF16)批 1、4、32 时每步要读多少专家权重?

参考答案

(a) 64×(1−(60/64)n)64 \times (1 - (60/64)^n):批 1 是 4,批 16 是 41.2,批 64 是 63.0——批 64 时几乎全部被碰到,每个专家平均只分到 4 个 token。(b) 每个专家 1.76 亿参数 × 2 字节 = 0.35 GB。批 1 碰到 2 个:每层 0.70 GB,32 层 22.6 GB;批 4 碰到 5.47 个:61.7 GB;批 32 碰到 8 个:90.2 GB(就是全部专家权重)。批从 1 到 4,读的专家权重涨了 2.7 倍,而计算只涨 4 倍——“读字节”很快饱和在总参数上。

练习 2(每个专家的批)。Qwen3-235B(128 选 8),BF16 权重,A100(312 TFLOPS、2 TB/s)。按 MegaScale-Infer 的公式,整批至少要多大,专家的 GEMM 才 compute-bound?如果显存只允许每步 400 个序列,FFN 的理论利用率是多少?

参考答案

每个专家要 F/B=156F/B = 156 个 token,整批要 156×128/8=2496156 \times 128/8 = 2496。批 400 时每个专家平均分到 400×8/128=25400 \times 8/128 = 25 个 token,理论利用率 min⁡(25/156,1)=16%\min(25/156, 1) = 16\%。要提高它,要么把多个副本的请求汇到同一组专家(专家并行 / 解耦),要么换用更”粗”的专家(E/kE/k 更小)。

练习 3(all-to-all 的账)。DeepSeek-V3 decode,专家并行组 64 块卡,每块卡 128 个 token,隐藏维度 7168,每个 token 选 8 个专家,dispatch 用 FP8、combine 用 BF16,网卡有效带宽 40 GB/s、固定延迟 5 µs。(a) 每块卡每层 dispatch 和 combine 各发多少字节、要多久?(b) 58 个 MoE 层合计多久?(c) 这对 TBT 意味着什么?

参考答案

(a) 每块卡发出 128×8×63/64=1008128 \times 8 \times 63/64 = 1008 份 token 向量:dispatch 1008×7168×1=7.231008 \times 7168 \times 1 = 7.23 MB,约 5+181=1865 + 181 = 186 µs;combine 14.45 MB,约 366 µs。(b) 58×(0.186+0.366)≈3258 \times (0.186 + 0.366) \approx 32 ms(已含每次 5 µs 的固定延迟)。(c) 如果 TBT 的目标是 30–50 ms,通信几乎吃掉了全部预算,必须和计算重叠(双微批交替,代码 03 演示 2),或者减少跨节点的流量(把 token 先在节点内按目标节点去重再发——DeepEP 的做法)。

练习 4(ping-pong)。Ta=0.4T_a = 0.4 ms、Te=0.5T_e = 0.5 ms、Tc=0.3T_c = 0.3 ms、L=40L = 40。(a) 按论文的条件 (3),最少要几个微批?(b) 用本讲的闭式公式算 mm = 2、3、4 时的迭代时间和”每个微批的平均产出间隔”。(c) 论文的条件和实际情况一致吗?

参考答案

(a) m≥2(1+0.3/0.5)=3.2m \ge 2(1 + 0.3/0.5) = 3.2,即 4 个。(b) R=0.4+0.5+0.6=1.5R = 0.4 + 0.5 + 0.6 = 1.5。m=2m = 2:C=max⁡(0.8,1.0,0.6,1.5)=1.5C = \max(0.8, 1.0, 0.6, 1.5) = 1.5,T=1.5+39×1.5+1×0.5=60.5T = 1.5 + 39 \times 1.5 + 1 \times 0.5 = 60.5 ms(公式 (5) 给 41.0,低估了);m=3m = 3:C=max⁡(1.2,1.5,0.9,1.5)=1.5C = \max(1.2, 1.5, 0.9, 1.5) = 1.5,T=1.5+58.5+1.0=61.0T = 1.5 + 58.5 + 1.0 = 61.0 ms;m=4m = 4:C=2.0C = 2.0,T=1.5+78+1.5=81.0T = 1.5 + 78 + 1.5 = 81.0 ms。每个微批的平均间隔 T/mT/m:30.3、20.3、20.25 ms。(c) 实际上 m=3m = 3 已经让专家节点忙满(3Te=1.5=R3 T_e = 1.5 = R),m=4m = 4 几乎没有吞吐收益、TBT 却从 61 ms 涨到 81 ms。论文的条件 (3) 用 2(Tf+Tc)=1.62(T_f + T_c) = 1.6 代替 R=1.5R = 1.5,在 Ta≠TeT_a \ne T_e 时偏保守。

练习 5(思考:为什么推理不丢 token)。训练系统长期用容量因子丢 token 而模型质量没怎么受影响,为什么推理系统却坚持 dropless?如果一定要给专家的计算设上限(例如为了 CUDA Graph 的固定形状),你会怎么设计?

参考答案(要点)

训练时丢 token 只是一种正则化式的噪声,梯度会补偿,而且训练不要求单个样本的确定性;推理时丢 token 直接改变某个用户请求的输出,而且结果取决于同批的其他请求,不可复现、难以调试,还违反”优化不改变输出”的约定。要固定形状又不丢 token:按最坏情况留缓冲区(浪费,代码 01 演示 3 的有效载荷只有一半左右),或者准备几档容量、按本批实际负载选一档(CUDA Graph 捕获多个形状),溢出时回退到不用 Graph 的路径;FlashMoE 式的常驻内核则从根本上绕开了固定形状的要求。

练习 6(Helix 的显存与通信)。Llama-405B(126 层,8 个 KV 头,头维度 128,FP4 的 KV),上下文 100 万。(a) 一个序列的 KV 有多大?(b) TP = 16 与 TPA × KVP = 8 × 2 时,每块卡要为每个序列存多少 KV?(c) 8 × 2、批 16 时,每块卡每层 all-to-all 发多少字节?在 NVL72 上(有效 720 GB/s、固定延迟 3 µs)要多久?

参考答案

(a) 126×2×8×128×0.5×106=129126 \times 2 \times 8 \times 128 \times 0.5 \times 10^6 = 129 GB。(b) TP = 16 超过 KV 头数,每块卡存一个完整的 KV 头:129/8=16.1129/8 = 16.1 GB;8 × 2 时 KV 头切 8 份、序列再切 2 份:8.06 GB。同样 16 块卡,Helix 每块卡能多放一倍的序列。(c) 每块卡持有自己那组头(16384/8=204816384/8 = 2048 个隐藏维度)的部分输出,发出其中的一半:16×2048×0.5×216 \times 2048 \times 0.5 \times 2 字节 = 32.8 KB(加上很少的 LSE),约 3+0.05=3.053 + 0.05 = 3.05 µs——几乎全是固定延迟,和上下文长度无关。

练习 7(冗余专家)。8 个专家放在 4 块卡上,长期负载为 [8, 4, 2, 2, 1, 1, 1, 1]。(a) 连续放置(每卡 2 个)时最忙的卡是平均的几倍?(b) 每卡 3 个专家位(共 12 个位,4 个给副本),按 §3.3 的贪心算法放置,结果是什么?

参考答案

(a) 各卡负载 [12, 4, 2, 2],平均 5,最忙的是 2.4 倍。(b) 副本数:每次给”每副本负载”最大的专家加一份——专家 0(8)先加到 2 份(每份 4),这时专家 0 和专家 1 都是 4,按编号先给专家 0 加到 3 份(每份 2.67),再给专家 1 加到 2 份(每份 2),最后给专家 0 加到 4 份(每份 2)。副本按负载从大到小用 LPT 放到最闲的卡上,结果每块卡负载都是 5(例如 {专家 3, 专家 0 的 1/4, 专家 7}、{专家 2, 专家 0 的 1/4, 专家 6}、{专家 1 的 1/2, 专家 0 的 1/4, 专家 5}、{专家 1 的 1/2, 专家 0 的 1/4, 专家 4}),最忙 / 平均 = 1.0。可以用 place_redundant(np.array([8,4,2,2,1,1,1,1.]), 4, 3) 核对。

练习 8(思考:prefill 和 decode 的专家并行组为什么不一样大)。DeepSeek 的 prefill 用 32 块卡的专家并行组,decode 用 144–320 块。用本讲的分析解释。

参考答案(要点)

prefill 一次处理成千上万个 token,即使只有 32 块卡,每个专家也能分到几百上千个 token,专家已经 compute-bound,再扩大专家并行组只会增加 all-to-all 的代价和负载不均的影响(prefill 的负载波动大,§3.4)。decode 每个请求每步只有 1 个 token,单卡的批受 KV 显存限制(几十到两百个序列),只有把很多卡的请求汇到同一组专家,每个专家才能分到足够的 token;同时专家分得越散,每卡省下的显存越多、能放的序列越多(代码 03 演示 2)。代价是 all-to-all 跨节点,必须用双微批交替掩盖。

练习 9(思考:MegaScale-Infer 与 Helix 的共同点)。两篇论文一个面向 MoE 的吞吐 / 成本,一个面向长上下文的延迟,看起来没什么关系。它们在”怎样对待注意力和 FFN”上有什么共同点?又有什么根本不同?

参考答案(要点)

共同点:都承认注意力和 FFN 的资源特性不同(注意力受 KV 的容量和带宽限制,FFN 受权重读取和批大小限制),都让两者用不同的并行方式,都用一次针对性的通信衔接两个阶段,都用流水线(ping-pong / HOP-B)把这次通信藏起来。不同:MegaScale-Infer 在空间上解耦——注意力和专家在不同的 GPU(甚至不同型号的 GPU)上,可以独立扩展,追求吞吐和成本;Helix 在时间上解耦——同一组 GPU 在一层之内先当注意力的 KVP × TPA 网格、再当 FFN 的 TP 网格,追求单个请求的低延迟。前者每层两次跨节点通信,后者依赖大 NVLink 域的低延迟集合通信。


11. 延伸阅读 · Further reading

本讲的必读材料

  • R. Zhu, Z. Jiang, C. Jin, P. Wu, C. A. Stuardo, D. Wang, X. Zhang, H. Zhou, H. Wei, Y. Cheng, J. Xiao, X. Zhang, L. Liu, H. Lin, L.-W. Chang, J. Ye, X. Yu, X. Liu, X. Jin, X. Liu. MegaScale-Infer: Serving Mixture-of-Experts at Scale with Disaggregated Expert Parallelism. arXiv:2504.02263,2025。
  • N. Bhatia, A. More, R. Borkar, T. Mitra, R. Matas, R. Zhao, M. Golub, D. Mudigere, B. Pharris, B. Rouhani. Helix Parallelism: Rethinking Sharding Strategies for Interactive Multi-Million-Token LLM Decoding. arXiv:2507.07120,2025。

选读

  • N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, J. Dean. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. ICLR 2017。
  • O. J. Aimuyo, B. Oh, R. Singh. FlashMoE: Fast Distributed MoE in a Single Kernel. NeurIPS 2025。

相关工作

  • D. Lepikhin et al. GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. ICLR 2021(容量因子与专家并行的训练系统)。
  • W. Fedus, B. Zoph, N. Shazeer. Switch Transformers. JMLR 2022(top-1 路由与容量因子)。
  • DeepSeek-AI. DeepSeek-V3 Technical Report. 2024(细粒度专家 + 共享专家、无辅助损失的负载均衡、推理部署)。
  • A. Agrawal et al. Medha: Efficiently Serving Multi-Million Context Length LLM Inference Requests Without Approximations. 2025(KV 并行,Helix 的对照)。
  • DeepSeek 开源的 DeepEP、EPLB、DeepGEMM 仓库及其说明文档(2025 年 2 月)。

背景

  • 第 4 讲:注意力状态与 ⊕、split-KV(Helix 的数学基础)。
  • 第 6 讲:all-to-all、张量并行时 KV 头的复制、α–β 模型。
  • 第 7 讲:prefill / decode 解耦与异构部署(MegaScale-Infer 的前身思想)。
  • 前置讲 P5:NVLink / NVSwitch、InfiniBand / RDMA。

12. 下一讲预告 · Next lecture

到这里,KV cache 一直被当作”活在 GPU 显存里、请求结束就丢掉”的东西。可在多轮对话、agent、RAG 里,同样的长上下文会被反复使用;一个请求的 KV 动辄几 GB,重新 prefill 既慢又贵,而 GPU 显存装不下所有用户的历史。第 9 讲:上下文管理读两篇论文:Mooncake(月之暗面 Kimi 的服务平台)以 KV cache 为中心设计整个架构——用 CPU 内存和 SSD 组成分布式的 KV 池,调度器在”复用已有的 KV”和”负载均衡”之间权衡,过载时提前拒绝请求;CacheGen 把 KV cache 当作一种需要压缩、流式传输的数据,利用它在 token 之间的局部性和各层敏感度的差异做差分、量化和熵编码,按网络带宽自适应。选读的 Strata 讨论分层上下文缓存的 I/O 效率。我们会做一个分层缓存的模拟器(“加载 vs 重算”的决策),并在 tiny_llm 的 KV 上真实地做 CacheGen 式的压缩,测量它对输出的影响。

← 返回课程目录