FlashAttention and the Kernel Era

It is May 2022. The Transformer has been the default architecture for sequence modeling for five years, and its self-attention layer still costs memory and compute that grow with the square of sequence length, on accelerators whose arithmetic throughput has been climbing faster than the memory bandwidth that feeds it. A team at Stanford publishes a kernel that computes exactly the same attention output as before — no approximation, no dropped terms — and reports it running several times faster by changing nothing but the order in which numbers move between a GPU’s memory tiers. Two follow-on papers, one a year later and one two years later, repeat the trick on newer hardware, and by the middle of the decade “run FlashAttention” is close to a synonym for “run attention” at all.

Where We Were

By 2022, the Transformer covered in Language Models, Chapter 5 was the settled architecture behind essentially every frontier language model, and its self-attention layer had a well-known structural cost: computing attention over a sequence of length $N$ requires an $N \times N$ matrix of pairwise scores, so both the compute and the memory needed grow quadratically in $N$. Language Models, Chapter 4 covers where the attention mechanism itself came from; this chapter starts from the assumption that the mechanism is fixed and asks a narrower question: given that formula, how do you actually run it fast on real hardware.

The standard implementation of that period made the quadratic cost worse than the arithmetic alone required. A naive attention kernel computes the full $N \times N$ score matrix, writes it out to the accelerator’s main memory, reads it back to apply softmax, writes the result out again, reads it back once more to multiply by the value matrix, and so on — each of those intermediate $N \times N$ matrices round-tripping through the slowest tier of the memory hierarchy. Chapter 4 covers that hierarchy and the roofline model that classifies a kernel as compute-bound or memory-bound depending on how much arithmetic it does per byte moved; standard attention, because it re-reads and re-writes matrices whose size grows with $N^2$ while doing comparatively little arithmetic per element, sits on the memory-bound side of that line on every accelerator built since GPUs gained the fast, small on-chip memory described there. The gap this chapter addresses had already widened by the time it needed addressing: Chapter 4 shows peak compute throughput on NVIDIA accelerators growing faster than HBM bandwidth across every generation from Volta to Blackwell, which means an algorithm that does not deliberately manage its memory traffic loses more ground with every new chip, not less.

Two lines of prior work had already tried to escape the quadratic cost by changing the algorithm rather than the implementation. Sparse-attention variants restrict which token pairs get scored at all, and linear-attention methods reformulate attention as a recurrence with cost linear in $N$ — the latter covered in full in Chapter 18. Both routes reduce the FLOP count, and both routes are approximations: they change what attention computes, not merely how it is computed, and by the early 2020s neither had displaced exact softmax attention as the default in frontier models, in part because approximation quality varied by task in ways practitioners found hard to predict. The gap that remained open was a kernel that computed the same attention output, to numerical precision, while spending less time waiting on memory — and by 2022, with the NVIDIA A100 (Chapter 20, May 2020) two years into deployment and sequence lengths climbing in pursuit of longer documents and codebases, that gap had real cost attached to it.

The Breakthrough

FlashAttention (Dao et al., 2022) reframed the problem the same way the roofline model in Chapter 4 frames every kernel-design question: not “how do we do less arithmetic” but “how do we move fewer bytes.” The algorithm never materializes the full $N \times N$ score or probability matrix in the accelerator’s high-bandwidth memory (HBM) at all. Instead, it splits the query, key, and value matrices into blocks small enough to fit in the GPU’s fast on-chip SRAM, loads one block of each into SRAM, computes the partial attention output for that block, and folds the result into a running output using an incrementally updated softmax — detailed under The Mathematics below — before moving to the next block. Only the final output, of size $N \times d$ rather than $N \times N$, is written back to HBM. On the backward pass, rather than storing the large intermediate matrices computed on the forward pass, FlashAttention recomputes them on demand from the blocks already in SRAM — an application of the same trade-off, recompute instead of store, that gradient checkpointing (Chapter 10) makes across whole layers rather than within one.

The result was an exact attention computation — the same output, to floating-point precision, as the naive implementation — that used $O(N)$ memory instead of $O(N^2)$ and moved far fewer bytes through HBM per unit of arithmetic performed. Measured end to end, FlashAttention trained BERT-large 15% faster than the fastest known MLPerf 1.1 training implementation at the time, trained GPT-2 up to 3x faster than a well-optimized HuggingFace baseline, and ran 2.4x faster on Long-Range Arena benchmarks. Because the memory savings let models train on longer sequences than had previously fit, quality moved too: models trained with FlashAttention reached 0.7 better perplexity on GPT-2 and 6.4 accuracy points better on a long-document classification task than shorter-context baselines, and a block-sparse extension of the same kernel produced the first better-than-chance results on the Path-X (sequence length 16K, 61.4% accuracy) and Path-256 (sequence length 64K, 63.1% accuracy) long-range benchmarks.

FlashAttention-2 (Dao, 2023) kept the same no-materialization, no-approximation strategy and rewrote the work distribution inside it. The original kernel left some of a GPU’s parallel thread blocks under-used for certain problem shapes and split work between warps (small groups of threads) in a way that required extra reads and writes to shared memory; FlashAttention-2 parallelizes across the sequence-length dimension in addition to batch and attention-head dimensions, and repartitions work between warps to cut those extra shared-memory round trips. Dao reports FlashAttention-2 reaching 50%–73% of an A100’s theoretical peak FLOPs/s, against only 25%–40% for the original kernel, and an end-to-end GPT-style training throughput of up to 225 TFLOPs/s per A100 — 72% model FLOPs utilization — roughly double FlashAttention’s own training speed.

FlashAttention-3 (Shah et al., 2024) targeted a specific hardware generation: the NVIDIA H100 (Chapter 20, announced March 2022), whose Hopper architecture added asynchronous data-movement units and low-precision matrix-multiply hardware that the first two FlashAttention kernels did not exploit. FlashAttention-3 overlaps computation and data movement through warp specialization — dedicating some warps to memory transfers and others to arithmetic so neither has to wait idle on the other — and interleaves the block-wise matrix multiplications with the block-wise softmax computation so the two run concurrently rather than in sequence. It also adds an FP8 low-precision path with block-level quantization and an incoherent-processing technique to control the additional rounding error that low precision would otherwise introduce. FlashAttention-2, run unmodified on an H100, had reached only about 35% of the chip’s FP16 peak; FlashAttention-3 reached 740 TFLOPs/s in FP16 (75% utilization) and close to 1.2 PFLOPs/s in FP8, a 1.5x–2.0x speedup over FlashAttention-2 on the same hardware, while measuring 2.6x lower numerical error than a baseline FP8 attention implementation.

The Mathematics

Scaled dot-product attention, in the notation Language Models, Chapter 5 uses to derive it, is:

$$O = \mathrm{softmax}\!\left(\frac{QK^{\top}}{\sqrt{d}}\right)V$$

Symbol	Meaning
$N$	sequence length, number of tokens
$d$	head dimension, columns per query/key/value row
$Q, K, V$	query, key, value matrices, each $N \times d$
$QK^{\top}$	$N \times N$ matrix of pairwise attention scores
$O$	attention output, $N \times d$

The FlashAttention paper’s central result is not about FLOPs at all — it is a bound on how many times data has to cross the boundary between HBM and the GPU’s fast on-chip SRAM, of size $M$:

$$\Theta(Nd + N^2) \ \text{HBM accesses (standard attention)}, \qquad \Theta\!\left(\frac{N^2 d^2}{M}\right) \ \text{HBM accesses (FlashAttention)}$$

Symbol	Meaning
$N$	sequence length
$d$	head dimension
$M$	size of the GPU’s on-chip SRAM, in floating-point elements

Inside each block, FlashAttention needs the softmax over the full row of scores without ever holding that full row in memory. It gets there with an online softmax: as each new block of scores $S^{(j)}$ arrives, the running maximum and running (unnormalized) sum are updated together, rescaling the sum so it stays consistent with the new maximum:

$$m^{\text{new}} = \max\!\big(m, \max(S^{(j)})\big), \qquad l^{\text{new}} = l \cdot e^{\,m - m^{\text{new}}} + \sum_{k} e^{\,S^{(j)}_k - m^{\text{new}}}$$

Symbol	Meaning
$m$	running maximum score seen across blocks processed so far
$l$	running (unnormalized) sum of exponentials, rescaled at each step
$S^{(j)}$	the block of raw attention scores just computed
$m^{\text{new}}, l^{\text{new}}$	updated running maximum and sum after folding in block $j$

Model FLOPs utilization, the metric FlashAttention-2 and FlashAttention-3 report their gains in, is simply the fraction of a chip’s peak arithmetic throughput a real run achieves:

$$\mathrm{MFU} = \frac{\text{measured FLOPs/s}}{\text{chip's peak FLOPs/s}}$$

What It Enabled

The three kernels did not change what attention computes; they changed how close a real run gets to the hardware’s own ceiling, on successive hardware generations.

	Standard attention	FlashAttention, 2022	FlashAttention-2, 2023	FlashAttention-3, 2024
Attention matrix materialized in HBM	yes, full $N \times N$	no	no	no
Memory use	$O(N^2)$	$O(N)$	$O(N)$	$O(N)$
Reported gain	baseline	up to 3x faster (GPT-2), 15% faster (BERT-large)	~2x faster than FlashAttention; 50–73% of A100 peak FLOPs/s	1.5–2.0x faster than FlashAttention-2 on H100
Reported utilization	25–40% of A100 peak (naive kernel)	~25–40% of A100 peak	up to 72% MFU on A100 (225 TFLOPs/s)	75% of H100 FP16 peak (740 TFLOPs/s); ~1.2 PFLOPs/s FP8
Target hardware	any GPU	A100-era	A100-era	Hopper H100, FP8

Once an exact-attention kernel stopped costing quadratic memory, sequence lengths that had been impractical became routine, feeding directly into the long-context techniques — extended position encodings and attention variants built to use a large context window efficiently — covered in Language Models, Chapter 16. At serving time, the same IO-aware thinking carried over from training to inference: keeping a growing key-value cache small and moving it efficiently through the memory hierarchy is the central problem Chapter 17 covers, and FlashAttention-style kernels are the computational substrate underneath serving-side techniques such as PagedAttention (2023) that manage that cache’s memory layout, detailed in full there. Mixture-of-experts models (Chapter 13) route each token through only a fraction of a model’s feed-forward parameters but still run full dense attention over every token, so an efficient attention kernel remains on the critical path even as the rest of the model gets sparser. And because linear-attention and state-space alternatives (Chapter 18) are judged in large part against how fast exact attention has become, FlashAttention reset the bar those approximations have to clear before an efficiency argument alone is enough to displace standard attention.

Key Papers

Paper	Authors	Year	Link	Why it matters
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness	Dao et al.	2022	arXiv:2205.14135	First exact attention kernel to avoid materializing the full $N \times N$ score matrix in HBM, cutting memory to $O(N)$ and training time by up to 3x with no approximation
FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning	Dao	2023	arXiv:2307.08691	Redistributed parallelism across the sequence dimension and between warps, roughly doubling FlashAttention’s speed and reaching 72% model FLOPs utilization on A100
FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision	Shah et al.	2024	arXiv:2407.08608	Exploited Hopper-specific asynchrony and FP8 to reach 75% of H100’s FP16 peak and near-1.2-PFLOPs/s FP8 throughput, with controlled low-precision error

PM Lens