fix: fix the wrong return value type of draft model runner (#18105)

Co-authored-by: Liangsheng Yin <lsyincs@gmail.com>

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()

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:

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()