Clean up follow-ups for eagle hidden dim clean up (#29546)

python/sglang/srt/managers/scheduler.py

@@ -1124,16 +1124,11 @@ class Scheduler(
         )
 
         # todo: should we fix this when enabling mtp or it doesn't matter since we only enable mtp in decode node thus we don't transfer draft kvs between P and D?
-        draft_token_to_kv_pool, model_config = kv_cache_builder.get_draft_kv_pool(
+        draft_token_to_kv_pool = kv_cache_builder.get_draft_kv_pool(
             draft_worker=self.draft_worker,
             spec_algorithm=self.spec_algorithm,
             server_args=self.server_args,
         )
-        # Default to the target model_config so the MetadataBuffers branches
-        # below can always access it; overridden by the draft model_config
-        # when this node runs a spec module.
-        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

python/sglang/srt/mem_cache/kv_cache_builder.py

@@ -52,17 +52,17 @@ def get_draft_kv_pool(
     spec_algorithm: SpeculativeAlgorithm,
     server_args: ServerArgs,
 ):
-    """Return (draft_token_to_kv_pool, draft_model_config) for the current
+    """Return the draft token-to-KV pool for the current draft worker,
-    draft worker, or (None, None) when no draft KV pool is available."""
+    or None when no draft KV pool is available."""
     if draft_worker is None or spec_algorithm.is_ngram():
-        return None, None
+        return None
 
     # V2 workers nest the draft runner under `.draft_worker`.
     if server_args.enable_multi_layer_eagle:
         draft_runner = draft_worker.draft_worker.draft_runner_list[0]
     else:
         draft_runner = draft_worker.draft_worker.draft_runner
-    return draft_runner.token_to_kv_pool, draft_runner.model_config
+    return draft_runner.token_to_kv_pool
 
 
 def maybe_register_hicache_draft(
@@ -78,7 +78,7 @@ def maybe_register_hicache_draft(
     if not enable_hierarchical_cache:
         return
 
-    draft_kv_pool, _ = get_draft_kv_pool(
+    draft_kv_pool = get_draft_kv_pool(
         draft_worker=draft_worker,
         spec_algorithm=spec_algorithm,
         server_args=server_args,

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

@@ -27,6 +27,7 @@ Backend selection comes from cuda_graph_config.prefill:
 
 from __future__ import annotations
 
+import inspect
 import logging
 import warnings
 from typing import TYPE_CHECKING, Dict, Optional, Union
@@ -191,13 +192,6 @@ 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,
         )
 
@@ -294,6 +288,10 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                 language_model.model, "layers"
             ):
                 self.layer_model = language_model.model
+                params = list(inspect.signature(self.layer_model.forward).parameters)
+                self._input_embeds_arg_idx = (
+                    params.index("input_embeds") if "input_embeds" in params else None
+                )
             else:
                 raise RuntimeError(
                     f"BCG could not resolve inner layer_model on "
@@ -880,6 +878,8 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                 shape_key = ShapeKey(size=self._static_num_tokens)
                 static_n = self._static_num_tokens
 
+                ie_idx = self._input_embeds_arg_idx
+
                 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)
@@ -891,8 +891,8 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                     # 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:
+                        if ie is None and ie_idx is not None and len(args) > ie_idx:
-                            ie = args[3]
+                            ie = args[ie_idx]
                         if ie is not None:
                             self.buffer_registry.get_slot("input_embeds").slice_for(
                                 1, static_n