NVIDIA/kvpress

何を解決するか

長文コンテキストを持つ大規模言語モデル（LLM）を展開するのは高コストです。キーバリュー（KV）キャッシュはトークン数に比例して増大するため、例えば Llama 3.1-70B を float16 で 100 万トークン処理すると、最大 330GB のメモリが必要になることがあります。KVPress は KV キャッシュの圧縮により、これらのメモリ要件を削減し、デコード速度を向上させます。

仕組み

KVPress は、プレフィルフェーズまたはデコードフェーズ中に定期的に KV ペアを削除または統合する「プレス」と呼ばれるさまざまな圧縮手法を実装しています。Hugging Face の transformers ライブラリと統合され、トークン化やチャットテンプレートを処理するカスタム KVPressTextGenerationPipeline を提供しています。

圧縮は以下の戦略により実現されます：

• スコアベースの削除：重要度スコア（例：ランダム、逆ノルム、またはアテンション重み）を使用して、最も重要度の低い KV ペアを削除します。
• 構造的圧縮：特定のキャッシュパターンを維持します。たとえば、層間でピラミッド型のサイズを維持する、または初期および最近のトークンのみを保持する（StreamingLLM）。
• 次元削減：チャネルアテンションスコアに基づいてキーの次元を圧縮します。
• 高度な技術：CUR分解、コンテキスト再構成（KVzip）、または削除されたトークンを生存する隣接トークンに統合する手法を使用します。

対象ユーザー

長文コンテキスト LLM の効率化に取り組む研究者や開発者で、KV キャッシュ圧縮手法を実装・ベンチマーク・展開するための標準化されたフレームワークが必要な方々。

特徴

• 豊富なライブラリ：SnapKV、StreamingLLM、KVzip など、学習不要の圧縮手法を多数サポート。
• 柔軟なパイプライン：transformers パイプラインをカスタマイズし、簡単な統合と評価を可能に。
• 二段階圧縮：プレフィルフェーズおよびデコードフェーズ（DecodingPress を通じて）での圧縮をサポート。
• 量子化対応：QuantizedCache と互換性があり、さらなるメモリ削減を可能に。
• ベンチマークツール：CLI とノートブックを提供し、精度、ピークメモリ使用量、合計時間の改善を測定可能。