KV Cache

A key-value cache stores attention keys and values from earlier tokens in an autoregressive decoder. When the next token arrives, the model can reuse that state instead of projecting the same prefix again.

The saving is in avoiding repeated work on past tokens. The current query still reads the relevant cached keys and values, so attention work and cache memory continue to grow with the context.

In a causal decoder, token $t$ attends to tokens $1..t$. When you later generate token $t+1$, tokens $1..t$ are unchanged, so their keys and values are identical to what you already computed. Recomputing them is pure waste. The KV cache stores them; generation appends the new token's K and V and reuses everything else.

The cache trades repeated computation for stored state. It grows with context length and must be read during attention. Decoding can become memory-bandwidth-bound, depending on batch size, hardware, and the attention implementation.

Without a cache, a naive decoder reprojects the full prefix. For dense $d \times d$ key/value projections, step $t$ takes $O(t d^2)$ projection work. Summing over $n$ generated positions gives $O(n^2 d^2)$ for those projections alone.

With a cache, each step projects only the new token ($O(d^2)$ under the same assumption). Its query still attends to the cached keys and values, requiring $O(t d)$ attention work. Caching removes repeated prefix projection; it does not make attention or memory use constant in context length.

Memory cost of the cache for one sequence:

The factor of 2 is for K and V; count KV heads, which may be fewer than query heads. The formula assumes equally sized key/value heads and excludes allocation overhead. Multiply by batch size for equal-length sequences. The cache can dominate memory for long contexts. Two standard mitigations shrink it by sharing keys/values across query heads: Multi-Query Attention (MQA) uses a single K/V head for all query heads; Grouped-Query Attention (GQA) uses a few: cutting cache size (and bandwidth) by the head-sharing factor with a quality tradeoff that depends on the model and training.

Key takeaway: Cache reuse assumes unchanged weights, prefix, positional scheme, and inference settings. Dense decoding reads growing cached state and can be bandwidth-bound, but weight reads, batching, arithmetic, and allocation also affect serving performance.

This inference-only example receives already projected Q/K/V. Repeated torch.cat makes allocation simple but copies the cache; production implementations usually preallocate or page the storage. Create a fresh instance for each independent sequence.

1import torch
2from torch import Tensor
3
4
5class KVCacheAttention:
6    """Single-head causal attention with an incremental KV cache."""
7
8    def __init__(self) -> None:
9        self.k_cache: Tensor | None = None      # (seq, d)
10        self.v_cache: Tensor | None = None
11
12    @torch.no_grad()
13    def step(self, q_t: Tensor, k_t: Tensor, v_t: Tensor) -> Tensor:
14        # q_t, k_t, v_t: (1, d) for the single new token
15        self.k_cache = k_t if self.k_cache is None else torch.cat([self.k_cache, k_t], dim=0)
16        self.v_cache = v_t if self.v_cache is None else torch.cat([self.v_cache, v_t], dim=0)
17        d = q_t.size(-1)
18        scores = (q_t @ self.k_cache.T) / d**0.5     # (1, seq_so_far): no future to mask
19        weights = torch.softmax(scores, dim=-1)
20        return weights @ self.v_cache                # (1, d)
21
22
23attn = KVCacheAttention()
24d = 64
25for _ in range(5):                                   # generate 5 tokens
26    out = attn.step(torch.randn(1, d), torch.randn(1, d), torch.randn(1, d))
27assert attn.k_cache.shape == (5, d) and out.shape == (1, d)

1. Conceptual: What does the KV cache store and why is caching it correct? (Past tokens' keys and values: they don't change as new tokens are generated in a causal decoder, so they can be reused.)
2. Implementation: How does the cache change the per-step cost of decoding? (From recomputing K/V over the whole prefix each step to computing only the new token's projections plus attention over the cache.)
3. Applied: Write the memory cost of the cache and name what makes it grow. (2 · layers · KV heads · d_head · seq_len · bytes: grows linearly with sequence length and batch size.)
4. Systems-level: Why can autoregressive decoding become memory-bandwidth-bound? (At small batch sizes, reading weights and a growing KV cache can cost more time than the arithmetic. The limiting resource depends on the workload and hardware.)
5. Failure modes: How do MQA and GQA reduce the KV cache? (They share K/V heads across query heads (MQA uses one K/V head, GQA a few) shrinking cache size and bandwidth with model-dependent quality effects.)

Without looking: explain what the KV cache stores and why reuse is valid, give the memory-cost formula, and explain when decoding can become memory-bound. Then name two techniques that shrink the cache. Check against Stages 1–3.

Sources and further reading

• Shazeer, Fast Transformer Decoding: incremental decoding, cache bandwidth, and multi-query attention.
• Ainslie et al., GQA: shared key/value heads and evaluated quality tradeoffs.