Phase 1.1 · OptionalOptionalModule 18 of 21

KV Cache

Without caching, generating token 100 requires recomputing attention for all 99 previous tokens. KV cache stores each layer's K and V matrices, appending only the new token's K/V each step. Critical for production inference speed.

Instead of re-reading an entire book from page 1 every time you want to write the next sentence, you keep bookmarks (cache) at each page you've already read and only read the new page.

Visual Workflows

Start here — study each diagram, then use + / − to zoom if needed.

100%

Loading diagram...

Scroll inside the frame · use + / − to zoom

Key Takeaways

1.KV cache stores previously computed Key and Value tensors during autoregressive generation — avoiding redundant recomputation and making inference O(n) per token instead of O(n²). During decode: for each new token, compute Q, K, V for that token only.
2.Append new K, V to cached K, V from all previous tokens. Attention: Q_new @ [K_cache; K_new]ᵀ → softmax → weighted sum of [V_cache; V_new].
3.Memory cost: 2 × num_layers × num_heads × seq_len × d_k × bytes_per_element. For Llama-70B at 4096 context: ~10GB just for KV cache.
4.Optimizations: Multi-Query Attention (MQA), Grouped-Query Attention (GQA) reduce KV heads, PagedAttention (vLLM) manages cache like virtual memory.

Real Example

Scenario

Generating 100 tokens without cache: 1+2+3+...+100 = 5050 attention computations. With KV cache: 100 (one per new token). ~50× speedup for long sequences.

What you would do

In Transformer & ML Foundations, apply KV Cache to this scenario: Generating 100 tokens without cache: 1+2+3+. Identify the inputs, run the technique, validate the output, and note one thing you would monitor in production.

Practice Task

Open the Code Walkthrough below and run it locally. Change one parameter related to KV Cache (e.g. model, temperature, top_k, or tool name), observe the difference in output, and write 2–3 sentences explaining what changed.

Code Walkthrough

Highlighted lines show where KV Cache happens in the code.

KV Cache

1import torch  # import dependencies2
3# Simplified KV cache concept4class KVCache:  # define a data structure or component5    def __init__(self):  # define a reusable function6        self.keys = []7        self.values = []8
9    def update(self, k, v):  # define a reusable function10        self.keys.append(k)11        self.values.append(v)12        return torch.cat(self.keys, dim=2), torch.cat(self.values, dim=2)  # return the result13
14cache = KVCache()  # key line for KV Cache15# Prefill: process prompt tokens, populate cache16# Decode: for each new token17for step in range(5):18    q = torch.randn(1, 8, 1, 64)  # query for new token only19    k = torch.randn(1, 8, 1, 64)  # key for new token only20    v = torch.randn(1, 8, 1, 64)  # value for new token only21    all_k, all_v = cache.update(k, v)  # key line for KV Cache22  # attention: q @ all_k.T -> softmax -> @ all_v

Cheat Sheet

Quick recap

quick ref

•Cache K,V per layer per token
•Prefill: batch all prompt tokens
•Decode: one token at a time
•GQA: shared KV heads
•PagedAttention: vLLM
•Memory ∝ seq_len × layers × heads

Common Mistakes

✕Not using KV cache in production — inference is unusably slow
✕Underestimating KV cache memory for long contexts
✕Confusing KV cache with model weight caching
✕Forgetting cache invalidation when context changes mid-conversation