FlashAttention

Self-attention is a handful of matrix multiplies — so why is it the slowest part of a transformer at long context? The surprise: on a modern GPU, attention is memory-bound, not compute-bound. The arithmetic is cheap; the bottleneck is shuffling a giant intermediate matrix to and from slow memory. FlashAttention fixes exactly that — without changing the math.

The real bottleneck: HBM vs SRAM

A GPU has two very different memories:

• HBM (high-bandwidth memory) — the tens of GB you think of as “GPU memory.” Large, but “only” ~2–3 TB/s.
• SRAM — tiny on-chip memory next to the compute units, ~20 MB, but roughly 10× the bandwidth of HBM.

Compute has outpaced memory for years, so for an operation like attention the chip finishes the math and then waits on memory. The metric that matters is how many bytes you move through HBM, not how many FLOPs you do.

Standard attention is wasteful by this metric. For a sequence of length N, it:

1. computes the scores S = QKᵀ — an N×N matrix — and writes it to HBM,
2. reads N×N back to apply softmax, writes N×N again,
3. reads N×N once more to multiply by V.

That N×N matrix is the problem. At N = 32k it’s multiple gigabytes per head, and it makes three round-trips through slow memory. The dot products are trivial; the data movement is everything.

Watch the tiling

Standard attention lights up the entire N×N matrix — it lives in HBM and is read back twice. FlashAttention streams one K/V block-column through SRAM at a time; only the current tile is ever resident, and the full matrix is never stored — so HBM cost stays O(N) instead of exploding as O(N²):

The online softmax — why it’s exact

The catch: softmax needs the whole row to normalise (you divide by the sum over all keys). If you only ever see one tile of keys at a time, how can you get the exact softmax? The online softmax trick: keep a running max m and running denominator ℓ, and rescale the accumulated output whenever a later tile reveals a larger value. When the last tile is done, the result is bit-for-bit the same as computing softmax over the full row at once.

That’s the key point most people miss: FlashAttention is not an approximation. Sparse and linear attention change the math to get speed; FlashAttention computes the identical softmax attention — it’s just IO-aware about how it touches memory. Prove it to yourself:

import numpy as np
rng = np.random.default_rng(0)

N, d = 8, 4
Q = rng.standard_normal((1, d))     # one query row, for clarity
K = rng.standard_normal((N, d))
V = rng.standard_normal((N, d))
scores = (Q @ K.T / np.sqrt(d))[0]  # N scores

# --- standard softmax over the full row at once ---
w = np.exp(scores - scores.max()); w /= w.sum()
out_full = w @ V

# --- online softmax: stream K/V in tiles of 2, keep running (m, l, acc) ---
m, l, acc = -np.inf, 0.0, np.zeros(d)
for t in range(0, N, 2):                 # one tile at a time
    s = scores[t:t+2]
    m_new = max(m, s.max())
    alpha = np.exp(m - m_new)            # rescale old state to the new max
    p = np.exp(s - m_new)
    l = l * alpha + p.sum()
    acc = acc * alpha + p @ V[t:t+2]
    m = m_new
out_online = acc / l

print("max abs difference:", np.abs(out_full - out_online).max())
# ~1e-16 — the tiled result is the exact same softmax attention.

max abs difference: 1.1102230246251565e-16

Don’t store it — recompute it

The forward pass never keeps the N×N matrix, so the backward pass can’t read it back either. FlashAttention simply recomputes the needed tiles during the backward pass. That sounds wasteful, but recomputing in SRAM is far cheaper than the HBM round-trips it avoids — a classic compute-for-memory trade. The upshot: training memory drops from O(N²) to O(N), which is what let context windows jump from 2k to 128k+.

In one breath

• Attention is memory-bound, not compute-bound: the cost is moving the N×N score matrix through slow HBM, not the FLOPs.
• FlashAttention never writes that matrix to HBM — it tiles the computation through fast on-chip SRAM and carries a running softmax (max m, denominator ℓ) from tile to tile.
• The online softmax rescales the accumulated output when a later tile reveals a bigger score, so the result is bit-for-bit exact — not an approximation.
• The backward pass recomputes tiles instead of storing them, dropping training memory from O(N²) to O(N) — what made 128k-token context practical.
• You get it for free via F.scaled_dot_product_attention and every serving stack; FA-2 improved warp partitioning, FA-3 added FP8 and Hopper async.

Check yourself

What to remember

• Attention is memory-bound: the cost is moving the N×N score matrix through HBM, not the FLOPs.
• FlashAttention keeps that matrix out of HBM entirely — tiling through fast SRAM plus an online softmax that rescales running state.
• It’s exact, not approximate; the backward pass recomputes tiles instead of storing them, cutting memory from O(N²) to O(N).
• You get it for free via scaled_dot_product_attention and every modern serving stack — it’s why 128k-token context is practical.