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

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)

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

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

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

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

@@ -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:
+                if self.kv_cache_layout == "vectorized_5d":
-                    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":
                     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
+                    k_shape, v_shape = self._kv_buffer_shapes()
-                    # The padded slot 0 is used for writing dummy outputs from padded tokens.
                     self.k_buffer = [
-                        torch.zeros(
+                        torch.zeros(k_shape, dtype=self.store_dtype, device=self.device)
-                            (self.size + self.page_size, self.head_num, self.head_dim),
-                            dtype=self.store_dtype,
-                            device=self.device,
-                        )
                         for _ in range(self.layer_num)
                     ]
                     self.v_buffer = [
-                        torch.zeros(
+                        torch.zeros(v_shape, dtype=self.store_dtype, device=self.device)
-                            (
-                                self.size + self.page_size,
-                                self.head_num,
-                                self.v_head_dim,
-                            ),
-                            dtype=self.store_dtype,
-                            device=self.device,
-                        )
                         for _ in range(self.layer_num)
                     ]
 
-        self.k_data_ptrs = torch.tensor(
+    # -- post-capture VA backing (opt-in; overridable per layout) --------------
-            [x.data_ptr() for x in self.k_buffer],
+
-            dtype=torch.uint64,
+    def _build_kv_buffer_descs(self):
-            device=self.device,
+        """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(
+        descs = []
-            [x.data_ptr() for x in self.v_buffer],
+        for prefix, shape in (("k", k_shape), ("v", v_shape)):
-            dtype=torch.uint64,
+            row_bytes = int(np.prod(shape[1:])) * itemsize
-            device=self.device,
+            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]
+        tensors = self._pd_registerable_tensors()
-        # layer_num x [page_num, page_size, head_num, head_dim]
+        ptrs = [t.data_ptr() for t in tensors]
-        kv_data_ptrs = [
+        lens = [
-            self._get_key_buffer(i).data_ptr()
+            d.final_span_bytes(self.size, self.page_size) for d in self._kv_buffer_descs
-            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)
         ]
-        kv_data_lens = [
+        item_lens = [d.item_len_bytes(self.page_size) for d in self._kv_buffer_descs]
-            self._get_key_buffer(i).nbytes
+        return ptrs, lens, item_lens
-            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
 
     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

@@ -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

@@ -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 (
+from sglang.srt.distributed.parallel_state import get_world_group
-    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 * (
+        slack_gb = pre_model_load_memory * (1 - self.mem_fraction_static)
-            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
+        config = self._config_from_budget(available_bytes)
-
-        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.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

@@ -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
+                    activation_tokens = max(self.max_prefill_tokens, 2048)
-            # For decode cuda graphs (skip on prefill-only nodes)
+                # Constant meta data (e.g., from attention backend) + activation slack.
-            if (
+                reserved_mem = 512
-                self.disaggregation_mode != "prefill"
+                reserved_mem += activation_tokens * 1.5
-                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
-
+                reserved_mem += self.reserve_for_graph_mb()
-            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
-
                 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.
-            # DeepEP all-to-all buffers captured in the decode graph are real
+                reserved_mem += self.reserve_for_deepep_a2a_mb()
-            # 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
 
             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

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