vLLM x Mooncake による大規模なエージェンティック・ワークロードのサービング

エージェンティック・ワークロードは大規模な再利用可能なプレフィックスを生成する

エージェンティック・ワークロードは、再利用性の高い長いコンテキストを生成するため、プレフィックス・キャッシングが極めて重要になります。 Codex/SWE‑bench Pro のトレースによると、30ターン目までにコンテキスト長は約 80K トークンに達し、180K トークンを超えることもありますが、各ターンで追加される新しいトークンはわずか数百から数千程度です。トレースあたりのターン数の中央値が 33 ターンの 610 個のトレース全体において、データセットは以下の特性を示しています：

94.2% のキャッシュヒット率
131:1 の入力対出力トークン比
1ターンあたり平均約 2,242 トークンのコンテキスト増加
1トレースあたりのコンテキスト増加の中央値は 12K から 80K トークン
ターン間の遅延は中央値 5.2s から P99 で 81.4s の範囲これらの数値は、入力の大部分が再利用可能なプレフィックスであることを裏付けており、キャッシングによって冗長な prefill を排除できます。

Mooncake Store による分散 KV キャッシュプールがプレフィックス共有を解決

Mooncake Store を統合することで、ヒット率を向上させ、インスタンス間での再利用を可能にする共有 KV キャッシュが実現します。 Mooncake は、KV キャッシュの転送と分散ストレージのためのオープンソースライブラリです。 vLLM はすでに MooncakeConnector を介して、prefill-decode 分離（disaggregation）に Mooncake を使用しています。新しい統合では、複数の vLLM インスタンスが Mooncake クライアントを組み込み、クラスター全体の Mooncake Store マスターに接続する分散 KV キャッシュプールを構築します。マスターは KV ブロックのメタデータ、サービスディスカバリ、クライアントのヘルスチェックを管理し、ワーカーは RDMA を介して GPU HBM と分散 DRAM/SSD プール間で KV ブロックを転送します。スケジューラ側では、vLLM がプロンプトのトークンブロックをハッシュ化し、一致する KV キャッシュブロックを Mooncake マスターに問い合わせ、その結果を使用してスケジューリングをガイドします。ワーカー側では、各 GPU ワーカーが Mooncake クライアントを実行し、GPU KV キャッシュメモリを RDMA バッファとして登録し、バックグラウンド I/O スレッドを介してデータを移動させます。

設計のハイライト: SM-free GPUDirect RDMA、非同期転送、MultiConnector

実装では、ゼロコピーのために GPUDirect RDMA を使用し、転送をバックグラウンド I/O スレッドで実行し、MultiConnector を介した PD 分離（disaggregation）に対応しています。 KV ブロックの移動には GPUDirect RDMA を活用することで、ステージングバッファと SM の消費を回避し、複数の RNIC をプールして帯域幅を高めることができます。すべての RDMA 操作は専用のバックグラウンド I/O スレッドで実行されるため、転送パスは完全に非同期となり、メインの CPU パスでのストールを防ぎます。 MultiConnector インターフェースはサブコネクタを連鎖させます。prefill インスタンスは KV ブロックを分散プールに保存し、キャッシュヒットのためにそれらを回収でき、decode インスタンスは prefill から可視化できるようにプールにブロックを書き込みます。

Codex トレースにおけるパフォーマンス結果とスケーリング

現実的なエージェンティック・トレースにおいて、キャッシュヒット率は 1.7% から 92.2% に跳ね上がり、スループットは 3.8 倍、TTFT は 46 倍低減、レイテンシは 8.6 倍低減しました。60 GB200 GPU へのスケーリングにおいても、95% 以上のヒット率とほぼ線形なスループットの増加を維持しています。 Codex エージェンティック・トレース実験（1P1D、12 GB200 GPU）では、分散 KV キャッシュプールによりスループットが 3.8 倍向上し、P50 TTFT が 46 倍短縮され、エンドツーエンドのレイテンシが 8.6 倍削減されました。これらの利点は、キャッシュヒット率が 1.7%（システムプロンプトのみをキャッシュ）から 92.2%（ほぼすべてのプレフィックスをキャッシュ）に上昇したことによるものです。スケーリングについては、Codex ワークロードから派生した合成データセットを使用し、ノード間でラウンドロビン・ルーティングを行いました。システムはあらゆるスケールで 95% を超えるキャッシュヒット率を維持し、12 から 60 GB200 GPU にかけてスループットをほぼ線形にスケールさせ、分散プールがパフォーマンスを低下させるインスタンスをまたぐミスを防ぐことを実証しました。

今後の課題

計画されている拡張には、分散ディスクオフロード、ハイブリッドモデルのサポート、キャッシュ認識ルーティング、およびデータパスの最適化が含まれます。今後の作業では、より大きなキャッシュ容量のために、ストレージ階層を NVMe SSD や分散ファイルシステムに拡張する予定です。混合アテンションメカニズムを持つハイブリッドモデルのサポートが追加され、レイヤーごとに異なるキャッシング戦略を適用できるようになります。キャッシュ認識ルーティングは、リクエストルーターを KV キャッシュプールと共同設計し、関連するプレフィックスを保持しているインスタンスにターンを転送するようにすることで、分散プールにフォールバックする前にローカルヒットを最大化します。データパスの最適化では、NVIDIA のマルチノード NVLink や、prefill と decode インスタンスからの DualPath のような同時 KV ロードを調査し、総帯域幅を向上させます。

謝辞

vLLM Mooncake Store の統合は、vLLM-Ascend での先行研究から着想を得たものです。初期実装を行った Ant Group の Chao Lei 氏、およびエージェンティック・トレースと分析を行った Inferact の Zijing Liu 氏に感謝いたします。また、Approaching.AI の Jiahao Lu、Zuoyuan Zhang、Zihan Tang、Ke Yang 氏、Huawei の Pengbo Zhao、Fuqiao Duan、Tianyu Xu 氏、Alibaba Cloud Computing の Tianchen Ding、Xuchun Shang、Xingrui Yi、Teng Ma 氏、Ant Group の Yunxiao Ning、Dejiang Zhu、Shoujian Zheng 氏、および 9#AISoft の Feng Ren 氏に技術的なフィードバックへの謝意を表します。広範な vLLM および Mooncake コミュニティのサポート、ならびに密接な協力関係にある Inferact チームに感謝いたします。