InferCeptとは?AIエージェントのツール呼び出し中にKVキャッシュを最適管理する推論基盤

InferCeptは、AIエージェントが検索や外部APIを呼び出して推論が中断する際、KVキャッシュの保持・CPU退避・再計算を選び分けるLLMサービング技術です。仕組み、実験結果、エージェント開発への応用を解説します。

参考文献

InferCept: Efficient Intercept Support for Augmented Large Language Model Inference

Reyna Abhyankar, Zijian He, Vikranth Srivatsa, Hao Zhang, Yiying Zhang

論文を見る

今回の論文

今回取り上げるのは、Reyna Abhyankarらによる「InferCept: Efficient Intercept Support for Augmented Large Language Model Inference」です。2024年に国際会議ICML 2024で発表され、Proceedings of Machine Learning Research(PMLR 235、pp. 81–95)で公開されました。URLは https://proceedings.mlr.press/v235/abhyankar24a.html 、研究分野はLLM推論基盤・AIエージェントのサービングです。

選んだ理由は、ツールを使うAIエージェントを本番運用すると必ずぶつかる「外部API待ち」の問題を、モデルの賢さではなく推論基盤から解いているためです。検索、計算、画像生成、人との対話の最中に、過去文脈を捨てるべきか、GPUに残すべきか、退避すべきかを定量的に判断する考え方は、エージェント開発のコストと応答性に直結します。

どんな技術か

InferCeptは、LLMが外部ツールや人の応答を待つために推論が中断される場面を効率よく扱うサービングフレームワークです。例えばエージェントが検索APIを呼び、結果を受け取ってから回答を続けるとします。この間、過去の会話・ツール呼び出しを再び読むためのKVキャッシュをどう扱うかが問題になります。

単純に捨てればGPUメモリは空きますが、再開時に長い文脈を再計算しなければなりません。GPUに残せば再計算は不要ですが、待機中の要求がメモリを占有します。InferCeptは、この二択にCPUメモリへの退避も加え、待機時間、文脈長、GPUの混雑度を基に、要求ごとに「保持・退避・分割再計算」を選びます。論文ではこの方針をmin-waste interceptionと呼びます。

課題

従来のLLM推論サーバーは、連続してテキストを生成する単独の要求を主な前提にしています。しかしAIエージェントでは、生成の途中でツール呼び出しが発生し、数ミリ秒の計算APIから人間の入力待ちまで、停止時間が大きく変わります。

  • 既存の「中断したらKVキャッシュを破棄する」方式では、再開のたびにそれまでの文脈をprefillし直します。論文の測定では、この不要な再計算がモデルforward時間の37〜40%を占めました。
  • 逆にKVキャッシュをGPUに保持し続けると、数秒〜数十秒待つ画像生成やチャットで、実行していない要求がGPUメモリを占有します。他の利用者の要求を同時に処理できなくなります。
  • CPUへswapすればGPUメモリは回収できますが、GPU-CPU間の転送が前景の推論を止めると、今度は帯域がボトルネックになります。論文の混合ワークロードでは、素朴なSwapはGPU資源の26%を無駄にし、全処理時間の25%以上がswap待ちでした。

これはRAG、ブラウザ操作、コード実行、社内API連携のあるAIシステムで現実に起こります。モデルの出力品質を変えなくても、ツールを一度呼ぶたびに長い履歴を計算し直せば、TTFT(最初のトークンまでの時間)、同時接続数、GPUコストが悪化します。

用語解説

KVキャッシュ
Transformerが過去トークンのattention計算で作ったKeyとValueを保存するメモリです。再利用すれば過去文脈を読み直さずに次トークンを生成できます。本論文は、中断中のKVキャッシュをどこに置くかを最適化します。
interception(推論中断)
LLMのデコードが外部ツール、人、環境からの結果待ちで一時停止することです。単なるリクエスト完了ではなく、同じ文脈で再開するため、通常のLLMサーバーの前提を崩します。
prefillとdecode
prefillは入力文脈全体を処理してKVキャッシュを作る段階、decodeは1トークンずつ生成する段階です。キャッシュを捨てた要求の再開では長いprefillが再発し、InferCeptが削減したい負担になります。
swap
GPUメモリ上のKVキャッシュをCPUメモリへ移し、必要になれば戻す操作です。容量は節約できますが、PCIeなどの転送帯域と待ち時間を消費するため、無制限に実行すると逆効果になります。
GPU saturation point
GPUの演算器を十分に使い切れる1イテレーション当たりのクエリトークン数です。InferCeptはこの上限まで、通常のdecodeと再計算を混ぜることで、長い再計算が一度に推論を塞ぐことを避けます。

技術の仕組み

InferCeptはvLLM上に実装され、通常のPagedAttentionによるメモリ管理を利用しながら、中断要求用のスケジューラ、資源浪費の推定、分割処理を追加します。重要なのは、どの方式も常に正解ではないと捉えている点です。

1. 三つの選択肢を「失われるGPU資源」で比較する

中断した要求には三つの選択肢があります。DiscardはKVキャッシュを破棄して再開時に文脈を再計算します。PreserveはGPUに残します。SwapはCPUへ退避します。InferCeptは、文脈トークン数、1トークンのKVメモリ量、forward時間、推定中断時間、同時実行中の他要求のメモリ量から、それぞれが消費するGPU時間×メモリ量を見積もります。

例えば短い検索API待ちなら、再計算するよりPreserveが有利です。一方で人の返信待ちのように長い中断は、GPUに残すコストが大きくなるためSwapまたはDiscardが候補になります。ただし実行時間が未知・変動する場合は、呼び出し開始からの経過時間を推定中断時間として毎イテレーション更新します。混合ワークロードでは、この推定でも中断時間を正確に知るoracleの93%の性能に達しました。

2. swapを推論の裏でパイプライン化する

素朴なSwapは、KVキャッシュ全体を一度にコピーするため、通常の推論を待たせます。InferCeptはモデルの各層のswapをパイプライン段として扱い、カーネル起動、データ転送、通常のforwardを重ねます。

さらに、swap-inとswap-outを複数イテレーションへ分割します。オフライン計測でバッチサイズごとのforward時間を取得し、その時間に隠せる転送トークン数をswap limitとして計算します。1回で運ぶ量をこの上限に抑えることで、転送のために前景のデコードを止めない設計です。各イテレーションでは、CPU空き容量、GPU空き容量、swap-in/out予算を同時に満たす量を選びます。

3. 再計算も一括で行わず、decodeと混ぜる

Discardを選んだ要求は、過去文脈を再計算する必要があります。しかし長い文脈を一度にprefillすると、そのイテレーションが重くなり、他の要求の生成まで遅れます。そこでInferCeptは再計算をチャンクに分けます。

GPU saturation pointをS、現在decode中のトークン数をRとすると、概念的には最大S - Rトークンぶんの再計算を混ぜます。decodeは少数のクエリ、再計算は多くのクエリを使うため、両者を混ぜるとメモリ上限内でGPU演算器を使いやすくなります。論文の分析では、分割再計算は再計算自身によるメモリ浪費を一括再計算の半分にできます。

4. 要求横断で優先順位を付ける

実運用では中断要求が複数あります。InferCeptは各要求についてPreserveと分割Discardのうち小さい浪費量を求め、その値が大きい順に、swap予算が尽きるまでCPUへ退避します。残りはPreserveまたはDiscardに割り当てます。

再開後の要求には、実行中、swapから戻す待ち、通常待ちの三つのキューを持ちます。通常待ちでは元の到着時刻に従うFCFSで公平性を保ち、swap-inキューには別の帯域予算を与えます。この「中断した一要求だけ」の局所最適化で終えず、全要求のGPU資源と帯域を配分する点がInferCeptの核です。

実験と結果

論文は、算術ツール(GSM8K-XL)、知識QA(MultiHopQA Wikipedia)、仮想環境(ALFWorld)、チャット(ShareGPT)、画像生成(Stable Diffusion呼び出し)、TTS(Bark呼び出し)という六種類の中断を混ぜたワークロードで評価しました。短時間の自動APIと、人や別モデルを待つ長時間APIを混ぜている点が実運用に近い設定です。

比較条件と指標

GPT-J 6BをA100 1枚、Vicuna 13BをA100 1枚と2枚のtensor parallel、Llama 3 70BをA100 4枚で動かしました。比較対象は、無変更のvLLM(Discard)、到着順を改善したImprovedDiscard、常時Preserve、素朴なSwapです。中断時間そのものは全方式で同じため、エンドツーエンド遅延から除外し、出力トークンで正規化した遅延、完了要求数/秒、TTFTを測定しました。

混合ワークロードでの改善

6Bモデルでは、同じ正規化遅延でvLLMの最大1.6倍の到着率を処理し、同じ要求率での正規化遅延は1.9〜5.7倍低下しました。最小TTFTを保ったまま処理できる到着率も1.7倍でした。分解評価では、InferCeptはベースラインより再計算由来のGPU浪費を60%以上、swap由来を96%削減しています。

13Bを2GPUで動かす設定では、vLLMより最大1.8倍高い到着率を処理し、正規化遅延は1.6〜10倍改善しました。Llama 3 70B・4GPUでは、同じ要求率で1.3〜12倍低い正規化遅延、同じTTFTで2.4倍高い負荷を達成しています。モデル、GPU、負荷で値は変わりますが、中断が多いほどキャッシュの扱いが全体の性能を左右することを示します。

単一APIでも効果を確認

QAのみ、チャットのみのワークロードでも確認し、正規化遅延でvLLMを最大2.3倍(QA)、1.9倍(チャット)上回りました。QAは短いAPI呼び出しが多いため、PreserveがDiscardより有利な場面が多く、改善が大きくなりました。つまり「ツール呼び出しがあるから必ずCPU退避」ではなく、停止時間の性質を利用する必要があります。

何に使える?

InferCeptの本質は、ツール利用そのものではなく、再開する生成の文脈を失わずに、待機資源を減らすことです。以下のようなシステムで特に応用しやすいでしょう。

検索・社内APIを多段で呼ぶRAGエージェント

商品在庫、顧客情報、規程、Web検索を順に呼ぶエージェントでは、ツール結果のたびに会話履歴と取得済み根拠を持ったまま再開します。短いAPIならキャッシュを保持し、遅いAPIや人の承認待ちでは退避する方針にすると、長いエージェント履歴の再prefillを減らせます。

人の承認を挟む業務ワークフロー

見積もり作成、申請チェック、サポート返信案では、LLMが下書きを出し、人が修正・承認して再開する流れがあります。人待ち中にGPUキャッシュを固定すると同時利用者を圧迫するため、経過時間に応じて退避・破棄を選ぶ設計が有効です。再開が遅い要求だけをCPUに退避できれば、対話の継続性とサーバーの混雑緩和を両立できます。

画像・音声・コード実行を併用するマルチモーダルエージェント

画像生成、音声合成、サンドボックスでのコード実行は、LLMの次トークン生成より長く不規則な待ち時間を作ります。InferCeptは画像とTTSを評価に含めており、別モデルを呼ぶ間のKVキャッシュ運用を考える土台になります。ただし実装には、ツール呼び出しを確実に中断イベントとしてサーバーへ渡す連携が必要です。

開発や事業へのヒント

この論文から得られる実務上のヒントは、「エージェントの遅さ」をLLMの生成速度だけで測らないことです。外部待ちと文脈の再利用を観測すると、モデルを替えずに改善できる余地が見つかります。

中断をリクエスト終了として扱わない

ツール呼び出し後の再開を新規チャット要求として実装すると、簡単ですが長い履歴ほど再計算が増えます。ツールコールID、再開トークン、KVキャッシュの所有権を推論基盤で追跡し、「一時停止・再開」という状態遷移を持たせると、将来の最適化余地を残せます。

ツール別に待ち時間の分布を計測する

ツールを「速い・遅い」と固定的に分類するだけでは不十分です。論文のQAのようにネットワーク待ちは変動します。API種別、文脈長、中断時間、再開率、GPUメモリ量をログに取り、p50だけでなくp95/p99を見てください。短時間で再開する要求は保持、長く不確実な要求は退避というルールに根拠を持たせられます。

小規模でも「再計算トークン」をKPIにする

大規模な独自サーバーがなくても、ツール呼び出しの後に何トークンを再prefillしたか、要求当たりのKVキャッシュ滞留時間、TTFTを計測できます。再計算トークンが多い経路を減らすだけでも、LLM APIの入力課金や自前GPUの待ち行列を抑えられます。まずはキャッシュのTTL、会話の要約、ツール結果の圧縮と組み合わせて可視化するのが現実的です。

限界

InferCeptはvLLMを基盤にした研究実装であり、すべての推論エンジンやクラウドAPIでそのまま使えるわけではありません。KVキャッシュの退避・復元、GPUとCPUの転送、ツール中断のフックを制御できるインフラが必要です。

また、最適判断は中断時間、文脈長、GPU帯域、同時要求数の推定に依存します。予想外に長い外部API待ち、CPUメモリ不足、急激な負荷変動では、プロファイルに基づく予算が外れる可能性があります。論文の経過時間ベース推定はoracleの93%でしたが、残りの差は実運用でのSLOに影響し得ます。

評価は六種類の合成・推定を含むワークロードと、2024年時点のモデル・A100環境です。現代の推論エンジン、NVLink構成、prefix cache、disaggregated serving、ツールの並列実行との相互作用は個別検証が必要です。さらに、キャッシュをCPUへ移すことは性能の問題であって、テナント分離、暗号化、保持期限などのデータ保護要件を満たす仕組みを別途設計しなければなりません。

よくある質問

Q. InferCeptはモデルの回答精度を上げますか?

A. 直接は上げません。同じモデル・同じ文脈で再開する際のGPUメモリ、再計算、転送を最適化する推論基盤です。得られる主な効果は、混雑時の遅延、TTFT、処理可能な同時要求数の改善です。

Q. 外部ツールが短時間なら、KVキャッシュは常にGPUに残すべきですか?

A. 多くの場合は有利ですが、常にではありません。短い待機でも、文脈が非常に長い、同時に多数の要求が止まる、GPUメモリが逼迫する場合は他要求への影響が大きくなります。InferCeptはこの保持コストと再計算コストを要求単位で比較します。

Q. CPUへのswapはいつ有効ですか?

A. GPUに保持するには長すぎるが、捨てて再計算するには文脈が長い要求で候補になります。ただし転送を一括で行うと推論を止めるため、論文のようにforwardと重ね、イテレーションごとの帯域予算を設けることが重要です。

Q. RAGのprefix cacheとInferCeptは競合しますか?

A. 目的が異なります。prefix cacheは複数要求に共通するプロンプト先頭の計算を再利用し、InferCeptは一つの中断された要求を再開するためのKVキャッシュを管理します。実装上は同じGPUメモリを使うため、容量配分と退避方針を統合して評価する必要があります。

Q. 小さなAIアプリでも導入価値はありますか?

A. ツール呼び出しが少なく、短い会話だけなら専用のサービング改修は過剰かもしれません。一方、長い会話履歴、複数の社内API、人の承認待ちがあり、同時利用者が増えているなら、まず再prefill量と待機中KVキャッシュを測ってください。そこが大きければ、同じ発想をTTLやCPU退避の運用ルールへ段階的に活かせます。

今日の学び

この論文は、AIエージェントが外部ツールや人を待つたびに、文脈の再計算またはGPUメモリの占有が起きる課題を扱いました。InferCeptは、KVキャッシュの保持・CPU退避・分割再計算を、待機時間と資源浪費の推定に基づいて選びます。開発では、ツール呼び出しを単なるAPI遅延として見るのではなく、推論の「中断・再開」として観測し、再計算トークンと待機メモリを最適化対象にすることが重要です。

関連記事