[Feature] Prefill Context Parallel (PCP) basic support

[pisceskkk]

Purpose

This PR, splited from full PR #26864, adds the basic supports for the Prefill Context Parallelism (PCP) feature, which corresponds to DCP. For specific implementation details, please refer to the RFC #25749.

TL;DR: PCP enhances long-sequence inference capabilities by partitioning the sequence dimension during the prefill stage.

The current implementation primarily includes the following changes:

• Modified files such as block_tables.py to extend the KV cache storage based on DCP & PCP;
• Added a communication group pcp_group for PCP;
• Added necessary command-line arguments to control parallelism for PCP. Temporarily disabled PCP parameters until backend support is complete, then re-enable;
• Added PCP-related parameters to the attention backend prototype class;

CC @LookAround0301 @FENP @LucasWilkinson

[chatgpt-codex-connector]

💡 Codex Review

https://github.com/vllm-project/vllm/blob/14870a720161273d7583493580f81e14ab45199f/vllm/engine/arg_utils.py#L762-L772
Avoid duplicate --data-parallel-size argument registration

The CLI now calls parallel_group.add_argument("--data-parallel-size", …) twice in a row. argparse rejects duplicate option strings, so EngineArgs.add_cli_args() will raise ArgumentError: conflicting option string(s): --data-parallel-size before any command line can be parsed. This prevents vLLM from starting at all. One of the two registrations should be removed or renamed.

---

https://github.com/vllm-project/vllm/blob/14870a720161273d7583493580f81e14ab45199f/vllm/v1/worker/gpu_worker.py#L726-L737
Pass new PCP argument to FusedMoEParallelConfig.make

FusedMoEParallelConfig.make now requires a pcp_size_ positional argument, but the call in update_moe_modules still passes only tp_size_ and dp_size_. Any MoE model will hit this code path and raise TypeError: make() missing 1 required positional argument: 'pcp_size_' when the worker adjusts MoE modules. Update the invocation to include the prefill-context size (or provide a default) so MoE models can initialize.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[gemini-code-assist]

Code Review

This pull request introduces basic support for Prefill Context Parallelism (PCP), aligning with the RFC. The changes are consistently applied across the codebase, including updates to parallel configuration, KV cache management, and attention backends. Renaming dcp_kv_cache_interleave_size to cp_kv_cache_interleave_size generalizes the context parallelism KV cache interleaving logic to support both decode and prefill context parallelism. Compatibility checks are in place to temporarily disable PCP for certain features like full CUDA graphs and hybrid attention, indicating a phased rollout of full support. The integration appears thorough and well-considered for the initial support phase.

[luccafong]

thanks for the PR, can you share a few combinations of PCP x DCP x TP in your summary?

[heroes999]

can you share a few combinations of PCP x DCP x TP in your summary?

We usually set PCP=2, and TP=num_devices/PCP. For the MLA model, we set DCP=TP, and for the GQA model, we set DCP=TP/num_kv_heads, to fully utilize the KVCache. Some tests indicate that setting a larger PCP can further optimize TTFT, but currently, our testing primarily focuses on the case where PCP=2.

If we set PCP=2, TP=8 but without DCP, what's the KVCache split strategy? Input prompt segmented between ranks, but decoding output tokens containing copies?

[pisceskkk]

If we set PCP=2, TP=8 but without DCP, what's the KVCache split strategy? Input prompt segmented between ranks, but decoding output tokens containing copies?

Different PCP ranks contain split sequence dimensions and distinct KV caches, while copies may exist within the same PCP rank (depending on MLA or GQA configurations, specifically the relationship between num_kv_heads and tp_size).

In your example, assuming we are using an MLA model, all KV caches belonging to PCP0 are identical. The KV caches of PCP0 and PCP1 are different, following the same distribution pattern as DCP, where they are stored in interleave-style on cp_kvcache_interleave_size.

[heroes999]

If we set PCP=2, TP=8 but without DCP, what's the KVCache split strategy? Input prompt segmented between ranks, but decoding output tokens containing copies?

Different PCP ranks contain split sequence dimensions and distinct KV caches, while copies may exist within the same PCP rank (depending on MLA or GQA configurations, specifically the relationship between num_kv_heads and tp_size).

In your example, assuming we are using an MLA model, all KV caches belonging to PCP0 are identical. The KV caches of PCP0 and PCP1 are different, following the same distribution pattern as DCP, where they are stored in interleave-style on cp_kvcache_interleave_size.

Does it mean in decoding phase，pcp acts similarly to dcp, which requires gather q, calc local attn, correct lse..., sharding new generated tokens' kv cache interleavely?

[heroes999]

28988 , it seems dcp&pcp are merged to cp concept, so they do the same thing in decoding phase?

[pisceskkk]

If we set PCP=2, TP=8 but without DCP, what's the KVCache split strategy? Input prompt segmented between ranks, but decoding output tokens containing copies?

Different PCP ranks contain split sequence dimensions and distinct KV caches, while copies may exist within the same PCP rank (depending on MLA or GQA configurations, specifically the relationship between num_kv_heads and tp_size).
In your example, assuming we are using an MLA model, all KV caches belonging to PCP0 are identical. The KV caches of PCP0 and PCP1 are different, following the same distribution pattern as DCP, where they are stored in interleave-style on cp_kvcache_interleave_size.

Does it mean in decoding phase，pcp acts similarly to dcp, which requires gather q, calc local attn, correct lse..., sharding new generated tokens' kv cache interleavely?

#28988 , it seems dcp&pcp are merged to cp concept, so they do the same thing in decoding phase?

yes, in decoding phase, they act similarly for kvcache. If I remember correctly, there are only two differences: one is that dcp introduce an all-gather op for Q in head dim, and another is that dcp need a reduce-scatter after update output using lse.