[Bugfix] Support MTP speculative decoding with pipeline parallelism on the PD producer side

vllm/config/model.py

@@ -1198,8 +1198,15 @@ def verify_with_parallel_config(
             self._verify_with_expert_parallelism()
 
         pipeline_parallel_size = parallel_config.pipeline_parallel_size
-        if pipeline_parallel_size > 1 and not self.registry.is_pp_supported_model(
-            self.architectures, self
+        # Speculative drafters (Eagle, MTP, etc.) are loaded locally on the
+        # last pipeline stage rather than partitioned across all PP ranks.
+        # Skip the PP support check for these models since they run with an
+        # effective PP size of 1.
+        is_local_drafter = self.runner_type == "draft"
+        if (
+            pipeline_parallel_size > 1
+            and not is_local_drafter
+            and not self.registry.is_pp_supported_model(self.architectures, self)
         ):
             raise NotImplementedError(
                 "Pipeline parallelism is not supported for this model. "

vllm/config/vllm.py

@@ -844,6 +844,35 @@ def _verify_kv_transfer_compat(self) -> None:
             "expandable_segments is automatically disabled)."
         )
 
+    def _validate_spec_decode_pp_config(self) -> None:
+        """Validate speculative decoding + pipeline parallelism combinations.
+
+        MTP-style speculative decoding with PP > 1 is only supported on the
+        prefill (producer) side of PD-disaggregated deployments. The decode
+        (consumer) side must use PP=1 because the draft verification loop
+        cannot coordinate draft token propagation across pipeline stages.
+        """
+        if (
+            self.speculative_config is None
+            or self.speculative_config.method != "mtp"
+            or self.parallel_config.pipeline_parallel_size <= 1
+        ):
+            return
+
+        if (
+            self.kv_transfer_config is not None
+            and self.kv_transfer_config.is_kv_producer
+        ):
+            return
+
+        raise ValueError(
+            "MTP speculative decoding with pipeline_parallel_size > 1 is only "
+            "supported on the prefill (producer) side of PD-disaggregated "
+            "deployments (kv_role='kv_producer'). The decode (consumer) side "
+            "must use pipeline_parallel_size=1; combine data parallelism with "
+            "MTP instead."
+        )
+
     def __post_init__(self):
         """Verify configs are valid & consistent with each other."""
 
@@ -1538,6 +1567,7 @@ def has_blocked_weights():
             if "-quant_fp8" not in custom_ops:
                 custom_ops.append("+quant_fp8")
 
+        self._validate_spec_decode_pp_config()
         self._verify_kv_transfer_compat()
         # Log the custom passes that are enabled
         self.compilation_config.pass_config.log_enabled_passes()

vllm/model_executor/models/deepseek_eagle3.py

@@ -321,9 +321,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             base_vocab_size = getattr(self.config, "vocab_size", None)
             self.config.draft_vocab_size = base_vocab_size
 
-        target_layer_num = vllm_config.model_config.get_num_layers(
-            vllm_config.parallel_config
-        )
+        target_layer_num = vllm_config.model_config.get_total_num_hidden_layers()
 
         # Store target layer count in draft config
         self.config.target_layer_count = target_layer_num

vllm/model_executor/models/llama_eagle3.py

@@ -297,9 +297,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         if getattr(self.config, "draft_vocab_size", None) is None:
             base_vocab_size = getattr(self.config, "vocab_size", None)
             self.config.draft_vocab_size = base_vocab_size
-        target_layer_num = vllm_config.model_config.get_num_layers(
-            vllm_config.parallel_config
-        )
+        target_layer_num = vllm_config.model_config.get_total_num_hidden_layers()
 
         # Store target layer count in draft config for
         # proper layer_types indexing in draft models

vllm/model_executor/models/qwen3_5_mtp.py

@@ -129,19 +129,17 @@ def forward(
         inputs_embeds: torch.Tensor | None = None,
         spec_step_idx: int = 0,
     ) -> torch.Tensor:
-        if get_pp_group().is_first_rank:
-            if inputs_embeds is None:
-                inputs_embeds = self.embed_input_ids(input_ids)
-            assert hidden_states.shape[-1] == inputs_embeds.shape[-1]
-            inputs_embeds = self.pre_fc_norm_embedding(inputs_embeds)
-            hidden_states = self.pre_fc_norm_hidden(hidden_states)
-            hidden_states = torch.cat([inputs_embeds, hidden_states], dim=-1)
-            hidden_states = self.fc(hidden_states)
-            residual = None
-        else:
-            assert intermediate_tensors is not None
-            hidden_states = intermediate_tensors["hidden_states"]
-            residual = intermediate_tensors["residual"]
+        # The MTP drafter is loaded locally on the last PP stage and always
+        # combines token embeddings with the target hidden states.
+        # It does not consume PP intermediate tensors from previous stages.
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_input_ids(input_ids)
+        assert hidden_states.shape[-1] == inputs_embeds.shape[-1]
+        inputs_embeds = self.pre_fc_norm_embedding(inputs_embeds)
+        hidden_states = self.pre_fc_norm_hidden(hidden_states)
+        hidden_states = torch.cat([inputs_embeds, hidden_states], dim=-1)
+        hidden_states = self.fc(hidden_states)
+        residual = None
 
         current_step_idx = spec_step_idx % self.num_mtp_layers
         hidden_states, residual = self.layers[current_step_idx](

vllm/v1/worker/gpu_model_runner.py

@@ -6805,9 +6805,14 @@ def initialize_metadata_builders(
         self.calculate_reorder_batch_threshold()
 
         # Initialize drafter attention backend
-        if self.speculative_config and (
-            self.speculative_config.use_eagle()
-            or self.speculative_config.uses_draft_model()
+        # Drafter is only loaded on the last PP rank, so skip on other ranks.
+        if (
+            self.speculative_config
+            and get_pp_group().is_last_rank
+            and (
+                self.speculative_config.use_eagle()
+                or self.speculative_config.uses_draft_model()
+            )
         ):
             assert isinstance(
                 self.drafter,
@@ -6858,9 +6863,14 @@ def _check_and_update_cudagraph_mode(
         )
 
         # Initialize drafter's cudagraph dispatcher if using spec decode.
-        if self.speculative_config and (
-            self.speculative_config.use_eagle()
-            or self.speculative_config.uses_extract_hidden_states()
+        # Drafter is only loaded on the last PP rank, so skip on other ranks.
+        if (
+            self.speculative_config
+            and get_pp_group().is_last_rank
+            and (
+                self.speculative_config.use_eagle()
+                or self.speculative_config.uses_extract_hidden_states()
+            )
         ):
             assert isinstance(
                 self.drafter,