KV Cache Compression

Quantization reduces the numerical precision of the cached key and value tensors. Instead of storing them in full 16-bit (FP16/BF16) or 32-bit (FP32) floating-point format, they are compressed to 8-bit integers (INT8) or 8-bit floating-point (FP8).

• Mechanism: A quantization function maps a range of high-precision values to a smaller set of discrete low-precision values, often using a scaling factor and zero-point.
• Impact: Reduces the cache's memory footprint by 50% (FP16→INT8) or 75% (FP32→INT8), directly increasing the maximum possible context length. This also lowers memory bandwidth pressure during the attention computation.
• Challenge: Requires careful calibration to minimize the loss in model accuracy (perplexity increase) due to the approximation error. Techniques like per-tensor or per-channel quantization and smoothing are used to mitigate this.

Pruning involves selectively removing entries from the KV cache that are deemed less important for future token generation. It is a lossy compression technique that trades some model quality for significant memory savings.

• Spatial Pruning: Drops entire tokens from the cache. Policies include:
  • StreamingLLM's Attention Sink: Retains the first few tokens (the "sink") and a sliding window of recent tokens.
  • H2O (Heavy-Hitter Oracle): Keeps tokens that have historically received high attention scores.
  • Random Eviction: A simple baseline for comparison.
• Structural Pruning: Drops specific dimensions or heads within the cached tensors. For example, head pruning removes entire attention heads from the cache if they are found to be redundant.
• Use Case: Essential for enabling infinite-length generation with models trained on finite context windows, or for operating within strict, fixed memory budgets.

Low-Rank Approximation exploits the inherent structure and redundancy in the key and value matrices by representing them using a compressed, factorized form. The core idea is that the high-dimensional KV cache can often be approximated by the product of smaller matrices.

• Mechanism: Techniques like singular value decomposition (SVD) are applied to project the keys and values into a lower-dimensional subspace. Instead of caching the full [sequence_length, d_model] tensor, the system caches a smaller [sequence_length, k] projection matrix and a [k, d_model] basis matrix, where k << d_model.
• Benefit: Provides a mathematically grounded, lossy compression method. The compression factor is tunable via the chosen rank k, offering a direct trade-off between memory savings and approximation fidelity.
• Application: More computationally intensive to apply dynamically than pruning or quantization, and is often researched as a method to compress cached states between decoding steps or for long-term context storage.

Selective Caching is a policy-based technique where the system decides what to cache in the first place, rather than caching all computed keys and values. It prevents non-essential data from entering the cache, acting as a form of proactive compression.

• Layer Skipping: In deep transformer models, keys and values from certain intermediate layers may be skipped and not cached if analysis shows they contribute less to the final output.
• Token Importance Gating: A lightweight learned or heuristic gating function evaluates each token as it is processed in the prefill phase. Only tokens that pass an importance threshold have their KV states written to the long-term cache.
• Benefit: Reduces the working set size of the cache from the moment of creation, improving memory efficiency and potentially reducing write amplification. It is often combined with other techniques like pruning for a multi-stage compression pipeline.

This approach compresses the KV cache by fundamentally changing the transformer's attention mechanism. Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) are model architecture changes that reduce the number of unique key and value heads.

• Multi-Query Attention (MQA): All query heads share a single key head and a single value head. This drastically shrinks the KV cache size by a factor equal to the number of attention heads (e.g., 32x).
• Grouped-Query Attention (GQA): A generalization of MQA where query heads are divided into groups, and each group shares a single key and value head. This provides a tunable trade-off: an 8-group GQA offers an 8x cache reduction with less quality degradation than full MQA.
• Impact: These are pre-training techniques. Models like Llama 2 and 3 use GQA. The compression is "baked in," requiring no runtime algorithm, making it one of the most effective methods for production inference systems.

Offloading is a systems-level compression technique that treats high-cost GPU memory as a cache for a larger, slower memory hierarchy. It moves portions of the KV cache from GPU RAM to CPU RAM or even NVMe storage.

• Mechanism: A caching policy (e.g., LRU) decides which blocks of the KV cache reside in fast GPU memory. "Cold" blocks are evicted to slower, higher-capacity CPU memory or SSD. When needed again, they are fetched back, causing increased I/O latency.
• Technology Enabler: Relies on Unified Virtual Memory (UVM) frameworks like CUDA UVM or high-speed interconnifts like NVLink to manage the coherent address space between GPU and CPU.
• Use Case: Critical for supporting extremely long context windows (e.g., 1M+ tokens) that far exceed available GPU memory. It explicitly trades latency for capacity, making it suitable for offline or non-latency-sensitive batch processing of long documents.

A primary compression technique that reduces the numerical precision of stored key and value tensors. Instead of standard 16-bit (FP16/BF16) or 32-bit (FP32) formats, the cache is stored in 8-bit (INT8/FP8) or 4-bit formats.

• Mechanism: Applies post-training quantization (PTQ) or quantization-aware training (QAT) to the cache activations.
• Impact: Can reduce the KV cache memory footprint by 50% or more, directly enabling longer contexts or larger batch sizes.
• Trade-off: Introduces potential quality degradation that must be managed via calibration and selective quantization of sensitive layers.

Architectural modifications that structurally reduce the size of the KV cache by sharing key and value projections across multiple query heads.

• Multi-Query Attention (MQA): All query heads share a single key head and a single value head. This creates a minimal cache but can impact model quality.
• Grouped-Query Attention (GQA): A tunable middle ground where queries are grouped; each group shares a single key and value head. It offers a favorable trade-off, used in models like Llama 2 and 3.
• Cache Reduction: These methods directly shrink the KV cache size by a factor equal to the number of query heads per shared key/value head, reducing memory bandwidth pressure.

A compression technique that removes entries deemed non-essential from the KV cache based on a scoring heuristic.

• Common Policies: Include attention score-based pruning (dropping key-value pairs with low aggregate attention) and magnitude-based pruning.
• Dynamic vs. Static: Can be applied statically after the prefill phase or dynamically during the decode phase.
• Objective: To maintain a fixed-size "working set" of the most relevant context, allowing the model to operate within a target memory budget for arbitrarily long sequences, albeit with potential information loss.

An attention pattern that imposes a hard, fixed-length limit on the cache by having each token only attend to the W tokens that immediately precede it.

• Bounded Cache: The KV cache size becomes O(batch_size * W) instead of O(batch_size * sequence_length), enabling linear-time complexity.
• Use Case: Fundamental to models like Mistral 7B and Llama 3, which use a 128K context window but a sliding window of 4K or 8K tokens for efficient long-context inference.
• Trade-off: The model cannot leverage information outside the immediate window, which may be unsuitable for tasks requiring long-range dependency.

A method for enabling infinite-length generation with models trained on finite windows, relying on the attention sink phenomenon.

• Attention Sink: The initial tokens (e.g., the first 4) of a sequence receive disproportionately high attention scores. StreamingLLM framework keeps these sink tokens in the KV cache permanently.
• Mechanism: It maintains a cache composed of the sink tokens plus the most recent tokens (sliding window), recycling the rest. This provides a stable numerical anchor for attention calculation.
• Outcome: Allows models to generate text far beyond their trained context window without retraining, though recall of mid-range information is limited.

A memory management algorithm, not a compression technique, that enables highly efficient utilization of the KV cache memory space.

• Analogy: Manages the cache in fixed-size blocks (pages), similar to virtual memory in an OS.
• Benefit: Eliminates internal fragmentation caused by variable-length sequences in continuous batching, allowing near-100% memory utilization.
• Synergy with Compression: PagedAttention's block-based management is highly compatible with compression techniques like quantization, which can be applied at the page level. It is the core innovation behind the vLLM inference server.