前一篇线性注意力从 FLOPs 账本上砍 $n^2$，代价是换掉 softmax 核。FlashAttention 走了另一条路：核一个数都不改， attention 还是标准的 $\mathrm{softmax}(QK^\top)V$，只是把计算顺序重排，让 $n\times n$ 的中间矩阵从头到尾不落地 HBM。它的出发点是一个容易被忽略的事实：GPU 上跑 attention，慢在搬数据，不在算。
先看硬件账本：HBM vs SRAM
以 A100 为例（FlashAttention 论文 Tri Dao et al. 2022 的测试环境）：HBM 容量 40GB、带宽约 1.5–2.0 TB/s；片上 SRAM 每流处理器约 192KB、带宽约 19TB/s。两者的带宽差接近两个数量级，而容量差是六个数量级。注意力恰好是那种"中间结果远大于输入"的算子：输入 $Q,K,V$ 各 $n\times d$，朴素实现却要把 $n\times n$ 的 $S$ 和 $P$ 各写读一遍。论文里有个 IO 复杂度分析：朴素 attention 的 HBM 访问次数是 $\Theta(nd + n^2)$，当 $n \gg d$（比如 $n=8\text{k}, d=64$ per head）时 $n^2$ 项主导，每元素一次读写的成本远超一次乘加。
把三次矩阵乘拆成访存看更直观：$S=QK^\top$ 要把 $S$ 写出；$P=\mathrm{softmax}(S)$ 要把 $S$ 读入、$P$ 写出；$O=PV$ 要把 $P$ 读入。一个 $n^2$ 矩阵在 HBM 上三进三出，这就是要消灭的开销。
分块：让中间矩阵留在片上
FlashAttention 的调度很直白：把 $Q$ 切成 $B_r$ 行一块，$K,V$ 切成 $B_c$ 列一块，块大小选到能同时驻留 SRAM 的程度（A100 上 $B_r\times B_c$ 的 $S_{ij}$ 块大约 64×64 或 128×64）。外层循环遍历 query 块，内层循环逐个载入 key/value 块，在片上完成 $S_{ij}=Q_iK_j^\top$、局部 softmax、累乘 $P_{ij}V_j$，只有累加结果 $O_i$ 最终写回 HBM。SRAM 容量 $M$ 个元素时，HBM 访问降到 $\Theta(n^2 d^2 / M)$——$M$ 是常量，所以渐近上仍是 $n^2$，但常数项被 $M$ 除掉了，这就是 2–4 倍实测加速的来源。
唯一的技术难点是：softmax 需要全局的行最大值做数值稳定（减去 max 防溢出），分块之后"全局 max"在内层循环跑到一半时还不知道。答案是在线 softmax。
在线 softmax：边算边修正
回忆稳定 softmax 的形式：对一行 $x$，$\mathrm{softmax}(x)_i = e^{x_i - m} / \sum_j e^{x_j - m}$，$m=\max_j x_j$。分块场景下维护两个 running 量：当前见过的行最大值 $m$ 和按 $m$ 归一的行和 $l=\sum_j e^{x_j-m}$。处理完新块 $S_{ij}$ 后，三个量同步更新：
$$ \begin{aligned} m^{\text{new}} &= \max\big(m,\ \mathrm{rowmax}(S_{ij})\big),\\ l^{\text{new}} &= e^{m - m^{\text{new}}}\, l \;+\; \mathrm{rowsum}\big(e^{S_{ij} - m^{\text{new}}}\big),\\ O^{\text{new}} &= O\, e^{m - m^{\text{new}}} \;+\; e^{S_{ij} - m^{\text{new}}} V_j. \end{aligned} $$
推导就一行：旧累加 $\sum e^{x-m}V$ 里每一项都按旧 max $m$ 归一，新 max 更大时，等价于把旧和乘上 $e^{m-m^{\text{new}}}$ 再补新块的贡献。注意旧块贡献 $O$ 要整体重缩放，不是只处理新块——漏掉这一步是在线 softmax 最常见的 bug，结果不会 NaN，只是 quietly 错，且错得离正确值很近。两个常识性的退化检查：只有一块时 $m^{\text{new}}=m$、缩放因子为 1，退化为标准 softmax；新块 max 没超旧 max 时同样因子为 1。
整个内层循环的额外开销是每个 query 块维护两个长度为 $B_r$ 的向量，访存可忽略；重缩放本身是 $B_r\times d_v$ 的一次逐行乘法，和 GEMM 比是小头。
反向传播：用重算换显存
前向不存 $n^2$ 矩阵，反向却需要 $P$ 和 $S$ 来算梯度（$\partial L/\partial V = P^\top dO$、$\partial L/\partial S = (dO\,V^\top)\odot P$ 再减行和项）。FlashAttention 的做法是从 $Q,K,V,O$ 和保存的 $m,l$ 在片上重算 $P$：块重载回 SRAM，重跑一遍前向的块级计算，材料化出局部 $P_{ij}$ 就立刻消费掉。重算的成本是多一遍前向的 FLOPs，换来的是反向激活显存从 $\Theta(n^2)$ 降到 $\Theta(nd)$——训练时配合梯度检查点，长序列终于塞得进卡了。这也是 activation recomputation 这个思想在 kernel 级别的同款应用：算力换显存，而 GPU 上算力恰恰是最不缺的。
小结
FlashAttention 的核心论断可以浓缩成一句：IO -aware 的算法设计。它不改变注意力的数学形式（和线性注意力、稀疏注意力有本质区别），只改变执行顺序，用在线 softmax 解决分块与全局归一的矛盾，用重算解决反向传播的记忆需求。后续所有高效注意力 kernel——FlashAttention-2/3、PagedAttention 里的 attention、各种推理框架的 fused kernel——都是在这个"少搬数据"的框架里做工程细化。