FlashAttention computes exact attention faster and with memory that grows linearly in the sequence length; it trained GPT-2 with 1K-token sequences 3× faster and BERT-large 15% faster end to end than the baselines it compared against (Dao et al., 2022).
A newer notebook covers this concept. This older page keeps its summary, math and sources.
How to study this pagePredict · Check · Generalize · Continue
PredictSay what you expect.
Before you use the demo in the newer notebook, say what you expect to change.
CheckCheck against the math and the papers.
Compare what the notebook demo shows with the key equation here and with the papers cited below.
GeneralizeName what stays the same.
Identify the part of the idea that survives a change of notation or model, so you can recognize it in another paper or system.
ContinueContinue in the newer notebook.
The newer notebook covers the same concept and has an interactive demo.
Why It Matters for Modern Models
FlashAttention computes exact attention faster and with memory that grows linearly in the sequence length; it trained GPT-2 with 1K-token sequences 3× faster and BERT-large 15% faster end to end than the baselines it compared against (Dao et al., 2022).
Because memory no longer grows as N², longer contexts became practical: FlashAttention gave the first Transformers to do better than chance on the Path-X benchmark (16K tokens), and its block-sparse version did the same on Path-256 (64K tokens) (Dao et al., 2022).
Published training recipes use memory-efficient exact attention of this kind; LLaMA, for example, used an implementation inspired by Rabe & Staats (2021) with the backward pass of Dao et al. (2022) (Touvron et al., 2023).
Questions to Think About
Why can an algorithm that performs more arithmetic, by recomputing attention in the backward pass, still be faster than one that stores and rereads the attention matrix?
How can softmax be computed in one pass over blocks by keeping a running maximum and a running sum, and why must earlier partial outputs be rescaled when the maximum changes?
Interactive Demo
The interactive demo for this concept is in the newer notebook.
If you prefer intuition first, try the demo in the newer notebook, then come back for the math.
Key Equation
extra memory: O(N) instead of O(N2)
Standard attention over a sequence of N tokens computes an N×N matrix of scores, writes it to the GPU's main memory (HBM), and reads it back for the softmax and the product with V:
S=dQK⊤,P=softmax(S),O=PV
Moving S and P in and out of HBM costs memory traffic that grows as N2, and on modern GPUs most of the operations in a transformer are limited by such memory traffic rather than by arithmetic (Dao et al., 2022).
FlashAttention computes exactly the same output without ever storing S or P:
split Q, K and V into blocks small enough to fit in the fast on-chip memory (SRAM);
for each block of queries, loop over the blocks of keys and values, computing scores and a partial output inside SRAM;
combine the partial results with the online softmax: keep a running maximum and a running sum of exponentials for each query, and rescale the partial output whenever the maximum increases (Milakov & Gimelshein, 2018).
In the backward pass it recomputes the blocks instead of reading a stored P, so it performs more arithmetic than standard attention yet runs faster, because it moves far less data. The extra memory it needs grows as O(N) (the running statistics) instead of O(N2). Dao (2024) summarizes FlashAttention's speed-up as 2–4× over optimized baselines, and FlashAttention-2 is about twice as fast again.
Key Papers
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness