diff --git a/docs/advanced_features/nsa_prefill_cp_hicache_per_layer_backup_plan.md b/docs/advanced_features/nsa_prefill_cp_hicache_per_layer_backup_plan.md index 2e81de2ab..5652fbfbd 100644 --- a/docs/advanced_features/nsa_prefill_cp_hicache_per_layer_backup_plan.md +++ b/docs/advanced_features/nsa_prefill_cp_hicache_per_layer_backup_plan.md @@ -226,6 +226,18 @@ all_reduce per node victim in a tight host-eviction loop all_reduce every scheduler tick when no completion prefix advanced ``` +Current correctness note: CP host reservation now synchronizes +`required_host_slots` with `all_reduce(MAX)` before the host-eviction retry. +This forces every rank into the same reserve/evict/retry branch and avoids +collective mismatches when one rank is host-full and another rank reserves +successfully. It is intentionally a coarse slow-path collective, not a +per-layer collective, but it can become a performance cost when host pressure +is frequent because every reservation failure pays at least one rank-wide MAX +sync and retry failures pay a second one. Treat this as a correctness guard to +be amortized later with batched reservation epochs, deterministic host +watermarks, or less frequent proactive host eviction; do not move it into the +per-layer data path. + ## Per-Layer Backup Data Plane The existing host pool exposes `load_to_device_per_layer()` but not a matching @@ -415,6 +427,13 @@ Refactor host full handling so reserve failure returns required local physical slots, host eviction frees only valid host-only victims, and reservation retries once. Add logs/counters for backup skipped due to host pressure. +The current first implementation uses rank-wide +`MAX(required_host_slots)` synchronization before eviction/retry so all CP ranks +observe the same capacity-pressure branch. This fixes correctness under +rank-skewed host pressure, but it is a known performance risk if host-full +becomes common. Follow-up optimization should reduce how often this slow path +is reached rather than adding more collectives around it. + ### Phase 7: Split safety Implement defer/requeue-on-pending-split for nodes with pending backup. Add @@ -436,10 +455,14 @@ paths: ## Open Questions / Decisions -1. **Layer completion hook.** The chosen semantic hook is layer-end after that - layer's KV store is ordered. The exact code location still needs - implementation-time confirmation. A next-layer-start hook is a fallback only - if it preserves the same ordering and adds an explicit final-layer trigger. +1. **Layer completion hook.** The chosen semantic hook is explicit model + layer-end after that layer's target/draft KV and NSA indexer writes have + been enqueued. Store-side notifiers are intentionally not the ownership + boundary because MLA, NSA indexer, TAI fused store, Triton store, and + zero-local paths can bypass one another. The first implementation hooks + `DeepseekV2DecoderLayer.forward()` and the TBO `op_comm_postprocess_layer()` + path; TBO children only notify from subbatch 1 so a layer is not backed up + after only the first child has written its KV. 2. **Pending split behavior.** The chosen first pass is defer/requeue the request that would split a node with pending backup. Do not split the in-flight backup op, including for future `bs > 1`. @@ -466,9 +489,8 @@ Implemented in the first P4-P8 pass: - pending backup split attempts are deferred instead of mutating the radix tree; - request admission marks deferred prefix matches and keeps those requests in the waiting path for a later retry; -- KV pools expose a no-op-safe layer-store notifier. NSA defers notification - until `set_index_k_scale_buffer()` so the indexer payload is ordered after the - KV payload. +- KV pools expose a no-op-safe explicit layer-end notifier. Store methods and + fused store kernels do not drive backup progress. Important limitation: the current `sgl_kernel.kvcacheio` surface has H2D per-layer page-first kernels and D2H all-layer page-first kernels, but not a @@ -503,14 +525,16 @@ This covers the GLM5/NSA production shape: MHA page-first remains fail-fast until it has a dedicated kernel. -The current radix insertion path creates host reservations after request KV has -already been materialized. Therefore `submit_write_cp_per_layer()` performs a -safe catch-up submission of all layers through the per-layer API for current -call sites. The layer-store notifier path is also wired: +The first production forward-overlap path reserves host slots after +`prepare_for_extend()` and before the model forward, then calls +`submit_write_cp_per_layer(..., catch_up_all_layers=False)`. Radix insertion +only attaches that prepared reservation; it does not replay an all-layer +catch-up copy. The model layer-end hook is the forward progress driver: ```text submit_write_cp_per_layer(..., catch_up_all_layers=False) - -> on_layer_kv_stored(layer_id, source="target" | "draft") + -> DeepseekV2 layer-end calls token_to_kv_pool.notify_layer_end_for_backup(layer_id) + -> HiCacheController.on_layer_end(layer_id, source="target" | "draft") -> submit_write_cp_layer(reservation, layer_id, submit_target=..., submit_draft=...) -> final ack only after all target/draft layers are submitted ``` @@ -521,6 +545,32 @@ payloads remain invisible until the final write ack commits separately so an early target hook cannot copy draft KV before draft has stored the same layer. +The previous store-notifier attempt was rejected as the primary correctness +boundary. It was too easy for fused paths to skip it: NSA fused indexer store +bypassed `set_index_k_scale_buffer()`, MLA can use `SGLANG_CP_SHARED_KV_FUSED_MLA_STORE=1`, +and zero-local paths do not necessarily touch the same setter. Missing one +store notifier left the prepared backup in `pending_layer_writes` forever: the +first request could finish prefill without a final write ack, and a second +identical request could hit the pending node and wedge the scheduler. The +layer-end hook makes backup progress independent of the chosen store kernel; +store-side hooks should remain absent or debug-only. + +One more correctness/performance edge was observed with repeated EAGLE requests. +Scheduler prefix matching and final radix insertion can have different effective +lengths because matching is based on the schedulable prefix while insertion uses +the final page-aligned key. With EAGLE bigram keys this can leave +`req.cache_protected_len` shorter than the prefix that insertion will find in the +tree. If early backup reserves `[cache_protected_len, insertion_key_len)` blindly, +the system copies a duplicate page, then final insertion discovers that no new +node was created and rolls the prepared backup back with `reason=insert_miss`. +This is safe but wastes host reservation, per-layer callbacks, and write-stream +work on repeated prompts. + +The early backup path now probes the existing radix prefix for the final +insertion key without splitting the tree, and only reserves the suffix beyond +`max(cache_protected_len, existing_insertion_prefix_len)`. This preserves the +cold-request behavior while avoiding duplicate prepared backups on exact repeats. + ## Summary Per-layer backup should be implemented as a data-plane refinement, not a radix diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index e5b95b2d0..440bebd1b 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -361,7 +361,7 @@ class HiCacheController: self.mem_pool_device.register_layer_transfer_counter(self.layer_done_counter) if hasattr(self.mem_pool_device, "register_layer_backup_notifier"): self.mem_pool_device.register_layer_backup_notifier( - lambda layer_id: self.on_layer_kv_stored(layer_id, source="target") + lambda layer_id: self.on_layer_end(layer_id, source="target") ) self.draft_mem_pool_host = None self.draft_mem_pool_device = None @@ -431,10 +431,10 @@ class HiCacheController: ) if hasattr(draft_mem_pool_device, "register_layer_backup_notifier"): draft_mem_pool_device.register_layer_backup_notifier( - lambda layer_id: self.on_layer_kv_stored(layer_id, source="draft") + lambda layer_id: self.on_layer_end(layer_id, source="draft") ) - def on_layer_kv_stored(self, layer_id: int, source: str = "target") -> None: + def on_layer_end(self, layer_id: int, source: str = "target") -> None: if not self.pending_layer_writes: return if source not in ("target", "draft"): @@ -1236,7 +1236,7 @@ class HiCacheController: Current radix insertion calls write_backup after the request KV has already been materialized. For that path, catch up by submitting all layers immediately through the per-layer API. Future early reservations can use - catch_up_all_layers=False and rely on on_layer_kv_stored(). + catch_up_all_layers=False and rely on explicit model layer-end hooks. """ state = self._get_or_create_layer_write_state(reservation) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index ec52db19b..ca7920732 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -644,6 +644,7 @@ class Req(ReqDllmMixin): self.last_host_backup_node: Any = None self.host_hit_length = 0 self.prefix_match_deferred_by_pending_backup = False + self.cp_hicache_prepared_backup = None # Tokens loaded from storage backend (L3) during prefetch for this request self.storage_hit_length = 0 # The node to lock until for swa radix tree lock ref diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index be59d995f..9984c961c 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2457,6 +2457,11 @@ class Scheduler( ) new_batch.prepare_for_extend() + if self.enable_hierarchical_cache and hasattr( + self.tree_cache, "prepare_write_backup_for_req" + ): + for req in can_run_list: + self.tree_cache.prepare_write_backup_for_req(req) # Record prefill stats for logging after forward new_batch.prefill_stats = PrefillStats.from_adder( diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 563566cca..dea7cec91 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -60,6 +60,7 @@ class InsertParams: # General chunked: bool = False priority: int = 0 + cp_hicache_prepared_backup: Optional[object] = None @dataclasses.dataclass diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 93f258f75..0629ea8ac 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -17,6 +17,7 @@ from sglang.srt.environ import envs from sglang.srt.managers.cache_controller import ( HiCacheController, HiCacheWriteFailure, + HiCacheWriteReservation, PrefetchOperation, ) from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator @@ -44,6 +45,7 @@ from sglang.srt.mem_cache.radix_cache import ( RadixKey, TreeNode, compute_node_hash_values, + page_align_keys, split_node_hash_value, ) from sglang.srt.mem_cache.utils import convert_to_bigram_key @@ -246,6 +248,15 @@ class PendingHiCacheBackup: locked: bool = True +@dataclass +class PreparedCpHiCacheBackup: + node_id: int + reservation: HiCacheWriteReservation + metadata: CpHiCacheNodeMetadata + logical_len: int + attached: bool = False + + class HiCachePendingBackupSplit(Exception): def __init__(self, node: TreeNode): self.node = node @@ -1010,6 +1021,260 @@ class HiRadixCache(RadixCache): self.dec_node_lock_ref(node) return node + def _sync_cp_write_required_host_slots(self, result) -> int: + """Return the max host slots any CP rank needs before write backup. + + CP ranks must make the reserve/evict/retry decision collectively. A + local successful reservation is not enough: another CP rank may be full + and enter host eviction, whose implementation contains collective + all_reduce calls. If successful ranks skip that branch they can instead + enter unrelated collectives (for example disagg polling), causing a + process-group mismatch. + """ + + required_host_slots = ( + int(result.required_host_slots) + if isinstance(result, HiCacheWriteFailure) + else 0 + ) + if ( + getattr(self, "tp_group", None) is None + or getattr(self, "tp_world_size", 1) <= 1 + ): + return required_host_slots + + required = torch.tensor(required_host_slots, dtype=torch.int64, device="cpu") + torch.distributed.all_reduce( + required, op=torch.distributed.ReduceOp.MAX, group=self.tp_group + ) + return int(required.item()) + + def _release_cp_write_reservation_for_retry(self, result, node_id: int) -> None: + if isinstance(result, HiCacheWriteFailure): + return + logger.info( + "[HiCache-write] release local CP reservation before collective retry: node_id=%d owned_positions=%d", + node_id, + result.metadata.owned_positions.numel(), + ) + self.cache_controller.evict_cp_host(result.metadata) + + def _reserve_write_cp_indices_collectively( + self, device_indices: torch.Tensor, node_id: int + ): + result = self.cache_controller.reserve_write_cp( + device_indices=device_indices, + node_id=node_id, + ) + required_host_slots = self._sync_cp_write_required_host_slots(result) + if required_host_slots <= 0: + return result + + self._release_cp_write_reservation_for_retry(result, node_id) + self._evict_host_for_physical_slots( + required_host_slots, + synchronize_across_ranks=getattr(self, "tp_world_size", 1) > 1, + ) + logger.info( + "[HiCache-write] write_backup CP retry after host eviction: node_id=%d needed_slots=%d", + node_id, + required_host_slots, + ) + result = self.cache_controller.reserve_write_cp( + device_indices=device_indices, + node_id=node_id, + ) + required_host_slots = self._sync_cp_write_required_host_slots(result) + if required_host_slots <= 0: + return result + + self._release_cp_write_reservation_for_retry(result, node_id) + logger.info( + "[HiCache-write] write_backup CP FAILED after collective retry: node_id=%d len=%d needed_slots=%d", + node_id, + len(device_indices) if device_indices is not None else 0, + required_host_slots, + ) + return HiCacheWriteFailure(required_host_slots=required_host_slots) + + def _reserve_write_cp_collectively(self, node: TreeNode): + return self._reserve_write_cp_indices_collectively(node.value, node.id) + + def _attach_prepared_cp_backup( + self, node: TreeNode, prepared: PreparedCpHiCacheBackup + ) -> None: + if prepared.attached: + raise RuntimeError( + f"CP HiCache prepared backup node_id={prepared.node_id} was attached twice" + ) + if len(node.value) != prepared.logical_len: + raise RuntimeError( + f"Prepared CP HiCache backup length mismatch for node_id={prepared.node_id}: " + f"node_len={len(node.value)} prepared_len={prepared.logical_len}" + ) + node.id = prepared.node_id + self.ongoing_write_through[node.id] = node + self.inc_node_lock_ref(node) + self.pending_host_backups[node.id] = PendingHiCacheBackup( + node=node, + metadata=prepared.metadata, + logical_len=prepared.logical_len, + submitted=True, + locked=True, + ) + prepared.attached = True + logger.info( + "[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(), + len(self.pending_host_backups), + ) + + def _rollback_prepared_cp_backup( + self, prepared: Optional[PreparedCpHiCacheBackup], reason: str + ) -> None: + if prepared is None or prepared.attached: + return + if prepared.node_id in self.cache_controller.pending_layer_writes: + raise RuntimeError( + f"Cannot rollback prepared CP HiCache backup node_id={prepared.node_id} " + f"while per-layer writes are still pending; reason={reason}" + ) + + retained_acks = [] + removed_ack = False + for ack in self.cache_controller.ack_write_queue: + if prepared.node_id not in ack.node_ids: + retained_acks.append(ack) + continue + ack.finish_event.synchronize() + remaining_node_ids = [ + node_id for node_id in ack.node_ids if node_id != prepared.node_id + ] + if remaining_node_ids: + retained_acks.append(ack._replace(node_ids=remaining_node_ids)) + removed_ack = True + if removed_ack: + self.cache_controller.ack_write_queue = retained_acks + logger.info( + "[HiCache-write] rollback unattached prepared CP backup: node_id=%d logical_len=%d reason=%s", + prepared.node_id, + prepared.logical_len, + reason, + ) + self.cache_controller.evict_cp_host(prepared.metadata) + + def _probe_existing_radix_prefix_len_no_split(self, key: RadixKey) -> int: + """Return the already-present radix prefix length without mutating the tree. + + `prepare_write_backup_for_req` runs before forward and must reserve host + slots only for KV that will become a *new* radix node at insertion time. + The request's `cache_protected_len` is based on scheduler prefix-match + semantics, which can be shorter than the prefix that a final insertion + will find (for example EAGLE/bigram plus page alignment). Probing the + full insertion key here avoids backing up duplicate pages that insertion + will later reject and roll back. + """ + + if self.disable or len(key) == 0: + return 0 + + key, _ = self.maybe_bigram_convert(key) + if self.page_size != 1: + page_aligned_len = len(key) // self.page_size * self.page_size + key = key[:page_aligned_len] + + if len(key) == 0: + return 0 + + node = getattr(self, "root_node", None) + if node is None: + return 0 + child_key = self.get_child_key_fn(key) + total_prefix_len = 0 + + while len(key) > 0 and child_key in node.children: + child = node.children[child_key] + prefix_len = self.key_match_fn(child.key, key) + if prefix_len <= 0: + break + + total_prefix_len += prefix_len + if prefix_len < len(child.key): + break + + node = child + key = key[prefix_len:] + if len(key): + child_key = self.get_child_key_fn(key) + + return total_prefix_len + + def prepare_write_backup_for_req(self, req) -> None: + if ( + self.disable + or not self._uses_cp_hicache + or self.cache_controller.write_policy == "write_back" + or getattr(req, "is_chunked", 0) > 0 + or getattr(req, "cp_hicache_prepared_backup", None) is not None + ): + return + + token_ids = req.fill_ids + keys = convert_to_bigram_key(token_ids) if self.is_eagle else token_ids + keys = page_align_keys(keys, self.page_size) + existing_prefix_len = self._probe_existing_radix_prefix_len_no_split( + RadixKey( + keys, + getattr(req, "extra_key", None), + is_bigram=self.is_eagle, + ) + ) + start = min(max(req.cache_protected_len, existing_prefix_len), len(keys)) + if len(keys) <= start: + return + + kv_indices = self.req_to_token_pool.req_to_token[ + req.req_pool_idx, start : len(keys) + ].to(dtype=torch.int64, copy=True) + if len(kv_indices) == 0: + return + + node_id = TreeNode.counter + TreeNode.counter += 1 + result = self._reserve_write_cp_indices_collectively(kv_indices, node_id) + if isinstance(result, HiCacheWriteFailure): + logger.info( + "[HiCache-write] prepare CP backup skipped after host reservation failure: node_id=%d rid=%s len=%d", + node_id, + getattr(req, "rid", ""), + len(kv_indices), + ) + return + + try: + self.cache_controller.submit_write_cp_per_layer( + result, catch_up_all_layers=False + ) + except Exception: + self.cache_controller.evict_cp_host(result.metadata) + raise + + req.cp_hicache_prepared_backup = PreparedCpHiCacheBackup( + node_id=node_id, + reservation=result, + metadata=result.metadata, + logical_len=len(kv_indices), + ) + logger.info( + "[HiCache-write] prepared CP per-layer backup before forward: node_id=%d rid=%s logical_len=%d owned_positions=%d", + node_id, + getattr(req, "rid", ""), + len(kv_indices), + result.metadata.owned_positions.numel(), + ) + def _node_host_len(self, node: TreeNode) -> int: if self._uses_cp_hicache: return node.host_len @@ -1029,25 +1294,7 @@ class HiRadixCache(RadixCache): write_back, ) if self._uses_cp_hicache: - result = self.cache_controller.reserve_write_cp( - device_indices=node.value, - node_id=node.id, - ) - if isinstance(result, HiCacheWriteFailure): - required_host_slots = result.required_host_slots - self._evict_host_for_physical_slots( - required_host_slots, - synchronize_across_ranks=getattr(self, "tp_world_size", 1) > 1, - ) - logger.info( - "[HiCache-write] write_backup CP retry after host eviction: node_id=%d needed_slots=%d", - node.id, - required_host_slots, - ) - result = self.cache_controller.reserve_write_cp( - device_indices=node.value, - node_id=node.id, - ) + result = self._reserve_write_cp_collectively(node) if isinstance(result, HiCacheWriteFailure): logger.info( "[HiCache-write] write_backup CP FAILED (host full): node_id=%d len=%d", @@ -1191,6 +1438,7 @@ class HiRadixCache(RadixCache): if not finish_event.query(): break finish_count += 1 + local_finish_count = finish_count queue_size = torch.tensor(finish_count, dtype=torch.int, device="cpu") if self.tp_world_size > 1: # synchronize TP workers to make the same update to radix cache @@ -1201,11 +1449,11 @@ class HiRadixCache(RadixCache): ) finish_count = int(queue_size.item()) - logger.info( + logger.debug( "[HiCache-write] writing_check: ongoing=%d ack_queue=%d local_finished=%d sync_finished=%d tp_size=%d", len(self.ongoing_write_through), ack_queue_len, - finish_count if self.tp_world_size <= 1 else queue_size.item(), + local_finish_count, finish_count, getattr(self, "tp_world_size", 1), ) @@ -2264,6 +2512,7 @@ class HiRadixCache(RadixCache): value = params.value chunked = params.chunked priority = params.priority + prepared_cp_backup = params.cp_hicache_prepared_backup if priority is None: priority = 0 @@ -2325,7 +2574,18 @@ class HiRadixCache(RadixCache): child_key = self.get_child_key_fn(key) if len(key): - new_node = TreeNode(priority=priority) + use_prepared_cp_backup = ( + prepared_cp_backup is not None + and getattr(prepared_cp_backup, "logical_len", -1) == len(value) + ) + new_node = TreeNode( + id=( + prepared_cp_backup.node_id + if use_prepared_cp_backup + else None + ), + priority=priority, + ) new_node.parent = node new_node.key = key new_node.value = value.clone() @@ -2342,7 +2602,10 @@ class HiRadixCache(RadixCache): self._record_store_event(new_node) if self.cache_controller.write_policy != "write_back": - self._inc_hit_count(new_node, chunked) + if use_prepared_cp_backup: + self._attach_prepared_cp_backup(new_node, prepared_cp_backup) + else: + self._inc_hit_count(new_node, chunked) return InsertResult(prefix_len=total_prefix_length) def release_aborted_request(self, rid: str): diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 8f6302ed0..65688fcd3 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -695,7 +695,6 @@ class KVCache(abc.ABC): # default state for optional layer-wise transfer control self.layer_transfer_counter = None self.layer_backup_notifiers = [] - self.layer_backup_notify_after_indexer = False # for disagg with nvlink self.enable_custom_mem_pool, self.custom_mem_pool, _ = ( @@ -750,11 +749,10 @@ class KVCache(abc.ABC): def register_layer_backup_notifier(self, notifier): self.layer_backup_notifiers.append(notifier) - def _notify_layer_kv_stored(self, layer_id: int, source: str = "kv"): - if self.layer_backup_notify_after_indexer and source != "indexer": - return + def notify_layer_end_for_backup(self, layer_id: int): + local_layer_id = layer_id - self.start_layer for notifier in self.layer_backup_notifiers: - notifier(layer_id) + notifier(local_layer_id) def get_cpu_copy(self, indices): raise NotImplementedError() @@ -1047,7 +1045,6 @@ class MHATokenToKVPool(KVCache): alt_stream=self.alt_stream, same_kv_dim=self.same_kv_dim, ) - self._notify_layer_kv_stored(layer_id - self.start_layer) def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): if envs.SGLANG_NATIVE_MOVE_KV_CACHE.get(): @@ -1595,7 +1592,6 @@ class MLATokenToKVPool(KVCache): ) else: self.kv_buffer[layer_id - self.start_layer][loc] = cache_k - self._notify_layer_kv_stored(layer_id - self.start_layer) def set_mla_kv_buffer( self, @@ -1637,7 +1633,6 @@ class MLATokenToKVPool(KVCache): cache_k_nope, cache_k_rope, ) - self._notify_layer_kv_stored(layer_id - self.start_layer) def get_mla_kv_buffer( self, @@ -1859,7 +1854,6 @@ class NSATokenToKVPool(MLATokenToKVPool): use_nsa=True, override_kv_cache_dim=override_dim, ) - self.layer_backup_notify_after_indexer = True # self.index_k_dtype = torch.float8_e4m3fn # self.index_k_scale_dtype = torch.float32 self.index_head_dim = index_head_dim @@ -1966,7 +1960,6 @@ class NSATokenToKVPool(MLATokenToKVPool): index_buf_accessor.SetKAndS.execute( pool=self, buf=buf, loc=loc, index_k=index_k, index_k_scale=index_k_scale ) - self._notify_layer_kv_stored(layer_id - self.start_layer, source="indexer") def get_state_buf_infos(self): data_ptrs = [ diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index f06ab7803..d440257d2 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -489,22 +489,40 @@ class RadixCache(BasePrefixCache): keys = page_align_keys(keys, self.page_size) values = kv_indices[: len(keys)].to(dtype=torch.int64, copy=True) radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle) + prepared_cp_backup = getattr(req, "cp_hicache_prepared_backup", None) # Radix Cache takes one ref in memory pool if is_insert: priority = getattr(req, "priority", 0) or 0 result = self.insert( - InsertParams(key=radix_key, value=values, priority=priority) + InsertParams( + key=radix_key, + value=values, + priority=priority, + cp_hicache_prepared_backup=prepared_cp_backup, + ) ) + if ( + prepared_cp_backup is not None + and not getattr(prepared_cp_backup, "attached", False) + and hasattr(self, "_rollback_prepared_cp_backup") + ): + self._rollback_prepared_cp_backup(prepared_cp_backup, "insert_miss") new_prefix_len = result.prefix_len # Free the duplicates that were already in the tree self.token_to_kv_pool_allocator.free( kv_indices[req.cache_protected_len : new_prefix_len] ) else: + if prepared_cp_backup is not None and hasattr( + self, "_rollback_prepared_cp_backup" + ): + self._rollback_prepared_cp_backup(prepared_cp_backup, "no_insert") self.token_to_kv_pool_allocator.free( kv_indices[req.cache_protected_len : len(keys)] ) + if prepared_cp_backup is not None: + req.cp_hicache_prepared_backup = None # free the unaligned tail self.token_to_kv_pool_allocator.free(kv_indices[len(keys) :]) @@ -527,6 +545,7 @@ class RadixCache(BasePrefixCache): keys = page_align_keys(keys, self.page_size) values = kv_indices[: len(keys)].to(dtype=torch.int64, copy=True) radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle) + prepared_cp_backup = getattr(req, "cp_hicache_prepared_backup", None) # Radix Cache takes one ref in memory pool result = self.insert( @@ -535,8 +554,17 @@ class RadixCache(BasePrefixCache): value=values, chunked=chunked, priority=getattr(req, "priority", 0) or 0, + cp_hicache_prepared_backup=prepared_cp_backup, ) ) + if ( + prepared_cp_backup is not None + and not getattr(prepared_cp_backup, "attached", False) + and hasattr(self, "_rollback_prepared_cp_backup") + ): + self._rollback_prepared_cp_backup(prepared_cp_backup, "insert_miss") + if prepared_cp_backup is not None: + req.cp_hicache_prepared_backup = None new_prefix_len = result.prefix_len self.token_to_kv_pool_allocator.free( diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 17744eb2d..07b30ae98 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -1620,6 +1620,22 @@ class DeepseekV2DecoderLayer(nn.Module): and layer_id % self.config.moe_layer_freq == 0 ) + def _notify_cp_hicache_layer_end( + self, + forward_batch: ForwardBatch, + tbo_subbatch_index: Optional[int] = None, + ) -> None: + if ( + getattr(forward_batch, "tbo_parent_token_range", None) is not None + and tbo_subbatch_index != 1 + ): + return + notifier = getattr( + forward_batch.token_to_kv_pool, "notify_layer_end_for_backup", None + ) + if notifier is not None: + notifier(self.layer_id) + def forward( self, positions: torch.Tensor, @@ -1707,6 +1723,7 @@ class DeepseekV2DecoderLayer(nn.Module): hidden_states, residual, forward_batch ) + self._notify_cp_hicache_layer_end(forward_batch) return hidden_states, residual def op_comm_prepare_attn( @@ -1761,6 +1778,9 @@ class DeepseekV2DecoderLayer(nn.Module): state.pop("residual_after_comm_pre_mlp"), state.forward_batch, ) + self._notify_cp_hicache_layer_end( + state.forward_batch, tbo_subbatch_index=state.tbo_subbatch_index + ) output = dict( positions=state.positions, diff --git a/test/registered/unit/managers/test_hicache_controller_cp.py b/test/registered/unit/managers/test_hicache_controller_cp.py index c13e4f02d..9c529d24d 100644 --- a/test/registered/unit/managers/test_hicache_controller_cp.py +++ b/test/registered/unit/managers/test_hicache_controller_cp.py @@ -175,6 +175,10 @@ class FakeDevicePool: def register_layer_backup_notifier(self, notifier): self.layer_backup_notifiers.append(notifier) + def notify_layer_end_for_backup(self, layer_id): + for notifier in self.layer_backup_notifiers: + notifier(layer_id) + class TestPageFirstPerLayerBackupTaiKernel(CustomTestCase): def test_mla_page_first_per_layer_backup_uses_tai_lf_pf_kernel(self): @@ -677,12 +681,12 @@ class TestHiCacheControllerCPWrite(CustomTestCase): self.assertEqual(host_pool.layer_backups, []) self.assertEqual(controller.ack_write_queue, []) - controller.on_layer_kv_stored(0) + allocator.device_pool.notify_layer_end_for_backup(0) self.assertEqual(host_pool.layer_backups[0][2], 0) self.assertEqual(controller.ack_write_queue, []) - controller.on_layer_kv_stored(1) + allocator.device_pool.notify_layer_end_for_backup(1) self.assertEqual([x[2] for x in host_pool.layer_backups], [0, 1]) self.assertEqual(len(controller.ack_write_queue), 1) @@ -708,13 +712,13 @@ class TestHiCacheControllerCPWrite(CustomTestCase): reservation = controller.reserve_write_cp(logical_locs, node_id=85) controller.submit_write_cp_per_layer(reservation, catch_up_all_layers=False) - controller.on_layer_kv_stored(0, source="target") + allocator.device_pool.notify_layer_end_for_backup(0) self.assertEqual([x[2] for x in host_pool.layer_backups], [0]) self.assertEqual(draft_host_pool.layer_backups, []) self.assertEqual(controller.ack_write_queue, []) - controller.on_layer_kv_stored(0, source="draft") + draft_device_pool.notify_layer_end_for_backup(0) self.assertEqual([x[2] for x in draft_host_pool.layer_backups], [0]) self.assertEqual(len(controller.ack_write_queue), 1) diff --git a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py index 89a4624f0..34d66a066 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -91,14 +91,20 @@ for _schema in ( raise from sglang.srt.managers.cache_controller import ( + HiCacheAck, HiCacheWriteFailure, HiCacheWriteReservation, ) -from sglang.srt.mem_cache.base_prefix_cache import EvictParams, MatchPrefixParams +from sglang.srt.mem_cache.base_prefix_cache import ( + EvictParams, + InsertParams, + MatchPrefixParams, +) from sglang.srt.mem_cache.hiradix_cache import ( CpHiCacheNodeMetadata, HiRadixCache, PendingHiCacheBackup, + PreparedCpHiCacheBackup, _compute_shared_hicache_token_capacities, ) from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode @@ -544,18 +550,25 @@ class FakeReserveWriteController: self.results = list(results) self.reservations = [] self.submitted = [] + self.submit_kwargs = [] self.evicted_host_indices = [] def reserve_write_cp(self, device_indices, priority=None, node_id=-1): self.reservations.append((device_indices.clone(), node_id)) result = self.results.pop(0) - return result(device_indices) if callable(result) else result + if not callable(result): + return result + try: + return result(device_indices, node_id=node_id) + except TypeError: + return result(device_indices) def submit_write_cp_all_layer(self, reservation): self.submitted.append(reservation) - def submit_write_cp_per_layer(self, reservation): + def submit_write_cp_per_layer(self, reservation, **kwargs): self.submitted.append(reservation) + self.submit_kwargs.append(kwargs) def evict_cp_host(self, metadata): self.evicted_host_indices.append(metadata.host_indices.clone()) @@ -883,6 +896,221 @@ class TestHiRadixCacheCPBackup(CustomTestCase): self.assertIn(node.id, cache.pending_host_backups) self.assertEqual(len(cache.cache_controller.submitted), 1) + def test_write_backup_rolls_back_local_success_when_peer_needs_host_eviction(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.tp_world_size = 2 + cache.tp_group = object() + cache.cache_controller = FakeReserveWriteController( + [ + lambda device_indices: make_write_reservation( + device_indices, node_id=130, host_start=90 + ), + lambda device_indices: make_write_reservation( + device_indices, node_id=130, host_start=100 + ), + ] + ) + evictions = [] + cache._evict_host_for_physical_slots = lambda required, synchronize_across_ranks=False: ( + evictions.append((required, synchronize_across_ranks)) or required + ) + cache.ongoing_write_through = {} + cache.pending_host_backups = {} + cache.inc_node_lock_ref = lambda node: None + + reduce_values = iter([4, 0]) + + def fake_all_reduce(tensor, op=None, group=None): + tensor.fill_(next(reduce_values)) + + node = TreeNode() + node.id = 130 + node.value = torch.arange(16, dtype=torch.int64) + + with patch("torch.distributed.all_reduce", side_effect=fake_all_reduce): + backed_len = cache.write_backup(node) + + self.assertEqual(backed_len, 16) + # The first local success must be released before all ranks enter the + # collective host eviction/retry branch. Otherwise CP ranks can diverge: + # successful ranks submit backup while failing ranks enter eviction + # all_reduce, which matches the observed Gloo 4-vs-1 mismatch. + self.assertEqual( + cache.cache_controller.evicted_host_indices[0].tolist(), + list(range(90, 106)), + ) + self.assertEqual(evictions, [(4, True)]) + self.assertEqual(len(cache.cache_controller.submitted), 1) + self.assertEqual( + cache.cache_controller.submitted[0].metadata.host_indices.tolist(), + list(range(100, 116)), + ) + + def test_insert_attaches_prepared_cp_backup_without_catchup_copy(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.root_node = TreeNode() + cache.root_node.key = RadixKey([]) + cache.evictable_size_ = 0 + cache.get_child_key_fn = lambda key: key.token_ids[0] + cache.key_match_fn = lambda lhs, rhs: 0 + cache.maybe_bigram_convert = lambda key, value: (key, value) + cache.is_eagle = False + cache.enable_storage = False + cache.enable_kv_cache_events = False + cache._update_leaf_status = lambda node: None + cache._update_host_leaf_status = lambda node: None + cache._record_store_event = lambda node: None + cache.ongoing_write_through = {} + cache.pending_host_backups = {} + locked_nodes = [] + cache.inc_node_lock_ref = lambda node: locked_nodes.append(node) + cache.cache_controller = types.SimpleNamespace(write_policy="write_through") + + value = torch.arange(16, dtype=torch.int64) + reservation = make_write_reservation(value, node_id=131, host_start=110) + prepared = PreparedCpHiCacheBackup( + node_id=131, + reservation=reservation, + metadata=reservation.metadata, + logical_len=16, + ) + + result = cache.insert( + InsertParams( + key=RadixKey(list(range(16))), + value=value, + cp_hicache_prepared_backup=prepared, + ) + ) + + node = cache.root_node.children[0] + self.assertEqual(result.prefix_len, 0) + self.assertEqual(node.id, 131) + self.assertTrue(prepared.attached) + self.assertIs(cache.ongoing_write_through[131], node) + self.assertIs(cache.pending_host_backups[131].node, node) + self.assertIs(cache.pending_host_backups[131].metadata, reservation.metadata) + self.assertEqual(cache.pending_host_backups[131].logical_len, 16) + self.assertEqual(locked_nodes, [node]) + + def test_prepare_write_backup_for_req_registers_before_forward_without_catchup( + self, + ): + cache = HiRadixCache.__new__(HiRadixCache) + cache.disable = False + cache._uses_cp_hicache = True + cache.is_eagle = False + cache.page_size = 1 + cache.req_to_token_pool = types.SimpleNamespace( + req_to_token=torch.arange(16, dtype=torch.int64).view(1, 16) + ) + cache.cache_controller = FakeReserveWriteController( + [ + lambda device_indices, node_id: make_write_reservation( + device_indices, node_id=node_id, host_start=130 + ) + ] + ) + + req = types.SimpleNamespace( + rid="rid-prepare", + fill_ids=list(range(16)), + cache_protected_len=0, + req_pool_idx=0, + is_chunked=0, + cp_hicache_prepared_backup=None, + ) + + cache.prepare_write_backup_for_req(req) + + prepared = req.cp_hicache_prepared_backup + self.assertIsNotNone(prepared) + self.assertEqual(prepared.logical_len, 16) + self.assertEqual(len(cache.cache_controller.submitted), 1) + self.assertIs(cache.cache_controller.submitted[0], prepared.reservation) + self.assertEqual( + cache.cache_controller.submit_kwargs, + [{"catch_up_all_layers": False}], + ) + + def test_prepare_write_backup_for_req_skips_existing_insert_prefix(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache.disable = False + cache._uses_cp_hicache = True + cache.is_eagle = False + cache.page_size = 1 + cache.root_node = TreeNode() + cache.root_node.key = RadixKey([]) + cache.root_node.children = {} + cache.get_child_key_fn = lambda key: key.token_ids[0] + def key_match(lhs, rhs): + matched = 0 + for left, right in zip(lhs.token_ids, rhs.token_ids): + if left != right: + break + matched += 1 + return matched + + cache.key_match_fn = key_match + cache.maybe_bigram_convert = lambda key, value=None: (key, value) + cache.req_to_token_pool = types.SimpleNamespace( + req_to_token=torch.arange(16, dtype=torch.int64).view(1, 16) + ) + cache.cache_controller = FakeReserveWriteController([]) + + existing = TreeNode() + existing.id = 201 + existing.parent = cache.root_node + existing.key = RadixKey(list(range(16))) + existing.value = torch.arange(16, dtype=torch.int64) + cache.root_node.children[0] = existing + + req = types.SimpleNamespace( + rid="rid-existing-prefix", + fill_ids=list(range(16)), + extra_key=None, + cache_protected_len=8, + req_pool_idx=0, + is_chunked=0, + cp_hicache_prepared_backup=None, + ) + + cache.prepare_write_backup_for_req(req) + + self.assertIsNone(req.cp_hicache_prepared_backup) + self.assertEqual(cache.cache_controller.reservations, []) + self.assertEqual(cache.cache_controller.submitted, []) + + def test_rollback_unattached_prepared_cp_backup_removes_orphan_ack(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + value = torch.arange(4, dtype=torch.int64) + reservation = make_write_reservation(value, node_id=132, host_start=120) + prepared = PreparedCpHiCacheBackup( + node_id=132, + reservation=reservation, + metadata=reservation.metadata, + logical_len=4, + ) + evicted = [] + + class ReadyEvent: + def synchronize(self): + pass + + cache.cache_controller = types.SimpleNamespace( + pending_layer_writes={}, + ack_write_queue=[HiCacheAck(ReadyEvent(), ReadyEvent(), [132])], + evict_cp_host=lambda metadata: evicted.append(metadata), + ) + + cache._rollback_prepared_cp_backup(prepared, "test") + + self.assertEqual(cache.cache_controller.ack_write_queue, []) + self.assertEqual(evicted, [reservation.metadata]) + def test_write_backup_cp_retry_failure_leaves_node_device_only(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True @@ -1641,5 +1869,102 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase): self.assertEqual(cache.ongoing_load_back, {}) +class TestCPHiCacheLayerBackupNotifications(CustomTestCase): + def _call_store_index_fast_path(self, out_loc): + from sglang.srt.layers.attention.nsa import nsa_indexer + + notifications = [] + fused_calls = [] + + class FakePool: + page_size = 64 + start_layer = 0 + + def get_index_k_with_scale_buffer(self, layer_id): + return torch.empty((1,), dtype=torch.uint8) + + def notify_layer_kv_stored_for_backup(self, layer_id, source="kv"): + notifications.append((layer_id, source)) + + forward_batch = types.SimpleNamespace(token_to_kv_pool=FakePool()) + indexer = types.SimpleNamespace() + + def fake_fused_store(key, buf, loc, page_size): + fused_calls.append((key, buf, loc.clone(), page_size)) + + with ( + patch.object(nsa_indexer, "_is_cuda", True), + patch.object(nsa_indexer, "_is_fp8_fnuz", False), + patch.object(nsa_indexer, "can_use_nsa_fused_store", return_value=True), + patch.object( + nsa_indexer, "fused_store_index_k_cache", side_effect=fake_fused_store + ), + ): + nsa_indexer.Indexer._store_index_k_cache( + indexer, + forward_batch, + layer_id=5, + key=torch.empty((max(out_loc.numel(), 1), 128), dtype=torch.float32), + out_loc_override=out_loc, + ) + + return notifications, fused_calls + + def test_nsa_indexer_fused_store_does_not_notify_cp_hicache_layer_backup(self): + notifications, fused_calls = self._call_store_index_fast_path( + torch.tensor([1, 2, 3], dtype=torch.int64) + ) + + self.assertEqual(notifications, []) + self.assertEqual(len(fused_calls), 1) + self.assertEqual(fused_calls[0][2].tolist(), [1, 2, 3]) + + def test_nsa_indexer_empty_store_does_not_notify_cp_hicache_layer_backup(self): + notifications, fused_calls = self._call_store_index_fast_path( + torch.empty((0,), dtype=torch.int64) + ) + + self.assertEqual(notifications, []) + self.assertEqual(fused_calls, []) + + def test_cp_shared_zero_local_index_store_does_not_notify_layer_backup(self): + from sglang.srt.layers.attention.nsa import nsa_indexer + + notifications = [] + + class FakePool: + page_size = 64 + + def notify_layer_kv_stored_for_backup(self, layer_id, source="kv"): + notifications.append((layer_id, source)) + + forward_batch = types.SimpleNamespace(token_to_kv_pool=FakePool()) + indexer = types.SimpleNamespace(nsa_enable_prefill_cp=True) + + with ( + patch.object(nsa_indexer, "nsa_use_prefill_cp", return_value=True), + patch.object( + nsa_indexer, + "get_cp_shared_kv_local_out_cache_loc", + return_value=torch.empty((0,), dtype=torch.int64), + ), + patch.object( + nsa_indexer, + "get_cp_shared_kv_local_physical_out_cache_loc", + side_effect=AssertionError("must not require physical locs"), + ), + ): + handled = nsa_indexer.Indexer._store_cp_shared_local_index_k_cache( + indexer, + forward_batch=forward_batch, + layer_id=6, + local_key=torch.empty((0, 128), dtype=torch.float32), + act_quant=None, + ) + + self.assertTrue(handled) + self.assertEqual(notifications, []) + + if __name__ == "__main__": unittest.main()