From f355fdd39e96b640a044f8033d46cd4632c5bc49 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Wed, 27 May 2026 09:50:47 +0800 Subject: [PATCH] Overlap CP HiCache backup without exposing partial host state CP shared KV with HiCache and EAGLE needs host backup to overlap forward while keeping radix visibility synchronous. The change reserves host slots before forward, drives target and draft backup from explicit layer-end hooks, and commits host visibility only after the final target/draft ack. It also probes the final insertion prefix before early reservation so repeated EAGLE prompts do not prepare duplicate suffix backups that later rollback as insert_miss. Constraint: CP ranks use independent shared-KV pools, so target/draft host state must remain atomically visible at the radix boundary. Constraint: Fused MLA and NSA store paths can bypass store-side notifier hooks, so layer end is the safer backup progress boundary. Rejected: Store-side backup notifier as the primary trigger | fused store and zero-local paths made notifier coverage fragile. Rejected: Reserve from cache_protected_len alone | EAGLE bigram/page alignment can make final insertion find a longer existing prefix and force duplicate rollback work. Confidence: medium Scope-risk: moderate Directive: Do not add per-layer CP collectives here; keep radix state synchronous and data transfer asynchronous/local-event driven. Tested: local git diff --check Tested: local py_compile for touched CP HiCache/cache-controller/deepseek/test files Tested: remote pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/managers/test_hicache_controller_cp.py -q (115 passed, 5 warnings) Not-tested: full GLM5 ETE server rerun after this commit --- ...refill_cp_hicache_per_layer_backup_plan.md | 74 +++- .../sglang/srt/managers/cache_controller.py | 8 +- python/sglang/srt/managers/schedule_batch.py | 1 + python/sglang/srt/managers/scheduler.py | 5 + .../sglang/srt/mem_cache/base_prefix_cache.py | 1 + python/sglang/srt/mem_cache/hiradix_cache.py | 309 ++++++++++++++-- python/sglang/srt/mem_cache/memory_pool.py | 13 +- python/sglang/srt/mem_cache/radix_cache.py | 30 +- python/sglang/srt/models/deepseek_v2.py | 20 ++ .../managers/test_hicache_controller_cp.py | 12 +- .../mem_cache/test_cp_hicache_metadata.py | 331 +++++++++++++++++- 11 files changed, 747 insertions(+), 57 deletions(-) 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()