[PCP] Support prefill context parallelism with data parallelism (#57075)
Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: QiuChunshuo <qiuchunshuo@huawei.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Lucas Wilkinson committed
Sep 23, 2026 at 22:32 UTC
e6dc16cebd57cabc55c1b3503ce16faaf6717d21
15 files changed
+268
-26
.buildkite/test_areas/misc.yaml
+2
@@ -518,6 +518,7 @@ steps:
518
- vllm/v1/
519
- tests/test_envs.py
520
- tests/test_outputs.py
521
+ - tests/test_pcp_dp.py
522
- tests/test_pooling_params.py
523
- tests/test_ray_env.py
524
- tests/test_sampling_params.py
@@ -538,6 +539,7 @@ steps:
539
- python3 standalone_tests/lazy_imports.py
540
- pytest -v -s test_envs.py
541
- pytest -v -s test_outputs.py
542
+ - pytest -v -s test_pcp_dp.py
543
- pytest -v -s test_pooling_params.py
544
- pytest -v -s test_ray_env.py
545
- pytest -v -s test_sampling_params.py
tests/test_pcp_dp.py
new
+161
@@ -0,0 +1,161 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from types import SimpleNamespace
4
+
5
+import pytest
6
+import torch
7
+
8
+from vllm.config import ParallelConfig
9
+from vllm.distributed.device_communicators.all2all import AgRsAll2AllManager
10
+from vllm.forward_context import DPMetadata
11
+from vllm.v1.attention.ops.pcp import maybe_gather_mla_latent_cache_inputs
12
+
13
+
14
+@pytest.mark.parametrize(
15
+ "pcp_size,sp_size,enable_ep,expected",
16
+ [
17
+ (1, 1, True, [5, 7]),
18
+ (1, 2, True, [3, 3, 4, 4]),
19
+ (2, 1, False, [10, 14]),
20
+ (2, 1, True, [5, 5, 7, 7]),
21
+ (2, 2, True, [3, 3, 3, 3, 4, 4, 4, 4]),
22
+ ],
23
+)
24
+def test_dispatch_sizes_expand_pcp_before_tp(pcp_size, sp_size, enable_ep, expected):
25
+ config = ParallelConfig(
26
+ distributed_executor_backend="mp",
27
+ data_parallel_size=2,
28
+ prefill_context_parallel_size=pcp_size,
29
+ tensor_parallel_size=sp_size,
30
+ enable_expert_parallel=enable_ep,
31
+ )
32
+ metadata = DPMetadata.make(config, 5, torch.tensor([5, 7]))
33
+ with metadata.sp_local_sizes(sp_size, pcp_size=pcp_size, use_ep=enable_ep) as sizes:
34
+ assert sizes == expected
35
+ assert metadata.local_sizes is None
36
+ assert metadata.num_tokens_across_dp_cpu.tolist() == [5, 7]
37
+
38
+
39
+@pytest.mark.parametrize(
40
+ "dp_size,pcp_size,tp_size,use_ep,is_sp,expected",
41
+ [
42
+ (2, 2, 1, True, False, "ep"),
43
+ (2, 2, 2, True, True, "ep"),
44
+ (2, 2, 2, False, False, "dp"),
45
+ (2, 1, 2, True, False, "dp"),
46
+ (1, 2, 2, True, False, "pcp"),
47
+ (2, 2, 2, True, False, None),
48
+ ],
49
+)
50
+def test_dispatch_reuses_existing_groups(
51
+ dp_size, pcp_size, tp_size, use_ep, is_sp, expected, monkeypatch
52
+):
53
+ groups = {
54
+ "dp": SimpleNamespace(world_size=dp_size),
55
+ "pcp": SimpleNamespace(world_size=pcp_size),
56
+ "ep": object(),
57
+ }
58
+ for name, group in groups.items():
59
+ monkeypatch.setattr(
60
+ f"vllm.distributed.device_communicators.all2all.get_{name}_group",
61
+ lambda group=group: group,
62
+ )
63
+ manager = AgRsAll2AllManager.__new__(AgRsAll2AllManager)
64
+ manager.dp_world_size = dp_size
65
+ manager.tp_group = SimpleNamespace(world_size=tp_size)
66
+ manager.use_ep = use_ep
67
+ if expected is None:
68
+ with pytest.raises(AssertionError, match="requires sequence-parallel MoE"):
69
+ manager._get_comm_group(is_sp)
70
+ else:
71
+ assert manager._get_comm_group(is_sp) is groups[expected]
72
+
73
+
74
+@pytest.mark.parametrize("enable_ep", [False, True])
75
+def test_ag_rs_dispatch_and_combine_use_dp_pcp_sizes(monkeypatch, enable_ep):
76
+ calls = []
77
+ local_tokens = [20, 21] if enable_ep else [20, 21, 22, 23]
78
+
79
+ class FakeGroup:
80
+ world_size = 4 if enable_ep else 2
81
+ rank_in_group = 2 if enable_ep else 1
82
+
83
+ def all_gatherv(self, tensors, dim, sizes):
84
+ calls.append(("gather", sizes))
85
+ return [torch.tensor([10, 11, 20, 21, 22, 23]) for _ in tensors]
86
+
87
+ def reduce_scatterv(self, tensor, dim, sizes):
88
+ calls.append(("scatter", sizes))
89
+ return tensor[2 : 2 + len(local_tokens)]
90
+
91
+ config = ParallelConfig(
92
+ distributed_executor_backend="mp",
93
+ data_parallel_size=2,
94
+ data_parallel_rank=1,
95
+ prefill_context_parallel_size=2,
96
+ enable_expert_parallel=enable_ep,
97
+ )
98
+ metadata = DPMetadata.make(config, 2, torch.tensor([1, 2]))
99
+ manager = AgRsAll2AllManager.__new__(AgRsAll2AllManager)
100
+ manager.dp_world_size = 2
101
+ manager.use_ep = enable_ep
102
+
103
+ manager.tp_group = SimpleNamespace(world_size=1)
104
+ for name in ("dp", "ep"):
105
+ monkeypatch.setattr(
106
+ f"vllm.distributed.device_communicators.all2all.get_{name}_group",
107
+ FakeGroup,
108
+ )
109
+ monkeypatch.setattr(
110
+ "vllm.distributed.device_communicators.all2all.get_pcp_group",
111
+ lambda: SimpleNamespace(world_size=2),
112
+ )
113
+ monkeypatch.setattr(
114
+ "vllm.distributed.device_communicators.all2all.get_forward_context",
115
+ lambda: SimpleNamespace(dp_metadata=metadata),
116
+ )
117
+ with metadata.sp_local_sizes(1, pcp_size=2, use_ep=enable_ep):
118
+ hidden_states, _, _ = manager.dispatch(
119
+ torch.tensor(local_tokens),
120
+ torch.ones(len(local_tokens)),
121
+ torch.zeros(len(local_tokens)),
122
+ )
123
+ combined = manager.combine(hidden_states)
124
+ assert combined.tolist() == local_tokens
125
+ sizes = [1, 1, 2, 2] if enable_ep else [2, 4]
126
+ assert calls == [("gather", sizes), ("scatter", sizes)]
127
+
128
+
129
+@pytest.mark.parametrize("slots", [[3, 4], [3, 4, -1, -1, -1, -1]])
130
+def test_decode_cache_write_ignores_dp_padding(slots):
131
+ kv = torch.arange(12).reshape(6, 2)
132
+ pe = torch.arange(6).reshape(6, 1, 1)
133
+ slots = torch.tensor(slots)
134
+ cache_kv, cache_pe, cache_slots = maybe_gather_mla_latent_cache_inputs(
135
+ kv, pe, slots, num_decode_tokens=2, use_pcp=True
136
+ )
137
+ torch.testing.assert_close(cache_kv, kv[:2])
138
+ torch.testing.assert_close(cache_pe, pe[:2])
139
+ torch.testing.assert_close(cache_slots, slots[:2])
140
+
141
+
142
+def test_expanded_slot_mapping_keeps_pcp_prefill_padding(monkeypatch):
143
+ calls = []
144
+
145
+ def all_gather(tensor, dim):
146
+ calls.append(tensor.shape[0])
147
+ return torch.cat((tensor, tensor), dim=dim)
148
+
149
+ monkeypatch.setattr(
150
+ "vllm.v1.attention.ops.pcp.get_pcp_group",
151
+ lambda: SimpleNamespace(world_size=2, all_gather=all_gather),
152
+ )
153
+ kv = torch.zeros(3, 2) # Two decodes followed by one PCP padding row.
154
+ pe = torch.zeros(3, 1, 1)
155
+ slots = torch.tensor([3, 4, 8, 3, 4, -1])
156
+ cache_kv, _, cache_slots = maybe_gather_mla_latent_cache_inputs(
157
+ kv, pe, slots, num_decode_tokens=2, use_pcp=True
158
+ )
159
+ assert calls == [1, 1]
160
+ assert cache_kv.shape == (4, 2)
161
+ assert cache_slots.tolist() == [3, 4, 8, -1]
tests/v1/worker/test_gpu_autoregressive_speculator.py
+44
@@ -19,6 +19,7 @@ from vllm.model_executor.models.mistral_large_3_eagle import (
19
)
20
from vllm.v1.attention.backends import flash_attn as flash_attn_module
21
from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
22
+from vllm.v1.attention.backends.utils import split_decodes_and_prefills
23
from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
24
from vllm.v1.worker.gpu.spec_decode import speculator as base_spec_module
25
from vllm.v1.worker.gpu.spec_decode.autoregressive import speculator as spec_module
@@ -69,6 +70,49 @@ class _TextOnlyDraftModel(torch.nn.Module):
70
raise AssertionError("embed_input_ids should not be called during loading")
71
72
73
+@pytest.mark.parametrize("cg_mode", [CUDAGraphMode.NONE, CUDAGraphMode.FULL])
74
+def test_pcp_draft_metadata_keeps_graph_padding_in_decode(cg_mode):
75
+ def build(common_prefix_len, common_attn_metadata):
76
+ return split_decodes_and_prefills(
77
+ common_attn_metadata,
78
+ decode_threshold=1,
79
+ require_uniform=True,
80
+ treat_short_extends_as_decodes=False,
81
+ )
82
+
83
+ speculator = object.__new__(_TestSpeculator)
84
+ speculator.arange_np = torch.arange(5, dtype=torch.int32).numpy()
85
+ speculator.max_model_len = speculator.draft_max_seq_len = 32
86
+ speculator.draft_is_prefilling = torch.zeros(4, dtype=torch.bool)
87
+ speculator.input_buffers = SimpleNamespace(
88
+ query_start_loc=torch.tensor([0, 1, 2, 2, 2], dtype=torch.int32),
89
+ seq_lens=torch.tensor([11, 21, 0, 0], dtype=torch.int32),
90
+ )
91
+ speculator.block_tables = SimpleNamespace(
92
+ cp_size=1,
93
+ input_block_tables=[torch.zeros(4, 1, dtype=torch.int32)],
94
+ slot_mappings=torch.tensor([[10, 20, -1, -1]]),
95
+ )
96
+ speculator.kv_cache_config = SimpleNamespace(kv_cache_groups=[object()])
97
+ speculator.attn_groups = [
98
+ [
99
+ SimpleNamespace(
100
+ get_metadata_builder=lambda _: SimpleNamespace(build=build),
101
+ layer_names=["draft"],
102
+ )
103
+ ]
104
+ ]
105
+ num_reqs_padded = 4 if cg_mode == CUDAGraphMode.FULL else 2
106
+ metadata = speculator._build_uniform_attn_metadata(
107
+ batch_desc=BatchExecutionDescriptor(cg_mode, 4, num_reqs_padded),
108
+ num_reqs=2,
109
+ num_query_per_req=1,
110
+ seq_lens_cpu_upper_bound=torch.tensor([10, 20], dtype=torch.int32),
111
+ step=1,
112
+ )
113
+ assert metadata["draft"] == (num_reqs_padded, 0, num_reqs_padded, 0)
114
+
115
+
116
def _mock_base_model_load(monkeypatch):
117
monkeypatch.setattr(
118
base_spec_module,
tests/v1/worker/test_gpu_pcp_manager.py
+12
@@ -209,6 +209,18 @@ def test_partition_padding_is_derived_from_batch_descriptor(
209
assert local_batch.num_reqs_after_padding == expected_reqs
210
211
212
+def test_dummy_draft_does_not_reuse_previous_graph_batch():
213
+ manager, _ = _make_capture_manager(torch.ones((4, 2), dtype=torch.int32))
214
+ dummy_batch = InputBatch.make_dummy(1, 4, manager.input_buffers)
215
+ manager.draft_prefill_batch = replace(dummy_batch)
216
+
217
+ manager.prepare_draft_prefill(dummy_batch, dummy_batch.input_ids)
218
+
219
+ assert (
220
+ manager.get_draft_input_buffers(manager.input_buffers) is manager.input_buffers
221
+ )
222
+
223
+
224
def test_capture_uses_pcp_persistent_inputs():
225
manager, _ = _make_capture_manager(torch.ones((4, 2), dtype=torch.int32))
226
vllm/distributed/device_communicators/all2all.py
+7
@@ -56,11 +56,17 @@ class AgRsAll2AllManager(All2AllManagerBase):
56
57
def __init__(self, cpu_group, tcp_store_group=None):
58
super().__init__(cpu_group, tcp_store_group)
59
+ self.use_ep = get_current_vllm_config().parallel_config.enable_expert_parallel
60
61
def _get_comm_group(self, is_sequence_parallel: bool) -> Any:
62
if is_sequence_parallel:
63
return get_ep_group()
64
if self.dp_world_size > 1:
65
+ if self.use_ep and get_pcp_group().world_size > 1:
66
+ assert self.tp_group.world_size == 1, (
67
+ "DP+PCP with TP>1 requires sequence-parallel MoE inputs"
68
+ )
69
+ return get_ep_group()
70
return get_dp_group()
71
return get_pcp_group()
72
@@ -72,6 +78,7 @@ class AgRsAll2AllManager(All2AllManagerBase):
78
assert dp_metadata is not None
79
sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
80
assert sizes is not None
81
+ assert len(sizes) == comm_group.world_size
82
return sizes
83
84
def dispatch_router_logits(
vllm/distributed/elastic_ep/standby_state.py
+7
-3
@@ -6,6 +6,7 @@ from vllm.distributed.device_communicators.pynccl import defer_comm_warmup_on_ro
6
from vllm.distributed.parallel_state import (
7
_init_stateless_group,
8
_node_count,
9
+ get_pcp_group,
10
get_pp_group,
11
get_tp_group,
12
get_world_group,
@@ -72,12 +73,13 @@ def create_standby_groups(
73
_STANDBY_WORLD_NODE_COUNT = _node_count(_STANDBY_WORLD.tcp_store_group)
74
75
tp_size = get_tp_group().world_size
76
+ pcp_size = get_pcp_group().world_size
77
pp_size = get_pp_group().world_size
78
79
all_ranks = torch.arange(new_world_size_across_dp).reshape(
78
- -1, new_dp_size, pp_size, tp_size
80
+ -1, new_dp_size, pp_size, pcp_size, tp_size
81
)
80
- standby_dp_ranks = all_ranks.transpose(1, 3).reshape(-1, new_dp_size).unbind(0)
82
+ standby_dp_ranks = all_ranks.transpose(1, 4).reshape(-1, new_dp_size).unbind(0)
83
standby_dp_ranks = [x.tolist() for x in standby_dp_ranks]
84
85
# Deferred to commit so the warm-up runs while the engine is paused.
@@ -87,7 +89,9 @@ def create_standby_groups(
89
)
90
91
standby_ep_ranks = (
90
- all_ranks.transpose(1, 2).reshape(-1, new_dp_size * tp_size).unbind(0)
92
+ all_ranks.transpose(1, 2)
93
+ .reshape(-1, new_dp_size * pcp_size * tp_size)
94
+ .unbind(0)
95
)
96
standby_ep_ranks = [x.tolist() for x in standby_ep_ranks]
97
_STANDBY_EP = _init_stateless_group(
vllm/forward_context.py
+14
-3
@@ -58,8 +58,17 @@ class BatchDescriptor:
58
59
60
def _compute_sp_num_tokens(
61
- num_tokens_across_dp_cpu: torch.Tensor, sequence_parallel_size: int
61
+ num_tokens_across_dp_cpu: torch.Tensor,
62
+ sequence_parallel_size: int,
63
+ pcp_size: int = 1,
64
+ use_ep: bool = True,
65
) -> list[int]:
66
+ if pcp_size > 1:
67
+ num_tokens_across_dp_cpu = (
68
+ num_tokens_across_dp_cpu.repeat_interleave(pcp_size)
69
+ if use_ep
70
+ else num_tokens_across_dp_cpu * pcp_size
71
+ )
72
sp_tokens = (
73
num_tokens_across_dp_cpu + sequence_parallel_size - 1
74
) // sequence_parallel_size
@@ -98,12 +107,14 @@ class DPMetadata:
107
return DPMetadata(num_tokens_across_dp_cpu)
108
109
@contextmanager
101
- def sp_local_sizes(self, sequence_parallel_size: int):
110
+ def sp_local_sizes(
111
+ self, sequence_parallel_size: int, pcp_size: int = 1, use_ep: bool = False
112
+ ):
113
"""Context manager for setting self.local_sizes. Same as self.chunked_sizes
114
but without any chunking.
115
"""
116
self.local_sizes = _compute_sp_num_tokens(
106
- self.num_tokens_across_dp_cpu, sequence_parallel_size
117
+ self.num_tokens_across_dp_cpu, sequence_parallel_size, pcp_size, use_ep
118
)
119
try:
120
yield self.local_sizes
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+5
-1
@@ -657,7 +657,11 @@ class MoERunner(MoERunnerInterface):
657
"""
658
ctx = get_forward_context()
659
return (
660
- ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size)
660
+ ctx.dp_metadata.sp_local_sizes(
661
+ self.moe_config.sp_size,
662
+ pcp_size=self.moe_config.pcp_size,
663
+ use_ep=self.moe_config.use_ep,
664
+ )
665
if ctx.dp_metadata
666
else nullcontext()
667
)
vllm/platforms/cuda.py
-6
@@ -329,12 +329,6 @@ class CudaPlatformBase(Platform):
329
parallel_config = vllm_config.parallel_config
330
model_config = vllm_config.model_config
331
332
- if (
333
- parallel_config.prefill_context_parallel_size > 1
334
- and parallel_config.data_parallel_size > 1
335
- ):
336
- raise ValueError("PCP does not support data parallelism on CUDA yet.")
337
-
332
if parallel_config.worker_cls == "auto":
333
parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker"
334
vllm/platforms/rocm.py
-6
@@ -934,12 +934,6 @@ class RocmPlatform(Platform):
934
compilation_config = vllm_config.compilation_config
935
parallel_config = vllm_config.parallel_config
936
937
- if (
938
- parallel_config.prefill_context_parallel_size > 1
939
- and parallel_config.data_parallel_size > 1
940
- ):
941
- raise ValueError("PCP does not support data parallelism on ROCm yet.")
942
-
937
if (
938
compilation_config.cudagraph_mode.has_full_cudagraphs()
939
and parallel_config.prefill_context_parallel_size > 1
vllm/v1/attention/backends/utils.py
+1
-1
@@ -828,7 +828,7 @@ def split_decodes_and_prefills(
828
(query_lens == query_lens[0]) | (query_lens == 0)
829
):
830
return num_reqs, 0, num_tokens, 0 # all decodes
831
- is_prefill = query_lens != query_lens[0]
831
+ is_prefill = (query_lens != query_lens[0]) & (query_lens != 0)
832
else:
833
is_prefill = query_lens > decode_threshold
834
vllm/v1/attention/ops/pcp.py
+9
-2
@@ -18,8 +18,15 @@ def _gather_prefill_cache_inputs(
18
assert all(tensor.shape[0] == local_num_tokens for tensor in tensors)
19
assert 0 <= num_decode_tokens <= local_num_tokens
20
21
- if num_decode_tokens == local_num_tokens:
22
- return tensors, slot_mapping[:num_decode_tokens]
21
+ # Replicated draft decodes use unexpanded slot mappings, even with DP padding.
22
+ if (
23
+ num_decode_tokens == local_num_tokens
24
+ or slot_mapping.shape[0] <= local_num_tokens
25
+ ):
26
+ return (
27
+ tuple(tensor[:num_decode_tokens] for tensor in tensors),
28
+ slot_mapping[:num_decode_tokens],
29
+ )
30
31
pcp_group = get_pcp_group()
32
gathered_prefills = tuple(
vllm/v1/worker/gpu/pcp_manager.py
+1
@@ -754,6 +754,7 @@ class PCPManager:
754
def prepare_draft_prefill(
755
self, input_batch: InputBatch, input_ids: torch.Tensor
756
) -> None:
757
+ self.draft_prefill_batch = None
758
if input_batch is not self._global_batch or self._local_batch is None:
759
return
760
local_batch = self._local_batch
vllm/v1/worker/gpu/spec_decode/speculator.py
+1
-1
@@ -364,7 +364,7 @@ class DraftModelSpeculator(BaseSpeculator):
364
kv_cache_config=self.kv_cache_config,
365
causal=causal,
366
seq_lens_cpu_upper_bound=draft_seq_lens_cpu_upper_bound,
367
- is_prefilling=self.draft_is_prefilling[:num_reqs],
367
+ is_prefilling=self.draft_is_prefilling[:num_reqs_padded],
368
)
369
return attn_metadata
370
vllm/v1/worker/gpu_worker.py
+4
-3
@@ -406,13 +406,14 @@ class Worker(WorkerBase):
406
if dp_local_rank is None:
407
dp_local_rank = self.parallel_config.data_parallel_index
408
409
- tp_pp_world_size = (
409
+ tp_pcp_pp_world_size = (
410
self.parallel_config.pipeline_parallel_size
411
+ * self.parallel_config.prefill_context_parallel_size
412
* self.parallel_config.tensor_parallel_size
413
)
414
414
- # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK
415
- self.local_rank += dp_local_rank * tp_pp_world_size
415
+ # DP_LOCAL_RANK * TP_PCP_PP_WORLD_SIZE + TP_LOCAL_RANK
416
+ self.local_rank += dp_local_rank * tp_pcp_pp_world_size
417
418
# Publish the logical-to-physical mapping for topology queries
419
# such as NIC affinity and P2P checks.