Lecture 11 · 内核优化技术：FlashAttention、flash-decoding 与 megakernel · Kernel Optimization Techniques

对应 Maps to：Stanford CS349E / EE339 Efficient ML Inference at Scale（Autumn 2026），原课第 12 次课 · Thu Oct 29 · Kernel Optimization Techniques（Priyanka Raina）。阅读材料：T. Dao, D. Y. Fu, S. Ermon, A. Rudra, C. Ré, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness（NeurIPS 2022）；B. Spector, J. Juravsky, S. Sul, O. Dugan, D. Lim, D. Fu, S. Arora, C. Ré（Hazy Research），Look Ma, No Bubbles! Designing a Low-Latency Megakernel for Llama-1B（博客，2025-05-27）。 依据说明：课程网站只公开了课表和阅读材料，没有 slides。本讲以两篇阅读材料为主线，补充了 FlashAttention-2（Dao 2023）、FlashAttention-3（Shah et al. 2024）的摘要结论和 flash-decoding 博客（Dao、Haziza、Massa、Sizov，PyTorch 博客 2023）的实测数据。补充材料是我选的，不代表课堂上实际讲了什么。 配套代码 Code：code/lecture11/（numpy；一个共享模块 + 5 个脚本，每个都在几秒内跑完） 先修：第 10 讲（分块、共享内存、Tensor Core、融合的 softmax）、前置讲 P3 §5（安全 softmax、在线 softmax 与它的合并）、第 2 讲（注意力、GQA、KV cache）。 预计阅读时间：3 小时（含运行代码）

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

0.1 学习目标 Learning goals

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

1. 用”显存读写次数”（IO）而不是 FLOPs 分析注意力，说出标准写法为什么受限于带宽。
2. 推导在线 softmax 的换算因子，写出 FlashAttention 前向的分块算法；解释两种循环顺序（论文算法 1 和 FlashAttention-2）各自反复读写什么、各自的并行性。
3. 复述 FlashAttention 的 IO 复杂度 $\Theta(N^2 d^2 / M)$ 和下界，并用代码的计数验证；说清它在什么条件下省得多、什么条件下省得少。
4. 推导反向传播里的 $D_i = \sum_j dP_{ij} P_{ij} = dO_i \cdot O_i$，解释”重算 $P$ 比存储 $P$ 便宜”。
5. 解释 decode 时 FlashAttention 为什么占不满 GPU，flash-decoding 怎样沿 KV 的长度切分并用 log-sum-exp 精确合并；用一个 CTA × 波次的时间模型复现 flash-decoding 博客的实测趋势。
6. 说出一次批 1 的 decode 里”内核之间的空隙”由哪几部分组成，megakernel 怎样消灭它们，以及 CUDA Graphs、PDL 能做到哪一步。

0.2 本讲地图 Map

§1  两类内核的瓶颈：搬得太多、启动得太多
§2  IO 感知：标准注意力为什么慢
§3  在线 softmax：一遍扫完 K、V                                    ← 代码 01
§4  FlashAttention 前向：分块、两种循环顺序、IO 复杂度；FA-2、FA-3    ← 代码 02
§5  反向传播：重算；因果掩码                                        ← 代码 03
§6  decode：flash-decoding                                          ← 代码 04
§7  融合的一般规律
§8  megakernel：消灭内核之间的空隙                                   ← 代码 05
§9  现状（2025–2026）
§10–14  小结、常见误区、练习（含答案）、延伸阅读、下一讲预告

0.3 本讲的代码

flash_lab.py 用 numpy 真的实现了标准注意力、FlashAttention 的前向（两种循环顺序）和反向、flash-decoding 的分段与合并，全部和标准写法比较到 $10^{-15}$，并数出每种算法读写显存的元素个数（“显存”是一个概念：算法里明确写回、读回的数据）。decode 注意力的速度和 megakernel 的收益用时间模型估计，参数和假设都写在代码里。

1. 两类内核的瓶颈 · Two kinds of kernel bottlenecks

第 10 讲的矩阵乘法是”算得多、搬得少”的内核（强度上限 683 FLOP/字节），难点是把 Tensor Core 喂饱。推理里还有两类内核，难点不一样：

1. 搬得太多：中间结果比输入大得多。注意力的输入是 $Q, K, V$（各 $N \times d$），中间的分数矩阵 $S = QK^\top$ 是 $N \times N$——$N = 8192$、$d = 128$ 时，一个头的 $S$ 是输入的 64 倍。如果把 $S$ 写回显存再读出来，时间就被这一来一回占满。FlashAttention 让 $S$ 从不离开片上存储。
2. 启动得太多：每个内核只有几微秒的活。批 1 的 decode 一次前向要跑上百个内核，每个之间都有启动、等第一批数据、等最后一个块的空隙。megakernel 把整个前向写成一个内核。

两件事背后是同一个原则：数据搬运和同步是贵的，运算是便宜的（第 1 讲）。所以 FlashAttention 宁可多算（反向时重算 $S$）也不多搬；megakernel 宁可在每个 SM 上跑一个解释器，也不让 GPU 空着等下一个内核。

2. IO 感知：标准注意力为什么慢 · IO-awareness

注意力 $O = \mathrm{softmax}(QK^\top / \sqrt{d})\, V$。标准的写法（FlashAttention 论文的算法 0）分三步，每步一个（或几个）内核：

S = Q Kᵀ / √d          读 Q、K（各 N×d），写 S（N×N）
P = softmax(S)          读 S，写 P（N×N）
O = P V                 读 P、V，写 O（N×d）

显存读写 = $4Nd + 4N^2$ 个元素，$\Theta(Nd + N^2)$。运算 = $4N^2 d$（两个矩阵乘法）。算术强度 ≈ $4N^2 d / (4N^2 \times 2\text{ 字节}) = d/2$ FLOP/字节——$d = 64$ 时只有 32，远低于 A100 的临界强度（约 200）。所以标准注意力受限于带宽，而且大部分带宽花在搬 $N \times N$ 的中间矩阵上，不是搬输入。

FlashAttention 论文把这种分析叫 IO 感知（IO-aware）：衡量一个算法时，不只数 FLOPs，还数它在存储层次之间搬了多少数据。这正是第 9 讲”访问次数 × 每次的代价”的思路。GPU 上两层就够了：显存（HBM，几十 GB、几 TB/s）和片上 SRAM（每个 SM 一百多 KB、带宽高一个数量级）。

问题于是变成：能不能不把 $S$、$P$ 写回显存？ 难点在 softmax——每个输出元素需要整行分数的最大值和总和。

3. 在线 softmax：一遍扫完 K、V · Online softmax

代码：01_online_softmax.py

3.1 换算因子

前置讲 P3 §5.3 讲过在线 softmax：扫一遍数据，同时维护目前见过的最大值 $m$ 和 $\ell = \sum_j e^{x_j - m}$。来了一块新数据（块内最大值 $\tilde m$）：

旧的 $\ell$ 是”按旧的 $m$ 缩放”的，乘以 $\alpha = e^{m - m_{\text{new}}}$ 就换算成”按新的 $m$ 缩放”。FlashAttention 把这个技巧推广到连同乘 $V$：再维护一个 $d$ 维的累加器

扫完以后 $o = \mathrm{acc} / \ell$。一个 query 对 $S$ 个 key 的注意力，只需要一遍读 $K$、$V$，不需要存长度 $S$ 的分数向量。

代码 01 演示 1 对比三种算法（$S = 4096$、$d = 64$、每块 256）：存下整行分数（3 遍）、在线 softmax（2 遍：第一遍求 $m$、$\ell$，第二遍重算分数并乘 $V$）、在线 softmax × V（1 遍），结果相差 $10^{-16}$ 量级。一遍算法的前几块：

第 1 块的最大值更大，$m$ 变大，之前累加的 $\ell$ 和 acc 整体乘 0.4087 换算——$\ell$ 反而变小了，因为它现在是按更大的 $m$ 计算的。

3.2 局部状态可以任意合并

把一行切成任意几段，每段算出局部状态 $(m_s, \ell_s, o_s)$，再两两合并：

这个合并运算满足结合律和交换律（代码 01 演示 2：长度 3000 的一行随机切成 12 段，按顺序、树形、乱序合并，和一次算完相差 $5 \times 10^{-16}$）。这给了并行的自由：不同的 SM 可以算同一行的不同段（flash-decoding，§6），不同的 GPU 可以算序列的不同部分（ring attention）。

3.3 推迟归一化

FlashAttention 论文的算法 1 每处理一块就把输出重新归一化（除以新的 $\ell$）。FlashAttention-2 改成只在最后除一次：中间只把累加器乘 $\alpha$。为什么要在乎这一点点运算？因为 exp、max、除法这些”非矩阵乘法”的运算不能用 Tensor Core：A100 的 FP16 矩阵乘法峰值 312 TFLOPS，非矩阵乘法的 FP32 只有 19.5 TFLOPS，每次运算贵 16 倍（FlashAttention-2 论文的说法）。代码 01 演示 3（一个头，$N = 4096$、$d = 128$）：非矩阵乘法的运算只有矩阵乘法的 1.6%（每步归一化）和 1.0%（推迟归一化），按 16 倍折算，相当于矩阵乘法时间的 25% 和 16%——不可忽略。

3.4 精度

真实的内核里：$m$、$\ell$、累加器都用 FP32；$P = e^{S - m}$ 转成 FP16 / BF16 后再和 $V$ 做第二个矩阵乘法（Tensor Core 的输入只能是低精度），累加仍是 FP32。代码 01 演示 4（$S = 8192$）：

$\ell$ 绝不能用低精度累加：部分和变大以后，后面加的小数被”吞掉”（前置讲 P3 §4.3 的停滞）。减最大值也不只是为了 FP32：FP16 的 $e^x$ 在 $x \approx 11.09$ 就溢出了。

4. FlashAttention 前向 · The forward pass

代码：02_flash_forward.py

4.1 分块

把 §3 用到 $N$ 个 query 上：$Q$ 切成 $T_r = \lceil N/B_r \rceil$ 块（每块 $B_r$ 行），$K$、$V$ 切成 $T_c = \lceil N/B_c \rceil$ 块。片上同时放得下 $Q_i$、$K_j$、$V_j$、$O_i$ 四块，就能在片上算出 $S_{ij} = Q_i K_j^\top$、更新 $(m_i, \ell_i, O_i)$，$S_{ij}$ 用完就扔。论文的块大小（片上 SRAM 能放 $M$ 个元素）：

“4d”来自四块都要放下。论文的定理 1：FlashAttention 算的是精确的注意力（不是近似），运算量 $O(N^2 d)$，除输入输出外只需要 $O(N)$ 的额外显存（每行的 $m$、$\ell$）。

代码 02 演示 1：两种循环顺序 × 带不带因果掩码 × 长度是不是块的倍数（包括 $N = 1$、$N_k = 1000$ 的 decode 情形），都和标准写法相差不到 $10^{-15}$。

4.2 两种循环顺序

FlashAttention-2 的顺序还带来了沿序列长度的并行：FA1 只按 (批, 头) 分配线程块，批 × 头数小于 SM 数时（长序列训练常常这样）GPU 占不满；FA2 每个 $Q$ 块一个 CTA，序列长就有足够多的块。FA2 的第三个改进是块内的分工：一个 CTA 的 4 个 warp 各负责 $Q_i$ 的一部分行（“split-Q”），而不是各负责 $K_j$ 的一部分（“split-K”，最后要通过共享内存合并），省掉了 warp 之间的通信。论文报告比 FlashAttention 快约 2 倍，前向达到 A100 理论峰值的 50%–73%。

4.3 IO 复杂度

论文的定理 2：标准注意力的显存访问是 $\Theta(Nd + N^2)$，FlashAttention 是 $\Theta(N^2 d^2 M^{-1})$；命题 3：对所有 $M \in [d, Nd]$，任何精确注意力算法都做不到 $o(N^2 d^2 / M)$——FlashAttention 在这个意义下最优。

直觉：算法 1 的顺序里，$Q$、$O$ 被读写 $T_c = N / B_c \approx 4Nd/M$ 次，每次 $Nd$ 个元素，合计 $\approx 4N^2 d^2 / M$。

代码 02 演示 2 真的跑一遍、数读写（一个头，$d = 64$，$M$ = 100 KB = 51,200 个 FP16 元素 → $B_c = 200$、$B_r = 64$）：

几点值得注意：

• 计数和定理 2 一致：读写 ÷ $(N^2 d^2 / M)$ 在算法 1 顺序下是 12.6–19、FA-2 顺序下是 25–28，几乎不随 $N$ 变。
• 省多少取决于 $M / d^2$。$d = 64$、$M \approx 5$ 万时，前向的读写只比”最少读写的标准注意力”少 3–4 倍，不是几十倍。SRAM 越大、头维度越小越省（代码 02 演示 3：$N = 4096$ 时 SRAM 从 16 KB 到 256 KB，算法 1 顺序的读写从 207 MB 降到 14 MB）。
• 论文的公式让 $B_r = \min(B_c, d) = 64$，对 FA-2 顺序不利（$K$、$V$ 被读 $N/64$ 次，读写是算法 1 顺序的 2 倍）。FlashAttention-2 实际常用 $B_r = 128$：读写和算法 1 顺序相当，同时换来了并行性。真实的块大小由寄存器、共享内存和 Tensor Core 的形状共同决定（第 10 讲 §8）。

4.4 论文图 2 的设置

论文用 GPT-2 medium 的形状（$N = 1024$、$d = 64$、16 个头、批 64，A100）测了前向 + 反向。代码 02 演示 4 按同样的形状数我们的算法：

我们的读写是”每个算法最少要读写的元素”：论文测到的标准实现（PyTorch）还有单独的掩码、dropout、softmax 内核，每个都把 $N \times N$ 的矩阵再读写一遍，所以 40.3 GB 比我们的 25.4 GB 多。论文的 GFLOPs 比我们按”批 64 × 16 头”算的小一个数量级，论文没有说明它的统计范围，我们没能对上绝对值；两者的比值接近（FlashAttention 多做 13%–17% 的运算）。结论一致：运算多了一点，读写少了很多，时间少了几倍。

4.5 FlashAttention-3：Hopper 上的异步

FlashAttention-2 在 H100 上只用到约 35% 的峰值（FlashAttention-3 论文的说法）：Hopper 的 Tensor Core 太快，同步的”读 → 算”跟不上。FlashAttention-3（Shah et al. 2024）用了第 8 讲 §8 和第 10 讲 §8 的全部 Hopper 特性：

1. warp 专门化 + TMA：producer warp 只发 TMA 读 $K_j$、$V_j$，consumer warpgroup 只做 WGMMA，搬运和计算完全重叠；
2. 矩阵乘法和 softmax 交错（ping-pong）：一个 warpgroup 在做 softmax 的 exp 时，另一个在做下一块的矩阵乘法——把 §3.3 那 16 倍贵的非矩阵乘法藏在矩阵乘法后面；
3. FP8：块量化 + incoherent processing（用随机正交变换打散离群值，第 3 讲）。

论文报告比 FlashAttention-2 快 1.5–2.0 倍，FP16 达到 740 TFLOPS（75% 的利用率），FP8 接近 1.2 PFLOPS，FP8 的数值误差比基线 FP8 注意力低 2.6 倍。

5. 反向传播：重算比存储便宜 · The backward pass

代码：03_flash_backward.py

推理不需要反向传播，但 FlashAttention 的反向是”重算 vs 存储”这个系统原则最干净的例子，而且训练的长上下文能力完全依赖它。

5.1 公式

前向 $S = QK^\top \cdot s$（$s = 1/\sqrt d$），$P = \mathrm{softmax}(S)$，$O = PV$。给定 $dO$：

其中 $D_i = \sum_j dP_{ij} P_{ij}$（softmax 的雅可比矩阵乘向量）。直接算 $D_i$ 要整行 $P$。FlashAttention 用一个恒等式：

只要 $O$ 和 $dO$ 的一行，可以在反向开始前一次算好。

标准写法要用前向存下的 $P$（$N \times N$）。FlashAttention 只存每行一个 $L_i = m_i + \log \ell_i$（logsumexp），反向时在片上重算 $S_{ij}$，再 $P_{ij} = e^{S_{ij} - L_i}$——一次 exp 就恢复出归一化后的 $P$，不需要 $m$、$\ell$ 分开存。

代码 03 演示 1：FlashAttention 的反向和标准反向的梯度相差不到 $10^{-15}$，和有限差分相差约 $10^{-9}$（中心差分本身的误差）。

5.2 代价与收益

• 显存（代码 03 演示 2；32 个注意力矩阵、FP16）：标准写法为反向存 $P$，$N = 32\text{K}$ 时每层 68.7 GB——一块 80 GB 的 GPU 连一层都存不下；FlashAttention 存 $L$，4.2 MB。
• 运算：前向 2 次矩阵乘法；标准反向 4 次；FlashAttention 反向 5 次（多一次重算 $S$）。前向 + 反向合计是标准写法的 $7/6 \approx 1.167$ 倍。
• 读写（代码 03 演示 3；一个头，$d = 128$，块 128×128，所有读写都算成显存）：只比”最少读写的标准写法”少约 1.8 倍，而且不随 $N$ 变——FlashAttention 的反向对每个 $(i, j)$ 块都要读 $Q_i$、$dO_i$ 并读写 $dQ_i$，这部分约 $4N^2$ 个元素，和标准写法的 $N \times N$ 矩阵（约 $11N^2$）同一个量级。实际测到的加速更大：标准实现还有更多的内核；被反复读的 $Q_i$、$dO_i$ 块在很多 CTA 之间共享，大部分在 L2 里命中——这一项我们的计数没有算，要用 profiler 才能确定。

“重算比存储便宜”在本课里反复出现：激活值重算（gradient checkpointing）、KV cache 的重算 vs 换出（第 7 讲）、megakernel 里 RMSNorm 的重复计算（§8）。共同的理由：算力便宜，带宽和容量贵。

5.3 因果掩码：跳过一半的块

因果掩码让分数矩阵的一半是 $-\infty$。整块都被掩码的 $(i, j)$ 不读、不算；只有对角线上的块要在块内加掩码。$T = N / B$ 个块行时，算的块数是 $T(T+1)/2$，占比 $\frac12 + \frac{1}{2T}$（代码 03 演示 4：$N = 4096$，块 64 / 128 / 256 时 FLOPs 是不带掩码的 0.508 / 0.516 / 0.531，和公式一致）。对角线块里约一半是算了但不用的，块越大浪费越多。

6. decode：flash-decoding · Attention at decode time

代码：04_flash_decoding.py

6.1 问题：占不满 GPU

decode 时每个序列每步只有 1 个新 query。FlashAttention 沿 query 切块、每块一个 CTA，于是 CTA 数 = 批 × 头数：批 1、16 个头只有 16 个 CTA，A100 的 108 个 SM 只用了 15%。更糟的是，每个 CTA 的 Tensor Core 仍按 128 行的 $Q$ 块算，其中只有 1 行是真的 query——flash-decoding 博客说”批 1 时 FlashAttention 用不到 GPU 的 1%“，指的是这两个因素相乘。而长上下文时批往往最小（KV cache 占满了显存，第 1 讲 §4.4），偏偏这时注意力是 decode 的大头。

6.2 沿 KV 的长度再切一刀

flash-decoding（Dao、Haziza、Massa、Sizov，2023）分三步：

1. 把 $K$、$V$ 切成若干段（只是视图，不搬数据）；
2. 每个 (序列, KV 头, 段) 一个 CTA，用 FlashAttention 算这一段的输出 $o_s$（段内归一化），外加每行一个数 $\mathrm{lse}_s = \log \sum_{\text{段}} e^{\text{分数}}$；
3. 一个小内核合并：$\mathrm{lse} = \log \sum_s e^{\mathrm{lse}_s}$，$o = \sum_s e^{\mathrm{lse}_s - \mathrm{lse}}\, o_s$。

这就是 §3.2 的合并，换成了 lse 的写法。代码 04 演示 1：切成 1、7、64、1000 段，和直接计算相差 $10^{-16}$ 量级。GQA（第 2 讲）：16 个 query 头共享 2 组 $K$、$V$ 时，把同一组的 8 个 query 当成一个 8 行的 $Q$ 块，一次读 $K$、$V$ 服务 8 个头——flash-decoding 的实现就是这样把 GQA 的头”打包”进矩阵乘法的 M 维的。

6.3 时间模型 vs 实测

博客在 A100 上做了微基准（批 1–256、序列 256–128K，16 个 query 头、2 个 KV 头、$d = 128$，即 CodeLlama-34B 在 4 卡上的形状）。我们写了一个”CTA × 波次”的模型（flash_lab.decode_attention_time）：

• 每个 CTA 的时间 = max(补齐后的运算量 / (每 SM 的 Tensor Core 峰值 × 效率)，读的字节 / 每 SM 能拉的带宽)；
• 每 SM 能拉的带宽 = min(约 80 GB/s（按 Little 定律：在途约 64 KB、延迟约 0.8 µs），显存带宽 / 同时活跃的 SM 数)；
• 总时间 = 波数 × 每个 CTA 的时间 + 每个内核约 5 µs 的固定开销；
• FlashAttention-2：每个 (序列, query 头) 一个 CTA，$Q$ 块按 128 行算；flash-decoding：每个 (序列, KV 头, 段) 一个 CTA，8 个头打包、补到 64 行。

这些参数是事先按 Little 定律和 FlashAttention-2 论文（前向达到峰值的 50%–73%，取 0.7）定的，没有按实测调。

（完整的 10 行见代码输出。）模型对 FlashAttention-2 的估计在实测的 ±12% 以内，并且告诉我们为什么：批大、序列短时受限于带宽——每个 query 头单独读一遍它那组的 $K$、$V$（读了 8 倍的数据）；批小、序列长时受限于计算——16 个 CTA 在 16 个 SM 上做了 128 倍的无用功（补齐的 127 行）。flash-decoding 几乎不随序列长度变：这张表里批 × 序列长度是常数，只要 CTA 填满了 GPU，时间就是”读一遍 KV + 两个内核的开销”。博客说注意力最多快 50 倍；表里最大的实测倍数是 36（批 1、64K）和 43（128K）。

6.4 段数与形状

段数太少，CTA 填不满 SM；段数太多，CTA 超过一波、每段太短，固定开销和合并的代价占比变大。批 64、序列 2K 时，不切只有 128 个 CTA（比 108 个 SM 多一点，第二波几乎空着，151 µs），切 4–16 段降到约 100 µs——波次量化（前置讲 P2 §7.3）和并行度一样重要。真实的实现按批大小、长度和 SM 数自动选段数。

H100 上的 Llama-3.1-8B（32 个 query 头、8 个 KV 头）：一层的 KV 在 128K 时 537 MB，按满带宽也要 160 µs，32 层就是 5 ms。flash-decoding 让 decode 注意力接近”读一遍 KV”的下限（模型 170 µs，FlashAttention-2 式批 1 时 1.64 ms）；再往下只能减少 KV 的字节：GQA / MLA（第 2 讲）、KV 量化（第 3 讲）、稀疏注意力。第 7 讲的分页 KV 需要注意力内核支持”按块表间接寻址”——vLLM 的 PagedAttention 内核、FlashInfer 都在 flash-decoding 的框架上做了这件事。

7. 融合的一般规律 · Fusion in general

FlashAttention 是融合的极致例子，但推理里到处都是融合：

规律：

1. 融合省的是中间结果的往返，对受限于带宽的算子（逐元素、归一化、softmax）收益最大（第 10 讲 §10.3 的 softmax：流量少 4 倍）。
2. 归约是融合的障碍：一个元素的输出依赖一整行（softmax、LayerNorm）时，要么把整行放进片上，要么找到可以分块合并的形式（在线 softmax、Welford 方差）。
3. 融合有代价：融合后的内核要同时放下所有算子的数据，寄存器和共享内存的压力变大、占用率下降；有时还要重算（RMSNorm 的 RMS 每个块各算一次）。
4. 谁来融合：手写（FlashAttention、cuBLASLt 的 epilogue）、模板库（CUTLASS 的 epilogue visitor）、编译器（第 12 讲：TorchInductor 把逐元素、归约、矩阵乘法的 epilogue 自动融合）。PyTorch 2 论文的消融实验里，去掉融合和内联后，Inductor 生成的代码反而比 eager 模式慢（第 12 讲 §6）。

8. megakernel：消灭内核之间的空隙 · Megakernels

代码：05_megakernel.py

8.1 问题：一百多个内核，每个只跑几微秒

Hazy Research 的博客研究的是最极端的情形：Llama-3.2-1B、批 1 的 decode——聊天机器人”打字”时的延迟。批 1 时每个权重只用一次，强度约 1 FLOP/字节，完全受限于带宽。理想时间 = 读一遍所有权重：

可现有的引擎（vLLM、SGLang，都已经用了 CUDA Graphs 和 torch.compile）在 H100 上最多用到约 50% 的带宽。一次前向被切成约 100 个内核，每个算子只读几 MB 到几十 MB：一个 8 MB 的 O 投影按满带宽只要 2.5 µs，和一次内核启动的开销同一个量级。博客列出了三种空隙：

1. 启动与收尾：同一个 stream 上，前一个内核的所有块结束，下一个才能开始。博客测到 H100 上一个空内核在 stream 上约 2.1 µs，用 CUDA Graphs 也要约 1.3 µs。
2. 序幕（load latency）：内核开始后，要等第一批权重从显存到达，几千个周期里 GPU 什么也不做。
3. 尾部（tail effects）：最后一波块没占满 SM。博客的例子：B200 有 148 个 SM，512 个块 = 3 × 148 + 68，最后一波 80 个 SM 空闲。

Hopper 的 PDL（Programmatic Dependent Launch）让下一个内核在前一个结束前就开始执行序幕，但同步仍然太粗：用 cudaGridDependencySynchronize 等待的是前一个内核的全部结果——注意力要等所有的 Q、K、V 都算完，哪怕它只需要其中一个头。

博客的估计：每层 7 个内核 × 16 层 × 每个约 5 µs 的停顿，H100 上每秒最多约 770 次前向。

8.2 设计

megakernel 把整个前向写成一个常驻的内核：

• 解释器：每个 SM 执行一串”指令”，所有指令用同一个 CUDA 模板实现。指令序列（调度）在 Python 里事先排好，可以复用几百次前向。指令集（每层 5 种 + 最后 1 种）：RMSNorm + QKV + RoPE；注意力；注意力的归约（长序列时的 flash-decoding 合并）；O 投影 + 残差；RMSNorm + gate/up + SiLU；down 投影 + 残差；最后的 RMSNorm + lm_head。
• 共享内存分页：H100 共享内存的前 213 KB 分成 13 页、每页 16 KiB。指令显式地申请、释放页；一条指令释放的页立刻交给下一条指令，下一条指令的权重可以提前开始读——权重不依赖上一个算子的结果，只有激活依赖。
• 细粒度同步：显存里放一个计数器数组，初值为 0。每条指令完成时把对应的计数器加一；开始前等它依赖的计数器到达目标值。down 投影被切成 4 块，各有自己的计数器，每块只等自己那部分输入——比 PDL 的”等前一个内核全部结束”细得多。

博客报告：H100 上用到 78% 的带宽、每次前向不到 1 ms，比 vLLM 快约 2.5 倍、比 SGLang 快 1.5 倍以上；B200 上不到 680 µs，比 vLLM 快 3.5 倍以上。B200 上约 600 µs 的时间分解：存取激活和检查一致性约 250 µs（每次读写约 500 ns）、RMSNorm 和矩阵乘向量约 200 µs、等权重约 30 µs（四成在 lm_head）、warp 同步约 40 µs、设置等约 80 µs。

8.3 时间线模型

代码 05 写了一个时间线模型（flash_lab.timeline）：多内核的写法里，每个内核 = 启动间隙 + 序幕（假设 1 µs，只付一次）+ 按”每 SM 一份带宽”读完它那份权重（块数不是 SM 数的整数倍时有尾部）；几乎不读数据的小内核（单独的 RMSNorm、RoPE、残差）按 2 µs 算。megakernel 里，权重在共享内存的缓冲允许的范围内连续地读，只有激活的依赖要等（上一个算子结束 + 1 µs 的同步），每条指令再加 0.5 µs 的固定开销。带宽效率按 0.9。启动间隙用博客的数字，其余是假设。

时间的分解（7 个内核 + CUDA Graphs）：读权重 739 µs + 带宽效率损失 82 µs + 启动间隙 148 µs + 序幕 81 µs + 尾部 49 µs + 小内核 66 µs。空隙合计约 350 µs——和博客”每个内核约 5 µs × 112 个”的估计（约 560 µs）同一个量级。

模型比实测乐观：多内核的 63% 高于 vLLM、SGLang 实测的不到 50%（真实的引擎还有 CPU 上的调度、采样）；megakernel 的 88% 高于博客的 78%（模型没算矩阵乘向量本身的时间和每条指令读写激活的开销）。但结论不依赖具体参数（代码 05 演示 4）：序幕、小内核、同步、指令开销、带宽效率各自变化时，megakernel 都比多内核快 1.3–1.5 倍、带宽利用率高 20–29 个百分点。B200 上带宽是 H100 的 2.4 倍，读权重的时间短了，固定的空隙占比更大：模型里多内核只有 46%，megakernel 优势 1.8 倍——和博客”B200 上比 vLLM 快 3.5 倍以上（H100 上 2.5 倍）“的方向一致。硬件越快，软件的空隙越显眼。

8.4 尾部效应

块数恰好比 SM 数多一点最糟。常见的缓解：持久化内核（块数 = SM 数，每个块循环处理多份工作）、split-K、改块大小。megakernel 天然是”持久化”的：每个 SM 一直在跑，没有”一个内核的最后一波”。

8.5 megakernel 的代价

• 工程复杂度：所有算子要写进同一个模板，资源（寄存器、共享内存页）要在指令之间手工管理；调度要事先排好。博客的实现基于 ThunderKittens。
• 只对”空隙占比大”的情形有用：批大、序列长时每个内核的活多，空隙占比小，megakernel 的收益随之变小；那时瓶颈回到 Tensor Core 和带宽本身。
• 和其他手段的关系：CUDA Graphs 去掉 CPU 的启动开销，但去不掉内核之间的依赖等待；PDL 让序幕和前一个内核重叠，但同步太粗；投机解码（第 6 讲）和批处理（第 5 讲）则是另一个方向——让每个内核的活变多。

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

• 注意力内核：FlashAttention-3（Hopper）之后，Blackwell 上的新版本利用第五代 Tensor Core 和 Tensor Memory；FlashInfer、vLLM / SGLang 的内核在 flash-decoding 的框架上支持分页 KV、前缀共享（级联注意力）、投机解码的树形掩码；DeepSeek 开源了 MLA 的 decode 内核 FlashMLA。长上下文的 decode 注意力已经接近”读一遍 KV”的下限，进一步的收益来自减少 KV 的字节（MLA、KV 量化、稀疏注意力）。
• megakernel 和”持久化”的思路在扩散：除了 Hazy 的 Llama 演示，也有把整个模型编译成一个内核的编译器研究，以及把通信（all-reduce）融合进计算内核的做法。批 1 的低延迟推理（聊天、智能体的工具调用循环，第 15 讲）是它们的主要动机。
• 内核由谁来写：FlashAttention 这类”算法 + 硬件”深度耦合的内核仍然主要靠专家手写；编译器（第 12–13 讲）能自动做大部分逐元素和归约的融合；LLM 写内核（第 14 讲）在 KernelBench 上开始能做出一些融合，但 FlashAttention 级别的算法重排仍然很难。

10. 小结 · Summary

1. 注意力的瓶颈是 $N \times N$ 中间矩阵的显存往返，不是 FLOPs。IO 感知的分析：标准 $\Theta(Nd + N^2)$，FlashAttention $\Theta(N^2 d^2 / M)$，后者在 $M \in [d, Nd]$ 上是最优的。
2. 在线 softmax 的换算因子 $\alpha = e^{m - m_{\text{new}}}$ 让输出累加器可以”边扫边算”；局部状态 $(m, \ell, o)$ 的合并满足结合律，可以任意切分、并行、合并。
3. FlashAttention 前向：$Q$、$K$、$V$ 分块，$S$ 只在片上存在。论文算法 1 的顺序反复读写 $Q$、$O$；FlashAttention-2 的顺序反复读 $K$、$V$，但不同 query 块完全独立，可以沿序列并行，并推迟归一化以减少非矩阵乘法的运算。
4. 省多少取决于 $M / d^2$：$d = 64$、100 KB SRAM 时前向读写只少 3–4 倍；真实实现的加速还来自少了大量的辅助内核。
5. 反向不存 $P$，只存 $L$，在片上重算——多 1/6 的运算，省 $N^2$ 的显存；$D_i = dO_i \cdot O_i$ 是关键的恒等式。因果掩码跳过 $\frac12 - \frac{1}{2T}$ 的块。
6. decode 时 query 只有 1 个，FlashAttention 的 CTA 太少、每个 CTA 的 Tensor Core 大部分在算补齐的 0。flash-decoding 沿 KV 的长度切分、用 lse 合并，让时间接近”读一遍 KV”。一个 CTA × 波次的模型能复现实测的趋势（FA2 在 ±12% 以内）。
7. 批 1 的 decode 有一百多个内核，每个之间的启动、序幕、尾部合起来可以占到总时间的三分之一以上。megakernel 用常驻的解释器、共享内存分页预取权重、计数器做细粒度依赖，把带宽利用率从约 50% 提高到 78%（Hazy 在 H100 上的实测）。硬件越快，软件的空隙越显眼。

11. 常见误区 · Common pitfalls

1. “FlashAttention 是近似注意力”——它算的是精确的注意力，和标准写法相差只是浮点舍入（§4.1）。
2. “FlashAttention 减少了运算量”——不带掩码时前向的运算一样，反向还多 1/6；它减少的是显存读写（§4.3、§5.2）。
3. “FlashAttention 总能省几十倍的读写”——省多少取决于 $M/d^2$；$d = 64$、100 KB SRAM 时前向只少 3–4 倍（§4.3）。
4. “在线 softmax 的 $\ell$ 可以用 FP16 累加”——会停滞，误差到 $10^{-1}$ 量级（§3.4）。
5. “decode 时用 FlashAttention 就够了”——批小时它只能用上很少的 SM，而且每个 CTA 的 Tensor Core 大部分在算补齐的行；长上下文要 flash-decoding（§6.1）。
6. “段切得越多越好”——超过一波的 CTA 会带来波次量化，固定开销和合并的代价也变大（§6.4）。
7. “用了 CUDA Graphs 就没有启动开销了”——Graphs 去掉了 CPU 端的开销，但内核之间的依赖等待、序幕和尾部还在（博客：每个内核仍约 1.3 µs，§8.1）。
8. “megakernel 对所有推理都有用”——它解决的是”每个内核的活太少”；批大、序列长时收益随之变小（§8.5）。

12. 练习 · Exercises

练习 1（在线 softmax） 一行分数分两块到达：$[1, 3]$ 和 $[2, 5]$。按在线 softmax 写出每块之后的 $m$ 和 $\ell$，并和一次算完比较。

答案：第一块：$m = 3$，$\ell = e^{-2} + e^{0} = 1.1353$。第二块：$\tilde m = 5$，$m = 5$，$\alpha = e^{3 - 5} = 0.1353$，$\ell = 1.1353 \times 0.1353 + (e^{-3} + e^{0}) = 0.1537 + 1.0498 = 1.2034$。一次算完：$e^{-4} + e^{-2} + e^{-3} + 1 = 1.2034$。

练习 2（标准注意力的读写） 一个头，$N = 4096$、$d = 64$，FP16。按算法 0 数读写的字节。

答案：$4Nd + 4N^2 = 4 \times 262{,}144 + 4 \times 16{,}777{,}216 = 68{,}157{,}440$ 个元素 × 2 = 136.3 MB（代码 02 演示 2 的数）。

练习 3（FlashAttention 的读写） 同上，$M$ = 51,200 个元素，算法 1 的顺序（$B_c = 200$、$B_r = 64$）。读写多少？

答案：$T_c = \lceil 4096/200 \rceil = 21$。$K$、$V$ 各读一次 $2Nd$；每个 $j$：读 $Q$、$O$、$\ell$、$m$（$2Nd + 2N$），写 $O$、$\ell$、$m$（$Nd + 2N$）；再加初始化写 $Nd + 2N$。合计 $524{,}288 + 21 \times 532{,}480 + 21 \times 270{,}336 + 270{,}336 = 17{,}653{,}760$ 个元素 = 35.3 MB，标准的 1/3.9。

练习 4（因果掩码） $N = 4096$，块 128×128。因果掩码下要算的块占几成？

答案：$T = 32$，$\frac{T(T+1)/2}{T^2} = \frac12 + \frac{1}{64} = 0.516$（代码 03 演示 4）。

练习 5（为反向存什么） 32 个注意力矩阵、$N = 32768$、FP16：标准写法为反向存的 $P$ 多大？FlashAttention 存的 $L$（FP32）多大？

答案：$32 \times 32768^2 \times 2 = 68.7$ GB；$32 \times 32768 \times 4 = 4.2$ MB。

练习 6（$D_i$ 的恒等式） 证明 $\sum_j dP_{ij} P_{ij} = dO_i \cdot O_i$。

答案：$dP = dO\, V^\top$，所以 $dP_{ij} = dO_i \cdot v_j$。$\sum_j P_{ij}(dO_i \cdot v_j) = dO_i \cdot \sum_j P_{ij} v_j = dO_i \cdot O_i$。

练习 7（非矩阵乘法为什么贵） A100 的 FP16 矩阵乘法峰值 312 TFLOPS，FP32 非矩阵乘法 19.5 TFLOPS。如果注意力里非矩阵乘法的运算占矩阵乘法的 1%，它的时间大约占矩阵乘法时间的几成？

答案：$1\% \times 312 / 19.5 = 16\%$（代码 01 演示 3 里推迟归一化的那一行）。

练习 8（decode 时的补齐） 批 1、64K 的序列、16 个 query 头、$d = 128$。FlashAttention-2 每个 CTA 按 128 行的 $Q$ 块算两次矩阵乘法，运算量多少？A100 每 SM 的 Tensor Core 峰值约 2.89 TFLOPS，按 70% 的效率要多久？

答案：$4 \times 128 \times 65536 \times 128 = 4.29$ GFLOP，$4.29\text{e}9 / (2.89\text{e}12 \times 0.7) = 2.12$ ms。16 个 CTA 并行在 16 个 SM 上，所以总时间约 2.1 ms——博客实测 2.3 ms（代码 04 的模型：2.13 ms）。真正有用的只有 1/128。

练习 9（选段数） 批 1、64K、2 个 KV 头，A100 有 108 个 SM。为什么在我们的模型里切 32 段（64 个 CTA）比切 64 段（128 个 CTA）快？

答案：64 个 CTA 时每个 CTA 可以拉 $1555 / 64 = 24$ GB/s（低于单 SM 的上限 80），64 个 SM 已经能用满显存带宽，一波完成；128 个 CTA 超过 108 个 SM，要两波，第二波只有 20 个 CTA，却要付一整波的时间。

练习 10（decode 注意力的下限） Llama-3.1-8B（8 个 KV 头、$d = 128$、32 层），批 1、128K 的上下文，BF16。读一遍 KV cache 要多少字节、H100 上至少多久？

答案：每层 $2 \times 131072 \times 8 \times 128 \times 2 = 537$ MB，32 层 17.2 GB，按 3.35 TB/s 至少 5.1 ms——比读一遍 16 GB 的权重（4.8 ms）还长。

练习 11（megakernel 的理想时间） Llama-3.2-1B 有 2.48 GB 的 BF16 权重。H100 上批 1 的 decode 最快每秒几次？如果多内核的写法每层 7 个内核、每个内核多花 4 µs，带宽利用率是多少？

答案：$3.35\text{e}12 / 2.48\text{e}9 = 1350$ 次（0.74 ms）。$7 \times 16 + 2 = 114$ 个内核 × 4 µs = 456 µs，利用率 $0.74 / (0.74 + 0.456) = 62\%$（还没算带宽效率的损失）。

练习 12（尾部） 192 个块在 132 个 SM 上，每个块的工作量相同。SM 的平均利用率是多少？把每块的行数从 16 改成 24（块数变成 128）呢？

答案：两波，$192 / (2 \times 132) = 73\%$；128 个块一波，$128 / 132 = 97\%$。块变大反而更快——只要块数不超过 SM 数太多。

13. 延伸阅读 · Further reading

阅读材料

• T. Dao, D. Y. Fu, S. Ermon, A. Rudra, C. Ré, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022. https://arxiv.org/abs/2205.14135
• B. Spector et al.（Hazy Research），Look Ma, No Bubbles! Designing a Low-Latency Megakernel for Llama-1B, 2025. https://hazyresearch.stanford.edu/blog/2025-05-27-no-bubbles

更多

• T. Dao, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, 2023. https://arxiv.org/abs/2307.08691
• J. Shah, G. Bikshandi, Y. Zhang, V. Thakkar, P. Ramani, T. Dao, FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision, 2024. https://arxiv.org/abs/2407.08608
• T. Dao, D. Haziza, F. Massa, G. Sizov, Flash-Decoding for long-context inference, PyTorch 博客 2023. https://pytorch.org/blog/flash-decoding/
• M. Milakov, N. Gimelshein, Online normalizer calculation for softmax, 2018（在线 softmax 的出处）。
• Z. Ye et al., FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving, 2025（分页 KV、级联注意力、JIT 生成注意力变体）。
• B. Spector et al., ThunderKittens: Simple, Fast, and Adorable AI Kernels, 2024（megakernel 的实现基础）。

本仓库

• 前置讲 P3 §5：安全 softmax、log-sum-exp、在线 softmax 的合并。
• 第 7 讲：分页 KV 与注意力内核的块表。
• 第 10 讲：分块、Tensor Core 的片段、融合的 softmax、启动顺序与波次。
• CS349D 讲义里也有 FlashAttention 和 flash-decoding 的讨论（侧重推理系统）：../../CS349D/README.md。

14. 下一讲预告 · Next lecture

这两讲的内核都是手写的：每一个融合、每一种分块都要专家花几周到几个月。能不能让编译器来做？第 12 讲讲 ML 编译器的前端：怎样从一段 PyTorch / JAX 程序里抓出计算图（tracing vs 分析字节码，TorchDynamo 的 guard 和 graph break），图上的中间表示长什么样，图级优化（常量折叠、公共子表达式、代数化简、算子融合、内存规划）怎样自动完成本讲 §7 那张表里的大部分融合，以及动态形状和 CUDA Graphs 带来的问题。PyTorch 2 的论文会告诉我们：Inductor 的大部分加速正来自融合——关掉融合和内联，编译后的代码反而比不编译还慢。

→ 第 12 讲：ML 编译器（前端）