Flash Attention

Most explanations of flash attention start with tiling. That's backwards, and it's why they don't land.

Tiling a matmul is easy and obvious — you've probably written one. Tiling attention looks impossible, and the reason is softmax: it needs a denominator summed over the whole row before it can emit any single output value. You cannot process the first block of keys and be finished with it, because a key you haven't looked at yet might have the largest score, which changes the normalization of everything you already computed.

So the interesting question isn't "how do we tile attention." It's "how do we compute softmax without ever seeing the whole row at once?" Answer that and tiling is a corollary. Don't answer it and no amount of tiling helps.

Here's the other thing worth stating immediately, because it's the most common misconception in this entire subject: flash attention is exact. It is not an approximation, not a sparse pattern, not a low-rank projection. It produces bit-comparable output to a naive implementation, up to floating-point reassociation. Every technique in chapter 13 that does approximate attention lost to this one, and that's exactly why.

The problem, quantified#

Standard attention, as written in the paper everyone read first:

Implemented literally, that's three kernels, and between them $S$ and $P$ — each $N \times N$ per head — get written to HBM and read back.

For Llama 3 8B at $N = 8192$ with 32 heads, in bf16, one layer:

8.6 GB of HBM traffic, per layer, per forward pass, to compute something whose inputs ($Q$, $K$, $V$) total 101 MB — an 85× amplification, all of it intermediates. And from chapter 1 we know what that means: this is a memory-bound kernel sitting on a machine that wants 295 FLOPs per byte. The FLOPs were never the problem. The problem is that $S$ and $P$ exist in HBM at all.

So don't put them there. But softmax says you have to. Unless —

Online softmax#

Take the standard numerically-stable softmax over a vector $x \in \mathbb{R}^N$:

The subtraction of $m$ is the standard trick to keep $e^{x_i}$ from overflowing. It also looks like exactly the thing that forces a full pass before you can start.

But watch what happens if you've processed a prefix of the vector and know its max $m^{\text{old}}$ and sum $\ell^{\text{old}}$, and now a new block arrives with its own max $m^{\text{new}}$ and local sum. The combined max is $m' = \max(m^{\text{old}}, m^{\text{new}})$, and every previously-accumulated term was scaled by the wrong max. Fix it by rescaling:

That's it. That's the whole idea.Milakov & Gimelshein, Online normalizer calculation for softmax, 2018 — two years before anyone applied it to attention. Rabe & Staats, Self-attention Does Not Need $O(n^2)$ Memory, 2021, made the attention connection; Dao et al. made it fast. The correction factor $e^{m^{\text{old}} - m'}$ retroactively re-normalizes everything you accumulated under a stale maximum. You never needed the whole row — you needed a running max, a running sum, and the discipline to rescale when the max moves.

Now extend it from the denominator to the output itself. The unnormalized output accumulator $O$ gets the same treatment:

and at the very end you divide once by $\ell$. Each block of keys and values is visited exactly once, $S$ and $P$ never exist in full, and the answer is exact.

In code, stripped to the mechanism:

def flash_attention_reference(Q, K, V, block_size=256):
    """One query block against all key blocks. Mirrors the kernel's
    inner loop; a real implementation keeps every tile in SRAM."""
    N, d = K.shape
    O = np.zeros((Q.shape[0], d), dtype=np.float32)
    m = np.full((Q.shape[0], 1), -np.inf)   # running max
    l = np.zeros((Q.shape[0], 1))           # running sum
 
    for start in range(0, N, block_size):
        Kj, Vj = K[start:start + block_size], V[start:start + block_size]
 
        S = Q @ Kj.T / np.sqrt(d)           # this block only — never stored
        m_new = np.maximum(m, S.max(axis=-1, keepdims=True))
 
        # Rescale the accumulated state to the new maximum, then add.
        correction = np.exp(m - m_new)
        P = np.exp(S - m_new)
 
        l = correction * l + P.sum(axis=-1, keepdims=True)
        O = correction * O + P @ Vj
        m = m_new
 
    return O / l                            # single normalization at the end

Two properties of that loop are worth dwelling on. The state carried between iterations is $(O, m, \ell)$ — size $O(N d)$ and $O(N)$, never $O(N^2)$. And the loop body touches $K_j$ and $V_j$ exactly once, so each key block crosses the HBM boundary once per query block rather than being implied by a materialized $N \times N$ intermediate.

Now tiling makes sense#

With online softmax in hand, the kernel design is mechanical. An H100 SM has 228 KB of shared memory; an A100 has 192 KB. The tile sizes are chosen so that $Q_i$, $K_j$, $V_j$, and the accumulator all fit in SRAM simultaneously — roughly $B_c = \lceil M / 4d \rceil$ for SRAM size $M$ and head dimension $d$. Then:

• Outer loop over query blocks $Q_i$ (in FlashAttention-2; version 1 had the loops the other way, which mattered).
• Inner loop over key/value blocks: load $K_j, V_j$ into SRAM, compute the block's scores, update $(O, m, \ell)$ in registers.
• Write $O_i$ and the statistics $(m_i, \ell_i)$ to HBM once.

HBM traffic drops from $\Theta(Nd + N^2)$ to $\Theta(N^2 d^2 M^{-1})$.Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022. The IO-complexity analysis in §3.2 is the paper's real contribution — it proves the tiling is optimal up to constants. With $d = 128$ and $M \approx 200$ KB, $d^2/M$ is around $1/12$, and the $N^2$ term shrinks by that factor while the quadratic-in-memory term disappears entirely.

The FLOPs are unchanged. Every multiply-add the naive version does, flash does too — it does slightly more, because of the rescaling. What changed is the denominator of the arithmetic intensity, which moved the kernel from far left of the ridge point to the right of it. That's the whole win, and it's a pure chapter 1 argument.

The backward pass#

This is where flash attention and chapter 3 turn out to be the same idea wearing different clothes.

Backward needs $P$ to compute the gradients. $P$ is $N \times N$ and we deliberately never stored it. So flash recomputes it — from $Q$, $K$, and the saved per-row statistics $(m, \ell)$, which are only $O(N)$.

That is textbook activation checkpointing, applied at the granularity of a single fused kernel. And it wins on both axes at once, which is unusual: recomputing $P$ from SRAM-resident tiles is faster than reading a materialized $P$ back from HBM, because the machine is bandwidth-starved and FLOPs are the cheap resource. You save memory and time simultaneously.

This is the concrete reason the previous chapter warned you not to stack selective recomputation on top of flash attention. Flash already did it, better.

Versions, and what each one actually changed#

The three papers are frequently cited interchangeably. They aren't.

FlashAttention (2022) established IO-awareness and the tiled algorithm. Roughly 2–4× faster than the standard implementation, and the first time long-context training was memory-feasible.

FlashAttention-2 (2023) is an occupancy and work-partitioning fix, not a new algorithm. Version 1 spent too much time on non-matmul operations — the rescaling — which is expensive because tensor cores do matmul FLOPs roughly 16× faster than the CUDA cores do everything else. FA2 reduced the number of rescalings, swapped the loop order so query blocks are the outer loop (letting each warp own a slice of output with no inter-warp communication), and parallelized over sequence length as well as batch and heads. About 2× over FA1, reaching 50–73% of theoretical peak on A100.Dao, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, 2023.

FlashAttention-3 (2024) is Hopper-specific. It exploits three things FA2 couldn't: TMA for asynchronous bulk memory transfer, warp specialization (dedicated producer warps fetching tiles while consumer warps compute), and software pipelining that overlaps the softmax — which runs on slow CUDA cores — with the matmuls on tensor cores. Plus fp8 support with incoherent processing to control outlier error. Around 1.5–2× over FA2 on H100, roughly 75% utilization in bf16.Shah et al., FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision, NeurIPS 2024.

The trajectory is worth noting: FA1 was an algorithmic insight, FA2 and FA3 were increasingly hardware-specific engineering. That's what a mature optimization looks like.

Where it doesn't help#

Short sequences. At $N = 512$, attention is a small fraction of block FLOPs and the $N^2$ intermediates fit comfortably. Flash still helps a little; it isn't the difference between working and not.

It doesn't reduce FLOPs. Attention is still $O(N^2 d)$ of compute. Flash makes that compute run near peak instead of near bandwidth — it fixes the constant, not the asymptote. At 1M context the quadratic term dominates your FLOP budget no matter how efficiently you execute it, which is why sparse and linear attention research didn't stop.

It doesn't shrink the KV cache. During decode, the cache is read from HBM in full regardless of how attention is computed. Flash attention is a training and prefill win; the decode bottleneck is chapter 5's problem, and a completely different one.

Head dimension limits. Implementations support specific head dims (64, 128, 256). Unusual architectures fall back to the slow path, which is worth knowing before you design one.