diff --git a/python/sglang/bench_one_batch.py b/python/sglang/bench_one_batch.py index 2a5164299..b1bc2ae03 100644 --- a/python/sglang/bench_one_batch.py +++ b/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 diff --git a/python/sglang/srt/dllm/algorithm/low_confidence.py b/python/sglang/srt/dllm/algorithm/low_confidence.py index 2f7b8f2b0..c8912b47c 100644 --- a/python/sglang/srt/dllm/algorithm/low_confidence.py +++ b/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( - forward_batch, pp_proxy_tensors=None - ) + out = model_runner.forward(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( - forward_batch, pp_proxy_tensors=None - ) + out = model_runner.forward(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 diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index c7d0ea51b..3467069bd 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/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 ) diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index bed034246..0a132c709 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/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: diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index a8cc32440..4ab0240dd 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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 diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index b84c8ed72..db25c5066 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/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: diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index ce68e3cdd..ab0b2d6a3 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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[