KV Cache Explained: Efficient Attention for LLM Generation

When you ask a language model to complete the sentence "The capital of France is," it generates tokens one at a time: first "Paris," then perhaps a comma, then additional context. At each step, the model computes attention over all previous tokens to decide what comes next. Without optimization, this means the model recomputes the same attention calculations for "The," "capital," "of," "France," "is" every single time it generates a new token. For a 100-token response, that's nearly 5,000 redundant computations just for the prompt tokens. This cost does not just slow things down a little; it makes practical deployment nearly impossible.

The key-value cache, commonly called the KV cache, eliminates this redundancy by storing intermediate attention computations and reusing them across generation steps. This optimization reduces the computational complexity of autoregressive generation from quadratic to linear, making inference with large language models feasible at production scale. Think of the KV cache as the model's notepad: instead of re-reading the entire novel every time it needs to write the next sentence, the model keeps its reading notes and consults them directly. The notes grow with each sentence written, but the act of re-reading the full novel from scratch never happens.

Understanding the KV cache is necessary for reasoning about modern LLM serving. Paged attention, continuous batching, speculative decoding, and grouped-query attention all manage or extend the KV cache more effectively. Before you can reason about why these techniques matter, you need to understand what the cache is, why it works, what it costs, and where it breaks down.

This chapter begins by dissecting exactly where redundancy arises during autoregressive generation, building up from the attention formula itself. We then formalize how caching eliminates that redundancy, examine the per-layer and per-head structure of the cache in real architectures, and work through the memory arithmetic that determines whether a deployment is feasible. A fully worked numerical example shows each computation step-by-step before we move into a complete PyTorch implementation. We close with an honest examination of the limitations that have motivated years of follow-on research.

As we explored in Part XIII, decoder-only transformers generate text one token at a time. At each step, the model takes all tokens generated so far (including the prompt) and predicts the next token. The computational core of this process is the self-attention mechanism, which allows each token to gather information from all tokens that came before it. To understand why caching gives such large benefits, we must first examine exactly how attention operates during generation and identify where the redundant work occurs. The redundancy is not an accident or a design flaw; it is an inherent consequence of computing context-aware representations at every position.

Recall from Part XIII that self-attention computes queries, keys, and values from the input. The basic insight behind these three components is that attention works like an information retrieval system: queries stand for what a token is looking for, keys stand for what each token offers, and values contain the actual information to be retrieved. Think of it like a library search: the query is your search term, the keys are the catalog entries for every book on the shelf, and the values are the books themselves. When you search, you compare your query against every catalog entry and retrieve the books whose entries best match. We express these three components through linear projections:

where:

• $X$: the input embedding matrix, containing the vector representations of all tokens in the sequence
• $W_Q, W_K, W_V$: the learned projection matrices that transform inputs into query, key, and value subspaces, each capturing a different aspect of the token's meaning
• $Q, K, V$: the resulting matrices, where $Q$ contains vectors representing what each token seeks, $K$ contains vectors representing what each token can be matched against, and $V$ contains the actual content to be aggregated

Each of these projection matrices is learned during training to extract the most useful representations for the attention mechanism. Once training is complete, these matrices remain fixed during inference. This means that given the same input token, the projections will always produce the same key and value vectors. This determinism is the property that makes caching possible: if the input and the weight matrix are unchanged, the output cannot change either.

The attention mechanism then computes a weighted sum of values based on the compatibility between queries and keys. This compatibility is measured through dot products, which capture how well a query's "question" matches each key's "answer." The full attention formula is:

where:

• $Q, K, V$: the query, key, and value matrices computed from the input
• $K^T$: the transpose of the key matrix, arranged so that the dot product between queries and keys can be computed efficiently through matrix multiplication
• $d_k$: the dimension of the key vectors, used for scaling to prevent the dot products from growing too large
• $\text{softmax}$: the function that converts raw compatibility scores into probability weights summing to 1

The scaling factor $\sqrt{d_k}$ is necessary. Without this normalization, the dot products between high-dimensional vectors can become very large in magnitude. When these large values pass through the softmax function, the result becomes extremely peaked, meaning almost all attention weight concentrates on a single token. This causes gradients to vanish for all other positions, making training unstable. The square root scaling keeps the variance of the dot products roughly constant regardless of the dimension, and we explored this phenomenon in depth in Part XIII.

During generation, let's trace what happens at each step to understand precisely where the wasted work accumulates. Suppose we're generating a response to a 10-token prompt. The model must produce one token at a time, with each new token depending on everything that came before it.

Step 1 (Prefill): The model processes all 10 prompt tokens simultaneously. It computes the $Q$, $K$, $V$ matrices, each of shape $(10, d_k)$ or $(10, d_v)$ depending on the specific projection. The attention mechanism allows each prompt token to attend to all previous prompt tokens, establishing the initial context representation.

Step 2: After generating the first new token, we now have 11 tokens total. Without any optimization, the model would need to process the entire 11-token sequence from scratch. This means computing $Q$, $K$, $V$ for all 11 tokens, even though 10 of them are unchanged from the previous step.

Step 3: With 12 tokens total, we compute $Q$, $K$, $V$ for all 12 tokens, once again recomputing the projections for the original prompt tokens that haven't changed at all.

Notice the wasteful pattern: at step $t$, we recompute the keys and values for the first 10 prompt tokens, even though nothing about them has changed since step 1. The same tokens, passing through the same fixed projection matrices $W_K$ and $W_V$, produce the same $K$ and $V$ vectors every single time. This redundancy stems from a basic property of the attention mechanism: the key and value for a given token depend only on that token's embedding and the projection weights, not on what comes after it in the sequence. The causal mask ensures that earlier tokens never attend to later ones, so later tokens can never influence the keys and values of earlier tokens.

This redundancy has concrete costs that grow sharply with sequence length. For a model with $L$ layers, $H$ attention heads, and head dimension $d_h$, generating $T$ tokens requires computing $K$ and $V$ projections repeatedly at every step. The total number of projection operations across all generation steps sums to:

where:

• $T$: the final sequence length, representing the total number of tokens processed including both the prompt and generated tokens
• $t$: the current step number, which also is the sequence length at that particular step
• $\sum_{t=1}^{T} t$: the sum of the arithmetic progression $1 + 2 + 3 + \ldots + T$, capturing the total number of token-projection operations across all steps

This quadratic scaling has severe implications for practical deployment. If generating 100 tokens requires $\frac{100 \times 101}{2} = 5,050$ projection operations, then generating 1,000 tokens requires $\frac{1000 \times 1001}{2} = 500,500$ operations, a 100-fold increase in work for only a 10-fold increase in output length. This contrasts with the ideal case where each token's keys and values are computed only once, requiring $T$ projection operations. The KV cache reaches exactly this ideal.

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:

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:

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:

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:

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.

To make the abstract formulas concrete, let's trace through a minimal numerical example from start to finish. We will use a toy model with $d_{\text{model}} = 4$, one attention head ($H = 1$), and $d_k = d_v = 4$. The prompt is three tokens: "The cat sat." We will then generate one additional token and track every vector and matrix involved.

Suppose the three prompt tokens produce input embeddings (after the embedding layer and any positional encoding):

For the projection matrices, we use simplified illustrative weights. In practice these are dense learned matrices, but here we treat them as identity-like for clarity:

This means the key and value for each token are just that token's embedding, and the query is half the embedding. The simplification does not change the structural argument at all.

We compute the keys and values for all three tokens simultaneously:

We store these in the cache immediately:

(In this toy example the keys and values are equal since $W_K = W_V = I$; in real models they differ entirely.)

Suppose the model samples token 4 from the distribution produced after the prefill phase, and that token 4 has embedding $x_4 = [0.5, 0.9, 0.3, 0.6]$. To predict token 5, we need to compute attention from position 4 over all previous positions.

Step 1: Compute projections for token 4 only.

The key insight is that we compute exactly three vectors ($q_4$, $k_4$, $v_4$) for this step, not $4 \times 3 = 12$ vectors. The keys and values for tokens 1 through 3 remain exactly as computed during prefill.

Step 2: Append $k_4$ and $v_4$ to the cache.

Step 3: Compute attention scores for token 4's query against all four cached keys.

The raw dot products are:

Step 4: Scale by $\sqrt{d_k} = \sqrt{4} = 2.0$.

Step 5: Apply softmax to get attention weights.

First compute the exponentials:

Sum: $1.451 + 1.380 + 1.346 + 1.459 = 5.636$

Normalized weights:

Notice that token 4's own key (position 4) receives slightly more weight than the others, which is typical for recent tokens in causal attention.

Step 6: Compute the attention output as a weighted sum of cached values.

So $o_4 \approx [0.627, 0.703, 0.515, 0.458]$.

What did we compute in this decode step? Three vector projections (for token 4 only), four dot products, a softmax over four values, and four weighted additions. Without the cache, we would have recomputed all projections for tokens 1, 2, and 3 as well. That is 9 additional projection operations for a sequence of length 4. For a sequence of length 100, the savings would be 297 additional projections avoided. For length 1000, the savings grow to 2997 projections avoided in that single step alone.

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:

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.

Understanding KV cache memory consumption is needed for deployment planning, as the cache often becomes the primary memory bottleneck in production systems. Unlike model weights that have a fixed size once the model is loaded, the cache grows dynamically during generation and can easily exceed the model weights in memory usage for long sequences or large batches. A system that can comfortably load and run a 7B model may be unable to serve multiple concurrent users at long context lengths, not because the model is too large, but because the caches are.

To derive the memory requirements, we must account for all the tensors that constitute the complete cache. For a single sequence of length $T$, we need to store keys and values for every layer and every head. The basic memory formula is:

where:

• $2$: accounts for storing both keys and values, as each requires a separate tensor
• $L$: the number of transformer layers, each maintaining its own independent cache
• $H$: the number of attention heads per layer
• $d_h$: the dimension of each head, determining the size of individual key and value vectors
• $T$: the sequence length in tokens, the dynamic factor that grows during generation
• $\text{bytes per element}$: the memory size required for a single floating-point number (e.g., 2 bytes for FP16, 4 bytes for FP32)

Since the total model dimension is typically expressed as $d_{\text{model}} = H \times d_h$, we can simplify this formula by substituting:

where:

• $d_{\text{model}}$: the total model dimension, equal to $H \times d_h$
• $L$: the number of layers in the model
• $T$: the sequence length
• $\text{bytes per element}$: the numerical precision in bytes

This simplified formula reveals an important insight: the cache memory scales with the product of model depth and width, multiplied by sequence length. Doubling any of these factors doubles the memory requirement. This linear relationship is the cache's great advantage over the naive quadratic scheme, but it still produces very large absolute values for modern models.

To make these formulas concrete, let's calculate the cache requirements for a widely deployed model. The LLaMA 2 7B model has the following specifications:

• $L = 32$ layers
• $d_{\text{model}} = 4096$
• Typically stored in FP16 (2 bytes per element)

For a single sequence of length $T = 2048$, we can compute the cache size step by step:

A single 2048-token sequence requires a full gigabyte just for the KV cache. This is separate from the memory needed for model weights and activations, along with the rest of the inference overhead.

For a batch of 8 sequences at the 4096-token context length, the memory requirement grows proportionally:

For comparison, the model weights themselves require about 14 GB in FP16. This means that for a batch of 8 long-context requests, the KV cache consumes more memory than the entire model. At longer context lengths or larger batch sizes, the cache can substantially exceed the model size, making memory management the necessary challenge for deployment.

Larger models and longer contexts exacerbate the memory pressure, creating significant challenges for deployment at scale. The following table shows KV cache sizes for various model configurations, illustrating how the requirements grow across different architectural choices:

The cache memory scales linearly with both model size (via the $L \times d_{\text{model}}$ product) and context length. This linear scaling in context length, combined with the quadratic scaling of the attention computation itself, makes long-context inference particularly challenging. A model like GPT-4 with a 128K context window requires hundreds of gigabytes per sequence, necessitating distributed systems and advanced memory management strategies. Understanding this table should give you a visceral sense of why every modern LLM serving system devotes significant engineering effort to cache optimization.

Efficient cache management involves necessary decisions about memory allocation strategies, handling variable sequence lengths across concurrent requests, and managing batched generation effectively. These operational considerations often determine whether a deployment can reach its throughput and latency targets. Getting the compute right is necessary but not sufficient; you also need to get the memory management right.

Think of cache management like hotel room allocation. Model weights are the building itself: fixed overhead that is always there. Each user request is a guest who needs a room (cache memory) for the duration of their stay (generation). Some guests stay for 10 minutes (short generations), others for hours (long generations). If you pre-assign the biggest possible room to everyone, you waste capacity. If you assign rooms dynamically but poorly, you get fragmentation where no contiguous block is large enough for the next guest even though total free memory is sufficient.

Static allocation pre-allocates cache tensors for the maximum sequence length at the start of generation. This approach reserves all memory upfront, avoiding the overhead of repeated memory allocation during the generation process. Static allocation is preferred when the maximum sequence length is known in advance, when memory fragmentation must be avoided, and when consistent, predictable latency is required. The downside is waste: a request that only generates 100 tokens uses the same amount of pre-allocated memory as one that generates 4,096 tokens.

Dynamic allocation grows the cache as needed, typically by allocating larger tensors and copying existing data when the current allocation is exhausted. This saves memory for shorter sequences but introduces allocation overhead and potential fragmentation. Modern frameworks often use a hybrid approach, allocating in chunks of 128 or 256 tokens at a time to balance these tradeoffs. The chunk size is a compromise between allocation frequency (smaller chunks mean more frequent allocations) and waste (larger chunks mean more potential waste at the end).

The management challenge compounds when serving multiple requests concurrently. A request that finishes early frees its cache memory, but that freed memory may be fragmented: a 500-token hole here, a 200-token hole there. A new incoming request wanting 1,000 tokens cannot use these fragments even though total free memory is sufficient. This is the classical memory fragmentation problem, and it is why vLLM's Paged Attention, which we will cover in the next chapter, attracted such attention: it applies the operating system's paging approach to eliminate fragmentation entirely.

When generating for multiple sequences simultaneously, cache management becomes more complex because different sequences may have different lengths. The typical approaches each involve different tradeoffs:

• Padding: Allocate cache based on the longest sequence in the batch, padding shorter sequences. Simple to implement and preserves batched matrix operations, but wasteful when sequence lengths vary materially. A batch where one sequence has 4,096 tokens and five others have 50 tokens wastes nearly all the allocated memory for those five.
• Separate caches: Maintain independent cache tensors per sequence. Avoids wasted memory but complicates the attention computation and prevents batched matrix operations, hurting GPU utilization.
• Paged attention: Allocate cache in fixed-size blocks and track which blocks belong to which sequence using a block table (in effect a page table). This approach lets efficient memory utilization with variable-length sequences and is the basis of the vLLM serving system.

When a sequence reaches the model's maximum context length, the system must decide how to proceed. Three main strategies exist, each with different tradeoffs:

Truncation simply stops accepting new tokens or drops the oldest tokens from the sequence. The model loses access to early context, which can cause coherence problems in long conversations.

Sliding window maintains only the most recent $W$ tokens, discarding older ones. This is the approach used in Mistral's architecture, as we discussed in Part XXIX. The advantage is a constant memory footprint; the disadvantage is that the model cannot reference information from the distant past.

Attention sinks keeps initial tokens (which empirically receive disproportionate attention regardless of content, the "sink" phenomenon we explored in Part XVIII) plus a window of recent tokens. This preserves the model's numerical stability while using a bounded cache.

Each strategy has tradeoffs between memory usage, coherence over long conversations, and computational complexity. Production systems often need to support multiple strategies and allow users or operators to choose based on their application requirements.

Let's build a simple KV cache implementation to see how these concepts work in practice. The implementation below prioritizes clarity over production-level efficiency, but it captures all the needed logic.

A necessary property of KV caching is that it must produce identical outputs to non-cached attention. This guarantee makes caching a valid optimization rather than an approximation. If caching changed the output even slightly, every generated token would be slightly different, and those small errors would compound over the course of a generation, creating wildly different sequences. Let's verify this mathematically and empirically.

The algebraic argument is straightforward. In standard full-sequence attention, the output for position $t$ in a sequence of length $T$ is:

where $K[1:t]$ denotes the submatrix of $K$ containing only the first $t$ rows. In cached attention, we have:

where $K_{\text{cache}}$ at step $t$ contains exactly the first $t$ rows of $K$. Since $k_i = x_i W_K$ is computed identically whether we process the full sequence at once or token-by-token, the two expressions are algebraically identical. The numerical values in $K_{\text{cache}}$ are bitwise identical to the corresponding rows of $K$, so the outputs are numerically equivalent up to floating-point precision.

The numerical equivalence confirms that caching is purely an optimization: it changes how we compute the result, not what we compute. The tiny residual difference (typically around $10^{-7}$) shows floating-point rounding from the different order of operations, not a systematic error.

Let's measure actual memory consumption for different configurations:

Understanding each parameter's role helps you reason about deployment decisions when configuring inference servers:

• d_model: The dimensionality of the model's hidden states. This appears directly in the memory formula and in the projection computation cost.
• num_heads: The number of attention heads. Along with head_dim, it determines how the model dimension is partitioned. More heads mean more granular attention patterns but the same total cache size (since $H \times d_h = d_{\text{model}}$).
• head_dim: The dimension of each attention head, equal to $d_{\text{model}} / \text{num\_heads}$. Smaller head dimensions mean more heads for the same model width.
• max_seq_len: The maximum sequence length the cache can store. This determines the upper bound on context length and the maximum cache memory for a given model.
• batch_size: The number of sequences processed simultaneously. Larger batches improve GPU utilization and throughput but multiply cache memory linearly.

KV caching introduces several challenges that drive ongoing research in inference optimization. These limitations directly motivate major inference optimization techniques that have emerged since 2020, from grouped-query attention to paged attention to speculative decoding.

The most basic limitation is that memory consumption scales linearly with sequence length. While this is better than the quadratic scaling we would face without caching, it still creates a hard constraint on context length and batch size. A server with 80 GB of GPU memory might fit the model weights comfortably at 14 GB, but may be unable to serve multiple long-context requests concurrently because the caches fill the remaining 66 GB faster than expected. This tension between throughput and context length is basic to LLM deployment and not solvable by any single optimization. The KV cache memory is the single most important factor determining how many requests a given GPU can serve simultaneously. Every percentage reduction in cache memory directly translates to higher concurrency and lower cost per query.

Memory fragmentation becomes increasingly problematic as the number of concurrent requests grows. When different sequences in a batch have different lengths, naive approaches either waste memory through padding or sacrifice batching efficiency. Consider a batch where one sequence is actively using 3,000 tokens of cache and another finished early and freed 2,000 tokens. A new incoming request wanting 2,500 tokens cannot use the freed space because it is not contiguous. The server either rejects the request, queues it until more contiguous space is available, or wastes memory by padding smaller requests to the maximum. Production systems must carefully manage cache allocation to maximize GPU utilization. The Paged Attention chapter (coming next) addresses this with a memory management approach inspired by operating system virtual memory, dividing the cache into fixed-size pages that can be assigned non-contiguously.

The cache must persist across the entire generation process. Unlike model weights that are read-only during inference, the KV cache is constantly growing and being read from. This means the cache cannot be easily offloaded to CPU during generation without incurring significant latency for memory transfers (CPU-GPU bandwidth is typically 10-30 GB/s, while GPU memory bandwidth is 1-3 TB/s, a gap of two orders of magnitude). Systems that need to handle many concurrent requests must carefully orchestrate which caches are active on GPU at any moment. Some systems implement "cache swapping" that moves inactive caches to CPU memory between generation steps, accepting the bandwidth cost in exchange for higher concurrency. This tradeoff is only viable when GPU processing time per step is long enough to hide the transfer latency.

A fourth, less obvious limitation is that the decode phase becomes memory-bandwidth-bound rather than compute-bound. During decode, each step processes a single token's worth of activations through the model's projection matrices. Modern GPUs have large FLOPs but moderate memory bandwidth. The weight matrices for a large model contain billions of parameters, and at each decode step the model must read in effect all of them from DRAM to compute a single output token. For LLaMA-2 7B, this means reading ~14 GB of weights and up to several GB of cache per generated token. The GPU's arithmetic units are starved waiting for memory, resulting in low hardware utilization. This is why batching is so important: serving 32 concurrent requests at once amortizes the weight-reading cost over 32 tokens, sharply improving utilization. The cache then becomes the bottleneck that limits how large a batch can be.

Grouped-Query Attention (GQA), which we discussed in Part XXIX, directly addresses the cache size problem by having multiple query heads share a single key-value head. For example, if 8 query heads share 1 KV head, the cache is 8 times smaller while the model's representational capacity for queries is unchanged. LLaMA-2 70B uses GQA with a 2:1 reduction (8 KV heads for 64 query heads), cutting the cache size by 8 times. This architectural choice was motivated specifically by the memory constraints we have described here: making the cache smaller is the most direct way to let larger batches or longer contexts.

The KV cache is a basic optimization that makes autoregressive generation practical. By storing and reusing the key and value projections from attention, we avoid recomputing the same values at every generation step. This reduces the computational overhead from quadratic to linear in the number of generated tokens, and the reduction is not an approximation but an algebraically exact reformulation of the same computation.

The key concepts covered in this chapter are:

• Redundancy problem: Without caching, each generation step recomputes $K$ and $V$ for all previous tokens, yielding $O(T^2)$ total projection operations for a $T$-token generation
• Cache structure: Separate $K$ and $V$ tensors per layer per head, growing with sequence length, storing projections for all processed tokens
• Prefill vs decode: The initial prompt processing (prefill) populates the cache in a single parallel pass, while subsequent generation steps (decode) each process only one token at a time
• Memory requirements: Cache size scales as $2 \times L \times d_{\text{model}} \times T \times \text{bytes}$, easily reaching multiple gigabytes for long sequences with modern models
• Cache management: Decisions about static vs dynamic allocation, batching strategies, and context length limits materially impact deployment efficiency
• Limitations: Memory pressure, fragmentation, bandwidth bottlenecks, and the memory-bandwidth-bound decode phase motivate a detailed ecosystem of follow-on optimizations

Understanding KV cache mechanics is needed for working with modern LLMs, as cache memory often becomes the primary constraint on throughput and context length. The next chapter examines KV cache memory arithmetic in greater detail, including how grouped-query attention changes the calculation. Following that, Paged Attention introduces an operating-system-inspired approach to cache memory management that lets the vLLM serving system's large throughput improvements. After that, we explore techniques for compressing the cache through quantization and eviction to extend context lengths further, all building directly on the foundation established here.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about KV cache optimization in transformer inference.