ElemedexA pocket encyclopedia of AI architectures

Glossary (TM) › Architectural parts

TM092

KV cache

Also: key–value cache

In Spanish: Caché KV

Memory in which a generative model stores the attention keys and values of tokens already processed, so as not to recompute them for each new token.

An autoregressive generation model generates one token at a time: it appends the last one to the text and runs everything through the network again (Alammar & Grootendorst, 2024, pp. 75–76). But thanks to the causal mask, the keys and values of earlier tokens do not change when a new one is added. So they are stored: at each step only the new token's query, key and value are computed, and the query is compared with the stored keys. That is the KV cache (Alammar & Grootendorst, 2024, pp. 83–84). In a test by Alammar & Grootendorst (2024), generating 100 tokens with a small model took 4.5 seconds with the cache and 21.8 without it (2024, p. 85).

The price is memory, and it grows with context length. For each token, a key and a value must be stored at every layer: roughly 2 × layers × d numbers. For GPT-3 (96 layers, d = 12,288; (Brown et al., 2020, sec. 2.1)) that comes to ≈ 2.4 million numbers per token: ≈ 4.7 MB at 16 bits, and ≈ 9.7 GB for a 2,048-token window (our own rough arithmetic). This is why the cache, not only the weights, limits how long the context window can be and how many users a server can handle at once.

How it is slimmed down. Multi-query attention makes all heads share a single key and value; grouped-query attention (used in Llama 2 and 3) shares them across groups of heads (Alammar & Grootendorst, 2024, pp. 98–100). Pope et al. (2023) show that sharing keys and values allows contexts up to 32 times longer. State-space models such as Mamba avoid the problem at the root: their memory does not grow.

A common misunderstanding. The "memory" of a conversation with ChatGPT (InstructGPT) is not something the model remembers: language models are stateless between calls, and to make them "remember" the history has to be fed back to them (Alammar & Grootendorst, 2024, p. 209). The KV cache is just a shortcut to avoid redoing arithmetic within a single generation.

Sources