[Memory] Borrow CUDA graph pool storage for EAGLE sampling (#35375)

Co-authored-by: cctry <cctry@fb.com>

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

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 <cstddef>
+#include <cstdint>
+#include <mutex>
+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<std::mutex> 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<std::mutex> lk(g_mu); if (a) g_align=a; }}
+size_t bumparena_cursor_{sfx}(void){{
+  std::lock_guard<std::mutex> 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<std::mutex> 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<void*>(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<std::mutex> lk(g_mu); return g_freed_bytes; }}
+void bumparena_free_{sfx}(void* ptr, size_t size, int device, void* stream){{
+  std::lock_guard<std::mutex> 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."""
 

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)

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 <cstddef>
-#include <cstdint>
-#include <mutex>
-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<std::mutex> lk(g_mu); g_base=b; g_cursor=0; }}
-void kvarena_set_reserved_{sfx}(size_t r){{ std::lock_guard<std::mutex> lk(g_mu); g_reserved=r; }}
-void kvarena_set_align_{sfx}(size_t a){{ std::lock_guard<std::mutex> lk(g_mu); if (a) g_align=a; }}
-size_t kvarena_cursor_{sfx}(void){{ std::lock_guard<std::mutex> lk(g_mu); return g_cursor; }}
-void* kvarena_malloc_{sfx}(size_t size, int device, void* stream){{
-  std::lock_guard<std::mutex> 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<void*>(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._stub = BumpArenaStub()
-        self._fn_set_base(ctypes.c_void_p(self.base))
+        self._stub.set_extents([(self.base, self.reserved)])
-        self._fn_set_reserved(ctypes.c_size_t(self.reserved))
+        self._stub.set_align(self.granularity)
-        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()
         # 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:

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:

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:

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

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:
+        with graph_pool_capture_scope():
-            entry = self._capture(key, pixel_values, grid_thws, grid_thw_list, metadata)
+            try:
-        except Exception:
+                entry = self._capture(
-            self.failed_keys.add(key)
+                    key, pixel_values, grid_thws, grid_thw_list, metadata
-            logger.exception(
+                )
-                "Kimi-K3 ViT CUDA graph capture failed for key=%s; "
+            except Exception:
-                "using eager fallback",
+                self.failed_keys.add(key)
-                key,
+                logger.exception(
-            )
+                    "Kimi-K3 ViT CUDA graph capture failed for key=%s; "
-            outputs, _ = self._run_eager(pixel_values, grid_thws, grid_thw_list)
+                    "using eager fallback",
-            return outputs
+                    key,
+                )
+                outputs, _ = self._run_eager(pixel_values, grid_thws, grid_thw_list)
+                return outputs
 
-        self.graphs[key] = entry
+            self.graphs[key] = entry
-        entry.input_buffer.copy_(pixel_values)
+            entry.input_buffer.copy_(pixel_values)
-        entry.graph.replay()
+            entry.graph.replay()
         return list(entry.outputs)

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
+        # These full-vocabulary matrices are consumed by the sampling kernel
-            accept_index=accept_index,  # mutable
+        # within this step. Returned tensors were allocated before the scope,
-            accept_token_num=num_correct_drafts,  # mutable
+        # so the next CUDA graph replay may safely reclaim these borrowed bytes.
-            candidates=candidates,
+        with borrow_graph_pool(user="EAGLE probability borrow"):
-            # kwarg LHS retained as `retrive_*` to match sgl_kernel op schema.
+            expanded_temperature = torch.repeat_interleave(
-            retrive_index=verify_input.retrieve_index,
+                sampling_info.temperatures, verify_input.draft_token_num, dim=0
-            retrive_next_token=verify_input.retrieve_next_token,
+            )  # (bs * num_draft_tokens, 1)
-            retrive_next_sibling=verify_input.retrieve_next_sibling,
+
-            uniform_samples=coins,
+            target_probs = F.softmax(
-            uniform_samples_for_final_sampling=coins_for_final_sampling,
+                next_token_logits / expanded_temperature, dim=-1
-            target_probs=target_probs,
+            )  # (bs * num_draft_tokens, vocab_size)
-            draft_probs=draft_probs,
+            maybe_detect_nan(target_probs, "v2 verify: target_probs after softmax")
-            threshold_single=get_spec().speculative_accept_threshold_single,
+            if sampling_info.need_top_k_sampling:
-            threshold_acc=get_spec().speculative_accept_threshold_acc,
+                target_probs = top_k_renorm_prob(
-            deterministic=True,
+                    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

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):

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()