chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Apply default stream to priority 0 in scheduling. (#16438)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-03-03 00:05:27 -08:00
committed by GitHub
parent 07b8d763ef
commit 7a2d3df96f
2 changed files with 39 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

69
python/sglang/srt/managers/scheduler.py
View File

@@ -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()

4
python/sglang/srt/managers/scheduler_pp_mixin.py
View File

@@ -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,