NVIDIA KVPress: Toolkit for KV Cache Compression in Long-Context LLMs

NVIDIA has introduced KVPress, a Python toolkit designed to compress the Key-Value (KV) cache in Large Language Models (LLMs), enabling the efficient processing of long context windows. By reducing the memory overhead associated with storing intermediate attention results, KVPress allows models to handle larger sequences—such as those required for in-context retrieval and extended reasoning—without requiring prohibitive amounts of GPU memory.

The KV Cache Memory Bottleneck

In autoregressive LLMs, text generation occurs token by token. To avoid recalculating the representations of all previous tokens for every new token generated, models use a KV Cache to store the keys (K) and values (V) from attention layers. While this optimizes computation, the cache size scales linearly with the context window.

Memory consumption for the KV cache is determined by the following formula:

$$\text{Size} \left(\text{KV}\right) = 2 \times \text{precision} \times n_{layers} \times n_{heads} \times d \times n_{tokens}$$

For a model like Llama 3-70B running in bfloat16 precision with a 1M token context, the KV cache alone requires approximately 327.6 GB. When combined with the 140 GB required for model weights, the total memory demand reaches roughly 470 GB, meaning the KV cache accounts for approximately 70% of the total memory usage.

KVPress: A Modular Compression Toolkit

KVPress provides a suite of "presses"—compression algorithms that prune less important KV pairs to reduce the memory footprint. These presses are integrated into the model's attention layers via forward hooks and target the pre-filling phase, compressing the cache when it is at its largest.

Available Compression Techniques

KVPress implements several state-of-the-art pruning methods, including:

  • KnormPress: Prunes KV pairs based on the lowest key-value norm.
  • SnapKVPress: Prunes KV pairs associated with low attention weights for the most recent queries.
  • ExpectedAttentionPress: A new, unpublished technique created by the KVPress authors that prunes KV pairs based on the lowest expected attention weight for future queries.

KVPress is designed to be modular, allowing researchers to easily extend the toolkit with new methods and integrate it with other memory-saving techniques, such as KV Cache Quantization available in the transformers library.

Performance and Memory Impact

Compressing the KV cache reduces peak memory usage and improves generation speed. For Llama 3.1 8B in bfloat16 with a 128k context length, applying KVPress with a 50% compression ratio results in the following improvements:

  • Memory Reduction: Peak memory usage drops from 45GB to 37GB.
  • Speed Increase: Decoding speed increases from 11 tokens per second to 17 tokens per second on an A100 GPU.

Benchmarking and Accuracy Trade-offs

KVPress includes a CLI for benchmarking compression techniques on standard long-context datasets, including RULER, InfiniteBench, and Loogle.

Benchmarks on the RULER dataset (4k context length) indicate that the most effective combination is AdaKVPress and ExpectedAttentionPress. However, the toolkit highlights a critical trade-off: as the compression ratio increases, model accuracy generally declines. This underscores the need for continued research into compression algorithms that minimize the impact on model performance.

Sources