KV Cache and Memory Management

KV Cache and Memory Management

When a language model generates text, it does not recompute everything from scratch for each new token. It uses a clever optimization called the KV cache, which stores intermediate results and dramatically speeds up generation. Understanding the KV cache is essential for understanding memory usage during inference.

KV stands for Key and Value, which are the same keys and values we discussed in the attention mechanism article. When the model processes a token, it computes keys and values for that token. These are needed again for every subsequent token because each new token needs to attend to all previous tokens. Instead of recomputing them, the model caches them.

The KV cache grows as the conversation continues. For each new token generated, the model adds new keys and values to the cache. For a long conversation with a 70B model, the KV cache can grow to several gigabytes. This is why long conversations consume more memory than single turn interactions.

There are several optimization techniques for the KV cache. Multi-Query Attention and Grouped Query Attention reduce the cache size by sharing keys and values across multiple attention heads. Techniques like KV cache quantization compress the cached values to use less memory. Some inference engines also support cache eviction, removing old tokens from the cache when it gets too large.

The KV cache is also why context size matters so much for memory. A 128K context with a 70B model can require tens of gigabytes just for the KV cache, on top of the model weights. This is why running models with very long contexts requires so much VRAM, and why efficient KV cache management is an active area of research.