[Bugfix] Account for PCP in multi-node world size validation (#55111)

Signed-off-by: DebugSy <532151398@qq.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

tests/engine/test_arg_utils.py

@@ -464,6 +464,25 @@ def test_attention_config():
         engine_args.create_engine_config()
 
 
+def test_multi_node_world_size_includes_pcp(monkeypatch):
+    """PCP expands the process world size, so the --nnodes divisibility check
+    must include it. Without this, TP=1/PCP=2 over 2 nodes computes a world
+    size of 1 and the launch is rejected before the engine starts."""
+    import vllm.config.vllm
+
+    # PCP requires the V2 model runner, which is gated on Triton.
+    monkeypatch.setattr(vllm.config.vllm, "HAS_TRITON", True)
+
+    engine_args = EngineArgs(
+        model="facebook/opt-125m",
+        tensor_parallel_size=1,
+        prefill_context_parallel_size=2,
+        nnodes=2,
+    )
+    vllm_config = engine_args.create_engine_config()
+    assert vllm_config.parallel_config.world_size == 2
+
+
 def test_prefix_cache_default():
     parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
     args = parser.parse_args([])

vllm/engine/arg_utils.py

@@ -2121,21 +2121,19 @@ class EngineArgs:
             )
         inferred_data_parallel_rank = 0
         if self.nnodes > 1:
-            world_size = (
-                self.data_parallel_size
-                * self.pipeline_parallel_size
-                * self.tensor_parallel_size
-            )
             world_size_within_dp = (
-                self.pipeline_parallel_size * self.tensor_parallel_size
+                self.pipeline_parallel_size
+                * self.tensor_parallel_size
+                * self.prefill_context_parallel_size
             )
+            world_size = self.data_parallel_size * world_size_within_dp
             if world_size % self.nnodes != 0:
                 raise ValueError(
                     "Invalid data-parallel launch options: "
                     f"`--nnodes {self.nnodes}` must evenly divide the total "
                     f"world size ({world_size}). Adjust `--nnodes`, "
-                    "`--data-parallel-size`, `--pipeline-parallel-size`, or "
-                    "`--tensor-parallel-size`."
+                    "`--data-parallel-size`, `--pipeline-parallel-size`, "
+                    "`--tensor-parallel-size`, or `--prefill-context-parallel-size`."
                 )
             if not 0 <= self.node_rank < self.nnodes:
                 raise ValueError(