From aeced950ec2d70f186d0b293ae8a1689f915a26b Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Sun, 8 Nov 2026 21:56:13 +0000
Subject: [PATCH 01/10] [Transformers backend] Support EAGLE3 aux hidden states
 with pipeline parallelism

Relay aux hidden states between PP ranks the same way as native Llama, keep
the input embedding on the last rank when the drafter may share it, and record
aux hidden states from the real decoder layer class on ranks whose last layer
is missing.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 .../e2e/spec_decode/eagle/test_eagle3_pp.py   | 11 ++++---
 .../models/transformers/base.py               | 32 +++++++++++++++----
 .../models/transformers/multimodal.py         |  3 ++
 .../gpu/spec_decode/eagle/eagle3_utils.py     |  8 ++++-
 4 files changed, 42 insertions(+), 12 deletions(-)

diff --git a/tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py b/tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py
index c29666c97179..3d6fb47c1565 100644
--- a/tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py
+++ b/tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pytest
 import torch
 
 from tests.utils import multi_gpu_test
@@ -19,9 +20,10 @@
 ACCEPTANCE_TOLERANCE = 0.95
 
 
-def _run(pp_size: int) -> float:
+def _run(pp_size: int, model_impl: str) -> float:
     llm = LLM(
         model=MODEL,
+        model_impl=model_impl,
         tensor_parallel_size=1,
         pipeline_parallel_size=pp_size,
         max_model_len=512,
@@ -49,10 +51,11 @@ def _run(pp_size: int) -> float:
 
 
 @multi_gpu_test(num_gpus=4)
-def test_eagle3_pipeline_parallel_acceptance():
-    baseline = _run(1)
+@pytest.mark.parametrize("model_impl", ["vllm", "transformers"])
+def test_eagle3_pipeline_parallel_acceptance(model_impl: str):
+    baseline = _run(1, model_impl)
     for pp_size in (2, 4):
-        parallel = _run(pp_size)
+        parallel = _run(pp_size, model_impl)
         assert parallel >= baseline * ACCEPTANCE_TOLERANCE, (
             f"PP={pp_size} acceptance regressed: {parallel:.3f} < "
             f"{baseline:.3f} * {ACCEPTANCE_TOLERANCE}"
diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py
index 46cd9f517c93..094443c38b8e 100644
--- a/vllm/model_executor/models/transformers/base.py
+++ b/vllm/model_executor/models/transformers/base.py
@@ -51,6 +51,7 @@
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.layers.fused_moe import MoERunner
 from vllm.model_executor.models.interfaces import (
+    EagleModelMixin,
     SupportsEagle,
     SupportsEagle3,
     SupportsLoRA,
@@ -84,6 +85,7 @@
     WeightsMapper,
     make_empty_intermediate_tensors_factory,
     maybe_prefix,
+    spec_decode_needs_target_embed,
 )
 from vllm.sequence import IntermediateTensors
 
@@ -114,8 +116,10 @@ class Base(
     SupportsPP,
     SupportsEagle,
     SupportsEagle3,
+    EagleModelMixin,
 ):
     hf_to_vllm_mapper: WeightsMapper
+    supports_aux_hidden_states_over_pp = True
 
     # TODO transformers will have a util to get these
     embedding_modules = {"embed_tokens": "input_embeddings"}
@@ -182,6 +186,7 @@ def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
         # Create weight name to module qualname mapper
         self._create_hf_to_vllm_mapper()
         # Remove layers not on this pipeline parallel rank
+        self.start_layer, _ = self._pp_indices()
         self.pipeline_parallel()
         # Skip the device sync of dynamic rope where it can never rescale
         disable_noop_dynamic_rope(self.model.get_decoder(), vllm_config)
@@ -426,8 +431,10 @@ def pipeline_parallel(self):
 
         # Layers before module list
         for name in names[:module_list_idx]:
-            if self.pp_group.is_first_rank or (
-                self._get_tie_word_embeddings() and self.pp_group.is_last_rank
+            if (
+                self.pp_group.is_first_rank
+                or (self._get_tie_word_embeddings() and self.pp_group.is_last_rank)
+                or spec_decode_needs_target_embed(self.vllm_config)
             ):
                 continue
             # attrsetter in case the module is nested (e.g. "text_model.embed_tokens")
@@ -545,7 +552,8 @@ def _recursive_replace(module: nn.Module, prefix: str):
                     and len(module) == self.text_config.num_hidden_layers
                 ):
                     # Populate Eagle3 attrs
-                    self._target_class = type(child_module)
+                    if not isinstance(child_module, PPMissingLayer):
+                        self._target_class = type(child_module)
                     layer_name = qual_name.removeprefix("model.")
                     self._layer_names[int(child_name)] = layer_name
                     # MTP weights should not be loaded into the base model
@@ -784,6 +792,7 @@ def forward(
             assert intermediate_tensors is not None
             input_ids = None
             inputs_embeds = intermediate_tensors["hidden_states"]
+        remote_aux = self.collect_remote_aux_hidden_states(intermediate_tensors)
 
         # Add batch dimension before entering Transformers model
         if input_ids is not None and input_ids.ndim == 1:
@@ -812,13 +821,21 @@ def forward(
 
         # Remove batch dimension after exiting Transformers model
         hidden_states = outputs[0][0, ...]
+        aux_hidden_states = []
         if self._output_aux_hidden_states_kwargs:
-            aux_hidden_states = [x[0][0, ...] for x in outputs[1:]]
+            # Layers on other PP ranks capture nothing
+            aux_hidden_states = [x[0][0, ...] for x in outputs[1:] if x]
 
         if not self.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._output_aux_hidden_states_kwargs and len(aux_hidden_states) > 0:
+        aux_hidden_states = remote_aux + aux_hidden_states
+        if aux_hidden_states:
             return hidden_states, aux_hidden_states
         return hidden_states
 
@@ -850,8 +867,9 @@ def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
         if self.model._can_record_outputs is None:
             self.model._can_record_outputs = {}
 
+        self._set_aux_hidden_state_layers(layers)
         target_class = self._target_class
-        for layer in layers:
+        for layer in self.aux_hidden_state_layers:
             # layer - 1 because we want the input to the layer
             layer_name = self._layer_names[layer - 1]
             layer_key = f"aux_hidden_state_{layer}"
diff --git a/vllm/model_executor/models/transformers/multimodal.py b/vllm/model_executor/models/transformers/multimodal.py
index f772e3e85c42..968605d2172c 100644
--- a/vllm/model_executor/models/transformers/multimodal.py
+++ b/vllm/model_executor/models/transformers/multimodal.py
@@ -985,6 +985,9 @@ def get_language_model(self) -> torch.nn.Module:
         bases = [b for b in bases if b is not MultiModalMixin]
 
         class LanguageModel(*bases):  # type: ignore[misc]
+            # `reserve_aux_intermediate_tensor_slots` needs `config.hidden_size`
+            supports_aux_hidden_states_over_pp = False
+
             def __init__(self, multimodal_model):
                 # Don't call super().__init__() to avoid re-initialization
                 self.__dict__.update(multimodal_model.__dict__)
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 86377496d395..68e13bd3408e 100644
--- a/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py
+++ b/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py
@@ -7,7 +7,11 @@
 
 from vllm.config import SpeculativeConfig
 from vllm.logger import init_logger
-from vllm.model_executor.models.interfaces import SupportsEagle3, supports_eagle3
+from vllm.model_executor.models.interfaces import (
+    EagleModelMixin,
+    SupportsEagle3,
+    supports_eagle3,
+)
 
 logger = init_logger(__name__)
 
@@ -41,6 +45,8 @@ def _inner_decoder(model: nn.Module) -> nn.Module | None:
         parent_ref = model.get_language_model()
     elif hasattr(model, "language_model"):
         parent_ref = model.language_model
+    if isinstance(parent_ref, EagleModelMixin):
+        return parent_ref
     return getattr(parent_ref, "model", None)
 
 

From 8623288b8f2c2caf2bcf8c25bbef20b4a26b9093 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Mon, 9 Nov 2026 11:14:14 +0000
Subject: [PATCH 02/10] Support aux hidden state layer 0 (moved from #66413)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 .../test_extraction.py                        | 38 +++++++++++++++++++
 vllm/model_executor/models/interfaces.py      |  6 ++-
 .../models/transformers/base.py               | 15 ++++++++
 .../extract_hidden_states_config.py           |  2 +
 4 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py b/tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py
index db8272887fcb..38b817f0251e 100644
--- a/tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py
+++ b/tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py
@@ -358,3 +358,41 @@ def test_extract_hidden_states_tp2():
             len(layer_ids),
             hidden_size,
         )
+
+
+@pytest.mark.parametrize("model_impl", ["vllm", "transformers"])
+@create_new_process_for_each_test()
+def test_extract_layer_zero_is_input_embeddings(tmp_path, model_impl):
+    """Layer 0 is the input to the first decoder layer, i.e. the input embeddings."""
+    from transformers import AutoModelForCausalLM
+
+    model = "meta-llama/Llama-3.2-1B-Instruct"
+    llm = LLM(
+        model=model,
+        model_impl=model_impl,
+        speculative_config={
+            "method": "extract_hidden_states",
+            "num_speculative_tokens": 1,
+            "draft_model_config": {
+                "hf_config": {"eagle_aux_hidden_state_layer_ids": [0, 1]}
+            },
+        },
+        kv_transfer_config={
+            "kv_connector": "ExampleHiddenStatesConnector",
+            "kv_role": "kv_producer",
+            "kv_connector_extra_config": {"shared_storage_path": str(tmp_path)},
+        },
+        max_model_len=256,
+        enforce_eager=True,
+        gpu_memory_utilization=0.4,
+    )
+    hidden_size = llm.llm_engine.model_config.get_hidden_size()
+    sampling_params = SamplingParams(max_tokens=1, temperature=0.0)
+    output = llm.generate(["Hello world"], sampling_params)[0]
+    expected_shape = (len(output.prompt_token_ids), 2, hidden_size)
+    token_ids, hidden_states = get_and_check_output(output, expected_shape)
+
+    hf_model = AutoModelForCausalLM.from_pretrained(model, dtype=hidden_states.dtype)
+    with torch.no_grad():
+        embeddings = hf_model.get_input_embeddings()(token_ids)
+    torch.testing.assert_close(hidden_states[:, 0, :], embeddings)
diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py
index 34a6fc6e02f9..ea24e358c886 100644
--- a/vllm/model_executor/models/interfaces.py
+++ b/vllm/model_executor/models/interfaces.py
@@ -1658,7 +1658,11 @@ def _maybe_add_hidden_state(
         residual: torch.Tensor | None,
     ) -> list[torch.Tensor]:
         if layer_idx in self.aux_hidden_state_layers:
-            value = hidden_states + residual if residual is not None else hidden_states
+            if residual is None:
+                # The first layer updates its residual (the embeddings) in place
+                value = hidden_states.clone()
+            else:
+                value = hidden_states + residual
             aux_hidden_states.append(value)
         return aux_hidden_states
 
diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py
index 094443c38b8e..a987c13a3126 100644
--- a/vllm/model_executor/models/transformers/base.py
+++ b/vllm/model_executor/models/transformers/base.py
@@ -164,6 +164,8 @@ def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
         """Mapping from layer index to layer name for Eagle3."""
         self._output_aux_hidden_states_kwargs: dict[str, bool] = {}
         """Kwargs to pass to model forward for Eagle3 aux hidden states."""
+        self._aux_embeddings = False
+        """Whether layer 0 (the input embeddings) is an Eagle3 aux hidden state."""
 
         if self.quant_config:
             quant_method_name = self.quant_config.get_name()
@@ -809,6 +811,13 @@ def forward(
         if input_ids is not None and inputs_embeds is not None:
             input_ids = None
 
+        aux_embeddings = None
+        if self._aux_embeddings and self.pp_group.is_first_rank:
+            if inputs_embeds is None:
+                inputs_embeds = self.embed_input_ids(input_ids)
+                input_ids = None
+            aux_embeddings = inputs_embeds[0, ...]
+
         outputs = self.model(
             input_ids=input_ids,
             inputs_embeds=inputs_embeds,
@@ -825,6 +834,8 @@ def forward(
         if self._output_aux_hidden_states_kwargs:
             # Layers on other PP ranks capture nothing
             aux_hidden_states = [x[0][0, ...] for x in outputs[1:] if x]
+        if aux_embeddings is not None:
+            aux_hidden_states.insert(0, aux_embeddings)
 
         if not self.pp_group.is_last_rank:
             return IntermediateTensors(
@@ -869,7 +880,11 @@ def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
 
         self._set_aux_hidden_state_layers(layers)
         target_class = self._target_class
+        # The input to layer 0 is captured in `forward`
+        self._aux_embeddings = 0 in self.aux_hidden_state_layers
         for layer in self.aux_hidden_state_layers:
+            if layer == 0:
+                continue
             # layer - 1 because we want the input to the layer
             layer_name = self._layer_names[layer - 1]
             layer_key = f"aux_hidden_state_{layer}"
diff --git a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states_config.py b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states_config.py
index 159827e644c0..20fe85143ed5 100644
--- a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states_config.py
+++ b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states_config.py
@@ -37,6 +37,8 @@ def configure(cls, spec: "SpeculativeConfig") -> None:
             hf_config = {}
 
         spec.draft_model_config = copy.copy(spec.target_model_config)
+        # The target's `model_impl` doesn't apply to the vLLM-only draft model
+        spec.draft_model_config.model_impl = "auto"
         spec.draft_model_config.hf_config = ExtractHiddenStatesConfig(
             spec.draft_model_config.hf_config, **hf_config
         )

From daf5f8087d5ca39f7524d3f7853eb09ca0f6f01b Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Mon, 9 Nov 2026 15:59:36 +0000
Subject: [PATCH 03/10] Use more prompts in the Eagle3 PP acceptance test

Four 32-token prompts let one greedy path that changes under compilation
swing the backend's acceptance by 20%.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py b/tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py
index 3d6fb47c1565..95cdd09f7c00 100644
--- a/tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py
+++ b/tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py
@@ -10,12 +10,15 @@
 
 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:",
+# Enough tokens that one greedy path changing under compilation can't move the mean
+TOPICS = ["the ocean", "a robot", "Rome", "volcanoes", "chess", "bread", "jazz", "bees"]
+TEMPLATES = [
+    "Write a short story about {}.",
+    "Explain {} to a child.",
+    "List five facts about {}:",
+    "Q: What is the history of {}?\nA:",
 ]
+PROMPTS = [template.format(topic) for topic in TOPICS for template in TEMPLATES]
 
 ACCEPTANCE_TOLERANCE = 0.95
 
@@ -39,7 +42,7 @@ def _run(pp_size: int, model_impl: str) -> float:
     try:
         llm.generate(
             PROMPTS,
-            SamplingParams(temperature=0.0, max_tokens=32, ignore_eos=True),
+            SamplingParams(temperature=0.0, max_tokens=64, ignore_eos=True),
         )
         acceptance = compute_acceptance_len(llm.get_metrics())
         assert acceptance > 1

From 67651da99e5a010baaa545d1e353c192e7cb70da Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Mon, 9 Nov 2026 17:18:21 +0000
Subject: [PATCH 04/10] Check aux layer 0 directly instead of caching it

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 vllm/model_executor/models/transformers/base.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py
index a987c13a3126..de30ba793089 100644
--- a/vllm/model_executor/models/transformers/base.py
+++ b/vllm/model_executor/models/transformers/base.py
@@ -164,8 +164,6 @@ def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
         """Mapping from layer index to layer name for Eagle3."""
         self._output_aux_hidden_states_kwargs: dict[str, bool] = {}
         """Kwargs to pass to model forward for Eagle3 aux hidden states."""
-        self._aux_embeddings = False
-        """Whether layer 0 (the input embeddings) is an Eagle3 aux hidden state."""
 
         if self.quant_config:
             quant_method_name = self.quant_config.get_name()
@@ -812,7 +810,7 @@ def forward(
             input_ids = None
 
         aux_embeddings = None
-        if self._aux_embeddings and self.pp_group.is_first_rank:
+        if 0 in self.aux_hidden_state_layers and self.pp_group.is_first_rank:
             if inputs_embeds is None:
                 inputs_embeds = self.embed_input_ids(input_ids)
                 input_ids = None
@@ -880,9 +878,8 @@ def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
 
         self._set_aux_hidden_state_layers(layers)
         target_class = self._target_class
-        # The input to layer 0 is captured in `forward`
-        self._aux_embeddings = 0 in self.aux_hidden_state_layers
         for layer in self.aux_hidden_state_layers:
+            # The input to layer 0 is captured in `forward`
             if layer == 0:
                 continue
             # layer - 1 because we want the input to the layer

From 1055620b1bc808f8faf06f245591cf4bc89be2d9 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Mon, 9 Nov 2026 17:20:10 +0000
Subject: [PATCH 05/10] Use walrus for the combined aux hidden states

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 vllm/model_executor/models/transformers/base.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py
index de30ba793089..b8c17f91a1ff 100644
--- a/vllm/model_executor/models/transformers/base.py
+++ b/vllm/model_executor/models/transformers/base.py
@@ -843,8 +843,7 @@ def forward(
                 }
             )
 
-        aux_hidden_states = remote_aux + aux_hidden_states
-        if aux_hidden_states:
+        if aux_hidden_states := remote_aux + aux_hidden_states:
             return hidden_states, aux_hidden_states
         return hidden_states
 

From 2ec9dd4aab141bcaeb54e9996d3e401b88231829 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Mon, 9 Nov 2026 17:24:29 +0000
Subject: [PATCH 06/10] Record aux layer 0 inside the compiled graph

Computing the embeddings in eager code handed the piecewise CUDA graphs a
fresh tensor each step, causing an illegal memory access on replay.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 .../test_extraction.py                        |  1 -
 .../models/transformers/base.py               | 26 +++++++++----------
 2 files changed, 12 insertions(+), 15 deletions(-)

diff --git a/tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py b/tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py
index 38b817f0251e..a18673235bbc 100644
--- a/tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py
+++ b/tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py
@@ -383,7 +383,6 @@ def test_extract_layer_zero_is_input_embeddings(tmp_path, model_impl):
             "kv_connector_extra_config": {"shared_storage_path": str(tmp_path)},
         },
         max_model_len=256,
-        enforce_eager=True,
         gpu_memory_utilization=0.4,
     )
     hidden_size = llm.llm_engine.model_config.get_hidden_size()
diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py
index b8c17f91a1ff..d2bf0b47bb81 100644
--- a/vllm/model_executor/models/transformers/base.py
+++ b/vllm/model_executor/models/transformers/base.py
@@ -809,13 +809,6 @@ def forward(
         if input_ids is not None and inputs_embeds is not None:
             input_ids = None
 
-        aux_embeddings = None
-        if 0 in self.aux_hidden_state_layers and self.pp_group.is_first_rank:
-            if inputs_embeds is None:
-                inputs_embeds = self.embed_input_ids(input_ids)
-                input_ids = None
-            aux_embeddings = inputs_embeds[0, ...]
-
         outputs = self.model(
             input_ids=input_ids,
             inputs_embeds=inputs_embeds,
@@ -832,8 +825,10 @@ def forward(
         if self._output_aux_hidden_states_kwargs:
             # Layers on other PP ranks capture nothing
             aux_hidden_states = [x[0][0, ...] for x in outputs[1:] if x]
-        if aux_embeddings is not None:
-            aux_hidden_states.insert(0, aux_embeddings)
+        # The input embedding isn't called when the caller passes `inputs_embeds`
+        caller_embeds = inputs_embeds if self.pp_group.is_first_rank else None
+        if 0 in self.aux_hidden_state_layers and caller_embeds is not None:
+            aux_hidden_states.insert(0, caller_embeds[0, ...])
 
         if not self.pp_group.is_last_rank:
             return IntermediateTensors(
@@ -876,13 +871,16 @@ def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
             self.model._can_record_outputs = {}
 
         self._set_aux_hidden_state_layers(layers)
-        target_class = self._target_class
+        embeddings = self.model.get_input_embeddings()
         for layer in self.aux_hidden_state_layers:
-            # The input to layer 0 is captured in `forward`
+            # The input to a layer is the output of the module before it
             if layer == 0:
-                continue
-            # layer - 1 because we want the input to the layer
-            layer_name = self._layer_names[layer - 1]
+                target_class = type(embeddings)
+                modules = self.model.named_modules()
+                layer_name = next(n for n, m in modules if m is embeddings)
+            else:
+                target_class = self._target_class
+                layer_name = self._layer_names[layer - 1]
             layer_key = f"aux_hidden_state_{layer}"
             aux_hidden_state_i = OutputRecorder(target_class, layer_name=layer_name)
             self.model._can_record_outputs[layer_key] = aux_hidden_state_i

From e46d1be8859e9e48b68887e6835f694cf400b7e2 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Mon, 9 Nov 2026 17:30:09 +0000
Subject: [PATCH 07/10] Only look up the input embedding for aux layer 0

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 vllm/model_executor/models/transformers/base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py
index d2bf0b47bb81..e2cd0c9c007c 100644
--- a/vllm/model_executor/models/transformers/base.py
+++ b/vllm/model_executor/models/transformers/base.py
@@ -871,10 +871,10 @@ def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
             self.model._can_record_outputs = {}
 
         self._set_aux_hidden_state_layers(layers)
-        embeddings = self.model.get_input_embeddings()
         for layer in self.aux_hidden_state_layers:
             # The input to a layer is the output of the module before it
             if layer == 0:
+                embeddings = self.model.get_input_embeddings()
                 target_class = type(embeddings)
                 modules = self.model.named_modules()
                 layer_name = next(n for n, m in modules if m is embeddings)

From e0b777d9f1cca7bdaa01ec02c90b046f0cb4516f Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Mon, 9 Nov 2026 17:35:51 +0000
Subject: [PATCH 08/10] Clarify why the extraction draft resets model_impl

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 vllm/v1/worker/gpu/spec_decode/extract_hidden_states_config.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states_config.py b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states_config.py
index 20fe85143ed5..0fd9629fcafe 100644
--- a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states_config.py
+++ b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states_config.py
@@ -37,7 +37,7 @@ def configure(cls, spec: "SpeculativeConfig") -> None:
             hf_config = {}
 
         spec.draft_model_config = copy.copy(spec.target_model_config)
-        # The target's `model_impl` doesn't apply to the vLLM-only draft model
+        # `ExtractHiddenStatesModel` is native, ignore the target's `model_impl`
         spec.draft_model_config.model_impl = "auto"
         spec.draft_model_config.hf_config = ExtractHiddenStatesConfig(
             spec.draft_model_config.hf_config, **hf_config

From 4b4321d8bc556c8bad1476cdf3ef3859d23c60c7 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Mon, 9 Nov 2026 17:44:48 +0000
Subject: [PATCH 09/10] Halve the Eagle3 PP test prompts

16 prompts x 64 tokens keep every PP size within 3% of PP=1 for both impls.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py b/tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py
index 95cdd09f7c00..3c47b597e733 100644
--- a/tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py
+++ b/tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py
@@ -11,7 +11,7 @@
 MODEL = "meta-llama/Llama-3.2-1B-Instruct"
 DRAFT = "nm-testing/Llama3_2_1B_speculator.eagle3"
 # Enough tokens that one greedy path changing under compilation can't move the mean
-TOPICS = ["the ocean", "a robot", "Rome", "volcanoes", "chess", "bread", "jazz", "bees"]
+TOPICS = ["the ocean", "a robot", "Rome", "volcanoes"]
 TEMPLATES = [
     "Write a short story about {}.",
     "Explain {} to a child.",

From ee9d7b2e065299c1baecf29851352edf87cddb7d Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Mon, 9 Nov 2026 17:57:54 +0000
Subject: [PATCH 10/10] Set end_layer for the centralized aux hidden state
 relay

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 vllm/model_executor/models/transformers/base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py
index 9f70cab7b3e7..484dd8e45883 100644
--- a/vllm/model_executor/models/transformers/base.py
+++ b/vllm/model_executor/models/transformers/base.py
@@ -192,7 +192,7 @@ def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
         # Create weight name to module qualname mapper
         self._create_hf_to_vllm_mapper()
         # Remove layers not on this pipeline parallel rank
-        self.start_layer, _ = self._pp_indices()
+        self.start_layer, self.end_layer = self._pp_indices()
         self.pipeline_parallel()
         # Skip the device sync of rope updates where `inv_freq` can be fixed
         maybe_disable_dynamic_rope_update(self.model.get_decoder(), vllm_config)