Lecture 4 · 分布式推理：切分矩阵与 Transformer 的并行 · Distributed Serving

对应 Maps to：Stanford CS349E / EE339 Efficient ML Inference at Scale（Autumn 2026），Lecture 4 · Thu Oct 1 · Distributed Serving（Mark Horowitz）。阅读材料：J. Austin et al., How to Scale Your Model（Google DeepMind, 2025）第 2 章（All About TPUs）、第 3 章（Sharded Matrices and How to Multiply Them）、第 7 章（All About Transformer Inference）；R. Pope et al., Efficiently Scaling Transformer Inference（MLSys 2023）。原课作业 3（Tensor Parallelism，10/12–10/19）的主题落在本讲。 依据说明：课程网站只公开了课表和阅读材料，没有 slides。本讲以阅读材料为主线，另外用到 Megatron-LM（Shoeybi et al. 2019，张量并行的切法）和前置讲 P5（链路、拓扑、集合通信）。补充材料是我选的，不代表课堂上实际讲了什么。 配套代码 Code：code/lecture04/（numpy；1 个共享模块 + 4 个脚本，每个都在 1 秒内跑完） 预计阅读时间：3–4 小时（含运行代码）

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

0.1 学习目标 Learning goals

学完这一讲，你应该能够：

1. 说清楚推理为什么要用多块芯片（放得下、更快地读权重、更高的吞吐），以及推理的并行和训练有什么不同（memory-bound、小批量、对延迟敏感）。
2. 用切分记号（$A[I_X, J]$）描述一个张量怎样分布在设备上，并对”切分矩阵乘法的四种情形”说出每一种需要什么通信：不需要、all-gather、all-reduce / reduce-scatter，或者先 all-gather 变成别的情形。
3. 解释 Megatron 张量并行为什么”先按列切、再按行切”，为什么每层恰好两次 all-reduce；知道 KV cache 怎样跟着头被切开，以及 GQA / MQA 的 KV 组不够分时会发生什么。
4. 复述 Pope et al. 的主要结论：1D vs 2D weight-stationary（芯片数超过 16 时 2D 更省）、weight-gathered（prefill 时搬权重比搬激活便宜）、MQA 按批切（最大上下文长 32–64 倍）。
5. 给一个模型做部署规划：张量并行度、批大小、精度怎样决定延迟和成本；为什么小批量 decode 对通信延迟最敏感；为什么张量并行要留在一个 NVLink 域里；什么时候用流水线并行、什么时候多开副本。

0.2 先修与本讲范围 Prerequisites & scope

0.3 本讲地图 Map

§1  为什么要分布式推理：放得下、读得快、吞吐高；推理和训练的并行有什么不同
§2  硬件：TPU 的 ICI 与 torus，GPU 的 NVLink 域与 InfiniBand
§3  切分记号与四种情形；集合通信的价目表；网络 roofline                      ← 代码 01
§4  Megatron 张量并行：先列后行；每层两次 all-reduce；KV cache 与 GQA；lm_head    ← 代码 02
§5  Pope et al.：1D / 2D weight-stationary；weight-gathered；MQA 按批切；PaLM 540B ← 代码 03
§6  延迟、吞吐与成本：并行度 × 批大小；前沿；切多少份；跨机器                    ← 代码 04
§7  其他并行方式：流水线、专家、序列
§8  现状（2025–2026）
§9–13  小结、常见误区、练习（含答案）、延伸阅读、下一讲预告

0.4 记号 Notation

1. 为什么要分布式推理 · Why distribute inference

1.1 三个理由

第 1 讲 §4.4 已经看到：Llama-3.1-70B 的 BF16 权重（141 GB）放不进一块 H100（80 GB）。所以第一个理由是放得下。但即使放得下，还有两个理由：

1. 放得下（容量）：权重 + KV cache 要装进显存。第 3 讲的量化能缓解，但 405B、671B 的模型怎样都要多块卡。
2. 读得快（延迟）：decode 每一步要把权重读一遍。8 块卡一起读，每块只读 1/8，一步的时间理论上降到 1/8（第 1 讲 §4.2 的”读权重下界”被 $n$ 块卡的总带宽摊薄）。这是推理里张量并行最主要的用途。
3. 吞吐高：多块卡、多份副本，同时服务更多请求。

1.2 推理的并行和训练不一样

How to Scale Your Model 第 7 章特别强调，训练里好用的并行方式，到了 decode 很多都不能用：

• FSDP（完全切分的数据并行）不能用：它在每一层计算前把权重 all-gather 过来。decode 本来就是”读权重”受限的，再把权重通过比 HBM 慢一个数量级的网络搬一遍，只会更慢。decode 要搬激活，不要搬权重。
• 数据并行没意义：它复制参数，并不能让每一步读参数更快。与其做数据并行，不如直接多起几个模型副本（§6.4）。
• decode 没有序列可切：每一步每个序列只有一个新 token，序列并行无从谈起（prefill 可以，§7.3）。

剩下的主要就是各种模型并行：张量并行（切矩阵）、流水线并行（切层）、专家并行（切专家）。本讲的重点是张量并行。

2. 硬件：芯片怎样连在一起 · How chips are connected

详细的链路和拓扑见 P5 §2–§3。这里只回顾和本讲有关的数字。

TPU（How to Scale Your Model 第 2 章）：芯片之间用 ICI（inter-chip interconnect）直接相连，组成 2D 或 3D 的 torus（每一维首尾相接成环）。TPU v5e 一个 pod 是 16 × 16 = 256 块芯片的 2D torus，每条 ICI 链路单向约 45 GB/s；TPU v4 / v5p 是 3D torus。Pope et al. 用的 TPU v4：每块芯片 BF16 275 TFLOP/s、32 GiB HBM、1200 GB/s HBM 带宽、270 GB/s 互连带宽。torus 的好处是规模大（几千块芯片在一个 ICI 域里），代价是远处的芯片要走好几跳。

GPU：一台机器里 8 块 H100 通过 NVSwitch 两两全连（NVLink 4：每块卡单向 450 GB/s）；机器之间用 InfiniBand（每块卡一张 400 Gb/s 的网卡，约 50 GB/s）。所以 GPU 的世界有两层：NVLink 域内快、跨机器慢一个数量级。GB200 / GB300 NVL72 把 NVLink 域扩大到 72 块卡（第 1 讲 §5.7）。

一个贯穿全讲的比值是网络临界强度 = 峰值算力 / 链路带宽（第 1 讲 §3.10）：H100 + NVLink 约 2,198 FLOP/byte。每通过网络传 1 字节，至少要做这么多次运算，网络才不会成为瓶颈。

3. 切分记号与四种情形 · Sharding notation and the four cases

代码：01_sharded_matmuls.py，共享模块 shard_lab.py。8 块”虚拟 H100”挂在 NVLink 交换机上，每块持有真实的 numpy 分片，集合通信用 P5 的 comm_sim（真实搬运 + α–β 计时）。

3.1 记号

$A[I_X, J]$ 表示：$A$ 的第一维 $I$ 被切到网格轴 $X$ 上，每块设备只有 $I / |X|$ 行；$J$ 没有下标，每块设备都有完整的 $J$。代码 01 演示 1 用 Llama-3.1-70B 的 MLP 上投影 $W[D = 8192, F = 28672]$ 举例（8 块设备，BF16）：

这个记号同时说明了两件事：每块设备存多少，以及做矩阵乘法时需要什么通信。

3.2 四种情形

矩阵乘法 $C = A \cdot B$ 里，$J$ 是被求和的维度（contracting dimension）。How to Scale Your Model 第 3 章按”$J$ 有没有被切、怎么切”把所有情形分成四类。代码 01 演示 2 在 8 块虚拟 H100 上把每种情形都做了一遍（模拟尺寸：$I = 512$ 个 token，$J = 8192$，$K = 28672$，BF16；实际数组每维缩小 32 倍，通信字节按比例放大）：

逐一看：

• 情形 1：被求和的维度没被切，每块设备用自己的那几行做本地乘法，一个字节都不用传。
• 情形 2：只有 $A$ 的 $J$ 被切，$B$ 是完整的。最直接的做法是先把 $A$ all-gather 回来，但这样每块设备都做了全部运算（算力没被分摊）。
• 情形 3：两边的 $J$ 切法一致。每块设备算一个部分和 $C[I, K]\{U_X\}$（“还没加起来”），再 all-reduce。运算被分摊了，代价是传输输出 $C$——这里 $C$ 比 $A$ 大 3.5 倍、all-reduce 又是 all-gather 的两倍，通信是情形 2 的 7 倍。如果下一步本来就要按 $K$ 切着用，就只做 reduce-scatter（半价），得到 $C[I, K_X]$。
• 情形 4：两个”不被求和”的维度切在同一条轴上，每块设备只有 $I$ 的一段和 $K$ 的一段，只能算出 $C$ 的对角块——$C[I_X, K_X]$ 不是一个合法的切法。先 all-gather 其中一个，变回情形 1。这里 all-gather 的是权重 $B$（411 MB）；如果改为 all-gather 激活 $A$，只要传 7 MB。能搬激活就别搬权重。

3.3 集合通信的价目表

代码 01 演示 4（8 块 H100，NVLink，α = 2 µs）：

• 大消息时 all-gather 和 reduce-scatter 一样贵（约 $\frac{n-1}{n} \cdot V / W$），all-reduce = reduce-scatter + all-gather，是两倍（Pope et al. 附录 A.1 也是这么建模的）。
• 小消息时都被延迟 $\alpha$ 主导：ring 要 $2(n-1)$ 步，递归倍增只要 $\log_2 n$ 步，快得多。张量并行 decode 的 all-reduce 恰好是小消息（§6.1）。
• How to Scale Your Model 还给出：双向环上的 all-to-all 只要 all-gather 的约 1/4。

于是四种情形的”价格”一目了然：情形 2 / 4 付一次 all-gather，情形 3 付一次 all-reduce（或半价的 reduce-scatter），情形 1 免费。

3.4 网络 roofline：切开被求和的维度值不值

情形 3 里，每块设备的运算是 $2BDF/n$，ring all-reduce 的带宽项约 $2\frac{n-1}{n} \cdot (2BF \text{ 字节}) / W$（输出的大小，和 $D$ 无关）。两者之比：

H100 + NVLink 时，要 $D / (n-1) \gtrsim 4{,}396$ 才是计算更慢。代码 01 演示 3（$B = 512$）：Llama-3.1-70B（$D = 8192$）切 2 份时计算更慢（121.6 vs 69.2 µs），切 8 份就是通信更慢（30.4 vs 142.2 µs）；连 $D = 16384$ 的模型切 8 份也是通信更慢。$B = 16$ 时还要加上延迟：每次 all-reduce 至少 $2(n-1)\alpha$，计算却只有几微秒——decode 的张量并行几乎总是被通信的延迟主导。

这说明了一件看似矛盾的事：按 roofline，张量并行在 decode 时”通信受限”；可它仍然值得做，因为 decode 的计算本来就不是瓶颈，读权重才是（§6）。

4. Megatron 张量并行 · Megatron tensor parallelism

代码：02_megatron_tensor_parallel.py。shard_lab.shard_params_tp 把第 2 讲的模型切给 $n$ 块虚拟设备，tp_forward 做张量并行的前向。

4.1 MLP：先列后行

一个 SwiGLU MLP 是两步矩阵乘法：$h \to (W_\text{gate}, W_\text{up}) \to$ 逐元素 $\to W_\text{down}$。Megatron-LM 的技巧：

• 第一步的矩阵按列切（按输出维 $F$）：$W_\text{up}[D, F_X]$。输入 $h[B, D]$ 是复制的——这是情形 1，不需要通信。每块设备得到中间值的 $F/n$ 列 $m[B, F_X]$；SwiGLU 的逐元素运算也在本地完成。
• 第二步的矩阵按行切（按输入维 $F$）：$W_\text{down}[F_X, D]$。$m[B, F_X] \cdot W_\text{down}[F_X, D]$ 正好是情形 3——本地部分和 + 一次 all-reduce。

两个矩阵之间一个字节都不用传，整个 MLP 只在最后 all-reduce 一个 $[B, D]$。代码 02 演示 3 比较了三种切法（Llama-3.1-70B，8 块设备）：

“先行后列”要在中间 all-reduce 一个 $[B, F]$（$F = 3.5D$，大 3.5 倍）；“两次都按列”要 all-gather 整个中间值。两者都要两次集合通信，字节是 Megatron 的 2.25–4 倍。

4.2 注意力：按头切

注意力是同一个道理：$W_q, W_k, W_v$ 按头切（列），每块设备完整地算自己那几个头的注意力（头和头之间本来就独立），$W_o$ 按行切，最后 all-reduce。所以 Megatron 张量并行的全部通信就是：每层两次 all-reduce（注意力一次、MLP 一次），和设备数无关。词嵌入、RMSNorm 的缩放系数是复制的，每块设备各算一遍 RMSNorm（重复但便宜）。

4.3 正确性：和单卡一模一样

代码 02 演示 1 把第 2 讲的小模型切到 1、2、4、8 块虚拟设备上做 prefill：

差别只来自 all-reduce 时求和的顺序不同（浮点加法不满足结合律，P3 §4.2）。每块设备的参数没有降到 $1/n$：词嵌入、RMSNorm 是复制的，$n > G$ 时 K、V 投影也被复制（§4.4）。这个小模型词表占比大，复制的部分显得多。

作业 3 的检验方法就是这样：切开以后的输出必须和单卡逐元素相等（差在浮点舍入量级）。不然就是切法写错了。

4.4 KV cache 跟着头切开——直到 KV 组不够分

代码 02 演示 2 用张量并行做带 KV cache 的贪心生成（3 个序列，提示 12 个 token，生成 16 个），生成的 token 和单卡完全相同。每块设备只存、只读自己那几个头所属的 KV 组：

（小模型：8 个 query 头、2 组 KV。）设备数 ≤ KV 组数时，张量并行同时摊薄了权重和 KV 的带宽、容量。设备数 > KV 组数时，一组 KV 要服务分到不同设备上的几个 query 头，只能每块设备各存一份——GQA / MQA 和张量并行的冲突。Llama-3.1-70B 有 8 组 KV：张量并行度 ≤ 8 时 KV 完全切开，16 时每组要存 2 份。MQA 只有 1 组，切几份就复制几份。§5.3 讲 Pope 的解法。

4.5 lm_head 和序列并行

lm_head 通常按词表切（每块设备 $V/n$ 个词的分数），最后 all-gather logits。注意它的大小：Llama-3.1-70B、$B = 64$ 时 logits 有 16.42 MB，而一次层内 all-reduce 的消息只有 1.05 MB。只需要下一个 token 时，可以在每块设备上先取本地的最大值或 top-$k$，只交换候选。

序列并行（Korthikanti et al. 2022）是 Megatron 的一个改进：把每层的 all-reduce 拆成 reduce-scatter + all-gather，中间的 RMSNorm、残差加法在按序列切开的激活上做。通信量不变（all-reduce = RS + AG），省掉了 norm 的重复计算和激活的重复存储，主要用于训练和长 prefill。

5. Pope et al.：大模型推理怎样分区 · Partitioning PaLM for inference

代码：03_pope_partitioning.py

5.1 背景：一个解析模型

Pope et al. 在最多 256 块 TPU v4 上给 PaLM（8B、62B、540B）做推理。他们的方法论值得学：先写一个简单的解析模型（每种分区方式的计算、访存、通信时间），用它在”延迟 vs 成本”的目标下挑选分区方式，再做底层优化逼近模型。成本用”芯片·秒 / token”衡量（= 芯片数 × 每步时间 / 这一步处理的 token 数），它直接正比于钱，反比于 MFU。

记号沿用切分的下标：$BLE_{xyz}$ 表示形状 $[B, L, E]$ 的张量，最后一维被切成 $X \times Y \times Z$ 份（TPU v4 的 3D torus 正好三条轴）。

5.2 前馈层：1D vs 2D weight-stationary

1D weight-stationary（Megatron 式）：权重按 $F$ 切到全部 $n$ 块芯片上，激活要在每层前后 all-gather / reduce-scatter 一整份 $BLE$。通信时间

和芯片数无关：芯片越多，读权重和计算越快，通信却不变，迟早成为瓶颈。

2D weight-stationary：权重在 $E$ 和 $F$ 两个方向都切（$E_x F_{yz}$），每块芯片拿一个近似方形的块。激活每次只沿一部分轴聚合，交替在 $E$ 和 $F$ 方向上进行，任何时候都不需要完整复制的激活。附录 A.2.1 的推导：$T = \frac{2BL}{\text{带宽}}\left(\frac{E}{X} + \frac{F}{YZ}\right)$，选最优的 $X$ 得

两者相等在 $n = 4F/E$：$F = 4E$ 时 $n = 16$。论文说”芯片数超过 16 时 2D 更省”。代码 03 演示 1（PaLM 540B：$E = 18432$、$F = 73728$、118 层，decode $B = 512$，BF16 激活，只算带宽项）：

$n = 64$ 时 1D 的通信已经接近计算、是读权重的两倍；2D 只有一半，而且还能继续随 $\sqrt{n}$ 下降。通信时间能不能随芯片数下降，决定了加芯片还能不能降低延迟。

5.3 prefill：weight-gathered

当一步里的 token 很多时（prefill），激活比权重还大。这时反过来：激活不动，搬权重（weight-gathered）。权重从 2D 的布局出发，在 $N$ 块芯片之间 all-gather，激活在其余方向聚合。附录 A.2.2：

weight-stationary 的通信 $\propto BL$，weight-gathered $\propto \sqrt{BL}$，两者相等在 $BL = F/4$。代码 03 演示 2（论文图 3 的设置：$E = 16384$、$F = 65536$、64 块芯片）：

decode 的批（几百个 token）远在交叉点左边，用 2D weight-stationary；prefill 动辄几十万个 token，换成 weight-gathered。同一套权重布局两种都能用，只是运行时选不同的通信方式。token 越多通信占比越小，论文在 prefill 时达到了 76% 的 MFU。

5.4 注意力：MQA 按批切

PaLM 用多查询注意力（MQA，第 2 讲 §6.1）：只有一组 K/V。如果像 MHA 那样按头切，每块芯片都需要这唯一的一组 K/V，只能各存一份——MQA 省下的显存全被复制吃掉了。Pope 的做法：注意力部分按批切——每块芯片负责一部分序列，存这些序列的完整 KV，不复制。代价是一次 all-to-all，把 Q 从”按头切”换成”按批切”，算完再换回去。

代码 03 演示 3 复现了论文表 1（64 块 TPU v4，30% 的显存给 KV）：

按批切让最大上下文长了 32–64 倍。$B = 256$、上下文 8192 时，每块芯片读 KV 的时间从 211 ms（按头切）降到 3.3 ms，多出来的 all-to-all 只有约 115 µs（只算带宽项）。decode 时 Q 只有一个 token，比几千个 token 的 KV 小几个数量级——花一点通信，省大量 KV 读取。prefill 时 Q 很长，就不划算了，论文 prefill 时仍按头切。

5.5 其他技巧与端到端结果

• 并行的注意力 / 前馈层（PaLM 的结构）：$y = x + \text{Attn}(\text{LN}(x)) + \text{FFN}(\text{LN}(x))$，两个子层可以共享一次 all-reduce。论文测得串行结构的延迟高 14%。
• Looped CollectiveEinsum：把集合通信拆成小块，和矩阵乘法交错进行，让通信藏在计算后面。论文说底层优化带来了约 1.4 倍的提升。
• int8 权重（第 3 讲）：只量化权重，省读权重的时间。

PaLM 540B 在 64 块 TPU v4 上：$B = 64$ 时每 token 28.5 ms（int8 权重）、36.9 ms（bf16）。代码 03 演示 4 算了 roofline 下界：

实测是下界的 2.4–3.4 倍：下界没算每次集合通信的启动延迟（118 层 × 每层若干次）、通信和计算没能完全重叠、小运算的开销、芯片达不到峰值带宽。roofline 告诉你”不可能更快”，不告诉你”能做到多快”。论文还报告：成本在批大小约 512 时最低，最低延迟比批大小 512 时低约 3 倍。

6. 延迟、吞吐与成本 · Latency, throughput, and cost

代码：04_latency_throughput_cost.py。Llama-3.1-70B，H100，上下文 4096；每步 = max(计算, 访存) + 160 次 all-reduce（不重叠，偏保守）。

6.1 张量并行度 × 批大小

代码 04 演示 1（BF16）：

（“一步”按两种 all-reduce 算法中更快的那个算；BF16 时 $n = 2$ 只放得下 $B = 1$。）

• 张量并行降低单用户延迟：$B = 1$ 时，$n$ 从 2 到 8，读权重从 21 ms 降到 5.2 ms。
• 小消息的集合通信算法很关键：$n = 8$ 时 ring all-reduce 160 次一共 4.5 ms，和读权重差不多；递归倍增只要 1 ms。（这和 P5 §7.1 的结论一致。）
• 批变大，通信的占比下降（46% → 28%），但通信时间本身在涨：$B = 256$ 时每次 all-reduce 的消息有 4 MB。小批量、低延迟的部署对通信延迟最敏感。

6.2 延迟–成本前沿

把精度（BF16 / FP8）、张量并行度（1–8）、批大小（1–256）的所有组合都算一遍，代码 04 演示 2 找出”没有别的组合能同时更快、更便宜”的点：

（节选；完整的前沿有 9 个点。）沿着前沿往右主要是加大批：$B$ 从 1 到 128，每个 token 的成本降了约 38 倍，每个用户只慢了约 3.4 倍。最便宜的点比最快的点便宜约 46 倍。前沿上全是 FP8（字节减半）。“芯片·时 / 百万 token”乘以每卡每小时的价格，就是每百万 token 的芯片成本。

没有”最好”的配置，只有”在你的延迟要求下最便宜”的配置。Pope et al. 图 1 画的就是这样的前沿。

6.3 张量并行能切多少份

How to Scale Your Model 第 7 章给了一个简单的上限。decode 是 memory-bound 的：每层读权重 $2DF / (Y \cdot W_\text{hbm})$，每层通信约 $2BD / W_\text{link}$。通信超过读权重时

代码 04 演示 3：H100 + NVLink 的 $\beta = 7.4$，Llama-3.1-70B（$F = 28672$）在 $B = 32$ 时上限约 120 份，$B = 128$ 时 30 份，$B = 512$ 时约 8 份。批越小，能切的份数越多——小批量本来就该多切，用更多卡的 HBM 带宽换延迟。一台 H100 机器只有 8 块卡，这个带宽上限几乎总是大于 8；真正的限制是延迟（每次 all-reduce 几微秒）和 NVLink 域的大小。

这和训练不同：训练要 compute-bound，张量并行只能切到 $F / 2200$ 左右（4–8 份）；decode 是 memory-bound，可以切得更细。

6.4 跨机器：TP vs PP vs 副本

一台机器不够时（16 块 H100，2 台机器，一共服务 128 个序列，BF16），代码 04 演示 4：

• TP16 让每块卡读的权重和 KV 再少一半，可 160 次 all-reduce 都要走网卡（跨机器一次 51 µs，机器内 20 µs），一共约 8 ms，把省下的全吃掉了——甚至比一台机器扛全部 128 个序列还慢。张量并行最好留在一个 NVLink 域里。
• 流水线并行（TP8 × PP2）：每台机器放 40 层，两个微批交替流过两台机器。延迟和吞吐与”两个副本”几乎一样——工作量相同。机器之间只传一个 $[64, d]$ 的激活（约 1 MB），很便宜。
• 两者的区别在显存：流水线每块卡只存半个模型，省下的空间可以放更多 KV。所以流水线的用处是模型一台机器放不下、或者想把显存留给 KV 的时候；放得下时，多开副本最简单（没有任何跨机器依赖）。

7. 其他并行方式 · Other forms of parallelism

7.1 流水线并行

把层切成几段放在不同的芯片上，激活依次往下传。§6.4 已经看到它的特点：单个请求的路径不变短（每个 token 仍要依次经过所有段），吞吐靠”多个微批同时在不同段上”来提高；段与段之间只传激活，通信很少，适合跨机器。代价是气泡（bubble）：微批不够多时，有的段在等。decode 时每个 token 都要走完全部段才能开始下一个 token，所以要有足够多的并发请求才能把流水线填满。

7.2 专家并行

MoE 模型（第 3 讲 §7）把专家分散到不同的芯片上，每块芯片只放几个专家。token 被路由到哪个专家，就通过 all-to-all 发到那块芯片，算完再发回来。它解决了第 3 讲 §7.6 的问题：每个专家分到的 token 来自整个集群的批，“有效批”变大了。代价是每层两次 all-to-all，跨机器时很贵（P5 §7.4 算过：MoE 的 all-to-all 跨机器比机器内慢约 10 倍）。DeepSeek-V3 的部署把专家并行和按批切的注意力组合在一起，第 5 讲会再提到。

7.3 序列并行与 prefill

prefill 有很长的序列，可以把序列切开：每块芯片算一段 token 的 Q，K/V 在芯片之间轮转（ring attention）。How to Scale Your Model 第 7 章的建议：prefill 时”张量并行切到网络的上限，然后做序列并行”。decode 没有序列可切。

8. 现状（2025–2026）· State of the field

• 单机 TP8 是默认：vLLM、SGLang、TensorRT-LLM 部署 70B 级模型时，通常在一台 8 卡机器内做张量并行（FP8 时常常 TP4 甚至 TP2 就够），多台机器之间开副本。
• NVLink 域变大：GB200 / GB300 NVL72 把 72 块 GPU 放进一个 NVLink 域，跨 72 块卡的张量并行 / 专家并行变得可行；TPU 一直是几百到几千块芯片的 ICI 域。
• MoE 推理以专家并行为主：DeepSeek-V3 / R1 这样的大 MoE 模型用”注意力数据并行（按批切）+ 专家并行”的组合（DeepEP 等通信库专门优化 all-to-all），并且和 prefill / decode 解耦一起使用（第 5 讲）。
• 通信与计算重叠成为标配：类似 Looped CollectiveEinsum 的思路在 GPU 上以”计算–通信融合内核”（把 all-reduce 和矩阵乘法写进同一个内核）的形式出现，目标正是 §5.5 里实测和下界之间的那几倍差距。

9. 小结 · Summary

1. 推理用多块芯片有三个理由：放得下、读得快（单用户延迟）、吞吐高。decode 是 memory-bound 的，所以训练里的 FSDP、数据并行在 decode 里没用——搬激活，不搬权重。
2. 切分矩阵乘法的四种情形：被求和的维度没切——免费；只有一边切了——all-gather；两边切法一致——部分和 + all-reduce / reduce-scatter；两个非求和维度切在同一轴上——不合法，先 all-gather 一个。all-reduce = reduce-scatter + all-gather；小消息被延迟主导。
3. Megatron 张量并行：先按列、再按行切，中间不通信；每层两次 all-reduce；KV cache 跟着头切开，KV 组不够分时要复制。切开后的输出必须和单卡逐元素相等。
4. Pope et al.：1D weight-stationary 的通信不随芯片数下降，2D 按 $1/\sqrt{n}$ 下降（$n > 4F/E$ 时更省）；prefill 时 token 多到一定程度（$BL > F/4$），搬权重比搬激活便宜；MQA 要按批切才省 KV（最大上下文长 32–64 倍）。roofline 下界和实测差 2–3 倍。
5. 部署是一条前沿：张量并行降低延迟，批处理降低成本，FP8 两头都帮；小批量 decode 对通信延迟最敏感（集合通信算法要选对）；张量并行留在 NVLink 域内，跨机器用流水线或副本。

10. 常见误区 · Common pitfalls

1. “8 块卡就快 8 倍。” 读权重快了 8 倍，但每层两次 all-reduce 的延迟不随卡数下降，ring 算法还随卡数变慢（§6.1）。
2. “decode 时张量并行是通信受限的，所以不该做。” 按 roofline 它确实”通信比计算慢”，但 decode 的瓶颈是读权重，张量并行把读权重摊薄了，仍然划算（§3.4、§6.3）。
3. “FSDP 在训练里好用，推理也可以用。” 它每层把权重搬一遍，decode 本来就是读权重受限，只会更慢（§1.2）。
4. “切分只要算对就行，用哪种切法无所谓。” 同一个 MLP 的三种切法，通信量差 2–4 倍（§4.1）；情形 4 里 all-gather 权重还是激活，差 50 多倍（§3.2）。
5. “GQA 的 KV 和张量并行没关系。” 张量并行度超过 KV 组数时，KV 要被复制，KV 的显存和带宽优势就打了折扣（§4.4）。
6. “MQA 的 KV 只有一组，所以总是省显存。” 按头切时每块芯片都要存一份，比 MHA 还差（§5.4）。
7. “跨机器的张量并行只是慢一点。” 跨机器的 all-reduce 慢好几倍，一层两次、几十上百层累加起来，足以让 TP16 比 TP8 还慢（§6.4）。
8. “流水线并行能降低延迟。” 它不缩短单个请求的路径；它的好处是显存（每块卡只存一部分层）和跨机器的低通信量（§6.4、§7.1）。
9. “roofline 的下界就是能达到的速度。” PaLM 540B 实测是下界的 2.4–3.4 倍（§5.5）。
10. “成本最低的配置就是最好的。” 成本最低的点每个用户最慢；部署要先定延迟要求，再在前沿上找最便宜的点（§6.2）。

11. 练习 · Exercises

带 💻 的题可以用本讲的模块验证。参考答案在每题后面。

练习 1（切分记号）。8 块设备，激活 $A[B = 64, D = 8192]$ 按 $A[B, D_X]$ 存，权重 $W[D, F = 28672]$ 按 $W[D_X, F]$ 存（BF16）。每块设备各存多少字节？$C = AW$ 属于哪种情形？需要什么通信？

$A$：$64 \times 1024 \times 2 = 131$ KB；$W$：$1024 \times 28672 \times 2 = 58.7$ MB。被求和的维度 $D$ 在两边都按 $X$ 切——情形 3：每块设备算部分和 $C[B, F]\{U_X\}$，再 all-reduce（得到复制的 $C[B, F]$）或 reduce-scatter（得到 $C[B, F_X]$）。

练习 2（情形判断）。$A[I_X, J] \cdot B[J_X, K]$——$A$ 的 $I$ 和 $B$ 的 $J$ 切在同一条轴上。可以直接算吗？怎样做最省？

不能直接算：设备 $d$ 只有 $A$ 的第 $d$ 段行（但 $J$ 完整），和 $B$ 的第 $d$ 段 $J$，它们的 $J$ 对不上。把 $B$ all-gather 回 $B[J, K]$，变成情形 1，得到 $C[I_X, K]$。如果 $B$ 是权重、$A$ 是激活，也可以考虑先把 $A$ 重新切成 $A[I, J_X]$（all-to-all），变成情形 3——哪种更便宜取决于 $A$、$B$、$C$ 的大小。

练习 3（每层的通信）💻。Llama-3.1-8B，张量并行度 4，decode $B = 32$，BF16。每次 all-reduce 的消息多大？每块设备每一步一共发送多少字节（ring）？用 tp_decode_step 比较 ring 和递归倍增的通信时间。

消息 $32 \times 4096 \times 2 = 262{,}144$ 字节。ring 每次每块设备发 $2 \cdot \frac{3}{4} \cdot 262{,}144 \approx 393$ KB，32 层 × 2 次 = 64 次，一共约 25.2 MB。通信时间：ring 0.82 ms，递归倍增 0.33 ms；一步的下界分别是 3.23 ms 和 2.73 ms（读权重和 KV 2.40 ms）。

练习 4（1D vs 2D 的交叉点）。由 $T_\text{2D} = 4BL\sqrt{EF/n}/\text{带宽}$ 和 $T_\text{1D} = 2BLE/\text{带宽}$，求两者相等时的 $n$。Llama-3.1-70B（$F/E = 3.5$）是多少？

$4\sqrt{EF/n} = 2E \Rightarrow n = 4F/E$。$F = 4E$ 时 16；Llama-3.1-70B 是 $4 \times 3.5 =$ 14。一台 8 卡机器内 1D（Megatron）就够了；几十块卡以上才值得 2D。

练习 5（MQA 的切法）💻。仿照代码 03 演示 3，算 PaLM 540B 在 64 块 TPU v4、批大小 256 时三种注意力切法的最大上下文长度。

MHA 666；MQA 按头切 333；MQA 按批切 21,326。批大小翻倍，最大上下文减半；三者的比例不变。

练习 6（集合通信算法）。为什么张量并行 decode 时递归倍增比 ring 快，而 prefill 时 ring 更好？用 α–β 公式解释。

ring：$2(n-1)\alpha + 2\frac{n-1}{n} \cdot \frac{V}{W}$；递归倍增：$\log_2 n \cdot (\alpha + V/W)$。decode 的消息 $V$ 只有几十 KB 到 1 MB，$\alpha$ 项主导，递归倍增的步数少（$n = 8$ 时 3 步 vs 14 步）。prefill 的消息几十到上百 MB，带宽项主导，ring 每一步只传 $V/n$，总字节约 $2V$，而递归倍增每一步都传完整的 $V$，总共 $\log_2 n \cdot V$（代码 01 演示 4：64 MB 时 289 µs vs 453 µs）。

练习 7（能切多少份）。Llama-3.1-8B（$F = 14336$），$B = 64$，H100（$\beta = 7.44$）。按 $Y > F/(B\beta)$，张量并行切到多少份时通信赶上读权重？实际部署 8B 模型通常用几卡？为什么？

$14336 / (64 \times 7.44) \approx$ 30 份。但 8B 的 BF16 权重只有 16 GB，一块卡就放得下；切开的唯一好处是延迟，而每次 all-reduce 的固定延迟（几微秒 × 64 次）很快就抵消了读权重省下的时间（一块 H100 读完 8B 只要约 4.5 ms）。所以 8B 通常单卡（或 TP2）部署，多开副本提高吞吐。

练习 8（weight-gathered 的交叉点）。PaLM 540B（$F = 73728$）的 prefill，一步里的 token 数超过多少时 weight-gathered 比 2D weight-stationary 省？一个 2048 token 的提示够吗？

$BL > F/4 =$ 18,432 个 token。单个 2048 token 的提示不够；把 9 个以上这样的提示放在同一批 prefill，才值得切换到 weight-gathered。

练习 9（成本）。FP8、$n = 8$、$B = 64$ 时每 token 7.88 ms。每个 token 占多少”芯片·秒”？每百万 token 多少芯片·时？如果每块 H100 每小时 2 美元，每百万 token 的芯片成本是多少？

$8 \times 7.88\text{ ms} / 64 = 0.985$ 芯片·毫秒 / token；乘 $10^6$ 再除以 3600 秒 = 0.274 芯片·时 / 百万 token（代码 04 演示 2）；× 2 美元 ≈ 0.55 美元 / 百万 token（只算芯片租金，且是 roofline 下界）。

练习 10（思考题）。DeepSeek-V3 用 MLA（第 2 讲 §6.4）：KV 是一个所有头共享的 576 维潜在向量。按头切的张量并行会遇到什么问题？结合 §5.4，你会怎样切它的注意力？

潜在向量是所有头共享的，和 MQA 的单组 KV 一样：按头切时每块卡都要存完整的潜在向量，KV 被复制 $n$ 份。所以应该像 Pope 对 MQA 那样按批切注意力（每块卡负责一部分序列的完整 KV），这正是 DeepSeek 部署里”注意力数据并行”的做法；MoE 部分再用专家并行，两者之间用 all-to-all 换布局。

12. 延伸阅读 · Further reading

本讲阅读材料

• J. Austin et al., How to Scale Your Model, Google DeepMind, 2025. 第 2 章 All About TPUs、第 3 章 Sharded Matrices and How to Multiply Them、第 7 章 All About Transformer Inference。https://jax-ml.github.io/scaling-book/
• R. Pope, S. Douglas, A. Chowdhery, J. Devlin, J. Bradbury, A. Levskaya, J. Heek, K. Xiao, S. Agrawal, J. Dean, Efficiently Scaling Transformer Inference, MLSys 2023（arXiv:2211.05102）。

本讲补充材料

• M. Shoeybi et al., Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, 2019（arXiv:1909.08053）。—— 先列后行的切法。
• V. Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models, 2022（arXiv:2205.05198）。—— 序列并行。
• S. Wang et al., Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models（ASPLOS 2023）。—— Looped CollectiveEinsum。
• How to Scale Your Model 第 8 章 Serving LLaMA 3-70B on TPUs。—— 本讲 §6 的 TPU 版本。

本仓库的相关讲义

• 前置讲 P5：链路、拓扑、all-reduce 算法、LLM 推理里的通信账。
• CS349D 第 6 讲：并行策略搜索与 Transformer 推理的分区（系统角度）。

13. 下一讲预告 · Next lecture

第 5 讲（Zain Asgar）讲解耦式推理，阅读材料是 DistServe 和 Gimlet Labs 关于”d-Matrix Corsair 上的投机解码”的博客。本讲一直在问”一个阶段怎样切到多块芯片上”；第 5 讲问的是另一个维度的问题：prefill 和 decode 的”性格”完全不同（第 1 讲 §4.6），为什么要让它们挤在同一组 GPU 上？ 把它们分到不同的机器、甚至不同种类的芯片上，各自选最合适的并行方式和批大小，代价是把 KV cache 从 prefill 机器搬到 decode 机器。我们会写一个离散事件模拟器，比较”共置”和”解耦”在不同延迟目标下的 goodput，并重现 DistServe 的放置搜索。