技術サマリー：ハイブリッドおよび再帰型 LLM サービングのためのスパース接頭辞キャッシング

1. 問題定義

既存の大規模言語モデル（LLM）サービングシステムは、共有接頭辞内のすべてのトークンのキー/バリュー（KV）状態を保存して重複したプリフィル計算を回避する「密な接頭辞キャッシング」に大きく依存しています。このアプローチはトランスフォーマーアーキテクチャでは標準的です。しかし、状態空間モデル（SSM）やアテンションと再帰層を組み合わせたハイブリッドアーキテクチャは構造的な非対称性を呈します。再帰層は $h_t = F(h_{t-1}, x_t)$ として状態を進化させるため、中間状態の全履歴を必要とするのではなく、単一の保存状態 $h_c$ から計算を再開できます。

再帰状態のトークンごとの密なキャッシングは、しばしば不要であり、かつ prohibitively 高価です。Qwen-3.5 のようなモデルにおいて、ヘッドあたりの再帰状態は $d_{head}^2$ に比例してスケーリングするのに対し、アテンション KV 保存は線形にスケーリングします。その結果、すべてのトークンの再帰状態を保存することは、アテンション KV キャッシュよりもはるかに多くのメモリを消費しますが、スパースなチェックポイントのセットを保存するよりも追加の利点を提供しません。

解決される核心的な問題は以下の通りです：「長さ $N$ の接頭辞に対して $M$ 個のチェックポイントスロットというメモリ制約が与えられたとき、将来の要求のオーバーラップ深度の特定の分布の下で、期待される再計算コストを最小化するために、チェックポイントをどこに配置すべきか？」この最適化は、キャッシュの受け入れ/淘汰ポリシーとは区別され、単一の保持エントリ内でどのトークンをチェックポイント化するかを焦点とし、部分的な接頭辞のオーバーラップに対する有用性を最大化することを目指します。

2. 手法

理論的定式化

著者らは、この問題を直線上の「片側重み付き k-メジアン問題」として定式化しました。

• $C = \{c_1, \dots, c_M\}$ をチェックポイント位置の集合とします。
• 深度 $t$ まで接頭辞とオーバーラップする要求に対して、再利用可能な最も深いチェックポイントは $\ell(t; C) = \max(\{0\} \cup \{c \in C : c \le t\})$ です。
• コストは再計算するトークンの数であり、$r(t; C) = t - \ell(t; C)$ です。
• 目的は、確率 $p_t$ を持つオーバーラップ深度分布 $T$ が与えられたとき、期待再計算コスト $E[r(T; C)]$ を最小化することです。

アルゴリズム的アプローチ

1. 最適バランス配置（一様/最悪ケース）： 著者らは、一様なオーバーラップ分布または最悪ケースのシナリオにおいて、バランスの取れた間隔（均等に分散されたチェックポイント）が最適であることを証明しました。
2. 分布認識型動的計画法（DP）： 非一様で現実的なオーバーラップ分布に対して、著者らは正確な $O(NM)$ 動的計画法アルゴリズムを導出しました。
   • 再帰関係は以下の通りです：$dp[m, j] = \min_{1 \le s \le j} (dp[m-1, s-1] + w(s, j))$。ここで、$w(s, j)$ は深度 $s \dots j$ に対して $s$ を最も深いチェックポイントとして使用した場合のコストです。
   • これは、接頭辞和の単調性を利用した単調凸包トリックを用いて効率的に解かれます。
3. 分布ドリフトへの対応： 真のオーバーラップ分布は未知であり、過去のデータから推定されなければならないため、本論文は（補題 1 および定理 3 により）値関数がオーバーラップ分布に関してリプシッツ連続であることを理論的に保証しています。これにより、経験ヒストグラム（またはドリフトする分布に対する指数重み付きヒストグラム）を使用しても、有界な非最適性を持つ準最適スケジュールが得られることが保証されます。

実装戦略

• チェックポイント化： システムはスパースな位置で正確な再帰状態を保存します。キャッシュヒットが発生すると、システムは保存された最も深いチェックポイントから再開し、接尾辞を正確に再計算します。
• 完全性： 再帰更新が決定論的であるため、チェックポイントから再開して接尾辞を再生成することは、完全なプリフィルとビット単位で同一の出力を生み出します。
• ハイブリッド統合： ハイブリッドモデルの場合、この手法は完全なアテンション KV キャッシュ（密）を保存しますが、再帰層に対してのみスパースなチェックポイント化を適用します。

3. 主要な貢献

1. 理論的最適性： 一様オーバーラップおよび最悪ケースの条件下で、バランスの取れた間隔が最適であることを証明しました（定理 1）。
2. 正確なアルゴリズム： スパース接頭辞キャッシングを重み付き k-メジアン問題として定式化し、正確な $O(NM)$ 動的計画法ソリューションを提供しました（定理 2）。
3. 堅牢性の保証： 最適スケジュールが分布の誤指定に対して安定していることを実証しました。非最適性は、経験分布と真のオーバーラップ分布との間の総変動距離によって有界であり（定理 3）、指数再重み付けを用いて分布ドリフトのシナリオにも拡張されます（定理 4）。
4. 実証的検証： 現実のデータセット（QuALITY、NarrativeQA、システムプロンプト）で手法を検証し、固定予算のベースラインに対して一貫した改善を示しました。

4. 実験結果

データセットと設定

• データセット：
  • QuALITY: 単一ドキュメントに複数の質問が含まれる（全長に近い高いオーバーラップ）。
  • NarrativeQA: 広範なマルチモーダルオーバーラップを伴う類似構造。
  • システムプロンプト: ユーザークエリと組み合わせた実際のシステムプロンプト（短い長さでオーバーラップが集中）。
• ベースライン： 「キャッシュなし」、「KV のみ」、「バランス型」（均等間隔）、「ブロック」（固定ブロックサイズ）、「対数型」、および「$\sqrt{L}$ 間隔」戦略と比較しました。
• ハードウェア： NVIDIA RTX 2080 Super 上で Qwen-3.5-0.8B モデルを用いたプロトタイプ実行。

知見

• トークン節約： DP 最適スケジュールは、再帰作業削減率の観点から、決定論的ベースライン（バランス型、ブロック型、対数型）を常に凌駕します。
• 低予算の利点： パフォーマンスの差は、オーバーラップ分布が最も非一様となる低チェックポイント予算で最大となります。例えば、システムプロンプト（オーバーラップが短い場合）では、DP 手法は発散のピーク付近にチェックポイントを密に配置し、より少ないチェックポイントでブロックキャッシングと同じ実行時間を達成します。
• ウォールクロック時間： QuALITY およびシステムプロンプトにおいて、DP 導出スケジュールは代表的な層グループで測定可能なウォールクロックの高速化をもたらします。実行時間の削減は、節約されたトークン数と線形に相関します。
• オーバーヘッド： CPU から GPU へのチェックポイント読み込みのオーバーヘッドは、削減された再計算から得られる節約に比べて小さく、結果としてネットの実行時間改善となります。

5. 意義と主張

本論文は、スパース接頭辞キャッシングが、以前は利用不可能だったハイブリッドおよび再帰型 LLM サービングのための新たな設計点を解放すると主張しています。

• メモリ/計算のトレードオフ： 状態を正確に抽出・復元できるアーキテクチャにおいて、密な KV キャッシングでは提供できないメモリ使用量と再計算の間のトレードオフを提供します。
• 完全性： 近似 KV 圧縮技術とは異なり、この手法は再帰計算カーネルを変更することなく正確な出力を維持します。
• 適用性： これは、要求が実質的だが同一ではない接頭辞を共有するワークロード（異なる質問を伴う RAG、変化するシステムプロンプトなど）で最も有益です。厳密な追加専用チャットワークロード（最後の状態で十分である場合）には追加の利点は提供しませんが、そのような場合でもオーバーヘッドは発生しません。
• 将来の作業： 著者らは、生産ランタイムにおける効率的な状態抽出/復元 API の必要性、およびこの配置戦略と受け入れ/淘汰ポリシー（Marconi らなど）との統合という工学的課題を認識しています。また、現在の単一接頭辞の理論を一般的な接頭辞木に拡張できる可能性にも言及しています。