From 2ad9a243f560b557f75f922a1456435244fa7586 Mon Sep 17 00:00:00 2001 From: cctry Date: Tue, 7 Jul 2026 12:05:01 -0700 Subject: [PATCH] Size KV pool after CUDA graph capture (opt-in) (#30157) --- python/sglang/srt/environ.py | 3 + python/sglang/srt/managers/scheduler.py | 7 +- python/sglang/srt/mem_cache/allocator/base.py | 6 + python/sglang/srt/mem_cache/allocator/swa.py | 14 + python/sglang/srt/mem_cache/kv_vmm_backing.py | 424 ++++++++++++++++++ python/sglang/srt/mem_cache/memory_pool.py | 263 +++++++---- .../sglang/srt/mem_cache/swa_memory_pool.py | 20 + .../model_runner_kv_cache_mixin.py | 147 +++++- python/sglang/srt/server_args.py | 186 +++++--- .../mem_cache/test_post_capture_kv_sizing.py | 99 ++++ 10 files changed, 1013 insertions(+), 156 deletions(-) create mode 100644 python/sglang/srt/mem_cache/kv_vmm_backing.py create mode 100644 test/registered/mem_cache/test_post_capture_kv_sizing.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 8ce9ad1f6..44c467025 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -287,6 +287,9 @@ class Envs: # page tables (DP attn); paged backends like trtllm_mha consume it directly. SGLANG_USE_HND_KVCACHE = EnvBool(False) + # size the KV pool after CUDA-graph capture + SGLANG_ENABLE_POST_CAPTURE_KV_SIZING = EnvBool(False) + # Scheduler: memory leak test SGLANG_TEST_RETRACT = EnvBool(False) SGLANG_TEST_RETRACT_INTERVAL = EnvInt(3) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 03f8dd9eb..8557352e0 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -866,13 +866,16 @@ class Scheduler( self.init_tp_model_worker() self.maybe_init_draft_worker() - # Allocate KV cache pools for all workers. + # Prepare KV cache pools for all workers self.init_memory_pools() - # TODO: make memory profile consider cuda graph memory as well self.init_all_attention_backends() self.init_all_cuda_graphs() + model_runner = self.tp_worker.model_runner + if model_runner.token_to_kv_pool.post_capture_active: + model_runner.post_capture_resize_kv_pool() + # Dispatch the model worker if self.spec_algorithm.is_none(): self.model_worker = self.tp_worker diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index 949d92d73..7b6151dda 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -97,6 +97,12 @@ class BaseTokenToKVPoolAllocator(abc.ABC): def alloc_decode(self, *args, **kwargs): raise NotImplementedError("alloc_decode is only for paged allocator") + def resize(self, config) -> None: + self.size = config.max_total_num_tokens + if self.page_size > 1: + self.num_pages = config.max_total_num_tokens // self.page_size + self.clear() + @abc.abstractmethod def clear(self): raise NotImplementedError() diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index ebfec94b7..750683b8e 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -376,6 +376,20 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.full_attn_allocator.restore_state(state[0]) self.swa_attn_allocator.restore_state(state[1]) + def resize(self, config) -> None: + size_full = int(config.full_max_total_num_tokens) + size_swa = int(config.swa_max_total_num_tokens) + self._size_full = size_full + self._size_swa = size_swa + for alloc, sz in ( + (self.full_attn_allocator, size_full), + (self.swa_attn_allocator, size_swa), + ): + alloc.size = int(sz) + if self.page_size > 1: + alloc.num_pages = int(sz) // self.page_size + self.clear() + def clear(self): self.swa_attn_allocator.clear() self.full_attn_allocator.clear() diff --git a/python/sglang/srt/mem_cache/kv_vmm_backing.py b/python/sglang/srt/mem_cache/kv_vmm_backing.py new file mode 100644 index 000000000..6831755ae --- /dev/null +++ b/python/sglang/srt/mem_cache/kv_vmm_backing.py @@ -0,0 +1,424 @@ +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 + +if TYPE_CHECKING: + from sglang.srt.mem_cache.memory_pool import KvBufferDesc + +logger = logging.getLogger(__name__) + +_drv = None + + +def _driver(): + global _drv + if _drv is None: + from cuda.bindings import driver + + _drv = driver + return _drv + + +def _check(result, label: str): + drv = _driver() + err = result[0] if isinstance(result, tuple) else result + if err != drv.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"{label} failed: {err}") + return result[1] if isinstance(result, tuple) and len(result) > 1 else None + + +def align_up(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +def query_granularity(device_id: int) -> int: + """Minimum CUDA virtual-memory allocation granularity (bytes) for ``device_id``.""" + drv = _driver() + prop = drv.CUmemAllocationProp() + prop.type = drv.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED + prop.location.type = drv.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + prop.location.id = int(device_id) + return int( + _check( + drv.cuMemGetAllocationGranularity( + prop, + drv.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_MINIMUM, + ), + "cuMemGetAllocationGranularity", + ) + ) + + +# 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 arena instance and each instance loads its own .so, so +# multiple arenas per process (hybrid-SWA: full + swa) don't 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) + self._sfx = str(KvVmmArena._instance_count) + KvVmmArena._instance_count += 1 + drv = _driver() + with torch.cuda.device(self.device_id): + _check(drv.cuInit(0), "cuInit") + self._prop = drv.CUmemAllocationProp() + self._prop.type = drv.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED + self._prop.location.type = drv.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + self._prop.location.id = self.device_id + self.granularity = query_granularity(self.device_id) + self._access = drv.CUmemAccessDesc() + self._access.location.type = ( + drv.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + ) + self._access.location.id = self.device_id + self._access.flags = ( + drv.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE + ) + + self.reserved = self._align(reserve_bytes) + # Align the base to granularity so base + (granularity-aligned cursor) is + # always a valid cuMemMap address for per-buffer commit_range(). + self.base = int( + _check( + drv.cuMemAddressReserve(self.reserved, self.granularity, 0, 0), + "cuMemAddressReserve", + ) + ) + # commit_range bookkeeping: mapped VA -> (size, handle); committed bytes per offset. + self._ranges = {} + self._committed_by_offset = {} + 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() + # no_split so the caching allocator hands our bump pointers back verbatim. + self.pool = torch.cuda.MemPool(self._allocator, no_split=True) + logger.info( + "KvVmmArena[%s] ready: device=%d reserved=%.1f GiB granularity=%d KiB", + self._sfx, + self.device_id, + self.reserved / (1024**3), + self.granularity // 1024, + ) + + def _align(self, v: int) -> int: + return align_up(v, self.granularity) + + def _build_stub(self) -> ctypes.CDLL: + out_dir = os.path.join(tempfile.gettempdir(), "sgl_kv_vmm_arena") + 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, + ) + self._so_path = f"{out_dir}/{libname}.so" + lib = ctypes.CDLL(self._so_path) + self._fn_set_base = getattr(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 = getattr(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 = getattr(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 = getattr(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). + Maps one full handle per extension -- GB300 rejects partial-handle maps.""" + if self._closed: + raise RuntimeError("KvVmmArena.commit_range after close") + if offset % self.granularity != 0: + raise ValueError( + f"commit_range offset {offset} not granularity-aligned " + f"({self.granularity})" + ) + want = self._align(int(want_bytes)) + prev = self._committed_by_offset.get(offset, 0) + if want <= prev: + return + if offset + want > self.reserved: + raise RuntimeError( + f"commit_range [{offset}, {offset + want}) exceeds reservation " + f"{self.reserved}" + ) + drv = _driver() + add = want - prev + addr = self.base + offset + prev + with torch.cuda.device(self.device_id): + handle = _check(drv.cuMemCreate(add, self._prop, 0), "cuMemCreate") + try: + _check(drv.cuMemMap(addr, add, 0, handle, 0), "cuMemMap") + _check( + drv.cuMemSetAccess(addr, add, [self._access], 1), "cuMemSetAccess" + ) + except Exception: + # Roll back this failed extension; leave already-mapped ranges intact. + unmap = drv.cuMemUnmap(addr, add) + unmap = unmap[0] if isinstance(unmap, tuple) else unmap + rel = drv.cuMemRelease(handle) + rel = rel[0] if isinstance(rel, tuple) else rel + raise + self._ranges[addr] = (add, handle) + self._committed_by_offset[offset] = want + self._range_backed += add + + @property + def backed_bytes(self) -> int: + """Total physically-backed bytes (sum of scattered per-buffer ranges).""" + return self._range_backed + + @property + def cursor_bytes(self) -> int: + return int(self._fn_cursor()) + + def close(self) -> None: + if self._closed: + return + self._closed = True + drv = _driver() + try: + torch.cuda.synchronize() + except Exception as e: # pragma: no cover + logger.warning("KvVmmArena.close synchronize failed: %s", e) + for addr, (size, handle) in self._ranges.items(): + err = drv.cuMemUnmap(addr, size) + err = err[0] if isinstance(err, tuple) else err + if err != drv.CUresult.CUDA_SUCCESS: + logger.warning("cuMemUnmap range -> %s", err) + err = drv.cuMemRelease(handle) + err = err[0] if isinstance(err, tuple) else err + if err != drv.CUresult.CUDA_SUCCESS: + logger.warning("cuMemRelease range -> %s", err) + self._ranges.clear() + err = drv.cuMemAddressFree(self.base, self.reserved) + err = err[0] if isinstance(err, tuple) else err + if err != drv.CUresult.CUDA_SUCCESS: + logger.warning("cuMemAddressFree -> %s", err) + + +# torch's caching allocator hands the pluggable allocator whole large-pool segments +# (rounded up to >= ~20 MiB) per tensor, so reserve slack beyond the tight tensor sum. +# VA is free until committed, so this costs only address space, not GPU memory. +_PER_BUFFER_VA_SLACK = 32 << 20 + + +class _BufferSpec: + """Per-buffer placement + backing state inside the shared VA reservation.""" + + __slots__ = ("desc", "offset", "reserved_span", "aligned_reserved", "backed_to") + + def __init__( + self, + desc: KvBufferDesc, + offset: int, + reserved_span: int, + aligned_reserved: int, + ): + self.desc = desc + self.offset = offset # granularity-aligned arena offset of this buffer + self.reserved_span = reserved_span # logical (unaligned) tensor bytes + self.aligned_reserved = aligned_reserved # reserved span rounded to granularity + self.backed_to = 0 # bytes from offset currently backed + + +class KvVmmBufferOwner: + """Owns one ``KvVmmArena`` plus its incrementally-backed KV buffers. + + ``buffer_descs`` is an ordered list of ``KvBufferDesc``; the created ``torch.empty`` + tensors are exposed in the same order as ``self.tensors``. + """ + + def __init__( + self, + *, + device: str, + device_id: int, + store_dtype: torch.dtype, + page_size: int, + reserved_num_tokens: int, + buffer_descs: Sequence[KvBufferDesc], + ): + self.device = device + self.device_id = int(device_id) + self.store_dtype = store_dtype + self.page_size = int(page_size) + self._reserved_num_tokens = int(reserved_num_tokens) + self._final_num_tokens: Optional[int] = None + self._arena: Optional[KvVmmArena] = None + self._specs: List[_BufferSpec] = [] + self.tensors: List[torch.Tensor] = [] + + itemsize = store_dtype.itemsize + with torch.cuda.device(self.device_id): + gran = query_granularity(self.device_id) + reserved_spans = [d.reserved_span_bytes(itemsize) for d in buffer_descs] + aligned = [align_up(s, gran) for s in reserved_spans] + reserve_bytes = sum(a + _PER_BUFFER_VA_SLACK for a in aligned) + gran + self._arena = KvVmmArena(self.device_id, reserve_bytes=reserve_bytes) + assert self._arena.granularity == gran, (self._arena.granularity, gran) + + # NORMAL torch tensors through the arena MemPool; torch.empty never touches + # the unbacked tail. + with torch.cuda.use_mem_pool(self._arena.pool): + self.tensors = [ + torch.empty(d.shape, dtype=store_dtype, device=self.device) + for d in buffer_descs + ] + + specs: List[_BufferSpec] = [] + for desc, tensor, reserved_span, aligned_reserved in zip( + buffer_descs, self.tensors, reserved_spans, aligned + ): + if prod(tensor.shape) * itemsize != reserved_span: + raise RuntimeError( + f"buffer {desc.name!r} tensor bytes " + f"{prod(tensor.shape) * itemsize} != reserved span {reserved_span}" + ) + offset = tensor.data_ptr() - self._arena.base + if offset < 0 or offset % gran != 0: + raise RuntimeError( + f"buffer {desc.name!r} arena offset {offset} not " + f"granularity-aligned ({gran})" + ) + if offset + aligned_reserved > self._arena.reserved: + raise RuntimeError( + f"buffer {desc.name!r} [{offset}, {offset + aligned_reserved}) " + f"exceeds reservation {self._arena.reserved}" + ) + specs.append(_BufferSpec(desc, offset, reserved_span, aligned_reserved)) + self._specs = specs + + # Back one page so slot 0 is resident before capture: capture routes every + # dummy KV write to slot 0 (out_cache_loc is zeros). finalize() backs the rest. + self.ensure_prefix(self.page_size) + + for t in self.tensors: + assert ( + t.is_cuda and t.device.index == self.device_id + ), f"post-capture KV buffer landed on {t.device}, expected cuda:{self.device_id}" + + # -- backing -------------------------------------------------------------- + + @staticmethod + def _check_span(spec: _BufferSpec, span: int) -> int: + """Return ``span`` if it fits ``[0, reserved_span]``; raise otherwise.""" + span = int(span) + if not (0 <= span <= spec.reserved_span): + raise ValueError( + f"buffer {spec.desc.name!r}: span {span} outside " + f"[0, {spec.reserved_span}] (reserved tensor bytes)" + ) + return span + + def _back_spans(self, span_bytes: Sequence[int]) -> None: + """Back each buffer to (at least) ``span_bytes[i]``. An out-of-reservation + span is a descriptor bug: raise before committing anything, never clamp.""" + if self._arena is None: + raise RuntimeError("backing after close / before construction") + for spec, span in zip(self._specs, span_bytes): + self._check_span(spec, span) + gran = self._arena.granularity + for spec, span in zip(self._specs, span_bytes): + want = align_up( + int(span), gran + ) # <= aligned_reserved since span <= reserved + if want > spec.backed_to: + self._arena.commit_range(spec.offset, want) + spec.backed_to = want + + def ensure_prefix(self, num_tokens: int) -> None: + """Ensure the first ``num_tokens`` slots of every buffer are physically backed.""" + self._back_spans( + [s.desc.prefix_span_bytes(num_tokens, self.page_size) for s in self._specs] + ) + + def finalize(self, final_num_tokens: int) -> None: + """Back each buffer's final advertised span; set the final serving capacity.""" + final = int(final_num_tokens) + if not (self.page_size <= final <= self._reserved_num_tokens): + raise ValueError( + f"final_num_tokens={final} must satisfy page_size=" + f"{self.page_size} <= final <= reserved={self._reserved_num_tokens}" + ) + self._back_spans( + [s.desc.final_span_bytes(final, self.page_size) for s in self._specs] + ) + self._final_num_tokens = final + + # -- accessors / teardown ------------------------------------------------- + + @property + def backed_bytes(self) -> int: + return self._arena.backed_bytes if self._arena is not None else 0 + + def close(self) -> None: + self.tensors = [] + self._specs = [] + if self._arena is not None: + self._arena.close() + self._arena = None diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 87e2638c1..2a535af1c 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -25,6 +25,7 @@ from __future__ import annotations import abc import dataclasses import logging +import math from contextlib import contextmanager, nullcontext from dataclasses import dataclass, fields from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union @@ -52,6 +53,7 @@ from sglang.srt.layers.dcp import ( from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype, is_fp8_fnuz from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.mem_cache.allocator.mamba import MambaSlotAllocator +from sglang.srt.mem_cache.kv_vmm_backing import KvVmmBufferOwner from sglang.srt.mem_cache.layout.page_major import ( build_page_major_mamba_views, build_page_major_mha_views, @@ -1188,7 +1190,44 @@ def unwrap_write_loc(loc_info): return loc_info, None, None +class KvBufferDesc: + """Byte-span math for one KV buffer laid out as rows of ``row_bytes`` holding + ``tokens_per_row`` tokens each (a row = one token slot, or one whole page).""" + + __slots__ = ("name", "shape", "row_bytes", "tokens_per_row") + + def __init__(self, name: str, shape: tuple, *, row_bytes: int, tokens_per_row: int): + self.name = name + self.shape = tuple(shape) + self.row_bytes = int(row_bytes) + self.tokens_per_row = int(tokens_per_row) + + def _rows(self, num_tokens: int) -> int: + n = max(int(num_tokens), 0) + return (n + self.tokens_per_row - 1) // self.tokens_per_row + + def reserved_span_bytes(self, itemsize: int) -> int: + """Full upper-bound byte size of the buffer (its whole tensor).""" + return math.prod(self.shape) * itemsize + + def prefix_span_bytes(self, num_tokens: int, page_size: int) -> int: + """Bytes to back to make the first ``num_tokens`` tokens usable.""" + return self._rows(num_tokens) * self.row_bytes + + def final_span_bytes(self, num_tokens: int, page_size: int) -> int: + """Bytes of the final advertised span (adds the padded page). CEIL, not floor: + an unaligned count must still cover its partial last page (e.g. n=17, page=16 + -> 3 pages, not 2).""" + return self._rows(max(int(num_tokens), 0) + page_size) * self.row_bytes + + def item_len_bytes(self, page_size: int) -> int: + """Per-page transfer chunk (one page's worth of this buffer).""" + return (page_size // self.tokens_per_row) * self.row_bytes + + class KVCache(abc.ABC): + post_capture_active: bool = False + @abc.abstractmethod def __init__( self, @@ -1308,7 +1347,12 @@ class MHATokenToKVPool(KVCache): enable_alt_stream: bool = True, enable_kv_cache_copy: bool = False, kv_cache_layout: Optional[str] = None, + post_capture_active: bool = False, ): + if post_capture_active: + # Reserved upper bound only (unbacked VA): page-align UP so + # (size + page_size) % page_size == 0 holds for paged layouts. + size = (size + page_size - 1) // page_size * page_size super().__init__( size, page_size, @@ -1319,6 +1363,8 @@ class MHATokenToKVPool(KVCache): start_layer, end_layer, ) + self.post_capture_active = post_capture_active + self._post_capture_owner = None self.head_num = swa_head_num if swa_head_num is not None else head_num self.head_dim = swa_head_dim if swa_head_dim is not None else head_dim self.v_head_dim = ( @@ -1448,6 +1494,44 @@ class MHATokenToKVPool(KVCache): ) def _create_buffers(self): + if self.post_capture_active: + self._alloc_post_capture_buffers() + else: + self._create_buffers_normal() + self._kv_buffer_descs = self._build_kv_buffer_descs() + self.k_data_ptrs = torch.tensor( + [x.data_ptr() for x in self.k_buffer], + dtype=torch.uint64, + device=self.device, + ) + self.v_data_ptrs = torch.tensor( + [x.data_ptr() for x in self.v_buffer], + dtype=torch.uint64, + device=self.device, + ) + self.data_ptrs = torch.cat([self.k_data_ptrs, self.v_data_ptrs], dim=0) + self.data_strides = torch.tensor( + [ + np.prod(x.shape[1:]) * x.dtype.itemsize + for x in self.k_buffer + self.v_buffer + ], + device=self.device, + ) + + def _kv_buffer_shapes(self): + """(k_shape, v_shape)""" + if self.use_hnd: + return ( + (self.num_pages, self.head_num, self.page_size, self.head_dim), + (self.num_pages, self.head_num, self.page_size, self.v_head_dim), + ) + rows = self.size + self.page_size + return ( + (rows, self.head_num, self.head_dim), + (rows, self.head_num, self.v_head_dim), + ) + + def _create_buffers_normal(self): with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): with ( torch.cuda.use_mem_pool(self.custom_mem_pool) @@ -1455,28 +1539,7 @@ class MHATokenToKVPool(KVCache): else nullcontext() ): # The padded page (slot 0's page) absorbs dummy padded-token writes. - if self.use_hnd: - k_shape = ( - self.num_pages, - self.head_num, - self.page_size, - self.head_dim, - ) - v_shape = ( - self.num_pages, - self.head_num, - self.page_size, - self.v_head_dim, - ) - self.k_buffer = [ - torch.zeros(k_shape, dtype=self.store_dtype, device=self.device) - for _ in range(self.layer_num) - ] - self.v_buffer = [ - torch.zeros(v_shape, dtype=self.store_dtype, device=self.device) - for _ in range(self.layer_num) - ] - elif self.kv_cache_layout == "vectorized_5d": + if self.kv_cache_layout == "vectorized_5d": total_slots = self.size + self.page_size num_blocks = total_slots // self.page_size x = self._kv_vector_x @@ -1511,51 +1574,90 @@ class MHATokenToKVPool(KVCache): for _ in range(self.layer_num) ] else: - # [size, head_num, head_dim] for each layer - # The padded slot 0 is used for writing dummy outputs from padded tokens. + k_shape, v_shape = self._kv_buffer_shapes() self.k_buffer = [ - torch.zeros( - (self.size + self.page_size, self.head_num, self.head_dim), - dtype=self.store_dtype, - device=self.device, - ) + torch.zeros(k_shape, dtype=self.store_dtype, device=self.device) for _ in range(self.layer_num) ] self.v_buffer = [ - torch.zeros( - ( - self.size + self.page_size, - self.head_num, - self.v_head_dim, - ), - dtype=self.store_dtype, - device=self.device, - ) + torch.zeros(v_shape, dtype=self.store_dtype, device=self.device) for _ in range(self.layer_num) ] - self.k_data_ptrs = torch.tensor( - [x.data_ptr() for x in self.k_buffer], - dtype=torch.uint64, - device=self.device, + # -- post-capture VA backing (opt-in; overridable per layout) -------------- + + def _build_kv_buffer_descs(self): + """Per-buffer layout descriptors, k0..k(L-1) then v0..v(L-1). Drives both the + CUDA-VMM post-capture backing and PD-transfer registration + (get_contiguous_buf_infos). Override per layout.""" + itemsize = self.store_dtype.itemsize + # Derive from the real buffers when they exist (covers arbitrary layouts, + # e.g. vectorized_5d); fall back to _kv_buffer_shapes for the pre-allocation + # post-capture call, which only runs for NHD/HND. + if getattr(self, "k_buffer", None) and getattr(self, "v_buffer", None): + k_shape = tuple(self.k_buffer[0].shape) + v_shape = tuple(self.v_buffer[0].shape) + else: + k_shape, v_shape = self._kv_buffer_shapes() + # A row is a whole page when the leading dim is pages (hnd, vectorized_5d), + # a single token slot for the plain NHD [slots, ...] layout. + num_slots = self.size + self.page_size + tokens_per_row = ( + self.page_size if k_shape[0] * self.page_size == num_slots else 1 ) - self.v_data_ptrs = torch.tensor( - [x.data_ptr() for x in self.v_buffer], - dtype=torch.uint64, - device=self.device, - ) - self.data_ptrs = torch.cat([self.k_data_ptrs, self.v_data_ptrs], dim=0) - self.data_strides = torch.tensor( - [ - np.prod(x.shape[1:]) * x.dtype.itemsize - for x in self.k_buffer + self.v_buffer - ], + descs = [] + for prefix, shape in (("k", k_shape), ("v", v_shape)): + row_bytes = int(np.prod(shape[1:])) * itemsize + for layer in range(self.layer_num): + descs.append( + KvBufferDesc( + f"{prefix}{layer}", + shape, + row_bytes=row_bytes, + tokens_per_row=tokens_per_row, + ) + ) + return descs + + def _assign_post_capture_tensors(self, tensors): + """Map owner tensors (in ``_build_kv_buffer_descs`` order) to k/v_buffer.""" + self.k_buffer = tensors[: self.layer_num] + self.v_buffer = tensors[self.layer_num :] + + def _alloc_post_capture_buffers(self): + dev = torch.device(self.device) + device_id = dev.index if dev.index is not None else torch.cuda.current_device() + self._post_capture_owner = KvVmmBufferOwner( device=self.device, + device_id=device_id, + store_dtype=self.store_dtype, + page_size=self.page_size, + reserved_num_tokens=self.size, + buffer_descs=self._build_kv_buffer_descs(), ) + self._assign_post_capture_tensors(self._post_capture_owner.tensors) + + def finalize_backing(self, config) -> None: + """After capture+sizing: back the final span and set serving capacity. + ``config`` is a MemoryPoolConfig (duck-typed); each pool family reads the + fields it needs, so the finalizer stays pool-agnostic.""" + self._finalize_backing_tokens(config.max_total_num_tokens) + + def _finalize_backing_tokens(self, final_num_tokens: int) -> None: + """Token-count primitive shared by composite pools (e.g. SWA sub-pools).""" + self._post_capture_owner.finalize(final_num_tokens) + self.size = int(final_num_tokens) + + @property + def post_capture_backed_bytes(self) -> int: + return self._post_capture_owner.backed_bytes if self._post_capture_owner else 0 def _clear_buffers(self): del self.k_buffer del self.v_buffer + if self._post_capture_owner is not None: + self._post_capture_owner.close() + self._post_capture_owner = None def get_kv_size_bytes(self): assert hasattr(self, "k_buffer") @@ -1569,35 +1671,26 @@ class MHATokenToKVPool(KVCache): return k_size_bytes, v_size_bytes # for disagg + def _pd_registerable_tensors(self): + """Buffers to register for PD KV transfer, in ``_kv_buffer_descs`` order. + Override when the registerable storage differs from k/v_buffer.""" + return self.k_buffer + self.v_buffer + def get_contiguous_buf_infos(self): + """(ptrs, lens, item_lens) for PD KV transfer, derived from the descriptors. + ``lens`` is the final span at the CURRENT serving size -- for a post-capture + pool that is the physically-backed span, not the reserved VA upper bound.""" assert not self.use_hnd, ( "PD-disaggregation KV transfer assumes NHD slot-row layout; " "HND KV cache (SGLANG_USE_HND_KVCACHE) is not supported with disagg yet." ) - # layer_num x [seq_len, head_num, head_dim] - # layer_num x [page_num, page_size, head_num, head_dim] - kv_data_ptrs = [ - self._get_key_buffer(i).data_ptr() - for i in range(self.start_layer, self.start_layer + self.layer_num) - ] + [ - self._get_value_buffer(i).data_ptr() - for i in range(self.start_layer, self.start_layer + self.layer_num) + tensors = self._pd_registerable_tensors() + ptrs = [t.data_ptr() for t in tensors] + lens = [ + d.final_span_bytes(self.size, self.page_size) for d in self._kv_buffer_descs ] - kv_data_lens = [ - self._get_key_buffer(i).nbytes - for i in range(self.start_layer, self.start_layer + self.layer_num) - ] + [ - self._get_value_buffer(i).nbytes - for i in range(self.start_layer, self.start_layer + self.layer_num) - ] - kv_item_lens = [ - self._get_key_buffer(i)[0].nbytes * self.page_size - for i in range(self.start_layer, self.start_layer + self.layer_num) - ] + [ - self._get_value_buffer(i)[0].nbytes * self.page_size - for i in range(self.start_layer, self.start_layer + self.layer_num) - ] - return kv_data_ptrs, kv_data_lens, kv_item_lens + item_lens = [d.item_len_bytes(self.page_size) for d in self._kv_buffer_descs] + return ptrs, lens, item_lens def get_cpu_copy(self, indices, mamba_indices=None): assert not self.use_hnd, ( @@ -2382,6 +2475,7 @@ class HybridLinearKVPool(KVCache): # When provided (shared-KV-pool path), use this pool for the # full-attention layers instead of constructing one internally. full_kv_pool: Optional[KVCache] = None, + post_capture_active: bool = False, ): self.size = size self.dtype = dtype @@ -2418,6 +2512,9 @@ class HybridLinearKVPool(KVCache): # priority since they don't understand alternate layouts. TokenToKVPoolClass = full_kv_pool_class + post_capture_kwargs = ( + {"post_capture_active": True} if post_capture_active else {} + ) self.full_kv_pool = TokenToKVPoolClass( size=size, page_size=self.page_size, @@ -2428,6 +2525,7 @@ class HybridLinearKVPool(KVCache): device=device, enable_memory_saver=enable_memory_saver, enable_kv_cache_copy=enable_kv_cache_copy, + **post_capture_kwargs, ) else: TokenToKVPoolClass = MLATokenToKVPool @@ -2460,6 +2558,19 @@ class HybridLinearKVPool(KVCache): k_size, v_size = self.get_kv_size_bytes() self.mem_usage = (k_size + v_size) / GB + @property + def post_capture_active(self) -> bool: + return getattr(self.full_kv_pool, "post_capture_active", False) + + @property + def post_capture_backed_bytes(self) -> int: + return getattr(self.full_kv_pool, "post_capture_backed_bytes", 0) + + def finalize_backing(self, config) -> None: + # Only the attention KV is resized; the mamba state cache is fixed pre-capture. + self.full_kv_pool._finalize_backing_tokens(config.max_total_num_tokens) + self.size = int(config.max_total_num_tokens) + def get_kv_size_bytes(self): return self.full_kv_pool.get_kv_size_bytes() diff --git a/python/sglang/srt/mem_cache/swa_memory_pool.py b/python/sglang/srt/mem_cache/swa_memory_pool.py index c56b4238f..8cef5a735 100644 --- a/python/sglang/srt/mem_cache/swa_memory_pool.py +++ b/python/sglang/srt/mem_cache/swa_memory_pool.py @@ -86,6 +86,26 @@ class SWAKVPool(BaseSWAKVPool): f"SWAKVPool mem usage: {self.mem_usage:.2f} GB, swa size: {self.size_swa}, full size: {self.size}" ) + @property + def post_capture_active(self) -> bool: + """True iff the sub-pools took the post-capture VA-backed path (both share the flag).""" + return self.full_kv_pool.post_capture_active + + @property + def post_capture_backed_bytes(self) -> int: + """Physically-backed KV bytes across both sub-pools (post-capture only).""" + return ( + self.full_kv_pool.post_capture_backed_bytes + + self.swa_kv_pool.post_capture_backed_bytes + ) + + def finalize_backing(self, config) -> None: + """Back both sub-pools to their post-capture final sizes and record them.""" + self.full_kv_pool._finalize_backing_tokens(config.full_max_total_num_tokens) + self.swa_kv_pool._finalize_backing_tokens(config.swa_max_total_num_tokens) + self.size = int(config.full_max_total_num_tokens) + self.size_swa = int(config.swa_max_total_num_tokens) + def register_mapping(self, full_to_swa_index_mapping: torch.Tensor): self.full_to_swa_index_mapping = full_to_swa_index_mapping diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index ef8ea8558..3cc968cd1 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging import math -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional import torch @@ -15,9 +15,7 @@ from sglang.srt.configs.model_config import ( is_deepseek_v4, is_minimax_sparse, ) -from sglang.srt.distributed.parallel_state import ( - get_world_group, -) +from sglang.srt.distributed.parallel_state import get_world_group from sglang.srt.environ import envs from sglang.srt.layers.dp_attention import get_attention_tp_size from sglang.srt.mem_cache.allocator import ( @@ -49,9 +47,11 @@ from sglang.srt.mem_cache.memory_pool import ( ReqToTokenPool, ) from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool +from sglang.srt.model_executor.cuda_graph_config import Backend from sglang.srt.platforms import current_platform from sglang.srt.utils.common import ( get_available_gpu_memory, + get_device_memory_capacity, is_float4_e2m1fn_x2, is_hip, is_npu, @@ -112,9 +112,17 @@ class ModelRunnerKVCacheMixin: cpu_group=get_world_group().cpu_group, ) - rest_memory = available_gpu_memory - pre_model_load_memory * ( - 1 - self.mem_fraction_static - ) + slack_gb = pre_model_load_memory * (1 - self.mem_fraction_static) + if self.mambaish_config is not None and self.post_capture_kv_active: + # Mamba state is a fixed pre-capture allocation, so it can't ride the ~0 post-capture slack. + slack_gb = max( + slack_gb, + self.server_args.mamba_pre_capture_reserve_mb( + get_device_memory_capacity(self.device) + ) + / 1024, + ) + rest_memory = available_gpu_memory - slack_gb if self.mambaish_config is not None: rest_memory = self.handle_max_mamba_cache(rest_memory) @@ -349,6 +357,92 @@ class ModelRunnerKVCacheMixin: "attention, no HiSparse, and --kv-cache-dtype != fp4_e2m1." ) + @property + def post_capture_kv_active(self: ModelRunner) -> bool: + return ( + self.server_args.post_capture_kv_sizing_planned() + and current_platform.is_cuda() + and not self.is_draft_worker + ) + + def post_capture_resize_kv_pool(self: ModelRunner) -> None: + """Resize the KV pool after capture.""" + pool = self.token_to_kv_pool + torch.cuda.synchronize() + free_gb = get_available_gpu_memory( + self.device, + self.gpu_id, + distributed=get_world_group().world_size > 1, + cpu_group=get_world_group().cpu_group, + ) + headroom_gb = self.pre_model_load_memory * (1 - self.mem_fraction_static) + decode_cuda_graph_config = self.server_args.cuda_graph_config.decode + decode_max_bs = int(decode_cuda_graph_config.max_bs or 0) + running_requests = int(self.max_running_requests or decode_max_bs or 1) + eager_decode_gap = ( + self.server_args.disaggregation_mode != "prefill" + and decode_cuda_graph_config.backend != Backend.DISABLED + and decode_max_bs < running_requests + ) + if eager_decode_gap: + logger.warning( + "Post-capture KV sizing: decode CUDA graph max_bs=%d < " + "max_running_requests=%d; reserving activation headroom", + decode_max_bs, + running_requests, + ) + if eager_decode_gap or self.mambaish_config is not None: + headroom_gb = max( + headroom_gb, + self.server_args.mamba_pre_capture_reserve_mb( + get_device_memory_capacity(self.device) + ) + / 1024, + ) + budget_bytes = ( + int(max(0.0, free_gb - headroom_gb) * (1 << 30)) + + pool.post_capture_backed_bytes + ) + config = self._config_from_budget( + budget_bytes, cap_tokens=self.max_total_num_tokens + ) + pool.finalize_backing(config) + self.token_to_kv_pool_allocator.resize(config) + + # Set the new pool size + self.max_total_num_tokens = config.max_total_num_tokens + if self.is_hybrid_swa: + self.full_max_total_num_tokens = config.full_max_total_num_tokens + self.swa_max_total_num_tokens = config.swa_max_total_num_tokens + if self.memory_pool_config is not None: + self.memory_pool_config.max_total_num_tokens = config.max_total_num_tokens + self.memory_pool_config.full_max_total_num_tokens = ( + config.full_max_total_num_tokens + ) + self.memory_pool_config.swa_max_total_num_tokens = ( + config.swa_max_total_num_tokens + ) + if self.max_running_requests is not None: + # Re-calculate max_running_requests for the now smaller pool + capped_reqs = min( + self.max_running_requests, + self._resolve_max_num_reqs(config.max_total_num_tokens), + ) + if capped_reqs < self.max_running_requests: + logger.warning( + "Post-capture KV sizing: max_running_requests %d -> %d", + self.max_running_requests, + capped_reqs, + ) + self.max_running_requests = capped_reqs + if self.memory_pool_config is not None: + self.memory_pool_config.max_running_requests = capped_reqs + logger.info( + "Post-capture KV sizing: max_total_num_tokens=%d, free memory=%.2f GB", + config.max_total_num_tokens, + get_available_gpu_memory(self.device, self.gpu_id), + ) + def _init_unified_mamba_pools(self: ModelRunner, max_num_reqs: int): """Build the shared-KV-pool stack for a hybrid-Mamba model: one byte buffer split between the full-attn MHA KV pool and the @@ -792,6 +886,7 @@ class ModelRunnerKVCacheMixin: size_swa=self.swa_max_total_num_tokens, page_size=self.page_size, dtype=self.kv_cache_dtype, + post_capture_active=self.post_capture_kv_active, head_num=self.model_config.get_num_kv_heads( get_attention_tp_size() ), @@ -988,6 +1083,7 @@ class ModelRunnerKVCacheMixin: use_mla=self.use_mla_backend, start_layer=self.start_layer, full_kv_pool_class=mha_pool_class, + post_capture_active=self.post_capture_kv_active, **extra_args, ) else: @@ -1038,6 +1134,7 @@ class ModelRunnerKVCacheMixin: enable_kv_cache_copy=( self.server_args.speculative_algorithm is not None ), + post_capture_active=self.post_capture_kv_active, ) # Initialize token_to_kv_pool_allocator @@ -1292,6 +1389,28 @@ class ModelRunnerKVCacheMixin: self._init_pools() + def _config_from_budget( + self: ModelRunner, budget_bytes: int, *, cap_tokens: Optional[int] = None + ) -> MemoryPoolConfig: + """Turn a KV byte budget into a pool config via the configurator, re-applying + the external token constraints (user cap, page alignment, PP sync) and the + optional ``cap_tokens`` clamp.""" + # Local import avoids a pool_configurator import cycle. + from sglang.srt.model_executor.pool_configurator import ( + create_memory_pool_configurator, + ) + + configurator = create_memory_pool_configurator(self) + config = configurator.calculate_pool_sizes(budget_bytes, self.page_size) + max_tokens = self._apply_token_constraints(config.max_total_num_tokens) + if cap_tokens is not None: + max_tokens = min(max_tokens, cap_tokens) + if max_tokens != config.max_total_num_tokens: + config = configurator.calculate_pool_sizes_from_max_tokens( + max_tokens, self.page_size + ) + return config + def _resolve_memory_pool_config( self: ModelRunner, pre_model_load_memory: int ) -> MemoryPoolConfig: @@ -1301,21 +1420,11 @@ class ModelRunnerKVCacheMixin: ) available_bytes = self._profile_available_bytes(pre_model_load_memory) - page_size = self.server_args.page_size - - configurator = create_memory_pool_configurator(self) - config = configurator.calculate_pool_sizes(available_bytes, page_size) - - # Apply external constraints (user cap, page alignment, PP sync) - constrained = self._apply_token_constraints(config.max_total_num_tokens) - if constrained != config.max_total_num_tokens: - config = configurator.calculate_pool_sizes_from_max_tokens( - constrained, page_size - ) - + config = self._config_from_budget(available_bytes) config.max_running_requests = self._resolve_max_num_reqs( config.max_total_num_tokens ) + configurator = create_memory_pool_configurator(self) config = configurator.finalize_with_max_running_requests(config) config.mem_fraction_static = self.server_args.mem_fraction_static return config diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index fa03b6feb..63591ed6c 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -3643,67 +3643,35 @@ class ServerArgs: ) if self.mem_fraction_static is None: - # Constant meta data (e.g., from attention backend) - reserved_mem = 512 - # For activation slack - if self.disaggregation_mode == "decode": - # Decode nodes do no prefill; size activation to the decode batch. - running_requests = ( - self.max_running_requests or decode_cuda_graph_config.max_bs or 1 - ) - draft_tokens = self.speculative_num_draft_tokens or 1 - reserved_mem += max(running_requests * draft_tokens, 2048) * 1.5 - elif self.chunked_prefill_size > 0: - reserved_mem += max(self.chunked_prefill_size, 2048) * 1.5 + if self.post_capture_kv_sizing_planned(): + # Post-capture sizing measures free memory after graph capture, so + # skip the graph/activation reserve; keep only the floor + parallel slack. + reserved_mem = 512 + reserved_mem += self.tp_size * self.pp_size / 8 * 1024 else: - reserved_mem += max(self.max_prefill_tokens, 2048) * 1.5 - # For decode cuda graphs (skip on prefill-only nodes) - if ( - self.disaggregation_mode != "prefill" - and decode_cuda_graph_config.backend != Backend.DISABLED - ): - reserved_mem += decode_cuda_graph_config.max_bs * 2 - # Some adjustments for large parallel size - reserved_mem += self.tp_size * self.pp_size / 8 * 1024 - - if ( - self._resolved().enable_dp_attention - and self.disaggregation_mode != "prefill" - ): - # DP attention needs more padding for some operations - reserved_mem += decode_cuda_graph_config.max_bs * self.dp_size * 3 - - # DP attention uses much more memory for large cuda graph max bs, - # likely due to some inefficiencies in torch allocator or our implementation. - # So we need to reserve more memory. - if decode_cuda_graph_config.max_bs > 300: - reserved_mem += decode_cuda_graph_config.max_bs * self.dp_size * 1.5 - - # For prefill piecewise cuda graphs (skip on decode-only nodes) - if ( - self.disaggregation_mode != "decode" - and prefill_cuda_graph_config.backend != Backend.DISABLED - ): - if not self.use_mla_backend(): - # Only calculate the memory overhead for Non-Torch Memory use since the Torch Memory can be reused with Cuda Graph Capture - reserved_mem += len(prefill_cuda_graph_config.bs) * 8 + # Tokens the activation working set scales with (per serving mode). + if self.disaggregation_mode == "decode": + running_requests = ( + self.max_running_requests + or decode_cuda_graph_config.max_bs + or 1 + ) + draft_tokens = self.speculative_num_draft_tokens or 1 + activation_tokens = max(running_requests * draft_tokens, 2048) + elif self.chunked_prefill_size > 0: + activation_tokens = max(self.chunked_prefill_size, 2048) else: - # For MLA backend the memory overhead is much higher than expected with fa3 - reserved_mem += 1.5 * 1024 - - if gpu_mem is not None and gpu_mem > 60 * 1024: - reserved_mem = max(reserved_mem, 10 * 1024) - - # DeepEP all-to-all buffers captured in the decode graph are real - # extra allocations, so reserve them on top of the floor. - from sglang.srt.arg_groups.overrides import resolved_view - - if ( - self.disaggregation_mode != "prefill" - and decode_cuda_graph_config.backend != Backend.DISABLED - and resolved_view(self).moe_a2a_backend == "deepep" - ): - reserved_mem += 2 * 1024 + activation_tokens = max(self.max_prefill_tokens, 2048) + # Constant meta data (e.g., from attention backend) + activation slack. + reserved_mem = 512 + reserved_mem += activation_tokens * 1.5 + # Some adjustments for large parallel size + reserved_mem += self.tp_size * self.pp_size / 8 * 1024 + reserved_mem += self.reserve_for_graph_mb() + if gpu_mem is not None and gpu_mem > 60 * 1024: + reserved_mem = max(reserved_mem, 10 * 1024) + # Reserve headroom for DeepEP all-to-all buffers on top of the floor. + reserved_mem += self.reserve_for_deepep_a2a_mb() self.mem_fraction_static = ( round((gpu_mem - reserved_mem) / gpu_mem, 3) @@ -3725,6 +3693,106 @@ class ServerArgs: "Use environment variable SGLANG_SYMM_MEM_PREALLOC_GB_SIZE to change the prealloc size." ) + def post_capture_kv_sizing_planned(self) -> bool: + """Whether the mem_fraction heuristic may skip the graph reserve; must be + False for any config the runtime won't post-capture-size, else it gets an + under-reserved fraction (still-unsupported: MiniMax sparse).""" + # use_mla_backend is a method at args time but ModelRunner overwrites it + # with a bool on global_server_args (see the FIXME there) -- handle both. + use_mla = self.use_mla_backend + return ( + envs.SGLANG_ENABLE_POST_CAPTURE_KV_SIZING.get() + and self.device == "cuda" + and self.dcp_size == 1 + and not (use_mla() if callable(use_mla) else use_mla) + and not self.prefill_only_disable_kv_cache + and not self.enable_memory_saver + and envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL.get() is None + # Accurate sizing assumes graph-covered execution (graphs retain the + # activation workspace, so it is measured post-capture). An eager + # phase would pay activations outside the measurement: DP attention + # runs prefill eager internally, and an explicitly disabled phase + # backend runs eager -- keep those on the heuristic reserve. + and not self.enable_dp_attention + and ( + self.disaggregation_mode == "decode" + or self.cuda_graph_config.prefill.backend != Backend.DISABLED + ) + and ( + self.disaggregation_mode == "prefill" + or self.cuda_graph_config.decode.backend != Backend.DISABLED + ) + ) + + def mamba_pre_capture_reserve_mb(self, gpu_mem: Optional[float]) -> float: + # Realistic runtime reserve for the fixed (non-resizable) mamba state cache, + # which post-capture can't size from measured free memory. + if self.disaggregation_mode == "decode": + running_requests = ( + self.max_running_requests or self.cuda_graph_config.decode.max_bs or 1 + ) + activation_tokens = max( + running_requests * (self.speculative_num_draft_tokens or 1), 2048 + ) + elif self.chunked_prefill_size > 0: + activation_tokens = max(self.chunked_prefill_size, 2048) + else: + activation_tokens = max(self.max_prefill_tokens, 2048) + reserved_mem = ( + 512 + activation_tokens * 1.5 + self.tp_size * self.pp_size / 8 * 1024 + ) + if gpu_mem is not None and gpu_mem > 60 * 1024: + reserved_mem = max(reserved_mem, 10 * 1024) + return reserved_mem + + def reserve_for_graph_mb(self) -> float: + decode_cuda_graph_config = self.cuda_graph_config.decode + prefill_cuda_graph_config = self.cuda_graph_config.prefill + + reserved_mem = 0.0 + if ( + self.disaggregation_mode != "prefill" + and decode_cuda_graph_config.backend != Backend.DISABLED + ): + reserved_mem += decode_cuda_graph_config.max_bs * 2 + + if ( + self._resolved().enable_dp_attention + and self.disaggregation_mode != "prefill" + ): + # DP attention needs more padding for some operations, and much more for large + # cuda graph max bs (torch allocator / implementation inefficiencies). + reserved_mem += decode_cuda_graph_config.max_bs * self.dp_size * 3 + if decode_cuda_graph_config.max_bs > 300: + reserved_mem += decode_cuda_graph_config.max_bs * self.dp_size * 1.5 + + if ( + self.disaggregation_mode != "decode" + and prefill_cuda_graph_config.backend != Backend.DISABLED + ): + if not self.use_mla_backend(): + # Only non-torch memory is counted; torch memory is reused by cuda graph capture. + reserved_mem += len(prefill_cuda_graph_config.bs) * 8 + else: + # MLA backend overhead is much higher than expected with fa3. + reserved_mem += 1.5 * 1024 + + return reserved_mem + + def reserve_for_deepep_a2a_mb(self) -> float: + # DeepEP all-to-all buffers captured in the decode graph are real extra + # allocations, reserved on top of the floor. + from sglang.srt.arg_groups.overrides import resolved_view + + decode_cuda_graph_config = self.cuda_graph_config.decode + if ( + self.disaggregation_mode != "prefill" + and decode_cuda_graph_config.backend != Backend.DISABLED + and resolved_view(self).moe_a2a_backend == "deepep" + ): + return 2 * 1024 + return 0.0 + def _generate_decode_cuda_graph_batch_sizes(self, max_bs: int): """ Generate the list of batch sizes for CUDA graph capture based on max_bs. diff --git a/test/registered/mem_cache/test_post_capture_kv_sizing.py b/test/registered/mem_cache/test_post_capture_kv_sizing.py new file mode 100644 index 000000000..e4638297d --- /dev/null +++ b/test/registered/mem_cache/test_post_capture_kv_sizing.py @@ -0,0 +1,99 @@ +"""E2E guard for SGLANG_ENABLE_POST_CAPTURE_KV_SIZING. + +Post-capture KV sizing reserves the KV pool as CUDA VMM virtual memory, captures +CUDA graphs, then sizes and physically backs the pool from measured free memory. +This test launches a server with the feature enabled and asserts that: + 1. the post-capture sizing path actually ran (log line present, not a silent + no-op skip via post_capture_kv_sizing_planned), + 2. the pool was sized to a positive max_total_num_tokens, and + 3. gsm8k accuracy is unchanged vs. the default sizing path. +""" + +import os +import re +import unittest + +import requests + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.run_eval import run_eval +from sglang.test.test_utils import ( + DEFAULT_MODEL_NAME_FOR_TEST, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + SimpleNamespace, + popen_launch_server, +) + +# CI Registration +register_cuda_ci(est_time=240, stage="base-b", runner_config="1-gpu-large") + +STDOUT_FILENAME = "post_capture_kv_sizing_stdout.log" +STDERR_FILENAME = "post_capture_kv_sizing_stderr.log" + + +class TestPostCaptureKVSizing(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = DEFAULT_MODEL_NAME_FOR_TEST + cls.base_url = DEFAULT_URL_FOR_TEST + cls.stdout = open(STDOUT_FILENAME, "w") + cls.stderr = open(STDERR_FILENAME, "w") + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + env={**os.environ, "SGLANG_ENABLE_POST_CAPTURE_KV_SIZING": "1"}, + return_stdout_stderr=(cls.stdout, cls.stderr), + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + cls.stdout.close() + cls.stderr.close() + for f in (STDOUT_FILENAME, STDERR_FILENAME): + if os.path.exists(f): + os.remove(f) + + def _server_logs(self) -> str: + text = "" + for f in (STDOUT_FILENAME, STDERR_FILENAME): + if os.path.exists(f): + with open(f) as fh: + text += fh.read() + return text + + def test_post_capture_sizing_ran(self): + """The post-capture path must actually execute, not silently skip.""" + m = re.search( + r"Post-capture KV sizing: max_total_num_tokens=(\d+)", self._server_logs() + ) + self.assertIsNotNone( + m, + "Post-capture KV sizing log line not found; the feature was gated off " + "or the resize path did not run.", + ) + self.assertGreater(int(m.group(1)), 0) + + def test_server_info_pool_sized(self): + info = requests.get(f"{self.base_url}/server_info").json() + self.assertGreater(info["max_total_num_tokens"], 0) + + def test_gsm8k_accuracy(self): + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="gsm8k", + num_examples=500, + num_threads=1024, + ) + metrics = run_eval(args) + print(f"GSM8K Accuracy: {metrics['score']:.3f}") + self.assertGreater(metrics["score"], 0.80) + + +if __name__ == "__main__": + unittest.main()