Core ML

FlashAttention

Standard attention writes a huge score matrix to GPU RAM and reads it back for the softmax — FlashAttention tiles the computation so the score matrix never leaves the small, fast on-chip cache, making attention faster without changing the result by a single bit.

Why Does This Exist?

Take the attention complexity page's finding seriously: the $N \times N$ score matrix, for a sequence of length 4,096 and 32 heads, already measures tens of gigabytes per layer. That matrix has to be written to GPU memory (HBM) and then read back for the softmax — twice. On modern GPUs, HBM bandwidth is a more severe constraint than raw compute. A100 GPUs have roughly 312 TFLOPS of compute but only 2 TB/s of HBM bandwidth. Standard attention is memory-bound, not compute-bound, which means the attention arithmetic is fast but the data movement is the thing that costs time.

Dao et al. published FlashAttention in 2022 with a simple observation: A100 also has 20 MB of SRAM on-chip — much smaller than HBM but roughly 10× faster to access. If you could fit the relevant tiles of Q, K, and V into SRAM, run the attention math there, and accumulate the output without ever materialising the full $N \times N$ score matrix in HBM, you'd cut HBM reads and writes dramatically. The FLOPs don't change. The result doesn't change. Only the data movement changes.

That's FlashAttention. IO-aware, exact, and now the default implementation in PyTorch, JAX, and every serious LLM training stack.

Think of It Like This

Doing maths at your desk instead of photocopying pages to the library

Imagine scoring 1,000 exam papers by printing all 1,000 answer sheets, taking them to a filing room, reading each one for scoring, carrying the scores back, then summing them. That's standard attention: compute the scores, write them all to memory, read them back for softmax, read them again for the weighted sum.

FlashAttention is scoring small batches at your desk without leaving the room. You pull a block of answer sheets, score them, note the running total, put them back, pull the next block. The final total is identical. But you never needed the 1,000-sheet pile in the filing room at once.

How It Actually Works

The memory-bandwidth problem

Standard attention does four large HBM operations per layer: write $QK^\top$ (the $N \times N$ score matrix), read it for softmax, write the softmax weights, read them for the $\text{weights} \times V$ multiply. For N=4,096 and 32-bit floats, the score matrix alone is $4096 \times 4096 \times 32 \times 4$ bytes ≈ 2 GB per layer. Even at 2 TB/s bandwidth that's a millisecond per read or write, and there are four of them.

Tiling and online softmax

FlashAttention processes the sequence in tiles of size $B_r \times B_c$ chosen to fit in SRAM. For each tile:

Load a tile of $Q_i$ and $K_j$ from HBM into SRAM.
Compute the tile's raw scores $S_{ij} = Q_i K_j^\top / \sqrt{d_k}$ .
Maintain a running maximum $m$ and running sum $\ell$ of softmax numerators using the online softmax recurrence — updating both as each new tile arrives rather than waiting for all scores to be computed.
Load the corresponding $V_j$ tile, multiply by the current softmax weights, and accumulate into the output tile $O_i$ .
Once all K/V tiles are processed, normalise $O_i$ by the final $\ell$ and write back to HBM.

The $N \times N$ score matrix is never materialised. SRAM holds only the current tile. The online softmax recurrence, derived by Milakov and Gimelshein (2018), is what makes this numerically equivalent to computing the full softmax over all scores at once.

Recomputation in the backward pass

Backward pass needs the attention weights to compute gradients. Standard attention stores them during forward (another $O(N^2)$ memory cost). FlashAttention instead stores only the softmax statistics ( $m$ and $\ell$ per query row) and recomputes the attention weights from tiles during the backward pass. This trades extra compute for drastically lower memory, which is the right trade on modern hardware.

The result

Memory usage drops from $O(N^2)$ (the score matrix) to $O(N)$ (the output). FLOPs are the same. HBM reads and writes drop dramatically. Wall-clock speedups of 2–4× on A100 were measured for N=1,024–4,096, growing as N grows because the score matrix that standard attention has to copy around gets larger.

Show Me the Code

Verifying that a tiled attention pass produces the same output as standard attention — the whole point of "exact".

1import numpy as np2
3
4def standard_attention(Q, K, V):5    d = Q.shape[-1]6    scores = Q @ K.T / d**0.57    scores -= scores.max(-1, keepdims=True)      # numerical stability8    weights = np.exp(scores)9    weights /= weights.sum(-1, keepdims=True)10    return weights @ V11
12
13def tiled_attention(Q, K, V, tile=2):14    """Simplified tiled attention; accumulates output without full score matrix."""15    n, d = Q.shape16    O = np.zeros_like(Q)17    m = np.full(n, -np.inf)    # running row-wise max18    l = np.zeros(n)             # running row-wise softmax denominator19
20    for j in range(0, n, tile):21        Kj = K[j:j+tile]22        Vj = V[j:j+tile]23        for i in range(0, n, tile):24            Qi = Q[i:i+tile]25            s = Qi @ Kj.T / d**0.5          # tile scores26            m_new = np.maximum(m[i:i+tile], s.max(-1))27            e_old = np.exp(m[i:i+tile] - m_new)[:, None]28            e_s = np.exp(s - m_new[:, None])29            O[i:i+tile] = O[i:i+tile] * e_old + e_s @ Vj30            l[i:i+tile] = l[i:i+tile] * e_old.squeeze() + e_s.sum(-1)31            m[i:i+tile] = m_new32
33    O /= l[:, None]34    return O35
36
37rng = np.random.default_rng(7)38Q, K, V = (rng.standard_normal((8, 16)) for _ in range(3))39out_std = standard_attention(Q, K, V)40out_tiled = tiled_attention(Q, K, V, tile=2)41print("Max difference:", np.abs(out_std - out_tiled).max())42# Max difference: ~1e-15  (floating-point rounding only, not approximation)

The difference is floating-point rounding — the order of the operations differs slightly between tiled and standard, which shifts the last bits. There is no approximation.

Watch Out For

Assuming FlashAttention is an approximation

It's named differently from standard attention, it rewrites the computation, and it produces faster results — so it's easy to assume it's cutting a corner. It isn't. Given the same inputs, FlashAttention produces the same output as standard attention, to floating-point precision. The only thing it changes is the pattern of memory access. This is what the paper means by "exact": the algorithm is IO-aware, not approximate.

Expecting linear-time complexity

FlashAttention reduces memory from $O(N^2)$ to $O(N)$ and reduces HBM bandwidth dramatically. But it does not reduce FLOPs — the compute is still $O(N^2 d)$ , the same as standard attention. If your bottleneck is pure arithmetic throughput rather than memory bandwidth (rare on current hardware, but possible on specialised accelerators), FlashAttention won't help there.

The Quick Version

Standard attention writes the $N \times N$ score matrix to HBM and reads it back multiple times; on modern GPUs, this memory movement is the real bottleneck, not the arithmetic.
FlashAttention tiles Q, K, V to fit in fast SRAM, runs an online softmax recurrence to accumulate partial results, and never writes the full score matrix to HBM.
Memory drops from $O(N^2)$ to $O(N)$ ; FLOPs are unchanged; the result is numerically identical to standard attention.
Backward pass recomputes attention weights from tiles rather than storing them, trading a small amount of extra compute for large memory savings.
FlashAttention is the default implementation in PyTorch (F.scaled_dot_product_attention), JAX, and most training frameworks — you're almost certainly using it already.