NVIDIA/kvpress
LLM KV cache compression made easy
What it solves
長コンテキストの大規模言語モデル(LLM)のデプロイはコストが高いです。なぜならキー・バリュー(KV)キャッシュがトークン数に比例して線形に増大するからです。例えば、70B モデルが float16 で 100 万トークンを処理すると、最大で 330GB のメモリが必要になることがあります。KVPress は KV キャッシュを圧縮することでこのメモリオーバーヘッドを削減し、デコード速度を向上させます。
How it works
KVPress は「presses」― さまざまな重要度スコアに基づいて KV ペアを剪定またはマージする圧縮手法のライブラリを提供します。これらの presses はカスタム KVPressTextGenerationPipeline を介して Hugging Face transformers パイプラインに統合されます。
- Prefilling Compression:入力コンテキストの初期処理中にキャッシュを圧縮します。
- Decoding Compression:トークン生成中に定期的にキャッシュを圧縮し、目標サイズを維持します。
- 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:コンテキストマネージャまたはカスタム transformers パイプラインとして利用でき、導入が容易。
- Hybrid Strategies:
PrefillDecodingPressを通じて、プレフィル段階とデコード段階で異なる圧縮手法を組み合わせ可能。 - Evaluation Tools:CLI とノートブックが同梱されており、精度、ピークメモリ使用量、総時間の改善を測定できます。