KV Cache Lab
Why isn’t token 100 a hundred times more work than token 1?
An interactive companion to the lecture. One generated sentence, six ideas — causal attention, wasted recomputation, the KV cache, its memory price, prefill vs. decode latency, and the trade-off between them — each with something to drag, step through, or predict. Explore at your own pace by scrolling, or start the Guided tour for a full-screen, one-step-per-page class presentation.
- 6 ideas
- 5 predict & check questions
- 23 guided steps
- press F for fullscreen
Same sentence, two ways to generate it.
An LLM writes one token at a time. Press Step and watch how much work each approach does at every token. Both produce exactly the same text.
Recall the generation loop from the Decoding Lab: every new token is appended and the loop runs again. The KV cache is what keeps that loop cheap.
- AttentionEvery token looks back
- WasteThe past never changes
- KV cacheKeep it, reuse it
- MemoryWhat it costs
- LatencyPrefill, then decode
Every new token looks back at all of its history.
Token i may attend to tokens 1…i and nothing after. Each row of the score matrix is one query; each filled cell is one query–key score.
The context doubles from 8 to 16 tokens. The number of attention scores grows by about…
Attention work grows with the square of the context.
Old keys and values never change.
Without a cache, every step re-projects K and V for the whole prefix — even though token 3’s K and V are identical at step 3 and at step 12. Step through and click a token to see it.
At step 9, how many of the 8 earlier tokens have K/V that changed since step 8?
Without a cache, almost all of the work is repeated work.
Keep the past. Compute only the new row.
Store each token’s K and V once (hatched = reused). At each step project just the new token, append its K and V, and let the new query read all cached keys. The text is identical; the work is not.
With a KV cache, is the attention step itself free?
The cache skips old work — not the new token’s attention.
The cache trades compute for memory.
Every cached token stores a K and a V in every layer. Memory grows with context length and with the number of sequences (or beam paths) kept alive. Pick a model size and drag.
We double the context and double the batch. KV memory becomes…
Advanced: cache precision (bytes per number)
Quantising the cache to fewer bytes shrinks the bars above proportionally — one common way to fit longer contexts.
KV memory ∝ context length × active sequences.
Prefill once, then one token at a time.
The prompt is processed in parallel (prefill) and produces token 1. Every later token is a separate decode step. Three numbers describe a request: TTFT, inter-token latency (ITL), and end-to-end latency (E2E).
Does a KV cache make the first token (TTFT) arrive faster?
Toy timings with the right shapes (prefill ∝ prompt, cached step ≈ constant, uncached step ∝ context) — not benchmarks.
The cache speeds up decode (ITL, E2E) — TTFT still pays for the prompt.
Benefit and cost both grow with context.
Benefit and cost both grow with context length. Slide the context and watch the decode speed-up (teal) climb alongside the KV memory (purple) against the GPU’s limit.
Illustrative constants: a = 0.25 ms/token (parallel prefill), d₀ = 18 ms (reading the weights each step), d₁ = 0.0035 ms/token (reading the cache). Shapes, not benchmarks.
A KV cache turns a compute problem into a memory problem.
Decode is fast, but one sequence’s cache can take gigabytes — memory is the limit.
Memory multiplies with every active sequence; serving systems budget the cache like a resource.
Each beam path keeps its own history, so B paths cost about B× the cache.