[PCP][Spec Decode] Prefer sharded MTP with replicated draft fallback

tests/v1/worker/test_gpu_autoregressive_speculator.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from contextlib import nullcontext
+from dataclasses import dataclass
 from types import SimpleNamespace
 from unittest.mock import Mock
 
@@ -28,7 +29,10 @@
 from vllm.v1.worker.gpu.spec_decode.multi_module_mtp.speculator import (
     MultiModuleMTPSpeculator,
 )
-from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
+from vllm.v1.worker.gpu.spec_decode.speculator import (
+    DraftModelSpeculator,
+    _draft_attention_backend_supports_pcp,
+)
 
 
 class _TestSpeculator(AutoRegressiveSpeculator):
@@ -95,6 +99,8 @@ def _make_speculator(
     speculator = object.__new__(_TestSpeculator)
     speculator.supports_mm_inputs = False
     speculator.pcp_manager = None
+    speculator.target_pcp_size = 1
+    speculator.uses_pcp_draft_sharding = False
     speculator.vllm_config = None
     speculator.input_buffers = SimpleNamespace(
         input_ids=torch.arange(4),
@@ -105,6 +111,173 @@ def _make_speculator(
     return speculator
 
 
+def _draft_backend_config(
+    draft_backend=None,
+    target_backend=None,
+    backend_per_kind=None,
+):
+    draft_model_config = SimpleNamespace(
+        use_mla=False,
+        get_sliding_window=lambda: None,
+    )
+    return SimpleNamespace(
+        speculative_config=SimpleNamespace(
+            attention_backend=draft_backend,
+            draft_model_config=draft_model_config,
+        ),
+        attention_config=SimpleNamespace(
+            backend=target_backend,
+            backend_per_kind=backend_per_kind or {},
+        ),
+    )
+
+
+def _backend_with_pcp_support(supports_pcp: bool):
+    backend_cls = SimpleNamespace(supports_pcp=lambda: supports_pcp)
+    return SimpleNamespace(get_class=lambda: backend_cls)
+
+
+@dataclass
+class _ParallelConfig:
+    prefill_context_parallel_size: int = 4
+    data_parallel_size: int = 1
+    data_parallel_rank: int = 0
+
+
+@dataclass
+class _SpeculatorVllmConfig:
+    speculative_config: object
+    attention_config: object
+    parallel_config: _ParallelConfig
+    scheduler_config: object
+    model_config: object
+
+
+def test_auto_draft_backend_prefers_pcp_sharding():
+    assert _draft_attention_backend_supports_pcp(_draft_backend_config())
+
+
+@pytest.mark.parametrize("supports_pcp", [False, True])
+def test_explicit_draft_backend_controls_pcp_sharding(supports_pcp):
+    backend = _backend_with_pcp_support(supports_pcp)
+    config = _draft_backend_config(draft_backend=backend)
+
+    assert _draft_attention_backend_supports_pcp(config) is supports_pcp
+
+
+def test_draft_backend_per_kind_takes_precedence_for_pcp_sharding():
+    supported = _backend_with_pcp_support(True)
+    unsupported = _backend_with_pcp_support(False)
+    config = _draft_backend_config(
+        draft_backend=supported,
+        backend_per_kind={"full_attention": unsupported},
+    )
+
+    assert not _draft_attention_backend_supports_pcp(config)
+
+
+@pytest.mark.parametrize(
+    ("supports_pcp", "expected_draft_pcp_size"),
+    [(True, 4), (False, 1)],
+)
+def test_draft_config_prefers_sharding_and_falls_back_to_replication(
+    monkeypatch, supports_pcp, expected_draft_pcp_size
+):
+    monkeypatch.setattr(base_spec_module, "_target_feeds_hc_residual", lambda _: False)
+    draft_model_config = SimpleNamespace(
+        hf_config=SimpleNamespace(),
+        use_mla=False,
+        get_sliding_window=lambda: None,
+        get_hidden_size=lambda: 16,
+        get_vocab_size=lambda: 32,
+    )
+    speculative_config = SimpleNamespace(
+        method="mtp",
+        num_speculative_tokens=3,
+        draft_model_config=draft_model_config,
+        attention_backend=_backend_with_pcp_support(supports_pcp),
+        use_local_argmax_reduction=False,
+        draft_sample_method="greedy",
+        enable_adaptive_verification=False,
+    )
+    config = _SpeculatorVllmConfig(
+        speculative_config=speculative_config,
+        attention_config=SimpleNamespace(backend=None, backend_per_kind={}),
+        parallel_config=_ParallelConfig(),
+        scheduler_config=SimpleNamespace(
+            max_num_seqs=2,
+            max_num_batched_tokens=8,
+        ),
+        model_config=SimpleNamespace(
+            max_model_len=32,
+            dtype=torch.float32,
+            use_fp64_gumbel=False,
+        ),
+    )
+
+    speculator = _TestSpeculator(config, torch.device("cpu"))
+
+    assert speculator.target_pcp_size == 4
+    assert speculator.uses_pcp_draft_sharding is supports_pcp
+    assert (
+        speculator.vllm_config.parallel_config.prefill_context_parallel_size
+        == expected_draft_pcp_size
+    )
+    assert config.parallel_config.prefill_context_parallel_size == 4
+
+
+def test_replicated_pcp_prefill_builds_global_draft_metadata(monkeypatch):
+    speculator = object.__new__(_TestSpeculator)
+    slot_mappings = torch.tensor([[7, 8, 9]])
+    speculator.block_tables = SimpleNamespace(
+        gather_block_tables=Mock(),
+        compute_slot_mappings=Mock(return_value=slot_mappings),
+    )
+    speculator.kv_cache_config = object()
+    speculator._build_attn_metadata = Mock(return_value={"draft": "metadata"})
+    monkeypatch.setattr(
+        spec_module,
+        "build_slot_mappings_by_layer",
+        lambda mappings, kv_config: {"draft.layer": mappings},
+    )
+    input_batch = SimpleNamespace(
+        num_reqs=2,
+        idx_mapping=torch.tensor([3, 1]),
+        query_start_loc=torch.tensor([0, 2, 3]),
+        query_start_loc_np=torch.tensor([0, 2, 3]).numpy(),
+        positions=torch.tensor([5, 6, 9]),
+        seq_lens_cpu_upper_bound=torch.tensor([7, 10]),
+    )
+    batch_desc = BatchExecutionDescriptor(
+        cg_mode=CUDAGraphMode.PIECEWISE,
+        num_tokens=4,
+        num_reqs=4,
+    )
+
+    metadata, mappings_by_layer = speculator._prepare_replicated_pcp_prefill_attn(
+        input_batch, batch_desc
+    )
+
+    speculator.block_tables.gather_block_tables.assert_called_once_with(
+        input_batch.idx_mapping, num_reqs_padded=4
+    )
+    speculator.block_tables.compute_slot_mappings.assert_called_once_with(
+        input_batch.idx_mapping,
+        input_batch.query_start_loc,
+        input_batch.positions,
+        4,
+    )
+    speculator._build_attn_metadata.assert_called_once_with(
+        num_reqs=2,
+        batch_desc=batch_desc,
+        query_start_loc_np=input_batch.query_start_loc_np,
+        seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound,
+        step=0,
+    )
+    assert metadata == {"draft": "metadata"}
+    assert torch.equal(mappings_by_layer["draft.layer"], slot_mappings)
+
+
 @pytest.mark.parametrize(("hc_mult", "expected"), [(None, 64), (4, 256)])
 def test_speculator_uses_draft_model_hidden_size(monkeypatch, hc_mult, expected):
     # Qwen4Exp targets expose multi-stream HC residuals to the drafter.

vllm/v1/worker/gpu/model_runner.py

@@ -876,14 +876,36 @@ def _dummy_run(
                     ),
                 )
 
+            draft_hidden_states = hidden_states
+            assert draft_hidden_states is not None
+            draft_uses_pcp_sharding = (
+                isinstance(self.speculator, DraftModelSpeculator)
+                and self.speculator.uses_pcp_draft_sharding
+            )
+            if self.pcp_manager is not None and not draft_uses_pcp_sharding:
+                hidden_states, input_batch = self.pcp_manager.restore_for_sampling(
+                    hidden_states
+                )
+                if aux_hidden_states is not None:
+                    aux_hidden_states = [
+                        self.pcp_manager.restore_hidden_states(states)
+                        for states in aux_hidden_states
+                    ]
+
             # Let the target override the hidden state fed to the drafter
             # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The
             # target returns a persistent buffer sized at max_num_batched_tokens;
             # slice to the active token count that propose() expects.
             spec_hidden_states = hidden_states
             if hasattr(self.model, "get_mtp_target_hidden_states"):
                 pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
-                spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]]  # type: ignore[union-attr]
+                spec_hidden_states = pre_hc_hidden_states[
+                    : draft_hidden_states.shape[0]
+                ]
+                if self.pcp_manager is not None and not draft_uses_pcp_sharding:
+                    spec_hidden_states = self.pcp_manager.restore_hidden_states(
+                        spec_hidden_states
+                    )
             if isinstance(self.sampler, GPUWatermarkSampler):
                 self.speculator.prepare_watermarking(
                     self.sampler._get_contexts(input_batch.idx_mapping),
@@ -2128,10 +2150,20 @@ def sample_tokens(
             # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The
             # target returns a persistent buffer sized at max_num_batched_tokens;
             # slice to the active token count that propose() expects.
-            spec_hidden_states = draft_hidden_states
+            draft_uses_pcp_sharding = (
+                isinstance(self.speculator, DraftModelSpeculator)
+                and self.speculator.uses_pcp_draft_sharding
+            )
+            spec_hidden_states = (
+                draft_hidden_states if draft_uses_pcp_sharding else hidden_states
+            )
             if hasattr(self.model, "get_mtp_target_hidden_states"):
                 pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
                 spec_hidden_states = pre_hc_hidden_states[: draft_hidden_states.size(0)]
+                if self.pcp_manager is not None and not draft_uses_pcp_sharding:
+                    spec_hidden_states = self.pcp_manager.restore_hidden_states(
+                        spec_hidden_states
+                    )
             if isinstance(self.sampler, GPUWatermarkSampler):
                 self.speculator.prepare_watermarking(
                     self.sampler._get_contexts(input_batch.idx_mapping),

vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py

@@ -167,12 +167,17 @@ def capture(self) -> None:
             self.prefill_cudagraph_manager.init_breakable_cg_runner(self.model)
 
         self.on_prefill_begin(self.max_num_reqs)
+        capture_input_buffers = self.target_input_buffers
+        capture_attn_groups = self.target_attn_groups
+        if self.target_pcp_size > 1 and not self.uses_pcp_draft_sharding:
+            capture_input_buffers = self.input_buffers
+            capture_attn_groups = self.attn_groups
         self.prefill_cudagraph_manager.capture(
             self._prefill,
             self.model_state,
-            self.target_input_buffers,
+            capture_input_buffers,
             self.block_tables,
-            self.target_attn_groups,
+            capture_attn_groups,
             self.kv_cache_config,
             progress_bar_desc="Capturing prefill CUDA graphs",
         )
@@ -200,6 +205,37 @@ def capture(self) -> None:
         )
         self.on_multi_step_decode_end(self.max_num_reqs)
 
+    @torch.inference_mode()
+    def _prepare_replicated_pcp_prefill_attn(
+        self,
+        input_batch: InputBatch,
+        batch_desc: BatchExecutionDescriptor,
+    ) -> tuple[dict[str, Any] | None, dict[str, torch.Tensor]]:
+        """Build global draft metadata when the target metadata is sharded."""
+        num_reqs = input_batch.num_reqs
+        num_reqs_padded = batch_desc.num_reqs or num_reqs
+        self.block_tables.gather_block_tables(
+            input_batch.idx_mapping,
+            num_reqs_padded=num_reqs_padded,
+        )
+        slot_mappings = self.block_tables.compute_slot_mappings(
+            input_batch.idx_mapping,
+            input_batch.query_start_loc,
+            input_batch.positions,
+            batch_desc.num_tokens,
+        )
+        slot_mappings_by_layer = build_slot_mappings_by_layer(
+            slot_mappings, self.kv_cache_config
+        )
+        attn_metadata = self._build_attn_metadata(
+            num_reqs=num_reqs,
+            batch_desc=batch_desc,
+            query_start_loc_np=input_batch.query_start_loc_np,
+            seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound,
+            step=0,
+        )
+        return attn_metadata, slot_mappings_by_layer
+
     @torch.inference_mode()
     def propose(
         self,
@@ -270,7 +306,7 @@ def propose(
             self.max_num_reqs,
         )
 
-        if self.pcp_manager is not None:
+        if self.pcp_manager is not None and self.uses_pcp_draft_sharding:
             self.pcp_manager.prepare_draft_prefill(
                 input_batch,
                 self.input_buffers.input_ids[:num_tokens_padded],
@@ -279,7 +315,13 @@ def propose(
             if prefill is not None:
                 num_tokens = prefill.num_tokens
                 num_tokens_padded = prefill.num_tokens_after_padding
-        self.hidden_states[:num_tokens_padded].copy_(hidden_states)
+        self.hidden_states[:num_tokens].copy_(hidden_states[:num_tokens])
+        if num_tokens_padded > num_tokens:
+            self.hidden_states[num_tokens:num_tokens_padded].zero_()
+        if self.target_pcp_size > 1 and not self.uses_pcp_draft_sharding:
+            self.input_buffers.is_padding[:num_tokens_padded].copy_(
+                input_batch.is_padding[:num_tokens_padded]
+            )
 
         # When all requests are decoding (no true prefills), each has
         # num_speculative_steps + 1 tokens, enabling FULL graph replay.
@@ -307,6 +349,18 @@ def propose(
             else None
         )
 
+        prefill_attn_metadata: dict[str, Any] | None = attn_metadata
+        prefill_slot_mappings: dict[str, torch.Tensor] | None = slot_mappings
+        if self.target_pcp_size > 1 and not self.uses_pcp_draft_sharding:
+            if dummy_run and skip_attn_for_dummy_run:
+                prefill_attn_metadata, prefill_slot_mappings = None, None
+            else:
+                prefill_attn_metadata, prefill_slot_mappings = (
+                    self._prepare_replicated_pcp_prefill_attn(
+                        input_batch, prefill_batch_desc
+                    )
+                )
+
         self._prepare_eplb_forward(num_tokens)
 
         self.on_prefill_begin(num_reqs)
@@ -321,8 +375,8 @@ def propose(
             self._prefill(
                 num_reqs,
                 prefill_batch_desc.num_tokens,
-                attn_metadata,
-                slot_mappings,
+                prefill_attn_metadata,
+                prefill_slot_mappings,
                 num_tokens_across_dp=num_tokens_across_dp,
                 cudagraph_runtime_mode=prefill_batch_desc.cg_mode,
                 mm_inputs=mm_inputs,
@@ -403,8 +457,9 @@ def _run_model(
     ) -> tuple[torch.Tensor, torch.Tensor]:
         input_buffers: InputBatch | InputBuffers = self.input_buffers
         is_padding = None
-        if self.pcp_manager is not None:
+        if self.pcp_manager is not None and self.uses_pcp_draft_sharding:
             input_buffers = self.pcp_manager.get_draft_input_buffers(self.input_buffers)
+        if self.target_pcp_size > 1:
             is_padding = input_buffers.is_padding[:num_tokens]
         batch_descriptor = BatchDescriptor(num_tokens=num_tokens)
         with set_forward_context(
@@ -483,7 +538,7 @@ def _prefill(
             cudagraph_runtime_mode=cudagraph_runtime_mode,
             mm_inputs=mm_inputs,
         )
-        if self.pcp_manager is not None:
+        if self.pcp_manager is not None and self.uses_pcp_draft_sharding:
             last_hidden_states, hidden_states = self.pcp_manager.restore_draft_prefill(
                 last_hidden_states, hidden_states
             )