vLLM Decode Context Parallelism for Long Context Workloads

TL;DR

vLLMは、KVキャッシュをアテンションヘッド単位ではなくシーケンス次元でGPU間に分割する技術、Decode Context Parallelism (DCP) を実装しました。これにより、Grouped-Query Attention (GQA) や Multi-head Latent Attention (MLA) モデルにおけるKVキャッシュの重複を防ぎ、長文コンテキスト推論におけるリクエストの並列実行数（Concurrency）とスループットを大幅に向上させることができます。

長文コンテキスト推論におけるメモリの壁

標準的なTensor Parallelism (TP) は、KVキャッシュをアテンションヘッドごとに分割します。この手法は、GPUの数がKVヘッドの数を超えると、システムが複数のGPUにわたってKVキャッシュを複製しなければならなくなるため、メモリのボトルネックを生じさせます。

Grouped-Query Attention (GQA): TPはKVキャッシュを1つのGPUあたり1つのヘッドにまでしか分割できません。TPの数がKVヘッドの数を超えると、キャッシュは複製されます。
Multi-head Latent Attention (MLA): MLAは、KVをすべてのクエリヘッドで共有される単一の低ランク潜在ベクトルに圧縮します。実質的に1つのKVヘッドしか持たないため、潜在KVキャッシュはすべてのTPランクに完全に複製され、並列リクエストに利用可能なGPUメモリを著しく制限します。

この重複は重要なGPUメモリを消費し、同時実行リクエスト数を制限し、トークンあたりのコストを増大させます。

Decode Context Parallelism (DCP) の解説

Decode Context Parallelismは、KVキャッシュをシーケンス（コンテキスト）次元でGPU間に分割することで、メモリのボトルネックを解決します。各GPUがシーケンス全体の特定のヘッドのキャッシュの完全なコピーを保持するのではなく、各GPUが同じシーケンス内のトークン位置のチャンク（塊）を担当します。

例えば、200Kトークンのリクエストにおいて、DCPを使用する4つのGPUは、GPU 0がトークン0–50Kを保持し、GPU 1が50K–100Kを保持するといった形でワークロードを分割できます。

DCPの実行プロセス

DCPは、デコードフェーズ中の正確性を維持するために、特定の通信リズムに従います。

AllGather Q: アテンションには完全なクエリベクトルが必要なため、all-gather操作によってすべてのGPU上でクエリの完全なコピーを組み立てます。（MLAモデルの場合、ロード時にクエリ投影を複製するために VLLM_DCP_Q_REPLICATE=1 を使用して最適化できます）。
Compute: 各GPUは、集約されたクエリと、自身のローカルなKVキャッシュのスライス（断片）に対してアテンションを計算します。
AllGather + ReduceScatter: 部分的な結果とLog-Sum-Exp (LSE) 値は、AllGatherを介して共有されます。LSE値は、オンライン・ソフトマックス・トリックを用いて部分的な結果を再重み付けおよびマージするために使用され、ReduceScatterはそれらを合計しつつ、各GPUに該当するヘッド・スライスのみを返します。

パフォーマンス・ベンチマーク

vLLMは、8×B200ノード上で、NVFP4を使用したKimi K2.6において、エージェント的な長文コンテキスト・トレース（中央値入力 ~67kトークン）を用いて、DCPをベースラインのTP展開に対してテストしました。

スループットと同時実行数

DCPは、ベースラインのTPと比較して、大幅に高い同時実行数とスループットを維持します。

Baseline TP: 同時実行数64においてメモリ使用率が100%に達し、スループットは1,863 tok/s/GPU付近で頭打ちになります。
DCP: 同時実行数512までスケールし、KVメモリの82%のみを使用しながら、6,091 tok/s/GPUに達します。

シーケンス長に応じた安定性

コンテキスト長が増加しても、パフォーマンスは安定しています。200k+トークンの範囲では、DCPは高いスループットとインタラクティビティの境界を維持しますが、複製されたKVキャッシュを用いるベースラインは、メモリ不足（OOM）が発生し、スケールしません。

実装と使用方法

DCPは、--decode-context-parallel-size 引数によって有効化されます。

Multi-head Latent Attention (MLA) バックエンド

DeepSeek-V2, V3, R1, および Kimi K2.6 のようなモデルで使用されます。MLAは実質的に1つのKVヘッドを持つため、シーケンスはTPの全次数まで分割可能です。

制約: tensor_parallel_size >= decode_context_parallel_size かつ tensor_parallel_size % decode_context_parallel_size == 0。

Grouped-Query Attention (GQA) バックエンド

Qwen3-235B や Llama-family モデルで使用されます。TPはまずKVヘッドごとに分割し、DCPは残りの冗長なコピーを分割（シャード）します。

制約: (tensor_parallel_size // num_key_value_heads) >= decode_context_parallel_size かつ (tensor_parallel_size // num_key_value_heads) % decode_context_parallel_size == 0。

将来のロードマップ

vLLMは、以下の開発を通じてDCPを拡張する計画です。

よりきめ細かな並列化: 過剰なプロビジョニングを抑えるため、TPとDCPのサイズをより精密に制御できるようにします。
通信カーネル: 計算のオーバーラップを改善するため、単一およびマルチノード設定における、より優れたall-to-all (A2A) カーネルを開発します。
Speculative Decoding: レイテンシの利点（メリット）を維持するため、DCPをMTPおよびSpeculative Decodingと統合します。
Disaggregated Serving: Prefill/Decode (P/D) 分離（Disaggregation）のサポートを強化します。
Expanded Model Support: GLM-5.2, Kimi K3 へのサポートを拡張し、Prefill Context Parallelism (PCP) を開発します。