Module 05 · Chapter 15

Why Cache Keys and Values?

Compare recomputing the full sequence with reusing earlier attention results; show why a causal model may save keys and values and label their shapes.

Beginner13 min readUpdated Nov 8, 2026

Learning objectives

Demonstrate the catastrophic quadratic recomputation cost of naive token generation without a cache
Explain why causal attention mathematically guarantees that past Keys and Values never change
Distinguish why Keys and Values must be cached while Queries and Hidden States are discarded
Track the exact tensor shapes of K and V at every layer as sequence length grows
Calculate the computational savings factor (FLOPs reduction) achieved by the KV cache

Before you start

Prefill vs. Decode: A Worked Latency Example (/llm-foundations/prefill-decode-latency-worked-example/)

In the previous lesson, you saw that prefill and decode have dramatically different performance profiles. Prefill processes the whole prompt at once with large matrix multiplications. Decode processes one token at a time, reading the model weights from memory at every step.

But we skipped a crucial question: what actually happens inside the decoder during generation if we do not save anything between steps?

This lesson answers that question from first principles. You will see why naive autoregressive generation without caching recomputes the entire prompt and all past generated tokens at every step — an O(T²) cost per layer, O(T² × N) total FLOPs — and why the causal attention mask mathematically guarantees that past Key and Value vectors are strictly invariant to future tokens. You will also trace exactly what tensor shapes the KV cache holds at each layer, and compute the savings factor that makes practical LLM serving possible.

The Naive Generation Thought Experiment

Suppose you are building a text generator and you decide not to save anything between steps. You process each request as a fresh forward pass through the entire model.

Let the prompt have P tokens: [x₀, x₁, …, x₍ₚ₋₁₎].

Step-by-step without a cache

Step 1: Pass all P tokens through every layer. Predict token xₚ.
Step 2: You now have P + 1 tokens: [x₀, x₁, …, x₍ₚ₋₁₎, xₚ]. Pass all P + 1 tokens through all layers from scratch! Predict token x₍ₚ₊₁₎.
Step 3: Pass all P + 2 tokens from scratch! Predict token x₍ₚ₊₂₎.
Step t: Pass all P + t - 1 tokens from scratch!

Notice what happens: You are recomputing the exact same hidden representations, queries, keys, and values for tokens 0, 1, …, P-1 over and over again.

Total tokens processed across `G` generation steps

Total tokens = Σₜ₌₀ᴳ⁻¹ (P + t)
             = P × G + G(G - 1) / 2

This scales quadratically with generation length! If P = 1000 and G = 1000, you process roughly 1,500,000 token-equivalents instead of just 2,000.

FLOPs consequence

Recall from the previous lesson that a forward pass for one token costs approximately 2 × N FLOPs (where N is model parameter count). Without caching:

Total FLOPs (naive) ≈ 2 × N × [P × G + G(G - 1) / 2]

With a 7B parameter model (N = 7 × 10⁹), P = 500, G = 500:

Tokens evaluated = 500 × 500 + 499 × 500 / 2 = 374,750
Total FLOPs      = 2 × 374,750 × 7 × 10⁹ ≈ 5.25 × 10¹⁵ FLOPs = 5.25 PFLOPs

This is completely impractical. The KV cache exists to avoid this quadratic recomputation.

Why Causal Attention Permits Freezing Past Keys and Values

In a decoder-only autoregressive transformer, every attention layer uses a causal lower-triangular mask.

What “causal” means

At position j, the model can only attend to positions i ≤ j. It never sees or receives information from positions i > j.

Position:    0    1    2    3    4
Attends to:  ✓    ✓    ✓    ✓    ✓
             —    ✓    ✓    ✓    ✓
             —    —    ✓    ✓    ✓
             —    —    —    ✓    ✓
             —    —    —    —    ✓

The key mathematical consequence

When a new token t > j arrives later in the generation:

It cannot change the input at position j (the input is fixed)
It cannot change the activations at position j (causal mask blocks backward flow)
It cannot change the Q/K/V projections at position j (linear projections depend only on the input at j)

Therefore: The Key vector kⱼ = xⱼ @ Wₖ and Value vector vⱼ = xⱼ @ Wᵥ produced at position j are time-invariant. Once computed, they are set in stone.

You can compute kⱼ and vⱼ once, store them in high-speed GPU memory (the KV cache), and never touch them with a matrix multiplication again.

Why Don’t We Cache Queries?

The Query vector qⱼ = xⱼ @ W_q is used exclusively to ask: “What past keys does token j care about?”

Once token j has computed its attention output, its query qⱼ is never needed by any future token. Future token j+1 will bring its own query q₍ⱼ₊₁₎. Past queries can be safely deleted after their attention step completes.

Why Don’t We Cache Attention Outputs or Hidden States?

You might wonder: why not cache the attention output (the weighted sum of values) or the hidden state after the attention block?

Because the new query qₙₑ𝓌 must compute new attention scores against all past keys! The attention distribution dynamically changes with every new token. The softmax weights depend on the new query, so the blending of past values changes at every step.

The only two tensors needed to evaluate dynamic attention against an incoming query are:

Keys — to compute dot products q @ k_pastᵀ (attention scores)
Values — to compute the weighted sum softmax_weights @ v_past

Hence: KV Cache — we cache Keys and Values, nothing else.

Tensor Organization Across Layers

Each transformer layer l ∈ [0, L-1] maintains its own independent KV cache.

At layer l, the cache tensors have shape:

K_layer ∈ [B, H_kv, T_seq, d_head]
V_layer ∈ [B, H_kv, T_seq, d_head]

Where:

B = batch size
H_kv = number of key/value heads (for GQA, this is fewer than query heads)
T_seq = current sequence length (grows by 1 each decode step)
d_head = head dimension (d_model / H_q)

The cache grows along the sequence dimension T_seq at every decode step. At step t, each layer appends one new key and one new value vector per head per batch item.

Hands-On Exercise 1: Recomputation Cost vs. Cached Cost Hand Calculation

Problem: Compare the total FLOPs for generating 500 tokens after a 500-token prompt, with and without a KV cache.

Given:

Model: 7B parameters (N = 7 × 10⁹)
Prompt length: P = 500 tokens
Generation length: G = 500 tokens

Scenario A: Without KV Cache (Naive Recomputation)

Step 1: Total tokens evaluated across all steps

Tokens = Σₜ₌₀⁴⁹⁹ (500 + t)
       = 500 × 500 + (499 × 500) / 2
       = 250,000 + 124,750
       = 374,750 tokens

Step 2: Total FLOPs

FLOPs = 2 × Tokens × N
      = 2 × 374,750 × 7 × 10⁹
      = 5.2465 × 10¹⁵ FLOPs
      = 5,246.5 TFLOPs
      ≈ 5.25 PFLOPs

Scenario B: With KV Cache

Step 1: Prefill tokens (evaluated once)

Prefill tokens = P = 500

Step 2: Decode tokens (one per step)

Decode tokens = G = 500

Step 3: Total tokens evaluated

Total tokens = 500 + 500 = 1,000 tokens

Step 4: Total FLOPs

FLOPs = 2 × 1,000 × 7 × 10⁹
      = 1.4 × 10¹³ FLOPs
      = 14.0 TFLOPs

Step 5: Savings Ratio

Savings = Naive FLOPs / Cached FLOPs
        = 5,246.5 / 14.0
        = 374.75× fewer FLOPs!

Final Answer: The KV cache reduces compute by 374.75× — from 5.25 PFLOPs to 14 TFLOPs. This is why LLM serving is feasible at all.

Hands-On Exercise 2: Causal Invariance Trace on a 3-Token Toy Sequence

Problem: Verify step-by-step that causal attention makes past Keys and Values invariant.

Setup:

1 layer, 1 head, d_head = 2
Wₖ = [[1, 0], [0, 1]] (identity for clarity)
Wᵥ = [[1, 0], [0, 1]] (identity for clarity)
Causal mask: lower-triangular

Step 1: Process token `x₀ = [1.0, 0.0]`

k₀ = x₀ @ Wₖ = [1.0, 0.0] @ I = [1.0, 0.0]
v₀ = x₀ @ Wᵥ = [1.0, 0.0] @ I = [1.0, 0.0]
Cache after Step 1: K = [[1.0, 0.0]], V = [[1.0, 0.0]]

Step 2: Process token `x₁ = [0.0, 1.0]` (with cache)

k₁ = x₁ @ Wₖ = [0.0, 1.0] @ I = [0.0, 1.0]
v₁ = x₁ @ Wᵥ = [0.0, 1.0] @ I = [0.0, 1.0]

Key observation: k₀ and v₀ are retrieved from cache, NOT recomputed.
If we recomputed them: x₀ @ Wₖ = [1.0, 0.0] @ I = [1.0, 0.0] — identical!
Cache after Step 2: K = [[1.0, 0.0], [0.0, 1.0]], V = [[1.0, 0.0], [0.0, 1.0]]

Step 3: Process token `x₂ = [1.0, 1.0]` (with cache)

k₂ = x₂ @ Wₖ = [1.0, 1.0] @ I = [1.0, 1.0]
v₂ = x₂ @ Wᵥ = [1.0, 1.0] @ I = [1.0, 1.0]

Key observation: k₀, v₀, k₁, v₁ are retrieved from cache.
x₂ cannot back-propagate or influence representations at positions 0 or 1!
Cache after Step 3: K = [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]], V = [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]

Conclusion: At every step, previously cached keys and values are bitwise identical to what a full recomputation would produce. The causal mask mathematically guarantees this invariance.

Hands-On Exercise 3: KV Cache Tensor Dimension Tracker Across Layers

Problem: Trace exact tensor shapes as context grows from 10 to 11 tokens.

Given:

L = 32 layers
H_q = 32 query heads
H_kv = 8 key/value heads (Grouped-Query Attention)
d_head = 128
B = 1 (batch size)

Single Layer: Adding the 11th token

New key/value for the incoming token:

k_new ∈ [1, 8, 1, 128]
v_new ∈ [1, 8, 1, 128]

Layer cache BEFORE (T = 10):

K_layer: [1, 8, 10, 128] → 1 × 8 × 10 × 128 = 10,240 elements
V_layer: [1, 8, 10, 128] → 10,240 elements

Layer cache AFTER (T = 11):

K_layer: [1, 8, 11, 128] → 11,264 elements
V_layer: [1, 8, 11, 128] → 11,264 elements

New query for the incoming token:

q_new ∈ [1, 32, 1, 128]

GQA broadcasting:

8 KV heads shared across 32 query heads
Ratio: 32 / 8 = 4 query heads per KV head
Each KV head is read by 4 different query heads

Attention scores shape:

scores = q_new @ K_layerᵀ
       ∈ [1, 32, 1, 128] @ [1, 8, 128, 11]  (after broadcasting KV heads)
       → [1, 32, 1, 11]

→ 11 scores per query head (10 past + 1 current)

Full Model (All 32 Layers)

Total cache elements per tensor (K or V):

Per layer:  11,264 elements
All layers: 32 × 11,264 = 360,448 elements

Memory at FP16 (2 bytes/element):

360,448 × 2 bytes = 720,896 bytes ≈ 0.72 MB per tensor
Both K and V: 1.44 MB for the full model at T = 11

Summary Reference Table: What Is Cached vs. What Is Discarded

Tensor / State	Cached?	Reason
Keys (K)	✅ Yes	Needed by every future query to compute attention scores (`q @ kᵀ`)
Values (V)	✅ Yes	Needed by every future query to compute weighted sum (`softmax @ v`)
Queries (Q)	❌ No	Used once to attend to past; future tokens bring their own queries
Attention scores	❌ No	Recomputed per step because softmax weights change with each new query
Attention output	❌ No	Dynamic blend of past values; changes every step
Hidden states	❌ No	Output of attention + MLP; depends on dynamic attention weights
MLP intermediates	❌ No	Depend on hidden states, which change every step

Primary Source Citations

Vaswani et al. (2017), Attention Is All You Need — arXiv:1706.03762 (Scaled dot-product attention, causal masking, transformer architecture)
Shazeer (2019), Fast Transformer Decoding: One Write-Head is All You Need — arXiv:1911.02150 (Multi-query attention, KV cache mechanics, decoding efficiency)
Pope et al. (2022), Efficiently Scaling Transformer Inference — arXiv:2211.05102 (Prefill/decode distinction, compute vs memory bounds, KV cache mechanics, latency modeling)
Kwon et al. (2023), Efficient Memory Management for Large Language Model Serving with PagedAttention — arXiv:2309.06180 (PagedAttention, continuous batching, memory-efficient KV cache)

What Comes Next

Next lesson: Calculate KV Cache Memory by Hand — Derive the exact byte formula for KV cache memory across context length, layers, head configurations, batch size, and precision.
Previous lesson: Prefill vs. Decode: A Worked Latency Example
Track overview: Start Here: LLM Foundations