Keep graph-pool borrows on their allocation stream (#39180)

Co-authored-by: cctry <17473714+cctry@users.noreply.github.com>

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;

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.