NVIDIA/kvpress

LLM KV cache compression made easy

What it solves

部署长上下文的大型语言模型（LLM）成本高昂，因为键值（KV）缓存会随 token 数量线性增长。例如，一个 70B 模型在 float16 下处理 100 万个 token 可能需要高达 330GB 内存。KVPress 通过压缩 KV 缓存来降低此内存开销，并提升解码速度。

How it works

KVPress 提供一套「presses」——根据各种重要性分数剪枝或合并 KV 对的压缩方法。这些 presses 通过自定义的 KVPressTextGenerationPipeline 集成到 Hugging Face transformers pipeline 中。

  • Prefilling Compression：在输入上下文的初始处理阶段压缩缓存。
  • Decoding Compression：在 token 生成过程中定期压缩缓存，以保持目标大小。
  • Scorer-based Pruning：许多方法（如 KnormPress 或 SnapKVPress）会为 KV 对打分，并淘汰重要性最低的项。
  • Quantization Support：它可与 QuantizedCache 结合，通过低位精度（例如 4 位）进一步降低内存需求。

Who it’s for

从事长上下文 LLM 效率优化的研究人员和开发者，特别是想实现、基准测试或部署 KV 缓存压缩技术以减少 GPU 内存使用的群体。

Highlights

  • Extensive Library：实现多种无需训练的压缩方法，包括 StreamingLLM、SnapKV、PyramidKV 与 KVzip。
  • Flexible Integration：可作为 context manager 或自定义的 transformers pipeline 使用，方便快速采纳。
  • Hybrid Strategies：支持通过 PrefillDecodingPress 结合不同压缩方法于预填和解码阶段。
  • Evaluation Tools：提供 CLI 与 notebook，可测量准确率、峰值内存使用和总体时间提升。