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

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

Browse Source 
Co-authored-by: cctry <cctry@fb.com>
This commit is contained in:
cctry
2026-08-18 16:19:41 -07:00
committed by GitHub
co-authored by cctry
parent 7ebaa98f81
commit 37c09ff3d8
11 changed files with 783 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/compilation/cuda_piecewise_backend.py
+9 -2
View File
@@ -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
+160
View File
@@ -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
+2
View File
@@ -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
+7 -93
View File
@@ -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._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:
python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py
+5 -2
View File
@@ -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
+8 -2
View File
@@ -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
+179 -1
View File
@@ -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
+21 -15
View File
@@ -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)
python/sglang/srt/speculative/eagle_utils.py
+82 -68
View File
@@ -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