Optimizing LLM KV Cache Memory for Cost-Effective AI Inference

The KV Cache: A Core LLM Memory Component

The KV cache is a crucial working memory block in LLMs, distinct from model weights, that stores key and value vectors for every input token. This caching prevents redundant recomputation during token generation, significantly speeding up the attention mechanism. However, its size grows linearly with context length and batch size, making it a primary driver of GPU memory consumption and a bottleneck for long-context LLM inference.

• Prefill Phase (Compute-Bound): The model processes all input tokens in parallel, building the KV cache. This phase is limited by the GPU's arithmetic speed.
• Decoding Phase (Memory-Bound): The model generates output one token at a time. Each new token requires reading the *entire* cache from GPU memory, making this phase limited by memory bandwidth. Larger caches directly translate to slower and costlier generation.

Techniques for KV Cache Optimization

Optimizing the KV cache involves pushing against several factors that determine its size. These techniques can be broadly categorized into architectural changes during model training and serving-side optimizations for deployed models.

Architectural Optimizations (Model Training)

• Grouped-Query Attention (GQA): Allows multiple query heads to share a single key-value head, drastically reducing the number of stored key-value sets. This can cut cache size by factors like 8x (e.g., Llama 3 70B vs. older models). Multi-query attention (MQA) is an even more aggressive version but can impact quality.
• Multi-Head Latent Attention: Compresses key and value vectors into a smaller latent representation before caching, then expands them back on read. Offers significant memory savings (e.g., DeepSeek-V3), but adds computational overhead during serving, best for very large models and contexts.

Post-Training and Serving-Side Optimizations

• Quantization: Reduces the precision of stored key-value numbers (e.g., from 16-bit to 8-bit or 4-bit). This directly halves or quarters the cache size. 8-bit quantization often has minimal quality impact, while 4-bit may show measurable losses on demanding tasks.
• Eviction: Drops less relevant tokens from the cache, typically retaining a window of recent tokens and initial 'anchor' tokens. The challenge is predicting which tokens will be needed later, as aggressive eviction can lead to information loss and degraded performance on retrieval tasks.
• Paged Attention: Borrows from operating system memory management, splitting the KV cache into small, fixed-size blocks (pages) that can be allocated on demand and stored non-contiguously. This significantly reduces memory fragmentation (from 60-80% wasted space to <4%) and improves throughput (2-3x).
• Prefix Caching (Prompt Caching): Leverages paged attention to share identical initial prompt segments across multiple requests. If two requests start with the same text, they can share the same physical KV cache blocks, leading to substantial cost and latency reductions (50-90% savings for cached tokens) for workloads with repetitive prompts.