KV Cache › The KV Cache Solution

The insight behind KV caching is straightforward once we recognize where the redundancy lies: since each token's key and value vectors depend only on that token's representation and the fixed projection weights, we can compute them once and store them for reuse in all subsequent generation steps. The query vectors, by contrast, must be recomputed because they are used differently at each step. A query asks "what information should I attend to?" and the answer depends on the token's position in the evolving sequence. The key insight is that keys and values are determined entirely by the token that produced them, while queries are only real when there is a full, growing context to attend to.

This asymmetry between queries and key-values is basic to understanding the cache. When token 5 generates a query, it needs to attend to tokens 1 through 4. When token 10 generates a query, it needs to attend to tokens 1 through 9. The keys and values for tokens 1 through 4 are the same in both cases since they haven't changed. But the query's role is different at each step: it always asks about everything that came before, and "everything that came before" grows with each step. Think of queries as questions that depend on who is asking and when, while keys and values are stable answers that remain constant once a token has been processed.

During generation with KV caching, the process changes fundamentally from the naive approach. The two phases, prefill and decode, serve very different computational roles:

Prefill phase (Step 1): Process all prompt tokens at once, computing queries, keys, and values for the entire prompt. Store the computed $K$ and $V$ matrices in the cache, as these will be reused throughout generation. Use the queries to compute attention within the prompt, establishing the initial hidden states. This phase is compute-intensive but happens only once.

Decode phase (Steps 2+): For each new token, the process is efficient:

1. Compute $Q$, $K$, $V$ only for the single new token
2. Append the new $K$ and $V$ vectors to the end of the cache
3. Compute attention using the new $Q$ against all cached $K$ values
4. Weight all cached $V$ values by the resulting attention scores
5. Produce the output for this position

The important difference is that we only compute projections for one token per step (the newly generated one), while the attention computation can freely access all previously computed keys and values through the cache. This transforms the projection work from quadratic to linear in the total sequence length. The attention computation itself remains linear in the cache size at each step (since we compute $q_t K_{\text{cache}}^T$ for a single query), but that is unavoidable because the query must attend over all prior context.

Let's formalize this caching mechanism to understand exactly how it reaches the computational savings. The formalization reveals that caching is not an approximation but an algebraically equivalent reformulation of the attention computation. The outputs are bit-for-bit identical to non-cached attention; only the order and grouping of computations changes.

At generation step $t$, let $x_t$ denote the embedding of the newly generated token. This embedding captures all the information the model has about this token at the input to the attention layer. We project this single token's embedding to obtain its query, key, and value vectors:

$$q_t = x_t W_Q, \quad k_t = x_t W_K, \quad v_t = x_t W_V$$

where:

• $x_t$: the input embedding for the current token at step $t$, a vector of dimension $d_{\text{model}}$
• $W_Q, W_K, W_V$: the fixed projection matrices shared across all positions and all generation steps
• $q_t$: the query vector representing what the current token is seeking to attend to
• $k_t$: the key vector representing how this token can be matched by future queries
• $v_t$: the value vector containing the information this token contributes when attended to

Note the important distinction from the non-cached case: each of these is a single vector of dimension $d_k$ or $d_v$, not a full matrix covering all positions. We compute projections for exactly one token, not the entire sequence.

The cache is the model's memory of past tokens, maintaining the history of key and value vectors computed in all previous steps. At the beginning of step $t$, before processing the new token, the cache contains concatenated matrices:

$$\begin{aligned} K_{\text{cache}} &= [k_1; k_2; \ldots; k_{t-1}] \in \mathbb{R}^{(t-1) \times d_k} \\ V_{\text{cache}} &= [v_1; v_2; \ldots; v_{t-1}] \in \mathbb{R}^{(t-1) \times d_v} \end{aligned}$$

where:

• $K_{\text{cache}}, V_{\text{cache}}$: matrices storing the complete history of keys and values for all previous $t-1$ tokens
• $k_i, v_i$: the key and value vectors computed when token $i$ was first processed
• $[;]$: the concatenation operation along the sequence dimension, stacking vectors as rows
• $d_k, d_v$: the dimensions of the key and value vectors, often equal in practice

This cache is all the "memory" the attention mechanism has of the sequence so far. Each row corresponds to a token's contribution to the attention computation.

After computing the new token's projections, we update the cache by appending the new key and value vectors:

$$\begin{aligned} K_{\text{cache}} &\leftarrow [K_{\text{cache}}; k_t] \\ V_{\text{cache}} &\leftarrow [V_{\text{cache}}; v_t] \end{aligned}$$

where:

• $K_{\text{cache}}, V_{\text{cache}}$: the cached key and value matrices being extended with new entries
• $\leftarrow$: the assignment operator showing an in-place update to the stored cache state
• $k_t, v_t$: the newly computed key and value vectors being appended to preserve the sequence history

After this update, the cache contains keys and values for all $t$ tokens, ready for the attention computation.

With the cache updated, the model computes attention for the new token by having its query vector interact with the entire history:

$$\begin{aligned} a_t &= \text{softmax}\left(\frac{q_t K_{\text{cache}}^T}{\sqrt{d_k}}\right) \in \mathbb{R}^{1 \times t} \\ o_t &= a_t V_{\text{cache}} \in \mathbb{R}^{1 \times d_v} \end{aligned}$$

where:

• $a_t$: the attention weights for the current step, a vector of $t$ probabilities representing how much the new token should attend to each previous position
• $q_t$: the query vector for the current token, seeking needed information from the context
• $K_{\text{cache}}^T$: the transpose of the cached key matrix, shaped for efficient dot product computation with the query
• $V_{\text{cache}}$: the cached value matrix containing the actual content to be retrieved and aggregated
• $d_k$: the dimension of key vectors, giving the scaling factor for numerical stability
• $d_v$: the dimension of value vectors, determining the output size
• $o_t$: the final attention output for step $t$, a weighted combination of all cached values

The computation $q_t K_{\text{cache}}^T$ produces a vector of $t$ scores, one for each cached position. These scores measure how needed each previous token is to the current query. After scaling and applying softmax, these become proper attention weights that sum to 1. Finally, multiplying by $V_{\text{cache}}$ aggregates the cached values according to these weights, creating the output vector that captures all the needed information from the sequence history.

This formulation reaches efficiency because the matrix multiplications involve a single query vector against the cached matrices, requiring $O(t \cdot d)$ operations rather than $O(t^2 \cdot d)$ that would be needed to recompute attention from scratch at every step.