diff --git a/python/sglang/srt/mem_cache/cp_l3_slab_accessor.py b/python/sglang/srt/mem_cache/cp_l3_slab_accessor.py new file mode 100644 index 000000000..6c6d6ac99 --- /dev/null +++ b/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 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 new file mode 100644 index 000000000..d868c72ea --- /dev/null +++ b/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()