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