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，可測量準確度、峰值記憶體使用量與總體時間提升。