From 2ae46013ac461e1c934b4820735b711f84383f1b Mon Sep 17 00:00:00 2001 From: leavelet Date: Sat, 20 Jun 2026 23:30:41 +0000 Subject: [PATCH] L3 3.0: wire CpL3Store into HiRadixCache (construct + drain + reset/clear) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _maybe_init_cp_l3: build the 3 CpSharedL2SlabAccessor (target_kv always; draft_kv if draft pool; index_k if NSA) from the LIVE host pool slabs (verified attrs: host.allocator.mapping.mmap, layer_num/page_num/kv_cache_dim/dtype.itemsize; index_host_tensor_allocator + indexer attrs), then CpL3Store.from_config + connect (PLP gate). Single-slab only (default --cp-shared-l2-slab-size-gb 0); fail loud on the multi-slab group case (future extension, not a silent stub). _drain_l3_control_queues: per-tick CP-cpu-group MIN over the 2 ack-qsizes + drain MIN of each, idle early-return on the replicated has_inflight (R1 CRIT-3/4) — wired in check_hicache_events (gated on enable_cp_l3, NOT the enable_storage drain). reset() -> cp_l3_store.clear(); shutdown() -> close(). Default-off + behavior-neutral. Compiles. 3.0 CODE complete; flag-on smoke gated on cluster branch sync. Co-Authored-By: Claude Opus 4.8 (1M context) --- python/sglang/srt/mem_cache/hiradix_cache.py | 98 ++++++++++++++++++++ 1 file changed, 98 insertions(+) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index f5fd7fc1a..9b022ddd1 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -1191,6 +1191,94 @@ class HiRadixCache(RadixCache): self.eviction_strategy = CpReplicatedSLRUStrategy() + self._maybe_init_cp_l3(server_args, params) + + def _maybe_init_cp_l3(self, server_args, params) -> None: + """Construct the CP L3 disk store from the live host pool slabs (Model B: content-addressed, + rank-local I/O + the shared LMDB). Default-off; only when --enable-cp-l3.""" + if not self.enable_cp_l3: + return + from sglang.srt.mem_cache.cp_l3_config import CpL3Config + from sglang.srt.mem_cache.cp_l3_slab_accessor import ( + CpL3SlabLayout, + CpSharedL2SlabAccessor, + ) + from sglang.srt.mem_cache.cp_l3_store import CpL3Store + from sglang.srt.mem_cache.memory_pool_host import NSATokenToKVPoolHost + + cp_allocator = params.token_to_kv_pool_allocator + cp_rank, cp_size = int(cp_allocator.cp_rank), int(cp_allocator.cp_size) + cfg = CpL3Config.from_file(server_args.cp_l3_config) + + host = self.token_to_kv_pool_host + accessors = { + "target_kv": CpSharedL2SlabAccessor( + self._cp_l3_single_slab_mmap(host.allocator, "target_kv"), + 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, + ), + ) + } + 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"), + 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, + ), + ) + if isinstance(host, NSATokenToKVPoolHost): + accessors["index_k"] = CpSharedL2SlabAccessor( + self._cp_l3_single_slab_mmap(host.index_host_tensor_allocator, "index_k"), + 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, + ), + ) + self.cp_l3_store = CpL3Store.from_config( + cfg, cp_rank=cp_rank, cp_size=cp_size, accessors=accessors + ) + self.cp_l3_store.connect(cfg) + logger.info( + "[CP-L3] enabled rank=%d/%d payloads=%s disk_dir=%s", + cp_rank, cp_size, list(accessors), self.cp_l3_store.disk_dir, + ) + + @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).""" + 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." + ) + return mapping.mmap + + def _drain_l3_control_queues(self) -> None: + """Per-tick CP-cpu-group MIN over the L3 ack-queue sizes, then drain MIN of each (R1 CRIT-3/4). + Idle early-return: ongoing_spill/reload is a replicated quantity (3.1+ submits are lockstep on the + replicated plan), so an idle tick issues no collective. 3.0: always idle -> inert (3.1/3.2 act on acks).""" + store = self.cp_l3_store + if store is None or not store.has_inflight(): + return + qsizes = torch.tensor( + [store.ack_spill_qsize(), store.ack_reload_qsize()], dtype=torch.int + ) + if self.tp_world_size > 1: + self._cp_hicache_all_reduce( + qsizes, op=torch.distributed.ReduceOp.MIN, group=self.tp_group, tag="l3_queue_min" + ) + n_spill, n_reload = map(int, qsizes.tolist()) + store.drain_spill_acks(n_spill) + store.drain_reload_acks(n_reload) + def _cp_hicache_all_reduce(self, tensor, *, op, group, tag: str) -> None: start = time.perf_counter() torch.distributed.all_reduce(tensor, op=op, group=group) @@ -2303,6 +2391,11 @@ class HiRadixCache(RadixCache): self.detach_storage_backend() except Exception: logger.exception("Failed to detach storage backend on process shutdown.") + try: + if getattr(self, "cp_l3_store", None) is not None: + self.cp_l3_store.close() + except Exception: + logger.exception("Failed to close CP L3 store on process shutdown.") def _apply_storage_runtime_config( self, @@ -2709,6 +2802,9 @@ class HiRadixCache(RadixCache): TreeNode.counter = 0 TreeNode.last_access_counter = 0 self.cache_controller.reset() + if self.cp_l3_store is not None: + # flush_cache is idle-gated -> no in-flight L3 ops; clear() stops/drains/resets/restarts. + self.cp_l3_store.clear() self.token_to_kv_pool_host.clear() self.cache_controller.clear_draft_host_pool() # Clear per-request tracking dicts @@ -4975,6 +5071,8 @@ class HiRadixCache(RadixCache): self.loading_check() if self.enable_storage: self.drain_storage_control_queues() + if self.enable_cp_l3: + self._drain_l3_control_queues() if self.enable_storage_metrics: self.storage_metrics_collector.log_storage_metrics( self.cache_controller.storage_backend.get_stats()