[Transformers backend] Support EAGLE3 aux hidden states with pipeline parallelism and layer 0

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

@@ -14,12 +14,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"]
+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
 
@@ -29,11 +32,13 @@ def _run(
     model: str = MODEL,
     draft: str = DRAFT,
     max_model_len: int = 512,
-    max_tokens: int = 32,
+    max_tokens: int = 64,
     chat_prompts: list[Messages] | None = None,
+    model_impl: str = "auto",
 ) -> float:
     llm = LLM(
         model=model,
+        model_impl=model_impl,
         tensor_parallel_size=1,
         pipeline_parallel_size=pp_size,
         max_model_len=max_model_len,
@@ -64,10 +69,11 @@ def _run(
 
 
 @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=model_impl)
     for pp_size in (2, 4):
-        parallel = _run(pp_size)
+        parallel = _run(pp_size, model_impl=model_impl)
         assert parallel >= baseline * ACCEPTANCE_TOLERANCE, (
             f"PP={pp_size} acceptance regressed: {parallel:.3f} < "
             f"{baseline:.3f} * {ACCEPTANCE_TOLERANCE}"

tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py

@@ -358,3 +358,40 @@ 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,
+        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)

vllm/model_executor/models/interfaces.py

@@ -1667,7 +1667,11 @@ def _maybe_add_hidden_state(
             layer_idx >= self._aux_first_local_layer
             and 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
 
@@ -1721,6 +1725,9 @@ def _get_decoder_holder(model: object) -> object:
         parent_ref = model.get_language_model()
     elif hasattr(model, "language_model"):
         parent_ref = model.language_model
+    # The Transformers modeling backend's `.model` is the HF model, not the decoder
+    if isinstance(parent_ref, EagleModelMixin):
+        return parent_ref
     # A multimodal model that builds its decoder inside `_mark_language_model`
     # has get_language_model() return the decoder itself, with no `.model`.
     return getattr(parent_ref, "model", parent_ref)

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
 
     def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
         nn.Module.__init__(self)
@@ -188,6 +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.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)
@@ -432,8 +437,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")
@@ -551,7 +558,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
@@ -790,6 +798,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:
@@ -818,13 +827,24 @@ 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]
+        # 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({"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:
+        if aux_hidden_states := remote_aux + aux_hidden_states:
             return hidden_states, aux_hidden_states
         return hidden_states
 
@@ -856,10 +876,17 @@ def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
         if self.model._can_record_outputs is None:
             self.model._can_record_outputs = {}
 
-        target_class = self._target_class
-        for layer in layers:
-            # layer - 1 because we want the input to the layer
-            layer_name = self._layer_names[layer - 1]
+        self._set_aux_hidden_state_layers(layers)
+        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)
+            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

vllm/model_executor/models/transformers/multimodal.py

@@ -1083,6 +1083,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__)

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)
+        # `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
         )