From 01029c8df484e08851535fe21ab4e63e66d031d8 Mon Sep 17 00:00:00 2001 From: leavelet Date: Sat, 20 Jun 2026 13:59:39 +0000 Subject: [PATCH] 2b.1b: B1 collective-free pooled-L2 write-through (default-off) The B1 write-through: every CP rank runs the IDENTICAL deterministic CpSharedL2PageAllocator.reserve over the replicated event stream -> identical placement with NO broadcast/gather; commit rides the writing_check ReduceOp.MIN frontier (mark_object_committed) instead of a per-(rank,layer,payload) gather quorum. cache_controller.py: _reserve_write_cp_shared_l2 stripped to every-rank reserve (dropped preflight#2/rank0-gate/broadcast#3/adopt/missing-payloads; per-object contract kept MF3; idempotent abort-on-partial SF3); _submit_write_cp_layer_states drops the 3 gather-commit calls (per-layer D2H + all-payload done-gate MF1 + per-node ack kept); submit_write_cp_all_layer drops both fallback gathers (zero-owned keeps its ack MF2). hiradix_cache.py: 3-way merge of l2_pooling shared-L2 init (clean, 0-conflict; Phase-1 clock + E1 + reserve-budget preserved) builds the allocator/slab when the flag is on and passes it to HiCacheController WITHOUT collective fns. The #5 deadlock removed: shared_l2_capacity reserve-miss now SKIPS the backup (reactive, rank-uniform) instead of _evict_host_for_physical_slots(synchronize_across_ranks=True) -- the deterministic shared-pool evict is 2b.2. mark_object_committed wired at _commit_pending_backup (the MIN commit). New _cp_assert_placement_replicated (SGLANG_CP_HICACHE_PLACEMENT_ASSERT) in check_hicache_events: MIN/MAX-reduce placement_digest across the CP group, fail-loud on divergence (Theorem 1 runtime gate; rank-uniform entry, cannot deadlock). environ.py: SGLANG_CP_HICACHE_PLACEMENT_ASSERT (EnvBool, default off). Tests: 8-rank reserve-determinism (identical placement across 8 ranks + divergence-detectable) + full pool suite 91/91 + import smoke on g0033 syh-dev-new. TWO opus adversarial reviews both SHIP (cache_controller strip + hiradix wiring). NOTES: default-off (--enable-cp-shared-physical-l2-hicache); validated for write_through policy (prod default). Capacity-miss currently SKIPS (B1 shared-pool evict = 2b.2) so the flag must stay off in real runs until 2b.2. Dead gather island (3 fns + 3 imports, opus-confirmed inert) pending GC follow-up. Design+proof: docs_internal/cp_hicache_2b_pooled_l2_b1_design.md sec 2.4/9. Co-Authored-By: Claude Opus 4.8 (1M context) --- python/sglang/srt/environ.py | 4 + .../sglang/srt/managers/cache_controller.py | 1128 +++++++++++++++-- python/sglang/srt/mem_cache/hiradix_cache.py | 938 +++++++++++++- .../unit/mem_cache/test_cp_shared_l2_pool.py | 64 + 4 files changed, 1970 insertions(+), 164 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 7a30de4da..7f25fb37b 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -389,6 +389,10 @@ class Envs: # during bring-up so E1 measures the baseline (reactive drops, not waits); # flip on once validated (E3). Opt-in. SGLANG_CP_HICACHE_RESERVE_ADMISSION = EnvBool(False) + # Phase 2b (B1): runtime cross-rank assert that the pooled-L2 allocator placement + # is replicated (Theorem 1). When on, MIN/MAX-reduce placement_digest across the + # CP group each tick and fail loud on divergence. On in CI/bring-up, off in prod. + SGLANG_CP_HICACHE_PLACEMENT_ASSERT = EnvBool(False) # Mooncake KV Transfer SGLANG_MOONCAKE_CUSTOM_MEM_POOL = EnvStr(None) diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 2ae5d4a53..ea2c05a13 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -43,6 +43,17 @@ from sglang.srt.mem_cache.cp_shared_kv_layout import ( CpSharedKVLayout, pad_token_locs_to_page_boundary, ) +from sglang.srt.mem_cache.cp_shared_l2_pool import ( + PAYLOAD_DRAFT_KV, + PAYLOAD_INDEX_K, + PAYLOAD_TARGET_KV, + CpSharedL2NodeMetadata, + CpSharedL2PageAllocator, + broadcast_cp_shared_l2_decision, + cp_shared_l2_logical_token_indices, + gather_cp_shared_l2_commits, + gather_cp_shared_l2_preflight, +) from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool, NSATokenToKVPool from sglang.srt.mem_cache.page_index_utils import validate_page_aligned_token_indices from sglang.srt.utils import get_device_module @@ -336,7 +347,10 @@ class HiCacheWriteReservation: node_id: int = -1 priority: Optional[int] = None draft_host_indices: Optional[torch.Tensor] = None + index_host_indices: Optional[torch.Tensor] = None required_host_slots: int = 0 + owned_positions: Optional[torch.Tensor] = None + shared_l2_object_key: Optional[str] = None @dataclass @@ -348,9 +362,12 @@ class HiCacheLayerWriteState: layer_events: List[object] = field(default_factory=list) completed_target_layers: Set[int] = field(default_factory=set) completed_draft_layers: Set[int] = field(default_factory=set) + completed_index_layers: Set[int] = field(default_factory=set) host_indices: Optional[torch.Tensor] = None physical_device_indices: Optional[torch.Tensor] = None draft_host_indices: Optional[torch.Tensor] = None + index_host_indices: Optional[torch.Tensor] = None + index_layer_ids: tuple[int, ...] = () ack_appended: bool = False @@ -375,6 +392,12 @@ class HiCacheController: cp_shared_kv_layout: Optional[CpSharedKVLayout] = None, draft_mem_pool_host: Optional["HostKVCache"] = None, draft_mem_pool_device=None, + cp_shared_l2_page_allocator: Optional[CpSharedL2PageAllocator] = None, + cp_shared_l2_cpu_group=None, + cp_shared_l2_source_rank: int = 0, + cp_shared_l2_broadcast_fn=None, + cp_shared_l2_preflight_fn=None, + cp_shared_l2_commit_fn=None, ): self.tp_group = tp_group self.mem_pool_device_allocator = token_to_kv_pool_allocator @@ -385,6 +408,16 @@ class HiCacheController: mem_pool_device = mem_pool_device.full_kv_pool self.mem_pool_device = mem_pool_device self.cp_shared_kv_layout = cp_shared_kv_layout + self.cp_shared_l2_page_allocator = cp_shared_l2_page_allocator + self.cp_shared_l2_cpu_group = cp_shared_l2_cpu_group + self.cp_shared_l2_source_rank = int(cp_shared_l2_source_rank) + self.cp_shared_l2_broadcast_fn = cp_shared_l2_broadcast_fn + self.cp_shared_l2_preflight_fn = cp_shared_l2_preflight_fn + self.cp_shared_l2_commit_fn = cp_shared_l2_commit_fn + self._cp_shared_l2_stats = { + "cp_shared_l2_load_hits": 0, + "cp_shared_l2_load_misses": 0, + } self.has_draft = False self.mem_pool_device_draft = None @@ -441,6 +474,8 @@ class HiCacheController: self.write_queue: List[CacheOperation] = [] self.draft_load_queue: List[CacheOperation] = [] self.draft_write_queue: List[CacheOperation] = [] + self.index_load_queue: List[CacheOperation] = [] + self.index_write_queue: List[CacheOperation] = [] self.ack_load_queue: List[HiCacheAck] = [] self.ack_write_queue: List[HiCacheAck] = [] self.pending_layer_writes: Dict[int, HiCacheLayerWriteState] = {} @@ -813,6 +848,8 @@ class HiCacheController: self.load_queue.clear() self.draft_write_queue.clear() self.draft_load_queue.clear() + self.index_write_queue.clear() + self.index_load_queue.clear() self.write_buffer.clear() self.load_buffer.clear() self.ack_write_queue.clear() @@ -889,7 +926,11 @@ class HiCacheController: return host_indices def start_writing(self) -> None: - if len(self.write_queue) == 0 and len(self.draft_write_queue) == 0: + if ( + len(self.write_queue) == 0 + and len(self.draft_write_queue) == 0 + and len(self.index_write_queue) == 0 + ): return op = CacheOperation.merge_ops(self.write_queue) if self.write_queue else None @@ -898,7 +939,13 @@ class HiCacheController: if self.draft_write_queue else None ) - node_ids = op.node_ids if op is not None else draft_op.node_ids + index_op = ( + CacheOperation.merge_ops(self.index_write_queue) + if self.index_write_queue + else None + ) + primary_op = op or draft_op or index_op + node_ids = primary_op.node_ids if op is not None: host_indices, device_indices = self.move_indices(op, self.mem_pool_host) else: @@ -909,8 +956,15 @@ class HiCacheController: ) else: draft_host_indices = draft_device_indices = None + if index_op is not None: + index_host_indices, index_device_indices = self.move_indices( + index_op, self.mem_pool_host + ) + else: + index_host_indices = index_device_indices = None self.write_queue.clear() self.draft_write_queue.clear() + self.index_write_queue.clear() start_event = device_module.Event() finish_event = device_module.Event() @@ -919,9 +973,34 @@ class HiCacheController: with device_module.stream(self.write_stream): start_event.wait(self.write_stream) if op is not None: - self.mem_pool_host.backup_from_device_all_layer( + backup_target = self.mem_pool_host.backup_from_device_all_layer + if index_op is not None: + backup_target = getattr( + self.mem_pool_host, "backup_kv_from_device_all_layer", None + ) + if backup_target is None: + raise RuntimeError( + "CP shared physical L2 index_k write requires " + "backup_kv_from_device_all_layer on the target host pool" + ) + backup_target( self.mem_pool_device, host_indices, device_indices, self.io_backend ) + if index_op is not None: + backup_index = getattr( + self.mem_pool_host, "backup_indexer_from_device_all_layer", None + ) + if backup_index is None: + raise RuntimeError( + "CP shared physical L2 index_k write requires " + "backup_indexer_from_device_all_layer on the target host pool" + ) + backup_index( + self.mem_pool_device, + index_host_indices, + index_device_indices, + self.io_backend, + ) if draft_op is not None: self.draft_mem_pool_host.backup_from_device_all_layer( self.draft_mem_pool_device, @@ -941,6 +1020,10 @@ class HiCacheController: draft_host_indices.record_stream(self.write_stream) if draft_device_indices is not None and draft_device_indices.is_cuda: draft_device_indices.record_stream(self.write_stream) + if index_host_indices is not None and index_host_indices.is_cuda: + index_host_indices.record_stream(self.write_stream) + if index_device_indices is not None and index_device_indices.is_cuda: + index_device_indices.record_stream(self.write_stream) self._append_write_ack(HiCacheAck(start_event, finish_event, node_ids)) @@ -961,6 +1044,25 @@ class HiCacheController: event.record() self._append_write_ack(HiCacheAck(event, event, [node_id])) + def _remove_write_acks_for_node(self, node_id: int) -> bool: + retained_acks: List[HiCacheAck] = [] + removed = False + for ack in self.ack_write_queue: + if node_id not in ack.node_ids: + retained_acks.append(ack) + continue + remaining_node_ids = [nid for nid in ack.node_ids if nid != node_id] + if remaining_node_ids: + retained_acks.append(ack._replace(node_ids=remaining_node_ids)) + removed = True + if removed: + self.ack_write_queue = retained_acks + self._max_write_ack_node_id = max( + (nid for ack in retained_acks for nid in ack.node_ids), + default=-1, + ) + return removed + def _append_completed_load_ack(self, node_id: int) -> None: event = device_module.Event() event.record() @@ -1033,6 +1135,9 @@ class HiCacheController: priority: Optional[int] = None, node_id: int = -1, ) -> HiCacheWriteReservation | HiCacheWriteFailure: + if self.cp_shared_l2_page_allocator is not None: + return self._reserve_write_cp_shared_l2(device_indices, priority, node_id) + from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata if node_id >= 0 and self.node_has_undrained_write_ack(node_id): @@ -1160,6 +1265,259 @@ class HiCacheController: draft_host_indices=draft_host_indices, ) + def _cp_shared_l2_object_key(self, node_id: int) -> str: + return f"cp_hicache_node:{int(node_id)}" + + def _abort_cp_shared_l2_reservation( + self, reservation_or_metadata, *, fence_write_stream: bool = False + ) -> None: + object_key = getattr(reservation_or_metadata, "object_key", None) + if object_key is None: + object_key = getattr(reservation_or_metadata, "shared_l2_object_key", None) + allocator = self.cp_shared_l2_page_allocator + if allocator is not None and object_key: + if fence_write_stream: + self.write_stream.synchronize() + allocator.abort(object_key) + + def _release_cp_shared_l2_metadata(self, metadata) -> bool: + object_key = getattr(metadata, "object_key", None) + allocator = self.cp_shared_l2_page_allocator + if allocator is None or not object_key: + return False + return allocator.release(object_key) + + def cp_shared_l2_cache_report(self) -> dict[str, int]: + allocator = self.cp_shared_l2_page_allocator + if allocator is not None and hasattr(allocator, "stats"): + report = dict(allocator.stats()) + else: + report = { + "cp_shared_l2_pages_capacity": 0, + "cp_shared_l2_pages_used": 0, + "cp_shared_l2_objects_committed": 0, + "cp_shared_l2_objects_aborted": 0, + "cp_shared_l2_objects_evicted": 0, + } + report.update(self._cp_shared_l2_stats) + return report + + @staticmethod + def _reservation_owned_count(reservation: HiCacheWriteReservation) -> int: + owned_positions = getattr(reservation, "owned_positions", None) + if owned_positions is not None: + return int(owned_positions.numel()) + metadata = reservation.metadata + metadata_owned_positions = getattr(metadata, "owned_positions", None) + if metadata_owned_positions is not None: + return int(metadata_owned_positions.numel()) + return int(reservation.host_indices.numel()) + + def _cp_shared_l2_index_layer_ids(self) -> tuple[int, ...]: + active_layer_ids = getattr(self.mem_pool_host, "index_active_layer_ids", None) + if active_layer_ids is not None: + return tuple(int(layer_id) for layer_id in active_layer_ids) + index_layer_num = int( + getattr(self.mem_pool_host, "index_active_layer_num", 0) or 0 + ) + start_layer = int(getattr(self.mem_pool_host, "start_layer", 0) or 0) + return tuple(range(start_layer, start_layer + index_layer_num)) + + def _cp_shared_l2_expected_layers_by_payload( + self, required_payloads: List[str] + ) -> dict[str, object]: + expected_layers_by_payload = {} + if PAYLOAD_TARGET_KV in required_payloads: + expected_layers_by_payload[PAYLOAD_TARGET_KV] = range(self.layer_num) + if PAYLOAD_DRAFT_KV in required_payloads: + if self.draft_mem_pool_device is None: + raise RuntimeError( + "CP shared physical L2 draft payload requires draft KV device pool" + ) + draft_layer_num = int(getattr(self.draft_mem_pool_device, "layer_num", 0)) + if draft_layer_num <= 0: + raise RuntimeError( + "CP shared physical L2 draft payload requires positive draft layer count" + ) + expected_layers_by_payload[PAYLOAD_DRAFT_KV] = range(draft_layer_num) + if PAYLOAD_INDEX_K in required_payloads: + index_layer_ids = self._cp_shared_l2_index_layer_ids() + if not index_layer_ids: + raise RuntimeError( + "CP shared physical L2 index_k payload requires active NSA index layers" + ) + expected_layers_by_payload[PAYLOAD_INDEX_K] = index_layer_ids + return expected_layers_by_payload + + def _set_cp_shared_l2_object_commit_contract( + self, allocator, object_key: str, required_payloads: List[str] + ) -> None: + set_required = getattr(allocator, "set_object_required_payloads", None) + if set_required is None: + return + set_required( + object_key, + required_payloads, + expected_layers_by_payload=self._cp_shared_l2_expected_layers_by_payload( + required_payloads + ), + ) + + def _cp_shared_l2_index_payload_active(self) -> bool: + return bool(self._cp_shared_l2_index_layer_ids()) + + def _reserve_write_cp_shared_l2( + self, + device_indices: torch.Tensor, + priority: Optional[int] = None, + node_id: int = -1, + ) -> HiCacheWriteReservation | HiCacheWriteFailure: + allocator = self.cp_shared_l2_page_allocator + if allocator is None: + raise RuntimeError("shared L2 write reservation requested without allocator") + layout = self.cp_shared_kv_layout + if layout is None: + raise RuntimeError("shared L2 write reservation requires CP layout") + required_payloads = [PAYLOAD_TARGET_KV] + if self.has_draft_hicache: + required_payloads.append(PAYLOAD_DRAFT_KV) + if self._cp_shared_l2_index_payload_active(): + required_payloads.append(PAYLOAD_INDEX_K) + + page_size = self.page_size + logical_len = len(device_indices) + padded_device_indices = pad_token_locs_to_page_boundary( + device_indices, + page_size, + name="CP shared physical L2 HiCache write device_indices", + ) + padded_len = int(padded_device_indices.numel()) + num_pages = padded_len // page_size + object_key = self._cp_shared_l2_object_key(node_id) + # B1: no preflight gather -- node_has_undrained_write_ack reads the + # replicated ack_write_queue (node.id-keyed), so every CP rank decides + # identically without a collective. + if node_id >= 0 and self.node_has_undrained_write_ack(node_id): + return HiCacheWriteFailure(required_host_slots=0, reason="undrained_ack") + + object_ranges = None + try: + # B1: EVERY rank runs the identical deterministic reserve over the + # replicated free list -> identical (slab_id, base_page) with NO + # rank-0 gate / broadcast_decision / adopt. A capacity ValueError + # raises identically on all ranks (same free list), so the + # HiCacheWriteFailure is rank-uniform. + object_ranges = {} + try: + for payload_kind in required_payloads: + object_ranges[payload_kind] = allocator.reserve( + object_key, payload_kind, num_pages + ) + except ValueError as exc: + allocator.abort(object_key) # idempotent; drops any partial reserve + logger.debug( + "[CacheCtrl-write] reserve_write_cp shared L2 FAILED (B1): " + "node_id=%d logical_len=%d pages=%d payloads=%s error=%s", + node_id, + logical_len, + num_pages, + required_payloads, + exc, + ) + return HiCacheWriteFailure( + required_host_slots=padded_len, + reason="shared_l2_capacity", + ) + # MF3: reserve() only sets the allocator-wide default payload set; pin + # the actual per-object contract (required payloads + expected layers) + # so split_committed_object reads correct state. + self._set_cp_shared_l2_object_commit_contract( + allocator, object_key, required_payloads + ) + + object_range = object_ranges[PAYLOAD_TARGET_KV] + owned_mask = layout.owned_by_this_rank(padded_device_indices) + owned_positions = owned_mask.nonzero(as_tuple=True)[0].cpu() + if owned_positions.numel() == 0: + host_indices = torch.empty((0,), dtype=torch.int64) + physical_device_indices = torch.empty((0,), dtype=torch.int64) + else: + owned_logical_indices = padded_device_indices[owned_mask] + physical_device_indices = layout.logical_locs_to_physical( + owned_logical_indices + ) + host_indices = cp_shared_l2_logical_token_indices( + object_range, page_size, owned_positions + ).to(dtype=torch.int64) + + draft_host_indices = None + if PAYLOAD_DRAFT_KV in required_payloads: + draft_range = object_ranges[PAYLOAD_DRAFT_KV] + draft_host_indices = cp_shared_l2_logical_token_indices( + draft_range, page_size, owned_positions + ).to(dtype=torch.int64) + index_host_indices = None + if PAYLOAD_INDEX_K in required_payloads: + index_range = object_ranges[PAYLOAD_INDEX_K] + index_host_indices = cp_shared_l2_logical_token_indices( + index_range, page_size, owned_positions + ).to(dtype=torch.int64) + + metadata = CpSharedL2NodeMetadata( + logical_len=logical_len, + padded_len=padded_len, + page_size=page_size, + object_ranges=object_ranges, + required_payloads=tuple(required_payloads), + committed_payload_layers={payload: set() for payload in required_payloads}, + committed=False, + object_key=object_key, + ) + return HiCacheWriteReservation( + metadata=metadata, + host_indices=host_indices, + physical_device_indices=physical_device_indices, + node_id=node_id, + priority=priority, + draft_host_indices=draft_host_indices, + index_host_indices=index_host_indices, + required_host_slots=padded_len, + owned_positions=owned_positions, + shared_l2_object_key=object_key, + ) + except Exception: + if object_ranges is not None: + allocator.abort(object_key) + raise + + def _commit_cp_shared_l2_all_required_layers( + self, reservation: HiCacheWriteReservation + ) -> None: + metadata = reservation.metadata + required_payloads = tuple( + getattr(metadata, "required_payloads", (PAYLOAD_TARGET_KV,)) + ) + for layer_id in range(self.layer_num): + if PAYLOAD_TARGET_KV in required_payloads: + self._commit_cp_shared_l2_layer_if_needed( + reservation, PAYLOAD_TARGET_KV, layer_id + ) + if PAYLOAD_DRAFT_KV in required_payloads: + if self.draft_mem_pool_device is None: + raise RuntimeError( + "CP shared physical L2 draft commit requires draft KV device pool" + ) + draft_layer_num = int(getattr(self.draft_mem_pool_device, "layer_num", 0)) + for layer_id in range(draft_layer_num): + self._commit_cp_shared_l2_layer_if_needed( + reservation, PAYLOAD_DRAFT_KV, layer_id + ) + if PAYLOAD_INDEX_K in required_payloads: + for layer_id in self._cp_shared_l2_index_layer_ids(): + self._commit_cp_shared_l2_layer_if_needed( + reservation, PAYLOAD_INDEX_K, layer_id + ) + def submit_write_cp_all_layer(self, reservation: HiCacheWriteReservation) -> None: logger.warning( "[CP_HICACHE_FALLBACK][submit_write_cp_all_layer] " @@ -1167,14 +1525,18 @@ class HiCacheController: "node_id=%d logical_len=%d owned=%d physical=%d draft=%s", reservation.node_id, reservation.metadata.logical_len, - reservation.metadata.owned_positions.numel(), + self._reservation_owned_count(reservation), len(reservation.physical_device_indices), reservation.draft_host_indices is not None, ) if len(reservation.physical_device_indices) == 0: + # B1: zero-owned rank -- no D2H and NO gather-commit; still append the + # per-node ack so the writing_check MIN advances rank-uniformly with the + # owning ranks (MF2). The object commits via mark_object_committed at the + # MIN frontier, identically on every rank. self._append_completed_write_ack(reservation.node_id) logger.debug( - "[CacheCtrl-write] submit_write_cp_all_layer zero-owned ack: node_id=%d logical_len=%d", + "[CacheCtrl-write] submit_write_cp_all_layer zero-owned ack (B1): node_id=%d logical_len=%d", reservation.node_id, reservation.metadata.logical_len, ) @@ -1197,12 +1559,43 @@ class HiCacheController: reservation.priority, ) ) - self.start_writing() + if reservation.index_host_indices is not None: + self.index_write_queue.append( + CacheOperation( + reservation.index_host_indices, + reservation.physical_device_indices, + reservation.node_id, + reservation.priority, + ) + ) + try: + self.start_writing() + except Exception: + self.write_queue = [ + op for op in self.write_queue if reservation.node_id not in op.node_ids + ] + self.draft_write_queue = [ + op + for op in self.draft_write_queue + if reservation.node_id not in op.node_ids + ] + self.index_write_queue = [ + op + for op in self.index_write_queue + if reservation.node_id not in op.node_ids + ] + self._abort_cp_shared_l2_reservation( + reservation, fence_write_stream=True + ) + raise + # B1: NO gather-commit. The all-layer write enqueues its per-node ack via + # the write_queue/start_writing path; the object commits via the + # writing_check MIN frontier (mark_object_committed), not a quorum gather. logger.debug( - "[CacheCtrl-write] submit_write_cp_all_layer submitted: node_id=%d logical_len=%d owned=%d physical=%d draft=%s", + "[CacheCtrl-write] submit_write_cp_all_layer submitted (B1): node_id=%d logical_len=%d owned=%d physical=%d draft=%s", reservation.node_id, reservation.metadata.logical_len, - reservation.metadata.owned_positions.numel(), + self._reservation_owned_count(reservation), len(reservation.physical_device_indices), reservation.draft_host_indices is not None, ) @@ -1224,7 +1617,13 @@ class HiCacheController: if reservation.draft_host_indices is not None else 0 ) - total_layers = max(self.layer_num, draft_layer_num) + index_layer_ids = ( + self._cp_shared_l2_index_layer_ids() + if reservation.index_host_indices is not None + else () + ) + index_layer_num = (max(index_layer_ids) + 1) if index_layer_ids else 0 + total_layers = max(self.layer_num, draft_layer_num, index_layer_num) if total_layers <= 0: raise RuntimeError( f"Invalid CP HiCache per-layer backup layer count: {total_layers}" @@ -1235,27 +1634,53 @@ class HiCacheController: total_layers=total_layers, start_event=device_module.Event(), finish_event=device_module.Event(), + index_layer_ids=index_layer_ids, ) - if len(reservation.physical_device_indices) > 0: - op = CacheOperation( - reservation.host_indices, - reservation.physical_device_indices, - reservation.node_id, - reservation.priority, - ) - state.host_indices, state.physical_device_indices = self.move_indices( - op, self.mem_pool_host - ) - if reservation.draft_host_indices is not None: - draft_op = CacheOperation( - reservation.draft_host_indices, + try: + if len(reservation.physical_device_indices) > 0: + op = CacheOperation( + reservation.host_indices, reservation.physical_device_indices, reservation.node_id, reservation.priority, ) - state.draft_host_indices, _ = self.move_indices( - draft_op, self.draft_mem_pool_host + state.host_indices, state.physical_device_indices = self.move_indices( + op, self.mem_pool_host ) + if reservation.draft_host_indices is not None: + draft_op = CacheOperation( + reservation.draft_host_indices, + reservation.physical_device_indices, + reservation.node_id, + reservation.priority, + ) + state.draft_host_indices, _ = self.move_indices( + draft_op, self.draft_mem_pool_host + ) + if reservation.index_host_indices is not None: + index_op = CacheOperation( + reservation.index_host_indices, + reservation.physical_device_indices, + reservation.node_id, + reservation.priority, + ) + state.index_host_indices, _ = self.move_indices( + index_op, self.mem_pool_host + ) + else: + # Zero-owned CP ranks still participate in shared-L2 commit + # collectives. Preserve the empty payload tensors in the + # per-layer state so draft/index payloads are committed in the + # same rank-uniform order as ranks that performed a real D2H + # copy. The transfer lists below still skip them because the + # physical-device index tensor is empty. + state.host_indices = reservation.host_indices + state.physical_device_indices = reservation.physical_device_indices + state.draft_host_indices = reservation.draft_host_indices + state.index_host_indices = reservation.index_host_indices + except Exception: + self._abort_cp_shared_l2_reservation(reservation.metadata) + raise self.pending_layer_writes[reservation.node_id] = state state.start_event.record() return state @@ -1267,7 +1692,11 @@ class HiCacheController: or len(state.completed_draft_layers) >= self.draft_mem_pool_device.layer_num ) - return target_done and draft_done + index_done = ( + not state.index_layer_ids + or len(state.completed_index_layers) >= len(state.index_layer_ids) + ) + return target_done and draft_done and index_done @staticmethod def _record_tensor_on_stream(tensor: Optional[torch.Tensor], stream) -> None: @@ -1297,6 +1726,92 @@ class HiCacheController: reservation.draft_host_indices is not None, ) + def _commit_cp_shared_l2_layer_if_needed( + self, reservation: HiCacheWriteReservation, payload: str, layer_id: int + ) -> bool: + return self._commit_cp_shared_l2_layers_batched( + [reservation], payload, layer_id + ) + + def _commit_cp_shared_l2_layers_batched( + self, + reservations: List[HiCacheWriteReservation], + payload: str, + layer_id: int, + ) -> bool: + allocator = self.cp_shared_l2_page_allocator + if allocator is None: + return False + reservations_by_object_key: Dict[str, HiCacheWriteReservation] = {} + local_commits = [] + for reservation in reservations: + metadata = reservation.metadata + object_key = getattr(metadata, "object_key", None) + if object_key is None: + continue + reservations_by_object_key[object_key] = reservation + local_commits.append( + ( + object_key, + payload, + int(layer_id), + int(self.cp_shared_kv_layout.cp_rank), + ) + ) + if not local_commits: + return False + local_commits_tuple = tuple(local_commits) + if self.cp_shared_l2_commit_fn is not None: + gathered_commits = self.cp_shared_l2_commit_fn( + local_commits_tuple, + cp_cpu_group=self.cp_shared_l2_cpu_group, + rank=int(self.cp_shared_kv_layout.cp_rank), + ) + else: + gathered_commits = gather_cp_shared_l2_commits( + local_commits_tuple, + cp_cpu_group=self.cp_shared_l2_cpu_group, + ) + committed = False + for commit_object_key, commit_payload, commit_layer_id, commit_rank in ( + gathered_commits + ): + if commit_payload != payload: + continue + if allocator.get_range(commit_object_key, commit_payload) is None: + continue + committed = allocator.commit_layer( + commit_object_key, + commit_payload, + int(commit_layer_id), + int(commit_rank), + ) or committed + for object_key, reservation in reservations_by_object_key.items(): + metadata = reservation.metadata + getattr(metadata, "committed_payload_layers", {}).setdefault( + payload, set() + ).add(layer_id) + if allocator.is_committed(object_key): + metadata.committed = True + return committed + + def _abort_cp_shared_l2_states( + self, states: List[HiCacheLayerWriteState], *, fence_write_stream: bool = False + ) -> None: + seen_reservations = set() + for state in states: + reservation = state.reservation + reservation_id = id(reservation) + if reservation_id in seen_reservations: + continue + seen_reservations.add(reservation_id) + if getattr(reservation, "shared_l2_object_key", None) is None: + continue + self.pending_layer_writes.pop(reservation.node_id, None) + self._abort_cp_shared_l2_reservation( + reservation, fence_write_stream=fence_write_stream + ) + def _submit_write_cp_layer_states( self, states: List[HiCacheLayerWriteState], @@ -1304,6 +1819,7 @@ class HiCacheController: *, submit_target: bool = True, submit_draft: bool = True, + submit_index: bool = True, ) -> None: """Submit one layer for a group of reserved CP host backups. @@ -1346,12 +1862,19 @@ class HiCacheController: and layer_id < self.draft_mem_pool_device.layer_num and layer_id not in state.completed_draft_layers ] - if not target_states and not draft_states: + index_states = [ + state + for state in unique_states + if submit_index + and layer_id in state.index_layer_ids + and layer_id not in state.completed_index_layers + ] + if not target_states and not draft_states and not index_states: return active_states: List[HiCacheLayerWriteState] = [] seen_active_ids = set() - for state in target_states + draft_states: + for state in target_states + draft_states + index_states: state_id = id(state) if state_id in seen_active_ids: continue @@ -1375,72 +1898,132 @@ class HiCacheController: if state.physical_device_indices is not None and len(state.physical_device_indices) > 0 ] + index_transfer_states = [ + state + for state in index_states + if state.index_host_indices is not None + and state.physical_device_indices is not None + and len(state.physical_device_indices) > 0 + ] grouped_tensors: List[torch.Tensor] = [] final_states: List[HiCacheLayerWriteState] = [] - with device_module.stream(self.write_stream): - for state in active_states: - state.start_event.wait(self.write_stream) - layer_event.wait(self.write_stream) + try: + with device_module.stream(self.write_stream): + for state in active_states: + state.start_event.wait(self.write_stream) + layer_event.wait(self.write_stream) - if target_transfer_states: - target_host_indices = self._concat_layer_write_tensors( - [state.host_indices for state in target_transfer_states] - ) - target_device_indices = self._concat_layer_write_tensors( - [ - state.physical_device_indices - for state in target_transfer_states - ] - ) - self.mem_pool_host.backup_from_device_per_layer( - self.mem_pool_device, - target_host_indices, - target_device_indices, - layer_id, - self.io_backend, - ) - grouped_tensors.extend([target_host_indices, target_device_indices]) - if draft_transfer_states: - draft_host_indices = self._concat_layer_write_tensors( - [state.draft_host_indices for state in draft_transfer_states] - ) - draft_device_indices = self._concat_layer_write_tensors( - [ - state.physical_device_indices - for state in draft_transfer_states - ] - ) - self.draft_mem_pool_host.backup_from_device_per_layer( - self.draft_mem_pool_device, - draft_host_indices, - draft_device_indices, - layer_id, - self.io_backend, - ) - grouped_tensors.extend([draft_host_indices, draft_device_indices]) - for tensor in grouped_tensors: - self._record_tensor_on_stream(tensor, self.write_stream) + if target_transfer_states: + target_host_indices = self._concat_layer_write_tensors( + [state.host_indices for state in target_transfer_states] + ) + target_device_indices = self._concat_layer_write_tensors( + [ + state.physical_device_indices + for state in target_transfer_states + ] + ) + backup_target = self.mem_pool_host.backup_from_device_per_layer + if index_transfer_states: + backup_target = getattr( + self.mem_pool_host, "backup_kv_from_device_per_layer", None + ) + if backup_target is None: + raise RuntimeError( + "CP shared physical L2 index_k write requires " + "backup_kv_from_device_per_layer on the target host pool" + ) + backup_target( + self.mem_pool_device, + target_host_indices, + target_device_indices, + layer_id, + self.io_backend, + ) + grouped_tensors.extend([target_host_indices, target_device_indices]) - for state in target_states: - state.completed_target_layers.add(layer_id) - for state in draft_states: - state.completed_draft_layers.add(layer_id) + if index_transfer_states: + index_host_indices = self._concat_layer_write_tensors( + [state.index_host_indices for state in index_transfer_states] + ) + index_device_indices = self._concat_layer_write_tensors( + [ + state.physical_device_indices + for state in index_transfer_states + ] + ) + backup_index = getattr( + self.mem_pool_host, "backup_indexer_from_device_per_layer", None + ) + if backup_index is None: + raise RuntimeError( + "CP shared physical L2 index_k write requires " + "backup_indexer_from_device_per_layer on the target host pool" + ) + backup_index( + self.mem_pool_device, + index_host_indices, + index_device_indices, + layer_id, + self.io_backend, + ) + grouped_tensors.extend([index_host_indices, index_device_indices]) - final_states = [ - state - for state in active_states - if self._layer_write_state_done(state) and not state.ack_appended - ] - for state in final_states: - state.finish_event.record() - self._record_tensor_on_stream(state.host_indices, self.write_stream) - self._record_tensor_on_stream( - state.physical_device_indices, self.write_stream - ) - self._record_tensor_on_stream( - state.draft_host_indices, self.write_stream - ) + if draft_transfer_states: + draft_host_indices = self._concat_layer_write_tensors( + [state.draft_host_indices for state in draft_transfer_states] + ) + draft_device_indices = self._concat_layer_write_tensors( + [ + state.physical_device_indices + for state in draft_transfer_states + ] + ) + self.draft_mem_pool_host.backup_from_device_per_layer( + self.draft_mem_pool_device, + draft_host_indices, + draft_device_indices, + layer_id, + self.io_backend, + ) + grouped_tensors.extend([draft_host_indices, draft_device_indices]) + + for tensor in grouped_tensors: + self._record_tensor_on_stream(tensor, self.write_stream) + + # B1: NO per-(rank,layer,payload) gather-commit quorum. The object's + # committed/readable transition rides the writing_check ReduceOp.MIN + # frontier (mark_object_committed at _commit_pending_backup); here we + # only track per-layer completion so the per-node ack fires once all + # payloads x all layers have landed (_layer_write_state_done). + for state in target_states: + state.completed_target_layers.add(layer_id) + for state in draft_states: + state.completed_draft_layers.add(layer_id) + for state in index_states: + state.completed_index_layers.add(layer_id) + + final_states = [ + state + for state in active_states + if self._layer_write_state_done(state) and not state.ack_appended + ] + for state in final_states: + state.finish_event.record() + self._record_tensor_on_stream(state.host_indices, self.write_stream) + self._record_tensor_on_stream( + state.physical_device_indices, self.write_stream + ) + self._record_tensor_on_stream( + state.draft_host_indices, self.write_stream + ) + self._record_tensor_on_stream( + state.index_host_indices, self.write_stream + ) + except Exception: + self._abort_cp_shared_l2_states(active_states, fence_write_stream=True) + raise # Value fingerprint of the EXACT device KV being backed up, taken on the # DEFAULT stream (outside the write_stream block) so it reflects the @@ -1457,6 +2040,7 @@ class HiCacheController: *, submit_target: bool = True, submit_draft: bool = True, + submit_index: bool = True, ) -> None: """Submit one layer of a reserved CP host backup. @@ -1470,6 +2054,7 @@ class HiCacheController: layer_id, submit_target=submit_target, submit_draft=submit_draft, + submit_index=submit_index, ) def submit_write_cp_per_layer( @@ -1556,6 +2141,281 @@ class HiCacheController: ) return device_indices + def _shared_l2_current_page_owners(self, num_pages: int) -> List[int]: + layout = self.cp_shared_kv_layout + if layout is None: + raise RuntimeError("CP shared physical L2 load requires CP layout") + if num_pages < 0: + raise ValueError(f"num_pages must be non-negative, got {num_pages}") + if num_pages == 0: + return [] + logical_pages = torch.arange(1, num_pages + 1, dtype=torch.int64) + return [ + int(owner) + for owner in layout.owner_for_logical_pages(logical_pages).tolist() + ] + + def _is_cp_shared_l2_load_metadata(self, metadata) -> bool: + return hasattr(metadata, "object_ranges") and not hasattr( + metadata, "page_owners" + ) + + def _validate_shared_l2_load_metadata_payloads(self, meta, node) -> None: + required_payloads = set( + getattr(meta, "required_payloads", (PAYLOAD_TARGET_KV,)) + or (PAYLOAD_TARGET_KV,) + ) + range_payloads = set(getattr(meta, "object_ranges", {}).keys()) + payloads = required_payloads | range_payloads + extra_ranges = sorted(range_payloads - required_payloads) + if extra_ranges: + raise RuntimeError( + "CP shared physical L2 load metadata has non-required payload ranges: " + f"node_id={getattr(node, 'id', '?')} payloads={extra_ranges}" + ) + if PAYLOAD_INDEX_K in payloads and not self._cp_shared_l2_index_payload_active(): + raise RuntimeError( + "CP shared physical L2 load index_k payload requires active NSA index layers: " + f"node_id={getattr(node, 'id', '?')} payloads={sorted(payloads)}" + ) + if PAYLOAD_DRAFT_KV in payloads and not self.has_draft_hicache: + raise RuntimeError( + "CP shared physical L2 load unsupported draft_kv payload because draft " + f"HiCache is not attached: node_id={getattr(node, 'id', '?')}" + ) + unsupported = sorted( + payloads - {PAYLOAD_TARGET_KV, PAYLOAD_DRAFT_KV, PAYLOAD_INDEX_K} + ) + if unsupported: + raise RuntimeError( + "CP shared physical L2 load has unsupported payloads: " + f"node_id={getattr(node, 'id', '?')} payloads={unsupported}" + ) + + def load_cp_shared_l2( + self, nodes_to_load, node_id: int = -1 + ) -> Optional[torch.Tensor]: + start_time = time.perf_counter() + layout = self.cp_shared_kv_layout + if layout is None: + raise RuntimeError("CP shared physical L2 load requires CP layout") + + metas = [] + total_padded_len = 0 + total_pages = 0 + for node in nodes_to_load: + meta = node.cp_hicache + if meta is None: + raise RuntimeError( + f"load_cp_shared_l2 called with node {getattr(node, 'id', '?')} " + "that has no cp_hicache metadata" + ) + if not self._is_cp_shared_l2_load_metadata(meta): + raise RuntimeError( + "CP shared physical L2 load requires all nodes to carry " + f"shared L2 metadata: node_id={getattr(node, 'id', '?')}" + ) + self._validate_shared_l2_load_metadata_payloads(meta, node) + valid_len = int(getattr(meta, "logical_len", node.host_len)) + padded_len = int(getattr(meta, "padded_len", node.host_len)) + page_size = int(getattr(meta, "page_size", self.page_size)) + if page_size != self.page_size: + raise RuntimeError( + "CP shared physical L2 metadata page size mismatch: " + f"node_id={getattr(node, 'id', '?')} metadata={page_size} " + f"controller={self.page_size}" + ) + if valid_len != int(node.host_len): + raise RuntimeError( + "CP shared physical L2 load metadata length mismatch: " + f"node_id={getattr(node, 'id', '?')} host_len={node.host_len} " + f"logical_len={valid_len}" + ) + if padded_len < valid_len or padded_len % self.page_size != 0: + raise RuntimeError( + "CP shared physical L2 load metadata has invalid padded length: " + f"node_id={getattr(node, 'id', '?')} logical_len={valid_len} " + f"padded_len={padded_len} page_size={self.page_size}" + ) + object_ranges = getattr(meta, "object_ranges", {}) + object_range = object_ranges.get(PAYLOAD_TARGET_KV) + if object_range is None: + raise RuntimeError( + "CP shared physical L2 target payload range missing: " + f"node_id={getattr(node, 'id', '?')}" + ) + expected_range_pages = padded_len // self.page_size + if int(object_range.num_pages) != expected_range_pages: + raise RuntimeError( + "CP shared physical L2 target payload range size mismatch: " + f"node_id={getattr(node, 'id', '?')} " + f"range_pages={int(object_range.num_pages)} " + f"expected_pages={expected_range_pages} " + f"padded_len={padded_len} page_size={self.page_size}" + ) + for payload_kind in getattr(meta, "required_payloads", (PAYLOAD_TARGET_KV,)): + payload_range = object_ranges.get(payload_kind) + if payload_range is None: + raise RuntimeError( + "CP shared physical L2 required payload range missing: " + f"node_id={getattr(node, 'id', '?')} payload={payload_kind}" + ) + if int(payload_range.num_pages) != expected_range_pages: + raise RuntimeError( + "CP shared physical L2 payload range size mismatch: " + f"node_id={getattr(node, 'id', '?')} payload={payload_kind} " + f"range_pages={int(payload_range.num_pages)} " + f"expected_pages={expected_range_pages}" + ) + metas.append((node, meta, valid_len, padded_len, object_ranges)) + total_padded_len += padded_len + total_pages += padded_len // self.page_size + + current_page_owners = self._shared_l2_current_page_owners(total_pages) + alloc_for_shared = getattr( + self.mem_pool_device_allocator, "alloc_pages_for_shared_l2_load", None + ) + if alloc_for_shared is None: + raise RuntimeError( + "CP shared physical L2 load requires allocator " + "alloc_pages_for_shared_l2_load(current_page_owners)" + ) + device_indices = alloc_for_shared(current_page_owners) + if device_indices is None: + _cp_shared_kv_bs_gt1_cache_timing( + "hicache_load_cp_shared_l2_plan", + start_time, + "node_id=%d result=alloc_none pages=%d", + node_id, + total_pages, + ) + self._cp_shared_l2_stats["cp_shared_l2_load_misses"] += 1 + return None + + try: + if device_indices.numel() != total_padded_len: + raise RuntimeError( + "alloc_pages_for_shared_l2_load returned unexpected length: " + f"got {device_indices.numel()}, expected {total_padded_len}" + ) + + host_chunks = [] + draft_host_chunks = [] + index_host_chunks = [] + physical_chunks = [] + visible_chunks = [] + offset = 0 + for _node, _meta, valid_len, padded_len, object_ranges in metas: + node_device_indices = device_indices[offset : offset + padded_len] + offset += padded_len + visible_chunks.append(node_device_indices[:valid_len]) + if padded_len == 0: + continue + owned_mask = layout.owned_by_this_rank(node_device_indices) + owned_positions = owned_mask.nonzero(as_tuple=True)[0] + if owned_positions.numel() == 0: + continue + selected_logical_locs = node_device_indices[owned_mask] + physical_chunks.append( + layout.logical_locs_to_physical(selected_logical_locs) + ) + cpu_owned_positions = owned_positions.to(device="cpu", dtype=torch.int64) + target_range = object_ranges[PAYLOAD_TARGET_KV] + host_chunks.append( + cp_shared_l2_logical_token_indices( + target_range, self.page_size, cpu_owned_positions + ) + ) + draft_range = object_ranges.get(PAYLOAD_DRAFT_KV) + if draft_range is not None: + draft_host_chunks.append( + cp_shared_l2_logical_token_indices( + draft_range, self.page_size, cpu_owned_positions + ) + ) + index_range = object_ranges.get(PAYLOAD_INDEX_K) + if index_range is not None: + index_host_chunks.append( + cp_shared_l2_logical_token_indices( + index_range, self.page_size, cpu_owned_positions + ) + ) + + visible_device_indices = ( + torch.cat(visible_chunks) + if visible_chunks + else torch.empty( + (0,), dtype=device_indices.dtype, device=device_indices.device + ) + ) + + if not host_chunks: + empty_host = torch.empty((0,), dtype=torch.int64) + empty_device = torch.empty( + (0,), dtype=device_indices.dtype, device=device_indices.device + ) + self.load_queue.append(CacheOperation(empty_host, empty_device, node_id)) + if self.has_draft_hicache and any( + PAYLOAD_DRAFT_KV in getattr(meta, "required_payloads", ()) + for _node, meta, _valid_len, _padded_len, _ranges in metas + ): + self.draft_load_queue.append( + CacheOperation(empty_host.clone(), empty_device.clone(), node_id) + ) + if self._cp_shared_l2_index_payload_active() and any( + PAYLOAD_INDEX_K in getattr(meta, "required_payloads", ()) + for _node, meta, _valid_len, _padded_len, _ranges in metas + ): + self.index_load_queue.append( + CacheOperation(empty_host.clone(), empty_device.clone(), node_id) + ) + self._cp_shared_l2_stats["cp_shared_l2_load_hits"] += 1 + return visible_device_indices + + host_indices = torch.cat(host_chunks) + physical_device_indices = torch.cat(physical_chunks) + self._validate_cp_hicache_page_indices( + host_indices, physical_device_indices + ) + draft_host_indices = None + if draft_host_chunks: + draft_host_indices = torch.cat(draft_host_chunks) + self._validate_cp_hicache_page_indices( + draft_host_indices, physical_device_indices + ) + index_host_indices = None + if index_host_chunks: + index_host_indices = torch.cat(index_host_chunks) + self._validate_cp_hicache_page_indices( + index_host_indices, physical_device_indices + ) + self.load_queue.append( + CacheOperation(host_indices, physical_device_indices, node_id) + ) + if draft_host_indices is not None: + self.draft_load_queue.append( + CacheOperation(draft_host_indices, physical_device_indices, node_id) + ) + if index_host_indices is not None: + self.index_load_queue.append( + CacheOperation(index_host_indices, physical_device_indices, node_id) + ) + self._cp_shared_l2_stats["cp_shared_l2_load_hits"] += 1 + _cp_shared_kv_bs_gt1_cache_timing( + "hicache_load_cp_shared_l2_plan", + start_time, + "node_id=%d result=queued pages=%d host_indices=%d physical_indices=%d visible_indices=%d", + node_id, + total_pages, + int(host_indices.numel()), + int(physical_device_indices.numel()), + int(visible_device_indices.numel()), + ) + return visible_device_indices + except Exception: + self.mem_pool_device_allocator.free(device_indices) + raise + def load_cp(self, nodes_to_load, node_id: int = -1) -> Optional[torch.Tensor]: start_time = time.perf_counter() stage_start_time = start_time @@ -1584,6 +2444,8 @@ class HiCacheController: f"load_cp called with node {getattr(node, 'id', '?')} " "that has no cp_hicache metadata" ) + if not hasattr(meta, "page_owners"): + return self.load_cp_shared_l2(nodes_to_load, node_id=node_id) # page_owners is CPU int8; .tolist() returns list[int] directly. page_owners.extend(meta.page_owners.tolist()) record_stage("collect_page_owners") @@ -1803,7 +2665,11 @@ class HiCacheController: raise ValueError(f"Unsupported io backend") def start_loading(self) -> int: - if len(self.load_queue) == 0 and len(self.draft_load_queue) == 0: + if ( + len(self.load_queue) == 0 + and len(self.draft_load_queue) == 0 + and len(self.index_load_queue) == 0 + ): return -1 def begin_load_op(mem_pool_host, host_indices, device_indices) -> bool: @@ -1825,7 +2691,13 @@ class HiCacheController: if self.draft_load_queue else None ) - node_ids = op.node_ids if op is not None else draft_op.node_ids + index_op = ( + CacheOperation.merge_ops(self.index_load_queue) + if self.index_load_queue + else None + ) + primary_op = op or draft_op or index_op + node_ids = primary_op.node_ids if op is not None: host_indices, device_indices = self.move_indices(op, self.mem_pool_host) else: @@ -1836,8 +2708,15 @@ class HiCacheController: ) else: draft_host_indices = draft_device_indices = None + if index_op is not None: + index_host_indices, index_device_indices = self.move_indices( + index_op, self.mem_pool_host + ) + else: + index_host_indices = index_device_indices = None self.load_queue.clear() self.draft_load_queue.clear() + self.index_load_queue.clear() producer_event = self.layer_done_counter.events[producer_id] producer_event.start_event.record() @@ -1871,11 +2750,19 @@ class HiCacheController: draft_layer_num = ( self.draft_mem_pool_device.layer_num if draft_op is not None else 0 ) + index_layer_ids = ( + set(self._cp_shared_l2_index_layer_ids()) + if index_op is not None + else set() + ) + index_layer_num = (max(index_layer_ids) + 1) if index_layer_ids else 0 layer_loop_start = _cp_shared_kv_bs_gt1_cache_timing_start() - for i in range(max(self.layer_num, draft_layer_num)): + total_layers = max(self.layer_num, draft_layer_num, index_layer_num) + for i in range(total_layers): layer_submit_start = _cp_shared_kv_bs_gt1_cache_timing_start() submitted_target = False submitted_draft = False + submitted_index = False if draft_op is not None and i < draft_layer_num: if len(draft_host_indices) > 0: self.draft_mem_pool_host.load_to_device_per_layer( @@ -1888,7 +2775,19 @@ class HiCacheController: submitted_draft = True if op is not None and i < self.layer_num: if len(host_indices) > 0: - self.mem_pool_host.load_to_device_per_layer( + load_target = self.mem_pool_host.load_to_device_per_layer + if index_op is not None: + load_target = getattr( + self.mem_pool_host, + "load_kv_to_device_per_layer", + None, + ) + if load_target is None: + raise RuntimeError( + "CP shared physical L2 index_k load requires " + "load_kv_to_device_per_layer on the target host pool" + ) + load_target( self.mem_pool_device, host_indices, device_indices, @@ -1896,35 +2795,58 @@ class HiCacheController: self.io_backend, ) submitted_target = True - # Value fingerprint of what actually landed on device - # after H2D (on load_stream, in-order after the load). - # Compare to backup_kv_hash(node,layer): inequality = - # round-trip delivered wrong KV. node_ids is the merged - # op's nodes (single node in the c=1 rehit pass). + if index_op is not None and i in index_layer_ids: + if len(index_host_indices) > 0: + load_index = getattr( + self.mem_pool_host, + "load_indexer_to_device_per_layer", + None, + ) + if load_index is None: + raise RuntimeError( + "CP shared physical L2 index_k load requires " + "load_indexer_to_device_per_layer on the target host pool" + ) + load_index( + self.mem_pool_device, + index_host_indices, + index_device_indices, + i, + self.io_backend, + ) + submitted_index = True + if op is not None and i < self.layer_num: producer_event.complete(i) elif op is None and i < self.layer_num: producer_event.complete(i) _cp_shared_kv_bs_gt1_cache_timing_log( "submit_h2d_layer_per_call_slow", layer_submit_start, - "layer_id=%s target=%s draft=%s target_tokens=%s draft_tokens=%s", + "layer_id=%s target=%s draft=%s index=%s target_tokens=%s draft_tokens=%s index_tokens=%s", i, submitted_target, submitted_draft, + submitted_index, int(host_indices.numel()) if host_indices is not None else 0, int(draft_host_indices.numel()) if draft_host_indices is not None else 0, + int(index_host_indices.numel()) + if index_host_indices is not None + else 0, ) _cp_shared_kv_bs_gt1_cache_timing_log( "submit_h2d_layer_loop", layer_loop_start, - "layers=%s target_tokens=%s draft_tokens=%s", - max(self.layer_num, draft_layer_num), + "layers=%s target_tokens=%s draft_tokens=%s index_tokens=%s", + total_layers, int(host_indices.numel()) if host_indices is not None else 0, int(draft_host_indices.numel()) if draft_host_indices is not None else 0, + int(index_host_indices.numel()) + if index_host_indices is not None + else 0, ) # NOTE: We must save the host indices and device indices here, # this is because we need to guarantee that these tensors are @@ -1937,6 +2859,10 @@ class HiCacheController: draft_host_indices.record_stream(self.load_stream) if draft_device_indices is not None and draft_device_indices.is_cuda: draft_device_indices.record_stream(self.load_stream) + if index_host_indices is not None and index_host_indices.is_cuda: + index_host_indices.record_stream(self.load_stream) + if index_device_indices is not None and index_device_indices.is_cuda: + index_device_indices.record_stream(self.load_stream) finally: end_start = _cp_shared_kv_bs_gt1_cache_timing_start() if draft_load_op_started: @@ -1975,6 +2901,10 @@ class HiCacheController: if not backup_only: raise ValueError("Other eviction policies are not supported yet.") + if isinstance(metadata, CpSharedL2NodeMetadata): + self._release_cp_shared_l2_metadata(metadata) + return int(getattr(metadata, "padded_len", 0)) + draft_host_indices = getattr(metadata, "draft_host_indices", None) if self.has_draft_hicache and draft_host_indices is None: raise RuntimeError( diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 96ba46651..d9732e64d 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -66,34 +66,133 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -def _estimate_hicache_size_per_token(kv_cache) -> int: +def _estimate_nsa_index_size_per_token(kv_cache) -> int: + if not isinstance(kv_cache, NSATokenToKVPool): + raise ValueError( + f"NSA index size is only defined for NSATokenToKVPool; " + f"got {type(kv_cache).__name__}" + ) + indexer_size_per_token = ( + kv_cache.index_head_dim + + kv_cache.index_head_dim // kv_cache.quant_block_size * 4 + ) + index_layer_count = len( + getattr( + kv_cache, + "index_active_layer_ids", + range(kv_cache.layer_num), + ) + ) + return ( + indexer_size_per_token + * index_layer_count + * NSATokenToKVPool.index_k_with_scale_buffer_dtype.itemsize + ) + + +def _estimate_hicache_kv_size_per_token(kv_cache) -> int: dtype_size = kv_cache.store_dtype.itemsize if isinstance(kv_cache, MHATokenToKVPool): return kv_cache.head_dim * kv_cache.head_num * kv_cache.layer_num * dtype_size * 2 - if isinstance(kv_cache, NSATokenToKVPool): - indexer_size_per_token = ( - kv_cache.index_head_dim - + kv_cache.index_head_dim // kv_cache.quant_block_size * 4 - ) - index_layer_count = len( - getattr( - kv_cache, - "index_active_layer_ids", - range(kv_cache.layer_num), - ) - ) - return ( - kv_cache.kv_cache_dim * dtype_size * kv_cache.layer_num - + indexer_size_per_token - * index_layer_count - * NSATokenToKVPool.index_k_with_scale_buffer_dtype.itemsize - ) if isinstance(kv_cache, MLATokenToKVPool): - kv_cache_dim = kv_cache.kv_lora_rank + kv_cache.qk_rope_head_dim + kv_cache_dim = getattr( + kv_cache, + "kv_cache_dim", + kv_cache.kv_lora_rank + kv_cache.qk_rope_head_dim, + ) return kv_cache_dim * dtype_size * kv_cache.layer_num raise ValueError(f"Unsupported HiCache KV pool type: {type(kv_cache).__name__}") +def _estimate_hicache_size_per_token(kv_cache) -> int: + size_per_token = _estimate_hicache_kv_size_per_token(kv_cache) + if isinstance(kv_cache, NSATokenToKVPool): + size_per_token += _estimate_nsa_index_size_per_token(kv_cache) + return size_per_token + + +def _get_cp_shared_l2_rank_and_group(token_to_kv_pool_allocator): + cp_rank = getattr(token_to_kv_pool_allocator, "cp_rank", None) + cp_size = getattr(token_to_kv_pool_allocator, "cp_size", None) + if cp_rank is None or cp_size is None: + raise ValueError( + "[CP_SHARED_L2_FAILFAST][missing_cp_layout] " + "CP shared physical L2 HiCache requires an allocator with " + "explicit cp_rank/cp_size; refusing rank-local host allocation." + ) + + cp_group = None + cp_cpu_group = None + if int(cp_size) > 1: + try: + from sglang.srt.layers.dp_attention import get_attention_cp_group + + cp_group = get_attention_cp_group() + except Exception as exc: + raise ValueError( + "[CP_SHARED_L2_FAILFAST][missing_cp_group] " + "CP shared physical L2 HiCache requires an attention CP " + "group so rank 0 can broadcast one shared slab handle." + ) from exc + if cp_group is None: + raise ValueError( + "[CP_SHARED_L2_FAILFAST][missing_cp_group] " + "CP shared physical L2 HiCache requires a non-null attention " + "CP group for multi-rank CP." + ) + cp_cpu_group = getattr(cp_group, "cpu_group", cp_group) + if cp_cpu_group is None: + raise ValueError( + "[CP_SHARED_L2_FAILFAST][missing_cp_group] " + "CP shared physical L2 HiCache requires a usable CP CPU " + "group for multi-rank CP." + ) + return int(cp_rank), int(cp_size), cp_cpu_group + + +def _make_cp_shared_l2_host_tensor_allocator( + *, + server_args, + token_to_kv_pool_allocator, + payload_kind: str, + slab_name_suffix: str, + slabs_by_payload=None, + validate_production: bool = True, +): + from sglang.srt.mem_cache.memory_pool_host import ( + SharedHostTensorAllocator, + SharedHostTensorGroupAllocator, + ) + + cp_rank, _cp_size, cp_cpu_group = _get_cp_shared_l2_rank_and_group( + token_to_kv_pool_allocator + ) + slab_name = f"sglang-cp-shared-l2-{payload_kind}-{slab_name_suffix}" + slabs = tuple((slabs_by_payload or {}).get(payload_kind, ())) + common_kwargs = dict( + directory=server_args.cp_shared_l2_hugetlbfs_dir, + name=slab_name, + creator_rank=0, + validate_production=validate_production, + rank=cp_rank, + source_rank=0, + cp_cpu_group=cp_cpu_group, + ) + if len(slabs) > 1: + return SharedHostTensorGroupAllocator(slabs=slabs, **common_kwargs) + return SharedHostTensorAllocator(**common_kwargs) + + +def _cp_shared_l2_host_slab_suffix( + kv_cache, *, page_size: int, draft: bool = False +) -> str: + return ( + f"{'draft' if draft else 'target'}-" + f"layers{getattr(kv_cache, 'start_layer', 0)}-" + f"{getattr(kv_cache, 'end_layer', 0)}-" + f"ps{int(page_size)}" + ) + def _compute_shared_hicache_token_capacities( *, total_host_bytes: int, @@ -137,6 +236,113 @@ def _compute_shared_hicache_token_capacities( return target_tokens, draft_tokens +def _cp_shared_l2_bytes_per_page_by_payload( + *, + kv_cache, + page_size: int, + draft_token_to_kv_pool=None, +) -> Dict[str, int]: + from sglang.srt.mem_cache.cp_shared_l2_pool import ( + PAYLOAD_DRAFT_KV, + PAYLOAD_INDEX_K, + PAYLOAD_TARGET_KV, + ) + + page_size = int(page_size) + if page_size <= 0: + raise ValueError(f"page_size must be positive, got {page_size}") + bytes_per_page = { + PAYLOAD_TARGET_KV: _estimate_hicache_kv_size_per_token(kv_cache) * page_size + } + if draft_token_to_kv_pool is not None: + bytes_per_page[PAYLOAD_DRAFT_KV] = ( + _estimate_hicache_kv_size_per_token(draft_token_to_kv_pool) * page_size + ) + if isinstance(kv_cache, NSATokenToKVPool): + bytes_per_page[PAYLOAD_INDEX_K] = ( + _estimate_nsa_index_size_per_token(kv_cache) * page_size + ) + return bytes_per_page + + +def _cp_shared_l2_slab_pages_by_payload( + *, + server_args, + page_size: int, + bytes_per_page_by_payload: Dict[str, int], +) -> Dict[str, int]: + slab_size_gb = int(getattr(server_args, "cp_shared_l2_slab_size_gb", 0) or 0) + if slab_size_gb <= 0: + return {} + slab_size_bytes = slab_size_gb * 1000 * 1000 * 1000 + slab_pages: Dict[str, int] = {} + for payload_kind, bytes_per_page in bytes_per_page_by_payload.items(): + bytes_per_page = int(bytes_per_page) + if bytes_per_page <= 0: + raise ValueError(f"bytes_per_page for {payload_kind!r} must be positive") + if slab_size_bytes < bytes_per_page: + suggested_gb = math.ceil(bytes_per_page / 1e9) + raise ValueError( + "cp_shared_l2_slab_size_gb is smaller than one logical page for " + f"payload {payload_kind!r}: configured_bytes={slab_size_bytes} " + f"bytes_per_page={bytes_per_page}. Increase " + f"--cp-shared-l2-slab-size-gb to at least {suggested_gb}." + ) + slab_pages[payload_kind] = slab_size_bytes // bytes_per_page + return slab_pages + + +def _cp_shared_l2_pages_for_kv_cache_host( + *, + kv_cache, + server_args, + page_size: int, + host_size: int, + host_token_capacity: Optional[int], +) -> int: + if host_token_capacity is not None: + tokens = int(host_token_capacity) + return max(1, (tokens + int(page_size) - 1) // int(page_size)) + size_per_token = _estimate_hicache_size_per_token(kv_cache) + if int(host_size) > 0: + tokens = int(int(host_size) * 1e9 // size_per_token) + else: + tokens = int(kv_cache.size * server_args.hicache_ratio) + return max(1, tokens // int(page_size) + 1) + + +def _cp_shared_l2_pages_per_payload( + *, + token_to_kv_pool_host, + page_size: int, + draft_token_to_kv_pool=None, + draft_host_token_capacity: Optional[int] = None, +) -> Dict[str, int]: + from sglang.srt.mem_cache.cp_shared_l2_pool import ( + PAYLOAD_DRAFT_KV, + PAYLOAD_INDEX_K, + PAYLOAD_TARGET_KV, + ) + + target_pages = max( + 1, int(getattr(token_to_kv_pool_host, "size", 0)) // int(page_size) + ) + pages_per_payload = {PAYLOAD_TARGET_KV: target_pages} + if draft_token_to_kv_pool is not None: + draft_tokens = ( + int(draft_host_token_capacity) + if draft_host_token_capacity is not None + else int(getattr(token_to_kv_pool_host, "size", 0)) + ) + pages_per_payload[PAYLOAD_DRAFT_KV] = max(1, draft_tokens // int(page_size)) + if int(getattr(token_to_kv_pool_host, "index_active_layer_num", 0) or 0) > 0: + pages_per_payload[PAYLOAD_INDEX_K] = max( + 1, + int(getattr(token_to_kv_pool_host, "indexer_page_num", target_pages)), + ) + return pages_per_payload + + def _page_aligned_room_from_ratio(capacity: int, ratio: float, page_size: int) -> int: if page_size <= 0: raise ValueError(f"page_size must be positive, got {page_size}") @@ -376,6 +582,142 @@ class CpHiCacheCapacitySnapshot: ) +class _CpRunningHostCounts: + """Running per-owner host TOKEN counts for the CP-HiCache capacity snapshot. + + Equals, by construction, the full radix-walk result -- committed nodes and + pending backups, summed per CP owner lane, gated on draft presence -- but is + maintained incrementally at each state transition so the snapshot is + O(cp_size) rather than O(total-cached-pages) per write-admission. Created + EMPTY (production starts with an empty tree) and then maintained by the + accounting hooks; ``_cp_walk_capacity_counts`` is the reference the env-gated + drift verifier compares against, NOT a lazy-build source. + """ + + __slots__ = ( + "cp_size", + "committed_target", + "committed_draft", + "pending_target", + "pending_draft", + ) + + def __init__(self, cp_size: int): + self.cp_size = int(cp_size) + self.committed_target = [0] * self.cp_size + self.committed_draft = [0] * self.cp_size + self.pending_target = [0] * self.cp_size + self.pending_draft = [0] * self.cp_size + + def set_from_tuples( + self, committed_target, committed_draft, pending_target, pending_draft + ) -> None: + self.committed_target = list(committed_target) + self.committed_draft = list(committed_draft) + self.pending_target = list(pending_target) + self.pending_draft = list(pending_draft) + + @staticmethod + def _token_counts(metadata, cp_size: int) -> List[int]: + # owner_page_counts caches on the metadata object (and survives commit -- + # pending.metadata IS node.cp_hicache -- and split, which builds fresh + # objects with fresh caches); tokens == pages * page_size. + page_counts = metadata.owner_page_counts(cp_size) + if len(page_counts) != cp_size: + raise RuntimeError( + "CP HiCache running-count cp_size mismatch: " + f"metadata={len(page_counts)} expected={cp_size}" + ) + page_size = int(metadata.page_size) + return [int(c) * page_size for c in page_counts] + + def _add(self, target, draft, metadata) -> None: + counts = self._token_counts(metadata, self.cp_size) + for i, c in enumerate(counts): + target[i] += c + if getattr(metadata, "draft_host_indices", None) is not None: + for i, c in enumerate(counts): + draft[i] += c + + def _sub(self, target, draft, metadata) -> None: + counts = self._token_counts(metadata, self.cp_size) + has_draft = getattr(metadata, "draft_host_indices", None) is not None + # Underflow is a hard error, not a floor: the full walk can never go + # negative, so a negative here means a missed add or a double sub. + for i, c in enumerate(counts): + new_t = target[i] - c + if new_t < 0: + raise RuntimeError( + f"CP HiCache running token underflow at lane {i}: " + f"{target[i]} - {c} < 0" + ) + target[i] = new_t + if has_draft: + new_d = draft[i] - c + if new_d < 0: + raise RuntimeError( + f"CP HiCache running draft underflow at lane {i}: " + f"{draft[i]} - {c} < 0" + ) + draft[i] = new_d + + def enter_committed(self, metadata) -> None: + self._add(self.committed_target, self.committed_draft, metadata) + + def leave_committed(self, metadata) -> None: + self._sub(self.committed_target, self.committed_draft, metadata) + + def enter_pending(self, metadata) -> None: + self._add(self.pending_target, self.pending_draft, metadata) + + def leave_pending(self, metadata) -> None: + self._sub(self.pending_target, self.pending_draft, metadata) + + def reset(self) -> None: + n = self.cp_size + self.committed_target = [0] * n + self.committed_draft = [0] * n + self.pending_target = [0] * n + self.pending_draft = [0] * n + + def as_tuples(self): + return ( + tuple(self.committed_target), + tuple(self.committed_draft), + tuple(self.pending_target), + tuple(self.pending_draft), + ) + + def sanity_check(self) -> None: + """Cheap always-on invariants (O(cp_size)): non-negativity and + draft<=target lane-wise (draft counts are a subset of target counts for + the same metadata). Does NOT bound per-lane sums against pool .size -- + availability floors at zero (`_cp_sub_counts_floor_zero`) so used may + legitimately exceed capacity. The env-gated full-walk verifier is the + real drift net.""" + for name, vec in ( + ("committed_target", self.committed_target), + ("committed_draft", self.committed_draft), + ("pending_target", self.pending_target), + ("pending_draft", self.pending_draft), + ): + for i, v in enumerate(vec): + if v < 0: + raise RuntimeError( + f"CP HiCache running {name} lane {i} negative: {v}" + ) + for i in range(self.cp_size): + if ( + self.committed_draft[i] > self.committed_target[i] + or self.pending_draft[i] > self.pending_target[i] + ): + raise RuntimeError( + f"CP HiCache running draft>target lane {i}: " + f"cd={self.committed_draft[i]} ct={self.committed_target[i]} " + f"pd={self.pending_draft[i]} pt={self.pending_target[i]}" + ) + + @dataclass(frozen=True) class CpHiCacheEvictionPlan: victims: Tuple[TreeNode, ...] @@ -432,7 +774,14 @@ class HiRadixCache(RadixCache): *, host_size: Optional[int] = None, host_token_capacity: Optional[int] = None, + shared_l2_payload_kind: Optional[str] = None, + slabs_by_payload=None, ): + from sglang.srt.mem_cache.cp_shared_l2_pool import ( + PAYLOAD_DRAFT_KV, + PAYLOAD_INDEX_K, + PAYLOAD_TARGET_KV, + ) from sglang.srt.mem_cache.memory_pool_host import ( MHATokenToKVPoolHost, MLATokenToKVPoolHost, @@ -440,6 +789,39 @@ class HiRadixCache(RadixCache): ) host_size = server_args.hicache_size if host_size is None else host_size + host_tensor_allocator = None + index_host_tensor_allocator = None + if getattr(server_args, "enable_cp_shared_physical_l2_hicache", False): + payload_kind = shared_l2_payload_kind or PAYLOAD_TARGET_KV + suffix = _cp_shared_l2_host_slab_suffix( + kv_cache, + page_size=self.page_size, + draft=(payload_kind == PAYLOAD_DRAFT_KV), + ) + host_tensor_allocator = _make_cp_shared_l2_host_tensor_allocator( + server_args=server_args, + token_to_kv_pool_allocator=self._token_to_kv_pool_allocator, + payload_kind=payload_kind, + slab_name_suffix=suffix, + slabs_by_payload=slabs_by_payload, + ) + if isinstance(kv_cache, NSATokenToKVPool): + index_host_tensor_allocator = _make_cp_shared_l2_host_tensor_allocator( + server_args=server_args, + token_to_kv_pool_allocator=self._token_to_kv_pool_allocator, + payload_kind=PAYLOAD_INDEX_K, + slab_name_suffix=suffix, + slabs_by_payload=slabs_by_payload, + ) + logger.info( + "[CP_SHARED_L2] constructing shared physical host L2 pool: " + "hicache_size_gb=%s is_node_local_total=true payload=%s index_payload=%s slabs=%s numa_policy=%s", + server_args.hicache_size, + payload_kind, + index_host_tensor_allocator is not None, + {key: len(value) for key, value in (slabs_by_payload or {}).items()}, + getattr(server_args, "cp_shared_l2_numa_policy", "interleave_2m"), + ) if isinstance(kv_cache, MHATokenToKVPool): return MHATokenToKVPoolHost( kv_cache, @@ -449,6 +831,7 @@ class HiRadixCache(RadixCache): server_args.hicache_mem_layout, allocator_type=server_args.hicache_storage_backend, host_token_capacity=host_token_capacity, + host_tensor_allocator=host_tensor_allocator, ) if isinstance(kv_cache, NSATokenToKVPool): return NSATokenToKVPoolHost( @@ -459,6 +842,8 @@ class HiRadixCache(RadixCache): server_args.hicache_mem_layout, allocator_type=server_args.hicache_storage_backend, host_token_capacity=host_token_capacity, + host_tensor_allocator=host_tensor_allocator, + index_host_tensor_allocator=index_host_tensor_allocator, ) if isinstance(kv_cache, MLATokenToKVPool): return MLATokenToKVPoolHost( @@ -469,6 +854,7 @@ class HiRadixCache(RadixCache): server_args.hicache_mem_layout, allocator_type=server_args.hicache_storage_backend, host_token_capacity=host_token_capacity, + host_tensor_allocator=host_tensor_allocator, ) raise ValueError("HiRadixCache only supports MHA, MLA, and NSA KV pools") @@ -487,6 +873,8 @@ class HiRadixCache(RadixCache): self._server_args, host_size=0 if host_token_capacity is not None else None, host_token_capacity=host_token_capacity, + shared_l2_payload_kind="draft_kv", + slabs_by_payload=getattr(self, "_cp_shared_l2_slabs_by_payload", None), ) self.cache_controller.attach_draft_pool( draft_token_to_kv_pool, draft_token_to_kv_pool_host @@ -523,6 +911,7 @@ class HiRadixCache(RadixCache): "CP shared KV HiCache host integration requires NSATokenToKVPool." ) self.kv_cache = params.token_to_kv_pool_allocator.get_kvcache() + self._token_to_kv_pool_allocator = params.token_to_kv_pool_allocator draft_token_to_kv_pool = getattr(params, "draft_token_to_kv_pool", None) target_host_token_capacity = None draft_host_token_capacity = None @@ -553,11 +942,77 @@ class HiRadixCache(RadixCache): target_host_token_capacity, draft_host_token_capacity, ) + cp_shared_l2_slabs_by_payload = None + cp_shared_l2_pages_per_payload_planned = None + if getattr(server_args, "enable_cp_shared_physical_l2_hicache", False): + from sglang.srt.mem_cache.cp_shared_l2_pool import ( + PAYLOAD_DRAFT_KV, + PAYLOAD_INDEX_K, + PAYLOAD_TARGET_KV, + build_cp_shared_l2_slabs_by_payload, + ) + + target_host_size = ( + 0 if target_host_token_capacity is not None else server_args.hicache_size + ) + target_pages = _cp_shared_l2_pages_for_kv_cache_host( + kv_cache=self.kv_cache, + server_args=server_args, + page_size=self.page_size, + host_size=target_host_size, + host_token_capacity=target_host_token_capacity, + ) + pages_per_payload = {PAYLOAD_TARGET_KV: target_pages} + bytes_per_page_by_payload = _cp_shared_l2_bytes_per_page_by_payload( + kv_cache=self.kv_cache, + page_size=self.page_size, + draft_token_to_kv_pool=draft_token_to_kv_pool, + ) + if draft_token_to_kv_pool is not None: + draft_pages = _cp_shared_l2_pages_for_kv_cache_host( + kv_cache=draft_token_to_kv_pool, + server_args=server_args, + page_size=self.page_size, + host_size=( + 0 + if draft_host_token_capacity is not None + else server_args.hicache_size + ), + host_token_capacity=draft_host_token_capacity, + ) + pages_per_payload[PAYLOAD_DRAFT_KV] = draft_pages + if isinstance(self.kv_cache, NSATokenToKVPool): + pages_per_payload[PAYLOAD_INDEX_K] = target_pages + 1 + cp_shared_l2_pages_per_payload_planned = dict(pages_per_payload) + slab_pages_by_payload = _cp_shared_l2_slab_pages_by_payload( + server_args=server_args, + page_size=self.page_size, + bytes_per_page_by_payload=bytes_per_page_by_payload, + ) + cp_shared_l2_slabs_by_payload = build_cp_shared_l2_slabs_by_payload( + pages_per_payload, + slab_pages_by_payload=slab_pages_by_payload, + numa_policy=getattr( + server_args, "cp_shared_l2_numa_policy", "interleave_2m" + ), + ) + logger.info( + "[CP_SHARED_L2] planned shared physical L2 slabs: policy=%s counts=%s pages=%s", + getattr(server_args, "cp_shared_l2_numa_policy", "interleave_2m"), + { + key: len(value) + for key, value in cp_shared_l2_slabs_by_payload.items() + }, + pages_per_payload, + ) + + self._cp_shared_l2_slabs_by_payload = cp_shared_l2_slabs_by_payload self.token_to_kv_pool_host = self._create_token_to_kv_pool_host( self.kv_cache, server_args, host_size=0 if target_host_token_capacity is not None else None, host_token_capacity=target_host_token_capacity, + slabs_by_payload=cp_shared_l2_slabs_by_payload, ) self.tp_group = params.tp_cache_group @@ -583,12 +1038,39 @@ class HiRadixCache(RadixCache): self.load_cache_event = threading.Event() cp_shared_kv_layout = None + cp_shared_l2_page_allocator = None + cp_shared_l2_cpu_group = None if self._uses_cp_hicache: cp_shared_kv_layout = CpSharedKVLayout( page_size=self.page_size, cp_size=params.token_to_kv_pool_allocator.cp_size, cp_rank=params.token_to_kv_pool_allocator.cp_rank, ) + if getattr(server_args, "enable_cp_shared_physical_l2_hicache", False): + from sglang.srt.mem_cache.cp_shared_l2_pool import ( + PAYLOAD_TARGET_KV, + CpSharedL2PageAllocator, + ) + + _cp_rank, cp_size, cp_shared_l2_cpu_group = _get_cp_shared_l2_rank_and_group( + params.token_to_kv_pool_allocator + ) + if cp_shared_l2_pages_per_payload_planned is not None: + pages_per_payload = dict(cp_shared_l2_pages_per_payload_planned) + else: + pages_per_payload = _cp_shared_l2_pages_per_payload( + token_to_kv_pool_host=self.token_to_kv_pool_host, + page_size=self.page_size, + draft_token_to_kv_pool=draft_token_to_kv_pool, + draft_host_token_capacity=draft_host_token_capacity, + ) + cp_shared_l2_page_allocator = CpSharedL2PageAllocator( + pages_per_payload=pages_per_payload, + slabs_by_payload=cp_shared_l2_slabs_by_payload, + expected_ranks=range(cp_size), + expected_layers=range(self.kv_cache.layer_num), + required_payloads=(PAYLOAD_TARGET_KV,), + ) self.cache_controller = HiCacheController( params.token_to_kv_pool_allocator, self.token_to_kv_pool_host, @@ -605,6 +1087,8 @@ class HiRadixCache(RadixCache): pp_size=self.pp_size, enable_storage_metrics=self.enable_storage_metrics, cp_shared_kv_layout=cp_shared_kv_layout, + cp_shared_l2_page_allocator=cp_shared_l2_page_allocator, + cp_shared_l2_cpu_group=cp_shared_l2_cpu_group, ) if draft_token_to_kv_pool is not None: self.attach_draft_kv_pool( @@ -626,6 +1110,11 @@ class HiRadixCache(RadixCache): self.ongoing_write_through = {} # record CP host reservations/backups that are not request-visible yet. self.pending_host_backups: Dict[int, PendingHiCacheBackup] = {} + # CP-HiCache running host-capacity counts are lazily built from the tree + # on first access (see the _cp_counts property) and then maintained + # incrementally at every state transition, so _cp_host_capacity_snapshot + # is O(cp_size) rather than a full O(total-cached-pages) radix walk per + # write-admission. # record the node segments with ongoing load back self.ongoing_load_back = {} # record the ongoing prefetch requests @@ -840,6 +1329,133 @@ class HiRadixCache(RadixCache): stack.extend(getattr(node, "children", {}).values()) return nodes + # ----- B1: O(cp_size) running host-capacity counts ------------------- + # `_cp_counts` (an _CpRunningHostCounts) holds the per-owner committed / + # pending host TOKEN counts that `_cp_host_capacity_snapshot` used to derive + # by walking the whole radix tree (O(total cached pages)) on every + # write-admission. It is created EMPTY on first access (production starts + # with an empty tree) and then maintained incrementally at each state + # transition (enter/leave pending, commit, evict, split, reset), so the + # snapshot is O(cp_size). `_cp_walk_capacity_counts` recomputes the same + # result from scratch and backs the env-gated drift verifier (it is the + # correctness reference, not a build source). + + def _cp_walk_capacity_counts( + self, + ) -> Tuple[Tuple[int, ...], Tuple[int, ...], Tuple[int, ...], Tuple[int, ...]]: + """Authoritative (committed_target, committed_draft, pending_target, + pending_draft) per-owner token counts from the full radix walk + pending + dict. Runs the per-node owner-lane integrity assertion. This is the + correctness reference the env-gated drift verifier compares the + incrementally-maintained running counts against -- NOT a build source for + them (the running counts start empty and are maintained by the hooks).""" + cp_size = self._cp_hicache_cp_size() + ct = self._cp_zero_counts(cp_size) + cd = self._cp_zero_counts(cp_size) + pt = self._cp_zero_counts(cp_size) + pd = self._cp_zero_counts(cp_size) + pending = getattr(self, "pending_host_backups", {}) + pending_node_ids = set(pending.keys()) + for node in self._cp_walk_radix_nodes(): + metadata = getattr(node, "cp_hicache", None) + if metadata is None or getattr(node, "host_len", 0) <= 0: + continue + if self._is_cp_shared_l2_metadata(metadata): + continue + if getattr(node, "id", None) in pending_node_ids: + continue + counts = self._cp_assert_metadata_counts( + metadata, + context=f"capacity_committed node_id={getattr(node, 'id', '?')}", + ) + ct = self._cp_add_counts(ct, counts) + if getattr(metadata, "draft_host_indices", None) is not None: + cd = self._cp_add_counts(cd, counts) + for node_id, pending_backup in pending.items(): + if self._is_cp_shared_l2_metadata(pending_backup.metadata): + continue + counts = self._cp_assert_metadata_counts( + pending_backup.metadata, + context=f"capacity_pending node_id={node_id}", + ) + pt = self._cp_add_counts(pt, counts) + if getattr(pending_backup.metadata, "draft_host_indices", None) is not None: + pd = self._cp_add_counts(pd, counts) + return ct, cd, pt, pd + + @property + def _cp_counts(self) -> "_CpRunningHostCounts": + """Running host-capacity counts, created empty on first access and then + maintained incrementally at every state transition. + + Production starts with an empty radix tree, so zero-initialisation is + correct: every committed/pending entry is created only via commit / + split (committed) or attach / write_backup (pending), each of which + flows through the accounting hooks below. Only reached when CP HiCache + is enabled (the wrappers and the snapshot guard on _uses_cp_hicache). + Tests that construct nodes directly register them via the same hooks.""" + counts = self.__dict__.get("_cp_running_host_counts") + if counts is None: + counts = _CpRunningHostCounts(self._cp_hicache_cp_size()) + self.__dict__["_cp_running_host_counts"] = counts + return counts + + # CP accounting is a no-op unless CP HiCache is enabled; an absent + # _uses_cp_hicache (a cache built without __init__, e.g. a unit fixture or a + # non-CP path) counts as disabled. + @staticmethod + def _is_cp_shared_l2_metadata(metadata) -> bool: + return hasattr(metadata, "object_ranges") and not hasattr(metadata, "page_owners") + + @staticmethod + def _cp_reservation_owned_count(reservation) -> int: + owned_positions = getattr(reservation, "owned_positions", None) + if owned_positions is not None: + return int(owned_positions.numel()) + metadata = getattr(reservation, "metadata", None) + metadata_owned_positions = getattr(metadata, "owned_positions", None) + if metadata_owned_positions is not None: + return int(metadata_owned_positions.numel()) + return int(getattr(reservation, "host_indices").numel()) + + def _cp_account_enter_pending(self, metadata) -> None: + if getattr(self, "_uses_cp_hicache", False) and not self._is_cp_shared_l2_metadata(metadata): + self._cp_counts.enter_pending(metadata) + + def _cp_account_leave_pending(self, metadata) -> None: + if getattr(self, "_uses_cp_hicache", False) and not self._is_cp_shared_l2_metadata(metadata): + self._cp_counts.leave_pending(metadata) + + def _cp_account_enter_committed(self, metadata) -> None: + if getattr(self, "_uses_cp_hicache", False) and not self._is_cp_shared_l2_metadata(metadata): + self._cp_counts.enter_committed(metadata) + + def _cp_account_leave_committed(self, metadata) -> None: + if getattr(self, "_uses_cp_hicache", False) and not self._is_cp_shared_l2_metadata(metadata): + self._cp_counts.leave_committed(metadata) + + def _cp_account_leave_node_evict(self, node, metadata) -> None: + """Leave-accounting for an evicted node; metadata captured BEFORE the + caller nulls node.cp_hicache. Picks the bucket by pending membership so + the running total self-corrects; a pending eviction is provably + impossible (pending nodes are locked and skipped by both evictors and + the planner) so it also fails loud.""" + if metadata is None or not getattr(self, "_uses_cp_hicache", False): + return + if self._is_cp_shared_l2_metadata(metadata): + return + if getattr(node, "id", None) in getattr(self, "pending_host_backups", {}): + self._cp_counts.leave_pending(metadata) + raise RuntimeError( + "CP HiCache invariant violated: evicting a PENDING node " + f"id={getattr(node, 'id', None)} — should be impossible" + ) + self._cp_counts.leave_committed(metadata) + + def _cp_reset_running_counts(self) -> None: + if getattr(self, "_uses_cp_hicache", False): + self._cp_counts.reset() + def _cp_host_capacity_snapshot(self) -> CpHiCacheCapacitySnapshot: cp_size = self._cp_hicache_cp_size() target_capacity = tuple( @@ -854,34 +1470,30 @@ class HiRadixCache(RadixCache): if has_draft else None ) - committed_target = self._cp_zero_counts(cp_size) - committed_draft = self._cp_zero_counts(cp_size) - pending_target = self._cp_zero_counts(cp_size) - pending_draft = self._cp_zero_counts(cp_size) - pending = getattr(self, "pending_host_backups", {}) - pending_node_ids = set(pending.keys()) - for node in self._cp_walk_radix_nodes(): - metadata = getattr(node, "cp_hicache", None) - if metadata is None or getattr(node, "host_len", 0) <= 0: - continue - if getattr(node, "id", None) in pending_node_ids: - continue - counts = self._cp_assert_metadata_counts( - metadata, context=f"snapshot_committed node_id={getattr(node, 'id', '?')}" + if not self._uses_cp_hicache: + zero = self._cp_zero_counts(cp_size) + committed_target = committed_draft = pending_target = pending_draft = zero + else: + counts = self._cp_counts + counts.sanity_check() + committed_target, committed_draft, pending_target, pending_draft = ( + counts.as_tuples() ) - committed_target = self._cp_add_counts(committed_target, counts) - if getattr(metadata, "draft_host_indices", None) is not None: - committed_draft = self._cp_add_counts(committed_draft, counts) - - for node_id, pending_backup in pending.items(): - metadata = pending_backup.metadata - counts = self._cp_assert_metadata_counts( - metadata, context=f"snapshot_pending node_id={node_id}" - ) - pending_target = self._cp_add_counts(pending_target, counts) - if getattr(metadata, "draft_host_indices", None) is not None: - pending_draft = self._cp_add_counts(pending_draft, counts) + if envs.SGLANG_CP_HICACHE_VERIFY_SNAPSHOT.get(): + walk = self._cp_walk_capacity_counts() + running = ( + committed_target, + committed_draft, + pending_target, + pending_draft, + ) + if running != walk: + raise RuntimeError( + "CP HiCache running-count DRIFT vs full walk:\n" + f" running={running}\n" + f" walk={walk}" + ) return CpHiCacheCapacitySnapshot( target_capacity=target_capacity, @@ -1057,6 +1669,12 @@ class HiRadixCache(RadixCache): remaining_deficit=tuple(deficits), ) + def _cp_shared_l2_current_page_owners(self, num_pages: int) -> List[int]: + cp_size = self._cp_hicache_cp_size() + if num_pages < 0: + raise ValueError(f"num_pages must be non-negative, got {num_pages}") + return [int(page_index % cp_size) for page_index in range(num_pages)] + def _build_cp_load_back_plan( self, nodes_to_load: List[TreeNode], *, node_id: int ) -> CpLoadBackPlan: @@ -1069,6 +1687,8 @@ class HiRadixCache(RadixCache): ) page_owners: List[int] = [] + saw_shared_l2_metadata = False + saw_legacy_cp_metadata = False host_hit_len = 0 physical_hit_len = 0 for node in nodes_to_load: @@ -1078,6 +1698,15 @@ class HiRadixCache(RadixCache): "CP HiCache load-back missing cp_hicache metadata: " f"load_node_id={getattr(node, 'id', '?')} root_node_id={node_id}" ) + shared_l2_metadata = self._is_cp_shared_l2_metadata(metadata) + saw_shared_l2_metadata = saw_shared_l2_metadata or shared_l2_metadata + saw_legacy_cp_metadata = saw_legacy_cp_metadata or not shared_l2_metadata + if saw_shared_l2_metadata and saw_legacy_cp_metadata: + raise RuntimeError( + "CP HiCache load-back cannot mix shared physical L2 metadata " + "with legacy page_owners metadata in one load plan: " + f"root_node_id={node_id}" + ) node_host_len = int(self._node_host_len(node)) node_padded_len = int(getattr(metadata, "padded_len", node_host_len)) if node_padded_len % self.page_size != 0: @@ -1099,17 +1728,21 @@ class HiRadixCache(RadixCache): f"load_node_id={getattr(node, 'id', '?')} " f"host_len={node_host_len} padded_len={node_padded_len}" ) - page_owners.extend(int(owner) for owner in metadata.page_owners.tolist()) + if not shared_l2_metadata: + page_owners.extend(int(owner) for owner in metadata.page_owners.tolist()) host_hit_len += node_host_len physical_hit_len += node_padded_len expected_pages = physical_hit_len // self.page_size - if len(page_owners) != expected_pages: - raise RuntimeError( - "CP HiCache load-back page owner count mismatch: " - f"node_id={node_id} physical_hit_len={physical_hit_len} " - f"page_size={self.page_size} page_owners={len(page_owners)}" - ) + if saw_shared_l2_metadata: + page_owners = self._cp_shared_l2_current_page_owners(expected_pages) + else: + if len(page_owners) != expected_pages: + raise RuntimeError( + "CP HiCache load-back page owner count mismatch: " + f"node_id={node_id} physical_hit_len={physical_hit_len} " + f"page_size={self.page_size} page_owners={len(page_owners)}" + ) required, available, deficits, free_room_deficits = ( self._cp_load_back_owner_lane_stats(page_owners) @@ -1202,7 +1835,7 @@ class HiRadixCache(RadixCache): self, node: TreeNode, cp_size: int ) -> Tuple[int, ...]: metadata = getattr(node, "cp_hicache", None) - if metadata is not None: + if metadata is not None and not self._is_cp_shared_l2_metadata(metadata): return metadata.owner_page_counts(cp_size) value = getattr(node, "value", None) @@ -2065,6 +2698,7 @@ class HiRadixCache(RadixCache): self.prefetch_loaded_tokens_by_reqid.clear() if hasattr(self, "pending_host_backups"): self.pending_host_backups.clear() + self._cp_reset_running_counts() self.evictable_host_leaves.clear() self.pinned_size_ = 0 super().reset() @@ -2114,6 +2748,22 @@ class HiRadixCache(RadixCache): ) if controller is None or not getattr(controller, "has_draft_hicache", False): return True + if self._is_cp_shared_l2_metadata(metadata): + from sglang.srt.mem_cache.cp_shared_l2_pool import PAYLOAD_DRAFT_KV + + required_payloads = set(getattr(metadata, "required_payloads", ())) + object_ranges = getattr(metadata, "object_ranges", {}) + if ( + PAYLOAD_DRAFT_KV in required_payloads + and PAYLOAD_DRAFT_KV in object_ranges + ): + return True + raise RuntimeError( + "CP HiCache invariant violation: draft HiCache is attached but " + "shared physical L2 metadata is missing draft_kv payload: " + f"node_id={getattr(node, 'id', '?')} " + f"host_len={getattr(node, 'host_len', '?')} cp_rank={cp_rank}" + ) if getattr(metadata, "draft_host_indices", None) is None: raise RuntimeError( "CP HiCache invariant violation: draft HiCache is attached but " @@ -2143,13 +2793,34 @@ class HiRadixCache(RadixCache): return node.host_value is not None def _commit_pending_backup(self, node_id: int) -> TreeNode: - pending = self.pending_host_backups.pop(node_id) + pending = self.pending_host_backups[node_id] node = pending.node + # Leave pending BEFORE the pop, then enter committed AFTER the node owns + # the metadata, so a first-touch lazy rebuild of the running counts + # always sees a tree consistent with the single transition in flight + # (same metadata object; counts net-transfer pending -> committed). + self._cp_account_leave_pending(pending.metadata) + self.pending_host_backups.pop(node_id) node.host_len = pending.logical_len node.cp_hicache = pending.metadata node.host_value = None + self._cp_account_enter_committed(pending.metadata) if pending.locked: self.dec_node_lock_ref(node) + # B1: this commit is gated by the writing_check ReduceOp.MIN frontier, which + # is the all-ranks-done consensus. Mark the shared-L2 object committed here + # (replacing the per-(rank,layer,payload) gather quorum) so it transitions + # rank-uniformly. Guarded on object_ranges so it fires only for objects that + # were actually reserved into the pool (non-shared / skipped nodes no-op). + # B1 is validated for the write_through policy (the prod default); under + # write_back this commit is reached via the blocking drain rather than the + # MIN frontier -- still rank-uniform (deterministic registration), but the + # MIN-frontier reasoning above is the write_through path. + allocator = getattr(self.cache_controller, "cp_shared_l2_page_allocator", None) + if allocator is not None: + object_key = self.cache_controller._cp_shared_l2_object_key(node_id) + if allocator.object_ranges(object_key): + allocator.mark_object_committed(object_key) return node def _remove_undrained_write_acks(self, node_id: int) -> bool: @@ -2178,8 +2849,14 @@ class HiRadixCache(RadixCache): return removed def _rollback_pending_backup(self, node_id: int) -> TreeNode: - pending = self.pending_host_backups.pop(node_id) + pending = self.pending_host_backups[node_id] node = pending.node + # Leave pending BEFORE removing the entry (so a first-touch lazy rebuild + # of the running counts still sees this pending backup); unconditional + # because the entry existed regardless of whether node.cp_hicache + # currently aliases pending.metadata (pending/committed are exclusive). + self._cp_account_leave_pending(pending.metadata) + self.pending_host_backups.pop(node_id) # Cancel the half-submitted per-layer state (so later layer hooks # cannot write into the host slots evicted below) and scrub any # already-appended final ack before releasing the registration. @@ -2254,6 +2931,7 @@ class HiRadixCache(RadixCache): victim_ids.append(victim_id) self._record_remove_event(victim) + self._cp_account_leave_node_evict(victim, victim.cp_hicache) local_freed += self.cache_controller.evict_cp_host(victim.cp_hicache) victim.host_len = 0 victim.cp_hicache = None @@ -2278,6 +2956,28 @@ class HiRadixCache(RadixCache): *, admission_checked: bool = False, ): + if getattr( + self.cache_controller, "cp_shared_l2_page_allocator", None + ) is not None: + result = self.cache_controller.reserve_write_cp( + device_indices=device_indices, + node_id=node_id, + ) + if not isinstance(result, HiCacheWriteFailure): + return result + if getattr(result, "reason", "host_capacity") != "shared_l2_capacity": + return result + + # B1 (2b.1b): on a shared-L2 capacity miss, SKIP this backup (the + # existing reactive behavior; reserve-at-admission + background rotation + # [2a] keep this off the steady-state path). Do NOT call + # _evict_host_for_physical_slots(synchronize_across_ranks=True) here -- + # its `while len(heap) and not all_ranks_done()` loop is the #5 collective + # deadlock shape (per-rank len(heap)/wall-clock). The collective-free + # deterministic shared-pool eviction-and-retry is 2b.2; until it lands, + # capacity miss degrades to skip, never a collective. + return result + if not admission_checked: admission = self._cp_build_write_admission( device_indices, node_id=node_id, phase="initial" @@ -2373,12 +3073,13 @@ class HiRadixCache(RadixCache): submitted=True, locked=True, ) + self._cp_account_enter_pending(prepared.metadata) prepared.attached = True logger.debug( "[HiCache-write] attached prepared CP backup: node_id=%d logical_len=%d owned_positions=%d pending_backups=%d", node.id, prepared.logical_len, - prepared.metadata.owned_positions.numel(), + self._cp_reservation_owned_count(prepared.reservation), len(self.pending_host_backups), ) @@ -2740,6 +3441,7 @@ class HiRadixCache(RadixCache): node.value = None if node.cp_hicache is not None: + self._cp_account_leave_node_evict(node, node.cp_hicache) cache_controller = getattr(self, "cache_controller", None) if cache_controller is not None and hasattr(cache_controller, "evict_cp_host"): cache_controller.evict_cp_host(node.cp_hicache) @@ -2893,7 +3595,7 @@ class HiRadixCache(RadixCache): node_id, getattr(req, "rid", ""), len(kv_indices), - result.metadata.owned_positions.numel(), + self._cp_reservation_owned_count(result), ) def prepare_write_backup_for_req(self, req) -> None: if self.disable or not self._uses_cp_hicache: @@ -2934,7 +3636,11 @@ class HiRadixCache(RadixCache): return admission_checked = False - if len(candidates) > 1: + shared_l2_write = ( + getattr(self.cache_controller, "cp_shared_l2_page_allocator", None) + is not None + ) + if len(candidates) > 1 and not shared_l2_write: required = self._cp_zero_counts() for candidate in candidates: required = self._cp_add_counts( @@ -2978,7 +3684,10 @@ class HiRadixCache(RadixCache): def _node_host_evict_indices(self, node: TreeNode) -> torch.Tensor: if self._uses_cp_hicache: - return node.cp_hicache.host_indices + metadata = node.cp_hicache + if self._is_cp_shared_l2_metadata(metadata): + return torch.empty((0,), dtype=torch.int64) + return metadata.host_indices return node.host_value def write_backup(self, node: TreeNode, write_back=False): @@ -3019,6 +3728,10 @@ class HiRadixCache(RadixCache): submitted=True, locked=not write_back, ) + # Enter pending AFTER the dict insert and BEFORE the submit try: + # if submit raises, the except calls _rollback_pending_backup which + # does the matching leave_pending (net zero, no double-sub). + self._cp_account_enter_pending(result.metadata) try: self.cache_controller.submit_write_cp_per_layer(result) except Exception: @@ -3029,7 +3742,7 @@ class HiRadixCache(RadixCache): "[HiCache-write] write_backup CP SUBMITTED: node_id=%d logical_len=%d owned_positions=%d ongoing_writes=%d pending_backups=%d", node.id, len(node.value), - result.metadata.owned_positions.numel(), + self._cp_reservation_owned_count(result), len(self.ongoing_write_through), len(self.pending_host_backups), ) @@ -3670,15 +4383,26 @@ class HiRadixCache(RadixCache): host_indices = self._node_host_evict_indices(x) physical_count = len(host_indices) if host_indices is not None else 0 + is_shared_l2_metadata = self._is_cp_shared_l2_metadata(x.cp_hicache) self._record_remove_event(x) - if physical_count > 0: + if ( + self._uses_cp_hicache + and is_shared_l2_metadata + and hasattr(self.cache_controller, "evict_cp_host") + ): + num_evicted += self.cache_controller.evict_cp_host(x.cp_hicache) + elif physical_count > 0: if self._uses_cp_hicache and hasattr( self.cache_controller, "evict_cp_host" ): num_evicted += self.cache_controller.evict_cp_host(x.cp_hicache) else: num_evicted += self.cache_controller.evict_host(host_indices) + # Account before nulling (helper no-ops when not CP / metadata None); + # placed outside the physical_count gate so any committed teardown + # here is subtracted from the running counts. + self._cp_account_leave_node_evict(x, x.cp_hicache) x.host_len = 0 x.cp_hicache = None x.host_value = None @@ -4044,8 +4768,45 @@ class HiRadixCache(RadixCache): ongoing_after, ) + def _cp_assert_placement_replicated(self) -> None: + """B1 proof-obligation RUNTIME gate (Theorem 1). When + SGLANG_CP_HICACHE_PLACEMENT_ASSERT is on, MIN/MAX-reduce the pooled-L2 + allocator's placement_digest across the CP TP group; any cross-rank + divergence (a non-replicated reserve/release/commit) makes MIN != MAX -> + fail loud. Entry is gated ONLY on the global flag + tp_world_size + the + flag-implied allocator (all rank-uniform), never on per-rank state, so the + all_reduce can never be entered divergently (opus SF2). Off in prod. + """ + if not envs.SGLANG_CP_HICACHE_PLACEMENT_ASSERT.get(): + return + if self.tp_world_size <= 1: + return + allocator = getattr(self.cache_controller, "cp_shared_l2_page_allocator", None) + if allocator is None: + return + local = int(allocator.placement_digest(), 16) % (1 << 61) + lo = torch.tensor(local, dtype=torch.int64, device="cpu") + hi = torch.tensor(local, dtype=torch.int64, device="cpu") + self._cp_hicache_all_reduce( + lo, op=torch.distributed.ReduceOp.MIN, group=self.tp_group, + tag="placement_digest_min", + ) + self._cp_hicache_all_reduce( + hi, op=torch.distributed.ReduceOp.MAX, group=self.tp_group, + tag="placement_digest_max", + ) + if int(lo.item()) != local or int(hi.item()) != local: + raise RuntimeError( + "CP shared-L2 placement digest DIVERGED across ranks (B1 Theorem-1 " + f"violation): local={local} min={int(lo.item())} max={int(hi.item())}" + ) + def check_hicache_events(self): self.writing_check() + # B1 runtime invariant: placement is mutated by reserve (prepare) / release + # (evict) / mark_object_committed (the writing_check just above), so assert + # cross-rank replication here, right after the MIN commit frontier. + self._cp_assert_placement_replicated() self.loading_check() if self.enable_storage: self.drain_storage_control_queues() @@ -4471,11 +5232,58 @@ class HiRadixCache(RadixCache): child.value = child.value[split_len:].clone() if self._uses_cp_hicache: if self._node_backuped(child): - new_node.cp_hicache, child.cp_hicache = child.cp_hicache.split( - split_len - ) + _old_cp_meta = child.cp_hicache + # Leave committed for the pre-split metadata BEFORE reassigning + # child.cp_hicache, so a first-touch lazy rebuild of the running + # counts still sees the intact, in-tree child as committed. + # child is provably committed here (a pending split raises at the + # top of _split_node; _node_backuped requires host_len>0 & not + # pending). + self._cp_account_leave_committed(_old_cp_meta) + if self._is_cp_shared_l2_metadata(_old_cp_meta): + parent_object_key = self.cache_controller._cp_shared_l2_object_key( + new_node.id + ) + child_object_key = getattr(_old_cp_meta, "object_key", None) + if not child_object_key: + child_object_key = self.cache_controller._cp_shared_l2_object_key( + child.id + ) + split_pages = split_len // self.page_size + allocator = getattr( + self.cache_controller, "cp_shared_l2_page_allocator", None + ) + split_object = getattr(allocator, "split_committed_object", None) + if split_object is None: + raise RuntimeError( + "CP shared physical L2 split requires allocator " + "split_committed_object" + ) + split_object( + child_object_key, + split_pages_by_payload={ + payload_kind: split_pages + for payload_kind in _old_cp_meta.object_ranges + }, + parent_object_key=parent_object_key, + child_object_key=child_object_key, + ) + new_node.cp_hicache, child.cp_hicache = _old_cp_meta.split( + split_len, + parent_object_key=parent_object_key, + child_object_key=child_object_key, + ) + else: + new_node.cp_hicache, child.cp_hicache = _old_cp_meta.split( + split_len + ) new_node.host_len = split_len child.host_len = child.host_len - split_len + # Re-enter each half's OWN metadata object (count-neutral overall + # -- .split() partitions page_owners) so the later + # evict-subtraction uses the correct per-half counts. + self._cp_account_enter_committed(new_node.cp_hicache) + self._cp_account_enter_committed(child.cp_hicache) elif child.backuped: new_node.host_value = child.host_value[:split_len].clone() child.host_value = child.host_value[split_len:].clone() diff --git a/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py b/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py index e0abce0a1..931ac60c9 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py +++ b/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py @@ -2862,5 +2862,69 @@ class TestCpSharedL2MarkObjectCommitted(unittest.TestCase): self.assertFalse(a.is_committed("n")) +class TestCpSharedL2EightRankReserveDeterminism(unittest.TestCase): + """B1 (2b.1b) write-through end-to-end determinism: every CP rank runs the + IDENTICAL deterministic reserve over the replicated event stream, so all 8 + ranks reach identical placement with NO broadcast (the controller's + _reserve_write_cp_shared_l2 strips to exactly this loop). Models 8 ranks as 8 + independent allocators applying the same reserve/commit/release sequence -- + the runtime placement_digest assert is the production analog of this check.""" + + CP = 8 + + def make_rank_allocators(self, pages=64): + return [ + _cp_shared_l2_pool.CpSharedL2PageAllocator( + pages_per_payload={PAYLOAD_TARGET_KV: pages, PAYLOAD_DRAFT_KV: pages}, + slab_ids_by_payload={PAYLOAD_TARGET_KV: 10, PAYLOAD_DRAFT_KV: 11}, + expected_ranks=range(self.CP), + expected_layers=(0, 1), + required_payloads=(PAYLOAD_TARGET_KV,), + ) + for _ in range(self.CP) + ] + + def _reserve_on_all(self, allocs, node_id, payloads, num_pages): + # the B1 every-rank reserve loop, run identically on all ranks + for a in allocs: + for p in payloads: + a.reserve(f"cp_hicache_node:{node_id}", p, num_pages) + + def test_eight_ranks_reach_identical_placement(self): + allocs = self.make_rank_allocators() + for nid, (payloads, npages) in enumerate( + [ + ((PAYLOAD_TARGET_KV,), 3), + ((PAYLOAD_TARGET_KV, PAYLOAD_DRAFT_KV), 2), + ((PAYLOAD_TARGET_KV,), 5), + ] + ): + self._reserve_on_all(allocs, nid, payloads, npages) + self.assertEqual( + len({a.placement_digest() for a in allocs}), + 1, + "all 8 ranks must compute identical placement with no broadcast", + ) + + def test_commit_then_release_stays_identical(self): + allocs = self.make_rank_allocators() + self._reserve_on_all(allocs, 0, (PAYLOAD_TARGET_KV,), 4) + for a in allocs: # MIN-commit fires on every rank with the same object_key + a.mark_object_committed("cp_hicache_node:0") + self.assertEqual(len({a.placement_digest() for a in allocs}), 1) + for a in allocs: # deterministic evict on every rank + a.release("cp_hicache_node:0") + self.assertEqual(len({a.placement_digest() for a in allocs}), 1) + + def test_one_rank_divergence_is_detectable(self): + # A reserve that fires on one rank but not the others (a non-replicated + # trigger) makes the placement_digests diverge -- exactly what the runtime + # cross-rank assert catches. + allocs = self.make_rank_allocators() + self._reserve_on_all(allocs, 0, (PAYLOAD_TARGET_KV,), 3) + allocs[3].reserve("cp_hicache_node:1", PAYLOAD_TARGET_KV, 1) # rank 3 only + self.assertGreater(len({a.placement_digest() for a in allocs}), 1) + + if __name__ == "__main__": unittest.main()