[Spec][V2] Support MTP speculative decoding under pipeline parallelism

tests/v1/worker/test_pp_utils.py

@@ -0,0 +1,201 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the PPHandler sampled-token / draft-token relay under PP."""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.v1.worker.gpu import pp_utils
+from vllm.v1.worker.gpu.pp_utils import PPHandler
+
+requires_cuda = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="PPHandler drives a side CUDA stream"
+)
+
+
+def make_handler(
+    monkeypatch,
+    *,
+    is_last_rank: bool,
+    num_speculative_steps: int,
+    relay_draft_tokens: bool,
+    world_size: int = 2,
+) -> PPHandler:
+    """Build a real PPHandler with the PP group stubbed out."""
+    pp_group = SimpleNamespace(
+        is_last_rank=is_last_rank,
+        last_rank=world_size - 1,
+        world_size=world_size,
+        make_sibling_device_group=lambda group_desc: object(),
+    )
+    monkeypatch.setattr(pp_utils, "get_pp_group", lambda: pp_group)
+    return PPHandler(
+        max_num_reqs=8,
+        num_speculative_steps=num_speculative_steps,
+        device=torch.device("cuda"),
+        relay_draft_tokens=relay_draft_tokens,
+    )
+
+
+def record_broadcasts(monkeypatch) -> list[torch.Tensor]:
+    """Capture every tensor handed to the collective, in call order."""
+    calls: list[torch.Tensor] = []
+    monkeypatch.setattr(
+        torch.distributed, "broadcast", lambda t, src, group: calls.append(t)
+    )
+    return calls
+
+
+def make_input_batch(num_reqs: int = 3, *, needs_sample: bool = True):
+    # compute_need_sampled_mask only reads these fields. With needs_sample=False
+    # every request is already at max_seq_len, so no sample is needed next step.
+    return SimpleNamespace(
+        num_reqs=num_reqs,
+        num_computed_tokens_np=np.zeros(num_reqs, dtype=np.int32),
+        prefill_len_np=np.full(num_reqs, 4, dtype=np.int32),
+        num_scheduled_tokens=np.full(num_reqs, 4, dtype=np.int32),
+        max_seq_len_np=np.full(num_reqs, 100 if needs_sample else 1, dtype=np.int32),
+        idx_mapping=torch.arange(num_reqs, device="cuda"),
+        idx_mapping_np=np.arange(num_reqs, dtype=np.int32),
+    )
+
+
+def send_step(handler, input_batch, *, width: int, with_draft: bool):
+    num_reqs = input_batch.num_reqs
+    sampled = torch.zeros(num_reqs, width, dtype=torch.int64, device="cuda")
+    counts = torch.zeros(num_reqs, dtype=torch.int32, device="cuda")
+    handler.broadcast(sampled, counts, counts, input_batch)
+    if with_draft:
+        draft = torch.zeros(
+            num_reqs, handler.max_sample_len - 1, dtype=torch.int64, device="cuda"
+        )
+        handler.broadcast_draft(draft, input_batch)
+
+
+# ---------------------------------------------------------------------------
+# broadcast() pads so send/recv element counts match on every step
+# ---------------------------------------------------------------------------
+
+
+@requires_cuda
+@pytest.mark.parametrize("width,num_spec", [(1, 1), (1, 3), (2, 3)])
+def test_broadcast_pads_sampled_tokens_to_max_sample_len(monkeypatch, width, num_spec):
+    """The sampler emits width 1 on steps with no draft tokens (prefill, first
+    decode) and num_spec+1 only after rejection sampling. The receiver always
+    posts a max_sample_len buffer, so an unpadded send is a count mismatch."""
+    handler = make_handler(
+        monkeypatch,
+        is_last_rank=True,
+        num_speculative_steps=num_spec,
+        relay_draft_tokens=True,
+    )
+    calls = record_broadcasts(monkeypatch)
+    input_batch = make_input_batch()
+
+    send_step(handler, input_batch, width=width, with_draft=False)
+
+    sent_sampled = calls[0]
+    assert sent_sampled.shape == (input_batch.num_reqs, handler.max_sample_len)
+    # Placeholder columns are ignored by post_update, which advances each
+    # request by its own num_sampled count.
+    assert (sent_sampled[:, width:] == -1).all()
+    assert (sent_sampled[:, :width] == 0).all()
+
+
+# ---------------------------------------------------------------------------
+# Sender and receiver must post the same number of collectives per step
+# ---------------------------------------------------------------------------
+
+
+@requires_cuda
+def test_send_and_recv_op_counts_match_with_speculator(monkeypatch):
+    """With a speculator the step is three broadcasts: sampled, combined, draft."""
+    sender = make_handler(
+        monkeypatch, is_last_rank=True, num_speculative_steps=3, relay_draft_tokens=True
+    )
+    calls = record_broadcasts(monkeypatch)
+    send_step(sender, make_input_batch(), width=1, with_draft=True)
+    assert len(calls) == 3
+
+    receiver = make_handler(
+        monkeypatch,
+        is_last_rank=False,
+        num_speculative_steps=3,
+        relay_draft_tokens=True,
+    )
+    calls.clear()
+    assert receiver.receive(make_input_batch())
+    assert len(calls) == 3
+    assert calls[2].shape == (3, sender.max_sample_len - 1)
+
+
+@requires_cuda
+def test_send_and_recv_op_counts_match_without_speculator(monkeypatch):
+    """Diffusion LLMs set num_speculative_steps > 0 but have no speculator, so
+    the last rank never relays draft tokens. Gating the receiver's third recv on
+    num_speculative_steps instead of on the speculator hangs the non-last ranks
+    waiting for a broadcast that is never issued."""
+    sender = make_handler(
+        monkeypatch,
+        is_last_rank=True,
+        num_speculative_steps=3,
+        relay_draft_tokens=False,
+    )
+    calls = record_broadcasts(monkeypatch)
+    send_step(sender, make_input_batch(), width=1, with_draft=False)
+    assert len(calls) == 2
+
+    receiver = make_handler(
+        monkeypatch,
+        is_last_rank=False,
+        num_speculative_steps=3,
+        relay_draft_tokens=False,
+    )
+    calls.clear()
+    assert receiver.receive(make_input_batch())
+    assert len(calls) == 2
+    assert receiver.queue[-1].draft_tokens is None
+
+
+@requires_cuda
+def test_both_ranks_skip_when_no_request_needs_sampling(monkeypatch):
+    """The skip gate must be symmetric, or the ranks desynchronize."""
+    sender = make_handler(
+        monkeypatch, is_last_rank=True, num_speculative_steps=3, relay_draft_tokens=True
+    )
+    calls = record_broadcasts(monkeypatch)
+    send_step(sender, make_input_batch(needs_sample=False), width=1, with_draft=True)
+    assert calls == []
+
+    receiver = make_handler(
+        monkeypatch,
+        is_last_rank=False,
+        num_speculative_steps=3,
+        relay_draft_tokens=True,
+    )
+    calls.clear()
+    assert not receiver.receive(make_input_batch(needs_sample=False))
+    assert calls == []
+
+
+# ---------------------------------------------------------------------------
+# Relayed draft tokens survive the deferred consume
+# ---------------------------------------------------------------------------
+
+
+# ---------------------------------------------------------------------------
+# DeepSeekMTP under pipeline parallelism
+# ---------------------------------------------------------------------------
+
+
+def test_deepseek_mtp_passes_supports_pp_gate():
+    """DeepSeekMTP must pass the supports_pp() gate used at model resolution;
+    otherwise the engine refuses to build it under pipeline parallelism. The
+    gate covers both the SupportsPP MRO entry and the forward() signature."""
+    from vllm.model_executor.models.deepseek_mtp import DeepSeekMTP
+    from vllm.model_executor.models.interfaces import supports_pp
+
+    assert supports_pp(DeepSeekMTP)

vllm/model_executor/models/deepseek_mtp.py

@@ -37,9 +37,11 @@
     DeepseekV2MoE,
     _try_load_fp8_indexer_wk,
 )
+from .interfaces import SupportsPP
 from .utils import (
     get_pp_missing_layer_names,
     get_spec_layer_idx_from_weight_name,
+    make_empty_intermediate_tensors_factory,
     maybe_prefix,
 )
 
@@ -228,7 +230,7 @@ def compute_logits(
 
 
 @support_torch_compile
-class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
+class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts, SupportsPP):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
         self.config = vllm_config.model_config.hf_config
@@ -238,6 +240,12 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         )
         # Set MoE hyperparameters
         self.set_moe_parameters()
+        # PP support: the MTP draft runs only on the last PP stage (the runner gates
+        # drafter construction on get_pp_group().is_last_rank), so it never actually
+        # consumes PP intermediate tensors — but SupportsPP requires this factory.
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+            ["hidden_states", "residual"], self.config.hidden_size
+        )
 
     def set_moe_parameters(self):
         self.num_moe_layers = self.config.num_nextn_predict_layers
@@ -322,6 +330,24 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         for name, loaded_weight in weights:
             if "rotary_emb.inv_freq" in name:
                 continue
+            # Load the draft's own embed_tokens. The MTP module embeds the
+            # tokens it drafts via its own embed_tokens, but the checkpoint
+            # stores embed_tokens as a single top-level (shared/tied) weight
+            # whose spec_layer is None -- so the loop below skips it. Under
+            # pipeline parallelism the target model's embed_tokens is a
+            # PPMissingLayer on the draft's last stage, so the draft cannot
+            # borrow it and MUST load its own copy here, or it embeds tokens
+            # with uninitialized weights and produces garbage drafts.
+            if "embed_tokens" in name:
+                param = params_dict.get(name)
+                if param is not None:
+                    weight_loader = getattr(param, "weight_loader", None)
+                    if weight_loader is not None:
+                        weight_loader(param, loaded_weight)
+                    else:
+                        param.data.copy_(loaded_weight)
+                    loaded_params.add(name)
+                continue
             spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
             if spec_layer is None:
                 continue

vllm/model_executor/models/qwen3_5_mtp.py

@@ -36,6 +36,7 @@
 from .interfaces import (
     MultiModalEmbeddings,
     SupportsMultiModal,
+    SupportsPP,
     _require_is_multimodal,
 )
 from .utils import (
@@ -145,7 +146,13 @@ def forward(
         inputs_embeds: torch.Tensor | None = None,
         spec_step_idx: int = 0,
     ) -> torch.Tensor:
-        if get_pp_group().is_first_rank:
+        pp_group = get_pp_group()
+        if pp_group.is_first_rank or pp_group.is_last_rank:
+            # The drafter is only built on the last PP stage, so the target's
+            # hidden states are available locally there and the fc projection
+            # applies exactly as it does at PP=1. Without the last-rank case the
+            # draft would fall through to the intermediate-tensor path and
+            # project uninitialized state.
             if inputs_embeds is None:
                 inputs_embeds = self.embed_input_ids(input_ids)
             assert hidden_states.shape[-1] == inputs_embeds.shape[-1]
@@ -155,6 +162,7 @@ def forward(
             hidden_states = self.fc(hidden_states)
             residual = None
         else:
+            # Middle PP rank: use intermediate tensors from previous rank.
             assert intermediate_tensors is not None
             hidden_states = intermediate_tensors["hidden_states"]
             residual = intermediate_tensors["residual"]
@@ -208,7 +216,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         "hidden_states": 0,
     }
 )
-class Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal):
+class Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal, SupportsPP):
     packed_modules_mapping = {
         "qkv_proj": [
             "q_proj",
@@ -250,6 +258,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
 
         self.logits_processor = LogitsProcessor(config.vocab_size)
 
+        self.make_empty_intermediate_tensors = (
+            self.model.make_empty_intermediate_tensors
+        )
+
     def embed_input_ids(
         self,
         input_ids: torch.Tensor,
@@ -285,6 +297,12 @@ def forward(
         inputs_embeds: torch.Tensor | None = None,
         **kwargs: object,
     ):
+        if not get_pp_group().is_first_rank and intermediate_tensors is None:
+            intermediate_tensors = self.make_empty_intermediate_tensors(
+                batch_size=hidden_states.shape[0],
+                dtype=hidden_states.dtype,
+                device=hidden_states.device,
+            )
         hidden_states = self.model(
             input_ids, positions, hidden_states, intermediate_tensors, inputs_embeds
         )

vllm/models/deepseek_v32/nvidia/model.py

@@ -174,12 +174,13 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         # DSA is always sparse (has index_topk); allocate the shared top-k
         # buffer the indexer writes and the sparse MLA backend reads.
         self.is_v32 = True
-        topk_indices_buffer = torch.empty(
+        self.topk_indices_buffer = torch.empty(
             vllm_config.scheduler_config.max_num_batched_tokens,
             config.index_topk,
             dtype=torch.int32,
             device=self.device,
         )
+        topk_indices_buffer = self.topk_indices_buffer
 
         if get_pp_group().is_first_rank:
             self.embed_tokens = VocabParallelEmbedding(

vllm/v1/attention/backends/mla/flashattn_mla_sparse.py

@@ -202,6 +202,10 @@ def __init__(
             topk_indices_buffer=topk_indices_buffer,
             **mla_args,
         )
+        # Keep the indexer, not just the buffer the base class snapshots from it:
+        # under MTP+PP the indexer swaps its buffer between the target and draft
+        # passes, so forward_mqa has to re-read it rather than use the snapshot.
+        self._indexer = indexer
         assert self.topk_indices_buffer is not None, (
             "Indexer or topk_indices_buffer required for sparse MLA"
         )
@@ -221,8 +225,13 @@ def forward_mqa(
         q_nope, q_rope = q
         num_actual_toks = q_rope.shape[0]
 
-        assert self.topk_indices_buffer is not None
-        topk_indices = self.topk_indices_buffer[:num_actual_toks]
+        buf = (
+            self._indexer.topk_indices_buffer
+            if self._indexer is not None
+            else self.topk_indices_buffer
+        )
+        assert buf is not None, "topk_indices_buffer required for sparse MLA"
+        topk_indices = buf[:num_actual_toks]
         topk_indices, valid_counts = triton_convert_req_index_to_global_index(
             attn_metadata.req_id_per_token[:num_actual_toks],
             attn_metadata.block_table,

vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py

@@ -367,6 +367,10 @@ def __init__(
             topk_indices_buffer=topk_indices_buffer,
             **mla_args,
         )
+        # Keep the indexer, not just the buffer the base class snapshots from it:
+        # under MTP+PP the indexer swaps its buffer between the target and draft
+        # passes, so forward_mqa has to re-read it rather than use the snapshot.
+        self._indexer = indexer
 
         self._workspace_buffer: torch.Tensor | None = None
         self.bmm1_scale: float | None = None
@@ -389,8 +393,13 @@ def forward_mqa(
 
         num_actual_toks = q.shape[0]
 
-        assert self.topk_indices_buffer is not None
-        topk_indices = self.topk_indices_buffer[:num_actual_toks]
+        buf = (
+            self._indexer.topk_indices_buffer
+            if self._indexer is not None
+            else self.topk_indices_buffer
+        )
+        assert buf is not None, "topk_indices_buffer required for sparse MLA"
+        topk_indices = buf[:num_actual_toks]
 
         if self.dcp_world_size > 1:
             topk_indices_physical, seq_lens = triton_filter_and_convert_dcp_index(