Tiny extract ModelRunnerOutput (#15400)

python/sglang/bench_one_batch.py

@@ -373,7 +373,7 @@ def extend(reqs, model_runner):
     _maybe_prepare_mlp_sync_batch(batch, model_runner)
     model_worker_batch = batch.get_model_worker_batch()
     forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner)
-    logits_output, _ = model_runner.forward(forward_batch)
+    logits_output = model_runner.forward(forward_batch).logits_output
     next_token_ids = model_runner.sample(logits_output, forward_batch)
     return next_token_ids, logits_output.next_token_logits, batch
 
@@ -385,7 +385,7 @@ def decode(input_token_ids, batch, model_runner):
     _maybe_prepare_mlp_sync_batch(batch, model_runner)
     model_worker_batch = batch.get_model_worker_batch()
     forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner)
-    logits_output, _ = model_runner.forward(forward_batch)
+    logits_output = model_runner.forward(forward_batch).logits_output
     next_token_ids = model_runner.sample(logits_output, forward_batch)
     return next_token_ids, logits_output.next_token_logits
 

python/sglang/srt/dllm/algorithm/low_confidence.py

@@ -35,9 +35,8 @@ class LowConfidence(DllmAlgorithm):
             if torch.sum(mask_index).item() == 0:
                 break
 
-            logits_output, can_run_cuda_graph = model_runner.forward(
+            out = model_runner.forward(forward_batch, pp_proxy_tensors=None)
-                forward_batch, pp_proxy_tensors=None
+            logits_output, can_run_cuda_graph = out.logits_output, out.can_run_graph
-            )
 
             x = torch.argmax(logits_output.full_logits, dim=-1)
             p = torch.squeeze(
@@ -58,9 +57,8 @@ class LowConfidence(DllmAlgorithm):
 
             forward_batch.input_ids[transfer_index] = x[transfer_index]
 
-        logits_output, can_run_cuda_graph = model_runner.forward(
+        out = model_runner.forward(forward_batch, pp_proxy_tensors=None)
-            forward_batch, pp_proxy_tensors=None
+        logits_output, can_run_cuda_graph = out.logits_output, out.can_run_graph
-        )
 
         next_token_ids = forward_batch.input_ids[start:]
         return logits_output, next_token_ids, can_run_cuda_graph

python/sglang/srt/managers/scheduler_pp_mixin.py

@@ -614,7 +614,7 @@ class SchedulerPPMixin:
                 forward_batch = ForwardBatch.init_new(
                     model_worker_batch, self.tp_worker.model_runner
                 )
-                _, _ = self.tp_worker.model_runner.forward(
+                _ = self.tp_worker.model_runner.forward(
                     forward_batch=forward_batch, pp_proxy_tensors=pp_proxy
                 )
 

python/sglang/srt/managers/tp_worker.py

@@ -196,7 +196,7 @@ class BaseTpWorker(ABC):
 
     def forward_batch_embedding(self, model_worker_batch: ModelWorkerBatch):
         forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
-        logits_output, _ = self.model_runner.forward(forward_batch)
+        logits_output = self.model_runner.forward(forward_batch).logits_output
         embeddings = logits_output.embeddings
         return embeddings
 
@@ -397,11 +397,12 @@ class TpModelWorker(BaseTpWorker):
             return self._forward_batch_generation_dllm(forward_batch)
 
         if self.pp_group.is_last_rank:
-            logits_output, can_run_cuda_graph = self.model_runner.forward(
+            out = self.model_runner.forward(
                 forward_batch,
                 pp_proxy_tensors=pp_proxy_tensors,
                 skip_attn_backend_init=skip_attn_backend_init,
             )
+            logits_output, can_run_cuda_graph = out.logits_output, out.can_run_graph
             batch_result = GenerationBatchResult(
                 logits_output=logits_output,
                 can_run_cuda_graph=can_run_cuda_graph,
@@ -450,11 +451,12 @@ class TpModelWorker(BaseTpWorker):
 
             return batch_result
         else:
-            pp_proxy_tensors, can_run_cuda_graph = self.model_runner.forward(
+            out = self.model_runner.forward(
                 forward_batch,
                 pp_proxy_tensors=pp_proxy_tensors,
                 skip_attn_backend_init=skip_attn_backend_init,
             )
+            pp_proxy_tensors, can_run_cuda_graph = out.logits_output, out.can_run_graph
             return GenerationBatchResult(
                 pp_hidden_states_proxy_tensors=pp_proxy_tensors,
                 can_run_cuda_graph=can_run_cuda_graph,
@@ -469,9 +471,10 @@ class TpModelWorker(BaseTpWorker):
         else:
             model_worker_batch = batch.get_model_worker_batch(batch.seq_lens_cpu_cache)
 
-        logits_output, can_run_cuda_graph = self.model_runner.forward(
+        out = self.model_runner.forward(
             batch.split_forward_batch, split_forward_count=batch.split_forward_count
         )
+        logits_output, can_run_cuda_graph = out.logits_output, out.can_run_graph
         if logits_output:
             next_token_ids = self.model_runner.sample(logits_output, model_worker_batch)
         else:

python/sglang/srt/model_executor/model_runner.py

@@ -268,6 +268,12 @@ class RankZeroFilter(logging.Filter):
         return True
 
 
+@dataclass
+class ModelRunnerOutput:
+    logits_output: Union[LogitsProcessorOutput, PPProxyTensors]
+    can_run_graph: bool
+
+
 class ModelRunner:
     """ModelRunner runs the forward passes of the models."""
 
@@ -2726,7 +2732,7 @@ class ModelRunner:
         pp_proxy_tensors: Optional[PPProxyTensors] = None,
         reinit_attn_backend: bool = False,
         split_forward_count: int = 1,
-    ) -> Tuple[Union[LogitsProcessorOutput, PPProxyTensors], bool]:
+    ) -> ModelRunnerOutput:
         self.forward_pass_id += 1
 
         with get_global_expert_distribution_recorder().with_forward_pass(
@@ -2753,7 +2759,7 @@ class ModelRunner:
         pp_proxy_tensors: Optional[PPProxyTensors],
         reinit_attn_backend: bool = False,
         split_forward_count: int = 1,
-    ) -> Tuple[Union[LogitsProcessorOutput, PPProxyTensors], bool]:
+    ) -> ModelRunnerOutput:
         mode_check = (
             forward_batch.forward_mode.is_cpu_graph
             if self.device == "cpu"
@@ -2771,7 +2777,7 @@ class ModelRunner:
                 skip_attn_backend_init=skip_attn_backend_init,
                 pp_proxy_tensors=pp_proxy_tensors,
             )
-            return ret, can_run_graph
+            return ModelRunnerOutput(logits_output=ret, can_run_graph=can_run_graph)
 
         # For MLP sync
         if forward_batch.global_num_tokens_cpu is not None:
@@ -2819,7 +2825,7 @@ class ModelRunner:
         ):
             forward_batch.post_forward_mlp_sync_batch(ret)
 
-        return ret, can_run_graph
+        return ModelRunnerOutput(logits_output=ret, can_run_graph=can_run_graph)
 
     def _preprocess_logits(
         self, logits_output: LogitsProcessorOutput, sampling_info: SamplingBatchInfo

python/sglang/srt/speculative/eagle_worker.py

@@ -646,9 +646,9 @@ class EAGLEWorker(TpModelWorker):
             spec_info.hidden_states = hidden_states
 
             # Run forward
-            logits_output, _ = self.draft_model_runner.forward(
+            logits_output = self.draft_model_runner.forward(
                 forward_batch, skip_attn_backend_init=True
-            )
+            ).logits_output
             if self.server_args.enable_nan_detection:
                 detect_nan(logits_output)
             probs = torch.softmax(logits_output.next_token_logits, dim=-1)
@@ -946,7 +946,7 @@ class EAGLEWorker(TpModelWorker):
             model_worker_batch, self.draft_model_runner
         )
         forward_batch.return_logprob = False
-        logits_output, _ = self.draft_model_runner.forward(forward_batch)
+        logits_output = self.draft_model_runner.forward(forward_batch).logits_output
         if self.enable_nan_detection:
             detect_nan(logits_output)
         assert isinstance(forward_batch.spec_info, EagleDraftInput)
@@ -1023,9 +1023,9 @@ class EAGLEWorker(TpModelWorker):
                 self.draft_model_runner.attn_backend.init_forward_metadata(
                     forward_batch
                 )
-            logits_output, _ = self.draft_model_runner.forward(
+            logits_output = self.draft_model_runner.forward(
                 forward_batch, skip_attn_backend_init=True
-            )
+            ).logits_output
             self.capture_for_decode(logits_output, forward_batch.spec_info)
 
         if self.enable_nan_detection:

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -391,9 +391,9 @@ class EagleDraftWorker(BaseDraftWorker):
             spec_info.hidden_states = hidden_states
 
             # Run forward
-            logits_output, _ = self.draft_runner.forward(
+            logits_output = self.draft_runner.forward(
                 forward_batch, skip_attn_backend_init=True
-            )
+            ).logits_output
             if self.server_args.enable_nan_detection:
                 detect_nan(logits_output)
             probs = torch.softmax(logits_output.next_token_logits, dim=-1)
@@ -465,7 +465,7 @@ class EagleDraftWorker(BaseDraftWorker):
 
         # Run forward
         forward_batch = ForwardBatch.init_new(batch, self.draft_runner)
-        logits_output, _ = self.draft_runner.forward(forward_batch)
+        logits_output = self.draft_runner.forward(forward_batch).logits_output
 
         # Update spec_info for the next draft step
         probs = torch.softmax(logits_output.next_token_logits, dim=-1)
@@ -516,9 +516,9 @@ class EagleDraftWorker(BaseDraftWorker):
                 forward_batch
             )
         else:
-            draft_logits_output, _ = self.draft_runner.forward(
+            draft_logits_output = self.draft_runner.forward(
                 forward_batch, skip_attn_backend_init=True
-            )
+            ).logits_output
 
         # Reorganize the spec info for the next batch
         draft_logits_output.next_token_logits = draft_logits_output.next_token_logits[