STAGE 0 · 2·A · 3 NODES

FlashAttention

In How attention works you built the query·key score grid — one score for every pair of tokens. That grid is the problem: it grows with the square of the sequence. FlashAttention gets the exact same answer without ever storing it whole.

THE 3 STEPS

01The memory wall 02Stream it in tiles 03Share the keys & values

NODE 01 / 03

The wall: an N×N grid for N tokens

Attention scores every token against every other token. That's an N×N grid — and doubling the sequence quadruples it. This is the memory wall long context runs into.

One score per pair. N tokens means N×N scores — the grid you built in the attention lesson, now drawn to scale.
It grows quadratically. 2× the tokens is 4× the grid. The bar is the bytes to hold it for one head in fp16 (2 bytes each).
Past a point it won't fit. At long context the grid alone blows past a whole GPU's memory — and that's before you count the dozens of heads and layers.
This is the intermediate, not the model. The scores are thrown away after each step; storing them all at once is pure waste.

N×N SCORE GRID · MEMORY vs SEQUENCE

Drag the sequence length — the bar grows quadratically, and many heads and layers push it past a whole GPU.

SEQUENCE LENGTH64k tokens

Every token attends to every token, so the grid is N×N. The bytes to store it grow with the square of the sequence — the core scaling problem.

SCORE_MATRIX
SEQUENCE64k tokens
CELLS4.3B
MEMORY8.6 GB
GPU80 GB
The full N×N score grid for one attention head in fp16, against one H100's 80 GB of memory.

NODE 02 / 03

Never build the whole grid — stream it in tiles

FlashAttention computes the same attention, but block by block. It slides small tiles of the grid through the GPU's tiny fast memory (SRAM), keeping only a running summary — so the full N×N grid is never stored.

Left: naive. The whole grid sits in slow GPU memory (HBM) at once — O(N²) bytes.
Right: FlashAttention. Only one tile is live at a time, plus a running summary per row. Memory is O(N) — flat, no matter how long the sequence.
SRAM is fast but tiny (megabytes). A tile fits; the full grid never could — so streaming tiles is also faster, not just smaller.
The answer is identical. A running-max trick keeps the softmax exact as each tile arrives — same numbers, a fraction of the memory.

NAIVE vs FLASHATTENTION · SAME 32k SEQUENCE

Change the block size — more tiles or fewer, but FlashAttention's memory stays a sliver of the naive bar.

BLOCK SIZE2k · 16 tiles

The block size sets how big each tile is. Bigger tiles mean fewer passes but more SRAM per tile; either way the whole grid is never resident.

TILING · N = 32k
BLOCK2k tokens
TILES16
NAIVE2.1 GB
FLASH18.1 MB
Peak memory for one head over a 32k sequence: the whole grid (naive) versus one live tile plus running row summaries (FlashAttention).

NODE 03 / 03

Let query heads share the keys and values

There's a second memory cost when a model answers: the KV cache — the keys and values of every past token, kept so they aren't recomputed. Giving every query head its own K/V is expensive, so heads share them.

MHA — one K/V set per head. The classic design: 8 query heads, 8 key/value heads. Biggest cache.
GQA — heads share in groups. Several query heads read one shared K/V. Here 8 query heads share 2 sets — a 4× smaller cache, almost no quality loss.
MQA — one K/V for all. Every query head reads a single K/V set: 8× smaller cache, the leanest option.
The bar is the KV cache for a 32k chat across 32 layers. Sharing is why long-context models stay affordable to run.

KV SHARING · QUERY HEADS → K/V HEADS

Cycle MHA → GQA → MQA — the wiring collapses onto fewer K/V heads and the cache bar shrinks 4× then 8×.

MHA (multi-head attention) — every query head has its own key/value head. Full quality, biggest KV cache.

KV_CACHE
SCHEMEMHA
Q HEADS8
K/V HEADS8
CACHE4.3 GB
The key/value cache held during generation — 32 layers, 32k tokens, fp16 — as query heads share fewer K/V heads.

EXPLAIN IT BACK

FlashAttention still computes a score for every pair of tokens. So how does it avoid the N×N memory blow-up?

NEXT: MIXTURE OF EXPERTS

Attention now scales to long context without the memory wall. Next, Mixture of Experts: the other half of a big model — routing each token to just a few expert sub-networks, so the model can be huge while only a slice of it runs per token.

← Back: long context Continue to Mixture of Experts →

Language: English