LABS · 14  /  INFERENCE

The Memory Tax.

A model generating text pays a hidden rent: for every token, in every layer, it stores a key and a value so it never recomputes the past. Watch that cache stack up, collapse it with grouped-query attention, slam into the VRAM wall that caps how many users a GPU serves — then quantize it back.

Act 1 · Prefill & decode
Watch the cache grow
Feed a prompt (prefill), then generate tokens one at a time (decode). Each token appends a key and value to the cache — for every layer. Nothing is recomputed; it just accumulates.
Sequence (KV cached per token)
0
Tokens cached
0
K/V vectors stored
across 32 layers
Recompute avoided
vs no-cache
Prefill processes the whole prompt in parallel; decode appends one token per step. Watch the "recompute avoided" climb — without the cache, step N would re-run attention over all N tokens.
Act 2 · Grouped-query attention
Collapse the cache with GQA
Standard attention (MHA) caches a separate key/value for every query head. GQA lets several query heads share one K/V head. Fewer K/V heads = smaller cache — the trick that made long context servable.
Query heads32
KV heads (GQA group)8

Query heads → shared KV heads

Cache reduction
vs full multi-head
4:1
Query : KV heads
— KB
Per token / seq
llama-70B geometry
GQA with 8 KV heads under 64 query heads is an 8× cache cut — for negligible quality loss. Drag KV heads to 1 for multi-query attention (maximum sharing), or up to match query heads for full multi-head.
Act 3 · The concurrency wall
Where the cache caps your users
Weights load once and everyone shares them. But every concurrent user holds their own KV cache for their whole context. Push context and batch up and watch the cache eat the free VRAM — then overflow it.
Context length / user8K
Concurrent users8
GPU:
weights 132 GB (shared)cache 0 GB
0
Max users that fit
at this context
0 GB
Cache demanded
0 GB
Free beside weights
Note what caps you: not compute, but whether another cache fits. This is why providers meter long context so aggressively — a 128K request can cost as much cache as dozens of short ones.
Act 4 · Quantize the cache
Claw the memory back
The cache doesn't have to be FP16. Store keys and values in fewer bits and it shrinks proportionally — buying more concurrent users from the same card, at some risk to quality on long generations.
Cache dtype:
Cache shrink
vs FP16
Users on 80 GB
70B, 32K context
safe
Quality risk
FP8 is the production sweet spot — roughly 2× the users, negligible quality loss. INT4 quarters the cache but its error compounds over long outputs, so validate it on your task before shipping.

Why the cache exists

A transformer generates one token at a time, and each new token must attend to every previous token. Done naively, generating token N re-runs attention over all N tokens — quadratic re-work that gets slower every step. The fix is the KV cache: after processing a token, keep its attention keys and values around. The next step only computes the new token's query and reuses every cached K and V. Generation becomes a linear append.

The trade

You trade compute for memory. The math per token drops enormously; in exchange you store, for every token × every layer × every sequence in flight, a key and a value. That storage is the KV cache — and it grows without mercy.

The size formula

Per token, the cache is 2 × layers × (hidden × kv_heads ÷ heads) × dtype_bytes — the 2 is keys plus values, and the kv_heads ÷ heads ratio is exactly what GQA shrinks. Multiply by context length for one sequence, and by concurrent sequences for the batch. For a 70B model at long context, this reaches tens of gigabytes — rivaling the weights.

Why it rules serving

Decode is memory-bound: the GPU's compute sits mostly idle while it streams weights and cache. So what stops you adding another user is never spare FLOPs — it's whether another full cache fits in the VRAM left beside the weights. The KV cache, not compute, is the concurrency ceiling, which is why every serving optimization (GQA, paged attention, quantized caches, prefix sharing) is really a cache-shrinking trick.

Check yourself

Four questions. The cache remembers if you don't.

Finished this one? 0 / 22 Labs done

Explore the topic

See this alongside everything else on the same subject — handbooks, system designs, challenges and tools, in one place.

More Labs