The KV Cache from First Principles

In Post 5.3, each decode step added one row to the attention computation. The query vector is new, but the keys and values for previous positions are unchanged. We can reuse them.

The cache grows linearly with sequence length. Across enough long requests, it can consume more GPU memory than the weights. PagedAttention, batching, and prefix caching all have to account for this growing memory demand.

Prerequisites

You'll want to be comfortable with:

• Attention mechanics from Post 4.1, specifically how Q, K, V matrices work
• The one-new-row insight from Post 5.3, that decode produces one query row per step
• Matrix multiplication shapes, enough to follow (1, d) @ (d, seq) = (1, seq)
• GPU memory basics from Post 5.1, just the idea that HBM is finite and precious

The redundancy

Say 100 positions have been processed and token 101 has just been sampled. Processing token 101 to predict token 102 computes its query and attends over positions 1 through 101, including itself.

To get the key vector at position 37, we take the layer input at position 37 and multiply it by $W_K$. To get the value at position 37, we multiply the same layer input by $W_V$. But we already did this when we were generating token 38. And again when generating token 39. And every step since. Under causal attention, that layer input hasn't changed when later tokens arrive. $W_K$ and $W_V$ haven't changed. The result is identical every time.

Count $T$ processing steps starting with a single position: step $t$ processes positions 1 through $t$. Without a cache, position 1 is computed $T$ times and position 2 is computed $T-1$ times. Across these $T$ one-position extensions, a naive implementation computes $T(T+1)/2$ position-level K/V projections instead of $T$. The redundant count is $T(T-1)/2$: about two million position-level projection evaluations at $T=2{,}000$, not two million separate GPU matmul launches.

With a cache, each position is computed exactly once. In this one-position-at-a-time example, avoided projection evaluations grow quadratically with sequence length while cache storage grows linearly.

Deriving the cache

At each decode step $t$, we need to:

1. Compute $q_t = x_t W_Q$, the query for the new position
2. Compute $k_t = x_t W_K$ and $v_t = x_t W_V$, the key and value for the new position
3. Attend: compute $q_t$ against all keys $K_{1:t}$, then take the weighted sum of all values $V_{1:t}$

Steps 1 and 2 read projection weights; step 3 reads cached keys and values. Which dominates depends on model shape, batch size, and context length.

The cache is just: keep those matrices around. After step $t$, we have $K_{1:t}$ and $V_{1:t}$ stored in GPU memory. At step $t+1$, we compute one new key $k_{t+1}$ and one new value $v_{t+1}$, append them to the cache, and attend using the full cached matrices. We never recompute a key or value we've already seen.

In math:

When model, positions, inputs, and cache precision match, reusing K/V gives the same attention output mathematically. Different kernel shapes may change floating-point rounding; quantized or compressed caches introduce separate approximations.

How much memory does this cost?

Every layer in the model maintains its own K and V cache. And every KV head within each layer has its own cache. The total memory is:

where $L$ is the number of layers, $n_{\text{kv\_heads}}$ is the number of KV heads per layer, $d_{\text{head}}$ is the head dimension, $T$ is the sequence length, and $b$ is the bytes per element (2 for fp16, 1 for int8).

The factor of 2 is for K and V, two separate caches.

Llama 3.1 70B has 80 layers, 8 KV heads (it uses grouped-query attention, which we'll get to), and a head dimension of 128. At a 32,768-token context in fp16:

That's per request, not per model. If you're serving 8 concurrent requests at 32k context, that's 85 GB of KV cache memory on top of however much the model weights themselves consume. An 80 GB H100 SXM has less memory than this aggregate cache alone.

And this is with GQA, which is already using far fewer KV heads than query heads. With 64 KV heads to match this model’s 64 query heads, a hypothetical MHA variant would use 8× as much: about 85.9 GB per request. The FP16 weights already require multiple 80 GB GPUs; the example counts total cache bytes before sharding.

The model weights are fixed. You load them once and they stay put. The KV cache is dynamic, growing with every token, and you need a separate one for every active request. Managing this growth limits how many requests a server can keep active.

Sharing key and value heads

In 2019, Noam Shazeer proposed: what if all query heads shared the same key and value projections?

In standard multi-head attention (MHA), each head gets its own $W_Q$, $W_K$, and $W_V$. If you have 32 heads, you have 32 sets of KV projections and 32 KV caches. Multi-query attention (MQA) keeps 32 independent query heads but uses just one shared key head and one shared value head. Every query head attends over the same keys and values.

Instead of one KV head per query head, you have one per layer. Holding the 70B example’s other dimensions fixed gives about 1.34 GB at 32,768 tokens.

Each query head still has its own $W_Q$, so head 0 and head 7 still compute different query vectors from the same input embedding. They still produce different attention patterns. They're just computing those patterns against the same set of keys, and pulling from the same set of values. Think of it as 32 different questions, all searched against the same lookup table.

The possible tradeoff is quality. When all query heads share the same keys and values, the model loses representational capacity. Different heads can no longer learn to attend to different aspects of the input through different K/V subspaces. The MQA and GQA papers evaluate this tradeoff for particular architectures and training setups; the size of a quality difference is not determined by head count alone.

Grouped-query attention: the compromise

In 2023, Ainslie et al. proposed a middle ground: grouped-query attention (GQA). Instead of one KV head shared across all query heads, or one KV head per query head, you use $G$ groups. Each group of query heads shares a single KV head.

If you have 64 query heads and 8 KV heads, each KV head is shared across 8 query heads. The cache is 8x smaller than MHA and 8x larger than MQA. The GQA paper found quality close to MHA in its evaluated models; Llama 2 70B and Llama 3 use GQA.

You can think of MHA, GQA, and MQA as points on a single continuum. MHA is GQA with $G = n_{\text{query\_heads}}$ groups (every query head gets its own KV head). MQA is GQA with $G = 1$ (all query heads share one KV head). Ainslie et al. tested intermediate group counts to measure the quality/cache tradeoff.

The Ainslie et al. paper showed that you can even convert an existing MHA model to GQA by mean-pooling the KV heads within each group and then fine-tuning for a small fraction of the original training budget.

Choose the KV head count when training or uptraining the model. Switching the diagram between MHA, GQA, and MQA shows memory arithmetic, not a free runtime conversion of a trained checkpoint.

Implementation

The implementation below appends one key and value row per step, then attends over the accumulated cache.

import torch
import torch.nn.functional as F
 
class KVCache:
    """Minimal KV cache for one attention layer."""
 
    def __init__(self):
        self.k: torch.Tensor | None = None  # (batch, heads, seq, d_head)
        self.v: torch.Tensor | None = None
 
    def update(self, k_new: torch.Tensor, v_new: torch.Tensor):
        """Append new K/V vectors to the cache."""
        if self.k is None:
            self.k = k_new
            self.v = v_new
        else:
            self.k = torch.cat([self.k, k_new], dim=2)  # concat along seq dim
            self.v = torch.cat([self.v, v_new], dim=2)
        return self.k, self.v
 
    @property
    def seq_len(self) -> int:
        return 0 if self.k is None else self.k.size(2)
 
def attention_with_kv_cache(
    q: torch.Tensor,        # (batch, heads, 1, d_head)  — single new query
    k_new: torch.Tensor,    # (batch, kv_heads, 1, d_head)
    v_new: torch.Tensor,    # (batch, kv_heads, 1, d_head)
    cache: KVCache,
) -> tuple[torch.Tensor, KVCache]:
    """
    Single decode step with KV cache.
    Returns the attention output and updated cache.
    """
    groups, remainder = divmod(q.size(1), k_new.size(1))
    if remainder or groups < 1:
        raise ValueError("query heads must be a multiple of KV heads")
    # Validate the head mapping before mutating the cache.
    k_full, v_full = cache.update(k_new, v_new)
    # k_full: (batch, kv_heads, seq_so_far, d_head)
 
    # Expand KV heads for this simple GQA demonstration. Optimized kernels
    # can share them without materializing repeated tensors.
    k_full = k_full.repeat_interleave(groups, dim=1)
    v_full = v_full.repeat_interleave(groups, dim=1)
    d_k = q.size(-1)
    # Score the new query against ALL cached keys
    scores = (q @ k_full.transpose(-2, -1)) / (d_k ** 0.5)
    # scores: (batch, heads, 1, seq_so_far)
 
    weights = F.softmax(scores, dim=-1)
    output = weights @ v_full
    # output: (batch, heads, 1, d_head)
 
    return output, cache
 
# --- Demo: decode 5 tokens ---
batch, heads, d_head = 1, 4, 64
cache = KVCache()
 
for step in range(5):
    x = torch.randn(batch, 1, heads * d_head)  # new token embedding
    # In a real model, these projections use learned W_Q, W_K, W_V
    q = x.view(batch, 1, heads, d_head).transpose(1, 2)
    k = x.view(batch, 1, heads, d_head).transpose(1, 2)
    v = x.view(batch, 1, heads, d_head).transpose(1, 2)
 
    out, cache = attention_with_kv_cache(q, k, v, cache)
    print(f"Step {step}: cache has {cache.seq_len} rows, "
          f"output shape {out.shape}")
 
# Step 0: cache has 1 rows, output shape torch.Size([1, 4, 1, 64])
# Step 1: cache has 2 rows, output shape torch.Size([1, 4, 1, 64])
# Step 2: cache has 3 rows, output shape torch.Size([1, 4, 1, 64])
# ...

The query is always shape (batch, heads, 1, d_head), representing one new query position. The cached keys and values grow by one row each step, so the score matrix is (batch, heads, 1, seq_so_far). The output is always one vector per head. This is the one-position query shape derived in Post 5.3.

A dynamic cache can use torch.cat, but repeated reallocation is expensive. Each append allocates a new tensor and copies the existing entries. A static-cache implementation can pre-allocate a buffer for the maximum sequence length and write into it with an index pointer:

PagedAttention (which we'll cover in Post 6.3) goes even further by managing cache memory in fixed-size blocks, like an operating system manages virtual memory. But the logical operation is always the same: append one row, attend over everything.

Prefill populates the cache

During prefill, the model computes K and V for the known prompt positions in batched matrix multiplications and stores those results in the cache.

Decode extends those cached prompt entries by one row whenever the model processes a newly selected token.

New prompt tokens need a forward pass before their K/V state can be reused. Prefix caching can skip work for an already computed, compatible prefix; new tool output or a rewritten conversation summary still needs processing.

For a dense model at small batch size, prefill often has more weight reuse than decode. A cache eliminates redundant computation, but each decode step still reads weights and cached K/V. Batch size, context length, kernels, and hardware determine the bottleneck. We covered that distinction in Post 5.3.

Misconceptions

"The KV cache is an approximation." An ordinary cache reuses K/V at their original precision. Quantizing or compressing those entries introduces an approximation; retaining them does not.

"The cache stores the attention weights." It stores the K and V projections, the raw key and value vectors. The attention weights (the softmax output) are computed fresh every step, because the query is new and the score distribution changes. Storing attention weights wouldn't help since they depend on the current query, which didn't exist at the previous step.

"Bigger model always means bigger cache." The cache depends on $L \times n_{\text{kv\_heads}} \times d_{\text{head}}$, not on total parameter count directly. A model with 80 layers and 8 KV heads (Llama 3.1 70B) can have a smaller per-request cache than a model with 32 layers and 32 KV heads, even if the second model has fewer total parameters. The number of KV heads can outweigh the difference in layer count.

What's next

The next post covers sampling strategies, where we shift from the internals of the model to what happens after the forward pass: how you actually pick the next token from the probability distribution, and why greedy, temperature, top-k, and nucleus sampling each change the output distribution in different ways.

Additional reading (and watching)

• Vaswani, A., et al. (2017). Attention Is All You Need. NeurIPS 2017.

• Shazeer, N. (2019). Fast Transformer Decoding: One Write-Head is All You Need. arXiv.

• Ainslie, J., et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP 2023.

• Pope, R., et al. (2022). Efficiently Scaling Transformer Inference. MLSys 2023.

• Kwon, W., et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP 2023.

• Hooper, C., et al. (2024). KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization. arXiv.