KV-Cache: Theory & Derivation

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. The Redundancy in Autoregressive Attention
5. KV-Cache: Formal Derivation
6. Memory Growth Analysis
7. Inference Latency with KV-Cache
8. Prefill vs Decode Phases
9. Multi-Layer KV-Cache
10. Common Pitfalls
11. Summary
12. Exercises

Learning Objectives

1. Identify the redundant computation in naive autoregressive attention.
2. Derive the KV-cache update rule and show it reduces per-step complexity from $O(T^2d)$ to $O(Td)$.
3. Compute KV-cache memory as a function of $L$, $H$, $d_k$, and $T$.
4. Distinguish prefill (compute-bound) from decode (memory-bound) phases.
5. Analyze how KV-cache size limits maximum batch size and context length.

Notation

• $\mathbf{K}_{\leq t} \in \mathbb{R}^{t \times d_k}$ — cached keys up to step $t$
• $\mathbf{V}_{\leq t} \in \mathbb{R}^{t \times d_v}$ — cached values up to step $t$
• $\mathbf{q}_t \in \mathbb{R}^{d_k}$ — query at current step $t$
• $L$ — number of layers, $H$ — number of heads

Core Intuition

In autoregressive generation, the model generates one token at a time. At step $t$, it needs to compute attention over all previous tokens $1, \ldots, t$. Without caching, this means recomputing all keys and values from scratch at every step — $O(t)$ redundant work per step, $O(T^2)$ total. The KV-cache stores previously computed keys and values, so each step only computes the new token's key/value and appends them.

The Redundancy in Autoregressive Attention

Naive approach at step $t$: Given all tokens $x_1, \ldots, x_t$:

1. Compute $\mathbf{Q}, \mathbf{K}, \mathbf{V} \in \mathbb{R}^{t \times d_k}$ from all $t$ tokens.
2. Compute attention: $\text{softmax}(\mathbf{QK}^T/\sqrt{d_k})\mathbf{V}$.
3. Take only the last row (the new token's output).

Problem: Steps 1 and 2 have already been done for tokens $1, \ldots, t-1$ at the previous step. We're recomputing $\mathbf{K}_{1:t-1}$ and $\mathbf{V}_{1:t-1}$ redundantly.

Total naive cost: $\sum_{t=1}^T O(t \cdot d) = O(T^2 d)$ — same as full-sequence attention.

KV-Cache: Formal Derivation

Key observation: The projections $\mathbf{k}_i = \mathbf{x}_i\mathbf{W}^K$ and $\mathbf{v}_i = \mathbf{x}_i\mathbf{W}^V$ depend only on token $i$, not on future tokens (due to causal masking).

Cache update at step $t$:

Attention at step $t$ (only for the new token):

Per-step cost: Computing $\mathbf{q}_t\mathbf{K}_{\leq t}^T$ is $O(t \cdot d_k)$. Over $H$ heads: $O(t \cdot d)$.

Total cost for generating $T$ tokens: $\sum_{t=1}^T O(td) = O(T^2d/2)$ — same total FLOPs but now incrementally distributed, avoiding $O(T^2)$ at each step.

Memory Growth Analysis

At step $t$, the KV-cache stores:

Factor 2: one for keys, one for values.

Example: LLaMA-2 70B ($L=80, H=8$ GQA groups, $d_k=128$, FP16):

For $T = 4096$ tokens: $320 \text{ KB} \times 4096 = 1.28 \text{ GB}$ per request.

For $T = 128K$ tokens: $320 \text{ KB} \times 128K = 40 \text{ GB}$ — may exceed GPU memory for a single request!

Inference Latency with KV-Cache

Without cache: At step $t$, load all $t$ token representations through $L$ layers. Cost: $O(t \cdot C_{\text{layer}})$ per step.

With cache: At step $t$, only process 1 new token through $L$ layers + one attention dot product against $t$ cached entries. Cost: $O(C_{\text{layer}} + t \cdot d)$ per step.

The KV read at each layer is:

At large $t$, this KV read dominates the per-step latency (memory-bandwidth-bound).

Prefill vs Decode Phases

Prefill phase (processing the prompt):

• All prompt tokens processed in parallel (one forward pass, batch of $T_{\text{prompt}}$ tokens).
• Compute-bound (large GEMM operations).
• Populates the KV-cache for all prompt tokens.

Decode phase (generating tokens one by one):

• Single token processed per step.
• Memory-bandwidth-bound (loading KV-cache + model weights for tiny batch).
• Latency: dominated by memory reads.

Time to first token (TTFT): Determined by prefill time. Inter-token latency (ITL): Determined by decode step time.

Multi-Layer KV-Cache

Each of the $L$ layers maintains its own KV-cache (keys and values are layer-specific since they depend on that layer's hidden state):

Important: KV-cache cannot be shared across layers because $\mathbf{K}^{(l)} = \mathbf{h}^{(l)}\mathbf{W}_K^{(l)}$ depends on the layer-$l$ representation.

However, with GQA/MQA, the number of distinct K/V heads is reduced:

• MHA: $L \times H \times d_k$ per token per K and V
• GQA: $L \times G \times d_k$ per token (where $G < H$)
• MQA: $L \times d_k$ per token

Common Pitfalls

Pitfall 1. Forgetting that KV-cache grows linearly with context. For a 128K context window, the cache alone can exceed the model weights in memory.

Pitfall 2. Not accounting for KV-cache in batch size planning. With $B$ concurrent requests: total KV memory = $B \times T \times \text{per-token cache}$. This often limits max batch size.

Pitfall 3. Assuming KV-cache is only relevant for generation. During prefill, the cache must be written — this write bandwidth can bottleneck short-prompt, long-generation workloads.

Summary

• KV-cache eliminates redundant key/value computation during autoregressive generation.
• Reduces per-step complexity from $O(T^2d)$ to $O(Td)$ but introduces linear memory growth.
• Cache size per token: $2LGd_k$ elements (with GQA).
• Prefill = parallel, compute-bound. Decode = sequential, memory-bound.
• KV-cache size often determines max context length and batch size.

Exercises

Exercise 1. Compute the KV-cache memory (in GB) for GPT-3 175B ($L=96, H=96, d_k=128$) at $T=2048$ in FP16.

Exercise 2. Derive the total FLOPs saved by using KV-cache vs naive recomputation for generating $T$ tokens.

Exercise 3. For a GPU with 80 GB memory and a model taking 40 GB, compute the maximum batch size $B$ at context length $T=4096$ with per-token cache of 320 KB.

Exercise 4. Explain why the decode phase is memory-bandwidth-bound while prefill is compute-bound.

Exercise 5. Compute the speedup from using GQA ($G=8$) vs MHA ($H=32$) in terms of KV-cache memory and attention compute during decode.