FutureMLS-Lab/OSCAR

OSCAR: Offline Spectral Covariance-Aware Rotation for 2-bit KV Cache Quantization

何を解決するか

OSCARは、大規模言語モデルにおけるKV(キー・バリュー)キャッシュの高いメモリオーバーヘッドに対処します。これは、しばしばコンテキスト長やスループットを制限する原因となります。他の2ビット手法が通常失敗するような、複雑な推論やコーディングタスクにおいても、元の16ビット(BF16)精度に近い精度を維持しながら、KVキャッシュの2ビット量子化(INT2)を可能にします。

仕組み

OSCARは、モデルが実際に注意機構(attention mechanism)をどのように使用しているかを分析するオフラインキャリブレーションプロセスを使用します。小さなデータセット上で活性化(activations)を取得し、キーとバリューの「注意に配慮した」共分散構造を推定します。これにより、各レイヤーごとの直交回転とクリッピングしきい値を導出し、量子化プロセスをモデルが精度のために最も必要とする方向に合わせます。

品質をさらに保護するために、ミックスプレシジョン戦略を採用しています:初期トークンの小さな「シンク」と最近のトークンウィンドウはBF16で保持され、キャッシュの大部分はINT2で保存されます。これは、高性能を実現するための統合ミックスプレシジョンFlash-Attentionカーネルによって実装されています。

対象ユーザー

  • LLMデプロイャー:限られたハードウェア上で、コンテキストを大幅に延長したり、より大きなバッチサイズをサポートするために、KVキャッシュメモリを約8倍削減したい方。
  • ローカルLLMユーザー:コンシューマーハードウェア(例:MacBook)上でモデルを実行している方で、長文コンテキストで大規模モデル(例:Qwen3-32B)を動かしたい方。
  • ML研究者:KVキャッシュ量子化や効率的な推論エンジン開発に取り組んでいる開発者。

主な特徴

  • 極端な圧縮:BF16と比較して、KVキャッシュメモリを約8倍削減。
  • パフォーマンス向上:大バッチサイズでは最大7倍のスループット向上、バッチサイズ1では最大3倍の高速化。
  • フレームワーク統合:SGLangおよびllama.cpp(特定のフォーク経由)に組み込み済み。
  • 高忠実度:ナイーブな2ビット量子化が通常失敗する推論ベンチマーク(GPQA、HumanEval)で、BF16に近い品質を維持。
  • 回転のコレクション:Hugging Faceに事前計算されたキャリブレーション済み回転を提供し、キャリブレーションフェーズをスキップ可能。

関連

  • Dispatch
  • プロジェクト
  • プロジェクト
  • プロジェクト
  • Dispatch