L3 3.0: wire CpL3Store into HiRadixCache (construct + drain + reset/clear)

_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) <noreply@anthropic.com>
This commit is contained in:
2026-06-20 23:30:41 +00:00
parent 81a1b72be6
commit 2ae46013ac

View File

@@ -1191,6 +1191,94 @@ class HiRadixCache(RadixCache):
self.eviction_strategy = CpReplicatedSLRUStrategy() 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: def _cp_hicache_all_reduce(self, tensor, *, op, group, tag: str) -> None:
start = time.perf_counter() start = time.perf_counter()
torch.distributed.all_reduce(tensor, op=op, group=group) torch.distributed.all_reduce(tensor, op=op, group=group)
@@ -2303,6 +2391,11 @@ class HiRadixCache(RadixCache):
self.detach_storage_backend() self.detach_storage_backend()
except Exception: except Exception:
logger.exception("Failed to detach storage backend on process shutdown.") 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( def _apply_storage_runtime_config(
self, self,
@@ -2709,6 +2802,9 @@ class HiRadixCache(RadixCache):
TreeNode.counter = 0 TreeNode.counter = 0
TreeNode.last_access_counter = 0 TreeNode.last_access_counter = 0
self.cache_controller.reset() 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.token_to_kv_pool_host.clear()
self.cache_controller.clear_draft_host_pool() self.cache_controller.clear_draft_host_pool()
# Clear per-request tracking dicts # Clear per-request tracking dicts
@@ -4975,6 +5071,8 @@ class HiRadixCache(RadixCache):
self.loading_check() self.loading_check()
if self.enable_storage: if self.enable_storage:
self.drain_storage_control_queues() self.drain_storage_control_queues()
if self.enable_cp_l3:
self._drain_l3_control_queues()
if self.enable_storage_metrics: if self.enable_storage_metrics:
self.storage_metrics_collector.log_storage_metrics( self.storage_metrics_collector.log_storage_metrics(
self.cache_controller.storage_backend.get_stats() self.cache_controller.storage_backend.get_stats()