Lecture 6 · 并行：从 FlexFlow 的策略搜索到 Transformer 推理的分区 · Parallelism

对应 Maps to：Stanford CS349D AI Inference Infrastructure（Spring 2026），#2 Wed Apr 1 Parallelism。必读：Z. Jia, M. Zaharia, A. Aiken, Beyond Data and Model Parallelism for Deep Neural Networks（FlexFlow，SysML 2019）；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）。选读：LoongServe（Wu 等，SOSP 2024）、Alpa（Zheng 等，OSDI 2022）、Arctic Inference with Shift Parallelism（Snowflake，2025）。 依据说明：原课没有公开 slides。本讲依据两篇必读论文的全文；选读论文中，Arctic Inference 读了全文，LoongServe 和 Alpa 读了摘要与引言（§5 对它们的介绍限于这些部分）。原课把”并行”放在第 2 次课，因为项目的第一个里程碑是数据并行 / 张量并行；本讲义把它挪到单卡引擎之后（见 SYLLABUS §0.1）。§6”之后的发展”是我补充的。 配套代码 Code：code/lecture06/（numpy + 标准库；5 个脚本 + 1 个模块 parallel_sim.py：ring 集合通信、α–β 通信模型、tiny_llm 的张量并行与流水线并行、Pope 等人的 1D / 2D 前馈层分区；每个脚本 1 秒左右） 预计阅读时间：3–4 小时（含运行代码）

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

0.1 学习目标 Learning goals

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

1. 说出推理需要多卡的三个理由（放不下、降延迟、提吞吐），以及每种理由对应什么并行方式。
2. 用 α–β 模型估算一次集合通信的时间，解释 ring all-reduce 为什么”带宽最优”、为什么 decode 时它受限于启动延迟。
3. 解释 Megatron 张量并行的”列切 + 行切”为什么每层只需要 2 次 all-reduce；知道 KV 头少于 GPU 数时 KV 会被复制。
4. 比较张量并行、数据并行（多副本）、流水线并行在延迟、吞吐、通信、显存上的取舍，知道流水线并行不缩短单个请求的延迟。
5. 说清 FlexFlow 的三个想法：SOAP 搜索空间、基于任务图的执行模拟器、MCMC 搜索；知道它为什么比”数据并行或模型并行”二选一好。
6. 推导 Pope 等人的前馈层通信量：1D weight-stationary 约 $2BLE$、2D 约 $8BLE/\sqrt{n}$、weight-gathered 约 $4E\sqrt{BLF}/\sqrt{n}$；解释为什么 decode 用 weight-stationary、大批 prefill 用 weight-gathered；解释多查询注意力为什么要按批切。
7. 理解”低延迟和高吞吐的最优配置不同”，以及 LoongServe、Arctic 这类动态切换并行方式的系统要解决什么问题。

0.2 先修与本讲范围 Prerequisites & scope

0.3 本讲地图 Map

§1  推理为什么要并行：放不下、降延迟、提吞吐；通信的代价（α–β、ring all-reduce）  ← 代码 01
§2  三种基本并行：数据并行、张量并行（Megatron）、流水线并行；Alpa 的两层分类      ← 代码 02、03
§3  FlexFlow：SOAP 搜索空间、执行模拟器、MCMC 搜索、评测与论文卡片                  ← 代码 05
§4  Pope 等人：前馈层的 1D / 2D / weight-gathered 布局、多查询注意力按批切、
    PaLM 540B 的结果与论文卡片                                                     ← 代码 04、03
§5  选读：Alpa、LoongServe、Arctic Inference
§6  之后的发展
§7–11  小结、常见误区、练习、延伸阅读、下一讲预告

0.4 记号 Notation

0.5 关于数字 · About the numbers

本讲的数字有三类：(1) 虚拟设备上的真实计算（代码 02、04）：并行之后的输出与单卡核对，通信量由 CommLog 逐字节数出来，并与公式核对；(2) 纸上 GPU + α–β 通信模型（代码 01、03）：时间来自第 1 讲的 hw_model，通信用 α–β 模型，不与计算重叠（偏保守），互连的延迟 α 是粗略的假设值；(3) 迷你 FlexFlow（代码 05）：一个 4 个算子、4 块 GPU 的人造问题，任务时间用 roofline 公式代替实测。论文里的数字写成”论文报告”。

1. 推理为什么要并行 · Why parallelize inference

1.1 三个理由

放不下。 Llama-3-70B 的 BF16 权重有 141 GB，一块 80 GB 的 H100 放不下；Llama-3.1-405B 有约 812 GB，一台 8 卡 H100 机器（640 GB）都放不下。权重放得下还不够，KV cache 也要显存：代码 03 演示 1 算出，70B 用 2 块 H100 做张量并行时每块卡要放 70 GB 权重，只剩 1 GB 多给 KV cache，上下文 1024 的请求最多放 8 个——能跑，但几乎没法批处理。

降延迟。 batch 小时 decode 受限于读权重（第 1 讲）。把权重分到 $n$ 块卡上，每块卡只读 $1/n$，读权重的时间就降到 $1/n$。代码 03 演示 1：70B 在 2、4、8 块 H100 上，batch = 1 时每步读显存的时间是 26 ms、13 ms、6.5 ms。prefill 是 compute-bound 的，多卡同样能分担计算、降低 TTFT。

提吞吐。 两条路：一是让每个副本放下更大的批（权重分出去了，剩下的显存给 KV cache）；二是多放几个副本，各自服务不同的请求。

不同的理由对应不同的并行方式。这一讲的核心问题就是：给定模型、硬件和目标（延迟还是吞吐），怎样切？

1.2 切开就要通信

把一个模型切到多块卡上，卡之间就要交换数据。通信的代价用 α–β 模型 描述（前置讲 P5）：传一条 $D$ 字节的消息，时间约为 $\alpha + D/\beta$，其中 $\alpha$ 是启动延迟（发起一次传输的固定开销），$\beta$ 是带宽。消息小时 $\alpha$ 主导，消息大时 $D/\beta$ 主导。

不同的互连差别巨大。代码 01 用的数量级参数（有效带宽取公开规格的约 80%；延迟是粗略的假设值）：

节点内的 NVLink 比跨节点的 InfiniBand 快一个数量级，比普通以太网快两个数量级。并行方式要和这个层次结构相匹配：通信多的放在快链路上，通信少的才能跨慢链路（§2.3 Alpa 的核心观察）。

1.3 ring all-reduce

推理里最常见的集合通信是 all-reduce：$K$ 个设备各有一份同样形状的数组，结束时每个设备都拿到所有数组之和。张量并行每层都要做（§2.2）。

朴素的做法——每个设备把数据发给所有人——每个设备要发 $(K-1) \cdot D$ 字节，随设备数线性增长。ring all-reduce 把它拆成两步：

1. reduce-scatter：把数组切成 $K$ 块，设备排成一个环。$K - 1$ 步里，每个设备每步把一块”部分和”发给下一个设备，对方加到自己的同一块上。结束时设备 $i$ 拥有总和的第 $i$ 块。
2. all-gather：再走 $K - 1$ 步，把每个设备的完整块沿环复制给所有人。

代码 01 演示 1 用”这一块已经加进了哪些设备的数据”打印了 $K = 4$ 时 reduce-scatter 的每一步：3 步之后，设备 $i$ 的块 $i$ 恰好含有 $\{0, 1, 2, 3\}$ 全部数据。每一步每个设备只发 1 块，所有链路同时在忙，没有谁是瓶颈。

于是每个设备一共只发送 $2 \cdot \frac{K-1}{K} \cdot D$ 字节，几乎与设备数无关——这就是”带宽最优”。代码 01 演示 2 用 CommLog 实测（$D$ = 67 MB）：$K$ = 2、4、8、16 时每个设备发送 67.11 MB、100.66 MB、117.44 MB、125.83 MB，与公式完全一致；朴素做法在 $K = 16$ 时要发 1.01 GB。

但”与设备数无关”说的是字节数。ring 要走 $2(K-1)$ 步，每一步都要付一次 $\alpha$：

代码 01 演示 3 算了 8 卡一次 all-reduce 的时间（Llama-3-70B，$d_{model}$ = 8192，BF16）：

decode 时消息只有几十 KB，时间几乎全是启动延迟；70B 有 80 层、每层 2 次 all-reduce，一次 forward 160 次，batch 1 在 NVLink 上累计约 4.49 ms——和 70B 在 8 卡上读一遍权重（约 6.5 ms）是同一个量级。prefill 时消息有几十 MB，带宽成为主角，跨节点的 InfiniBand 上每次就要 1.5 ms。这两个数字解释了本讲大部分的结论。

2. 三种基本并行 · The three basic forms

2.1 数据并行：推理里就是”多放几个副本”

训练里的数据并行需要每一步同步梯度；推理里没有梯度，数据并行就是部署多个完整的副本，请求分给不同的副本，副本之间完全不通信。

它的优点是简单、没有通信开销、吞吐随副本数线性增长；缺点是每个副本都要放下整个模型（加上 KV cache），而且不降低单个请求的延迟。模型一块卡放不下时，数据并行必须和其他并行方式组合（每个副本本身是多卡的）。

2.2 张量并行：Megatron 的列切 + 行切

**张量并行（tensor parallelism，TP）**把每一层的矩阵切开，每块卡算一部分。关键是怎样切才能让通信最少。Megatron-LM（Shoeybi 等，2019）发现：对两个连续的矩阵乘法 $Y = \sigma(X W_1) W_2$，

• 把 $W_1$ 按列切成 $n$ 份：每块卡算出 $X W_1$ 的一部分列，$\sigma$（逐元素）可以在本地做；
• 把 $W_2$ 按行切成对应的 $n$ 份：每块卡用自己那部分列乘 $W_2$ 的对应行，得到最终结果的一个部分和；
• 一次 all-reduce 把部分和加起来。

两个矩阵乘法之间不需要任何通信。Transformer 的一层恰好有两组这样的结构：

• 注意力：$W_q, W_k, W_v$ 按列切——每块卡负责几个头，在本地算这几个头的注意力；$W_o$ 按行切，一次 all-reduce。
• MLP：$W_{gate}, W_{up}$ 按列切（中间维 $d_{ff}$ 切成 $n$ 段），$W_{down}$ 按行切，一次 all-reduce。

所以每层恰好 2 次 all-reduce，与 TP 度无关。代码 02 在 tiny_llm 上实现了这种切法（lm_head 按词表切，最后一次 all-gather 拼出完整的 logits）：

• TP = 2、4 时 logits 与单卡的最大差是 $1.1 \times 10^{-6}$（部分和的加法顺序变了），贪心生成的 80 个 token 完全相同；TP = 1 时逐位相同。和第 4 讲的 ⊕ 合并一样：并行在数学上精确，在浮点上只是误差范围内一致。
• 每次 forward 恰好 4 次 all-reduce（2 层 × 2）加 1 次 all-gather；通信量与 token 数成正比，TP = 4 时 prefill 54 个 token 每次 all-reduce 每个设备发 20.74 KB，decode 1 个 token 只发 384 B。

KV cache 怎么切。 注意力按头切，每块卡只存自己那几个头的 KV。但如果 KV 头比 GPU 少（GQA、MQA），KV 头不够分：每块卡的 query 头都需要一个完整的 KV 头，同一个 KV 头只能复制到几块卡上。代码 02 演示 3：

跟着复制的还有 $W_k$、$W_v$ 这两个投影矩阵：TP = 4 时三种模型每块卡的权重完全一样（103.17 KB）——KV 头少带来的节省在切分中消失了。Llama-3-70B 有 8 个 KV 头，TP 超过 8 时也会遇到这个问题。Pope 等人的解法是 decode 时按批切注意力（§4.3）。

2.3 张量并行的极限：通信不随卡数减少

张量并行降低了每块卡要读的权重和要做的计算，但每层 2 次 all-reduce 的次数不变，每次的消息大小也不变（每块卡都要拿到完整的激活），ring 的步数还随卡数增加。代码 03 演示 1 把这一点量化了（Llama-3-70B，H100，NVLink，上下文 1024）：

batch = 1 时延迟随 TP 下降，但通信从 0.65 ms 涨到 4.51 ms，TP = 8 时已经占了一步的约四成；batch = 64 时，每块卡每秒产出的 token 数随 TP 增加而下降。prefill 也一样（代码 03 演示 3，2048 个 token）：TP = 2、4、8 时 TTFT 是 222.53 ms、128.03 ms、82.71 ms，但”卡·秒”成本从 445 涨到 662，扩展效率降到 67%；如果 8 卡的环跨了节点（瓶颈是 InfiniBand），通信涨到 246 ms，TTFT 反而变成 298 ms。

所以张量并行通常只在一台机器的 NVLink 范围内做，而且 TP 度不是越大越好：追求延迟用更多卡，追求单位成本的吞吐用更少的卡——这正是 Pope 等人反复强调的取舍（§4）。

2.4 数据并行 vs 张量并行：同样 8 块卡

代码 03 演示 2 比较了一台 8 卡 H100 跑 Llama-3-8B 的四种配置（上下文 2048；每个副本的批大小受 KV 显存限制，最多 512）：

几个和直觉不完全一致的地方：

• 单请求延迟最低的是 TP=4，不是 TP=8：8B 模型在 8 块卡上每块只剩 2 GB 权重要读，省下的读显存时间抵不过 64 次 all-reduce 的启动延迟。prefill 是 compute-bound 的，TP=8 最快。
• 满载吞吐最高的是 TP=2 × DP=4，不是 DP=8。我原本预期”完全不通信”的 DP=8 最高，但每个 DP 副本都要放一整份 16 GB 的权重，剩下的 KV 显存只够 208 个上下文 2048 的序列，批不够大；TP=2 把权重分到两块卡上，每个副本能放 476 个序列。
• TP=8 吞吐最低，但满载时每个请求的延迟最低（11 ms 对 27 ms）。

这正是 Arctic Inference 论文开头描述的困境（§5.3）：延迟优化的部署用 TP=8，吞吐优化的部署用 TP=1、DP=8，两者 KV 布局不同，很难在线切换。

2.5 流水线并行：不缩短延迟，只提高吞吐

流水线并行（pipeline parallelism，PP）按层切：第 1 块卡负责前 1/n 的层，第 2 块卡负责接下来的 1/n……前一级算完，把激活点对点发给下一级。代码 02 演示 4：tiny_llm 切成 2 级，输出与单卡逐位相同（同样的运算、同样的顺序），prefill 54 个 token 只发送一次 13.82 KB 的激活；同一次 prefill 张量并行（TP = 2）每个设备要发 55.38 KB。

流水线并行的通信少得多（每 $L/n$ 层才发一次，而且是点对点），所以跨节点时常用流水线并行。它的代价在别处：

1. 单个批的延迟不会缩短：一个批必须依次走完所有级。
2. 要有多个批同时在流水线里才能让各级都忙。decode 时如果只有一个批，就只有一级在干活，其余都在等——这就是”气泡”（bubble）。
3. 推理特有的不均衡：每个微批里 prefill 和 decode 的比例不同，各级的时间不同，产生额外的气泡。Sarathi-Serve（第 7 讲）指出，连 Orca 的迭代级调度也没能消除它。

代码 03 演示 4 用 Llama-3.1-405B（约 812 GB，两台 8×H100）比较了跨节点的几种切法（上下文 1024）：

• 跨节点的张量并行每层 2 次 all-reduce 都要过 InfiniBand，延迟比”节点内 TP + 节点间 PP”高；满流水时后者的吞吐是前者的 2.5 倍以上。Sarathi-Serve 论文实测 Falcon-180B：跨节点 TP-8 的 decode 延迟是”节点内 TP-4 + 节点间 PP-2”的 2 倍以上，比我们模型里的差距（约 1.26 倍）大——跨节点小消息的实际启动延迟通常远大于我们假设的 5 µs。
• TP=4 × PP=4 每个批要走 4 级、付 4 次固定开销，延迟最高，但满流水时吞吐最高。
• 气泡：TP=8 × PP=2、batch 32，如果流水线里同时只有 1 个批，GPU 只有 50% 在忙，吞吐减半。

2.6 一个统一的视角：算子内 vs 算子间（Alpa）

Alpa（OSDI 2022）把各种并行重新分成两类：

• 算子内并行（intra-operator）：把一个算子沿某些维度切开分给多个设备——数据并行（切批）、张量并行（切参数）都属于这一类。设备利用率高，但每次切分和合并都要通信。
• 算子间并行（inter-operator）：把模型切成若干段（stage），流水线式地放在不同设备组上——流水线并行属于这一类。只在相邻段之间通信，量小，但有调度带来的空闲。

Alpa 的关键观察是：这两类并行的带宽需求不同，正好对应集群的层次结构——机器内的设备之间带宽高，适合算子内并行；机器之间带宽低，适合算子间并行。于是 Alpa 把集群分成若干”设备网格”（mesh），网格内做算子内并行、网格间做流水线，分两层自动求解（详见 §5.1）。前面几个实验得出的”节点内 TP、节点间 PP”，正是这个原则的特例。

除了这三种，还有两种推理里常用的并行：序列 / 上下文并行（把一个长序列的 token 切到多块卡上，§5.2 的 LoongServe、§5.3 的 Arctic 都用到）和 MoE 模型的专家并行（第 8 讲）。

3. FlexFlow：在巨大的策略空间里搜索 · FlexFlow

3.1 问题：只会”数据并行或模型并行”

2018 年前后的深度学习框架（TensorFlow、PyTorch、Caffe2、MXNet）默认用数据并行，把模型并行当作”手动指定设备”的备选。数据并行对参数少、计算多的算子（卷积）很好，对参数多的算子（全连接层）很差——每一步都要同步大量参数。模型并行省掉了参数同步，却要在算子之间搬数据，也没有算子内部的并行。专家设计的混合策略（例如卷积层用数据并行、全连接层用模型并行的”one weird trick”）好一些，但依赖专家，换一个模型或集群就要重新设计。已有的自动化方法只探索其中一个方向：REINFORCE 用强化学习学设备分配（只有算子间并行），OptCNN 优化单个算子内部的切法（只适用于线性的算子图）。

FlexFlow 的主张是：定义一个包含所有这些策略的更大的搜索空间，然后自动地在里面找。

3.2 SOAP 搜索空间

FlexFlow 把并行化一个算子定义为切分它的输出张量，每个设备负责输出的一个不相交的子集。可以切的维度分四类（SOAP）：

• Sample（样本）：切批维——就是数据并行；
• Attribute（属性）：切样本内部、不涉及参数的维度，例如图像的高宽、序列的长度；
• Parameter（参数）：切需要切分模型参数的维度，例如全连接层的输出通道——就是张量并行 / 模型并行；
• Operation（算子）：不同的算子放在不同的设备上并行执行——算子间并行。

一个算子的并行配置给出每个可切维度的并行度和每个子任务放在哪个设备上；一个并行策略为每个算子各选一个配置。数据并行、模型并行、专家设计的策略都是这个空间里的特例；空间的大小随算子数指数增长。

3.3 执行模拟器

空间这么大，不可能每个策略都在真实硬件上跑一遍（REINFORCE 就是这么做的，要 160 个节点跑 12–27 小时）。FlexFlow 用一个执行模拟器预测一个策略要多久，它依赖四个假设（论文 §5）：

• A1：每个任务的执行时间可预测、方差小，与输入数据的内容无关（稠密矩阵运算满足）；
• A2：在带宽为 $b$ 的连接上传 $s$ 字节需要 $s/b$ 的时间；
• A3：每个设备按先来先服务处理分给它的任务（GPU 就是这样）；
• A4：运行时开销可以忽略。

模拟分两步。先建任务图：每个算子按配置拆成若干任务；生产者和消费者任务在不同设备上时，插入一个通信任务，放在连接两个设备的”链路设备”上（把链路也当作设备，计算和通信才能重叠）。每类任务的执行时间只需对每种输入大小实测一次，之后复用——一个神经机器翻译模型有几百个算子，却只有四种不同的算子。然后运行完整模拟算法（Algorithm 1）：用一个按”就绪时间”排序的优先队列，依次取出任务，开始时间 = max(就绪时间, 它的设备上一个任务的结束时间)，更新后继任务的就绪时间，最后所有任务结束时间的最大值就是预测的执行时间。

由于搜索每次只改一个算子的配置，FlexFlow 还设计了增量模拟（delta simulation，Algorithm 2）：只重算时间线上受影响的那部分任务，结果与完整模拟相同，快 2.2–6.9 倍，设备越多加速越大。

3.4 MCMC 搜索

有了模拟器当”神谕”，搜索就变成最小化预测时间。找最优策略是 NP 难的，FlexFlow 用 Markov 链蒙特卡洛（MCMC）：令 $p(S) \propto \exp(-\beta \cdot \text{cost}(S))$，用 Metropolis–Hastings 算法——每次随机选一个算子，把它的配置换成随机的另一种；新策略更好就接受，更差就以 $\exp(\beta(\text{cost}(S) - \text{cost}(S^*)))$ 的概率接受。它大多数时候像贪心搜索，偶尔”往上走”以跳出局部最优。起点用数据并行、专家策略和随机策略；搜索预算用完、或者一半预算内没有改进时停止。

3.5 评测

论文在两个 GPU 集群（4 节点 × 4 块 P100，节点内 NVLink；16 节点 × 4 块 K80，PCIe 连接不对称）上评测了 6 个模型（AlexNet、Inception-v3、ResNet-101 和三个 RNN）（论文报告）：

• 相对数据并行和专家设计的策略，训练吞吐提高 1.3–3.3 倍（ResNet-101 上 FlexFlow 找到的接近数据并行，收益很小）；在 64 块 K80 上训练 NMT，每次迭代快 1.7–2.4 倍、数据传输少 2–5.5 倍。Inception-v3 端到端训练时间比 TensorFlow 少 38%。
• 相对 REINFORCE 快 3.4–3.8 倍，搜索只要 14–40 秒；相对 OptCNN 快 1.2–1.6 倍（差别来自算子间并行）。
• 模拟器：预测与实测的相对误差都在 30% 以内，而且保持了不同策略之间的快慢顺序——对搜索来说这才是关键。
• 搜索质量：在两个搜索空间约 $10^{11}$ 的小网络上，FlexFlow 找到了全局最优（穷举分别花了 0.8 和 18 小时）；在大网络上，找到的策略都是局部最优。
• 案例：Inception-v3 在 4 块 P100 上，最优策略对关键路径上的算子用算子内并行、对不同分支用算子间并行，参数同步开销降低 75%、每次迭代快 12%；NMT 中参数多但计算少的嵌入层放在少数几块卡上，参数多且计算重的 softmax 层按通道切。

3.6 迷你 FlexFlow

代码 05 实现了一个缩小版：4 个算子（词嵌入 → 两个并行的塔 → 输出头），4 块 GPU，(0,1) 和 (2,3) 之间是快链路、跨对的是慢链路。每个算子 11 种配置（按行切、按参数切、两者混合，放在 1、2 或 4 块卡上），共 14,641 种策略。和论文的一个区别：FlexFlow 优化的是训练，这里改成推理的前向，并给每块卡加了显存上限（64 MB）——数据并行要复制权重，放不下；参数切分省显存，但要搬激活。

放得下的 3,307 种策略里，执行时间从 1.07 ms 到 6.39 ms 不等（中位数 3.29 ms）。全部参数切分时，GPU 只忙了 0.95 ms，跨对的慢链路却忙了 2.88 ms。最优策略在一个算子内部混合两种切法（2 路按行 × 2 路按参数），并让参数切分恰好落在快链路连着的那一对卡上：两个按行切出来的组各自独立，组内两块卡分摊权重、只在快链路上交换激活。这种策略既不是数据并行，也不是张量并行——正是 SOAP 空间里”混合并行”的意义。

MCMC 搜索（β = 1/(0.1 ms)）从”全部参数切分”或随机起点出发，6 条链都在 1000 次评估之内找到了穷举的最优，最快 106 次、最慢 591 次，每条链只看了 180–362 种不同的策略。

3.7 为什么推理系统也关心 FlexFlow

FlexFlow 是一个训练系统，但它的三个想法在推理里处处可见：

• 把并行策略当作搜索问题：Alpa 自动化了训练和推理的并行；DistServe（第 7 讲）枚举 prefill 和 decode 实例的并行配置；Splitwise 搜索 prompt 机器和 token 机器的数量。
• 用模拟器代替实测：DistServe 用模拟器估计 SLO 达标率并二分查找最大请求率（与实测误差不到 2%），Splitwise 用事件驱动的集群模拟器，Sarathi-Serve 用 Vidur 选择 token 预算。它们都依赖 FlexFlow 的假设 A1：神经网络的执行时间高度可预测。（第 1 讲的 Clockwork 也建立在同一个观察上。）
• SpecInfer（第 5 讲）就建在 FlexFlow 之上，用它的运行时实现张量并行 + 流水线并行。

4. Pope 等人：Transformer 推理的分区 · Efficiently Scaling Transformer Inference

FlexFlow 用通用的搜索解决一般的网络；Pope 等人走了另一条路：针对 Transformer 推理，用简单的解析模型推导出最优的分区方式，让用户能直观地理解取舍，而不是黑箱搜索。他们的实验平台是 TPU v4（每芯片 275 TFLOPS 的 BF16、32 GiB HBM、1200 GB/s 显存带宽、270 GB/s 的芯片间互连，3D 环面拓扑）和 PaLM 系列模型。

4.1 推理的代价：显存时间、计算时间、通信时间

论文把一次 forward 的时间拆成三部分：

• 显存时间：把权重和 KV cache 从 HBM 读进计算单元。batch 小、序列短时读权重占主导；batch 和序列长度都大时（例如 2048 个 token × batch 512）读 KV cache 占主导。
• 计算时间：$N$ 个参数的模型每个 token 约 $2N$ 次矩阵乘法运算；注意力的运算量相对较小，但它的 KV cache 占了大量显存容量和带宽。
• 通信时间：芯片之间交换激活或权重。

前两项都随芯片数近似线性下降，通信时间下降得慢甚至不降，所以芯片越多，通信越成为瓶颈。论文还指出两个极端场景：要求最低延迟时，要用很多芯片、小 batch，代价是 MFU 低、每 token 成本高；离线推理不在乎延迟时，应该用大 batch 提高 MFU。一个直观的例子：500B 以上的多头注意力模型，batch 512、上下文 2048 时 KV cache 有 3 TB，是参数的 3 倍。

衡量成本时，论文用”芯片·秒每 token”：$n_{chips} \times$ 时间 ÷ token 数，它与 MFU 成反比。

4.2 前馈层的三种布局

1D weight-stationary（Megatron）。 权重按 $F$ 维切到所有芯片上（$EF_{xyz}$、$F_{xyz}E$），权重不动，激活在芯片之间移动：一次 all-gather 让每个芯片拿到完整的输入激活，乘本地的权重分片，再一次 reduce-scatter 把输出切回去。每个芯片的通信量约 $2BLE$ 个元素，与芯片数无关——芯片越多，通信占比越高。

2D weight-stationary。 把每个 $E \times F$ 的权重矩阵沿 $E$ 和 $F$ 两个方向切，让每个芯片拿到一块”接近正方形”的权重（例如 $E$ = 1024、$F$ = 4096、64 个芯片：$E$ 切 4 份、$F$ 切 16 份，每块 256 × 256）。在连续的矩阵乘法之间交替地在两个轴上聚合激活，每个芯片的激活从不完整复制。通信量约 $2BL(E/X + F/YZ)$；$F = 4E$ 时最优取 $X = 0.5\sqrt{n}$、$YZ = 2\sqrt{n}$，得到

随 $\sqrt{n}$ 下降。代价是每层 4 次集合通信而不是 2 次，所以论文说当 $\sqrt{n} > F/E$（通常 $n > 16$）时 2D 才更省。

代码 04 演示 1 在虚拟设备上真实运行了两种布局（$E$ = 64、$F$ = 256、$BL$ = 32，输出都与单卡核对），用 CommLog 数出每个芯片的通信量：

实测值与带 $(K-1)/K$ 因子的精确公式完全相等。论文的近似公式省掉了这个因子：$K$ 大时无所谓，但 $n = 4$ 时最优的 $X = 1$，$x$ 方向的通信根本不存在，近似公式却还算了 $F/YZ$ 那一项，把 2D 高估到 8,192（实测 3,072，与 1D 相同）；$n = 16$ 时近似公式说两者持平，精确值是 0.73。设备数小时要用精确公式。

weight-gathered。 batch 很大时，激活比权重还大，与其搬激活，不如搬权重、让激活不动：每层开始前把权重 all-gather 到 $N$ 个芯片上，激活按批切。权重的通信量是 $2EF \cdot N/n$，激活的是 $2BLE/N$；取 $N = \sqrt{BL \cdot n / F}$ 时总通信量最小：

weight-stationary 的通信量与 token 数 $BL$ 成正比，weight-gathered 只与 $\sqrt{BL}$ 成正比，所以 token 数足够大时后者更省。根据 $N$ 取 $X$、$XY$ 或 $XYZ$，有 X-、XY-、XYZ-weight-gathered 三种变体。代码 04 演示 2 用论文图 3 的参数（$E$ = 16384、$F$ = 65536、$X = Y = Z = 4$，64 个芯片）算出每个芯片的通信量（BF16）：

分界点分别在约 16,384、65,536、1,048,576 个 token，与论文图 3 的趋势一致。

于是论文的策略是：decode（每批 token 数小）用 2D weight-stationary；prefill 按这一批的 token 数在 weight-stationary 和几种 weight-gathered 之间选。几种布局共享同一种权重切分（$E_x F_{yz}$），切换时不需要搬权重。论文图 7：PaLM 540B 在 64 个芯片上做 prefill，batch 从 2048 个 token 增大到 100 万个 token 时，最优布局从 2D weight-stationary 换成 weight-gathered，最高达到 76% 的 MFU。

4.3 注意力：多查询注意力要按批切

多查询注意力（multiquery attention，MQA）让所有 query 头共用一个 K、V 头，KV cache 缩小 $n_{heads}$ 倍（GQA 是它和多头注意力之间的折中）。可是它去掉了一个可以并行切分的维度：如果像多头注意力那样按头切，唯一的 KV 头就得复制到每个芯片上，省下的显存全丢了（§2.2 的问题的极端情况）。

论文的解法：decode 时把注意力按批切——Q、K、V 都按批维切成 $n$ 份，每个芯片只存 $1/n$ 的序列的 KV。代价是要用 all-to-all 把 Q、K、V 从”按头切”的布局换成”按批切”、算完再换回来。decode 时每个序列只有 1 个新 token，Q、K、V 是很小的张量，而 KV cache 有上千个 token：花在小张量上的通信换来大张量上的显存时间，非常划算。prefill 时情况相反（很多 query 共用同一份 K、V，读 KV 不是瓶颈），所以 prefill 仍然按头切。

代码 04 演示 4 复现了论文表 1（PaLM 540B，64 个芯片，每芯片 32 GiB，30% 给 KV cache；多头变体按论文的做法把头数补齐到 64、$d_{head}$ 缩到 128，使注意力的参数量不变）：

按头切的多查询注意力能放下的上下文反而只有多头的一半（$d_{head}$ 是 256），按批切则长 32 倍以上。一层 decode（batch 512、上下文 2048）里，按头切每个芯片要读 1.07 GB 的 KV（约 1.12 ms），按批切只读 16.78 MB（约 17 µs），额外的 all-to-all 每个芯片只发约 532 KB（约 6.5 µs）。论文图 8：上下文越长，优化后的布局越快；多查询注意力加上按批切，支持到 8192–32768 个 token（batch 512 和 128），注意力只占总时间的 8%–31%。

4.4 其他优化

• 并行的注意力 / 前馈层（PaLM 的结构：两者都从同一个归一化后的输入算出、再相加）：每层只有一个 LayerNorm；前馈层的输入矩阵可以和 $W_Q$ 融合、输出矩阵可以和 $W_O$ 融合，每层少一次 all-reduce。换回串行结构，decode 的延迟高 14%。
• 通信与计算重叠：用 Looped CollectiveEinsum 把大部分 reduce-scatter 和 all-gather 藏在矩阵乘法后面，比最初的编译器自动分区快约 1.4 倍。（我们的代码 03 假设通信完全不重叠，所以通信的占比偏高。）
• int8 权重量化：低 batch 时读权重占主导，int8 把它减半；PaLM 540B 在 batch 64 时每 token 28.5 ms（BF16 是 36.9 ms）。
• 补齐头数：PaLM 540B 的 48 个头补到 64 个，多了 18B 参数、损失 3% MFU，但在 64+ 个芯片上切得更好，收益更大。

4.5 结果

PaLM 540B 在 64 个 TPU v4 上的几个配置（论文表 2，上下文 2048；prefill 指处理 2048 个输入 token，decode 指生成 64 个 token）：

论文强调的几点：

• 低延迟和高吞吐用的配置不同。低延迟时 prefill 用 batch 1、decode 用 batch 64（batch 增大到 64 对 decode 延迟影响很小，MFU 却高得多）；作者指出，这可以通过”一个 batch 1 的 prefill 服务器流水线式地喂给一个 batch 64 的 decode 服务器”实现——这已经是第 7 讲解耦的雏形。
• 对话场景：处理 64 个新输入 token、查阅 1920 个 token 的缓存历史、生成 64 个 token，共 1.9 秒；离线场景：1984 个输入 + 64 个输出，整体 FLOPS 效率 73%。
• 与 FasterTransformer 比较（以 MFU 归一化硬件差异）：FasterTransformer 的 32 路张量并行最高只有 33% MFU，16 路是 46%，说明它在 32 路时被通信卡住；Pope 等人的 2D 布局在 64 路时仍有 44% MFU。
• 低 batch 的延迟随模型大小大约按平方根增长：大模型读的权重多，但能切到更多芯片上才被通信卡住。

我们的代码 03 在 GPU 上得到了同样的定性结论：TP 越大，单请求延迟越低（直到通信占主导），每块卡的吞吐越低；8B 模型在 8 卡上，延迟最低的配置（TP=4）和吞吐最高的配置（TP=2 × DP=4）不同。

5. 选读：Alpa、LoongServe、Arctic Inference · Optional readings

5.1 Alpa：自动化的两层并行（OSDI 2022）

Alpa 自动生成统一了数据、算子、流水线并行的执行计划。它的出发点是 §2.6 的两层分类：算子内并行（切张量的某些维）和算子间并行（把模型切成段做流水线），并观察到它们的带宽需求恰好对应集群的层次结构。Alpa 把集群划分成若干设备网格（网格内的设备之间带宽高），把计算图划分成若干段，分别在两层上近似最优地求解（论文用整数线性规划求算子内的切分、用动态规划求段的划分），再由运行时协调两层的执行。论文报告：Alpa 生成的计划在 Megatron 等系统专门设计的模型上与手工调优相当甚至更好，而且能推广到结构异构、没有手工方案的模型。Alpa 主要面向训练，但”按带宽层次分配并行方式”的原则在推理部署中同样适用。

5.2 LoongServe：弹性序列并行（SOSP 2024）

长上下文模型（例如支持 100 万 token 的模型）让请求之间、以及同一个请求的 prefill 和 decode 之间，资源需求差异巨大：论文举例，一个 100 万 token 的输入在 LWM 模型上光 KV cache 就要 488 GB。静态的并行配置无法适应：要么为最长的请求配置，浪费在短请求上；要么把 GPU 分组、每组负责一段长度范围或一个阶段，组之间的显存无法共享（碎片），切换阶段时还要迁移全部 KV。

LoongServe 提出弹性序列并行（ESP）：每一个迭代都重新决定每个请求的并行度。例如 prefill 时让一个长请求用上所有 GPU 快速算完，进入较轻的 decode 后降低并行度、释放资源给别的请求。为了让伸缩没有额外开销，它在 prefill 时”主动缩容”（复用 prefill 本身的通信完成 KV 的重新分布），在 decode 时用”多主 decode”（不迁移已有的 KV，把部分注意力的通信与计算重叠），并以单个 token 为粒度在实例之间管理 KV，消除碎片。论文报告，吞吐比分块 prefill 最多高 3.85 倍、比 prefill / decode 解耦最多高 5.81 倍。

5.3 Arctic Inference：Shift Parallelism（Snowflake，2025）

Arctic Inference 是一个 vLLM 插件，它要解决 §2.4 的困境：真实的推理负载是突发的，延迟（TTFT、TPOT）和吞吐不能兼得——延迟优化的部署用 TP=8，吞吐优化的部署用 TP=1、DP=8，两者 KV cache 的布局不兼容，在线切换要大量搬数据，很多团队只好部署两套。

Shift Parallelism 在两种模式之间按负载切换：小批时用张量并行（最低的 TPOT）；大批时用 Arctic 序列并行（Ulysses 式：把输入序列切到各 GPU 上，注意力处按头重新分布），它避免了张量并行逐 token 的通信，又能充分利用 GPU。关键是 KV cache 布局不变：只要 SP × TP = GPU 总数，两种模式下每块 GPU 负责的注意力头完全相同，所以可以在相邻两次 forward 之间无缝切换；再通过仔细地映射 TP 的 rank，让小 TP 需要的参数分片恰好包含在大 SP 所需的分片里。

论文报告（Llama 3.3 70B，8×H200，结合了它的推测解码和 SwiftKV）：相对吞吐优化的部署，请求完成快 3.4 倍、吞吐高 1.06 倍；相对延迟优化的部署，吞吐高 1.7 倍、请求完成快 1.28 倍。在动态负载下，中位 TTFT 从 1355 ms 降到 148 ms，中位 TPOT 从 83 ms 降到 51 ms，最大吞吐只比最好的配置低不到 10%（75.5K → 69.1K token/s）。

LoongServe 和 Arctic 的共同点：并行方式不再是部署时一次定死的，而是随负载、随阶段在线调整——FlexFlow 的静态搜索在推理里不够用了。

6. 之后的发展 · Since then

（这一节是我的补充，不是原课内容。）

• 更大的 NVLink 域：GB200 NVL72 把 72 块 GPU 放在同一个 NVLink 域里，“张量并行只能在 8 卡内做”的限制放宽了，Pope 等人为 TPU 大环面设计的 2D 布局这类思路在 GPU 上也有了用武之地。
• 专家并行成为 MoE 推理的主角：DeepSeek-V3 等 MoE 模型的部署把注意力和专家用不同的并行方式（例如注意力用数据并行、专家用专家并行），中间用 all-to-all 交换 token——第 8 讲详细讨论。这和 Pope 等人”注意力按批切”的思想一脉相承。
• 上下文并行 / 序列并行：Ring Attention 等方法把超长序列切到多块卡上，按环传递 KV 块；它们和第 4 讲的 ⊕ 合并是同一个数学（分段算注意力状态再合并）。
• 推理引擎里的并行选项：vLLM、SGLang、TensorRT-LLM 都支持 TP、PP、DP、EP 的组合；Dynamo（第 7 讲）等框架让 prefill 和 decode 使用不同的并行配置。
• 自动并行编译器：GSPMD / JAX 的分片注解让用户只写”哪个张量按哪个维度切”，编译器自动插入集合通信——Pope 等人的实现就基于这种分片表示。

7. 小结 · Summary

1. 推理需要多卡的三个理由：放不下（70B 的 BF16 权重 141 GB）、降延迟（每块卡只读 1/n 的权重）、提吞吐（腾出 KV 显存放更大的批，或多放副本）。不同理由对应不同的并行方式。
2. 通信的代价用 α–β 模型描述。ring all-reduce 每个设备只发约 2D 字节（与设备数几乎无关），但要走 2(K − 1) 步。decode 时消息小、受限于启动延迟，prefill 时消息大、受限于带宽；节点内 NVLink 比跨节点快一个数量级（代码 01）。
3. Megatron 张量并行：列切 + 行切，每层 2 次 all-reduce；输出与单卡在浮点误差内一致（代码 02）。KV 头少于 GPU 数时 KV 和 $W_k$、$W_v$ 被复制。通信不随卡数减少，TP 度越大，每卡吞吐越低（代码 03）。
4. 数据并行（多副本）不通信但不降延迟；流水线并行通信少、适合跨节点，但不缩短单个批的延迟、需要多个批填满流水线。8 卡跑 8B 时，延迟最低（TP=4）和吞吐最高（TP=2 × DP=4）的配置不同；跨节点时”节点内 TP + 节点间 PP”优于跨节点 TP（代码 03）。Alpa 把这概括为：算子内并行放在高带宽设备之间，算子间并行跨低带宽链路。
5. FlexFlow：SOAP 策略空间 + 基于任务图的执行模拟器 + MCMC 搜索，比数据 / 模型并行快最多 3.8 倍（论文报告）。迷你 FlexFlow 里，最优策略是算子内的行 × 参数混合切分，比”全部参数切分”快 3.4 倍，MCMC 在几百次评估内找到它（代码 05）。
6. Pope 等人：1D weight-stationary 通信 ≈ $2BLE$（不随芯片数下降）、2D ≈ $8BLE/\sqrt{n}$、weight-gathered ≈ $4E\sqrt{BLF}/\sqrt{n}$；decode 用 2D WS，大批 prefill 用 weight-gathered（76% MFU）；多查询注意力 decode 时按批切，上下文长 32 倍以上（代码 04 复现了论文表 1）。低延迟和高吞吐的最优配置不同——这引出了第 7 讲的解耦和 §5 的动态并行。

8. 常见误区 · Common pitfalls

1. “卡越多越快。” 读权重和计算随卡数下降，通信不降、甚至因为步数增加而上升。8B 模型在 8 卡上 TP=8 的单请求延迟反而高于 TP=4（代码 03 演示 2）。
2. “ring all-reduce 的时间与设备数无关。” 字节数几乎无关，步数是 2(K − 1)；decode 时消息小，时间主要是启动延迟，随设备数线性增长。
3. “数据并行完全不通信，所以吞吐一定最高。” 每个副本都要放一整份权重，KV 显存可能不够，批放不大（代码 03 演示 2：DP=8 的吞吐低于 TP=2 × DP=4）。
4. “流水线并行能降低延迟。” 不能。一个批必须依次走完所有级，每级还有固定开销；流水线只提高吞吐，而且要有足够多的批填满它（代码 03 演示 4）。
5. “张量并行可以跨节点随便做。” 每层两次 all-reduce 在 InfiniBand 上非常慢（代码 03 演示 3：8 卡跨节点时 prefill 的通信从 31 ms 涨到 246 ms）。跨节点优先用流水线并行。
6. “GQA / MQA 的 KV 节省在任何并行方式下都成立。” 按头切、KV 头比 GPU 少时，KV（和 $W_k$、$W_v$）被复制，节省消失（代码 02 演示 3）。要么按批切注意力（Pope 等人），要么控制 TP 度不超过 KV 头数。
7. “论文的近似通信公式在任何规模都准确。” 它们省掉了 $(K-1)/K$；设备数小时误差很大，例如 $n = 4$ 时 2D 布局的近似公式高估了 2.7 倍（代码 04 演示 1）。
8. “找到一个最好的并行配置，部署之后就不用管了。” 推理负载在变（批大小、输入长度、prefill / decode 比例），最优配置也在变。这正是 LoongServe、Arctic 和第 7 讲的 Dynamo Planner 要解决的问题。

9. 练习 · Exercises

练习 2、3、5 的参考答案里的数字是用 parallel_sim.py 和 hw_model.py 的函数实际算出来的。

练习 1（Megatron 的切法）。一个 MLP：$Y = \text{ReLU}(X W_1) W_2$，$X$ 是 $T \times d$，$W_1$ 是 $d \times 4d$，$W_2$ 是 $4d \times d$，切到 $n$ 块卡上。(a) 如果 $W_1$ 按行切、$W_2$ 按列切，会出什么问题？(b) 按 Megatron 的切法，需要几次通信、每次多少数据？

(a) $W_1$ 按行切意味着每块卡只有 $X$ 的一部分列乘 $W_1$ 的一部分行，得到的是 $XW_1$ 的部分和；ReLU 不是线性的，$\text{ReLU}(a + b) \ne \text{ReLU}(a) + \text{ReLU}(b)$，所以必须先 all-reduce 求和、再做 ReLU，中间多一次通信。(b) $W_1$ 按列切、$W_2$ 按行切：ReLU 在本地做，最后一次 all-reduce，数据量 $T \times d$ 个元素（每块卡发约 $2(n-1)/n \cdot T d$ 个）。这就是”列切在前、行切在后”的道理。

练习 2（all-reduce 的时间）。用 t_all_reduce 算：Llama-3-8B（$d_{model}$ = 4096）在 4 块 H100（NVLink）上做张量并行，decode batch 1 时一次 all-reduce 要多久？一次 forward（32 层）的通信总时间呢？和读权重的时间（16 GB / 4 卡）比如何？

消息 $4096 \times 2 = 8192$ 字节。$t = 2 \times 3 \times 2\,\mu s + 2 \times \frac{3}{4} \times 8192 / (360 \times 10^9) \approx 12.03\,\mu s$，几乎全是启动延迟。32 层 × 2 次 = 64 次 ≈ 0.77 ms（再加 lm_head 的一次 all-gather，很小）。读权重：16.06 GB ÷ 4 ÷ (3.35 TB/s × 0.8) ≈ 1.50 ms。通信已经是读权重时间的一半，所以 8B 模型用 TP=8 的收益很小（代码 03 演示 2）。

练习 3（KV 显存的复制）。Llama-3-70B 有 8 个 KV 头（每个头 128 维，80 层，BF16）。按头切做张量并行，TP = 8 和 TP = 16 时每块卡上每个 token 的 KV 是多少字节？如果改成按批切（TP = 16 的 16 块卡各负责 1/16 的序列）呢？

全模型每 token 的 KV = $2 \times 80 \times 8 \times 128 \times 2$ = 327,680 字节。TP = 8：每卡 1 个 KV 头，40,960 字节。TP = 16：KV 头不够分，每卡仍存 1 个头（复制 2 倍），还是 40,960 字节——多出来的 8 块卡没有带来任何 KV 显存的节省。按批切：每卡只存 1/16 的序列的全部 KV，平均每 token 20,480 字节。

练习 4（流水线的气泡）。PP = 4，每级时间相同，流水线里同时有 $M$ 个批。(a) $M$ = 1、2、4、8 时 GPU 的忙碌比例是多少？(b) 为什么 decode 时 $M$ 不容易做大？(c) 每个批里 prefill 和 decode 的比例不同，会怎样？

(a) $\min(M, 4)/4$：25%、50%、100%、100%。(b) 每个批 decode 一步之后要等它的下一个 token 才能做下一步；要让 4 级都忙，至少要有 4 个独立的批轮流进入，这要求有足够多的请求，而且每个批变小了（批越小，decode 的效率越低）。(c) 含 prefill 的批在每一级都比纯 decode 的批慢得多，后面的级要等它，前面的级算完又要等后面腾出来，产生额外的气泡。Sarathi-Serve 用分块 prefill 让每个批的计算量接近，从而减少气泡（第 7 讲）。

练习 5（weight-gathered 的分界点）。用 §4.2 的公式，在 $E$ = 8192、$F$ = 28672（Llama-3-70B 的 $d_{model}$、$d_{ff}$）、$n$ = 64 时，2D weight-stationary（$X = 4$、$YZ = 16$）和 X-weight-gathered（$N = 4$）的分界点在多少个 token？这时的 prefill 是什么规模？

2D WS：$2BL(E/4 + F/16) = 2BL(2048 + 1792) = 7680\,BL$；X-WG：$2EF \cdot 4/64 + 2BLE/4 = 29{,}360{,}128 + 4096\,BL$。相等时 $3584\,BL = 29{,}360{,}128$，$BL \approx 8192$ 个 token——相当于 4 个 2048 长的提示一起 prefill。所以在 64 个芯片上，只要一批 prefill 超过几千个 token，搬权重就比搬激活省；decode（每个序列 1 个 token）远在分界点以下。

练习 6（多查询注意力：all-to-all 划算吗）。decode 时，batch $B$、上下文 $C$、$H$ 个 query 头、1 个 KV 头、$d_{head}$ = $d$、$n$ 个芯片。按批切要做两次 all-to-all（交换 Q、K、V，再换回输出），每次大约交换 $B(H+2)d$ 个元素；省下的是每个芯片少读 $(1 - 1/n) \cdot B C \cdot 2d$ 个元素的 KV。什么条件下划算？

比较数量级：通信约 $B(H + 2) d$，省下的读取约 $2BCd$（$n$ 大时）。只要 $C \gg H/2$（上下文远长于头数的一半），省下的读取就远大于通信量；再考虑互连带宽（TPU v4 的 270 GB/s）低于 HBM 带宽（1200 GB/s）约 4.4 倍，条件大约是 $C \gg 2H$。PaLM 540B 的 $H = 64$，上下文几百以上就划算了（代码 04 演示 4：上下文 2048 时 6.5 µs 换 1.1 ms）。prefill 时每个序列有很多 query，KV 的读取被分摊，这个不等式不成立，所以 prefill 按头切。

练习 7（迷你 FlexFlow）。在代码 05 里把每块 GPU 的显存上限从 64 MB 改成 200 MB（权重全部复制也放得下），你预期最优策略会变成什么？为什么？（提示：推理的前向里，数据并行需要通信吗？）

预期最优策略变成”全部按行切 4 路”（纯数据并行）：每块卡处理 1/4 的 token，所有算子的输入都在本地，完全不需要通信，计算也平均分配。这说明 FlexFlow 式搜索的价值取决于约束：推理时数据并行唯一的代价是显存（复制权重），显存够时它就是最优的；训练时数据并行还要同步梯度，这才是 FlexFlow 论文里参数多的层不该用数据并行的原因。可以自己改常数验证（改的是脚本的副本，不要改课程文件）。

练习 8（思考：为什么 prefill 和 decode 想要不同的并行方式）。结合 §2.3 和 §4.5，说说 prefill 和 decode 各自喜欢什么样的张量并行度、批大小和布局。如果它们必须在同一组 GPU 上用同一种配置，会损失什么？

prefill 是 compute-bound 的、受 TTFT 约束：喜欢较大的 TP（分担计算、降低延迟）、小 batch（batch 1 已经能用满算力）、token 多时用 weight-gathered。decode 是 memory-bound 的、受 TPOT 约束：喜欢大 batch（分摊读权重）、TP 度够用即可（太大则通信和固定开销占比高）、注意力按批切。放在同一组 GPU 上，就只能选一个折中的配置，而且两种工作会互相干扰（长 prefill 让 decode 卡顿）。Pope 等人已经提出”batch 1 的 prefill 服务器喂给 batch 64 的 decode 服务器”；第 7 讲的 DistServe、Splitwise、Dynamo 把这个想法做成了完整的系统。

10. 延伸阅读 · Further reading

本讲的必读论文

• Z. Jia, M. Zaharia, A. Aiken. Beyond Data and Model Parallelism for Deep Neural Networks. SysML 2019（FlexFlow）。
• 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。

选读

• L. Zheng, Z. Li, H. Zhang et al. Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning. OSDI 2022。
• B. Wu, S. Liu, Y. Zhong, P. Sun, X. Liu, X. Jin. LoongServe: Efficiently Serving Long-Context Large Language Models with Elastic Sequence Parallelism. SOSP 2024。
• S. Rajbhandari, M. Hidayetoglu, A. Qiao, Y. Wang, J. Yang, J. Rasley, M. Wyatt, Y. He. Arctic Inference with Shift Parallelism: Fast and Efficient Open Source Inference System for Enterprise AI. 2025。

相关工作

• M. Shoeybi et al. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. 2019（列切 + 行切）。
• Y. Xu et al. GSPMD: General and Scalable Parallelization for ML Computation Graphs. 2021（Pope 等人使用的分片表示）。
• N. Shazeer. Fast Transformer Decoding: One Write-Head is All You Need. 2019（多查询注意力）。
• Y. Huang et al. GPipe. 2019；D. Narayanan et al. PipeDream. 2019（流水线并行与气泡）。

背景

• 前置讲 P5（分布式与通信基础）：互连、α–β 模型、集合通信算法、并行策略的概念。

11. 下一讲预告 · Next lecture

这一讲反复出现一个结论：prefill 和 decode 想要的东西不一样——prefill 受限于算力、在意 TTFT，decode 受限于带宽、在意 TPOT；它们喜欢的批大小、并行度、布局都不同。可在第 2 讲的连续批处理里，它们挤在同一批、同一组 GPU 上：一个长 prefill 插进来，所有正在 decode 的请求都要等它（第 2 讲代码 04 已经看到 ITL 的尖峰）。第 7 讲：Prefill 与 Decode 读三篇论文和一篇技术博客：Sarathi-Serve 把长 prefill 切成小块，和 decode 拼在一起，用一个 token 预算让每一步的时间可控（“不停顿的批”）；DistServe 干脆把 prefill 和 decode 放到不同的 GPU 上，各自选择并行方式和资源，以 goodput 为目标搜索放置方案；Splitwise 进一步让两个阶段用不同代际的硬件；NVIDIA Dynamo 把解耦做成了一个生产框架。我们会在 tiny_llm 上验证分块 prefill 与整段 prefill 输出一致，并在模拟器里比较”共置 / 分块 / 解耦”在不同 SLO 下的 goodput。