A modern GPU can do arithmetic far faster than it can fetch the numbers to work on, so most Transformer operations are limited by bytes moved between memory levels, not by floating-point work — and the levels differ by three orders of magnitude in both speed and size.
This is the fact the paper is built on, and it is worth looking at to scale before anything else. On-chip SRAM is roughly ten times faster than HBM and about two thousand times smaller. Any algorithm that keeps its working set in SRAM wins; any algorithm that spills to HBM pays.
A100 numbers as quoted in the paper. Bar length is bandwidth; box area is capacity.
The ratio that matters: SRAM is ~13× the bandwidth of HBM but holds 0.05% as much. So the design question for any GPU kernel is not "how few FLOPs?" but "how much work can I finish while the data is up there?"
Standard attention computes an N×N score matrix, writes it to HBM, reads it back to softmax it, writes it again, and reads it once more to multiply by V — so its memory traffic grows quadratically with sequence length even though the answer, an N×d output, does not.
Softmax, masking and dropout are all cheap arithmetic on an enormous array; the GPU spends its time carrying that array back and forth. Follow the round trips below — the numbers are the paper's measurements for one GPT-2-medium attention layer.
Left: the N×N grid standard attention actually builds. Right: the same grid, swept by one tile. The meters below are the paper's measured totals for one GPT-2-medium layer.
Sparse and low-rank attention methods cut the FLOP count to near linear, but most showed no wall-clock speedup and never displaced standard attention — because they attacked arithmetic while the machine was paying for memory traffic.
That sentence is the observation the whole paper turns on, and it is a methodological one rather than a technical one: the field had settled on a proxy metric. The same computation, re-scheduled so it never spills, turns out to beat approximations that do asymptotically less arithmetic — a result that is only surprising if you were counting FLOPs.
Compute the same attention output — bit-for-bit the same mathematics, no approximation — while never writing the N×N matrices to HBM at all, using memory linear in sequence length.
Two obstacles stand in the way, and the rest of the paper is those two obstacles being removed. First: softmax needs a whole row before it can normalise anything. Second: the backward pass wants the very matrix we refuse to store.
Split a matrix multiply into blocks and you can add the partial results at the end; that is why tiling is standard practice. Softmax breaks the pattern, because every output depends on a sum over the entire row — including the entries you have not loaded yet. You cannot normalise until you have seen everything, and seeing everything is exactly what we are trying to avoid.
One query, N keys. Each output is a ratio whose denominator runs across the whole row.
The denominator is the problem. If a later block contains a larger score than anything seen so far, both the max used for numerical stability and the sum change — and every weight computed so far was computed against the wrong ones.
Keep two running statistics per row — the maximum seen so far and the running sum of exponentials — and when a new block arrives with a larger maximum, rescale everything already accumulated by a single factor e^{m_old − m_new}; the result is not an approximation of the full softmax, it is identical to it.
This is the whole idea, and it is easier to trust once you have watched it happen. Below, one query row meets its keys four at a time. Every step prints its own arithmetic: the new block's scores, the new maximum, the correction factor applied to everything already accumulated, and the running output. The bar at the bottom tracks the distance between the running answer and the answer a full-row softmax would give.
N = 12 keys in blocks of 4, d = 4, real numbers throughout. This is Algorithm 1 lines 9–12 for a single row.
Notice what never exists: a row of 12 weights. At any moment the kernel holds one block of scores, two scalars, and the running output — that is the entire working set, and it does not grow with N. The correction factor is where the exactness comes from: rescaling by e^{m_old − m_new} is algebraically the same as having known the true maximum from the start.
Step to a block whose maximum beats everything before it, and watch the correction factor drop below 1. What would happen to the running sum if you skipped that rescale — and why does that make it a correctness bug rather than a precision one?
Scale that row up to the whole matrix and you have Algorithm 1: an outer loop over blocks of K and V, an inner loop over blocks of Q, and O, ℓ, m updated in place. The block sizes are not tuned by hand — they fall out of one constraint, that Q, K, V and the score tile must fit in SRAM at once.
Step through the beats. The panel on the right is the SRAM occupancy at that moment; everything else is sitting in HBM, untouched.
Algorithm 1 line 1: Bc = ⌈M / 4d⌉, Br = min(⌈M / 4d⌉, d). Move the two knobs that the hardware and the model actually give you.
The formula is just accounting: four arrays of one block each must live in SRAM simultaneously. Its consequence is the paper's IO complexity — a bigger SRAM means fewer, larger tiles, which means fewer passes over the data, which is why the traffic scales as N²d²/M rather than N².
The backward pass needs S and P — the very matrices we refused to store. FlashAttention keeps only the output O and the two statistics (m, ℓ), a few numbers per row, and rebuilds the score tiles on-chip when the gradient needs them. It does strictly more arithmetic than the standard backward pass and still finishes sooner, which is the paper's thesis stated as a single design decision.
Same gradient, same result — different bill.
This is also why the trick is not just gradient checkpointing. Checkpointing trades speed for memory; here the recomputation is the speedup, because the thing it avoids is an HBM round trip rather than a matmul.
Treat reads and writes between SRAM and HBM as the quantity to minimise — then the algorithm, the block sizes and even the decision to recompute all follow from one budget, and the paper can prove the resulting traffic is optimal up to constants.
Stated as complexity: standard attention moves Θ(Nd + N²) words; FlashAttention moves Θ(N²d²M⁻¹). With A100 numbers (d = 64, M = 192 KB) that ratio is roughly 9×, and the paper proves no exact attention algorithm can do asymptotically better across all SRAM sizes — so this is not a clever implementation, it is the floor.
Three layers: microbenchmarks of the attention op alone (runtime and memory vs sequence length), end-to-end training speed on BERT and GPT-2 against the MLPerf and Megatron records, and model-quality results that are only reachable because longer sequences became affordable.
Because the output is exact, quality can only change by way of what longer context enables — which is why the quality claims are all "same model, more sequence". The probe figures here are our own runs on one A100, measuring the two things a reader should not have to take on faith: that memory really is linear, and that the answer really is identical.
On GPT-2-medium attention the paper trades 13% more arithmetic for 9× less HBM traffic and comes out 5.7× faster, with memory linear in sequence length — and our own A100 probe reproduces the shape of both curves.
One attention layer of GPT-2 medium on an A100, three quantities.
Read it left to right: FLOPs went up, traffic went down 9×, and the clock followed the traffic. Any story in which FLOPs predict runtime has to explain this row away.
The claim that separates this paper from every approximate attention method. We checked it rather than trusting it.
Same models, same initialisation, same final quality — only the attention kernel changed.
The BERT number is against NVIDIA's own MLPerf 1.1 record, which is the honest comparison: not against a naive baseline, but against the best-tuned implementation of the day.
Against exact baselines (PyTorch, Megatron) FlashAttention simply wins — up to 3× on common lengths; against approximate methods it wins below about 1K tokens despite doing asymptotically more arithmetic, and beyond that the honest answer is that Linformer and friends eventually get faster.
That crossover is the paper being candid, and it is also the seam its own block-sparse variant is aimed at: apply the same IO-aware treatment to a sparse pattern and it outruns every approximate method the authors could find, at every length they tested.
It made sequence lengths that used to be impossible merely expensive — Path-X at 16K became solvable for the first time, GPT-2 could train at 4K faster than Megatron trained it at 1K, and within two years the kernel was the default path in every major framework.
The deeper effect is on how systems papers are written. "Count the bytes, not the FLOPs" is now the standard opening move for accelerating a memory-bound operation, and the same recipe — fuse, tile to SRAM, keep the reduction online, recompute rather than store — reappears in KV-cache kernels, in linear attention, and in the paper's own successors.
Every variant needs a hand-written CUDA kernel that does not transfer across GPU architectures, the technique is single-GPU only as published, and the approach helps exactly to the degree an operation is memory-bound — which attention at large batch or with a short sequence may not be.
The authors are explicit about the first one: they want a compiler that takes a high-level description and emits an IO-aware kernel, in the way Halide did for image processing. Until then, each new attention variant is a rewrite by a specialist.
The paper asked for IO-aware everything and for multi-GPU versions; what arrived was FlashAttention-2 (better work partitioning, ~2× again) and FlashAttention-3 (Hopper asynchrony and FP8), plus the paged-KV kernels that made long-context serving practical.
The through-line is that each successor optimises a different level of the same hierarchy: v1 fixed HBM traffic, v2 fixed how work was split across warps, v3 fixed overlap between memory movement and tensor cores. The IO-aware lens kept producing bottlenecks after the original one was gone.
The kernel is open source
(Dao-AILab/flash-attention) and, more
usefully, it is already inside PyTorch — scaled_dot_product_attention dispatches to
it, which is exactly how our probes call it.