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 <omx@oh-my-codex.dev>

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

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)

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)