Memoizing Context via KV Caching
The Key-Value (KV) cache is a mechanism to bypass redundant calculations. In the attention layers of a transformer, each token embedding is transformed into three vectors: Query (Q), Key (K), and Value (V). When predicting a new token, the Query vector of the current token must be compared against the Key vectors of all previous tokens to determine the attention weights. The resulting weights are then used to compute a weighted sum of the Value vectors.
By storing the computed Keys and Values for every token generated so far in a cache, the system only needs to perform the forward pass for the single newly generated token. The Attention layer simply retrieves the cached Keys and Values for the entire history, concatenates them with the current token's Key and Value, and proceeds with the computation. This transforms the complexity of each generation step from quadratic to linear relative to sequence length.
Trade-offs and Memory Management
While KV caching drastically improves inference latency, it introduces significant memory overhead. Each token adds new entries to the cache, consuming VRAM. As the context window grows, the memory footprint increases linearly. Managing this cache effectively is a critical engineering challenge:
- Memory Fragmentation: Traditional contiguous buffers for caches often lead to memory fragmentation, limiting the number of concurrent requests.
- Paged Attention: Advanced implementations partition the cache into blocks, similar to virtual memory in operating systems, allowing non-contiguous storage of context blocks to minimize fragmentation.
- Quantization: Reducing the precision of cached tensors (e.g., from FP16 to INT8) can double the effective cache capacity at the cost of slight numerical degradation.
Engineering Takeaways
For engineers scaling production inference, the KV cache is the primary bottleneck for concurrency. Maximizing throughput requires balancing the sequence length against the number of active requests the GPU memory can sustain. Effective cache eviction policies, block-based memory management, and precision reduction are essential tools for optimizing LLM serving infrastructure.
