What is a KV cache?

What does a KV cache store?

A KV cache is the saved set of key and value vectors that a transformer's attention layers compute for every token it has already processed. The PagedAttention paper (Kwon et al., 2023) describes how, at each attention layer, a token's hidden state is projected into a query, a key and a value vector. Generating a new token requires comparing its query against the key vectors of all prior tokens. It then blends their value vectors.

Recomputing those keys and values from scratch at every step would mean reprocessing the whole sequence each time a single new token is generated. Instead, the cache keeps the key and value vectors from earlier tokens in memory. Each new decoding step then only needs to compute the key and value for the newest token and reuse the rest. Shazeer (2019) frames this as the reason incremental decoding is memory-bandwidth bound: the model must reload those cached tensors at every step.

Why does it dominate memory at inference?

The cache grows with every token generated, and it grows across every layer and every attention head at once. The PagedAttention paper gives a concrete figure: for a 13-billion-parameter OPT model, the KV cache for a single token takes 800 kilobytes. That figure comes from two vectors, key and value, times a 5,120-dimension hidden state, times 40 layers, times 2 bytes per value in 16-bit precision. Because OPT can generate sequences up to 2,048 tokens long, one request's cache can reach 1.6 gigabytes.

On a 13-billion-parameter model served on an NVIDIA A100 with 40 gigabytes of memory, the PagedAttention paper found that roughly 65% of memory holds the static model weights. Close to 30% holds the KV cache, leaving only a small remainder for activations. Weights are fixed and activations are small, so the cache is what limits how many requests a server can batch together, and therefore its throughput.

Multi-query and grouped-query attention: what's the difference?

Multi-query attention, proposed by Shazeer (2019), has every query head share a single set of key and value vectors instead of each head keeping its own. This shrinks the cache by a factor equal to the number of heads, which is what consumes memory bandwidth during incremental decoding. On WMT 2014 English-German translation, Shazeer's multi-query model scored 28.5 BLEU with beam-4 decoding, matching or beating the multi-head baseline. It also cut the decoder's incremental inference cost from 46 microseconds per token to 3.8.

Grouped-query attention, introduced by Ainslie et al. (2023), sits between the two: query heads are split into groups, and each group shares one key and value head. Grouped-query attention with one group is identical to multi-query attention; with as many groups as heads, it is identical to ordinary multi-head attention. On T5-XXL, Ainslie et al. report that grouped-query attention with 8 groups reached an average score of 47.1 across summarisation, translation and question-answering benchmarks, close to multi-head attention's 47.2. It ran in 0.28 seconds per sample versus multi-head's 1.51 seconds.

Ainslie et al. also show that multi-head checkpoints can be converted to multi-query or grouped-query attention cheaply, through a process they call uptraining. Mean-pooling the existing key and value projection heads into fewer heads, then retraining for about 5% of the original pre-training compute, recovers most of the quality lost from the conversion.

How does paged memory fix KV cache waste?

Before PagedAttention, serving systems stored each request's KV cache as one contiguous block of memory, sized for the request's maximum possible length. The PagedAttention paper found this wasted memory three ways: space reserved for tokens not yet generated, internal fragmentation when a request ends up shorter than its reserved maximum, and external fragmentation from mismatched chunk sizes. Across their experiments, only 20.4% to 38.2% of allocated KV cache memory held actual token data.

PagedAttention, built by Kwon et al. (2023) into a serving engine called vLLM, borrows the idea of paged virtual memory from operating systems. It splits each request's cache into fixed-size blocks that need not sit next to each other in physical memory, tracked through a per-request block table. This works much as an OS maps a program's logical pages to scattered physical ones. Blocks are allocated only as new tokens arrive, which the paper reports brings KV cache waste close to zero.

Because blocks are addressed indirectly, separate sequences can also point to the same physical block. The paper reports that this lets parallel sampling and beam search share the KV cache of a common prompt, saving up to 55% of memory in beam search. A copy-on-write scheme protects blocks once a sequence needs to diverge.

What does this mean for serving throughput?

The PagedAttention paper reports that vLLM improved throughput by 2 to 4 times over FasterTransformer and Orca, two prior serving systems, at the same latency, measured across various models and workloads. The authors say the gain grows with longer sequences, larger models and more complex decoding algorithms such as beam search. These are precisely the cases where KV cache waste under the old contiguous-allocation approach was worst.

None of this changes how much computation a transformer does; it changes how much of a GPU's memory and bandwidth that computation wastes. For anyone choosing a serving stack, the attention variant a model was trained with sets a hard ceiling on batch size. So does the memory manager the server uses.