Expose a capacity check for graph-pool borrows (#39178)

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

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

@@ -288,6 +288,16 @@ _PRECARVE_SMALL_RESERVE_BYTES = 32 << 20
 _PRECARVE_GRANULARITY = 2 << 20
 
 
+def graph_pool_borrow_can_fit(nbytes: int) -> bool:
+    """Whether one free run fits the payload plus allocator padding and reserve."""
+    if nbytes <= 0:
+        return False
+    required = (
+        nbytes + _PRECARVE_GRANULARITY - 1
+    ) // _PRECARVE_GRANULARITY * _PRECARVE_GRANULARITY + _PRECARVE_SMALL_RESERVE_BYTES
+    return graph_pool_borrow_largest_run() >= required
+
+
 def _precarve_run_segments(runs: list[tuple[int, int]]) -> None:
     """Seed coalescible segments on the stream that will allocate borrows."""
     for _, run_bytes in runs:

test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py

@@ -76,6 +76,25 @@ class TestGraphPoolBorrow(CustomTestCase):
             self.assertEqual(pool.graph_pool_borrow_largest_run(), 0)
             self.assertEqual(snapshot.call_count, 3)
 
+    def test_borrow_capacity_includes_rounding_and_small_pool_reserve(self):
+        for payload, largest_run, expected in (
+            (-1, 64 << 20, False),
+            (0, 64 << 20, False),
+            (1, 0, False),
+            (1, 32 << 20, False),
+            (1, 34 << 20, True),
+            (2 << 20, 34 << 20, True),
+            ((2 << 20) + 1, 34 << 20, False),
+            ((2 << 20) + 1, 36 << 20, True),
+        ):
+            with (
+                self.subTest(payload=payload, largest_run=largest_run),
+                patch.object(
+                    pool, "graph_pool_borrow_largest_run", return_value=largest_run
+                ),
+            ):
+                self.assertEqual(pool.graph_pool_borrow_can_fit(payload), expected)
+
     def test_graph_replay_fails_during_active_pool_borrow(self):
         graph = Mock()
         backend = object.__new__(FullCudaGraphBackend)