第 18 讲 · FlashAttention 与高效解码 · Efficient Attention (FlashAttention) & Efficient Decoding

对应 Maps to：CMU 10-423 / 10-623 / 10-723 Generative AI（Aran Nayebi、Matt Gormley），Lecture 18 · Mon 23-Mar · Flash Attention / Efficient decoding strategies（slides：lecture18-efficient）。 阅读材料：Milakov & Gimelshein (2018), Online normalizer calculation for softmax；Dao et al. (2022), FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness；Dao (2023), FlashAttention-2；Ivanov et al. (2021), Data Movement Is All You Need；Kwon et al. (2023), Efficient Memory Management for Large Language Model Serving with PagedAttention（vLLM）；Leviathan et al. (2023), Fast Inference from Transformers via Speculative Decoding；Chen et al. (2023), Accelerating Large Language Model Decoding with Speculative Sampling。 依据说明：本讲按原课第 18 讲公开的 slides 展开：为什么关心 FlashAttention（精确注意力，关键指标是运行时间）；背景一：矩阵乘法的分块，以及”softmax 不满足结合律”；背景二：普通 softmax、安全 softmax、在线 softmax（“三遍变两遍，表面上 1.33 倍、实际上 1.3 倍”）；FlashAttention：GPU 的存储层次、Transformer 训练是访存受限的、算子融合、标准注意力的读写、两个关键想法（分块和重计算）、两层网络的重计算、结果；高效解码：虚拟内存、PagedAttention、推测解码。所有实验是本讲义自己设计的，不代表课堂上讲了什么。 配套代码 Code：code/lecture18/（5 个脚本 + 2 个模块；每个都在十几秒以内） 预计阅读时间：2.5 小时（含运行代码）

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

0.1 学习目标 Learning goals

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

1. 解释为什么矩阵乘法可以分块、softmax 不能直接分块，以及在线 softmax 的 $(m, \ell)$ 怎样补上这一环。
2. 写出 FlashAttention 前向的循环：每块 $K_j, V_j$、每块 $Q_i$ 要算什么、要修正什么、往显存写什么。
3. 数出标准注意力和 FlashAttention 读写显存（HBM）的次数，说明 FlashAttention 省的是常数倍的读写、真正”从平方变成线性”的是额外显存；说出它需要 $M \gg d^2$。
4. 解释重计算（激活检查点）怎样用运算换显存，FlashAttention 的反向为什么只需要每行一个数 $L_i$。
5. 说出 PagedAttention 借用了虚拟内存的哪几个概念，分页为什么能让同一块显存装下多一倍以上的请求。
6. 写出推测解码的接受 / 拒绝规则，证明它的输出分布和目标模型完全相同，用 $\alpha$、$\gamma$、$c$ 估计加速比。

0.2 先修与本讲范围 Prerequisites & scope

0.3 本讲地图 Map

§1  为什么关心 FlashAttention：存储层次，Transformer 访存受限，分块回顾        ← 代码 01
§2  在线 softmax：三种写法，(m, ℓ, o) 的合并                                    ← 代码 01
§3  FlashAttention 的前向：分块 + 在线 softmax                                  ← 代码 02
§4  读写多少：IO 复杂度、SRAM 大小、因果掩码、显存                              ← 代码 02
§5  重计算与反向                                                                ← 代码 03
§6  运算更多，却更快：roofline 估计                                             ← 代码 03
§7  高效解码 I：PagedAttention                                                  ← 代码 04
§8  高效解码 II：推测解码                                                       ← 代码 05
§9–13  小结、误区、练习、延伸阅读、下一讲

本讲的”显存”和”SRAM”是记账的模型：CPU 上没有可编程的片上内存。flash_lab.py 把 FlashAttention 的循环原样实现一遍（float64，可以和 PyTorch 逐位比较），在每一次”从 HBM 读 / 往 HBM 写”时记账——访问了哪些数和真实 GPU 上的实现一致，耗时再用前置讲 P3 的 roofline 模型估计。解码部分的 decode_lab.py 用真实的小模型（第 2 讲的农场 GPT 做目标模型）跑推测解码。

1. 为什么关心 FlashAttention · Why care

1.1 它是精确的，关键指标是时间

slides 开门见山：“这个算法做的是精确的注意力，所以困惑度和模型质量没有任何下降；关键指标是运行时间。“FlashAttention（Dao et al., 2022）“是近年 ML 里影响最大的想法之一——虽然很多人用着它却不知道”。今天 PyTorch 的 F.scaled_dot_product_attention、Hugging Face 的 attn_implementation="flash_attention_2"、vLLM 背后都是它或它的后代。

它和第 19 讲要讲的稀疏、低秩等近似注意力是两条路：近似注意力减少计算量，但改变了模型、常常还”拿不到实际的加速”（论文摘要原话）；FlashAttention 不改变任何数学，只改变计算的顺序和数据放在哪里。

1.2 GPU 的存储层次

slides 的图（取自 FlashAttention 论文图 1）：

SRAM 比 HBM 快一个数量级，却小好几个数量级。前置讲 P3 §3 讲过：一个运算的耗时是 $\max(\text{FLOPs}/\text{算力},\ \text{字节}/\text{带宽})$；A100 的平衡点约 153 FLOP/字节——每从 HBM 读一个字节，要做 150 次以上运算才”不浪费”。

1.3 Transformer 的训练是访存受限的

slides 引用 Ivanov et al. (2021) 对 BERT 的剖析：“矩阵乘法占了 99% 的 FLOPs，却只占 61% 的运行时间”——剩下的 39% 时间花在 softmax、LayerNorm、dropout、残差加法这些 FLOPs 很少、却要把整个张量从 HBM 读一遍写一遍的运算上。“大量时间浪费在 GPU 上搬数据，而不是计算。”

对策是前置讲 P3 §5 的算子融合。slides 画了两个版本：

• 版本 A：一层一层地算——把这一层的输入从 HBM 搬进 SRAM、算一点、把输出写回 HBM，下一层再搬进来……
• 版本 B：把原始输入搬进 SRAM 一次，在片上连做一串运算，只把最后的输出写回 HBM。

slides：“标准注意力恰好就是版本 A”：$S = QK^\top$ 写回 HBM，读出来做 softmax 得到 $P$ 再写回 HBM，再读 $P$ 和 $V$ 算 $O$。$S$ 和 $P$ 都是 $N \times N$ 的。能不能像版本 B 那样，把三步融合成一步、$S$ 和 $P$ 永远不落地？难点在于 SRAM 放不下 $N \times N$——必须分块。

1.4 回顾：矩阵乘法为什么可以分块（代码 01 第 1、2 部分）

slides 的例子：$9 \times 9$ 的矩阵切成 $3 \times 3$ 的块，

每个输出块是若干个”块乘积”的和。于是可以把一对输入块搬进快速存储、相乘、加到输出块 $Z$ 上（$Z = Z + XY$），循环三次。01 第 1 部分验证这个分解（和整块相乘差 $4.4\times10^{-16}$），并用前置讲 P3 的公式数读写：$1024 \times 1024$ 的矩阵，块边长 1（不复用）读写 21.5 亿个数，块边长 128 只要 1783 万个——每个数搬进来后被用了 120 次。

能这样做，靠的是加法满足结合律：分几次加、按什么顺序加都一样。slides 接着说：“要是能直接对自注意力分块就好了。可惜，矩阵乘法里的加法满足结合律，自注意力里的 softmax 不满足！”

01 第 2 部分：一行 12 个分数分成 3 块、各自做 softmax 再拼起来，结果的和是 3.000 而不是 1.000，和正确结果最大相差 0.530。每一块的分母只是自己那几项的和——softmax 的分母把整行绑在了一起。

2. 在线 softmax · Online softmax

2.1 三种写法（代码 01 第 3 部分）

slides 引用 Milakov & Gimelshein (2018) 的三个算法：

朴素 softmax：$y_i = e^{x_i} / \sum_j e^{x_j}$。第一遍读 $x$ 求分母，第二遍读 $x$、写 $y$。

安全 softmax：“为了避免 $e$ 的大次幂溢出，用安全 softmax——每个深度学习框架都这样实现”：

分子分母同乘 $e^{-m}$，比值不变；最大的指数变成 $e^0 = 1$，不会溢出。代价是三遍：求 $m$、求分母、写 $y$。

在线 softmax：一遍同时求出最大值和分母。读到第 $j$ 个元素时：

第二个式子的第一项是修正：最大值变大了，之前累加的每一项都要换到新的基准上，乘上 $e^{m_{j-1} - m_j} \le 1$。

01 逐个元素地数访存：

$e^{1000}$ 在 float32 里是 inf，$e^{12} \approx 1.6\times10^5$ 也超过了 float16 的最大值 65504（前置讲 P1）——朴素写法在混合精度训练里根本不能用。在线写法和朴素写法一样只访存 3 次，又和安全写法一样不溢出。$4/3 \approx 1.33$：slides 说的”表面上 1.33 倍的加速，由于减少了带宽需求，实际上也快了 1.3 倍”，就是因为 softmax 访存受限，访存次数几乎决定了耗时。

2.2 一步一步（代码 01 第 4 部分）

一行 $[1, 3, 2, 6, 4, 5]$：

每一行的 $d_j$ 都等于 $\sum_{k \le j} e^{x_k - m_j}$（代码逐行核对）。最大值变大时（读到 3、6）整体打一个折，换到新基准；最大值不变时修正因子是 1。

2.3 从”一个一个读”到”一块一块读”：$(m, \ell, o)$（代码 01 第 5 部分）

在线 softmax 的修正规则不只对单个元素成立，对一整块也成立。对一个查询，一块分数 $s$ 和对应的值 $V$ 只需要记三样东西：

• $m$：块内最大值；
• $\ell = \sum_j e^{s_j - m}$：以 $m$ 为基准的分母；
• $o = \sum_j e^{s_j - m} v_j$：以 $m$ 为基准的未归一化输出（$d$ 个数）。

两块合并时先换到同一个基准 $m = \max(m_a, m_b)$ 再相加：

最后的输出是 $o / \ell$。这个合并运算满足结合律和交换律——01 把 64 个分数分成 8 块，从左到右、从右到左、两两配对、打乱顺序四种合并方式，和整行 $\mathrm{softmax}(s)V$ 的差都在 $2\times10^{-16}$–$3\times10^{-16}$；而 slides 里那个”作弊的 softmax”（只看一块、在块内归一化）差 1.43。

这就是全部的诀窍：softmax 本身不满足结合律，但带上 $(m, \ell)$ 之后的”加权求和”满足。从此注意力也可以像矩阵乘法一样分块计算了。（这种把不可分解的统计量改写成可合并的”充分统计量”的做法，数据库里叫 algebraic aggregation，论文的脚注也这样说。）

3. FlashAttention 的前向 · The forward pass

3.1 标准注意力（算法 0）

论文的算法 0（slides 的”Standard Attention”）：

1. 从 HBM 读 $Q, K$，算 $S = QK^\top$，把 $S$ 写回 HBM；
2. 读 $S$，算 $P = \mathrm{softmax}(S)$，把 $P$ 写回 HBM；
3. 读 $P, V$，算 $O = PV$，把 $O$ 写回 HBM。

3.2 FlashAttention（算法 1）

slides：“两个关键想法组合成 FlashAttention。两者都是成熟的想法，有趣的是怎样把它们用到注意力上：① 分块：一块一块地计算注意力权重，这样不必一次把所有东西都搬进 SRAM；② 重计算：永远不存完整的注意力矩阵，反向时需要哪一块就重新算哪一块。”

把 $Q$ 切成 $T_r$ 块（每块 $B_r$ 行），$K, V$ 切成 $T_c$ 块（每块 $B_c$ 行）。HBM 里放着输出 $O$（初值 0）、每行的 $\ell$（初值 0）和 $m$（初值 $-\infty$）：

对每一块 K_j、V_j（搬进 SRAM）：
    对每一块 Q_i（连同 O_i、ℓ_i、m_i 搬进 SRAM）：
        S_ij = Q_i K_jᵀ / √d                     （B_r × B_c，只在片上）
        m̃ = 行最大值(S_ij)；P̃ = exp(S_ij − m̃)；ℓ̃ = 行和(P̃)
        m_new = max(m_i, m̃)；ℓ_new = e^{m_i − m_new} ℓ_i + e^{m̃ − m_new} ℓ̃
        O_i ← (ℓ_i e^{m_i − m_new} O_i + e^{m̃ − m_new} P̃ V_j) / ℓ_new
        把 O_i、ℓ_new、m_new 写回 HBM

$O_i$ 的更新就是 §2.3 的合并：先乘 $\ell_i$ 撤销旧的归一化、换到新基准，再加上新的一块，再除以新的分母。$S_{ij}$ 和 $\tilde P$ 只在片上存在，用完就丢。块大小是论文算法 1 第 1 行的 $B_c = \lceil M/4d \rceil$、$B_r = \min(\lceil M/4d\rceil, d)$，$M$ 是 SRAM 能放的数的个数——保证 $K_j, V_j, Q_i, O_i$ 和 $S_{ij}$ 同时放得下。

（slides 的图里写着 $O_i = P_{ij}V_j$、$O_4 = \text{cheating\_softmax}(Q_4K_2^\top)V_2$——只用一块是”作弊”，§2.3 的 1.43 就是作弊的代价；修正项让它变成精确的。）

3.3 结果一样吗（代码 02 第 1 部分）

$N = 256$、$d = 32$，float64：

差都是 float64 的舍入：FlashAttention 不是近似，只是换了计算的顺序。它顺便留下每行的 logsumexp $L_i = m_i + \log \ell_i$（$N$ 个数），反向要用（§5）。

3.4 FlashAttention-2 的循环顺序

slides 的分块图取自 FlashAttention-2（Dao, 2023）。它把两层循环对调：外层遍历 $Q$ 的块，内层遍历 $K, V$ 的块。这样 $Q_i$、$O_i$、$m_i$、$\ell_i$ 一直留在片上，内层循环结束才写回一次，并且只在最后除一次 $\ell$（少做非矩阵乘法的运算——GPU 上它们比矩阵乘法慢得多）；不同的 $Q$ 块互不依赖，可以分给不同的 SM 并行，长序列、小批量时也能占满 GPU。代价是 $K, V$ 要被每一块 $Q$ 读一遍。代码里的 flash_attention_v2 就是这个顺序。

4. 读写多少 · IO complexity

4.1 定理与实测（代码 02 第 2 部分）

论文定理 2：“标准注意力需要 $\Theta(Nd + N^2)$ 次 HBM 访问，FlashAttention 需要 $\Theta(N^2d^2M^{-1})$ 次。“代码逐项数出来的次数和闭式公式完全一致：

• 标准：$4N^2 + 4Nd$（$S$ 写、读，$P$ 写、读，再加 $Q, K, V, O$）；
• 算法 1：$K, V$ 各读一次，但每一块 $K_j$ 都要把全部 $Q$、$O$ 读一遍、$O$ 写一遍：主项 $3Nd \cdot T_c = 12N^2d^2/M$；
• FA-2 的顺序：$Q$ 读一次、$O$ 写一次，$K, V$ 各读 $T_r$ 次：主项 $2N^2d/B_r$。

$d = 64$、$M = 65536$（$B_c = 256$、$B_r = 64$），FA-2 用 $B_r = 128$、$B_c = 64$：

一个常被误解的地方：FlashAttention 的读写也随 $N^2$ 增长。$N$ 翻倍，三列都约变成 4 倍；省下的倍数趋近一个和 $N$ 无关的常数——$M/(3d^2) = 5.3$ 或 $2B_r/d = 4.0$。它没有把注意力的读写从平方变成线性，而是把平方项前面的系数压小了；真正从 $O(N^2)$ 变成 $O(N)$ 的是额外的显存（§4.4）。片上同时存放的数最多 $0.88M$，没有超过 SRAM 的容量。

4.2 SRAM 越大越好——而且必须足够大（代码 02 第 3 部分）

$N = 1024$、$d = 64$，改变 $M$：

$M$ 每翻一倍，读写约减半（论文图 2 中：“读写越少越快，直到别的因素成为瓶颈”）。更值得注意的是前两行：SRAM 太小时，FlashAttention 反而比标准做法读写得更多。$12N^2d^2/M < 4N^2$ 要求 $M > 3d^2 = 12{,}288$。论文的前提是”$d$（64–128）远小于 $M$（约 100 KB）“；头维度越大，需要的片上内存越多。这也解释了 FlashAttention 为什么必须针对每一代 GPU 的 SRAM 大小调块大小。

4.3 因果掩码：整块跳过（代码 02 第 4 部分）

完全落在对角线上方的块，所有键都在所有查询的”未来”，直接跳过——不读、不算；只有对角线上的块需要在块内加掩码。块数从 $T^2$ 降到 $T(T+1)/2$：

标准实现做不到这一点：它先把整个 $N \times N$ 算出来，再把一半盖成 $-\infty$。

4.4 显存：从 $O(N^2)$ 到 $O(N)$（代码 02 第 5 部分）

标准实现要为反向保存 $N \times N$ 的 $P$；FlashAttention 只存每行的 $L_i$。每个头（FP16）：

128K 的上下文，一层的一个头就要 32 GiB。这才是 FlashAttention 让长上下文训练（第 19 讲）成为可能的主要原因——slides 的结果页也写着：更长的上下文带来更好的模型（论文：GPT-2 困惑度好 0.7、长文档分类提高 6.4 个点，第一次在 Path-X（16K）和 Path-256（64K）上超过随机猜测）。

5. 重计算与反向 · Recomputation

5.1 slides 的两层网络（代码 03 第 1 部分）

slides 先用一个前馈网络说明重计算：

标准做法前向时存下 $z$，反向要用它：$\partial J/\partial W_2$ 需要 $z$，$\partial z/\partial(\cdot)$ 需要 $z(1 - z)$。重计算的做法：前向算完就”删掉 $z$“，反向需要时再算一遍 $z = \sigma(W_1x + b_1)$。03 手写两种反向，和自动求导都差 $5.6\times10^{-17}$——梯度一模一样，省下 $z$ 的显存，多做一次第一层的前向。

5.2 激活检查点（代码 03 第 2 部分）

推广到 $L$ 层：每 $k$ 层存一个”检查点”，反向时从检查点出发重算这一段。$L = 16$ 层、每层 $h = \tanh(hW)$：

同时存着 $L/k$ 个检查点加一段之内的 $k$ 份，$k \approx \sqrt{L}$ 时最少（Chen et al., 2016 的 $O(\sqrt{L})$ 显存）；代价是几乎每一层的前向多做一遍。梯度和全存时完全一样（差 0）。大模型训练普遍使用激活检查点（第 17 讲的显存账里，激活值是大头之一）。

5.3 FlashAttention 的反向（代码 03 第 3 部分）

记 $P = \mathrm{softmax}(S)$、$O = PV$，给定 $dO$：

其中 $D_i = \sum_j P_{ij}\,dP_{ij}$ 是 softmax 雅可比的那一项。两个技巧让反向不需要 $P$：

1. 用 $L$ 重算 $P$：$P_{ij} = \exp(Q_iK_j^\top/\sqrt d - L_i)$——任何一块都能由 $Q_i$、$K_j$ 和 $L_i$ 现算；
2. 用 $O$ 算 $D$：$D_i = \sum_j P_{ij}(dO_i \cdot V_j) = dO_i \cdot \sum_j P_{ij}V_j = dO_i \cdot O_i$——只要 $O$ 和 $dO$，$N$ 个数。

03 和 PyTorch 的自动求导比，$dQ, dK, dV$ 的差都在 $10^{-15}$ 量级（有无因果掩码都一样）。

6. 运算更多，却更快 · More FLOPs, less time

6.1 FLOPs 和读写（代码 03 第 4 部分）

单个头、$N = 1024$、$d = 64$，前向 + 反向：

多出的 FLOPs 来自反向重算 $QK^\top$：反向从 $8N^2d$ 变成 $10N^2d$。论文图 2 左（GPT-2 medium，$N = 1024$、$d = 64$、16 个头、批 64，A100 实测）是同样的形状：

6.2 用 roofline 估一估（代码 03 第 5 部分）

按论文图 2 的设置（1024 个头）、A100 BF16，把标准实现的每个核函数和 FlashAttention 的两个核函数放到 P3 的 roofline 上：

估计快 3.8 倍（实测 5.7 倍）。一个意外：$d = 64$ 时连 $S = QK^\top$ 这个矩阵乘法都是访存受限的——它只做 $2N^2d$ 次运算，却要写出 $N^2$ 个数，算术强度 57 远低于平衡点 153；softmax、掩码就更不用说了。

在我们的模型里 FlashAttention 也还是访存受限：我们把”每个 $Q$ 块各读一遍 $K, V$“全部算成 HBM 读写。真实 GPU 上同一个头的几个 $Q$ 块同时在不同 SM 上算，重复读的 $K, V$ 很大一部分会命中 L2 缓存，所以这个模型对 FlashAttention 是偏保守的；它也没有算 dropout 和固定开销。可比的是结论：HBM 读写决定了运行时间，FLOPs 多一点无关紧要。slides 的结果页列出论文的端到端数字：BERT-large（序列 512）比 MLPerf 1.1 的训练纪录快 15%，GPT-2（序列 1K）快 3 倍，Long-Range Arena（1K–4K）快 2.4 倍。

7. 高效解码 I：PagedAttention

7.1 背景：虚拟内存

slides 先回顾操作系统的虚拟内存：每个进程看到自己的一段连续的虚拟地址空间；操作系统用每个进程的页表把虚拟页映射到物理页；物理页可以散落在内存各处，甚至换到磁盘上。好处：进程之间隔离、可以共享、内存利用率高；处理保护问题、减少碎片。

7.2 问题

slides：“PagedAttention = GPU 上注意力的虚拟内存。“服务一个 LLM 时，每个请求都有一份随 token 数增长的 KV 缓存（第 3 讲 §4）。普通系统给每个请求分配一整段连续的显存，造成三种浪费：

• 预留：为以后才生成的 token 提前占着位置；
• 内部碎片：按最大长度预留，但请求提前结束，剩下的永远用不到；
• 外部碎片：请求有长有短、结束时间不同，空闲显存被切成小段，总量够、连续的不够。

显存利用率低 → 能同时跑的请求少 → 解码访存受限时，吞吐量（每秒 token 数）几乎和批大小成正比（P3 §6.2），所以吞吐量低。

7.3 解决办法

slides：把 KV 缓存切成块（类似页），每块固定放若干个 token；块不必连续存放，用一张块表（类似页表）把逻辑块映射到物理块；注意力核函数按块计算；不需要的块回收再用，只在用到时才分配。

04 第 1 部分：12 个物理块、每块 4 个 token，3 个请求轮流生成。请求 B 的块表是 1 → 4 → 5（物理块不相邻），按块表取回的 K 完全正确；逐块计算的注意力（每块用 §2.3 的 $(m, \ell, o)$ 合并——和 FlashAttention 是同一个技巧）和连续存放时差 $10^{-16}$。浪费只发生在每个请求的最后一块，最多 block_size − 1 格。

7.4 能多装多少（代码 04 第 2 部分）

600 个请求（提示平均 203、输出平均 376 个 token，最长 2048），KV 缓存能放 65,536 个 token（Llama-3 8B 每个 token 的 KV 是 128 KiB，共 8 GiB）。每一步所有在跑的请求各生成一个 token；统计”还有请求在排队”时（显存是瓶颈的时候）显存的去向：

预留最大长度时，只有 28% 的 KV 显存真的存着 token——vLLM 论文实测已有系统是 20.4%–38.2%，分页后 96.3%，和我们的 97.9% 很接近。即使”事先知道输出长度”（现实中做不到），连续分配也要为以后才写的 token 提前占位，还留下 13% 的外部碎片。分页让平均批大小从 32 变成 134，估计吞吐量是”预留最大长度”的 2.2 倍（vLLM 论文：吞吐量提高 2–4 倍）。代价是偶尔的抢占：显存满了时把最晚来的请求赶出去、以后重算它的 KV（206 次）。

7.5 块多大、共享（代码 04 第 3、4 部分）

块越小，最后一块浪费越少（块 1：99.4%；块 256：79.6%）；但块越小，块表越长、注意力核函数每次只能读一小块，GPU 的并行度差。我们的模拟只算显存不算核函数效率，所以总显得块越小越好；vLLM 默认 16，是两者之间的折中。

分页还让共享变得自然。同一个 200 token 的提示生成 4 个各 100 token 的回答（并行采样）：不共享要 76 个块，共享提示要 40 个（省 47%）。200 不是 16 的倍数，提示的最后一块被 4 个回答共享、引用计数为 4；前 3 个回答第一次往里写时各复制一份（写时复制，copy-on-write），引用计数降到 1，最后一个回答直接在原块上写——和操作系统里 fork 之后的内存页一模一样。束搜索共享得更多（vLLM 论文：最多省 55%）。

8. 高效解码 II：推测解码 · Speculative decoding

8.1 算法

解码时每生成一个 token，都要把整个大模型的权重从显存读一遍（P3 §6.2）：70B 模型 BF16 权重 141 GB，H100 上读一遍要 42 ms；一次喂 5 个 token 多出的计算只要 0.71 ms。检验 5 个 token 和生成 1 个 token 几乎一样快——推测解码就是利用这一点。

slides 的算法：

1. 用大的目标模型预填充提示 $[x_1, \ldots, x_p]$；
2. 用小的草稿模型快速提出 $n$ 个后续 token $[x_{p+1}, \ldots, x_{p+n}]$；
3. 把这些草稿 token 一次喂给目标模型，检验接受或拒绝，保留前 $k \le n$ 个；
4. $p = p + k$，重复 2、3。

优点：“减少每个输出 token 需要调用大模型的次数，提高吞吐量”；“保持和标准解码相同的输出分布”。

检验的规则（Leviathan et al., 2023；Chen et al., 2023）：草稿第 $i$ 个 token $x$ 来自草稿分布 $q$，目标分布是 $p$。

• 以概率 $\min\!\left(1, p(x)/q(x)\right)$ 接受 $x$，看下一个；
• 第一次拒绝时，从 $\mathrm{norm}(\max(0, p - q))$ 采一个 token 代替它，本轮结束；
• $\gamma$ 个都接受了，再从 $p_{\gamma+1}$ 采一个”奖励 token”（目标模型那次前向本来就算出了它）。

8.2 为什么分布不变（代码 05 第 1、2 部分）

看第一个输出 token。$x$ 被接受的概率是 $q(x)\min(1, p(x)/q(x)) = \min(p(x), q(x))$；被拒绝的总概率是 $1 - \alpha$，其中 $\alpha = \sum_y \min(p(y), q(y))$；拒绝后输出 $x$ 的概率是 $\frac{p(x) - \min(p(x), q(x))}{1 - \alpha}$。相加：

直觉：$q$ 给多了的地方（$q > p$）按 $p/q$ 打折；$q$ 给少了的地方（$q < p$），靠拒绝后的重采样补齐，而残差分布 $\max(0, p - q)$ 恰好只补这部分。05 第 1 部分在 6 个 token 的随机分布上逐项展示，输出和 $p$ 差 $1.1\times10^{-16}$；$\alpha = 0.373$ 就是两个分布”重叠”的面积。

多步时，第一个 token 分布对了，以它为条件的下一个 token 也是同样的论证（归纳）。05 第 2 部分不做证明，而是枚举全部随机性（草稿的采样、每次接受 / 拒绝、重采样、奖励 token），精确地算出前 4 个 token 的分布（词表 4、$\gamma = 3$，共 256 种）：

“拒绝后从 $p$ 重采样”听起来很合理，却会让 $q$ 已经给多了的 token 再被抽中，分布偏向草稿。规则里的每一部分都不能省。

8.3 真实的小模型（代码 05 第 3 部分）

目标 = 第 2 讲的农场 GPT（11 万参数）。先在真实模型上做同样的精确检验（2-gram 草稿、$\gamma = 2$）：前 3 个 token 的 175,616 种组合，和目标模型差 $6.9\times10^{-18}$。然后换不同的草稿模型，各采样 3000 个句子（$\gamma = 4$、温度 1）：

• 大模型的调用次数：3-gram 做草稿时每句只调用 3.7 次目标模型（逐个采样要 11.4 次）。一个只有 6 千参数、训练一个 epoch 的 1 层 GPT 的 $\alpha$ 最高（0.91），比 3-gram 还好——草稿模型不必强，只要在”容易的 token”上和目标模型一致。
• 实测略低于公式：公式假设序列无限长、每个位置的 $\alpha$ 相同；我们的句子在 </s> 处结束，最后一轮往往用不满 $\gamma$ 个草稿（连”目标模型自己”都只有 4.35，而不是 5）。
• 分布不变：平均对数概率和两次逐个采样的平均值相差 −2.1 到 +0.2 个标准误（都在 ±3 以内），语法正确的比例（0.988–0.990）、长度也一样。草稿再差（均匀分布），输出的分布也不变——只是变慢：几乎每个草稿都被拒绝，还白白多算了草稿。

8.4 贪心解码（代码 05 第 4 部分）

温度为 0 时规则变成”草稿 = $\arg\max p$ 就接受，否则换成 $\arg\max p$“。16 个提示（“<s> the 某个词”）、3-gram 草稿：16 / 16 个输出和目标模型自己的贪心解码逐字相同。（顺便看到第 3 讲 §6.6 的复读：贪心生成了 “the fox beside the foxes beside the foxes beside the foxes chases the hens at night .”）

8.5 $\gamma$ 取多大、能快多少（代码 05 第 5 部分）

记 $c$ = 草稿跑一次的时间 / 目标模型跑一次的时间。Leviathan et al. 的定理 3.8：总耗时的期望加速比

换算到 Llama-3 70B（目标）+ Llama-3 8B（草稿）、H100、批大小 1：两者都访存受限，$c \approx 8/70 = 0.114$。

$\alpha$ 越高，$\gamma$ 可以取得越大；$\alpha$ 低时草稿太长只是浪费。还要注意前提：批很大时解码不再访存受限，检验多个 token 不再”免费”，推测解码的收益会变小。

9. 小结 · Summary

1. 注意力的瓶颈是读写 HBM，不是计算：$d = 64$ 时连 $QK^\top$ 都访存受限，softmax、掩码更是。
2. softmax 不满足结合律，但带上 $(m, \ell)$ 的加权求和满足——在线 softmax 让注意力可以分块计算。
3. FlashAttention = 分块 + 在线 softmax + 重计算：$S$、$P$ 只在片上存在；输出精确相同；读写少一个常数倍（需要 $M > 3d^2$），额外显存从 $O(N^2)$ 变成 $O(N)$；因果掩码可以整块跳过。
4. 反向用 $L_i$ 重算 $P$、用 $D_i = dO_i \cdot O_i$ 代替整行求和；FLOPs 多 13–17%，时间却快几倍。
5. PagedAttention：KV 缓存分块、块表映射、按需分配、引用计数 + 写时复制；显存里真正存 token 的比例从约 28% 提高到约 98%，同一块显存能同时跑的请求多了几倍。
6. 推测解码：小模型草拟、大模型一次检验；规则 $\min(1, p/q)$ + 残差 $\mathrm{norm}(\max(0, p-q))$ 保证输出分布完全不变；加速比由 $\alpha$、$\gamma$、$c$ 决定。

10. 常见误区 · Common pitfalls

• “FlashAttention 是一种近似注意力。” 不是，结果和标准注意力差 $10^{-16}$（float64 下）；它只改变计算顺序和数据的位置。
• “FlashAttention 把注意力从 $O(N^2)$ 变成了 $O(N)$。” 计算量仍是 $O(N^2d)$，读写也仍随 $N^2$ 增长（只是系数小了几倍）；变成 $O(N)$ 的是额外显存。
• “FlashAttention 少做了运算，所以快。” 它多做了运算（反向重算 $QK^\top$，FLOPs 多 13%）；快是因为 HBM 读写少了 9 倍。
• “块越小越省 SRAM，所以越好。” 块太小（$M < 3d^2$）时读写反而比标准做法多（我们的：0.3×）。
• “推测解码会改变模型的输出，所以只适合对质量不敏感的场景。” 按正确的规则，输出分布和目标模型完全相同（精确枚举差 $10^{-17}$）；变的只有速度。
• “草稿模型越大越好。” 越大 $\alpha$ 越高，但 $c$ 也越大；我们的 6 千参数小 GPT 就有 $\alpha = 0.91$。
• “分页会让注意力变慢很多。” 注意力本来就是一块一块读的（FlashAttention），按块表找块只是多一次间接寻址。

11. 练习 · Exercises

练习 1（在线 softmax 的正确性）。 用归纳法证明：在线 softmax 读完全部 $N$ 个元素后，$d_N = \sum_{k=1}^{N} e^{x_k - m_N}$，其中 $m_N = \max_k x_k$。

$j = 1$：$d_1 = e^{x_1 - m_1} = 1$，成立。假设 $d_{j-1} = \sum_{k<j} e^{x_k - m_{j-1}}$，则 $d_j = d_{j-1}e^{m_{j-1} - m_j} + e^{x_j - m_j} = \sum_{k<j} e^{x_k - m_{j-1} + m_{j-1} - m_j} + e^{x_j - m_j} = \sum_{k \le j} e^{x_k - m_j}$。所以修正因子 $e^{m_{j-1} - m_j}$ 的作用正是把旧的和换到新基准上。01 第 4 部分的最后一列逐行核对了这个等式。

练习 2（合并两块）。 证明 §2.3 的合并结果 $o/\ell$ 等于把两块拼起来一起算的 $\mathrm{softmax}(s)V$。

设 $m = \max(m_a, m_b)$ 是两块合起来的最大值。$\ell_a e^{m_a - m} = \sum_{j \in a} e^{s_j - m_a}e^{m_a - m} = \sum_{j\in a} e^{s_j - m}$，同理 $b$；相加就是 $\sum_{j} e^{s_j - m}$，即合起来的分母（以 $m$ 为基准）。$o$ 同理是 $\sum_j e^{s_j - m}v_j$。所以 $o/\ell = \sum_j \frac{e^{s_j - m}}{\sum_k e^{s_k - m}}v_j = \mathrm{softmax}(s)V$。

练习 3（数读写）。 $N = 4096$、$d = 128$，FA-2 的顺序取 $B_r = 128$。(a) 标准注意力（前向）读写多少个数？(b) FA-2 呢？(c) 省了几倍？和 $2B_r/d$ 比较。

(a) $4N^2 + 4Nd = 67{,}108{,}864 + 2{,}097{,}152 = 69{,}206{,}016$。(b) $Nd + 2Nd\cdot(N/B_r) + Nd + N = 524{,}288 + 33{,}554{,}432 + 524{,}288 + 4{,}096 = 34{,}607{,}104$。(c) 2.0 倍，正好是 $2B_r/d = 2$。$d$ 变大时同样的块大小省得更少——要用更大的块（更多 SRAM）才能保持倍数。（用 flash_lab.io_formula 可以验证。）

练习 4（因果掩码）。 序列切成 $T$ 块，因果注意力要算多少块？$T$ 很大时省多少？

第 $i$ 块查询需要第 $1..i$ 块键，共 $\sum_{i=1}^T i = T(T+1)/2$ 块，比例 $(T+1)/(2T) \to 1/2$。02 第 4 部分：$T = 32$（$N = 2048$、块 64）时 FLOPs 之比 0.516 = 33/64。

练习 5（推测解码的分布）。 证明 §8.2 的等式，并说明：如果拒绝后从 $p$ 而不是 $\mathrm{norm}(\max(0, p-q))$ 采样，第一个 token 的分布是什么？

等式见 §8.2。若拒绝后从 $p$ 采样：$P(x) = \min(p(x), q(x)) + (1 - \alpha)p(x)$。对 $q(x) \ge p(x)$ 的 token，这是 $p(x)(2 - \alpha) > p(x)$——草稿给多了的 token 被多抽了；对 $q < p$ 的 token 是 $q(x) + (1-\alpha)p(x) < p(x)$。总和仍是 1，但分布偏向草稿。05 第 2 部分：全变差 0.18。

练习 6（加速比）。 (a) $\alpha = 0.8$、$\gamma = 4$ 时每次调用目标模型平均产生几个 token？(b) $c = 0.1$ 时加速比是多少？(c) $\alpha = 0.6$、$c = 0.05$ 时最好的 $\gamma$ 是多少？

(a) $(1 - 0.8^5)/0.2 = 3.36$。(b) $3.36/(4 \times 0.1 + 1) = 2.40$。(c) 逐个试 $\gamma$：$\gamma = 3$ 时 1.89，$\gamma = 4$ 时 1.92（最好），之后下降。（用 decode_lab.walltime_speedup 验证。）

练习 7（PagedAttention 的碎片）。 块大小 16，100 个请求的长度在 1–2048 之间均匀分布。(a) 分页时平均每个请求浪费几格？(b) 按最大长度 2048 预留时平均浪费几格？

(a) 只有最后一块可能没写满，平均约 $(16 - 1)/2 = 7.5$ 格（长度对 16 取余近似均匀），100 个请求约 750 格。(b) 平均长度约 1024，浪费约 $2048 - 1024 = 1024$ 格，100 个请求约 10 万格——差两个数量级。

12. 延伸阅读 · Further reading

• 在线 softmax 与 FlashAttention：Milakov & Gimelshein (2018)；Dao et al. (2022)；Dao (2023), FlashAttention-2；Shah et al. (2024), FlashAttention-3（Hopper 上的异步与 FP8）；Rabe & Staats (2021), Self-attention Does Not Need $O(n^2)$ Memory。
• 访存与融合：Ivanov et al. (2021), Data Movement Is All You Need；Horace He, Making Deep Learning Go Brrrr From First Principles（slides 的版本 A / B 图）。
• 重计算：Chen et al. (2016), Training Deep Nets with Sublinear Memory Cost；Korthikanti et al. (2022)，选择性重计算。
• 服务系统：Kwon et al. (2023), vLLM / PagedAttention；Yu et al. (2022), Orca（连续批处理）；Zheng et al. (2024), SGLang（前缀缓存 RadixAttention）。
• 推测解码：Leviathan et al. (2023)；Chen et al. (2023)；Cai et al. (2024), Medusa；Li et al. (2024), EAGLE。

13. 下一讲预告 · Next lecture

第 19 讲：长上下文 LLM。上下文从 GPT-2 的 1024 个 token 涨到今天的 100 万以上。FlashAttention 让 $O(N^2)$ 的显存变成 $O(N)$，但计算仍是平方的、一块 GPU 的显存也有上限。下一讲看三条路：近似注意力（稀疏、滑动窗口、空洞）——便宜，但会丢掉什么？扩展位置编码——一个在 4K 上训练的模型，怎样用很少的数据学会 128K？环形注意力——把 FlashAttention 的分块拓展到很多块 GPU 上，让能处理的上下文随 GPU 的数量线性增长。我们会在小模型上做”大海捞针”测试。