Flash Attention

June 2025

FlashAttention is a way to compute the same attention as a Transformer textbook, without ever writing the giant N×N score matrix to GPU memory. That one change is why long-context LLMs and VLMs became practical.

What attention actually computes

For a sequence of length N and head dimension d, each head has queries, keys, and values Q, K, V of shape N×d:

S = QK^T → P = softmax(S) → O = PV

The math is simple. The problem is S and P: both are N×N. At N = 8,192 that is about 64 million numbers per head. At N = 32,768 it is a billion. Training and serving spend most of their time shipping that matrix to and from memory, not multiplying numbers.

The bottleneck is HBM, not FLOPs

A GPU has a small, fast on-chip SRAM and a large, slow HBM (the “GPU memory” you see in nvidia-smi). On an A100, SRAM is on the order of tens of megabytes; HBM is tens of gigabytes, but much higher latency.

Standard attention:

Reads Q and K from HBM, writes the full S back to HBM.
Reads S, writes P to HBM.
Reads P and V, writes O.

Those round trips dominate. Attention is memory-bound: the Tensor Cores sit idle waiting for HBM. You cannot “just add more FLOPs” to fix it.

Keep the N×N matrix in SRAM

SRAM cannot hold a full 8k×8k matrix. It can hold small tiles. FlashAttention cuts Q, K, and V into blocks that fit on-chip:

Load one block of Q into SRAM.
Loop over blocks of K and V.
For each pair of blocks, compute a slice of scores, softmax, and a partial output entirely in SRAM.
Write only the final output block O back to HBM.

The full score matrix still exists mathematically. It never exists as a tensor in HBM. Memory for attention drops from O(N²) to O(N).

Online softmax: why tiling is not trivial

Softmax is a global operation: each row is exp(s_i) / Σ exp(s_j). If you only see one tile of a row, you do not know the sum — or the maximum, which you need for numerical stability.

The fix is online softmax. For each query row you keep two running statistics while you stream key tiles:

the running maximum of the scores, and
the running sum of exp(score − max).

When a new tile arrives, you rescale the partial output you already have (because the max and the denominator changed), then add the new tile’s contribution. After the last key tile, the output is exactly the same as full softmax — not an approximation.

What you get in practice

Faster: fewer HBM round trips, often 2–4× on long sequences.
Leaner: you can raise context length without the attention matrix blowing GPU memory.
Exact: same numerics as standard attention, so checkpoints and downstream metrics stay comparable.

FlashAttention-2 improves how work is split across warps (less non-matmul overhead). FlashAttention-3 targets newer NVIDIA hardware (Hopper) with asynchrony and lower-precision Tensor Cores. The idea is unchanged: do attention where the data already is.

That is also why it shows up in inference stacks (TensorRT-LLM, vLLM, Triton) and not only in training papers. Long prompts and KV-cache-heavy decoding hit the same HBM wall.