[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

vllm/vllm/engine/arg_utils.py

Lines 762 to 772 in 14870a7

	parallel_group.add_argument(
	"--data-parallel-size", "-dp", **parallel_kwargs["data_parallel_size"]
	)
	parallel_group.add_argument(
	"--prefill-context-parallel-size",
	"-pcp",
	**parallel_kwargs["prefill_context_parallel_size"],
	)
	parallel_group.add_argument(
	"--data-parallel-size", "-dp", **parallel_kwargs["data_parallel_size"]
	)

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.

---

vllm/vllm/v1/worker/gpu_worker.py

Lines 726 to 737 in 14870a7

	def update_moe_modules(moe_modules: list[FusedMoE], num_local_experts: int):
	assert all(
	module.moe_config.num_local_experts == num_local_experts
	for module in moe_modules
	), "All MoE modules must have the same number of experts"
	for module in moe_modules:
	module.moe_config.num_experts = num_local_experts * new_ep_size
	module.global_num_experts = module.moe_config.num_experts
	module.moe_parallel_config = FusedMoEParallelConfig.make(
	tp_size_=get_tp_group().world_size,
	dp_size_=get_dp_group().world_size,
	vllm_parallel_config=parallel_config,

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?

[LucasWilkinson]

Overall looks pretty good to me; running the CI, left a couple of final nits/comments

[LucasWilkinson]

LGTM