Loading...
Back to Concepts
Efficient Attention2022
FlashAttention
Fast and Memory-Efficient Exact Attention with IO-Awareness
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré
Read the Paper on arXivPaper Overview
FlashAttention (Dao, Fu, Ermon, Rudra & Ré — Stanford and SUNY Buffalo, NeurIPS 2022) is best understood as a correction to what the field had been optimizing. Self-attention is quadratic in sequence length, so for five years the response was to make it cheaper in FLOPs — sparse approximations, low-rank approximations, and combinations of the two. Many of these reduced compute to linear or near-linear in . Many of them were also, on real hardware, no faster.
The paper's diagnosis is that FLOPs were the wrong currency. On an A100, on-chip SRAM runs at roughly TB/s while HBM runs at – TB/s, and compute throughput has outpaced memory throughput for years. Most operations in a Transformer are therefore bottlenecked by memory accesses, not arithmetic. An algorithm that cuts FLOPs while still writing an matrix to HBM has optimized the term that was not binding.
So FlashAttention changes the objective to IO-awareness: count the reads and writes between HBM and SRAM, and minimize those. The mechanism is tiling — process attention in blocks small enough to live in SRAM — which runs immediately into the obstacle that softmax couples every column in a row, so the normalizer is not known until the whole row is seen. The resolution is an online softmax that carries two running statistics and rescales as it goes, producing the exact same output as standard attention rather than an approximation of it.
The payoff is unusual in that nothing is traded away. Attention computation itself runs 7.6x faster, memory footprint drops from quadratic to linear (up to 20x less), and the backward pass gets faster despite doing more FLOPs. BERT-large trains 15% faster than the MLPerf 1.1 record; GPT-2 up to 3x faster than HuggingFace. And because 64K-token sequences now fit in memory, block-sparse FlashAttention produced the first sequence model to beat chance on Path-256.
Chapter Roadmap
Click any topic to jump in
1
The GPU Memory Hierarchy
SRAM is ~10x faster than HBM and ~4000x smaller. Compute has outpaced memory, so attention is bottlenecked by data movement, not arithmetic.
measured against
2
What Standard Attention Costs
Three kernels, each round-tripping an N×N matrix through HBM. The FLOPs are fine; the traffic is the problem.
restructured by
3
Tiling and the Online Softmax
Softmax couples the whole row, so it resists tiling. Carrying two running statistics (m, ℓ) makes it decomposable — exactly, not approximately.
made trainable by
4
Recomputation and Kernel Fusion
Throw away the N×N matrices and rebuild them in SRAM during the backward pass. More FLOPs, less traffic, net faster.
analyzed as
Exact bound and approximate extension
5
IO Complexity and Optimality
Θ(N²d²M⁻¹) HBM accesses versus Θ(Nd + N²), with a matching lower bound proving no exact algorithm does asymptotically better.
6
Block-Sparse FlashAttention
Skip the masked-out blocks entirely and the IO cost scales directly with the sparsity — the one part of the paper that is approximate.
validated by
7
Results
15% over the MLPerf BERT record, 3x on GPT-2, 20x less attention memory — and the first Transformer to beat chance on Path-X.
Every claim in this paper rests on one hardware fact, so it is worth stating precisely. An A100 GPU has 40–80GB of HBM with bandwidth 1.5–2.0 TB/s, and 192KB of on-chip SRAM per streaming multiprocessor across 108 SMs — roughly 20MB in total — at an estimated 19 TB/s. Behind both sits CPU DRAM: over a terabyte, at 12.8 GB/s.
The shape of that hierarchy is the whole story. SRAM is an order of magnitude faster than HBM but many orders of magnitude smaller. Each rung down buys capacity and pays for it in bandwidth, and the ratios are steep enough that where your data sits matters more than how much arithmetic you do to it.
GPU kernels have a fixed rhythm: load inputs from HBM into registers and SRAM, compute, write outputs back to HBM. An operation's speed therefore depends on its arithmetic intensity — work done per byte moved. Matrix multiplication has high intensity and is compute-bound, which is why it gets the attention and the tensor cores. But softmax, dropout, and masking are elementwise: one or two operations per element loaded. They are memory-bound, and on modern GPUs, where compute speed has out-paced memory speed, they dominate runtime despite contributing almost nothing to the FLOP count.
This is the gap the paper walks through. Approximate-attention methods were evaluated in FLOPs, cut FLOPs substantially, and many of them do not display wall-clock speedup against standard attention and have not gained wide adoption. The measure had come apart from the thing being measured. IO-awareness — explicitly accounting for reads and writes between levels of memory — is a well-established principle in database joins, image processing, and numerical linear algebra; the contribution here is bringing it to attention.
Key Points
1
A100 HBM: 40–80GB at 1.5–2.0 TB/s. A100 SRAM: 192KB per SM across 108 SMs (~20MB total) at an estimated 19 TB/s. CPU DRAM: >1TB at 12.8 GB/s
2
SRAM is an order of magnitude faster than HBM but many orders of magnitude smaller — the trade that makes tiling both necessary and possible
3
Every kernel follows the same pattern: load from HBM → compute in SRAM/registers → write back to HBM. Data movement is the tax on every step
4
Compute speed has out-paced memory speed, so most operations in Transformers are bottlenecked by memory accesses rather than arithmetic
5
Matmul is compute-bound; softmax, dropout and masking are memory-bound. Matmul dominates the FLOP count while the elementwise operations dominate the runtime
6
The paper's core critique: approximate attention methods reduce compute to linear or near-linear in yet often show no wall-clock speedup, because they optimize FLOPs and ignore IO
7
IO-awareness is not new — it is standard in database joins, image processing and numerical linear algebra. The contribution is applying it to attention