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) <noreply@anthropic.com>
This commit is contained in:
Claude Opus 4.8
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user