KV Cache › Cache Structure

The KV cache must be maintained separately for each layer and each attention head in the transformer architecture. This requirement arises because different layers and heads learn different attention patterns and have different projection weights. Layer 1's keys and values are fundamentally different from layer 32's, as each layer captures different levels of abstraction: early layers tend to track syntactic patterns like word boundaries and part-of-speech relationships, while later layers tend to encode more semantic, task-needed information. Similarly, within a single layer, head 1 might learn to track coreference (which pronoun refers to which entity), while head 4 tracks positional relationships, each requiring its own separate cache.

Think of the cache as a filing cabinet with $L$ drawers (one per layer), and within each drawer, $H$ hanging folders (one per head). Every hanging folder contains two documents: the key history and the value history. When the model processes a new token, it opens every drawer, pulls out all the folders, appends one line to each key document and one line to each value document, then puts everything back. The output depends on reading every folder in the needed drawer at every layer.

For a transformer with $L$ layers and $H$ attention heads per layer, the complete cache consists of $2 \times L \times H$ tensors: one $K$ cache and one $V$ cache for each attention head at each layer. This combinatorial structure means that even modest models require tracking many separate cache tensors. For LLaMA-2 7B with 32 layers and 32 heads, this is $2 \times 32 \times 32 = 2,048$ individual tensors per sequence in the batch.

In practice, deep learning frameworks typically consolidate this organization into two tensors per layer, using an additional dimension to index across heads:

$$\begin{aligned} K_{\text{cache}}^{(l)} &\in \mathbb{R}^{B \times H \times T \times d_h} \\ V_{\text{cache}}^{(l)} &\in \mathbb{R}^{B \times H \times T \times d_h} \end{aligned}$$

where:

• $K_{\text{cache}}^{(l)}, V_{\text{cache}}^{(l)}$: the key and value cache tensors for layer $l$, containing caches for all heads
• $B$: the batch size, letting multiple sequences to be processed in parallel
• $H$: the number of attention heads, each with its own cached key-value pairs
• $T$: the current sequence length, growing as generation proceeds
• $d_h$: the head dimension, typically $d_{\text{model}} / H$

Some implementations combine keys and values into a single tensor of shape $(B, H, T, 2, d_h)$ or $(B, 2, H, T, d_h)$ for improved memory locality. This layout keeps each position's key and value adjacent in memory, which can improve cache efficiency during the attention computation. However, the logical structure remains the same regardless of the physical memory layout.

Understanding how the cache integrates with the transformer's layer-by-layer computation clarifies why the optimization is both correct and efficient. During the decode phase, when generating token by token, each transformer layer performs the following sequence of operations:

1. Receives the hidden state for only the new token: $h_t^{(l)} \in \mathbb{R}^{B \times 1 \times d_{\text{model}}}$. This is the representation of the new token as computed by the previous layer.
2. Projects this hidden state to obtain $q_t$, $k_t$, $v_t$ for the new token only.
3. Retrieves the cached $K_{\text{cache}}^{(l)}$ and $V_{\text{cache}}^{(l)}$ containing all previously computed keys and values for this layer.
4. Appends the new $k_t$ and $v_t$ to the cache, extending the history by one position.
5. Computes attention between the new query $q_t$ and the full cache, letting the new token to attend to all previous positions.
6. Passes the attention output through the feed-forward network, which processes each position independently.
7. Returns the hidden state for the new token: $h_t^{(l+1)}$, ready for the next layer.

Notice that operations 1, 2, 6, and 7 all operate on just a single position. The feed-forward network, layer normalization, and residual connections process only the new token's hidden state, making them computationally trivial during decoding. The attention computation in step 5 is the only operation that touches the full sequence length, and even there, the work is linear in $t$ rather than quadratic because we compute attention for just one query position.

This also explains why the KV cache only caches keys and values, not queries: queries are consumed immediately within the layer to compute attention, and a query for position $t$ can never be needed again because future positions $t', t' > t$ will generate their own queries using their own embeddings. The query is inherently step-local; the key and value are inherently step-global.

An important subtlety is that causal masking interacts with the cache in a way that is often glossed over. During prefill, the model processes the full prompt as a matrix operation and applies a causal mask so that position $i$ cannot attend to positions $j > i$. During decode, however, the new token at position $t$ can attend to all prior positions 1 through $t-1$ without any masking restriction, because all those positions precede it in the sequence. The decode phase therefore applies no mask at all for the single-token query case. This simplification is valid and saves computation: there is no mask matrix to construct or apply when the query has length 1.