AWSがvLLMに双方向KVキャッシュ転送、マルチターンでTTFT最大3.3倍改善

vLLMに双方向KV転送を追加

AWSはMachine Learningブログで、vLLMに双方向KVキャッシュ転送（bi-directional KV transfer）の仕組みを追加したと報告しました。マルチターンのエージェント型ワークロードで、会話が長くなるほどTTFTが伸びる問題を抑える狙いです。AWSはまずTrainiumインスタンスクラスタとEFA接続の環境で構築・検証し、アクセラレータ非依存の機能としてvLLMにcontributeしたとしています。今回公開された性能データは、EFA接続のAmazon EC2 P5en GPUクラスタ上で取得したものです。

対象となるのは、prefillとdecodeを別ノードに分けるP-D分離（Prefill-decode disaggregation）構成です。prefillノードがプロンプトを処理してKVキャッシュを構築し、decodeノードがそのKVブロックを受け取って応答を生成します。従来はこの転送がprefillからdecodeへの一方向に限られていました。

一方向転送が生む2つの無駄

一方向の転送には2つの非効率があると記事は説明します。第一に、応答生成後のdecodeノードはプロンプトと生成済みトークンの両方のKVエントリを保持しますが、次のターンでprefillノードはそれを取得できず、前回応答分のKVを再計算します。

第二に、セッション状態を保持するのはdecodeノード側で、prefillノードは保持しません。そのためprefillノードでキャッシュが破棄されやすく、フルキャッシュミスになるとシステム指示・過去のプロンプト・ツール結果・モデル応答を含む会話履歴全体のKVを再計算することになります。会話が伸びるほどこの冗長な計算がアクセラレータの容量を消費し、TTFTを押し上げます。

双方向化では、decodeノードが生成後に再利用可能なKVブロックを保持し、そのメタデータをストリーミング応答で返します。ルーターがそれを次のprefillリクエスト用にキャッシュし、prefillノードはローカルキャッシュ一致を適用した後に再利用可能なKVを確認します。再利用トークン数がしきい値を満たせば非同期受信をキューに入れ、残りのサフィックスを計算する前にD→PのNIXL READを実行します。転送自体はNIXL READで、P→DとD→Pに共通の仕組みを使います。

ルーターと2つの調整パラメータ

ターンをまたいだKV再利用をまとめるのが、KVキャッシュ対応ルーター（KV-cache-aware router）です。クライアント・prefill・decodeノードの間に置く軽量なステートフル制御プレーンで、クライアントが渡すconversation_idとdecodeのkv_transfer_paramsの対応をメモリ上に保持します。保持するのはブロック位置のメタデータだけでKVテンソルは持ちません。またメタデータは応答ストリームから取得し、decodeノードへ別途問い合わせることはないとしています。vLLMのGitHubには、この記事で使ったフローを実装したサンプルプロキシが公開されているということです。

調整用のパラメータは2つあります。decoder_kv_blocks_ttlはブロック保持期間でデフォルト480秒、固定値で更新されません。ターン間隔の想定に合わせて調整し、KV再利用とアクセラレータメモリ使用量のバランスを取ることになります。もう1つのkv_recompute_thresholdはデフォルト64トークンで、計算対象のトークン数がこれ未満ならprefillで再計算し、以上ならdecodeからKVを取得します。短い再利用プレフィックスでは、相互接続と調整のオーバーヘッドが再計算時間を上回り得るためです。

P5enクラスタでの性能データ

評価はAmazon EC2 P5enインスタンスをEFAで接続した1P1D構成（prefill 1台、decode 1台）で、Qwen3-32Bとzai-org/GLM-5-FP8を使いました。prefillノードのローカルキャッシュが破棄された状況や、別のprefillノードが選ばれた状況を再現するため、prefill側のプレフィックスキャッシュは無効化しています（--no-enable-prefix-caching）。

機能OFFではprefillが毎ターン拡大した文脈全体を再計算するため、TTFTが急上昇します。Qwen3-32Bは2kで158msから24kで778msへ、GLM-5-FP8は2kで500msから20kで2,334msへ上昇しました。ONではdecodeが保持するKVを再利用し、Qwen3-32Bは122msから460msへ、GLM-5-FP8は358msから702msへと抑えられています。最長プロンプト時点で、Qwen3-32Bは約1.6倍、GLM-5-FP8は最大3.3倍のTTFT削減です。

内訳を見ると、双方向転送はKV転送時間を追加します（Qwen3-32Bで約30〜130ms、GLM-5-FP8で約29〜110ms）。ただし削減されるprefill計算の方が大きく、prefill実行時間の合計はONでQwen3-32Bが128〜414ms、GLM-5-FP8が316〜541msとなり、再計算のみの経路（Qwen3-32B 153〜744ms、GLM-5-FP8 500〜2,331ms）を下回ります。トークンあたりの再計算コストが大きいGLM-5-FP8の方が差が大きい、と記事は説明しています。

現時点での留意点

性能値は特定の構成での結果であり、ハードウェア・モデル・ワークロード・ネットワーク構成によって変わると明記されています。今回の数値はEFA接続のP5en GPUクラスタのもので、Trainiumでの検証結果が今回の記事のデータとして示されているわけではありません。

TTLが固定で更新されない点も運用上の論点です。記事は、ハートビートで更新されるリースをdecode側にも広げれば、会話がアイドルになった時点でブロックを回収できるようになる、という次の改善方向を示しています。ルーターはサンプルプロキシとして提供されている段階で、本番運用には自環境向けの適応が必要とみられます。