【Transformer 与注意力机制】42｜FlashAttention：注意力计算的硬件级重写
内容提要
FlashAttention是一种IO感知的精确注意力算法，通过分块和在线softmax避免物化完整注意力矩阵，减少HBM读写，提升效率并降低显存占用，虽FLOPs仍为O(n²)，但优化了硬件执行路径，使长序列训练更可行。
延伸解读
FlashAttention 不是近似，而是精确重排
很多人误以为 FlashAttention 是稀疏或线性注意力，但文章明确指出它计算的是标准 softmax attention 的精确结果，只是通过分块和在线 softmax 改变了计算顺序，避免物化完整的注意力矩阵。它不改变 attention 的语义，因此可以作为 drop-in 替换，无需重训模型。理解这一点有助于避免在技术选型时混淆不同优化路径。
瓶颈在内存读写，而非 FLOPs
文章强调，标准 attention 的朴素实现会产生巨大的中间矩阵，导致频繁的 HBM 读写，成为实际性能瓶颈。FlashAttention 通过 IO-aware 设计，将数据尽量留在片上 SRAM，减少 HBM 流量，从而提升吞吐并降低显存占用。这提醒我们，在评估模型性能时，不能只看 FLOPs，还要考虑内存带宽和中间状态的存储开销。
FlashAttention 与稀疏/线性注意力互补
FlashAttention 优化的是标准 attention 的执行效率，而稀疏或线性注意力则改变 attention 的数学形式或近似目标，以降低计算复杂度。文章指出，两者并非替代关系，而是互补：在长上下文场景中，可先用 FlashAttention 将标准 attention 训练到尽可能长的序列，若仍不够，再引入架构层面的改动。理解这一边界有助于合理选择优化策略。
Q&A
FlashAttention 是近似注意力吗？
不是。FlashAttention 计算的是标准 softmax attention 的精确结果，只是改变了计算和内存访问顺序，因此是 exact attention。
FlashAttention 的核心思想是什么？
FlashAttention 通过分块（tiling）和在线 softmax，避免物化完整的注意力矩阵，减少 HBM 读写，从而提升效率并降低显存占用。
为什么标准 attention 的瓶颈常常不是 FLOPs 而是内存读写？
因为标准 attention 的中间矩阵（如 S 和 P）很大，需要频繁写入和读取 HBM，而 GPU 的 HBM 访问代价远高于片上 SRAM，导致内存 I/O 成为瓶颈。
FlashAttention 与稀疏注意力、线性注意力有什么区别？
FlashAttention 不改变 attention 语义，仍计算标准 softmax attention，只是优化了内存访问；而稀疏注意力和线性注意力改变了 attention 模式或数学形式，以降低计算复杂度。
FlashAttention 如何节省显存？
FlashAttention 避免将完整的注意力矩阵 S 和 P 写回 HBM，forward 阶段保存的中间状态显著减少；backward 阶段通过重算部分注意力来节省显存，用少量计算换取显存和 I/O 的减少。
FlashAttention 能消除长上下文瓶颈吗？
不能。FlashAttention 降低了显存和 I/O 成本，但 attention 关系数量仍随序列长度平方增长，因此长上下文瓶颈被缓解，但并未消失。
FlashAttention-2 和 FlashAttention-3 的主要改进方向是什么？
FlashAttention-2 优化了并行性和 work partitioning，在 A100 上接近 GEMM 利用率；FlashAttention-3 面向 Hopper 架构，利用 TMA 异步搬 tile 和 FP8 低精度，进一步提升性能。
FlashAttention 对训练和推理分别有什么好处？
训练时减少显存占用和 HBM 流量，支持更长序列和更大 batch；推理时 prefill 阶段同样受益，但 decode 阶段瓶颈转向 KV Cache 读写。
