NVIDIA/kvpress

KVPress – 長文LLM向けKVキャッシュ圧縮

概要

• トランスフォーマー系言語モデルのキー・バリュー（KV）キャッシュをプリフィリング段階（および必要に応じてデコーディング段階）でプルーニングまたは圧縮するアルゴリズム（プレス）を幅広く実装したPythonライブラリです。KVキャッシュを削減することでメモリ使用量が減り、非常に長いコンテキスト（例：100万トークン）での推論が高速化される可能性があります。

重要性

• 現代のLLMでは、KVキャッシュは処理トークン数に比例して増加します。70Bモデルを100万トークンのドキュメントで実行する場合、300GB以上のGPUメモリが必要になることがあります。KVPressを使うと、研究者や開発者はモデルコードを書き換えることなく、最先端の圧縮手法を多数試すことができます。

主な特長

• プラグ＆プレイのパイプライン – KVPressTextGenerationPipeline はHugging Faceパイプライン（kv-press-text-generation）として登録されます。
• 豊富なプレスカタログ – 単純なランダムプルーニングから、高度なアテンションベースや行列分解の手法（例：KnormPress、SnapKVPress、KVzipPress、CURPress）まで、30以上の既製手法を提供。
• プリフィルとデコーディングの圧縮 – DecodingPress および関連ラッパーにより、トークン生成中の定期的な圧縮が可能。
• 構成可能 – ラッパープレス（AdaKVPress、ComposedPress、ChunkPress など）を使って戦略を積み重ねたり組み合わせたりできます。
• トレーニング不要 – リストされたすべてのプレスは追加のモデル微調整なしで動作します。
• 量子化キャッシュのサポート – Hugging FaceのQuantizedCache（例：optimum-quantoによる4ビット）と連携。
• マルチGPU対応 – accelerate のデバイスマップと互換性があります。
• ベンチマークとリーダーボード – 評価CLIとノートブックに加え、RULER 4kトークンベンチマークでの公開HF Spaceリーダーボードを提供。

インストール

pip install kvpress               # 基本インストール
# または、開発用にオプションの依存関係（評価、フラッシュアテンション）を含める場合
uv sync --extra eval --extra flash-attn

基本的な使用例

from transformers import pipeline
from kvpress import ExpectedAttentionPress

pipe = pipeline(
    "kv-press-text-generation",
    model="Qwen/Qwen3-8B",
    device_map="auto",
    dtype="auto",
)

context = "A very long text you want to compress once and for all"
question = "\nA question about the compressed context"
press = ExpectedAttentionPress(compression_ratio=0.5)

answer = pipe(context, question=question, press=press)["answer"]
print(answer)

プレスは、モデルが質問に注意を向ける前にコンテキストトークンのKVキャッシュを圧縮するため、異なる圧縮率を比較できます。

評価

• 各プレスを長文コンテキストデータセットで実行し、精度、メモリ削減、高速化を報告するCLIを提供。
• リポジトリには、対話型プロファイリング用のノートブック（speed_and_memory.ipynb、wikipedia_demo.ipynb）が含まれています。
• RULER 4kトークンベンチマークの結果は、READMEにリンクされているHugging Faceリーダーボードに表示されます。

ドキュメントとリソース

• Colabノートブックデモ、HF Spaceデモ、ブログ記事でライブラリを紹介。
• 各プレスには、ソースファイルと対応するarXiv論文へのリンクがあり、深く学べます。
• 「新しいプレス」ノートブックは、カスタム圧縮手法の追加方法をガイドします。

ライセンス

• Apache 2.0（READMEのバッジを参照）。

引用

@article{devoto2025expectedattention,
  title={Expected Attention: KV Cache Compression by Estimating Attention from Future Queries Distribution},
  author={Devoto, Alessio and Jeblick, Maximilian and J{\'e}gou, Simon},
  journal={arXiv preprint arXiv:2510.00636},
  year={2025},
  url={https://arxiv.org/abs/2510.00636}
}

入手先

• GitHub: https://github.com/NVIDIA/kvpress
• PyPI: kvpress
• Hugging Face Space: https://huggingface.co/spaces/nvidia/kvpress