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)