The KV Cache

In the last lesson (L207) we said decode is memory-bandwidth-bound — each token must stream the model's weights and "the KV cache" out of memory. This lesson is about that cache. It's not a minor optimization: the KV cache is the single thing that makes generation fast enough to be usable at all, and it's simultaneously the biggest memory consumer in LLM serving — the thing that decides how long a context and how large a batch you can afford.

Generation is autoregressive: the model writes one token at a time, each one attending back over every token before it (recall L7). The naive way to do that re-reads the entire story every time it adds a word. The KV cache is the fix — and understanding it explains most of why inference costs what it costs.

In this lesson:

• The problem — why naive decoding is O(n²) and gets slower with every token
• The fix — cache the keys and values of past tokens so each step processes one token, not the whole prefix (O(n))
• Why this is the root of the memory-bound behavior from L207
• The cost — the KV-cache memory formula, real GB numbers, and why it caps batch size and context length
• Shrinking it — MQA, GQA, MLA, and quantization — what every modern model does

Scope: this is the cache itself. How it powers continuous batching for throughput is L209; PagedAttention shows up here as memory management but its batching payoff is L209. Speculative decoding (L210) and vLLM/TGI (L211) build on this. And the KV cache (inside one request) is not the same as prompt caching (reusing a prefix across requests) — that's a cost feature in L214. We'll flag the difference, not cover it here.

Start from how attention works. To produce the next token, the model computes a query (Q) for the current position and compares it against a key (K) for every previous token; the match scores weight a sum of the values (V). Because generation is causal, token t attends over the K and V of all t tokens so far.

Now watch what happens if you implement decoding naively — just call the model on the sequence so far, repeatedly:

Every step reprocesses the entire growing prefix from scratch. Generating n tokens costs 6 + 7 + 8 + … ≈ O(n²) total work, and — worse for UX — each token is slower than the last (TPOT keeps rising; recall L207). For a 2,000-token answer that's millions of redundant token-passes. It's like re-reading the whole book from page one every time you write a single new word.

The waste has a specific shape: you are recomputing the K and V of tokens you already processed a moment ago.

Here's the insight that fixes it: a past token's K and V never change. They're computed from that token (and its position) alone — they don't depend on anything that comes after. So recomputing them every step is pure waste. Compute them once, store them, and reuse them. That store is the KV cache.

With the cache, decoding splits cleanly into the two phases from L207:

• Prefill: run the whole prompt through once, in parallel, and save every token's K and V into the cache. (This is the work behind TTFT.)
• Decode: for each new token, compute K/V/Q for just that one token, append its K and V to the cache, and attend its Q over the cached K/V. Then repeat.

Each decode step is now constant work — one token — instead of reprocessing the prefix. Generating n tokens drops from O(n²) to O(n). In pseudocode the difference is stark:

This is exactly what use_cache=True does in libraries like Hugging Face Transformers, and what every serving engine (vLLM, TGI, SGLang) does under the hood — you rarely write the loop yourself, but this is the mechanism you're paying for. The cache turns an impossible quadratic into a linear stream of single-token steps.

The KV cache also explains the deepest fact from L207 — why decode is memory-bandwidth-bound. Look at what a single decode step actually does: it computes one token's worth of math (tiny), but to do it, it must read the entire model's weights and the entire KV cache out of memory. The compute is small; the bytes moved are huge. That's low arithmetic intensity — the GPU's math units sit mostly idle (~20-40% utilization) waiting on memory.

Prefill is compute-bound (it builds the cache by crunching the whole prompt in parallel). Decode is memory-bound (it reads the cache, one token at a time). The KV cache is the object being streamed — so its size directly drives your TPOT. A bigger cache (longer context) means more bytes to move per token, means a slower decode. Memory footprint and decode speed are the same problem.

This is why the rest of the lesson is about cache size — it's not just a capacity question, it's a speed question too.

How big does the cache get? It's deterministic. You store a K and a V vector for every token, in every layer, for every KV head:

KV bytes = 2 × layers × kv_heads × head_dim × seq_len × batch × bytes_per_element

(The leading 2 is because you store both K and V.) The thing to notice: it grows linearly with sequence length and with batch size — both of which you want to make large. Plug in a Llama-3.1-70B-class model (80 layers, grouped-query attention with 8 KV heads, head_dim 128, BF16 = 2 bytes) and run it:

Output:

per token:   0.33 MB
  4,096 tokens ->    1.3 GB
 32,768 tokens ->   10.7 GB
131,072 tokens ->   42.9 GB
128K @ MHA (64 heads) -> 344.0 GB

Sit with those numbers. At 128K context the KV cache is ~43 GB — that can be larger than a quantized copy of the model's own weights, and it's per concurrent request (multiply by batch size). On an 80 GB GPU, the weights plus a few long-context requests' caches is all your memory — gone.

This is the real reason you can't just "turn up the context length" or "batch more users": the KV cache is the budget. It sets the hard ceiling that batching strategies (L209) spend their cleverness trying to raise. And notice the last line: full multi-head attention would need ~344 GB for the same context — which brings us to how modern models make the cache fit.

Since the cache is the budget, modern models are designed to make each token's K/V smaller. The big lever is how many KV heads you store (the kv_heads term in the formula):

GQA is the modern default because it sits in the sweet spot: query heads share K/V in small groups, giving ~8× smaller cache at quality indistinguishable from full MHA (and you can uptrain an MHA model into GQA with ~5% of pretraining compute). MQA pushes to a single KV head for maximum savings at a small quality cost. MLA (DeepSeek) takes a different route — instead of sharing heads, it compresses what's stored into a low-rank latent, beating even GQA. On top of any of these, KV-cache quantization (store the cache in FP8/INT8 instead of BF16) gives another 2-4×.

There's a second, orthogonal idea: don't shrink the cache, manage its memory better. Naively you'd reserve a contiguous block for each request's maximum length — wasting most of it. PagedAttention (the technique behind vLLM) instead stores the cache in fixed-size pages, allocated on demand like an operating system's virtual memory — no fragmentation, no over-reservation, so you fit far more sequences in the same VRAM. That is what unlocks high-throughput batching — which is the whole of L209. Here, just hold the idea: shrink the cache (GQA/MQA/MLA/quant), then pack it efficiently (paging).

Make it concrete. Below, step through generation and flip the cache OFF vs ON, then switch the attention scheme. Watch the work per step and the memory cost react:

Three things to take away from playing with it:

• Cache OFF turns every cell red — each step re-derives the whole prefix, and the decode-work-so-far counter grows like n². Cache ON keeps the past tokens green (read for free) and computes only the amber new token — O(n), flat work per step. That gap is the lesson.
• The price is the cache itself: it grows every token, and the @128K readout shows it in GB.
• Flip MHA → GQA → MQA → MLA and the @128K cache collapses from ~344 GB to a few GB — the difference between "impossible on one GPU" and "fits comfortably." That's why these schemes exist.

Reason through these, then use the lab (and the calculator above) to confirm:

1. Predict: with the cache ON, why doesn't "decode work this step" grow as you generate more tokens?
2. You double your context length from 64K to 128K. Roughly what happens to (a) the KV-cache memory and (b) the per-token decode time, and why are those the same answer?
3. A model switches from MHA to GQA with 8 KV heads (down from 64). About how much smaller is the KV cache, and what roughly happens to quality?
4. Your serving box OOMs (out of memory) when a few users send very long prompts, even though the model weights fit fine. What's eating the memory, and name two levers to fix it.
5. A teammate says "let's enable the KV cache to make our prompt caching work across requests." What's the confusion?

→ (1) Because past tokens' K/V are already stored — each step computes only the one new token and reads the rest from the cache. Work per step is constant (O(n)). (2) Both double: the cache is linear in sequence length, and since decode is memory-bound, the per-token time tracks how many cache bytes you must read — memory size and decode speed are the same problem. (3) ~8× smaller cache (you store 8 KV heads instead of 64), with quality ≈ MHA — that's why GQA is the default. (4) The KV cache (it grows with context × batch and can exceed the weights). Levers: a smaller-cache attention scheme (GQA/MQA/MLA), KV quantization (FP8/INT8), PagedAttention to pack memory, shorter context, or fewer concurrent long requests. (5) They're different things: the KV cache is the within-request store that accelerates decode (always on); prompt caching reuses a prefix's KV across requests to save cost (L214). Turning on one doesn't give you the other.

• "The KV cache stores the model's answers/outputs." No — it stores the keys and values of past tokens within a single generation, so they aren't recomputed each step.
• "KV cache = prompt caching." Different. KV cache = within-request decode accelerator (always on). Prompt caching = reusing a prefix's KV across requests for cost (L214).
• "Caching saves memory." The opposite — it spends memory to save compute. The cache is one of the largest memory consumers in serving and can exceed the model weights.
• "Longer context is mainly a compute cost." It's mainly a memory cost: the KV cache grows linearly with context and caps your batch size — and, since decode is memory-bound, it slows every token too.
• "GQA/MQA are just speed tricks." They're KV-cache-size reductions (fewer KV heads). The speed and the long-context feasibility both come from a smaller cache; GQA keeps quality ≈ MHA.
• "The KV cache speeds up prefill." No — prefill builds the cache. The cache speeds up decode, by avoiding reprocessing the prefix on every token.
• "Just raise the context window / batch size." Both are gated by KV-cache memory. You raise them by shrinking the cache (GQA/MQA/MLA/quant) or paging it (L209), not by wishing.

• The KV cache is why decode is feasible. Past tokens' K/V never change, so store them once and reuse them — each decode step processes one token instead of the whole prefix. O(n²) → O(n).
• Prefill builds the cache; decode reads and appends to it. That read-the-whole-cache-every-token pattern is why decode is memory-bandwidth-bound (L207) — so cache size drives TPOT.
• The cache is a memory monster: 2 × layers × kv_heads × head_dim × seq_len × batch × bytes. For a 70B GQA model that's ~0.33 MB/token → ~43 GB at 128K context, often bigger than the weights, and per request.
• It caps context length and batch size — it's the budget every serving optimization fights over.
• Shrink it: MHA → GQA (8×, quality ≈ MHA, the default) → MQA (64×, small quality hit) → MLA (DeepSeek latent, ~3-5× over GQA); plus KV quantization (2-4×). Then page it (PagedAttention) to pack memory.
• Don't confuse it with prompt caching (cross-request prefix reuse, L214).
• Next — L209: Batching & Continuous Batching — now that you know the cache is the memory budget, see how continuous batching and PagedAttention pack many requests onto one GPU to turn that budget into throughput.