Minor code cleanup / improvement for PREBUILT_EXTEND mode (#12948)

python/sglang/srt/disaggregation/decode.py

@@ -588,7 +588,7 @@ class DecodePreallocQueue:
         #       the extend batch is not in any queue, so we need to explicitly add the tokens slots here
         if (
             self.scheduler.last_batch
-            and self.scheduler.last_batch.forward_mode.is_prebuilt_extend()
+            and self.scheduler.last_batch.forward_mode.is_prebuilt()
         ):
             allocatable_tokens -= self.num_reserved_decode_tokens * len(
                 self.scheduler.last_batch.reqs
@@ -850,11 +850,14 @@ class SchedulerDisaggregationDecodeMixin:
             self.launch_batch_sample_if_needed(batch_result)
             self.last_batch = batch
 
-    def _run_batch_prebuilt_extend(
+    def _run_batch_prebuilt(
         self: Scheduler, batch: ScheduleBatch
     ) -> GenerationBatchResult:
         if batch.inner_idle_batch is not None:
-            return self.run_batch(batch.inner_idle_batch)
+            idle_batch = batch.inner_idle_batch
+            # Reset the inner idle batch to avoid reusing it.
+            batch.inner_idle_batch = None
+            return self.run_batch(idle_batch)
 
         return GenerationBatchResult()
 
@@ -864,7 +867,7 @@ class SchedulerDisaggregationDecodeMixin:
         """Create fake completed prefill if possible and merge with running batch"""
         # Merge the prefill batch into the running batch
         last_batch = self.last_batch
-        if last_batch and last_batch.forward_mode.is_prebuilt_extend():
+        if last_batch and last_batch.forward_mode.is_prebuilt():
             # chunked prefill doesn't happen in decode instance.
             assert self.chunked_req is None
             # Filter finished batches.
@@ -947,8 +950,8 @@ class SchedulerDisaggregationDecodeMixin:
         )
 
         # construct fake completed prefill
-        new_batch.prepare_for_prebuilt_extend()
+        new_batch.prepare_for_prebuilt()
-        new_batch.process_prebuilt_extend(self.server_args, self.model_config)
+        new_batch.process_prebuilt(self.server_args, self.model_config)
 
         return new_batch
 

python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py

@@ -20,13 +20,13 @@ if TYPE_CHECKING:
 
 class ScheduleBatchDisaggregationDecodeMixin:
 
-    def prepare_for_prebuilt_extend(self: ScheduleBatch):
+    def prepare_for_prebuilt(self: ScheduleBatch):
         """
         Prepare a prebuilt extend by populate metadata
         Adapted from .prepare_for_extend().
         """
 
-        self.forward_mode = ForwardMode.PREBUILT_EXTEND
+        self.forward_mode = ForwardMode.PREBUILT
         reqs = self.reqs
         input_ids = [r.fill_ids[len(r.prefix_indices) :] for r in reqs]
         extend_num_tokens = sum(len(ids) for ids in input_ids)
@@ -100,7 +100,7 @@ class ScheduleBatchDisaggregationDecodeMixin:
             self.model_config.vocab_size,
         )
 
-    def process_prebuilt_extend(
+    def process_prebuilt(
         self: ScheduleBatch, server_args: ServerArgs, model_config: ModelConfig
     ):
         """Assign the buffered last input id to schedule batch"""

python/sglang/srt/managers/scheduler.py

@@ -1964,8 +1964,8 @@ class Scheduler(
                 req.time_stats.prefill_start_time = current_time
 
         # Place holder handling for pd-disagg decode event loop
-        if batch.forward_mode.is_prebuilt_extend():
+        if batch.forward_mode.is_prebuilt():
-            return self._run_batch_prebuilt_extend(batch)
+            return self._run_batch_prebuilt(batch)
 
         # Run forward
         if self.is_generation:
@@ -2099,8 +2099,8 @@ class Scheduler(
             trace_slice_batch(RequestStage.DECODE_LOOP, batch.reqs)
         elif batch.forward_mode.is_extend():
             self.process_batch_result_prefill(batch, result)
-        elif batch.forward_mode.is_prebuilt_extend():
+        elif batch.forward_mode.is_prebuilt():
-            self.process_batch_result_prebuilt_extend(batch)
+            self.process_batch_result_prebuilt(batch)
         elif batch.forward_mode.is_idle():
             if self.enable_overlap:
                 if result.copy_done is not None:

python/sglang/srt/managers/scheduler_dp_attn_mixin.py

@@ -74,25 +74,21 @@ def _update_gather_batch(
     mlp_sync_info: MLPSyncBatchInfo,
     require_mlp_tp_gather: bool,
 ):
-    local_batch = batch
-
     # TODO: handle the case when moe_dense_tp_size != 1
     if not require_mlp_tp_gather:
-        local_batch.global_num_tokens = [mlp_sync_info.num_tokens]
+        batch.global_num_tokens = [mlp_sync_info.num_tokens]
-        local_batch.global_num_tokens_for_logprob = [
+        batch.global_num_tokens_for_logprob = [mlp_sync_info.num_tokens_for_logprob]
-            mlp_sync_info.num_tokens_for_logprob
-        ]
     else:
-        local_batch.global_num_tokens = mlp_sync_info.global_num_tokens
+        batch.global_num_tokens = mlp_sync_info.global_num_tokens
-        local_batch.global_num_tokens_for_logprob = (
+        batch.global_num_tokens_for_logprob = (
             mlp_sync_info.global_num_tokens_for_logprob
         )
-    local_batch.is_extend_in_batch = mlp_sync_info.is_extend_in_batch
+    batch.is_extend_in_batch = mlp_sync_info.is_extend_in_batch
-    local_batch.tbo_split_seq_index = mlp_sync_info.tbo_split_seq_index
+    batch.tbo_split_seq_index = mlp_sync_info.tbo_split_seq_index
-    local_batch.global_forward_mode = mlp_sync_info.global_forward_mode
+    batch.global_forward_mode = mlp_sync_info.global_forward_mode
 
     # Check forward mode for cuda graph
-    local_batch.can_run_dp_cuda_graph = mlp_sync_info.can_cuda_graph
+    batch.can_run_dp_cuda_graph = mlp_sync_info.can_cuda_graph
 
 
 def prepare_mlp_sync_batch_raw(
@@ -107,7 +103,7 @@ def prepare_mlp_sync_batch_raw(
     offload_tags: set[str],
 ):
     # Check if other DP workers have running batches
-    if local_batch is None or local_batch.forward_mode.is_prebuilt_extend():
+    if local_batch is None or local_batch.forward_mode.is_prebuilt():
         num_tokens = 0
         num_tokens_for_logprob = 0
     elif local_batch.forward_mode.is_decode():
@@ -131,7 +127,7 @@ def prepare_mlp_sync_batch_raw(
     can_cuda_graph = (
         local_batch is None
         or local_batch.forward_mode.is_decode_or_idle()
-        or local_batch.forward_mode.is_prebuilt_extend()
+        or local_batch.forward_mode.is_prebuilt()
     ) and not disable_cuda_graph
 
     is_extend_in_batch = local_batch.forward_mode.is_extend() if local_batch else False
@@ -165,16 +161,13 @@ def prepare_mlp_sync_batch_raw(
     )
 
     need_idle_batch = max(mlp_sync_info.global_num_tokens) > 0
-    if need_idle_batch and local_batch is None:
+    if need_idle_batch:
-        local_batch = get_idle_batch()
-
         batch_to_gather = local_batch
-    if need_idle_batch and local_batch.forward_mode.is_prebuilt_extend():
+        if local_batch is None:
+            batch_to_gather = local_batch = get_idle_batch()
+        elif local_batch.forward_mode.is_prebuilt():
             # NOTE: for prebuilt batch, we add an inner idle batch to run MLP sync
-        local_batch.inner_idle_batch = get_idle_batch()
+            batch_to_gather = local_batch.inner_idle_batch = get_idle_batch()
-        batch_to_gather = local_batch.inner_idle_batch
-
-    if local_batch is not None:
         _update_gather_batch(batch_to_gather, mlp_sync_info, require_mlp_tp_gather)
 
     return local_batch

python/sglang/srt/managers/scheduler_output_processor_mixin.py

@@ -42,15 +42,14 @@ class SchedulerOutputProcessorMixin:
     We put them into a separate file to make the `scheduler.py` shorter.
     """
 
-    def process_batch_result_prebuilt_extend(self: Scheduler, batch: ScheduleBatch):
+    def process_batch_result_prebuilt(self: Scheduler, batch: ScheduleBatch):
         assert self.disaggregation_mode == DisaggregationMode.DECODE
-        for req in batch.reqs:
         for req in batch.reqs:
             req.check_finished()
             if req.finished():
-                    req.time_stats.forward_entry_time = (
+                req.time_stats.forward_entry_time = req.time_stats.completion_time = (
-                        req.time_stats.completion_time
+                    time.perf_counter()
-                    ) = time.perf_counter()
+                )
                 trace_slice_end(
                     RequestStage.DECODE_QUICK_FINISH,
                     req.rid,

python/sglang/srt/model_executor/forward_batch_info.py

@@ -81,7 +81,7 @@ class ForwardMode(IntEnum):
 
     # Used in disaggregated decode worker
     # Represent a batch of requests having their KV cache ready to start decoding
-    PREBUILT_EXTEND = auto()
+    PREBUILT = auto()
 
     # Split Prefill for PD multiplexing
     SPLIT_PREFILL = auto()
@@ -152,8 +152,8 @@ class ForwardMode(IntEnum):
             and not self.is_draft_extend()
         )
 
-    def is_prebuilt_extend(self):
+    def is_prebuilt(self):
-        return self == ForwardMode.PREBUILT_EXTEND
+        return self == ForwardMode.PREBUILT
 
 
 @total_ordering

python/sglang/srt/two_batch_overlap.py

@@ -79,7 +79,7 @@ def compute_split_seq_index(
     elif forward_mode.is_target_verify() or forward_mode.is_decode():
         assert token_num_per_seq is not None
         return (num_tokens // token_num_per_seq) // 2
-    elif forward_mode.is_idle() or forward_mode.is_prebuilt_extend():
+    elif forward_mode.is_idle() or forward_mode.is_prebuilt():
         assert num_tokens == 0
         return 0
     else:
@@ -381,7 +381,7 @@ class TboDPAttentionPreparer:
                 or local_batch.forward_mode.is_decode()
             ):
                 num_tokens = local_batch.batch_size() * token_num_per_seq
-            elif local_batch.forward_mode.is_prebuilt_extend():
+            elif local_batch.forward_mode.is_prebuilt():
                 num_tokens = 0
             else:
                 num_tokens = local_batch.extend_num_tokens