Keep CP HiCache backup state replicated on radix nodes

Rank-local CP HiCache async write queues can drain at different times across ranks. Radix structural decisions that consult those local queues can diverge, so backup-pending state now lives on the replicated radix node and uses a separate backup lock ref from request locks. Commit and rollback clear the replicated marker, and split/prune/hit-count probes use the marker rather than local pending dictionaries.

Constraint: CP ranks must make identical radix structural decisions even when local async write acknowledgements drain at different times

Rejected: Drain write acknowledgements before every split/prune decision | still rank-local and adds hot-path synchronization

Rejected: Treat backup lock_ref as a normal request lock | hides the distinction between request protection and backup-pending protection

Confidence: high

Scope-risk: moderate

Directive: Do not derive CP HiCache split/prune/write-through decisions from pending_host_backups or ongoing_write_through without a replicated node marker

Tested: local py_compile python/sglang/srt/mem_cache/hiradix_cache.py python/sglang/srt/mem_cache/radix_cache.py

Tested: local git diff --check

Tested: remote cjy-glm5-new pytest replicated marker test plus targeted CP HiCache metadata subset, 10 passed

Not-tested: full test_cp_hicache_metadata.py because current branch has unrelated fixture drift around enable_cp_l3/cp_shared_l2_page_allocator/running-count setup

Co-authored-by: OmX <omx@oh-my-codex.dev>

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -3524,6 +3524,54 @@ class HiRadixCache(RadixCache):
             self, "pending_host_backups", {}
         )
 
+    def _node_backup_outstanding_replicated(self, node: TreeNode) -> bool:
+        return bool(getattr(node, "cp_backup_pending", False))
+
+    def _node_request_lock_ref(self, node: TreeNode) -> int:
+        return max(
+            0,
+            int(getattr(node, "lock_ref", 0) or 0)
+            - int(getattr(node, "cp_backup_lock_ref", 0) or 0),
+        )
+
+    def _mark_cp_backup_pending(self, node: TreeNode, *, locked: bool) -> None:
+        if self._node_backup_outstanding_replicated(node):
+            raise RuntimeError(
+                "CP HiCache backup marker was attached twice: "
+                f"node_id={getattr(node, 'id', '?')} "
+                f"cp_backup_lock_ref={getattr(node, 'cp_backup_lock_ref', '?')} "
+                f"lock_ref={getattr(node, 'lock_ref', '?')}"
+            )
+        node.cp_backup_pending = True
+        if locked:
+            self.inc_node_lock_ref(node)
+            node.cp_backup_lock_ref = (
+                int(getattr(node, "cp_backup_lock_ref", 0) or 0) + 1
+            )
+
+    def _clear_cp_backup_pending(
+        self, node: TreeNode, *, locked: bool, reason: str
+    ) -> None:
+        if not self._node_backup_outstanding_replicated(node):
+            raise RuntimeError(
+                "CP HiCache backup marker was cleared without being pending: "
+                f"node_id={getattr(node, 'id', '?')} reason={reason} "
+                f"cp_backup_lock_ref={getattr(node, 'cp_backup_lock_ref', '?')} "
+                f"lock_ref={getattr(node, 'lock_ref', '?')}"
+            )
+        if locked:
+            cp_backup_lock_ref = int(getattr(node, "cp_backup_lock_ref", 0) or 0)
+            if cp_backup_lock_ref <= 0:
+                raise RuntimeError(
+                    "CP HiCache backup lock marker underflow: "
+                    f"node_id={getattr(node, 'id', '?')} reason={reason} "
+                    f"cp_backup_lock_ref={cp_backup_lock_ref} "
+                    f"lock_ref={getattr(node, 'lock_ref', '?')}"
+                )
+            node.cp_backup_lock_ref = cp_backup_lock_ref - 1
+            self.dec_node_lock_ref(node)
+        node.cp_backup_pending = False
+
     def _node_host_write_ready(self, node: TreeNode) -> bool:
         return not self._node_host_write_pending(node)
 
@@ -3550,8 +3598,7 @@ class HiRadixCache(RadixCache):
         node.cp_hicache = pending.metadata
         node.host_value = None
         self._cp_account_enter_committed(pending.metadata)
-        if pending.locked:
-            self.dec_node_lock_ref(node)
+        self._clear_cp_backup_pending(node, locked=pending.locked, reason="commit")
         # B1: this commit is gated by the writing_check ReduceOp.MIN frontier, which
         # is the all-ranks-done consensus. Mark the shared-L2 object committed here
         # (replacing the per-(rank,layer,payload) gather quorum) so it transitions
@@ -3616,8 +3663,7 @@ class HiRadixCache(RadixCache):
         if node.cp_hicache is pending.metadata:
             node.cp_hicache = None
             node.host_len = 0
-        if pending.locked:
-            self.dec_node_lock_ref(node)
+        self._clear_cp_backup_pending(node, locked=pending.locked, reason="rollback")
         return node
 
     def _evict_cp_host_for_write_admission(
@@ -3905,7 +3951,7 @@ class HiRadixCache(RadixCache):
             )
         node.id = prepared.node_id
         self.ongoing_write_through[node.id] = node
-        self.inc_node_lock_ref(node)
+        self._mark_cp_backup_pending(node, locked=True)
         self.pending_host_backups[node.id] = PendingHiCacheBackup(
             node=node,
             metadata=prepared.metadata,
@@ -4248,9 +4294,9 @@ class HiRadixCache(RadixCache):
         while stack:
             current = stack.pop()
             if (
-                current.lock_ref > 0
+                self._node_request_lock_ref(current) > 0
                 or current.host_ref_counter > 0
-                or self._node_host_write_pending(current)
+                or self._node_backup_outstanding_replicated(current)
             ):
                 return True
             stack.extend(current.children.values())
@@ -4274,17 +4320,15 @@ class HiRadixCache(RadixCache):
         return
 
     def _cp_node_split_still_pending(self, node: TreeNode) -> bool:
-        if not self._uses_cp_hicache or not self._node_host_write_pending(node):
+        if not self._uses_cp_hicache:
             return False
-        self._cp_drain_write_acks_before_pending_split()
-        return self._node_host_write_pending(node)
+        return self._node_backup_outstanding_replicated(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(
@@ -4636,8 +4680,7 @@ class HiRadixCache(RadixCache):
                 return 0
 
             self.ongoing_write_through[node.id] = node
-            if not write_back:
-                self.inc_node_lock_ref(node)
+            self._mark_cp_backup_pending(node, locked=not write_back)
             self.pending_host_backups[node.id] = PendingHiCacheBackup(
                 node=node,
                 metadata=result.metadata,
@@ -4726,7 +4769,7 @@ class HiRadixCache(RadixCache):
             return
         node.hit_count += 1
 
-        if self._uses_cp_hicache and self._node_host_write_pending(node):
+        if self._uses_cp_hicache and self._node_backup_outstanding_replicated(node):
             return
 
         if not self._node_backuped(node):
@@ -6122,9 +6165,16 @@ class HiRadixCache(RadixCache):
     def _split_node(self, key: RadixKey, child: TreeNode, split_len: int):
         if (
             self._uses_cp_hicache
-            and self._node_host_write_pending(child)
+            and self._node_backup_outstanding_replicated(child)
         ):
             raise HiCachePendingBackupSplit(child)
+        if self._uses_cp_hicache and getattr(child, "cp_backup_lock_ref", 0):
+            raise RuntimeError(
+                "CP HiCache split found backup lock without pending marker: "
+                f"node_id={getattr(child, 'id', '?')} "
+                f"cp_backup_lock_ref={getattr(child, 'cp_backup_lock_ref', '?')} "
+                f"lock_ref={getattr(child, 'lock_ref', '?')}"
+            )
         # child node split into new_node -> child
         new_node = TreeNode(priority=child.priority)
         new_node.children = {self.get_child_key_fn(key[split_len:]): child}
@@ -6132,6 +6182,9 @@ class HiRadixCache(RadixCache):
         new_node.lock_ref = child.lock_ref
         new_node.pin_expiry = child.pin_expiry
         new_node.pin_ttl = child.pin_ttl
+        if self._uses_cp_hicache:
+            new_node.cp_backup_pending = False
+            new_node.cp_backup_lock_ref = 0
         # If child is pinned, new parent inherits a host_ref_counter hold
         if child.pin_expiry > 0:
             new_node.host_ref_counter += 1

python/sglang/srt/mem_cache/radix_cache.py

@@ -229,6 +229,12 @@ class TreeNode:
         self.host_value: Optional[torch.Tensor] = None
         self.host_len: int = 0
         self.cp_hicache = None
+        # CP HiCache per-layer backup state that must stay replicated on the
+        # radix node itself.  Do not derive radix structural decisions from
+        # rank-local async write queues; those queues can drain at different
+        # times on different CP ranks.
+        self.cp_backup_pending: bool = False
+        self.cp_backup_lock_ref: int = 0
         # store hash values of each pages
         self.hash_value: Optional[List[str]] = None
         # priority for priority-aware eviction