KV Cache — Module 11.1

How This Lesson Fits the Module & Volume

Decoding strategies choose the next token; they do not make generation cheap. Autoregressive Transformers (Vol. 10) would recompute attention over the entire prefix every step if naïve. The KV cache stores past keys and values so each new token only computes its own K/V and attends against the cache.

This is the main systems trick behind interactive inference. Without it, chat latency grows harshly with context length.

Learning Objectives

By the end of this lesson, students should be able to:

Explain why naïve autoregressive decoding recomputes past K/V.
Define the KV cache and what is stored per layer / head.
Distinguish prefill (prompt) from decode (one token at a time).
Estimate memory growth with sequence length and layers.
Implement a minimal cached attention step in PyTorch.
Relate caching to causal masking and next-token loops.

Definition

A KV cache is a per-layer store of key and value tensors for tokens already processed. At generation step t, the model computes Q/K/V only for the new token, concatenates the new K/V onto the cache, and runs attention with Q_t against K_{1:t}, V_{1:t}. Past tokens’ K/V are reused, not recomputed.

Prefill vs Decode

Prefill

Run full prompt; fill cache

Decode

One new token

Append KV

Grow cache by 1

Repeat

Until stop / max len

Phase	Tokens processed	Cache action	Bottleneck feel
Prefill	Full prompt length T_p	Write K/V for all positions	Compute-heavy
Decode	1 new token per step	Append one K/V slice	Memory-bandwidth heavy
Beam (optional)	B hypotheses	Often B caches or shared tricks	Memory × B

With vs Without Cache

No cache

Forward over all t tokens every step.
Recomputes past K/V repeatedly.
Simple but slow for long contexts.

With KV cache

Forward mainly on the new token.
Attention reads cached past K/V.
Memory grows with context.

Systems notes

Cache dtype / quantization matter.
Paged attention manages huge caches.
Multi-user serving shares GPU carefully.

Code: Minimal Cached Attention Step

import math
import torch
import torch.nn.functional as F

def attn_with_cache(q, k_new, v_new, k_cache=None, v_cache=None):
    # q, k_new, v_new: (B, h, 1, d) for a single new token
    if k_cache is None:
        k = k_new
        v = v_new
    else:
        k = torch.cat([k_cache, k_new], dim=2)  # (B, h, T, d)
        v = torch.cat([v_cache, v_new], dim=2)
    scores = (q @ k.transpose(-2, -1)) / math.sqrt(q.size(-1))
    weights = scores.softmax(dim=-1)
    out = weights @ v
    return out, k, v

B, h, d = 1, 2, 8
# Prefill-like: first token
q = k = v = torch.randn(B, h, 1, d)
out, k_cache, v_cache = attn_with_cache(q, k, v)
# Decode: second token reuses cache
q2 = k2 = v2 = torch.randn(B, h, 1, d)
out2, k_cache, v_cache = attn_with_cache(q2, k2, v2, k_cache, v_cache)
print(out2.shape, k_cache.shape)  # (1,2,1,8) and (1,2,2,8)

Strengths and Tradeoffs

Strengths

Turns O(T²) repeated work into incremental updates.
Essential for interactive token streaming.
Compatible with all decoding strategies above.

Tradeoffs

Cache memory can dominate GPU RAM.
Incorrect cache indexing causes subtle wrong generations.
Beam / speculative decoding complicate cache layouts.

Common Misconception

“The KV cache stores the generated text tokens.” It stores key and value activations inside each attention layer—continuous tensors—not the token ID string. The token IDs live in the growing input list; the cache is an acceleration structure for attention.

Related module pages: Inference, Context Window, Hidden State, Key, Value, Masked Attention.

Knowledge Check

Short Answer: What two tensors does the cache store? Answer: Keys and values (per layer/head).
True/False: Without a cache, past K/V are recomputed every decode step. Answer: True.
Multiple Choice: Prefill typically: (a) fills the cache for the prompt, (b) deletes attention, (c) trains the tokenizer. Answer: (a).
Short Answer: At decode time, what length is the new Q usually? Answer: One new position (sequence length 1 for that step).
True/False: The KV cache stores detokenized strings. Answer: False.
Multiple Choice: Cache memory grows mainly with: (a) context length × layers, (b) learning rate, (c) CSS themes. Answer: (a).
Short Answer: Why is causal masking still required conceptually? Answer: Tokens must not attend to the future; cache only stores past positions.
True/False: Top-p sampling replaces the need for a KV cache. Answer: False—orthogonal concerns.
Multiple Choice: Decode phase is often limited by: (a) memory bandwidth, (b) only disk fonts, (c) BLEU. Answer: (a).
Short Answer: Which lecture assembles the full generation loop? Answer: Inference.

Key Takeaways

KV cache reuses past keys/values so each step is incremental.
Prefill writes the prompt cache; decode appends one slice at a time.
Huge win for latency; cache memory is the price.
Orthogonal to sampling vs beam—both can use a cache.
Next: Inference (module capstone).

Trainer’s Guide

Hands-on idea: Time a tiny model generating 64 tokens with and without cache (re-forward full sequence each step); plot wall time vs length.

Discussion prompt: If GPU RAM is tight, would you shorten context, quantize the cache, or reduce batch size first?

Recap: The KV cache makes autoregressive generation efficient by reusing past attention keys and values. Continue with Inference.