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>
This commit is contained in:
125
test/registered/unit/mem_cache/test_cp_l3_slab_accessor.py
Normal file
125
test/registered/unit/mem_cache/test_cp_l3_slab_accessor.py
Normal file
@@ -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()
|
||||
Reference in New Issue
Block a user