1.3 KiB
1.3 KiB
FlashAttention
IO感知的精确注意力优化，由 Dao 等 2022 年提出，是注意力计算效率的最大单次突破。
核心洞见
传统注意力实现的瓶颈不是计算（FLOPs），而是GPU 内存层级之间的数据传输（IO）。标准实现需要多次读写 HBM（高带宽内存），而 HBM 带宽远低于 SRAM。
关键创新
- Tiling: 将注意力矩阵分块计算，每块保持在快速 SRAM 中
- Recomputation: 反向传播时重新计算 softmax 而非存储中间结果，节省内存
- IO-Aware: 算法设计以最小化 HBM↔SRAM 数据传输为核心目标
性能收益
- 速度: 比标准注意力快 2-4x
- 内存: 内存占用量从 O(n²) 降至 O(n)
- 精度: 数值精确（非近似），无精度损失
版本演进
- FlashAttention-1 (2022): Tiling + Recomputation
- FlashAttention-2 (2023): 更好的并行化和 work partitioning
- FlashAttention-3 (2024): 异步计算 + 低精度 (FP8)
相关概念
- flash-attention-3 — 最新版本
- kv-cache-bottleneck — KV 缓存瓶颈（FlashAttention 不直接解决但互补）
- sparse-attention-patterns — 稀疏注意力也可结合 FlashAttention
- llm-attention-survey-2026 — 综述参考