From a1189068faa0a9f0db3ea851b1fe9713d2afeeca Mon Sep 17 00:00:00 2001 From: Yingchun Lai Date: Mon, 9 Feb 2026 12:51:35 +0800 Subject: [PATCH] fix: fix the wrong return value type of draft model runner (#18105) Co-authored-by: Liangsheng Yin --- python/sglang/srt/managers/tp_worker.py | 4 ++-- .../srt/speculative/multi_layer_eagle_worker.py | 17 ++++++++++++----- .../speculative/multi_layer_eagle_worker_v2.py | 4 ++-- 3 files changed, 16 insertions(+), 9 deletions(-) diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index 37416ba8b..7a334f628 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -16,7 +16,7 @@ from __future__ import annotations import logging from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, List, Optional import torch @@ -236,7 +236,7 @@ class TpModelWorker(BaseTpWorker): self.token_to_kv_pool_allocator = token_to_kv_pool_allocator # MTP model runners - self.model_runner_list = [] + self.model_runner_list: List[ModelRunner] = [] self._init_model_config() self._init_model_runner() diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker.py b/python/sglang/srt/speculative/multi_layer_eagle_worker.py index 305eff355..21fc97777 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker.py @@ -14,7 +14,7 @@ import logging import time -from typing import List, Optional, Tuple +from typing import TYPE_CHECKING, List, Optional, Tuple import torch @@ -56,6 +56,9 @@ from sglang.srt.speculative.spec_utils import ( ) from sglang.srt.utils import empty_context, get_available_gpu_memory, is_cuda, is_npu +if TYPE_CHECKING: + from sglang.srt.model_executor.model_runner import ModelRunner + _is_npu = is_npu() if is_cuda(): @@ -226,7 +229,7 @@ class MultiLayerEagleWorker(TpModelWorker): f"Capture draft extend cuda graph end. Time elapsed: {time.perf_counter() - tic:.2f} s. mem usage={(before_mem - after_mem):.2f} GB. avail mem={after_mem:.2f} GB." ) - def mtp_model_runner(self, layer_id: int): + def mtp_model_runner(self, layer_id: int) -> ModelRunner: return self.model_runner_list[layer_id] def forward_batch_generation(self, batch: ScheduleBatch) -> GenerationBatchResult: @@ -613,7 +616,9 @@ class MultiLayerEagleWorker(TpModelWorker): topk_p_list = [] topk_index_list = [] for step in range(self.speculative_num_steps): - logits_output, _ = self.mtp_model_runner(step).forward(forward_batch) + logits_output = ( + self.mtp_model_runner(step).forward(forward_batch).logits_output + ) if self.enable_nan_detection: detect_nan(logits_output) probs = torch.softmax(logits_output.next_token_logits, dim=-1) @@ -718,8 +723,10 @@ class MultiLayerEagleWorker(TpModelWorker): self.mtp_model_runner(step).attn_backend.init_forward_metadata( forward_batch ) - logits_output, _ = self.mtp_model_runner(step).forward( - forward_batch, skip_attn_backend_init=True + logits_output = ( + self.mtp_model_runner(step) + .forward(forward_batch, skip_attn_backend_init=True) + .logits_output ) if self.enable_nan_detection: 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 8320eaa7f..a0432e2cf 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -45,7 +45,7 @@ from sglang.srt.speculative.spec_utils import ( from sglang.srt.utils.common import empty_context, fast_topk if TYPE_CHECKING: - from sglang.srt.model_executor.model_runner import ModelRunnerOutput + from sglang.srt.model_executor.model_runner import ModelRunner, ModelRunnerOutput logger = logging.getLogger(__name__) @@ -125,7 +125,7 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker): ) # Alias for better readability - self.draft_runner_list = self.draft_worker.model_runner_list + self.draft_runner_list: List[ModelRunner] = self.draft_worker.model_runner_list self.init_lm_head()