FlashAttention
The same attention maths, done in tiles so the big score table never touches slow memory.
Read first:Self-attention
Step 1 of 6· Illustrative grid; A100 figures and GPT-2 measurements are from the paper, traffic counts follow its two algorithms, and the softmax maths is real
Attention compares every token with every other, producing a table of scores that grows with the square of the text length. A GPU has a large but slower main memory and a tiny, much faster memory on the chip. The standard way of computing attention writes the whole score table to slow memory, reads it back, then does the same with its softmaxed copy; that traffic, not the arithmetic, is what takes the time. FlashAttention splits the work into tiles small enough to stay on the chip, keeps a running softmax so the pieces add up exactly, and never stores the full table. Same answer, far less waiting.
Why it matters for your product
FlashAttention helped make long contexts practical: attention's extra memory now grows in step with the context instead of with its square, and it runs faster without changing a single answer. The broader lesson for any team tuning a model: on modern GPUs, moving data is often the real bottleneck, so check whether a slow step is waiting on memory before paying for more compute.
For engineersShow the maths
m_new = max(m, m̃), ℓ_new = e^(m − m_new) · ℓ + e^(m̃ − m_new) · ℓ̃
m and ℓ are the running maximum and sum of e^(score − m) for one row; m̃ and ℓ̃ are the same for the newly arrived tile. Rescale both to the new maximum and add.
Worked example: Row 1, 0.5, 2, 1.5 | 3, 0.2, 1, 2.5. After tile 1: m = 2, ℓ = 2.198. Tile 2 has m̃ = 3 and ℓ̃ = 1.803, so ℓ = 2.198 × e⁻¹ + 1.803 = 2.611, exactly the full-row sum.
HBM accesses: standard Θ(N·d + N²), FlashAttention Θ(N²·d² / M)
N is the sequence length, d the head size and M the size of SRAM. For typical d of 64 to 128 and M around 100 KB, d² is many times smaller than M, so FlashAttention touches HBM far less.
Worked example: N = 4,096: the score table alone is N² ≈ 16.8 million numbers, 33.6 MB at 2 bytes, for every head in every layer.
FlashAttention makes exact attention faster and leaner but still does N² work, so doubling the context still quadruples the compute. Its gains depend on kernels tuned for specific GPUs: the first version reached only 25 to 40% of an A100's peak, which FlashAttention-2 raised to 50 to 73%.
The words you will hear
HBM
Primary sources
- 01FlashAttention: Fast and Memory-Efficient Exact Attention with IO-AwarenessDao et al. · 2022
- 02FlashAttention-2: Faster Attention with Better Parallelism and Work PartitioningDao · 2023
- 03Online normalizer calculation for softmaxMilakov, Gimelshein · 2018
Connected ideas
Inside the transformer
Self-attention
How every word looks at every other word to work out what it means here.
ExploreGenerating text
KV cache
Why a model keeps the keys and values of earlier tokens instead of recomputing them, and what that memory costs.
ExploreGenerating text
Prefill and decode
Why the first word of an answer takes a moment and the rest stream out at a steady pace.
ExploreInside the transformer
Beyond attention
Why attention gets expensive on long text, and the designs that trade exact recall for linear cost.
ExploreWant this working on your data?
We design and build the systems these ideas power: retrieval, agents, voice and the models behind them. Start with a free discovery call.