@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d87a440f88

[Core][MRV2] Support eagle3 spec decode with pipeline parallel (#50514)

Signed-off-by: Yongqin Wang <yongqinwang@roblox.com> Co-authored-by: Cursor <noreply@cursor.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>

yongqinwang-cmd committed Sep 4, 2026 at 20:45 UTC d87a440f88e28e5b37f9b1e22ce214d0426d5352
27 files changed +582 -152
tests/model_executor/test_qwen3_omni.py
+3 -3
@@ -347,6 +347,7 @@ def test_dspark_shares_target_embedding_with_smaller_draft_vocabulary():
347 vllm_config = SimpleNamespace(
348 speculative_config=SimpleNamespace(
349 draft_model_config=draft_model_config,
350 + draft_parallel_config=SimpleNamespace(),
351 attention_backend=None,
352 kv_cache_dtype=None,
353 ),
@@ -362,9 +363,8 @@ def test_dspark_shares_target_embedding_with_smaller_draft_vocabulary():
363
364 with (
365 patch.object(dspark_utils, "replace", side_effect=fake_replace),
365 - patch.object(
366 - dspark_utils,
367 - "get_pp_group",
366 + patch(
367 + "vllm.v1.worker.gpu.spec_decode.eagle.utils.get_pp_group",
368 return_value=SimpleNamespace(world_size=1),
369 ),
370 patch(
tests/models/kimi_k3/test_aux_attn_res_stream.py
+14 -16
@@ -157,25 +157,23 @@ def test_last_layer_on_the_final_rank_uses_the_output_aggregation(
157 assert recorder[0].kwargs["num_blocks"] == 99
158
159
160 -def test_last_layer_of_a_non_final_stage_falls_back(recorder, monkeypatch):
161 - """The consumer lives on the next rank and the output aggregation only
162 - exists on the last one, so there is nothing here to mix against.
163 -
164 - This is the case that would otherwise reach for weights this rank never
165 - constructs. The forward guard is `layer_idx + 1 < end_layer`, where
166 - `end_layer` is the rank's own exclusive bound from `get_pp_indices`, so a
167 - `PPMissingLayer` is unreachable by construction -- the fallback below is
168 - what makes that true rather than merely likely.
169 - """
160 +def test_aux_layer_at_non_final_pp_boundary_is_rejected(monkeypatch):
161 + model = k3_model.KimiLinearModel.__new__(k3_model.KimiLinearModel)
162 + torch.nn.Module.__init__(model)
163 + model.use_attn_res = True
164 + model.end_layer = 72
165 + monkeypatch.setattr(
166 + "vllm.distributed.parallel_state.model_parallel_is_initialized", lambda: False
167 + )
168 _set_last_rank(monkeypatch, False)
171 - prefix_sum = torch.tensor([3.0, 4.0])
172 -
173 - got = _call(
174 - _stub_model(enabled=True), END_LAYER - 1, prefix_sum, None, torch.zeros(2)
169 + monkeypatch.setattr(
170 + k3_model.KimiLinearModel,
171 + "_aux_attn_res_stream",
172 + property(lambda self: True),
173 )
174
177 - torch.testing.assert_close(got, prefix_sum)
178 - assert not recorder, "no weights exist on this rank to mix against"
175 + with pytest.raises(ValueError, match="Auxiliary layer 72"):
176 + model._set_aux_hidden_state_layers((3, 24, 48, 72, 90))
177
178
179 def test_pending_mlp_output_is_folded_in_rather_than_passed_as_delta(
tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py new
+59
@@ -0,0 +1,59 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch
4 +
5 +from tests.utils import multi_gpu_test
6 +from tests.v1.e2e.spec_decode.utils import compute_acceptance_len
7 +from vllm import LLM, SamplingParams
8 +from vllm.distributed import cleanup_dist_env_and_memory
9 +
10 +MODEL = "meta-llama/Llama-3.2-1B-Instruct"
11 +DRAFT = "nm-testing/Llama3_2_1B_speculator.eagle3"
12 +PROMPTS = [
13 + "The capital of France is",
14 + "2 + 2 equals",
15 + "In one word, the color of the sky is",
16 + "Q: If a train travels 60 miles in 1.5 hours, what is its average speed?\nA:",
17 +]
18 +
19 +ACCEPTANCE_TOLERANCE = 0.95
20 +
21 +
22 +def _run(pp_size: int) -> float:
23 + llm = LLM(
24 + model=MODEL,
25 + tensor_parallel_size=1,
26 + pipeline_parallel_size=pp_size,
27 + max_model_len=512,
28 + gpu_memory_utilization=0.45,
29 + disable_log_stats=False,
30 + compilation_config={"cudagraph_mode": "FULL_AND_PIECEWISE"},
31 + speculative_config={
32 + "method": "eagle3",
33 + "model": DRAFT,
34 + "num_speculative_tokens": 3,
35 + },
36 + )
37 + try:
38 + llm.generate(
39 + PROMPTS,
40 + SamplingParams(temperature=0.0, max_tokens=32, ignore_eos=True),
41 + )
42 + acceptance = compute_acceptance_len(llm.get_metrics())
43 + assert acceptance > 1
44 + return acceptance
45 + finally:
46 + del llm
47 + torch.accelerator.empty_cache()
48 + cleanup_dist_env_and_memory()
49 +
50 +
51 +@multi_gpu_test(num_gpus=4)
52 +def test_eagle3_pipeline_parallel_acceptance():
53 + baseline = _run(1)
54 + for pp_size in (2, 4):
55 + parallel = _run(pp_size)
56 + assert parallel >= baseline * ACCEPTANCE_TOLERANCE, (
57 + f"PP={pp_size} acceptance regressed: {parallel:.3f} < "
58 + f"{baseline:.3f} * {ACCEPTANCE_TOLERANCE}"
59 + )
tests/v1/worker/test_eagle3_aux_hidden_states_pp.py new
+26
@@ -0,0 +1,26 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from types import SimpleNamespace
4 +
5 +import pytest
6 +
7 +from vllm.model_executor.models.interfaces import EagleModelMixin
8 +from vllm.model_executor.models.mimo import MiMoModel
9 +from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
10 + verify_supports_aux_hidden_states_over_pp,
11 +)
12 +
13 +
14 +def test_aux_layers_are_sorted_and_deduplicated():
15 + model = EagleModelMixin()
16 + model._set_aux_hidden_state_layers((48, 3, 90, 24, 48))
17 + assert model.aux_hidden_state_layers == (3, 24, 48, 90)
18 +
19 +
20 +def test_mimo_does_not_inherit_aux_hidden_state_pp_support():
21 + inner = MiMoModel.__new__(MiMoModel)
22 + target = SimpleNamespace(model=inner)
23 +
24 + assert not inner.supports_aux_hidden_states_over_pp
25 + with pytest.raises(ValueError, match="does not support eagle3"):
26 + verify_supports_aux_hidden_states_over_pp(target, "eagle3")
tests/v1/worker/test_spec_decode_embed_sharing_pp.py new
+110
@@ -0,0 +1,110 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from types import SimpleNamespace
4 +
5 +import pytest
6 +import torch
7 +import torch.nn as nn
8 +
9 +from vllm.model_executor.models.utils import (
10 + PPMissingLayer,
11 + spec_decode_needs_target_embed,
12 +)
13 +from vllm.v1.worker.gpu.spec_decode.eagle import utils as eagle_utils
14 +
15 +VOCAB, HIDDEN = 32, 8
16 +
17 +
18 +def _fake_pp(world_size: int, is_last_rank: bool = True):
19 + return lambda: SimpleNamespace(
20 + world_size=world_size,
21 + is_last_rank=is_last_rank,
22 + is_first_rank=world_size == 1,
23 + )
24 +
25 +
26 +def _inner(embed: nn.Module | None) -> nn.Module:
27 + inner = nn.Module()
28 + if embed is not None:
29 + inner.embed_tokens = embed
30 + return inner
31 +
32 +
33 +def _embed(fill: float | None = None) -> nn.Embedding:
34 + embed = nn.Embedding(VOCAB, HIDDEN)
35 + if fill is not None:
36 + with torch.no_grad():
37 + embed.weight.fill_(fill)
38 + return embed
39 +
40 +
41 +@pytest.mark.parametrize("draft_embed", ["loaded", "unset"])
42 +def test_drafter_without_own_embedding_gets_the_targets(monkeypatch, draft_embed):
43 + monkeypatch.setattr(eagle_utils, "get_pp_group", _fake_pp(2))
44 + target_embed = _embed()
45 + draft_inner = _inner(_embed() if draft_embed == "loaded" else None)
46 + if draft_embed == "unset":
47 + draft_inner.embed_tokens = None
48 + draft = SimpleNamespace(has_own_embed_tokens=False)
49 +
50 + eagle_utils.maybe_share_target_embed(draft, draft_inner, _inner(target_embed))
51 +
52 + assert draft_inner.embed_tokens is target_embed
53 +
54 +
55 +def test_missing_target_embedding_raises_instead_of_running_on_garbage(monkeypatch):
56 + monkeypatch.setattr(eagle_utils, "get_pp_group", _fake_pp(2))
57 + draft_inner = _inner(_embed())
58 + draft = SimpleNamespace(has_own_embed_tokens=False)
59 +
60 + with pytest.raises(RuntimeError, match="needs the target input embedding"):
61 + eagle_utils.maybe_share_target_embed(
62 + draft, draft_inner, _inner(PPMissingLayer())
63 + )
64 +
65 +
66 +def test_drafter_with_distinct_weights_keeps_them(monkeypatch):
67 + monkeypatch.setattr(eagle_utils, "get_pp_group", _fake_pp(2))
68 + draft_embed = _embed(fill=1.0)
69 + draft_inner = _inner(draft_embed)
70 + draft = SimpleNamespace(has_own_embed_tokens=True)
71 +
72 + eagle_utils.maybe_share_target_embed(draft, draft_inner, _inner(_embed(fill=2.0)))
73 +
74 + assert draft_inner.embed_tokens is draft_embed
75 +
76 +
77 +def test_mtp_style_drafter_is_left_alone_under_pp(monkeypatch):
78 + monkeypatch.setattr(eagle_utils, "get_pp_group", _fake_pp(2))
79 + draft_embed = _embed()
80 + draft_inner = _inner(draft_embed)
81 +
82 + eagle_utils.maybe_share_target_embed(nn.Module(), draft_inner, _inner(_embed()))
83 +
84 + assert draft_inner.embed_tokens is draft_embed
85 +
86 +
87 +@pytest.mark.parametrize(
88 + "method,pp_size,is_last_rank,expected",
89 + [
90 + ("eagle", 2, True, True),
91 + ("eagle3", 2, True, True),
92 + ("dflash", 2, True, True),
93 + ("dspark", 2, True, True),
94 + ("eagle3", 1, True, False),
95 + ("eagle3", 2, False, False),
96 + ("mtp", 2, True, False),
97 + (None, 2, True, False),
98 + ],
99 +)
100 +def test_target_embedding_provisioning(
101 + monkeypatch, method, pp_size, is_last_rank, expected
102 +):
103 + monkeypatch.setattr(
104 + "vllm.distributed.parallel_state.get_pp_group",
105 + _fake_pp(pp_size, is_last_rank),
106 + raising=True,
107 + )
108 + speculative_config = None if method is None else SimpleNamespace(method=method)
109 + config = SimpleNamespace(speculative_config=speculative_config)
110 + assert spec_decode_needs_target_embed(config) is expected
vllm/config/speculative.py
+1 -1
@@ -1699,7 +1699,7 @@ class SpeculativeConfig:
1699 This is mostly a copy of the target parallel config, except the tp_size.
1700 """
1701 draft_parallel_config = ParallelConfig(
1702 - pipeline_parallel_size=target_parallel_config.pipeline_parallel_size,
1702 + pipeline_parallel_size=1,
1703 tensor_parallel_size=speculative_draft_tensor_parallel_size,
1704 distributed_executor_backend=target_parallel_config.distributed_executor_backend,
1705 max_parallel_loading_workers=target_parallel_config.max_parallel_loading_workers,
vllm/config/vllm.py
-6
@@ -2590,12 +2590,6 @@ class VllmConfig:
2590 ):
2591 unsupported.append("parallel drafting for EAGLE speculative decoding")
2592
2593 - if (
2594 - speculative_config.method == "eagle3"
2595 - and self.parallel_config.pipeline_parallel_size > 1
2596 - ):
2597 - unsupported.append("EAGLE3 with pipeline parallelism")
2598 -
2593 if self.parallel_config.use_ubatching:
2594 unsupported.extend(self._get_dbo_unsupported_features())
2595
vllm/model_executor/models/deepseek_eagle3.py
+1 -3
@@ -285,9 +285,7 @@ class Eagle3DeepseekV2ForCausalLM(LocalArgmaxMixin, DeepseekV2ForCausalLM):
285 base_vocab_size = getattr(self.config, "vocab_size", None)
286 self.config.draft_vocab_size = base_vocab_size
287
288 - target_layer_num = vllm_config.model_config.get_num_layers(
289 - vllm_config.parallel_config
290 - )
288 + target_layer_num = vllm_config.model_config.get_total_num_hidden_layers()
289
290 # Store target layer count in draft config
291 self.config.target_layer_count = target_layer_num
vllm/model_executor/models/interfaces.py
+58 -1
@@ -3,6 +3,7 @@
3
4 import asyncio
5 import weakref
6 +from bisect import bisect_right
7 from collections.abc import (
8 AsyncGenerator,
9 Callable,
@@ -1572,10 +1573,36 @@ class LocalArgmaxMixin:
1573
1574
1575 class EagleModelMixin:
1576 + start_layer: int
1577 aux_hidden_state_layers: tuple[int, ...] = ()
1578 + supports_aux_hidden_states_over_pp: ClassVar[bool] = False
1579 + AUX_HIDDEN_STATE_KEY: ClassVar[str] = "aux_hidden_states_"
1580 + _aux_slot_base_cached: int = 0
1581 + _aux_upstream_total_cached: int = 0
1582
1583 def _set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
1578 - self.aux_hidden_state_layers = layers
1584 + self.aux_hidden_state_layers = tuple(sorted(set(layers)))
1585 + self._aux_slot_base_cached = 0
1586 + self._aux_upstream_total_cached = 0
1587 + self._cache_aux_pp_layout()
1588 +
1589 + def _cache_aux_pp_layout(self) -> None:
1590 + from vllm.distributed.parallel_state import (
1591 + get_pp_group,
1592 + model_parallel_is_initialized,
1593 + )
1594 +
1595 + if not model_parallel_is_initialized():
1596 + return
1597 + pp = get_pp_group()
1598 + if pp.world_size < 2:
1599 + return
1600 + if not pp.is_first_rank:
1601 + self._aux_slot_base_cached = bisect_right(
1602 + self.aux_hidden_state_layers, self.start_layer
1603 + )
1604 + if pp.is_last_rank:
1605 + self._aux_upstream_total_cached = self._aux_slot_base_cached
1606
1607 def _maybe_add_hidden_state(
1608 self,
@@ -1589,6 +1616,36 @@ class EagleModelMixin:
1616 aux_hidden_states.append(value)
1617 return aux_hidden_states
1618
1619 + def pack_local_aux_hidden_states(
1620 + self, aux_hidden_states: list[torch.Tensor]
1621 + ) -> dict[str, torch.Tensor]:
1622 + if not aux_hidden_states:
1623 + return {}
1624 + base = self._aux_slot_base_cached
1625 + return {
1626 + f"{self.AUX_HIDDEN_STATE_KEY}{base + i}": t
1627 + for i, t in enumerate(aux_hidden_states)
1628 + }
1629 +
1630 + def collect_remote_aux_hidden_states(
1631 + self, intermediate_tensors: "IntermediateTensors | None"
1632 + ) -> list[torch.Tensor]:
1633 + total = self._aux_upstream_total_cached
1634 + if total == 0:
1635 + return []
1636 +
1637 + assert intermediate_tensors is not None
1638 + out: list[torch.Tensor] = []
1639 + for i in range(total):
1640 + key = f"{self.AUX_HIDDEN_STATE_KEY}{i}"
1641 + if key not in intermediate_tensors.tensors:
1642 + raise RuntimeError(
1643 + f"Missing {key} from PP intermediate tensors: "
1644 + f"{sorted(intermediate_tensors.tensors)}"
1645 + )
1646 + out.append(intermediate_tensors[key])
1647 + return out
1648 +
1649
1650 @runtime_checkable
1651 class SupportsEagle(SupportsEagleBase, Protocol):
vllm/model_executor/models/laguna_dflash.py
+1 -3
@@ -245,9 +245,7 @@ class DFlashLagunaForCausalLM(nn.Module, SupportsEagle3):
245 raise ValueError("Laguna DFlash config requires `draft_vocab_size`.")
246 self.has_own_embed_tokens = False
247 self.has_own_lm_head = False
248 - target_layer_num = vllm_config.model_config.get_num_layers(
249 - vllm_config.parallel_config
250 - )
248 + target_layer_num = vllm_config.model_config.get_total_num_hidden_layers()
249 self.config.target_layer_count = target_layer_num
250 target_vocab_size = vllm_config.model_config.get_vocab_size()
251 if self.config.draft_vocab_size != target_vocab_size:
vllm/model_executor/models/llama.py
+21 -5
@@ -75,6 +75,7 @@ from .utils import (
75 make_empty_intermediate_tensors_factory,
76 make_layers,
77 maybe_prefix,
78 + spec_decode_needs_target_embed,
79 )
80
81
@@ -354,6 +355,7 @@ class LlamaModel(nn.Module, EagleModelMixin):
355 ".up_proj": (".gate_up_proj", 1),
356 }
357 )
358 + supports_aux_hidden_states_over_pp = True
359
360 def __init__(
361 self,
@@ -372,8 +374,10 @@ class LlamaModel(nn.Module, EagleModelMixin):
374
375 self.vocab_size = config.vocab_size
376
375 - if get_pp_group().is_first_rank or (
376 - config.tie_word_embeddings and get_pp_group().is_last_rank
377 + if (
378 + get_pp_group().is_first_rank
379 + or (config.tie_word_embeddings and get_pp_group().is_last_rank)
380 + or spec_decode_needs_target_embed(vllm_config)
381 ):
382 self.embed_tokens = VocabParallelEmbedding(
383 self.vocab_size,
@@ -418,9 +422,16 @@ class LlamaModel(nn.Module, EagleModelMixin):
422 hidden_states = intermediate_tensors["hidden_states"]
423 residual = intermediate_tensors["residual"]
424
421 - aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
425 + remote_aux = self.collect_remote_aux_hidden_states(intermediate_tensors)
426 +
427 + aux_hidden_states: list[torch.Tensor] = []
428 + if get_pp_group().is_first_rank:
429 + self._maybe_add_hidden_state(
430 + aux_hidden_states, self.start_layer, hidden_states, residual
431 + )
432 for idx, layer in enumerate(
423 - islice(self.layers, self.start_layer, self.end_layer)
433 + islice(self.layers, self.start_layer, self.end_layer),
434 + start=self.start_layer,
435 ):
436 hidden_states, residual = layer(
437 positions, hidden_states, residual, **extra_layer_kwargs
@@ -431,11 +442,16 @@ class LlamaModel(nn.Module, EagleModelMixin):
442
443 if not get_pp_group().is_last_rank:
444 return IntermediateTensors(
434 - {"hidden_states": hidden_states, "residual": residual}
445 + {
446 + "hidden_states": hidden_states,
447 + "residual": residual,
448 + **self.pack_local_aux_hidden_states(aux_hidden_states),
449 + }
450 )
451
452 hidden_states, _ = self.norm(hidden_states, residual)
453
454 + aux_hidden_states = remote_aux + aux_hidden_states
455 if len(aux_hidden_states) > 0:
456 return hidden_states, aux_hidden_states
457 return hidden_states
vllm/model_executor/models/llama_eagle3.py
+1 -3
@@ -314,9 +314,7 @@ class Eagle3LlamaForCausalLM(_Eagle3LlamaForCausalLMBase):
314 if getattr(self.config, "draft_vocab_size", None) is None:
315 base_vocab_size = getattr(self.config, "vocab_size", None)
316 self.config.draft_vocab_size = base_vocab_size
317 - target_layer_num = vllm_config.model_config.get_num_layers(
318 - vllm_config.parallel_config
319 - )
317 + target_layer_num = vllm_config.model_config.get_total_num_hidden_layers()
318
319 # Store target layer count in draft config for
320 # proper layer_types indexing in draft models
vllm/model_executor/models/mimo.py
+2
@@ -55,6 +55,8 @@ logger = init_logger(__name__)
55 }
56 )
57 class MiMoModel(Qwen2Model):
58 + supports_aux_hidden_states_over_pp = False
59 +
60 def forward(
61 self,
62 input_ids: torch.Tensor | None,
vllm/model_executor/models/qwen2.py
+21 -5
@@ -74,6 +74,7 @@ from .utils import (
74 make_empty_intermediate_tensors_factory,
75 make_layers,
76 maybe_prefix,
77 + spec_decode_needs_target_embed,
78 )
79
80
@@ -330,6 +331,7 @@ class Qwen2Model(nn.Module, EagleModelMixin):
331 ".up_proj": (".gate_up_proj", 1),
332 }
333 )
334 + supports_aux_hidden_states_over_pp = True
335
336 def __init__(
337 self,
@@ -360,8 +362,10 @@ class Qwen2Model(nn.Module, EagleModelMixin):
362 self.quant_config = quant_config
363 self.vocab_size = config.vocab_size
364
363 - if get_pp_group().is_first_rank or (
364 - config.tie_word_embeddings and get_pp_group().is_last_rank
365 + if (
366 + get_pp_group().is_first_rank
367 + or (config.tie_word_embeddings and get_pp_group().is_last_rank)
368 + or spec_decode_needs_target_embed(vllm_config)
369 ):
370 self.embed_tokens = VocabParallelEmbedding(
371 config.vocab_size,
@@ -412,9 +416,16 @@ class Qwen2Model(nn.Module, EagleModelMixin):
416 hidden_states = intermediate_tensors["hidden_states"]
417 residual = intermediate_tensors["residual"]
418
415 - aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
419 + remote_aux = self.collect_remote_aux_hidden_states(intermediate_tensors)
420 +
421 + aux_hidden_states: list[torch.Tensor] = []
422 + if get_pp_group().is_first_rank:
423 + self._maybe_add_hidden_state(
424 + aux_hidden_states, self.start_layer, hidden_states, residual
425 + )
426 for idx, layer in enumerate(
417 - islice(self.layers, self.start_layer, self.end_layer)
427 + islice(self.layers, self.start_layer, self.end_layer),
428 + start=self.start_layer,
429 ):
430 hidden_states, residual = layer(positions, hidden_states, residual)
431 self._maybe_add_hidden_state(
@@ -423,11 +434,16 @@ class Qwen2Model(nn.Module, EagleModelMixin):
434
435 if not get_pp_group().is_last_rank:
436 return IntermediateTensors(
426 - {"hidden_states": hidden_states, "residual": residual}
437 + {
438 + "hidden_states": hidden_states,
439 + "residual": residual,
440 + **self.pack_local_aux_hidden_states(aux_hidden_states),
441 + }
442 )
443
444 hidden_states, _ = self.norm(hidden_states, residual)
445
446 + aux_hidden_states = remote_aux + aux_hidden_states
447 if len(aux_hidden_states) > 0:
448 return hidden_states, aux_hidden_states
449
vllm/model_executor/models/qwen3_dflash.py
+1 -3
@@ -694,9 +694,7 @@ class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):
694 self.config = self.draft_model_config.hf_config
695 if getattr(self.config, "draft_vocab_size", None) is None:
696 self.config.draft_vocab_size = getattr(self.config, "vocab_size", None)
697 - target_layer_num = vllm_config.model_config.get_num_layers(
698 - vllm_config.parallel_config
699 - )
697 + target_layer_num = vllm_config.model_config.get_total_num_hidden_layers()
698 self.model = self.model_cls(
699 vllm_config=vllm_config,
700 prefix=maybe_prefix(prefix, "model"),
vllm/model_executor/models/qwen3_dspark.py
+1 -3
@@ -182,9 +182,7 @@ class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM):
182 self.config = self.draft_model_config.hf_config
183 if getattr(self.config, "draft_vocab_size", None) is None:
184 self.config.draft_vocab_size = getattr(self.config, "vocab_size", None)
185 - target_layer_num = vllm_config.model_config.get_num_layers(
186 - vllm_config.parallel_config
187 - )
185 + target_layer_num = vllm_config.model_config.get_total_num_hidden_layers()
186 self.model = Qwen3DSparkModel(
187 vllm_config=vllm_config,
188 prefix=maybe_prefix(prefix, "model"),
vllm/model_executor/models/qwen3_eagle3.py
+1 -3
@@ -288,9 +288,7 @@ class Eagle3Qwen3ForCausalLM(Qwen3ForCausalLM):
288 if getattr(self.config, "draft_vocab_size", None) is None:
289 base_vocab_size = getattr(self.config, "vocab_size", None)
290 self.config.draft_vocab_size = base_vocab_size
291 - target_layer_num = vllm_config.model_config.get_num_layers(
292 - vllm_config.parallel_config
293 - )
291 + target_layer_num = vllm_config.model_config.get_total_num_hidden_layers()
292
293 # Store target layer count in draft config for
294 # proper layer_types indexing in draft models
vllm/model_executor/models/utils.py
+16
@@ -835,6 +835,22 @@ class PPMissingLayer(torch.nn.Identity):
835 return args[0] if args else next(iter(kwargs.values()))
836
837
838 +def spec_decode_needs_target_embed(vllm_config: VllmConfig) -> bool:
839 + """Whether the last PP rank needs the target input embedding."""
840 + from vllm.distributed.parallel_state import get_pp_group
841 +
842 + speculative_config = vllm_config.speculative_config
843 + if speculative_config is None or speculative_config.method not in (
844 + "eagle",
845 + "eagle3",
846 + "dflash",
847 + "dspark",
848 + ):
849 + return False
850 + pp = get_pp_group()
851 + return pp.world_size > 1 and pp.is_last_rank
852 +
853 +
854 def make_layers(
855 num_hidden_layers: int,
856 layer_fn: LayerFn,
vllm/models/deepseek_v4/nvidia/model.py
+12 -2
@@ -67,6 +67,7 @@ from vllm.model_executor.models.utils import (
67 is_pp_missing_parameter,
68 make_layers,
69 maybe_prefix,
70 + spec_decode_needs_target_embed,
71 )
72 from vllm.model_executor.utils import set_weight_attrs
73 from vllm.models.common.ops.sequence_parallel import (
@@ -1270,6 +1271,8 @@ class DeepseekV4DecoderLayer(nn.Module):
1271
1272
1273 class DeepseekV4Model(nn.Module, EagleModelMixin):
1274 + supports_aux_hidden_states_over_pp = True
1275 +
1276 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1277 super().__init__()
1278
@@ -1305,7 +1308,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
1308 dtype=torch.int32,
1309 )
1310
1308 - if get_pp_group().is_first_rank:
1311 + if get_pp_group().is_first_rank or spec_decode_needs_target_embed(vllm_config):
1312 self.embed_tokens = VocabParallelEmbedding(
1313 config.vocab_size,
1314 config.hidden_size,
@@ -1416,6 +1419,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
1419 input_ids = sp_shard(input_ids)
1420
1421 residual, post_mix, res_mix = None, None, None
1422 + remote_aux = self.collect_remote_aux_hidden_states(intermediate_tensors)
1423 aux_hidden_states: list[torch.Tensor] = []
1424 final_aux_recon: torch.Tensor | None = None # avoid duplicate mhc_post call
1425 for idx, layer in enumerate(
@@ -1450,7 +1454,12 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
1454 )
1455
1456 if not get_pp_group().is_last_rank:
1453 - return IntermediateTensors({"hidden_states": hidden_states})
1457 + return IntermediateTensors(
1458 + {
1459 + "hidden_states": hidden_states,
1460 + **self.pack_local_aux_hidden_states(aux_hidden_states),
1461 + }
1462 + )
1463
1464 if self.use_sequence_parallel:
1465 hidden_states = sp_all_gather(hidden_states)[:full_num_tokens]
@@ -1468,6 +1477,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
1477 self.hc_eps,
1478 )
1479 hidden_states = self.norm(hidden_states)
1480 + aux_hidden_states = remote_aux + aux_hidden_states
1481 if len(aux_hidden_states) > 0:
1482 return hidden_states, aux_hidden_states
1483 return hidden_states
vllm/models/kimi_k3/nvidia/dspark_mla.py
+1 -3
@@ -421,9 +421,7 @@ class K3DSparkForCausalLM(nn.Module):
421 assert vllm_config.speculative_config is not None
422 self.draft_model_config = vllm_config.speculative_config.draft_model_config
423 self.config = self.draft_model_config.hf_config
424 - target_layer_num = vllm_config.model_config.get_num_layers(
425 - vllm_config.parallel_config
426 - )
424 + target_layer_num = vllm_config.model_config.get_total_num_hidden_layers()
425 self.model = K3DSparkModel(
426 vllm_config=vllm_config,
427 start_layer_id=target_layer_num,
vllm/models/kimi_k3/nvidia/model.py
+26 -33
@@ -85,6 +85,7 @@ from vllm.model_executor.models.utils import (
85 is_pp_missing_parameter,
86 make_layers,
87 maybe_prefix,
88 + spec_decode_needs_target_embed,
89 )
90 from vllm.model_executor.models.vision import is_vit_use_data_parallel
91 from vllm.models.common.ops.sequence_parallel import (
@@ -1121,6 +1122,8 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant):
1122 "fused_qkv_a_g_proj": ["q_a_proj", "kv_a_proj_with_mqa", "g_proj"],
1123 }
1124
1125 + supports_aux_hidden_states_over_pp = True
1126 +
1127 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1128 super().__init__()
1129
@@ -1145,7 +1148,7 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant):
1148 vllm_config, self.use_sequence_parallel
1149 )
1150
1148 - if get_pp_group().is_first_rank:
1151 + if get_pp_group().is_first_rank or spec_decode_needs_target_embed(vllm_config):
1152 self.embed_tokens = VocabParallelEmbedding(
1153 config.vocab_size,
1154 config.hidden_size,
@@ -1228,15 +1231,18 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant):
1231
1232 def _set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
1233 super()._set_aux_hidden_state_layers(layers)
1234 + if self.use_attn_res and self._aux_attn_res_stream:
1235 + pp = get_pp_group()
1236 + if not pp.is_last_rank and self.end_layer in self.aux_hidden_state_layers:
1237 + raise ValueError(
1238 + f"Auxiliary layer {self.end_layer} cannot end a non-final PP "
1239 + "stage when VLLM_KIMI_K3_AUX_ATTN_RES_STREAM=1"
1240 + )
1241 if self.use_attn_res:
1232 - # Emitted once, at configuration time. Which layers are tapped and
1233 - # which convention is in force are the two things you need to
1234 - # confirm from a running process, and neither is recoverable from
1235 - # the served output.
1242 logger.info_once(
1243 "Kimi-K3 aux hidden capture: layers=%s mode=%s "
1244 "(VLLM_KIMI_K3_AUX_ATTN_RES_STREAM=%d)",
1239 - layers,
1245 + self.aux_hidden_state_layers,
1246 "attn_res_stream" if self._aux_attn_res_stream else "prefix_only",
1247 int(self._aux_attn_res_stream),
1248 )
@@ -1252,25 +1258,8 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant):
1258 pending_mlp_out: torch.Tensor | None,
1259 block_residual: torch.Tensor,
1260 ) -> torch.Tensor:
1255 - """Auxiliary feature tapped after ``layer_idx`` under AttnRes.
1256 -
1257 - The wire between layers only carries the current block's running prefix;
1258 - the committed blocks live in the bank. The value the next consumer
1259 - actually reads is the pre-norm AttnRes mixture over
1260 - ``bank[:num_blocks] + prefix``, which is what the DFlash drafters were
1261 - trained against. ``attn_res`` with no delta, no block write and no
1262 - output norm computes exactly that and leaves both the prefix and the
1263 - bank untouched.
1264 -
1265 - Folding the pending MLP output into the prefix rather than passing it as
1266 - ``delta`` is deliberate: the kernel writes an applied delta back into
1267 - the prefix in place, which would double-add it into the live residual
1268 - stream.
1269 - """
1261 + """Return the AttnRes stream after ``layer_idx``."""
1262 prefix = prefix_sum if pending_mlp_out is None else prefix_sum + pending_mlp_out
1271 - # `use_attn_res` is what constructs the norm and projection weights this
1272 - # reads; without it there is no mixture to compute and the attribute
1273 - # lookups below would raise.
1263 if not (self._aux_attn_res_stream and self.use_attn_res):
1264 return prefix
1265
@@ -1280,17 +1269,11 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant):
1269 score_proj = consumer.self_attention_res_proj
1270 num_blocks = consumer.prev_valid_blocks
1271 elif get_pp_group().is_last_rank:
1283 - # Nothing downstream but the model's own output-side aggregation.
1272 score_norm = self.output_attn_res_norm
1273 score_proj = self.output_attn_res_proj
1274 num_blocks = self.num_attn_res_blocks
1275 else:
1288 - # Last layer of a non-final pipeline stage: the consumer lives on
1289 - # the next rank and the output-side aggregation only exists on the
1290 - # last one, so there is nothing here to mix against. Falling back
1291 - # to the running prefix keeps the tap defined rather than reaching
1292 - # for weights this rank does not construct.
1293 - return prefix
1276 + raise RuntimeError("Auxiliary AttnRes capture crossed a PP boundary")
1277
1278 return attn_res(
1279 prefix,
@@ -1338,9 +1321,14 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant):
1321 hidden_states = sp_shard(hidden_states)
1322 assert residual is None, "Currently, SP is not supported with PP"
1323
1324 + remote_aux = self.collect_remote_aux_hidden_states(intermediate_tensors)
1325 +
1326 # sharded aux hidden states when sp is enabled
1327 aux_hidden_states: list[torch.Tensor] = []
1343 - if self.start_layer in self.aux_hidden_state_layers:
1328 + if (
1329 + get_pp_group().is_first_rank
1330 + and self.start_layer in self.aux_hidden_state_layers
1331 + ):
1332 if self.use_attn_res or residual is None:
1333 aux_hidden_states.append(hidden_states)
1334 else:
@@ -1391,7 +1379,11 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant):
1379 if prefix_sum is not None:
1380 hidden_states = hidden_states + prefix_sum
1381 return IntermediateTensors(
1394 - {"hidden_states": hidden_states, "residual": residual}
1382 + {
1383 + "hidden_states": hidden_states,
1384 + "residual": residual,
1385 + **self.pack_local_aux_hidden_states(aux_hidden_states),
1386 + }
1387 )
1388
1389 if self.use_attn_res:
@@ -1428,6 +1420,7 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant):
1420
1421 # NOTE: the final norm is applied in compute_logits instead of here, so
1422 # the MTP draft model receives the pre-norm hidden states.
1423 + aux_hidden_states = remote_aux + aux_hidden_states
1424 if aux_hidden_states:
1425 return hidden_states, aux_hidden_states
1426 return hidden_states
vllm/v1/worker/gpu/model_runner.py
+20 -7
@@ -152,6 +152,7 @@ from vllm.v1.worker.gpu.spec_decode.adaptive_verification import (
152 )
153 from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
154 set_eagle3_aux_hidden_state_layers,
155 + verify_supports_aux_hidden_states_over_pp,
156 )
157 from vllm.v1.worker.gpu.spec_decode.rejection_sampler import (
158 RejectionSampler,
@@ -280,11 +281,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
281 ):
282 # Drafting may require auxiliary hidden states from target model outputs
283 self.use_aux_hidden_state_outputs = True
283 - if self.use_pp:
284 - raise ValueError(
285 - f"{self.speculative_config.method} with pipeline parallel "
286 - "is not supported."
287 - )
284
285 # Draft tokens propagation - for spec-dec + struct outputs.
286 self.draft_tokens_handler = DraftTokensHandler(self.device)
@@ -407,6 +403,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
403 if self.use_aux_hidden_state_outputs:
404 assert self.speculative_config is not None
405 set_eagle3_aux_hidden_state_layers(self.model, self.speculative_config)
406 + if self.use_pp:
407 + assert self.speculative_config.method is not None
408 + verify_supports_aux_hidden_states_over_pp(
409 + self.model, self.speculative_config.method
410 + )
411 + assert self.pp_handler is not None
412 + self.pp_handler.configure_aux_hidden_state_relay(self.model)
413 if isinstance(self.speculator, DraftModelSpeculator):
414 with use_workspace_lane(self._draft_workspace_lane):
415 self.speculator.load_model(self.model)
@@ -1043,7 +1046,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1046 # For non-last PP ranks, update decode requests with sampler output from
1047 # the prior step in which they were scheduled (pp_size steps ago).
1048 if self.pp_handler is not None:
1046 - outputs = self.pp_handler.get_prev_sampled_outputs()
1049 + outputs = self.pp_handler.get_prev_sampled_outputs(
1050 + self.req_states.draft_tokens
1051 + )
1052 if outputs is not None:
1053 self.postprocess_sampled(**outputs)
1054
@@ -1879,7 +1884,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1884
1885 if not self.is_last_pp_rank:
1886 # Non-last PP rank: return IntermediateTensors for sending.
1882 - return output_intermediate_tensors
1887 + assert output_intermediate_tensors is not None
1888 + assert self.pp_handler is not None
1889 + return self.pp_handler.relay_aux_hidden_states(
1890 + model_inputs["intermediate_tensors"], output_intermediate_tensors
1891 + )
1892 return None
1893
1894 @torch.inference_mode()
@@ -2036,6 +2045,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2045 input_batch,
2046 self.req_states.draft_tokens[input_batch.idx_mapping],
2047 )
2048 + if self.pp_handler is not None:
2049 + self.pp_handler.broadcast_drafts(
2050 + self.req_states.draft_tokens, input_batch
2051 + )
2052
2053 # Post-step KV connector related operations.
2054 kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
vllm/v1/worker/gpu/pp_utils.py
+76 -2
@@ -10,6 +10,7 @@ import torch
10
11 from vllm.distributed.parallel_state import get_pp_group
12 from vllm.platforms import current_platform
13 +from vllm.sequence import IntermediateTensors
14 from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
15 from vllm.v1.worker.gpu.input_batch import InputBatch
16
@@ -30,6 +31,7 @@ class PendingRecv:
31 # Snapshot of slot generation counters at receive time, used to
32 # detect requests aborted since then.
33 gen_at_receive_np: np.ndarray # [num_reqs]
34 + draft_tokens: torch.Tensor | None = None # [num_reqs, num_speculative_steps]
35
36
37 def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None:
@@ -61,6 +63,7 @@ class PPHandler:
63 self.is_last_rank = get_pp_group().is_last_rank
64 self.last_rank = get_pp_group().last_rank
65 self.max_sample_len = num_speculative_steps + 1
66 + self.num_speculative_steps = num_speculative_steps
67 self.device = device
68 self.main_stream = torch.cuda.current_stream(device)
69 self.broadcast_stream = torch.cuda.Stream(device)
@@ -82,11 +85,37 @@ class PPHandler:
85 self.broadcast_group = get_pp_group().make_sibling_device_group(
86 group_desc="pp_broadcast"
87 )
88 + self.aux_hidden_state_relay_keys: tuple[str, ...] = ()
89
90 def on_req_idx_freed(self, req_idx: int) -> None:
91 self.req_idx_gen_np[req_idx] += 1
92
89 - def get_prev_sampled_outputs(self) -> dict[str, torch.Tensor] | None:
93 + def configure_aux_hidden_state_relay(self, model: torch.nn.Module) -> None:
94 + from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
95 + aux_hidden_state_relay_keys,
96 + )
97 +
98 + self.aux_hidden_state_relay_keys = aux_hidden_state_relay_keys(model)
99 +
100 + def relay_aux_hidden_states(
101 + self,
102 + intermediate_tensors: IntermediateTensors | None,
103 + output_intermediate_tensors: IntermediateTensors,
104 + ) -> IntermediateTensors:
105 + if not self.aux_hidden_state_relay_keys:
106 + return output_intermediate_tensors
107 + assert intermediate_tensors is not None
108 + return IntermediateTensors(
109 + output_intermediate_tensors.tensors
110 + | {
111 + key: intermediate_tensors[key]
112 + for key in self.aux_hidden_state_relay_keys
113 + }
114 + )
115 +
116 + def get_prev_sampled_outputs(
117 + self, draft_tokens_to_update: torch.Tensor | None = None
118 + ) -> dict[str, torch.Tensor] | None:
119 """Consume the entry from pp_size steps ago and wait for its recv event,
120 then filter out entries whose request was freed since `receive`.
121 """
@@ -112,6 +141,18 @@ class PPHandler:
141 idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device)
142
143 self.main_stream.wait_event(slot.event)
144 + if slot.draft_tokens is not None and draft_tokens_to_update is not None:
145 + draft_tokens = slot.draft_tokens
146 + draft_idx_mapping = slot.idx_mapping
147 + if exclude_mask.any():
148 + keep = ~exclude_mask
149 + keep_t = torch.as_tensor(keep, device=self.device)
150 + draft_tokens = draft_tokens[keep_t]
151 + draft_idx_mapping = async_copy_to_gpu(
152 + slot.idx_mapping_np[keep], device=self.device
153 + )
154 + draft_tokens_to_update[draft_idx_mapping] = draft_tokens
155 +
156 return dict(
157 sampled_tokens=slot.sampled_tokens,
158 num_sampled=slot.num_sampled,
@@ -119,6 +160,21 @@ class PPHandler:
160 idx_mapping=idx_mapping,
161 )
162
163 + def broadcast_drafts(
164 + self, draft_tokens: torch.Tensor, input_batch: InputBatch
165 + ) -> None:
166 + """Broadcast draft proposals so non-last ranks can embed real token ids."""
167 + assert self.is_last_rank
168 + if compute_need_sampled_mask(input_batch) is None:
169 + return
170 + with torch.cuda.stream(self.broadcast_stream):
171 + self.broadcast_stream.wait_stream(self.main_stream)
172 + send = draft_tokens[input_batch.idx_mapping].contiguous()
173 + torch.distributed.broadcast(
174 + send, src=self.last_rank, group=self.broadcast_group
175 + )
176 + send.record_stream(self.broadcast_stream)
177 +
178 def receive(self, input_batch: InputBatch) -> bool:
179 """Returns True iff sampled tokens need to be gathered from *all*
180 requests in the batch."""
@@ -145,12 +201,25 @@ class PPHandler:
201 torch.distributed.broadcast(
202 combined, src=self.last_rank, group=self.broadcast_group
203 )
204 + draft_tokens = None
205 + if self.num_speculative_steps > 0:
206 + draft_tokens = torch.empty(
207 + num_reqs,
208 + self.num_speculative_steps,
209 + dtype=torch.int64,
210 + device=self.device,
211 + )
212 + torch.distributed.broadcast(
213 + draft_tokens, src=self.last_rank, group=self.broadcast_group
214 + )
215 event = self.broadcast_stream.record_event()
216 num_sampled, num_rejected = combined.unbind(dim=0)
217 # Must record_stream since these were allocated on broadcast stream but
218 # later used on the main stream.
219 sampled_tokens.record_stream(self.main_stream)
220 combined.record_stream(self.main_stream)
221 + if draft_tokens is not None:
222 + draft_tokens.record_stream(self.main_stream)
223 self.queue[-1] = PendingRecv(
224 event,
225 sampled_tokens,
@@ -160,6 +229,7 @@ class PPHandler:
229 input_batch.idx_mapping_np,
230 need_sampled_mask,
231 gen_at_receive_np,
232 + draft_tokens,
233 )
234 return bool(need_sampled_mask.all())
235
@@ -182,8 +252,12 @@ class PPHandler:
252
253 with torch.cuda.stream(self.broadcast_stream):
254 self.broadcast_stream.wait_stream(self.main_stream)
255 + send_tokens = torch.nn.functional.pad(
256 + sampled_token_ids,
257 + (0, self.max_sample_len - sampled_token_ids.shape[-1]),
258 + )
259 torch.distributed.broadcast(
186 - sampled_token_ids.contiguous(),
260 + send_tokens.contiguous(),
261 src=self.last_rank,
262 group=self.broadcast_group,
263 )
vllm/v1/worker/gpu/spec_decode/dflash/utils.py
+2 -13
@@ -3,11 +3,11 @@
3 import torch.nn as nn
4
5 from vllm.config import VllmConfig, replace
6 -from vllm.distributed.parallel_state import get_pp_group
6 from vllm.model_executor.model_loader import get_model
7 from vllm.v1.worker.gpu.spec_decode.eagle.utils import (
8 _should_share,
9 get_target_lm_head,
10 + maybe_share_target_embed,
11 )
12
13
@@ -54,18 +54,7 @@ def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
54 target_inner = getattr(target_language_model, "model", target_language_model)
55 draft_inner = dflash_model.model
56
57 - # Skip embedding sharing under PP — each rank owns its own embedding.
58 - if get_pp_group().world_size == 1:
59 - target_embed = getattr(target_inner, "embed_tokens", None) or getattr(
60 - target_inner, "embedding", None
61 - )
62 - draft_embed = getattr(draft_inner, "embed_tokens", None)
63 - if target_embed is not None and _should_share(
64 - dflash_model, "has_own_embed_tokens", draft_embed, target_embed
65 - ):
66 - if draft_embed is not None:
67 - del draft_inner.embed_tokens
68 - draft_inner.embed_tokens = target_embed
57 + maybe_share_target_embed(dflash_model, draft_inner, target_inner)
58
59 target_lm_head = get_target_lm_head(target_model, target_language_model)
60 draft_lm_head = getattr(dflash_model, "lm_head", None)
vllm/v1/worker/gpu/spec_decode/dspark/utils.py
+4 -16
@@ -4,7 +4,6 @@
4 import torch.nn as nn
5
6 from vllm.config import ModelConfig, VllmConfig, replace
7 -from vllm.distributed.parallel_state import get_pp_group
7 from vllm.logger import init_logger
8 from vllm.v1.attention.backends.registry import AttentionBackendEnum
9
@@ -43,6 +42,7 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
42 from vllm.v1.worker.gpu.spec_decode.eagle.utils import (
43 _should_share,
44 get_target_lm_head,
45 + maybe_share_target_embed,
46 )
47
48 draft_attention_backend = _resolve_dspark_attention_backend(
@@ -53,6 +53,7 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
53
54 draft_vllm_config = replace(
55 vllm_config,
56 + parallel_config=speculative_config.draft_parallel_config,
57 attention_config=replace(
58 vllm_config.attention_config,
59 use_non_causal=dflash_has_any_non_causal(draft_model_config.hf_config),
@@ -76,9 +77,6 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
77 vllm_config=draft_vllm_config, model_config=draft_model_config
78 )
79
79 - if get_pp_group().world_size != 1:
80 - raise NotImplementedError("DSpark does not support pipeline parallelism.")
81 -
80 target_language_model = (
81 target_model.get_language_model()
82 if hasattr(target_model, "get_language_model")
@@ -88,18 +86,8 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
86 draft_inner = draft_model.model
87 target_vocab_size = vllm_config.model_config.get_vocab_size()
88
91 - target_embed = getattr(target_inner, "embed_tokens", None)
92 - draft_embed = getattr(draft_inner, "embed_tokens", None)
93 - if (
94 - target_embed is not None
95 - and draft_model_config.get_vocab_size() <= target_vocab_size
96 - and _should_share(
97 - draft_model, "has_own_embed_tokens", draft_embed, target_embed
98 - )
99 - ):
100 - if draft_embed is not None:
101 - del draft_inner.embed_tokens
102 - draft_inner.embed_tokens = target_embed
89 + if draft_model_config.get_vocab_size() <= target_vocab_size:
90 + maybe_share_target_embed(draft_model, draft_inner, target_inner)
91
92 target_lm_head = get_target_lm_head(target_model, target_language_model)
93 draft_lm_head = getattr(draft_model, "lm_head", None)
vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py
+67 -2
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from typing import cast
4
5 +import torch
6 import torch.nn as nn
7
8 from vllm.config import SpeculativeConfig
@@ -17,8 +18,6 @@ def set_eagle3_aux_hidden_state_layers(
18 ) -> None:
19 if not supports_eagle3(model):
20 raise RuntimeError("Model does not support EAGLE3 interface")
20 - # mypy may infer the class-level overload for supports_eagle3.
21 - # Narrow explicitly to the runtime protocol instance.
21 if isinstance(model, type):
22 raise RuntimeError("Expected model instance for EAGLE3 configuration")
23 eagle3_model = cast(SupportsEagle3, model)
@@ -30,6 +29,72 @@ def set_eagle3_aux_hidden_state_layers(
29 aux_layers = eagle3_model.get_eagle3_default_aux_hidden_state_layers()
30 logger.info("Using Eagle3 auxiliary layers from model: %s", aux_layers)
31 eagle3_model.set_aux_hidden_state_layers(aux_layers)
32 + reserve_aux_intermediate_tensor_slots(model)
33 +
34 +
35 +def _inner_decoder(model: nn.Module) -> nn.Module | None:
36 + parent_ref = model
37 + if hasattr(model, "get_language_model"):
38 + parent_ref = model.get_language_model()
39 + elif hasattr(model, "language_model"):
40 + parent_ref = model.language_model
41 + return getattr(parent_ref, "model", None)
42 +
43 +
44 +def verify_supports_aux_hidden_states_over_pp(model: nn.Module, method: str) -> None:
45 + inner = _inner_decoder(model)
46 + if not getattr(inner, "supports_aux_hidden_states_over_pp", False):
47 + raise ValueError(
48 + f"{type(model).__name__} does not support {method} with "
49 + "pipeline parallelism"
50 + )
51 +
52 +
53 +def aux_hidden_state_relay_keys(model: nn.Module) -> tuple[str, ...]:
54 + from vllm.distributed.parallel_state import get_pp_group
55 +
56 + pp = get_pp_group()
57 + if pp.world_size < 2 or pp.is_first_rank or pp.is_last_rank:
58 + return ()
59 + inner = _inner_decoder(model)
60 + assert inner is not None
61 + return tuple(
62 + f"{inner.AUX_HIDDEN_STATE_KEY}{i}" for i in range(inner._aux_slot_base_cached)
63 + )
64 +
65 +
66 +def reserve_aux_intermediate_tensor_slots(model: nn.Module) -> None:
67 + from vllm.distributed.parallel_state import (
68 + get_pp_group,
69 + model_parallel_is_initialized,
70 + )
71 +
72 + if not model_parallel_is_initialized():
73 + return
74 + pp = get_pp_group()
75 + if pp.world_size < 2 or pp.is_first_rank:
76 + return
77 + inner = _inner_decoder(model)
78 + if inner is None or not getattr(inner, "supports_aux_hidden_states_over_pp", False):
79 + return
80 +
81 + num_aux_states = inner._aux_slot_base_cached
82 + if num_aux_states == 0:
83 + return
84 +
85 + key = inner.AUX_HIDDEN_STATE_KEY
86 + hidden_size = inner.config.hidden_size
87 + make_empty = model.make_empty_intermediate_tensors
88 +
89 + def make_empty_with_aux(batch_size, dtype, device):
90 + tensors = make_empty(batch_size, dtype, device)
91 + for i in range(num_aux_states):
92 + tensors[f"{key}{i}"] = torch.zeros(
93 + (batch_size, hidden_size), dtype=dtype, device=device
94 + )
95 + return tensors
96 +
97 + model.make_empty_intermediate_tensors = make_empty_with_aux
98
99
100 def get_eagle3_aux_layers_from_config(
vllm/v1/worker/gpu/spec_decode/eagle/utils.py
+37 -19
@@ -7,6 +7,7 @@ from vllm.config import VllmConfig, replace
7 from vllm.distributed.parallel_state import get_pp_group
8 from vllm.lora.layers.base import BaseLayerWithLoRA
9 from vllm.model_executor.model_loader import get_model
10 +from vllm.model_executor.models.utils import PPMissingLayer
11
12
13 def _should_share(eagle: nn.Module, flag: str, draft, target) -> bool:
@@ -33,6 +34,41 @@ def get_target_lm_head(target_model: nn.Module, target_language_model: nn.Module
34 )
35
36
37 +def maybe_share_target_embed(
38 + draft_model: nn.Module, draft_inner: nn.Module, target_inner: nn.Module
39 +) -> None:
40 + """Share the target input embedding with the drafter when needed."""
41 + target_embed = getattr(target_inner, "embed_tokens", None) or getattr(
42 + target_inner, "embedding", None
43 + )
44 + if isinstance(target_embed, PPMissingLayer):
45 + target_embed = None
46 + # The drafter does not use the target's LoRA adapter.
47 + if isinstance(target_embed, BaseLayerWithLoRA):
48 + target_embed = target_embed.base_layer
49 + draft_embed = getattr(draft_inner, "embed_tokens", None)
50 +
51 + if get_pp_group().world_size > 1 and not hasattr(
52 + draft_model, "has_own_embed_tokens"
53 + ):
54 + return
55 +
56 + if target_embed is None:
57 + if hasattr(draft_inner, "embed_tokens") and not getattr(
58 + draft_model, "has_own_embed_tokens", False
59 + ):
60 + raise RuntimeError(
61 + f"{type(draft_model).__name__} needs the target input embedding, "
62 + "but it is unavailable on this PP stage"
63 + )
64 + return
65 +
66 + if _should_share(draft_model, "has_own_embed_tokens", draft_embed, target_embed):
67 + if draft_embed is not None:
68 + del draft_inner.embed_tokens
69 + draft_inner.embed_tokens = target_embed
70 +
71 +
72 def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module:
73 from vllm.compilation.backends import set_model_tag
74
@@ -70,25 +106,7 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod
106 target_inner = target_language_model.model
107 draft_inner = eagle_model.model
108
73 - # Skip embedding sharing under PP — each rank owns its own embedding.
74 - if get_pp_group().world_size == 1:
75 - target_embed = getattr(target_inner, "embed_tokens", None) or getattr(
76 - target_inner, "embedding", None
77 - )
78 - # If the target's embedding is LoRA-wrapped, share the underlying base
79 - # layer. The draft is not part of the LoRA adapter; sharing the wrapper
80 - # would make the draft run the LoRA embedding kernel with the target's
81 - # punica metadata (sized for the target's token count), causing an
82 - # out-of-bounds GPU access during multi-step draft decode.
83 - if isinstance(target_embed, BaseLayerWithLoRA):
84 - target_embed = target_embed.base_layer
85 - draft_embed = getattr(draft_inner, "embed_tokens", None)
86 - if target_embed is not None and _should_share(
87 - eagle_model, "has_own_embed_tokens", draft_embed, target_embed
88 - ):
89 - if draft_embed is not None:
90 - del draft_inner.embed_tokens
91 - draft_inner.embed_tokens = target_embed
109 + maybe_share_target_embed(eagle_model, draft_inner, target_inner)
110
111 target_lm_head = get_target_lm_head(target_model, target_language_model)
112 draft_lm_head = getattr(eagle_model, "lm_head", None)