[Core][PCP] Support data parallelism with MRV2 PCP

tests/entrypoints/openai/test_dp_supervisor.py

@@ -84,6 +84,7 @@ def _make_unit_args(**overrides) -> argparse.Namespace:
         "node_rank": 1,
         "tensor_parallel_size": 1,
         "pipeline_parallel_size": 1,
+        "prefill_context_parallel_size": 1,
         "uvicorn_log_level": "info",
         "shutdown_timeout": 5.0,
     }
@@ -119,6 +120,7 @@ def _make_args(**overrides) -> argparse.Namespace:
         node_rank=0,
         tensor_parallel_size=1,
         pipeline_parallel_size=1,
+        prefill_context_parallel_size=1,
         uvicorn_log_level="warning",
         shutdown_timeout=0.0,
     )
@@ -176,6 +178,22 @@ def test_build_multi_port_external_lb_child_args_sets_external_rank_server():
     assert child_args.api_server_count == 1
 
 
+def test_build_multi_port_external_lb_child_args_assigns_pcp_devices():
+    args = _make_unit_args(
+        data_parallel_size=2,
+        data_parallel_size_local=2,
+        data_parallel_start_rank=0,
+        node_rank=0,
+        prefill_context_parallel_size=2,
+    )
+
+    child_0 = _build_vllm_dp_server_args(args, local_rank=0)
+    child_1 = _build_vllm_dp_server_args(args, local_rank=1)
+
+    assert child_0.device_ids == [0, 1]
+    assert child_1.device_ids == [2, 3]
+
+
 def test_run_vllm_dp_server_uses_python_server_by_default(monkeypatch):
     calls: list[str] = []
 

tests/test_config.py

@@ -1532,6 +1532,17 @@ def test_needs_dp_coordination(
     assert vllm_config.needs_dp_coordinator == expected_needs_coordinator
 
 
+def test_parallel_config_supports_dp_with_pcp():
+    parallel_config = ParallelConfig(
+        data_parallel_size=2,
+        prefill_context_parallel_size=2,
+        enable_expert_parallel=True,
+    )
+
+    assert parallel_config.world_size == 2
+    assert parallel_config.world_size_across_dp == 4
+
+
 def test_renderer_num_workers_with_mm_cache():
     """Disallow renderer_num_workers > 1 when mm processor cache is enabled,
     since neither cache type is thread-safe."""

vllm/config/parallel.py

@@ -505,8 +505,6 @@ def _validate_parallel_config(self) -> Self:
         tp = self.tensor_parallel_size
         pcp = self.prefill_context_parallel_size
         dcp = self.decode_context_parallel_size
-        if pcp > 1 and self.data_parallel_size > 1:
-            raise ValueError("PCP does not support data parallelism yet.")
         if pcp == 1:
             # DCP reuses the TP ranks when PCP is disabled.
             if tp % dcp != 0:

vllm/distributed/device_communicators/all2all.py

@@ -49,23 +49,41 @@ class AgRsAll2AllManager(All2AllManagerBase):
     def __init__(self, cpu_group, tcp_store_group=None):
         super().__init__(cpu_group, tcp_store_group)
 
-    def _get_comm_group(self, is_sequence_parallel: bool) -> Any:
+    def _get_comm_groups(self, is_sequence_parallel: bool) -> list[Any]:
         if is_sequence_parallel:
-            return get_ep_group()
+            return [get_ep_group()]
+
+        # Compose the token-sharding axes and reverse them during combine.
+        groups = []
         if self.dp_world_size > 1:
-            return get_dp_group()
-        return get_pcp_group()
+            groups.append(get_dp_group())
+        if get_pcp_group().world_size > 1:
+            groups.append(get_pcp_group())
+        return groups
 
     def _get_sizes(self, num_local_tokens: int, comm_group: Any) -> list[int]:
-        if self.dp_world_size == 1:
+        if comm_group is get_pcp_group():
             return [num_local_tokens] * comm_group.world_size
 
         dp_metadata = get_forward_context().dp_metadata
         assert dp_metadata is not None
+        if comm_group is get_dp_group():
+            return dp_metadata.num_tokens_across_dp_cpu.tolist()
         sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
         assert sizes is not None
         return sizes
 
+    def _all_gatherv(
+        self, tensors: list[torch.Tensor], is_sequence_parallel: bool
+    ) -> list[torch.Tensor]:
+        for dist_group in self._get_comm_groups(is_sequence_parallel):
+            sizes = self._get_sizes(tensors[0].shape[0], dist_group)
+            assert sizes[dist_group.rank_in_group] == tensors[0].shape[0]
+            gathered = dist_group.all_gatherv(tensors, dim=0, sizes=sizes)
+            assert isinstance(gathered, list)
+            tensors = gathered
+        return tensors
+
     def dispatch_router_logits(
         self,
         hidden_states: torch.Tensor,
@@ -77,21 +95,13 @@ def dispatch_router_logits(
         | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
     ):
         """
-        Gather hidden_states and router_logits from all dp ranks.
+        Gather hidden_states and router_logits from all token-sharded ranks.
         """
-        dist_group = self._get_comm_group(is_sequence_parallel)
-        sizes = self._get_sizes(hidden_states.shape[0], dist_group)
-        assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]
-
         tensors_to_gather = [hidden_states, router_logits]
         if extra_tensors is not None:
             tensors_to_gather.extend(extra_tensors)
 
-        gathered_tensors = dist_group.all_gatherv(
-            tensors_to_gather,
-            dim=0,
-            sizes=sizes,
-        )
+        gathered_tensors = self._all_gatherv(tensors_to_gather, is_sequence_parallel)
 
         if extra_tensors is not None:
             return (gathered_tensors[0], gathered_tensors[1], gathered_tensors[2:])
@@ -109,21 +119,13 @@ def dispatch(
         | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
     ):
         """
-        Gather hidden_states and router_logits from all dp ranks.
+        Gather hidden states and routing outputs from all token-sharded ranks.
         """
-        dist_group = self._get_comm_group(is_sequence_parallel)
-        sizes = self._get_sizes(hidden_states.shape[0], dist_group)
-        assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]
-
         tensors_to_gather = [hidden_states, topk_weights, topk_ids]
         if extra_tensors is not None:
             tensors_to_gather.extend(extra_tensors)
 
-        gathered_tensors = dist_group.all_gatherv(
-            tensors_to_gather,
-            dim=0,
-            sizes=sizes,
-        )
+        gathered_tensors = self._all_gatherv(tensors_to_gather, is_sequence_parallel)
 
         hidden_states = gathered_tensors[0]
         topk_weights = gathered_tensors[1]
@@ -138,14 +140,16 @@ def combine(
         self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
     ) -> torch.Tensor:
         """
-        Reduce-scatter hidden_states across all dp ranks.
+        Reduce-scatter hidden states back to the local token shard.
         """
-        dist_group = self._get_comm_group(is_sequence_parallel)
-        sizes = self._get_sizes(
-            hidden_states.shape[0] // dist_group.world_size,
-            dist_group,
-        )
-        hidden_states = dist_group.reduce_scatterv(hidden_states, dim=0, sizes=sizes)
+        for dist_group in reversed(self._get_comm_groups(is_sequence_parallel)):
+            sizes = self._get_sizes(
+                hidden_states.shape[0] // dist_group.world_size,
+                dist_group,
+            )
+            hidden_states = dist_group.reduce_scatterv(
+                hidden_states, dim=0, sizes=sizes
+            )
         return hidden_states
 
     def destroy(self):

vllm/entrypoints/openai/dp_supervisor.py

@@ -137,7 +137,11 @@ def _build_device_ids(args: argparse.Namespace, local_rank: int) -> list[int | s
     var (e.g. CUDA_VISIBLE_DEVICES), so integer IDs must stay env-relative
     here rather than being translated to physical IDs.
     """
-    devices_per_rank = args.tensor_parallel_size * args.pipeline_parallel_size
+    devices_per_rank = (
+        args.tensor_parallel_size
+        * args.pipeline_parallel_size
+        * args.prefill_context_parallel_size
+    )
     start = local_rank * devices_per_rank
     stop = start + devices_per_rank
     device_ids = getattr(args, "device_ids", None)

vllm/forward_context.py

@@ -59,13 +59,16 @@ class BatchDescriptor:
 
 
 def _compute_sp_num_tokens(
-    num_tokens_across_dp_cpu: torch.Tensor, sequence_parallel_size: int
+    num_tokens_across_dp_cpu: torch.Tensor,
+    sequence_parallel_size: int,
+    pcp_size: int,
 ) -> list[int]:
     sp_tokens = (
         num_tokens_across_dp_cpu + sequence_parallel_size - 1
     ) // sequence_parallel_size
 
-    sp_tokens = sp_tokens.repeat_interleave(sequence_parallel_size)
+    # PCP counts are already rank-local, so replicate rather than divide them.
+    sp_tokens = sp_tokens.repeat_interleave(sequence_parallel_size * pcp_size)
     return sp_tokens.tolist()
 
 
@@ -99,13 +102,15 @@ def make(
         return DPMetadata(num_tokens_across_dp_cpu)
 
     @contextmanager
-    def sp_local_sizes(self, sequence_parallel_size: int):
+    def sp_local_sizes(self, sequence_parallel_size: int, pcp_size: int = 1):
         """
         Context manager for setting self.local_sizes. Same as self.chunked_sizes
         but without any chunking.
         """
         self.local_sizes = _compute_sp_num_tokens(
-            self.num_tokens_across_dp_cpu, sequence_parallel_size
+            self.num_tokens_across_dp_cpu,
+            sequence_parallel_size,
+            pcp_size,
         )
         try:
             yield self.local_sizes

vllm/model_executor/layers/attention/mla_attention.py

@@ -253,6 +253,7 @@
     LayerNameType,
     _encode_layer_name,
     _resolve_layer_name,
+    async_tensor_h2d,
     direct_register_custom_op,
     is_quantized_kv_cache,
     kv_cache_dtype_str_to_dtype,
@@ -1342,6 +1343,8 @@ class ChunkedContextMetadata:
         workspace: torch.Tensor
         token_to_seq: torch.Tensor
         chunk_total_token: list[int]
+        token_context_lens: torch.Tensor
+        max_context_chunk: int
 
         # for mla DCP
         padded_local_chunk_seq_lens: list[list[int]] | None = None
@@ -1567,6 +1570,13 @@ def build_mla_chunked_context_metadata(
         token_to_seq = torch.repeat_interleave(req_indices, chunk_seq_lens[i])
         token_to_seq_cpu[i, : token_to_seq.shape[0]] = token_to_seq
 
+    prefill_query_lens_cpu = (
+        prefill_query_start_loc_cpu[1:] - prefill_query_start_loc_cpu[:-1]
+    )
+    token_context_lens = async_tensor_h2d(
+        torch.repeat_interleave(context_lens_cpu, prefill_query_lens_cpu), device
+    )
+
     prefill_tokens_with_context = prefill_query_start_loc_cpu[
         num_prefills_with_context
     ].item()
@@ -1629,6 +1639,8 @@ def build_mla_chunked_context_metadata(
             token_to_seq=token_to_seq_cpu.to(device, non_blocking=True),
             chunk_total_token=chunk_total_token.tolist(),
             workspace=chunked_prefill_workspace,
+            token_context_lens=token_context_lens,
+            max_context_chunk=max_context_chunk,
             prefill_tokens_with_context=prefill_tokens_with_context,
             padded_local_chunk_seq_lens=padded_local_chunk_seq_lens.tolist(),
             local_context_lens_allranks=local_context_lens_allranks.tolist(),
@@ -1651,6 +1663,8 @@ def build_mla_chunked_context_metadata(
             token_to_seq=token_to_seq_cpu.to(device, non_blocking=True),
             chunk_total_token=chunk_total_token,
             workspace=chunked_prefill_workspace,
+            token_context_lens=token_context_lens,
+            max_context_chunk=max_context_chunk,
             prefill_tokens_with_context=prefill_tokens_with_context,
         )
 
@@ -2238,6 +2252,11 @@ def _compute_prefill_context(
                     v=v,
                 )
             )
+            chunk_start = i * prefill_metadata.chunked_context.max_context_chunk
+            attn_softmax_lse.masked_fill_(
+                prefill_metadata.chunked_context.token_context_lens <= chunk_start,
+                float("-inf"),
+            )
 
             if output is None:
                 output = attn_output
@@ -2388,6 +2407,11 @@ def _context_parallel_compute_prefill_context(
                     v=v,
                 )
             )
+            chunk_start = i * prefill_metadata.chunked_context.max_context_chunk
+            attn_softmax_lse.masked_fill_(
+                prefill_metadata.chunked_context.token_context_lens <= chunk_start,
+                float("-inf"),
+            )
 
             if output is None:
                 output = attn_output

vllm/model_executor/layers/fused_moe/runner/moe_runner.py

@@ -605,7 +605,9 @@ def _sequence_parallel_context(self):
         """
         ctx = get_forward_context()
         return (
-            ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size)
+            ctx.dp_metadata.sp_local_sizes(
+                self.moe_config.sp_size, self.moe_config.pcp_size
+            )
             if ctx.dp_metadata
             else nullcontext()
         )

vllm/v1/worker/gpu/model_runner.py

@@ -1246,6 +1246,11 @@ def execute_model(
                 block_tables = None
                 slot_mappings = None
 
+        if self.pcp_manager is not None and self.dp_size > 1:
+            # The earlier DP sync used the global, pre-PCP token count.
+            # Recompute MoE metadata from the rank-local batch.
+            num_tokens_across_dp = None
+
         attn_metadata = None
         slot_mappings_by_layer = None
         if not (dummy_run and skip_attn_for_dummy_run):

vllm/v1/worker/gpu_worker.py

@@ -310,13 +310,14 @@ def init_device(self):
                 if dp_local_rank is None:
                     dp_local_rank = self.parallel_config.data_parallel_index
 
-                tp_pp_world_size = (
+                tp_pp_pcp_world_size = (
                     self.parallel_config.pipeline_parallel_size
                     * self.parallel_config.tensor_parallel_size
+                    * self.parallel_config.prefill_context_parallel_size
                 )
 
-                # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK
-                self.local_rank += dp_local_rank * tp_pp_world_size
+                # DP_LOCAL_RANK * TP_PP_PCP_WORLD_SIZE + LOCAL_RANK
+                self.local_rank += dp_local_rank * tp_pp_pcp_world_size
 
             # Publish the logical-to-physical mapping for topology queries
             # such as NIC affinity and P2P checks.