[2205.14135] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
FlashAttention 的提出，核心在於降低Transformer 注意力計算中的 記憶體讀寫成本（IO）。這一改善，使其在多數 GPU 上能獲得 2–4× 以上的速度提升。然而，一個有趣的實驗現象是：不同 GPU 上的 speedup 並不一致。例如：
- RTX 3090：最高可達 4.5× 加速
- A100：約 2–4×
- NVIDIA T4：明顯較低
一、HBM 記憶體頻寬：速度差異的主要來源
Transformer 注意力運算有一個重大瓶頸：
要讀寫一個 N×N 的注意力矩陣，這使得標準 Attention 的時間主要被「記憶體頻寬」決定，而非計算量。
以下是論文中的 GPU 比較結果：
GPUHBM 帶寬FlashAttention SpeedupRTX 3090約 900 GB/s2.5× – 4.5×（最高）A100約 1.5 TB/s2× – 4×T4約 300 GB/s最小
HBM 帶寬越低，FlashAttention 的加速比反而越高。
為什麼會如此？
因為標準 Attention 是 IO-bound（記憶體頻寬限制） 的計算，當 HBM 速度不夠快時：
- 標準 Attention 特別慢
- FlashAttention 因大量降低 IO 需求，相對快非常多
因此 RTX 3090 雖然 FLOPs 不如 A100，但因其記憶體頻寬較低，反而更能顯現 FlashAttention 的優勢。
二、SRAM（片上快取）大小：決定 FlashAttention 的極限效益
FlashAttention 的核心技巧之一是 tiling（分塊）。所有注意力計算會被切成小區塊，在 GPU 的 SRAM（L2 cache、shared memory、register file）中完成，幾乎不從 HBM 讀寫 NxN 交互矩陣。
因此：
- SRAM 越大 → 每次可處理越大的 block → IO reuse 越高 → speedup 越大
這也是為什麼：
- A100 SRAM 大 → 表現穩定優秀
- RTX 3090 SRAM 也大 → speedup 甚至比 A100 更高
- T4 SRAM 小 → block 必須切得更碎 → speedup 明顯下降
在 T4 的實驗中，block size 的縮小使得 FlashAttention 無法充分利用 SRAM，導致加速比下降，儘管它仍比標準 Attention 省記憶體。
三、FLOPs（計算能力）並非 Speedup 的主因素
很多人直覺認為 GPU 算力越高，FlashAttention 越快。但實際上：
- 標準 Attention 不是 Compute-bound（計算瓶頸）
- 而是 Memory-bound（記憶體瓶頸）
FlashAttention 雖然額外做了一些重算（recomputation），增加少量 FLOPs，但這些計算量遠遠小於省下的 HBM IO，因此：
- GPU FLOPs 高低，對 speedup 不造成主要差異
- 決定性因素仍是 HBM 與 SRAM
這解釋了為什麼：
- A100 的 FLOPs 遠高於 3090
- 但 FlashAttention 在 3090 反而更快
因為：
速度決定因素不是算力，而是 IO 密度（IO intensity）能否大幅降低。
綜合分析：FlashAttention Speedup 的三要素模型
可以將加速比概念化為：
其中：
- 標準 Attention IO 成本：
- FlashAttention IO 成本：
因此 speedup 取決於：
① HBM 帶寬越小 → Speedup 越高
（標準 Attention 更慢）
② SRAM 越大 → Speedup 越高
（FlashAttention block 越大）
③ FLOPs 幾乎不影響 speedup
（不是計算瓶頸）
結論：FlashAttention 的價值在於「IO-aware AI 計算」
FlashAttention 的成功，象徵著深度學習效能的瓶頸已經從「算力不足」轉向：
「記憶體頻寬不足 — Memory wall」。
因此其加速效果會隨不同 GPU 的記憶體子系統而變動，是可以預期的結果。
這也揭示了未來 AI 加速的方向：
- 軟體層面：IO-aware 演算法（如 FlashAttention / FlashDecoding）
- 硬體層面：更大 SRAM、更高帶寬、更靠近計算單元的記憶體架構
FlashAttention 不只是一個快速 attention kernel，更是 深度學習從 FLOPs 時代走向 IO 時代 的重要里程碑。
資料來源
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and memory-efficient exact attention with IO-awareness. arXiv:2205.14135. https://arxiv.org/abs/2205.14135
