From c2d25ff591df29aa2d451ae526f3f2b92f5fc737 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Tue, 2 Jun 2026 04:51:39 +0800 Subject: [PATCH] Prevent CP HiCache pending-backup splits from killing chunked prefill Chunked prefill can revisit a sub-page CP HiCache tail while a per-layer backup is still in flight. The old insert path split the radix node first and only then tried to prune the stale tail, so an unprunable pending backup raised after tree mutation and propagated to the scheduler. This makes split/prune atomic from the radix-tree perspective: drain completed write acks only on the conflict path, preflight pending/unprunable state before split, and return a deferred insert result when the backup is still in flight. Unfinished requests keep their KV ownership for transfer/release instead of rematching or freeing pages under a stale tree state. Constraint: CP HiCache backup metadata is node/page owned and cannot be repartitioned while per-layer D2H is pending Rejected: Split the pending backup node | would require repartitioning in-flight backup metadata and host reservations Rejected: Delete the stale tail unconditionally | risks freeing device/host pages still owned by pending backup Confidence: high Scope-risk: moderate Directive: Do not move stale-tail prune after split without a preflight; pending backup split conflicts must remain non-mutating Tested: remote g0034 container py_compile for mem_cache files Tested: remote g0034 PYTHONPATH=python pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py => 115 passed Tested: local git diff --check Not-tested: full chunked prefill ETE after service restart Co-authored-by: OmX --- ..._prefill_cp_page_aligned_cache_contract.md | 15 ++ .../sglang/srt/mem_cache/base_prefix_cache.py | 1 + python/sglang/srt/mem_cache/hiradix_cache.py | 74 ++++++++- python/sglang/srt/mem_cache/radix_cache.py | 69 ++++++--- .../mem_cache/test_cp_hicache_metadata.py | 144 ++++++++++++++++++ 5 files changed, 279 insertions(+), 24 deletions(-) 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