目录
前面学 Attention 时，我们写下了一个很简洁的公式：
\[\operatorname{Attention}(Q,K,V)=\operatorname{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V\]公式只占一行，照着它一步一步计算，却可能在 GPU 上反复搬运一张很大的中间表。FlashAttention 的核心不是改掉这个公式，而是改变计算顺序：尽量在片上算完当前的一小块，不把整张注意力矩阵写回显存。[1]
两种数据流：公式相同，搬运方式不同。点击查看原始 SVG
从一张注意力矩阵说起
假设一段输入有 $T$ 个 Token。每个 Query 都要与 Key 比较，分数形成一张 $T\times T$ 的表。朴素实现先算 $QK^\top$，把分数写到 GPU 显存；读出来做 Softmax，再写回；接着又读出来与 $V$ 相乘。
如果 $T=4096$，仅一张单头、单样本的 FP16 分数表就有 $4096^2$ 个数，约 32 MiB。这还没算其他中间结果、多个头和批次。关键不只是“放不放得下”：写出、再读回这张表，也要花时间。
显存大，不等于访问快
可以把 GPU 的 HBM 想成仓库：容量大，但从中取货有搬运成本。片上 SRAM 像工作台：空间小，取用却快。这个比喻只用于理解存储层级，不代表真实 GPU 只有两层内存。
如果每做一步都把整张表送回仓库，计算单元就会等数据。FlashAttention 设计时把这种读写成本也算进去，所以论文称它为 IO-aware（考虑数据搬运的算法）。[1]
| 存储位置 | 在这篇文章里的作用 | 直观特点 |
|---|---|---|
| HBM（GPU 显存） | 保存输入、输出及较大的数据 | 容量大，往返搬运有成本 |
| SRAM（片上存储） | 暂存正在计算的小块 | 容量小，访问快 |
把整张表改为分块计算
FlashAttention 将 $Q$、$K$、$V$ 划成能放进片上存储的块。对一块 Query，依次载入 Key 和 Value 的小块，计算局部分数，更新当前输出，然后再处理下一块。已经处理过的局部分数无需组成完整的 $T\times T$ 表写回 HBM。
这里要留意：它仍然会比较需要比较的 Query 和 Key。常见的精确 FlashAttention 没有把平方量级的注意力计算变成线性；省下的是大块中间矩阵的存储与搬运。论文另讨论了 block-sparse 近似变体，那是另一回事。[1]
Softmax 分块后怎么保持结果不变
真正容易卡住的是这里：Softmax 要知道一整行所有分数，分块时后面的分数还没见到，怎么先算前面的块？答案是保存一小组可更新的统计量，而不是保存整行的所有中间分数。
以一行分数 $[1,2,3]$ 为例。先看 $[1,2]$：当前最大值是 $m=2$，指数和是 $\ell=e^{1-2}+e^{2-2}$。当分数 $3$ 到来，最大值变为 $m’=3$。原先那一部分相对新最大值要乘 $e^{2-3}$，因此：
\[\ell'=e^{m-m'}\ell+e^{3-m'}=e^{-1}(e^{-1}+1)+1\]输出向量的加权和也按同样比例缩放，再加上新块对 $V$ 的贡献。所有块处理完后，用累计的指数和做归一化。这就是在线 Softmax 的直觉：旧块不用重新存一遍，但遇到更大的分数时必须重新标定旧贡献。对同一数学定义，它得到精确 Attention 的结果；真实计算仍会有浮点舍入差异，不能把“精确”理解为逐位完全相同。[1]
动手看看中间结果去了哪里
这里的大小只用于比较一张 FP16 $T\times T$ 表；真实显存峰值还取决于批量、头数、精度、内核实现和训练阶段。交互展示的是数据流差异，不是性能基准测试。
反向传播为什么还要重算
训练时需要求梯度。朴素做法可以存下大量前向中间结果，反向时直接读取。FlashAttention 选择保存更少的信息，反向时从 $Q,K,V$ 和少量统计量重算需要的局部分数。这会多做一些算术，但减少 HBM 读写和显存占用；在合适的 GPU 与输入形状下，整体训练反而更快。[1]
更快的边界在哪里
FlashAttention 优化的是一种具体的计算实现，不是新的 Attention 定义。它不保证任何 GPU、任何序列长度都能按同一个倍数提速；内核版本、硬件、精度、头维度和掩码方式都会影响收益。FlashAttention-2 又进一步改进了工作划分与并行度，说明“减少 IO”之后，如何把 GPU 算力用满仍然重要。[2]
记住三条边界就够了：精确 Attention 不等于逐位相同；省中间显存不等于消灭全部显存开销；更少 IO 不等于理论计算量从 $O(T^2)$ 变成 $O(T)$。
约 90 秒讲解视频
参考文献（References）
- Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Ré, C. “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.” NeurIPS, 2022. arXiv:2205.14135.↩
- Dao, T. “FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning.” ICLR, 2024. arXiv:2307.08691.↩
第十六课复习总图
分块、在线 Softmax、重算与边界，一页回顾。点击查看原始 SVG
下一课进入生成：Temperature、Top-K 与 Top-P 怎样改变模型选出的下一个 Token？