KV Cache Explained

Table of Contents

1. Learning Objectives
2. Prerequisites
3. Notation
4. Core Intuition
5. Autoregressive Generation
6. The Redundancy Problem
7. KV Cache Mechanism
8. Prefill vs. Decode Phases
9. Memory Analysis
10. Multi-Query and Grouped-Query Attention
11. Worked Examples
12. Connection to the Broader Curriculum
13. Common Pitfalls and Misconceptions
14. Research Perspective
15. Summary of Takeaways
16. Exercises

Learning Objectives

After reading this chapter, you should be able to:

1. Describe autoregressive generation as iterative Self-Attention over growing sequences.
2. Explain why naive recomputation costs $O(t^2)$ per step $t$.
3. Define the KV cache and reduce per-token cost to $O(t)$ attention over cached keys/values.
4. Compute KV cache memory: $2 \times B \times L \times H \times n \times d_k \times \text{bytes}$.
5. Compare MHA, MQA, and GQA cache requirements.

Prerequisites

• Self-Attention — Q/K/V projections
• RoPE — position encoding in cached K

Notation

• $t, n$ — Current decode step and total sequence length
• $\mathbf{K}_{\mathrm{cache}}, \mathbf{V}_{\mathrm{cache}}$ — Cached key and value tensors
• $L, H, d_k$ — Layers, heads, and key dimension
• $B$ — Batch size
• $M_{\mathrm{KV}}$ — Total KV cache memory
• $p$ — Prompt (prefill) length

Core Intuition

LLMs generate text one token at a time. At step $t$, Self-Attention needs all previous keys and values. Without caching, recomputing K/V for tokens $1, \ldots, t$ at every step wastes $O(t^2)$ work across generation.

The KV cache stores computed $\mathbf{k}_i, \mathbf{v}_i$ for $i < t$; each new step computes only the new token's Q, K, V and attends over the cache. This is essential for practical inference latency.

Series context. Volume III, Chapter 15 (Inference).

Autoregressive Generation

Definition 1 (Autoregressive Model).

Definition 2 (Generation Step). At step $t$, compute logits from $x_{1:t}$, sample $x_{t+1}$.

Each step runs full Transformer forward pass on sequence of length $t$.

The Redundancy Problem

Proposition 1 (Naive Cost). At decode step $t$, computing $\mathbf{K}, \mathbf{V} \in \mathbb{R}^{t \times d_k}$ for all positions costs $O(t \cdot d)$ per layer. Summing over generation steps $t = 1, \ldots, n$ yields total attention-related work $O(n^2 d)$ per layer.

Proof. At step $t$, a full forward pass processes a sequence of length $t$. Key and value projections require $O(t \cdot d \cdot d_k) = O(t d)$ operations per layer (with $d_k = O(d)$). Summing:

Observation. Keys and values for positions $1, \ldots, t-1$ are identical to their values at step $t-1$ — weights are frozen during inference and prior token representations do not change. Recomputing them at every step is pure redundancy.

KV Cache Mechanism

Definition 3 (KV Cache). Store per layer:

Decode step for new token $x_{t+1}$:

1. Compute $\mathbf{q}_{t+1}, \mathbf{k}_{t+1}, \mathbf{v}_{t+1}$ (single token only)
2. Append: $\mathbf{K}_{\text{cache}} \leftarrow [\mathbf{K}_{\text{cache}}; \mathbf{k}_{t+1}]$
3. Attention: $\text{softmax}(\mathbf{q}_{t+1} \mathbf{K}_{\text{cache}}^T / \sqrt{d_k}) \mathbf{V}_{\text{cache}}$

Theorem 1. Per decode step: $O(t \cdot d_k)$ attention vs. $O(t^2)$ for full recomputation matrix operations at length $t$.

Prefill vs. Decode Phases

Definition 4 (Prefill). Process prompt $x_{1:p}$ in parallel — compute and fill entire cache at once. Cost: $O(p^2 d)$ — one-time.

Definition 5 (Decode). Generate tokens $p+1, \ldots, n$ one at a time using cache. Cost per token: $O(t \cdot d)$ where $t$ grows.

Proposition 2. Prefill is compute-bound (parallel); decode is memory-bandwidth-bound (sequential, cache-heavy).

Memory Analysis

Definition 6 (KV Cache Memory).

where $B$ = batch size, $L$ = layers, $H$ = heads, $n$ = sequence length, factor 2 for K and V.

Example. Llama 70B: $L=80$, $H=64$, $d_k=128$, $n=4096$, FP16:

Often exceeds weight memory for long contexts.

Multi-Query and Grouped-Query Attention

Definition 7 (MQA). Share single $\mathbf{K}, \mathbf{V}$ across all heads; only Q is per-head. Cache reduced by factor $H$.

Definition 8 (GQA). $G$ KV head groups; each group of $H/G$ query heads shares KV. Interpolates MHA ↔ MQA.

Used in Llama 2/3, PaLM — critical for long-context inference with Flash Attention.

Worked Examples

Example 1: Cache Growth

After 1000 tokens, cache stores 1000 key vectors per head per layer.

Example 2: MQA Savings

64-head MHA → MQA: 64× KV memory reduction.

Connection to the Broader Curriculum

• Self-Attention — mechanism
• RoPE — apply rotation when writing to cache
• Flash Attention — IO-efficient attention
• Quantization — INT8/FP8 KV cache

Common Pitfalls and Misconceptions

Pitfall 1: KV cache memory scales linearly with context — limits long sequences.

Pitfall 2: RoPE position must be correct when reading cached K.

Pitfall 3: Prefill and decode have different bottlenecks.

Pitfall 4: Batch inference multiplies cache by $B$.

Research Perspective

The KV cache is implicit in any efficient autoregressive implementation but was formalized in the context of transformer inference scaling. Multi-Query Attention (Shazeer, 2019) demonstrated that sharing key and value heads across query heads reduces cache bandwidth with minimal quality loss. Grouped-Query Attention (Ainslie et al., 2023) interpolates between MHA and MQA and is now standard in Llama 2/3 and related architectures. PagedAttention (Kwon et al., 2023) applies virtual-memory paging to KV storage, eliminating fragmentation in batched serving. Quantized and compressed KV caches (INT8, FP8) remain active research for extending context length under fixed memory budgets.

Summary of Takeaways

• Cache — Store K, V per layer
• Decode — $O(t)$ per step vs. $O(t^2)$ naive
• Memory — (3)
• MQA/GQA — Reduce cache by grouping

Next: RAG

Exercises

Exercise 1. Derive total naive cost $O(n^2 d)$ over $n$ generation steps by summing per-step costs as in equation (2).

Exercise 2. Compute (4) for given model config.

Exercise 3. MQA vs. MHA memory for Llama 7B.

Exercise 4. Prefill vs. decode bottleneck analysis.

Exercise 5. RoPE indexing with cache at position $t$.

Exercise 6. PagedAttention conceptually.

Exercise 7. Batch size effect on (3).

Exercise 8. Connect to Flash Attention memory complexity.