vLLMのプッシュ型KV転送、分離推論におけるレイテンシを削減

Amazon Web Servicesは、オープンソースの推論エンジン「vLLM」向けに、プッシュ型のキーバリュー（KV）転送メカニズムを導入しました。このアップデートは、プリフィルとデコードを分離したアーキテクチャを対象としており、最初のトークン生成までの時間（TTFT）を大幅に短縮し、トークン生成時のスループット向上を目指しています。本機能は、Amazon EC2 P5enインスタンス上でQwen3-32Bモデルを用いて検証されました。

何が起きたか

大規模言語モデル（LLM）の推論には、入力プロンプトを処理してKVキャッシュを構築する「プリフィル」と、そのキャッシュを使用して出力トークンを生成する「デコード」という2つの異なるフェーズがあります。分離サービングでは、これらのフェーズが別々のエンジンインスタンスで実行され、独立したスケーリングと最適化が可能になります。しかし、この分離により、生成開始前にプリフィルインスタンスからデコードインスタンスへKVキャッシュを転送する必要があり、この転送の効率が全体のレイテンシにとって極めて重要となります。

従来、vLLMではNVIDIA Inference Xfer Library（NIXL）コネクタを介したプル方式が採用されていました。このデフォルトモードでは、デコードインスタンスはプリフィルの完了を待ち、ルーターからメタデータを受信し、メモリを割り当てた後、データを取得します。これらの手順はプリフィル完了後に順次実行されるため、クリティカルパスに大きな遅延が生じます。新しいプッシュ型コネクタは、調整作業をプリフィル計算自体と重複させることで、このフローを変更します。

AWSエンジニアは、Remote Direct Memory Access（RDMA）をサポートするAWS Elastic Fabric Adapter（EFA）で接続されたクラスター上で、この機能を開発・テストしました。当初はAWS Trainiumインスタンス向けに構築されましたが、この貢献はアクセラレーターに依存せず、GPUでも動作します。ベンチマークによると、高負荷時にはプッシュモードにより、プルモードと比較してTTFTが最大3倍短縮され、さらに1トークンあたりの出力時間（TPOT）も約30％改善することが示されています。

仕組み

プッシュ型アーキテクチャは、デコードインスタンスがあらかじめターゲットメモリブロックを登録できるようにすることで、直列化によるボトルネックを解消します。プリフィルの完了を待ってからデータ受信の準備をするのではなく、デコード側はローカルブロックを割り当てた瞬間にプリフィルインスタンスへ登録通知を送信します。これにより、ソースデータがまだ準備できていない段階で、宛先が確定している状態になります。

プリフィル側では、専用のバックグラウンドスレッドが転送準備を担当します。プリフィル計算が完了し、デコーダーからの登録が届くと、メインエンジンループはKVブロックをこのバックグラウンドスレッドに引き渡します。スレッドはNIXL接続を確立し、ブロックIDをマッピングした後、NIXL WRITE操作を発行して、データをデコーダーの登録済みメモリへ直接コピーします。これにより、ディスクリプターの準備や転送送信といった重い処理が、両方のインスタンスのレイテンシに敏感なメインループから切り離されます。

システムはタイミングの差異を処理するためにランデブープロトコルを使用します。登録がプリフィル完了より先に届いた場合、プロデューサーはそれを保持します。逆にプリフィルが先に完了した場合、登録が届くまでブロックを保持します。どちらのイベントが後に発生するかによって、即時転送がトリガーされます。デコーダーは各プロデューサーランクからの完了通知をカウントし、期待されるすべてのデータが到着してから初めてトークン生成に進みます。

主要な詳細

• 性能向上: 高いリクエストレートでの1P1Dデプロイメントにおいて、Qwen3-32Bに対するTTFTの最大3倍の短縮と、TPOTの30％の改善。
• テストハードウェア: EFAを搭載したAmazon EC2 P5enインスタンス。GPUDirect RDMAを使用し、GPU間での直接メモリ転送を実現。
• メカニズム: デコードインスタンスによるNIXL READではなく、プリフィルインスタンスが開始するNIXL WRITE操作を使用。
• スレッディング: プリフィルワーカー上のバックグラウンドスレッドが転送セットアップを処理し、メインモデル実行ループのスタール（停止）を防ぐ。
• トレードオフ: プッシュモードではプロセスの早い段階でデコードインスタンスを選択する必要があるため、プルモードと比較して負荷感知ルーティングの柔軟性が低下する可能性がある。
• 設定: ユーザーは、デフォルトのNixlConnectorの代わりに、プリフィルサーバーとデコードサーバーの両方をNixlPushConnectorに設定することで、この機能を有効にできます。

なぜ重要なのか

高スループットのLLMサービスを構築するエンジニアにとって、分離サービングはリソース使用量を効率的に管理するための一般的な戦略です。しかし、ノード間で巨大なKVキャッシュを移動するオーバーヘッドは、慎重に扱わなければその利点を打ち消してしまいます。従来のプルモデルでは、デコードインスタンスがメタデータを待機し、その後データを取得するという同期待ち時間が生じ、この期間中インスタンスはアイドル状態になります。コンテキスト長が増加し、リクエストレートが高まるにつれて、この遅延は顕著になり、ユーザーが体感するレイテンシに直接影響を与えます。

調整作業をバックグラウンドスレッドへ移行し、プリフィル計算と重複させることで、プッシュモデルはデータが利用可能になった瞬間にデコードインスタンスが受信準備を整えられるようにします。これにより、クリティカルパス上のアイドル時間が減少し、ユーザーへの応答速度が向上します。また、デコードワーカーが転送セットアップタスクに煩わされることを防ぎ、高い並行性下でも一貫したトークン生成速度を維持できるようにします。

この開発は、AIサービングにおける低レベルのインフラ最適化の重要性を浮き彫りにします。モデルが大規模化し、コンテキストが長くなるにつれて、計算ノード間のデータ移動効率は、行列演算そのものの速度と同様に重要になります。AWSや同様のRDMA対応インフラ上でvLLMを展開するエンジニアは、パフォーマンスを最大化するためにこのオプションを検討すべきです。

できること

• 現在のvLLMデプロイメントを評価し、ワークロードに対して分離プリフィル・デコードサービングが適切かどうかを確認してください。
• ステージング環境で現実的なトラフィックパターンを用いてNixlPushConnectorをテストし、TTFTの改善度を測定してください。
• AWS EFAなど、直接メモリアクセス機能を活用するために、インフラがRDMAをサポートしていることを確認してください。
• ルーター設定において、早期のデコードインスタンス選択と動的な負荷分散ニーズとのトレードオフを監視してください。
• キャッシングが転送コストを隠蔽する可能性があるため、ベンチマーク中はプレフィックスキャッシングを無効にして、転送パフォーマンスを隔離してください。
• プルモードからプッシュモードへ切り替えるために必要な具体的な設定フラグについては、vLLMドキュメントを確認してください。