diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index ad3d2b57d..e99da5581 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index df824ac2c..2508c1fcd 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/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 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 cd5fac5c4..6283d14af 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/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( diff --git a/test/registered/unit/mem_cache/test_cp_hicache_replicated_pending_marker.py b/test/registered/unit/mem_cache/test_cp_hicache_replicated_pending_marker.py new file mode 100644 index 000000000..e2c094c72 --- /dev/null +++ b/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()