Let ModelRunner take InputMetadata as input, instead of ScheduleBatch (#1541)

python/sglang/srt/managers/schedule_batch.py

@@ -29,7 +29,7 @@ from sglang.srt.constrained.jump_forward import JumpForwardMap
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
 from sglang.srt.mem_cache.chunk_cache import ChunkCache
 from sglang.srt.mem_cache.memory_pool import BaseTokenToKVPool, ReqToTokenPool
-from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.model_executor.forward_batch_info import ForwardMode, InputMetadata
 from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
 from sglang.srt.sampling.sampling_params import SamplingParams
 from sglang.srt.server_args import ServerArgs
@@ -511,6 +511,9 @@ class ScheduleBatch:
         self.extend_logprob_start_lens_cpu = [r.extend_logprob_start_len for r in reqs]
         self.sampling_info = SamplingBatchInfo.from_schedule_batch(self, vocab_size)
 
+    def get_input_metadata(self):
+        return InputMetadata.from_schedule_batch(self)
+
     def mix_with_running(self, running_batch: "ScheduleBatch"):
         self.forward_mode = ForwardMode.MIXED
         running_bs = running_batch.batch_size()

python/sglang/srt/managers/scheduler.py

@@ -575,8 +575,9 @@ class Scheduler:
         if self.is_generation:
             # Forward and sample the next tokens
             if batch.extend_num_tokens != 0:
+                input_metadata = batch.get_input_metadata()
                 logits_output, next_token_ids = self.tp_worker.forward_batch_generation(
-                    batch
+                    input_metadata, batch
                 )
                 batch.sampling_info.penalizer_orchestrator.cumulate_output_tokens(
                     next_token_ids
@@ -640,7 +641,8 @@ class Scheduler:
                     )
         else:
             assert batch.extend_num_tokens != 0
-            embeddings = self.tp_worker.forward_batch_embedding(batch)
+            input_metadata = batch.get_input_metadata()
+            embeddings = self.tp_worker.forward_batch_embedding(input_metadata)
 
             # Check finish conditions
             for i, req in enumerate(batch.reqs):
@@ -769,7 +771,10 @@ class Scheduler:
         batch.prepare_for_decode()
 
         # Forward and sample the next tokens
-        logits_output, next_token_ids = self.tp_worker.forward_batch_generation(batch)
+        input_metadata = batch.get_input_metadata()
+        logits_output, next_token_ids = self.tp_worker.forward_batch_generation(
+            input_metadata, batch
+        )
         batch.sampling_info.penalizer_orchestrator.cumulate_output_tokens(
             next_token_ids
         )

python/sglang/srt/managers/tp_worker.py

@@ -21,6 +21,7 @@ import logging
 from sglang.srt.configs.model_config import ModelConfig
 from sglang.srt.hf_transformers_utils import get_processor, get_tokenizer
 from sglang.srt.managers.io_struct import UpdateWeightReqInput
+from sglang.srt.model_executor.forward_batch_info import InputMetadata
 from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import broadcast_pyobj, is_multimodal_model, set_random_seed
@@ -105,13 +106,13 @@ class ModelTpWorker:
             self.random_seed,
         )
 
-    def forward_batch_generation(self, batch):
-        logits_output = self.model_runner.forward(batch)
+    def forward_batch_generation(self, input_metadata: InputMetadata, batch):
+        logits_output = self.model_runner.forward(input_metadata)
         next_token_ids = self.model_runner.sample(logits_output, batch)
         return logits_output, next_token_ids
 
-    def forward_batch_embedding(self, batch):
-        logits_output = self.model_runner.forward(batch)
+    def forward_batch_embedding(self, input_metadata: InputMetadata):
+        logits_output = self.model_runner.forward(input_metadata)
         embeddings = logits_output.embeddings.tolist()
         return embeddings