chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

CP HiCache: multi-slab host KV cache to respect cudaHostRegister per-call ceiling

Browse Source 
A single cudaHostRegister over a >~1 TB host buffer fails with
cudaErrorMemoryAllocation on B300 (hard per-call ceiling: 512 GiB OK, 1024 GiB
FAIL), crashing the hicache_size=1600 (~1.5 TB CP shared-L2 slab) prefill at
startup. The registration cannot simply be chunked: a memcpy (cudaMemcpyBatchAsync,
the CP-L2 H2D/D2H transfer) fails with cudaErrorInvalidValue when its host range
straddles a registration boundary (verified empirically on b300-049).

Fix: physically split the host cache into multiple page-aligned slabs, each <= a
safe single-registration size (default 480 GiB, env SGLANG_CP_HICACHE_MAX_SLAB_GB),
reusing the existing SharedHostTensorGroupAllocator + per-slab transfer splitting
(_host_transfer_segments). Each slab is one whole registration and no transfer
crosses a boundary; small configs (hicache_size<=400) stay single-slab unchanged.

- memory_pool_host.py: add cp_hicache_max_single_register_bytes() + the fail-loud
  _check_single_cuda_host_register_size guard; revert the (transfer-unsafe)
  registration chunking back to one cudaHostRegister per buffer/slab.
- hiradix_cache.py: _cp_shared_l2_slab_pages_by_payload auto-caps each payload's
  slab <= the ceiling so large caches auto-split.
- cp_l3_slab_accessor.py: CpSharedL2SlabAccessor is now slab-count-aware
  (CpL3SlabSpan + per-slab dispatch via global_base_page; per-slab layer stride);
  _cp_l3_slab_spans rewires _maybe_init_cp_l3 off the single-slab assumption. L3
  disk slabs / slot pool / LMDB index / GC are content-addressed and unchanged.

Tests: multi-slab accessor incl. a non-circular torch.frombuffer-layout check; a
real allocate_group + _cp_l3_slab_spans roundtrip; slab-cap auto-split; L3 store
cross-slab spill/reload. Reviewed by 3 adversarial agents, no correctness bugs.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-06-23 15:33:10 +00:00
parent 0025b0e819
commit dbc5ebbaa0
6 changed files with 555 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

125
python/sglang/srt/mem_cache/cp_l3_slab_accessor.py
View File

@@ -57,67 +57,130 @@ class CpL3SlabLayout:
return cls(n_layers=n_active_layers, page_num=indexer_page_num, slice_bytes=indexer_page_stride_size)
class CpSharedL2SlabAccessor:
"""Gather/scatter one global page's strided layer slices to/from a contiguous buffer."""
@dataclass(frozen=True)
class CpL3SlabSpan:
"""One physical shared-L2 slab inside a (possibly multi-slab) payload.
def __init__(self, slab_mmap, layout: CpL3SlabLayout):
self._mm = slab_mmap
self._lo = layout
if len(slab_mmap) < layout.total_bytes:
``mmap`` is this slab's own byte buffer; the slab owns global pages
``[global_base_page, global_base_page + num_pages)``. The host KV cache is
split into multiple slabs when its total registration would exceed the
per-call ``cudaHostRegister`` ceiling, so a global page must be dispatched to
its owning slab and addressed with that slab's *own* layer stride
(``num_pages * slice_bytes``)."""
global_base_page: int
num_pages: int
mmap: object
def __post_init__(self):
if self.num_pages <= 0 or self.global_base_page < 0:
_l3_failfast(
f"slab mmap {len(slab_mmap)} bytes < layout total {layout.total_bytes} "
f"(n_layers={layout.n_layers} page_num={layout.page_num} slice={layout.slice_bytes})"
f"CpL3SlabSpan requires num_pages>0 and global_base_page>=0; got "
f"global_base_page={self.global_base_page} num_pages={self.num_pages}"
)
class CpSharedL2SlabAccessor:
"""Gather/scatter one global page's strided layer slices to/from a contiguous buffer.
Supports a payload split across multiple physical slabs (each its own mmap).
A global page is dispatched to its owning slab via ``global_base_page`` and
addressed with that slab's per-slab layer stride ``num_pages * slice_bytes``.
A single-slab payload is just one span with ``global_base_page == 0``."""
def __init__(self, slabs, *, n_layers: int, slice_bytes: int):
if int(n_layers) <= 0 or int(slice_bytes) <= 0:
_l3_failfast(
f"CpSharedL2SlabAccessor requires positive dims; got "
f"n_layers={n_layers} slice_bytes={slice_bytes}"
)
spans = sorted(slabs, key=lambda s: int(s.global_base_page))
if not spans:
_l3_failfast("CpSharedL2SlabAccessor requires at least one slab span")
self._n_layers = int(n_layers)
self._slice_bytes = int(slice_bytes)
# Spans must tile [0, total_pages) contiguously, and each mmap must back
# its full layer-strided extent.
expected_base = 0
for s in spans:
if int(s.global_base_page) != expected_base:
_l3_failfast(
f"CpSharedL2SlabAccessor slab spans not contiguous: expected "
f"global_base_page={expected_base}, got {s.global_base_page} "
f"(spans must tile [0,total) with no gaps/overlap)"
)
need = self._n_layers * int(s.num_pages) * self._slice_bytes
if len(s.mmap) < need:
_l3_failfast(
f"slab mmap {len(s.mmap)} bytes < required {need} "
f"(n_layers={self._n_layers} num_pages={s.num_pages} slice={self._slice_bytes})"
)
expected_base += int(s.num_pages)
self._spans = tuple(spans)
self._total_pages = expected_base
@classmethod
def from_single_mmap(cls, slab_mmap, layout: "CpL3SlabLayout") -> "CpSharedL2SlabAccessor":
"""Build a one-slab accessor from a single mmap + layout (the historical /
single-slab path; also used by tests)."""
return cls(
[CpL3SlabSpan(0, layout.page_num, slab_mmap)],
n_layers=layout.n_layers,
slice_bytes=layout.slice_bytes,
)
@property
def page_blob_bytes(self) -> int:
return self._lo.page_blob_bytes
return self._n_layers * self._slice_bytes
@property
def n_layers(self) -> int:
return self._lo.n_layers
return self._n_layers
@property
def page_num(self) -> int:
return self._lo.page_num
def total_pages(self) -> int:
return self._total_pages
def _slice_offset(self, layer: int, global_page: int) -> int:
return layer * self._lo.layer_stride_bytes + global_page * self._lo.slice_bytes
def _check_page(self, global_page: int) -> None:
if global_page < 0 or global_page >= self._lo.page_num:
_l3_failfast(f"global_page {global_page} out of range [0,{self._lo.page_num})")
def _locate(self, global_page: int):
"""Return (mmap, local_page, per_slab_layer_stride_bytes) for a global page."""
if global_page < 0 or global_page >= self._total_pages:
_l3_failfast(f"global_page {global_page} out of range [0,{self._total_pages})")
for s in self._spans:
base = s.global_base_page
if base <= global_page < base + s.num_pages:
return s.mmap, global_page - base, s.num_pages * self._slice_bytes
_l3_failfast(f"global_page {global_page} not covered by slab spans")
def gather_into(self, global_page: int, dst, dst_offset: int = 0) -> int:
"""Copy this page's N_layers strided slices contiguously into ``dst[dst_offset:]``. Returns bytes
written (page_blob_bytes). ``dst`` is typically the aligned disk-slot buffer at HEADER_BYTES."""
self._check_page(global_page)
sb = self._lo.slice_bytes
blob = self._lo.page_blob_bytes
mm, local_page, layer_stride = self._locate(global_page)
sb = self._slice_bytes
blob = self.page_blob_bytes
if len(dst) < dst_offset + blob:
_l3_failfast(f"dst too small: need {dst_offset + blob}, have {len(dst)}")
dview = memoryview(dst)
for layer in range(self._lo.n_layers):
off = self._slice_offset(layer, global_page)
dview[dst_offset + layer * sb : dst_offset + (layer + 1) * sb] = self._mm[off : off + sb]
for layer in range(self._n_layers):
off = layer * layer_stride + local_page * sb
dview[dst_offset + layer * sb : dst_offset + (layer + 1) * sb] = mm[off : off + sb]
return blob
def gather(self, global_page: int) -> bytes:
"""Convenience: return the contiguous page blob as bytes."""
buf = bytearray(self._lo.page_blob_bytes)
buf = bytearray(self.page_blob_bytes)
self.gather_into(global_page, buf, 0)
return bytes(buf)
def scatter_from(self, global_page: int, src, src_offset: int = 0) -> int:
"""Write a contiguous page blob (``src[src_offset:]``) back into this page's N_layers strided slices.
Returns bytes written. Inverse of gather_into (same canonical layer order)."""
self._check_page(global_page)
sb = self._lo.slice_bytes
blob = self._lo.page_blob_bytes
mm, local_page, layer_stride = self._locate(global_page)
sb = self._slice_bytes
blob = self.page_blob_bytes
if len(src) < src_offset + blob:
_l3_failfast(f"src too small: need {src_offset + blob}, have {len(src)}")
sview = memoryview(src)
for layer in range(self._lo.n_layers):
off = self._slice_offset(layer, global_page)
self._mm[off : off + sb] = sview[src_offset + layer * sb : src_offset + (layer + 1) * sb]
for layer in range(self._n_layers):
off = layer * layer_stride + local_page * sb
mm[off : off + sb] = sview[src_offset + layer * sb : src_offset + (layer + 1) * sb]
return blob

103
python/sglang/srt/mem_cache/hiradix_cache.py
View File

@@ -273,22 +273,43 @@ def _cp_shared_l2_slab_pages_by_payload(
page_size: int,
bytes_per_page_by_payload: Dict[str, int],
) -> Dict[str, int]:
from sglang.srt.mem_cache.memory_pool_host import (
cp_hicache_max_single_register_bytes,
)
ceiling = cp_hicache_max_single_register_bytes()
slab_size_gb = int(getattr(server_args, "cp_shared_l2_slab_size_gb", 0) or 0)
if slab_size_gb <= 0:
return {}
slab_size_bytes = slab_size_gb * 1000 * 1000 * 1000
# AUTO: split each payload into slabs no larger than ONE cudaHostRegister
# call can pin. A single >~1 TiB registration OOMs, and a registration
# cannot be chunked (a transfer memcpy cannot cross a registration
# boundary), so large host caches must be physically multi-slab. Payloads
# that already fit in one slab stay single (build_cp_shared_l2_slabs_by_payload
# collapses slab_pages >= total_pages to one slab).
slab_size_bytes = ceiling
else:
slab_size_bytes = slab_size_gb * 1000 * 1000 * 1000
if slab_size_bytes > ceiling:
logger.warning(
"cp_shared_l2_slab_size_gb=%d GB exceeds the safe single-"
"cudaHostRegister ceiling (%.0f GiB); capping to the ceiling.",
slab_size_gb,
ceiling / (1024**3),
)
slab_size_bytes = ceiling
slab_pages: Dict[str, int] = {}
for payload_kind, bytes_per_page in bytes_per_page_by_payload.items():
bytes_per_page = int(bytes_per_page)
if bytes_per_page <= 0:
raise ValueError(f"bytes_per_page for {payload_kind!r} must be positive")
if slab_size_bytes < bytes_per_page:
suggested_gb = math.ceil(bytes_per_page / 1e9)
suggested_gb = math.ceil(bytes_per_page / (1024**3))
raise ValueError(
"cp_shared_l2_slab_size_gb is smaller than one logical page for "
f"payload {payload_kind!r}: configured_bytes={slab_size_bytes} "
f"bytes_per_page={bytes_per_page}. Increase "
f"--cp-shared-l2-slab-size-gb to at least {suggested_gb}."
"CP shared-L2 single-registration slab budget is smaller than one "
f"logical page for payload {payload_kind!r}: slab_bytes={slab_size_bytes} "
f"bytes_per_page={bytes_per_page}. One logical page must fit in one "
f"cudaHostRegister; raise SGLANG_CP_HICACHE_MAX_SLAB_GB to at least "
f"{suggested_gb}."
)
slab_pages[payload_kind] = slab_size_bytes // bytes_per_page
return slab_pages
@@ -1240,35 +1261,48 @@ class HiRadixCache(RadixCache):
cp_rank, cp_size = int(cp_allocator.cp_rank), int(cp_allocator.cp_size)
cfg = CpL3Config.from_file(server_args.cp_l3_config)
def _build_accessor(allocator, layout, payload_kind):
# One accessor over N physical slabs (N==1 for a single slab); the
# layout supplies n_layers + slice_bytes (slab-count invariant), the
# spans supply each slab's global page range + its own mmap.
return CpSharedL2SlabAccessor(
self._cp_l3_slab_spans(allocator, layout.page_num, payload_kind),
n_layers=layout.n_layers,
slice_bytes=layout.slice_bytes,
)
host = self.token_to_kv_pool_host
accessors = {
"target_kv": CpSharedL2SlabAccessor(
self._cp_l3_single_slab_mmap(host.allocator, "target_kv"),
"target_kv": _build_accessor(
host.allocator,
CpL3SlabLayout.for_mla(
layer_num=host.layer_num, page_num=host.page_num,
page_size=self.page_size, kv_cache_dim=host.kv_cache_dim,
itemsize=host.dtype.itemsize,
),
"target_kv",
)
}
draft = getattr(self, "draft_token_to_kv_pool_host", None)
if draft is not None:
accessors["draft_kv"] = CpSharedL2SlabAccessor(
self._cp_l3_single_slab_mmap(draft.allocator, "draft_kv"),
accessors["draft_kv"] = _build_accessor(
draft.allocator,
CpL3SlabLayout.for_mla(
layer_num=draft.layer_num, page_num=draft.page_num,
page_size=self.page_size, kv_cache_dim=draft.kv_cache_dim,
itemsize=draft.dtype.itemsize,
),
"draft_kv",
)
if isinstance(host, NSATokenToKVPoolHost):
accessors["index_k"] = CpSharedL2SlabAccessor(
self._cp_l3_single_slab_mmap(host.index_host_tensor_allocator, "index_k"),
accessors["index_k"] = _build_accessor(
host.index_host_tensor_allocator,
CpL3SlabLayout.for_index(
n_active_layers=host.index_active_layer_num,
indexer_page_num=host.indexer_page_num,
indexer_page_stride_size=host.indexer_page_stride_size,
),
"index_k",
)
self.cp_l3_store = CpL3Store.from_config(
cfg, cp_rank=cp_rank, cp_size=cp_size, accessors=accessors
@@ -1282,16 +1316,45 @@ class HiRadixCache(RadixCache):
# cover PURE-L2-without-L3 too, not just this L3 path.)
@staticmethod
def _cp_l3_single_slab_mmap(allocator, payload_kind: str):
"""Reach the single shared-L2 slab mmap (default --cp-shared-l2-slab-size-gb 0). Fail loud on the
multi-slab group case (a future extension; needs per-slab accessors + base_page->slab conversion)."""
def _cp_l3_slab_spans(allocator, total_pages: int, payload_kind: str):
"""Per-slab page ranges + mmaps for an L3 accessor.
Handles both the single-slab allocator (one span over the whole payload)
and the multi-slab group allocator (one span per physical slab, carrying
each slab's global_base_page/num_pages from its slab_info). L3 addresses
host pages by *global* page id, so a multi-slab payload must dispatch each
page to its owning slab -- which these spans + CpSharedL2SlabAccessor do.
"""
from sglang.srt.mem_cache.cp_l3_slab_accessor import CpL3SlabSpan
from sglang.srt.mem_cache.memory_pool_host import (
SharedHostTensorGroupAllocator,
)
if isinstance(allocator, SharedHostTensorGroupAllocator):
spans = []
for slab_info, sub in zip(allocator.slabs, allocator.allocators):
mapping = getattr(sub, "mapping", None)
if mapping is None or getattr(mapping, "mmap", None) is None:
raise RuntimeError(
f"[CP_L3_FAILFAST] {payload_kind}: group slab "
f"{int(slab_info.slab_id)} has no live mapping for L3 access."
)
spans.append(
CpL3SlabSpan(
int(slab_info.global_base_page),
int(slab_info.num_pages),
mapping.mmap,
)
)
return spans
mapping = getattr(allocator, "mapping", None)
if mapping is None or getattr(mapping, "mmap", None) is None:
raise NotImplementedError(
f"[CP_L3_FAILFAST] {payload_kind}: L3 requires a single shared-L2 slab "
f"(--cp-shared-l2-slab-size-gb 0); multi-slab group accessors are not yet supported."
raise RuntimeError(
f"[CP_L3_FAILFAST] {payload_kind}: shared-L2 slab has no live "
f"mapping for L3 access."
)
return mapping.mmap
return [CpL3SlabSpan(0, int(total_pages), mapping.mmap)]
def _drain_l3_control_queues(self) -> None:
"""Per-tick CP-cpu-group MIN over the three L3 ack-queue sizes [gather, durable, reload], then drain

54
python/sglang/srt/mem_cache/memory_pool_host.py
View File

@@ -2,6 +2,7 @@ import abc
import bisect
import heapq
import logging
import os
import threading
import weakref
from collections import defaultdict
@@ -860,6 +861,8 @@ class SharedHostTensorGroupAllocator:
tensor = view.tensor
ptr = tensor.data_ptr()
nbytes = int(tensor.numel()) * int(tensor.element_size())
# Each slab is capped at the per-call ceiling upstream; guard anyway.
_check_single_cuda_host_register_size(nbytes)
registered = False
register_error = None
try:
@@ -927,6 +930,10 @@ def alloc_with_host_register(
if pin_memory:
ptr = buffer.data_ptr()
nbytes = buffer.numel() * buffer.element_size()
# A single cudaHostRegister cannot exceed the per-call ceiling, and it
# cannot be chunked (memcpy can't cross a registration boundary) -- large
# host caches must be split into multiple slabs upstream.
_check_single_cuda_host_register_size(nbytes)
registered = False
register_error = None
try:
@@ -982,6 +989,53 @@ def _cuda_host_unregister(ptr: int) -> None:
pass
def cp_hicache_max_single_register_bytes() -> int:
"""Largest host region we will hand to ONE ``cudaHostRegister`` call.
A single ``cudaHostRegister`` over a very large region fails with
``cudaErrorMemoryAllocation`` ("out of memory") above a hard per-call size
ceiling (empirically between 512 GiB and 1 TiB on B300, driver 590.48.01),
even with abundant host RAM and hugepages. The limit is per *call*, but a
single contiguous registration cannot be chunked either: a ``cudaMemcpy``
(incl. ``cudaMemcpyBatchAsync``, used by the CP-L2 H2D/D2H transfers) whose
host range straddles a registration boundary fails with
``cudaErrorInvalidValue``. So the host KV cache is physically split into
multiple slabs each <= this size (each one whole registration; transfers
split per slab and never cross a boundary). Default 480 GiB stays under the
512-GiB-proven-OK ceiling while keeping ``hicache_size<=400`` slabs single.
Override via ``SGLANG_CP_HICACHE_MAX_SLAB_GB``.
"""
raw = os.environ.get("SGLANG_CP_HICACHE_MAX_SLAB_GB", "").strip()
if raw:
try:
gb = float(raw)
except ValueError:
gb = 0.0
if gb > 0:
return int(gb * (1024**3))
return 480 * (1024**3)
def _check_single_cuda_host_register_size(nbytes: int) -> None:
"""Fail loud if a single registration exceeds the per-call ceiling.
CP shared-L2 never trips this (slabs are auto-capped at the ceiling in
``hiradix_cache._cp_shared_l2_slab_pages_by_payload``); a non-CP HiCache host
buffer larger than the ceiling gets an actionable error instead of the
cryptic ``cudaErrorMemoryAllocation``.
"""
ceiling = cp_hicache_max_single_register_bytes()
if int(nbytes) > ceiling:
raise RuntimeError(
"[CP_HICACHE_FAILFAST][host_register_too_large] single cudaHostRegister of "
f"{int(nbytes) / 1024**3:.1f} GiB exceeds the per-call ceiling "
f"({ceiling / 1024**3:.0f} GiB). A single cudaHostRegister over a >~1 TiB region "
"fails on B300, and a memcpy cannot cross a registration boundary. For CP shared-L2 "
"this should not happen (slabs are auto-capped); for non-CP HiCache reduce "
"--hicache-size. Override the ceiling via SGLANG_CP_HICACHE_MAX_SLAB_GB."
)
def alloc_with_pin_memory(
dims,
dtype: torch.dtype,