From 18cc55dc0b2f6a5133bdde61f0ac8c02c8e3d17d Mon Sep 17 00:00:00 2001 From: cctry Date: Sat, 12 Sep 2026 19:24:57 -0700 Subject: [PATCH] Expose a capacity check for graph-pool borrows (#39178) Co-authored-by: cctry <17473714+cctry@users.noreply.github.com> --- .../srt/model_executor/runner_utils/pool.py | 10 ++++++++++ .../runner_utils/test_graph_pool_borrow.py | 19 +++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/python/sglang/srt/model_executor/runner_utils/pool.py b/python/sglang/srt/model_executor/runner_utils/pool.py index 2b6ed2949..1388c5d1d 100644 --- a/python/sglang/srt/model_executor/runner_utils/pool.py +++ b/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: 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 3f39e3fb2..0650e60b7 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 @@ -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)