Rename InputMetadata -> ForwardBatch (#1543)

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, InputMetadata
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 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,8 +511,8 @@ 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 get_forward_batch(self):
+        return ForwardBatch.from_schedule_batch(self)
 
     def mix_with_running(self, running_batch: "ScheduleBatch"):
         self.forward_mode = ForwardMode.MIXED

python/sglang/srt/managers/schedule_policy.py

@@ -32,7 +32,7 @@ from sglang.srt.mem_cache.radix_cache import TreeNode
 CLIP_MAX_NEW_TOKENS = int(os.environ.get("SGLANG_CLIP_MAX_NEW_TOKENS", "4096"))
 
 
-class SchedulerPolicy:
+class SchedulePolicy:
     def __init__(self, policy: str, tree_cache: BasePrefixCache):
         if tree_cache.disable and policy in ["lpm", "dfs-weight"]:
             # LPM and DFS-weight is meaningless when the tree cache is disabled.

python/sglang/srt/managers/scheduler.py

@@ -50,8 +50,8 @@ from sglang.srt.managers.schedule_batch import (
     Req,
     ScheduleBatch,
 )
-from sglang.srt.managers.scheduler_policy import PrefillAdder, SchedulerPolicy
-from sglang.srt.managers.tp_worker import ModelTpWorker
+from sglang.srt.managers.schedule_policy import PrefillAdder, SchedulePolicy
+from sglang.srt.managers.tp_worker import TpModelWorker
 from sglang.srt.mem_cache.chunk_cache import ChunkCache
 from sglang.srt.mem_cache.radix_cache import RadixCache
 from sglang.srt.server_args import PortArgs, ServerArgs
@@ -134,7 +134,7 @@ class Scheduler:
         )
 
         # Launch a tensor parallel worker
-        self.tp_worker = ModelTpWorker(
+        self.tp_worker = TpModelWorker(
             gpu_id=gpu_id,
             tp_rank=tp_rank,
             server_args=server_args,
@@ -179,7 +179,7 @@ class Scheduler:
                 disable=server_args.disable_radix_cache,
             )
         self.tree_cache_metrics = {"total": 0, "hit": 0}
-        self.policy = SchedulerPolicy(self.schedule_policy, self.tree_cache)
+        self.policy = SchedulePolicy(self.schedule_policy, self.tree_cache)
 
         # Init running status
         self.waiting_queue: List[Req] = []
@@ -575,9 +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()
+                forward_batch = batch.get_forward_batch()
                 logits_output, next_token_ids = self.tp_worker.forward_batch_generation(
-                    input_metadata, batch
+                    forward_batch, batch
                 )
                 batch.sampling_info.penalizer_orchestrator.cumulate_output_tokens(
                     next_token_ids
@@ -641,8 +641,8 @@ class Scheduler:
                     )
         else:
             assert batch.extend_num_tokens != 0
-            input_metadata = batch.get_input_metadata()
-            embeddings = self.tp_worker.forward_batch_embedding(input_metadata)
+            forward_batch = batch.get_forward_batch()
+            embeddings = self.tp_worker.forward_batch_embedding(forward_batch)
 
             # Check finish conditions
             for i, req in enumerate(batch.reqs):
@@ -771,9 +771,9 @@ class Scheduler:
         batch.prepare_for_decode()
 
         # Forward and sample the next tokens
-        input_metadata = batch.get_input_metadata()
+        forward_batch = batch.get_forward_batch()
         logits_output, next_token_ids = self.tp_worker.forward_batch_generation(
-            input_metadata, batch
+            forward_batch, batch
         )
         batch.sampling_info.penalizer_orchestrator.cumulate_output_tokens(
             next_token_ids

python/sglang/srt/managers/tp_worker.py

@@ -21,7 +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.forward_batch_info import ForwardBatch
 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
@@ -29,7 +29,9 @@ from sglang.srt.utils import broadcast_pyobj, is_multimodal_model, set_random_se
 logger = logging.getLogger(__name__)
 
 
-class ModelTpWorker:
+class TpModelWorker:
+    """A tensor parallel model worker."""
+
     def __init__(
         self,
         gpu_id: int,
@@ -106,13 +108,13 @@ class ModelTpWorker:
             self.random_seed,
         )
 
-    def forward_batch_generation(self, input_metadata: InputMetadata, batch):
-        logits_output = self.model_runner.forward(input_metadata)
+    def forward_batch_generation(self, forward_batch: ForwardBatch, batch):
+        logits_output = self.model_runner.forward(forward_batch)
         next_token_ids = self.model_runner.sample(logits_output, batch)
         return logits_output, next_token_ids
 
-    def forward_batch_embedding(self, input_metadata: InputMetadata):
-        logits_output = self.model_runner.forward(input_metadata)
+    def forward_batch_embedding(self, forward_batch: ForwardBatch):
+        logits_output = self.model_runner.forward(forward_batch)
         embeddings = logits_output.embeddings.tolist()
         return embeddings