Fix EAGLE draft hidden dim extraction and centralize spec helpers (#29464)

python/sglang/srt/managers/scheduler.py

@@ -243,6 +243,7 @@ from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
 from sglang.srt.server_args import PortArgs, ServerArgs, get_global_server_args
 from sglang.srt.session.session_controller import SessionController
 from sglang.srt.speculative.dflash_utils import validate_dflash_request
+from sglang.srt.speculative.eagle_utils import get_draft_recurrent_hidden_state_spec
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.utils import (
     DynamicGradMode,
@@ -1120,6 +1121,17 @@ class Scheduler(
         if model_config is None:
             model_config = self.model_config
 
+        if self.spec_algorithm.carries_draft_hidden_states():
+            # `draft_runner` aliases `draft_runner_list[0]` in the multi-layer
+            # worker, so a single accessor covers both shapes.
+            draft_runner = self.draft_worker.draft_worker.draft_runner
+            disagg_hidden_size, disagg_hidden_states_dtype = (
+                get_draft_recurrent_hidden_state_spec(draft_runner)
+            )
+        else:
+            disagg_hidden_size = 16  # minimal padding size for RDMA
+            disagg_hidden_states_dtype = torch.float32
+
         if (
             self.disaggregation_mode == DisaggregationMode.DECODE
         ):  # *2 for the headroom.
@@ -1129,16 +1141,8 @@ class Scheduler(
             )
             self.disagg_metadata_buffers = MetadataBuffers(
                 buffer_size,
-                hidden_size=(
+                hidden_size=disagg_hidden_size,
-                    model_config.spec_hidden_size
+                hidden_states_dtype=disagg_hidden_states_dtype,
-                    if self.spec_algorithm.carries_draft_hidden_states()
-                    else 16  # minimal padding size for RDMA
-                ),
-                hidden_states_dtype=(
-                    model_config.dtype
-                    if self.spec_algorithm.carries_draft_hidden_states()
-                    else torch.float32
-                ),
                 custom_mem_pool=self.token_to_kv_pool_allocator.get_kvcache().maybe_get_custom_mem_pool(),
             )
 
@@ -1182,16 +1186,8 @@ class Scheduler(
             )
             self.disagg_metadata_buffers = MetadataBuffers(
                 buffer_size,
-                hidden_size=(
+                hidden_size=disagg_hidden_size,
-                    model_config.spec_hidden_size
+                hidden_states_dtype=disagg_hidden_states_dtype,
-                    if self.spec_algorithm.carries_draft_hidden_states()
-                    else 16  # minimal padding size for RDMA
-                ),
-                hidden_states_dtype=(
-                    model_config.dtype
-                    if self.spec_algorithm.carries_draft_hidden_states()
-                    else torch.float32
-                ),
                 custom_mem_pool=self.token_to_kv_pool_allocator.get_kvcache().maybe_get_custom_mem_pool(),
             )
 

python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py

@@ -234,7 +234,9 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         if model_runner.spec_algorithm.is_speculative():
             if self.model_runner.is_draft_worker:
                 # Draft workers can use TARGET_VERIFY mode.
-                if not self.model_runner.spec_algorithm.is_dflash():
+                if (
+                    not self.model_runner.spec_algorithm.supports_target_verify_for_draft()
+                ):
                     raise RuntimeError("This should not happen")
             self.capture_forward_mode = ForwardMode.TARGET_VERIFY
             self.num_tokens_per_bs = (

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -75,6 +75,7 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph impo
 from sglang.srt.model_executor.runner_utils.buffers import (
     PrefillInputBuffers,
 )
+from sglang.srt.speculative.eagle_utils import get_draft_input_from_target_hidden_dim
 from sglang.srt.utils import (
     get_available_gpu_memory,
     get_bool_env_var,
@@ -190,6 +191,13 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             hidden_size=self.model_runner.model_config.hidden_size,
             embed_dtype=self.model_runner.dtype,
             enable_mamba_track=self.mamba_track_enabled,
+            # Register the multimodal input_embeds slot for every prefill
+            # backend (default True). The slot is only added when is_multimodal,
+            # so text-only models are unaffected. Both tc_piecewise (outer MM
+            # wrapper passes composed input_embeds as an argument) and breakable
+            # (captures the input_embeds path; general_mm_embed_routine fills the
+            # slot) need it, otherwise the captured graph re-embeds input_ids and
+            # drops the scattered vision embeddings.
             source=self.buffers,
         )
 
@@ -237,24 +245,15 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
 
         # Static hidden_states buffer giving the captured graph a stable
         # address; load_batch refreshes it from live spec_info at replay.
-        # Draft consumes the aux-concatenated hidden states from the target
+        # Draft consumes aux-concatenated hidden states from the target
-        # (e.g. EAGLE3 stacks 3 target layers), so read the dim from the
+        # (e.g. EAGLE3 stacks 3 target layers), so capture the pre-reduction
-        # draft model's input fc when available; fall back to the per-layer
+        # width when the draft model exposes it.
-        # dim for arches without an fc projection.
         if (
             isinstance(self.backend, BreakableCudaGraphBackend)
             and model_runner.is_draft_worker
             and model_runner.spec_algorithm.is_eagle()
         ):
-            from sglang.srt.speculative.eagle_utils import get_draft_hidden_dim
+            hidden_dim = get_draft_input_from_target_hidden_dim(model_runner)
-
-            inner = getattr(model_runner.model, "model", model_runner.model)
-            fc = getattr(inner, "fc", None)
-            hidden_dim = (
-                fc.in_features
-                if fc is not None and hasattr(fc, "in_features")
-                else get_draft_hidden_dim(model_runner)
-            )
             with torch.device(self.device):
                 self.static_draft_hidden_states = torch.zeros(
                     (self.max_num_tokens, hidden_dim),
@@ -392,15 +391,18 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         )
         set_is_extend_in_batch(False)
 
-        with forward_context(
+        with (
+            forward_context(
                 ForwardContext(attn_backend=self.model_runner.attn_backend)
-        ), set_tc_piecewise_forward_context(
+            ),
+            set_tc_piecewise_forward_context(
                 forward_batch,
                 self.attention_layers,
                 self.quant_config,
                 self.moe_layers,
                 self.moe_fusions,
                 dsa_indexers=self.dsa_indexers,
+            ),
         ):
             if self.layer_model is not None:
                 return self.layer_model.forward(
@@ -876,8 +878,25 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                 # static_forward_batch. The outer's logits_processor /
                 # pooler then runs on top with live multi-req metadata.
                 shape_key = ShapeKey(size=self._static_num_tokens)
+                static_n = self._static_num_tokens
 
                 def replay_layer_forward(*args, **layer_kwargs):
+                    # The captured BCG graph reads activations from the static
+                    # input_embeds slot. The outer model.forward (run eagerly)
+                    # passes the live embeddings into layer_model.forward as the
+                    # 4th positional arg (or input_embeds kwarg): for multimodal
+                    # batches these are the composed text+vision embeds, for
+                    # text-only batches they are get_input_embeddings()(input_ids).
+                    # Copy them into the slot before replay so the graph sees the
+                    # current request's embeddings (mirrors main's BCG closure).
+                    if self.buffer_registry.has_slot("input_embeds"):
+                        ie = layer_kwargs.get("input_embeds")
+                        if ie is None and len(args) > 3:
+                            ie = args[3]
+                        if ie is not None:
+                            self.buffer_registry.get_slot("input_embeds").slice_for(
+                                1, static_n
+                            ).copy_(ie[:static_n])
                     return self.backend.replay(
                         shape_key, static_forward_batch, **kwargs
                     )
@@ -885,9 +904,11 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                 original_layer_forward = self.layer_model.forward
                 self.layer_model.forward = replay_layer_forward
                 try:
-                    with forward_context(
+                    with (
+                        forward_context(
                             ForwardContext(attn_backend=self.model_runner.attn_backend)
-                    ), set_tc_piecewise_forward_context(
+                        ),
+                        set_tc_piecewise_forward_context(
                             static_forward_batch,
                             self.attention_layers,
                             self.quant_config,
@@ -896,6 +917,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                             dsa_indexers=self.dsa_indexers,
                             num_tokens=static_num_tokens,
                             raw_num_tokens=raw_num_tokens,
+                        ),
                     ):
                         output = self.model_runner.model.forward(
                             static_forward_batch.input_ids,
@@ -909,9 +931,11 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                 # TC_PIECEWISE path. backend.replay calls the compiled
                 # outer model.forward directly (torch.compile handles
                 # multi-req via bs-invariant FX-traced kernels).
-                with forward_context(
+                with (
+                    forward_context(
                         ForwardContext(attn_backend=self.model_runner.attn_backend)
-                ), set_tc_piecewise_forward_context(
+                    ),
+                    set_tc_piecewise_forward_context(
                         static_forward_batch,
                         self.attention_layers,
                         self.quant_config,
@@ -920,6 +944,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                         dsa_indexers=self.dsa_indexers,
                         num_tokens=static_num_tokens,
                         raw_num_tokens=raw_num_tokens,
+                    ),
                 ):
                     output = self.backend.replay(
                         self._static_num_tokens, static_forward_batch, **kwargs

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -34,6 +34,7 @@ from sglang.srt.model_executor.runner_backend_utils import (
 )
 from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
 from sglang.srt.speculative.eagle_info import EagleDraftInput
+from sglang.srt.speculative.eagle_utils import get_draft_recurrent_hidden_state_spec
 from sglang.srt.utils import (
     require_attn_tp_gather,
     require_gathered_buffer,
@@ -187,11 +188,13 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
                 if self.model_runner.server_args.speculative_use_rejection_sampling
                 else None
             )
-            _hidden_size = EagleDraftInput.hidden_size_for(self.eagle_worker)
+            _hidden_size, _hidden_dtype = get_draft_recurrent_hidden_state_spec(
+                model_runner
+            )
             hidden_states = (
                 torch.zeros(
                     (self.max_bs, _hidden_size),
-                    dtype=EagleDraftInput.dtype_for(self.eagle_worker),
+                    dtype=_hidden_dtype,
                 )
                 if _hidden_size is not None
                 else None

python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py

@@ -33,6 +33,7 @@ from sglang.srt.model_executor.runner_backend_utils import (
     CUDA_GRAPH_CAPTURE_FAILED_MSG,
 )
 from sglang.srt.speculative.eagle_info import EagleDraftExtendInput
+from sglang.srt.speculative.eagle_utils import get_draft_input_from_target_hidden_dim
 from sglang.srt.speculative.spec_utils import fast_topk
 from sglang.srt.utils import (
     is_hip,
@@ -152,11 +153,19 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
             positions = torch.zeros((self.max_num_token,), dtype=torch.int64)
             mrope_positions = torch.zeros((3, self.max_num_token), dtype=torch.int64)
 
-            _hidden_size = EagleDraftExtendInput.hidden_size_for(self.eagle_worker)
+            # Width and dtype both come from the draft `model_runner` so the
+            # source stays consistent (the draft dtype matches the target dtype
+            # that produced these hidden states).
+            _hidden_dtype = model_runner.model_config.dtype
+            _hidden_size = (
+                None
+                if self.eagle_worker.speculative_algorithm.is_standalone()
+                else get_draft_input_from_target_hidden_dim(model_runner)
+            )
             hidden_states = (
                 torch.zeros(
                     (self.max_num_token, _hidden_size),
-                    dtype=EagleDraftExtendInput.dtype_for(self.eagle_worker),
+                    dtype=_hidden_dtype,
                 )
                 if _hidden_size is not None
                 else None

python/sglang/srt/speculative/eagle_info.py

@@ -14,18 +14,6 @@ from sglang.srt.speculative.spec_info import SpecInput, SpecInputType
 logger = logging.getLogger(__name__)
 
 
-def _draft_runner_of(worker):
-    """Draft model_runner accessor across worker shapes.
-
-    v2 draft workers (`EagleDraftWorker` and subclasses) expose the draft
-    model_runner as `draft_runner`; fall back to `model_runner` for workers
-    that run the draft model directly.
-    """
-    return (
-        worker.draft_runner if hasattr(worker, "draft_runner") else worker.model_runner
-    )
-
-
 @dataclass
 class EagleVerifyInput(SpecInput):
     draft_token: torch.Tensor
@@ -191,22 +179,6 @@ class EagleDraftInput(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) -> Optional[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). Returns None when the draft architecture doesn't
-        consume the field (e.g., STANDALONE)."""
-        if worker.speculative_algorithm.is_standalone():
-            return None
-        return _draft_runner_of(worker).model_config.spec_hidden_size
-
-    @classmethod
-    def dtype_for(cls, worker) -> Optional[torch.dtype]:
-        if worker.speculative_algorithm.is_standalone():
-            return None
-        return _draft_runner_of(worker).model_config.dtype
-
     @classmethod
     def create_idle_input(
         cls,
@@ -351,39 +323,6 @@ 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) -> Optional[int]:
-        """Extend-phase `hidden_states` width: target's `spec_hidden_size`,
-        widened to `num_aux * target_hidden` for EAGLE-3 aux mode. Returns
-        None when the draft architecture doesn't consume the field
-        (e.g., STANDALONE)."""
-        if worker.speculative_algorithm.is_standalone():
-            return None
-        target_cfg = worker.target_worker.model_runner.model_config
-        if not (
-            worker.speculative_algorithm.is_eagle3()
-            and worker.eagle_use_aux_hidden_state
-        ):
-            return target_cfg.spec_hidden_size
-
-        hf_config = target_cfg.hf_config
-
-        # `num_aux` resolution: explicit attr > eagle_config layer_ids > default 3.
-        num_aux = getattr(hf_config, "num_aux_hidden_states", None)
-        if num_aux is None:
-            eagle_config = getattr(hf_config, "eagle_config", None) or {}
-            layer_ids = eagle_config.get("eagle_aux_hidden_state_layer_ids")
-            num_aux = len(layer_ids) if layer_ids else 3
-
-        target_hidden = getattr(hf_config, "target_hidden_size", target_cfg.hidden_size)
-        return target_hidden * num_aux
-
-    @classmethod
-    def dtype_for(cls, worker) -> Optional[torch.dtype]:
-        if worker.speculative_algorithm.is_standalone():
-            return None
-        return worker.target_worker.model_runner.model_config.dtype
-
     @classmethod
     def create_idle_input(
         cls,

python/sglang/srt/speculative/eagle_utils.py

@@ -267,21 +267,55 @@ def verify_tree_greedy_func(
     return predicts, accept_index, accept_token_num
 
 
-def get_draft_hidden_dim(model_runner: ModelRunner) -> int:
+def get_draft_input_from_target_hidden_dim(model_runner: ModelRunner) -> int:
-    """Derive the hidden dimension of target hidden states fed to the draft model."""
+    """Width of the target hidden states fed into the draft model.
-    hf_config = model_runner.model_config.hf_config
+
-    eagle_config = getattr(hf_config, "eagle_config", {})
+    This is the single source of truth and is derived entirely from config: for
-    use_aux = eagle_config.get("use_aux_hidden_state", False)
+    EAGLE3 aux mode the draft consumes `num_aux` concatenated target layers
+    (each `target_hidden_size` wide); every other arch consumes the per-layer
+    `spec_hidden_size`.
+
+    Do NOT read this off a draft projection's `in_features` (e.g. an `fc`
+    layer): that width is arch-specific.
+
+    Note: read entirely from the *draft* `model_runner`'s config. The non-aux
+    branch assumes the draft's `spec_hidden_size` equals the target hidden width
+    fed to the draft (true for standard EAGLE, where the draft mirrors the
+    target hidden size); aux mode reads the explicit `target_hidden_size`.
+    """
+    model_config = model_runner.model_config
+    hf_config = model_config.hf_config
+    eagle_config = getattr(hf_config, "eagle_config", None) or {}
+    get_eagle_config = (
+        eagle_config.get
+        if isinstance(eagle_config, dict)
+        else lambda key, default=None: getattr(eagle_config, key, default)
+    )
+    use_aux = get_eagle_config("use_aux_hidden_state", True)
     spec_algorithm = model_runner.spec_algorithm
 
-    if spec_algorithm is not None and spec_algorithm.is_eagle3() and use_aux:
+    if not (spec_algorithm is not None and spec_algorithm.is_eagle3() and use_aux):
-        base = getattr(hf_config, "target_hidden_size", None)
+        return model_config.spec_hidden_size
-        if base is None:
+
-            base = model_runner.model_config.hidden_size
+    target_hidden = getattr(hf_config, "target_hidden_size", None)
-        layer_ids = eagle_config.get("eagle_aux_hidden_state_layer_ids", [])
+    if target_hidden is None:
-        num_aux = max(len(layer_ids), 1)
+        target_hidden = model_config.hidden_size
-        return base * num_aux
+    num_aux = getattr(hf_config, "num_aux_hidden_states", None)
-    return model_runner.model_config.spec_hidden_size
+    if num_aux is None:
+        layer_ids = get_eagle_config("eagle_aux_hidden_state_layer_ids", None)
+        if layer_ids is None:
+            layer_ids = getattr(hf_config, "eagle_aux_hidden_state_layer_ids", None)
+        num_aux = len(layer_ids) if layer_ids else 3
+    return target_hidden * num_aux
+
+
+def get_draft_recurrent_hidden_state_spec(
+    model_runner: ModelRunner,
+) -> tuple[Optional[int], Optional[torch.dtype]]:
+    """Return hidden_states width/dtype carried between draft decode steps."""
+    if model_runner.spec_algorithm.is_standalone():
+        return None, None
+    return model_runner.model_config.spec_hidden_size, model_runner.model_config.dtype
 
 
 def eagle_prepare_for_verify(

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -70,6 +70,7 @@ from sglang.srt.speculative.eagle_utils import (
     build_tree_kernel_efficient,
     eagle_prepare_for_verify,
     eagle_sample,
+    get_draft_recurrent_hidden_state_spec,
     organize_draft_results,
     per_step_draft_out_cache_loc,
 )
@@ -188,14 +189,6 @@ class EagleDraftWorker(EagleDraftWorkerBase):
 
         # Alias for better readability
         self.draft_runner = self.draft_worker.model_runner
-        self.eagle_use_aux_hidden_state = False
-        if self.speculative_algorithm.is_eagle3():
-            eagle_config = getattr(
-                self.draft_runner.model_config.hf_config, "eagle_config", {}
-            )
-            self.eagle_use_aux_hidden_state = eagle_config.get(
-                "use_aux_hidden_state", True
-            )
         # Reuse the first draft step's NSA/DSA indexer topk across the rest;
         # topk == 1 only (select_top_k_tokens reorders rows, desyncing indices).
         self.index_share_for_mtp_iteration = (
@@ -247,16 +240,20 @@ class EagleDraftWorker(EagleDraftWorkerBase):
                 )
 
     def init_attention_backends(self):
-        with self.draft_tp_context(
+        with (
-            self.draft_runner.tp_group
+            self.draft_tp_context(self.draft_runner.tp_group),
-        ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
+            speculative_moe_backend_context(),
+            speculative_moe_a2a_backend_context(),
+        ):
             self.draft_worker.init_attention_backends()
             self.init_attention_backend()
 
     def init_cuda_graphs(self):
-        with self.draft_tp_context(
+        with (
-            self.draft_runner.tp_group
+            self.draft_tp_context(self.draft_runner.tp_group),
-        ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
+            speculative_moe_backend_context(),
+            speculative_moe_a2a_backend_context(),
+        ):
             self.draft_worker.init_cuda_graphs(capture_decode_cuda_graph=False)
             if check_cuda_graph_backend(Phase.PREFILL, Backend.BREAKABLE):
                 self.draft_runner.init_prefill_cuda_graph(force_for_draft_worker=True)
@@ -1126,10 +1123,13 @@ class EAGLEWorkerV2(BaseSpecWorker):
                     if self.speculative_algorithm.is_standalone()
                     else CaptureHiddenMode.LAST
                 )
+                hidden_size, hidden_dtype = get_draft_recurrent_hidden_state_spec(
+                    self.draft_worker.draft_runner
+                )
                 batch.spec_info = EagleDraftInput.create_idle_input(
                     device=self.device,
-                    hidden_size=EagleDraftInput.hidden_size_for(self.draft_worker),
+                    hidden_size=hidden_size,
-                    dtype=EagleDraftInput.dtype_for(self.draft_worker),
+                    dtype=hidden_dtype,
                     topk=self.topk,
                     capture_hidden_mode=capture_mode,
                     vocab_size=self.target_worker.model_config.vocab_size,
@@ -1249,11 +1249,13 @@ class EAGLEWorkerV2(BaseSpecWorker):
         next_draft_input.topk_index = torch.zeros(
             (bs, self.topk), dtype=torch.int64, device=device
         )
-        hidden_size = EagleDraftInput.hidden_size_for(self.draft_worker)
+        hidden_size, hidden_dtype = get_draft_recurrent_hidden_state_spec(
+            self.draft_worker.draft_runner
+        )
         if hidden_size is not None:
             next_draft_input.hidden_states = torch.zeros(
                 (bs, hidden_size),
-                dtype=EagleDraftInput.dtype_for(self.draft_worker),
+                dtype=hidden_dtype,
                 device=device,
             )
 

python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py

@@ -57,6 +57,7 @@ from sglang.srt.model_executor.runner_backend_utils import (
     CUDA_GRAPH_CAPTURE_FAILED_MSG,
 )
 from sglang.srt.speculative.eagle_info import EagleDraftExtendInput
+from sglang.srt.speculative.eagle_utils import get_draft_input_from_target_hidden_dim
 from sglang.srt.speculative.spec_utils import fast_topk
 from sglang.srt.utils import (
     get_available_gpu_memory,
@@ -517,11 +518,12 @@ class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner:
 
     def _allocate_buffers(self) -> MultiLayerEagleDraftExtendInputBuffers:
         runner = next(r for r in self.runners if r is not None)
+        model_runner = runner.model_runner
         max_bs = self.max_bs
         num_tokens_per_bs = self.num_tokens_per_bs
         max_num_token = max_bs * num_tokens_per_bs
-        hidden_size = EagleDraftExtendInput.hidden_size_for(self.eagle_worker)
+        hidden_size = get_draft_input_from_target_hidden_dim(model_runner)
-        dtype = EagleDraftExtendInput.dtype_for(self.eagle_worker)
+        dtype = model_runner.model_config.dtype
         vocab_size = self._vocab_size()
 
         seq_lens_cpu = torch.full((max_bs,), self.seq_len_fill_value, dtype=torch.int32)

python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py

@@ -53,6 +53,7 @@ from sglang.srt.speculative.eagle_utils import (
     build_tree_kernel_efficient,
     eagle_prepare_for_verify,
     eagle_sample,
+    get_draft_recurrent_hidden_state_spec,
 )
 from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner import (
     MultiLayerEagleMultiStepDraftExtendCudaGraphRunner,
@@ -155,8 +156,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
 
         # Alias for better readability
         self.draft_runner_list: List[ModelRunner] = self.draft_worker.model_runner_list
-        # Match `EagleDraftWorker.draft_runner` so `_draft_runner_of(self)` works
+        # Match `EagleDraftWorker.draft_runner` for generic draft-runner access.
-        # for the EagleDraftInput shape classmethods.
         self.draft_runner: ModelRunner = self.draft_runner_list[0]
 
         # Chain-style MTP: each step propagates its own output hidden states to the
@@ -186,15 +186,17 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
         self.init_lm_head()
 
     def init_attention_backends(self):
-        with self.draft_tp_context(
+        with (
-            self.draft_runner_list[0].tp_group
+            self.draft_tp_context(self.draft_runner_list[0].tp_group),
-        ), speculative_moe_backend_context():
+            speculative_moe_backend_context(),
+        ):
             super().init_attention_backends()
 
     def init_cuda_graphs(self):
-        with self.draft_tp_context(
+        with (
-            self.draft_runner_list[0].tp_group
+            self.draft_tp_context(self.draft_runner_list[0].tp_group),
-        ), speculative_moe_backend_context():
+            speculative_moe_backend_context(),
+        ):
             super().init_cuda_graphs()
 
     def mtp_model_runner(self, step: int):
@@ -734,10 +736,13 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
                     if self.speculative_algorithm.is_standalone()
                     else CaptureHiddenMode.LAST
                 )
+                hidden_size, hidden_dtype = get_draft_recurrent_hidden_state_spec(
+                    self.draft_worker.draft_runner
+                )
                 batch.spec_info = EagleDraftInput.create_idle_input(
                     device=self.device,
-                    hidden_size=EagleDraftInput.hidden_size_for(self.draft_worker),
+                    hidden_size=hidden_size,
-                    dtype=EagleDraftInput.dtype_for(self.draft_worker),
+                    dtype=hidden_dtype,
                     topk=self.topk * self.speculative_num_steps,
                     capture_hidden_mode=capture_mode,
                 )