vLLMの双方向KV転送が多ターンエージェントのTTFTを短縮

概要

エージェント型の対話では、モデルが推論し、ツールを呼び出し、その結果を取り込むたびに会話履歴が伸びていく。新しいターンでは、システム指示、過去のユーザー入力、ツールの結果、以前のモデル出力が再び送信されることが多い。Prefill-Decode、つまりP-D分離型の推論基盤では、Prefillノードがプロンプトを処理し、その後Decodeノードが出力を生成する。過去のKV状態をPrefill側で利用できなければ、長くなった履歴を繰り返し計算することになる。

vLLMが紹介した双方向KV転送は、この重複を減らすための仕組みだ。従来の分離構成では、KVは主にPrefillからDecodeへ渡される。新しい方式では、生成後もDecodeノードが利用可能なKVブロックを保持し、次のターンでPrefillノードへ返せる。

主なポイント

• ターンをまたいでKVを再利用する。 初回は従来通り、PrefillがプロンプトのKVを計算し、Decodeがそれを読み取る。生成終了後、Decodeは履歴と応答に対応するブロックを保持し、次回はPrefillが再利用できる部分を読み込み、新しい末尾だけを計算する。
• ルーターはメタデータだけを扱う。 KV対応ルーターはconversation_idを使って、Decodeから返されたブロック位置や有効期限の情報を管理する。KVテンソルそのものはルーターを通らず、NIXLによってノード間で直接転送される。
• 保持期間と再計算閾値がある。 通常、Decodeは生成完了時にリクエストのブロックを解放する。双方向モードでは一定時間保持し、説明された設定ではdecoder_kv_blocks_ttlの既定値が480秒となっている。Prefillはまずローカルのプレフィックスキャッシュを確認し、再利用可能な量がkv_recompute_thresholdに達した場合に非同期のリモート読み出しを開始する。
• ノードの役割は変わらない。 Prefillは新しいプロンプト部分のKVを計算し、Decodeは既存のキャッシュブロックを提供する。PからD、DからPの両方向でNIXL READを利用するが、各ノードは役割に応じたメタデータ処理を行う。

性能と意味

報告されたテストは、AWS Elastic Fabric Adapterで接続したAmazon EC2 P5en上の1P1D構成で実施された。双方向転送により、Qwen3-32BではTTFTが最大1.6倍、zai-org/GLM-5-FP8では最大3.3倍改善した。実際の効果はワークロードと設定によって変わるため、すべての環境で同じ結果になるわけではない。

この方式の重要性は、単に一回の遅延を下げることだけではない。生成済みのKVをそのターン限りの一時データにせず、次のターンで使える状態としてPrefillとDecodeの間に保持する。これにより、長い会話での重複したPrefill計算を抑え、アクセラレーターの計算資源を新しい入力の処理に振り向けやすくなる。

一方、運用面では新しい管理も必要になる。状態を持つKV対応ルーター、安定した会話ID、期限切れ処理、リモートブロックが利用できない場合の再計算経路が求められる。TTL、再計算閾値、ローカルキャッシュの命中率、ネットワーク性能、セッション配置によって、リモート転送の利点は変動する。

この機能は当初、AWS TrainiumとEFAのクラスターで構築・検証され、その後アクセラレーターに依存しない機能としてvLLMへ提供された。今回の結果はEFA接続GPUクラスターのものであり、同じ考え方が異なるアクセラレーター環境にも適用できることを示している。

出典：vLLM Blog