Unify forward output datastructure (#11124)

python/sglang/srt/speculative/ngram_worker.py

@@ -7,7 +7,7 @@ from sgl_kernel.speculative import reconstruct_indices_from_tree_mask
 
 from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.managers.tp_worker import TpModelWorker
-from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.model_executor.forward_batch_info import ForwardBatchOutput, ForwardMode
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.speculative.cpp_ngram.ngram_cache import NgramCache
 from sglang.srt.speculative.ngram_utils import NgramVerifyInput
@@ -207,17 +207,18 @@ class NGRAMWorker:
             batch_tokens.append(put_ids)
         self.ngram_cache.batch_put(batch_tokens)
 
-    def forward_batch_speculative_generation(self, batch: ScheduleBatch):
+    def forward_batch_generation(self, batch: ScheduleBatch) -> ForwardBatchOutput:
         self._prepare_for_speculative_decoding(batch)
         model_worker_batch = batch.get_model_worker_batch()
-        bid = model_worker_batch.bid
         num_accepted_tokens = 0
 
         if model_worker_batch.forward_mode.is_target_verify():
-            logits_output, _, can_run_cuda_graph = (
-                self.target_worker.forward_batch_generation(
-                    model_worker_batch, skip_sample=True
-                )
+            forward_batch_output = self.target_worker.forward_batch_generation(
+                model_worker_batch, skip_sample=True
+            )
+            logits_output, can_run_cuda_graph = (
+                forward_batch_output.logits_output,
+                forward_batch_output.can_run_cuda_graph,
             )
             verify_input = model_worker_batch.spec_info
             logits_output, next_token_ids, num_accepted_tokens = verify_input.verify(
@@ -227,14 +228,18 @@ class NGRAMWorker:
             batch.forward_mode = ForwardMode.DECODE
 
         else:
+            forward_batch_output = self.target_worker.forward_batch_generation(
+                model_worker_batch
+            )
             logits_output, next_token_ids, can_run_cuda_graph = (
-                self.target_worker.forward_batch_generation(model_worker_batch)
+                forward_batch_output.logits_output,
+                forward_batch_output.next_token_ids,
+                forward_batch_output.can_run_cuda_graph,
             )
 
-        return (
-            logits_output,
-            next_token_ids,
-            bid,
-            num_accepted_tokens,
-            can_run_cuda_graph,
+        return ForwardBatchOutput(
+            logits_output=logits_output,
+            next_token_ids=next_token_ids,
+            num_accepted_tokens=num_accepted_tokens,
+            can_run_cuda_graph=can_run_cuda_graph,
         )