Zefan-Cai/KVCache-Factory

何を解決するか

KVCache-Factory は、長文脈 Large Language Model (LLM) 推論における Key-Value (KV) キャッシュに関連するメモリのボトルネックに対処します。文脈長が増加するにつれて、KV キャッシュは線形に増大し、膨大な GPU メモリを消費し、デコードを遅くします。このプロジェクトは、モデルのパフォーマンスを損なうことなく、このメモリ使用量を削減するためのさまざまな戦略を実装・比較・評価できる統合フレームワークを提供します。

仕組み

このプロジェクトは、複数の KV キャッシュ最適化技術を統合した「実験場」として機能し、単一の評価インターフェースで動作します。以下の複数のカテゴリの手法をサポートしています：

圧縮と除外：重要度の低いトークンを削除（例：StreamingLLM、H2O、SnapKV、PyramidKV）。
検索：現在のクエリに最も関連するトークンのみを選択（例：Quest、L2Norm）。
マージ：複数の KV ペアを1つに集約（例：CAM）。
量子化：キャッシュされた値の精度を低下（例：KIVI、KVQuant、GEAR）。
オフロード：KV キャッシュを CPU メモリに移動し、GPU に1つずつストリーミングして戻すことで、GPU メモリのオーバーフローを防ぎながら完全な精度を維持（HeadInfer）。

LongBench、RULER、および「Haystack に針を探す」などの標準ベンチマーク用のランナーを含み、これらの手法が精度と遅延に与える影響を測定できます。

対象ユーザー

このツールは、LLMの効率性に取り組む AI リサーチャーおよび開発者向けに設計されており、特に長文脈ウィンドウのパフォーマンス最適化や、さまざまな KV キャッシュ管理戦略のベンチマークを検討している人にとって有用です。

特徴

統合インターフェース：PyramidKV、SnapKV、H2O、StreamingLLM などの幅広いベースラインをサポート。
多様な最適化戦略：除外、検索、マージ、量子化、および損失なしのオフロードを1つの場所で統合。
包括的なベンチマーク：LongBench、RULER、Needle-in-a-haystack の評価に対応。
ハードウェアの柔軟性：大規模モデル（例：Llama-3-70B）向けのマルチGPU推論をサポートし、FlashAttention v2 と SDPA の両方のパスを提供。
可視化ツール：注意機構のパターンを可視化するユーティリティを備えており、異なる圧縮手法がモデルの注目点にどのように影響するかを理解できます。