第 21 讲 · 为现代硬件优化注意力：FlashAttention 1–4 · Optimizing Attention for Modern Hardware

对应 Maps to：CMU 11-868 / 11-968 Large Language Model Systems（Spring 2026，Lei Li），Lecture 21 · Wed 4/1 · Guest Lecture by Tri Dao: Optimizing Attention for Modern Hardware（slides：llmsys-21-FlashAttention_tridao2026.4.pdf，63 页）。 阅读材料：Dao et al. (2022) FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness；Dao (2023) FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning；Shah et al. (2024) FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision；Zadouri et al. (2026) FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling（MLSys）。 依据说明：slides 覆盖：长序列建模的动机、注意力的背景、近似注意力、“瓶颈在读写”的观察、GPU 存储层次、分块的两次尝试与 softmax 重新缩放、重计算、FlashAttention 的结果；硬件的变化与 FlashAttention-2 在 H100 上只有 35–40% 的利用率；FlashAttention-3（WGMMA / TMA、exp 与矩阵乘的周期账、乒乓调度、组内重叠、FP8 与不相干处理、持久内核、因果注意力的 LPT 调度、BF16 / FP8 的基准）；FlashAttention-4（不对称扩展、多项式模拟 exp、跳过大部分重新缩放、反向的 2-CTA MMA、性能、CuTe-DSL 与 FlexAttention）；推理（Flash-Decoding、GQA 打包）与总结。slides 对 FlashAttention-2 只有一句话，§3 的内容取自 FlashAttention-2 论文；§2.3–2.4 的 IO 复杂度与下界、§5.2–5.3 的多项式与阈值细节取自对应论文。本讲义的”会记账的 HBM”、所有读写量的计数与公式核对、屋顶线估算、反向与 dropout 的实现、warp 分工与 SM 利用率的模型、列表调度的时间线、FP8 与多项式的数值实验，都是本讲义自己设计的，不代表课堂上讲了什么；本机没有 GPU，GPU 上的时间、周期都是按标称规格的模型估算。 配套代码 Code：code/lecture21/（5 个脚本 + flash.py；每个都在 3 秒以内） 预计阅读时间：3 小时（含运行代码）

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

0.1 学习目标 Learning goals

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

1. 用”读写了多少字节”而不是”算了多少 FLOPs”来分析注意力，算出标准实现的 HBM 读写量，解释为什么它是访存受限的。
2. 说清楚 FlashAttention 的两个关键技术——分块（配合在线 softmax 的重新缩放）和重计算——各自解决了什么，推导 Θ(N²d²/M) 的 IO 复杂度，理解它为什么是渐近最优的。
3. 推导 FlashAttention 的反向：为什么只存每行一个 logsumexp 就够，D = rowsum(dO ∘ O) 从哪来，为什么”多算 17% 的 FLOPs 反而快好几倍”。
4. 说出 FlashAttention-2 的三条改进（更少的非矩阵乘运算、沿序列维并行、warp 之间 split-Q），并用 SM 利用率的模型解释它们的效果。
5. 解释为什么到了 H100、B200，exp 成了和矩阵乘一样贵的东西，以及 FlashAttention-3（异步、乒乓调度、FP8 + 不相干处理）和 FlashAttention-4（多项式模拟 exp、有条件的重新缩放）怎样应对。
6. 说明推理时的注意力为什么要另外优化（Flash-Decoding、GQA 打包）。

0.2 先修与本讲范围 Prerequisites & scope

第 13 讲已经从 TPU / Pallas 的角度讲过在线 softmax 和 Flash Attention 的前向内核。这一讲换到 GPU，并且往两头走：往”理论”走——IO 复杂度、下界、反向；往”硬件”走——同一个算法在 A100、H100、B200 上的瓶颈怎样一代代转移。

0.3 本讲地图 Map

§1  动机：长序列、近似注意力为什么没成功、瓶颈在读写                              ← 代码 01
§2  FlashAttention：分块 + 重新缩放、IO 复杂度与下界、因果掩码、反向与重计算          ← 代码 02、03
§3  FlashAttention-2：更少的非矩阵乘运算、沿序列并行、split-Q                        ← 代码 04
§4  FlashAttention-3：Hopper 的异步、exp 的周期账、乒乓调度、FP8、持久内核与 LPT      ← 代码 05、04
§5  FlashAttention-4：Blackwell 的不对称扩展、多项式 exp、有条件的重新缩放             ← 代码 05
§6  推理时的注意力：Flash-Decoding、GQA 打包                                          ← 代码 04
§7 小结   §8 常见误区   §9 练习   §10 延伸阅读   §11 下一讲预告

0.4 记号 Notation

N：序列长度；d：头维度；Q、K、V、O：[N, d]（一个头；批和多头只是外面再套循环）；S = QKᵀ/√d、P = softmax(S)：[N, N]；M：片上 SRAM 能放下的元素个数；Br、Bc：Q 块、K/V 块的行数；Tr = ⌈N/Br⌉、Tc = ⌈N/Bc⌉；m、ℓ：每行到目前为止的最大值、以 m 为基准的指数和；L = m + log ℓ：每行的 logsumexp。读写量的单位是”元素”，FP16/BF16 乘以 2 就是字节。

1. 动机：瓶颈不在算，在搬 · Motivation

1.1 长序列：能力、真实性、新领域

slides 开头用三句话讲为什么要长序列：新能力（读懂整本书、整部剧本、整个代码库）、缩小与现实的差距（视觉里更高的分辨率带来更好的结果）、打开新领域（时间序列、音频、视频、医学影像天然就是几百万步的序列）。而注意力的计算量和显存都随 N² 增长：slides 第 3 页的图里，上下文长度翻倍，训练速度明显下降，再长就干脆跑不动。

01 号脚本第 4 部分算了一笔账：FP16 下，一个头、一条序列的 N × N 分数矩阵，N = 2048 时 8 MiB，N = 8192 时 128 MiB，N = 128K 时 32 GiB——一层 32 个头就是 1 TiB。训练时反向还要用到 P，标准实现在长序列上不只是慢，而是根本放不下。这就是 2022 年以前，大多数模型的上下文停在 2K–4K 的一个重要原因。

1.2 近似注意力为什么没有流行起来

2020–2021 年有大量”高效 Transformer”：稀疏注意力、低秩近似（Linformer）、核方法（Performer）、哈希（Reformer）……它们把 FLOPs 从 N² 降到 N log N 甚至 N。slides 第 6 页的总结是：它们拿质量换速度，而且很多在实际的墙钟时间上并不比标准注意力快。

为什么 FLOPs 少了却不快？因为它们优化的是错误的指标。slides 第 7 页的核心观察只有一句话：最大的开销在搬运比特。标准实现要把 N × N 的 S、P 在慢速的 GPU 显存里反复读写。于是 slides 提出的问题是：有没有一个又快、又省显存、又精确的注意力算法？

1.3 GPU 的存储层次

slides 第 8 页（引用 Horace He 的博客 Making Deep Learning Go Brrrr）画了 GPU 计算的三步：输入从 HBM（显存）出发 → 搬到计算单元和 SRAM 里计算 → 结果写回 HBM。FlashAttention 论文给的 A100 数字：HBM 40–80 GB、带宽 1.5–2.0 TB/s；SRAM 每个 SM 192 KB、108 个 SM、带宽约 19 TB/s——SRAM 快一个数量级，但小了五个数量级。

“IO 感知（IO-aware）“的意思就是：分析算法时把”在哪一级存储之间搬了多少”算进去，而不是只数 FLOPs。

1.4 标准实现的读写账

标准注意力（FlashAttention 论文的算法 0）是三个内核：

内核 1：读 Q、K，算 S = QKᵀ，写 S          读 2Nd，写 N²
内核 2：读 S，算 P = softmax(S)，写 P       读 N²，写 N²
内核 3：读 P、V，算 O = PV，写 O           读 N² + Nd，写 Nd
                                          合计 4Nd + 4N²

01 号脚本第 1 部分用一块”会记账的 HBM”（flash.HBM，每次读写按元素计数）跑了一遍：N = 256、d = 64 时计数 327,680，与公式逐元素一致，其中 S、P 的读写占 80%。这些全是中间结果——S 写下去只是为了马上被读回来。

再看算术强度（第 2 部分，FP16，d = 128，N = 4096）：

三种 GPU 的平衡点（标称算力 ÷ 带宽）：A100 153、H100 295、B200 281 FLOPs/字节。softmax 内核低了两个数量级；连两个矩阵乘也只有约 d，因为它们的输出（或输入）是 N × N 的矩阵，搬它的字节和计算量同阶。三个内核全是访存受限的。而且平衡点一代比一代高——同样的写法，在新 GPU 上”访存受限”得更厉害。

1.5 一个对照：在 CPU 上会怎样？

01 号脚本第 5 部分在本机 CPU 上用 NumPy 跑了标准实现和分块实现（d = 64，块 256）：

峰值内存的差别在任何硬件上都成立（O(N²) vs O(N)）；时间却没有变快。粗算一下，N = 8192 时标准实现搬了约 1 GiB，用 120 ms，相当于 9 GB/s，远低于这台机器的内存带宽——CPU 的瓶颈是 6700 万次 exp 的计算，不是搬数据。这个反例说明：“IO 感知”是相对于硬件的算力 / 带宽之比而言的。GPU 的张量核心让矩阵乘快了两个数量级，带宽才成了瓶颈。

2. FlashAttention：分块与重计算 · Tiling and recomputation

2.1 两个难点

想让 S、P 不出片上存储，就得分块计算。slides 第 9 页列出两个难点和对应的办法：

2.2 分块的两次尝试

slides 第 10–12 页的推导很值得一步步走。把 Q 切块很容易（每行的输出只依赖这一行的查询）；难的是 K、V。假设 K、V 切成两块：

第一次尝试：每块各算 A⁽ⁱ⁾ = exp(S⁽ⁱ⁾)、各自的 O⁽ⁱ⁾ = A⁽ⁱ⁾V⁽ⁱ⁾/ℓ⁽ⁱ⁾，再相加。错了——每块用的是自己的分母 ℓ⁽ⁱ⁾，而我们要的是整行的 ℓ = ℓ⁽¹⁾ + ℓ⁽²⁾。

第二次尝试：算完第二块后，把第一块的结果重新缩放：O⁽²⁾ = (ℓ⁽¹⁾/ℓ⁽²⁾)·O⁽¹⁾ + A⁽²⁾V⁽²⁾/ℓ⁽²⁾，这里的 ℓ⁽²⁾ 是累计到第二块为止的分母。分母对了，而且全程只在 SRAM 里做局部计算。再加上减去最大值（数值稳定），就是第 13 讲的在线 softmax：每来一块，

$m_{\text{new}} = \max(m, \tilde m),\quad \ell_{\text{new}} = e^{m - m_{\text{new}}}\ell + e^{\tilde m - m_{\text{new}}}\tilde\ell,\quad O \leftarrow \frac{\ell\, e^{m - m_{\text{new}}} O + e^{\tilde m - m_{\text{new}}} \tilde P V}{\ell_{\text{new}}}$

（m̃、ℓ̃、P̃ 是这一块自己的最大值、指数和与 exp(S − m̃)。）slides 第 13 页的动画就是这个过程：按块读入 → 在片上算这一块的贡献 → 缩放后更新输出。

02 号脚本第 1 部分：FlashAttention（论文算法 1）与参考实现的差，float64 下 8.9e-16（无掩码）/ 1.0e-15（因果）——它是精确注意力，不是近似，差别只来自加法顺序。float32 下误差与标准实现同一量级（2.8e-7 vs 4.0e-7）。

2.3 IO 复杂度：Θ(N²d²/M)

论文算法 1 的循环结构：外层循环 K/V 块 j，把 Kj、Vj 读进 SRAM；内层循环 Q 块 i，把 Qi 和到目前为止的 Oi、ℓi、mi 读进来，更新后写回。块大小取 Bc = ⌈M/4d⌉、Br = min(⌈M/4d⌉, d)，让 Kj、Vj、Qi、Oi 都放得进 SRAM。

数一数读写：K、V 各读一次（2Nd）；每个 K/V 块都把整个 Q、O 读一遍、O 写一遍，共 Tc 次——主项是 3Nd·Tc = 3Nd·N/(M/4d) = 12·N²d²/M。02 号脚本第 2 部分（N = 1024，d = 64）：

几个要点：

• 计数与闭式公式逐元素一致；M 每翻一倍，读写量的主项减半。比值在 12.4–15.3 之间，M 大时低阶项占比变大。
• M 太小时 FlashAttention 反而读写更多（第一行）：它省掉了 N × N 的中间矩阵，代价是 Q、O 被反复读写 Tc 次。只有 M 远大于 d² 时才划算。GPU 的 SRAM（A100 一个 SM 192 KB ≈ 10 万个 FP16）满足这一点。
• 论文的常数 4 是渐近意义上的：按论文的块大小，Kj、Vj、Qi、Oi 加上分数块的峰值是 M 的 0.88–1.26 倍。真实内核按寄存器和共享内存分别排布。

第 3 部分用公式看缩放关系（基准 N = 8192、d = 64、M = 98304）：N 翻倍 → ×3.83（N² 项）；d 翻倍 → ×3.79（d² 项，标准实现几乎不变）；M 翻倍 → ×0.52。

2.4 下界：不可能对所有 M 都做得更好

论文的命题 3：不存在一个精确注意力算法，对所有 M ∈ [d, Nd] 都只用 o(N²d²/M) 次 HBM 访问。 直观的理由看 M = Nd 这一端（SRAM 放得下全部输入）：N²d²/M = Nd，而任何算法至少要把 Q、K、V 读一遍——Θ(Nd) 已经是最少了。所以在这个”两级存储”的模型里，FlashAttention 的 IO 是渐近最优的。

这里有一个与真实 GPU 的差别值得说清楚（01 号脚本第 3 部分）。论文的模型里 HBM 之上只有 SRAM；如果每个 Q 块都真的从 HBM 重读 K、V，FlashAttention-2（§3）的读写约为 2N²·d/Br，与标准实现的 4N² 之比约 d/(2Br)——Br = d = 128 时只省一半（N = 4096：132 → 66 MiB）。但同一个头的所有 Q 块读的是同一份 K、V，它们只有 2Nd 个元素（N = 16K、d = 128 时 8 MiB），远小于 A100 的 40 MB、H100 的 50 MB L2 缓存，同时运行的线程块在 L2 里共享它们。于是 HBM 实际上只需读 Q、K、V、写 O 各一次。按这个情形做屋顶线估算（d = 128，N = 4096）：

FlashAttention 从访存受限变成了计算受限；平衡点越高的 GPU，省下的越多。

2.5 因果掩码：整块跳过

因果注意力里，第 i 个查询只能看到前 i 个键。分块之后，对角线右上方的块整块都被遮住，可以直接跳过——不算、也不读。02 号脚本第 4 部分（N = 1024，块 64）：16 × 16 = 256 个块只算 136 个（53.1%），矩阵乘 FLOPs 从 268.4 降到 142.6 MFLOP。标准实现做不到这一点：它先把整个 N × N 的 S 算出来，再往一半的位置填 −inf。对角线上的块只遮住一部分，在块内按全局行号、列号现算掩码，不存掩码矩阵（第 13 讲的 Splash Attention 把这个思路推广到任意稀疏掩码）。

2.6 反向：只存 logsumexp，重算 P

slides 第 14 页：前向存下 softmax 的归一化信息（大小 N），反向时从 SRAM 里的 Q、K、V 很快地重新算出注意力矩阵。

具体地（03 号脚本用的是 FlashAttention-2 论文的写法）：前向只存 O 和每行的 L = m + log ℓ。反向时，对每一块重新算 S = QᵢKⱼᵀ/√d，然后 P = exp(S − Lᵢ) 一步算出——L 已经是整行的 logsumexp，不需要再做在线 softmax。然后：

$dV_j \mathrel{+}= P^\top dO_i,\quad dP = dO_i V_j^\top,\quad dS = P \circ (dP - D_i),\quad dK_j \mathrel{+}= dS^\top Q_i,\quad dQ_i \mathrel{+}= dS\,K_j$

（省略了 1/√d。）D 是什么？softmax 的反向是 dSᵢⱼ = Pᵢⱼ(dPᵢⱼ − Σₖ Pᵢₖ dPᵢₖ)，最后那一项要整行的 P 和 dP，分块时拿不到。但 dPᵢₖ = dOᵢ·Vₖ，所以

$\sum_k P_{ik}\, dP_{ik} = dO_i \cdot \sum_k P_{ik} V_k = dO_i \cdot O_i$

只要前向的输出 O 就能算。反向先用一个小内核算 D = rowsum(dO ∘ O)，主循环里每块直接用。03 号脚本第 2 部分核对了这个恒等式（差 1.3e-15）。

正确性（第 1 部分，N = 512，d = 64）：dQ、dK、dV 与 PyTorch autograd 的差都在 1e-15 量级（因果时最大 7.1e-15）。读写量里没有任何 N × N 的项。

为什么 dQ 用原子加：反向的外层循环是 K/V 块（每个 K/V 块一个线程块），dKⱼ、dVⱼ 只属于这个线程块，在片上累加、最后写一次；但 dQᵢ 要收集所有 K/V 块的贡献，只能用原子加往 HBM 里累加（计数 524,288 = Tc × 2Nd）。如果反过来按 Q 块并行，dK、dV 两个都要原子加，量翻倍（04 号脚本第 4 部分）。

省了多少显存（第 3 部分）：GPT-2 medium（批 64 × 16 个头）N = 1024 时，存 P 每层要 2 GiB，存 L 只要 4 MiB；N = 16K 时 512 GiB vs 64 MiB。这与第 16、18 讲的激活重计算是同一个思想——用计算换显存——不同的是它只重算注意力内部最便宜的部分，而且不增加 HBM 读写。

2.7 多算了 FLOPs，为什么反而快

03 号脚本第 4 部分：前向 4N²d、反向 8N²d；FlashAttention 反向多算一次 QKᵀ（2N²d），合计多 16.7%。但读写量从 O(N²) 变成了 O(Nd)：N = 1024、d = 64 时，算术强度从 33 升到 551 FLOPs/字节（假设 K/V 的重复读命中 L2），越过了所有 GPU 的平衡点。slides 第 14 页的实测（GPT-2 medium，A100）：

GPU 上 FLOP 便宜、搬运贵。 slides 第 15 页的总结：2–4 倍加速、10–20 倍显存节省，而且没有任何近似。

2.8 dropout：存种子，不存掩码

训练时注意力还有 dropout。标准实现要把 N × N 的掩码存给反向；FlashAttention 只存随机数种子和偏移：第 (i, j) 块的掩码只由（种子, i, j）决定，前向、反向需要时各自重新生成（真实实现用 Philox 这种”计数器式”随机数生成器，任何线程都能直接算出第几个随机数）。03 号脚本第 5 部分：前向按 Q 块、反向按 K/V 块循环，掩码仍然一致，输出和三个梯度与”存了完整掩码”的参考实现差 ≤ 1.3e-15。

3. FlashAttention-2：更好的并行与分工 · Parallelism and work partitioning

FlashAttention 在 A100 上只达到理论峰值的 25–40%（FlashAttention-2 论文），而优化好的矩阵乘能到 80–90%。FlashAttention-2 的三条改进都不改变 IO 复杂度，改的是”GPU 上的工作怎么分”。

3.1 换循环顺序

FlashAttention-2 把循环换过来：外层循环 Q 块，内层循环 K/V 块。每个 Q 块从头到尾独立完成：Oᵢ、mᵢ、ℓᵢ 全程留在片上，只在最后写一次。02 号脚本第 5 部分：Br = Bc 时读写量少 1.53 倍。但两者的主项依赖不同的块大小——FA1 是 Q、O 被读写 N/Bc 次，FA2 是 K、V 被读 N/Br 次——Br = 64、Bc = 256 时 FA2 反而多。所以 FA2 用较大的 Q 块（128），并依赖 L2 吸收 K、V 的重复读取。

换顺序真正的收益在下面两条。

3.2 更少的非矩阵乘运算

A100 的张量核心有 312 TFLOP/s，普通的 FP32 单元只有 19.5 TFLOP/s——FlashAttention-2 论文说每个非矩阵乘 FLOP 贵 16 倍。所以虽然 exp、max、缩放这些运算只占 FLOPs 的 2–3%，按这个比例算却占了 28–31% 的时间（04 号脚本第 1 部分，d = 64）。

FlashAttention-2 的改法：

1. 延迟归一化：Oᵢ 累加的是未归一化的 Σ e^{S − m}V，每步只乘一个 α = e^{m_old − m_new}，除以 ℓ 只在最后做一次（FA1 每步都要两次缩放加一次除法）。
2. 只存 logsumexp L 一个数给反向，而不是 m、ℓ 两个。

按本讲义的计数规则，d = 64 时非矩阵乘运算从 7.41M 降到 6.39M（−14%）——省得不多。

3.3 沿序列维并行

FlashAttention 只按”批 × 头”并行：一个线程块负责一个头的全部计算。长序列训练时，为了放得下，批往往很小（甚至是 1）。04 号脚本第 2 部分（每个线程块占一个 SM、按波次执行的简单模型）：

FlashAttention-2 把每个 Q 块作为一个独立的线程块（换循环顺序的直接结果），线程块数乘上 N/Br。最后一行还显示了波次量化（wave quantization）：256 个线程块在 108 个 SM 上要 3 波，第三波只有 40 个，利用率 79%；同样的形状在 H100（132 个 SM）上恰好两波，97%。

3.4 一个线程块里 warp 怎么分工：split-K vs split-Q

一个线程块通常有 4 个 warp。FlashAttention 让它们切 K/V（split-K）：4 个 warp 共享同一组 Q 行、各算 1/4 的 K/V 列，得到 4 份部分结果 (m, ℓ, o)，必须写进共享内存、同步、再合并。FlashAttention-2 改成切 Q（split-Q）：每个 warp 管 1/4 的 Q 行、扫过全部 K/V，它的结果就是最终结果，不需要交换。04 号脚本第 3 部分：两种分工的结果相同（只差舍入），split-K 每个 Q 块要往返 540,672 字节的共享内存外加栅栏同步，split-Q 为 0。

3.5 结果

FlashAttention-2 论文：比 FlashAttention 快约 2 倍，在 A100 上达到理论峰值的 50–73%；端到端训练 GPT 式模型时每块 A100 约 225 TFLOP/s（72% 的模型 FLOPs 利用率）。slides 第 17 页：“FlashAttention-2 在 A100 上高度优化，达到 70% 的利用率”——可是在 H100 上只有 35–40%。下一节讲为什么。

4. FlashAttention-3：Hopper 的异步与低精度 · Asynchrony and low precision on Hopper

4.1 新硬件，新指令

slides 第 18–19 页列出 H100 的两条新指令：

• WGMMA（warpgroup 矩阵乘，4 个 warp 一起发出）：吞吐更高，而且是异步的。老的 mma.sync 在 H100 上只能达到峰值的 2/3。
• TMA（张量内存加速器）：专门把数据从全局内存搬到共享内存的硬件单元，异步，而且省寄存器（不用每个线程算地址）。

两条指令都是异步的：线程发出它们之后可以去做别的事。FlashAttention-2 是为同步的 A100 写的，没有利用这一点。FlashAttention-3 的思路：让 softmax 和矩阵乘、搬运和计算全部重叠起来。

4.2 为什么要重叠：exp 的周期账

slides 第 20 页算了一个块（头维度 128，块 128 × 128）的账：

• FP16 WGMMA：两个矩阵乘共 2 × 2 × 128 × 128 × 128 = 8.4 MFLOP，每 SM 每周期 4096 FLOP → 2048 周期；
• MUFU.EX2（特殊函数单元的 exp2）：128 × 128 = 16K 次，每周期 16 次 → 1024 周期。

exp 要花矩阵乘一半的时间！ 05 号脚本第 1 部分把它写成一个公式：一个块的 exp 周期 / 矩阵乘周期 = (张量核心每周期 FLOPs) / (4·d·每周期 exp 数)，与块大小无关：

由规格反推，H100 的 exp 吞吐只有 3.86 T/s（FlashAttention-3 论文：“特殊函数只有 3.9 TFLOPS”），矩阵乘却有 989 TFLOP/s。张量核心每代翻倍，MUFU 停在每 SM 每周期 16 个。到了 FP8（张量核心吞吐再翻倍）或 B200，exp 和矩阵乘一样要 1024 周期。

4.3 乒乓调度与组内流水线

slides 第 21–22 页的两种重叠：

• warpgroup 之间（乒乓）：两个 warpgroup 各算各的 Q 块，用栅栏（bar.sync）强制它们交替——一个做 softmax 时，另一个正好用张量核心做矩阵乘。
• warpgroup 之内（两级流水线）：在一个 warpgroup 里，让第 j 次迭代的 softmax 与第 j − 1 次迭代的第二个矩阵乘（PV）重叠。代价是寄存器：下一个 QKᵀ 的累加器和当前 PV 的操作数要同时活着。

05 号脚本第 2 部分用一个资源受限的列表调度模型（张量核心、MUFU 各是一个互斥的资源）模拟了这几种排法：

顺序执行时的甘特图里，张量核心在每次 softmax 期间都闲着（利用率 = 2G/(2G + S)）；乒乓之后两条时间线互相填满。slides 的实测是：交给 warp 调度器自己重叠 580 TFLOP/s → 乒乓 640 → 再加组内流水线 670。真实收益比模型小——模型里没有寄存器压力、同步开销和其他非矩阵乘运算。B200 那一行是 §5 的伏笔：即使完美重叠，MUFU 也已经满载，没有余量。

4.4 FP8 与不相干处理

FP8 张量核心让 WGMMA 的吞吐再翻一倍，代价是精度（slides 第 23 页）。LLM 的激活里有离群值（第 19 讲），一个大值就决定了整个张量的缩放系数，其余数值挤在 E4M3 很窄的几档里。FlashAttention-3 的两个办法：

• 逐块量化：每个块一个缩放系数，而不是整个张量一个。
• 不相干处理（incoherent processing）：把 Q、K 都乘上同一个随机正交矩阵 J（随机符号的 Hadamard 矩阵）。因为 JJᵀ = I，(QJ)(KJ)ᵀ = QKᵀ，注意力分数不变；但每个输出维度都是所有输入维度的 ±1 加权和，离群值被”摊”开了。Hadamard 变换可以用快速算法在 O(d log d) 内完成，还能和 RoPE 之类的逐元素运算融合。

slides 第 24 页：在”正态分布 + 0.1% 的大值”的数据上，不相干处理把量化误差降低 2.6 倍。05 号脚本第 3 部分用同样的数据设置（N = 1024，d = 128）：

（BF16 对照 7.30e-04。）旋转把 |Q| 的最大值 / RMS 从 26.9 降到 5.0，Q、K 的误差降低 2.1 倍；但全部 FP8 时，V 和 P 的量化误差不受旋转影响，总误差只降低 1.6 倍。逐块量化在旋转之后反而略差。这里没有复现出 2.6 倍——离群值的幅度、基线的定义、P 的处理都可能不同——方向是一致的。V 为什么不旋转？PV 的输出上没有另一半能把旋转约掉（除非在输出端再乘一次 Jᵀ）。

FP8 还有一个工程上的麻烦：FP8 的 WGMMA 要求操作数在内存里按特定的维度连续存放，V 的布局不对，FlashAttention-3 在内核里就地转置 V（slides 第 18 页”layout conformance”）。

4.5 持久内核与因果注意力的负载均衡

张量核心太快了，每个任务开头（读 Q、读第一块 K/V）和结尾（写 O、写 L）的延迟变得不可忽略。持久内核（slides 第 25–26 页）：启动的线程块数固定等于 SM 数，每个线程块自己从任务列表里取活，用异步的 TMA 写回让上一个任务的收尾和下一个任务的开头重叠——670 → 700 TFLOP/s。

有了”自己取活”，就可以决定取活的顺序。因果注意力里第 i 个 Q 块只需扫前 i 个 K/V 块，工作量从 1 到 N/Br 不等。slides 第 27–46 页的例子：2 个批 × 4 个 Q 块（工作量 1、2、3、4），3 个 SM。按自然顺序分配，最长的块恰好排在最后；改用最长处理时间优先（LPT）——先分最长的任务，每次交给最空闲的 SM。04 号脚本第 6 部分复现了这个例子：自然顺序 [9, 5, 6]（完成时间 9），LPT [7, 7, 6]（完成时间 7，等于下限 ⌈20/3⌉）。放大到 H100 的 132 个 SM，批 1 × 16 头、N = 8192 时自然顺序的效率 86.6%，LPT 100%；Q 块的个数远多于 SM 时自然顺序也有 99%——LPT 在任务数与 SM 数同一量级时最重要。实测因果注意力 670 → 730 TFLOP/s。slides 的提醒：打乱顺序时小心别把 L2 缓存挤爆（同一个头的 Q 块共用 K/V），要配合 L2 swizzling。

4.6 结果

slides 第 47–49 页（H100，CUDA 12.8）：BF16 比 FlashAttention-2 快 1.8–2.2 倍；因果注意力 730–750 TFLOP/s，约等于矩阵乘的速度；BF16 最高 840 TFLOP/s；FP8 最高 1.3 PFLOP/s。（论文发表时的数字是 FP16 最高 740 TFLOP/s、约 75% 利用率。）

5. FlashAttention-4：Blackwell 的不对称扩展 · Asymmetric hardware scaling

5.1 瓶颈又转移了

slides 第 50 页的标题是”不对称扩展的兴起”：Blackwell 的矩阵乘单元又快了一倍（FlashAttention-4 论文：B200 每 SM 每周期 8192 FLOP，H100 是 4096；整卡 2.25 vs 1 PFLOP/s），但 exp 单元和共享内存的速度没变（MUFU 仍是每 SM 每周期 16 个，共享内存每 SM 每周期 128 字节）。于是：前向被 exp 卡住，反向被共享内存的流量卡住。

论文给出的一个块（M = N = d = 128）的周期账：前向矩阵乘 1024、exp 1024、共享内存 768；反向（1-CTA）矩阵乘 2560、共享内存 3328——共享内存比矩阵乘多约 30%。

slides 第 51 页列出前向的三个算法改动：Q 块的乒乓流水线；用多项式模拟 exp；一种新的在线 softmax，跳过 90% 的重新缩放。

5.2 用 FMA 单元算 exp

MUFU 每 SM 每周期只能算 16 个 exp，而普通的 FMA（乘加）单元每周期有 128 个。能不能用乘加算 exp？可以——这是数学库的老办法（Cody-Waite 区间约化 + 多项式）：

$2^x = 2^{\lfloor x\rfloor}\cdot 2^{f},\quad f = x - \lfloor x\rfloor \in [0, 1)$

2^f 用一个低次多项式近似（Horner 法：几次乘加）；乘 2^⌊x⌋ 不用乘法——把整数 ⌊x⌋ 直接加到 float32 的指数字段上（比特操作）。05 号脚本第 4 部分用 Lawson 算法拟合了 p(0) = 1 的极小极大多项式：

关键在第三列：P 最后要舍入成 BF16 才能喂给张量核心，而 BF16 只有 8 个有效二进制位，相对误差可达 2⁻⁸ ≈ 3.9e-3。3 次多项式舍入到 BF16 后，99% 与精确值完全相同，其余也只差 1 个 ULP；最大误差就是 BF16 自身的舍入误差——更高的次数买不到更多精度。放进整个注意力里（N = 512），P 保持 float32 时 3 次多项式的误差（1.4e-4）比 torch.exp（3.3e-6）大得多；但 P 舍入到 BF16 后两者都是 3.1e-3，多项式的误差被淹没了。

FlashAttention-4 并不把所有 exp 都交给多项式：论文说每行 10%–25% 的 exp 用多项式算，其余仍用 MUFU，两种单元并行工作，比例按块的配置调。（slides 提到多项式用的是切比雪夫（Chebyshev）近似。）

5.3 有条件的重新缩放

在线 softmax 每当某块让行最大值变大，就要把 Oᵢ、ℓᵢ 乘上 α = e^{m_old − m_new}——这又是一批非矩阵乘运算，还在关键路径上。FlashAttention-4 的观察：基准 m 不必是真正的最大值（与第 13 讲的 Max Logit Estimate 同一个道理）。Oᵢ 和 ℓᵢ 只要一直以同一个 m 为基准累加，最后 O = Oᵢ/ℓᵢ 时这个基准就被约掉了。所以只在最大值增加超过一个阈值 τ 时才更新 m 并缩放；否则继续用旧的 m，代价是 e^{s − m} 可以大于 1，最大到 2^τ。论文取 τ = 8（以 2 为底，即 256 倍）。一个 warp 里只要有一个线程需要缩放，整个 warp 一起缩放。

05 号脚本第 5 部分（N = 2048，块 128，float32）：

随机数据上最大值很快就稳定了，τ = 8 时只在每个 Q 块的第一块缩放（必须的那一次），与 slides 说的”跳过 90%“一致；误差不变。τ 也不能太大：分数涨到约 316 的数据上，τ = 200 时 float32 的累加器溢出，出现 31,152 个 inf / NaN（float32 的最大值约 2^128）。τ = 8 留下了充足的余量。

5.4 反向：2-CTA 与 TMEM

反向的瓶颈是共享内存：每个块要反复从共享内存读 Q、K、V、dO、dS 喂给张量核心。Blackwell 有两个新东西（slides 第 52 页、论文）：每个 SM 256 KB 的张量内存（TMEM），矩阵乘的累加器直接放在里面；2-CTA MMA，两个 SM（CTA）合作做一个 M = 256 的矩阵乘，各自只需要读一半的操作数，共享内存流量减少，dQ 的原子加也变少（论文：1-CTA 时共享内存比矩阵乘多约 30%，2-CTA 时只多约 5%）。代价是两个 CTA 之间精细的同步与数据搬运。

5.5 结果与 CuTe-DSL

• 前向：最高约 1600 TFLOP/s（slides 第 53 页），注意力和矩阵乘一样快；论文：B200 上 BF16 最高 1613 TFLOP/s（71% 利用率），比 cuDNN 9.13 快 1.3 倍、比 Triton 快 2.7 倍。新版 cuDNN 也吸收了 FlashAttention-4 的很多优化。
• 反向：比没有 Blackwell 优化的 FlashAttention-2 基线快 4 倍（slides 第 54 页）。
• 实现语言：CuTe-DSL——嵌在 Python 里的内核语言，论文说编译比 C++ 模板快 20–30 倍。FlashAttention-4 的块稀疏、掩码抽象让 PyTorch 的 FlexAttention 快了 2.7–3.0 倍（slides 第 55 页）。

6. 推理时的注意力 · Attention for inference

6.1 Flash-Decoding：沿 K/V 切开

解码时，每个请求每一步只有一个（或几个）新查询，上下文却可能有 128K（slides 第 58 页）。FlashAttention-2 的并行方式”每个 Q 块一个线程块”失效了：批 1 × 32 头只有 32 个线程块，H100 只用上 24% 的 SM（04 号脚本第 5 部分）。而解码注意力是访存受限的——要把整个 KV 缓存读一遍——用不上的 SM 也就用不上它们那份带宽。

Flash-Decoding（split-KV）：把 K/V 沿序列切成若干段，每段一个线程块各算部分结果 (m, ℓ, o)，再用一个小内核按 logsumexp 合并。合并规则就是 §2.2 的重新缩放，与分段数无关地精确（脚本里 1–64 段的误差都在 1e-16 量级），额外读写的部分结果相对 KV 缓存可以忽略。切几段是一个调度参数：脚本仿照 flash-attn 的启发式，取”波次效率不低于最佳 85%“的最小分段数——批 1 × 32 头切 4 段，利用率 97%；批大时不切。

6.2 GQA 打包

Hopper 的 WGMMA 在 M 维上至少 64 行（slides 第 59 页）。解码时一个头只有 1 个查询，64 行的块只用上 1/64。GQA / MQA 里，共用同一组 K/V 的几个查询头可以叠在一起当作一个块的不同行：K/V 只读一次，块的利用率乘以组的大小（每组 8 个查询头、查询长度 1 时 1.6% → 12.5%；查询长度 16 时 100%）。FlashAttention-2 对”1 个查询”做了这件事（只是一次 reshape），FlashAttention-3 推广到任意的查询长度。第 24 讲的 PagedAttention 会在此基础上再加一层：KV 缓存不连续地存放在一个个块里。

7. 小结 · Summary

1. 注意力的瓶颈在读写，不在计算。 标准实现的读写量是 4Nd + 4N²，三个内核的算术强度都低于 GPU 的平衡点。FLOPs 少的近似注意力没有解决真正的问题。
2. FlashAttention = 分块 + 重新缩放 + 重计算。 读写量 Θ(N²d²/M)，在两级存储的模型里渐近最优；额外显存从 O(N²) 降到 O(N)；结果精确。反向只存 logsumexp，D = rowsum(dO ∘ O) 让 softmax 的反向不需要整行；多算 17% 的 FLOPs，换来快好几倍。
3. FlashAttention-2：同样的算法，更好的分工。 换循环顺序带来沿序列维的并行；延迟归一化减少非矩阵乘运算；split-Q 消除 warp 之间的交换。
4. FlashAttention-3：用好异步。 exp 在 H100 上要花矩阵乘一半的时间，所以要让 softmax 与矩阵乘重叠（乒乓、组内流水线）；FP8 需要逐块量化与 Hadamard 旋转来压住离群值；持久内核 + LPT 调度处理因果注意力的负载不均。
5. FlashAttention-4：瓶颈随硬件转移。 B200 上 exp 和矩阵乘一样慢、共享内存成为反向的瓶颈；用 FMA 单元的多项式分担 exp（BF16 的精度只需要 3 次多项式）、用有条件的重新缩放省掉大部分缩放、用 2-CTA 与 TMEM 减少共享内存流量。
6. 推理另有一套： Flash-Decoding 沿 K/V 切开找并行度，GQA 打包填满矩阵乘的块。

贯穿四代的一句话：算法要随着硬件的平衡点一起设计——先找到这一代硬件上真正的瓶颈（HBM 带宽 → 并行度与非矩阵乘运算 → 异步与 exp 吞吐 → exp 与共享内存），再改算法绕过它。

8. 常见误区 · Common pitfalls

1. “FlashAttention 是一种近似注意力。” 不是。它与标准实现在数学上完全等价，float64 下只差舍入（02 号脚本第 1 部分）。
2. “FlashAttention 减少了计算量。” 前向的 FLOPs 一样，反向反而多一次 QKᵀ（+17%）。它减少的是 HBM 读写。
3. “分块之后在任何硬件上都会更快。” 在 CPU 上就不会（01 号脚本第 5 部分）；SRAM 太小时（M 接近 d²）读写量甚至比标准实现还多（02 号脚本第 2 部分）。IO 感知的收益取决于硬件的算力 / 带宽之比和 SRAM 大小。
4. “反向需要前向的注意力矩阵，所以 FlashAttention 省不了训练显存。” 反向用 L 和 Q、K 在片上重算 P；softmax 反向需要的那一项可以用 rowsum(dO ∘ O) 得到。
5. “在线 softmax 的 m 必须是真正的最大值。” 只要 exp 不溢出、不全部下溢，任何基准都会在最后被约掉。这正是 Max Logit Estimate（第 13 讲）和有条件的重新缩放（§5.3）的根据——但阈值必须有界。
6. “FlashAttention-2 的提速来自更少的读写。” 读写量的差别取决于块大小，Br 小时 FA2 甚至读写更多。它的提速主要来自并行度和工作划分。
7. “FP8 注意力的误差主要来自 P。” 在有离群值的数据上，Q、K 的量化误差才是大头（05 号脚本第 3 部分），所以不相干处理只作用在 Q、K 上就很有效。
8. “多项式 exp 精度不够。” 它只需要达到 P 最终要舍入到的 BF16 精度；3 次多项式舍入后与硬件 exp 几乎处处一致。
9. “解码时用训练用的注意力内核就行。” 解码只有 1 个查询，按 Q 块并行的内核只能用上一小部分 SM；需要 split-KV 和 GQA 打包。

9. 练习 · Exercises

1. N = 2048、d = 64 时，标准注意力前向的 HBM 读写是多少个元素？其中 N × N 矩阵占多少比例？
2. A100 一个 SM 的 SRAM 按 192 KB、FP16 计，M ≈ 98304 个元素。d = 128 时按论文的公式 Bc、Br 是多少？N = 8192 时 FlashAttention（算法 1）读写量的主项是多少个元素？与标准实现比呢？
3. 证明两段部分结果的合并（§6.1、第 13 讲 §5.1）满足结合律，并说明这对 Flash-Decoding 意味着什么。
4. 不看讲义，推导 Σₖ Pᵢₖ dPᵢₖ = dOᵢ·Oᵢ。如果注意力带 dropout（O = (P ∘ Z/(1 − p))V），这个恒等式还成立吗？
5. 数一数：一个头的注意力前向 + 反向，标准实现和 FlashAttention 各要多少 FLOPs？FlashAttention 多出的那部分是什么？
6. 为什么 FlashAttention-2 的反向外层循环是 K/V 块，而前向外层循环是 Q 块？
7. H100 上头维度 d = 96 时，一个块的 exp 周期 / 矩阵乘周期是多少？B200 上呢？
8. FP8 注意力里，为什么只旋转 Q、K 而不旋转 V？如果一定要旋转 V，要额外做什么？
9. 有条件的重新缩放里，Oᵢ、ℓᵢ 用 float32 累加，序列长度 N。τ（以 2 为底）最多能取多大才保证 ℓᵢ 不溢出？
10. 6 个任务按 [1, 2, 3, 3, 4, 5] 的顺序到达，2 个 worker。按到达顺序贪心分配（每次交给最空闲的 worker）和按 LPT 分配，完成时间各是多少？下限是多少？
11. 解码时批 2 × 8 个 KV 头、上下文 32K，在 132 个 SM 上，不切分时 SM 的利用率是多少？切成几段比较合适？

参考答案

1. 4Nd + 4N² = 4 × 2048 × 64 + 4 × 2048² = 524,288 + 16,777,216 = 17,301,504 个元素；N × N 部分占 97.0%。
2. Bc = ⌈98304/512⌉ = 192，Br = min(192, 128) = 128。Tc = ⌈8192/192⌉ = 43，主项 3Nd·Tc = 3 × 8192 × 128 × 43 ≈ 1.35 亿个元素；标准实现 4Nd + 4N² ≈ 2.73 亿个元素，约 2 倍。d = 128 时 d²/M 不够小，所以真实内核还依赖 L2 吸收重复读（§2.4）。
3. 把 (m, ℓ, o) 看成表示”Σ e^{s}·v”和”Σ e^{s}“的一对数（都以 e^m 为单位）：合并就是把两对数换到同一个单位再相加。加法满足结合律，换单位不改变所表示的值，所以合并满足结合律（也满足交换律）。因此 Flash-Decoding 可以任意切段、任意顺序合并（甚至两两树状合并），结果只差舍入。
4. dPᵢₖ = dOᵢ·Vₖ，所以 Σₖ Pᵢₖ dPᵢₖ = dOᵢ·(Σₖ PᵢₖVₖ) = dOᵢ·Oᵢ。带 dropout 时，对 S 求导需要的是 Σₖ Pᵢₖ·(∂L/∂Pᵢₖ)，而 ∂L/∂Pᵢₖ = (Zᵢₖ/(1 − p))·dOᵢ·Vₖ，于是 Σₖ Pᵢₖ (Zᵢₖ/(1 − p)) dOᵢ·Vₖ = dOᵢ·Oᵢ（O 正是带 dropout 的输出）——仍然成立（03 号脚本第 5 部分的实现就用了它）。
5. 标准：前向 4N²d（QKᵀ、PV），反向 8N²d（dV、dP、dQ、dK 各 2N²d），合计 12N²d。FlashAttention：反向再算一次 QKᵀ，合计 14N²d（+16.7%）。多出的是重计算 S（以及 exp），用来代替从 HBM 读回 P。
6. 反向要产出 dQ、dK、dV。按 K/V 块并行时 dK、dV 各属于一个线程块、在片上累加，只有 dQ 需要跨线程块的原子加；按 Q 块并行则 dK、dV 都要原子加，量翻倍（04 号脚本第 4 部分）。前向只有 O 一个输出，按 Q 块并行时没有任何原子加。
7. H100：4096 / (4 × 96 × 16) = 0.67；B200：8192 / (4 × 96 × 16) = 1.33——exp 比矩阵乘还慢。
8. (QJ)(KJ)ᵀ = QJJᵀKᵀ = QKᵀ：旋转在乘积里被约掉了。PV 的乘积里只有 V 一个出现 J：P(VJ) = OJ，输出被旋转了，必须在最后再乘一次 Jᵀ 才能得到 O。可以做（额外一次 O(d log d) 的变换），只是 FlashAttention-3 没有这样做。
9. ℓᵢ 是至多 N 个不超过 2^τ 的数之和，ℓᵢ ≤ N·2^τ。要求 N·2^τ < 2^128，即 τ < 128 − log₂N。N = 2^20 时 τ < 108。论文取 8，余量很大（还要考虑 Oᵢ 的分量可能更大，以及 BF16 的 P 本身的范围）。
10. 按到达顺序：1 → A，2 → B，3 → A（4），3 → B（5），4 → A（8），5 → B（10），完成时间 10——最长的任务排在最后，没有别的任务能和它”凑齐”。LPT 按 [5, 4, 3, 3, 2, 1]：5 → A，4 → B，3 → B（7），3 → A（8），2 → B（9），1 → A（9），完成时间 9。下限 ⌈18/2⌉ = 9，LPT 达到了下限。
11. 不切分：2 × 8 = 16 个线程块，16/132 ≈ 12%。按”波次效率”的思路，切 8 段得到 128 个线程块，一波 97%；切 33 段得到 528 个线程块，4 波整好 100%，但合并的开销更大。像 flash-attn 那样取”效率不低于最佳 85% 的最小段数”，答案是 8 段。

10. 延伸阅读 · Further reading

• FlashAttention 系列：Dao et al. (2022)；Dao (2023)；Shah et al. (2024)；Zadouri et al. (2026)。源码：github.com/Dao-AILab/flash-attention（FlashAttention-4 的前向在 flash_attn/cute/flash_fwd_sm100.py）。NVIDIA CUTLASS 的 examples/77_blackwell_fmha 是另一个 Blackwell 上的注意力实现。
• IO 复杂度的经典：Aggarwal & Vitter (1988) The Input/Output Complexity of Sorting and Related Problems——FlashAttention 下界证明的思路来源。
• 在线 softmax：Milakov & Gimelshein (2018) Online normalizer calculation for softmax；Rabe & Staats (2021) Self-attention Does Not Need O(n²) Memory。
• Flash-Decoding：Dao, Haziza, Massa, Sizov (2023) 的博客 Flash-Decoding for long-context inference。
• Horace He：Making Deep Learning Go Brrrr From First Principles（slides 第 8 页引用的博客，讲计算、带宽、开销三种瓶颈）。
• 近似注意力的综述：Tay et al. Long Range Arena 与 Efficient Transformers: A Survey。

11. 下一讲预告 · Next lecture

FlashAttention 让一次注意力算得又快又省；可是一个线上服务同时面对成千上万个请求，它们长度不一、到达时间随机、很多还共享相同的前缀（系统提示、少样本示例）。第 22 讲进入LLM 服务：SGLang 的请求调度器怎样做连续批处理（为什么按”迭代”而不是按”请求”组批），RadixAttention 怎样用一棵基数树在请求之间复用 KV 缓存，以及缓存感知的调度为什么能把命中率提高好几倍。