diff --git a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md index 6266a305f..785d987a0 100644 --- a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md +++ b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md @@ -5021,3 +5021,18 @@ Implication: - Production LPF speedup requires changing allocation/residency policy, not only tensor layout: allocate compact/extent-aware host pages for each backed node or add a transfer-time gather/staging path. + +### C118 — 2026-06-02 Chunked prefill can hit CP HiCache stale-tail split while backup is pending + +Finding: + +- Remote crash at `process_batch_result_disagg_prefill -> cache_unfinished_req -> insert -> _cp_prune_stale_tail_after_page_floor` raised `HiCachePendingBackupSplit` for a stale tail while CP HiCache backup was pending. +- The current insert path performs `_split_node()` first and only then calls `_cp_prune_stale_tail_after_page_floor()`. If the stale tail subtree is unprunable because a backup lock/ref is pending, the exception is raised after the radix tree has already been mutated. +- `_split_node()` only checks `pending_host_backups`; the prune guard checks both `ongoing_write_through` and `pending_host_backups` via `_node_host_write_pending()`. Therefore an ongoing-only write can pass the split precheck and fail during prune. +- `match_prefix()` catches `HiCachePendingBackupSplit` and defers scheduling, but `insert()` / `cache_unfinished_req()` do not convert this condition into a non-fatal deferred insertion. Under chunked prefill the conflict can surface after forward, where requeueing the request is no longer possible. + +Required correction: + +- Make stale-tail split/prune atomic from the radix tree perspective: preflight pending/unprunable state before `_split_node()` and drain completed write acks once on the conflict path. +- If the stale tail is still pending, do not mutate the tree and do not crash the scheduler. Defer/skip this cache insertion attempt while keeping request KV ownership intact for transfer/release. +- Do not delete the stale tail or split an in-flight backup node. Per-layer backup ownership remains node/page based and cannot be repartitioned mid-flight. diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 66e09128e..7aa847f4d 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -70,6 +70,7 @@ class InsertResult: prefix_len: int mamba_exist: bool = False + pending_backup_deferred_node: Any = None @dataclasses.dataclass diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 56b85f5c4..e340f971e 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -2191,6 +2191,51 @@ class HiRadixCache(RadixCache): stack.extend(current.children.values()) return False + def _cp_drain_write_acks_before_pending_split(self) -> None: + """Drain already-finished CP write acks before declaring split blocked. + + This is intentionally only called on the rare stale-tail/pending-split + path. It may run the final write-visibility collective when a completed + ack is available, but it avoids adding a collective to the normal insert + hot path. + """ + + if not self._uses_cp_hicache or not getattr( + self, "ongoing_write_through", None + ): + return + self.writing_check() + + def _cp_node_split_still_pending(self, node: TreeNode) -> bool: + if not self._uses_cp_hicache or not self._node_host_write_pending(node): + return False + self._cp_drain_write_acks_before_pending_split() + return self._node_host_write_pending(node) + + def _cp_stale_tail_prune_still_blocked(self, stale_tail: TreeNode) -> bool: + if not self._uses_cp_hicache or not self._cp_subtree_has_unprunable_state( + stale_tail + ): + return False + self._cp_drain_write_acks_before_pending_split() + return self._cp_subtree_has_unprunable_state(stale_tail) + + def _cp_defer_insert_for_pending_backup_split( + self, node: TreeNode, prefix_len: int, *, reason: str + ) -> InsertResult: + self._warn_cp_hicache_fallback( + "insert_deferred_pending_backup_split", + reason, + node_id=getattr(node, "id", None), + prefix_len=prefix_len, + ongoing=len(getattr(self, "ongoing_write_through", {})), + pending=len(getattr(self, "pending_host_backups", {})), + ) + return InsertResult( + prefix_len=prefix_len, + pending_backup_deferred_node=node, + ) + def _cp_prune_stale_tail_after_page_floor( self, parent: TreeNode, stale_tail: TreeNode ) -> None: @@ -2212,7 +2257,7 @@ class HiRadixCache(RadixCache): f"parent_id={getattr(parent, 'id', None)} " f"tail_id={getattr(stale_tail, 'id', None)}" ) - if self._cp_subtree_has_unprunable_state(stale_tail): + if self._cp_stale_tail_prune_still_blocked(stale_tail): raise HiCachePendingBackupSplit(stale_tail) self._cp_delete_stale_tail_subtree(stale_tail) @@ -3683,10 +3728,7 @@ class HiRadixCache(RadixCache): if prefix_len <= 0: break if prefix_len < len(child.key): - if ( - self._uses_cp_hicache - and child.id in getattr(self, "pending_host_backups", {}) - ): + if self._cp_node_split_still_pending(child): raise HiCachePendingBackupSplit(child) if ( self._uses_cp_hicache @@ -3703,6 +3745,11 @@ class HiRadixCache(RadixCache): ) if prefix_len == 0: break + if ( + prune_stale_tail_after_split + and self._cp_stale_tail_prune_still_blocked(child) + ): + raise HiCachePendingBackupSplit(child) new_node = self._split_node(child.key, child, prefix_len) if prune_stale_tail_after_split: self._cp_prune_stale_tail_after_page_floor(new_node, child) @@ -3724,7 +3771,7 @@ class HiRadixCache(RadixCache): def _split_node(self, key: RadixKey, child: TreeNode, split_len: int): if ( self._uses_cp_hicache - and child.id in getattr(self, "pending_host_backups", {}) + and self._node_host_write_pending(child) ): raise HiCachePendingBackupSplit(child) # child node split into new_node -> child @@ -3833,6 +3880,21 @@ class HiRadixCache(RadixCache): ) if prefix_len <= 0: break + if self._cp_node_split_still_pending(node): + return self._cp_defer_insert_for_pending_backup_split( + node, + total_prefix_length, + reason="node_pending_backup", + ) + if ( + prune_stale_tail_after_split + and self._cp_stale_tail_prune_still_blocked(node) + ): + return self._cp_defer_insert_for_pending_backup_split( + node, + total_prefix_length, + reason="stale_tail_unprunable", + ) new_node = self._split_node(node.key, node, prefix_len) if prune_stale_tail_after_split: self._cp_prune_stale_tail_after_page_floor(new_node, node) diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 6765bd548..d9b0d2b14 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -554,24 +554,42 @@ class RadixCache(BasePrefixCache): 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._free_kv_indices_range( - kv_indices, - req.cache_protected_len, - new_prefix_len, - include_partial_tail=self._should_free_cp_exact_match_tail( - start=req.cache_protected_len, - end=new_prefix_len, - key_len=len(keys), - ), - ) + if getattr(result, "pending_backup_deferred_node", None) is not None: + 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, "pending_backup_split_deferred" + ) + self._free_kv_indices_range( + kv_indices, + req.cache_protected_len, + len(keys), + include_partial_tail=True, + ) + else: + 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._free_kv_indices_range( + kv_indices, + req.cache_protected_len, + new_prefix_len, + include_partial_tail=self._should_free_cp_exact_match_tail( + start=req.cache_protected_len, + end=new_prefix_len, + key_len=len(keys), + ), + ) else: if prepared_cp_backup is not None and hasattr( self, "_rollback_prepared_cp_backup" @@ -631,6 +649,21 @@ class RadixCache(BasePrefixCache): cp_hicache_prepared_backup=prepared_cp_backup, ) ) + if getattr(result, "pending_backup_deferred_node", None) is not None: + 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, "unfinished_pending_backup_split_deferred" + ) + req.cp_hicache_prepared_backup = None + if len(values) < len(kv_indices): + req.prefix_indices = torch.cat([values, kv_indices[len(values) :]]) + else: + req.prefix_indices = values + return if ( prepared_cp_backup is not None and not getattr(prepared_cp_backup, "attached", False) 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 630f9c180..b4565abda 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -1612,6 +1612,96 @@ class TestHiRadixCacheCPBackup(CustomTestCase): self.assertEqual(allocator.freed, []) self.assertEqual(req.cache_protected_len, 6) + def test_cache_unfinished_req_deferred_insert_preserves_request_kv(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache.disable = False + cache._uses_cp_hicache = True + cache.is_eagle = False + cache.page_size = 4 + deferred_node = TreeNode() + + class Pool: + req_to_token = torch.arange(8, dtype=torch.int64).view(1, 8) + + def write(self, index, values): + raise AssertionError("deferred insert must not rewrite req pool") + + cache.req_to_token_pool = Pool() + cache.token_to_kv_pool_allocator = RecordingTokenAllocator() + cache.insert = lambda params: types.SimpleNamespace( + prefix_len=1, + pending_backup_deferred_node=deferred_node, + ) + cache.match_prefix = lambda params: (_ for _ in ()).throw( + AssertionError("deferred insert must not rematch") + ) + cache._free_kv_indices_range = lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("deferred unfinished insert must keep request KV") + ) + old_last_node = TreeNode() + req = types.SimpleNamespace( + fill_ids=list(range(6)), + req_pool_idx=0, + extra_key=None, + cache_protected_len=1, + last_node=old_last_node, + prefix_indices=torch.empty((0,), dtype=torch.int64), + cp_hicache_prepared_backup=None, + ) + + cache.cache_unfinished_req(req, chunked=True) + + self.assertEqual(req.cache_protected_len, 1) + self.assertIs(req.last_node, old_last_node) + self.assertEqual(req.prefix_indices.tolist(), list(range(6))) + self.assertEqual(cache.token_to_kv_pool_allocator.freed, []) + + def test_cache_unfinished_req_deferred_insert_rolls_back_unattached_backup(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache.disable = False + cache._uses_cp_hicache = True + cache.is_eagle = False + cache.page_size = 4 + deferred_node = TreeNode() + rollbacks = [] + + class Pool: + req_to_token = torch.arange(8, dtype=torch.int64).view(1, 8) + + cache.req_to_token_pool = Pool() + cache.token_to_kv_pool_allocator = RecordingTokenAllocator() + cache.insert = lambda params: types.SimpleNamespace( + prefix_len=1, + pending_backup_deferred_node=deferred_node, + ) + cache.match_prefix = lambda params: (_ for _ in ()).throw( + AssertionError("deferred insert must not rematch") + ) + cache._free_kv_indices_range = lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("deferred unfinished insert must keep request KV") + ) + cache._rollback_prepared_cp_backup = lambda prepared, reason: rollbacks.append( + (prepared, reason) + ) + prepared = types.SimpleNamespace(attached=False) + req = types.SimpleNamespace( + fill_ids=list(range(6)), + req_pool_idx=0, + extra_key=None, + cache_protected_len=1, + last_node=TreeNode(), + prefix_indices=torch.empty((0,), dtype=torch.int64), + cp_hicache_prepared_backup=prepared, + ) + + cache.cache_unfinished_req(req, chunked=False) + + self.assertEqual( + rollbacks, [(prepared, "unfinished_pending_backup_split_deferred")] + ) + self.assertIsNone(req.cp_hicache_prepared_backup) + self.assertEqual(req.prefix_indices.tolist(), list(range(6))) + def test_prepare_write_backup_for_req_skips_existing_insert_prefix(self): cache = HiRadixCache.__new__(HiRadixCache) cache.disable = False @@ -3011,6 +3101,60 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase): self.assertEqual(new_tail.key.token_ids, [4]) self.assertEqual(new_tail.value.tolist(), [4]) + def test_cp_insert_defers_stale_tail_prune_while_backup_pending(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.disable = False + cache.is_eagle = False + cache.page_size = 4 + cache.pending_host_backups = {} + cache.ongoing_write_through = {} + cache.cache_controller = types.SimpleNamespace( + has_draft_hicache=False, + write_policy="write_through", + ack_write_queue=[], + ) + cache.tp_world_size = 1 + cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4]) + cache.key_match_fn = lambda child_key, key: _key_match_paged( + child_key, key, page_size=4 + ) + cache.maybe_bigram_convert = lambda key, value=None: (key, value) + cache._update_leaf_status = lambda node: None + cache._update_host_leaf_status = lambda node: None + cache._inc_hit_count = lambda *args, **kwargs: None + cache._record_store_event = lambda node: None + cache.evictable_size_ = 0 + cache.enable_storage = False + cache.enable_kv_cache_events = False + root = TreeNode() + root.key = RadixKey([]) + root.value = torch.empty((0,), dtype=torch.int64) + root.children = {} + cache.root_node = root + node = TreeNode() + node.id = 152 + node.parent = root + node.key = RadixKey(list(range(6))) + node.value = torch.arange(6, dtype=torch.int64) + node.host_len = 0 + node.cp_hicache = None + root.children[(0, 1, 2, 3)] = node + cache.ongoing_write_through[node.id] = node + + result = cache.insert( + InsertParams( + key=RadixKey(list(range(10))), + value=torch.arange(10, dtype=torch.int64), + ) + ) + + self.assertEqual(result.prefix_len, 0) + self.assertIs(result.pending_backup_deferred_node, node) + self.assertIs(root.children[(0, 1, 2, 3)], node) + self.assertEqual(node.key.token_ids, list(range(6))) + self.assertEqual(node.value.tolist(), list(range(6))) + def test_cp_insert_extends_from_page_boundary_after_exact_valid_tail(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True