KV cache

The KV cache is where a language model keeps its intermediate results while it writes an answer. The model writes piece by piece, and each new piece depends on all the text before it. The cache stores the work already done on that text, so each step only processes the newest piece. Answers come faster, but the cache takes memory, so longer conversations mean fewer users per server.

How does it work?

In the attention mechanism of a transformer, every token gets three vectors: a query, a key and a value (Vaswani et al., 2017). To predict the next token, the model compares the query of the last token with the keys of all earlier tokens and uses the scores to weigh their values. A token only looks back, so its key and value never change once computed. According to the Hugging Face documentation, they can therefore be stored and reused: each step computes a key and a value only for the new token, in every layer.

Serving a request has two phases (Kwon et al., 2023). In the prompt phase, the model processes the whole prompt at once and fills the cache. In the generation phase, it adds one token per step. Those steps cannot run in parallel, and they take most of the time a request needs. The new token is still compared with every stored key, so the attention cost per step grows linearly with the length of the text instead of quadratically.

How much memory does it take?

The cache holds a key and a value per token for every layer. Kwon et al. calculate it for the 13-billion-parameter OPT model: 2 (key and value) × 5,120 (hidden size) × 40 (layers) × 2 bytes makes 800 KB per token, or up to 1.6 GB for a request of 2,048 tokens. On a 40 GB NVIDIA A100 with a 13B model, about 65% of the memory held the weights and close to 30% the cache. The cache grows with every token and every request, so it limits how many requests a GPU can handle at once. Models such as Mamba keep a state of fixed size instead (Lahoti et al., 2026).

How can the cache be made smaller?

• Share keys and values between attention heads. Step-by-step decoding is often slow because the large key and value tensors are loaded over and over (Shazeer, 2019). In his multi-query attention (MQA), all heads share one key and one value, which makes decoding much faster with a minor loss of quality. In grouped-query attention (GQA), groups of heads share them, with quality close to standard multi-head attention at a speed comparable to MQA (Ainslie et al., 2023). Meta used GQA for the 34B and 70B versions of Llama 2.
• Manage the memory in pages. Older serving systems reserved space for each request's maximum length, and in Kwon et al.'s measurements only 20.4% to 38.2% of the cache memory held actual token data. Their PagedAttention, the basis of vLLM, stores the cache in small blocks with near-zero waste, and reached 2 to 4 times the throughput of FasterTransformer and Orca at the same latency.
• Keep fewer tokens. H2O (Zhang et al., 2023) keeps the recent tokens and the few heavy hitters that contribute most to attention. With 20% heavy hitters, it reached up to 29 times the throughput of DeepSpeed Zero-Inference and Hugging Face Accelerate.
• Use fewer bits. vLLM can store the cache in 8 bits (FP8 KV cache), a form of quantization.
• Reuse it. When requests start with the same text, such as a long system prompt, vLLM can reuse the cache of that shared start (automatic prefix caching). API providers offer something similar as prompt caching.

Example

Suppose a company runs OPT-13B on one 40 GB A100 for its customer service chat. At 800 KB per token, a conversation of 2,000 tokens takes about 1.6 GB of cache. The roughly 12 GB for the cache (close to 30% of 40 GB) holds about seven such conversations at once. An 8-bit cache halves the 800 KB per token and doubles that number.