AI glossary
KV Cache
The KV cache (key-value cache) is a store of the key and value vectors, computed by a transformer’s self-attention layers for every token generated so far in a sequence. It is kept in memory so these vectors do not need to be recomputed at every new generation step.
When you run inference on a large language model, the system does not process the entire history of the conversation all at once for every new word it predicts. Instead, it relies on this cache to maintain context efficiently. Without it, the model would have to recalculate the attention weights for every single previous token every time it generates a new one, which would make the process prohibitively slow.
Why it is needed
In autoregressive generation, a language model produces one token at a time. Each new token’s prediction depends on attending to all previous tokens in the sequence. This means the model must understand the relationship between the current word and every word that came before it.
Without caching, generating each new token would require reprocessing the entire sequence from scratch through every attention layer. You would repeat a large amount of computation that had already been done for previous tokens. For long contexts, this redundancy becomes a major bottleneck.
With a KV cache, only the newest token’s query vector needs to be computed and compared against the already-cached keys and values of all previous tokens. This avoids redundant computation and significantly speeds up generation. The model essentially remembers the “context” of the conversation so it doesn’t have to re-read the entire history each time it speaks.
How it works
Inside each transformer attention layer, the model computes a query, key, and value vector for the current token. The key and value vectors for every token processed so far are stored in the KV cache. There is one entry per attention layer per token.
When generating the next token, the model computes only that token’s query vector. It then attends this query against all the cached keys and values, producing the output for that step. Finally, the model appends the new token’s own key and value vectors to the cache for future steps.
This process is fundamental to how the attention mechanism functions in sequence modeling. Each new token enriches the cache, allowing subsequent tokens to attend to the full context without recalculating the contributions of earlier tokens. The Transformer architecture relies heavily on this efficient storage pattern to scale to large context windows.
The memory cost
The KV cache grows linearly with the length of the sequence and with the size of the model. Specifically, it scales with the number of layers, the number of attention heads, and the dimension of each head. This cache must be kept in GPU memory alongside the model’s own weights.
For long-context models handling very large inputs, the KV cache can become the dominant consumer of GPU memory during inference. In many cases, it uses more memory than the model’s weights themselves. This limits how many requests can be served in parallel (the batch size) on a given amount of hardware.
If you are running multiple users on a single GPU, the KV cache is often the first resource to run out, not the model parameters. This is why understanding kv cache memory behavior is critical for deploying models at scale.
Reducing the memory cost
Several techniques exist to shrink the KV cache and reduce its memory footprint. These methods often involve architectural changes or optimizations to how the cache is stored.
Multi-query attention (MQA)
Multi-query attention (MQA) has all query heads share a single set of key and value heads. This sharply reduces the size of the cache compared to standard multi-head attention. The trade-off is a potential cost to model quality, as sharing keys and values reduces the expressiveness of the attention mechanism.
Grouped-query attention (GQA)
Grouped-query attention (GQA) is a middle ground. Query heads are divided into groups, and each group shares one set of key and value heads. This balances the memory savings of MQA against the quality of standard multi-head attention. GQA is used in models such as Llama 2, Llama 3, and Mistral. It offers a practical compromise for many production environments.
Quantization
The cache itself can also be stored at lower numeric precision (quantized) to reduce its memory footprint further. By reducing the precision of the stored key and value vectors, you save space without changing the model architecture. This is often combined with other techniques like quantization of the model weights.
Limitations
Techniques that shrink the KV cache, such as MQA and GQA, generally involve a trade-off against model quality. They must be designed into the model’s architecture from the start. This means they cannot always be applied to an already-trained model without retraining or fine-tuning.
Because the cache must be kept for the full length of an active generation, very long conversations or documents can still exhaust available GPU memory even with these optimizations. The linear growth means that doubling the context length roughly doubles the memory required for the cache.
For extremely long sequences, other strategies like speculative decoding or sliding window attention might be necessary to manage memory usage effectively. Understanding these limits helps you choose the right model configuration for your hardware constraints.
FAQ
What is a KV cache?
It is a memory store that holds the key and value vectors for all previously generated tokens in a sequence. This allows the model to avoid recomputing attention for past tokens during autoregressive generation.
Why does KV cache consume so much GPU memory?
The cache size grows linearly with the sequence length and the model’s depth. For long contexts, the cache can exceed the memory usage of the model weights, limiting how many parallel requests a GPU can handle.
What is the difference between MQA and GQA?
Multi-query attention (MQA) shares one set of keys and values across all query heads, maximizing memory savings but potentially reducing quality. Grouped-query attention (GQA) divides heads into groups that share keys/values, offering a balance between memory efficiency and model performance.
Can I reduce KV cache size after training?
You can reduce memory usage by storing the cache at lower precision (quantization). However, architectural changes like MQA or GQA usually require the model to be trained or fine-tuned with those specific attention heads to work correctly.