FlashAttention Explained: Why Modern LLMs Run Faster

TL;DR: FlashAttention computes exact attention without ever writing the full N×N attention matrix to GPU memory. By treating attention as a memory-movement problem rather than a FLOP-counting one, it trains GPT-2 up to 3× faster, beats the BERT-large MLPerf 1.1 record by 15%, and makes 4× longer context cheaper than the previous 1× baseline. It is the reason "long context" stopped being a research curiosity and became a default feature.

The problem

Self-attention is the engine of the Transformer, and its cost grows quadratically with sequence length N — in both time and memory. For years that quadratic wall was the thing standing between models and long documents, high-resolution images, and large context windows. The natural response was to attack the FLOP count: sparse attention, low-rank approximations, linear-attention variants. Many of these reduce the asymptotic compute on paper.

The uncomfortable finding that motivates this paper is that most of them deliver little or no wall-clock speedup. The reason is that standard attention is memory-bound, not compute-bound. On modern GPUs, arithmetic throughput has outpaced memory bandwidth for years, so what actually limits attention is the volume of reads and writes between the large-but-slow HBM (high-bandwidth memory) and the small-but-fast on-chip SRAM. A textbook attention kernel materializes the full N×N scores matrix S = QKᵀ and the softmax probabilities P in HBM, then reads them back — O(N²) traffic across the slowest link in the system. Cutting FLOPs while leaving that traffic in place optimizes the wrong number.

The authors' argument is that attention should be made IO-aware: explicitly account for and minimize movement between memory tiers, exactly the way classical numerical linear algebra is designed around the memory hierarchy.

How it works

FlashAttention produces the same output as standard attention — no approximation — but never builds the N×N matrix in HBM. Two ideas carry the method.

Tiling. Q, K, and V are split into blocks. The kernel loads a block from HBM into SRAM, does the attention math on-chip, and accumulates into the output incrementally. The hard part is softmax: it normalizes across an entire row, which couples every column together, so you seemingly can't process the row in pieces. FlashAttention sidesteps this by folding in the online-softmax formulation of Milakov and Gimelshein (2018), which the paper cites: for each query row it keeps two running statistics — the running maximum and the running sum of exponentials — and rescales the partial output each time a new block of keys and values arrives. The result is a numerically stable, mathematically exact softmax computed block-by-block, without the row ever existing in full anywhere.

Recomputation. Backpropagation normally needs the intermediate attention matrix. Instead of stashing that quadratic blob, FlashAttention stores only the per-row softmax statistics (the max and sum — O(N) values) and recomputes the needed attention blocks on-chip during the backward pass. That means more arithmetic. But since the operation is memory-bound, trading FLOPs for far fewer HBM accesses is still a net win, and it drops memory from quadratic to linear.

The whole thing ships as a single fused CUDA kernel, so the intermediates live and die in SRAM rather than round-tripping through HBM.

The paper backs this with an IO-complexity analysis: FlashAttention needs Θ(N²d²/M) HBM accesses (d is head dimension, M is SRAM size, for d ≤ M ≤ Nd) versus Θ(Nd + N²) for standard attention — a large reduction for realistic d and M. They also prove a lower bound: no exact attention algorithm can asymptotically beat that HBM-access count across all SRAM sizes. The paper is careful here — it calls the result optimality "for a range of SRAM sizes" and leaves a fully parameterized lower bound as future work. A block-sparse variant extends the same IO-aware treatment to sparse attention for further gains.

Results

All numbers below are as reported in the paper. Training results are averaged over runs on 8× Nvidia A100 GPUs; kernel benchmarks are single-A100.

At the kernel level the paper reports the attention operation running up to 3× faster than standard PyTorch attention across common sequence lengths (128–2K), with a typical 2–4× range measured on an A100 (batch 8, head dim 64, 12 heads) and a peak of 7.6× on the GPT-2 configuration (Figure 1, right). Memory is also up to 20× more efficient than exact-attention baselines, and the block-sparse variant is faster still. The headline isn't just speed: the longer-context numbers show the speedup converting into model quality, because sequences that were previously out of reach become trainable.

Limitations and context

A few honest caveats sit alongside the wins.

The speedups come from a hand-written, fused CUDA kernel. IO-aware kernels are hard to write, tied to specific GPU architectures, and tuned to a particular SRAM size, since the block size is set by M — the paper itself flags the need for higher-level abstractions so people don't hand-roll these per device. The analysis and implementation target the GPU HBM/SRAM hierarchy specifically; optimal block sizes are hardware-dependent. The optimality result is also scoped to a single GPU. The paper notes that attention is typically split across 4–8 GPUs on a node, which adds inter-GPU transfer as another layer of the IO hierarchy — an analysis it leaves to future work.

The original v1 kernel also parallelizes only over batch size and number of heads, so when their product is small — long sequences with a small batch, for instance — it leaves streaming multiprocessors idle; FlashAttention-2 reports v1 reaching only 25–40% of the A100's theoretical maximum FLOPs/s. That gap is precisely what FlashAttention-2 later addressed by restructuring the work partitioning between thread blocks and warps, roughly doubling throughput to 50–73% of peak.

And the central honest point: FlashAttention changes how attention is computed, not what. It does not reduce the asymptotic O(N²d) FLOP count of full attention. It makes that exact computation dramatically cheaper in wall-clock time and memory, but the quadratic compute is still there — only the block-sparse variant touches it, and only by introducing sparsity (i.e., approximation).

Why it matters

FlashAttention turned into foundational infrastructure. By making exact attention memory-linear and substantially faster, it directly enabled the long context windows that current models depend on and lowered training and inference cost across the board. Its ideas were absorbed into the training and serving stacks that came after it — adoption that sits outside the paper's own claims and isn't something the paper measures.

Its more durable contribution may be conceptual: it reframed efficiency research around IO-awareness and the memory hierarchy rather than FLOP counting alone, and it spawned a lineage — FlashAttention-2 (better parallelism and work partitioning) and FlashAttention-3 (exploiting newer Hopper-class hardware features, including FP8). When people say modern LLMs "just run faster," a large part of that story is IO-aware kernels like this one.

References

• Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. arXiv:2205.14135 — https://arxiv.org/abs/2205.14135
• Tri Dao. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. 2023. arXiv:2307.08691 — https://arxiv.org/abs/2307.08691
• Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. 2024. arXiv:2407.08608 — https://arxiv.org/abs/2407.08608