From dbc5ebbaa0968f0fe6ef4e86d5ad8ea61770cce2 Mon Sep 17 00:00:00 2001 From: leavelet Date: Tue, 23 Jun 2026 15:33:10 +0000 Subject: [PATCH] CP HiCache: multi-slab host KV cache to respect cudaHostRegister per-call ceiling 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) --- .../srt/mem_cache/cp_l3_slab_accessor.py | 125 ++++++++++---- python/sglang/srt/mem_cache/hiradix_cache.py | 103 +++++++++--- .../sglang/srt/mem_cache/memory_pool_host.py | 54 ++++++ .../mem_cache/test_cp_l3_slab_accessor.py | 107 +++++++++++- .../unit/mem_cache/test_cp_l3_store.py | 67 +++++++- .../mem_cache/test_cp_shared_l2_slab_cap.py | 159 ++++++++++++++++++ 6 files changed, 555 insertions(+), 60 deletions(-) create mode 100644 test/registered/unit/mem_cache/test_cp_shared_l2_slab_cap.py diff --git a/python/sglang/srt/mem_cache/cp_l3_slab_accessor.py b/python/sglang/srt/mem_cache/cp_l3_slab_accessor.py index 51b7c53bd..1cf43936c 100644 --- a/python/sglang/srt/mem_cache/cp_l3_slab_accessor.py +++ b/python/sglang/srt/mem_cache/cp_l3_slab_accessor.py @@ -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 diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 313e3e8b6..bf02baa8a 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -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 diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index ba4362a9d..65e76fd50 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -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, diff --git a/test/registered/unit/mem_cache/test_cp_l3_slab_accessor.py b/test/registered/unit/mem_cache/test_cp_l3_slab_accessor.py index d868c72ea..dc3ac87c3 100644 --- a/test/registered/unit/mem_cache/test_cp_l3_slab_accessor.py +++ b/test/registered/unit/mem_cache/test_cp_l3_slab_accessor.py @@ -7,6 +7,8 @@ import types import unittest from pathlib import Path +import torch + _REPO_ROOT = Path(__file__).resolve().parents[4] _MEM = _REPO_ROOT / "python" / "sglang" / "srt" / "mem_cache" @@ -52,10 +54,28 @@ def _make_slab(): return mm, lo +def _make_multislab(page_counts, *, fill=True): + """Build len(page_counts) slabs tiling [0, sum(page_counts)) in layer_page_first + layout (per-slab layer stride = num_pages * SLICE), returning CpL3SlabSpan list.""" + spans = [] + base = 0 + for num in page_counts: + mm = mmap.mmap(-1, N_LAYERS * num * SLICE) + if fill: + per_slab_layer_stride = num * SLICE + for layer in range(N_LAYERS): + for local in range(num): + off = layer * per_slab_layer_stride + local * SLICE + mm[off:off + SLICE] = _slice_pattern(layer, base + local) + spans.append(a.CpL3SlabSpan(base, num, mm)) + base += num + return spans + + class TestAccessor(unittest.TestCase): def test_gather_concatenates_layers_in_order(self): mm, lo = _make_slab() - acc = a.CpSharedL2SlabAccessor(mm, lo) + acc = a.CpSharedL2SlabAccessor.from_single_mmap(mm, lo) self.assertEqual(acc.page_blob_bytes, N_LAYERS * SLICE) for page in range(PAGE_NUM): expect = b"".join(_slice_pattern(layer, page) for layer in range(N_LAYERS)) @@ -63,7 +83,7 @@ class TestAccessor(unittest.TestCase): def test_gather_into_with_offset(self): mm, lo = _make_slab() - acc = a.CpSharedL2SlabAccessor(mm, lo) + acc = a.CpSharedL2SlabAccessor.from_single_mmap(mm, lo) dst = bytearray(64 + lo.page_blob_bytes) n = acc.gather_into(1, dst, 64) self.assertEqual(n, lo.page_blob_bytes) @@ -73,9 +93,9 @@ class TestAccessor(unittest.TestCase): def test_gather_scatter_roundtrip(self): src_mm, lo = _make_slab() - acc_src = a.CpSharedL2SlabAccessor(src_mm, lo) + acc_src = a.CpSharedL2SlabAccessor.from_single_mmap(src_mm, lo) dst_mm = mmap.mmap(-1, lo.total_bytes) - acc_dst = a.CpSharedL2SlabAccessor(dst_mm, lo) + acc_dst = a.CpSharedL2SlabAccessor.from_single_mmap(dst_mm, lo) for page in range(PAGE_NUM): blob = acc_src.gather(page) acc_dst.scatter_from(page, blob, 0) @@ -86,16 +106,16 @@ class TestAccessor(unittest.TestCase): def test_scatter_from_offset(self): src_mm, lo = _make_slab() - acc_src = a.CpSharedL2SlabAccessor(src_mm, lo) + acc_src = a.CpSharedL2SlabAccessor.from_single_mmap(src_mm, lo) dst_mm = mmap.mmap(-1, lo.total_bytes) - acc_dst = a.CpSharedL2SlabAccessor(dst_mm, lo) + acc_dst = a.CpSharedL2SlabAccessor.from_single_mmap(dst_mm, lo) framed = bytearray(64) + bytearray(acc_src.gather(2)) # simulate header+payload slot buffer acc_dst.scatter_from(2, framed, 64) self.assertEqual(acc_dst.gather(2), acc_src.gather(2)) def test_bounds_fail_loud(self): mm, lo = _make_slab() - acc = a.CpSharedL2SlabAccessor(mm, lo) + acc = a.CpSharedL2SlabAccessor.from_single_mmap(mm, lo) with self.assertRaises(ValueError): acc.gather(PAGE_NUM) # page out of range with self.assertRaises(ValueError): @@ -107,7 +127,78 @@ class TestAccessor(unittest.TestCase): lo = a.CpL3SlabLayout(n_layers=N_LAYERS, page_num=PAGE_NUM, slice_bytes=SLICE) small = mmap.mmap(-1, lo.total_bytes - 16) with self.assertRaises(ValueError): - a.CpSharedL2SlabAccessor(small, lo) + a.CpSharedL2SlabAccessor.from_single_mmap(small, lo) + + # --- multi-slab (host cache split across N physical slabs) --- + + def test_multislab_gather_dispatches_per_slab(self): + spans = _make_multislab([2, 2]) # pages 0-1 in slab0, pages 2-3 in slab1 (global_base 2) + acc = a.CpSharedL2SlabAccessor(spans, n_layers=N_LAYERS, slice_bytes=SLICE) + self.assertEqual(acc.total_pages, PAGE_NUM) + for page in range(PAGE_NUM): + expect = b"".join(_slice_pattern(layer, page) for layer in range(N_LAYERS)) + self.assertEqual(acc.gather(page), expect) # slab1 pages addressed via global_base + per-slab stride + + def test_multislab_roundtrip_uneven_split(self): + src = _make_multislab([1, 3]) # slab1 base=1, per-slab layer strides differ (1*SLICE vs 3*SLICE) + dst = _make_multislab([1, 3], fill=False) + acc_src = a.CpSharedL2SlabAccessor(src, n_layers=N_LAYERS, slice_bytes=SLICE) + acc_dst = a.CpSharedL2SlabAccessor(dst, n_layers=N_LAYERS, slice_bytes=SLICE) + for page in range(PAGE_NUM): + acc_dst.scatter_from(page, acc_src.gather(page), 0) + for page in range(PAGE_NUM): + self.assertEqual(acc_dst.gather(page), acc_src.gather(page)) + + def test_multislab_matches_single_slab_bytes(self): + # A page in slab1 must gather identically whether the payload is one slab or two. + single_mm, lo = _make_slab() + acc_single = a.CpSharedL2SlabAccessor.from_single_mmap(single_mm, lo) + acc_multi = a.CpSharedL2SlabAccessor( + _make_multislab([2, 2]), n_layers=N_LAYERS, slice_bytes=SLICE + ) + for page in range(PAGE_NUM): + self.assertEqual(acc_multi.gather(page), acc_single.gather(page)) + + def test_multislab_noncontiguous_spans_fail(self): + mm0 = mmap.mmap(-1, N_LAYERS * 2 * SLICE) + mm1 = mmap.mmap(-1, N_LAYERS * 2 * SLICE) + with self.assertRaises(ValueError): # gap: slab1 base 3 != 0+2 + a.CpSharedL2SlabAccessor( + [a.CpL3SlabSpan(0, 2, mm0), a.CpL3SlabSpan(3, 2, mm1)], + n_layers=N_LAYERS, + slice_bytes=SLICE, + ) + + def test_multislab_slab_mmap_too_small_fails(self): + spans = _make_multislab([2, 2]) + spans[1] = a.CpL3SlabSpan(2, 2, mmap.mmap(-1, N_LAYERS * 2 * SLICE - 8)) + with self.assertRaises(ValueError): + a.CpSharedL2SlabAccessor(spans, n_layers=N_LAYERS, slice_bytes=SLICE) + + def test_accessor_stride_matches_real_torch_frombuffer_view(self): + """Non-circular layout check: write each slab's pages through a real + torch.frombuffer(mmap).view(n_layers, num_pages, slice) -- exactly how the + production per-slab tensor is C-contiguously laid out (page_dim sized to the + SLAB's num_pages) -- and read them back through the accessor. A bug where + the accessor used the TOTAL page_num (instead of per-slab num_pages) for the + layer stride would read the wrong bytes here and fail.""" + page_counts = [2, 3] # uneven -> per-slab layer strides differ (2*SLICE vs 3*SLICE) + spans, base = [], 0 + for num in page_counts: + mm = mmap.mmap(-1, N_LAYERS * num * SLICE) + # Write via the production-style C-contiguous view, NOT the accessor's formula. + t = torch.frombuffer(mm, dtype=torch.uint8).view(N_LAYERS, num, SLICE) + for layer in range(N_LAYERS): + for local in range(num): + t[layer, local] = torch.frombuffer( + bytearray(_slice_pattern(layer, base + local)), dtype=torch.uint8 + ) + spans.append(a.CpL3SlabSpan(base, num, mm)) + base += num + acc = a.CpSharedL2SlabAccessor(spans, n_layers=N_LAYERS, slice_bytes=SLICE) + for page in range(sum(page_counts)): + expect = b"".join(_slice_pattern(layer, page) for layer in range(N_LAYERS)) + self.assertEqual(acc.gather(page), expect) def test_factory_mla_dims(self): lo = a.CpL3SlabLayout.for_mla(layer_num=78, page_num=100, page_size=64, kv_cache_dim=576, itemsize=1) diff --git a/test/registered/unit/mem_cache/test_cp_l3_store.py b/test/registered/unit/mem_cache/test_cp_l3_store.py index c821d7354..21d6d0757 100644 --- a/test/registered/unit/mem_cache/test_cp_l3_store.py +++ b/test/registered/unit/mem_cache/test_cp_l3_store.py @@ -65,7 +65,24 @@ def _make_slab_and_accessor(): for page in range(PAGE_NUM): off = layer * lo.layer_stride_bytes + page * SLICE mm[off:off + SLICE] = _pattern(layer, page) - return mm, acc_mod.CpSharedL2SlabAccessor(mm, lo) + return mm, acc_mod.CpSharedL2SlabAccessor.from_single_mmap(mm, lo) + + +def _make_multislab_accessor(page_counts): + """Build a multi-slab target_kv accessor tiling [0, sum(page_counts)) -- each slab + its own mmap with per-slab layer stride. Returns ([(base,num,mm)...], accessor).""" + spans, slabs, base = [], [], 0 + for num in page_counts: + mm = mmap.mmap(-1, N_LAYERS * num * SLICE) + stride = num * SLICE + for layer in range(N_LAYERS): + for local in range(num): + off = layer * stride + local * SLICE + mm[off:off + SLICE] = _pattern(layer, base + local) + spans.append(acc_mod.CpL3SlabSpan(base, num, mm)) + slabs.append((base, num, mm)) + base += num + return slabs, acc_mod.CpSharedL2SlabAccessor(spans, n_layers=N_LAYERS, slice_bytes=SLICE) def _wait_ack(qsize_fn, n=1, timeout=5.0): @@ -301,5 +318,53 @@ class TestCpL3Store(unittest.TestCase): self.assertFalse(self.store.has_inflight()) +class TestCpL3StoreMultiSlab(unittest.TestCase): + """The store must work when the host cache is split across multiple physical slabs: + the accessor dispatches each global page to its owning slab (per-slab layer stride).""" + + def setUp(self): + self._td = tempfile.TemporaryDirectory() + self.slabs, self.acc = _make_multislab_accessor([4, 4]) # pages 0-3 | 4-7 + cfg = cfg_mod.CpL3Config.from_dict({ + "backend": "posix", + "require_plp": False, + "index_map_gb": 0.05, + "disks": [{"path": os.path.join(self._td.name, "disk0"), "budget_gb": 0.02}], + }) + self.store = store_mod.CpL3Store.from_config( + cfg, cp_rank=0, cp_size=1, accessors={"target_kv": self.acc}) + self.store.connect(cfg) + + def tearDown(self): + self.store.close() + self._td.cleanup() + + def _zero_page(self, global_page): + for base, num, mm in self.slabs: + if base <= global_page < base + num: + local, stride = global_page - base, num * SLICE + for layer in range(N_LAYERS): + off = layer * stride + local * SLICE + mm[off:off + SLICE] = bytes(SLICE) + return + + def test_spill_reload_roundtrip_page_in_second_slab(self): + page = 5 # slab1: global_base 4, local 1 + orig = self.acc.gather(page) + pages = {"target_kv": [(page, _h(page))]} + op_id = self.store.submit_spill("obj-ms", pages, last_access=1) + self.assertTrue(_wait_ack(self.store.ack_durable_qsize, 1)) + self.store.drain_gather_acks(self.store.ack_gather_qsize()) + self.assertEqual( + self.store.drain_durable_acks(self.store.ack_durable_qsize()), [(op_id, True)] + ) + self._zero_page(page) + self.assertNotEqual(self.acc.gather(page), orig) + self.store.submit_reload("obj-ms", pages) + self.assertTrue(_wait_ack(self.store.ack_reload_qsize, 1)) + self.store.drain_reload_acks(self.store.ack_reload_qsize()) + self.assertEqual(self.acc.gather(page), orig) # byte-exact reload into slab1 + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_cp_shared_l2_slab_cap.py b/test/registered/unit/mem_cache/test_cp_shared_l2_slab_cap.py new file mode 100644 index 000000000..9bbc9648a --- /dev/null +++ b/test/registered/unit/mem_cache/test_cp_shared_l2_slab_cap.py @@ -0,0 +1,159 @@ +"""Unit tests for the CP shared-L2 per-payload slab-size cap. + +A single cudaHostRegister cannot exceed a per-call ceiling (and cannot be +chunked), so large host caches are split into multiple slabs each <= the +ceiling. `_cp_shared_l2_slab_pages_by_payload` derives the per-payload slab page +count; this verifies the auto-cap, env override, explicit-size capping, and the +one-page-too-big fail-loud. +""" +import os +import tempfile +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.srt.mem_cache.cp_l3_slab_accessor import CpSharedL2SlabAccessor +from sglang.srt.mem_cache.cp_shared_l2_pool import ( + PAYLOAD_TARGET_KV, + CpSharedL2SlabInfo, + build_cp_shared_l2_slabs_by_payload, +) +from sglang.srt.mem_cache.hiradix_cache import ( + HiRadixCache, + _cp_shared_l2_slab_pages_by_payload, +) +from sglang.srt.mem_cache.memory_pool_host import ( + SharedHostTensorAllocator, + SharedHostTensorGroupAllocator, + cp_hicache_max_single_register_bytes, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +# hiradix_cache imports sgl_kernel (needs libcuda to import); no GPU is used. +register_cuda_ci(est_time=1, suite="stage-b-test-1-gpu-small") + +GiB = 1024**3 + + +def _args(slab_size_gb=0): + return SimpleNamespace(cp_shared_l2_slab_size_gb=slab_size_gb) + + +def _slab_pages(server_args, bytes_per_page_by_payload): + return _cp_shared_l2_slab_pages_by_payload( + server_args=server_args, + page_size=64, + bytes_per_page_by_payload=bytes_per_page_by_payload, + ) + + +class TestCpSharedL2SlabCap(CustomTestCase): + def test_auto_caps_each_payload_at_ceiling(self): + ceiling = cp_hicache_max_single_register_bytes() + out = _slab_pages(_args(0), {"target_kv": 10 * GiB, "index_k": 4 * GiB}) + self.assertEqual(out["target_kv"], ceiling // (10 * GiB)) + self.assertEqual(out["index_k"], ceiling // (4 * GiB)) + # Each slab's registration bytes stay <= the ceiling. + self.assertLessEqual(out["target_kv"] * 10 * GiB, ceiling) + + def test_env_override_changes_cap(self): + with patch.dict(os.environ, {"SGLANG_CP_HICACHE_MAX_SLAB_GB": "256"}): + out = _slab_pages(_args(0), {"target_kv": 8 * GiB}) + self.assertEqual(out["target_kv"], (256 * GiB) // (8 * GiB)) # 32 + + def test_explicit_slab_size_capped_at_ceiling(self): + ceiling = cp_hicache_max_single_register_bytes() + # 4000 GB >> ceiling -> capped down to the ceiling. + out = _slab_pages(_args(4000), {"target_kv": 4 * GiB}) + self.assertEqual(out["target_kv"], ceiling // (4 * GiB)) + + def test_explicit_small_slab_size_respected(self): + # 100 GB < ceiling (~515 GB) -> honoured as-is. + out = _slab_pages(_args(100), {"target_kv": 1 * GiB}) + self.assertEqual(out["target_kv"], (100 * 1000**3) // (1 * GiB)) + + def test_page_larger_than_budget_fails_loud(self): + with patch.dict(os.environ, {"SGLANG_CP_HICACHE_MAX_SLAB_GB": "1"}): # 1 GiB budget + with self.assertRaises(ValueError): + _slab_pages(_args(0), {"target_kv": 2 * GiB}) # one page > budget + + def test_auto_cap_yields_multiple_slabs_for_large_payload(self): + """B3: a realistic payload + the auto-cap must build >1 page-aligned slab, + each <= ceiling, tiling [0, total) contiguously (so the group allocator path + and per-slab transfer splitting actually engage).""" + ceiling = cp_hicache_max_single_register_bytes() + bpp = 1 * GiB # 1 GiB per logical page + cap_pages = ceiling // bpp + total_pages = cap_pages * 3 + 5 # spans ~3.x slabs + slab_pages = _slab_pages(_args(0), {PAYLOAD_TARGET_KV: bpp}) # auto-cap + slabs = build_cp_shared_l2_slabs_by_payload( + {PAYLOAD_TARGET_KV: total_pages}, slab_pages + )[PAYLOAD_TARGET_KV] + self.assertGreater(len(slabs), 1) + base = 0 + for s in slabs: + self.assertLessEqual(s.num_pages * bpp, ceiling) # each slab <= one registration + self.assertEqual(s.global_base_page, base) # contiguous tiling + base += s.num_pages + self.assertEqual(base, total_pages) # exact + + +class TestCpL3SlabSpansWiring(CustomTestCase): + """B2: _cp_l3_slab_spans (the production wiring) + the accessor read against a + REAL SharedHostTensorGroupAllocator allocate_group tensor (not a hand-filled toy).""" + + def test_group_spans_and_accessor_roundtrip(self): + n_layers, trailing = 3, 2 # int16 -> slice_bytes = trailing * 2 = 4 + slice_bytes = trailing * 2 + infos = ( + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=0, global_base_page=0, num_pages=3), + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=1, global_base_page=3, num_pages=2), + ) + with tempfile.TemporaryDirectory() as tmp: + allocator = SharedHostTensorGroupAllocator( + slabs=infos, directory=tmp, name="target-kv", creator_rank=0, + validate_production=False, + ) + group = allocator.allocate_group( + (n_layers, 5, trailing), dtype=torch.int16, device="cpu", page_dim=1 + ) + # Write a distinct value per (layer, global_page) via the REAL per-slab tensors. + for view in group.slab_views: + b = int(view.slab_info.global_base_page) + for layer in range(n_layers): + for local in range(int(view.slab_info.num_pages)): + view.tensor[layer, local] = layer * 100 + (b + local) + # Production wiring derives spans from the live allocator: + spans = HiRadixCache._cp_l3_slab_spans(allocator, 5, "target_kv") + self.assertEqual( + [(s.global_base_page, s.num_pages) for s in spans], [(0, 3), (3, 2)] + ) + acc = CpSharedL2SlabAccessor(spans, n_layers=n_layers, slice_bytes=slice_bytes) + for gp in range(5): + vals = torch.frombuffer( + bytearray(acc.gather(gp)), dtype=torch.int16 + ).view(n_layers, trailing) + for layer in range(n_layers): + self.assertTrue(bool((vals[layer] == layer * 100 + gp).all())) + group.release_tensor_view_for_close() + group.close() + + def test_single_slab_spans_branch(self): + with tempfile.TemporaryDirectory() as tmp: + alloc = SharedHostTensorAllocator( + directory=tmp, name="single", creator_rank=0, validate_production=False + ) + alloc.allocate((3, 4, 2), dtype=torch.int16, device="cpu") + spans = HiRadixCache._cp_l3_slab_spans(alloc, 4, "target_kv") + self.assertEqual(len(spans), 1) + self.assertEqual((spans[0].global_base_page, spans[0].num_pages), (0, 4)) + self.assertIs(spans[0].mmap, alloc.mapping.mmap) + alloc.release_tensor_view_for_close() + alloc.close() + + +if __name__ == "__main__": + unittest.main()