From de1d1e0af2d6ee26fbaba52ed536a4d028449573 Mon Sep 17 00:00:00 2001 From: leavelet Date: Sat, 20 Jun 2026 23:22:10 +0000 Subject: [PATCH] L3 3.0: CpL3Store orchestrator (async spill/reload, object-granular acks) Ties config -> per-rank disk slabs + slot pools + shared LMDB index + slab accessors. Background spill/reload threads (off the scheduler tick, storage-template shape); object-granular acks (one per object after all owned pages durable; zero-owned ranks ack in lockstep) via ack queues + has_inflight, so the caller's CP-cpu-group MIN-drain frees the same objects on every rank. Spill = gather->O_DIRECT write->durable index (data->fdatasync->index-commit ordering) with content-hash dedup; reload = index lookup->read_into->scatter (verify-on-read). free_object frees owned slots+index (3.3 eviction consumes); clear() is the flush_cache hook (stop/drain/reset/restart). from_config splits the disk budget across ranks-on-disk x payloads (equal slot count). PLP gate at connect(). Model B: store is content-addressed + rank-local I/O; only the LMDB is shared. End-to-end test (spill/exists/evict/reload byte-exact/dedup/free/clear) 5/5 (venv lmdb). + accessor n_layers/page_num accessors. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../srt/mem_cache/cp_l3_slab_accessor.py | 8 + python/sglang/srt/mem_cache/cp_l3_store.py | 349 ++++++++++++++++++ .../unit/mem_cache/test_cp_l3_store.py | 167 +++++++++ 3 files changed, 524 insertions(+) create mode 100644 python/sglang/srt/mem_cache/cp_l3_store.py create mode 100644 test/registered/unit/mem_cache/test_cp_l3_store.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 6c6d6ac99..51b7c53bd 100644 --- a/python/sglang/srt/mem_cache/cp_l3_slab_accessor.py +++ b/python/sglang/srt/mem_cache/cp_l3_slab_accessor.py @@ -73,6 +73,14 @@ class CpSharedL2SlabAccessor: def page_blob_bytes(self) -> int: return self._lo.page_blob_bytes + @property + def n_layers(self) -> int: + return self._lo.n_layers + + @property + def page_num(self) -> int: + return self._lo.page_num + def _slice_offset(self, layer: int, global_page: int) -> int: return layer * self._lo.layer_stride_bytes + global_page * self._lo.slice_bytes diff --git a/python/sglang/srt/mem_cache/cp_l3_store.py b/python/sglang/srt/mem_cache/cp_l3_store.py new file mode 100644 index 000000000..6e70e9a9b --- /dev/null +++ b/python/sglang/srt/mem_cache/cp_l3_store.py @@ -0,0 +1,349 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 +"""CP HiCache L3 — the content-addressed disk store (orchestrator). + +`CpL3Store` ties the pieces together for one CP rank: per-(payload,assigned-disk) ``CpL3DiskSlab`` + rank-local +``CpL3SlotPool``, the shared per-node LMDB index, and the host-slab accessors. Spill/reload run on dedicated +background threads (off the scheduler tick, mirroring the mainline storage template); completion is reported +**object-granular** (one ack per object after ALL this rank's owned pages are durable; a zero-owned rank acks +in lockstep so the cross-rank MIN-drain frees the same objects). The CP-cpu-group MIN over the ack-queue sizes +is done by the caller (hiradix `_drain_l3_control_queues`) — this store just exposes the queues + the drain +helpers + ``has_inflight`` (the replicated idle early-return). + +Model B (R1): radix tracks RAM-residency; this store is the durable content-addressed tier (key = per-page +content hash). They meet at spill (on radix eviction) + reload (on a radix miss). L3 eviction (3.3) rides the +index's replicated clock. The store is rank-local for I/O (own disk); only the LMDB index is shared per node. +""" + +from __future__ import annotations + +import os +import queue +import threading +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Sequence, Tuple + +from sglang.srt.mem_cache.cp_l3_config import CpL3Config, probe_disk_plp +from sglang.srt.mem_cache.cp_l3_disk import ( + CP_L3_HEADER_BYTES, + compute_crc, + content_hash_to_bytes, + _l3_failfast, +) +from sglang.srt.mem_cache.cp_l3_index import CpL3IndexEntry, CpL3MetaStore +from sglang.srt.mem_cache.cp_l3_posix import CpL3DiskSlab +from sglang.srt.mem_cache.cp_l3_slab_accessor import CpSharedL2SlabAccessor + +# (global_page, content_hash_hex) for one of this rank's owned pages of an object +OwnedPage = Tuple[int, str] + + +@dataclass +class L3SpillOp: + op_id: int + object_key: str + owned_pages: Dict[str, List[OwnedPage]] # payload_kind -> this rank's owned (slab_page, hash) + last_access: int # replicated logical clock for the index entries + + def is_empty(self) -> bool: + return not any(self.owned_pages.values()) + + +@dataclass +class L3ReloadOp: + op_id: int + object_key: str + owned_pages: Dict[str, List[OwnedPage]] # payload_kind -> this rank's owned (dest_slab_page, hash) + + def is_empty(self) -> bool: + return not any(self.owned_pages.values()) + + +@dataclass +class _Ack: + op_id: int + ok: bool + + +class CpL3Store: + def __init__( + self, + *, + cp_rank: int, + cp_size: int, + disk_dir: str, + disk_id: int, + index: CpL3MetaStore, + slabs: Dict[str, CpL3DiskSlab], + pools, # Dict[str, CpL3SlotPool] + accessors: Dict[str, CpSharedL2SlabAccessor], + ): + self.cp_rank = cp_rank + self.cp_size = cp_size + self.disk_dir = disk_dir + self.disk_id = disk_id + self.index = index + self.slabs = slabs + self.pools = pools + self.accessors = accessors + self.payloads: Tuple[str, ...] = tuple(slabs.keys()) + + self._op_counter = 0 + self._spill_q: "queue.Queue[L3SpillOp]" = queue.Queue() + self._reload_q: "queue.Queue[L3ReloadOp]" = queue.Queue() + self._ack_spill_q: "queue.Queue[_Ack]" = queue.Queue() + self._ack_reload_q: "queue.Queue[_Ack]" = queue.Queue() + self._ongoing_spill: set = set() + self._ongoing_reload: set = set() + self._stop = threading.Event() + self._spill_thread: Optional[threading.Thread] = None + self._reload_thread: Optional[threading.Thread] = None + self._started = False + + # ---------- construction ---------- + @classmethod + def from_config( + cls, + cfg: CpL3Config, + *, + cp_rank: int, + cp_size: int, + accessors: Dict[str, CpSharedL2SlabAccessor], + ) -> "CpL3Store": + from sglang.srt.mem_cache.cp_l3_disk import CpL3SlotPool, slot_bytes_for_page + + cfg.validate(cp_size) + disk_id = cfg.disk_for_rank(cp_rank, cp_size) + disk = cfg.disks[disk_id] + # split this disk's budget across the ranks mapped to it (graceful sharing) and across payloads + ranks_here = sum(1 for r in range(cp_size) if cfg.disk_for_rank(r, cp_size) == disk_id) + per_rank_budget = disk.budget_bytes // max(1, ranks_here) + payloads = tuple(accessors.keys()) + sum_slot_bytes = sum(slot_bytes_for_page(accessors[pk].page_blob_bytes) for pk in payloads) + if sum_slot_bytes <= 0: + _l3_failfast("no payloads / zero slot bytes") + num_slots = max(1, per_rank_budget // sum_slot_bytes) # equal slot count across payloads + os.makedirs(disk.path, exist_ok=True) + slabs: Dict[str, CpL3DiskSlab] = {} + pools = {} + for pk in payloads: + page_bytes = accessors[pk].page_blob_bytes + path = os.path.join(disk.path, f"l3_r{cp_rank}_{pk}.slab") + slabs[pk] = CpL3DiskSlab(path, pk, page_bytes, num_slots).open() + pools[pk] = CpL3SlotPool(num_slots) + index = CpL3MetaStore(cfg.index_dir, cfg.index_map_bytes, durable=True) + return cls( + cp_rank=cp_rank, cp_size=cp_size, disk_dir=disk.path, disk_id=disk_id, + index=index, slabs=slabs, pools=pools, accessors=accessors, + ) + + def connect(self, cfg: Optional[CpL3Config] = None) -> None: + """Start the bg threads. If cfg.require_plp, fail loud on a non-PLP disk (durability gate).""" + if cfg is not None and cfg.require_plp: + is_plp, median_us = probe_disk_plp(self.disk_dir) + if not is_plp: + _l3_failfast( + f"L3 disk {self.disk_dir} is not power-loss-protected (fdatasync median " + f"{median_us:.1f}us > threshold); durability is unsafe. Set require_plp=false only if " + f"you accept data loss on power failure." + ) + self._start_threads() + + def _start_threads(self) -> None: + if self._started: + return + self._stop.clear() + self._spill_thread = threading.Thread(target=self._spill_loop, daemon=True, name="cp-l3-spill") + self._reload_thread = threading.Thread(target=self._reload_loop, daemon=True, name="cp-l3-reload") + self._spill_thread.start() + self._reload_thread.start() + self._started = True + + # ---------- query (scheduler thread; lockless) ---------- + def exists_prefix(self, page_keys: Sequence[str], payload_kinds: Sequence[str]) -> int: + keys = [content_hash_to_bytes(h) for h in page_keys] + return self.index.exists_prefix(keys, payload_kinds) + + # ---------- submit (scheduler thread; non-blocking) ---------- + def _next_op_id(self) -> int: + self._op_counter += 1 + return self._op_counter + + def submit_spill(self, object_key: str, owned_pages: Dict[str, List[OwnedPage]], last_access: int) -> int: + op = L3SpillOp(self._next_op_id(), object_key, owned_pages, last_access) + self._ongoing_spill.add(op.op_id) + self._spill_q.put(op) + return op.op_id + + def submit_reload(self, object_key: str, owned_pages: Dict[str, List[OwnedPage]]) -> int: + op = L3ReloadOp(self._next_op_id(), object_key, owned_pages) + self._ongoing_reload.add(op.op_id) + self._reload_q.put(op) + return op.op_id + + def has_inflight(self) -> bool: + return bool(self._ongoing_spill) or bool(self._ongoing_reload) + + # ---------- completion harvest (scheduler thread; caller does the CP-group MIN) ---------- + def ack_spill_qsize(self) -> int: + return self._ack_spill_q.qsize() + + def ack_reload_qsize(self) -> int: + return self._ack_reload_q.qsize() + + def drain_spill_acks(self, n: int) -> List[Tuple[int, bool]]: + out = [] + for _ in range(n): + try: + ack = self._ack_spill_q.get_nowait() + except queue.Empty: + break + self._ongoing_spill.discard(ack.op_id) + out.append((ack.op_id, ack.ok)) + return out + + def drain_reload_acks(self, n: int) -> List[Tuple[int, bool]]: + out = [] + for _ in range(n): + try: + ack = self._ack_reload_q.get_nowait() + except queue.Empty: + break + self._ongoing_reload.discard(ack.op_id) + out.append((ack.op_id, ack.ok)) + return out + + # ---------- background workers ---------- + def _spill_loop(self) -> None: + while not self._stop.is_set(): + try: + op = self._spill_q.get(block=True, timeout=0.2) + except queue.Empty: + continue + ok = True + try: + self._spill_object(op) + except Exception as exc: # never kill the thread; report failure -> caller keeps L2 (no defer-free) + ok = False + import logging + + logging.getLogger(__name__).error("L3 spill op %s failed: %r", op.op_id, exc) + self._ack_spill_q.put(_Ack(op.op_id, ok)) + + def _reload_loop(self) -> None: + while not self._stop.is_set(): + try: + op = self._reload_q.get(block=True, timeout=0.2) + except queue.Empty: + continue + ok = True + try: + ok = self._reload_object(op) + except Exception as exc: + ok = False + import logging + + logging.getLogger(__name__).error("L3 reload op %s failed: %r", op.op_id, exc) + self._ack_reload_q.put(_Ack(op.op_id, ok)) + + def _spill_object(self, op: L3SpillOp) -> None: + """Gather+write this rank's owned pages of the object, then durably commit their index entries. + Zero-owned ranks no-op (the ack alone keeps the MIN-drain in lockstep).""" + touched: set = set() + with self.index.write_batch() as wb: + for pk, pages in op.owned_pages.items(): + if not pages: + continue + slab, pool, accessor = self.slabs[pk], self.pools[pk], self.accessors[pk] + for (slab_page, h_hex) in pages: + h = content_hash_to_bytes(h_hex) + if self.index.exists(h, pk): + continue # content-hash dedup: already durable in L3 + slot = pool.alloc() + if slot < 0: + _l3_failfast( + f"L3 {pk} slab full on rank {self.cp_rank} (evict-to-fit must run before spill)" + ) + buf = slab.new_aligned_buffer() + try: + accessor.gather_into(slab_page, buf, CP_L3_HEADER_BYTES) + crc = compute_crc(memoryview(buf)[CP_L3_HEADER_BYTES:CP_L3_HEADER_BYTES + slab.page_bytes]) + slab.fill_header(buf, h, accessor.n_layers, crc) + slab.write_aligned(slot, buf) + finally: + buf.close() + wb.put(h, pk, CpL3IndexEntry( + disk_id=self.disk_id, file_id=0, slot_idx=slot, page_bytes=slab.page_bytes, + crc=crc, last_access=op.last_access, hit_count=0, + )) + touched.add(pk) + # data -> fdatasync (durable) -> index commit (on __exit__, durable) : the crash-safe ordering + for pk in touched: + self.slabs[pk].fdatasync() + # (write_batch __exit__ committed the index durably here) + + def _reload_object(self, op: L3ReloadOp) -> bool: + """Read this rank's owned pages from L3 and scatter into the reserved L2 dest pages. Returns False if + any owned page is absent/corrupt (the prefix is then not fully reloadable on this rank).""" + ok = True + for pk, pages in op.owned_pages.items(): + if not pages: + continue + slab, accessor = self.slabs[pk], self.accessors[pk] + for (dest_page, h_hex) in pages: + h = content_hash_to_bytes(h_hex) + entry = self.index.get(h, pk) + if entry is None: + ok = False + continue + buf = slab.new_aligned_buffer() + try: + if slab.read_into(entry.slot_idx, buf, expect_hash=h): + accessor.scatter_from(dest_page, buf, CP_L3_HEADER_BYTES) + else: + ok = False + finally: + buf.close() + return ok + + # ---------- L3 eviction support (3.3 consumes; here = the rank-local free) ---------- + def free_object(self, owned_pages: Dict[str, List[OwnedPage]]) -> None: + """Free this rank's owned disk slots + index entries for an evicted object (selection is replicated + by the caller via the index's replicated clock).""" + with self.index.write_batch() as wb: + for pk, pages in owned_pages.items(): + for (_page, h_hex) in pages: + h = content_hash_to_bytes(h_hex) + entry = self.index.get(h, pk) + if entry is not None: + self.pools[pk].free(entry.slot_idx) + wb.delete(h, pk) + + # ---------- lifecycle ---------- + def clear(self) -> None: + """flush_cache hook: stop+drain threads, reset slot pools + index, restart. Idle-gated by the caller.""" + self._stop_threads() + for q in (self._spill_q, self._reload_q, self._ack_spill_q, self._ack_reload_q): + with q.mutex: + q.queue.clear() + self._ongoing_spill.clear() + self._ongoing_reload.clear() + for pool in self.pools.values(): + pool.reset() + self.index.clear() + self._start_threads() + + def _stop_threads(self) -> None: + if not self._started: + return + self._stop.set() + for t in (self._spill_thread, self._reload_thread): + if t is not None: + t.join(timeout=3) + self._started = False + + def close(self) -> None: + self._stop_threads() + for slab in self.slabs.values(): + slab.close() + self.index.close() diff --git a/test/registered/unit/mem_cache/test_cp_l3_store.py b/test/registered/unit/mem_cache/test_cp_l3_store.py new file mode 100644 index 000000000..e5179c33a --- /dev/null +++ b/test/registered/unit/mem_cache/test_cp_l3_store.py @@ -0,0 +1,167 @@ +"""End-to-end unit test for the CP HiCache L3 store (cp_l3_store.CpL3Store), single rank (cp_size=1). + +Spill -> ack drain -> exists_prefix -> simulate L2 eviction (zero slab pages) -> reload (byte-exact) -> +free_object -> clear. Requires lmdb (run with the l3venv). Uses a synthetic host slab + temp-dir config. +""" + +import importlib.util +import mmap +import os +import sys +import tempfile +import time +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_l3(): + 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 + for mod in ("cp_l3_disk", "cp_l3_config", "cp_l3_index", "cp_l3_posix", "cp_l3_slab_accessor", "cp_l3_store"): + _load_file(f"sglang.srt.mem_cache.{mod}", _MEM / f"{mod}.py") + return ( + sys.modules["sglang.srt.mem_cache.cp_l3_store"], + sys.modules["sglang.srt.mem_cache.cp_l3_config"], + sys.modules["sglang.srt.mem_cache.cp_l3_slab_accessor"], + ) + + +store_mod, cfg_mod, acc_mod = _load_l3() + +N_LAYERS, PAGE_NUM, SLICE = 3, 8, 16 + + +def _h(i): + return f"{i:064x}" # 64 hex chars == 32 bytes + + +def _pattern(layer, page): + return bytes(((layer * 41 + page * 13 + k) & 0xFF) for k in range(SLICE)) + + +def _make_slab_and_accessor(): + lo = acc_mod.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] = _pattern(layer, page) + return mm, acc_mod.CpSharedL2SlabAccessor(mm, lo) + + +def _wait_ack(qsize_fn, n=1, timeout=5.0): + t0 = time.time() + while qsize_fn() < n and time.time() - t0 < timeout: + time.sleep(0.005) + return qsize_fn() >= n + + +class TestCpL3Store(unittest.TestCase): + def setUp(self): + self._td = tempfile.TemporaryDirectory() + self.mm, self.acc = _make_slab_and_accessor() + cfg = cfg_mod.CpL3Config.from_dict({ + "backend": "posix", + "require_plp": False, # temp FS / buffered + "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 test_spill_exists_reload_roundtrip(self): + # spill pages 2 and 5 + pages = {"target_kv": [(2, _h(2)), (5, _h(5))]} + orig2 = self.acc.gather(2) + orig5 = self.acc.gather(5) + self.store.submit_spill("obj-A", pages, last_access=10) + self.assertTrue(_wait_ack(self.store.ack_spill_qsize, 1)) + acks = self.store.drain_spill_acks(self.store.ack_spill_qsize()) + self.assertEqual(acks, [(1, True)]) + self.assertFalse(self.store.has_inflight()) + + # exists_prefix: consecutive present count + self.assertEqual(self.store.exists_prefix([_h(2), _h(5)], ["target_kv"]), 2) + self.assertEqual(self.store.exists_prefix([_h(2), _h(99)], ["target_kv"]), 1) + self.assertEqual(self.store.exists_prefix([_h(99)], ["target_kv"]), 0) + + # simulate L2 eviction: zero out pages 2 and 5 in the slab + for p in (2, 5): + for layer in range(N_LAYERS): + off = layer * (PAGE_NUM * SLICE) + p * SLICE + self.mm[off:off + SLICE] = bytes(SLICE) + self.assertNotEqual(self.acc.gather(2), orig2) + + # reload from L3 -> scatter back into the slab + self.store.submit_reload("obj-A", pages) + self.assertTrue(_wait_ack(self.store.ack_reload_qsize, 1)) + racks = self.store.drain_reload_acks(self.store.ack_reload_qsize()) + self.assertEqual(racks, [(2, True)]) + self.assertEqual(self.acc.gather(2), orig2) # byte-exact round-trip through disk + self.assertEqual(self.acc.gather(5), orig5) + + def test_zero_owned_spill_acks_in_lockstep(self): + # a rank that owns no pages of the object still must ack (empty op) + op_id = self.store.submit_spill("obj-empty", {"target_kv": []}, last_access=1) + self.assertTrue(_wait_ack(self.store.ack_spill_qsize, 1)) + self.assertEqual(self.store.drain_spill_acks(1), [(op_id, True)]) + + def test_dedup_skip_existing(self): + pages = {"target_kv": [(3, _h(3))]} + self.store.submit_spill("o1", pages, last_access=1) + self.assertTrue(_wait_ack(self.store.ack_spill_qsize, 1)) + self.store.drain_spill_acks(1) + free_after_first = self.store.pools["target_kv"].num_free + # spill the SAME hash again -> dedup, no new slot consumed + self.store.submit_spill("o1b", pages, last_access=2) + self.assertTrue(_wait_ack(self.store.ack_spill_qsize, 1)) + self.store.drain_spill_acks(1) + self.assertEqual(self.store.pools["target_kv"].num_free, free_after_first) + + def test_free_object_releases_slots_and_index(self): + pages = {"target_kv": [(1, _h(1)), (4, _h(4))]} + before = self.store.pools["target_kv"].num_free + self.store.submit_spill("o", pages, last_access=1) + self.assertTrue(_wait_ack(self.store.ack_spill_qsize, 1)) + self.store.drain_spill_acks(1) + self.assertEqual(self.store.pools["target_kv"].num_free, before - 2) + self.store.free_object(pages) + self.assertEqual(self.store.pools["target_kv"].num_free, before) + self.assertEqual(self.store.exists_prefix([_h(1)], ["target_kv"]), 0) + + def test_clear_resets(self): + self.store.submit_spill("o", {"target_kv": [(0, _h(0))]}, last_access=1) + self.assertTrue(_wait_ack(self.store.ack_spill_qsize, 1)) + self.store.drain_spill_acks(1) + self.store.clear() + self.assertEqual(self.store.exists_prefix([_h(0)], ["target_kv"]), 0) + self.assertEqual(self.store.pools["target_kv"].num_free, self.store.pools["target_kv"].num_slots) + self.assertFalse(self.store.has_inflight()) + + +if __name__ == "__main__": + unittest.main()