FlashAttention 与 Online Softmax

FlashAttention = 对标准注意力做 GPU 访存优化:通过 SRAM 切块(Tiling)与 Online Softmax 增量归一化,把 HBM 读写从 $O(N^2)$ 降到 $O(N^2 / M)$(大 $N$ 时近似 $O(N)$,$M$ 为 SRAM 块容量)。标准注意力 $S = \frac{QK^\top}{\sqrt{d}}$、softmax、$O = SV$,$N \times N$ 的 $S$ 矩阵需写回 HBM 再读回——这是 $O(N^2)$ 的访存瓶颈;FlashAttention 把 $Q/K/V$ 切成块常驻 SRAM,逐块计算局部 softmax 并用运行最大值 + 重缩放合并: 处理元素 $x_i$ 时,$m_{new} = \max(m, x_i)$,$l_{new} = l \cdot e^{m - m_{new}} + e^{x_i - m_{new}}$,输出按 $e^{m - m_{new}}$ 重缩放,最终 $\frac{1}{l}$ 归一化,数学结果与标准 softmax 完全一致。

长序列注意力(4K/32K/1M 上下文)训练与推理的标准 kernel;面试高频“FlashAttention 为什么快”“Online Softmax 是什么”“FA-1/2/3 区别”;也是 PD 分离、KV 量化等推理优化的底层支撑。

标准注意力受 HBM 带宽支配:长序列下 $S$ 矩阵 $O(N^2)$ 读写主导耗时,且显存 $O(N^2)$ 存不下。FlashAttention-1(2022)用切块 + Online Softmax + 反传重算(不存 $S$)把 HBM 访问降一个数量级、显存从 $O(N^2)$ 降到 $O(N)$,64K 序列上比 PyTorch eager 快 2-4 倍;FA-2(2023)做多头并行、把 $N^2$ 的 softmax 开销移出并行域,约再快 2 倍;FA-3(2024)用 Hopper 的 warp specialization + TMA 异步拷贝,并把 FP8 引入 kernel。

推导标准注意力的 HBM 瓶颈: 为什么 $S = \frac{QK^\top}{\sqrt{d}}$ 与 softmax 之间的 $N \times N$ 矩阵往返是 $O(N^2)$ 访存;长序列(如 128K)下为什么带宽主导?

写出 Online Softmax 增量公式: $m_{new} = \max(m, x_i)$、$l_{new} = l \cdot e^{m - m_{new}} + e^{x_i - m_{new}}$;为什么需要 running max + 重缩放而不能直接累加 $e^{x_i}$?

Tiling 细节: 块大小受什么限制(每 SM SRAM 如 H100 约 228KB)?HBM 访问如何从 $O(N^2)$ 降到 $O(N^2 / M)$;causal mask 如何在 kernel 内处理?

FA-1/2/3 演进: 反传重算为何把显存从 $O(N^2)$ 降到 $O(N)$;FA-2 的多头并行与 upcast 优化;FA-3 的 warp specialization/TMA 异步与 FP8 各解决什么?

FlashAttention 与 KV cache/PagedAttention 的关系: kernel 层优化为何不与推理引擎层优化冲突;为什么 decode 阶段带宽瓶颈(逐 token 读权重)不能靠 FA 解决?