Paper
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré
arXiv:2205.14135 · 0▲ · cs.LG
View on arXiv →Premium readers get an interactive explainer for this paper — a figure you can poke at, not just read.
What it is
FlashAttention is an exact attention algorithm that restructures the computation to avoid writing the large N x N attention matrix to GPU high bandwidth memory (HBM). It uses tiling to compute softmax block by block in fast on-chip SRAM and recomputes the attention matrix during the backward pass instead of storing it, all fused into a single CUDA kernel.
Why it matters
Attention was memory-bound, not compute-bound, so most of the runtime was spent shuffling data between HBM and SRAM rather than doing math. This makes attention faster and cuts memory from quadratic to linear in sequence length, which lets you train Transformers on much longer contexts (up to 64K) without running out of memory or wall-clock budget.
Practical takeaway
You can drop in FlashAttention as an exact replacement for standard attention (same numerical results, same training curves) to get 2-4x speedups and roughly 10-20x less attention memory. It has since become the default attention kernel in PyTorch and most training/inference stacks, so if you use a recent framework you are likely already benefiting from it.
Key result
7.6x speedup on the attention computation itself for GPT-2, driven by cutting HBM reads/writes from 40.3 GB to 4.4 GB (even though FLOPs rose from 66.6 to 75.2 GFLOPs due to recomputation); end-to-end gains are more modest, for example 15% faster BERT-large training vs the MLPerf 1.1 record and up to 3x vs HuggingFace GPT-2.
Subscribe
Get the next issue.
Free. One email a week. Unsubscribe any time: no account, no dark patterns.