FlashAttention：IO 感知的精确注意力
一句话定义
FlashAttention 通过分块（tiling）与 online softmax 把注意力的计算全部留在 GPU 片上高速缓存（SRAM）内完成，从不把 n×n 注意力矩阵写回显存，做到结果精确不变的前提下把显存从 O(n²) 降到 O(n)、速度提升数倍。
为什么重要
它是 2022 年后所有长上下文能力与 prefill 速度的底层功臣：没有它，4k 以上上下文的 prefill 要么慢一个量级、要么 O(n²) 显存直接放不下。它是"算法不变、只改计算组织方式就拿到数量级收益"的典范案例，也是理解"访存意识（IO-aware）"设计思想的最佳教材。
前置知识
- softmax 的性质（可分解为 max 与 log-sum-exp 两步）；kp-003（prefill 计算密集）。
- GPU 存储层级：SRAM（快、小，约几十 MB）与 HBM（慢、大，几十 GB）。
核心概念
- 标准注意力的浪费：先算 S=QKᵀ（n×n）写回 HBM，再 softmax 写回，再乘 V——n² 级中间矩阵在慢速 HBM 上往返两次，瓶颈是 IO 而非 FLOPs。
- 分块计算（tiling）：把 Q、K、V 切成小块，逐块在 SRAM 内算部分注意力结果。
- online softmax：块与块之间维护运行最大值与运行和，看到新块时对已积累的输出做一次代价很小的重缩放——数学上与全局 softmax 严格等价。
- 精确而非近似：与标准注意力结果数值一致（浮点误差级），这是它区别于各类"线性注意力/稀疏近似"的根本点。
- FlashAttention-2：减少非矩阵乘操作占比、沿序列维并行、改善 warp 分工，在 A100 上进一步接近算力上限。
原理与机制
标准注意力的浪费不在计算而在搬运：n×n 中间矩阵要在慢速 HBM 上往返两次。FlashAttention 用 tiling 把 Q/K/V 切成能装进 SRAM 的小块，块内算完注意力再换下一块，中间量只写 O(n) 的输出；online softmax 用"运行最大值 + 运行和 + 重缩放"保证跨块累加与全局 softmax 严格等价——所以它是精确算法，收益来自 IO 复杂度下降而非计算量减少。
公式与模型
IO 复杂度（HBM 读写字符数）：
M 为 SRAM 大小。d（头维度）通常 64–128，故 FlashAttention 的 HBM 往返可低一个数量级；显存占用从 O(n²) 降到 O(n)——32k 上下文下，n² 矩阵在 FP16 下超过 2GB/头，而 Flash 只需存 KV 本身。
直观类比
改卷子不必把全班所有卷子铺满操场（写回 HBM）再逐份批改，而是每次抽一小摞在桌上（SRAM）批完记总分，再换下一摞——桌面小，但没人要求把操场铺满。
实例或案例
prefill 8k 上下文的实测典型值：朴素实现受 HBM 往返限制，注意力部分耗时是 FlashAttention 的 3–7 倍；且朴素实现 O(n²) 显存使 16k+ 训练/推理直接不可行。工程上无需手动启用——vLLM/SGLang/TensorRT-LLM 默认即用，需要关注的是"自定义注意力算子是否也遵循 IO-aware"。
常见误区
- "FlashAttention 是近似/有损注意力"：不，结果精确等价，只是计算组织方式不同。
- "它减少了计算量"：FLOPs 几乎不变，减的是 HBM IO 与显存占用——所以对带宽受限的 decode 帮助小于对 prefill 的帮助。
- "任何模型都能直接替换"：依赖算子实现与硬件架构支持（SM 数量、SRAM 大小），老架构收益打折。
自测题
- FlashAttention 省的是什么资源？不省什么？
要点：省 HBM IO 与 O(n²) 显存；FLOPs 基本不变。
- online softmax 如何保证与全局 softmax 等价？
要点：维护运行 max 与运行和，新块到来时对累积结果按比例重缩放。
- 为什么 FlashAttention 对 prefill 的收益比 decode 更明显？
要点：prefill 序列长、注意力占比高且 compute-bound，O(n²) IO 是主要浪费；decode 的瓶颈更多在权重读取。
与其他知识点的关系
与 kp-013 组合决定长上下文可行性；kp-016 的 GQA 减 KV 体积、Flash 减 IO，两者互补；kp-024 的 PD 分离里 prefill 实例的效率依赖本页技术。
延伸阅读
《FlashAttention》（Dao 等，2022）第 3 节的 IO 分析是全领域"访存意识"设计思想的代表；FlashAttention-2 论文对并行化的讨论适合工程向读者。