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 804d421b1..5ba7f697a 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/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: 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 e8f6f5bfa..ca3d65a4a 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 @@ -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() ) diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index ecafdeb3b..14c17c9ea 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/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, diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index c6fd83636..010831fd3 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/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