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

Panduan membaca

Masalah dan Motivasi

Transformer memiliki kompleksitas waktu dan memori kuadratik terhadap panjang sekuens $N$. Meskipun metode approximate attention mampu memangkas jumlah FLOP teoretis, metode tersebut sering kali gagal memberikan percepatan wall-clock karena komputasi GPU modern utamanya terhambat oleh transfer memori (IO), bukan operasi aritmetika.

Pada attention standar, komputasi $\mathbf{S}=\mathbf{Q}{\mathbf{K}}^{\top }$, $\mathbf{P}=\mathrm{softmax}(\mathbf{S})$, dan $\mathbf{O}=\mathbf{PV}$ menulis matriks perantara berukuran $N\times N$ ke High Bandwidth Memory (HBM) yang lambat lalu membacanya kembali berulang kali. Hal ini menimbulkan $\Theta (Nd+N^2)$ akses HBM dan jejak memori $O(N^2)$, yang membatasi panjang context window secara praktis.

Ide Utama

• IO-Awareness: Memperhitungkan hierarki memori GPU secara eksplisit: SRAM on-chip yang cepat namun berukuran kecil (misalnya, 192 KB/SM dengan bandwidth ~19 TB/s pada A100) versus HBM yang besar namun lambat (40–80 GB dengan bandwidth 1,5–2,0 TB/s).
• Tiling (Komputasi Softmax Per Blok): Membagi input $\mathbf{Q},\mathbf{K},\mathbf{V}$ ke dalam blok-blok yang muat seluruhnya di SRAM. Dengan mencatat statistik running row-max dan normalisasi softmax, reduksi softmax dihitung secara bertahap antar-blok tanpa mematerialisasi matriks penuh $N\times N$ di HBM.
• Komputasi Ulang pada Backward Pass: Menghindari penyimpanan matriks perhatian perantara $N\times N$ untuk backpropagation. Sebagai gantinya, hanya output $\mathbf{O}$ dan statistik softmax $(\ell ,m)$ yang disimpan, lalu $\mathbf{S}$ dan $\mathbf{P}$ dihitung ulang secara on-chip blok demi blok di dalam SRAM selama backward pass.
• Kernel Fusion: Memfusikan operasi perkalian matriks, masking, softmax, dan dropout ke dalam satu kernel CUDA tunggal untuk mengeliminasi transfer bolak-balik yang redundan ke HBM.

Metode (beserta Persamaan Utama)

Untuk $\mathbf{Q},\mathbf{K},\mathbf{V}\in {\mathbb{R}}^{N\times d}$, FlashAttention menentukan ukuran blok $B_c=\lceil M/(4d)\rceil$ dan $B_r=\min (\lceil M/(4d)\rceil ,d)$, di mana $M$ adalah kapasitas SRAM.

Pelacakan Softmax Inkremental

Untuk vektor yang dipecah menjadi sub-vektor $x^{(1)}$ dan $x^{(2)}$, pembaruan statistik berjalannya adalah: $$m(x)=\max (m(x^{(1)}),m(x^{(2)})),$$ $$\ell (x)=e^{m(x^{(1)})-m(x)}\ell (x^{(1)})+e^{m(x^{(2)})-m(x)}\ell (x^{(2)}).$$

Untuk blok keys/values ke-$j$ dan blok queries ke-$i$:

1. Hitung skor tanpa normalisasi pada SRAM: ${\mathbf{S}}_{ij}={\mathbf{Q}}_i{\mathbf{K}}_j^{\top }\in {\mathbb{R}}^{B_r\times B_c}$.
2. Hitung statistik lokal blok: ${\tilde{m}}_{ij}=\mathrm{rowmax}({\mathbf{S}}_{ij})$, ${\tilde{\mathbf{P}}}_{ij}=\exp ({\mathbf{S}}_{ij}-{\tilde{m}}_{ij})$, dan ${\tilde{\ell }}_{ij}=\mathrm{rowsum}({\tilde{\mathbf{P}}}_{ij})$.
3. Perbarui normalisasi global berjalan: $$m_i^{\mathrm{new}}=\max (m_i,{\tilde{m}}_{ij}),{\ell }_i^{\mathrm{new}}=e^{m_i-m_i^{\mathrm{new}}}{\ell }_i+e^{{\tilde{m}}_{ij}-m_i^{\mathrm{new}}}{\tilde{\ell }}_{ij}.$$
4. Perbarui blok output berjalan ${\mathbf{O}}_i\in {\mathbb{R}}^{B_r\times d}$: $${\mathbf{O}}_i\leftarrow \mathrm{diag}({\ell }_i^{\mathrm{new}}{)}^{-1}\left(\mathrm{diag}({\ell }_i)e^{m_i-m_i^{\mathrm{new}}}{\mathbf{O}}_i+e^{{\tilde{m}}_{ij}-m_i^{\mathrm{new}}}{\tilde{\mathbf{P}}}_{ij}{\mathbf{V}}_j\right).$$

Analisis Kompleksitas

• Kompleksitas IO: Attention standar membutuhkan $\Theta (Nd+N^2)$ akses HBM. FlashAttention membutuhkan $\Theta (N^2{d}^2{M}^{-1})$ akses HBM. Penulis membuktikan bahwa tidak ada algoritma exact attention yang secara asimtotik dapat mencapai $o(N^2{d}^2{M}^{-1})$ akses untuk semua $M\in [d,Nd]$.
• FLOPs dan Memori: FlashAttention melakukan $O(N^{2}d)$ FLOPs (dengan sedikit tambahan FLOPs pada backward pass akibat komputasi ulang) dan hanya memerlukan memori tambahan $O(N)$ di luar input dan output.

Ekstensi: Block-Sparse FlashAttention

Membatasi evaluasi hanya pada blok-blok bukan-nol berdasarkan mask sparsity yang telah ditentukan $\mathbf{M}\in \{0,1{\}}^{N/B_r\times N/B_c}$ menghasilkan kompleksitas IO sebesar $\Theta (Nd+N^2{d}^2{M}^{-1}s)$, di mana $s$ merepresentasikan proporsi blok bukan-nol.

Hasil (beserta Angka Terpenting)

• Percepatan Pelatihan:
  • BERT-large (panjang sekuens 512): Melampaui rekor kecepatan pelatihan MLPerf 1.1 sebesar 15% (17,4 menit vs. 20,0 menit pada $8\times$ A100).
  • GPT-2 (panjang sekuens 1K): Meraih percepatan hingga 3,5x dibanding HuggingFace dan hingga 1,8x dibanding Megatron-LM.
  • Long-Range Arena (LRA): Memberikan percepatan 2,4x dibanding attention standar; Block-sparse FlashAttention meraih percepatan 2,8x.
• Kapabilitas Konteks Lebih Panjang:
  • GPT-2 dengan konteks 4K: Berlatih lebih cepat daripada Megatron-LM pada konteks 1K sembari meningkatkan perplexity validasi sebesar 0,7 poin (17,5 vs. 18,2).
  • Klasifikasi Dokumen Panjang: Penskalaan konteks dari 512 ke 16K meningkatkan micro-$F_1$ sebesar 4,3 poin pada MIMIC-III (52,8 ke 57,1) dan hingga 8,5 poin pada ECtHR (72,2 ke 80,7 pada 8K).
  • Path-X dan Path-256: FlashAttention merupakan Transformer pertama yang mampu menyelesaikan Path-X (panjang sekuens 16K) dengan akurasi 61,4%. Block-sparse FlashAttention meraih akurasi 63,1% pada Path-256 (panjang sekuens 64K).

Keterbatasan dan Pertanyaan Terbuka

• Beban Rekayasa CUDA: Mengimplementasikan kernel IO-aware membutuhkan pemrograman CUDA tingkat rendah (low-level) khusus daripada sekadar menggunakan abstraksi PyTorch standar atau compiler, sehingga portabilitas antar-arsitektur perangkat keras menjadi tantangan tersendiri.
• Multi-GPU Attention: Jadwal IO optimal untuk satu node tidak langsung otomatis dapat diterapkan pada distributed multi-GPU attention melalui interkoneksi (NVLink, InfiniBand).
• Batas Bawah Exact Attention: Meskipun terbukti optimal untuk rentang $M$ secara umum, pembuktian batas bawah kompleksitas terparameterisasi (parameterized complexity) yang detail terhadap $M$ masih menjadi masalah terbuka.