Add dynamic Eagle3 auxiliary layer configuration from speculative config

Implement _get_eagle3_aux_layers_from_config() helper method to extract
auxiliary layer IDs from the draft model's speculative config. The GPU
model runner now prefers config-specified layers over model defaults,
with fallback to model's get_eagle3_aux_hidden_state_layers() when not
configured.

Changes:
- Refactor auxiliary layer setup with early return pattern for errors
- Add config extraction with proper error handling
- Log only when using non-default layer configuration
- Enable dynamic layer configuration per deployment

Signed-off-by: rahul-tuli <[email protected]>
Signed-off-by: Rahul Tuli <[email protected]>

Author: rahul-tuli

vllm/v1/worker/gpu_model_runner.py

@@ -2943,15 +2943,29 @@ def load_model(self, eep_scale_up: bool = False) -> None:
                 logger.info("Loading drafter model...")
                 self.drafter.load_model(self.model)
             if self.use_aux_hidden_state_outputs:
-                if supports_eagle3(self.model):
-                    self.model.set_aux_hidden_state_layers(
-                        self.model.get_eagle3_aux_hidden_state_layers()
-                    )
-                else:
+                if not supports_eagle3(self.model):
                     raise RuntimeError(
                         "Model does not support EAGLE3 interface but "
                         "aux_hidden_state_outputs was requested"
                     )
+
+                # Try to get auxiliary layers from speculative config,
+                # otherwise use model's default layers
+                aux_layers = (
+                    self._get_eagle3_aux_layers_from_config()
+                    or self.model.get_eagle3_aux_hidden_state_layers()
+                )
+
+                if (
+                    aux_layers
+                    != self.model.get_eagle3_aux_hidden_state_layers()
+                ):
+                    logger.info(
+                        "Using auxiliary layers from speculative config: %s",
+                        aux_layers,
+                    )
+
+                self.model.set_aux_hidden_state_layers(aux_layers)
             time_after_load = time.perf_counter()
         self.model_memory_usage = m.consumed_memory
         logger.info(
@@ -3006,6 +3020,32 @@ def load_model(self, eep_scale_up: bool = False) -> None:
                     self.model, self.vllm_config, CUDAGraphMode.NONE, self.device
                 )
 
+    def _get_eagle3_aux_layers_from_config(self) -> Optional[tuple[int, ...]]:
+        """Extract Eagle3 auxiliary layer IDs from speculative config.
+
+        Returns:
+            Tuple of layer indices if found in draft model config,
+            None otherwise.
+        """
+        if not (self.speculative_config
+                and self.speculative_config.draft_model_config):
+            return None
+
+        try:
+            hf_config = self.speculative_config.draft_model_config.hf_config
+            if not hasattr(hf_config, 'eagle_aux_hidden_state_layer_ids'):
+                return None
+
+            layer_ids = hf_config.eagle_aux_hidden_state_layer_ids
+            if layer_ids and isinstance(layer_ids, (list, tuple)):
+                return tuple(layer_ids)
+        except Exception as e:
+            logger.warning(
+                "Failed to read auxiliary layers from speculative config: %s",
+                e)
+
+        return None
+
     def reload_weights(self) -> None:
         assert getattr(self, "model", None) is not None, (
             "Cannot reload weights before model is loaded."