diff --git a/python/sglang/srt/compilation/cuda_piecewise_backend.py b/python/sglang/srt/compilation/cuda_piecewise_backend.py index a6ba34e74..f4192b4bc 100644 --- a/python/sglang/srt/compilation/cuda_piecewise_backend.py +++ b/python/sglang/srt/compilation/cuda_piecewise_backend.py @@ -18,6 +18,10 @@ from sglang.srt.compilation.compile_phase import ( is_in_torch_compile_warmup, ) from sglang.srt.compilation.weak_ref_tensor import weak_ref_tensors +from sglang.srt.model_executor.runner_utils.pool import ( + graph_pool_capture_scope, + graph_pool_replay_scope, +) from sglang.srt.utils.common import print_warning_once logger = logging.getLogger(__name__) @@ -186,7 +190,9 @@ class CUDAPiecewiseBackend: stack.enter_context(patch("gc.collect", lambda: None)) stack.enter_context(patch("torch.cuda.empty_cache", lambda: None)) # mind-exploding: carefully manage the reference and memory. - with torch.cuda.graph(cudagraph, pool=self.graph_pool, stream=stream): + with graph_pool_capture_scope(), torch.cuda.graph( + cudagraph, pool=self.graph_pool, stream=stream + ): # `output` is managed by pytorch's cudagraph pool output = entry.runnable(*args) if self.is_last_graph: @@ -218,5 +224,6 @@ class CUDAPiecewiseBackend: "Input addresses for cudagraphs are different during replay." f" Expected {entry.input_addresses}, got {new_input_addresses}" ) - entry.cudagraph.replay() + with graph_pool_replay_scope(): + entry.cudagraph.replay() return entry.output diff --git a/python/sglang/srt/cuda_vmm_utils.py b/python/sglang/srt/cuda_vmm_utils.py index 79c14a762..b0d873e2f 100644 --- a/python/sglang/srt/cuda_vmm_utils.py +++ b/python/sglang/srt/cuda_vmm_utils.py @@ -337,6 +337,166 @@ def align_down(value: int, alignment: int) -> int: return int(value) // alignment * alignment +# Bump allocator over caller-provided extents: malloc first-fits an extent and +# hands back base+cursor, bounded by each extent's RESERVED size (not any +# committed watermark) so upper-bound tensors can be allocated before physical +# commit. Allocations are aligned so VMM users can commit each pointer at its +# own VA range (cuMemMap requires it; GB300 rejects partial-handle maps). +# Symbols are SUFFIXED per (process, arena instance) and each instance loads its +# own .so, so neither multiple arenas per process nor co-located engine +# processes sharing the tempdir clobber each other. +def _bump_arena_stub_source(sfx: str) -> str: + return f""" +#include +#include +#include +extern "C" {{ +enum {{ BUMPARENA_MAX_EXTENTS = 64 }}; +static uintptr_t g_bases[BUMPARENA_MAX_EXTENTS]; +static size_t g_reserved[BUMPARENA_MAX_EXTENTS]; +static size_t g_cursors[BUMPARENA_MAX_EXTENTS]; +static size_t g_num_extents = 0; +static size_t g_freed_bytes = 0; +static size_t g_align = 512; +static std::mutex g_mu; +static size_t align_up(size_t v, size_t a){{ return (v + a - 1) / a * a; }} +void bumparena_set_extents_{sfx}(const uintptr_t* bases, const size_t* sizes, size_t n){{ + std::lock_guard lk(g_mu); + if (n > BUMPARENA_MAX_EXTENTS) n = BUMPARENA_MAX_EXTENTS; + g_num_extents = n; + g_freed_bytes = 0; + for (size_t i = 0; i < n; ++i) {{ + g_bases[i] = bases[i]; + g_reserved[i] = sizes[i]; + g_cursors[i] = 0; + }} +}} +void bumparena_set_align_{sfx}(size_t a){{ std::lock_guard lk(g_mu); if (a) g_align=a; }} +size_t bumparena_cursor_{sfx}(void){{ + std::lock_guard lk(g_mu); + size_t total = 0; + for (size_t i = 0; i < g_num_extents; ++i) total += g_cursors[i]; + return total; +}} +void* bumparena_malloc_{sfx}(size_t size, int device, void* stream){{ + std::lock_guard lk(g_mu); + for (size_t i = 0; i < g_num_extents; ++i) {{ + size_t need = g_cursors[i] + align_up(size, g_align); + if (need > g_reserved[i]) continue; + void* p = reinterpret_cast(g_bases[i] + g_cursors[i]); + g_cursors[i] = need; + return p; + }} + return 0; // no extent fits -- surfaces as an allocator OOM +}} +size_t bumparena_freed_{sfx}(void){{ std::lock_guard lk(g_mu); return g_freed_bytes; }} +void bumparena_free_{sfx}(void* ptr, size_t size, int device, void* stream){{ + std::lock_guard lk(g_mu); + g_freed_bytes += size; +}} +}} +""" + + +class BumpArenaStub: + """JIT-built pluggable bump allocator over caller-provided device VA extents. + + ``malloc`` first-fits an extent and hands out ``base + cursor``; ``free`` + is a no-op. Plain ``torch.empty`` can thus be placed on externally managed + storage by wrapping ``allocator`` in a ``torch.cuda.MemPool``. + ``set_extents`` re-points the arena and resets every cursor, letting one + stub serve successive region sets. + """ + + MAX_EXTENTS = 64 # mirrors BUMPARENA_MAX_EXTENTS in the stub source + + # Per-instance suffix -> isolated allocator symbols/state (see _bump_arena_stub_source). + _instance_count = 0 + + def __init__(self): + # Unique per (process, instance): the stub .so lives in a host-shared + # tempdir, so co-located engine processes must not build the same-named + # .so (they race and one loads a half-relinked copy -> undefined symbol + # crash). + self.sfx = f"{os.getpid()}_{BumpArenaStub._instance_count}" + BumpArenaStub._instance_count += 1 + self._lib = self._build() + from torch.cuda.memory import CUDAPluggableAllocator + + self.allocator = CUDAPluggableAllocator( + self._so_path, + f"bumparena_malloc_{self.sfx}", + f"bumparena_free_{self.sfx}", + ).allocator() + + def _build(self) -> ctypes.CDLL: + import torch.utils.cpp_extension + + # Per-stub build dir: load_inline writes every caller's source to the + # same main.cpp inside build_directory, so any sharing (across + # co-located engine processes under the host tempdir, or across arenas + # within one process) can compile another stub's source and link a .so + # missing this stub's symbols. One dir per stub means no shared ninja + # scratch or .so, ever. + out_dir = os.path.join(tempfile.gettempdir(), "sgl_bump_arena", self.sfx) + os.makedirs(out_dir, exist_ok=True) + libname = f"sgl_bump_arena_stub_{self.sfx}" + torch.utils.cpp_extension.load_inline( + name=libname, + cpp_sources=_bump_arena_stub_source(self.sfx), + with_cuda=False, # pure arithmetic -- no nvcc, no CUDA headers + is_python_module=False, + verbose=False, + build_directory=out_dir, + no_implicit_headers=True, + ) + self._so_path = f"{out_dir}/{libname}.so" + lib = ctypes.CDLL(self._so_path) + self._fn_set_extents = lib[f"bumparena_set_extents_{self.sfx}"] + self._fn_set_extents.argtypes = [ + ctypes.POINTER(ctypes.c_void_p), + ctypes.POINTER(ctypes.c_size_t), + ctypes.c_size_t, + ] + self._fn_set_extents.restype = None + self._fn_set_align = lib[f"bumparena_set_align_{self.sfx}"] + self._fn_set_align.argtypes = [ctypes.c_size_t] + self._fn_set_align.restype = None + self._fn_cursor = lib[f"bumparena_cursor_{self.sfx}"] + self._fn_cursor.argtypes = [] + self._fn_cursor.restype = ctypes.c_size_t + self._fn_freed = lib[f"bumparena_freed_{self.sfx}"] + self._fn_freed.argtypes = [] + self._fn_freed.restype = ctypes.c_size_t + return lib + + def set_extents(self, extents: List[tuple]) -> None: + """Register ``(base, nbytes)`` extents (first-fit order) and reset + every bump cursor.""" + if len(extents) > BumpArenaStub.MAX_EXTENTS: + raise ValueError( + f"{len(extents)} extents exceed BUMPARENA_MAX_EXTENTS " + f"({BumpArenaStub.MAX_EXTENTS})" + ) + n = len(extents) + bases = (ctypes.c_void_p * n)(*(base for base, _ in extents)) + sizes = (ctypes.c_size_t * n)(*(nbytes for _, nbytes in extents)) + self._fn_set_extents(bases, sizes, ctypes.c_size_t(n)) + + def set_align(self, nbytes: int) -> None: + self._fn_set_align(ctypes.c_size_t(nbytes)) + + @property + def cursor_bytes(self) -> int: + return int(self._fn_cursor()) + + @property + def freed_bytes(self) -> int: + """Bytes handed back through ``free`` since the last ``set_extents`` + -- nonzero means someone (empty_cache) released this arena's segments.""" + return int(self._fn_freed()) + + class VmmReservation: """Own a VA reservation, its mappings, and their teardown order.""" diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index e9acf9578..f668210d9 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1189,6 +1189,8 @@ class Envs: # Guards CUDA graph executable dedup via cudaGraphExecUpdate. SGLANG_ENABLE_CUDA_GRAPH_DEDUP = EnvBool(False) SGLANG_MEMORY_SAVER_CUDA_GRAPH = EnvBool(False) + # Reuse wholly-free graph-pool segments for step-local eager allocations. + SGLANG_ENABLE_GRAPH_POOL_BORROW = EnvBool(False) # Eager forward wraps the ForwardBatch's own tensors instead of copying them # into the CUDA graph buffer registry (no per-iter device-to-device copy). SGLANG_EAGER_INPUT_NO_COPY = EnvBool(False) diff --git a/python/sglang/srt/mem_cache/kv_vmm_backing.py b/python/sglang/srt/mem_cache/kv_vmm_backing.py index c720cc69a..3f76a6179 100644 --- a/python/sglang/srt/mem_cache/kv_vmm_backing.py +++ b/python/sglang/srt/mem_cache/kv_vmm_backing.py @@ -1,17 +1,13 @@ from __future__ import annotations -import ctypes import logging -import os -import tempfile from math import prod from typing import TYPE_CHECKING, List, Optional, Sequence import torch -import torch.utils.cpp_extension -from torch.cuda.memory import CUDAPluggableAllocator from sglang.srt.cuda_vmm_utils import ( + BumpArenaStub, VmmReservation, align_up, allocation_handle_type_name, @@ -25,58 +21,14 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -# Bump allocator: hands back base+cursor, bounded by the RESERVED size (not the -# committed watermark) so upper-bound tensors can be allocated before physical -# commit. Allocations are granularity-aligned so each pointer can be committed at -# its own VA range (cuMemMap requires it; GB300 rejects partial-handle maps). -# Symbols are SUFFIXED per (process, arena instance) and each instance loads its -# own .so, so neither multiple arenas per process (hybrid-SWA: full + swa) nor -# co-located engine processes sharing the tempdir clobber each other. -def _stub_source(sfx: str) -> str: - return f""" -#include -#include -#include -extern "C" {{ -static uintptr_t g_base = 0; -static size_t g_cursor = 0; -static size_t g_reserved = 0; -static size_t g_align = 512; -static std::mutex g_mu; -static size_t align_up(size_t v, size_t a){{ return (v + a - 1) / a * a; }} -void kvarena_set_base_{sfx}(uintptr_t b){{ std::lock_guard lk(g_mu); g_base=b; g_cursor=0; }} -void kvarena_set_reserved_{sfx}(size_t r){{ std::lock_guard lk(g_mu); g_reserved=r; }} -void kvarena_set_align_{sfx}(size_t a){{ std::lock_guard lk(g_mu); if (a) g_align=a; }} -size_t kvarena_cursor_{sfx}(void){{ std::lock_guard lk(g_mu); return g_cursor; }} -void* kvarena_malloc_{sfx}(size_t size, int device, void* stream){{ - std::lock_guard lk(g_mu); - size_t need = g_cursor + align_up(size, g_align); - if (need > g_reserved) return 0; // never exceed the reserved VA range - void* p = reinterpret_cast(g_base + g_cursor); - g_cursor = need; - return p; -}} -void kvarena_free_{sfx}(void* ptr, size_t size, int device, void* stream){{}} -}} -""" - - _DEFAULT_RESERVE_BYTES = 256 * (1024**3) # 256 GiB virtual; free until committed class KvVmmArena: """One device's CUDA virtual-memory reservation exposed as a ``torch.cuda.MemPool``.""" - # Per-instance suffix source -> isolated allocator symbols/state (see _stub_source). - _instance_count = 0 - def __init__(self, device_id: int, reserve_bytes: int = _DEFAULT_RESERVE_BYTES): self.device_id = int(device_id) - # Unique per (process, arena instance): the stub .so lives in a host-shared - # tempdir, so co-located engine processes must not build the same-named .so - # (they race and one loads a half-relinked copy -> undefined symbol crash). - self._sfx = f"{os.getpid()}_{KvVmmArena._instance_count}" - KvVmmArena._instance_count += 1 with torch.cuda.device(self.device_id): prop = make_device_allocation_prop(self.device_id) self.handle_type = prop.requestedHandleTypes @@ -96,19 +48,15 @@ class KvVmmArena: self._range_backed = 0 self._closed = False - self._lib = self._build_stub() - self._fn_set_base(ctypes.c_void_p(self.base)) - self._fn_set_reserved(ctypes.c_size_t(self.reserved)) - self._fn_set_align(ctypes.c_size_t(self.granularity)) - self._allocator = CUDAPluggableAllocator( - self._so_path, f"kvarena_malloc_{self._sfx}", f"kvarena_free_{self._sfx}" - ).allocator() + self._stub = BumpArenaStub() + self._stub.set_extents([(self.base, self.reserved)]) + self._stub.set_align(self.granularity) # no_split so the caching allocator hands our bump pointers back verbatim. - self.pool = torch.cuda.MemPool(self._allocator, no_split=True) + self.pool = torch.cuda.MemPool(self._stub.allocator, no_split=True) logger.info( "KvVmmArena[%s] ready: device=%d reserved_va=%.1f GiB " "granularity=%d KiB handle_type=%s", - self._sfx, + self._stub.sfx, self.device_id, self.reserved / (1024**3), self.granularity // 1024, @@ -118,40 +66,6 @@ class KvVmmArena: def _align(self, v: int) -> int: return align_up(v, self.granularity) - def _build_stub(self) -> ctypes.CDLL: - # Per-arena build dir: load_inline writes every caller's source to the same - # main.cpp inside build_directory, so any sharing (across co-located engine - # processes under the host tempdir, or across arenas within one process) - # can compile another arena's source and link a .so missing this arena's - # symbols. One dir per stub means no shared ninja scratch or .so, ever. - out_dir = os.path.join(tempfile.gettempdir(), "sgl_kv_vmm_arena", self._sfx) - os.makedirs(out_dir, exist_ok=True) - libname = f"sgl_kv_vmm_arena_stub_{self._sfx}" - torch.utils.cpp_extension.load_inline( - name=libname, - cpp_sources=_stub_source(self._sfx), - with_cuda=False, # pure arithmetic — no nvcc, no CUDA headers - is_python_module=False, - verbose=False, - build_directory=out_dir, - no_implicit_headers=True, - ) - self._so_path = f"{out_dir}/{libname}.so" - lib = ctypes.CDLL(self._so_path) - self._fn_set_base = lib[f"kvarena_set_base_{self._sfx}"] - self._fn_set_base.argtypes = [ctypes.c_void_p] - self._fn_set_base.restype = None - self._fn_set_reserved = lib[f"kvarena_set_reserved_{self._sfx}"] - self._fn_set_reserved.argtypes = [ctypes.c_size_t] - self._fn_set_reserved.restype = None - self._fn_set_align = lib[f"kvarena_set_align_{self._sfx}"] - self._fn_set_align.argtypes = [ctypes.c_size_t] - self._fn_set_align.restype = None - self._fn_cursor = lib[f"kvarena_cursor_{self._sfx}"] - self._fn_cursor.argtypes = [] - self._fn_cursor.restype = ctypes.c_size_t - return lib - def commit_range(self, offset: int, want_bytes: int) -> None: """Back ``[base+offset, base+offset+want_bytes)`` (monotonic per offset). ``offset`` must be granularity-aligned (the bump allocator guarantees it). @@ -189,7 +103,7 @@ class KvVmmArena: @property def cursor_bytes(self) -> int: - return int(self._fn_cursor()) + return self._stub.cursor_bytes def close(self) -> None: if self._closed: diff --git a/python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py index c4fc86532..f12bef028 100644 --- a/python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py +++ b/python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py @@ -41,6 +41,8 @@ from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import ) from sglang.srt.model_executor.runner_utils.pool import ( get_or_create_global_graph_memory_pool, + graph_pool_capture_scope, + graph_pool_replay_scope, ) from sglang.srt.utils import get_bool_env_var from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter @@ -126,7 +128,7 @@ class BreakableCudaGraphBackend(DedupedCudaGraphMixin, BaseCudaGraphBackend): size = shape_key.size if self._shared_output_buffer is None: self._shared_output_buffer = self._alloc_full_buffer(warmup_out, size) - with BreakableCUDAGraphCapture( + with graph_pool_capture_scope(), BreakableCUDAGraphCapture( cuda_graph=graph, pool=self._pool, stream=self._capture_stream, @@ -245,7 +247,8 @@ class BreakableCudaGraphBackend(DedupedCudaGraphMixin, BaseCudaGraphBackend): static_forward_batch: ForwardBatch, **kwargs, ) -> Any: - self._graphs[shape_key].replay() + with graph_pool_replay_scope(): + self._graphs[shape_key].replay() return self._outputs[shape_key] def cleanup(self) -> None: diff --git a/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py index ff2b64910..aa5b9f3db 100644 --- a/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py +++ b/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py @@ -32,6 +32,8 @@ from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import ( ) from sglang.srt.model_executor.runner_utils.pool import ( get_or_create_global_graph_memory_pool, + graph_pool_capture_scope, + graph_pool_replay_scope, ) from sglang.srt.utils import get_bool_env_var from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter @@ -125,7 +127,10 @@ class FullCudaGraphBackend(BaseCudaGraphBackend): else: graph_ctx = self._device_module.graph - with graph_ctx(cuda_graph=graph, pool=self._pool, stream=self._capture_stream): + with ( + graph_pool_capture_scope(), + graph_ctx(cuda_graph=graph, pool=self._pool, stream=self._capture_stream), + ): out = forward_fn() if profiler is not None: @@ -147,7 +152,8 @@ class FullCudaGraphBackend(BaseCudaGraphBackend): static_forward_batch: ForwardBatch, **kwargs, ) -> Any: - self._graphs[shape_key].replay() + with graph_pool_replay_scope(): + self._graphs[shape_key].replay() return self._outputs[shape_key] def cleanup(self) -> None: diff --git a/python/sglang/srt/model_executor/runner_utils/pool.py b/python/sglang/srt/model_executor/runner_utils/pool.py index f0fd4cca4..7dd64dc8f 100644 --- a/python/sglang/srt/model_executor/runner_utils/pool.py +++ b/python/sglang/srt/model_executor/runner_utils/pool.py @@ -18,9 +18,50 @@ sharing one pool reserves only the larger phase's capture footprint. from __future__ import annotations -from typing import Any, Optional +import logging +from contextlib import contextmanager +from typing import Any, Iterator, Optional +import torch + +from sglang.srt.cuda_vmm_utils import BumpArenaStub +from sglang.srt.environ import envs from sglang.srt.runtime_context import get_resources +from sglang.srt.utils import is_cuda + +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 + + +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 + logger.info("Graph pool borrow disabled: %s", reason) + + +def set_graph_pool_borrow_runs(runs: list[tuple[int, int]]) -> None: + """Use fixed graph-storage extents instead of snapshots of the shared pool. + + This supports graph storage whose addresses are managed externally but + remain stable for the process lifetime. Registering an empty list disables + borrowing. + """ + global _borrow_static_runs + _borrow_static_runs = sorted(runs, key=lambda run: run[1], reverse=True)[ + : BumpArenaStub.MAX_EXTENTS + ] + logger.info( + "Graph pool borrow runs pinned: runs=%d free=%d", + len(_borrow_static_runs), + sum(nbytes for _, nbytes in _borrow_static_runs), + ) def get_global_graph_memory_pool() -> Optional[Any]: @@ -38,3 +79,140 @@ def get_or_create_global_graph_memory_pool(device_module: Any) -> Any: if resources.graph_memory_pool is None: resources.graph_memory_pool = device_module.graph_pool_handle() return resources.graph_memory_pool + + +def graph_pool_borrow_enabled() -> bool: + if ( + _borrow_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 + return get_global_graph_memory_pool() is not None + + +@contextmanager +def graph_pool_user_scope(user: str) -> Iterator[None]: + global _active_graph_pool_user + # Graph replay silently overwrites aliases of its allocator-free blocks. + if _active_graph_pool_user is not None: + raise RuntimeError( + f"graph pool already has live user {_active_graph_pool_user!r}; " + f"cannot use it for {user!r}" + ) + _active_graph_pool_user = user + try: + yield + finally: + _active_graph_pool_user = None + + +@contextmanager +def graph_pool_replay_scope() -> Iterator[None]: + if not graph_pool_borrow_enabled(): + yield + return + with graph_pool_user_scope("CUDA graph"): + yield + + +@contextmanager +def graph_pool_capture_scope() -> Iterator[None]: + if not graph_pool_borrow_enabled(): + yield + return + with graph_pool_user_scope("CUDA graph"): + # Capture re-carves the pool's free space: the borrow extents go stale, + # so retire the borrow pool before capturing. + _teardown_borrow_pool() + yield + + +def find_free_graph_pool_runs(pool_id: Any) -> list[tuple[int, int]]: + """Return contiguous inactive runs inside wholly-free segments, largest first.""" + runs: list[tuple[int, int]] = [] + for segment in torch.cuda.memory_snapshot(pool_id, include_traces=False): + if segment["allocated_size"] != 0: + continue + run_address = 0 + run_bytes = 0 + for block in segment["blocks"]: + if block["state"] == "inactive": + if run_bytes == 0: + run_address = block["address"] + run_bytes += block["size"] + continue + if run_bytes: + runs.append((run_address, run_bytes)) + run_bytes = 0 + if run_bytes: + runs.append((run_address, run_bytes)) + runs.sort(key=lambda run: run[1], reverse=True) + return runs + + +def _teardown_borrow_pool() -> None: + """Retire the persistent borrow pool after draining deferred frees.""" + global _borrow_mem_pool + if _borrow_mem_pool is None: + return + # 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 + + +@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 + 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 + if not graph_pool_borrow_enabled(): + yield + return + with graph_pool_user_scope(user): + if _borrow_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 + or _borrow_stub.cursor_bytes > _borrow_extents_total // 2 + ): + # Rebuild if empty_cache released the arena's segments or if + # unresolved deferred frees consumed half of the extents. + _teardown_borrow_pool() + if _borrow_mem_pool is None: + if _borrow_stub is None: + _borrow_stub = BumpArenaStub() + if _borrow_static_runs is not None: + runs = _borrow_static_runs + else: + runs = find_free_graph_pool_runs(get_global_graph_memory_pool())[ + : BumpArenaStub.MAX_EXTENTS + ] + _borrow_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) + logger.info( + "Graph pool borrow extents: runs=%d free=%d", + len(runs), + _borrow_extents_total, + ) + with torch.cuda.use_mem_pool(_borrow_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 diff --git a/python/sglang/srt/multimodal/kimi_k3_vit_cuda_graph_runner.py b/python/sglang/srt/multimodal/kimi_k3_vit_cuda_graph_runner.py index 3002a7578..358062e61 100644 --- a/python/sglang/srt/multimodal/kimi_k3_vit_cuda_graph_runner.py +++ b/python/sglang/srt/multimodal/kimi_k3_vit_cuda_graph_runner.py @@ -11,6 +11,8 @@ import torch from sglang.srt.model_executor.runner_utils.pool import ( get_or_create_global_graph_memory_pool, + graph_pool_capture_scope, + graph_pool_replay_scope, ) if TYPE_CHECKING: @@ -157,7 +159,8 @@ class KimiK3ViTCudaGraphRunner: entry = self.graphs.get(key) if entry is not None: entry.input_buffer.copy_(pixel_values) - entry.graph.replay() + with graph_pool_replay_scope(): + entry.graph.replay() return list(entry.outputs) max_seqlen = max(t * h * w for t, h, w in grid_thw_list) @@ -193,19 +196,22 @@ class KimiK3ViTCudaGraphRunner: total_tokens=pixel_values.shape[0], dtype=pixel_values.dtype, ) - try: - entry = self._capture(key, pixel_values, grid_thws, grid_thw_list, metadata) - except Exception: - self.failed_keys.add(key) - logger.exception( - "Kimi-K3 ViT CUDA graph capture failed for key=%s; " - "using eager fallback", - key, - ) - outputs, _ = self._run_eager(pixel_values, grid_thws, grid_thw_list) - return outputs + with graph_pool_capture_scope(): + try: + entry = self._capture( + key, pixel_values, grid_thws, grid_thw_list, metadata + ) + except Exception: + self.failed_keys.add(key) + logger.exception( + "Kimi-K3 ViT CUDA graph capture failed for key=%s; " + "using eager fallback", + key, + ) + outputs, _ = self._run_eager(pixel_values, grid_thws, grid_thw_list) + return outputs - self.graphs[key] = entry - entry.input_buffer.copy_(pixel_values) - entry.graph.replay() + self.graphs[key] = entry + entry.input_buffer.copy_(pixel_values) + entry.graph.replay() return list(entry.outputs) diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index 662ca8091..28cf4610c 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -19,6 +19,7 @@ from sglang.srt.mem_cache.allocation_sizing import ( get_alloc_reserve_per_decode, page_aligned_decode_alloc_lens, ) +from sglang.srt.model_executor.runner_utils.pool import borrow_graph_pool from sglang.srt.runtime_context import get_parallel, get_spec from sglang.srt.utils import ( is_cpu, @@ -768,79 +769,92 @@ def eagle_sample( use_rejection_sampling = get_spec().speculative_use_rejection_sampling - # Apply temperature and get target probs - expanded_temperature = torch.repeat_interleave( - sampling_info.temperatures, verify_input.draft_token_num, dim=0 - ) # (bs * num_draft_tokens, 1) - - target_probs = F.softmax( - next_token_logits / expanded_temperature, dim=-1 - ) # (bs * num_draft_tokens, vocab_size) - maybe_detect_nan(target_probs, "v2 verify: target_probs after softmax") - if sampling_info.need_top_k_sampling: - target_probs = top_k_renorm_prob( - target_probs, - torch.repeat_interleave( - sampling_info.top_ks, verify_input.draft_token_num, dim=0 - ), - ) # (bs * num_draft_tokens, vocab_size) - maybe_detect_nan(target_probs, "v2 verify: target_probs after top_k_renorm") - if sampling_info.need_top_p_sampling: - target_probs = top_p_renorm_prob( - target_probs, - torch.repeat_interleave( - sampling_info.top_ps, verify_input.draft_token_num, dim=0 - ), - ) - maybe_detect_nan(target_probs, "v2 verify: target_probs after top_p_renorm") - target_probs = target_probs.reshape(bs, verify_input.draft_token_num, -1) - draft_probs = ( - verify_input.draft_probs - if use_rejection_sampling - else torch.zeros_like(target_probs) - ) - # Defense-in-depth behind the spec_hook startup allowlist: validate the - # actual kernel inputs (catches draft_probs plumbing regressions or a - # startup guard bypassed by a worker subclass) before the Triton kernel. - if use_rejection_sampling and ( - draft_probs is None or draft_probs.shape[-1] != target_probs.shape[-1] - ): - raise ValueError( - "Rejection sampling requires a target-vocab draft proposal " - "distribution; the current speculative algorithm/draft worker " - "does not produce one (draft_probs missing or vocab-mismatched)." - ) - - coins, coins_for_final_sampling = _verify_coins( - sampling_info=sampling_info, - seq_lens=batch.seq_lens, - draft_token_num=verify_input.draft_token_num, - candidates=candidates, - device=device, - ) - sampling_fn = ( chain_speculative_sampling_triton if use_rejection_sampling else tree_speculative_sampling_target_only ) - sampling_fn( - predicts=predict, # mutable - accept_index=accept_index, # mutable - accept_token_num=num_correct_drafts, # mutable - candidates=candidates, - # kwarg LHS retained as `retrive_*` to match sgl_kernel op schema. - retrive_index=verify_input.retrieve_index, - retrive_next_token=verify_input.retrieve_next_token, - retrive_next_sibling=verify_input.retrieve_next_sibling, - uniform_samples=coins, - uniform_samples_for_final_sampling=coins_for_final_sampling, - target_probs=target_probs, - draft_probs=draft_probs, - threshold_single=get_spec().speculative_accept_threshold_single, - threshold_acc=get_spec().speculative_accept_threshold_acc, - deterministic=True, - ) + + # These full-vocabulary matrices are consumed by the sampling kernel + # within this step. Returned tensors were allocated before the scope, + # so the next CUDA graph replay may safely reclaim these borrowed bytes. + with borrow_graph_pool(user="EAGLE probability borrow"): + expanded_temperature = torch.repeat_interleave( + sampling_info.temperatures, verify_input.draft_token_num, dim=0 + ) # (bs * num_draft_tokens, 1) + + target_probs = F.softmax( + next_token_logits / expanded_temperature, dim=-1 + ) # (bs * num_draft_tokens, vocab_size) + maybe_detect_nan(target_probs, "v2 verify: target_probs after softmax") + if sampling_info.need_top_k_sampling: + target_probs = top_k_renorm_prob( + target_probs, + torch.repeat_interleave( + sampling_info.top_ks, verify_input.draft_token_num, dim=0 + ), + ) # (bs * num_draft_tokens, vocab_size) + maybe_detect_nan( + target_probs, "v2 verify: target_probs after top_k_renorm" + ) + if sampling_info.need_top_p_sampling: + target_probs = top_p_renorm_prob( + target_probs, + torch.repeat_interleave( + sampling_info.top_ps, verify_input.draft_token_num, dim=0 + ), + ) + maybe_detect_nan( + target_probs, "v2 verify: target_probs after top_p_renorm" + ) + target_probs = target_probs.reshape(bs, verify_input.draft_token_num, -1) + draft_probs = ( + verify_input.draft_probs + if use_rejection_sampling + else torch.zeros_like(target_probs) + ) + # Defense-in-depth behind the spec_hook startup allowlist: validate + # the actual kernel inputs before the Triton kernel. + if use_rejection_sampling and ( + draft_probs is None or draft_probs.shape[-1] != target_probs.shape[-1] + ): + raise ValueError( + "Rejection sampling requires a target-vocab draft proposal " + "distribution; the current speculative algorithm/draft worker " + "does not produce one (draft_probs missing or vocab-mismatched)." + ) + + coins, coins_for_final_sampling = _verify_coins( + sampling_info=sampling_info, + seq_lens=batch.seq_lens, + draft_token_num=verify_input.draft_token_num, + candidates=candidates, + device=device, + ) + sampling_fn( + predicts=predict, # mutable + accept_index=accept_index, # mutable + accept_token_num=num_correct_drafts, # mutable + candidates=candidates, + # kwarg LHS retained as `retrive_*` to match sgl_kernel op schema. + retrive_index=verify_input.retrieve_index, + retrive_next_token=verify_input.retrieve_next_token, + retrive_next_sibling=verify_input.retrieve_next_sibling, + uniform_samples=coins, + uniform_samples_for_final_sampling=coins_for_final_sampling, + target_probs=target_probs, + draft_probs=draft_probs, + threshold_single=get_spec().speculative_accept_threshold_single, + threshold_acc=get_spec().speculative_accept_threshold_acc, + deterministic=True, + ) + del ( + expanded_temperature, + target_probs, + draft_probs, + coins, + coins_for_final_sampling, + ) # Sync sampling results across TP ranks: different GPUs may # produce slightly different target_probs due to floating-point diff --git a/test/registered/spec/eagle/test_spec_eagle.py b/test/registered/spec/eagle/test_spec_eagle.py index 27fa45d43..21ac87fc2 100644 --- a/test/registered/spec/eagle/test_spec_eagle.py +++ b/test/registered/spec/eagle/test_spec_eagle.py @@ -33,7 +33,10 @@ _KITS = ( class _Core(Eagle3Base): - env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),) + env_overrides = ( + (envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1), + (envs.SGLANG_ENABLE_GRAPH_POOL_BORROW, 1), + ) class TestEagle3Overlap(_Core, *_KITS): 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 new file mode 100644 index 000000000..56f6f97c5 --- /dev/null +++ b/test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py @@ -0,0 +1,306 @@ +"""CUDA graph-pool borrowing allocator and lifetime regression tests.""" + +import unittest +from contextlib import contextmanager +from types import SimpleNamespace +from unittest.mock import MagicMock, Mock, patch + +import torch + +from sglang.srt.environ import envs +from sglang.srt.model_executor.runner_backend.full_cuda_graph_backend import ( + FullCudaGraphBackend, +) +from sglang.srt.model_executor.runner_utils import pool +from sglang.srt.speculative import eagle_utils +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=60, stage="base-b", runner_config="1-gpu-small") + + +class TestGraphPoolBorrow(CustomTestCase): + def setUp(self): + super().setUp() + self._reset_borrow_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 + + def test_graph_replay_fails_during_active_pool_borrow(self): + graph = Mock() + backend = object.__new__(FullCudaGraphBackend) + backend._graphs = {"shape": graph} + backend._outputs = {"shape": None} + snapshot = [ + { + "allocated_size": 0, + "blocks": [{"state": "inactive", "address": 4096, "size": 4096}], + } + ] + + 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(pool.torch.cuda, "MemPool"), + patch.object(pool.torch.cuda, "use_mem_pool"), + patch.object(pool.torch.cuda, "memory_snapshot", return_value=snapshot), + pool.borrow_graph_pool(user="test"), + ): + with self.assertRaisesRegex( + RuntimeError, "graph pool already has live user" + ): + backend.replay("shape", None) + + graph.replay.assert_not_called() + + def test_external_graph_storage_can_disable_borrowing(self): + 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)), + ): + self.assertTrue(pool.graph_pool_borrow_enabled()) + pool.disable_graph_pool_borrow("graph storage is externally managed") + self.assertFalse(pool.graph_pool_borrow_enabled()) + + def test_eagle_non_greedy_probabilities_use_borrow_scope(self): + state = {"active": False, "users": []} + + @contextmanager + def tracking_borrow(*, user): + self.assertFalse(state["active"]) + state["active"] = True + state["users"].append(user) + try: + yield + finally: + state["active"] = False + + import torch.nn.functional as F + + real_softmax = F.softmax + + def checked_softmax(*args, **kwargs): + self.assertTrue(state["active"]) + return real_softmax(*args, **kwargs) + + def fake_sampling(**kwargs): + self.assertTrue(state["active"]) + kwargs["predicts"].fill_(3) + kwargs["accept_index"].fill_(0) + kwargs["accept_token_num"].fill_(1) + + verify_input = SimpleNamespace( + draft_token_num=2, + draft_token=torch.tensor([1, 2], dtype=torch.int32), + max_tree_depth=2, + tree_topk=1, + retrieve_index=torch.zeros((1, 2), dtype=torch.int32), + retrieve_next_token=torch.zeros((1, 2), dtype=torch.int32), + retrieve_next_sibling=torch.zeros((1, 2), dtype=torch.int32), + draft_probs=None, + ) + sampling_info = SimpleNamespace( + acc_additive_penalties=None, + acc_scaling_penalties=None, + logit_bias=None, + is_all_greedy=False, + temperatures=torch.ones((1, 1)), + need_top_k_sampling=False, + need_top_p_sampling=False, + sampling_seed=None, + ) + batch = SimpleNamespace( + device="cpu", + seq_lens=torch.tensor([4], dtype=torch.int32), + sampling_info=sampling_info, + forward_mode=SimpleNamespace(is_idle=lambda: False), + ) + logits_output = SimpleNamespace(next_token_logits=torch.randn((2, 8))) + spec_config = SimpleNamespace( + speculative_use_rejection_sampling=False, + speculative_accept_threshold_single=1.0, + speculative_accept_threshold_acc=1.0, + ) + tp_group = SimpleNamespace(world_size=1) + + with ( + patch.object(eagle_utils, "borrow_graph_pool", tracking_borrow), + patch.object(eagle_utils, "get_spec", return_value=spec_config), + patch("torch.nn.functional.softmax", side_effect=checked_softmax), + patch( + "sglang.srt.layers.dp_attention.is_dp_attention_enabled", + return_value=False, + ), + patch("sglang.srt.distributed.get_tp_group", return_value=tp_group), + patch( + "sgl_kernel.tree_speculative_sampling_target_only", + side_effect=fake_sampling, + ), + ): + predict, accept_lens, accept_index = eagle_utils.eagle_sample( + verify_input, batch, logits_output + ) + + self.assertEqual(state["users"], ["EAGLE probability borrow"]) + self.assertFalse(state["active"]) + self.assertTrue(torch.equal(predict, torch.full_like(predict, 3))) + self.assertTrue(torch.equal(accept_lens, torch.full_like(accept_lens, 2))) + self.assertTrue(torch.equal(accept_index, torch.zeros_like(accept_index))) + + @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") + def test_borrowed_allocations_land_on_free_graph_pool_runs(self): + """Borrowed blocks span free runs, recycle, and add no reservation.""" + handle = torch.cuda.graph_pool_handle() + graph = torch.cuda.CUDAGraph() + x = torch.zeros(8, device="cuda") + stream = torch.cuda.Stream() + with torch.cuda.stream(stream), torch.cuda.graph( + graph, pool=handle, stream=stream + ): + # Two capture-only transients become disjoint free graph-pool runs. + transient_a = torch.empty(48 << 20, dtype=torch.uint8, device="cuda") + transient_b = torch.empty(24 << 20, dtype=torch.uint8, device="cuda") + y = x + 1 + del transient_a, transient_b + torch.cuda.synchronize() + + device_id = torch.cuda.current_device() + reserved_before = torch.cuda.memory_reserved(device_id) + 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) + largest_run_bytes = runs[0][1] + + def on_a_run(tensor): + start = tensor.data_ptr() + end = start + tensor.nbytes + return any( + address <= start and end <= address + nbytes + for address, nbytes in runs + ) + + with pool.borrow_graph_pool(user="test"): + # Together these exceed the largest run, forcing first-fit to + # use more than one captured extent. + a = torch.empty(40 << 20, dtype=torch.uint8, device="cuda") + b = torch.empty(20 << 20, dtype=torch.uint8, device="cuda") + self.assertGreater(a.nbytes + b.nbytes, largest_run_bytes) + self.assertTrue(on_a_run(a) and on_a_run(b)) + self.assertTrue( + a.data_ptr() + a.nbytes <= b.data_ptr() + or b.data_ptr() + b.nbytes <= a.data_ptr() + ) + + recycled_address = a.data_ptr() + del a + c = torch.empty(40 << 20, dtype=torch.uint8, device="cuda") + self.assertEqual(c.data_ptr(), recycled_address) + del b, c + + # Captures retire the persistent borrow pool. Its storage aliases + # existing graph-pool runs, so the reserved footprint is unchanged. + pool._teardown_borrow_pool() + + self.assertEqual(torch.cuda.memory_reserved(device_id), reserved_before) + del graph, y + + @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.""" + handle = torch.cuda.graph_pool_handle() + graph = torch.cuda.CUDAGraph() + x = 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(64 << 20, dtype=torch.uint8, device="cuda") + y = x + 1 + del transient + torch.cuda.synchronize() + + runs = pool.find_free_graph_pool_runs(handle) + self.assertTrue(runs) + 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()) + with pool.borrow_graph_pool(user="test"): + borrowed = torch.empty(16 << 20, dtype=torch.uint8, device="cuda") + self.assertTrue( + any( + address <= borrowed.data_ptr() < address + nbytes + for address, nbytes in runs + ) + ) + del borrowed + + del graph, y + + @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") + def test_cross_stream_borrow_frees_resolve_before_pointer_reuse(self): + """Deferred record_stream frees must not collide on the next borrow.""" + handle = torch.cuda.graph_pool_handle() + graph = torch.cuda.CUDAGraph() + x = 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(128 << 20, dtype=torch.uint8, device="cuda") + y = x + 1 + del transient + torch.cuda.synchronize() + + side = torch.cuda.Stream() + 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") + with torch.cuda.stream(side): + widened = borrowed.to(torch.int32) + borrowed.record_stream(side) + del borrowed, widened + # 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() + + del graph, y + + +if __name__ == "__main__": + unittest.main()