技術概要：OSCAR - 2 ビット KV キャッシュ量子化のためのオフラインスペクトル共分散感知回転

1. 問題定義

大規模言語モデル（LLM）における長文脈推論は、キー・バリュー（KV）キャッシュのメモリ占有量によって次第にボトルネック化しています。文脈長、バッチサイズ、モデルの深さが拡大するにつれ、キャッシュは GPU メモリを大量に消費し、スループットとスケーラビリティを制限します。KV キャッシュを超低精度（例：INT2）に圧縮することは 8 倍のメモリ削減をもたらしますが、重大な精度上の課題に直面します。すなわち、KV 活性化には深刻なチャネルごとの外れ値が存在することです。

低ビット量子化の下では、これらの外れ値が量子化スケールを支配し、大多数の「正常な」値を少数の有効な量子化レベルに押し込め、アテンションの品質を著しく劣化させます。既存の回転ベースの手法（例：アダマール変換）は、チャネル間でのエネルギー再分配によってこれらの外れ値を平滑化しようと試みます。しかし、これらの手法はデータ非依存であり、生キャッシュ分布を平滑化しますが、下流のアテンション機構が実際に消費する特定の方向とは整合しません。4 つの量子化レベルのみが存在する INT2 においては、この不一致により単純な回転手法は破綻し、推論タスクにおいてほぼゼロの精度をもたらします。

2. 手法：OSCAR

著者は、オフラインでアテンション感知の共分散構造を推定し、固定された回転とクリッピング閾値を導出する手法であるOSCAR（Offline Spectral Covariance-Aware Rotation）を提案します。従来の研究が生のユークリッド再構成誤差（$\|K - \hat{K}\|_F$）を最小化するのに対し、OSCAR は下流のアテンションスコアと出力における歪みを最小化します。

2.1 理論的動機

本論文は、最適な回転ターゲットはアテンション計算自体の統計から導出されるべきであると主張します。

• キー（$K$）の場合: アテンション対数（logits）の歪みは、単にキー共分散 $K^\top K$ ではなく、クエリ共分散 $Q^\top Q$ によって制御されます。
• バリュー（$V$）の場合: アテンション出力の歪みはアテンションスコア $S$ によって重み付けされ、ターゲット共分散は $V^\top S^\top S V$ となります。

2.2 アルゴリズム設計

OSCAR は 2 つのフェーズで動作します。

A. オフライン較正フェーズ:

1. 共分散推定: 少量の較正データセットを用いて、手法は経験的なターゲット共分散を計算します。
   • $C_Q = \frac{1}{N} \sum q_n^\top q_n$（クエリ感知キーターゲット）。
   • $C_S = \frac{1}{N} V^\top S^\top S V$（スコア感知バリューターゲット）。
2. 回転導出: 手法はこれらの行列に対して固有値分解を行い、最適な固有基底（$U_Q$ および $U_S$）を求めます。
3. 合成回転の構築: 最終的な回転行列（$R_K, R_V$）は、3 つの成分の積として構築されます。
   • $U_Q/U_S$: チャネルを重要度方向に整合させるアテンション感知固有基底。
   • $H_{Had}$: チャネル間での重要度指標を均等化し、外れ値のエネルギーを拡散させるアダマール変換。
   • $P_{br}$: 量子化グループ内のバランスの取れたダイナミックレンジを確保するため、高分散チャネルと低分散チャネルを交互に配置するビット反転置換。
   • 数式：$R_K = U_Q H_{Had} P_{br}$ および $R_V = U_S H_{Had} P_{br}$。
4. クリッピング閾値: 回転された較正活性化からレイヤーごとのクリッピング閾値を推定し、外れ値を制御します。

B. オンラインサービングフェーズ:

• 混合精度レイアウト: システムは、ページングされた KV キャッシュを維持し、「シンク」トークン（初期）と「最近」トークンを完全な BF16 精度で保持します。歴史的な中間文脈は、固定された OSCAR 回転とクリッピングを適用した後、INT2 で保存されます。
• カーネル統合: 著者は SGLang サービングフレームワーク用のカスタム融合 Triton カーネルを開発しました。このカーネルは、INT2 履歴のアンパッキング、量子化解除、逆回転、およびアテンション計算を処理し、ページングアテンションやプレフィックスキャッシングとの互換性を確保しつつ、追加のメモリトラフィックを導入しません。

3. 主要な貢献

1. 欠落したターゲットの特定: 本論文は、一般的な回転は外れ値を散らばらせますが、下流のアテンションと整合しないため INT2 では失敗することを特定しました。回転ターゲットは生キャッシュ再構成ではなく、アテンション統計（$Q^\top Q$ および $V^\top S^\top S V$）によって誘発される必要があります。
2. OSCAR フレームワーク: オフライン較正を用いてアテンション感知回転を導出する、完全に展開可能な INT2 KV キャッシュ量子化システム。理論的解析により、これらのターゲットが固定誤差代理仮定の下で最適であることが示されています。
3. 生産準備完了システム: ページング KV キャッシュレイアウトとプレフィックスキャッシングをサポートする、SGLang に統合されたカスタム INT2 アテンションカーネル。これは、高いスループットを維持しつつ、現代のサービングフレームワーク（SGLang、vLLM）とのシームレスな統合を実現します。

4. 実験結果

著者は、推論モデル（Qwen3-4B、Qwen3-8B、Qwen3-32B）と最先端モデル（GLM-4.7、358B パラメータ）を評価し、5 つのベンチマーク（AIME25、GPQA、HumanEval、LiveCodeBench、MATH500）および長文脈検索（RULER-NIAH、最大 128k）で検証を行いました。

• 精度維持:

  • Qwen3-4B-Thinking および Qwen3-8B において、OSCAR は BF16 に対する精度の差をそれぞれ3.78ポイントおよび1.42ポイントにまで縮小します。対照的に、単純な回転を用いた INT2 はほぼゼロの精度に崩壊します。
  • Qwen3-32B および GLM-4.7 において、OSCAR は BF16 と実質的に同等（0.02〜0.27 ポイント以内）を維持し、他の INT2 ベースラインは著しい劣化を被ります。
  • 長文脈検索（RULER-NIAH）において、OSCAR は 128k トークンまで堅牢な性能を維持しますが、単純な INT2 やアダマール変換のみの手法は、はるかに短い長さ（例：16k）で崩壊します。

• システム性能:

  • メモリ: BF16 と比較して KV キャッシュのメモリ占有量を約8 倍削減します。
  • スループット: 同じメモリ予算下で、大規模バッチサイズにおいて最大7 倍のスループット向上を実現します。
  • レイテンシ: メモリ帯域幅のオーバーヘッドが削減されるため、BF16 に対してバッチサイズ 1 のデコーディングを最大3 倍高速化します。
  • スケーラビリティ: 単一の H100 上で 100k トークン入力を伴う最大 28 の同時リクエストを処理する能力を実証しました。これは BF16 や INT4 ベースラインがメモリ不足に陥る領域です。

5. 意義と主張

本論文は、OSCAR が超低ビット（INT2）KV キャッシングの理論的な可能性と実用的な展開可能性の間のギャップを成功裡に埋めたと主張しています。回転ターゲットを生データ再構成からアテンション感知共分散へと転換することで、OSCAR は単純な量子化によって失われた精度を回復しつつ、2 ビット保存の巨大なメモリおよびスループット上の恩恵を保持します。

著者は、彼らのアプローチが単なるアルゴリズムの改善ではなく、システムレベルの解決策であると強調しています。ページングアテンションやプレフィックスキャッシングと統合するカスタムカーネルを提供することで、彼らは INT2 量子化が、現代の LLM サービングが依存する効率機能（プレフィックスの再利用など）を犠牲にすることなく、生産サービングフレームワーク（SGLang など）で使用可能であることを実証しました。この研究は、長文脈推論において、成功する量子化の鍵は、圧縮戦略をアテンション機構の特定の統計的ニーズに整合させることにあることを示唆しています。