Autoregressive Generation › Key-Value Caching

The key insight behind KV caching is that during autoregressive generation, the attention keys and values for positions 1 through $t-1$ don't change when we generate token $t$. Only the new token at position $t$ introduces new keys and values. Instead of recomputing everything, we cache the keys and values from previous positions and only compute the new entries. This changes the scaling: generation that was quadratically expensive becomes linearly expensive, making practical deployment of large language models feasible.

To understand why the keys and values are stable, recall how they are computed. Keys and values are linear projections of the token embeddings at each layer: $k_i = W_K h_i$ and $v_i = W_V h_i$, where $h_i$ is the hidden state at position $i$. In a decoder-only model, the hidden state at position $i$ depends only on tokens at positions $1$ through $i$ due to the causal mask. When we generate token $t+1$, the hidden state at position $i$ (for $i \leq t$) does not change, because the causal mask prevents position $i$ from attending to position $t+1$. This is the important invariant: adding new tokens to the right of the sequence cannot affect the representations of tokens already in the sequence. Therefore, all their keys and values remain valid from the previous step.

Think of the KV cache as a library where each book represents the compressed meaning of one input token as "seen" by each attention layer. Once a book has been written and shelved, it never changes; you only ever add new books to the right end of the shelf. When the model attends from the new token to all prior context, it reads from the existing library without re-reading or re-processing any prior books. The library grows by exactly one book per generation step, keeping the marginal cost of each step constant rather than growing.

To understand why KV caching helps, we first need to understand what happens in standard attention. In multi-head attention, for a sequence of length $T$, we compute QKV projections from the input, then use them to compute a weighted sum:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

where:

• $Q \in \mathbb{R}^{T \times d_k}$: the query matrix containing a query vector for each of the $T$ positions
• $K \in \mathbb{R}^{T \times d_k}$: the key matrix containing a key vector for each position
• $V \in \mathbb{R}^{T \times d_v}$: the value matrix containing a value vector for each position
• $d_k$: the dimension of each key (and query) vector, typically 64 in GPT-2
• $QK^T \in \mathbb{R}^{T \times T}$: the attention score matrix where entry $(i, j)$ measures how much position $i$ attends to position $j$
• $\sqrt{d_k}$: a scaling factor that prevents the dot products from growing too large, which would push softmax into regions with vanishing gradients
• $\text{softmax}(\cdot)$: normalizes each row of the attention scores to sum to 1, creating attention weights

The formula works in three steps: (1) compute attention scores by taking dot products between queries and keys, (2) normalize scores with softmax to get attention weights, and (3) compute weighted sums of values using those weights.

During generation, we only need the query at the last position to predict the next token. But without caching, we must recompute the entire $K$ and $V$ matrices at every step because the attention mechanism requires all previous keys and values to compute the softmax normalization correctly.

With KV caching, we maintain running matrices $K_{\text{cache}}$ and $V_{\text{cache}}$ that accumulate the keys and values for all tokens processed so far. When generating token $t$:

1. Compute only $q_t$, $k_t$, $v_t$ for the new token (a single vector each)
2. Append $k_t$ to $K_{\text{cache}}$ and $v_t$ to $V_{\text{cache}}$ (cache grows by one row)
3. Compute attention using $q_t$ (single query) against the full cache

The attention computation for the new token becomes:

$$\text{Attention}(q_t, K_{\text{cache}}, V_{\text{cache}}) = \text{softmax}\left(\frac{q_t K_{\text{cache}}^T}{\sqrt{d_k}}\right)V_{\text{cache}}$$

where:

• $q_t \in \mathbb{R}^{1 \times d_k}$: the query vector for the single new token at position $t$
• $K_{\text{cache}} \in \mathbb{R}^{t \times d_k}$: the cached key matrix containing keys for all $t$ positions seen so far
• $V_{\text{cache}} \in \mathbb{R}^{t \times d_v}$: the cached value matrix containing values for all $t$ positions
• $q_t K_{\text{cache}}^T \in \mathbb{R}^{1 \times t}$: a single row of attention scores (the new token attending to all previous positions)
• $\sqrt{d_k}$: the same scaling factor as before

The key insight is that $q_t K_{\text{cache}}^T$ produces just a single row of attention scores rather than a full $T \times T$ matrix. We compute $t$ dot products instead of $t^2$, reducing the per-step computation from $O(T^2)$ to $O(T)$. Over the full generation of $T$ tokens, this changes the total cost from $O(T^3)$ to $O(T^2)$, a dramatic improvement for long sequences.

The cache has grown from 10 to 11 positions after processing a single new token. Notice that the input to this step was just one token, not the entire sequence of 11. This is the efficiency gain: we only compute the new keys and values while reusing all previously cached entries.

The shapes reveal the broader efficiency principle. After the initial prompt, each generation step only processes a single token, but the cache grows to include all previous positions. The attention computation uses the single-token query against the full key cache. The computational cost per step is now $O(T)$ (one query dot-producted with $T$ cached keys), rather than $O(T^2)$ (the full attention matrix). Summed over $T$ generation steps, total cost is $O(T^2)$ rather than $O(T^3)$, a qualitative improvement that becomes more dramatic the longer the sequence grows.

The KV cache is per-layer. Each transformer layer maintains its own cache of keys and values, because each layer produces different projections. A model with 12 layers maintains 12 separate key caches and 12 separate value caches. This is why the memory formula for the KV cache includes the number of layers $L$ as a multiplicative factor: cache memory scales with both context length and model depth.

KV caching trades memory for computation. We gain speed by storing intermediate results, but this requires additional GPU memory that grows with sequence length. Understanding this trade-off is needed for deploying models with long context windows. In practice, KV cache memory often becomes the primary constraint on serving large language models: you might have enough GPU memory to load the model weights, but run out of memory when trying to maintain caches for many concurrent users with long conversations.

For a model with $L$ layers, $H$ attention heads per layer, head dimension $d$, batch size $B$, and sequence length $T$, the total cache memory requirement is:

$$\text{Cache memory} = 2 \times L \times B \times H \times d \times T \times \text{sizeof(dtype)}$$

where:

• $2$: accounts for storing both keys and values (two tensors per layer)
• $L$: number of transformer layers (each layer has its own cache)
• $B$: batch size (number of sequences being generated in parallel)
• $H$: number of attention heads per layer
• $d$: dimension of each head (typically $d_{\text{model}} / H$)
• $T$: current sequence length (grows during generation)
• $\text{sizeof(dtype)}$: bytes per element (4 for float32, 2 for float16/bfloat16)

The product $H \times d$ equals the model's hidden dimension $d_{\text{model}}$, so we can also write this as $2 \times L \times B \times d_{\text{model}} \times T \times \text{sizeof(dtype)}$. The key insight is that cache memory grows linearly with both model depth $L$ and sequence length $T$.

For GPT-2 Small (12 layers, 12 heads, 64-dimensional heads, giving $d_{\text{model}} = 768$) with float16 precision:

The custom CachedAttention implementation above reveals the mechanics, but in practice you will use the Hugging Face transformers library, which handles cache management automatically. The library's implementation is more complete: it handles multi-head attention, positional encodings, the interaction between the cache and the causal mask for mixed-length inputs, and GPU memory management. Understanding the principles from the custom implementation allows you to reason about what the library is doing under the hood, even when you do not implement it yourself.

Let's implement efficient generation with KV caching using the Hugging Face transformers library, which handles cache management automatically: