FlashAttention: Fast and Memory-Efficient Exact Attention

Paper revolusioner Tri Dao et al. (Stanford, 2022) yang mengatasi bottleneck memori terbesar dalam arsitektur Transformer.

Masalah Utama Self-Attention Standar

Attention standar menghitung: $S = Q K^T \in \mathbb{R}^{N \times N}$ $P = \text{softmax}(S) \in \mathbb{R}^{N \times N}$ $O = P V \in \mathbb{R}^{N \times d}$

Matriks atensi $S$ dan $P$ berukuran $N \times N$ (kuadratik terhadap panjang konteks). Pada context 32k token, matriks intermediate ini memakan memori VRAM puluhan Gigabyte dan membebani transfer I/O HBM.

Tiga Inovasi FlashAttention

1. Tiling (Block-based Computation): Membagi matriks $Q, K, V$ menjadi blok-blok kecil yang pas di dalam SRAM/Shared Memory (SRAM GPU berkapasitas ~100–200 KB per SM).
2. Online Softmax: Menghitung softmax secara inkremental per-blok tanpa perlu mengalokasikan seluruh matriks $N \times N$ di VRAM.
3. Recomputasi Cerdas: Pada fase backward pass training, atensi dihitung ulang dari blok $Q, K, V$ di SRAM alih-alih menyimpan matriks atensi $N \times N$ di VRAM.

Evolusi FlashAttention

• FlashAttention-1 (2022): 2x–4x lebih cepat dibanding attention standar, memori O(N) linear.
• FlashAttention-2 (2023): Optimasi pemisahan kerja antar-Warp dan pereduksian sinkronisasi non-matriks, mendekati 73% peak theoretical TFLOPS.
• FlashAttention-3 (2024): Eksploitasi Tensor Cores Asinkron dan komputasi presisi FP8 di arsitektur NVIDIA Hopper (H100).