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

@@ -2746,13 +2746,23 @@ 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("Model loading took %.4f GiB and %.6f seconds",
@@ -2803,6 +2813,32 @@ def load_model(self, eep_scale_up: bool = False) -> None:
                 self.model = UBatchWrapper(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."