[Memory] Size the CUDA graph pool from warmup measurements and fix graph-pool borrowing (#36911)

Co-authored-by: cctry <cctry@fb.com>

python/sglang/srt/model_executor/model_runner.py

@@ -253,6 +253,14 @@ elif current_platform.is_out_of_tree():
 logger = logging.getLogger(__name__)
 
 
+@dataclass(frozen=True)
+class SamplingPrewarmResult:
+    """Memory requirements observed while pre-warming a sampling path."""
+
+    sampling_input_bytes: int = 0
+    sampling_headroom_bytes: int = 0
+
+
 def _prefill_cuda_graph_allows_context_parallel(
     prefill_runner, forward_batch: ForwardBatch
 ) -> bool:
@@ -377,6 +385,7 @@ class ModelRunner:
         self.draft_model_idx = draft_model_idx
         self.enable_hisparse = get_memory().enable_hisparse
         self._sampling_observer: Optional[SamplingObserver] = None
+        self.sampling_prewarm_result = SamplingPrewarmResult()
 
         self.init_startup_observability()
 
@@ -1058,6 +1067,11 @@ class ModelRunner:
             n_prepared,
         )
 
+    def prewarm_sampling(self) -> SamplingPrewarmResult:
+        """Warm the sampling path after graph initialization."""
+        self.sampling_prewarm_result = SamplingPrewarmResult()
+        return self.sampling_prewarm_result
+
     def init_cuda_graphs(self, capture_decode_cuda_graph: bool = True):
         capture = capture_cuda_graphs(
             model_runner=self, capture_decode_cuda_graph=capture_decode_cuda_graph

python/sglang/srt/model_executor/model_runner_components/kv_pool_runtime.py

@@ -11,6 +11,7 @@ from sglang.srt.configs.hybrid_arch import mambaish_config
 from sglang.srt.distributed import get_world_group
 from sglang.srt.mem_cache.kv_cache_configurator import mm_runtime_reservation_gb
 from sglang.srt.model_executor.cuda_graph_config import Backend
+from sglang.srt.model_executor.runner_utils.pool import graph_pool_borrow_enabled
 from sglang.srt.platforms import current_platform
 from sglang.srt.runtime_context import (
     get_disagg,
@@ -85,6 +86,13 @@ def compute_post_capture_kv_resize(
             )
             / 1024,
         )
+    if not graph_pool_borrow_enabled():
+        # Borrowing serves the sampling temporaries out of idle graph storage;
+        # without it they need real headroom the KV pool must not claim.
+        headroom_gb = max(
+            headroom_gb,
+            model_runner.sampling_prewarm_result.sampling_headroom_bytes / (1 << 30),
+        )
     mm_reservation_gb = mm_runtime_reservation_gb(
         is_multimodal=model_runner.model_config.is_multimodal,
         mm_feature_transport=get_mm().mm_feature_transport,

python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py

@@ -40,6 +40,7 @@ from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import
     enable_breakable_cuda_graph,
 )
 from sglang.srt.model_executor.runner_utils.pool import (
+    GraphPoolPrecarve,
     get_or_create_global_graph_memory_pool,
     graph_pool_capture_scope,
     graph_pool_replay_scope,
@@ -77,6 +78,7 @@ class BreakableCudaGraphBackend(DedupedCudaGraphMixin, BaseCudaGraphBackend):
         self._capture_stream: Optional[torch.cuda.Stream] = None
         self._debug_eager = debug_eager
         self._shared_output_buffer: Optional[Any] = None
+        self._precarve = GraphPoolPrecarve()
         self._memory_saver_adapter: Optional[Any] = TorchMemorySaverAdapter.create(
             enable=enable_memory_saver
             and get_bool_env_var("SGLANG_MEMORY_SAVER_CUDA_GRAPH")
@@ -117,7 +119,8 @@ class BreakableCudaGraphBackend(DedupedCudaGraphMixin, BaseCudaGraphBackend):
         for _ in range(2):
             self._device_module.synchronize()
             self._tp_group.barrier()
-            warmup_out = forward_fn()
+            with self._precarve.measure():
+                warmup_out = forward_fn()
             if post_warmup_hook is not None:
                 post_warmup_hook()
 
@@ -134,6 +137,7 @@ class BreakableCudaGraphBackend(DedupedCudaGraphMixin, BaseCudaGraphBackend):
             stream=self._capture_stream,
             barrier_fn=self._tp_group.barrier,
         ):
+            self._precarve.mint()
             out = captured_fn()
             out_rows = self._output_rows(out, size)
             self._copy_output_to_buffer(out, self._shared_output_buffer, out_rows)

python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py

@@ -31,6 +31,7 @@ from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
     BaseCudaGraphBackend,
 )
 from sglang.srt.model_executor.runner_utils.pool import (
+    GraphPoolPrecarve,
     get_or_create_global_graph_memory_pool,
     graph_pool_capture_scope,
     graph_pool_replay_scope,
@@ -64,6 +65,7 @@ class FullCudaGraphBackend(BaseCudaGraphBackend):
         self._device_module = cuda_graph_runner.device_module
         self._tp_group = cuda_graph_runner.model_runner.tp_group
         self._capture_stream: Optional[torch.cuda.Stream] = None
+        self._precarve = GraphPoolPrecarve()
         self._memory_saver_adapter: Optional[Any] = TorchMemorySaverAdapter.create(
             enable=enable_memory_saver
             and get_bool_env_var("SGLANG_MEMORY_SAVER_CUDA_GRAPH")
@@ -107,7 +109,8 @@ class FullCudaGraphBackend(BaseCudaGraphBackend):
         for _ in range(2):
             self._device_module.synchronize()
             self._tp_group.barrier()
-            forward_fn()
+            with self._precarve.measure():
+                forward_fn()
             if profiler is not None:
                 profiler.step()
             if post_warmup_hook is not None:
@@ -131,6 +134,7 @@ class FullCudaGraphBackend(BaseCudaGraphBackend):
             graph_pool_capture_scope(),
             graph_ctx(cuda_graph=graph, pool=self._pool, stream=self._capture_stream),
         ):
+            self._precarve.mint()
             out = forward_fn()
 
         if profiler is not None:

python/sglang/srt/model_executor/runner_utils/pool.py

@@ -46,6 +46,7 @@ def disable_graph_pool_borrow(reason: str) -> None:
     """Disable borrowing when graph storage is managed outside the shared pool."""
     global _borrow_disabled_reason
     _borrow_disabled_reason = reason
+    _teardown_borrow_pool()
     logger.info("Graph pool borrow disabled: %s", reason)
 
 
@@ -93,6 +94,43 @@ def get_or_create_global_graph_capture_stream() -> Any:
     return get_stream(_CAPTURE_STREAM_NAME)
 
 
+class GraphPoolPrecarve:
+    """Pre-carve the memory pool to reduce fragmentation."""
+
+    def __init__(self) -> None:
+        self.nbytes = 0
+        self.minted = False
+
+    @contextmanager
+    def measure(self) -> Iterator[None]:
+        """Wrap one eager warmup. the last one before ``mint`` sets the size."""
+        if self.minted or not envs.SGLANG_ENABLE_GRAPH_POOL_PRECARVE.get():
+            yield
+            return
+        torch.cuda.synchronize()
+        # Shrink the cache first so the warmup's reserved growth is its own
+        # footprint. Reserved (not allocated) is the stat to use: allocated
+        # peak is the live-byte sum and undershoots by exactly the packing
+        # holes the carved span has to absorb.
+        torch.cuda.empty_cache()
+        torch.cuda.reset_peak_memory_stats()
+        base = torch.cuda.memory_stats()["reserved_bytes.all.current"]
+        yield
+        torch.cuda.synchronize()
+        self.nbytes = torch.cuda.memory_stats()["reserved_bytes.all.peak"] - base
+
+    def mint(self) -> None:
+        """Pre-allocate the space"""
+        if self.minted:
+            return
+        self.minted = True
+        if self.nbytes <= 0:
+            return
+        span = torch.empty(self.nbytes, dtype=torch.uint8, device="cuda")
+        del span
+        logger.info("Graph pool pre-carved: %.2f GB", self.nbytes / 2**30)
+
+
 def graph_pool_borrow_enabled() -> bool:
     if (
         _borrow_disabled_reason is not None
@@ -143,11 +181,11 @@ def graph_pool_capture_scope() -> Iterator[None]:
 
 
 def find_free_graph_pool_runs(pool_id: Any) -> list[tuple[int, int]]:
-    """Return contiguous inactive runs inside wholly-free segments, largest first."""
+    """Return contiguous inactive runs across all pool segments, largest first;
+    live and pending-free blocks break runs, so a run never overlaps live data.
+    """
     runs: list[tuple[int, int]] = []
     for segment in torch.cuda.memory_snapshot(pool_id, include_traces=False):
-        if segment["allocated_size"] != 0:
-            continue
         run_address = 0
         run_bytes = 0
         for block in segment["blocks"]:
@@ -195,16 +233,18 @@ def borrow_graph_pool(user: str) -> Iterator[None]:
             # Return completed cross-stream frees to the cache. The allocator
             # processes their events on a later allocation.
             torch.empty(1, device="cuda")
-            if (
-                _borrow_stub.freed_bytes
-                or _borrow_stub.cursor_bytes > _borrow_extents_total // 2
-            ):
-                # Rebuild if empty_cache released the arena's segments or if
-                # unresolved deferred frees consumed half of the extents.
+            if _borrow_stub.freed_bytes:
+                # Rebuild if the caching allocator released an arena segment.
+                # A high cursor alone means the cache owns reusable segments;
+                # rebuilding discards them and can turn the next large borrow
+                # into a fragmented cold-allocation OOM.
                 _teardown_borrow_pool()
         if _borrow_mem_pool is None:
             if _borrow_stub is None:
                 _borrow_stub = BumpArenaStub()
+                # Runs are sorted largest first, so first fit would carve every
+                # small allocation out of the run a probability matrix needs.
+                _borrow_stub.set_best_fit(True)
             if _borrow_static_runs is not None:
                 runs = _borrow_static_runs
             else: