大規模言語モデル(LLM)の推論において、プレフィックスキャッシングは重要なレイテンシ最適化技術です。しかし、既存のシステムは各トークンのキー・バリュー(KV)キャッシュが高密度に再利用されることを前提としています。状態空間モデル(SSM)やリカレント層は問題の構造を変えます。これらのモデルは、単一の保存された状態から再開できるため、全トークン履歴を必要としません。この非対称性により、無再利用と高密度キャッシュの間の新しい設計点が開かれます。つまり、スパースなチェックポイント位置に正確なリカレント状態を保存し、キャッシュヒット時に最も深い保存チェックポイントから再開し、残りのサフィックスを正確に再計算します。
本論文では、スパースプレフィックスキャッシングを、オーバーラップ深度の分布下でのチェックポイント配置として形式化し、正確なO(NM)動的計画法を提案します。複数のリクエストが有意なプレフィックスを共有する(例えば、単一の長い文書について異なる質問をする)ユースケースにおいて、本手法は実世界データ上の標準ヒューリスティックによってトレースされるパレートフロンティアを一貫して改善することを示します。QuALITYおよびSystem Promptsにおいて、分布認識型配置は測定されたレイヤーグループパレートフロンティア上のすべての固定予算ベースラインを凌駕し、最も強力なヒューリスティック(ブロックキャッシング)と一致またはそれを上回りながら、通常はかなり少ないチェックポイントを使用します。最大の利得は、オーバーラップ分布が最も不均一な低チェックポイント予算で得られます。
本手法は、多くのリクエストが保持されたキャッシュエントリ内で実質的だが同一ではないプレフィックスを共有する場合に最も適しています。正確な出力を保持し、リカレント計算自体を変更せず、新しいリカレント更新カーネルを必要としません。隠れ状態を正確に抽出および復元できるリカレント/SSM層に適用可能であり、ハイブリッドモデルでは既存のKVキャッシュ圧縮技術と組み合わせることができます。
さらに、この手法の実装は比較的簡単です。なぜなら、基盤となるリカレント計算を変更する必要がなく、チェックポイント位置での隠れ状態の保存と復元のみが必要だからです。これにより、既存の推論システムへの統合が容易になります。論文の実験では、少ないチェックポイント予算でも、キャッシュミス時の計算量を大幅に削減し、全体的な推論レイテンシを低減できることが示されています。これは、共通のプレフィックスを持つ多数のリクエストを処理するサービス(マルチターン対話、文書QAなど)にとって特に価値があります。