diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index eca270231..59359d5f7 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1107,9 +1107,15 @@ class Scheduler( self.init_all_cuda_graphs() model_runner = self.tp_worker.model_runner - with torch.get_device_module(model_runner.device).stream( + device_module = torch.get_device_module(model_runner.device) + self.schedule_stream = None if use_mlx() else device_module.Stream(priority=0) + # Match run_batch / _pp_launch_batch so warmup allocations stay reusable. + forward_stream = ( model_runner.forward_stream - ): + if self.enable_overlap or self.ps.pp_size > 1 or use_mlx() + else self.schedule_stream + ) + with device_module.stream(forward_stream): if self.draft_worker is None: model_runner.prewarm_sampling() else: @@ -1847,7 +1853,6 @@ class Scheduler( def run_event_loop(self) -> None: """Run the scheduler's event loop. - Sets up the schedule stream and dispatches to the appropriate event loop. The event loop blocks until shutdown. """ # Engine init (graph capture, warmups) is done; from here on any @@ -1862,10 +1867,9 @@ class Scheduler( dispatch_event_loop(self) return - self.schedule_stream = self.device_module.Stream(priority=0) if self.device == "cpu": self.schedule_stream.synchronize = lambda: None # No-op for CPU - elif is_cuda() or _is_hip: + elif (is_cuda() or _is_hip) and (self.enable_overlap or self.ps.pp_size > 1): # CUDA/HIP streams come from a fixed round-robin pool. Redraw if this # stream aliases forward_stream, which would eliminate scheduler # overlap. Only CUDA/HIP streams expose a ``cuda_stream`` handle; diff --git a/python/sglang/srt/model_executor/runner_utils/pool.py b/python/sglang/srt/model_executor/runner_utils/pool.py index 1388c5d1d..5acce066f 100644 --- a/python/sglang/srt/model_executor/runner_utils/pool.py +++ b/python/sglang/srt/model_executor/runner_utils/pool.py @@ -45,6 +45,7 @@ class GraphPoolBorrowState: active_user: Optional[str] = None stub: Optional[BumpArenaStub] = None mem_pool: Optional[torch.cuda.MemPool] = None + stream: Optional[torch.cuda.Stream] = None disabled_reason: Optional[str] = None static_runs: Optional[list[tuple[int, int]]] = None check_pending: bool = False @@ -279,6 +280,7 @@ def _teardown_borrow_pool() -> None: torch.cuda.synchronize() torch.empty(1, device="cuda") state.mem_pool = None + state.stream = None _PRECARVE_MIN_RUN_BYTES = 64 << 20 @@ -312,6 +314,8 @@ def _precarve_run_segments(runs: list[tuple[int, int]]) -> None: def borrow_graph_pool(user: str) -> Iterator[None]: """Route this thread's torch allocations onto the graph pool's free runs. + All borrows must use the stream that first creates the borrow pool, so + the caching allocator can reuse its pre-carved segments. Tensors allocated inside must be released before the next graph replay, which rewrites their bytes; the next replay (or pool teardown) raises if any are still referenced. An allocation no run can hold raises the @@ -322,7 +326,13 @@ def borrow_graph_pool(user: str) -> Iterator[None]: yield return with graph_pool_user_scope(user): + stream = torch.cuda.current_stream() if state.mem_pool is not None: + if stream != state.stream: + raise RuntimeError( + "Graph-pool borrow must use the stream that created the borrow pool: " + f"expected {state.stream}, got {stream}" + ) # Return completed cross-stream frees to the cache. The allocator # processes their events on a later allocation. torch.empty(1, device="cuda") @@ -346,6 +356,7 @@ def borrow_graph_pool(user: str) -> Iterator[None]: # stream-ordered deferred frees remain allocator-managed. Capture # retires it because capture changes the underlying free extents. state.mem_pool = torch.cuda.MemPool(state.stub.allocator) + state.stream = stream with torch.cuda.use_mem_pool(state.mem_pool): _precarve_run_segments(runs) # Only growth beyond the precarve is worth another log line. diff --git a/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py b/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py index 196ed3c55..9a59499b8 100644 --- a/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py +++ b/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py @@ -17,6 +17,7 @@ maybe_stub_sgl_kernel() from sglang.srt.configs.device_config import DeviceConfig from sglang.srt.configs.load_config import LoadConfig, LoadFormat from sglang.srt.configs.model_config import ModelImpl +from sglang.srt.managers.scheduler import Scheduler from sglang.srt.managers.tp_worker import TpModelWorker from sglang.srt.model_executor.cuda_graph_config import Backend, CudaGraphConfig from sglang.srt.model_executor.model_runner import ModelRunner @@ -728,9 +729,9 @@ class TestStartupWeightLoadSchedulerRouting(CustomTestCase): reset_context() self.addCleanup(reset_context) - def _scheduler(self, worker, trace, *, mode, draft_worker=None): - from sglang.srt.managers.scheduler import Scheduler - + def _scheduler( + self, worker, trace, *, mode, draft_worker=None, enable_overlap=True, pp_size=1 + ): # The schedule reads the mode from the bags, so the test states it by # publishing a record rather than by standing one in. reset_context() @@ -739,6 +740,8 @@ class TestStartupWeightLoadSchedulerRouting(CustomTestCase): role="scheduler", ) scheduler = Scheduler.__new__(Scheduler) + scheduler.enable_overlap = enable_overlap + scheduler.ps = SimpleNamespace(pp_size=pp_size) scheduler.init_tp_model_worker = lambda: setattr(scheduler, "tp_worker", worker) scheduler.maybe_init_draft_worker = lambda: setattr( scheduler, "draft_worker", draft_worker @@ -748,7 +751,9 @@ class TestStartupWeightLoadSchedulerRouting(CustomTestCase): scheduler.init_all_cuda_graphs = lambda: trace.append("capture") return scheduler - def _run_startup(self, mode, *, use_draft_worker=False): + def _run_startup( + self, mode, *, use_draft_worker=False, enable_overlap=True, pp_size=1, mlx=False + ): trace = [] worker = _SchedulerWorker(trace, post_capture_active=True) draft_worker = ( @@ -764,6 +769,14 @@ class TestStartupWeightLoadSchedulerRouting(CustomTestCase): trace, mode=mode, draft_worker=draft_worker, + enable_overlap=enable_overlap, + pp_size=pp_size, + ) + schedule_stream = object() + expected_stream = ( + worker.model_runner.forward_stream + if enable_overlap or pp_size > 1 or mlx + else schedule_stream ) class StreamContext: @@ -774,7 +787,7 @@ class TestStartupWeightLoadSchedulerRouting(CustomTestCase): trace.append("stream_exit") def stream_context(stream): - self.assertIs(stream, worker.model_runner.forward_stream) + self.assertIs(stream, expected_stream) return StreamContext() def stop_after_startup(): @@ -783,6 +796,7 @@ class TestStartupWeightLoadSchedulerRouting(CustomTestCase): scheduler.spec_algorithm = SimpleNamespace(is_none=stop_after_startup) with ( + patch("sglang.srt.managers.scheduler.use_mlx", return_value=mlx), patch( "sglang.srt.managers.scheduler.get_exec", return_value=SimpleNamespace( @@ -794,12 +808,15 @@ class TestStartupWeightLoadSchedulerRouting(CustomTestCase): ), patch( "sglang.srt.managers.scheduler.torch.get_device_module", - return_value=SimpleNamespace(stream=stream_context), + return_value=SimpleNamespace( + stream=stream_context, Stream=lambda priority: schedule_stream + ), ), self.assertRaisesRegex(RuntimeError, "stop after startup"), ): scheduler.init_model_worker() + self.assertIs(scheduler.schedule_stream, None if mlx else schedule_stream) worker.model_runner.post_capture_resize_kv_pool.assert_called_once_with( draft_runners=(worker.model_runner,) if use_draft_worker else () ) @@ -819,6 +836,18 @@ class TestStartupWeightLoadSchedulerRouting(CustomTestCase): ], ) + def test_sampling_warmup_uses_the_serving_stream(self): + for enable_overlap, pp_size, mlx in ( + (False, 1, False), + (False, 2, False), + (True, 1, False), + (False, 1, True), + ): + with self.subTest(enable_overlap=enable_overlap, pp_size=pp_size, mlx=mlx): + self._run_startup( + "serial", enable_overlap=enable_overlap, pp_size=pp_size, mlx=mlx + ) + def test_overlap_starts_before_capture_and_finalizes_after(self): self.assertEqual( self._run_startup("overlap"), diff --git a/test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py b/test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py index 0650e60b7..33f09db0d 100644 --- a/test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py +++ b/test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py @@ -114,6 +114,7 @@ class TestGraphPoolBorrow(CustomTestCase): patch.object(self.state, "mem_pool", None), patch.object(pool.torch.cuda, "MemPool"), patch.object(pool.torch.cuda, "use_mem_pool"), + patch.object(pool.torch.cuda, "current_stream"), patch.object(pool.torch.cuda, "memory_snapshot", return_value=snapshot), pool.borrow_graph_pool(user="test"), ): @@ -127,15 +128,18 @@ class TestGraphPoolBorrow(CustomTestCase): def test_high_cursor_keeps_reusable_cached_segments(self): stub = MagicMock(cursor_bytes=600, freed_bytes=0) mem_pool = MagicMock() + stream = object() with ( patch.object(pool, "graph_pool_borrow_enabled", return_value=True), patch.object(self.state, "stub", stub), patch.object(self.state, "mem_pool", mem_pool), + patch.object(self.state, "stream", stream), patch.object(self.state, "extents_total", 1000), patch.object(pool, "_teardown_borrow_pool") as teardown, patch.object(pool.torch, "empty"), patch.object(pool.torch.cuda, "use_mem_pool"), + patch.object(pool.torch.cuda, "current_stream", return_value=stream), ): with pool.borrow_graph_pool(user="test"): pass @@ -282,9 +286,35 @@ class TestGraphPoolBorrow(CustomTestCase): self.assertEqual(c.data_ptr(), recycled_address) del b, c + self.assertEqual(self.state.stream, torch.cuda.current_stream()) + with ( + torch.cuda.stream(stream), + self.assertRaisesRegex( + RuntimeError, "stream that created the borrow pool" + ), + ): + with pool.borrow_graph_pool(user="wrong stream"): + self.fail("cross-stream borrowing must fail before allocating") + self.assertIsNone(self.state.active_user) + with pool.borrow_graph_pool(user="same stream"): + reused = torch.empty(40 << 20, dtype=torch.uint8, device="cuda") + self.assertEqual(reused.data_ptr(), recycled_address) + del reused + # Captures retire the persistent borrow pool. Its storage aliases # existing graph-pool runs, so the reserved footprint is unchanged. pool._teardown_borrow_pool() + self.assertIsNone(self.state.stream) + with torch.cuda.stream(stream), pool.borrow_graph_pool(user="new pool"): + self.assertEqual(self.state.stream, stream) + reused = torch.empty(40 << 20, dtype=torch.uint8, device="cuda") + self.assertTrue(on_a_run(reused)) + del reused + pool._teardown_borrow_pool() + with pool.graph_pool_replay_scope(): + graph.replay() + torch.cuda.synchronize() + self.assertTrue(torch.equal(y, torch.ones_like(y))) self.assertEqual(torch.cuda.memory_reserved(device_id), reserved_before) del graph, y