第 5 讲 · 案例：Transformer、注意力与 FlashAttention · Transformers, Attention, FlashAttention

对应 Maps to：CMU 15-779（Fall 2025，Zhihao Jia），Lecture 5 · Fri 9/12 · Case Study: Transformers, Attention, FlashAttention。 必读论文：A. Vaswani et al., Attention Is All You Need（NeurIPS 2017）；T. Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness（NeurIPS 2022）。选读：A. Radford et al., Language Models are Unsupervised Multitask Learners（GPT-2，2019）。本讲还引用：T. Dao, FlashAttention-2（2023）；J. Shah et al., FlashAttention-3（2024）；T. Dao, D. Haziza, F. Massa, G. Sizov, Flash-Decoding for long-context inference（博客，2023）；M. Milakov & N. Gimelshein, Online normalizer calculation for softmax（2018）。 依据说明：本讲以原课 slides 05-transformers-attention.pdf 的主线组织：注意力机制、自注意力与多头；GPU 内存层次复习；安全 softmax 与在线 softmax；分块；FlashAttention-2 前向；反向的重计算；线程块级与 warp 级并行；推理的预填充与解码两阶段、KV 缓存；Flash-Decoding。论文中的数字引自论文摘要和正文（只做转述）。LikesWorld 小模型、全部模拟实验和 H100 上的估算都是本讲义自己设计的，不代表课堂上讲了什么。 配套代码 Code：code/lecture05/（5 个脚本 + 共享模块 tiny_transformer.py、attention_kernels.py；用 numpy 和 torch CPU 版；第一次运行要训练小模型约 1 分钟，之后每个脚本 ≤ 6 秒） 预计阅读时间：4–5 小时（含运行代码）

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

0.1 学习目标

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

1. 说清楚注意力”按内容查找”的直觉，写出缩放点积注意力、因果掩码、多头注意力的公式，并解释为什么要除以 $\sqrt{d}$。
2. 算出注意力的 FLOPs 和中间矩阵的大小，解释为什么标准注意力在 GPU 上受显存带宽限制，而不是受计算限制。
3. 推导安全 softmax 和在线 softmax，证明 $(m, \ell)$ 摘要的合并满足结合律，并说明它为什么是”分块”和”并行切分”的数学基础。
4. 写出 FlashAttention 前向的分块算法，算出它的显存读写量（$4Nd + 4N^2$ vs $2Nd + 2Nd \cdot N/B_r + N$），并解释 FlashAttention-2 相对第一版改了什么（循环顺序、并行维度、warp 分工）。
5. 解释反向传播为什么用重计算而不是存 $P$，推导 $D_i = dO_i \cdot O_i$ 这个关键恒等式，并用”$d$ / 拐点”判断重算是否划算。
6. 区分推理的预填充和解码两个阶段，解释 KV 缓存、解码阶段注意力的算术强度为什么永远约等于 1，以及 Flash-Decoding 为什么要沿 K/V 方向切分。

0.2 先修与本讲范围

本讲不假设你学过 Transformer。如果你已经熟悉，可以跳过 §2 的前半部分。CMU11-868 第 6 讲（Transformer 从零实现）和第 21 讲（FlashAttention 1–4）覆盖了相近内容，侧重点不同，可以互相参考。

0.3 本讲地图

§1  为什么拿注意力做案例：矩阵乘法 + 归约的组合
§2  注意力机制：直觉、公式、多头、Transformer 块        ← 代码 01（演示 1–3）
§3  注意力为什么慢：FLOPs、N×N 矩阵与访存               ← 代码 01（演示 4）
§4  softmax 的三个版本：朴素、安全、在线；合并运算 ⊕     ← 代码 02
§5  FlashAttention 前向：分块 + 在线 softmax，IO 分析    ← 代码 03（演示 1–2）
§6  实现细节：块大小、因果掩码、从 FA1 到 FA2            ← 代码 03（演示 3–4）
§7  反向传播：用重计算代替存储                           ← 代码 04
§8  推理：预填充与解码、KV 缓存、Flash-Decoding           ← 代码 05
§9  论文精读：Attention Is All You Need；FlashAttention
§10 小结  §11 常见误区  §12 练习  §13 延伸阅读  §14 下一讲预告

0.4 记号

1. 为什么拿注意力做案例 · Why attention

第 4 讲把 GPU 优化的两个基本零件讲透了：分块矩阵乘法（计算密集，关键是复用）和并行归约（访存密集，关键是不要比读一遍更慢）。注意力恰好是这两者的组合：

两次矩阵乘法中间夹着一个按行的 softmax，而 softmax 需要两个按行的归约（最大值、求和）。

把两个零件分别优化到极致，组合起来并不快。原因是”组合”本身：标准实现把第一次矩阵乘法的结果 $S$（$N \times N$）写回显存，softmax 读出来再写回 $P$，第二次矩阵乘法再读 $P$。中间结果的大小是 $N^2$，而输入输出只有 $Nd$。当 $N \gg d$ 时（$N$ = 4096、$d$ = 128 时 $N/d = 32$），绝大部分时间花在搬运中间结果上。

FlashAttention 的贡献就是跨越算子边界的融合：把三个内核合成一个，让 $S$、$P$ 永远不离开片上内存。难点在于 softmax 需要”整行”信息，而一行有 $N$ 个数，放不进片上内存。解决办法是在线 softmax：一个只需要两个数的”行摘要”，可以一块一块地更新。

这个案例为什么重要？它是”算法设计要考虑 IO”（IO-awareness）的范例。FlashAttention 没有减少 FLOPs（反向还增加了），也没有做任何近似，却在当时的 GPU 上把注意力加速了数倍。它也是后面几讲的起点：第 6 讲的 warp 特化以 FlashAttention-3 为例，第 8 讲用 Triton 写它，第 11 讲的 Mirage 能自动发现它，第 16 讲的 PagedAttention 是它的推理变体。

2. 注意力机制 · The attention mechanism

2.1 背景：为什么需要注意力

在 Transformer 之前，处理序列的主流模型是循环神经网络（RNN）：它从左到右读，每读一个词就更新一个固定大小的”记忆”向量。这有两个问题：

• 信息瓶颈：整个上文都要压进一个固定大小的向量。句子越长，早期的信息越容易被冲掉。
• 无法并行：第 $t$ 步必须等第 $t-1$ 步算完。GPU 上千个核心只能一步一步地等。

注意力的想法是：不要把上文压缩成一个向量，而是让每个位置在需要时直接”回头看”上文的任何位置，按相关程度取用。 2017 年的 Transformer 论文更进一步：完全去掉循环，只用注意力（标题 Attention Is All You Need 就是这个意思）。这样一来，所有位置可以同时计算——这正是 GPU 擅长的。

2.2 直觉：注意力是一次”软查找”

想象一个字典：每个条目有一个键（key）和一个值（value）。普通的查找是”找到和查询完全匹配的键，返回它的值”。注意力是”软”版本：

1. 拿查询（query） $q$ 和每个键 $k_j$ 算一个相似度 $s_j = q \cdot k_j$；
2. 把相似度变成一组加起来等于 1 的权重：$p_j = \text{softmax}(s)_j$；
3. 返回所有值的加权平均：$o = \sum_j p_j v_j$。

如果某个键和查询特别匹配，它的权重接近 1，结果就近似于”硬查找”。整个过程可微，所以 $q$、$k$、$v$ 由网络自己学。

代码 01 的演示让你亲眼看到这一点。我们造了一个迷你世界 LikesWorld（tiny_transformer.py）：

<bos> eve likes tea . dan likes jam . hal likes figs . gus likes plums . q dan jam . q gus plums . <eos>
      └──────────────────── 事实 ──────────────────────┘ └──────── 问答 ────────┘

“q dan” 之后的词就是 dan 喜欢的东西。要答对，模型必须在上文里找到 “dan likes ___“，把空格里的词抄过来。一个 2 层、4 个头、8.5 万个参数的 Transformer 训练约 1 分钟，测试集上答对 100.0%；而”从上文出现过的物品里随机挑一个”只有 29.2%。

把注意力权重打印出来（演示 2），两层的分工一目了然：

• 第 0 层头 3：在每个事实的物品位置（如 “jam”）上，平均把 95% 的权重放在”往前数第 2 个词”——也就是这个物品的主人 “dan”。效果是：每个物品的向量里”记住了自己的主人”。
• 第 1 层头 0：在被问的人名位置（“q dan” 里的 “dan”）上，平均把 91% 的权重放在正确的物品上。它用”dan”作为 query，去匹配那些”记住了主人是 dan”的物品 key，再把物品本身作为 value 抄出来。

这就是两步查找：“先让物品记住主人，再按主人找物品”。演示 3 把每个头分别关掉：关掉第 0 层头 3，准确率从 100% 掉到 24.0%（低于瞎猜）；关掉第 1 层头 0 掉到 49.9%；而有 4 个头关掉后准确率仍在 99% 以上。第 1 层头 3 也在做查找（平均权重 57%），关掉它掉到 82.4%，和头 0 互为备份。

一个训练上的插曲（代码 README 的”计划 vs 实际”）：最初的问法是 “q dan ? a jam”，被问的人名和答案之间隔了两个词。同一个模型怎么训练都停在 27–31%，正好是瞎猜的水平——它学会了”答案是上文的某个物品”，却没学会”是哪一个”。只要让人名紧挨着答案，2000 步就能学会。注意力能做的事，不等于小模型在有限的训练里一定能学会；“让信息只差一跳”对学习非常重要。

2.3 公式：缩放点积注意力

把所有 query 排成矩阵 $Q \in \mathbb{R}^{N \times d}$，key 和 value 同理，一次算完所有位置：

为什么除以 $\sqrt{d}$？ 假设 $q$、$k$ 的每一维独立、均值 0、方差 1，那么 $q \cdot k = \sum_{t=1}^{d} q_t k_t$ 的方差是 $d$，标准差 $\sqrt{d}$。$d = 128$ 时点积的典型大小在 ±11 左右，softmax 会被最大的那一项完全主导（几乎是 one-hot），梯度接近 0，训练卡住。除以 $\sqrt{d}$ 把方差拉回 1。Transformer 论文正是这样解释的。

自注意力（self-attention）：$Q$、$K$、$V$ 都来自同一个序列 $X$，$Q = XW_Q$、$K = XW_K$、$V = XW_V$。每个位置既是查询者，也是被查询者。

因果掩码（causal mask）：生成文本时，第 $i$ 个位置只能看到位置 $\le i$ 的词（不能偷看答案）。做法是把 $S$ 中 $j > i$ 的位置设成 $-\infty$，softmax 后权重就是 0。$S$ 的上三角整个被屏蔽，这为 §6.2 的”跳过一半的块”埋下伏笔。

2.4 多头注意力

一个头只能产生一种”相似度”。实际的句子需要同时关注多种关系：语法、指代、位置……多头注意力把 $d_{\text{model}}$ 维切成 $H$ 份，每份 $d = d_{\text{model}}/H$ 维，各自独立地做注意力，再拼起来乘一个输出矩阵 $W_O$：

LikesWorld 的例子里，“记住主人”和”按主人找物品”就是两个不同的头做的事。从系统的角度看，多头的好处是：每个头完全独立，$H$ 个头天然可以并行（GPU 上每个 (批, 头) 至少一个线程块）。

2.5 Transformer 块与现代变体

一个 Transformer 块 = 注意力子层 + 前馈（FFN）子层，各带一个残差连接和归一化。本讲的 tiny_transformer.py 采用 Llama 风格的”现代”写法：

x = x + Attention(RMSNorm(x))        ← 注意力里用 RoPE 编码位置
x = x + SwiGLU_FFN(RMSNorm(x))

• Pre-norm（先归一化再进子层）：比原论文的 post-norm 训练更稳定。
• RMSNorm：只除以均方根，不减均值，比 LayerNorm 少一次归约。
• SwiGLU：$W_2(\text{SiLU}(W_1 x) \odot W_3 x)$，Llama 的 FFN。
• RoPE（旋转位置编码）：按位置把 $q$、$k$ 的每两维旋转一个角度，使 $q \cdot k$ 只依赖于两个位置的差。本讲义的小模型用可学习的绝对位置嵌入时学不会 LikesWorld（README 记录），换成 RoPE 后就可以了。

从系统的角度看，除了注意力，块里的其余部分都是”每个位置独立”的矩阵乘法和逐元素运算——它们的代价随 $N$ 线性增长。只有注意力让位置之间发生交互，代价随 $N^2$ 增长。

3. 注意力为什么慢 · Why attention is slow

3.1 FLOPs：随 $N^2$ 增长

对一个头，$QK^\top$ 和 $PV$ 各需要 $2N^2d$ 次浮点运算（乘加算 2 次），合计 $4N^2d$；因果注意力只需要算一半。块里的其他部分（4 个投影 + FFN）是 $O(N d_{\text{model}}^2)$。

代码 01 演示 4 算了 7B 规模的一层（$d_{\text{model}} = 4096$、32 头 × 128 维、FFN 11008，因果）：

短序列时注意力的计算量微不足道；128K 上下文时它占 73%。

3.2 真正的问题：$N \times N$ 的中间结果

FLOPs 只是故事的一半。上表最后一列才是关键：$N$ = 32K 时，一层所有头的 $S$ 加起来 64 GiB，占 H100 显存的 86%。即使放得下，搬运它的时间也是灾难。

标准实现是三个内核：

内核 1：S = QKᵀ/√d        读 Q、K（2Nd），写 S（N²）
内核 2：P = softmax(S)     读 S（N²），写 P（N²）
内核 3：O = PV             读 P、V（N² + Nd），写 O（Nd）

即使有理想的缓存（每个字节只从显存读一次），显存读写也是 $4Nd + 4N^2$ 个元素。$N \gg d$ 时 $4N^2$ 占绝对多数。

用第 1 讲的算术强度看：内核 1 每读写一个元素做 $2N^2d / (2Nd + N^2) \approx 2d$ 次运算——$d = 128$ 时约 256 FLOP/元素，FP16 下约 128 FLOP/字节，低于 H100 的拐点 295。内核 2（softmax）几乎不计算，算术强度约 1。三个内核全部受访存限制。 01 演示 4 的最后一条观察给出了量级：$N$ = 32K 时这一层注意力光读写就要约 82 ms，矩阵乘法本身（连被屏蔽的一半也算上）只要约 18 ms。

3.3 GPU 内存层次复习

为什么”写回显存”这么贵？回顾第 3 讲的内存层次（H100 SXM 标称值）：

SRAM 比 HBM 快一个数量级，但全卡的共享内存合计只有约 30 MB（132 × 228 KB），比 80 GB 的显存小三个多数量级。所以优化的核心是：把数据留在 SRAM 里，尽量少碰 HBM。 分块矩阵乘法（第 4 讲）已经这样做了——但它的输出 $S$ 还是要写回 HBM，因为下一个内核（softmax）需要完整的行。要打破这个限制，必须改造 softmax 本身。

4. softmax 的三个版本 · Three versions of softmax

4.1 朴素 softmax 与溢出

直接按定义算会溢出。代码 02 演示 1：float32 能表示的最大数约 $3.4 \times 10^{38}$，$e^x$ 在 $x > 88.7$ 时就溢出成 inf；float16 最大 65,504，$x > 11.09$ 就溢出。inf / inf = nan，整行结果全部作废。注意力分数超过 11 并不罕见。

4.2 安全 softmax

分子分母同乘 $e^{-m}$（$m = \max_j x_j$），数学上结果不变：

现在所有指数都 $\le e^0 = 1$，永远不会溢出。代价是多读一遍：先找最大值，再求和，再归一化——三遍。代码 02 演示 2 在 4096 个数上逐元素计数：朴素版读输入 2 遍，安全版 3 遍。

4.3 在线 softmax

能不能在一遍里同时求最大值和求和？难点是：求和时要用最终的最大值 $m$ 做基准，但读到一半时还不知道最终的 $m$。

在线 softmax（Milakov & Gimelshein，2018）的办法是：用当前已知的最大值做基准，最大值变了就把旧的和”换基准”。维护两个数 $(m, \ell)$，读到新元素 $x$ 时：

$e^{m - m_{\text{new}}}$ 就是”换基准”的系数：如果最大值没变，它等于 1；如果变大了，旧的和要相应缩小。读完一遍，$(m, \ell)$ 就是最终的最大值和以它为基准的指数和，第二遍直接归一化。总共两遍，和朴素版一样，又不会溢出。

需要真的”换基准”的次数很少：对随机顺序的数据，第 $k$ 个元素成为新最大值的概率是 $1/k$，期望次数 $\sum_{k=2}^{N} 1/k \approx \ln N$。演示 2 里 4096 个数只换了 10 次（期望 7.9）。

4.4 合并运算 ⊕：在线 softmax 的真正力量

把”一段数”的摘要定义为 $(m, \ell)$，两段摘要的合并为：

逐个元素的在线更新只是”一段一个元素”的特例。这个运算满足结合律和交换律（证明见练习 3）。这意味着：

• 可以按任意分组计算：把一行切成几块，每块算出摘要，再合并——这就是 FlashAttention；
• 可以并行计算：各块同时算，最后树形合并——这就是第 4 讲的并行归约，也是 Flash-Decoding（§8.4）。

代码 02 演示 3 在 10000 个数上用四种方式合并：从左到右逐个合并、切成 7 段长短不一的块、树形两两合并（14 层）、7 段打乱顺序。与 torch.logsumexp 的差最大 $3.6 \times 10^{-15}$（float64）。float32 下顺序不同会差在最后几位（逐个合并与一次性计算差 $7.6 \times 10^{-6}$，在 20 左右的数值上），这就是 FlashAttention 和标准实现结果”数学上相同、但不逐位相同”的原因。

4.5 从 softmax 到注意力输出

注意力还要算 $o = \sum_j p_j v_j$。把摘要扩展成三元组 $(m, \ell, \mathbf{o})$，其中 $\mathbf{o} = \sum_j e^{s_j - m} v_j$ 是未归一化的加权和，合并规则相同：

全部合并完后，输出 $= \mathbf{o} / \ell$。代码 02 演示 4 打印了一行 query 流过 4 个 K/V 块的过程：

最后一行与标准注意力的结果 [−0.2106, −0.1905] 一致（差 $1.1 \times 10^{-16}$）。注意第 2 块：新出现的 +0.78 超过了旧最大值，旧的 $\ell$ 和 $\mathbf{o}$ 都乘了 0.676。

这一节的关键结论：softmax 的分母原本需要”整行算完”才知道；在线 softmax 让每一块只需要一个很小的摘要。于是一整行 $N$ 个分数不需要同时存在——$S$ 和 $P$ 不必写回显存。

5. FlashAttention 前向 · FlashAttention forward pass

5.1 算法

把在线注意力和第 4 讲的分块组合起来（这里写 FlashAttention-2 的循环顺序，§6.3 解释为什么）：

for 每个 Q 块 Q_i（B_r 行）：                 ← 每个线程块负责一个，彼此独立、并行
    从 HBM 读 Q_i 到 SRAM
    m = −∞, ℓ = 0, acc = 0                    ← 留在片上（寄存器）
    for 每个 K/V 块 K_j, V_j（B_c 行）：
        从 HBM 读 K_j, V_j 到 SRAM
        S_ij = Q_i K_jᵀ / √d                  ← B_r × B_c，只在片上存在
        m_new = max(m, rowmax(S_ij))
        P̃ = exp(S_ij − m_new)                ← 未归一化
        ℓ   = ℓ · exp(m − m_new) + rowsum(P̃)
        acc = acc · exp(m − m_new) + P̃ V_j
        m = m_new
    O_i = acc / ℓ；L_i = m + log ℓ            ← 写回 HBM，只写一次

attention_kernels.py 里的 flash_forward 就是这段伪代码的 PyTorch 版本，在 float64 下与标准注意力差 $10^{-16}$ 量级。

为什么是精确的？ 每一步都只是在线 softmax 的合并规则（§4.5），没有丢掉任何项、没有近似。它和标准注意力的区别仅在于运算顺序。

5.2 写成 CUDA 风格的内核（代码 03）

为了看清”在 GPU 上到底省了什么”，代码 03 用第 3 讲的 gpusim 把两种做法都写成内核，逐条访存指令计数。

标准注意力（3 个内核）：

• 内核 ①、③ 是第 4 讲的 16×16 共享内存分块矩阵乘法（threadIdx.x 对应输出的列，保证合并写；共享内存按 17 列 padding）；
• 内核 ② 是按行 softmax：一个线程块负责一行，两次块内树形归约（最大值、求和），用的是第 4 讲的 v3 顺序寻址。

FlashAttention（1 个内核）：每个线程块负责 $B_r$ 行 query，块内每个线程负责一行：

1. 合作加载 Q 块到共享内存（合并访存），每个线程把自己那一行抄进寄存器；
2. 对每个 K/V 块：合作加载 $K_j$、$V_j$（float4 向量加载）→ __syncthreads() → 每个线程算自己那一行的 $B_c$ 个分数（所有线程读同一个 $K_j[c]$，是广播）→ 在线 softmax 更新 → 累加 $P̃ V_j$ → __syncthreads()；
3. 输出经共享内存转置后合并写回 $O$，同时写回每行的 $L$。

这是一个教学简化版。真实的 FlashAttention 用 Tensor Core 的 MMA 指令，一个 warp 一起算十几行，但显存读写的结构完全一样：Q 读一次、O 写一次、K/V 每个 Q 块读一遍、$S$ 和 $P$ 从不出现在显存里。

5.3 数一数（代码 03 演示 1）

$N = 256$、$d = 64$、单个头、FP32：

两种做法与 float64 参考的差都在 $4 \times 10^{-7}$ 左右（float32 的舍入量级）。Flash 版的”显存”里只有 Q、K、V、O、L 五个数组——根本没有 S 和 P。

模拟器的四个计数与闭式公式逐字节相等：

理想缓存下差 5.0 倍，差别全在 $S$ 和 $P$。

（顺带一提：这两个演示内核里 Flash 的共享内存轮数也更少——71,680 对 616,448。这是因为”一个线程独占一行”时一条 float4 广播读取能喂 4 次乘加，而 16×16 分块矩阵乘法每次乘加要读 2 次共享内存。这是两个内核写法的差别，不是 FlashAttention 的本质。）

5.4 IO 复杂度

$S$ 和 $P$ 消失后，剩下的显存读写主要是”K/V 被每个 Q 块读一遍”：$2Nd \cdot N/B_r$。块能开多大？Q 块、K 块、V 块都要放进 SRAM：$B_r \approx M/d$ 量级（$M$ 是 SRAM 能放的元素数）。于是

典型值 $d = 64$–$128$、$M \approx 10^5$ 时 $d^2/M < 1$，FlashAttention 的访问量远小于标准实现。FlashAttention 论文证明了这一点，还证明了一个下界：在一定的 SRAM 大小范围内，任何精确注意力算法都不可能在渐近意义上做得比它更好。

5.5 在 H100 上估算（代码 03 演示 2）

模拟的是 FP32、CUDA core、$N = 256$；但公式已经被逐字节验证，可以外推。$N = 4096$、$d = 128$、32 头、FP16、Tensor Core 峰值 989 TFLOP/s、HBM 3.35 TB/s：

标准做法合计约 1,334 µs（含 3 次启动开销），Flash 约 282 µs，下限估算快 4.7 倍。即使 K/V 每次都从 HBM 读，Flash 也快 2.0 倍。真实 GPU 介于两者之间：同时在跑的线程块往往属于少数几个头，每个头的 K/V 只有 2 MiB，能留在 50 MB 的 L2 里。

这是下限，不是实测。 FlashAttention-3 论文报告，FlashAttention-2 在 H100 上只达到峰值的约 35%：exp 等非矩阵运算、warp 之间的同步、加载和计算之间的等待都会吃掉时间。FlashAttention-3 用 warp 特化和异步流水线把 FP16 提到约 740 TFLOP/s（约 75%）——这是第 6 讲的主题。

6. 实现细节 · Implementation details

6.1 块大小：被 SRAM 和寄存器夹住

$B_r$ 越大，K/V 被重复读的遍数 $N/B_r$ 越少；但 Q 块、K 块、V 块都要放进共享内存。为了让加载和计算重叠，K/V 通常还要双缓冲（一份在算，一份在读）。代码 03 演示 3（$d = 128$，FP16）：

还有寄存器的约束：每行的累加器 acc 有 $d$ 个数。真实内核里一个 warp 负责一段行，acc 分散在 32 个线程的寄存器里；$d = 256$ 时寄存器压力很大，这也是大 head dim 的 FlashAttention 更难写的原因。块太大还会降低占用率（第 3 讲）。所以块大小是一个多约束的权衡——第 9 讲的自动调优就是用来搜这种参数的。

6.2 因果掩码：整块跳过

因果注意力中 $S$ 的上三角全被屏蔽。分块后，一个 (Q 块, K/V 块) 对有三种情况：

• 完全在下三角：正常计算，不需要逐元素的掩码；
• 完全在上三角（K/V 块的第一个 key 比 Q 块的最后一个 query 还靠后）：整块跳过，连 K/V 都不用读；
• 跨对角线：计算，并对块内逐元素屏蔽。

$T$ 个块时，需要计算的块对是 $T(T+1)/2$ 个，比例 $(T+1)/(2T)$，块越多越接近一半。代码 03 在模拟器上实测：$N = 256$、$B = 64$（$T = 4$）时计算 10/16 个块对，K/V 的显存读取从 524,288 降到 327,680 字节；$T = 32$ 时比例为 51.6%。

注意这个”跳过”在 GPU 上要做得聪明：在 FA2 的循环顺序里，一个线程块（一个 Q 块）的内层循环提前结束，整个块一起跳出，不会造成 warp 发散。但不同的 Q 块工作量不同（越靠后的块要算越多），会造成负载不均，需要靠调度来缓解（例如让工作量大的块先启动，第 3 讲的调度器会把后到的小块填进空出来的 SM）。

6.3 从 FlashAttention 到 FlashAttention-2：谁在外层？

FlashAttention 第一版的循环顺序与上面相反：外层是 K/V 块，内层是 Q 块。

• 好处：K、V 只读一次；
• 坏处：每处理一个 K/V 块，都要把所有 Q 块读一遍，把 $O_i$、$\ell_i$、$m_i$ 读出来、更新、写回去。

代码 03 演示 4（$N = 4096$、$d = 128$、$B_r = B_c = 128$，一个头）：FA1 的显存读写约 51.9 M 个元素，FA2 约 34.6 M（K/V 命中 L2 时只有 2.1 M）。FA2 把 $O_i$ 留在寄存器里，直到这个 Q 块全部算完。

更重要的是并行度。FA1 的外层循环有依赖（每个 K/V 块都要更新同一批 $O_i$），只能按 (批, 头) 并行：批 1 × 32 头时只有 32 个线程块，而 H100 有 132 个 SM。FA2 的外层 Q 块彼此独立，每个都能分给一个线程块：同样的配置能开 1,024 个线程块。批小、序列长——这正是长上下文训练和推理的常见情况——时，这一点决定了 GPU 能不能被占满。

FlashAttention-2 论文总结了三条改进：

1. 减少非矩阵运算：例如不在每一步都用 $\ell$ 归一化输出，而是最后除一次（上面伪代码里的 acc / ℓ）；只存 logsumexp $L$ 而不是分开存 $m$ 和 $\ell$。FlashAttention-2 论文举的例子：A100 上 FP16 矩阵乘法的峰值是 312 TFLOP/s，非矩阵乘法的 FP32 运算只有 19.5 TFLOP/s，每个非矩阵运算相当于 16 个矩阵乘法运算，所以要尽量少做。
2. 沿序列长度并行：就是上面说的外层 Q 块独立。
3. 改进 warp 之间的分工（§6.4）。

论文报告：比第一版快约 2 倍，在 A100 上达到理论峰值的 50–73%；端到端训练 GPT 风格的模型达到每块 A100 225 TFLOP/s（约 72% 的模型 FLOPs 利用率）。

6.4 warp 级的分工：split-K vs split-Q

一个线程块里有几个 warp（通常 4 个），怎样把一个 Q 块的工作分给它们？

• split-K（FA1 的做法）：所有 warp 共享同一个 $Q_i$，各自负责 $K_j$、$V_j$ 的一部分列。问题是每个 warp 只算出了部分的 $S$ 行、部分的 $P̃ V$，必须通过共享内存交换、同步、相加，才能得到正确的行最大值和输出。
• split-Q（FA2 的做法）：所有 warp 共享同一个 $K_j$、$V_j$，各自负责 $Q_i$ 的一部分行。每一行的全部计算都在同一个 warp 里完成，warp 之间不需要任何通信。

代码 03 的模拟内核”每个线程负责一行”，是 split-Q 的极端形式：所有线程读同一个 $K_j[c]$（广播），之间不交换任何数据。这和 §4.4 的结合律是同一个道理：split-K 是把一行切成几段、最后合并；split-Q 干脆不切行。能不切就不切，切了就要付通信的代价。

7. 反向传播：用重计算代替存储 · Backward pass with recomputation

7.1 标准反向需要什么

训练时要算 $dQ$、$dK$、$dV$。链式法则给出：

（$dS$ 那一项来自 softmax 的雅可比矩阵 $\text{diag}(p) - pp^\top$。）这里需要 $P$。标准做法是前向时把 $P$（$N \times N$）存下来——又回到了 $N^2$ 的显存和读写。代码 04 演示 2：32 个头、FP16，$N$ = 64K 时一层存 $P$ 要 256 GiB。

7.2 FlashAttention 的做法：只存 $L$，反向时重算 $P$

前向时只存 $O$ 和每行一个数 $L_i = m_i + \log \ell_i$。反向时：

$L_i$ 已经包含了整行的归一化信息，所以重算出来的 $P$ 直接就是归一化的，不需要再做一次在线 softmax。存 $L$ 只要 $N$ 个数：$N$ = 64K、32 头时 8 MiB，比存 $P$ 小 32,768 倍。

这就是第 2 讲 §7 的重计算（rematerialization）：不存中间结果，需要时重算。

7.3 关键恒等式：$D_i = dO_i \cdot O_i$

$dS$ 的公式里有一个”整行的和” $\sum_k P_{ik}\, dP_{ik}$。看起来又需要整行同时在场。但：

$D_i$ 只用到 $dO$ 和 $O$，可以在反向开始前一次算好，每行一个数。于是 $dS_{ij} = P_{ij}(dP_{ij} - D_i)$ 可以逐块计算。代码 04 演示 1 验证：两种算法差 $10^{-15}$ 量级。

7.4 反向的循环顺序与原子加

FlashAttention-2 的反向以 K/V 块为外层：每个线程块负责一个 $K_j, V_j$，把 $dK_j$、$dV_j$ 累加在片上，最后写一次。内层遍历所有 Q 块，每次算出对 $dQ_i$ 的贡献——但 $dQ_i$ 会被所有 K/V 块各加一次，而这些 K/V 块属于不同的线程块，只能用原子加（atomicAdd）写到显存里。代码 04 的计数：$N = 300$、$d = 32$、$B_c = 48$ 时 dQ 原子加 67,200 个元素 $= \lceil N/B_c \rceil \cdot N \cdot d = 7 \times 300 \times 32$。

和自动微分对照（float64，$N = 300$，$B_r = 64$、$B_c = 48$）：$dQ$、$dK$、$dV$ 的最大差 $\le 5.3 \times 10^{-15}$（含因果掩码）。

7.5 多算 FLOPs，反而更快

代码 04 演示 3（$N = 4096$、$d = 128$、32 头、FP16）：

多算 25% 的反向 FLOPs，显存读写降到 1/29，下限估算快 3.3 倍。

7.6 一个和 $N$ 无关的判据

什么时候重算比读回更划算？重算 $S$ 的一个元素要 $2d$ 次运算，读回 $P$ 的一个元素要 2 字节（FP16）：

和 $N$ 无关！代码 04 演示 4：

比值小于 1 就是重算更便宜。$d = 128$ 时四代 GPU 都小于 1。随着拐点升高（算力涨得比带宽快，第 1 讲的”内存墙”），重算越来越划算。这还没算上”存不下”的问题——比值大于 1 时，重算依然可能因为显存而必须做。

7.7 接进模型训练（代码 04 演示 5）

把 flash_forward + flash_backward 包成 torch.autograd.Function（前向只保存 q、k、v、O、L），作为 tiny_transformer 的注意力实现，从同一个初始化训练 30 步：损失最大差 $4.8 \times 10^{-7}$，参数最大差 $1.7 \times 10^{-6}$——只是 float32 舍入顺序不同。自己写的 FlashAttention 可以直接拿来训练模型。

8. 推理：预填充、解码与 Flash-Decoding · Inference

8.1 两个阶段

LLM 生成文本分两个阶段：

• 预填充（prefill）：把整个提示词一次性送进模型，算出每一层的 K、V，并得到第一个输出 token 的分布。这一步处理 $N$ 个 token，和训练的前向一样。
• 解码（decode）：一次生成一个 token。每一步只有一个新 token 进入模型，它要和之前所有 token 做注意力。

8.2 KV 缓存

解码第 $t$ 步时，前面 $t - 1$ 个 token 的 K、V 和上一步算的完全一样（因果掩码保证前面的 token 看不到后面的）。把它们存下来，每步只算新 token 的 Q、K、V，就是 KV 缓存。

代码 05 演示 1 在 tiny_transformer 上验证：60 个提示词，用缓存和每步重算整个序列，生成结果全部相同（logits 最大差 $1.9 \times 10^{-5}$，是浮点舍入）；平均 19 个 token 的提示词生成 8 个 token，计算量差 6.9 倍。不缓存时总计算是 $O(n^2)$ 个 token 的前向，缓存后是 $O(n)$。

一个真实的 bug（代码 README”计划 vs 实际”第 2 条）：第一版代码里，每一层用 cache.length 现算新 token 的起始位置，但 cache.length 看的是第 0 层的缓存——第 1 层执行时第 0 层已经把新 token 拼进去了，起始位置多算了 $T$。结果预填充时第 1 层的因果掩码失效，60 个提示词里只有 52 个结果相同。KV 缓存是一个优化，优化必须和原始计算逐位等价；“用 / 不用缓存对照”这种测试能抓住这类 bug。

代价是显存：每个 token 每层要存 $2 \times d_{\text{model}}$ 个数。7B 模型（32 层、$d_{\text{model}} = 4096$、FP16）每个 token 0.5 MiB；批 128、上下文 2048 时整个模型的 KV 缓存要 128 GiB，比 H100 的显存还大。KV 缓存的管理是第 16 讲（PagedAttention）的主题。

8.3 预填充 vs 解码：算术强度天差地别

代码 05 演示 2，7B 规模的一层在 H100 上（FP16 拐点 295）：

几个结论：

• 预填充是计算受限的：每个权重被 2048 个 token 复用。
• 解码批 1 时一切都是访存受限的：每读一个权重只用一次（算术强度约 1），每个 token 的成本是预填充的约 300 倍。
• 增大批次能救线性层（同一份权重被批里所有序列共用，算术强度约等于批大小），救不了注意力：每个序列读的是自己的 K/V，没有任何复用，算术强度永远约 1。批 1 时注意力只占一层时间的 8%，批 128 时占 91%。

这张表解释了为什么推理系统的研究集中在解码阶段：连续批处理（第 16 讲）提高线性层的复用；PagedAttention 让更多序列的 KV 缓存放进显存；推测解码（第 17 讲）一次验证多个 token，把解码变得像预填充；GQA / MQA 让多个 query 头共享一份 K/V，直接提高注意力的算术强度（练习 7）。

8.4 Flash-Decoding：沿 K/V 方向切分

解码时 FlashAttention 有一个新问题：并行度不够。FA2 按 (批, 头, Q 块) 开线程块，但解码时每个序列只有 1 个 query，Q 块只有一个。批 1、32 个头时只有 32 个线程块，H100 的 132 个 SM 只用上 24%。Flash-Decoding 博客指出：解码时只要批大小小于 SM 的个数（A100 是 108），GPU 就有大量 SM 空闲；批 1 时 FlashAttention 用到的 GPU 不到 1%。

Flash-Decoding（Tri Dao、Daniel Haziza、Francisco Massa、Grigory Sizov，2023）的做法：

1. 把 K/V 沿序列方向切成 $S$ 段（只是张量的视图，不复制数据）；
2. 每段由一个线程块独立计算部分结果，同时记下这一段的 logsumexp；
3. 用一个很小的归约内核合并各段：按各段的 logsumexp 重新加权。

第 2–3 步就是 §4.4 的 ⊕ 运算。attention_kernels.flash_decode 实现了它；代码 05 演示 3 把它接进 tiny_transformer 的解码步，切 2、4、7 段时生成结果 60/60 全部相同。

代码 05 演示 4 用一个简化模型估算并行度的影响（批 1、32 头、$d = 128$、FP16；假设每个 SM 最多拿到总带宽的 1/132——真实 GPU 上单个 SM 能拿到的更多，所以这里的差距是偏大的上限）：

段数也不是越多越好。上下文 32K 时：$S$ = 4 是 169 µs，$S$ = 5 时 160 个块放不进一波（132 个 SM），要排两波，反而变成 268 µs——第 3 讲的波次量化。另外还多一个合并内核，上下文很短时这点开销不划算。

Flash-Decoding 博客在 CodeLlama-34B 上实测：对很长的序列，生成速度最多快 8 倍；注意力这一步本身比 FlashAttention 最多快 50 倍。我们的估算模型给不出这些数字（它只建模了 SM 利用率），但方向一致：解码阶段，问题不是读写量，而是能不能同时把足够多的 SM 用来读。

9. 论文精读 · Paper reading

9.1 Attention Is All You Need（Vaswani et al., NeurIPS 2017）

1. 这篇论文要解决什么问题？ 当时最好的序列到序列模型（机器翻译）基于 RNN 或 CNN 加注意力。RNN 的计算沿时间步串行，难以并行，长距离依赖要经过很多步才能传递。

2. 核心想法是什么？ 完全去掉循环和卷积，只用注意力（加上逐位置的前馈网络）构建编码器和解码器。关键组件：缩放点积注意力、多头注意力、位置编码（正弦函数）、残差 + LayerNorm。论文第 4 节专门比较了自注意力、循环层、卷积层的三项指标：每层的计算复杂度（自注意力 $O(n^2 d)$）、串行操作数（自注意力 $O(1)$，循环层 $O(n)$）、任意两个位置之间的最长路径长度（自注意力 $O(1)$）。

3. 怎样验证？ WMT 2014 英德翻译 28.4 BLEU，比之前最好的结果（包括集成模型）高 2 BLEU 以上；英法翻译单模型 41.8 BLEU，在 8 块 GPU 上训练 3.5 天，训练成本只是之前最好模型的一小部分。

4. 从系统的角度看，这篇论文改变了什么？ 它把序列建模的核心计算变成了大矩阵乘法——这正是 GPU 和 Tensor Core 最擅长的。“串行操作数 $O(1)$“这一列，从系统角度看就是”可以完全并行”。代价是 $O(n^2)$ 的计算和内存，当时 $n$ 只有几百，不是问题；当上下文长度到了 32K、128K 时，它成了本讲的主角。可以说 Transformer 赢在和硬件匹配，这是”模型–系统–硬件协同设计”的典型例子（第 1 讲）。

5. 局限与后续？ $n^2$ 的代价催生了大量”高效注意力”的近似方法（稀疏、低秩、线性注意力），但大多数在实际的墙钟时间上并没有更快——因为它们减少了 FLOPs，却没有减少显存读写。这正是 FlashAttention 的出发点。

9.2 FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness（Dao et al., NeurIPS 2022）

1. 这篇论文要解决什么问题？ Transformer 处理长序列又慢又费显存。已有的近似注意力方法减少了计算量，但常常没有带来墙钟时间上的加速。作者认为缺少的原则是 IO 感知：考虑 GPU 各级内存之间的读写。

2. 核心想法是什么？ 用分块（tiling）把注意力的计算组织成在 SRAM 里完成的小块，用 softmax 的分解（即在线 softmax 的合并规则）让分块后的结果精确等于原来的结果；反向时用重计算代替存储 $N \times N$ 的中间矩阵。实现上把整个注意力融合成一个 CUDA 内核。

3. 理论贡献？ IO 复杂度分析：标准注意力需要 $\Theta(Nd + N^2)$ 次 HBM 访问，FlashAttention 需要 $\Theta(N^2 d^2 M^{-1})$ 次；并证明在一个 SRAM 大小范围内，没有任何精确注意力算法能在所有 $M$ 上渐近地改进它。本讲的代码 03 用模拟器验证了这两个公式在小规模上的精确形式。

4. 怎样验证？ 论文摘要报告：BERT-large（序列 512）端到端训练比 MLPerf 1.1 的训练速度记录快 15%；GPT-2（序列 1K）快 3 倍；Long-Range Arena（序列 1K–4K）快 2.4 倍。因为显存省下来了，可以训练更长的上下文：GPT-2 的困惑度提高 0.7，长文档分类提高 6.4 个点；首次在 Path-X（序列 16K）上超过随机水平（61.4%），在 Path-256（序列 64K）上达到 63.1%。论文还提出块稀疏的 FlashAttention，比已有的近似注意力方法都快。

5. 局限与后续？ 论文自己指出：需要为每种注意力变体手写 CUDA 内核，工程量大；IO 分析只考虑了单块 GPU。后续工作沿着三个方向走：FlashAttention-2（2023，本讲 §6.3）改进并行和 warp 分工；FlashAttention-3（2024，第 6 讲）用 Hopper 的异步特性重叠计算和数据移动；用 Triton（第 8 讲）或编译器（第 11 讲 Mirage）自动生成这类内核，减少手写的负担。Flash-Decoding（§8.4）则把同样的数学用到推理的解码阶段。

9.3 两篇论文之间的联系

Transformer 论文为”可并行”付出了 $O(n^2)$ 的代价，并在复杂度表里把它写得清清楚楚；五年后的 FlashAttention 指出，这个 $O(n^2)$ 在计算上可以接受，但在显存读写上不必付出。两篇论文都是”算法和硬件特性匹配”的例子：前者让算法适应 GPU 的并行性，后者让算法适应 GPU 的内存层次。读这两篇论文时，可以特别注意它们各自的”复杂度表”——Transformer 比较的是 FLOPs 和串行步数，FlashAttention 比较的是 HBM 访问次数。选择什么作为代价模型，决定了你能看到什么优化机会。

10. 小结 · Summary

1. 注意力是一次”软查找”：query 和 key 的点积决定权重，value 的加权平均是输出。多头 = 多个并行的独立查找。在 LikesWorld 上，2 层的小模型学会了”先让物品记住主人，再按主人找物品”。
2. 标准注意力受显存带宽限制：中间结果 $S$、$P$ 是 $N \times N$ 的，读写 $4N^2$ 个元素，而输入输出只有 $4Nd$。
3. 安全 softmax 防溢出，但要多读一遍；在线 softmax 用”换基准”把最大值和求和合进一遍。$(m, \ell)$ 摘要的合并 ⊕ 满足结合律，可以任意分组、并行计算。
4. FlashAttention = 分块 + 在线 softmax + 一个融合的内核：$S$、$P$ 永远不离开片上内存。IO 复杂度从 $\Theta(Nd + N^2)$ 降到 $\Theta(N^2 d^2/M)$。它是精确的，不是近似。
5. FlashAttention-2 把 Q 块放到外层：$O$ 留在寄存器里，不同的 Q 块可以在不同的 SM 上并行；warp 之间按行切分（split-Q），不需要通信。
6. 反向传播只存 $L$（每行一个数），重算 $P$。恒等式 $D_i = dO_i \cdot O_i$ 让 softmax 的反向可以分块。重算 vs 读回的比值是 $d$ / 拐点，与 $N$ 无关；现代 GPU 上重算更便宜。
7. 推理分预填充（计算受限）和解码（访存受限）。批处理能救解码的线性层，救不了注意力（算术强度永远约 1）。Flash-Decoding 沿 K/V 方向切分，解决批小时 SM 利用率低的问题。

11. 常见误区 · Common misconceptions

• “FlashAttention 是一种近似注意力。” 不是。它和标准注意力数学上完全等价，只是计算顺序不同；数值差异只来自浮点舍入（代码 02、03、04 都验证了）。
• “FlashAttention 快是因为减少了计算量。” 正好相反：前向的 FLOPs 不变，反向还多了 25%。它快是因为减少了显存读写，而注意力是访存受限的。
• “FlashAttention 省显存是因为 K、V 压缩了。” 不是。省的是 $N \times N$ 的中间矩阵 $S$、$P$：前向不写回，反向不存储。Q、K、V、O 的大小不变。
• “在线 softmax 每读一个数都要乘一次缩放系数，很慢。” 只有最大值变大时系数才不等于 1；随机数据上这种情况大约只发生 $\ln N$ 次。分块后缩放是按块进行的，代价被一整块分摊。
• “块越大越好。” 块越大 K/V 重复读得越少，但受共享内存、寄存器和占用率的限制；还要给双缓冲留空间。这是需要调优的参数。
• “解码时增大批次就能解决访存受限。” 只对线性层有效。注意力部分每个序列读自己的 KV 缓存，算术强度与批大小无关。
• “KV 缓存只是一个小优化。” 它把生成的计算量从 $O(n^2)$ 降到 $O(n)$，但把问题从”算得多”变成了”存得多、读得多”——KV 缓存的显存管理是 LLM 服务系统的核心问题之一。
• “Flash-Decoding 是一个新的数学方法。” 它和 FlashAttention 用的是同一个合并规则，区别只在于沿哪个维度切分、怎样分配给线程块。

12. 练习 · Exercises（含参考答案）

练习 1　一个 13B 模型有 40 层、40 个头、每头 128 维。序列长度 8192、FP16，一层所有头的 $S$ 矩阵有多大？

答案：$40 \times 8192^2 \times 2$ 字节 $= 40 \times 67{,}108{,}864 \times 2 \approx 5.4 \times 10^9$ 字节 = 5 GiB。整个模型 40 层如果都同时存着，是 200 GiB。

练习 2　设 $q, k \in \mathbb{R}^d$ 的每一维独立、均值 0、方差 1。证明 $q \cdot k$ 的方差是 $d$，并说明除以 $\sqrt{d}$ 的作用。

答案：$q \cdot k = \sum_t q_t k_t$。每一项 $q_t k_t$ 的均值 $E[q_t]E[k_t] = 0$，方差 $E[q_t^2 k_t^2] = E[q_t^2]E[k_t^2] = 1$；$d$ 项独立，方差相加得 $d$。除以 $\sqrt{d}$ 后方差为 1，softmax 的输入保持在一个不会饱和的范围内，梯度不会消失。

练习 3　证明合并运算 $(m_a, \ell_a) \oplus (m_b, \ell_b)$ 满足结合律。

答案：一段摘要 $(m, \ell)$ 代表的量是 $Z = e^m \ell = \sum_{j \in \text{段}} e^{s_j}$（数学上，不考虑溢出）。合并结果的 $e^m \ell = e^m(\ell_a e^{m_a - m} + \ell_b e^{m_b - m}) = Z_a + Z_b$，而 $m = \max(m_a, m_b)$。所以 ⊕ 在”代表的量”上就是 $(\max, +)$，两者都满足结合律和交换律。$(m, \ell)$ 只是 $Z$ 的一种不会溢出的表示：给定 $Z$ 和 $m$，$\ell = Z e^{-m}$ 唯一确定。

练习 4　用 io_flash 的公式算 $N = 4096$、$d = 128$、$B_r = 128$ 时 FlashAttention 前向的显存读写（元素数，不考虑 L2）。和代码 03 演示 4 的数字对照。

答案：$2Nd + 2Nd \cdot N/B_r + N = 2 \times 524{,}288 + 2 \times 524{,}288 \times 32 + 4096 = 1{,}048{,}576 + 33{,}554{,}432 + 4{,}096 \approx 34.6$ M，与演示 4 的 FA2 一致。其中 97% 是 K/V 的重复读取，所以 K/V 能否命中 L2 影响很大。

练习 5　因果注意力分成 $T \times T$ 个块对，需要计算多少个？$T = 8$ 时占多少？为什么不是恰好一半？

答案：下三角加对角线：$T(T+1)/2$，比例 $(T+1)/(2T)$。$T = 8$ 时 36/64 = 56.25%。多出来的是对角线上的 $T$ 个块：它们一半被屏蔽，但仍要整块计算（再逐元素屏蔽）。$T \to \infty$ 时比例趋于 1/2。

练习 6　按 §7.6 的判据，V100 上 head dim 多大时”存 $P$ 读回”比”重算 $S$“更省时间？为什么即使在这种情况下 FlashAttention 仍然选择重算？

答案：比值 $d$ / 拐点 > 1，即 $d > 139$。但存 $P$ 要 $N^2$ 的显存，长序列时根本放不下（练习 1）；而且读回 $P$ 之外，标准反向还要读写 $dP$、$dS$ 等 $N^2$ 大小的矩阵（§7.5 的 $7N^2$）。重算让整个反向都在片上完成，节省的远不止”读回 $P$“这一项。

练习 7　解码批 1 时，多头注意力（MHA）每个 query 头有自己的 K/V 头。**分组查询注意力（GQA）**让 $g$ 个 query 头共享一个 K/V 头。解码时注意力的算术强度变成多少？Llama-2-70B 用 64 个 query 头、8 个 K/V 头，算术强度和 KV 缓存大小各变成 MHA 的几倍？

答案：每读一份 K/V（$2 \cdot t \cdot d$ 个元素），要和 $g$ 个 query 做注意力，FLOPs 是 $g \times 4td$，算术强度约 $g$ FLOP/字节（FP16）。Llama-2-70B 的 $g = 64/8 = 8$：算术强度提高 8 倍，KV 缓存缩小到 1/8。这是用模型结构的改动解决系统问题的例子。

练习 8　批 2、32 个头、H100（132 个 SM）、上下文 64K 做 Flash-Decoding。按代码 05 的估算模型，切几段最好？切 3 段会怎样？

答案：线程块数 $= 2 \times 32 \times S = 64S$。$S = 2$ 时 128 个块，一波放得下，每个块读 1/128 的数据；$S = 3$ 时 192 个块，要排两波，每波的块读 1/192 的数据，总时间约 $2/192 = 1/96$ 份，比 $S = 2$ 的 $1/128$ 份还慢。所以 $S = 2$ 最好（或者让块数接近 132 的整数倍）。

13. 延伸阅读 · Further reading

• T. Dao, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning（2023）——本讲 §6.3–§6.4 的来源。
• J. Shah, G. Bikshandi, Y. Zhang, V. Thakkar, P. Ramani, T. Dao, FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision（2024）——第 6 讲的主要案例。
• T. Dao, D. Haziza, F. Massa, G. Sizov, Flash-Decoding for long-context inference（Stanford CRFM / PyTorch 博客，2023）。
• M. Milakov, N. Gimelshein, Online normalizer calculation for softmax（2018）——在线 softmax 的原始论文。
• M. Rabe, C. Staats, Self-attention Does Not Need O(n²) Memory（2021）——与 FlashAttention 同时期、思路相近的工作，从内存而不是 IO 的角度出发。
• J. Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints（2023）——练习 7。
• CMU11-868 第 21 讲：FlashAttention 1–4 的更多实现细节（FP8、Blackwell 上的多项式 exp）。

14. 下一讲预告 · Next lecture

本讲的 FlashAttention 估算是”计算受限、下限约 282 µs”，但 FlashAttention-2 在 H100 上实际只用到峰值的 35% 左右。差在哪里？答案在 GPU 内部：同一组线程既要发出加载指令、又要等数据、又要做矩阵乘法、又要算 exp——Tensor Core 经常在等。第 6 讲《高级 CUDA：warp 特化》会用一个离散事件的 SM 模拟器，一步步看软件流水线、cp.async、Hopper 的 TMA、warp 特化（一部分 warp 专门搬数据，另一部分专门计算）和乒乓调度怎样把 Tensor Core 的利用率提上去——FlashAttention-3 正是这些技巧的集大成者。