4.2 FlashAttention I: IO-Awareness & The Online Softmax¶

In The Attention Mechanism From Scratch we wrote the canonical formula and the canonical loop:

That implementation is correct and it is slow — not because the arithmetic is expensive, but because of where the numbers live. The standard implementation builds the full \(N \times N\) score matrix \(S = QK^\top/\sqrt{d}\) in GPU main memory (HBM), reads it back to apply softmax, reads it again to multiply by \(V\). For a sequence of \(N = 8192\) tokens, that intermediate matrix is 64 million entries per head per layer, and every byte of it has to travel across the slowest wire in the machine. The matrix multiplies that everyone worries about are not the bottleneck. The traffic to and from HBM is.

FlashAttention (Dao, Fu, Ermon, Rudra, Ré, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, 2022) is the algorithm that removed that bottleneck. Its central trick is to never materialize the \(N \times N\) matrix at all — to compute attention in tiles small enough to live in on-chip SRAM, fusing the three steps into a single pass. To make that possible it needs one beautiful piece of numerical machinery: the online softmax, a way to compute a numerically-stable softmax incrementally, while streaming the inputs, without ever seeing all of them at once. That algorithm — running maxima, running denominators, and a correction factor that retroactively rescales partial results — is the heart of this chapter, and one of the most elegant ideas in all of systems-for-ML.

This chapter is a marquee chapter. We will derive the online softmax from scratch, prove it gives exactly the same answer as the naive version (FlashAttention is exact, not an approximation), build the full forward and backward passes in NumPy so you can run them, and then do the IO-complexity accounting that explains why it is faster. It builds directly on The Roofline Model & Performance Engineering and GPU Architecture & The Memory Hierarchy; FlashAttention 2 & 3 then takes the same core and squeezes out the remaining inefficiencies.

The memory bottleneck: why naive attention is IO-bound¶

Let us be precise about what the standard kernel does. Take a single attention head with sequence length \(N\) and head dimension \(d\) (think \(d = 64\) or \(128\)). The inputs \(Q, K, V\) are each \(N \times d\). The reference computation is three operations:

import numpy as np

def attention_naive(Q, K, V):
    # Q, K, V: (N, d) arrays already on the device
    d = Q.shape[-1]
    S = Q @ K.T / np.sqrt(d)          # (N, N)  -- the score matrix, WRITTEN to HBM
    P = softmax_rows(S)               # (N, N)  -- read S, write P to HBM
    O = P @ V                         # (N, d)  -- read P, read V, write O
    return O

def softmax_rows(S):
    S = S - S.max(axis=-1, keepdims=True)   # numerical stability: subtract row max
    e = np.exp(S)
    return e / e.sum(axis=-1, keepdims=True)

The arithmetic is two matrix multiplies (\(QK^\top\) and \(PV\)), each \(O(N^2 d)\) floating-point operations, plus an \(O(N^2)\) softmax. On a modern GPU those FLOPs run on the tensor cores at, on the order of, hundreds of teraFLOP/s. So compute is cheap. What is expensive is the data movement, and the data movement is dominated by the two \(N \times N\) intermediates \(S\) and \(P\).

Recall the memory hierarchy from GPU Architecture & The Memory Hierarchy. A GPU has a small amount of extremely fast on-chip SRAM (the shared memory / L1 cache of each streaming multiprocessor, on the order of 100–200 KB per SM, with aggregate bandwidth in the tens of TB/s) and a large amount of much slower off-chip HBM (tens of GB at, on the order of, 2–3.5 TB/s). The ratio matters: SRAM is roughly an order of magnitude faster per byte than HBM, and there is far less of it.

The naive kernel makes the worst possible use of this hierarchy. Look at the HBM traffic:

  operation        reads from HBM            writes to HBM       bytes (fp16)
  ---------------   -----------------------   -----------------   --------------
  S = Q K^T         Q (N·d), K (N·d)          S (N·N)             ~2·N² + 4·N·d
  P = softmax(S)    S (N·N)                   P (N·N)             ~4·N²
  O = P V           P (N·N), V (N·d)          O (N·d)             ~2·N² + 4·N·d
  ---------------------------------------------------------------------------
  TOTAL HBM traffic dominated by the N² terms:               Θ(N²)  elements

Every \(N \times N\) matrix is born in HBM, dies in HBM, and is dragged across the bus in between. The total HBM traffic scales as \(\Theta(N^2)\) — quadratic in sequence length, the same order as the FLOPs, but the bytes are moving over a channel that is an order of magnitude slower than the one the FLOPs run on. Plug that into the roofline picture from The Roofline Model: the arithmetic intensity (FLOPs per byte of HBM traffic) is low and roughly constant in \(N\), which puts attention firmly on the memory-bound (bandwidth-limited) side of the roofline. The hardware’s tensor cores sit mostly idle, starved, waiting for HBM.

There is a second, even harsher cost: memory capacity. The \(S\) and \(P\) matrices are \(O(N^2)\) in size. At \(N = 8192\) and fp16, a single \(N \times N\) matrix is \(8192^2 \times 2 \approx 134\) MB — per head, per layer. With dozens of heads and dozens of layers, and a backward pass that wants to keep \(P\) around for the gradient, you blow through HBM capacity long before you run out of compute. This is the wall that capped context lengths for years.

The diagnosis is now sharp. Attention is slow because (1) it moves \(\Theta(N^2)\) bytes across the slow HBM bus, and (2) it allocates \(\Theta(N^2)\) bytes of HBM it does not fundamentally need. Both pathologies share one cause: the algorithm materializes the full score matrix in HBM. If we could compute the exact same output without ever writing \(S\) or \(P\) to HBM — keeping the working set in SRAM — both problems vanish. The obstacle to doing that is the softmax, because softmax over a row needs to see the whole row (to normalize) before it can produce any output. The online softmax dissolves that obstacle.

The online softmax: computing a normalized sum in one streaming pass¶

Forget attention for a moment and consider the bare numerical problem. We have a long vector of logits \(x_1, x_2, \dots, x_N\) and we want the softmax weights \(p_i = e^{x_i} / \sum_j e^{x_j}\) — but in practice we never want \(p\) alone; in attention we want the weighted sum \(\sum_i p_i v_i\) for value vectors \(v_i\). The question: can we compute this while streaming the \(x_i\) (and \(v_i\)) one block at a time, using \(O(1)\) extra state, instead of making two full passes (one to find the normalizer, one to accumulate)?

Why the naive online attempt fails¶

Softmax must be computed in a numerically stable way. The naive expression \(e^{x_i}/\sum_j e^{x_j}\) overflows the moment any \(x_i\) exceeds ~88 (for fp32, \(e^{88} \approx 10^{38}\) is near the float max) or underflows to zero for very negative logits. The standard fix is the max-subtraction trick: subtract the maximum logit \(m = \max_j x_j\) before exponentiating,

which is algebraically identical (the \(e^{-m}\) cancels top and bottom) but keeps every exponent \(\le 0\), so \(e^{x_i - m} \in (0, 1]\) — no overflow, and the largest term is exactly 1. The catch: \(m\) depends on the whole vector. If we are streaming, we do not know the global max until we have seen everything. That is exactly the dependency that forces the naive kernel to write \(S\) to HBM and read it back.

The key idea: a running max and a correction factor¶

The online softmax breaks the dependency with a retroactive rescale. We process the logits in blocks. We keep two running statistics:

• \(m\) — the maximum logit seen so far,
• \(\ell\) — the running sum of exponentials, \(\sum e^{x_j - m}\), relative to the current running max \(m\),

and (for attention) an unnormalized output accumulator \(o = \sum e^{x_j - m}\, v_j\), also relative to \(m\). When a new block arrives and contains a logit larger than the current \(m\), the running max jumps from \(m_{\text{old}}\) to \(m_{\text{new}}\). Every quantity we accumulated was scaled by \(e^{-m_{\text{old}}}\); to put it on the new scale we must multiply it by

Because \(m_{\text{new}} \ge m_{\text{old}}\), this correction factor \(\alpha = e^{m_{\text{old}} - m_{\text{new}}} \in (0, 1]\) shrinks the old accumulators down onto the new, larger scale. That single multiply is the entire trick. It lets us “change our mind” about the normalizer after the fact, in \(O(1)\) work, every time we discover a larger logit.

Here is the recurrence in full. Suppose after processing some prefix we hold \((m, \ell, o)\), and a new block contributes logits \(\{x_j\}\) with values \(\{v_j\}\) (in attention these are one tile of \(K, V\)). Let the block’s local max be \(\tilde m = \max_j x_j\).

After the final block, the true softmax-weighted output is \(o / \ell\) — we divide by the denominator exactly once, at the very end. Notice what each line does: \(m_{\text{new}}\) updates the max; \(\alpha\) rescales the already-accumulated sum and output down onto the new max; then we add the new block’s contributions, exponentiated against the new max so they are on the same scale.

A from-scratch implementation and a correctness proof¶

Let us make this concrete and runnable. The function below computes a streaming, numerically-stable softmax-weighted sum and we assert it equals the two-pass reference.

import numpy as np

def online_softmax_weighted_sum(x, V, block_size=4):
    """
    Compute  sum_i softmax(x)_i * V[i]   in a single streaming pass,
    processing the logits x in blocks, using only O(d) running state.

    x : (N,)   logits
    V : (N, d) value vectors
    returns o/ell : (d,)  the softmax-weighted average of the rows of V
    """
    N, d = V.shape
    m = -np.inf            # running max of logits seen so far
    ell = 0.0              # running sum of exp(x - m)
    o = np.zeros(d)        # running sum of exp(x - m) * v, UNnormalized

    for start in range(0, N, block_size):
        xb = x[start:start + block_size]          # block of logits
        Vb = V[start:start + block_size]          # block of values

        m_block = xb.max()                        # local max of this block
        m_new = max(m, m_block)                   # updated running max

        alpha = np.exp(m - m_new)                 # correction for OLD state
        # exp(m - m_new) with m = -inf at the very first block -> exp(-inf) = 0,
        # which correctly zeroes the (empty) initial accumulators.

        p = np.exp(xb - m_new)                    # block weights vs new max, in (0,1]
        ell = alpha * ell + p.sum()               # rescale old sum, add new
        o   = alpha * o   + p @ Vb                # rescale old output, add new
        m = m_new

    return o / ell

# ---- correctness check against the honest two-pass softmax ----
rng = np.random.default_rng(0)
N, d = 37, 8
x = rng.normal(0, 5, size=N)        # wide spread -> stresses the max-subtraction
V = rng.normal(0, 1, size=(N, d))

# reference: stable softmax then weighted sum
p_ref = np.exp(x - x.max()); p_ref /= p_ref.sum()
o_ref = p_ref @ V

o_online = online_softmax_weighted_sum(x, V, block_size=4)
print("max abs error:", np.abs(o_online - o_ref).max())   # ~1e-16, machine epsilon

Run it: the maximum absolute error is on the order of \(10^{-16}\) — floating-point round-off, nothing more. The online algorithm is bit-for-bit equivalent up to round-off, not an approximation. Here is why, by induction on the blocks.

Claim. After processing the first \(k\) blocks (covering indices in a set \(I_k\)), the state satisfies, exactly,

Base case (\(k=0\), no blocks): \(m = -\infty\), and the empty sums are \(0\), matching \(\ell = o = 0\). ✓

Inductive step. Assume the claim holds after \(k\) blocks. Process block \(k{+}1\) with indices \(J\) and local max \(\tilde m\). The update sets \(m_{\text{new}} = \max(m, \tilde m) = \max_{i \in I_k \cup J} x_i\), which is correct. For the denominator, the old sum was relative to the old \(m\); multiplying by \(\alpha = e^{m - m_{\text{new}}}\) re-bases it:

and adding \(\sum_{j \in J} e^{x_j - m_{\text{new}}}\) gives \(\ell_{\text{new}} = \sum_{i \in I_k \cup J} e^{x_i - m_{\text{new}}}\), exactly the claim for \(k{+}1\). The identical argument holds for \(o\) with the \(v_i\) carried along. ∎

So at the end \(m = \max_i x_i\), \(\ell = \sum_i e^{x_i - m}\), and \(o = \sum_i e^{x_i - m} v_i\), hence \(o/\ell = \sum_i \frac{e^{x_i - m}}{\sum_j e^{x_j - m}} v_i = \sum_i p_i v_i\). The streaming computation equals the reference. The whole point is that we paid only \(O(d)\) state and one extra multiply per block to avoid ever holding all \(N\) logits at once. The numerical stability is preserved at every step because every exponent \(x_j - m_{\text{new}} \le 0\), so we never exponentiate a positive number — the running max guarantees it.

FlashAttention forward: tiling + online softmax fused into one pass¶

Now we lift the online softmax from a single row to the full attention matrix. The observation that makes the row-wise online softmax into FlashAttention: every output row is an independent online softmax over the key/value blocks. Row \(i\) of the output is \(\sum_j p_{ij} v_j\) where \(p_{ij} = \operatorname{softmax}_j(S_{ij})\), and we can compute that row by streaming over column-blocks of \(K, V\) exactly as above — running max \(m_i\), running denominator \(\ell_i\), running output \(o_i\) — all in SRAM. We do this for a block of query rows at a time, so that the \(Q\) tile, the \(K\)/\(V\) tiles, and the accumulators all fit on-chip simultaneously.

The structure is a double loop. The outer loop is over blocks of query rows (\(B_r\) rows at a time). The inner loop is over blocks of key/value rows (\(B_c\) columns at a time). For each \((i, j)\) tile we load \(Q_i, K_j, V_j\) from HBM into SRAM, compute the small \(B_r \times B_c\) score tile, update the running statistics for the \(B_r\) rows, and move on. The score tile is born and consumed entirely in SRAM; it is never written to HBM. That is the whole game.

Here is a faithful, runnable NumPy implementation of the forward pass. It mirrors the structure of the real CUDA kernel — the same loops, the same statistics, the same rescale — just without the GPU. (Causal masking and the actual SRAM tiling are noted in comments; we will add the mask in the code that follows.)

import numpy as np

def flash_attention_forward(Q, K, V, Br=32, Bc=32, causal=False):
    """
    Exact attention via tiling + online softmax. Mirrors the FlashAttention-1
    forward pass. No N×N matrix is ever fully materialized: the largest
    intermediate is one Br×Bc score tile.

    Q, K, V : (N, d)
    returns:
      O : (N, d)         attention output
      L : (N,)           logsumexp statistic per row  (= m + log ell), saved for backward
    """
    N, d = Q.shape
    scale = 1.0 / np.sqrt(d)

    O = np.zeros((N, d))
    L = np.zeros(N)                      # row logsumexp, needed by the backward pass

    # OUTER LOOP over query blocks (each maps to a CUDA thread block / program)
    for i0 in range(0, N, Br):
        i1 = min(i0 + Br, N)
        Qi = Q[i0:i1]                    # (br, d)  -- stays in SRAM for the whole inner loop
        br = i1 - i0

        # Per-row running statistics, initialized for "nothing seen yet"
        m_i  = np.full(br, -np.inf)      # running max logit per row
        l_i  = np.zeros(br)              # running denominator per row
        O_i  = np.zeros((br, d))         # running UNnormalized output per row

        # INNER LOOP over key/value blocks
        for j0 in range(0, N, Bc):
            j1 = min(j0 + Bc, N)
            Kj = K[j0:j1]                # (bc, d)
            Vj = V[j0:j1]                # (bc, d)

            # 1) score tile in SRAM (Br×Bc) -- the only "big" intermediate, and it is tiny
            Sij = (Qi @ Kj.T) * scale    # (br, bc)

            if causal:
                # mask out keys j > query i  (positions are absolute indices)
                rows = np.arange(i0, i1)[:, None]
                cols = np.arange(j0, j1)[None, :]
                Sij = np.where(cols <= rows, Sij, -np.inf)

            # 2) online-softmax update for these br rows over this block
            m_block = Sij.max(axis=1)               # (br,) local max of the tile
            m_new   = np.maximum(m_i, m_block)       # (br,) updated running max
            # guard rows that are entirely -inf (e.g. fully-masked causal blocks)
            m_new   = np.where(np.isneginf(m_new), 0.0, m_new)

            P = np.exp(Sij - m_new[:, None])         # (br, bc) tile weights, in [0,1]
            alpha = np.exp(m_i - m_new)              # (br,) correction for old state
            alpha = np.where(np.isneginf(m_i), 0.0, alpha)

            l_i = alpha * l_i + P.sum(axis=1)        # rescale + accumulate denominator
            O_i = alpha[:, None] * O_i + P @ Vj      # rescale + accumulate output
            m_i = m_new

        # 3) finalize: normalize once, and save the logsumexp for backward
        l_safe = np.where(l_i == 0.0, 1.0, l_i)      # avoid 0/0 for fully-masked rows
        O[i0:i1] = O_i / l_safe[:, None]
        L[i0:i1] = m_i + np.log(l_safe)              # logsumexp = m + log(sum exp(.-m))

    return O, L


# ---- verify against the reference, with and without causal masking ----
def reference_attention(Q, K, V, causal=False):
    d = Q.shape[-1]
    S = (Q @ K.T) / np.sqrt(d)
    if causal:
        N = Q.shape[0]
        mask = np.tril(np.ones((N, N), bool))
        S = np.where(mask, S, -np.inf)
    S = S - S.max(axis=1, keepdims=True)
    P = np.exp(S); P /= P.sum(axis=1, keepdims=True)
    return P @ V

rng = np.random.default_rng(1)
N, d = 100, 16
Q = rng.normal(size=(N, d)); K = rng.normal(size=(N, d)); V = rng.normal(size=(N, d))

for causal in (False, True):
    O_flash, _ = flash_attention_forward(Q, K, V, Br=32, Bc=32, causal=causal)
    O_ref = reference_attention(Q, K, V, causal=causal)
    print(f"causal={causal}: max abs error = {np.abs(O_flash - O_ref).max():.2e}")
# both print ~1e-15: exact up to floating-point round-off

A few details worth pausing on.

• The accumulators are the entire memory footprint per query block. For a block of \(B_r\) rows we hold \(m_i\) (\(B_r\) floats), \(\ell_i\) (\(B_r\) floats), and \(O_i\) (\(B_r \times d\) floats) — and the transient \(B_r \times B_c\) score tile. Choose \(B_r, B_c\) so all of this fits in SRAM. Nothing scales with \(N\). The peak extra HBM allocation across the whole kernel is just \(O\) (size \(N \times d\)) and \(L\) (size \(N\)) — linear in \(N\), never quadratic.

• The logsumexp \(L_i = m_i + \log \ell_i\) is the one statistic we save. It is the log of the softmax denominator for row \(i\), computed in a stable way. We will see in the next section that the backward pass needs exactly this scalar per row to reconstruct the softmax weights on the fly — so storing \(L\) (size \(N\)) lets us recompute \(P\) (size \(N^2\)) in backward instead of storing it. This is the recomputation trade we make.

• Causal masking is free. With a causal mask, query \(i\) only attends to keys \(j \le i\). In the tiled loop that means the outer query-block \(i\) can simply skip all key-blocks \(j\) that lie entirely above the diagonal (\(j_0 > i_1\)), and only the diagonal block needs the elementwise triangular mask. This roughly halves the work for long sequences — a major reason causal FlashAttention is so fast — and the real kernel exploits it explicitly.

The backward pass: recomputation instead of storage¶

Training needs gradients. The standard backward pass for attention stores the \(N \times N\) probability matrix \(P\) from the forward pass so it can compute \(\partial \mathcal{L}/\partial Q\), \(\partial \mathcal{L}/\partial K\), \(\partial \mathcal{L}/\partial V\). But \(P\) is exactly the \(O(N^2)\) object we worked so hard not to materialize. FlashAttention’s answer is recomputation (a.k.a. activation checkpointing, see Memory-Efficient Training): do not store \(P\); instead, recompute each \(P\) tile from \(Q, K\) and the saved logsumexp \(L\) during the backward pass, again block by block in SRAM. We trade extra FLOPs (recomputing the scores) for not touching HBM with an \(N^2\) matrix. Because attention is memory-bound, that trade is a clear win — the “wasted” recomputation FLOPs run on otherwise-idle tensor cores.

First, the math. With \(S = QK^\top/\sqrt{d}\), \(P = \operatorname{softmax}_{\text{row}}(S)\), and \(O = PV\), and given the upstream gradient \(\mathrm{d}O = \partial \mathcal{L} / \partial O\), the chain rule gives:

The softmax Jacobian (per row, since softmax is applied row-wise) turns \(\mathrm{d}P\) into \(\mathrm{d}S\). For a single softmax row \(p = \operatorname{softmax}(s)\) with upstream \(\mathrm{d}p\), the standard result is

In attention, \(\mathrm{d}p_{ij} = (\mathrm{d}O\,V^\top)_{ij}\) and \(\sum_k p_{ik}\,\mathrm{d}p_{ik} = \sum_k p_{ik}(\mathrm{d}O_i \cdot V_k) = \mathrm{d}O_i \cdot (\sum_k p_{ik} V_k) = \mathrm{d}O_i \cdot O_i\). Define the per-row scalar

which is just the row-wise dot product of the output gradient with the output — a cheap \(O(Nd)\) quantity computable up front. Then

The crucial enabling fact: we can reconstruct \(P_{ij}\) in any tile without storing it, because \(P_{ij} = \exp(S_{ij} - L_i)\), where \(S_{ij} = (Q_i K_j^\top)/\sqrt d\) is recomputed and \(L_i\) is the saved logsumexp. The logsumexp absorbs both the max-subtraction and the normalization into one number — \(\exp(S_{ij} - L_i)\) is already the normalized probability, numerically stable. This is precisely why the forward pass bothered to save \(L\).

import numpy as np

def flash_attention_backward(Q, K, V, O, L, dO, Br=32, Bc=32, causal=False):
    """
    Backward pass for FlashAttention. Recomputes P tile-by-tile from (Q,K,L)
    instead of storing the N×N matrix. Returns dQ, dK, dV.

    O, L are the forward outputs; dO is the upstream gradient (N, d).
    """
    N, d = Q.shape
    scale = 1.0 / np.sqrt(d)

    dQ = np.zeros_like(Q)
    dK = np.zeros_like(K)
    dV = np.zeros_like(V)

    # Per-row scalar D_i = sum_c dO_{ic} O_{ic}  (rowwise dot of dO and O)
    D = np.sum(dO * O, axis=1)                    # (N,)

    for i0 in range(0, N, Br):
        i1 = min(i0 + Br, N)
        Qi, dOi, Li, Di = Q[i0:i1], dO[i0:i1], L[i0:i1], D[i0:i1]
        dQi = np.zeros_like(Qi)

        for j0 in range(0, N, Bc):
            j1 = min(j0 + Bc, N)
            Kj, Vj = K[j0:j1], V[j0:j1]

            # RECOMPUTE the score tile and the probability tile (no storage)
            Sij = (Qi @ Kj.T) * scale             # (br, bc)
            if causal:
                rows = np.arange(i0, i1)[:, None]
                cols = np.arange(j0, j1)[None, :]
                Sij = np.where(cols <= rows, Sij, -np.inf)
            Pij = np.exp(Sij - Li[:, None])       # = softmax weights, since L is logsumexp

            # gradients flowing through this tile
            dV[j0:j1] += Pij.T @ dOi              # dV = P^T dO   (accumulate over query blocks)
            dPij = dOi @ Vj.T                      # (br, bc)  dP = dO V^T
            dSij = Pij * (dPij - Di[:, None])      # softmax-Jacobian: P*(dP - D)
            dQi   += (dSij @ Kj) * scale           # dQ = dS K / sqrt(d)
            dK[j0:j1] += (dSij.T @ Qi) * scale     # dK = dS^T Q / sqrt(d)

        dQ[i0:i1] = dQi

    return dQ, dK, dV


# ---- gradient check against autograd-style finite differences ----
def reference_attention_loss(Q, K, V, causal=False):
    O = reference_attention(Q, K, V, causal=causal)   # defined earlier
    return O

rng = np.random.default_rng(2)
N, d = 48, 8
Q = rng.normal(size=(N, d)); K = rng.normal(size=(N, d)); V = rng.normal(size=(N, d))
causal = True

O, L = flash_attention_forward(Q, K, V, Br=16, Bc=16, causal=causal)
dO = rng.normal(size=(N, d))                            # arbitrary upstream grad
dQ, dK, dV = flash_attention_backward(Q, K, V, O, L, dO, Br=16, Bc=16, causal=causal)

# numerical gradient of  scalar = sum(dO * O)  w.r.t. Q, via central differences
def num_grad(param, idx):
    eps = 1e-5
    p = param.copy(); orig = p[idx]
    p[idx] = orig + eps; Op = reference_attention_loss(*( (p,K,V) if param is Q else
                                  (Q,p,V) if param is K else (Q,K,p)), causal=causal)
    fp = np.sum(dO * Op)
    p[idx] = orig - eps; Om = reference_attention_loss(*( (p,K,V) if param is Q else
                                  (Q,p,V) if param is K else (Q,K,p)), causal=causal)
    fm = np.sum(dO * Om)
    return (fp - fm) / (2 * eps)

err = max(abs(num_grad(Q, (3, 2)) - dQ[3, 2]),
          abs(num_grad(K, (5, 1)) - dK[5, 1]),
          abs(num_grad(V, (7, 4)) - dV[7, 4]))
print("max grad-check error:", err)     # ~1e-7: analytic backward matches finite differences

The gradient check passes to \(\sim 10^{-7}\), the expected precision of central differences in fp64. The backward pass is genuinely exact. Two practical notes mirror the real kernel:

• \(\mathrm{d}V\) and \(\mathrm{d}K\) accumulate across query blocks (every query attends to a given key), so in a parallel implementation they require atomic adds or a separate reduction pass. FlashAttention-1 handles this by looping keys in the outer loop for the backward pass (the opposite of the forward), so each key block’s \(\mathrm{d}K_j, \mathrm{d}V_j\) is finalized in one go while \(\mathrm{d}Q\) is accumulated with atomics. The math above is loop-order-agnostic; the implementation chooses an order that minimizes atomics. A practical consequence: because floating-point addition is not associative and atomics land in nondeterministic order, the FlashAttention backward pass is run-to-run nondeterministic by default. If you are bisecting a training divergence and need bitwise reproducibility, set deterministic=True in the flash-attn API (it uses a slower deterministic reduction) or torch.use_deterministic_algorithms(True), and expect to pay for it.

• Recomputation cost is small and overlaps. The backward pass redoes the \(QK^\top\) matmul (the score tile) it could have stored. That is one extra \(O(N^2 d)\) matmul — but matmuls run on tensor cores that the memory-bound kernel was leaving idle anyway, and we save the \(O(N^2)\) HBM read/write of the stored \(P\). On the memory-bound side of the roofline, trading FLOPs for bytes is exactly the right direction.

Wiring it into PyTorch autograd¶

A forward and a backward function are not yet a usable layer. PyTorch’s autograd does not know that our tiled kernel implements attention, and if we let it trace the tile loop it would build a graph over every intermediate — reintroducing exactly the \(O(N^2)\) storage we removed. The bridge is torch.autograd.Function: we register the pair as one opaque op, declare what it saves, and hand back gradients ourselves. This is the same interface the real flash-attn package and every Triton attention kernel use, and it is worth writing once so nothing about the training loop stays a black box.

import torch

class FlashAttentionFn(torch.autograd.Function):
    """Turn the tiled kernel into a first-class PyTorch op.

    The body calls the NumPy routines above; a production version swaps them
    for a Triton or CUDA kernel. What matters is the *contract*: forward saves
    only (Q, K, V, O, L) -- O(N*d) state, never an NxN tensor -- and backward
    rebuilds every probability tile from L. That contract is what lets
    FlashAttention drop into any training loop.
    """

    @staticmethod
    def forward(ctx, Q, K, V, causal=False, Br=32, Bc=32):
        npy = lambda t: t.detach().cpu().double().numpy()
        O_np, L_np = flash_attention_forward(npy(Q), npy(K), npy(V), Br, Bc, causal)
        O = torch.as_tensor(O_np).to(Q)          # .to(Q) matches dtype AND device
        L = torch.as_tensor(L_np).to(Q)
        ctx.save_for_backward(Q, K, V, O, L)     # note: no NxN tensor is saved
        ctx.hparams = (causal, Br, Bc)           # non-tensors go on ctx directly
        return O

    @staticmethod
    def backward(ctx, dO):
        Q, K, V, O, L = ctx.saved_tensors
        causal, Br, Bc = ctx.hparams
        npy = lambda t: t.detach().cpu().double().numpy()
        dQ, dK, dV = flash_attention_backward(
            npy(Q), npy(K), npy(V), npy(O), npy(L), npy(dO), Br, Bc, causal)
        back = lambda a: torch.as_tensor(a).to(Q)
        # one return value per forward input; None for the non-tensor arguments
        return back(dQ), back(dK), back(dV), None, None, None

flash_attention = FlashAttentionFn.apply

# ---- check the op against PyTorch's own attention AND its autograd ----
torch.manual_seed(0)
N, d = 96, 16
Q, K, V = (torch.randn(N, d, dtype=torch.float64, requires_grad=True) for _ in range(3))

flash_attention(Q, K, V, True).square().sum().backward()   # our op
g_flash = [t.grad.clone() for t in (Q, K, V)]
for t in (Q, K, V):
    t.grad = None

torch.nn.functional.scaled_dot_product_attention(       # PyTorch's reference op
    Q, K, V, is_causal=True).square().sum().backward()
g_ref = [t.grad.clone() for t in (Q, K, V)]

print("max grad error vs PyTorch autograd:",
      max((a - b).abs().max().item() for a, b in zip(g_flash, g_ref)))   # ~1e-15

Three details generalize to any custom kernel you write. save_for_backward is for tensors only — it participates in autograd’s reference bookkeeping and version-counter checks — while non-tensor configuration is stashed on ctx as ordinary attributes. backward must return exactly as many values as forward took arguments, None for those that need no gradient. And the whole point of the exercise is visible in one line: the saved set is \((Q, K, V, O, L)\), all \(O(Nd)\) or smaller, which is why swapping this op into a transformer block turns attention’s activation memory from quadratic to linear. Replacing the NumPy body with a real GPU kernel is the subject of Writing GPU Kernels with Triton; the autograd.Function wrapper around it stays exactly as written here.

IO-aware analysis: counting the bytes¶

We have claimed FlashAttention is faster because it moves fewer bytes. Let us prove it with an IO-complexity argument — the kind of analysis the paper is named for, and the kind an interviewer loves.

The model: HBM is large and slow; SRAM has size \(M\) (in elements) and is fast. We count HBM accesses (reads + writes), because that is what dominates runtime for a memory-bound kernel. The inputs \(Q, K, V\) are \(N \times d\); assume \(d \le M\) (a few head-dimensions of vectors fit in SRAM, which is always true in practice).

Naive attention. It writes and reads the \(N \times N\) matrices \(S\) and \(P\). Even with perfect overlap, \(S\) is written once and read once, \(P\) is written once and read once: that is \(\Theta(N^2)\) HBM accesses, plus \(\Theta(Nd)\) for \(Q,K,V,O\). Total:

FlashAttention. Choose block sizes so that one \(Q\) tile, one \(K\) tile, one \(V\) tile, and the accumulators all fit in SRAM. With SRAM size \(M\), the column block can be \(B_c = \Theta(M/d)\) and the row block \(B_r = \Theta(M/d)\) as well (bounded so the \(B_r \times B_c\) score tile \(\le M\)). Now count: the outer loop runs \(N / B_r\) times; for each query block, the inner loop streams all of \(K\) and \(V\) from HBM once — that is \(\Theta(Nd)\) bytes per query block. So:

The last step uses \(B_r = \Theta(M/d)\). Compare the two: FlashAttention’s HBM traffic is smaller than naive’s \(\Theta(N^2)\) by a factor of \(\Theta(M/d^2)\). With SRAM on the order of \(M \sim 10^5\) elements and \(d^2 \sim 10^4\), that is an order-of-magnitude reduction in HBM traffic — and it is provably (the paper shows) optimal up to constants over all algorithms that compute exact attention by this tiling. The bigger the SRAM, the fewer passes over \(K, V\), the less traffic. This is the formal content of “IO-awareness”: the algorithm’s cost is dominated by HBM accesses, so we design the algorithm to minimize them, not to minimize FLOPs.

                  HBM accesses        SRAM scratch        extra HBM memory
  ---------------  -----------------   ----------------    ----------------
  Naive attention   Θ(N²)              O(Br·Bc) but it      Θ(N²)  (S and P)
                                       still spills S,P
  FlashAttention    Θ(N² d² / M)       O(Br·Bc) tile +      Θ(N)   (O and L)
                                       O(Br·d) accumulators

  M = SRAM size (elements). For typical M and d, FlashAttention moves ~M/d²
  fewer bytes — about an order of magnitude — and uses O(N) instead of O(N²)
  scratch memory. Same FLOPs, same answer, far fewer HBM trips.

Two consequences fall out of this accounting:

1. It is exact, and the speedup is “free.” FlashAttention computes the identical function as naive attention (we proved it above, twice). The speedup comes purely from data movement, not from any approximation, sparsity, or low-rank trick. This is what distinguishes it from earlier “efficient attention” methods (Linformer, Performer, sparse attention — see Beyond Attention: SSMs, Mamba, RWKV & Linear Attention) that traded accuracy for speed. FlashAttention trades nothing.

2. Memory drops from \(O(N^2)\) to \(O(N)\). The only persistent extra allocations are \(O\) (\(N \times d\)) and \(L\) (\(N\)). This linear scaling is what unlocked long-context training and inference; it is the foundation under nearly every long-context model shipped since 2022, and under the prefill phase of The Anatomy of LLM Inference.

Practical use: you almost never write this yourself¶

You now understand the algorithm well enough to implement it. In practice you call it — this is the single line every training loop in this book depends on, including Stack-100M’s, whose attention block calls it directly (The Stack-100M Architecture) and whose 8K-context extension is only affordable because of the \(O(N^2) \to O(N)\) memory drop (The Pretraining Run). PyTorch ships a fused, hardware-dispatched attention that uses a FlashAttention backend when shapes and dtypes allow:

import torch
import torch.nn.functional as F

# q, k, v : (batch, heads, seq_len, head_dim), on CUDA, in fp16/bf16
q = torch.randn(2, 16, 8192, 128, device="cuda", dtype=torch.bfloat16)
k = torch.randn_like(q)
v = torch.randn_like(q)

# Single fused call. PyTorch dispatches to a FlashAttention kernel when eligible:
# contiguous, supported head_dim, fp16/bf16, etc. is_causal applies the triangular mask
# WITHOUT materializing it (the kernel skips upper-triangular key blocks, as we discussed).
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)   # (2, 16, 8192, 128)

# You can inspect / force which backend is used:
from torch.nn.attention import sdpa_kernel, SDPBackend
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    out = F.scaled_dot_product_attention(q, k, v, is_causal=True)

Things that silently disable the FlashAttention backend and fall back to a slower math path — worth memorizing because they show up in profiling:

• Unsupported head_dim. Older kernels support head dims up to 128 (later up to 256). A head dim of 160, say, may fall back.
• fp32 inputs. FlashAttention kernels target fp16/bf16. Pass fp32 and you get the math fallback.
• An additive float mask with awkward shape, or a mask that the kernel cannot fuse, can force the materialized path. Prefer is_causal=True over hand-built masks when possible.
• Non-contiguous tensors or unsupported strides.

The official flash-attn package (Dao et al.) exposes the kernels directly, with more knobs than SDPA surfaces: sliding-window attention, ALiBi slopes, dropout fused into the kernel, deterministic backward, and a variable-length API that avoids padding waste.

# pip install flash-attn --no-build-isolation      # compiles CUDA kernels; needs nvcc
from flash_attn import flash_attn_func, flash_attn_varlen_func

# GOTCHA: flash-attn takes (batch, seqlen, heads, head_dim) -- the seq and head
# axes are TRANSPOSED relative to SDPA's (batch, heads, seqlen, head_dim).
qf, kf, vf = (t.transpose(1, 2).contiguous() for t in (q, k, v))    # from the SDPA layout
out = flash_attn_func(qf, kf, vf, causal=True)                     # (B, S, H, D)
out = flash_attn_func(qf, kf, vf, causal=True, window_size=(1023, 0))  # sliding window

# Packed/ragged batches: concatenate documents with NO padding and pass the
# cumulative sequence-length offsets, so no cross-document score is ever computed.
out = flash_attn_varlen_func(q_packed, k_packed, v_packed,
                             cu_seqlens_q=cu, cu_seqlens_k=cu,     # int32, len B+1
                             max_seqlen_q=max_len, max_seqlen_k=max_len,
                             causal=True)

flash_attn_varlen_func and its cu_seqlens convention are what Megatron-LM, TRL, and most production pretraining loops use for packed sequences — see Chat Templates, Data Formatting & Sequence Packing and, for the document-masking version of the same problem, Mid-Training: Quality Annealing, Long-Context Extension & Capability Injection. One level up, HF Transformers picks the backend for you with a single argument, and it is worth setting explicitly rather than trusting the default:

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",   # "sdpa" (default) | "eager" | "flex_attention"
)

"eager" is the materialized-\(N^2\) path — correct, useful for debugging or when you need the attention probabilities returned, and never what you want for a real run. For writing your own fused attention variants (custom masks, custom biases) the modern path is Writing GPU Kernels with Triton, where a FlashAttention-style kernel is the canonical tutorial example, or PyTorch’s flex_attention, which compiles a user-supplied score_mod/mask_mod into an FA-style kernel.

Putting it together: the mental model¶

Step back and hold the whole thing at once. Standard attention is correct but pathological on a GPU: it writes a quadratic intermediate to the slow memory and reads it back twice, so it is bandwidth-starved and capacity-limited. The single obstacle to fixing this is that softmax normalization seems to need the whole row before it can emit anything. The online softmax dissolves that obstacle with a running maximum, a running denominator, and a correction factor \(e^{m_{\text{old}} - m_{\text{new}}}\) that retroactively re-bases partial results whenever a bigger logit shows up — letting us compute a numerically-stable, exactly-correct softmax-weighted sum in one streaming pass over blocks. Wrap that streaming softmax in a tile loop that keeps every working block in SRAM, normalize once at the end, and save a single logsumexp per row so the backward pass can recompute the probabilities instead of storing them. The result moves an order of magnitude fewer bytes, uses linear instead of quadratic memory, and computes precisely the same function. No approximation. That combination — an elegant numerical recurrence married to a hardware-aware data layout — is why FlashAttention is one of the most consequential systems papers of the LLM era, and why “IO-awareness” is now a design principle, not just a kernel.

Further reading¶

• Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré — FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022). The original paper; read it for the IO-complexity theorem and the forward/backward algorithm boxes.
• Tri Dao — FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (2023). The follow-up on loop ordering, non-matmul FLOP reduction, and occupancy; covered in FlashAttention 2 & 3.
• Maxim Milakov and Natalia Gimelshein — Online normalizer calculation for softmax (2018). The streaming softmax recurrence that FlashAttention builds on.
• Markus N. Rabe and Charles Staats — Self-attention Does Not Need \(O(n^2)\) Memory (2021). An earlier memory-reduction-by-recomputation result for attention.
• Ashish Vaswani et al. — Attention Is All You Need (2017). The attention operator FlashAttention accelerates; the from-scratch version is in The Attention Mechanism From Scratch.
• The Dao-AILab/flash-attention GitHub repository — production CUDA kernels, the flash_attn Python API, and the varlen/packed-sequence entry points.
• The OpenAI Triton tutorials — the fused-attention example is the canonical place to see a FlashAttention-style kernel written in a high-level GPU language; see Writing GPU Kernels with Triton.

Exercises¶

1. Naive attention and FlashAttention execute the same number of floating-point operations, yet FlashAttention is far faster on a real GPU. In one or two paragraphs, explain (a) why the FLOP count is not what determines the wall-clock time here, and (b) precisely which step in the naive kernel is responsible for the slowdown. Then state what would happen to the relative speedup if you ran attention on a hypothetical GPU whose HBM bandwidth equalled its SRAM bandwidth.

2. Stream the logits \(x = [3,\,1,\,6,\,2]\) through the online-softmax recurrence with block size 2 (block one is \([3,1]\), block two is \([6,2]\)). Carry the running max \(m\) and running denominator \(\ell\) by hand, showing the correction factor \(\alpha\) explicitly at the second block. Then (a) confirm your final \(\ell\) matches the two-pass reference \(\sum_i e^{x_i - \max_j x_j}\), and (b) report the row logsumexp \(L = m + \log \ell\) that the forward pass would save for the backward pass. Use \(e^{-2}=0.135335\), \(e^{-3}=0.049787\), \(e^{-4}=0.018316\), \(e^{-5}=0.006738\).

3. Consider a single-head forward pass at sequence length \(N = 16{,}384\) and head dimension \(d = 128\), in fp16 (2 bytes/element); the saved logsumexp \(L\) is fp32 (4 bytes/element).

(a) How many bytes does the naive score matrix \(S\) occupy?

(b) FlashAttention’s only persistent extra allocations are the output \(O\) (\(N \times d\)) and the per-row statistic \(L\) (\(N\)). How many bytes is that, and what is the ratio of (a) to (b)?

© The chapter gives FlashAttention’s HBM traffic as \(\Theta(N^2 d^2 / M)\) versus naive’s \(\Theta(N^2)\), i.e. a reduction factor of \(M/d^2\) where \(M\) is SRAM size in elements. If an SM has \(M = 131{,}072\) elements of usable SRAM, what is the reduction factor at this \(d\)?

4. Modify the chapter’s flash_attention_forward to implement sliding-window causal attention: query \(i\) attends only to keys \(j\) with \(i - W < j \le i\) (a causal window of width \(W\)). Write the changed inner loop, exploit the window to skip key-blocks that fall entirely outside it (both future blocks and blocks past the left edge), and verify against a materialized reference. Keep the chapter’s -inf guards so fully-masked rows do not produce nan.

import numpy as np

def flash_attention_forward_swa(Q, K, V, W, Br=32, Bc=32):
    """Sliding-window causal attention: query i attends to keys i-W < j <= i."""
    N, d = Q.shape
    scale = 1.0 / np.sqrt(d)
    O = np.zeros((N, d))
    L = np.zeros(N)

    for i0 in range(0, N, Br):
        i1 = min(i0 + Br, N)
        Qi = Q[i0:i1]
        br = i1 - i0

        m_i = np.full(br, -np.inf)
        l_i = np.zeros(br)
        O_i = np.zeros((br, d))

        for j0 in range(0, N, Bc):
            j1 = min(j0 + Bc, N)

            # --- window-aware block skipping ---
            if j0 > i1 - 1:          # block entirely in the future -> and all later blocks too
                break
            if j1 - 1 < i0 - W + 1:  # block entirely left of the window for every row
                continue

            Kj = K[j0:j1]
            Vj = V[j0:j1]
            Sij = (Qi @ Kj.T) * scale

            rows = np.arange(i0, i1)[:, None]
            cols = np.arange(j0, j1)[None, :]
            keep = (cols <= rows) & (cols > rows - W)      # causal AND window
            Sij = np.where(keep, Sij, -np.inf)

            # --- online-softmax update (unchanged from the chapter) ---
            m_block = Sij.max(axis=1)
            m_new = np.maximum(m_i, m_block)
            m_new = np.where(np.isneginf(m_new), 0.0, m_new)

            P = np.exp(Sij - m_new[:, None])
            alpha = np.exp(m_i - m_new)
            alpha = np.where(np.isneginf(m_i), 0.0, alpha)

            l_i = alpha * l_i + P.sum(axis=1)
            O_i = alpha[:, None] * O_i + P @ Vj
            m_i = m_new

        l_safe = np.where(l_i == 0.0, 1.0, l_i)
        O[i0:i1] = O_i / l_safe[:, None]
        L[i0:i1] = m_i + np.log(l_safe)

    return O, L


# ---- reference and verification ----
def reference_swa(Q, K, V, W):
    d = Q.shape[-1]; N = Q.shape[0]
    S = (Q @ K.T) / np.sqrt(d)
    rows = np.arange(N)[:, None]; cols = np.arange(N)[None, :]
    keep = (cols <= rows) & (cols > rows - W)
    S = np.where(keep, S, -np.inf)
    S = S - S.max(axis=1, keepdims=True)
    P = np.exp(S); P /= P.sum(axis=1, keepdims=True)
    return P @ V

rng = np.random.default_rng(7)
N, d, W = 100, 16, 24
Q = rng.normal(size=(N, d)); K = rng.normal(size=(N, d)); V = rng.normal(size=(N, d))
O_flash, _ = flash_attention_forward_swa(Q, K, V, W, Br=32, Bc=32)
O_ref = reference_swa(Q, K, V, W)
print("max abs error:", np.abs(O_flash - O_ref).max())   # ~1e-15

5. The backward pass reconstructs each probability tile as \(P_{ij} = \exp(S_{ij} - L_i)\) instead of storing \(P\), using the saved logsumexp \(L_i = m_i + \log \ell_i\).

(a) Prove that this reconstruction is already normalized, i.e. that \(\sum_j P_{ij} = 1\) for each row \(i\) (over the row’s visible keys), so no extra division is needed in backward.

(b) At \(N = 16{,}384\), \(d = 128\), fp16, quantify the HBM memory that recomputation saves (storing \(L\) instead of \(P\)) and identify the compensating cost. Reuse numbers from Exercise 3 where useful.