Move aux hidden state layer 0 support to vllm-project#66546

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Author: hmellor

tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py

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

vllm/model_executor/models/interfaces.py

@@ -1658,11 +1658,7 @@ def _maybe_add_hidden_state(
         residual: torch.Tensor | None,
     ) -> list[torch.Tensor]:
         if layer_idx in self.aux_hidden_state_layers:
-            if residual is None:
-                # The first layer updates its residual (the embeddings) in place
-                value = hidden_states.clone()
-            else:
-                value = hidden_states + residual
+            value = hidden_states + residual if residual is not None else hidden_states
             aux_hidden_states.append(value)
         return aux_hidden_states
 

vllm/model_executor/models/transformers/base.py

@@ -160,8 +160,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()
@@ -804,13 +802,6 @@ def forward(
         if input_ids is not None and inputs_embeds is not None:
             input_ids = None
 
-        aux_embeddings = None
-        if self._aux_embeddings:
-            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,
@@ -823,16 +814,13 @@ 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:]]
-        if aux_embeddings is not None:
-            aux_hidden_states.insert(0, aux_embeddings)
 
         if not self.pp_group.is_last_rank:
             return IntermediateTensors({"hidden_states": hidden_states})
 
-        if aux_hidden_states:
+        if self._output_aux_hidden_states_kwargs and len(aux_hidden_states) > 0:
             return hidden_states, aux_hidden_states
         return hidden_states
 
@@ -865,11 +853,7 @@ def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
             self.model._can_record_outputs = {}
 
         target_class = self._target_class
-        # The input to layer 0 is captured in `forward`
-        self._aux_embeddings = 0 in layers
         for layer in 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}"

vllm/v1/worker/gpu/spec_decode/extract_hidden_states_config.py

@@ -37,8 +37,6 @@ 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
         )