🤖 machine learning

Efficient Serving for Dynamic Agent Workflows with Prediction-based KV-Cache Management

本論文は、LLM ワークフローにおけるエージェント呼び出しを予測に基づき動的に将来予測し、GPU メモリ内で高ポテンシャルの KV キャッシュエントリをインテリジェントに管理・保持する予測ベースのシステム PBKV を導入し、これにより動的および静的ワークフローの両方において既存のベースラインに対して大幅な高速化を達成するものである。

原著者： Haoyu Zheng, Fangcheng Fu, Jia Wu, Binhang Yuan, Yongqiang Zhang, Hao Wang, Yuanyuan Zhu, Xiao Yan, Jiawei Jiang

公開日 2026-05-08

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Haoyu Zheng, Fangcheng Fu, Jia Wu, Binhang Yuan, Yongqiang Zhang, Hao Wang, Yuanyuan Zhu, Xiao Yan, Jiawei Jiang

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

忙しい高級レストランのキッチンにいると想像してください。このキッチンでは、**大規模言語モデル（LLM）**がシェフ長であり、エージェントは複雑な注文を完了するために連携する専門の調理場（グリル、サラダバー、パティシエなど）です。

シェフが料理を作る際、彼らは「メンタルコンテキスト」（材料、レシピの手順、料理の現在の状態）を構築します。コンピュータ用語では、これをKV-Cacheと呼びます。これはシェフの準備ステーションのようなものです：次の手順で同じ刻んだ玉ねぎや同じソースが必要であれば、シェフはそれらをすべて再び刻んだり混ぜたりしたくありません。すでにそこにあるものを掴むだけです。これにより、莫大な時間が節約されます。

しかし、キッチンには限られたカウンタースペース（GPU メモリ）しかありません。すべての準備ステーションを永遠にセットアップしておくことはできません。新しい注文のスペースを作るために、何を捨てていいかを決めなければなりません。

問題：キッチンの「当てっこゲーム」

過去、キッチンでは**LRU（Least Recently Used：最も最近使用されていない）**という単純なルールが用いられていました。「あるステーションがしばらく触られていなければ、それを消去する」というものです。

欠点： 動的なレストランにおいて、あるステーションはシェフが顧客と話したり、配達を待ったりしている間、一時的にアイドル状態になるかもしれませんが、その後、すぐに再び必要になることがあります。LRU はそれが静かだったという理由だけでそれを捨ててしまい、シェフを最初からやり直す（「再プリフィル」）ことを強要します。これは遅く、コストがかかります。

もう一つの手法であるKVFlowは、キッチンマネージャーがすべてのステーションの正確な順序を事前に知っている（例：「グリル → サラダ → パティシエ」）と仮定していました。

欠点： 現実は厄介です。時にはサラダシェフがレタスが傷んでいることに気づき、冷蔵庫に戻らなければならない（「リトライループ」）ことがあります。時にはパティシエが顧客の反応に基づいて新しいデザートを作ることを決めることがあります。順序は静的ではなく、会話に応じて変化します。計画が変わると、KVFlow は混乱します。

解決策：PBKV（「予言者」キッチンマネージャー）

著者らは、PBKV（Prediction-Based KV-Cache Management：予測ベースの KV-Cache 管理）と呼ばれる新しいシステムを構築しました。「最後に忙しかったのは誰か」や「昨日の計画は何だったか」に基づいて推測するのではなく、PBKV は近い未来を覗き見ることができる超スマートなキッチンマネージャーのように機能します。

以下は、シンプルなアナロジーを用いた PBKV の仕組みです。

1. 水晶玉（予測器）

PBKV は「水晶玉」（機械学習モデル）を持っており、現在の注文と類似した注文の履歴を見て、*「次の 3 ステップでシェフが必要とするステーションは何か？」*を予測します。

仕組み： それは単に次のステップを推測するだけ（それは間違っている可能性があります）ではなく、次の3 つのステップを一度に推測します。「メニューの構造」（レストランの一般的な流れ）と「特定の注文の詳細」（顧客が実際に何を言ったか）を組み合わせます。
アナロジー： 顧客が「ハンバーガーが欲しいけど、もしかしたらチーズを多めにして」と言った場合、マネージャーは「まずグリル、次にチーズ・ステーション、もし考えが変わればパン・ステーション」と予測します。

2. 「引退」ステーションのルール（ライフサイクル認識型エヴィクション）

PBKV には黄金律があります：料理が完了したら、そのステーションを即座に消去する。

アナロジー： 客が支払いをして去った場合、そのテーブルを「最も最近使用されていない」ルールが適用されるまで待たず、今すぐ片付けて新しい客のスペースを作ります。PBKV は自動的に「完了した」ワークフローを特定し、他の何よりも先にそのメモリ領域を回収します。これは大きな勝利です。なぜなら、それらのステーションが再び必要になることは 100% ないからです。

3. 「スコアカード」システム（先読みスコアリング）

まだアクティブなステーションについては、PBKV は単に推測するのではなく、スコアを計算します。

アナロジー： すべての準備ステーションに、「今後数分以内にキッチン内のどのシェフがこれを必要とする可能性が高いか？」に基づいてスコアが付けられていると想像してください。
「チーズ・ステーション」がすぐに 5 つの異なる注文で必要になると予測されれば、そのスコアは高く、カウンターに残ります。
「スペシャルソース」がキャンセルされる可能性のある 1 つの注文でのみ必要であれば、そのスコアは低く、カウンタースペースを節約するために裏の冷蔵庫（ホストメモリ）へ移動されます。

4. 「安全」な事前注文（保守的プリフェッチング）

時には、マネージャーが近い将来あるステーションが必要になると思い、シェフが要求する前に材料を冷蔵庫からカウンターに出そうとします。

アナロジー： PBKV はここで非常に保守的です。材料をカウンターに出すのは、以下の条件が満たされる場合のみです：
1. すでに空いているカウンタースペースがあること。
2. カウンターから価値のあるものを追い出す必要がないこと。
3. 主要な配送トラック（PCIe バンド幅）を妨げないこと。
理由： もしマネージャーが間違えて間違った材料を出した場合、スペースを作るために正しい材料を捨てなければならないかもしれません。それは災難です。PBKV は、「確信が持てない場合は、価値のあるものでギャンブルはしない」と言います。

結果：より速いキッチン

著者らは PBKV を 3 つの異なる「キッチン」（ワークロード）でテストしました：

事実確認： エージェントが情報を検証するワークフロー。
コード作成： エージェントがコードを書き、修正するワークフロー（しばしばループ/リトライを含む）。
財務分析： 比較用の静的ワークフロー。

発見：

従来の「最後に使用された」ルール（LRU）と比較して、PBKV は動的タスクにおいてキッチンを1.85 倍高速化しました。
「準備ステーション」（キャッシュ）を LRU より2.55 倍多く「準備完了」の状態に保ちました。
静的タスクにおける以前の最良システム（KVFlow）と比較しても、PBKV は1.26 倍高速でした。

セーフティネット

この論文はまた、「水晶玉」（予測器）が間違いを犯してもシステムがクラッシュしないことを数学的に証明しています。システムは優雅に劣化します。

アナロジー： マネージャーの予測が少し外れても、システムは単にカウンターを管理する安全で標準的な方法にフォールバックします。マネージャーが予測を試みなかった場合よりも状況が悪化することはありません。

まとめ

PBKVは AI ワークフローのための賢いキッチンマネージャーです。一時的に静かだったからといって有用なツールを捨てたり、盲目的に推測したりすることをやめます。代わりに、先を見通し、完了した作業を即座に片付け、安全な場合のみ物事を移動させます。その結果、複雑で変化するタスクを立ち往生させることなく処理する、はるかに高速で効率的な AI システムが実現します。

技術概要：予測ベースの KV キャッシュ管理（PBKV）による動的エージェントワークフローの効率的なサービング

1. 問題定義

大規模言語モデル（LLM）に基づくマルチエージェントシステムは、複雑なタスクを、協働するエージェントのシーケンス（例：プランナー、コーダー、テスター）に分解します。これらのエージェントは膨大なコンテキストを共有するため、キー/バリュー（KV）キャッシュの再利用に大きな機会が存在します。しかし、既存の推論エンジン（例：vLLM、SGLang）は通常、時間的局所性を前提とした「最も最近使用されなかった（LRU）」ポリシーに依存しています。エージェントワークフローにおいて、キャッシュの再利用は時間的近接性ではなく、ワークフローの構造的性質によって駆動されます。キャッシュエントリは、再利用される前に、ツール呼び出しやループ反復中に複数のエージェント呼び出しにわたってアイドル状態のままになることがあり、これにより LRU がそれを早期に破棄してしまうのです。

KVFlow などの代替アプローチは、エビクションを「実行までのステップ数」に基づいて決定するために、静的で事前に定義されたエージェントステップグラフ（グローバルな DAG）を仮定します。この仮定は、呼び出されるエージェントの順序がランタイムのコンテキストに依存する動的ワークフロー（例：テスターが条件付きでリトライループをトリガーする、またはリトリーバーが新しいサブクエリを生成する）においては機能しません。

核心的な課題は二重です：

予測可能性: エージェントワークロードは LLM のデコーディングに起因して本質的に確率的であり、将来のエージェント呼び出しを予測することが困難です。
誤りのコスト: 予測誤差のコストは不均衡に高くなります。マルチエージェント環境でのキャッシュミスは、ワークフロー全体にわたって蓄積された数万トークンの再プリフィルを強いる可能性があります。

本論文は問いかけます：予測誤差に対して頑健でありながら、不完全な予測器から継続的に恩恵を受ける、動的マルチエージェントシステム向けの KV キャッシュ管理フレームワークをどのように設計できるか？

2. 手法：PBKV システム

著者は、将来のエージェント呼び出しを予測して既存の KV キャッシュエントリの再利用価値を評価するシステムである**PBKV（予測ベースの KV キャッシュ管理）**を提案します。PBKV は、GPU メモリとホストメモリ（HiCache を通じて）の 2 階層ストレージ階層上で動作し、以下の 3 つの主要コンポーネントで構成されます：

A. ワークフロー予測器

静的なグラフを仮定する代わりに、PBKV は補完的なシグナルを融合する予測器を採用し、各アクティブなワークフローに対して次の $K$ 個のエージェント呼び出しを予測します。

アーキテクチャ: GraphSAGEに基づき、予測器は帰納的であり、ランタイムにおいて未見のワークフロー接頭辞にも一般化できます。
シグナル融合:
1. トポロジー認識エージェント埋め込み（ $h_{cur}$ ）: オフライントレースから導出されたグローバル呼び出しグラフ $G$ における、現在のエージェントのアイデンティティとその近傍をエンコードします。
2. アテンションベースの履歴集約（ $h_{path}$ ）: スケールドドット積アテンションを用いてワークフロー接頭辞を要約し、異なる経路をたどって同じエージェントに到達するワークフローを区別します。
3. セマンティックシグナル（ $h_{txt}$ ）: 最後のプリフィルトークンのポストノーム隠れ状態を再利用してリクエスト固有の意図を捉え、同一の接頭辞を持つが目標が異なるリクエストを区別します。
出力: 予測器は、自己回帰的な誤差蓄積を回避するために、単一のフォワードパスで将来のエージェントに対する $K$ 個の確率分布を出力します。

B. 先読み KV キャッシュエビクション

PBKV は、単純な時間的ヒューリスティックに代わり、「先読み再利用スコア」によって駆動される階層的エビクションポリシーを採用します。

ライフサイクル認識リクレイメーション: 終了したワークフローに関連するキャッシュノード（「退役キャッシュ」）は決定論的に特定され、再利用ポテンシャルがゼロであるため最初にエビクトされます。これは予測の品質に依存しないベースラインの改善を提供します。
スコア駆動アクティブエビクション: アクティブなキャッシュについては、予測された再利用価値に基づいて連続的なスコアが計算されます：
$\text{Score}(c) = \sum_{k=1}^{K} \gamma^{k-1} \sum_{w \in \mathcal{W}_{act}(c)} s_w^{(k)} \cdot A_w(c) \cdot P_w^{(k)}$
ここで：
- $A_w(c)$ はアクセスインジケーター（ワークフロー $w$ のどのエージェントがノード $c$ にアクセスしたか）です。
- $P_w^{(k)}$ は、ステップ $k$ で $c$ にアクセスするエージェントを呼び出す予測確率です。
- $s_w^{(k)}$ は累積生存確率（ワークフローが終了していないこと）です。
- $\gamma < 1$ は、信頼性の低い長期的予測を割り引くための信頼性減衰因子です。
階層戦略: システムは、アクティブなキャッシュに触れる前に、スコアが 0 である退役キャッシュのエビクションを優先します。これにより、優雅な劣化が保証されます。予測器が不十分であっても、システムはライフサイクル認識から恩恵を受け、正確であればマルチステップの先読みを活用します。

C. 保守的プリフェッチ

転送レイテンシを隠蔽するために、PBKV は再利用される可能性が高いキャッシュをホストメモリから GPU に能動的に読み込みます。

保守的原則: プリフェッチは、それ以外ではアイドル状態の GPU スペース（空き領域＋退役キャッシュ）とアイドル状態の PCIe バンド幅（純粋なデコードバッチ中）に制限されます。
リスク回避: PBKV は、プリフェッチされたデータのためのスペースを作るためにアクティブなキャッシュをエビクトすることはありません。これにより、予測誤差がパフォーマンスをベースライン以下に低下させることなく、推測的な利益のために既知の価値あるキャッシュが置き換えられることを防ぎます。
ランキング: 候補は、厳格な予算下での貪欲ヒューリスティックを用いて、1 ステップ再利用価値（式 1）に基づいてランク付けされます。

3. 主要な貢献

形式化: 本論文は、マルチエージェントサービングにおける KV キャッシュ管理を、将来のアクセスパターンの予測という問題として形式化し、この設定における予測器の設計指針（マルチシグナル融合、マルチステップ先読み）を提供します。
PBKV フレームワーク: 著者は、階層的エビクション（決定論的ライフサイクルリクレイメーション＋確率的先読みスコアリング）と保守的プリフェッチを組み合わせたシステムを提案します。
理論的保証: 本論文は、PBKV のパフォーマンスが予測誤差に対してリプシッツ連続的に劣化することを証明しています。具体的には、エビクションコストの後悔は予測誤差によって線形的に有界であり、不完全な予測器であっても優雅な劣化が保証されます。
実装と評価: PBKV は HiCache を備えた SGLang 上で実装され、多様なワークロードと LLM において評価されました。

4. 実験結果

実験は、3 つのベンチマークで行われました：HoVer（事実検証、動的ループ）、SWE-bench（コード生成、リトライループ）、およびFinanceBench（静的ワークフロー）。使用されたモデルには Qwen3-14B と Qwen3-32B が含まれます。

動的ワークフロー:
- PBKV は、LRU と比較して、エンドツーエンドのレイテンシで最大1.85 倍の高速化、KV キャッシュヒット率で2.55 倍の改善を達成しました。
- 「ライフサイクル認識エビクション（PBKV-LAE）」単独でも、LRU に対してヒット率を 1.66 倍改善し、単純なライフサイクルシグナルの価値を実証しました。
- 完全なシステム（PBKV-HE ＋プリフェッチ）は、LRU が約 10% に低下する重いメモリ圧力下でも、高いヒット率（約 50-60%）を維持しました。
静的ワークフロー:
- 静的な FinanceBench ワークロードにおいて、PBKV は SOTA ベースラインであるKVFlowを、レイテンシで1.26 倍、ヒット率で1.39 倍上回りました。著者はこれを、PBKV のライフサイクル認識エビクション（アクティブなキャッシュの早期エビクションを防止）と、再利用スコアのワークフロー間集約に起因すると述べています。
頑健性:
- 感度分析により、より弱い予測器のバックボーン（例：R-GCN、マルコフモデル）であっても、PBKV は一貫して LRU と予測フリーのバリアント（PBKV-LAE）を上回ることが示されました。
- 保守的プリフェッチは安定した利益を提供しましたが、アクティブなキャッシュをエビクトする攻撃的なプリフェッチは、中程度の予測ノイズ下では逆効果となりました。

5. 意義と主張

本論文は、PBKV が、既存の静的仮定（KVFlow）または時間的ヒューリスティック（LRU）が機能しない動的エージェントワークフローのサービングにおける重要なギャップを埋めていると主張します。

最適性よりも頑健性: 著者は、PBKV がエージェントシステムの本質的な不確実性に対して頑健に設計されていることを強調しています。決定論的なガードレール（ライフサイクルリクレイメーション）と保守的な推測（アイドルリソースでのみプリフェッチ）を組み合わせることで、不完全な予測がパフォーマンスを害さないことを保証しています。
優雅な劣化: リプシッツ連続性の理論的証明は、PBKV のパフォーマンスが安定していることを確立しています。予測器の精度が向上すればパフォーマンスは線形的に向上しますが、精度が低くても壊滅的な失敗を引き起こすことはありません。
一般化可能性: 予測器はプラグ可能なモジュールとして設計されています。システムの設計原則（階層的エビクション、保守的プリフェッチ）は、使用される特定の予測器に依存しないため、このフレームワークは将来のワークフロー予測の進歩から恩恵を受けることができます。

著者は、PBKV が動的および静的ワークロードの両方において、LRU および最先端の静的認識ベースラインを大幅に上回ることを結論付け、マルチエージェントサービングにおける予測ベースの保守的キャッシュ管理の有効性を検証しました。