FlashAttention

FlashAttention 是 IO 感知的精确注意力实现:把 $Q K^T$ 分块 (tiling) 放进 SRAM 计算,配合 online softmax 维护运行最大值 $m_i$ 与归一化因子 $l_i$(合并新块时递归缩放:$m_{new} = max(m, m_{local})$),前向不把 $O(n^2)$ 注意力矩阵写回 HBM,反向用重算而非存储;动机是 HBM 与 SRAM 的带宽差距约 10-20 倍(A100: 40GB HBM vs 约 108KB SRAM)。

大模型训练与推理的事实标准实现(PyTorch SDPA、NVIDIA FA2/FA3、DeepSeek/Meta 训练栈),序列越长收益越大,长上下文训练几乎必用。

标准实现把注意力矩阵反复写读 HBM,IO 成为主要瓶颈;分块 + 重算把 HBM 访问从 $O(n^2 d)$ 降到 $O(\frac{n^2 d^2}{M})$($M$ 为 SRAM 大小),速度提升 2-4 倍、显存从 $O(n^2)$ 降到 $O(n)$,且结果与标准 softmax 逐位一致(不是近似)。

为什么说 attention 是 memory-bound?HBM 与 SRAM 的带宽/容量差距?

online softmax 如何合并分块统计量?(运行最大值 m 与归一化因子 l 的递归公式)

反向传播为什么重算而不存储?显存从 $O(n^2)$ 降到多少?

FlashAttention-2 相比 v1 改进了什么?(省略中间矩阵、减少非矩阵运算、约 2×)?

FlashAttention 为什么是精确的?与稀疏/低秩注意力有什么本质区别?