spec: centralize EagleDraft{,Extend}Input.hidden_states shape (#24926)

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -130,8 +130,8 @@ class EAGLEDraftCudaGraphRunner:
             topk_p = torch.zeros((self.max_bs, self.topk), dtype=torch.float32)
             topk_index = torch.zeros((self.max_bs, self.topk), dtype=torch.int64)
             hidden_states = torch.zeros(
-                (self.max_bs, self.model_runner.model_config.spec_hidden_size),
-                dtype=self.model_runner.dtype,
+                (self.max_bs, EagleDraftInput.hidden_size_for(self.eagle_worker)),
+                dtype=EagleDraftInput.dtype_for(self.eagle_worker),
             )
 
             if self.require_gathered_buffer:

python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py

@@ -132,35 +132,13 @@ class EAGLEDraftExtendCudaGraphRunner:
             positions = torch.zeros((self.max_num_token,), dtype=torch.int64)
             mrope_positions = torch.zeros((3, self.max_num_token), dtype=torch.int64)
 
-            if (
-                self.eagle_worker.speculative_algorithm.is_eagle3()
-                and self.eagle_worker.eagle_use_aux_hidden_state
-            ):
-                hidden_states = torch.zeros(
-                    (
-                        self.max_num_token,
-                        (
-                            self.model_runner.model_config.hf_config.target_hidden_size
-                            * 3
-                            if hasattr(
-                                self.model_runner.model_config.hf_config,
-                                "target_hidden_size",
-                            )
-                            else self.model_runner.model_config.hidden_size * 3
-                        ),
-                    ),
-                    dtype=self.model_runner.dtype,
-                )
-            else:
-                # Use target config: hidden_states carries target output.
-                target_cfg = self.eagle_worker.target_worker.model_runner.model_config
-                hidden_states = torch.zeros(
-                    (
-                        self.max_num_token,
-                        target_cfg.spec_hidden_size,
-                    ),
-                    dtype=target_cfg.dtype,
-                )
+            hidden_states = torch.zeros(
+                (
+                    self.max_num_token,
+                    EagleDraftExtendInput.hidden_size_for(self.eagle_worker),
+                ),
+                dtype=EagleDraftExtendInput.dtype_for(self.eagle_worker),
+            )
             self.seq_len_fill_value = (
                 self.model_runner.attn_backend.get_cuda_graph_seq_len_fill_value()
             )

python/sglang/srt/speculative/eagle_info.py

@@ -56,6 +56,18 @@ if is_cuda() or is_musa():
 logger = logging.getLogger(__name__)
 
 
+def _draft_runner_of(worker):
+    """Draft model_runner accessor that handles v1 / v2 worker naming.
+
+    v1 (`EAGLEWorker` and subclasses) exposes the draft model_runner as
+    `model_runner` (the worker itself runs the draft model);
+    v2 (`EagleDraftWorker` and subclasses) exposes it as `draft_runner`.
+    """
+    return (
+        worker.draft_runner if hasattr(worker, "draft_runner") else worker.model_runner
+    )
+
+
 @dataclass
 class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
     draft_token: torch.Tensor
@@ -698,6 +710,17 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin):
             )
             pt += extend_len
 
+    @classmethod
+    def hidden_size_for(cls, worker) -> int:
+        """Decode-phase `hidden_states` width: draft self-chain output
+        (draft model writes its own last hidden back via `capture_for_decode`
+        and the draft loop)."""
+        return _draft_runner_of(worker).model_config.spec_hidden_size
+
+    @classmethod
+    def dtype_for(cls, worker) -> torch.dtype:
+        return _draft_runner_of(worker).model_config.dtype
+
     @classmethod
     def create_idle_input(
         cls,
@@ -820,6 +843,24 @@ class EagleDraftExtendInput(SpecInput):
     def get_spec_adjust_token_coefficient(self) -> Tuple[int, int]:
         return self.num_tokens_per_req, self.num_tokens_for_logprob_per_req
 
+    @classmethod
+    def hidden_size_for(cls, worker) -> int:
+        """Extend-phase `hidden_states` width: target verify output (EAGLE
+        paper's "feature"). Widened to `target.hidden_size * 3` for EAGLE-3
+        aux mode (low/mid/high features fused into a 3k-dim vector, reduced
+        by draft's FC)."""
+        target_cfg = worker.target_worker.model_runner.model_config
+        if (
+            worker.speculative_algorithm.is_eagle3()
+            and worker.eagle_use_aux_hidden_state
+        ):
+            return target_cfg.hidden_size * 3
+        return target_cfg.spec_hidden_size
+
+    @classmethod
+    def dtype_for(cls, worker) -> torch.dtype:
+        return worker.target_worker.model_runner.model_config.dtype
+
     @classmethod
     def create_idle_input(
         cls,

python/sglang/srt/speculative/eagle_worker.py

@@ -727,8 +727,8 @@ class EAGLEWorker(TpModelWorker):
     def _draft_preprocess_idle(self, batch: ScheduleBatch):
         batch.spec_info = EagleDraftInput.create_idle_input(
             device=self.device,
-            hidden_size=self.model_config.spec_hidden_size,
-            dtype=self.model_config.dtype,
+            hidden_size=EagleDraftInput.hidden_size_for(self),
+            dtype=EagleDraftInput.dtype_for(self),
             topk=self.topk,
             capture_hidden_mode=CaptureHiddenMode.LAST,
         )
@@ -1130,17 +1130,10 @@ class EAGLEWorker(TpModelWorker):
             # All reqs finished this verify; swap to an idle ExtendInput.
             batch = batch.copy()
             batch.prepare_for_idle()
-            target_cfg = self.target_worker.model_runner.model_config
-            hidden_size = (
-                target_cfg.hidden_size * 3
-                if self.speculative_algorithm.is_eagle3()
-                and self.eagle_use_aux_hidden_state
-                else target_cfg.spec_hidden_size
-            )
             draft_extend_input = EagleDraftExtendInput.create_idle_input(
                 device=self.device,
-                hidden_size=hidden_size,
-                dtype=target_cfg.dtype,
+                hidden_size=EagleDraftExtendInput.hidden_size_for(self),
+                dtype=EagleDraftExtendInput.dtype_for(self),
                 capture_hidden_mode=CaptureHiddenMode.LAST,
             )
             batch.spec_info = draft_extend_input