L3 3.0: host-slab accessor (layer_page_first <-> contiguous blob)

CpSharedL2SlabAccessor (HIGH-1): memcpy-gather a page's N_layers strided slab slices into a
contiguous buffer (spill) + scatter a contiguous blob back into the strided slices (reload) —
the cheap RAM bridge that lets host stay layer_page_first (inference perf) while disk is
page-contiguous. Generic CpL3SlabLayout(n_layers, page_num, slice_bytes) with verified-from-source
factories: for_mla (layer_num,page_num,page_size,kv_cache_dim,itemsize -> 2.74 MiB/page) and
for_index (n_active_layers,indexer_page_num,indexer_page_stride -> 0.17/0.63 MiB/page). gather_into/
scatter_from take an offset so the spill writes straight after the 64B slot header (zero extra copy).
Pure (mmap+ints), no torch. 8/8 tests. Canonical layer 0..N-1 order shared by spill+reload.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

python/sglang/srt/mem_cache/cp_l3_slab_accessor.py

@@ -0,0 +1,115 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0
+"""CP HiCache L3 — host-slab byte accessor (the layer_page_first <-> contiguous-blob bridge).
+
+The CP shared-L2 host slab is ``layer_page_first`` (kept for inference perf): a page's bytes are N_layers
+*strided* slices, not contiguous. L3 wants one contiguous page blob on disk. This accessor does the
+**memcpy-gather** (spill: strided slab slices -> contiguous buffer) and **scatter** (reload: contiguous blob
+-> strided slab slices) — measured faster than vectored I/O under O_DIRECT, and the cheap RAM side of the
+spill/reload. The on-disk slice order (layer 0..N-1) is canonical: spill and reload use the same order.
+
+Verified layouts (memory_pool_host.py, GLM + layer_page_first):
+  * target_kv / draft_kv (MLA): slab tensor ``(layer_num, page_num, page_size, 1, kv_cache_dim)`` ->
+    slice(layer,page) at ``layer*(page_num*slice) + page*slice``, slice = ``page_size*kv_cache_dim*itemsize``.
+  * index_k (NSA):            ``(n_active_layers, indexer_page_num, 1, indexer_page_stride)`` ->
+    slice(slot,page) at ``slot*(indexer_page_num*stride) + page*stride``, slice = ``indexer_page_stride``.
+Both reduce to a generic ``(n_layers, layer_stride_bytes, slice_bytes)``. Pure (mmap + ints): no torch.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from sglang.srt.mem_cache.cp_l3_disk import _l3_failfast
+
+
+@dataclass(frozen=True)
+class CpL3SlabLayout:
+    n_layers: int            # layer_num (target/draft) or index_active_layer_num (index_k)
+    page_num: int            # slab pages along the page axis (for bounds + layer stride)
+    slice_bytes: int         # bytes of one (layer, page) contiguous slice
+
+    def __post_init__(self):
+        if self.n_layers <= 0 or self.page_num <= 0 or self.slice_bytes <= 0:
+            _l3_failfast(
+                f"CpL3SlabLayout requires positive dims; got n_layers={self.n_layers} "
+                f"page_num={self.page_num} slice_bytes={self.slice_bytes}"
+            )
+
+    @property
+    def layer_stride_bytes(self) -> int:
+        return self.page_num * self.slice_bytes
+
+    @property
+    def page_blob_bytes(self) -> int:
+        return self.n_layers * self.slice_bytes
+
+    @property
+    def total_bytes(self) -> int:
+        return self.n_layers * self.layer_stride_bytes
+
+    @classmethod
+    def for_mla(cls, *, layer_num: int, page_num: int, page_size: int, kv_cache_dim: int, itemsize: int):
+        return cls(n_layers=layer_num, page_num=page_num, slice_bytes=page_size * kv_cache_dim * itemsize)
+
+    @classmethod
+    def for_index(cls, *, n_active_layers: int, indexer_page_num: int, indexer_page_stride_size: int):
+        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."""
+
+    def __init__(self, slab_mmap, layout: CpL3SlabLayout):
+        self._mm = slab_mmap
+        self._lo = layout
+        if len(slab_mmap) < layout.total_bytes:
+            _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})"
+            )
+
+    @property
+    def page_blob_bytes(self) -> int:
+        return self._lo.page_blob_bytes
+
+    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 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
+        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]
+        return blob
+
+    def gather(self, global_page: int) -> bytes:
+        """Convenience: return the contiguous page blob as bytes."""
+        buf = bytearray(self._lo.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
+        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]
+        return blob

test/registered/unit/mem_cache/test_cp_l3_slab_accessor.py

@@ -0,0 +1,125 @@
+"""Unit tests for CP HiCache L3 slab accessor (cp_l3_slab_accessor): layer_page_first gather/scatter."""
+
+import importlib.util
+import mmap
+import sys
+import types
+import unittest
+from pathlib import Path
+
+_REPO_ROOT = Path(__file__).resolve().parents[4]
+_MEM = _REPO_ROOT / "python" / "sglang" / "srt" / "mem_cache"
+
+
+def _load_file(module_name, path):
+    spec = importlib.util.spec_from_file_location(module_name, path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_accessor():
+    for name, pkg_path in (
+        ("sglang", _REPO_ROOT / "python" / "sglang"),
+        ("sglang.srt", _REPO_ROOT / "python" / "sglang" / "srt"),
+        ("sglang.srt.mem_cache", _MEM),
+    ):
+        if name not in sys.modules:
+            mod = types.ModuleType(name)
+            mod.__path__ = [str(pkg_path)]
+            sys.modules[name] = mod
+    _load_file("sglang.srt.mem_cache.cp_l3_disk", _MEM / "cp_l3_disk.py")
+    return _load_file("sglang.srt.mem_cache.cp_l3_slab_accessor", _MEM / "cp_l3_slab_accessor.py")
+
+
+a = _load_accessor()
+
+N_LAYERS, PAGE_NUM, SLICE = 3, 4, 16
+
+
+def _slice_pattern(layer, page):
+    return bytes(((layer * 37 + page * 7 + k) & 0xFF) for k in range(SLICE))
+
+
+def _make_slab():
+    lo = a.CpL3SlabLayout(n_layers=N_LAYERS, page_num=PAGE_NUM, slice_bytes=SLICE)
+    mm = mmap.mmap(-1, lo.total_bytes)
+    for layer in range(N_LAYERS):
+        for page in range(PAGE_NUM):
+            off = layer * lo.layer_stride_bytes + page * SLICE
+            mm[off:off + SLICE] = _slice_pattern(layer, page)
+    return mm, lo
+
+
+class TestAccessor(unittest.TestCase):
+    def test_gather_concatenates_layers_in_order(self):
+        mm, lo = _make_slab()
+        acc = a.CpSharedL2SlabAccessor(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))
+            self.assertEqual(acc.gather(page), expect)
+
+    def test_gather_into_with_offset(self):
+        mm, lo = _make_slab()
+        acc = a.CpSharedL2SlabAccessor(mm, lo)
+        dst = bytearray(64 + lo.page_blob_bytes)
+        n = acc.gather_into(1, dst, 64)
+        self.assertEqual(n, lo.page_blob_bytes)
+        expect = b"".join(_slice_pattern(layer, 1) for layer in range(N_LAYERS))
+        self.assertEqual(bytes(dst[64:64 + lo.page_blob_bytes]), expect)
+        self.assertEqual(bytes(dst[:64]), bytes(64))  # header region untouched
+
+    def test_gather_scatter_roundtrip(self):
+        src_mm, lo = _make_slab()
+        acc_src = a.CpSharedL2SlabAccessor(src_mm, lo)
+        dst_mm = mmap.mmap(-1, lo.total_bytes)
+        acc_dst = a.CpSharedL2SlabAccessor(dst_mm, lo)
+        for page in range(PAGE_NUM):
+            blob = acc_src.gather(page)
+            acc_dst.scatter_from(page, blob, 0)
+        # dst now byte-identical to src for every page's slices
+        for page in range(PAGE_NUM):
+            self.assertEqual(acc_dst.gather(page), acc_src.gather(page))
+        self.assertEqual(bytes(dst_mm[:]), bytes(src_mm[:]))
+
+    def test_scatter_from_offset(self):
+        src_mm, lo = _make_slab()
+        acc_src = a.CpSharedL2SlabAccessor(src_mm, lo)
+        dst_mm = mmap.mmap(-1, lo.total_bytes)
+        acc_dst = a.CpSharedL2SlabAccessor(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)
+        with self.assertRaises(ValueError):
+            acc.gather(PAGE_NUM)            # page out of range
+        with self.assertRaises(ValueError):
+            acc.gather_into(0, bytearray(4))  # dst too small
+        with self.assertRaises(ValueError):
+            acc.scatter_from(0, b"\x00" * 4)  # src too small
+
+    def test_mmap_too_small_fails(self):
+        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)
+
+    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)
+        self.assertEqual(lo.slice_bytes, 64 * 576)                 # 36864 B/layer/page
+        self.assertEqual(lo.page_blob_bytes, 78 * 64 * 576)        # 2,875,392 = 2.74 MiB
+        self.assertEqual(lo.layer_stride_bytes, 100 * 64 * 576)
+
+    def test_factory_index_dims(self):
+        lo = a.CpL3SlabLayout.for_index(n_active_layers=21, indexer_page_num=100, indexer_page_stride_size=8448)
+        self.assertEqual(lo.page_blob_bytes, 21 * 8448)            # 177,408 = 0.17 MiB (GLM-5.2)
+        self.assertEqual(lo.layer_stride_bytes, 100 * 8448)
+
+
+if __name__ == "__main__":
+    unittest.main()