FlashAttention 为什么更快、更省显存？

　　标准 Attention 会产生 T×T 的分数矩阵。真正拖慢 GPU 的不只是乘法数量，还包括中间结果在 HBM 与片上 SRAM 之间反复读写。

FlashAttention 没有换掉 Attention 公式，它重新安排了同一套计算的数据流。

1、标准实现慢在哪里

朴素实现先算完整 QKᵀ，写入显存，再读取做 Softmax，之后又读取与 V 相乘。长序列下，中间矩阵很大，显存读写与峰值占用都很高。

2、GPU 内存有快慢层级

HBM 容量大但访问相对慢，片上 SRAM 容量小却快得多。高性能内核要尽量让已经搬进 SRAM 的数据多做计算，减少往返 HBM。

3、分块计算

FlashAttention 把 Q、K、V 切成能放进 SRAM 的小块。每次载入一部分，完成局部矩阵乘法与 Softmax 更新，再继续下一块。

它不把完整 T×T 分数矩阵写回 HBM，因此降低了中间显存和 IO。

4、Softmax 为什么能分块

Softmax 依赖一整行的最大值与分母。FlashAttention 使用在线算法，随着新块到来更新当前最大值和归一化和，并对先前结果做正确缩放。最终结果与普通精确 Attention 一致，而不是近似值。

5、Backward 为什么会重算

为了不保存巨大的中间分数矩阵，反向传播时可以根据 Q、K、V 和少量统计量重算部分结果。多做一些算术，换来更少的显存读写，整体反而可能更快。

6、它没有改变什么

FlashAttention 不会把 Attention 的理论 O(T²) 计算复杂度变成线性，也不会自动改变模型质量。它优化的是实际 GPU IO、内核融合与中间存储。硬件、形状、精度和实现版本都会影响收益。

容易混淆的三件事

1. FlashAttention 不是稀疏近似。
2. 省显存不等于线性复杂度。
3. 重算可能增加 FLOPs 但减少 IO。

记住这 5 件事

1. 结果是精确 Attention
2. 核心是 IO-aware
3. QKV 按块处理
4. 不保存完整分数矩阵
5. 计算量仍随 T² 增长

参考资料

• FlashAttention 论文

第十六课复习总图

　　下一阶段进入生成：Temperature、Top-K 与 Top-P 怎样改变模型选出的下一个 Token？