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

Size KV pool after CUDA graph capture (opt-in) (#30157)

Browse Source
This commit is contained in:
cctry
2026-07-07 12:05:01 -07:00
committed by GitHub
parent ead1e490b5
commit 2ad9a243f5
10 changed files with 1013 additions and 156 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+3
View File
@@ -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)
python/sglang/srt/managers/scheduler.py
+5 -2
View File
@@ -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
python/sglang/srt/mem_cache/allocator/base.py
+6
View File
@@ -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()
python/sglang/srt/mem_cache/allocator/swa.py
+14
View File
@@ -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()
python/sglang/srt/mem_cache/kv_vmm_backing.py
+424
View File
@@ -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 <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)
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
python/sglang/srt/mem_cache/memory_pool.py
+187 -76
View File
@@ -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,
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,
)
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,
)
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()
python/sglang/srt/mem_cache/swa_memory_pool.py
+20
View File
@@ -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
python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py
+127 -18
View File
@@ -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
python/sglang/srt/server_args.py
+118 -50
View File
@@ -3643,67 +3643,35 @@ class ServerArgs:
)
if self.mem_fraction_static is None:
# Constant meta data (e.g., from attention backend)
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
# For activation slack
reserved_mem += self.tp_size * self.pp_size / 8 * 1024
else:
# Tokens the activation working set scales with (per serving mode).
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
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
activation_tokens = max(running_requests * draft_tokens, 2048)
elif self.chunked_prefill_size > 0:
reserved_mem += max(self.chunked_prefill_size, 2048) * 1.5
activation_tokens = max(self.chunked_prefill_size, 2048)
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
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
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
else:
# For MLA backend the memory overhead is much higher than expected with fa3
reserved_mem += 1.5 * 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)
# 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
# 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.
test/registered/mem_cache/test_post_capture_kv_sizing.py
+99
View File
@@ -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()