FlashAttention

PagedAttention fixed how the KV cache is stored; FlashAttention fixes how attention is computed. The two are independent wins that stack. Where the last few pages were about packing memory, this one is about memory movement — and it is the cleanest illustration in the book of the roofline lesson from Part 1: a kernel can be slow not because it does too much math, but because it shoves too many bytes in and out of HBM. FlashAttention computes exactly the same attention as standard attention — no approximation — yet runs faster, uses less memory, and unlocks longer context. The trick is entirely about the bytes.

The problem: the n×n score matrix

Standard attention for a sequence of length n with head dimension d computes, per head:

S = Q · Kᵀ        (n×n)   ← scores for every token pair
P = softmax(S)    (n×n)
O = P · V         (n×d)

The villain is S. It has n × n entries, and the textbook implementation writes it to HBM, reads it back to apply softmax, writes it again, then reads it once more to multiply by V. That is several full passes of an n×n matrix across the slow memory bus.

How big is S? For n = 8192 in FP16:

S bytes = 8192 × 8192 × 2  =  134,217,728 bytes  ≈  128 MB   per head, per layer

128 MB shuttled across HBM per attention head per layer, for every forward pass. With dozens of heads and layers, attention’s HBM traffic — and its memory footprint — grows as O(n²). The actual multiply-adds are modest; the kernel is memory-bandwidth-bound, waiting on those reads and writes. And the O(n²) memory is precisely what makes long context blow up.

FLOPs are not the bottleneck here

Standard and FlashAttention do the same number of FLOPs. If FlashAttention is faster, it cannot be because it computes less — it is because it moves less. Keep that distinction sharp; it is the whole point of this page.

The fix: tile through SRAM, fuse the passes

A GPU has a tiny but enormous-bandwidth scratchpad — SRAM, on the order of ~100–200 KB per streaming multiprocessor, with bandwidth roughly an order of magnitude above HBM. FlashAttention is IO-aware: it restructures the computation so the n×n matrix never exists in HBM at all. Instead it streams tiles of Q, K, and V into SRAM, computes each block of the result there, and writes back only the final output O (which is merely n×d).

Standard:  build full S(n×n) in HBM → softmax → ×V        (O(n²) HBM traffic)

FlashAttention:  for each block Kj, Vj:          ← outer loop over K/V tiles
                   for each block Qi:            ← inner loop over Q tiles
                     load Qi, Kj, Vj into SRAM   (fast on-chip memory)
                     Sij = Qi · Kjᵀ              (small tile — lives in SRAM only)
                     update running max, running sum, and Oi   (online softmax)
                 only the final O(n×d) ever touches HBM

The whole S = QKᵀ, softmax, and ×V sequence is fused into one kernel: a score tile is produced, consumed, and discarded inside SRAM before the next tile is loaded. HBM never sees the n×n matrix.

Online softmax: the enabling trick

Softmax seems to need the whole row of scores at once — you must know the row maximum (for numerical stability) and the sum of exponentials before you can normalize. Tiling appears to break this, because you only have one block of the row in SRAM at a time.

The online softmax resolves it. As each new score tile arrives, FlashAttention keeps a running maximum m and a running denominator ℓ, and rescales the partial output it has accumulated so far whenever a larger value appears. After the last tile, the accumulated result is exactly the true softmax-weighted output — no approximation. This is why FlashAttention is mathematically identical to standard attention, not a quality trade.

What it buys, in the three currencies

• Memory: activation memory for attention drops from O(n²) to O(n) — the n×n matrix is never stored. This is what makes long context feasible at all.
• Latency: HBM accesses fall by a large factor (the n×n matrix is no longer read and written multiple times), so the bandwidth-bound kernel speeds up substantially — commonly 2–4× on the attention step for long sequences.
• Dollars: faster attention plus a smaller memory footprint means more sequence length and more batch fit per GPU, and each request finishes sooner — directly lowering dollar-per-token. Same FLOPs, fewer bytes moved, lower bill.

Memory footprint of attention scores
  standard:  O(n²)   ── grows quadratically, caps your max context
  flash:     O(n)    ── grows linearly, context limited by KV cache instead

Roofline payoff

This is the roofline model cashing out. Attention sat far under its compute roof, stalled on memory bandwidth. FlashAttention raises its arithmetic intensity (FLOPs per byte of HBM traffic) by keeping intermediates in SRAM, pushing the kernel toward the compute roof where the silicon is actually busy. No new math, just better data movement.

FlashAttention now ships inside vLLM, TensorRT-LLM, and PyTorch’s scaled-dot-product-attention, usually transparently. It cuts the cost per attention step. The next two pages attack a different axis — the number of expensive steps decode must take at all.

It actually happened — FlashAttention, three versions (2022–2024)

The original FlashAttention (Dao, Fu, Ermon, Rudra, Ré — NeurIPS 2022) argued that attention kernels were leaving most of their speed on the table not through wasted math but through wasted memory movement, and shipped an IO-aware exact kernel to prove it. FlashAttention-2 (Dao, 2023) reworked the GPU work-partitioning, cut non-matmul operations, and skipped fully-masked blocks in causal attention — together roughly doubling throughput over v1. FlashAttention-3 (2024) specialized for the Hopper generation (H100), overlapping the matmul and softmax pipelines and adding FP8 support. Across the line the contract never changed: exactly the same attention, no approximation, just fewer bytes crossing the HBM bus. It is the canonical proof that on memory-bound hardware, how you move data beats how much you compute.

The architect’s lens

FlashAttention is the cleanest “move fewer bytes, not fewer FLOPs” win in the book — the five questions:

• Why does it exist? Standard attention writes the n×n score matrix to HBM and reads it back several times (~128 MB per head, per layer at n=8,192); the kernel is memory-bandwidth-bound, slow from the bytes it shuttles, not the math it does.
• What problem does it solve? Attention’s O(n²) HBM traffic and O(n²) memory footprint — it tiles Q/K/V through on-chip SRAM, computes each score tile and discards it, and fuses QKᵀ → softmax → ×V into one kernel so the n×n matrix never touches HBM.
• What are the trade-offs? Almost none: it does the identical FLOPs and is lossless (online softmax keeps a running max/denominator), so the only price is kernel complexity — IO-aware tiling — paid once by the engine, not you.
• When should I avoid it? Effectively never for real sequence lengths; only at very small n is the n×n matrix cheap enough that the win shrinks. It ships transparently inside vLLM, TensorRT-LLM, and PyTorch SDPA.
• What breaks if I remove it? Attention reverts to materializing n×n in HBM — long context blows up memory (~128 MB per head, per layer at n = 8,192), the bandwidth-starved attention step runs 2–4× slower, and max context is capped by O(n²) memory instead of by the KV cache.

Check your understanding

1. What does standard attention materialize in HBM that FlashAttention never does, and how does its size grow with sequence length?
2. FlashAttention does the same number of FLOPs as standard attention. So where does its speedup come from?
3. What hardware feature does FlashAttention exploit by tiling, and how does its bandwidth compare to HBM?
4. Why is “online softmax” necessary for tiling, and does it change the numerical result?
5. Compute the size of the FP16 score matrix for one head at n = 4096, and state whether FlashAttention ever writes it to HBM.

Show answers

1. The n×n attention score matrix S = QKᵀ. Its size grows as O(n²) — quadratically in sequence length — which is why it dominates attention’s HBM traffic and memory.
2. From moving far fewer bytes: by fusing the passes and keeping score tiles in SRAM, it avoids reading and writing the n×n matrix to HBM multiple times. The kernel is memory-bandwidth-bound, so cutting HBM traffic cuts time even with identical FLOPs.
3. On-chip SRAM (the per-SM scratchpad), whose bandwidth is roughly an order of magnitude higher than HBM. Tiles are loaded into SRAM and the score tiles never leave it.
4. Softmax needs the row’s maximum and sum of exponentials, but tiling only exposes one block of the row at a time. Online softmax maintains a running max and denominator and rescales the partial output as new tiles arrive, yielding the exact (not approximate) softmax result.
5. 4096 × 4096 × 2 bytes = 33,554,432 bytes ≈ 32 MB per head. FlashAttention never writes it to HBM — it is produced and consumed tile-by-tile in SRAM.