Skip to content
LLM Inference
S12Batching & Scheduling·183 lines

FlashAttention

Attention never needs to materialise the score matrix. Tiling plus a running max and sum turns an O(N²) memory cost into O(N).

  • online softmax
  • tiling
  • rescaling
  • IO-awareness
  • paged variant

Motivation

The problem

Textbook attention computes the full score matrix S = qkᵀ, applies softmax to it, and multiplies by V. That matrix is N×N per head. At a context length of 8,192 with 32 heads in fp16, materialising it costs about 4 GB. That is for one request, in one layer.

Worse than the size is the traffic. The matrix is written to HBM, read back for the softmax, written again, then read a third time for the value multiply. Attention ends up bound by memory round trips rather than by arithmetic; the FLOPs are a small fraction of the time.

The obstacle is the softmax. It needs the maximum over the whole row for numerical stability, and it needs the sum over the whole row to normalise. Both look like they require seeing every score before producing any output.

Core idea

The solution

They do not. Softmax can be computed incrementally and exactly, in one pass, using a running maximum, a running sum, and a correction factor applied whenever the maximum moves.

online softmaxpython
def flash_attention(q, k, v, tile=64):
    d = q.shape[-1]
    m = -np.inf                        # running max
    l = 0.0                            # running sum of exp
    o = np.zeros(d)                    # running weighted sum of v

    for start in range(0, len(k), tile):
        s = (q @ k[start : start + tile].T) / np.sqrt(d)
        m_new = max(m, s.max())
        alpha = np.exp(m - m_new)      # rescale everything accumulated so far
        p = np.exp(s - m_new)
        l = alpha * l + p.sum()
        o = alpha * o + p @ v[start : start + tile]
        m = m_new

    return o / l

The trick is alpha. Everything accumulated so far was normalised against the old maximum; when a later tile contains a larger score, one multiplication corrects all of it at once. No tile is ever revisited, and the result is numerically equivalent to the full softmax rather than an approximation of it.

FlashAttention does not reduce FLOPs

It performs exactly the same arithmetic as standard attention. What it removes is the HBM traffic: the score matrix is created, consumed and discarded inside SRAM, roughly ten times faster than HBM on an H100. The 2–4× speedup comes entirely from not moving data, and the paper's title emphasises IO-awareness for that reason.

Mechanics

How it works

Diagramtiling, and the recurrence that makes it exact
THE ACTUAL BOTTLENECK: THE MEMORY HIERARCHYSRAM · ~33 MB~33 TB/sHBM · 80 GB~3.35 TB/sA ~10× bandwidth gap. Standard attentionwrites the N×N score matrix to HBM andreads it back twice, for softmax.At N=8192 that matrix is 128 MB per head.STANDARD ATTENTIONS = qkᵀwrite S to HBMsoftmax(S)write P, read againmemory: O(N²)HBM traffic dominatesThe arithmetic was neverthe problem. The roundtrips were.FLASH ATTENTIONload tile to SRAMscores + online softmaxaccumulate into onext tilememory: O(N)S never leaves SRAM2–4× faster, same resultTHE ONLINE SOFTMAX RECURRENCE — WHY STREAMING IS EXACT, NOT APPROXIMATEm_new = max(m_old, max(s_tile))α = exp(m_old − m_new)ℓ_new = α·ℓ_old + Σ exp(s_tile − m_new)o_new = α·o_old + Σ exp(s_tile − m_new)·vfinal output = o / ℓWhen a later tile contains a bigger score, everyvalue accumulated so far was normalised againstthe wrong maximum. α fixes all of them with onemultiply — no need to revisit any tile.The result is bit-comparable to full softmax.FlashAttention is not an approximation.The same recurrence lets you split one sequence across GPUs (ring attention) or across blocks (S06).
The memory hierarchy on the left is the motivation. The recurrence along the bottom is the algorithm, in full.

What this unlocks beyond speed

Once attention is a streaming recurrence over key tiles, several other things come almost for free:

  • PagedAttention (S06) — if you are already looping over tiles, the tiles need not be contiguous. Each tile can be a block fetched through the block table.
  • Ring attention — put different tiles on different GPUs, pass the running (m, ℓ, o) state around a ring, and one sequence can be longer than any single GPU's memory.
  • Flash-decoding — at batch 1 there is only one query, so the usual parallelism over queries vanishes. Split the keys across thread blocks instead and combine the partial (m, ℓ, o) states at the end. Same recurrence, applied across processors instead of across time.
memory, not O(N²)
O(N)
long context stops being a memory problem
wall-clock speedup
2–4×
entirely from avoided HBM round trips
not an approximation
exact
differs from full softmax only by float reassociation

How this composes with what you already built

A production decode kernel is FlashAttention's recurrence plus PagedAttention's indirection plus GQA's head sharing, all in one loop: for each block in the block table, load the KV head shared by this query's group, compute scores, update (m, ℓ, o). That single loop is most of what FlashInfer and vLLM's kernels do.

Explore

Try it

Real numbers, real recurrence. Step through the tiles and watch the running maximum, the correction factor and the output accumulator converge onto the exact answer, which is marked with green ticks.

Simulatoronline softmax, tile by tile
step 0 / 4
scores qᵀk/√d · 16 keys, tile 4
0.860.451.450.10————————————
  • not yet loaded — never in memory at once
  • current tile (in SRAM)
  • processed and discarded

Only the orange tile exists at any moment. The grey cells are never materialised, and that is the whole memory saving. FlashAttention does the same arithmetic as the naive version, which is why it is described as IO-aware rather than as a faster algorithm.

the online-softmax recurrence
tilem (running max)correction e^(m_old−m_new)ℓ (running sum)
press run

When a tile contains a larger score than anything seen before, the running maximum moves and every previously accumulated value is rescaled by e^(m_old−m_new), the orange correction below 1. That single multiply is what makes a streaming softmax exact rather than approximate.

output accumulator o / ℓ
o[0]0.952
o[1]-0.048
o[2]0.220
o[3]0.021

Green ticks mark the true attention output. Until the last tile the accumulator is wrong: a valid softmax over the keys seen so far, but not over all of them. That is fine, and it is why you can never early-exit a flash attention loop.

scores in memory
4 of 16
naive would hold
256
rescalings
0
max error vs exact
run to finish

Notice that the accumulator is wrong at every intermediate step: it is a correct softmax over the keys seen so far, which is not what you want. Only after the final tile does it match. Now set tile size to 1. The algorithm still works, does the same arithmetic, and performs a correction on nearly every step. Tile size only decides how much fits in SRAM.

Build it

Implementation

The paged variant is the one you ship. It is the same loop with a block table in front of the key access.

code/s12_flash_attention.py (excerpt)python
def paged_flash_attention(q, block_table, k_cache, v_cache, seq_len, block_size):
    d = q.shape[-1]
    m, l, o = -np.inf, 0.0, np.zeros(d)

    for logical, physical in enumerate(block_table):
        n = min(block_size, seq_len - logical * block_size)
        if n <= 0:
            break
        k_blk = k_cache[physical, :n]      # S06's indirection...
        v_blk = v_cache[physical, :n]
        s = (q @ k_blk.T) / np.sqrt(d)     # ...inside S12's recurrence

        m_new = max(m, s.max())
        alpha = np.exp(m - m_new)
        p = np.exp(s - m_new)
        l, o, m = alpha * l + p.sum(), alpha * o + p @ v_blk, m_new

    return o / l

exp(-inf − -inf) is NaN

On the first tile the running maximum is -inf, so alpha = exp(-inf - m_new) evaluates to 0, which is correct. But if the first tile is entirely masked out, m_new is also -inf and you get exp(nan). Every real kernel special-cases the first iteration or initialises the maximum to a large negative finite number. A fully-masked tile happens routinely with causal masking in a ragged batch.
Run it locally
Implements naive attention, tiled flash attention, paged flash attention and flash-decoding, then asserts all four agree to floating-point tolerance across random inputs, tile sizes and sequence lengths. Reports peak intermediate memory for each.
$ python code/s12_flash_attention.py
Expect: Agreement within ~1e-15, peak memory scaling as O(N²) for naive and O(tile) for flash, and a demonstration of the fully-masked-tile NaN alongside the guard that prevents it.

Only NumPy is required — setup instructions.

Production notes

In production

  • FlashAttention-2 / -3 — better work partitioning across warps and, in v3, asynchrony and FP8 support on Hopper.
  • FlashInfer — the kernel library vLLM and SGLang use; combines paging, GQA and flash in one implementation.
  • picoLM — the C implementation fuses attention into a single online-softmax pass, so no score array is ever written. The win is one pass over the KV with better cache behaviour. It is this chapter's recurrence in the wild, on a CPU.
  • PyTorch — scaled_dot_product_attention dispatches to a flash kernel automatically when the shapes and dtypes allow it.

Exercises

  1. 1
    Implement flash-decoding: split the keys across four "workers", run the recurrence independently on each, then combine their (m, ℓ, o) triples into one. The combination rule is the same rescaling you already wrote.
  2. 2
    Add causal masking to the tiled version and verify the output matches naive causal attention. Then make one tile fully masked and confirm you reproduce the NaN before fixing it.
  3. 3
    Measure peak memory for naive versus flash at N = 512, 2,048, 8,192 and confirm the quadratic-versus-constant scaling empirically.

Continue

What's next

The kernels are efficient. At small batch sizes, though, the GPU spends much of its time waiting for the CPU to tell it what to do next. S13 opens the decoding-acceleration layer with kernel fusion and CUDA graphs.

Check yourself

Questions

Answer before you look. A wrong answer tells you more than a right one, because the explanation points at the part of the chapter to reread.

Check yourselfQuestion 1 of 6

What does FlashAttention actually reduce?

Score 0/6