Below the API

The KV Cache

Basic Intermediate 2h Difficulty 4/5

Prerequisites 03, 04, III.08

★ The most important optimization in LLM inference — and the source of most of its problems.

1. What is it?#

During generation, the model stores the key and value tensors computed for every previous token, so it doesn’t have to recompute them.

WITHOUT CACHE — generating token 4:
  recompute K,V for tokens 1,2,3 (and 4)   ← wasteful: 1,2,3 haven't changed!

WITH CACHE — generating token 4:
  read K,V for tokens 1,2,3 from memory
  compute K,V for token 4 only
  append to the cache

Turns generation from O(N²) total work into O(N). In exchange, you must store 2 × L × n_kv_heads × head_dim × bytes per token per sequence — which becomes the binding constraint on how many users you can serve.

Diagram — One decode step with a KV cache#

flowchart TB
  T["New token t"] --> Q["Compute q, k, v for position t only"]
  Q --> AP["Append k and v"]
  AP --> KV[("KV cache<br/>K and V for positions 1..t<br/>per layer, per head")]
  KV --> ATT["Attention: q against all cached K<br/>weighted sum of cached V"]
  Q --> ATT
  ATT --> OUT["Hidden state, logits, token t+1"]
  OUT -.->|"next step"| T

  class KV,AP memory
  class Q,ATT compute
  class T,OUT neutral

2. Why does it exist?#

Because of a property of causal attention: the key and value for position i never change once computed.

K_i = x_i @ W_k     depends only on x_i (and its position, via RoPE)
V_i = x_i @ W_v     same

Token 4’s arrival doesn’t change token 2’s key. So recomputing it is pure waste.

Without the cache, generating N tokens costs:

Σ(i=1..N) [cost of processing i tokens] = O(N²)

For N=1000 with an 8B model, that’s ~500x more attention work than necessary. The KV cache is not an optimization; it’s the difference between viable and unviable.

3. Simple analogy#

Taking notes during a long meeting.

Each time someone speaks, you need to consider everything said so far. Without notes, you’d have to replay the entire meeting in your head before each new comment — impossible by hour two.

With notes, you glance at your summary and add one line. Fast.

The cost: the notebook grows. In a 3-hour meeting with 20 people, you need a lot of paper. And if you’re taking notes for 50 simultaneous meetings (batching), you need 50 notebooks, and you will run out of desk space long before you run out of pen.

Running out of desk space is exactly how LLM servers fail.

4. Tiny example#

Prove it matters, with one attention layer in plain Go:

// kvcache.go — one attention layer generating tokens, with and without a KV cache.
package main

import (
	"fmt"
	"math"
	"time"
)

const d = 256 // hidden size

type Mat [d][d]float32

func (m *Mat) apply(x *[d]float32) (y [d]float32) {
	for i := range y {
		var s float32
		for j := range x {
			s += m[i][j] * x[j]
		}
		y[i] = s
	}
	return y
}

var Wq, Wk, Wv Mat

// attend: the newest token's query against every key; returns the mixed value.
func attend(q *[d]float32, K, V [][d]float32) (out [d]float32) {
	w, sum := make([]float64, len(K)), 0.0
	for t := range K {
		var dot float32
		for i := range q {
			dot += q[i] * K[t][i]
		}
		w[t] = math.Exp(float64(dot) / 16)
		sum += w[t]
	}
	for t := range V {
		p := float32(w[t] / sum)
		for i := range out {
			out[i] += p * V[t][i]
		}
	}
	return out
}

func generate(n int, cached bool) (total, first, last time.Duration) {
	xs := [][d]float32{{1}} // the sequence so far (as hidden vectors)
	var K, V [][d]float32   // the KV cache
	for step := 0; step < n; step++ {
		t0 := time.Now()
		if cached {
			x := &xs[len(xs)-1] // compute K,V for the NEW token only, append
			K, V = append(K, Wk.apply(x)), append(V, Wv.apply(x))
		} else {
			K, V = K[:0], V[:0] // recompute K,V for EVERY token, every step
			for i := range xs {
				K, V = append(K, Wk.apply(&xs[i])), append(V, Wv.apply(&xs[i]))
			}
		}
		q := Wq.apply(&xs[len(xs)-1])
		xs = append(xs, attend(&q, K, V)) // the "next token"
		dt := time.Since(t0)
		total += dt
		if step == 0 {
			first = dt
		}
		last = dt
	}
	return
}

func main() {
	for i := range Wq {
		for j := range Wq[i] {
			Wq[i][j], Wk[i][j], Wv[i][j] = float32((i*7+j)%5-2)/40, float32((i+j*3)%7-3)/40, float32((i*3+j)%3-1)/40
		}
	}
	a, a0, aN := generate(400, false)
	b, b0, bN := generate(400, true)
	ms := func(t time.Duration) float64 { return t.Seconds() * 1e3 }
	fmt.Printf("no cache  : total %6.0f ms, first step %.2f ms, last step %.2f ms\n", ms(a), ms(a0), ms(aN))
	fmt.Printf("with cache: total %6.0f ms, first step %.2f ms, last step %.2f ms\n", ms(b), ms(b0), ms(bN))
	fmt.Printf("speedup: %.1fx\n", float64(a)/float64(b))
}

Typical result:

no cache  : total   5556 ms, first step 0.23 ms, last step 27.62 ms   ← per-token time GROWS
with cache: total     58 ms, first step 0.11 ms, last step 0.20 ms    ← per-token time is FLAT
speedup: 96.4x

This toy has only the part of a layer that the cache removes, so the speedup is dramatic. A real model has per-token work the cache cannot remove (the MLP, the attention read itself), so on a real 0.5B model you measure about 3x at 200 tokens, ~30x at 2,000 and ~300x at 20,000. The gap widens with length, which is exactly why it matters.

Note the shape of the no-cache curve: latency grows linearly with position, because you reprocess the whole sequence each time. If you ever see per-token latency growing linearly in your own system, you have a caching bug.

5. Technical explanation#

What is stored#

Per layer, per sequence:
  K: (n_kv_heads, seq_len, head_dim)
  V: (n_kv_heads, seq_len, head_dim)

Total per token per sequence:
  2 (K and V) × L (layers) × n_kv_heads × head_dim × bytes_per_element

Note: Q is not cached. You only need the query for the current token; past queries are never reused. This is a common point of confusion — the cache is K and V only, which is also why it’s called the KV cache.

The formula, with a worked example#

KV_bytes_per_token = 2 · L · n_kv_heads · head_dim · bytes

Llama-3-8B: L=32, n_kv_heads=8, head_dim=128, BF16 (2 bytes):

2 × 32 × 8 × 128 × 2 = 131,072 bytes = 128 KiB per token
Context     Per sequence      Batch 32       Batch 128
    1k         128 MB           4.0 GB         16 GB
    4k         512 MB          16.0 GB         64 GB
    8k        1024 MB          32.0 GB        128 GB   ← exceeds one H100
   32k        4096 MB         128.0 GB        512 GB
  128k         16.0 GB        512.0 GB          2 TB

At 128k context, ONE sequence needs 16 GB — the model is 16 GB too. File 06 does this arithmetic thoroughly; file 07 explores its consequences.

RoPE and what’s actually cached#

With rotary embeddings, the rotated keys are cached:

k_raw = x @ W_k
k_cached = apply_rope(k_raw, position)     ← position baked in

Consequences:

  1. You cannot reuse cached K at a different position without re-rotating. This constrains prefix caching (file 11): a shared prefix must be at the same absolute positions.
  2. Some systems cache pre-RoPE keys and apply RoPE at read time, trading compute for flexibility. DeepSeek’s MLA does something related.

V is not rotated, so it has no position dependence.

The four operations on a KV cache#

1. ALLOCATE    when a request is admitted (prefill length + expected generation)
2. WRITE       append new K,V each step
3. READ        the attention kernel reads all of it, every step
4. FREE        when the request completes or aborts

Each has a cost:

  • Allocate: must be fast and must not fragment (file 10).
  • Write: 128 KiB per token — trivial.
  • Read: the whole cache, every step — this is the bandwidth cost.
  • Free: must be immediate, or capacity leaks.

The read cost, quantified#

For batch B at context S:

KV bytes read per decode step = KV_per_token × S × B

Llama-3-8B, B=32, S=8192:

128 KiB × 8192 × 32 = 32 GB read per step
Model weights: 16 GB
→ KV reads are 2x the weight reads!

At long context and large batch, the KV cache dominates decode bandwidth, not the weights. This flips your optimization priorities entirely: GQA, KV quantization, and MLA become more valuable than weight quantization. Compute this crossover for your workload (file 06).

6. Under the hood#

A naive contiguous cache:

// Preallocate for the maximum: one contiguous slab per sequence.
type ContiguousKV struct {
	K, V   []float32 // len = layers * nKV * maxSeq * headDim
	layers, nKV, maxSeq, headDim int
}

func (c *ContiguousKV) at(layer, head, t int) int {
	return ((layer*c.nKV+head)*c.maxSeq + t) * c.headDim
}

// Write at step t
func (c *ContiguousKV) Write(layer, head, t int, k, v []float32) {
	i := c.at(layer, head, t)
	copy(c.K[i:i+c.headDim], k)
	copy(c.V[i:i+c.headDim], v)
}

// Read positions 0..t for one head: a single contiguous slice
func (c *ContiguousKV) Keys(layer, head, t int) []float32 {
	return c.K[c.at(layer, head, 0):c.at(layer, head, t+1)]
}

Problems with this, all of which PagedAttention fixes:

  1. You must preallocate for max_seq — a request that generates 50 tokens holds memory for 4096. Measured waste: 60-80%.
  2. Contiguous allocation fragments. After many variable-length requests, you have free memory but no contiguous block.
  3. No sharing. Two requests with the same system prompt each store their own copy.
  4. Growing a sequence requires reallocation if it exceeds its reservation.

File 10 is the fix.

7. Performance implications#

The KV cache is simultaneously:

  • the reason decode is affordable (O(N) not O(N²)),
  • the primary constraint on concurrency (memory),
  • a growing share of decode bandwidth (at long context),
  • the state that makes replicas stateful.

Capacity formula:

max_concurrent_sequences = (GPU_memory − weights − overhead) / (KV_per_token × avg_context)

For Llama-3-8B on one 80 GB H100:

(80 − 16 − 3) GB / (128 KiB × 4096) = 61 GB / 0.5 GB = 122 sequences at 4k
(80 − 16 − 3) GB / (128 KiB × 32768) = 61 / 4 = 15 sequences at 32k

8x more context → 8x fewer users. That’s the fundamental tradeoff of long context.

8. Production implications#

  • KV cache utilization is your primary capacity metric. Put it on the main dashboard. Autoscale on it. Alert on it.
  • Preallocate the KV pool at startup (vLLM’s gpu_memory_utilization, default 0.9) so steady-state allocation never fails.
  • Free blocks immediately on completion or abort. A leak here is a capacity leak.
  • Choose models with GQA/MQA when serving cost matters. An 8:1 GQA ratio is an 8x capacity multiplier.
  • Consider KV quantization (FP8 KV cache) for long-context workloads — Section VII.12.
  • Cap max context per request by tier. Long context is genuinely expensive; price and limit it.

9. Common mistakes#

Forgetting to enable the cache. O(N²) generation; symptom is per-token latency growing linearly.

Sizing capacity by weights only. KV usually exceeds them.

Preallocating for max context per request. 60-80% waste.

Not freeing on abort. Slow capacity leak until the server wedges.

Assuming the KV cache is small. At 128k context it’s larger than the model.

Caching Q. You don’t need it.

Reusing cached K at a different position with RoPE. Silently wrong outputs.

10. Hands-on exercise#

A. Prove the speedup. Run the section 4 benchmark. Plot per-token latency vs position for both. Extend to 1,000 tokens and re-measure the speedup.

B. Compute the formula. For five different open models, compute KV bytes per token. Rank them. Which is cheapest to serve at long context? Does it correlate with parameter count?

C. Find the crossover. For your model, compute the (batch, context) curve where KV read bytes equal weight read bytes. Plot it. Mark where your production traffic sits.

D. Measure real KV usage. On a running vLLM server, query the metrics endpoint for KV cache usage. Drive load until it saturates. What happens at 100%? (Preemption, or rejection — observe which.)

E. Implement a simple cache. Extend your generation loop from file 04 to maintain the KV cache manually (not using past_key_values). Verify it produces identical output. This is Project 05.

11. Interview questions#

  1. What is the KV cache and what problem does it solve? Give the complexity argument.
  2. Why is Q not cached?
  3. Write the KV size formula and compute it for a model you know.
  4. At what context length does the KV cache exceed the model weights for an 8B GQA model?
  5. Why does RoPE constrain KV cache reuse?
  6. What are the four operations on a KV cache and what does each cost?
  7. Your capacity dropped 50% after a change. KV usage is at 100%. What are three likely causes?
  8. Why does long context reduce concurrency proportionally?

12. Further reading#

  • [ESTABLISHED] Kwon et al., “PagedAttention” (SOSP 2023) §2-3 — the memory analysis
  • [ESTABLISHED] Pope et al., “Efficiently Scaling Transformer Inference” (2022)
  • [ESTABLISHED] Shazeer, “Fast Transformer Decoding” (2019) — why KV bandwidth dominates
  • Next: 06 — KV cache math
esc

↑↓ navigate ↵ open