Support DSpark MTP under pipeline parallelism on RDNA2

Three guards refused the drafter under PP -- the draft config inherited the
target's pipeline size, so a draft model without SupportsPP was rejected; the
V2 runner raised for eagle3/dflash/dspark because their aux taps may sit on an
earlier stage; and the DSpark loader aliases the target's embed_tokens, which
is a PPMissingLayer anywhere but the first stage. Past those, warmup deadlocked:
PPHandler.receive always allocates [num_reqs, max_sample_len] while the sampler
returns a single column when there are no drafts to verify, so broadcast sent a
narrower tensor, the ranks disagreed on the element count, and the stage that
had moved on wedged in its next device sync. Give the draft
pipeline_parallel_size=1 since it is built on the last stage only, let a model
declare that it carries its EAGLE3 aux taps across the handoff, build the
target's embedding on the last stage rather than draft with uninitialized
weights, pin one broadcast payload width on both sides and send the next step's
drafts in it through the filtered index mapping, and declare the AMD DSpark
head's draft_id_to_target_id. Decode TPOT 43.9 -> 35.2 ms and output throughput
20.0 -> 24.3 tok/s against the same runner without speculation, at 19.8% draft
acceptance.

Upstream adds pipeline-parallel support for these drafters in vllm-project#50514; drop this
commit when rebasing onto a tag that already contains it.

Author: yiminyuan

tests/v1/spec_decode/test_aux_hidden_state_pp.py

@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Aux hidden-state transport across pipeline stages.
+
+EAGLE3-style drafters (eagle3, dflash, dspark) consume auxiliary hidden states
+tapped from target layers. Under pipeline parallelism those layers can sit on an
+earlier stage than the drafter, which runs on the last rank, so the taps have to
+ride the intermediate-tensor handoff. Slot numbering is derived independently on
+every stage from the layer split, so these tests pin that the derivations agree
+-- a stage disagreeing by one silently feeds the drafter the wrong tap.
+"""
+
+import pytest
+import torch
+
+from vllm.model_executor.models.interfaces import EagleModelMixin
+from vllm.sequence import IntermediateTensors
+
+
+class _Config:
+    def __init__(self, num_hidden_layers: int):
+        self.num_hidden_layers = num_hidden_layers
+
+
+class _Model(EagleModelMixin):
+    supports_aux_hidden_states_over_pp = True
+
+    def __init__(self, num_hidden_layers: int, aux_layers: tuple[int, ...]):
+        self.config = _Config(num_hidden_layers)
+        self.aux_hidden_state_layers = aux_layers
+
+
+@pytest.mark.parametrize(
+    "start,end,aux_ids,is_first,expected",
+    [
+        # First stage also taps its own start layer, as the runner does.
+        (0, 4, (0, 2, 4), True, (0, 2, 4)),
+        (0, 4, (0, 2, 4), False, (2, 4)),
+        # A later stage only claims taps whose layer it owns.
+        (4, 8, (0, 2, 4), False, ()),
+        (4, 8, (5, 8), False, (5, 8)),
+        (4, 8, (), False, ()),
+    ],
+)
+def test_local_aux_tap_ids(start, end, aux_ids, is_first, expected):
+    assert EagleModelMixin.local_aux_tap_ids(start, end, aux_ids, is_first) == expected
+
+
+def test_every_tap_is_claimed_by_exactly_one_stage():
+    """The per-stage split must partition the taps, with none lost or doubled."""
+    num_layers, pp = 43, 2
+    aux = (0, 21, 42)
+    model = _Model(num_layers, aux)
+
+    from vllm.distributed.utils import get_pp_indices
+
+    claimed: list[int] = []
+    for rank in range(pp):
+        start, end = get_pp_indices(num_layers, rank, pp)
+        claimed.extend(
+            EagleModelMixin.local_aux_tap_ids(start, end, aux, rank == 0)
+        )
+    assert sorted(claimed) == sorted(aux)
+    assert len(claimed) == len(set(claimed))
+    # And the totals line up with the slot arithmetic.
+    assert sum(model._num_local_taps_on_rank(r, pp) for r in range(pp)) == len(aux)
+
+
+def test_slot_numbering_is_contiguous_and_agreed():
+    """Each rank's base is the count of taps produced by all earlier ranks."""
+    num_layers, pp = 43, 2
+    model = _Model(num_layers, (0, 21, 42))
+
+    bases = [model._aux_slot_base(r, pp) for r in range(pp)]
+    assert bases[0] == 0
+    for r in range(1, pp):
+        assert bases[r] == bases[r - 1] + model._num_local_taps_on_rank(r - 1, pp)
+
+
+def test_pack_local_aux_keys_by_global_slot():
+    model = _Model(43, (0, 21, 42))
+    model._aux_slot_base_cached = 2
+    taps = [torch.zeros(1), torch.ones(1)]
+
+    packed = model.pack_local_aux_for_last(taps)
+
+    assert sorted(packed) == ["aux_hidden_states_2", "aux_hidden_states_3"]
+    assert torch.equal(packed["aux_hidden_states_2"], taps[0])
+    assert torch.equal(packed["aux_hidden_states_3"], taps[1])
+    # Nothing to say when this stage owns no taps.
+    assert model.pack_local_aux_for_last([]) == {}
+
+
+def test_recv_remote_aux_returns_producer_order():
+    model = _Model(43, (0, 21, 42))
+    model._aux_upstream_total_cached = 2
+    a, b = torch.zeros(1), torch.ones(1)
+    it = IntermediateTensors(
+        {"aux_hidden_states_0": a, "aux_hidden_states_1": b, "hidden_states": a}
+    )
+
+    got = model.recv_remote_aux_from_producers(it)
+
+    assert len(got) == 2
+    assert torch.equal(got[0], a)
+    assert torch.equal(got[1], b)
+
+
+def test_recv_remote_aux_raises_rather_than_zero_filling():
+    """A missing slot must fail loudly; zeros would only cost acceptance."""
+    model = _Model(43, (0, 21, 42))
+    model._aux_upstream_total_cached = 2
+    it = IntermediateTensors({"aux_hidden_states_0": torch.zeros(1)})
+
+    with pytest.raises(RuntimeError, match="aux_hidden_states_1 missing"):
+        model.recv_remote_aux_from_producers(it)
+
+
+def test_no_upstream_taps_needs_no_intermediate_tensors():
+    """At PP=1, or when every tap is local, the receive side is a no-op."""
+    model = _Model(43, (0, 21, 42))
+    model._aux_upstream_total_cached = 0
+    assert model.recv_remote_aux_from_producers(None) == []

tests/v1/worker/test_pp_sampled_token_width.py

@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The PP sampled-token broadcast must agree on width between the two sides.
+
+`PPHandler.receive` allocates `[num_reqs, max_sample_len + draft_token_width]`
+unconditionally, while the sampler hands `broadcast` a single column whenever
+there were no draft tokens to verify -- which is the case throughout warmup.
+Broadcasting the narrower tensor makes the ranks disagree on the element count,
+and the collective then never completes: the last rank sails on while the other
+blocks in the next device sync with its GPU spinning.
+
+Without speculation `max_sample_len` is 1 and the two widths coincide, so this
+only bites with MTP enabled, which is why it is worth pinning here.
+"""
+
+import pytest
+import torch
+
+
+def _payload_width(max_sample_len: int, draft_token_width: int) -> int:
+    return max_sample_len + draft_token_width
+
+
+def _sender_payload(
+    sampled_token_ids: torch.Tensor,
+    max_sample_len: int,
+    draft_token_width: int,
+    draft_token_ids: torch.Tensor | None,
+) -> torch.Tensor:
+    """Mirrors PPHandler.broadcast's payload construction."""
+    payload = sampled_token_ids.new_zeros(
+        sampled_token_ids.shape[0], max_sample_len + draft_token_width
+    )
+    num_cols = min(sampled_token_ids.shape[1], max_sample_len)
+    payload[:, :num_cols].copy_(sampled_token_ids[:, :num_cols])
+    if draft_token_width:
+        assert draft_token_ids is not None
+        payload[:, max_sample_len:].copy_(draft_token_ids)
+    return payload
+
+
+@pytest.mark.parametrize("num_speculative_steps", [0, 1, 3, 7])
+@pytest.mark.parametrize("sampler_cols", [1, None])
+def test_sender_width_matches_receiver_allocation(num_speculative_steps, sampler_cols):
+    """A one-column sampler output must still broadcast at the agreed width."""
+    num_reqs = 4
+    max_sample_len = num_speculative_steps + 1
+    cols = sampler_cols if sampler_cols is not None else max_sample_len
+    sampled = torch.arange(num_reqs * cols, dtype=torch.int64).view(num_reqs, cols)
+
+    payload = _sender_payload(sampled, max_sample_len, 0, None)
+
+    # This is exactly what receive() allocates.
+    assert payload.shape == (num_reqs, _payload_width(max_sample_len, 0))
+    # The columns the sampler did provide survive unchanged.
+    assert torch.equal(payload[:, :cols], sampled[:, :cols])
+
+
+@pytest.mark.parametrize("num_speculative_steps", [1, 7])
+def test_draft_tokens_ride_the_same_payload(num_speculative_steps):
+    num_reqs = 3
+    max_sample_len = num_speculative_steps + 1
+    draft_width = num_speculative_steps
+    sampled = torch.full((num_reqs, 1), 5, dtype=torch.int64)
+    drafts = torch.arange(num_reqs * draft_width, dtype=torch.int64).view(
+        num_reqs, draft_width
+    )
+
+    payload = _sender_payload(sampled, max_sample_len, draft_width, drafts)
+
+    assert payload.shape == (num_reqs, _payload_width(max_sample_len, draft_width))
+    # Receiver splits at max_sample_len; both halves must come back intact.
+    assert torch.equal(payload[:, :1], sampled)
+    assert torch.equal(payload[:, max_sample_len:], drafts)
+
+
+def test_narrow_sampler_output_is_not_replicated_across_columns():
+    """Padding must be zeros, not a broadcast of the single sampled column.
+
+    `Tensor.copy_` broadcasts a [n, 1] source across a [n, k] destination, which
+    would silently fill every speculative slot with the bonus token. num_sampled
+    bounds what is read, but replicating real token ids into unused slots makes
+    any later off-by-one read plausible-looking garbage instead of an obvious 0.
+    """
+    num_reqs, max_sample_len = 2, 8
+    sampled = torch.full((num_reqs, 1), 7, dtype=torch.int64)
+
+    payload = _sender_payload(sampled, max_sample_len, 0, None)
+
+    assert torch.equal(payload[:, :1], sampled)
+    assert torch.equal(payload[:, 1:], torch.zeros(num_reqs, max_sample_len - 1,
+                                                   dtype=torch.int64))
+
+
+def test_stale_slot_rows_are_filtered_before_restore():
+    """Only still-valid rows may be scattered back into request state.
+
+    A pending PP entry is consumed pp_size steps after it is received; a request
+    can finish in that window and its state index be handed to a new request.
+    Restoring through the unfiltered mapping would write the finished request's
+    drafts into whoever owns the index now.
+    """
+    import numpy as np
+
+    idx_mapping_np = np.array([3, 1, 2], dtype=np.int32)
+    exclude_mask = np.array([False, True, False])  # row 1 finished
+    drafts = torch.tensor([[10, 11], [20, 21], [30, 31]], dtype=torch.int64)
+
+    valid_rows = np.flatnonzero(~exclude_mask)
+    update_indices = np.stack((valid_rows, idx_mapping_np[valid_rows]))
+    draft_rows = torch.from_numpy(update_indices[0]).to(torch.int64)
+    draft_idx_mapping = torch.from_numpy(update_indices[1]).to(torch.int64)
+    selected = drafts.index_select(0, draft_rows)
+
+    # Row 1 (state index 1) must not be written at all.
+    assert draft_idx_mapping.tolist() == [3, 2]
+    assert selected.tolist() == [[10, 11], [30, 31]]

vllm/config/speculative.py

@@ -1258,7 +1258,11 @@ def create_draft_parallel_config(
         This is mostly a copy of the target parallel config, except the tp_size.
         """
         draft_parallel_config = ParallelConfig(
-            pipeline_parallel_size=target_parallel_config.pipeline_parallel_size,
+            # The drafter is built on the last pipeline stage only, so it never
+            # spans stages and must not inherit the target's pipeline size --
+            # doing so makes verify_with_parallel_config reject every draft
+            # model that does not implement SupportsPP, which is all of them.
+            pipeline_parallel_size=1,
             tensor_parallel_size=speculative_draft_tensor_parallel_size,
             distributed_executor_backend=target_parallel_config.distributed_executor_backend,
             max_parallel_loading_workers=target_parallel_config.max_parallel_loading_workers,

vllm/model_executor/models/interfaces.py

@@ -1399,8 +1399,41 @@ def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor:
 class EagleModelMixin:
     aux_hidden_state_layers: tuple[int, ...] = ()
 
+    # Set by models that pack their own aux taps into the forward output so the
+    # runner can carry them across the pipeline handoff.
+    supports_aux_hidden_states_over_pp: ClassVar[bool] = False
+
+    AUX_HIDDEN_STATE_KEY: ClassVar[str] = "aux_hidden_states_"
+
+    # Resolved once at setup: get_pp_indices logs on an uneven split, and dynamo
+    # cannot trace logging from inside a forward.
+    _aux_slot_base_cached: int = 0
+    _aux_upstream_total_cached: int = 0
+
     def _set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
         self.aux_hidden_state_layers = layers
+        self._cache_aux_pp_layout()
+
+    def _cache_aux_pp_layout(self) -> None:
+        """Resolve this stage's aux slot numbering, off the forward path."""
+        from vllm.distributed.parallel_state import (
+            get_pp_group,
+            model_parallel_is_initialized,
+        )
+
+        # Models are also built outside a worker, where there is no PP group and
+        # nothing to forward.
+        if not model_parallel_is_initialized():
+            return
+        pp = get_pp_group()
+        if pp.world_size < 2:
+            return
+        self._aux_slot_base_cached = self._aux_slot_base(
+            pp.rank_in_group, pp.world_size
+        )
+        self._aux_upstream_total_cached = self._aux_slot_base(
+            pp.world_size - 1, pp.world_size
+        )
 
     def _maybe_add_hidden_state(
         self,
@@ -1414,6 +1447,90 @@ def _maybe_add_hidden_state(
             aux_hidden_states.append(value)
         return aux_hidden_states
 
+    @staticmethod
+    def local_aux_tap_ids(
+        start_layer: int,
+        end_layer: int,
+        aux_ids: tuple[int, ...],
+        is_first_rank: bool,
+    ) -> tuple[int, ...]:
+        """Tap ids this stage produces itself, excluding inherited ones."""
+        out: list[int] = []
+        if is_first_rank and start_layer in aux_ids:
+            out.append(start_layer)
+        for layer_idx in range(start_layer, end_layer):
+            if (layer_idx + 1) in aux_ids:
+                out.append(layer_idx + 1)
+        return tuple(out)
+
+    def _total_num_layers(self) -> int:
+        num_layers = getattr(getattr(self, "config", None), "num_hidden_layers", None)
+        if num_layers is None:
+            raise RuntimeError(
+                "aux-over-PP transport needs config.num_hidden_layers on the model"
+            )
+        return num_layers
+
+    def _num_local_taps_on_rank(self, rank: int, pp_world_size: int) -> int:
+        """Taps stage ``rank`` produces, derived without building that stage."""
+        from vllm.distributed.utils import get_pp_indices
+
+        start, end = get_pp_indices(self._total_num_layers(), rank, pp_world_size)
+        return len(
+            self.local_aux_tap_ids(
+                start, end, tuple(self.aux_hidden_state_layers), rank == 0
+            )
+        )
+
+    def _aux_slot_base(self, rank: int, pp_world_size: int) -> int:
+        """Global slot index of ``rank``'s first tap.
+
+        One slot per tap, ordered by producing rank. Every stage derives the same
+        numbering from the layer split, so the slots need no negotiation.
+        """
+        return sum(self._num_local_taps_on_rank(r, pp_world_size) for r in range(rank))
+
+    def pack_local_aux_for_last(
+        self, aux_hidden_states: list[torch.Tensor]
+    ) -> dict[str, torch.Tensor]:
+        """Expose this stage's taps to the runner, keyed by global slot.
+
+        Pure packing, so the forward stays capturable by a full CUDA graph.
+        """
+        if not aux_hidden_states:
+            return {}
+        base = self._aux_slot_base_cached
+        return {
+            f"{self.AUX_HIDDEN_STATE_KEY}{base + i}": t
+            for i, t in enumerate(aux_hidden_states)
+        }
+
+    def recv_remote_aux_from_producers(
+        self, intermediate_tensors: "IntermediateTensors | None"
+    ) -> list[torch.Tensor]:
+        """Collect earlier stages' taps on the last rank, in tap order.
+
+        The handoff has already landed them in the persistent buffer, so reading
+        fixed slots keeps the forward capturable by a full CUDA graph.
+        """
+        total = self._aux_upstream_total_cached
+        if total == 0:
+            return []
+
+        assert intermediate_tensors is not None
+        out: list[torch.Tensor] = []
+        for i in range(total):
+            key = f"{self.AUX_HIDDEN_STATE_KEY}{i}"
+            if key not in intermediate_tensors.tensors:
+                # Zero-filling here would silently cost acceptance rate.
+                raise RuntimeError(
+                    f"{key} missing from the last stage's intermediate tensors; "
+                    "the aux slots were not reserved (got "
+                    f"{sorted(intermediate_tensors.tensors)})"
+                )
+            out.append(intermediate_tensors[key])
+        return out
+
 
 @runtime_checkable
 class SupportsEagle(SupportsEagleBase, Protocol):