NVIDIA/kvpress

LLM KV cache compression made easy

What it solves

Deploying long-context Large Language Models (LLMs) is expensive because the key-value (KV) cache grows linearly with the number of tokens. For example, a 70B model handling 1 million tokens in float16 can require up to 330GB of memory. KVPress reduces this memory overhead and increases decoding speed by compressing the KV cache.

How it works

KVPress provides a library of "presses"—compression methods that prune or merge KV pairs based on various importance scores. These presses are integrated into the Hugging Face transformers pipeline via a custom KVPressTextGenerationPipeline.

  • Prefilling Compression: Compresses the cache during the initial processing of the input context.
  • Decoding Compression: Periodically compresses the cache during token generation to maintain a target size.
  • Scorer-based Pruning: Many methods (like KnormPress or SnapKVPress) assign scores to KV pairs and evict those with the lowest importance.
  • Quantization Support: It integrates with QuantizedCache to further reduce memory via lower-bit precision (e.g., 4-bit).

Who it’s for

Researchers and developers working on long-context LLM efficiency, specifically those looking to implement, benchmark, or deploy KV cache compression techniques to reduce GPU memory usage.

Highlights

  • Extensive Library: Implements a wide array of training-free compression methods including StreamingLLM, SnapKV, PyramidKV, and KVzip.
  • Flexible Integration: Works as a context manager or a custom transformers pipeline for easy adoption.
  • Hybrid Strategies: Supports combining different compression methods for the prefilling and decoding phases through PrefillDecodingPress.
  • Evaluation Tools: Includes a CLI and notebooks to measure accuracy, peak memory usage, and total time gains.