[Bugfix][MRV2][Spec Decode] Honor EAGLE3 single-state draft configuration (#61563)

Signed-off-by: Sandeep Agrawal <sandeepa@nvidia.com>
Co-authored-by: Kevin H. Luu <khluu000@gmail.com>

tests/v1/worker/test_eagle3_aux_hidden_states_pp.py

@@ -11,11 +11,36 @@ from vllm.model_executor.models.mimo import MiMoModel
 from vllm.models.minimax_m3.nvidia import model as minimax_m3
 from vllm.sequence import IntermediateTensors
 from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
+    requires_aux_hidden_states,
     reserve_aux_intermediate_tensor_slots,
     verify_supports_aux_hidden_states_over_pp,
 )
 
 
+@pytest.mark.parametrize(
+    ("method", "hf_config", "expected"),
+    [
+        ("eagle3", {"eagle_config": {"use_aux_hidden_state": False}}, False),
+        ("eagle3", {"eagle_config": {"use_aux_hidden_state": True}}, True),
+        ("eagle3", {"eagle_config": {}}, True),
+        ("eagle3", {"eagle_config": None}, True),
+        ("eagle3", {}, True),
+        ("dflash", {"eagle_config": {"use_aux_hidden_state": False}}, True),
+        ("dspark", {"eagle_config": {"use_aux_hidden_state": False}}, True),
+        ("extract_hidden_states", {}, True),
+        ("mtp", {}, False),
+    ],
+)
+def test_aux_hidden_state_requirement_respects_eagle3_config(
+    method, hf_config, expected
+):
+    spec_config = SimpleNamespace(
+        method=method,
+        draft_model_config=SimpleNamespace(hf_config=SimpleNamespace(**hf_config)),
+    )
+    assert requires_aux_hidden_states(spec_config) is expected
+
+
 def test_aux_layers_are_sorted_and_deduplicated():
     model = EagleModelMixin()
     model._set_aux_hidden_state_layers((48, 3, 90, 24, 48))

vllm/v1/worker/gpu/model_runner.py

@@ -156,6 +156,7 @@ from vllm.v1.worker.gpu.spec_decode.adaptive_verification import (
     resolve_adaptive_cudagraph_mode,
 )
 from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
+    requires_aux_hidden_states,
     set_eagle3_aux_hidden_state_layers,
     verify_supports_aux_hidden_states_over_pp,
 )
@@ -297,14 +298,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                     self.vllm_config, self.device, self.req_states
                 )
 
-            if self.speculative_config.method in (
-                "eagle3",
-                "dflash",
-                "dspark",
-                "extract_hidden_states",
-            ):
-                # Drafting may require auxiliary hidden states from target model outputs
-                self.use_aux_hidden_state_outputs = True
+            self.use_aux_hidden_state_outputs = requires_aux_hidden_states(
+                self.speculative_config
+            )
 
         # Draft tokens propagation - for spec-dec + struct outputs.
         self.draft_tokens_handler = DraftTokensHandler(self.device)

vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py

@@ -12,6 +12,16 @@ from vllm.model_executor.models.interfaces import SupportsEagle3, supports_eagle
 logger = init_logger(__name__)
 
 
+def requires_aux_hidden_states(spec_config: SpeculativeConfig) -> bool:
+    if spec_config.method == "eagle3":
+        eagle_config = (
+            getattr(spec_config.draft_model_config.hf_config, "eagle_config", None)
+            or {}
+        )
+        return eagle_config.get("use_aux_hidden_state", True)
+    return spec_config.method in ("dflash", "dspark", "extract_hidden_states")
+
+
 def set_eagle3_aux_hidden_state_layers(
     model: nn.Module,
     spec_config: SpeculativeConfig,