FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Tri Dao; Daniel Y. Fu; Stefano Ermon; Atri Rudra; Christopher Ré

論文 · arXiv 2205.14135

FlashAttention：具備 IO 感知的高速且高記憶體效率精確注意力機制

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré

發表於 2022-05-27cs.LG arXiv PDF

用 XYZNOT 閱讀（全文翻譯與 AI 講解）

三句話總結

🎯 標準精確注意力機制會在較慢的 GPU HBM 中實例化 $O (N^{2})$ 的注意力矩陣，導致其受限於記憶體頻寬，難以擴展至長上下文長度。
🔧 FlashAttention 透過將輸入資料分塊置於高速晶片上 SRAM、以增量方式更新 softmax，並在反向傳播時於晶片上重新計算注意力，在單一融合 CUDA 核心內完成精確注意力計算。
📈 FlashAttention 在 GPT-2 上實現高達 3 倍的端到端訓練加速，比 MLPerf 1.1 的 BERT-large 紀錄快 15%，並能將上下文長度擴展至 64K 個標元。

摘要

由於自注意力機制的計算時間與記憶體複雜度與序列長度成二次方關係，Transformer 在處理長序列時速度緩慢且極度消耗記憶體。近似注意力方法試圖透過犧牲模型品質來降低計算複雜度以解決此問題，但通常無法實現實際執行時間（wall-clock time）的加速。我們認為關鍵在於缺少了一項核心原則：讓注意力演算法具備 IO 感知（IO-aware）能力——即考量 GPU 各層級記憶體之間的讀取與寫入開銷。我們提出了 FlashAttention，這是一種具備 IO 感知的精確注意力演算法，利用分塊（tiling）技術來減少 GPU 高頻寬記憶體（HBM）與 GPU 晶片上 SRAM 之間的記憶體讀寫次數。我們分析了 FlashAttention 的 IO 複雜度，證明其所需的 HBM 存取次數少於標準注意力機制，且在多種 SRAM 大小下均達到漸近最優。我們亦將 FlashAttention 擴展至區塊稀疏注意力（block-sparse attention），實現了比現有任何近似注意力方法都更快速的近似演算法。FlashAttention 訓練 Transformer 的速度超越現有基準：相較於 MLPerf 1.1 的訓練速度紀錄，在 BERT-large（序列長度 512）上實現了 15% 的端到端實際時間加速；在 GPT-2（序列長度 1K）上實現了 3 倍加速；在 Long-Range Arena（序列長度 1K–4K）上實現了 2.4 倍加速。FlashAttention 與區塊稀疏 FlashAttention 讓 Transformer 能夠支援更長的上下文，從而帶來更高品質的模型（GPT-2 的困惑度改善 0.7，長文件分類提升 6.4 個百分點），並解鎖全新能力：成為首個在 Path-X 挑戰（序列長度 16K，準確率 61.4%）與 Path-256（序列長度 64K，準確率 63.1%）上表現優於隨機猜測的 Transformer。

英文原文摘要

Transformers are slow and memory-hungry on long sequences, since the time and memory complexity of self-attention are quadratic in sequence length. Approximate attention methods have attempted to address this problem by trading off model quality to reduce the compute complexity, but often do not achieve wall-clock speedup. We argue that a missing principle is making attention algorithms IO-aware -- accounting for reads and writes between levels of GPU memory. We propose FlashAttention, an IO-aware exact attention algorithm that uses tiling to reduce the number of memory reads/writes between GPU high bandwidth memory (HBM) and GPU on-chip SRAM. We analyze the IO complexity of FlashAttention, showing that it requires fewer HBM accesses than standard attention, and is optimal for a range of SRAM sizes. We also extend FlashAttention to block-sparse attention, yielding an approximate attention algorithm that is faster than any existing approximate attention method. FlashAttention trains Transformers faster than existing baselines: 15% end-to-end wall-clock speedup on BERT-large (seq. length 512) compared to the MLPerf 1.1 training speed record, 3 $\times$ speedup on GPT-2 (seq. length 1K), and 2.4 $\times$ speedup on long-range arena (seq. length 1K-4K). FlashAttention and block-sparse FlashAttention enable longer context in Transformers, yielding higher quality models (0.7 better perplexity on GPT-2 and 6.4 points of lift on long-document classification) and entirely new capabilities: the first Transformers to achieve better-than-chance performance on the Path-X challenge (seq. length 16K, 61.4% accuracy) and Path-256 (seq. length 64K, 63.1% accuracy).

論文解讀

問題與研究動機

Transformer 的時間與記憶體複雜度隨序列長度 $N$ 呈二次方增長。儘管近似注意力方法降低了理論上的 FLOP 數量，但它們往往無法帶來實際執行時間的加速，因為現代 GPU 計算的主要瓶頸在於記憶體傳輸（IO），而非算術運算。

在標準注意力機制中，計算 $S = Q K^{⊤}$ 、 $P = softmax (S)$ 以及 $O = PV$ 會將大小為 $N \times N$ 的中間矩陣寫入較慢的高頻寬記憶體（HBM），並多次重新讀取。這會產生 $Θ (N d + N^{2})$ 的 HBM 存取次數與 $O (N^{2})$ 的記憶體佔用，極大地限制了實際應用中的上下文長度。

核心概念

IO 感知（IO-Awareness）： 明確考量 GPU 的記憶體階層架構：高速但容量小的晶片上 SRAM（例如 A100 上每個 SM 有 192 KB，頻寬約 19 TB/s），對比容量大但速度較慢的 HBM（40–80 GB，頻寬 1.5–2.0 TB/s）。
分塊技術（Tiling，區塊 Softmax 計算）： 將輸入的 $Q, K, V$ 切分為可完全放入 SRAM 的小區塊。透過追蹤 softmax 即時的列最大值與正規化統計量，跨區塊以增量方式計算 softmax 歸約，無需在 HBM 中實例化完整的 $N \times N$ 矩陣。
反向傳播中的重計算（Recomputation in Backward Pass）： 避免為了反向傳播而儲存中間的 $N \times N$ 注意力矩陣。反之，僅儲存輸出 $O$ 與 softmax 統計量 $(ℓ, m)$ ，在反向傳播期間於 SRAM 內逐區塊重新計算 $S$ 和 $P$ 。
核心融合（Kernel Fusion）： 將矩陣乘法、遮罩（masking）、softmax 和 dropout 融合為單一 CUDA 核心（kernel），消除多餘的 HBM 來回傳輸開銷。

方法細節（含核心公式）

給定 $Q, K, V \in R^{N \times d}$ ，FlashAttention 設定區塊大小 $B_{c} = ⌈ M / (4 d)⌉$ 與 $B_{r} = min (⌈ M / (4 d)⌉, d)$ ，其中 $M$ 為 SRAM 容量。

增量 Softmax 追蹤

對於分解為子向量 $x^{(1)}$ 與 $x^{(2)}$ 的向量，其即時統計量更新公式如下： $m (x) = max (m (x^{(1)}), m (x^{(2)})),$ $ℓ (x) = e^{m (x^{(1)}) - m (x)} ℓ (x^{(1)}) + e^{m (x^{(2)}) - m (x)} ℓ (x^{(2)}) .$

針對 key/value 的第 $j$ 個區塊與 query 的第 $i$ 個區塊：

在 SRAM 上計算未正規化的分數： $S_{ij} = Q_{i} K_{j}^{⊤} \in R^{B_{r} \times B_{c}}$ 。
計算區塊局部統計量： $\tilde{m}_{ij} = rowmax (S_{ij})$ 、 $\tilde{P}_{ij} = exp (S_{ij} - \tilde{m}_{ij})$ ，以及 $\tilde{ℓ}_{ij} = rowsum (\tilde{P}_{ij})$ 。
更新即時全域正規化因子： $m_{i}^{new} = max (m_{i}, \tilde{m}_{ij}), ℓ_{i}^{new} = e^{m_{i} - m_{i}^{new}} ℓ_{i} + e^{\tilde{m}_{ij} - m_{i}^{new}} \tilde{ℓ}_{ij} .$
更新即時輸出區塊 $O_{i} \in R^{B_{r} \times d}$ ： $O_{i} \leftarrow diag (ℓ_{i}^{new})^{- 1} (diag (ℓ_{i}) e^{m_{i} - m_{i}^{new}} O_{i} + e^{\tilde{m}_{ij} - m_{i}^{new}} \tilde{P}_{ij} V_{j}) .$

複雜度分析

IO 複雜度： 標準注意力需要 $Θ (N d + N^{2})$ 次 HBM 存取。FlashAttention 僅需 $Θ (N^{2} d^{2} M^{- 1})$ 次 HBM 存取。作者證明，在所有 $M \in [d, N d]$ 的範圍內，沒有任何精確注意力演算法能在漸近意義上達到 $o (N^{2} d^{2} M^{- 1})$ 的存取量。
FLOP 與記憶體： FlashAttention 執行 $O (N^{2} d)$ 次 FLOP（由於反向傳播中的重計算，FLOP 略高），且除輸入與輸出外僅需 $O (N)$ 的額外記憶體空間。

擴展：區塊稀疏 FlashAttention（Block-Sparse FlashAttention）

若根據預定義的稀疏遮罩 $M \in {0, 1}^{N / B_{r} \times N / B_{c}}$ 僅計算非零區塊，可達到 $Θ (N d + N^{2} d^{2} M^{- 1} s)$ 的 IO 複雜度，其中 $s$ 為非零區塊的比例。

研究成果（關鍵數據）

訓練加速：
- BERT-large（序列長度 512）： 訓練速度比 MLPerf 1.1 的紀錄快 15%（在 8 顆 A100 上為 17.4 分鐘 vs. 20.0 分鐘）。
- GPT-2（序列長度 1K）： 比 HuggingFace 快高達 3.5 倍，比 Megatron-LM 快高達 1.8 倍。
- Long-Range Arena (LRA)： 相比標準注意力達到 2.4 倍加速；區塊稀疏 FlashAttention 達到 2.8 倍加速。
長上下文處理能力：
- 4K 上下文的 GPT-2： 訓練速度超越 1K 上下文的 Megatron-LM，同時驗證集困惑度改善 0.7 點（17.5 vs. 18.2）。
- 長文件分類： 將上下文從 512 擴展至 16K，在 MIMIC-III 上帶來 4.3 micro- $F_{1}$ 的提升（52.8 到 57.1），在 ECtHR 上達到高達 8.5 的增益（在 8K 長度下從 72.2 提升至 80.7）。
- Path-X 與 Path-256： FlashAttention 是首個解決 Path-X（序列長度 16K）的 Transformer，準確率達到 61.4%。區塊稀疏 FlashAttention 在 Path-256（序列長度 64K）上達到 63.1% 的準確率。

限制與開放研究問題

CUDA 工程實作開銷： 實作具備 IO 感知的核心需要撰寫客製化的底層 CUDA 程式碼，而非標準 PyTorch 或編譯器抽象層，這使得跨硬體架構的移植變得極具挑戰性。
多 GPU 注意力機制： 單節點上的最優 IO 排程無法直接套用到跨互連架構（NVLink、InfiniBand）的分散式多 GPU 注意力計算。
精確注意力下界： 雖然在一般的 $M$ 下已達到最優，但以 $M$ 為參數證明更細粒度的參數化複雜度下界仍是一個未解的開放問題。

關鍵術語

IO-awareness (IO 感知): 設計運算演算法時，明確考量硬體各記憶體層級之間（如高速晶片上 SRAM 與較慢的 HBM）讀取與寫入開銷的原則。
High Bandwidth Memory (HBM，高頻寬記憶體): GPU 的主要大型記憶體，具備高整體容量，但頻寬顯著低於晶片上的快取記憶體。
SRAM (靜態隨機存取記憶體): 現代 GPU 晶片上的高速暫存記憶體，位於每個串流多處理器（Streaming Multiprocessor, SM）內部。
Tiling (分塊技術): 透過將輸入矩陣切割為能放入 SRAM 的較小區塊，並在多次處理過程中間累加結果來重構計算流程的方法。
Selective gradient checkpointing (選擇性梯度檢查點): 在前向傳播期間避免儲存龐大的中間注意力矩陣，改為在反向傳播期間利用快取的 softmax 統計量於晶片上重新計算該矩陣的技術。
Kernel fusion (核心融合): 將多個獨立運算（矩陣乘法、遮罩、softmax、dropout）合併為單一 GPU 核心執行，以消除對 HBM 的重複讀寫。
Block-sparse FlashAttention (區塊稀疏 FlashAttention): FlashAttention 的擴展版本，僅計算由區塊稀疏遮罩所定義的非零區塊，從而依稀疏比例成正比地降低 IO 複雜度。
Path-X / Path-256: 來自 Long-Range Arena 的長序列視覺推理基準測試（序列長度分別為 16K 與 64K），先前的 Transformer 在該基準上均無法超越隨機猜測的準確率。

常見問題

FlashAttention 是否對注意力矩陣進行了近似計算？

沒有。FlashAttention 計算的是精確自注意力，在數值浮點精度範圍內，其輸出與標準注意力機制在數學上完全相同。

如果 FlashAttention 在反向傳播中計算了更多 FLOP，為什麼速度反而更快？

因為標準注意力機制的瓶頸在於記憶體頻寬（memory-bound）；在 HBM 與 SRAM 之間搬移 $N \times N$ 矩陣所消耗的實際時間，遠多於在高速 SRAM 上執行額外算術運算的時間。

FlashAttention 為反向傳播儲存了哪些輔助記憶體？

FlashAttention 僅儲存前向傳播的輸出 $O$ 、softmax 正規化因子 $ℓ$ 與 $m$ （僅需 $O (N)$ 的額外記憶體），以及用於在晶片上重現 dropout 遮罩的虛擬隨機數生成器狀態。

FlashAttention 與標準注意力在 IO 複雜度上有何差異？

標準注意力需要 $Θ (N d + N^{2})$ 次 HBM 存取，而 FlashAttention 僅需 $Θ (N^{2} d^{2} M^{- 1})$ 次存取，其中 $M$ 是 SRAM 的大小， $d$ 是注意力頭的維度。

本頁解讀由 AI 根據論文生成，可能有錯誤；細節請以原文為準。

用 XYZNOT 閱讀（全文翻譯與 AI 講解）