From f353cf47021b60d42d5692680b8c8674d740cf9b Mon Sep 17 00:00:00 2001 From: leavelet Date: Mon, 22 Jun 2026 21:35:27 +0000 Subject: [PATCH] Revert "Fix CP HiCache writing_check deadlock at the source: rank-replicated symmetric attach" This reverts commit 85585fcf4b62ed41b1f66cfb15a9352d841ed6c5. --- python/sglang/srt/mem_cache/hiradix_cache.py | 58 ++---- python/sglang/srt/mem_cache/radix_cache.py | 6 - .../test_cp_hicache_symmetric_attach.py | 185 ------------------ 3 files changed, 11 insertions(+), 238 deletions(-) delete mode 100644 test/registered/unit/mem_cache/test_cp_hicache_symmetric_attach.py diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 9754fe77c..a199edcd0 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -2967,19 +2967,6 @@ class HiRadixCache(RadixCache): self, "pending_host_backups", {} ) - def _node_backup_outstanding_replicated(self, node: TreeNode) -> bool: - """Rank-replicated 'backup outstanding' signal for the CP radix probe path. - - True from the rank-replicated prepared-backup attach until the MIN-committed - commit/rollback (see TreeNode.cp_backup_pending). Unlike - _node_host_write_pending (which reads the drain-timed ongoing_write_through/ - pending_host_backups dicts and so carries per-rank ack-drain timing), this is - identical across CP ranks at every sub-tick instant, so split/defer/prune - decisions stay rank-uniform. The radix probe path must read THIS; only - writing_check / eviction may read the drain-timed _node_host_write_pending. - """ - return bool(getattr(node, "cp_backup_pending", False)) - def _node_host_write_ready(self, node: TreeNode) -> bool: return not self._node_host_write_pending(node) @@ -3006,10 +2993,6 @@ class HiRadixCache(RadixCache): node.cp_hicache = pending.metadata node.host_value = None self._cp_account_enter_committed(pending.metadata) - # Clear the replicated pending marker at the MIN-committed edge, BEFORE - # releasing the backup lock (keeps marker and lock_ref consistent). After - # commit, node.host_len > 0 (_node_backuped) is the authoritative signal. - node.cp_backup_pending = False if pending.locked: self.dec_node_lock_ref(node) # B1: this commit is gated by the writing_check ReduceOp.MIN frontier, which @@ -3071,7 +3054,6 @@ class HiRadixCache(RadixCache): if node.cp_hicache is pending.metadata: node.cp_hicache = None node.host_len = 0 - node.cp_backup_pending = False if pending.locked: self.dec_node_lock_ref(node) return node @@ -3360,13 +3342,7 @@ class HiRadixCache(RadixCache): f"node_len={len(node.value)} prepared_len={prepared.logical_len}" ) node.id = prepared.node_id - if getattr(node, "cp_backup_pending", False): - raise RuntimeError( - f"CP HiCache double-attach for node_id={node.id}: " - "cp_backup_pending already set" - ) self.ongoing_write_through[node.id] = node - node.cp_backup_pending = True self.inc_node_lock_ref(node) self.pending_host_backups[node.id] = PendingHiCacheBackup( node=node, @@ -3673,15 +3649,10 @@ class HiRadixCache(RadixCache): stack = [node] while stack: current = stack.pop() - # Read only rank-replicated state: the backup-pending marker (covers the - # backup lock -- the sole rank-local lock_ref skew) and host_ref_counter - # (set by replicated pin/protect_host). Do NOT read lock_ref here: its only - # cross-rank skew is the backup lock, already covered by the marker, and a - # request-held lock is rank-uniform but irrelevant to splittability (the - # page-floor logic governs that). Reading it would re-leak per-rank timing. if ( - self._node_backup_outstanding_replicated(current) + current.lock_ref > 0 or current.host_ref_counter > 0 + or self._node_host_write_pending(current) ): return True stack.extend(current.children.values()) @@ -3705,19 +3676,17 @@ class HiRadixCache(RadixCache): return def _cp_node_split_still_pending(self, node: TreeNode) -> bool: - # Rank-replicated marker only (see _node_backup_outstanding_replicated): the - # drain-timed _node_host_write_pending here is what let CP ranks diverge on the - # attach/defer fork. Do NOT call _cp_drain_write_acks_before_pending_split (it is - # a deliberate no-op; draining acks from the radix path desyncs collectives). - if not self._uses_cp_hicache: + if not self._uses_cp_hicache or not self._node_host_write_pending(node): return False - return self._node_backup_outstanding_replicated(node) + 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: - # _cp_subtree_has_unprunable_state now reads only rank-replicated state, so a - # single read is rank-uniform; no ack drain (the helper is a deliberate no-op). - if not self._uses_cp_hicache: + 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( @@ -4069,7 +4038,6 @@ class HiRadixCache(RadixCache): return 0 self.ongoing_write_through[node.id] = node - node.cp_backup_pending = True if not write_back: self.inc_node_lock_ref(node) self.pending_host_backups[node.id] = PendingHiCacheBackup( @@ -4160,7 +4128,7 @@ class HiRadixCache(RadixCache): return node.hit_count += 1 - if self._uses_cp_hicache and self._node_backup_outstanding_replicated(node): + if self._uses_cp_hicache and self._node_host_write_pending(node): return if not self._node_backuped(node): @@ -5547,7 +5515,7 @@ class HiRadixCache(RadixCache): def _split_node(self, key: RadixKey, child: TreeNode, split_len: int): if ( self._uses_cp_hicache - and self._node_backup_outstanding_replicated(child) + and self._node_host_write_pending(child) ): raise HiCachePendingBackupSplit(child) # child node split into new_node -> child @@ -5562,10 +5530,6 @@ class HiRadixCache(RadixCache): new_node.host_ref_counter += 1 new_node.key = child.key[:split_len].compacted() new_node.hit_count = child.hit_count - # Propagate the replicated backup-pending marker. The guard above raises - # HiCachePendingBackupSplit before splitting a pending node, so child is - # non-pending here and both halves carry False -- copied for fail-loud symmetry. - new_node.cp_backup_pending = child.cp_backup_pending # split value and host value if exists if child.evicted: diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 5fd3d64fd..df824ac2c 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -229,12 +229,6 @@ class TreeNode: self.host_value: Optional[torch.Tensor] = None self.host_len: int = 0 self.cp_hicache = None - # Rank-replicated marker: True from the (rank-replicated) prepared-backup - # attach until the (MIN-committed) commit/rollback. The CP radix probe path - # reads THIS instead of the drain-timed ongoing_write_through/pending_host_backups - # dicts, so split/defer/prune decisions stay identical across CP ranks. - # See docs_internal/cp_hicache_symmetric_attach_design.md. - self.cp_backup_pending: bool = False # 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_symmetric_attach.py b/test/registered/unit/mem_cache/test_cp_hicache_symmetric_attach.py deleted file mode 100644 index f7a1770a8..000000000 --- a/test/registered/unit/mem_cache/test_cp_hicache_symmetric_attach.py +++ /dev/null @@ -1,185 +0,0 @@ -"""Regression: CP-HiCache radix probes must read the rank-replicated cp_backup_pending -marker, not the drain-timed ongoing_write_through / pending_host_backups dicts. - -The deadlock root cause: the attach/defer/prune/split probes read -`_node_host_write_pending` (membership in the dicts that are drained ASYNCHRONOUSLY at -per-rank ack completion), so `logical_len != len(value)` could diverge across CP ranks -> -asymmetric `write_backup` catch-up -> asymmetric ack-append -> the `ack_queue_len==0` skip -in `writing_check` splits the collective -> NCCL hang. The fix makes those probes read -`TreeNode.cp_backup_pending`, which is set at the rank-replicated prepared-backup attach -and cleared only at the MIN-committed commit/rollback -> rank-uniform by construction, so -the attach is symmetric and `ack_write_queue` stays symmetric. - -See docs_internal/cp_hicache_symmetric_attach_design.md. Run in the dev-cu13 container. -""" - -import sys -import types -import unittest -from types import SimpleNamespace - -import torch - -# Prefer the real sgl_kernel; fall back to a stub on CPU-only hosts. -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 - HiCachePendingBackupSplit, - HiRadixCache, - PendingHiCacheBackup, - PreparedCpHiCacheBackup, -) -from sglang.srt.mem_cache.radix_cache import TreeNode # noqa: E402 - - -def _cp_cache(): - c = HiRadixCache.__new__(HiRadixCache) - c._uses_cp_hicache = True - c.ongoing_write_through = {} - c.pending_host_backups = {} - return c - - -class TestRadixProbesReadReplicatedMarker(unittest.TestCase): - def test_split_still_pending_reads_marker_not_dicts(self): - c = _cp_cache() - node = TreeNode(id=42) - # In the drain-timed dicts but marker False -> probe must IGNORE the dicts. - c.ongoing_write_through[42] = node - c.pending_host_backups[42] = object() - self.assertFalse( - c._cp_node_split_still_pending(node), - "probe must read the replicated marker, not ongoing_write_through/pending", - ) - # Marker True, NOT in any dict -> probe must return True (reads the marker). - c.ongoing_write_through.clear() - c.pending_host_backups.clear() - node.cp_backup_pending = True - self.assertTrue(c._cp_node_split_still_pending(node)) - - def test_split_node_guard_raises_on_marker(self): - c = _cp_cache() - child = TreeNode(id=7) - child.cp_backup_pending = True - # The guard is the first statement in _split_node; it raises before touching - # the key / child structure, so a stub key is fine. - with self.assertRaises(HiCachePendingBackupSplit): - c._split_node(None, child, 1) - - def test_unprunable_ignores_lock_ref_reads_marker_and_pin(self): - c = _cp_cache() - node = TreeNode(id=9) - node.children = {} - # lock_ref is rank-local (backup lock) -> must be IGNORED now. - node.lock_ref = 5 - node.host_ref_counter = 0 - node.cp_backup_pending = False - self.assertFalse( - c._cp_subtree_has_unprunable_state(node), - "must not read lock_ref (its only cross-rank skew is the backup lock)", - ) - # host_ref_counter (replicated pin) still blocks. - node.host_ref_counter = 1 - self.assertTrue(c._cp_subtree_has_unprunable_state(node)) - # The replicated marker blocks. - node.host_ref_counter = 0 - node.cp_backup_pending = True - self.assertTrue(c._cp_subtree_has_unprunable_state(node)) - - def test_inc_hit_count_short_circuits_on_marker_not_dicts(self): - c = _cp_cache() - c.cache_controller = SimpleNamespace(write_policy="write_through") - c.write_through_threshold = 1 - reached_backuped_check = {"hit": False} - c._node_backuped = lambda n: reached_backuped_check.__setitem__("hit", True) or True - - node = TreeNode(id=3) - node.hit_count = 0 - # Marker True -> _inc_hit_count returns right after the marker check, before - # the _node_backuped/threshold path. - node.cp_backup_pending = True - c._inc_hit_count(node) - self.assertFalse( - reached_backuped_check["hit"], "marker True must short-circuit _inc_hit_count" - ) - # Marker False but node in the drain-timed dicts -> must NOT short-circuit - # (proves it reads the marker, not the dicts). - node.cp_backup_pending = False - c.ongoing_write_through[node.id] = node - c.pending_host_backups[node.id] = object() - c._inc_hit_count(node) - self.assertTrue( - reached_backuped_check["hit"], "dicts must be ignored; marker False -> proceed" - ) - - -class TestMarkerLifecycle(unittest.TestCase): - def test_set_at_attach_cleared_at_commit(self): - c = _cp_cache() - c.inc_node_lock_ref = lambda n: None - c.dec_node_lock_ref = lambda n: None - node = TreeNode(id=None) - node.value = torch.arange(4) - meta = SimpleNamespace(owned_positions=torch.arange(4)) - prepared = PreparedCpHiCacheBackup( - node_id=node.id, reservation=SimpleNamespace(), metadata=meta, logical_len=4 - ) - self.assertFalse(node.cp_backup_pending) - c._attach_prepared_cp_backup(node, prepared) - self.assertTrue(node.cp_backup_pending, "attach must set the marker") - - c.pending_host_backups[node.id] = PendingHiCacheBackup( - node=node, metadata=meta, logical_len=4, locked=True - ) - c._commit_pending_backup(node.id) - self.assertFalse(node.cp_backup_pending, "commit must clear the marker") - self.assertEqual(node.host_len, 4) - - def test_double_attach_fails_loud(self): - c = _cp_cache() - c.inc_node_lock_ref = lambda n: None - node = TreeNode(id=5) - node.value = torch.arange(2) - node.cp_backup_pending = True # already pending - prepared = PreparedCpHiCacheBackup( - node_id=5, reservation=SimpleNamespace(), metadata=SimpleNamespace(), logical_len=2 - ) - with self.assertRaises(RuntimeError): - c._attach_prepared_cp_backup(node, prepared) - - def test_cleared_on_rollback(self): - c = _cp_cache() - c.dec_node_lock_ref = lambda n: None - c.cache_controller = SimpleNamespace( - cancel_layer_write_state=lambda nid: None, - evict_cp_host=lambda m: None, - ack_write_queue=[], - ) - node = TreeNode(id=11) - meta = SimpleNamespace() - node.cp_hicache = meta - node.cp_backup_pending = True - c.pending_host_backups[node.id] = PendingHiCacheBackup( - node=node, metadata=meta, logical_len=3, locked=True - ) - c._rollback_pending_backup(node.id) - self.assertFalse(node.cp_backup_pending, "rollback must clear the marker") - - -if __name__ == "__main__": - unittest.main()