大模型算注意力为什么又快又省？聊聊FlashAttention

生产队队长

🧠 大模型一跑长上下文，显存就"爆"？罪魁祸首往往是那个看起来人畜无害的注意力计算。
🔍 问题出在哪
注意力要算一个 N×N 的矩阵，N 是序列长度。上下文 8K 时就是 6400 万个元素，中间结果全都要先写进显存、再读出来。

• ❌ 传统做法：先算出完整注意力矩阵，再过 softmax，最后乘 V——中间矩阵在显存里来回搬运
• ✔ 显存占用随序列长度平方增长，长文本直接撑爆
  

💡 FlashAttention 的思路：分块+在线 softmax
它不追求"一次算完"，而是把 Q、K、V 切成小块，在 GPU 片上高速缓存 SRAM 里逐块算完就走。

• 🔒 中间结果不落显存，显存占用从 O(N²) 降到 O(N)
• ⚡ 少了显存读写这个瓶颈，速度反而更快——计算不是最慢的，搬运才是
  

---

一句话总结：FlashAttention 快的秘诀不是算得猛，而是不把中间结果来回搬。

本质上它是一次 IO 优化：把数据搬运量砍掉一个数量级，同样的芯片就能跑更长的上下文。现在主流框架的注意力算子，底层几乎都是它的变体。
💬 互动一下
你训练或推理时被显存坑过吗？评论区聊聊你印象最深的一次 OOM 现场～