Lecture 8 · GPU 性能的第一性原理与 FlashAttention · GPU Performance & FlashAttention

对应 Maps to：Stanford CS324 Winter 2023 Advances in Foundation Models，1/11 客座报告 Tri Dao：FlashAttention（第 1 周，整门课的第二次课）。该次课日历上列出的阅读：Jay Alammar The Illustrated Transformer、Vaswani et al. Attention Is All You Need（2017）、Horace He Making Deep Learning Go Brrrr From First Principles（2022）、NVIDIA Deep Learning Performance Guide，以及报告本身的论文 Dao et al. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness（NeurIPS 2022）。Syllabus 的 “Training Methods and Infrastructure” 主题。 依据说明：课程网站没有提供 slides 或录像，报告的具体内容不可考。本讲的 §2–§4 依据 Horace He 的文章（三种”受限”、融合、开销）；§5–§6 依据 FlashAttention 论文（算法 1、附录 B 的反向算法、定理 2 与命题 3 的 IO 复杂度、§4 的实验）；§7 的 MFU / HFU 依据 PaLM 论文第 4 节与附录 B；§8”2023 年之后”依据 FlashAttention-2、FlashAttention-3 的论文摘要和我对后续发展的概括。硬件规格是厂商公开的数字，A100 的存储层次数字与 FlashAttention 论文 §2.1 一致。文献里的数字都是我在原文里读到的。 配套代码 Code：code/lecture08/（numpy + PyTorch；4 个脚本 + 1 个模块。代码 01 约 1 秒，02 约 15 秒（在你的机器上计时），03 约 2 秒，04 不到 1 秒） 预计阅读时间：3 小时（含运行代码）

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

0.1 学习目标 Learning goals

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

1. 用 Horace He 的框架说清一个深度学习程序的时间花在哪里：算力（compute）、显存带宽（memory bandwidth）、开销（overhead），并知道怎样判断自己处于哪一种。
2. 计算一个运算的算术强度（arithmetic intensity，FLOPs / 字节），用 roofline 模型估计它的最短用时，并解释为什么矩阵乘法通常受算力限制、逐元素运算和 softmax 受带宽限制、生成文本时的矩阵 × 向量受带宽限制。
3. 描述 GPU 的存储层次（片上 SRAM vs 显存 HBM）和融合（kernel fusion）为什么能加速，以及为什么融合在 CPU 上不总是有用。
4. 推导在线 softmax（online softmax）的更新公式，并解释它为什么让注意力可以分块计算。
5. 走一遍 FlashAttention 的前向（分块 + 在线 softmax）和反向（不存 $P$，重算），说明它为什么是精确算法、为什么 FLOPs 更多却更快。
6. 推导标准注意力 $\Theta(Nd + N^2)$ 和 FlashAttention $\Theta(N^2 d^2 / M)$ 的 IO 复杂度，并解释”线性”指的是显存占用而不是显存读写。
7. 复述 FlashAttention 的主要实验结果（BERT、GPT-2 的训练加速，Path-X / Path-256 上的长序列结果）。
8. 定义 MFU 与 HFU，用 PaLM 论文的公式重算 46.2% 和 57.8%，并解释为什么我们的玩具模型只有约 5%。
9. 了解 2023 年之后的发展：FlashAttention-2 / 3、推理时的注意力优化、长上下文，以及”硬件越来越快、数据越来越难搬”这条大趋势。

0.2 先修与本讲范围 Prerequisites & scope

如果你学过本目录下的 CS336，它的第 5 讲（GPU 与 TPU）和第 6 讲（kernel 与 Triton）从实现的角度讲得更深，会真的在 GPU 上写 kernel。本讲自成一体，侧重”为什么”：我们没有 GPU，所以用纸上的 roofline 计算和带计数器的 numpy 算法来理解 GPU，再在你自己的 CPU 上测一测同样的规律。

0.3 本讲地图 Map

§1  为什么一门基础模型的课要讲 GPU：算力账、效率就是钱                         ← 代码 01 演示 4
§2  三种"受限"：算力、带宽、开销（Horace He）；在你自己的机器上测一测          ← 代码 02
§3  roofline 与算术强度：拐点、常见运算的强度、推论                            ← 代码 01 演示 1、2
§4  存储层次与融合：SRAM vs HBM、kernel、融合；一个 Transformer 层的账本      ← 代码 01 演示 3、代码 02 演示 3
§5  FlashAttention：标准注意力为什么慢；在线 softmax；算法 1；反向时重算；
    IO 复杂度；精确 vs 近似；算法 ≠ kernel                                       ← 代码 03、04
§6  效果：更快的训练、更长的上下文
§7  MFU 与 HFU：怎样衡量"用满了多少"                                           ← 代码 01 演示 4、代码 02 演示 4、5
§8  2023 年之后：FlashAttention-2/3、推理、长上下文、硬件趋势
§9–13  小结、常见误区、练习、延伸阅读、下一讲预告

0.4 记号 Notation（只列本讲新增的，其余见第 1 讲 §0.4）

0.5 关于数字 · About the numbers

这一讲的数字有三种来源，可信度不同：

• 纸上计算（代码 01、04）：roofline 模型 + 公开的硬件规格 + 论文里的数字。结果是确定的，但它是下界——真实的 kernel 达不到峰值。我用它来判断”哪个运算是瓶颈""大约快几倍”，而不是预测毫秒数。
• 精确的数值验证（代码 03）：FlashAttention 与标准注意力逐数字对照，误差 1e-16 量级；同一台机器上重复运行逐字相同。
• 在你的机器上计时（代码 02）：每次运行都略有不同，换一台机器差别更大。输出里以 ⏱ 开头的行都是计时；讲义引用的是作者机器（Apple M4 Pro，PyTorch 单线程）上的一次运行，写作”约”。请以你自己运行的结果为准——规律相同，数字不同。

有几处结果和我写代码前的预想不一样，讲义按实际结果来写（完整清单在代码 README 的”关于实验结果”里），最值得注意的是：融合在 CPU 上对 GELU 几乎没用（§4.3），以及 FlashAttention 读写量的优势是”常数倍”而不是”从二次方到线性”（§5.5）。

1. 为什么一门基础模型的课要讲 GPU · Why systems matter

1.1 算力账：从 FLOPs 到天数

第 4 讲 §3.3 给每个大模型记了一笔账：PaLM 540B 的训练需要 $C \approx 6ND \approx 2.5 \times 10^{24}$ 次浮点运算。这个数字本身不告诉你要训练多久。要换算成时间，还需要三个数：

前两个是买来的：PaLM 用了 6144 块 TPU v4，每块 16 位峰值 275 TFLOP/s。最后一个——利用率——是工程做出来的。PaLM 论文报告的模型 FLOPs 利用率（MFU）是 46.2%；同一张表里，2020 年的 GPT-3 只有 21.3%。代码 01 演示 4 按 PaLM 的实测吞吐（每秒 238.3K 个 token）算出：780B 个 token 不间断地训练要 37.9 天。如果利用率只有 GPT-3 那么高，就要两倍多的时间——或者两倍多的芯片。

所以对基础模型来说，系统效率就是钱，也是”能不能做”。一个想法如果能让利用率从 21% 提高到 46%，它的价值相当于白送一倍的硬件。这一讲的主角 FlashAttention 正是这样的想法：它没有改变模型（结果逐数字相同），只改变了计算的顺序，就让注意力快了几倍，并让上下文长度从几千扩展到几万。

1.2 这次报告的背景

FlashAttention 来自 Stanford 的 Hazy Research（Chris Ré 的实验室，Ré 也是 CS324 的三位教师之一）。第 1 讲 §1.1 提到过，这个实验室把”基础模型本身”当成一个值得优化的系统。论文 2022 年 5 月出现在 arXiv 上；到 2023 年 1 月这次报告时，已经有不少开源训练框架和公司开始采用它，PyTorch 2.0（2023 年 3 月）随后把它集成进了标准的注意力接口。

这次课的阅读清单很有意思：前两篇是 Transformer 本身（图解 + 原论文），后两篇是 GPU 性能的入门（Horace He 的博客、NVIDIA 的性能指南），最后才是 FlashAttention。这个顺序本身就是一个提示：要理解 FlashAttention 为什么有效，先要理解 GPU 的时间花在哪里。本讲照着这个顺序来。

1.3 本讲的路线

时间花在哪里？          → §2 三种"受限"（算力 / 带宽 / 开销）
怎样定量判断？          → §3 roofline：算术强度 vs 拐点
带宽为什么是瓶颈？      → §4 存储层次：SRAM 很快但很小，HBM 很大但很慢；融合
注意力的问题在哪里？    → §5 N×N 矩阵要在 HBM 里来回搬
怎样不搬它？            → §5 分块 + 在线 softmax（前向），重算（反向）
效果怎样？怎样衡量？    → §6 训练加速、长上下文；§7 MFU / HFU

2. 三种”受限”：算力、带宽、开销 · Compute, memory, overhead

2.1 一个比喻：工厂、仓库和调度员

Horace He 用一个工厂的比喻来解释 GPU 上的时间去向：

• 算力（compute） 是工厂本身：机器在加工零件。对 GPU 来说就是真正在做浮点运算的时间。
• 带宽（memory bandwidth） 是运输：把原料从仓库运到工厂、把成品运回仓库。对 GPU 来说，是在显存（HBM）和计算单元之间搬数据的时间。
• 开销（overhead） 是其他一切：调度员决定”下一步加工什么”、填单子、打电话。对 PyTorch 来说，是 Python 解释器、框架选择该调用哪个 kernel、启动 kernel 的时间。

关键的一句话是：如果工厂的加工速度翻倍，而运输速度不变，产量不会翻倍——工厂会停下来等原料。这就是为什么要先弄清楚”现在受什么限制”：花力气优化不是瓶颈的那一项，毫无用处。

2.2 受算力限制：矩阵乘法与 tensor core

现代 GPU 上有专门做矩阵乘法的硬件（NVIDIA 叫 tensor core，TPU 叫 MXU）。在 A100 上，16 位矩阵乘法的峰值是 312 TFLOP/s；不走 tensor core 的普通运算（加法、exp、比较……）只有 19.5 TFLOP/s——差 16 倍（代码 01 演示 1）。

这意味着：几乎所有的计算力都在矩阵乘法上。一个运算如果不是矩阵乘法，就算它的 FLOPs 很少，它也跑在”慢车道”上。Horace He 引用了对 BERT 的一项测量（Ivanov et al. 2021, Data Movement Is All You Need）：非矩阵乘法的运算（归一化、逐元素运算）只占约 0.2% 的 FLOPs，但它们每秒完成的 FLOPs 比矩阵乘法少 250 倍到 700 倍——于是它们占用的时间远远超过 0.2%。

受算力限制时，能做的事情是：用 tensor core（确保运算是矩阵乘法、形状合适、用 16 位或更低的精度），或者——用 Horace He 的玩笑话说——“给 NVIDIA 更多的钱”。

2.3 受带宽限制：逐元素运算

一个逐元素运算，比如 y = x.cos()，对每个数只做一次运算，但要先从显存读入这个数（2 个字节）、算完再写回去（2 个字节）。A100 每秒能搬约 1.5 TB 的数据，却能做 312 万亿次矩阵运算；搬一个数的时间里，tensor core 能做几百次运算。所以这类运算的时间几乎全部花在搬数据上，计算单元大部分时间在等。

受带宽限制时，“少算”没用，要少搬：把几个运算合并成一个 kernel，中间结果不写回显存（§4 的融合），或者用更少的字节表示每个数（16 位、8 位）。

2.4 受开销限制：小运算与 Python

第三种情况在研究代码里很常见：张量很小，GPU 一眨眼就算完了，时间都花在”决定算什么”上。Horace He 给了一组数字：A100 每秒能做 312 万亿次运算，而 Python 每秒只能做约 3200 万次加法——Python 做一次运算的时间，A100 能做近一千万次。

GPU 的执行是异步的：CPU 只是把 kernel 放进队列就继续往下走，只要 CPU 排队的速度比 GPU 执行快，开销就被掩盖了。但如果每个 kernel 都很小，GPU 很快就把队列清空，转而等 CPU。

怎样判断自己是不是开销受限？ Horace He 的办法很简单：把 batch 加倍。如果工作量翻倍、时间只增加 10%，那时间主要不是花在计算或搬运上——是开销。解决办法是减少运算的个数：融合、用 torch.compile 之类的编译器、CUDA Graphs（把一串 kernel 录下来一次性重放）。

2.5 在你自己的机器上测一测

我们没有 GPU，但三种”受限”在 CPU 上同样存在。代码 02 在作者的机器（Apple M4 Pro，PyTorch 单线程）上测了一遍：

矩阵乘法（演示 1）：$n \times n$ 方阵相乘，$n = 8$ 时每次约 0.4 微秒、只有 2.4 GFLOP/s；$n = 1024$ 时约 3 TFLOP/s。小矩阵的时间几乎和 $n$ 无关——那是开销；大矩阵接近这台机器的矩阵峰值。

这个峰值本身就很说明问题：一个 CPU 核心做矩阵乘法能到约 3 TFLOP/s，比普通的逐元素运算快两个数量级——因为 macOS 的数学库（Accelerate）把矩阵乘法交给了芯片上专门的矩阵单元，和 GPU 的 tensor core 是同一个思路。在别的 CPU（例如 x86 服务器的单个核心）上，这个数通常是几十到一两百 GFLOP/s，但”大矩阵比小矩阵快几百倍”的规律相同。

带宽（演示 2）：原地乘一个常数 x.mul_(c)，每个数做 1 次乘法、搬 8 个字节。数组大时速度稳定在每秒约 130–150 GB——这就是这台机器的”显存带宽”。矩阵峰值除以带宽，这台机器的拐点约为 20 多 FLOPs / 字节（fp32；§3 会解释拐点）。

开销（演示 4）：第 1 讲的 base model 训练一步要调用 603 个底层运算（PyTorch 的 aten 算子：前向、反向、AdamW 对每个参数张量的更新……），这个数几乎不随 batch 变化。batch 从 1 加到 2，工作量翻倍，时间只变成 1.41 倍——按 Horace He 的标准，这已经是”开销占了相当一部分”；batch 到 8 以上，时间才基本和 batch 成正比。

宽度（演示 5）：同样 2 层、同样 batch，只把宽度 $d$ 从 32 加到 1024，每秒完成的 FLOPs 从约 66 G 涨到约 1580 G，从矩阵峰值的 2% 涨到一半左右。这就是**“矩阵越大，硬件用得越满”**——也是为什么第 4 讲 §3 说 6ND 只是”需要的计算量”，而花的时间还要除以利用率。

3. roofline 与算术强度 · Roofline and arithmetic intensity

3.1 一个运算至少要多久

把 §2 的直觉变成公式。一个运算要做 $F$ 次浮点运算、从显存读写 $B$ 个字节。硬件的峰值算力是 $P$（FLOP/s），带宽是 $W$（B/s）。计算和搬运可以同时进行，所以这个运算至少需要

哪一项大，这个运算就受谁限制。把它改写一下：令算术强度 $I = F / B$（每搬一个字节做多少次运算），那么每秒能完成的运算次数最多是

在对数坐标里画出来，像一个屋顶（roof）：

 达到的 FLOP/s（对数）
      │                     ________________  ← 峰值算力 P（平的屋顶）
      │                   ／
      │                 ／   受算力限制
      │               ／
      │             ／ ← 斜坡：I × W（受带宽限制）
      │           ／
      │         ／
      └──────────┬─────────────────────── 算术强度 I（对数）
               拐点 I* = P / W

拐点（ridge point）$I^* = P / W$ 是屋顶的拐角：强度低于它，运算的速度由带宽决定（在斜坡上）；高于它，由算力决定（在屋顶上）。这就是 roofline 模型（Williams et al. 2009）。它是一个下界模型——真实的 kernel 还有各种损耗——但用来判断”瓶颈在哪里”非常可靠。

3.2 拐点在升高

代码 01 演示 1 列出了几代加速器（16 位矩阵运算）的拐点：

从 V100 到 H100，矩阵峰值涨了 7.9 倍，显存带宽只涨了 3.7 倍。算力涨得比带宽快，拐点一代比一代高，越来越多的运算落到斜坡上。这个趋势至今没有改变，也是 FlashAttention 这类”少搬数据”的工作越来越重要的原因。

3.3 常见运算的强度

代码 01 演示 2 在 A100 上算了一组典型运算（16 位，每个数 2 字节；拐点约 200）：

几个值得记住的结论：

1. 方阵乘法的强度是 $n/3$：$2n^3$ 次运算，读写 $3n^2$ 个数、每个 2 字节。矩阵越大强度越高，$n$ 到几百才过拐点。
2. 训练和生成是两个世界。 训练时一次喂进几千个 token，矩阵乘法的强度上千；生成时每一步只算一个新 token，矩阵乘法退化成矩阵 × 向量，强度约 1——每个权重读进来只用一次。所以生成一个 token 的时间约等于”把整个模型从显存读一遍的时间”。这解释了推理系统的许多设计：把多个用户的请求攒成一个 batch、把权重量化成 8 位或 4 位（少读字节）、多查询注意力（MQA）缩小要读的 KV cache（第 10 讲）。
3. 逐元素运算、归一化、softmax 的强度只有 0.2–2，永远在斜坡的最底部。
4. 注意力本身也受带宽限制，尽管 $QK^\top$ 是矩阵乘法：每个头的维度 $d$ 只有 64–128，输出却是 $N \times N$ 的大矩阵，强度只有 60–114，低于拐点。再加上后面的缩放、掩码、softmax 都要把这个 $N \times N$ 矩阵读写一遍——标准注意力是一个受带宽限制的运算。这就是 FlashAttention 下手的地方。

3.4 精度也改变强度

强度的分母是字节，不是”个数”。同样的运算从 32 位换成 16 位，搬的字节减半，强度翻倍。这是低精度训练（第 10 讲）的好处之一：不只是 tensor core 在 16 位下更快，受带宽限制的运算也会快一倍。推理时把权重量化成 4 位，对受带宽限制的生成步骤来说，几乎就是 4 倍的加速。

4. 存储层次与融合 · Memory hierarchy and kernel fusion

4.1 GPU 的存储层次

“从显存读数据”具体是什么意思？一块 GPU 上有好几级存储，越靠近计算单元越快、越小：

                     容量               带宽                    位置
  寄存器            每个 SM 几百 KB     最快                   计算单元内部
  SRAM（shared      每个 SM 192 KB      约 19 TB/s（合计）      芯片上，每个 SM 一份
   memory / L1）    合计约 20 MB
  L2 缓存           40 MB               ——                     芯片上，所有 SM 共享
  HBM（显存）       40–80 GB            1.5–2.0 TB/s           芯片旁边的独立存储
  CPU 内存          TB 级               比 HBM 慢得多           通过 PCIe 连接

（A100 的数字；SRAM 和 HBM 的数字来自 FlashAttention 论文 §2.1，代码 01 演示 1 算出 108 个 SM × 192 KB = 20.25 MB，是 40 GB 显存的约 1/2000。）

SRAM 比 HBM 快一个数量级，但小三个数量级。 计算单元只能直接处理 SRAM 和寄存器里的数据；HBM 里的数据要先搬进来。CPU 上也是同一个结构：L1/L2 缓存对应 SRAM，内存（DRAM）对应 HBM。

4.2 一个 kernel 做了什么

在 PyTorch 里写 y = x.cos()，GPU 上发生的是：启动一个 kernel，每个 SM 从 HBM 读入一部分 x 到 SRAM / 寄存器，算 cos，把结果写回 HBM 里的 y。下一个运算 z = y.cos() 又是一个新的 kernel：从 HBM 读 y，算，写回 z。

每个 kernel 都从 HBM 读输入、往 HBM 写输出——这是 kernel 之间传递数据的唯一方式。于是 x.cos().cos() 要读写 HBM 4 次（读 x、写 y、读 y、写 z），而真正必要的只有 2 次（读 x、写 z）。

4.3 融合：中间结果不落地

融合（kernel fusion）就是把几个运算写进同一个 kernel：读进 x，在 SRAM / 寄存器里连做两次 cos，只把最后结果写回去。Horace He 的例子：融合后的 x.cos().cos() 读写 HBM 2 次，速度几乎和单独一个 x.cos() 一样。这也解释了一个常见的经验：在 GPU 上，GELU、SiLU 这些激活函数的”计算量”不同，用时却差不多——它们都受带宽限制，时间由读写决定。

融合不改变数学，只改变读写。框架里的许多”融合 kernel”就是这样来的：F.softmax（一个 kernel 完成 max、减、exp、求和、除）、F.layer_norm、带偏置的 GELU、BLOOM 训练用的”缩放 + 掩码 + softmax”融合 kernel（第 10 讲）。

代码 02 演示 3 在 CPU 上验证了一半、推翻了一半：

前两行符合预期：a*b+c 受带宽限制，加速比约等于”趟数之比” 6/4 = 1.5。后两行和 GPU 上的规律相反：在这台 CPU 上，GELU 融合几乎没用，x.cos().cos() 正好是 x.cos() 的两倍。原因是 CPU 上 tanh、cos 这类函数本身很贵（要算多项式近似，没有 GPU 那样的专用单元），这些运算在 CPU 上是受算力限制的——少读写几遍省不了多少。

这个”失败”其实是本讲最好的一课：同一个运算在不同硬件上可以落在 roofline 的不同一侧。优化之前先判断受什么限制；把 GPU 上的经验照搬到 CPU（或者反过来）是会出错的。

4.4 一个 Transformer 层的账本

现在把一个 Transformer 层拆开，逐个运算记账。代码 01 演示 3 用 roofline 模型估计了 GPT-2 medium 的一层（$d = 1024$，16 个头，$N = 1024$，batch 8）在 A100 上的前向，假设每个运算是一个独立的 kernel（PyTorch 的朴素写法）：

0.5% 的 FLOPs 占了 52% 的时间——逐元素运算和 softmax 几乎全是搬数据，其中缩放、加掩码、softmax 三个运算各占 15%，因为它们每一个都要把 $B \times H \times N \times N$ 的注意力矩阵读一遍、写一遍。这一层最多只能达到峰值的 34%。

如果把注意力的五步（$QK^\top$、缩放、掩码、softmax、$PV$）换成一个 FlashAttention kernel，再把 GELU、残差相加并进前一个矩阵乘法的”收尾”（epilogue），同一层的理论用时从 2.279 ms 降到 0.837 ms，快 2.72 倍，最多能达到峰值的比例升到 93%。

但同样的分析放到 GPT-3 175B 的一层（$d = 12288$，96 个头，$N = 2048$）上，线性层占了 97.2% 的 FLOPs、83.0% 的时间，融合只快 1.14 倍。模型越宽，注意力和逐元素运算越不重要；序列越长，它们越重要。 2022 年 FlashAttention 在 GPT-2 规模的模型上测出数倍的端到端加速，在超大模型上的收益要小一些；它真正改变格局的，是长序列（§6.2）。

5. FlashAttention

5.1 标准注意力为什么慢

回到第 3 讲 §3.2 的注意力。对一个头，$Q, K, V \in \mathbb{R}^{N \times d}$：

在 GPU 上的标准实现（FlashAttention 论文的”算法 0”）是三步，每步一个 kernel：

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

$S$ 和 $P$ 都是 $N \times N$。数一数读写 HBM 的元素个数：$2Nd + N^2$（第 1 步）$+ 2N^2$（第 2 步）$+ N^2 + 2Nd$（第 3 步）$= 4N^2 + 4Nd$。反向传播还要读前向存下的 $P$、再写几个 $N \times N$ 的梯度矩阵，又是 $7N^2 + 8Nd$（flash_tools.standard_io）。当 $N$ 比 $d$ 大得多（$N = 2048$、$d = 64$），读写几乎全是 $N^2$ 那几项。

这里有两个问题：

• 慢：§3.3 已经算过，这些步骤的强度都低于拐点，时间花在搬 $N \times N$ 矩阵上。
• 占显存：反向传播需要 $P$，所以前向结束后它必须一直留在显存里——每一层、每一个头、每一个序列都有一份 $N \times N$。$N$ 翻倍，这部分显存翻四倍。

为什么不直接分块？ 矩阵乘法很容易分块（第 9 讲的张量并行就是把它切开），但 softmax 不行：$P_{ij} = e^{S_{ij}} / \sum_k e^{S_{ik}}$，分母要用到一整行。而且实际实现必须用”安全 softmax”——先减去这一行的最大值再取 exp，否则会溢出（代码 03 演示 1：float64 下 $e^{1000}$ 溢出成 inf，结果全是 nan；16 位浮点数里 $e^{12}$ 就已经溢出，fp16 最大只能表示 65504）。“先找最大值”意味着必须先读完一整行，于是标准实现只好把 $S$ 整个算出来、存下来。

2022 年以前的主流思路是近似。 稀疏注意力（只算一部分位置）、低秩近似（Linformer）、核方法（Performer）、局部窗口（Longformer）……它们把计算量从 $O(N^2)$ 降到接近 $O(N)$，但改变了模型，质量常常下降；而且 FlashAttention 论文指出，其中很多方法在墙上时间上并不比标准注意力快——它们减少了 FLOPs，却没有减少读写，而瓶颈在读写。FlashAttention 的出发点是论文标题里的那个词：IO-aware——按”读写了多少”而不是”算了多少”来设计算法。

5.2 关键想法一：在线 softmax

softmax 真的需要先读完整行吗？不需要。Milakov & Gimelshein（2018）提出的在线 softmax 只需读一遍，边读边维护两个数：目前见过的最大值 $m$，和”相对于 $m$ 的”分母 $\ell = \sum_j e^{x_j - m}$。

读进新的一块 $x_{\text{blk}}$ 时：

第一项是关键：以前的 $\ell$ 是相对旧的最大值 $m$ 算的，现在基准变成了 $m_{\text{new}}$，每一项都要乘以 $e^{m - m_{\text{new}}}$（一个 ≤ 1 的数）“换算”过来。读完之后 $\mathrm{softmax}(x)_j = e^{x_j - m} / \ell$。

代码 03 演示 2 一步一步打印了这个过程（每次读 3 个数）：

x = [2.0, -1.0, 3.0, 0.5, 6.0, 1.0, -2.0, 5.5]
读到第几个数   新读入的块              m      ℓ
3             [2.0, -1.0, 3.0]        3.0    1.386195
6             [0.5, 6.0, 1.0]         6.0    1.079839   ← 出现了 6：旧 ℓ × e^(3−6)，再加新块
8             [-2.0, 5.5]             6.0    1.686705   ← 最大值没变，只累加

最后的 $\ell = 1.686705$ 正好等于 $\sum_j e^{x_j - 6}$，结果与”先读完整行”的安全 softmax 完全相同（差 0.0）。任何时候只需要两个数——这就是让注意力可以一块一块地算的那把钥匙。

5.3 算法 1：分块计算注意力

把在线 softmax 用到注意力上。注意力输出的第 $i$ 行是 $O_i = \sum_j P_{ij} V_j$，也是一个”按 softmax 权重加权”的求和——它可以和分母一起在线地更新。FlashAttention 论文的算法 1：

把 Q 切成 Tr 块（每块 Br 行），K、V 切成 Tc 块（每块 Bc 行）
在 HBM 里初始化 O = 0，ℓ = 0，m = −∞（每行一个）
for j = 1 … Tc:                         ← 外层：K、V 的块
    从 HBM 读 K_j、V_j 进 SRAM
    for i = 1 … Tr:                     ← 内层：Q 的块
        从 HBM 读 Q_i、O_i、ℓ_i、m_i 进 SRAM
        S_ij = τ Q_i K_jᵀ                 （Br × Bc，只在 SRAM 里）
        m̃ = rowmax(S_ij)，P̃ = exp(S_ij − m̃)，ℓ̃ = rowsum(P̃)
        m_new = max(m_i, m̃)
        ℓ_new = e^(m_i − m_new) ℓ_i + e^(m̃ − m_new) ℓ̃
        O_i ← diag(ℓ_new)⁻¹ ( diag(ℓ_i) e^(m_i − m_new) O_i + e^(m̃ − m_new) P̃ V_j )
        把 O_i、ℓ_i ← ℓ_new、m_i ← m_new 写回 HBM

$O_i$ 的更新式看起来复杂，拆开就是在线 softmax 的同一个逻辑：$O_i$ 一直保存”已经归一化”的结果；先乘回旧的分母 $\ell_i$ 变回”未归一化的和”，乘以 $e^{m_i - m_{\text{new}}}$ 换算到新基准，加上新块的贡献 $\tilde{P} V_j$，再除以新的分母。

$N \times N$ 的 $S$ 和 $P$ 从来没有完整地存在过——每次只有 $B_r \times B_c$ 的一小块，在 SRAM 里用完就扔。块的大小由 SRAM 的容量 $M$ 决定：论文取 $B_c = \lceil M / 4d \rceil$、$B_r = \min(\lceil M / 4d \rceil, d)$，让 $K_j$、$V_j$、$Q_i$、$O_i$ 和 $S_{ij}$ 能同时放进 SRAM。

代码 03 演示 3 用 numpy（float64）实现了这个算法（flash_tools.flash_attention_forward），在 $N = 256$、$d = 32$ 上和标准注意力对照：4 种块大小、有无因果掩码，最大差都在 1e-16 到 1e-15 量级——float64 的舍入误差。片上同时存在的数最多 2,304 到 32,768 个（取决于块大小），而一个 $N \times N$ 矩阵有 65,536 个；片上需要的存储只和块大小有关，和 $N$ 无关。

因果掩码（语言模型）还有一个额外的好处：如果一整块 $K_j$ 都在 $Q_i$ 的”未来”，这一块完全被遮住，可以直接跳过——不读、不算。

5.4 关键想法二：反向时重算

前向不存 $P$，反向传播怎么办？标准反向需要 $P$：

其中 $D_i = \sum_j P_{ij}\, dP_{ij}$。一个关键的恒等式是 $D_i = \sum_k dO_{ik} O_{ik}$——它只需要 $N \times d$ 的数，不需要 $P$。

FlashAttention 的做法是：前向只存输出 $O$ 和每一行的两个统计量 $m_i$、$\ell_i$（以及 dropout 用的随机数种子），反向时一块一块地重算 $S_{ij} = \tau Q_i K_j^\top$ 和 $P_{ij} = e^{S_{ij} - m_i} / \ell_i$，算完这一块的梯度贡献就扔掉。代价是多做一次 $QK^\top$ 的矩阵乘法。

这是一个典型的用计算换读写的交易。代码 03 演示 4 验证了重算的反向（flash_attention_backward，论文附录 B 的算法 4）和标准反向、和 PyTorch 自动求导三者给出同样的梯度（差在 1e-15 量级），然后算了这笔交易：一个头（$N = 1024$、$d = 64$）前向 + 反向的矩阵乘法 FLOPs，标准是 $12N^2 d$，FlashAttention 是 $14N^2 d$，多 17%。论文图 2 在 A100 上实测的是：

FLOPs 多了 13%，读写少了 9 倍，用时少了 5.7 倍。在 GPU 上，多算比多搬便宜——因为注意力在 roofline 的斜坡上，算力是闲着的。

（这个想法你可能在别处见过：激活重算，或者叫梯度检查点（gradient checkpointing），把整层的激活扔掉、反向时重算前向，用来省显存。第 9 讲会讲它。FlashAttention 的重算更精细：只重算注意力里那个 $N \times N$ 的部分，而且重算发生在 SRAM 里，所以不仅省显存，还省时间。）

5.5 IO 复杂度：到底少读写多少

论文的定理 2 给出了两种算法读写 HBM 的次数（$d \le M \le Nd$）：

FlashAttention 的式子从哪里来？外层循环每处理一块 $K_j, V_j$，内层都要把所有的 $Q$、$O$、$\ell$、$m$ 读一遍、写一遍——约 $3Nd$ 个数。外层一共 $T_c = N / B_c \approx 4Nd / M$ 块。乘起来，前向约 $12 N^2 d^2 / M$；反向每块要读写约 $5Nd$ 个数，前向 + 反向合计约 $32 N^2 d^2 / M$。

代码 04 演示 1 先确认：带计数器的 numpy 算法实际读写的次数，和这个公式（flash_tools.flash_io）逐个相等，包括 $N$ 除不尽块大小的情况。然后用公式算到几万的序列长度（演示 2，$d = 64$，$M = 98{,}304$——A100 一个 SM 的 192 KB 能放下的 16 位数个数）：

这张表纠正了一个流传很广的说法：“FlashAttention 让注意力从二次方变成线性。“读写量并没有变成线性——两种算法都随 $N^2$ 增长，差的是常数：标准算法是 $11N^2$，FlashAttention 是 $32 d^2 / M \cdot N^2 \approx 1.33 N^2$，比值趋近 8 倍左右（$N$ 小时只有 4 倍，因为 $O(Nd)$ 的项还占一部分）。变成线性的是显存占用（§6.2）。计算量也仍然是 $O(N^2 d)$——FlashAttention 是精确算法，不可能少算。

这个”常数”取决于 $d^2 / M$。演示 3 固定 $N = 4096$，改变片上存储 $M$：$M$ 每翻一倍，块大小翻一倍，$K/V$ 少扫一半遍数，读写近似减半。但 $M$ 太小时（16 KB 以下），FlashAttention 反而比标准注意力读写更多——$K/V$ 要被扫上百遍。分块的好处完全取决于片上存储够不够大；A100 每个 SM 192 KB，读写量比标准注意力少 5–10 倍。反过来，头维度 $d$ 越大，优势越小（练习 4：$d = 128$ 时只剩约 2 倍）。

论文还证明了一个下界（命题 3）：没有任何精确注意力算法能在所有 $M \in [d, Nd]$ 上渐近地比 $\Theta(N^2 d^2 / M)$ 更少地读写 HBM。在”精确”这个约束下，FlashAttention 是最优的。

5.6 用 roofline 对照论文

代码 04 演示 4 把读写量和 FLOPs 放进 roofline，估计 GPT-2 medium 的注意力（$N = 1024$、$d = 64$、16 个头）在 A100 上的用时。标准注意力有两种记账方式：论文算法 0 的”最简”写法，和更接近 PyTorch 朴素实现的写法（缩放、掩码、dropout 各是一个单独的 kernel，各多读写一遍 $N \times N$）。比值（标准 ÷ FlashAttention）：

论文没有给出这张图的 batch 和实现细节，所以只比较比值：读写量、用时的比值都落在我们两种记账之间，FLOPs 的比值也接近。roofline 模型给出的结论和实测一致：标准注意力受带宽限制，FlashAttention 受算力限制；省下的是搬数据的时间。

5.7 精确 vs 近似

“精确”在实践中意味着什么？代码 03 演示 5 把 PyTorch 版的分块因果注意力（flash_attention_torch，FlashAttention-2 的写法，§8.1）装进第 1 讲的 base model，只换注意力的核心计算，其余参数一个不改：

分块精确注意力的差在 1e-9 量级（float32 下求和顺序不同带来的舍入误差），同一段输入的 logits 差 9.3e-06：装上就能用，不用重新训练。这是 FlashAttention 被迅速采用的一个关键原因——它是一个可以直接替换的 kernel，不是一个新模型。

滑动窗口是一种真正的近似，代表了 2020–2021 年”高效 Transformer”的一大类方法。在这个迷你世界里，一行只有几十个字符，窗口 32 时几乎无损；窗口缩到 8，bit/字符多了 0.19；窗口 4 时变成原来的 4.7 倍。近似方法省了 FLOPs，代价是改变了模型。FlashAttention 选择了另一条路：FLOPs 一个不少（甚至多了 13%–17%），只减少读写，质量完全不变。

论文也给了近似方法一个位置：块稀疏 FlashAttention——只算一部分 $B_r \times B_c$ 的块，读写量和计算量都按非零块的比例下降（命题 4）。因果掩码就是最简单的例子：代码 04 演示 6 中，16 × 16 = 256 个块里只有 136 个需要算（比例 0.531），前向读写量降到全部块的 0.558 倍，结果与标准因果注意力的差 5.6e-16。更一般的稀疏模式（局部窗口 + 少数全局块）就是近似注意力了；论文报告，块稀疏 FlashAttention 在所有序列长度上都比当时已知的精确、稀疏和近似注意力实现更快——近似方法和 IO-aware 的 kernel 并不矛盾，可以叠加。

5.8 算法 ≠ kernel

最后一个容易误解的地方：如果你在 CPU 上计时代码 03 的 numpy 或 PyTorch 版分块注意力，它会比标准注意力慢得多——Python 的双重循环、每一块都是一个小矩阵乘法，正好落在 §2.4 的”开销受限”里。

FlashAttention 的速度不来自”分块”这个算法本身，而来自把整个算法写成一个融合的 CUDA kernel：每个线程块负责一部分 $Q$，把 $K$、$V$ 的块依次搬进 SRAM，在寄存器和 SRAM 里完成矩阵乘法、在线 softmax 和输出的更新，最后只把 $O$、$m$、$\ell$ 写回 HBM。所有中间结果都不离开芯片。我们的 numpy 版本用”显式的 load / store 计数”模拟了这件事，用来理解它为什么快；真正让它快起来，需要直接控制 GPU 存储层次的底层编程（CUDA，或者后来的 Triton、CUTLASS 等工具）。这也是为什么 FlashAttention 的每一代新版本都和新一代硬件绑在一起（§8）。

6. 效果：更快的训练，更长的上下文 · Results

6.1 训练加速

FlashAttention 论文报告了三组端到端的训练结果（都在 A100 上）：

• BERT-large（序列长度 512），训练到 MLPerf 1.1 的目标（掩码语言模型准确率 72.0%）：NVIDIA 在 MLPerf 1.1 的纪录是 20.0 ± 1.5 分钟，FlashAttention 17.4 ± 1.4 分钟——快 15%。BERT-large 的序列不长，注意力占比不大，所以端到端的提升有限，但这是在一个已经被高度优化过的基准上。

• GPT-2（序列长度 1024，8 × A100，OpenWebText）：

  模型	Hugging Face	Megatron-LM	FlashAttention
  GPT-2 small	9.5 天（1.0×）	4.7 天（2.0×）	2.7 天（3.5×）
  GPT-2 medium	21.0 天（1.0×）	11.5 天（1.8×）	6.9 天（3.0×）

  只看注意力这一个运算，论文图 1 报告了相对 PyTorch 实现 7.6 倍的加速。

• Long-Range Arena（长序列基准，序列长度 1K–4K）：比标准注意力快 2.4 倍，同时比许多近似注意力方法更快。

6.2 显存与长上下文

FlashAttention 更深远的影响在显存上。标准注意力为了反向传播，要把 $P$ 存下来；代码 04 演示 5 算了它的大小（batch 1、16 个头、16 位，每一层）：

$N = 65{,}536$ 时，一层的 $P$ 就超过了一张 80 GB 的卡。一个 24 层的模型还要再乘以 24。2022 年以前”上下文最多 2K 个 token”，很大程度上就是这个限制。FlashAttention 只存每行两个数，显存占用随 $N$ 线性增长——论文说它的显存效率比精确注意力的基线高到 20 倍。

这直接带来了新的能力：在 Long-Range Arena 的 Path-X（序列长度 16K）和 Path-256（64K）上，以前所有的 Transformer 都只能达到随机猜测的水平（要么放不下，要么学不会）；用 FlashAttention 的 Transformer 第一次超过了随机猜测，分别达到 61.4% 和 63.1%。论文还报告，把 GPT-2 small 的上下文从 1K 加长到 4K，用 FlashAttention 训练只要 3.6 天，比 Megatron-LM 训练 1K 的版本（4.7 天）还快 30%，困惑度从 18.2 降到 17.5；长文档分类上，MIMIC-III 用 16K 的长度比 512 高 4.3 个点。

6.3 为什么被这么快地采用

一个系统工作能在一年内成为默认选择，需要同时满足几个条件，FlashAttention 恰好都满足：

1. 精确：不改变模型，已有的训练配方、超参数、检查点都不用动（§5.7）。
2. 即插即用：只替换注意力的核心计算，接口和原来一样。PyTorch 2.0（2023 年 3 月）把它集成进了 F.scaled_dot_product_attention。
3. 又快又省：速度和显存同时改善，没有需要权衡的地方。
4. 开源：代码公开，任何人都能用。

7. MFU 与 HFU：怎样衡量”用满了多少” · Measuring utilization

7.1 定义

“利用率”听起来简单，其实有好几种算法。PaLM 论文第 4.1 节提出了一个现在被广泛使用的定义：

关键在于”模型 FLOPs”：它只数模型本身必须做的运算——前向 + 反向，按论文附录 B 的公式 $6N_p + 12 L H Q T$（$N_p$ 参数量，$L$ 层数，$H$ 头数，$Q$ 头维度，$T$ 序列长度；第二项是注意力的那部分，第 4 讲 §3.2 讲过 $6N$ 漏掉了它）。

为什么要强调”必须做的”？因为实际训练时硬件常常做额外的运算：为了省显存，把激活扔掉、反向时重算前向（第 9 讲的激活重算）。这些重算的 FLOPs 也让硬件忙着，但对”训练这个模型”没有贡献。把它们也算进去，就是 HFU（hardware FLOPs utilization）。同一次训练，HFU 总是 ≥ MFU。MFU 的好处是和实现无关：不管你用了什么技巧，只要每秒处理的 token 多，MFU 就高——它衡量的是”离理想有多远”。

7.2 重算 PaLM 的数字

代码 01 演示 4 用论文给出的数字重算了一遍。PaLM 540B：$N_p = 540.35\text{B}$，$L = 118$，$H = 48$，$Q = 256$，$T = 2048$；6144 块 TPU v4，每块峰值 275 TFLOP/s；batch 2048 时实测吞吐每秒 238.3K 个 token。

• 每个 token：$6N_p = 3.242 \times 10^{12}$，注意力 $12LHQT = 3.563 \times 10^{10}$，合计 $3.278 \times 10^{12}$（论文表 22：3.28 TFLOPs / token）。
• MFU（不算注意力）$= 238.3\text{K} \times 3.242 \times 10^{12} / (6144 \times 275 \times 10^{12}) =$ 45.7%；算上注意力 46.2%。
• 论文表 22 给出算上重算后每个 token 4.10 TFLOPs，HFU $=$ 57.8%。

三个数都和论文一致。4.10 / 3.28 = 1.25，说明 PaLM 只重算了一部分前向（全部重算的话是 8/6 ≈ 1.33，练习 9）。

几个大模型的训练效率放在一起（代码 01 演示 4；各论文的记账方式不完全相同，例如有的报告”每块 GPU 实测的 TFLOP/s”，可能包含重算）：

两年之内，大模型训练的利用率从约 1/5 提高到约一半——一部分来自更好的硬件互联和并行策略（第 9 讲），一部分来自 FlashAttention 这样的 kernel 优化。

7.3 利用率之外：有效训练时间

MFU 衡量的是”训练在跑的时候，用满了多少”。还有一层是”训练有多少时间在跑”。代码 01 演示 4 做了一个粗略的对比：按实测吞吐不间断地训练 780B 个 token，6144 块芯片需要 37.9 天；而 PaLM 论文附录 B 报告实际用了”6144 块芯片 1200 小时 + 3072 块芯片 336 小时（含停机和重复的步数）“，折合 6144 块芯片 57.0 天——是理想情况的 1.50 倍。多出来的部分包括：训练前期 batch 较小、吞吐较低（PaLM 的 batch 从 512 逐步加到 2048）；硬件故障后的重启；损失尖峰后回滚、重跑一部分步数。这些是第 10 讲的主题。

7.4 为什么我们的玩具模型只有约 5%

代码 02 演示 4、5 可以算一个”玩具版的 MFU”：每步的 FLOPs（PyTorch 的 FLOP 计数器数出来）÷ 用时 ÷ 这台机器的矩阵峰值。第 1 讲的 base model 在 batch ≥ 16 时稳定在约 5%；把宽度加到 1024，升到约一半。

原因正是本讲讲的三件事：base model 的矩阵只有 64 宽（强度低，§3.3）；LayerNorm、GELU、softmax 等逐元素运算占了大量时间（§4.4）；每步 600 多个运算的开销（§2.4）。小模型天生用不满硬件——这也是为什么”在小模型上做实验”和”在大模型上做实验”的成本比例，远小于它们参数量的比例。

8. 2023 年之后 · Since 2023

FlashAttention 发表后的三四年里，“IO-aware”这个思路从注意力扩展到了整个大模型系统。这里只列和本讲直接相关、比较确定的几条。

8.1 FlashAttention-2 与 FlashAttention-3

• FlashAttention-2（Dao，2023）：算法基本不变，改的是”怎样把工作分给 GPU 的计算单元”。三处改动：减少非矩阵乘法的运算（例如输出不在每一步都除以 $\ell$，而是最后统一除——代码 03 的 flash_attention_torch 就是这样写的）；即使只有一个头，也把计算分到多个线程块上并行，提高占用率；在一个线程块内部更好地给各个 warp 分工，减少通过共享内存的通信。结果比 FlashAttention 快约 2 倍，在 A100 上达到理论峰值的 50%–73%；端到端训练 GPT 类模型达到每块 A100 225 TFLOP/s（MFU 72%）。
• FlashAttention-3（Shah et al.，2024）：针对 H100。FlashAttention-2 在 H100 上只能用到约 35% 的峰值——新硬件的异步特性（Tensor Memory Accelerator、warp 专门化）需要重写 kernel。三个技巧：让数据搬运和计算重叠（warp 专门化）、把块内的矩阵乘法和 softmax 交错执行、用 FP8 低精度配合块量化。FP16 下快 1.5–2.0 倍，达到 740 TFLOP/s（75% 利用率）；FP8 接近 1.2 PFLOP/s，数值误差比 FP8 的基线注意力低 2.6 倍。

这条线索说明了一个规律：算法思想（分块、在线 softmax、重算）几年不变，但每一代硬件都要重写一遍 kernel。

8.2 推理：带宽受限的世界

§3.3 算过：生成时的矩阵 × 向量强度约 1，完全受带宽限制。2023 年以后，大量系统工作集中在推理上：

• KV cache 的管理：生成时每一步都要读之前所有 token 的 $K$、$V$。vLLM（2023）的 PagedAttention 借鉴操作系统的分页，减少 KV cache 的显存碎片，让同一张卡能同时服务更多请求。
• 缩小 KV cache：多查询注意力（MQA）和分组查询注意力（GQA）让多个查询头共享键 / 值头（第 3 讲 §9 提到，PaLM 在 2022 年已经用了 MQA，第 10 讲）。
• 长上下文的解码：一个查询对很长的 KV cache 做注意力时，标准的 FlashAttention 并行度不够；Flash-Decoding 一类方法把 KV cache 切开并行计算，再用在线 softmax 的合并公式（§5.2）把结果拼起来。
• 量化：把权重压到 8 位、4 位，直接减少每生成一个 token 要读的字节数。

8.3 长上下文

FlashAttention 解决了”$N \times N$ 矩阵放不下”的问题之后，上下文长度在两三年里涨了两个数量级：2023 年出现了 32K（GPT-4 的长版本）、100K（Anthropic 的 Claude）的商用模型，2023 年底到 2024 年出现了 128K 和 100 万 token 的上下文（例如 Gemini 1.5）。当一张卡放不下一个序列的激活时，还需要把序列切到多张卡上（序列并行、ring attention，第 9 讲）。显存问题之外，长上下文还带来了新的研究问题：模型能不能真正”用好”很长的上下文、位置编码怎样外推——这些不是本讲的主题。

8.4 编译器与硬件

• 手写融合 kernel 的门槛被降低了：Triton（2021 年起）让人用 Python 风格写 GPU kernel；PyTorch 2.0 的 torch.compile（2023）自动融合逐元素运算、减少开销；后来还出现了 FlexAttention 一类”写几行 Python 就能生成 FlashAttention 风格 kernel”的工具。
• 硬件继续沿着 §3.2 的趋势走：算力（尤其是低精度算力）涨得比带宽快。H100 之后的 GPU 进一步加强了 FP8、FP4 的支持。拐点越来越高，“少搬数据”越来越重要。

8.5 哪些结论仍然成立

仍然成立：roofline 和”先判断受什么限制”的方法；IO-aware 的设计思想（分块、融合、重算）；在线 softmax 作为把注意力切开的通用工具。被修正或补充的：“注意力是 Transformer 的瓶颈”只在序列长、模型窄时成立，对很宽的模型和 MoE 来说，矩阵乘法和通信（第 9 讲）更重要；“训练效率”之外，推理效率成了同样大的问题；“FlashAttention 是线性的”这种简化说法要小心——线性的是显存，不是读写和计算。

9. 小结 · Summary

1. 效率就是钱。 训练时间 = 需要的 FLOPs ÷（芯片数 × 峰值 × 利用率）。利用率从 21%（GPT-3）提高到 46%（PaLM），等于白送一倍的硬件。
2. 三种受限：算力（矩阵乘法、tensor core）、带宽（逐元素运算、softmax）、开销（小运算、Python）。优化之前先判断自己在哪一种。
3. roofline：$t \ge \max(F/P, B/W)$。算术强度 $I = F/B$ 低于拐点 $P/W$ 时受带宽限制。A100 的拐点约 200；方阵乘法的强度 $n/3$；逐元素运算 0.2–2；生成时的矩阵 × 向量约 1；一个注意力头 60–114。
4. 存储层次：SRAM 快但小（每个 SM 192 KB，约 19 TB/s），HBM 大但慢（40–80 GB，1.5–2.0 TB/s）。每个 kernel 都要读写 HBM；融合让中间结果不落地。在 GPT-2 medium 的一层里，0.5% 的 FLOPs 占了 52% 的时间。
5. 在线 softmax：只记当前最大值 $m$ 和分母 $\ell$，读到新块时换算到新基准——softmax 不必先读完整行。
6. FlashAttention 前向：$K/V$ 分块（外层）、$Q$ 分块（内层），在 SRAM 里算 $B_r \times B_c$ 的小块分数，在线更新输出；$N \times N$ 矩阵从不完整存在。反向：只存 $O$、$m$、$\ell$，重算 $P$；FLOPs 多 13%–17%，读写少约 9 倍，快 5.7 倍。
7. IO 复杂度：标准 $\Theta(Nd + N^2)$，FlashAttention $\Theta(N^2 d^2 / M)$，且这是精确算法的下界。读写仍是 $N^2$ 量级、差一个常数倍（$d = 64$ 时约 8 倍）；变成线性的是显存占用。
8. 精确：装进第 1 讲的 base model，bit/字符差 1e-9；近似方法（滑动窗口）省计算但改变模型。算法 ≠ kernel：速度来自一个融合的 CUDA kernel。
9. 效果：BERT-large 快 15%，GPT-2 快 3–3.5 倍（相对 Hugging Face）；Path-X / Path-256 首次超过随机猜测（61.4% / 63.1%）。
10. MFU / HFU：MFU 只数模型必须做的 FLOPs（$6N_p + 12LHQT$），HFU 还数重算。PaLM 46.2% / 57.8%。我们的玩具模型约 5%：小矩阵、逐元素运算、开销。

10. 常见误区 · Common pitfalls

1. “FLOPs 少就快。” 只在受算力限制时成立。受带宽限制的运算（逐元素、softmax、生成时的矩阵 × 向量），时间由读写决定。FlashAttention 的 FLOPs 比标准注意力多，却快几倍；很多近似注意力的 FLOPs 少得多，墙上时间却不快。
2. “99% 的 FLOPs 是矩阵乘法，所以只需要优化矩阵乘法。” 0.5% 的 FLOPs 可以占一半的时间（§4.4）。
3. “FlashAttention 是一种近似注意力 / 线性注意力。” 它是精确的，结果和标准注意力只差舍入误差；它的计算量和读写量都仍是 $N^2$ 量级。线性的是显存占用。
4. “FlashAttention 的速度来自分块这个算法。” 来自把整个算法写成一个融合的 kernel、中间结果留在 SRAM 里。用 Python 循环实现同样的分块，在 CPU 上会更慢（§5.8）。
5. “融合总是有用。” 只对受带宽限制的运算有用。在作者的 CPU 上，GELU、cos 是受算力限制的，融合几乎不省时间（§4.3）。
6. “GPU 利用率（nvidia-smi 显示 100%）就是 MFU。” nvidia-smi 的利用率只表示”有没有 kernel 在跑”，一个受带宽限制的 kernel 也算 100%。MFU 是”实际完成的有用 FLOPs ÷ 峰值”。
7. “HFU 高说明训练效率高。” 重算越多，HFU 越高，但有用的计算没有变多。比较不同的训练配置要看 MFU。
8. “硬件越新，带宽问题越小。” 恰恰相反：算力涨得比带宽快，拐点一代比一代高（V100 139 → H100 295），越来越多的运算受带宽限制。
9. “CPU 上测出来的规律可以直接搬到 GPU。” 规律（三种受限、roofline）相同，具体哪个运算在拐点哪一侧可能不同（§4.3）。

11. 练习 · Exercises

练习 1（算术强度） 一个注意力头的 $QK^\top$：$(4096 \times 64) \cdot (64 \times 4096)$，16 位。(a) 它的 FLOPs、读写字节数、算术强度各是多少？在 A100 上受什么限制？(b) 对比方阵乘法 $n = 4096$ 的强度 1365，为什么差这么多？

(a) FLOPs $= 2 \times 4096 \times 64 \times 4096 = 2{,}147{,}483{,}648$；读写 $2 \times (4096 \times 64 + 64 \times 4096 + 4096 \times 4096) = 34{,}603{,}008$ 字节；强度约 62.1，低于 A100 的拐点 201，受带宽限制。 (b) 方阵乘法的输入和输出一样大（都是 $n^2$），每个数参与 $n$ 次乘加；这里两个输入都很”瘦”（$4096 \times 64$），输出却是 $4096 \times 4096$ 的大矩阵——写这个输出就占了绝大部分字节，而每个输出元素只是一个长度 64 的点积。强度大约就是 $d$ 量级（$N d / (N + 2d)$），而 $d$ 只有 64。

练习 2（生成时攒 batch） §3.3 中 GPT-3 的 MLP 第一层（$12288 \to 49152$），一次处理 $T$ 个 token 时强度约等于 $T$（当 $T \ll d$）。在 A100（拐点约 200）上，大约要同时处理多少个 token，这个矩阵乘法才变成受算力限制？这对推理服务意味着什么？

用代码 01 里 matmul_cost 的公式算：$T = 128$ 时强度 126.4，$T = 200$ 时 196.0，$T = 224$ 时 219.0。所以大约要 200 多个 token 同时算，才能用满 tensor core。生成时每个请求每步只产生 1 个 token，所以推理服务要把两百个左右的请求攒成一个 batch，才能让这一层不受带宽限制。这就是推理系统里”连续批处理”（continuous batching）的动机；请求少的时候，生成速度由”读一遍权重”的时间决定，量化（每个权重少几个字节）直接有效。

练习 3（合并两个块） 在线 softmax 也可以合并两段分别处理过的结果。第一段 $x_1 = [1, 3]$，第二段 $x_2 = [5, 2]$。(a) 分别算出 $(m_1, \ell_1)$、$(m_2, \ell_2)$；(b) 写出合并公式并算出 $(m, \ell)$；(c) 和直接对 $[1, 3, 5, 2]$ 算的结果对比。这个”合并”在 §8.2 的哪个方法里用到了？

(a) $m_1 = 3$，$\ell_1 = e^{-2} + 1 = 1.135335$；$m_2 = 5$，$\ell_2 = 1 + e^{-3} = 1.049787$。 (b) $m = \max(m_1, m_2) = 5$，$\ell = \ell_1 e^{m_1 - m} + \ell_2 e^{m_2 - m} = 1.135335 \times e^{-2} + 1.049787 = 1.203438$。 (c) 直接算 $\sum e^{x_j - 5} = e^{-4} + e^{-2} + 1 + e^{-3} = 1.203438$，相同。合并公式是对称的，两段可以由两个不同的计算单元并行处理、最后再合并——Flash-Decoding 把很长的 KV cache 切成几段并行做注意力，最后就是用这个公式（加上对输出的同样换算）把结果拼起来。

练习 4（头维度的影响） 用 flash_tools.flash_io 和 standard_io，$N = 8192$、$M = 98{,}304$，分别算 $d = 64$ 和 $d = 128$ 时”标准 ÷ FlashAttention”的读写量比值。为什么 $d$ 越大，FlashAttention 的优势越小？

$d = 64$：比值 7.63。$d = 128$：块大小变成 $B_r = 128$、$B_c = 192$，$K/V$ 要扫 $T_c = 43$ 遍，标准 750,780,416 vs FlashAttention 371,228,672，比值只有 2.02。原因在 $\Theta(N^2 d^2 / M)$ 里的 $d^2$：$d$ 翻倍，每块 $K/V$ 占的 SRAM 翻倍，块的行数减半，$K/V$ 被扫的遍数翻倍，而每扫一遍要读写的 $Q$、$O$ 也翻倍——读写量变成 4 倍；标准算法的主项 $11N^2$ 和 $d$ 无关。所以 SRAM 越大、头维度越小，FlashAttention 越划算。（真实实现的块大小还受寄存器等限制，不完全按这个公式选；FlashAttention-2/3 通过更好的分工缓解了这个问题。）

练习 5（块大小的约束） 为什么论文要求 $B_r \le d$？如果取 $B_r = B_c = M / 4d$ 会怎样？

SRAM 里要同时放下 $K_j, V_j$（$2 B_c d$ 个数）、$Q_i, O_i$（$2 B_r d$）和分数块 $S_{ij}$（$B_r B_c$）。取 $B_c = \lceil M/4d \rceil$ 时 $K_j, V_j$ 约占 $M/2$；$B_r \le M/4d$ 让 $Q_i, O_i$ 至多占 $M/2$；关键是 $S_{ij}$：有了 $B_r \le d$，$B_r B_c \le d \cdot M/4d = M/4$。三部分都是 $M$ 的常数倍，论文的分析只需要这一点（真实实现会把常数调到正好放得下）。如果 $B_r$ 也取 $M/4d$，$S_{ij}$ 有 $(M/4d)^2$ 个数，当 $M/4d > 4d$（即 $M > 16 d^2$）时它一个就超过了 $M$，放不下。

练习 6（改一行代码，预测结果） 在代码 03 演示 5 里：(a) 把分块注意力的块大小改成 1（每次只处理一个位置），bit/字符会变吗？(b) 把滑动窗口改成 96（等于上下文长度），结果和原模型一样吗？(c) 改成 200 呢？

(a) 不会（除了 1e-9 量级的舍入误差）：块大小只改变计算顺序，不改变函数；块大小 1 时就是”一个一个位置地做在线 softmax”，会很慢，但结果相同。 (b) 一样：窗口 96 时每个位置能看到自己和前面最多 95 个位置，而序列最长 96，所以没有任何位置被额外遮住，和标准因果注意力完全相同（差别只来自舍入）。 (c) 同 (b)：窗口比序列还长，不起作用。这说明”近似”只有在窗口小于实际需要的上下文时才会改变模型。

练习 7（为什么不能直接分块 softmax） 一个同学说：“把 $S$ 按列切成两块，各自做 softmax，然后把两个结果拼起来就行。“错在哪里？怎样修正？

每一块各自做 softmax，分母只包含这一块的项，两块的概率各自加起来都是 1，拼起来一行的和是 2，而且两块的相对大小也不对（一块里的最大值和另一块里的最大值没有比较过）。修正方法就是在线 softmax：每块记下自己的 $(m, \ell)$ 和未归一化（或按自己的分母归一化）的加权和，最后按 $e^{m_k - m}$ 换算到统一的基准、用合并后的 $\ell$ 归一化（练习 3）。

练习 8（训练天数） OPT-175B 用 992 块 A100，论文报告每块 GPU 最高 147 TFLOP/s。按 $6N_pD$ 计算，训练 300B 个 token 至少要多少天？OPT 团队在训练日志里估计”在 1024 块 A100 上连续训练约 33 天”，两者为什么不同？

$C = 6 \times 175 \times 10^9 \times 300 \times 10^9 = 3.15 \times 10^{23}$ FLOPs；$3.15 \times 10^{23} / (992 \times 147 \times 10^{12})$ 秒 ≈ 25.0 天（1024 块 GPU 时约 24.2 天）。33 天更长，原因可能包括：147 TFLOP/s 是”最高”而不是平均；日志的估计里可能包含了注意力、重算等额外 FLOPs 和评测的时间。这个练习想说明：从”需要的 FLOPs”到”天数”，中间隔着利用率，而利用率本身的记账方式各不相同。（OPT 实际花了约两个月，原因是硬件故障和训练不稳定——第 10 讲。）

练习 9（HFU 与重算） 如果一次训练对所有层都做激活重算（反向前把整个前向再算一遍），每个 token 的硬件 FLOPs 是 $6N_p$ 的多少倍？若此时 MFU 是 46.2%，HFU 是多少？PaLM 的 4.10 / 3.28 说明了什么？

前向 $2N_p$ + 反向 $4N_p$ + 重算的前向 $2N_p$ $= 8N_p$，是 $6N_p$ 的 4/3 倍；HFU $= 46.2\% \times 4/3 = 61.6\%$。PaLM 的比值 $4.10 / 3.28 = 1.25 < 4/3$，说明它只重算了一部分前向（例如只重算部分层或部分运算），在显存和计算之间做了折中。

练习 10（GPT-3 的 $P$ 有多大） GPT-3 175B：96 层、每层 96 个头、上下文 2048、16 位。如果用标准注意力训练，一个序列的所有层、所有头的 $P$ 一共多大？和一张 80 GB 的 A100 比较。

每层 $96 \times 2048^2 \times 2$ 字节 $= 805{,}306{,}368$ 字节 = 768 MiB；96 层合计 72 GiB（约 77.3 GB）——一个序列的 $P$ 就快把一张 80 GB 的卡占满了，还没算参数、优化器状态和其他激活。实际训练会用激活重算、张量并行等办法（第 9 讲）把它分摊掉，代价是额外的计算和通信；FlashAttention 让这部分直接降到每行两个数。

练习 11（生成速度的上限） 一个 70B 参数的模型用 16 位权重，在一张 A100 80GB（2.039 TB/s）上做 batch = 1 的生成。忽略 KV cache 和一切开销，每秒最多能生成多少个 token？一个 7B 的模型呢？这个估计和 FLOPs 有关吗？

每生成一个 token 要把全部权重读一遍：70B × 2 字节 = 140 GB（顺便说，一张 80 GB 的卡放不下，这里只是假设）；140 GB / 2.039 TB/s ≈ 68.7 ms，即每秒最多约 14.6 个 token。7B 模型：每秒约 146 个 token。这个估计只和带宽、权重字节数有关，和 FLOPs 无关——这就是”生成受带宽限制”的含义，也是为什么 4 位量化（字节数变成 1/4）能让生成快接近 4 倍。

练习 12（概念） FlashAttention 对哪种训练收益最大：(a) 宽度 12288、序列 2048 的超大模型；(b) 宽度 1024、序列 8192 的中等模型？用 §4.4 的账本解释。

(b)。注意力相关的计算和读写 $\propto N^2 d_{\text{model}}$（每层），线性层 $\propto N d_{\text{model}}^2$，两者之比约 $N / d_{\text{model}}$。(a) 的比值 2048 / 12288 ≈ 0.17，注意力只是小头（代码 01：GPT-3 那层融合后只快 1.14 倍）；(b) 的比值 8，注意力是大头，而且标准注意力的 $P$ 会占大量显存。FlashAttention 对”长序列、窄模型”的收益最大。

12. 延伸阅读 · Further reading

原课阅读（1/11 FlashAttention）

• Jay Alammar. The Illustrated Transformer. 2018. jalammar.github.io/illustrated-transformer
• Vaswani et al. Attention Is All You Need. NeurIPS 2017. arXiv:1706.03762
• Horace He. Making Deep Learning Go Brrrr From First Principles. 2022. horace.io/brrr_intro.html——本讲 §2 的主要依据，强烈推荐全文阅读。
• NVIDIA. Deep Learning Performance Guide（GPU 性能背景、矩阵乘法的形状与 tensor core）. docs.nvidia.com/deeplearning/performance
• Dao, Fu, Ermon, Rudra, Ré. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. arXiv:2205.14135

补充

• Milakov & Gimelshein. Online Normalizer Calculation for Softmax. 2018. arXiv:1805.02867——在线 softmax 的出处。
• Ivanov et al. Data Movement Is All You Need: A Case Study on Optimizing Transformers. MLSys 2021. arXiv:2007.00072
• Williams, Waterman, Patterson. Roofline: An Insightful Visual Performance Model for Multicore Architectures. CACM 2009.
• Chowdhery et al. PaLM: Scaling Language Modeling with Pathways. 2022. arXiv:2204.02311——第 4.1 节与附录 B 定义了 MFU。
• Dao. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. 2023. arXiv:2307.08691
• Shah et al. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. 2024. arXiv:2407.08608
• 本目录 CS336 第 5 讲（GPU 与 TPU）、第 6 讲（kernel 与 Triton）：在真实 GPU 上实现这些想法。

13. 下一讲预告 · Next lecture

这一讲把一块芯片”用满”。但即使 MFU 达到 50%，一块 A100 要训练 PaLM 也需要几千年；而且 PaLM 的参数（16 位就是 1 TB 多）根本放不进一张 80 GB 的卡。

第 9 讲：分布式训练——数据、张量、流水线并行与去中心化训练（1/18 客座 Ce Zhang）回答：训练时显存都花在哪里（参数、梯度、优化器状态、激活——一个参数要 16 个字节）？上千张 GPU 怎样分工：每张卡算不同的数据（数据并行、ZeRO）、把一个矩阵切开（Megatron 的张量并行）、把层分给不同的卡（流水线并行与它的”气泡”）？如果这些 GPU 不在同一个数据中心，而是散落在世界各地、用慢速网络连接，还能训练吗？代码会模拟 ring all-reduce，把第 1 讲的 base model 切到多个”虚拟设备”上，验证每一种并行方式都和单机训练给出相同的结果。