diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 09231e1b8..f12843b74 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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=( - model_config.spec_hidden_size - 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 - ), + hidden_size=disagg_hidden_size, + hidden_states_dtype=disagg_hidden_states_dtype, 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=( - model_config.spec_hidden_size - 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 - ), + hidden_size=disagg_hidden_size, + hidden_states_dtype=disagg_hidden_states_dtype, custom_mem_pool=self.token_to_kv_pool_allocator.get_kvcache().maybe_get_custom_mem_pool(), ) diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index bfbe94774..ab7995237 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/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 = ( diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 43103b99d..0a9347609 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/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 - # (e.g. EAGLE3 stacks 3 target layers), so read the dim from the - # draft model's input fc when available; fall back to the per-layer - # dim for arches without an fc projection. + # Draft consumes aux-concatenated hidden states from the target + # (e.g. EAGLE3 stacks 3 target layers), so capture the pre-reduction + # width when the draft model exposes it. 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 - - 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) - ) + hidden_dim = get_draft_input_from_target_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( - ForwardContext(attn_backend=self.model_runner.attn_backend) - ), set_tc_piecewise_forward_context( - forward_batch, - self.attention_layers, - self.quant_config, - self.moe_layers, - self.moe_fusions, - dsa_indexers=self.dsa_indexers, + with ( + forward_context( + ForwardContext(attn_backend=self.model_runner.attn_backend) + ), + 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,17 +904,20 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): original_layer_forward = self.layer_model.forward self.layer_model.forward = replay_layer_forward try: - with forward_context( - ForwardContext(attn_backend=self.model_runner.attn_backend) - ), set_tc_piecewise_forward_context( - static_forward_batch, - self.attention_layers, - self.quant_config, - self.moe_layers, - self.moe_fusions, - dsa_indexers=self.dsa_indexers, - num_tokens=static_num_tokens, - raw_num_tokens=raw_num_tokens, + with ( + forward_context( + ForwardContext(attn_backend=self.model_runner.attn_backend) + ), + set_tc_piecewise_forward_context( + static_forward_batch, + self.attention_layers, + self.quant_config, + self.moe_layers, + self.moe_fusions, + 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,17 +931,20 @@ 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( - ForwardContext(attn_backend=self.model_runner.attn_backend) - ), set_tc_piecewise_forward_context( - static_forward_batch, - self.attention_layers, - self.quant_config, - self.moe_layers, - self.moe_fusions, - dsa_indexers=self.dsa_indexers, - num_tokens=static_num_tokens, - raw_num_tokens=raw_num_tokens, + with ( + forward_context( + ForwardContext(attn_backend=self.model_runner.attn_backend) + ), + set_tc_piecewise_forward_context( + static_forward_batch, + self.attention_layers, + self.quant_config, + self.moe_layers, + self.moe_fusions, + 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 diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index c2e324125..3c0ff9c54 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/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 diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index c755ea056..1394f8150 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/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 diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index 5efbc771f..926886054 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/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, diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index ffca4ab41..14f8ebb31 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/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: - """Derive the hidden dimension of target hidden states fed to the draft model.""" - hf_config = model_runner.model_config.hf_config - eagle_config = getattr(hf_config, "eagle_config", {}) - use_aux = eagle_config.get("use_aux_hidden_state", False) +def get_draft_input_from_target_hidden_dim(model_runner: ModelRunner) -> int: + """Width of the target hidden states fed into the draft model. + + This is the single source of truth and is derived entirely from config: for + 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: - base = getattr(hf_config, "target_hidden_size", None) - if base is None: - base = model_runner.model_config.hidden_size - layer_ids = eagle_config.get("eagle_aux_hidden_state_layer_ids", []) - num_aux = max(len(layer_ids), 1) - return base * num_aux - return model_runner.model_config.spec_hidden_size + if not (spec_algorithm is not None and spec_algorithm.is_eagle3() and use_aux): + return model_config.spec_hidden_size + + target_hidden = getattr(hf_config, "target_hidden_size", None) + if target_hidden is None: + target_hidden = model_config.hidden_size + num_aux = getattr(hf_config, "num_aux_hidden_states", None) + 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( diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index d0fed6405..0d158f429 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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( - self.draft_runner.tp_group - ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context(): + with ( + self.draft_tp_context(self.draft_runner.tp_group), + 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( - self.draft_runner.tp_group - ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context(): + with ( + self.draft_tp_context(self.draft_runner.tp_group), + 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), - dtype=EagleDraftInput.dtype_for(self.draft_worker), + hidden_size=hidden_size, + 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, ) diff --git a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py index 31ce5c826..0f8f047d7 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py +++ b/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) - dtype = EagleDraftExtendInput.dtype_for(self.eagle_worker) + hidden_size = get_draft_input_from_target_hidden_dim(model_runner) + 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) diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 6f4364f24..5361342c4 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/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 - # for the EagleDraftInput shape classmethods. + # Match `EagleDraftWorker.draft_runner` for generic draft-runner access. 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( - self.draft_runner_list[0].tp_group - ), speculative_moe_backend_context(): + with ( + self.draft_tp_context(self.draft_runner_list[0].tp_group), + speculative_moe_backend_context(), + ): super().init_attention_backends() def init_cuda_graphs(self): - with self.draft_tp_context( - self.draft_runner_list[0].tp_group - ), speculative_moe_backend_context(): + with ( + self.draft_tp_context(self.draft_runner_list[0].tp_group), + 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), - dtype=EagleDraftInput.dtype_for(self.draft_worker), + hidden_size=hidden_size, + dtype=hidden_dtype, topk=self.topk * self.speculative_num_steps, capture_hidden_mode=capture_mode, )