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

test/registered/unit/mem_cache/test_cp_hicache_metadata.py

@@ -1380,6 +1380,10 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
         cache.pending_host_backups[node.id] = PendingHiCacheBackup(
             node=node, metadata=metadata, logical_len=64
         )
+        node.cp_backup_pending = True
+        node.cp_backup_lock_ref = 1
+        node.lock_ref = 1
+        cache._cp_account_enter_pending(metadata)
 
         self.assertFalse(cache._node_backuped(node))
 
@@ -1412,6 +1416,10 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
         cache.pending_host_backups[node.id] = PendingHiCacheBackup(
             node=node, metadata=metadata, logical_len=64
         )
+        node.cp_backup_pending = True
+        node.cp_backup_lock_ref = 1
+        node.lock_ref = 1
+        cache._cp_account_enter_pending(metadata)
 
         cache._rollback_pending_backup(node.id)
 
@@ -1494,6 +1502,9 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
             page_size=1,
         )
         cache.ongoing_write_through[node.id] = node
+        node.cp_backup_pending = True
+        node.cp_backup_lock_ref = 1
+        node.lock_ref = 1
 
         cache._inc_hit_count(node)
 
@@ -2455,6 +2466,9 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
             )
         }
         cache.dec_node_lock_ref = lambda n: dec_locked.append(n.id)
+        node.cp_backup_pending = True
+        node.cp_backup_lock_ref = 1
+        node.lock_ref = 1
 
         finish_event = MagicMock()
         cancelled = []
@@ -3581,6 +3595,7 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
         cache.pending_host_backups[child.id] = PendingHiCacheBackup(
             node=child, metadata=metadata, logical_len=8
         )
+        child.cp_backup_pending = True
         root.children[0] = child
 
         result = cache.match_prefix(MatchPrefixParams(key=RadixKey([0, 1, 2, 9])))
@@ -4007,6 +4022,7 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
         node.cp_hicache = None
         root.children[(0, 1, 2, 3)] = node
         cache.ongoing_write_through[node.id] = node
+        node.cp_backup_pending = True
 
         result = cache.insert(
             InsertParams(

test/registered/unit/mem_cache/test_cp_hicache_replicated_pending_marker.py

@@ -0,0 +1,186 @@
+"""Regression coverage for CP HiCache backup-pending radix decisions.
+
+The radix split/prune/write-through probes must not branch on drain-timed,
+rank-local dictionaries such as ``ongoing_write_through`` or
+``pending_host_backups``.  Those dictionaries are advanced by local async write
+ack timing; using them for structural radix decisions lets CP ranks take
+different branches before the next collective.
+"""
+
+import sys
+import types
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+try:
+    import sgl_kernel  # noqa: F401
+except (ImportError, RuntimeError):
+    if "sgl_kernel" not in sys.modules:
+        stub = types.ModuleType("sgl_kernel")
+        stub.__path__ = []
+
+        def _getattr(name):
+            if name.startswith("__"):
+                raise AttributeError(name)
+            fn = lambda *a, **k: None
+            setattr(stub, name, fn)
+            return fn
+
+        stub.__getattr__ = _getattr
+        sys.modules["sgl_kernel"] = stub
+
+from sglang.srt.mem_cache.hiradix_cache import (  # noqa: E402
+    CpHiCacheNodeMetadata,
+    HiRadixCache,
+    PreparedCpHiCacheBackup,
+)
+from sglang.srt.mem_cache.radix_cache import TreeNode  # noqa: E402
+
+
+def _cp_cache():
+    cache = HiRadixCache.__new__(HiRadixCache)
+    cache._uses_cp_hicache = True
+    cache.ongoing_write_through = {}
+    cache.pending_host_backups = {}
+    return cache
+
+
+def _metadata(length: int) -> CpHiCacheNodeMetadata:
+    return CpHiCacheNodeMetadata(
+        logical_len=length,
+        owned_positions=torch.arange(length, dtype=torch.int64),
+        host_indices=torch.arange(length, dtype=torch.int64),
+        page_owners=torch.zeros(1, dtype=torch.int8),
+        page_size=length,
+    )
+
+
+def _reservation(metadata: CpHiCacheNodeMetadata) -> SimpleNamespace:
+    return SimpleNamespace(metadata=metadata)
+
+
+class TestCPHiCacheReplicatedPendingMarker(unittest.TestCase):
+    def test_split_probe_reads_marker_not_drain_timed_dicts(self):
+        cache = _cp_cache()
+        node = TreeNode(id=42)
+
+        cache.ongoing_write_through[node.id] = node
+        cache.pending_host_backups[node.id] = object()
+
+        self.assertFalse(
+            cache._cp_node_split_still_pending(node),
+            "split probe must ignore drain-timed local dicts when marker is clear",
+        )
+
+        cache.ongoing_write_through.clear()
+        cache.pending_host_backups.clear()
+        node.cp_backup_pending = True
+
+        self.assertTrue(
+            cache._cp_node_split_still_pending(node),
+            "split probe must honor the replicated node marker",
+        )
+
+    def test_inc_hit_count_reads_marker_not_drain_timed_dicts(self):
+        cache = _cp_cache()
+        cache.cache_controller = SimpleNamespace(write_policy="write_through")
+        cache.write_through_threshold = 1
+
+        node = TreeNode(id=7)
+        node.cp_backup_pending = True
+
+        def fail_if_reached(_node):
+            raise AssertionError("marker-pending node must not reach backuped check")
+
+        cache._node_backuped = fail_if_reached
+        cache._inc_hit_count(node)
+
+        reached_backuped_check = {"hit": False}
+
+        def record_backuped(_node):
+            reached_backuped_check["hit"] = True
+            return True
+
+        node.cp_backup_pending = False
+        cache.ongoing_write_through[node.id] = node
+        cache.pending_host_backups[node.id] = object()
+        cache._node_backuped = record_backuped
+
+        cache._inc_hit_count(node)
+
+        self.assertTrue(
+            reached_backuped_check["hit"],
+            "dict-pending alone must not short-circuit hit-count side effects",
+        )
+
+    def test_attach_commit_maintains_marker_and_backup_lock_ref(self):
+        cache = _cp_cache()
+        cache.inc_node_lock_ref = lambda node: setattr(
+            node, "lock_ref", node.lock_ref + 1
+        )
+        cache.dec_node_lock_ref = lambda node: setattr(
+            node, "lock_ref", node.lock_ref - 1
+        )
+        cache.cache_controller = SimpleNamespace(cp_shared_l2_page_allocator=None)
+
+        node = TreeNode()
+        node.value = torch.arange(4, dtype=torch.int64)
+        metadata = _metadata(4)
+        prepared = PreparedCpHiCacheBackup(
+            node_id=123,
+            reservation=_reservation(metadata),
+            metadata=metadata,
+            logical_len=4,
+        )
+
+        cache._attach_prepared_cp_backup(node, prepared)
+
+        self.assertTrue(node.cp_backup_pending)
+        self.assertEqual(node.cp_backup_lock_ref, 1)
+        self.assertEqual(node.lock_ref, 1)
+
+        cache._commit_pending_backup(node.id)
+
+        self.assertFalse(node.cp_backup_pending)
+        self.assertEqual(node.cp_backup_lock_ref, 0)
+        self.assertEqual(node.lock_ref, 0)
+        self.assertEqual(node.host_len, 4)
+
+    def test_rollback_clears_marker_and_backup_lock_ref(self):
+        cache = _cp_cache()
+        cache.inc_node_lock_ref = lambda node: setattr(
+            node, "lock_ref", node.lock_ref + 1
+        )
+        cache.dec_node_lock_ref = lambda node: setattr(
+            node, "lock_ref", node.lock_ref - 1
+        )
+        evicted = []
+        cache.cache_controller = SimpleNamespace(
+            ack_write_queue=[],
+            cancel_layer_write_state=lambda node_id: None,
+            evict_cp_host=lambda metadata: evicted.append(metadata) or 4,
+        )
+
+        node = TreeNode()
+        node.value = torch.arange(4, dtype=torch.int64)
+        metadata = _metadata(4)
+        prepared = PreparedCpHiCacheBackup(
+            node_id=124,
+            reservation=_reservation(metadata),
+            metadata=metadata,
+            logical_len=4,
+        )
+
+        cache._attach_prepared_cp_backup(node, prepared)
+        cache._rollback_pending_backup(node.id)
+
+        self.assertFalse(node.cp_backup_pending)
+        self.assertEqual(node.cp_backup_lock_ref, 0)
+        self.assertEqual(node.lock_ref, 0)
+        self.assertEqual(evicted, [metadata])
+
+
+if __name__ == "__main__":
+    unittest.main()