diff --git a/tests/model_executor/test_qwen3_omni.py b/tests/model_executor/test_qwen3_omni.py index 708460c92c84..dea87a122694 100644 --- a/tests/model_executor/test_qwen3_omni.py +++ b/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( diff --git a/tests/models/kimi_k3/test_aux_attn_res_stream.py b/tests/models/kimi_k3/test_aux_attn_res_stream.py index 7227ed6fd9ea..26cb76a55039 100644 --- a/tests/models/kimi_k3/test_aux_attn_res_stream.py +++ b/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( diff --git a/tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py b/tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py new file mode 100644 index 000000000000..c29666c97179 --- /dev/null +++ b/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}" + ) diff --git a/tests/v1/worker/test_eagle3_aux_hidden_states_pp.py b/tests/v1/worker/test_eagle3_aux_hidden_states_pp.py new file mode 100644 index 000000000000..b22da39c91c4 --- /dev/null +++ b/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") diff --git a/tests/v1/worker/test_spec_decode_embed_sharing_pp.py b/tests/v1/worker/test_spec_decode_embed_sharing_pp.py new file mode 100644 index 000000000000..1a1d96a6b4ef --- /dev/null +++ b/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 diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 49823bc03112..5a232d848069 100644 --- a/vllm/config/speculative.py +++ b/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, diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 25620e036752..a59e5bfc38f3 100644 --- a/vllm/config/vllm.py +++ b/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()) diff --git a/vllm/model_executor/models/deepseek_eagle3.py b/vllm/model_executor/models/deepseek_eagle3.py index 63673093439d..80c41501e6e9 100644 --- a/vllm/model_executor/models/deepseek_eagle3.py +++ b/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 diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 2404e44eb83d..a9a9bc245096 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/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): diff --git a/vllm/model_executor/models/laguna_dflash.py b/vllm/model_executor/models/laguna_dflash.py index 3c46a4cbb506..02af55c00033 100644 --- a/vllm/model_executor/models/laguna_dflash.py +++ b/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: diff --git a/vllm/model_executor/models/llama.py b/vllm/model_executor/models/llama.py index 7f2526f698a8..9e0f45388ab7 100644 --- a/vllm/model_executor/models/llama.py +++ b/vllm/model_executor/models/llama.py @@ -75,6 +75,7 @@ make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, + spec_decode_needs_target_embed, ) @@ -354,6 +355,7 @@ class LlamaModel(nn.Module, EagleModelMixin): ".up_proj": (".gate_up_proj", 1), } ) + supports_aux_hidden_states_over_pp = True def __init__( self, @@ -372,8 +374,10 @@ def __init__( self.vocab_size = config.vocab_size - if get_pp_group().is_first_rank or ( - config.tie_word_embeddings and get_pp_group().is_last_rank + if ( + get_pp_group().is_first_rank + or (config.tie_word_embeddings and get_pp_group().is_last_rank) + or spec_decode_needs_target_embed(vllm_config) ): self.embed_tokens = VocabParallelEmbedding( self.vocab_size, @@ -418,9 +422,16 @@ def forward( hidden_states = intermediate_tensors["hidden_states"] residual = intermediate_tensors["residual"] - aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual) + remote_aux = self.collect_remote_aux_hidden_states(intermediate_tensors) + + aux_hidden_states: list[torch.Tensor] = [] + if get_pp_group().is_first_rank: + self._maybe_add_hidden_state( + aux_hidden_states, self.start_layer, hidden_states, residual + ) for idx, layer in enumerate( - islice(self.layers, self.start_layer, self.end_layer) + islice(self.layers, self.start_layer, self.end_layer), + start=self.start_layer, ): hidden_states, residual = layer( positions, hidden_states, residual, **extra_layer_kwargs @@ -431,11 +442,16 @@ def forward( if not get_pp_group().is_last_rank: return IntermediateTensors( - {"hidden_states": hidden_states, "residual": residual} + { + "hidden_states": hidden_states, + "residual": residual, + **self.pack_local_aux_hidden_states(aux_hidden_states), + } ) hidden_states, _ = self.norm(hidden_states, residual) + aux_hidden_states = remote_aux + aux_hidden_states if len(aux_hidden_states) > 0: return hidden_states, aux_hidden_states return hidden_states diff --git a/vllm/model_executor/models/llama_eagle3.py b/vllm/model_executor/models/llama_eagle3.py index fab611cf317e..c3d956a946d1 100644 --- a/vllm/model_executor/models/llama_eagle3.py +++ b/vllm/model_executor/models/llama_eagle3.py @@ -314,9 +314,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): if getattr(self.config, "draft_vocab_size", None) is None: base_vocab_size = getattr(self.config, "vocab_size", None) self.config.draft_vocab_size = base_vocab_size - target_layer_num = vllm_config.model_config.get_num_layers( - vllm_config.parallel_config - ) + target_layer_num = vllm_config.model_config.get_total_num_hidden_layers() # Store target layer count in draft config for # proper layer_types indexing in draft models diff --git a/vllm/model_executor/models/mimo.py b/vllm/model_executor/models/mimo.py index 6de0c0e5967d..54151fb9dd5d 100644 --- a/vllm/model_executor/models/mimo.py +++ b/vllm/model_executor/models/mimo.py @@ -55,6 +55,8 @@ } ) class MiMoModel(Qwen2Model): + supports_aux_hidden_states_over_pp = False + def forward( self, input_ids: torch.Tensor | None, diff --git a/vllm/model_executor/models/qwen2.py b/vllm/model_executor/models/qwen2.py index 3820f5e39b42..7513936ebbe6 100644 --- a/vllm/model_executor/models/qwen2.py +++ b/vllm/model_executor/models/qwen2.py @@ -74,6 +74,7 @@ make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, + spec_decode_needs_target_embed, ) @@ -330,6 +331,7 @@ class Qwen2Model(nn.Module, EagleModelMixin): ".up_proj": (".gate_up_proj", 1), } ) + supports_aux_hidden_states_over_pp = True def __init__( self, @@ -360,8 +362,10 @@ def __init__( self.quant_config = quant_config self.vocab_size = config.vocab_size - if get_pp_group().is_first_rank or ( - config.tie_word_embeddings and get_pp_group().is_last_rank + if ( + get_pp_group().is_first_rank + or (config.tie_word_embeddings and get_pp_group().is_last_rank) + or spec_decode_needs_target_embed(vllm_config) ): self.embed_tokens = VocabParallelEmbedding( config.vocab_size, @@ -412,9 +416,16 @@ def forward( hidden_states = intermediate_tensors["hidden_states"] residual = intermediate_tensors["residual"] - aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual) + remote_aux = self.collect_remote_aux_hidden_states(intermediate_tensors) + + aux_hidden_states: list[torch.Tensor] = [] + if get_pp_group().is_first_rank: + self._maybe_add_hidden_state( + aux_hidden_states, self.start_layer, hidden_states, residual + ) for idx, layer in enumerate( - islice(self.layers, self.start_layer, self.end_layer) + islice(self.layers, self.start_layer, self.end_layer), + start=self.start_layer, ): hidden_states, residual = layer(positions, hidden_states, residual) self._maybe_add_hidden_state( @@ -423,11 +434,16 @@ def forward( if not get_pp_group().is_last_rank: return IntermediateTensors( - {"hidden_states": hidden_states, "residual": residual} + { + "hidden_states": hidden_states, + "residual": residual, + **self.pack_local_aux_hidden_states(aux_hidden_states), + } ) hidden_states, _ = self.norm(hidden_states, residual) + aux_hidden_states = remote_aux + aux_hidden_states if len(aux_hidden_states) > 0: return hidden_states, aux_hidden_states diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 0a481b38238d..ce3884a10eb8 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -694,9 +694,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.config = self.draft_model_config.hf_config if getattr(self.config, "draft_vocab_size", None) is None: self.config.draft_vocab_size = getattr(self.config, "vocab_size", None) - 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.model = self.model_cls( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"), diff --git a/vllm/model_executor/models/qwen3_dspark.py b/vllm/model_executor/models/qwen3_dspark.py index 2cee016124c7..7c37c6db368e 100644 --- a/vllm/model_executor/models/qwen3_dspark.py +++ b/vllm/model_executor/models/qwen3_dspark.py @@ -182,9 +182,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.config = self.draft_model_config.hf_config if getattr(self.config, "draft_vocab_size", None) is None: self.config.draft_vocab_size = getattr(self.config, "vocab_size", None) - 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.model = Qwen3DSparkModel( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"), diff --git a/vllm/model_executor/models/qwen3_eagle3.py b/vllm/model_executor/models/qwen3_eagle3.py index 5b4e5a7a57cd..39d0173b3ce9 100644 --- a/vllm/model_executor/models/qwen3_eagle3.py +++ b/vllm/model_executor/models/qwen3_eagle3.py @@ -288,9 +288,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): if getattr(self.config, "draft_vocab_size", None) is None: base_vocab_size = getattr(self.config, "vocab_size", None) self.config.draft_vocab_size = base_vocab_size - target_layer_num = vllm_config.model_config.get_num_layers( - vllm_config.parallel_config - ) + target_layer_num = vllm_config.model_config.get_total_num_hidden_layers() # Store target layer count in draft config for # proper layer_types indexing in draft models diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index f2da8a5109e3..84bc8a2d4b43 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -835,6 +835,22 @@ def forward(self, *args, **kwargs): return args[0] if args else next(iter(kwargs.values())) +def spec_decode_needs_target_embed(vllm_config: VllmConfig) -> bool: + """Whether the last PP rank needs the target input embedding.""" + from vllm.distributed.parallel_state import get_pp_group + + speculative_config = vllm_config.speculative_config + if speculative_config is None or speculative_config.method not in ( + "eagle", + "eagle3", + "dflash", + "dspark", + ): + return False + pp = get_pp_group() + return pp.world_size > 1 and pp.is_last_rank + + def make_layers( num_hidden_layers: int, layer_fn: LayerFn, diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 3247877db071..be4e87e77642 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -67,6 +67,7 @@ is_pp_missing_parameter, make_layers, maybe_prefix, + spec_decode_needs_target_embed, ) from vllm.model_executor.utils import set_weight_attrs from vllm.models.common.ops.sequence_parallel import ( @@ -1270,6 +1271,8 @@ def forward( class DeepseekV4Model(nn.Module, EagleModelMixin): + supports_aux_hidden_states_over_pp = True + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -1305,7 +1308,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): dtype=torch.int32, ) - if get_pp_group().is_first_rank: + if get_pp_group().is_first_rank or spec_decode_needs_target_embed(vllm_config): self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size, @@ -1416,6 +1419,7 @@ def forward( input_ids = sp_shard(input_ids) residual, post_mix, res_mix = None, None, None + remote_aux = self.collect_remote_aux_hidden_states(intermediate_tensors) aux_hidden_states: list[torch.Tensor] = [] final_aux_recon: torch.Tensor | None = None # avoid duplicate mhc_post call for idx, layer in enumerate( @@ -1450,7 +1454,12 @@ def forward( ) if not get_pp_group().is_last_rank: - return IntermediateTensors({"hidden_states": hidden_states}) + return IntermediateTensors( + { + "hidden_states": hidden_states, + **self.pack_local_aux_hidden_states(aux_hidden_states), + } + ) if self.use_sequence_parallel: hidden_states = sp_all_gather(hidden_states)[:full_num_tokens] @@ -1468,6 +1477,7 @@ def forward( self.hc_eps, ) hidden_states = self.norm(hidden_states) + aux_hidden_states = remote_aux + aux_hidden_states if len(aux_hidden_states) > 0: return hidden_states, aux_hidden_states return hidden_states diff --git a/vllm/models/kimi_k3/nvidia/dspark_mla.py b/vllm/models/kimi_k3/nvidia/dspark_mla.py index 5008015abdc8..97febda0c86e 100644 --- a/vllm/models/kimi_k3/nvidia/dspark_mla.py +++ b/vllm/models/kimi_k3/nvidia/dspark_mla.py @@ -421,9 +421,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: assert vllm_config.speculative_config is not None self.draft_model_config = vllm_config.speculative_config.draft_model_config self.config = self.draft_model_config.hf_config - 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.model = K3DSparkModel( vllm_config=vllm_config, start_layer_id=target_layer_num, diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 2fc20ecf2198..2115cdb71ea0 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -85,6 +85,7 @@ is_pp_missing_parameter, make_layers, maybe_prefix, + spec_decode_needs_target_embed, ) from vllm.model_executor.models.vision import is_vit_use_data_parallel from vllm.models.common.ops.sequence_parallel import ( @@ -1124,6 +1125,8 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant): "fused_qkv_a_g_proj": ["q_a_proj", "kv_a_proj_with_mqa", "g_proj"], } + supports_aux_hidden_states_over_pp = True + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -1148,7 +1151,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): vllm_config, self.use_sequence_parallel ) - if get_pp_group().is_first_rank: + if get_pp_group().is_first_rank or spec_decode_needs_target_embed(vllm_config): self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size, @@ -1231,15 +1234,18 @@ def make_empty_intermediate_tensors( def _set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: super()._set_aux_hidden_state_layers(layers) + if self.use_attn_res and self._aux_attn_res_stream: + pp = get_pp_group() + if not pp.is_last_rank and self.end_layer in self.aux_hidden_state_layers: + raise ValueError( + f"Auxiliary layer {self.end_layer} cannot end a non-final PP " + "stage when VLLM_KIMI_K3_AUX_ATTN_RES_STREAM=1" + ) if self.use_attn_res: - # Emitted once, at configuration time. Which layers are tapped and - # which convention is in force are the two things you need to - # confirm from a running process, and neither is recoverable from - # the served output. logger.info_once( "Kimi-K3 aux hidden capture: layers=%s mode=%s " "(VLLM_KIMI_K3_AUX_ATTN_RES_STREAM=%d)", - layers, + self.aux_hidden_state_layers, "attn_res_stream" if self._aux_attn_res_stream else "prefix_only", int(self._aux_attn_res_stream), ) @@ -1255,25 +1261,8 @@ def _capture_aux_hidden_stream( pending_mlp_out: torch.Tensor | None, block_residual: torch.Tensor, ) -> torch.Tensor: - """Auxiliary feature tapped after ``layer_idx`` under AttnRes. - - The wire between layers only carries the current block's running prefix; - the committed blocks live in the bank. The value the next consumer - actually reads is the pre-norm AttnRes mixture over - ``bank[:num_blocks] + prefix``, which is what the DFlash drafters were - trained against. ``attn_res`` with no delta, no block write and no - output norm computes exactly that and leaves both the prefix and the - bank untouched. - - Folding the pending MLP output into the prefix rather than passing it as - ``delta`` is deliberate: the kernel writes an applied delta back into - the prefix in place, which would double-add it into the live residual - stream. - """ + """Return the AttnRes stream after ``layer_idx``.""" prefix = prefix_sum if pending_mlp_out is None else prefix_sum + pending_mlp_out - # `use_attn_res` is what constructs the norm and projection weights this - # reads; without it there is no mixture to compute and the attribute - # lookups below would raise. if not (self._aux_attn_res_stream and self.use_attn_res): return prefix @@ -1283,17 +1272,11 @@ def _capture_aux_hidden_stream( score_proj = consumer.self_attention_res_proj num_blocks = consumer.prev_valid_blocks elif get_pp_group().is_last_rank: - # Nothing downstream but the model's own output-side aggregation. score_norm = self.output_attn_res_norm score_proj = self.output_attn_res_proj num_blocks = self.num_attn_res_blocks else: - # Last layer of a non-final pipeline stage: the consumer lives on - # the next rank and the output-side aggregation only exists on the - # last one, so there is nothing here to mix against. Falling back - # to the running prefix keeps the tap defined rather than reaching - # for weights this rank does not construct. - return prefix + raise RuntimeError("Auxiliary AttnRes capture crossed a PP boundary") return attn_res( prefix, @@ -1341,9 +1324,14 @@ def forward( hidden_states = sp_shard(hidden_states) assert residual is None, "Currently, SP is not supported with PP" + remote_aux = self.collect_remote_aux_hidden_states(intermediate_tensors) + # sharded aux hidden states when sp is enabled aux_hidden_states: list[torch.Tensor] = [] - if self.start_layer in self.aux_hidden_state_layers: + if ( + get_pp_group().is_first_rank + and self.start_layer in self.aux_hidden_state_layers + ): if self.use_attn_res or residual is None: aux_hidden_states.append(hidden_states) else: @@ -1394,7 +1382,11 @@ def forward( if prefix_sum is not None: hidden_states = hidden_states + prefix_sum return IntermediateTensors( - {"hidden_states": hidden_states, "residual": residual} + { + "hidden_states": hidden_states, + "residual": residual, + **self.pack_local_aux_hidden_states(aux_hidden_states), + } ) if self.use_attn_res: @@ -1431,6 +1423,7 @@ def forward( # NOTE: the final norm is applied in compute_logits instead of here, so # the MTP draft model receives the pre-norm hidden states. + aux_hidden_states = remote_aux + aux_hidden_states if aux_hidden_states: return hidden_states, aux_hidden_states return hidden_states diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 901bd687546c..9866724d1a00 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -152,6 +152,7 @@ ) from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( set_eagle3_aux_hidden_state_layers, + verify_supports_aux_hidden_states_over_pp, ) from vllm.v1.worker.gpu.spec_decode.rejection_sampler import ( RejectionSampler, @@ -280,11 +281,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): ): # Drafting may require auxiliary hidden states from target model outputs self.use_aux_hidden_state_outputs = True - if self.use_pp: - raise ValueError( - f"{self.speculative_config.method} with pipeline parallel " - "is not supported." - ) # Draft tokens propagation - for spec-dec + struct outputs. self.draft_tokens_handler = DraftTokensHandler(self.device) @@ -407,6 +403,13 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: if self.use_aux_hidden_state_outputs: assert self.speculative_config is not None set_eagle3_aux_hidden_state_layers(self.model, self.speculative_config) + if self.use_pp: + assert self.speculative_config.method is not None + verify_supports_aux_hidden_states_over_pp( + self.model, self.speculative_config.method + ) + assert self.pp_handler is not None + self.pp_handler.configure_aux_hidden_state_relay(self.model) if isinstance(self.speculator, DraftModelSpeculator): with use_workspace_lane(self._draft_workspace_lane): self.speculator.load_model(self.model) @@ -1038,7 +1041,9 @@ def update_pp_decode_requests(self): # For non-last PP ranks, update decode requests with sampler output from # the prior step in which they were scheduled (pp_size steps ago). if self.pp_handler is not None: - outputs = self.pp_handler.get_prev_sampled_outputs() + outputs = self.pp_handler.get_prev_sampled_outputs( + self.req_states.draft_tokens + ) if outputs is not None: self.postprocess_sampled(**outputs) @@ -1874,7 +1879,11 @@ def execute_model( if not self.is_last_pp_rank: # Non-last PP rank: return IntermediateTensors for sending. - return output_intermediate_tensors + assert output_intermediate_tensors is not None + assert self.pp_handler is not None + return self.pp_handler.relay_aux_hidden_states( + model_inputs["intermediate_tensors"], output_intermediate_tensors + ) return None @torch.inference_mode() @@ -2031,6 +2040,10 @@ def sample_tokens( input_batch, self.req_states.draft_tokens[input_batch.idx_mapping], ) + if self.pp_handler is not None: + self.pp_handler.broadcast_drafts( + self.req_states.draft_tokens, input_batch + ) # Post-step KV connector related operations. kv_connector_output = self.kv_connector.post_forward(finished_req_ids) diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py index cf52a6d3821e..33386c922270 100644 --- a/vllm/v1/worker/gpu/pp_utils.py +++ b/vllm/v1/worker/gpu/pp_utils.py @@ -10,6 +10,7 @@ from vllm.distributed.parallel_state import get_pp_group from vllm.platforms import current_platform +from vllm.sequence import IntermediateTensors from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu from vllm.v1.worker.gpu.input_batch import InputBatch @@ -30,6 +31,7 @@ class PendingRecv: # Snapshot of slot generation counters at receive time, used to # detect requests aborted since then. gen_at_receive_np: np.ndarray # [num_reqs] + draft_tokens: torch.Tensor | None = None # [num_reqs, num_speculative_steps] def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None: @@ -61,6 +63,7 @@ def __init__( self.is_last_rank = get_pp_group().is_last_rank self.last_rank = get_pp_group().last_rank self.max_sample_len = num_speculative_steps + 1 + self.num_speculative_steps = num_speculative_steps self.device = device self.main_stream = torch.cuda.current_stream(device) self.broadcast_stream = torch.cuda.Stream(device) @@ -82,11 +85,37 @@ def __init__( self.broadcast_group = get_pp_group().make_sibling_device_group( group_desc="pp_broadcast" ) + self.aux_hidden_state_relay_keys: tuple[str, ...] = () def on_req_idx_freed(self, req_idx: int) -> None: self.req_idx_gen_np[req_idx] += 1 - def get_prev_sampled_outputs(self) -> dict[str, torch.Tensor] | None: + def configure_aux_hidden_state_relay(self, model: torch.nn.Module) -> None: + from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( + aux_hidden_state_relay_keys, + ) + + self.aux_hidden_state_relay_keys = aux_hidden_state_relay_keys(model) + + def relay_aux_hidden_states( + self, + intermediate_tensors: IntermediateTensors | None, + output_intermediate_tensors: IntermediateTensors, + ) -> IntermediateTensors: + if not self.aux_hidden_state_relay_keys: + return output_intermediate_tensors + assert intermediate_tensors is not None + return IntermediateTensors( + output_intermediate_tensors.tensors + | { + key: intermediate_tensors[key] + for key in self.aux_hidden_state_relay_keys + } + ) + + def get_prev_sampled_outputs( + self, draft_tokens_to_update: torch.Tensor | None = None + ) -> dict[str, torch.Tensor] | None: """Consume the entry from pp_size steps ago and wait for its recv event, then filter out entries whose request was freed since `receive`. """ @@ -112,6 +141,18 @@ def get_prev_sampled_outputs(self) -> dict[str, torch.Tensor] | None: idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device) self.main_stream.wait_event(slot.event) + if slot.draft_tokens is not None and draft_tokens_to_update is not None: + draft_tokens = slot.draft_tokens + draft_idx_mapping = slot.idx_mapping + if exclude_mask.any(): + keep = ~exclude_mask + keep_t = torch.as_tensor(keep, device=self.device) + draft_tokens = draft_tokens[keep_t] + draft_idx_mapping = async_copy_to_gpu( + slot.idx_mapping_np[keep], device=self.device + ) + draft_tokens_to_update[draft_idx_mapping] = draft_tokens + return dict( sampled_tokens=slot.sampled_tokens, num_sampled=slot.num_sampled, @@ -119,6 +160,21 @@ def get_prev_sampled_outputs(self) -> dict[str, torch.Tensor] | None: idx_mapping=idx_mapping, ) + def broadcast_drafts( + self, draft_tokens: torch.Tensor, input_batch: InputBatch + ) -> None: + """Broadcast draft proposals so non-last ranks can embed real token ids.""" + assert self.is_last_rank + if compute_need_sampled_mask(input_batch) is None: + return + with torch.cuda.stream(self.broadcast_stream): + self.broadcast_stream.wait_stream(self.main_stream) + send = draft_tokens[input_batch.idx_mapping].contiguous() + torch.distributed.broadcast( + send, src=self.last_rank, group=self.broadcast_group + ) + send.record_stream(self.broadcast_stream) + def receive(self, input_batch: InputBatch) -> bool: """Returns True iff sampled tokens need to be gathered from *all* requests in the batch.""" @@ -145,12 +201,25 @@ def receive(self, input_batch: InputBatch) -> bool: torch.distributed.broadcast( combined, src=self.last_rank, group=self.broadcast_group ) + draft_tokens = None + if self.num_speculative_steps > 0: + draft_tokens = torch.empty( + num_reqs, + self.num_speculative_steps, + dtype=torch.int64, + device=self.device, + ) + torch.distributed.broadcast( + draft_tokens, src=self.last_rank, group=self.broadcast_group + ) event = self.broadcast_stream.record_event() num_sampled, num_rejected = combined.unbind(dim=0) # Must record_stream since these were allocated on broadcast stream but # later used on the main stream. sampled_tokens.record_stream(self.main_stream) combined.record_stream(self.main_stream) + if draft_tokens is not None: + draft_tokens.record_stream(self.main_stream) self.queue[-1] = PendingRecv( event, sampled_tokens, @@ -160,6 +229,7 @@ def receive(self, input_batch: InputBatch) -> bool: input_batch.idx_mapping_np, need_sampled_mask, gen_at_receive_np, + draft_tokens, ) return bool(need_sampled_mask.all()) @@ -182,8 +252,12 @@ def broadcast( with torch.cuda.stream(self.broadcast_stream): self.broadcast_stream.wait_stream(self.main_stream) + send_tokens = torch.nn.functional.pad( + sampled_token_ids, + (0, self.max_sample_len - sampled_token_ids.shape[-1]), + ) torch.distributed.broadcast( - sampled_token_ids.contiguous(), + send_tokens.contiguous(), src=self.last_rank, group=self.broadcast_group, ) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py index 562966e07eed..6226dea1c2db 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py @@ -3,11 +3,11 @@ import torch.nn as nn from vllm.config import VllmConfig, replace -from vllm.distributed.parallel_state import get_pp_group from vllm.model_executor.model_loader import get_model from vllm.v1.worker.gpu.spec_decode.eagle.utils import ( _should_share, get_target_lm_head, + maybe_share_target_embed, ) @@ -54,18 +54,7 @@ def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo target_inner = getattr(target_language_model, "model", target_language_model) draft_inner = dflash_model.model - # Skip embedding sharing under PP — each rank owns its own embedding. - if get_pp_group().world_size == 1: - target_embed = getattr(target_inner, "embed_tokens", None) or getattr( - target_inner, "embedding", None - ) - draft_embed = getattr(draft_inner, "embed_tokens", None) - if target_embed is not None and _should_share( - dflash_model, "has_own_embed_tokens", draft_embed, target_embed - ): - if draft_embed is not None: - del draft_inner.embed_tokens - draft_inner.embed_tokens = target_embed + maybe_share_target_embed(dflash_model, draft_inner, target_inner) target_lm_head = get_target_lm_head(target_model, target_language_model) draft_lm_head = getattr(dflash_model, "lm_head", None) diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index 177060747ab8..00b50912ef3b 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -4,7 +4,6 @@ import torch.nn as nn from vllm.config import ModelConfig, VllmConfig, replace -from vllm.distributed.parallel_state import get_pp_group from vllm.logger import init_logger from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -43,6 +42,7 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo from vllm.v1.worker.gpu.spec_decode.eagle.utils import ( _should_share, get_target_lm_head, + maybe_share_target_embed, ) draft_attention_backend = _resolve_dspark_attention_backend( @@ -53,6 +53,7 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo draft_vllm_config = replace( vllm_config, + parallel_config=speculative_config.draft_parallel_config, attention_config=replace( vllm_config.attention_config, 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 vllm_config=draft_vllm_config, model_config=draft_model_config ) - if get_pp_group().world_size != 1: - raise NotImplementedError("DSpark does not support pipeline parallelism.") - target_language_model = ( target_model.get_language_model() if hasattr(target_model, "get_language_model") @@ -88,18 +86,8 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo draft_inner = draft_model.model target_vocab_size = vllm_config.model_config.get_vocab_size() - target_embed = getattr(target_inner, "embed_tokens", None) - draft_embed = getattr(draft_inner, "embed_tokens", None) - if ( - target_embed is not None - and draft_model_config.get_vocab_size() <= target_vocab_size - and _should_share( - draft_model, "has_own_embed_tokens", draft_embed, target_embed - ) - ): - if draft_embed is not None: - del draft_inner.embed_tokens - draft_inner.embed_tokens = target_embed + if draft_model_config.get_vocab_size() <= target_vocab_size: + maybe_share_target_embed(draft_model, draft_inner, target_inner) target_lm_head = get_target_lm_head(target_model, target_language_model) draft_lm_head = getattr(draft_model, "lm_head", None) diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py index a37a4c83b3e5..0dca4588c332 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import cast +import torch import torch.nn as nn from vllm.config import SpeculativeConfig @@ -17,8 +18,6 @@ def set_eagle3_aux_hidden_state_layers( ) -> None: if not supports_eagle3(model): raise RuntimeError("Model does not support EAGLE3 interface") - # mypy may infer the class-level overload for supports_eagle3. - # Narrow explicitly to the runtime protocol instance. if isinstance(model, type): raise RuntimeError("Expected model instance for EAGLE3 configuration") eagle3_model = cast(SupportsEagle3, model) @@ -30,6 +29,72 @@ def set_eagle3_aux_hidden_state_layers( aux_layers = eagle3_model.get_eagle3_default_aux_hidden_state_layers() logger.info("Using Eagle3 auxiliary layers from model: %s", aux_layers) eagle3_model.set_aux_hidden_state_layers(aux_layers) + reserve_aux_intermediate_tensor_slots(model) + + +def _inner_decoder(model: nn.Module) -> nn.Module | None: + parent_ref = model + if hasattr(model, "get_language_model"): + parent_ref = model.get_language_model() + elif hasattr(model, "language_model"): + parent_ref = model.language_model + return getattr(parent_ref, "model", None) + + +def verify_supports_aux_hidden_states_over_pp(model: nn.Module, method: str) -> None: + inner = _inner_decoder(model) + if not getattr(inner, "supports_aux_hidden_states_over_pp", False): + raise ValueError( + f"{type(model).__name__} does not support {method} with " + "pipeline parallelism" + ) + + +def aux_hidden_state_relay_keys(model: nn.Module) -> tuple[str, ...]: + from vllm.distributed.parallel_state import get_pp_group + + pp = get_pp_group() + if pp.world_size < 2 or pp.is_first_rank or pp.is_last_rank: + return () + inner = _inner_decoder(model) + assert inner is not None + return tuple( + f"{inner.AUX_HIDDEN_STATE_KEY}{i}" for i in range(inner._aux_slot_base_cached) + ) + + +def reserve_aux_intermediate_tensor_slots(model: nn.Module) -> 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 or pp.is_first_rank: + return + inner = _inner_decoder(model) + if inner is None or not getattr(inner, "supports_aux_hidden_states_over_pp", False): + return + + num_aux_states = inner._aux_slot_base_cached + if num_aux_states == 0: + return + + key = inner.AUX_HIDDEN_STATE_KEY + hidden_size = inner.config.hidden_size + make_empty = model.make_empty_intermediate_tensors + + def make_empty_with_aux(batch_size, dtype, device): + tensors = make_empty(batch_size, dtype, device) + for i in range(num_aux_states): + tensors[f"{key}{i}"] = torch.zeros( + (batch_size, hidden_size), dtype=dtype, device=device + ) + return tensors + + model.make_empty_intermediate_tensors = make_empty_with_aux def get_eagle3_aux_layers_from_config( diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py index f1d8bd42e1d3..05f771ab8298 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py @@ -7,6 +7,7 @@ from vllm.distributed.parallel_state import get_pp_group from vllm.lora.layers.base import BaseLayerWithLoRA from vllm.model_executor.model_loader import get_model +from vllm.model_executor.models.utils import PPMissingLayer 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 ) +def maybe_share_target_embed( + draft_model: nn.Module, draft_inner: nn.Module, target_inner: nn.Module +) -> None: + """Share the target input embedding with the drafter when needed.""" + target_embed = getattr(target_inner, "embed_tokens", None) or getattr( + target_inner, "embedding", None + ) + if isinstance(target_embed, PPMissingLayer): + target_embed = None + # The drafter does not use the target's LoRA adapter. + if isinstance(target_embed, BaseLayerWithLoRA): + target_embed = target_embed.base_layer + draft_embed = getattr(draft_inner, "embed_tokens", None) + + if get_pp_group().world_size > 1 and not hasattr( + draft_model, "has_own_embed_tokens" + ): + return + + if target_embed is None: + if hasattr(draft_inner, "embed_tokens") and not getattr( + draft_model, "has_own_embed_tokens", False + ): + raise RuntimeError( + f"{type(draft_model).__name__} needs the target input embedding, " + "but it is unavailable on this PP stage" + ) + return + + if _should_share(draft_model, "has_own_embed_tokens", draft_embed, target_embed): + if draft_embed is not None: + del draft_inner.embed_tokens + draft_inner.embed_tokens = target_embed + + def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module: from vllm.compilation.backends import set_model_tag @@ -70,25 +106,7 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod target_inner = target_language_model.model draft_inner = eagle_model.model - # Skip embedding sharing under PP — each rank owns its own embedding. - if get_pp_group().world_size == 1: - target_embed = getattr(target_inner, "embed_tokens", None) or getattr( - target_inner, "embedding", None - ) - # If the target's embedding is LoRA-wrapped, share the underlying base - # layer. The draft is not part of the LoRA adapter; sharing the wrapper - # would make the draft run the LoRA embedding kernel with the target's - # punica metadata (sized for the target's token count), causing an - # out-of-bounds GPU access during multi-step draft decode. - if isinstance(target_embed, BaseLayerWithLoRA): - target_embed = target_embed.base_layer - draft_embed = getattr(draft_inner, "embed_tokens", None) - if target_embed is not None and _should_share( - eagle_model, "has_own_embed_tokens", draft_embed, target_embed - ): - if draft_embed is not None: - del draft_inner.embed_tokens - draft_inner.embed_tokens = target_embed + maybe_share_target_embed(eagle_model, draft_inner, target_inner) target_lm_head = get_target_lm_head(target_model, target_language_model) draft_lm_head = getattr(eagle_model, "lm_head", None)