一句话定义
FlashAttention 是一种 IO 感知的注意力实现：它不把完整的注意力矩阵写回显存，而是把计算切成小块，在片上高速缓存里完成，并用在线 softmax 逐步累积结果。
问题出在哪
标准注意力要先算出完整的 N×N 分数矩阵，再 softmax、再乘 Value。序列越长，这个中间矩阵越大，GPU 需要在显存和计算单元之间反复搬运数据。很多长上下文场景的瓶颈不是算力，而是显存带宽。注意力算子的算术强度低，数据搬运往往比乘加更先成为瓶颈。
关键机制
- 分块（tiling）：把 Q、K、V 切成能放进片上缓存的小块，逐块计算。
- 在线 softmax：不一次性拿到整行分数，而是边算边维护最大值和归一化因子，最后得到等价结果。
- 重计算：反向传播时不保存完整中间矩阵，用重算换显存。
- 数学结果与标准注意力一致，不是近似算法。
实际收益
- 显著降低显存占用，让更长序列和更大 batch 成为可能。
- 减少显存读写后，训练和推理速度通常都有提升。
- 与 KV Cache 配合时，主要优化的是注意力计算本身，而不是缓存管理。
- 短序列下收益可能不明显，长序列和高 batch 更能体现。
用在哪里
- 长上下文训练/推理、需要降低显存占用或提高吞吐的注意力实现。
- 短序列、算力受限或硬件不支持时，收益有限，未必值得替换基线实现。
常见误区
- FlashAttention 不是稀疏注意力，也不改变注意力的语义。
- 它不能单独解决 KV Cache 的显存增长问题。
- 是否加速取决于硬件、序列长度和实现版本，短序列下收益可能有限。