From 7a2d3df96fa16874a8bc1501ca113be659ae37f1 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Tue, 3 Mar 2026 00:05:27 -0800 Subject: [PATCH] Apply default stream to priority 0 in scheduling. (#16438) --- python/sglang/srt/managers/scheduler.py | 69 ++++++++++--------- .../sglang/srt/managers/scheduler_pp_mixin.py | 4 +- 2 files changed, 39 insertions(+), 34 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 23c83ba2f..d050b5c40 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -993,9 +993,6 @@ class Scheduler( def init_overlap(self): self.device_module = torch.get_device_module(self.device) - self.default_stream: CudaStream = self.device_module.current_stream() - if self.device == "cpu": - self.default_stream.synchronize = lambda: None # No-op for CPU self.forward_stream_ctx: CudaStreamContext = self.device_module.stream( self.forward_stream @@ -2325,7 +2322,7 @@ class Scheduler( future_indices = self.future_map.alloc_future_indices(bs) with self.forward_stream_ctx: - self.forward_stream.wait_stream(self.default_stream) + self.forward_stream.wait_stream(self.schedule_stream) self.future_map.resolve_future(model_worker_batch) with self.record_forward_metrics(batch): batch_result = self.model_worker.forward_batch_generation( @@ -2401,7 +2398,7 @@ class Scheduler( if self.enable_overlap: self.record_batch_in_overlap(model_worker_batch) with self.forward_stream_ctx: - self.forward_stream.wait_stream(self.default_stream) + self.forward_stream.wait_stream(self.schedule_stream) embeddings = self.tp_worker.forward_batch_embedding( model_worker_batch ) @@ -2439,7 +2436,7 @@ class Scheduler( return with self.forward_stream_ctx: - self.forward_stream.wait_stream(self.default_stream) + self.forward_stream.wait_stream(self.schedule_stream) _batch_result = batch_result.delay_sample_func() assert _batch_result is batch_result self.future_map.store_to_map(batch_result.future_indices, batch_result) @@ -3081,6 +3078,35 @@ class SenderWrapper: self.socket.send_pyobj(output) +def dispatch_event_loop(scheduler: Scheduler): + # Dispatch to the appropriate event loop based on the disaggregation mode + server_args = scheduler.server_args + disaggregation_mode: DisaggregationMode = scheduler.disaggregation_mode + if disaggregation_mode == DisaggregationMode.NULL: + if scheduler.enable_pdmux: + scheduler.event_loop_pdmux() + elif server_args.pp_size > 1: + scheduler.event_loop_pp() + elif scheduler.enable_overlap: + scheduler.event_loop_overlap() + else: + scheduler.event_loop_normal() + elif disaggregation_mode == DisaggregationMode.PREFILL: + if server_args.pp_size > 1: + scheduler.event_loop_pp_disagg_prefill() + elif scheduler.enable_overlap: + scheduler.event_loop_overlap_disagg_prefill() + else: + scheduler.event_loop_normal_disagg_prefill() + elif disaggregation_mode == DisaggregationMode.DECODE: + if server_args.pp_size > 1: + scheduler.event_loop_pp_disagg_decode() + elif scheduler.enable_overlap: + scheduler.event_loop_overlap_disagg_decode() + else: + scheduler.event_loop_normal_disagg_decode() + + def run_scheduler_process( server_args: ServerArgs, port_args: PortArgs, @@ -3174,32 +3200,11 @@ def run_scheduler_process( pipe_writer.send(result_dict) - # Dispatch to the appropriate event loop based on the disaggregation mode - disaggregation_mode: DisaggregationMode = scheduler.disaggregation_mode - if disaggregation_mode == DisaggregationMode.NULL: - if scheduler.enable_pdmux: - scheduler.event_loop_pdmux() - elif server_args.pp_size > 1: - scheduler.event_loop_pp() - elif scheduler.enable_overlap: - scheduler.event_loop_overlap() - else: - scheduler.event_loop_normal() - elif disaggregation_mode == DisaggregationMode.PREFILL: - if server_args.pp_size > 1: - scheduler.event_loop_pp_disagg_prefill() - elif scheduler.enable_overlap: - scheduler.event_loop_overlap_disagg_prefill() - else: - scheduler.event_loop_normal_disagg_prefill() - - elif disaggregation_mode == DisaggregationMode.DECODE: - if server_args.pp_size > 1: - scheduler.event_loop_pp_disagg_decode() - elif scheduler.enable_overlap: - scheduler.event_loop_overlap_disagg_decode() - else: - scheduler.event_loop_normal_disagg_decode() + scheduler.schedule_stream = scheduler.device_module.Stream(priority=0) + if scheduler.device == "cpu": + scheduler.schedule_stream.synchronize = lambda: None # No-op for CPU + with CudaStreamContext(scheduler.schedule_stream): + dispatch_event_loop(scheduler) except Exception: traceback = get_exception_traceback() diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 90b13576b..c39a8fc13 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -1039,7 +1039,7 @@ class SchedulerPPMixin: ) if not mbs[next_mb_id].forward_mode.is_prebuilt(): with self.copy_stream_ctx: - self.copy_stream.wait_stream(self.default_stream) + self.copy_stream.wait_stream(self.schedule_stream) batch_result = self._pp_prep_batch_result( mbs[next_mb_id], mb_metadata[next_mb_id], next_pp_outputs ) @@ -1057,7 +1057,7 @@ class SchedulerPPMixin: ): with torch.profiler.record_function("run_batch"): with self.forward_stream_ctx: - self.forward_stream.wait_stream(self.default_stream) + self.forward_stream.wait_stream(self.schedule_stream) result = self.run_batch(self.cur_batch, pp_proxy_tensors) mb_metadata[mb_id] = PPBatchMetadata( can_run_cuda_graph=result.can_run_cuda_graph,