TAO、ツールの返却予測をLLMエージェントの推論スケジューリングに組み込む

導入

LLMエージェントのサービスは、通常の連続生成とは異なります。セッションはGPU上でモデルを実行した後、分散CPU上の外部ツールを呼び出し、結果を待ってから再びモデル推論に戻ります。GPUのキューだけを見ているスケジューラでは、ツール待ちで実行できないセッションに資源を残したり、KVキャッシュの退避と再読み込みを繰り返したりする可能性があります。

TAO（Tool-Aware Joint Scheduling）は、ツール呼び出しをスケジューリングの外部にある一時停止ではなく、判断材料の一つとして扱います。重要なのは、どのセッションがGPUを必要としているかだけでなく、どのセッションが近く再開可能になるか、どのノードで効率よく再開できるかです。

TAOの仕組み

• ツール境界での中断。 セッションが外部ツールを呼び出すと、TAOはそのGPU実行を停止して容量を解放します。セッションは共有待機プールに移され、結果が戻った後に再びスケジューリング対象になります。
• 3要素による価値評価。 セッションとノードの組み合わせごとに、近い将来のツール返却準備度、ノード固有のKVキャッシュ信用度、トークン使用量を評価します。返却時刻は実行時ヒントと過去のツール実行時間から予測し、キャッシュの信用度は時間とともに減衰させます。
• 受け入れと配置の共同決定。 セッションを一つずつ貪欲に割り当てるのではなく、非同期の動的計画法で、GPU容量の制約下にある待機セッションの選択と配置を同時に決めます。報告された実行時オーバーヘッドは0.12%未満で、計画後の状態変化には再検証で対応します。

結果とトレードオフ

評価対象はDeepSeek-V4 Flash、Qwen3-32B、GLM-4.7で、モデル・ワークロード・同時実行数を組み合わせた12設定です。各指標で最も良かった評価済みベースラインと比べ、レイテンシは最大26.8%、平均TTFTは最大36.3%低下し、デコードスループットは最大29.5%向上しました。

一方で、TAOは過負荷時にすべての要求をすぐGPUへ入れるとは限りません。待機を意図的に許容することで、アクティブなKVキャッシュの追い出しや再計算を抑えます。AutomationBenchの128ワーカー構成では、平均TTFTはTAOが19.1秒、Dynamoが3.2秒でした。しかし全体タスクレイテンシは1417.2秒から563.1秒へ低下し、デコードスループットは482.4から1013.2 token/sへ向上しています。

意義

TAOが示すのは、エージェントサービスの文脈をGPUの空き容量だけで定義すべきではないという点です。外部ツールの時間的な振る舞い、KVキャッシュの局所性、GPU計算を一体として判断する必要があります。初期の配置が改善されれば、後続ターンでもキャッシュを保持しやすくなり、断片化やキューの混雑が連鎖的に抑えられる可能性があります。

ツール呼び出しが多く、負荷の変動が大きく、KVキャッシュ予算が限られるサービスでは特に参考になります。ただし、ツール時間の予測精度、ノードの異質性、TTFTとタスク完了時間の優先順位によって効果は変わります。実運用では、価値評価だけでなく既存の推論エンジンへ共同スケジューラを組み込む設計も重要になります。

出典：Hugging Face Daily Papers