[Core][MRV2] Support eagle3 spec decode with pipeline parallel

tests/model_executor/test_qwen3_omni.py

@@ -347,6 +347,7 @@ def test_dspark_shares_target_embedding_with_smaller_draft_vocabulary():
     vllm_config = SimpleNamespace(
         speculative_config=SimpleNamespace(
             draft_model_config=draft_model_config,
+            draft_parallel_config=SimpleNamespace(),
             attention_backend=None,
             kv_cache_dtype=None,
         ),
@@ -362,9 +363,8 @@ def fake_replace(config, **changes):
 
     with (
         patch.object(dspark_utils, "replace", side_effect=fake_replace),
-        patch.object(
-            dspark_utils,
-            "get_pp_group",
+        patch(
+            "vllm.v1.worker.gpu.spec_decode.eagle.utils.get_pp_group",
             return_value=SimpleNamespace(world_size=1),
         ),
         patch(

tests/models/kimi_k3/test_aux_attn_res_stream.py

@@ -157,25 +157,23 @@ def test_last_layer_on_the_final_rank_uses_the_output_aggregation(
     assert recorder[0].kwargs["num_blocks"] == 99
 
 
-def test_last_layer_of_a_non_final_stage_falls_back(recorder, monkeypatch):
-    """The consumer lives on the next rank and the output aggregation only
-    exists on the last one, so there is nothing here to mix against.
-
-    This is the case that would otherwise reach for weights this rank never
-    constructs. The forward guard is `layer_idx + 1 < end_layer`, where
-    `end_layer` is the rank's own exclusive bound from `get_pp_indices`, so a
-    `PPMissingLayer` is unreachable by construction -- the fallback below is
-    what makes that true rather than merely likely.
-    """
+def test_aux_layer_at_non_final_pp_boundary_is_rejected(monkeypatch):
+    model = k3_model.KimiLinearModel.__new__(k3_model.KimiLinearModel)
+    torch.nn.Module.__init__(model)
+    model.use_attn_res = True
+    model.end_layer = 72
+    monkeypatch.setattr(
+        "vllm.distributed.parallel_state.model_parallel_is_initialized", lambda: False
+    )
     _set_last_rank(monkeypatch, False)
-    prefix_sum = torch.tensor([3.0, 4.0])
-
-    got = _call(
-        _stub_model(enabled=True), END_LAYER - 1, prefix_sum, None, torch.zeros(2)
+    monkeypatch.setattr(
+        k3_model.KimiLinearModel,
+        "_aux_attn_res_stream",
+        property(lambda self: True),
     )
 
-    torch.testing.assert_close(got, prefix_sum)
-    assert not recorder, "no weights exist on this rank to mix against"
+    with pytest.raises(ValueError, match="Auxiliary layer 72"):
+        model._set_aux_hidden_state_layers((3, 24, 48, 72, 90))
 
 
 def test_pending_mlp_output_is_folded_in_rather_than_passed_as_delta(

tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+
+from tests.utils import multi_gpu_test
+from tests.v1.e2e.spec_decode.utils import compute_acceptance_len
+from vllm import LLM, SamplingParams
+from vllm.distributed import cleanup_dist_env_and_memory
+
+MODEL = "meta-llama/Llama-3.2-1B-Instruct"
+DRAFT = "nm-testing/Llama3_2_1B_speculator.eagle3"
+PROMPTS = [
+    "The capital of France is",
+    "2 + 2 equals",
+    "In one word, the color of the sky is",
+    "Q: If a train travels 60 miles in 1.5 hours, what is its average speed?\nA:",
+]
+
+ACCEPTANCE_TOLERANCE = 0.95
+
+
+def _run(pp_size: int) -> float:
+    llm = LLM(
+        model=MODEL,
+        tensor_parallel_size=1,
+        pipeline_parallel_size=pp_size,
+        max_model_len=512,
+        gpu_memory_utilization=0.45,
+        disable_log_stats=False,
+        compilation_config={"cudagraph_mode": "FULL_AND_PIECEWISE"},
+        speculative_config={
+            "method": "eagle3",
+            "model": DRAFT,
+            "num_speculative_tokens": 3,
+        },
+    )
+    try:
+        llm.generate(
+            PROMPTS,
+            SamplingParams(temperature=0.0, max_tokens=32, ignore_eos=True),
+        )
+        acceptance = compute_acceptance_len(llm.get_metrics())
+        assert acceptance > 1
+        return acceptance
+    finally:
+        del llm
+        torch.accelerator.empty_cache()
+        cleanup_dist_env_and_memory()
+
+
+@multi_gpu_test(num_gpus=4)
+def test_eagle3_pipeline_parallel_acceptance():
+    baseline = _run(1)
+    for pp_size in (2, 4):
+        parallel = _run(pp_size)
+        assert parallel >= baseline * ACCEPTANCE_TOLERANCE, (
+            f"PP={pp_size} acceptance regressed: {parallel:.3f} < "
+            f"{baseline:.3f} * {ACCEPTANCE_TOLERANCE}"
+        )

tests/v1/worker/test_eagle3_aux_hidden_states_pp.py

@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from types import SimpleNamespace
+
+import pytest
+
+from vllm.model_executor.models.interfaces import EagleModelMixin
+from vllm.model_executor.models.mimo import MiMoModel
+from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
+    verify_supports_aux_hidden_states_over_pp,
+)
+
+
+def test_aux_layers_are_sorted_and_deduplicated():
+    model = EagleModelMixin()
+    model._set_aux_hidden_state_layers((48, 3, 90, 24, 48))
+    assert model.aux_hidden_state_layers == (3, 24, 48, 90)
+
+
+def test_mimo_does_not_inherit_aux_hidden_state_pp_support():
+    inner = MiMoModel.__new__(MiMoModel)
+    target = SimpleNamespace(model=inner)
+
+    assert not inner.supports_aux_hidden_states_over_pp
+    with pytest.raises(ValueError, match="does not support eagle3"):
+        verify_supports_aux_hidden_states_over_pp(target, "eagle3")

tests/v1/worker/test_spec_decode_embed_sharing_pp.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm.model_executor.models.utils import (
+    PPMissingLayer,
+    spec_decode_needs_target_embed,
+)
+from vllm.v1.worker.gpu.spec_decode.eagle import utils as eagle_utils
+
+VOCAB, HIDDEN = 32, 8
+
+
+def _fake_pp(world_size: int, is_last_rank: bool = True):
+    return lambda: SimpleNamespace(
+        world_size=world_size,
+        is_last_rank=is_last_rank,
+        is_first_rank=world_size == 1,
+    )
+
+
+def _inner(embed: nn.Module | None) -> nn.Module:
+    inner = nn.Module()
+    if embed is not None:
+        inner.embed_tokens = embed
+    return inner
+
+
+def _embed(fill: float | None = None) -> nn.Embedding:
+    embed = nn.Embedding(VOCAB, HIDDEN)
+    if fill is not None:
+        with torch.no_grad():
+            embed.weight.fill_(fill)
+    return embed
+
+
+@pytest.mark.parametrize("draft_embed", ["loaded", "unset"])
+def test_drafter_without_own_embedding_gets_the_targets(monkeypatch, draft_embed):
+    monkeypatch.setattr(eagle_utils, "get_pp_group", _fake_pp(2))
+    target_embed = _embed()
+    draft_inner = _inner(_embed() if draft_embed == "loaded" else None)
+    if draft_embed == "unset":
+        draft_inner.embed_tokens = None
+    draft = SimpleNamespace(has_own_embed_tokens=False)
+
+    eagle_utils.maybe_share_target_embed(draft, draft_inner, _inner(target_embed))
+
+    assert draft_inner.embed_tokens is target_embed
+
+
+def test_missing_target_embedding_raises_instead_of_running_on_garbage(monkeypatch):
+    monkeypatch.setattr(eagle_utils, "get_pp_group", _fake_pp(2))
+    draft_inner = _inner(_embed())
+    draft = SimpleNamespace(has_own_embed_tokens=False)
+
+    with pytest.raises(RuntimeError, match="needs the target input embedding"):
+        eagle_utils.maybe_share_target_embed(
+            draft, draft_inner, _inner(PPMissingLayer())
+        )
+
+
+def test_drafter_with_distinct_weights_keeps_them(monkeypatch):
+    monkeypatch.setattr(eagle_utils, "get_pp_group", _fake_pp(2))
+    draft_embed = _embed(fill=1.0)
+    draft_inner = _inner(draft_embed)
+    draft = SimpleNamespace(has_own_embed_tokens=True)
+
+    eagle_utils.maybe_share_target_embed(draft, draft_inner, _inner(_embed(fill=2.0)))
+
+    assert draft_inner.embed_tokens is draft_embed
+
+
+def test_mtp_style_drafter_is_left_alone_under_pp(monkeypatch):
+    monkeypatch.setattr(eagle_utils, "get_pp_group", _fake_pp(2))
+    draft_embed = _embed()
+    draft_inner = _inner(draft_embed)
+
+    eagle_utils.maybe_share_target_embed(nn.Module(), draft_inner, _inner(_embed()))
+
+    assert draft_inner.embed_tokens is draft_embed
+
+
+@pytest.mark.parametrize(
+    "method,pp_size,is_last_rank,expected",
+    [
+        ("eagle", 2, True, True),
+        ("eagle3", 2, True, True),
+        ("dflash", 2, True, True),
+        ("dspark", 2, True, True),
+        ("eagle3", 1, True, False),
+        ("eagle3", 2, False, False),
+        ("mtp", 2, True, False),
+        (None, 2, True, False),
+    ],
+)
+def test_target_embedding_provisioning(
+    monkeypatch, method, pp_size, is_last_rank, expected
+):
+    monkeypatch.setattr(
+        "vllm.distributed.parallel_state.get_pp_group",
+        _fake_pp(pp_size, is_last_rank),
+        raising=True,
+    )
+    speculative_config = None if method is None else SimpleNamespace(method=method)
+    config = SimpleNamespace(speculative_config=speculative_config)
+    assert spec_decode_needs_target_embed(config) is expected

vllm/config/speculative.py

@@ -1692,7 +1692,7 @@ 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,
+            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/config/vllm.py

@@ -2590,12 +2590,6 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]:
             ):
                 unsupported.append("parallel drafting for EAGLE speculative decoding")
 
-            if (
-                speculative_config.method == "eagle3"
-                and self.parallel_config.pipeline_parallel_size > 1
-            ):
-                unsupported.append("EAGLE3 with pipeline parallelism")
-
         if self.parallel_config.use_ubatching:
             unsupported.extend(self._get_dbo_unsupported_features())
 

vllm/model_executor/models/deepseek_eagle3.py

@@ -285,9 +285,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/interfaces.py

@@ -3,6 +3,7 @@
 
 import asyncio
 import weakref
+from bisect import bisect_right
 from collections.abc import (
     AsyncGenerator,
     Callable,
@@ -1572,10 +1573,36 @@ def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
 
 class EagleModelMixin:
+    start_layer: int
     aux_hidden_state_layers: tuple[int, ...] = ()
+    supports_aux_hidden_states_over_pp: ClassVar[bool] = False
+    AUX_HIDDEN_STATE_KEY: ClassVar[str] = "aux_hidden_states_"
+    _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.aux_hidden_state_layers = tuple(sorted(set(layers)))
+        self._aux_slot_base_cached = 0
+        self._aux_upstream_total_cached = 0
+        self._cache_aux_pp_layout()
+
+    def _cache_aux_pp_layout(self) -> None:
+        from vllm.distributed.parallel_state import (
+            get_pp_group,
+            model_parallel_is_initialized,
+        )
+
+        if not model_parallel_is_initialized():
+            return
+        pp = get_pp_group()
+        if pp.world_size < 2:
+            return
+        if not pp.is_first_rank:
+            self._aux_slot_base_cached = bisect_right(
+                self.aux_hidden_state_layers, self.start_layer
+            )
+        if pp.is_last_rank:
+            self._aux_upstream_total_cached = self._aux_slot_base_cached
 
     def _maybe_add_hidden_state(
         self,
@@ -1589,6 +1616,36 @@ def _maybe_add_hidden_state(
             aux_hidden_states.append(value)
         return aux_hidden_states
 
+    def pack_local_aux_hidden_states(
+        self, aux_hidden_states: list[torch.Tensor]
+    ) -> dict[str, torch.Tensor]:
+        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 collect_remote_aux_hidden_states(
+        self, intermediate_tensors: "IntermediateTensors | None"
+    ) -> list[torch.Tensor]:
+        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:
+                raise RuntimeError(
+                    f"Missing {key} from PP intermediate tensors: "
+                    f"{sorted(intermediate_tensors.tensors)}"
+                )
+            out.append(intermediate_tensors[key])
+        return out
+
 
 @runtime_checkable
 class SupportsEagle(SupportsEagleBase, Protocol):

vllm/model_executor/models/laguna_dflash.py

@@ -245,9 +245,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             raise ValueError("Laguna DFlash config requires `draft_vocab_size`.")
         self.has_own_embed_tokens = False
         self.has_own_lm_head = False
-        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()
         self.config.target_layer_count = target_layer_num
         target_vocab_size = vllm_config.model_config.get_vocab_size()
         if self.config.draft_vocab_size != target_vocab_size: