diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index e9b9161a7..6625b1750 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/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.process_prebuilt_extend(self.server_args, self.model_config) + new_batch.prepare_for_prebuilt() + new_batch.process_prebuilt(self.server_args, self.model_config) return new_batch diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index eb40d4be1..b98467aba 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/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""" diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 369aa1410..39a323dd5 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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(): - return self._run_batch_prebuilt_extend(batch) + if batch.forward_mode.is_prebuilt(): + 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(): - self.process_batch_result_prebuilt_extend(batch) + elif batch.forward_mode.is_prebuilt(): + self.process_batch_result_prebuilt(batch) elif batch.forward_mode.is_idle(): if self.enable_overlap: if result.copy_done is not None: diff --git a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py b/python/sglang/srt/managers/scheduler_dp_attn_mixin.py index d180de410..13662cb37 100644 --- a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py +++ b/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] - local_batch.global_num_tokens_for_logprob = [ - mlp_sync_info.num_tokens_for_logprob - ] + batch.global_num_tokens = [mlp_sync_info.num_tokens] + batch.global_num_tokens_for_logprob = [mlp_sync_info.num_tokens_for_logprob] else: - local_batch.global_num_tokens = mlp_sync_info.global_num_tokens - local_batch.global_num_tokens_for_logprob = ( + batch.global_num_tokens = mlp_sync_info.global_num_tokens + 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 - local_batch.tbo_split_seq_index = mlp_sync_info.tbo_split_seq_index - local_batch.global_forward_mode = mlp_sync_info.global_forward_mode + batch.is_extend_in_batch = mlp_sync_info.is_extend_in_batch + batch.tbo_split_seq_index = mlp_sync_info.tbo_split_seq_index + 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: - local_batch = get_idle_batch() - - batch_to_gather = local_batch - if need_idle_batch and local_batch.forward_mode.is_prebuilt_extend(): - # 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 - - if local_batch is not None: + if need_idle_batch: + batch_to_gather = local_batch + 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 + batch_to_gather = local_batch.inner_idle_batch = get_idle_batch() _update_gather_batch(batch_to_gather, mlp_sync_info, require_mlp_tp_gather) return local_batch diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index 38f0ce0e8..8465db477 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -42,21 +42,20 @@ 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.completion_time - ) = time.perf_counter() - trace_slice_end( - RequestStage.DECODE_QUICK_FINISH, - req.rid, - thread_finish_flag=True, - ) - self.tree_cache.cache_finished_req(req) + req.check_finished() + if req.finished(): + req.time_stats.forward_entry_time = req.time_stats.completion_time = ( + time.perf_counter() + ) + trace_slice_end( + RequestStage.DECODE_QUICK_FINISH, + req.rid, + thread_finish_flag=True, + ) + self.tree_cache.cache_finished_req(req) # Note: Logprobs should be handled on the prefill engine. trace_slice_batch(RequestStage.DECODE_FAKE_OUTPUT, batch.reqs) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 98f129875..4c3a329a1 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/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): - return self == ForwardMode.PREBUILT_EXTEND + def is_prebuilt(self): + return self == ForwardMode.PREBUILT @total_ordering diff --git a/python/sglang/srt/two_batch_overlap.py b/python/sglang/srt/two_batch_overlap.py index 531331c6b..16fa3ef15 100644 --- a/python/sglang/srt/two_batch_overlap.py +++ b/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