技術概要：RDKV - KV キャッシュの結合エビクションと量子化のためのレート歪みビット割当

1. 問題定義

大規模言語モデル（LLM）は、長文脈推論においてメモリ帯域幅と容量という重大なボトルネックに直面している。キー・バリュー（KV）キャッシュのサイズはシーケンス長に比例して増加するため、各デコードステップでオフチップのハイ・バンドウィッド・メモリ（HBM）からオンチップ・メモリへ数十ギガバイトのデータを再読み出しする必要がある。これにより、推論は計算ボトルネックではなくメモリボトルネックとなる。

既存の圧縮戦略は、一般的に以下の 2 つの孤立したカテゴリに分類される：

1. エビクション（Eviction）： トークンやチャネルにスコアを付け、最も低いランクのものを完全に破棄する（バイナリの保持か破棄か）。これにより、中程度の重要性を持つトークンに対しては精度の階層が残されない。
2. 量子化（Quantization）： すべてのトークンを保持しつつ、ビット幅を削減する。これにより、無視できるようなトークンさえも保持され、メモリが浪費される。

最近のハイブリッド手法はこれらを組み合わせようとするが、多くの場合、ヒューリスティックなルーティング・スコアやオフライン検索によって決定された事前に割り当てられた比率に依存しており、エビクションと量子化を統一的な最適化問題として扱えていない。

2. 手法：RDKV

本論文は、KV キャッシュ圧縮をレート歪み問題として定式化するフレームワーク**RDKV（Rate-Distortion KV cache compression）**を提案する。この観点では、エビクションと量子化は別々の動作ではなく、単一のビット割当曲線の 2 つの端点として捉えられる。

2.1 理論的定式化

RDKV は、アテンション計算に対する除去または量子化によって引き起こされる摂動に基づき、各キャッシュ単位（V キャッシュ内のトークン、K キャッシュ内のチャネル）に対して歪み重みを導出する。

• V キャッシュ（トークン）： トークンの歪み重み $w_t$ は、エビクション後のアテンション分布の全変動距離である。これは、すべてのクエリ位置においてそのトークンが受け取ったアテンション・スコアの総和として計算される：$w_t = \sum_{\tau} a_{\tau,t}$。
• K キャッシュ（チャネル）： チャネルの歪み重み $w_c$ は、そのチャネルをゼロにすることで生じるログイトの偏差のスペクトルノルムである。これは、対応するクエリとキーの列のノルムの積として計算される：$w_c = \frac{1}{\sqrt{d}} \|Q_{:,c}\|_2 \|K_{:,c}\|_2$。

これらの重みは、レート歪み目的関数におけるコスト係数として機能する：
$$\min \sum_{u} w_u \sigma_u 2^{-b_u} \quad \text{s.t.} \quad \sum_{u} b_u \leq B$$
ここで、$b_u$ は単位 $u$ に割り当てられたビット幅、$B$ は総ビット予算である。

2.2 ビット割当パイプライン

割当は、離散的なハードウェア対応ビット幅 $\mathcal{B} = \{0, 2, 4, 8, 16\}$ に適応させた**逆ウォーター・フィルリング（ラグランジュ緩和）**によって解かれる。

1. ステージ 1（重み付け）： プリフィル段階の直後に、完全な非圧縮キャッシュを用いてトークン重み ($w_t$) とチャネル重み ($w_c$) を一度計算する。
2. ステージ 2（トークン割当）： V キャッシュに対して多項選択ナップサック問題（MCKP）を解く。トークンには $\{0, 2, 4, 8, 16\}$ からビット幅が割り当てられる。ビット幅 0 はエビクションに対応する。
3. ステージ 3（チャネル割当）： K キャッシュに対して MCKP を解くが、ステージ 2 で保持されたトークンのみを対象とする。チャネルには、トークンの V ビット幅とは独立してビット幅が割り当てられ、単一のトークンが混合精度の K 成分と V 成分を持つことを可能にする。

2.3 効率的なパックド・デコード（TriZone）

実際のメモリおよびレイテンシの節約を実現するため、RDKV はデ量子化をアテンションカーネルに融合させるTriZoneキャッシュ配置を導入し、処理前にデータを FP16 にアンパックする必要性を排除する。

• ゾーン A： 保持された K 行と V 行を含み、ビット幅が $\{2, 4, 8\}$ のパックド・オールド・キャッシュ・ストリーム。これらはビット幅ごとにグループ化され、均一なデ量子化経路を可能にする。
• ゾーン B： 16 ビット精度が割り当てられたトークン用の FP16 V 行。
• ゾーン C： 新規生成されたトークン用の FP16 K 行および V 行。

アテンションカーネルは、これら 3 つのゾーンを組み合わせることでログイトと出力を計算し、行列乗算にデ量子化を直接融合させることで、追加のメモリ帯域幅オーバーヘッドを排除する。

3. 主要な貢献

1. レート歪み定式化： 本論文は KV キャッシュ圧縮を統一的なレート歪み問題として再構成し、最適解が本質的にゼロレート（エビクション）と有限レート（量子化）の単位を混合することを証明した。これにより、ヒューリスティックな比率やオフライン検索の必要性が排除された。
2. 逆ウォーター・フィルリングによる離散割当： 著者は、アテンション動力学から導出された歪み重みに基づき、$\{0, 2, 4, 8, 16\}$ から最適なビット幅を割り当てる離散ナップサック問題として割当を実装し、ラグランジュ緩和によってこれを解いた。
3. TriZone パックド配置： デ量子化をアテンションカーネルに融合させるハードウェア対応の実装により、理論上のビット幅削減を実際のピークメモリおよびデコード・レイテンシの削減に変換した。

4. 実験結果

実験は、LongBench、RULER、InfiniteBench において、LLaMA-3.1-8B-Instruct、Mistral-7B、Qwen3-4B、LLaMA-2-13B、およびQwen2.5-72B に対して実施された。

• 精度： RDKV は、評価された最良のベースライン（Ada-SnapKV）を平均で**9.1%上回った。LongBench においては、キャッシュ保持率が2.48%のみであるにもかかわらず、フルキャッシュ精度の97.81%**を回復した。
• 長文脈パフォーマンス：
  • RULER： 128K 文脈において、RDKV は 66.95% の精度を達成し、最強のベースライン（AdaKV）の 58.43% を 8.5 ポイント上回った。
  • InfiniteBench： RDKV は平均スコア 39.46 で圧縮手法の中で 1 位となり、AdaKV を 1.40 ポイント上回った。
  • Needle-in-a-Haystack： RDKV は深さ全体でほぼ均一な検索精度（平均 0.99）を維持するのに対し、エビクションのみのベースラインは中間深度で著しく失敗した。
• 効率性：
  • レイテンシ： 128K 文脈において、RDKV はフルキャッシュの FlashAttention-2 デコードと比較して4.5 倍のデコード高速化を達成した。
  • メモリ： 128K 文脈において1.9 倍のピークメモリ削減を達成し、FullKV がメモリ不足（OOM）となる単一の A100 64GB GPU 上で 256K 文脈の実行に成功した。

5. 意義と主張

本論文は、エビクションと量子化を単一の最適化フレームワーク内で統合的に扱うことが、段階的または比率ベースのアプローチよりも優れていると主張する。アテンション計算に誘起される歪みから直接割当スコアを導出することで、RDKV はエビクションのみの手法が陥る「バイナリの罠」（中程度のトークンを破棄する）や、量子化のみの手法が引き起こす「浪費」（無視できるトークンを保持する）を回避する。

著者は、自らの手法が以下の点で優れていることを強調している：

• モデルの再トレーニングを必要としない。
• アーキテクチャに依存せず、異なるモデルサイズやファミリー間で一貫した性能を発揮する。
• レイテンシと精度の間のパレート最適なトレードオフを提供し、長文脈において両指標でベースラインを厳密に上回る。
• ハードウェア向けに離散化されたレート歪み理論の「連続的最適解」が、自然に混合ビット戦略を生み出し、現在の最先端のハイブリッド手法を上回ることを実証している。

本論文は、この統合アプローチが効率性を犠牲にすることなく精度を向上させ、リソース制約のあるハードウェアでの長文脈推論を可能にすると結論づけている。