FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré
cs.LG
2022-05-28
FlashAttention runs exact attention in SRAM tiles so the N×N matrix never hits HBM, cutting traffic about 9×. GPT-2 small training falls from 9.5 days to 2.7 days.
Self-attention spends time and memory quadratic in sequence length. Approximate attention pulls the FLOPs toward linear, and the wall-clock time often barely moves. On current GPUs many of those ops are limited by memory traffic. An A100's HBM runs at about 1.5-2 TB/s. On-chip SRAM is near 19 TB/s, but each streaming multiprocessor has only about 192KB.
The usual kernel writes S = QK^T to HBM, reads it back for softmax, writes P, then reads P with V for the second multiply. S and P are N by N. A typical GPT-2 head uses N = 1024 and d = 64, so that matrix dwarfs Q, K, and V. Masking and dropout touch it again.
FlashAttention returns the same softmax(QK^T)V as standard attention. What changes is the traffic between HBM and SRAM. Q, K, and V are split into blocks that fit on chip. The outer loop loads one block of K and V. The inner loop loads a block of Q, then multiplies, runs softmax, applies an optional mask and dropout, and multiplies by V, all in SRAM. Only the output block is written back. The N by N matrices are never materialized.
Softmax couples every key in a row, so separate blocks cannot be normalized and added as they stand. Each row keeps a running max m and a sum of exponentials ℓ. When the next block arrives, the gap between the old max and the new max rescales the previous output and the previous sum before the new block is folded in. After the last block, the result matches one full softmax.
Storing P for the backward pass would make memory quadratic again. The forward pass keeps the output O plus m and ℓ. Backward recomputes the attention block from Q, K, and V in SRAM, and regenerates the dropout mask from the saved RNG state instead of storing an N by N mask. FLOPs go up. HBM accesses go down. Forward and backward each run as one CUDA kernel.
Standard attention does Θ(Nd + N²) HBM accesses. FlashAttention does Θ(N²d²/M), where M is the SRAM size. For d of 64-128 and M around 100KB, d² is much smaller than M, and traffic drops by about an order of magnitude. The block-sparse variant skips all-zero blocks and multiplies the leading term by the nonzero-block fraction s. The experiments use a fixed butterfly pattern.
The speedup is fewer bytes moved, not fewer FLOPs. On GPT-2 medium (sequence 1024, head dim 64, 16 heads, batch 64), a forward-plus-backward pass on an A100 costs standard attention 66.6 GFLOPs, 40.3 GB of HBM traffic, and 41.7 ms. FlashAttention costs 75.2 GFLOPs, 4.4 GB, and 7.3 ms. Arithmetic is about 13% higher, traffic is about 9× lower, and the pass is about 5.7× faster. On the attention op alone versus PyTorch, Figure 1 reports up to 7.6×.
BERT and GPT-2 end-to-end times are measured on 8 A100s.
| Setup | Baseline | FlashAttention |
| BERT-large, seq 512, MLPerf 1.1 target 72.0% | Nvidia 20.0±1.5 min | 17.4±1.4 min, 15% faster |
| GPT-2 small, OpenWebText | HuggingFace 9.5 days, ppl 18.2 | 2.7 days, ppl 18.2, about 3.5× |
| GPT-2 medium | HuggingFace 21.0 days, ppl 14.2 | 6.9 days, ppl 14.3, about 3.0× |
| Versus Megatron-LM | small 4.7 days, medium 11.5 days | 2.7 and 6.9 days, about 1.7× |
| LRA, seq 1K-4K | standard Transformer, avg 59.3 | avg 59.8, 2.4× |
Quality moves when the context grows. GPT-2 small at context 4K reaches perplexity 17.5, which is 0.7 better than 18.2 at 1K, and finishes in 3.6 days, about 30% less time than Megatron at 1K (4.7 days). On long-document micro-F1, MIMIC-III goes from 52.8 at length 512 to 57.1 at 16K (+4.3). ECtHR goes from 72.2 at 512 to 80.7 at 8K (+8.5), then slips to 79.2 at 16K. The two gains average 6.4 points.
Path-X feeds a 128×128 black-and-white image one pixel at a time (sequence 16K) and asks whether two points are connected. Earlier Transformers either ran out of memory or stayed at chance. FlashAttention pretrains on Path-64, interpolates the positional embeddings, and scores 61.4%. The block-sparse variant reaches sequence 64K, with 63.1% on Path-256 and 56.0% on Path-X.
On one 40GB A100 with dropout and a padding mask, FlashAttention is up to 3× faster than PyTorch attention from sequence 128 to 2K. Memory grows linearly with N and is up to about 20× smaller than exact-attention baselines. Past roughly 512 to 1024 tokens, approximate methods such as Linformer start to win on wall-clock time. At 64K every baseline except Linformer is out of memory, and FlashAttention still uses about half of Linformer's memory. On LRA, block-sparse FlashAttention is 2.8× faster than standard attention at an average score of 59.6, close to dense attention.
At a fixed context length the math is unchanged. GPT-2 small lands at perplexity 18.2, the same figure as HuggingFace and Megatron. Medium lands at 14.3, matching Megatron. The memory and time that are saved get spent on longer context, so the quality change comes from reading more text.
Short sequences benefit too. BERT-large at length 512 still beats the MLPerf 1.1 single-node speed record by 15%. Block-sparse FlashAttention is another 2-4× faster than the dense kernel and reaches sequence 64K. Sparse attention was often missing a memory schedule.
Every new attention variant needs its own CUDA kernel. That sits below PyTorch, the engineering cost is high, and the kernel may not carry over to another GPU architecture. The IO argument stops at one GPU. Transfers between GPUs are another level of IO and sit outside the bound.
Calling the algorithm optimal for a range of SRAM sizes is stronger than the proof. The theorem shows that no exact algorithm improves the HBM bound by an asymptotic factor for every SRAM size M in [d, Nd]. The argument only uses M = Θ(Nd), where the bound collapses to reading the input at least once. A lower bound for one specific M is left open.
Exact FlashAttention still scales roughly with N². Past sequence length 1K, some approximate methods are already faster on the clock. Block sparsity needs a mask aligned to blocks. The butterfly pattern is fixed, not learned. LRA scores depend strongly on tuning, and the reproduced baselines beat the original leaderboard, so absolute scores should not be compared across papers. Path-256 at 63.1% beats Path-X at 61.4%, but Path-256 has relatively shorter paths, so the higher score is an easier task.
Rabe and Staats also tile attention so they do not store the N by N matrix, but they target peak memory. In the appendix their code matches standard attention or runs slightly slower. FlashAttention keeps a single running output and recomputes only the attention block on the backward pass, which cuts HBM traffic and runtime together.