The KV Cache from First Principles

Have you ever written a loop where you recompute the same thing every iteration, and then felt that specific kind of embarrassment when you realize you could have just stored the result? The KV cache is that insight, applied to attention.

In Post 5.3 we established the key thing about autoregressive decoding: at each step, the model generates one new token, which means one new row in the attention computation. The query vector is new. But the keys and values for every previous position? Those are exactly the same as they were last step. We already computed them. They haven't changed. Recomputing them is pure waste.

The KV cache is what you get when you take that observation seriously. You store the keys and values once, append a single row per step, and reuse everything you've already computed.

But the consequences of this simple idea end up dominating the entire inference stack. The cache grows linearly with sequence length. For large models at long contexts, it consumes more GPU memory than the model weights themselves. Every major innovation in Arc 6, things like PagedAttention, continuous batching, and prefix caching, exists because the KV cache is the scarce resource that has to be managed.

So let's derive it properly.

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.2, just the idea that HBM is finite and precious

The redundancy

Let me trace through what happens without a cache, to make the waste explicit.

Say we've generated 100 tokens so far and we're about to generate token 101. The model needs to compute attention for the new token. That means computing a query vector for position 101, and then attending over keys and values at all 101 positions.

Here's the thing. To get the key vector at position 37, we take the embedding at position 37 and multiply it by $W_K$. To get the value at position 37, we multiply the same embedding by $W_V$. But we already did this when we were generating token 38. And again when generating token 39. And every step since. The embedding at position 37 hasn't changed. $W_K$ and $W_V$ haven't changed. The result is identical every time.

Without a cache, generating a sequence of $T$ tokens means computing the key and value for position 1 a total of $T - 1$ times. Position 2 gets computed $T - 2$ times. The total number of redundant K/V computations is $\sum_{i=1}^{T-1} i = T(T-1)/2$. For a 2,000-token generation, that's roughly 2 million redundant matrix multiplications per layer, per head.

Once you see it written out like that, the solution becomes obvious.

With a cache, each position is computed exactly once. The savings grow quadratically with sequence length while the cache itself grows only 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 are cheap. They're single vector-matrix multiplies. Step 3 is where the work is, because we need the full $K$ and $V$ matrices going back to position 1.

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:

The attention output for position $t$ is mathematically identical whether you recompute everything from scratch or use the cache. There's no approximation, no quality tradeoff, just the same result with a lot less work.

How much memory does this cost?

Now here's where the simple idea starts to get scary. 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.

Let me run through a concrete example. Llama 3 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 H100 has 80 GB of HBM total. You can see the problem.

And this is with GQA, which is already using far fewer KV heads than query heads. With standard multi-head attention (32 KV heads instead of 8), that number would be 4x larger: around 43 GB per request. At which point you can barely fit one request alongside the model weights on a single GPU.

This is why I keep saying the KV cache is the scarce resource. 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 memory is the central problem of inference serving.

Multi-query attention: sharing everything

The memory problem was recognized early. In 2019, Noam Shazeer published a short, elegant paper with a simple proposal: 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.

The cache savings are immediate. Instead of 32 KV caches per layer, you have 1. For our Llama-scale example, the 43 GB cache drops to about 1.3 GB. You can serve dozens of concurrent requests comfortably.

I want to be precise about what's shared and what isn't. 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 catch is that quality drops. 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. In practice, MQA models tend to be slightly worse on quality benchmarks, especially on tasks that require fine-grained attention patterns.

Grouped-query attention: the compromise

In 2023, Ainslie et al. proposed the obvious 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. Quality is very close to full MHA, close enough that Llama 2 70B, Llama 3, and most large models released since 2023 use GQA as the default.

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). The 2023 insight was that you can slide the dial to some intermediate $G$ and get almost all the quality of MHA with a fraction of the cache.

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. Llama 2's 70B variant did exactly this.

GQA is one of those ideas that feels like it should have been discovered earlier. Once you frame the KV cache as a memory bottleneck (which Pope et al. did systematically in 2022), sharing heads within groups is the obvious lever. It took until 2023 for someone to do the careful ablation showing that the quality cost is genuinely small.

Implementation

Let me write a minimal version so you can see the mechanics. The math was already trivial; what's interesting is the append-and-reuse pattern, and the fact that the query is always exactly one row.

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.
    """
    # Append to cache
    k_full, v_full = cache.update(k_new, v_new)
    # k_full: (batch, kv_heads, seq_so_far, d_head)
 
    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])
# ...

A few things to notice. The query is always shape (batch, heads, 1, d_head), which is one row, always. 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. That's the one-new-row insight from Post 5.3 made concrete.

In production, nobody uses torch.cat to grow the cache. That would allocate a new tensor every step, which is terrible for performance. Real implementations 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, and the hydration problem

I've been talking about the decode phase, generating one token at a time. But there's also the prefill phase, where the model processes the user's prompt. During prefill, you're processing many tokens at once (the whole prompt), so you compute K and V for all prompt positions in a single batched matrix multiply. Those results go straight into the KV cache.

After prefill, the cache already contains the keys and values for every prompt token. Decode then picks up from there, appending one new row per generated token. The cache doesn't distinguish between "came from prefill" and "came from decode." It's all just rows of keys and values.

This act of populating the cache from a raw token sequence, before decode can start, has a name. I call it the hydration problem, and it's a handle we'll keep reusing throughout the series. Any time the model needs to "remember" a chunk of text it hasn't seen yet in its current session, those tokens have to be pushed through the full transformer so their K and V vectors land in the cache. If nothing is hydrated, the model has nothing to attend to. Prefill is one form of hydration. Prefix caching (Post 5.8) is a way to skip it when the rows already exist somewhere. And agent runtimes juggling tool outputs and compacted history are doing hydration logistics in disguise. The name is useful because the same shape of problem shows up at every layer above this one.

Prefix caching is the obvious consequence. If two requests share the same system prompt, the KV cache entries for that prompt are identical, so you can compute them once and reuse them across requests, skipping prefill for the shared prefix entirely. The cache doesn't care how its rows got there.

One subtlety: during prefill, the computation is compute-bound (lots of tokens, big matrix multiplies, high arithmetic intensity). During decode, it's memory-bound (one token, you're reading the entire weight matrix to produce a single output vector). The KV cache helps with decode by eliminating redundant compute, but it doesn't change the fundamental memory-boundedness of each decode step. We covered that distinction in Post 5.3.

Misconceptions

"The KV cache is an approximation." No. The attention output with a cache is bit-for-bit identical to recomputing from scratch. There's no quality loss, no approximation. It's pure memoization: same computation, same inputs, same outputs. You just store the result instead of redoing the work.

"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 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. GQA is the dominant factor.

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.