diff --git a/tests/v1/worker/test_eagle3_aux_hidden_states_pp.py b/tests/v1/worker/test_eagle3_aux_hidden_states_pp.py index f310c5a083..89bfc6a53d 100644 --- a/tests/v1/worker/test_eagle3_aux_hidden_states_pp.py +++ b/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)) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 6fc6c022a2..9cd7299642 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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) 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 6b0c5fe51d..e434e58a9e 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py +++ b/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,