KV caching from first principles

KV caching is the optimization that makes fast text generation possible. It trades computation for memory, and it is best understood from first principles: what a transformer computes during generation, which part of that work is redundant, and what the redundancy costs. This article builds the picture in order: the generation loop, the wasted work, the cache, its memory cost, and the modern ways of shrinking it.

Prerequisites

Five pieces need to be in place before the rest makes sense. The companion notes cover the first three in more depth; this section states what is needed and where to go for more.

1. Autoregressive generation. A language model does not produce a sentence in one pass. It predicts the next token, appends it to the input, and repeats. This loop is the reason caching exists; without it, there would be nothing to reuse. The transformer note describes the loop, and the NLP history note explains how next-token prediction became the paradigm.

2. Token embeddings and tensor shapes. Each token is mapped to a vector, so a sentence becomes a sequence of vectors. Grouping them gives a matrix of shape $[B, T, C]$: batch size $B$, sequence length $T$, and embedding dimension $C$. Nearly every operation in this article is a shape change on that matrix.

3. Attention and the Q, K, V split. The attention mechanism lets each token look at the tokens before it (the causal variant, which generation uses; the BERT geometry note covers the bidirectional encoder). For each token the model computes three vectors: a query (what this token is looking for), a key (what this token offers to match against), and a value (the content that is passed on). The embedding dimension splits across the attention heads, so with $H$ heads of dimension $d$ each, $C = H \cdot d$, and each token carries one key and one value per head. The worked example later uses a single head to keep the arithmetic small.

4. Matrix multiplication as a matching score. The dot product of a query with a key measures how much they agree, and attention is a weighted average of values, with weights derived from query-key agreement:

The scale factor $\sqrt{d_k}$ keeps the dot products from growing with the key dimension $d_k$.

5. Compute versus memory bandwidth. Arithmetic (FLOPs) is only part of what a GPU does; every number must also be moved in and out of memory, and moving bytes is often slower than multiplying them. Generation is frequently limited by the second rather than the first. The memory section returns to this in detail.

Generation is a loop

A transformer is, by default, stateless. A forward pass takes a sequence of tokens and produces a distribution over the next token; it remembers nothing between calls. Generation therefore re-feeds the whole sequence at every step.

The context grows by one token each step, and the model is called again on the full, longer sequence. This is correct but wasteful, and the waste is the subject of the next section.

The waste

At every forward pass, each layer projects every token into query, key, and value vectors, then attends. Consider the position of the past tokens. When the sequence is “The cat sat”, the model computes keys and values for “The”, “cat”, and “sat”. When the sequence becomes “The cat sat on”, it computes them again, and the results are identical, because the embeddings of the first three tokens did not change.

Recomputing unchanged values is reading the whole book again to reach page 200. The keys and values of the past are determined once and never change; only the newest token introduces new keys and values. That stability is not an accident of the example; it is guaranteed by the causal mask. A token’s representation at any layer depends only on the tokens up to itself, never on tokens that come after it, so appending a token cannot change anything already written. This is the property that makes a cache sound. The cost of the redundancy grows with the square of the sequence length, because each new token re-processes the entire past.

The fix: cache keys and values

The waste is avoidable by remembering. Instead of recomputing the keys and values of every past token at every step, compute each token’s keys and values once, store them, and reuse them. That stored block is the KV cache.

What matters is which vectors are cached and which are not. Keys and values represent the context, which is fixed once written. Queries represent what the current token is looking for. A query is used the moment it is computed, to score the new token against the cache, and it is never needed again, so there is nothing to reuse; queries are computed fresh every step and never stored. The cache holds keys and values only.

The consequence is a change in behavior. The naive model is a pure function of the whole context. The cached model is stateful: it takes the newest token, computes its keys and values, appends them to the cache, and attends over the full cache. The past is read from memory instead of recomputed.

Prefill and decode

Generation splits into two phases with different shapes and different costs.

Prefill: reading the prompt

The prompt arrives all at once, so the model can process every token of it in parallel. For a prompt of $T$ tokens with embedding dimension $C$, the input is a matrix of shape $[1, T, C]$, and the projection to keys and values produces matrices of the same shape in one large matrix multiply. This phase is compute-bound: it is a big parallel computation, which is why the model appears to pause for a moment before any text streams.

The keys and values produced here are the first entries of the cache.

Decode: generating one token at a time

Generation then proceeds one token at a time. Each step receives a single token, a matrix of shape $[1, 1, C]$, and computes its key and value vectors. These are appended to the cache, and the query attends over the entire cache, including the token just added.

The distinction is the whole point. Only the newest token is projected; everything before it is read from memory. The projection work per step no longer grows with the length of the past, while the naive approach does. The query still scores against every cached key, so attention work does grow, but that is a memory read rather than a recomputation.

A worked example

Concrete numbers make the shapes obvious. Take an embedding dimension $C = 4$ and three prompt tokens, each a 4-dimensional vector:

Using identity projection matrices for keys and values, so that $K = V = X$, the prefill step computes the key and value rows for all three tokens and stores them. The cache is this matrix.

The model predicts a fourth token, say “on”, with embedding $[0, 0, 1, 1]$. Its key and value are computed from this single vector and appended:

The query for “on” is $[0, 0, 1, 1]$. The attention scores are the dot products of this query with each cached key, divided by $\sqrt{d_k} = \sqrt{4} = 2$:

• “The”: $(0\cdot1 + 0\cdot0 + 1\cdot1 + 1\cdot0) / 2 = 0.5$
• “cat”: $(0\cdot0 + 0\cdot1 + 1\cdot0 + 1\cdot1) / 2 = 0.5$
• “sat”: $(0\cdot1 + 0\cdot1 + 1\cdot0 + 1\cdot0) / 2 = 0$
• “on”: $(0\cdot0 + 0\cdot0 + 1\cdot1 + 1\cdot1) / 2 = 1$

A softmax turns these scores into weights: $\text{softmax}([0.5, 0.5, 0, 1]) \approx [0.235, 0.235, 0.143, 0.387]$. The output is the weighted average of the four cached value rows:

Only “on” was projected during this step; the other three rows came straight from memory. Real models use distinct learned projections for keys and values, so the two differ; the identity projection here only keeps the arithmetic transparent.

The compute accounting

Without a cache, the $k$-th generation step re-projects all $k$ previous tokens, a total of $1 + 2 + \dots + T$ projections over a full generation, quadratic in the sequence length. With a cache, each step projects exactly one token, so projection work is linear. Attention also stops being recomputed for past positions, so the total attention work drops from cubic to quadratic. What remains quadratic is attention itself: each new query still scores against every cached key, so decode slows as the context grows. The point of the cache is that this residual cost is a memory read rather than a recomputation, which is exactly why generation ends up bandwidth-bound rather than compute-bound.

The memory wall

The tradeoff is now explicit: caching turned a compute problem into a memory problem. The cache grows by one key and one value per head, so two vectors per head per layer, across every layer and every key-value head.

The size per token follows a formula:

where $L$ is the number of layers, $H_{KV}$ the number of key-value heads, $d$ the head dimension, and $b$ the bytes per element. For a model with 32 layers, 8 key-value heads (a grouped-query setup, where several query heads share one key-value pair, defined below), head dimension 128, and 16-bit values (2 bytes):

A 4,096-token conversation therefore needs $128 \text{ KB} \times 4096 = 512$ MB of cache, and serving 32 such conversations needs 16 GB on top of the model weights. This is why long contexts are expensive even when the model itself fits in memory, and why the context window is often limited by the cache rather than by the model.

Generation is also memory-bandwidth-bound rather than compute-bound. In decode the input is a single token, so the matrix multiply per step is tiny, yet the whole weight matrix and the entire cache still have to be streamed from high-bandwidth memory. The GPU does little arithmetic per byte moved, so bandwidth, not FLOPs, sets the token rate. A large model can spend tens of milliseconds per token simply moving weights. This is the same reason serving throughput, covered in the transformer note , is so sensitive to cache size.

Modern optimizations

The cache is the dominant memory consumer, so the optimizations attack its size or its layout.

• Grouped-query attention (GQA) and multi-query attention (MQA) reduce $H_{KV}$ by sharing keys and values across query heads. The memory formula shrinks with $H_{KV}$, which is why most modern models use GQA.
• PagedAttention fixes fragmentation. Instead of one contiguous block per sequence, the cache is split into fixed-size pages that need not be adjacent, with a block table mapping logical positions to physical pages. Fragmentation waste drops sharply and batching improves.
• Multi-head latent attention (MLA), used by DeepSeek models, compresses all key-value heads into a single latent vector per token, cutting cache size by an order of magnitude. The inference section of the transformer note covers this recipe in detail.
• Prefix caching reuses cached keys and values across requests that share a prompt, so a repeated system prompt or chat prefix is computed once and served many times; it is the serving-side counterpart to the per-sequence cache.
• Quantization stores cache elements in lower precision, such as FP8 or FP4, reducing $b$ in the memory formula at some accuracy cost.

All four reduce the same term: the bytes per token in memory, which raises how many concurrent sequences fit on a GPU and how fast tokens can be generated.

Summary

The complete picture in one paragraph: generation is a loop, and a stateless model recomputes the past at every step. Keys and values of past tokens never change, so they are computed once and cached; queries are transient and recomputed every step. The prompt is processed in parallel in a compute-bound prefill, then generation proceeds one token at a time, appending new keys and values and attending over the whole cache. The price is memory: the cache grows linearly with context, across every layer and head, and generation is ultimately bounded by how fast that memory can be streamed. Modern systems shrink the cache with shared heads, paged allocation, latent compression, and lower precision.