FlashAttention 如何减少注意力的显存访问与运行时间？

图：A100 上的官方实现基准；回答时不要脱离图中硬件与输入条件复述倍数。来源：Dao-AILab/flash-attention，作者 FlashAttention contributors，许可 BSD-3-Clause。

考察意图

这道题不只考“分块”两个字。合格回答需要建立完整因果链：标准注意力为什么产生大量 HBM 往返；在线 Softmax 怎样跨块保持正确；为什么算术复杂度没有降阶却仍能加速；收益在哪些输入和硬件上可能不成立。

回答前自测

先合上答案，尝试写出：标准 attention 的三个大中间量、在线 Softmax 至少保存哪两个行统计量，以及 FlashAttention 对算术量和激活显存分别做了什么。

30 秒回答

FlashAttention 是标准全注意力的 IO-aware 实现。它把 Q、K、V 分块放进片上 SRAM，维护每行最大值和指数和来完成在线 Softmax，不把完整的 T×T 分数或概率矩阵写回 HBM。它没有消除所有 query-key 配对，算术复杂度仍约为 O(T²d)；主要收益是减少 HBM 读写和中间激活，延迟是否下降取决于序列长度、硬件和实际 backend。

90 秒回答

朴素实现通常先算 S=QKᵀ 并写回 HBM，再读取 S 做 Mask、Softmax，写回 P，最后读取 P 与 V 相乘。长序列时，S/P 是 T×T，物化和反复读写会成为显存与带宽瓶颈。

FlashAttention 将 Q 的行和 K/V 的列分块。每处理一个 K/V 块，就更新该 query 行目前的最大值 m、指数和 l 和加权输出；若新块出现更大分数，旧累计量按 exp(m_old-m_new) 重缩放。因此分块结果仍等价于全局稳定 Softmax，不需要保存完整注意力矩阵。Backward 再按块重算局部分数，用额外计算换取更少的激活和 HBM 流量。

它仍计算全注意力，所以不能说算术复杂度从二次降为线性。实际收益依赖序列长度、head dimension、dtype、Mask、GPU 和框架是否真的选中融合 backend；短序列和单 token Decode 的瓶颈可能不同。

评分点

• 指出瓶颈是 HBM 与片上存储之间的 IO，而不只是“显存不够”。
• 明确不物化完整 T×T 分数或概率矩阵。
• 能解释在线 Softmax 的 m/l 更新或旧累计量重缩放。
• 区分 O(T²d) 算术量、近线性的额外激活和真实墙钟时间。
• 提到 backward 重计算与适用条件。
• 不把它与稀疏注意力、线性注意力或 PagedAttention 混为一谈。

连续追问

1. 为什么不能对每个块单独 Softmax 再拼接？
2. 序列长度从 8K 到 16K，分数矩阵、QKV 和 KV Cache 分别怎样变化？
3. “精确注意力”为什么不保证 FP16 下逐 bit 相同？
4. FlashAttention-2 相比第一版主要改了模型数学还是 GPU 工作划分？
5. 如何确认 PyTorch 实际使用了 FlashAttention backend，而不是静默回退？
6. 为什么训练或 Prefill 的加速结论不能直接套到单 token Decode？

关联知识

• 完整原理：FlashAttention：分块、在线 Softmax 与 IO-aware 注意力
• 计算基线：缩放点积注意力
• 量级练习：8K 上下文的注意力中间量估算
• 方法辨析：注意力优化方法比较