minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

Scope graph-pool borrowing to the runtime and reduce fragmentation (#39177)

Browse Source 
Co-authored-by: cctry <17473714+cctry@users.noreply.github.com>
This commit is contained in:
cctry
2026-09-12 11:45:17 -07:00
committed by GitHub
co-authored by cctry
parent 2784a86062
commit 7ae4af8187
3 changed files with 270 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/runner_utils/pool.py
+143 -55
View File
@@ -21,6 +21,8 @@ from __future__ import annotations
import logging
from contextlib import contextmanager
from dataclasses import dataclass, field
from functools import cache
from typing import Any, Iterator, Optional
import torch
@@ -31,21 +33,50 @@ from sglang.srt.utils import is_cuda
from sglang.srt.utils.cuda_vmm_utils import BumpArenaStub
logger = logging.getLogger(__name__)
_active_graph_pool_user: Optional[str] = None
_borrow_stub: Optional[BumpArenaStub] = None
_borrow_mem_pool: Optional[torch.cuda.MemPool] = None
_borrow_disabled_reason: Optional[str] = None
_borrow_static_runs: Optional[list[tuple[int, int]]] = None
_borrow_extents_total = 0
_largest_logged_graph_pool_borrow = 0
_MIB = 1 << 20
_CAPTURE_STREAM_NAME = "cuda_graph_capture"
@dataclass(eq=False)
class GraphPoolBorrowState:
"""Mutable borrowing state and free-run cache for one runtime."""
active_user: Optional[str] = None
stub: Optional[BumpArenaStub] = None
mem_pool: Optional[torch.cuda.MemPool] = None
disabled_reason: Optional[str] = None
static_runs: Optional[list[tuple[int, int]]] = None
check_pending: bool = False
extents_total: int = 0
largest_logged_borrow: int = 0
snapshot_free_runs: Any = field(
default_factory=lambda: cache(find_free_graph_pool_runs), init=False, repr=False
)
def _get_graph_pool_borrow_state() -> GraphPoolBorrowState:
resources = get_resources()
if resources.graph_pool_borrow is None:
resources.graph_pool_borrow = GraphPoolBorrowState()
return resources.graph_pool_borrow
def _log_graph_pool_borrow_capacity(runs: list[tuple[int, int]]) -> None:
total_bytes = sum(nbytes for _, nbytes in runs)
largest_bytes = max((nbytes for _, nbytes in runs), default=0)
logger.info(
"Graph pool borrow capacity: %.1f MiB available, largest contiguous "
"region %.1f MiB, %d regions",
total_bytes / _MIB,
largest_bytes / _MIB,
len(runs),
)
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
_get_graph_pool_borrow_state().disabled_reason = reason
_teardown_borrow_pool()
logger.info("Graph pool borrow disabled: %s", reason)
@@ -57,17 +88,13 @@ def set_graph_pool_borrow_runs(runs: list[tuple[int, int]]) -> None:
remain stable for the process lifetime. Registering an empty list disables
borrowing.
"""
global _borrow_static_runs
state = _get_graph_pool_borrow_state()
static_runs = sorted(runs, key=lambda run: run[1], reverse=True)[
: BumpArenaStub.MAX_EXTENTS
]
_teardown_borrow_pool()
_borrow_static_runs = static_runs
logger.info(
"Graph pool borrow runs pinned: runs=%d free=%d",
len(_borrow_static_runs),
sum(nbytes for _, nbytes in _borrow_static_runs),
)
state.static_runs = static_runs
_log_graph_pool_borrow_capacity(state.static_runs)
def get_global_graph_memory_pool() -> Optional[Any]:
@@ -134,31 +161,32 @@ class GraphPoolPrecarve:
def graph_pool_borrow_enabled() -> bool:
state = _get_graph_pool_borrow_state()
if (
_borrow_disabled_reason is not None
state.disabled_reason is not None
or not envs.SGLANG_ENABLE_GRAPH_POOL_BORROW.get()
or not is_cuda()
):
return False
if _borrow_static_runs is not None:
return len(_borrow_static_runs) > 0
if state.static_runs is not None:
return len(state.static_runs) > 0
return get_global_graph_memory_pool() is not None
@contextmanager
def graph_pool_user_scope(user: str) -> Iterator[None]:
global _active_graph_pool_user
state = _get_graph_pool_borrow_state()
# Graph replay silently overwrites aliases of its allocator-free blocks.
if _active_graph_pool_user is not None:
if state.active_user is not None:
raise RuntimeError(
f"graph pool already has live user {_active_graph_pool_user!r}; "
f"graph pool already has live user {state.active_user!r}; "
f"cannot use it for {user!r}"
)
_active_graph_pool_user = user
state.active_user = user
try:
yield
finally:
_active_graph_pool_user = None
state.active_user = None
@contextmanager
@@ -167,6 +195,8 @@ def graph_pool_replay_scope() -> Iterator[None]:
yield
return
with graph_pool_user_scope("CUDA graph"):
if _get_graph_pool_borrow_state().check_pending:
_raise_on_live_borrows("graph replay")
yield
@@ -205,68 +235,126 @@ def find_free_graph_pool_runs(pool_id: Any) -> list[tuple[int, int]]:
return runs
def graph_pool_borrow_largest_run() -> int:
"""Largest contiguous free extent a single borrow can occupy, in bytes."""
if not graph_pool_borrow_enabled():
return 0
state = _get_graph_pool_borrow_state()
if state.static_runs is not None:
return state.static_runs[0][1]
runs = state.snapshot_free_runs(get_global_graph_memory_pool())
return runs[0][1] if runs else 0
def _raise_on_live_borrows(event: str) -> None:
"""Reject borrows that graph replay or pool teardown would overwrite."""
state = _get_graph_pool_borrow_state()
state.check_pending = False
if state.mem_pool is None:
return
live = sum(
block["size"]
for segment in torch.cuda.memory_snapshot(
state.mem_pool.id, include_traces=False
)
for block in segment["blocks"]
if block["state"] == "active_allocated"
)
if live:
raise RuntimeError(
f"Graph-pool borrow leak at {event}: {live} bytes still referenced"
)
def _teardown_borrow_pool() -> None:
"""Retire the persistent borrow pool after draining deferred frees."""
global _borrow_mem_pool
if _borrow_mem_pool is None:
state = _get_graph_pool_borrow_state()
state.snapshot_free_runs.cache_clear()
if state.mem_pool is None:
return
if state.check_pending:
_raise_on_live_borrows("borrow pool teardown")
# Borrowed blocks that saw cross-stream use can remain in event limbo.
# Synchronize, then drive allocator event processing before dropping the pool.
torch.cuda.synchronize()
torch.empty(1, device="cuda")
_borrow_mem_pool = None
state.mem_pool = None
_PRECARVE_MIN_RUN_BYTES = 64 << 20
# Left uncarved per run so 2 MiB small-pool segments keep a home.
_PRECARVE_SMALL_RESERVE_BYTES = 32 << 20
# The caching allocator rounds large segment requests up to 2 MiB.
_PRECARVE_GRANULARITY = 2 << 20
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:
seed = (
(run_bytes - _PRECARVE_SMALL_RESERVE_BYTES) // _PRECARVE_GRANULARITY
) * _PRECARVE_GRANULARITY
if seed >= _PRECARVE_MIN_RUN_BYTES:
torch.empty(seed, dtype=torch.uint8, device="cuda")
@contextmanager
def borrow_graph_pool(user: str) -> Iterator[None]:
"""Route this thread's torch allocations onto the graph pool's free runs.
Tensors allocated inside must not survive the current step: the next graph
replay may overwrite their bytes. An allocation no run can hold raises the
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
allocator's normal OOM. This is a no-op while borrowing is disabled.
"""
global _borrow_stub, _borrow_mem_pool, _borrow_extents_total
global _largest_logged_graph_pool_borrow
state = _get_graph_pool_borrow_state()
if not graph_pool_borrow_enabled():
yield
return
with graph_pool_user_scope(user):
if _borrow_mem_pool is not None:
if state.mem_pool is not 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:
# 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.
if state.stub.freed_bytes:
# The bump arena cannot reuse segments returned by empty_cache().
_teardown_borrow_pool()
if _borrow_mem_pool is None:
if _borrow_stub is None:
_borrow_stub = BumpArenaStub()
if state.mem_pool is None:
if state.stub is None:
state.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
state.stub.set_best_fit(True)
if state.static_runs is not None:
runs = state.static_runs
else:
runs = find_free_graph_pool_runs(get_global_graph_memory_pool())[
runs = state.snapshot_free_runs(get_global_graph_memory_pool())[
: BumpArenaStub.MAX_EXTENTS
]
_borrow_stub.set_extents(runs)
state.stub.set_extents(runs)
# Keep one caching layer across borrows so normal block reuse and
# stream-ordered deferred frees remain allocator-managed. Capture
# retires it because capture changes the underlying free extents.
_borrow_mem_pool = torch.cuda.MemPool(_borrow_stub.allocator)
_borrow_extents_total = sum(run_bytes for _, run_bytes in runs)
state.mem_pool = torch.cuda.MemPool(state.stub.allocator)
with torch.cuda.use_mem_pool(state.mem_pool):
_precarve_run_segments(runs)
# Only growth beyond the precarve is worth another log line.
state.largest_logged_borrow = state.stub.cursor_bytes
state.extents_total = sum(run_bytes for _, run_bytes in runs)
_log_graph_pool_borrow_capacity(runs)
logger.info(
"Graph pool borrow extents: runs=%d free=%d",
len(runs),
_borrow_extents_total,
"Graph pool borrow pre-carved: %.1f MiB",
state.stub.cursor_bytes / _MIB,
)
with torch.cuda.use_mem_pool(_borrow_mem_pool):
with torch.cuda.use_mem_pool(state.mem_pool):
yield
consumed = _borrow_stub.cursor_bytes
if consumed > _largest_logged_graph_pool_borrow:
logger.info("Graph pool borrow: consumed=%d", consumed)
_largest_logged_graph_pool_borrow = consumed
state.check_pending = True
consumed = state.stub.cursor_bytes
if consumed > state.largest_logged_borrow:
logger.info(
"Graph pool borrow high-water mark: %.1f MiB used of %.1f MiB "
"available",
consumed / _MIB,
state.extents_total / _MIB,
)
state.largest_logged_borrow = consumed
python/sglang/srt/runtime_context.py
+2
View File
@@ -58,6 +58,7 @@ from typing import TYPE_CHECKING, Any, Dict, Optional
import msgspec
if TYPE_CHECKING:
from sglang.srt.model_executor.runner_utils.pool import GraphPoolBorrowState
from sglang.srt.server_args import ServerArgs
logger = logging.getLogger(__name__)
@@ -607,6 +608,7 @@ class Resources(_FlagGroupBase):
# CUDA graph memory pool shared across the prefill and decode graph
# backends (created lazily by model_executor.runner_utils.pool).
graph_memory_pool: Any = None
graph_pool_borrow: GraphPoolBorrowState | None = None
# EPLB: per-process recorder and the publish-once location metadata
# (owning accessors live in sglang.srt.eplb).
expert_distribution_recorder: Any = None
test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py
+125 -34
View File
@@ -1,5 +1,6 @@
"""CUDA graph-pool borrowing allocator and lifetime regression tests."""
import contextlib
import unittest
from contextlib import contextmanager
from types import SimpleNamespace
@@ -17,32 +18,23 @@ from sglang.srt.speculative.dflash_worker_v2 import DFlashWorkerV2
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=13, stage="base-b", runner_config="1-gpu-small")
register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-small")
class TestGraphPoolBorrow(CustomTestCase):
def setUp(self):
super().setUp()
self._reset_borrow_state()
self.state = pool.GraphPoolBorrowState()
# ExitStack rather than enterContext, which is 3.11+.
stack = contextlib.ExitStack()
self.addCleanup(stack.close)
stack.enter_context(pool.get_resources().override(graph_pool_borrow=self.state))
def tearDown(self):
try:
if torch.cuda.is_available():
pool._teardown_borrow_pool()
torch.cuda.synchronize()
torch.cuda.empty_cache()
finally:
self._reset_borrow_state()
@staticmethod
def _reset_borrow_state():
pool._active_graph_pool_user = None
pool._borrow_stub = None
pool._borrow_mem_pool = None
pool._borrow_disabled_reason = None
pool._borrow_static_runs = None
pool._borrow_extents_total = 0
pool._largest_logged_graph_pool_borrow = 0
if torch.cuda.is_available():
pool._teardown_borrow_pool()
torch.cuda.synchronize()
torch.cuda.empty_cache()
def test_mixed_segment_runs_exclude_live_blocks(self):
"""A mixed segment's free runs are borrowable, but a returned run must
@@ -63,6 +55,27 @@ class TestGraphPoolBorrow(CustomTestCase):
runs = pool.find_free_graph_pool_runs((0, 1))
self.assertEqual(sorted(runs), [(0x1000, 4096), (0x3000, 4096)])
def test_free_run_snapshot_lifetime_follows_borrow_state(self):
runs = [{"blocks": [{"state": "inactive", "address": 0x1000, "size": 8192}]}]
with (
envs.SGLANG_ENABLE_GRAPH_POOL_BORROW.override(True),
patch.object(pool, "is_cuda", return_value=True),
patch.object(pool, "get_global_graph_memory_pool", return_value=(1, 2)),
patch.object(
pool.torch.cuda, "memory_snapshot", side_effect=[runs, [], runs]
) as snapshot,
):
self.assertEqual(pool.graph_pool_borrow_largest_run(), 8192)
self.assertEqual(pool.graph_pool_borrow_largest_run(), 8192)
snapshot.assert_called_once_with((1, 2), include_traces=False)
pool._teardown_borrow_pool()
self.assertEqual(pool.graph_pool_borrow_largest_run(), 0)
self.assertEqual(snapshot.call_count, 2)
with pool.get_resources().override(graph_pool_borrow=None):
self.assertEqual(pool.graph_pool_borrow_largest_run(), 8192)
self.assertEqual(pool.graph_pool_borrow_largest_run(), 0)
self.assertEqual(snapshot.call_count, 3)
def test_graph_replay_fails_during_active_pool_borrow(self):
graph = Mock()
backend = object.__new__(FullCudaGraphBackend)
@@ -78,10 +91,8 @@ class TestGraphPoolBorrow(CustomTestCase):
with (
envs.SGLANG_ENABLE_GRAPH_POOL_BORROW.override(True),
patch.object(pool, "get_global_graph_memory_pool", return_value=(1, 2)),
patch.object(
pool, "_borrow_stub", MagicMock(cursor_bytes=0, freed_bytes=0)
),
patch.object(pool, "_borrow_mem_pool", None),
patch.object(self.state, "stub", MagicMock(cursor_bytes=0, freed_bytes=0)),
patch.object(self.state, "mem_pool", None),
patch.object(pool.torch.cuda, "MemPool"),
patch.object(pool.torch.cuda, "use_mem_pool"),
patch.object(pool.torch.cuda, "memory_snapshot", return_value=snapshot),
@@ -100,9 +111,9 @@ class TestGraphPoolBorrow(CustomTestCase):
with (
patch.object(pool, "graph_pool_borrow_enabled", return_value=True),
patch.object(pool, "_borrow_stub", stub),
patch.object(pool, "_borrow_mem_pool", mem_pool),
patch.object(pool, "_borrow_extents_total", 1000),
patch.object(self.state, "stub", stub),
patch.object(self.state, "mem_pool", mem_pool),
patch.object(self.state, "extents_total", 1000),
patch.object(pool, "_teardown_borrow_pool") as teardown,
patch.object(pool.torch, "empty"),
patch.object(pool.torch.cuda, "use_mem_pool"),
@@ -126,7 +137,7 @@ class TestGraphPoolBorrow(CustomTestCase):
runs = [(0x1000, 4096), (0x2000, 8192)]
def reset_static_runs():
pool._borrow_static_runs = None
self.state.static_runs = None
with patch.object(
pool, "_teardown_borrow_pool", side_effect=reset_static_runs
@@ -134,7 +145,7 @@ class TestGraphPoolBorrow(CustomTestCase):
pool.set_graph_pool_borrow_runs(runs)
teardown.assert_called_once_with()
self.assertEqual(pool._borrow_static_runs, [(0x2000, 8192), (0x1000, 4096)])
self.assertEqual(self.state.static_runs, [(0x2000, 8192), (0x1000, 4096)])
def test_eagle_non_greedy_probabilities_do_not_borrow_graph_pool(self):
def fake_sampling(**kwargs):
@@ -221,7 +232,6 @@ class TestGraphPoolBorrow(CustomTestCase):
with (
envs.SGLANG_ENABLE_GRAPH_POOL_BORROW.override(True),
patch.object(pool, "get_global_graph_memory_pool", return_value=handle),
patch.object(pool, "_borrow_mem_pool", None),
):
runs = pool.find_free_graph_pool_runs(handle)
self.assertGreaterEqual(len(runs), 2)
@@ -260,6 +270,86 @@ class TestGraphPoolBorrow(CustomTestCase):
self.assertEqual(torch.cuda.memory_reserved(device_id), reserved_before)
del graph, y
@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
def test_borrow_recovers_from_arena_fragmentation(self):
handle = torch.cuda.graph_pool_handle()
graph = torch.cuda.CUDAGraph()
seed = torch.zeros(8, device="cuda")
stream = torch.cuda.Stream()
with (
torch.cuda.stream(stream),
torch.cuda.graph(graph, pool=handle, stream=stream),
):
transient = torch.empty(200 << 20, dtype=torch.uint8, device="cuda")
keep = seed + 1
del transient
torch.cuda.synchronize()
address, run_bytes = pool.find_free_graph_pool_runs(handle)[0]
self.assertEqual(run_bytes, 200 << 20)
with (
envs.SGLANG_ENABLE_GRAPH_POOL_BORROW.override(True),
patch.object(pool, "get_global_graph_memory_pool", return_value=handle),
):
# Unseeded 24/32/40 MiB segments strand 192 MiB before the 42 MiB request.
for rows in (1500, 2000, 2500, 2600):
with self.subTest(rows=rows), pool.borrow_graph_pool(user="test"):
first = torch.empty(
(rows, 4096), dtype=torch.float32, device="cuda"
)
second = torch.empty(
(rows, 4096), dtype=torch.float32, device="cuda"
)
self.assertTrue(
all(
address <= tensor.data_ptr()
and tensor.data_ptr() + tensor.nbytes <= address + run_bytes
for tensor in (first, second)
)
)
del first, second
pool._teardown_borrow_pool()
del graph, keep
@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
def test_replay_raises_when_borrowed_tensor_is_still_referenced(self):
"""Reject borrowed tensors that replay would silently overwrite."""
handle = torch.cuda.graph_pool_handle()
graph = torch.cuda.CUDAGraph()
seed = torch.zeros(8, device="cuda")
stream = torch.cuda.Stream()
with (
torch.cuda.stream(stream),
torch.cuda.graph(graph, pool=handle, stream=stream),
):
transient = torch.empty(48 << 20, dtype=torch.uint8, device="cuda")
keep = seed + 1
del transient
torch.cuda.synchronize()
with (
envs.SGLANG_ENABLE_GRAPH_POOL_BORROW.override(True),
patch.object(pool, "get_global_graph_memory_pool", return_value=handle),
):
with pool.borrow_graph_pool(user="leaky"):
leaked = torch.empty(1 << 20, device="cuda")
with self.assertRaisesRegex(
RuntimeError,
f"graph replay: {leaked.nbytes} bytes",
):
with pool.graph_pool_replay_scope():
pass
# A fresh borrow re-arms the replay check after releasing the leak.
del leaked
with pool.borrow_graph_pool(user="clean"):
released = torch.empty(1 << 20, device="cuda")
del released
with pool.graph_pool_replay_scope():
pass
pool._teardown_borrow_pool()
del graph, keep
@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
def test_static_borrow_runs_serve_without_a_pool_snapshot(self):
"""Fixed extents serve borrows without consulting the shared pool."""
@@ -281,8 +371,6 @@ class TestGraphPoolBorrow(CustomTestCase):
with (
envs.SGLANG_ENABLE_GRAPH_POOL_BORROW.override(True),
patch.object(pool, "get_global_graph_memory_pool", return_value=None),
patch.object(pool, "_borrow_static_runs", None),
patch.object(pool, "_borrow_mem_pool", None),
):
pool.set_graph_pool_borrow_runs(runs)
self.assertTrue(pool.graph_pool_borrow_enabled())
@@ -295,6 +383,7 @@ class TestGraphPoolBorrow(CustomTestCase):
)
)
del borrowed
pool._teardown_borrow_pool()
del graph, y
@@ -352,18 +441,20 @@ class TestGraphPoolBorrow(CustomTestCase):
with (
envs.SGLANG_ENABLE_GRAPH_POOL_BORROW.override(True),
patch.object(pool, "get_global_graph_memory_pool", return_value=handle),
patch.object(pool, "_borrow_mem_pool", None),
):
for _ in range(3):
with pool.borrow_graph_pool(user="test"):
borrowed = torch.empty(16 << 20, dtype=torch.uint8, device="cuda")
# Stream-keyed segments allow side-stream copies, not allocations.
sink = torch.empty_like(borrowed)
with torch.cuda.stream(side):
widened = borrowed.to(torch.int32)
sink.copy_(borrowed)
borrowed.record_stream(side)
del borrowed, widened
del borrowed, sink
# Regression: this used to fail with "Trying to free a pointer not
# allocated here" after a deferred free was re-issued too early.
torch.cuda.empty_cache()
pool._teardown_borrow_pool()
del graph, y