From 85585fcf4b62ed41b1f66cfb15a9352d841ed6c5 Mon Sep 17 00:00:00 2001 From: leavelet Date: Mon, 22 Jun 2026 19:42:41 +0000 Subject: [PATCH] Fix CP HiCache writing_check deadlock at the source: rank-replicated symmetric attach Replaces the interim ack_queue_len==0 band-aid with a structural fix so the perf skip is SAFE rather than removed. Root cause: the radix attach/defer/prune/split probes read `_node_host_write_pending` (membership in ongoing_write_through/pending_host_backups, which are drained ASYNCHRONOUSLY at per-rank ack completion), and `_cp_subtree_has_unprunable_state` also read the rank-local `lock_ref`. So per-rank ack-drain timing leaked into the match_prefix truncation -> cache_protected_len -> logical_len/len(value) -> the attach predicate. On some ranks the prepared backup attached, on others it was dropped to the synchronous write_backup catch-up (catch_up_all_layers draft=True), whose ack-append is per-rank. ack_write_queue then diverged across CP ranks, so writing_check's ack_queue_len==0 early-return diverged: a strict subset entered the writing_check_min all_reduce while the rest advanced to recv_requests' broadcast on the same tp_cpu_group (8 ranks; dp_size==1 resets enable_dp_attention to False) -> NCCL deadlock (prod hang: node_id=50713, all 8 GPUs 0% util). Fix: introduce TreeNode.cp_backup_pending, a rank-replicated marker set at the rank-replicated prepared-backup attach (_attach_prepared_cp_backup) and the guarded write_backup registration, cleared only at the MIN-committed commit (_commit_pending_backup) / rollback. Rewire the CP radix probes (_cp_node_split_still_pending, _split_node guard, _cp_subtree_has_unprunable_state, _inc_hit_count) to read this marker instead of the drain-timed dicts, and DROP the lock_ref read in the prune probe (its only cross-rank skew is the backup lock, now covered by the marker). _node_host_write_pending stays the drain-timed accessor used only by writing_check / eviction / the visibility check (_node_host_write_ready). With the radix-path reads replicated by construction, len(value)==logical_len on every rank -> the prepared backup always attaches (or rolls back) symmetrically -> the asymmetric catch-up never fires -> ack_write_queue is symmetric -> the existing ack_queue_len==0 skip is all-or-none (no per-tick no-op MIN, best perf). Correctness: the marker is cleared only at/after _commit_pending_backup, which runs only under the MIN frontier, so the MIN remains the SOLE host-visibility gate (no KV-corruption class reopened); the marker is strictly more conservative than the old dicts (stays True until commit). Fail-fasts: double-attach raises; _split_node propagates the marker. hi_mamba unchanged (gates on ongoing only, no async ack skip). Single-layer-draft (c3fc3ff752) preserved: the design touches only the attach decision; the draft notifier already fires post-MoE. Adds test_cp_hicache_symmetric_attach.py: the probes read the replicated marker not the dicts, the prune probe ignores lock_ref, and the marker lifecycle (attach->commit->rollback) + the double-attach guard. All pass on the fix; all fail on the pre-fix code. Design: docs_internal/cp_hicache_symmetric_attach_design.md. Pending: ETE (GSM8K + no-hang flood + TTFT) on the dev-cu13 container. Co-Authored-By: Claude Opus 4.8 (1M context) (cherry picked from commit 876a98177cd749b5a63623ba2ae8b868289b8e59) --- 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, 238 insertions(+), 11 deletions(-) create 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 ea7f88413..d621d6cd6 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -2792,6 +2792,19 @@ 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) @@ -2818,6 +2831,10 @@ 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 @@ -2879,6 +2896,7 @@ 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 @@ -3167,7 +3185,13 @@ 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, @@ -3474,10 +3498,15 @@ 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 ( - current.lock_ref > 0 + self._node_backup_outstanding_replicated(current) or current.host_ref_counter > 0 - or self._node_host_write_pending(current) ): return True stack.extend(current.children.values()) @@ -3501,17 +3530,19 @@ 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): + # 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: 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 - ): + # _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: 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( @@ -3863,6 +3894,7 @@ 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( @@ -3953,7 +3985,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): @@ -5337,7 +5369,7 @@ 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) # child node split into new_node -> child @@ -5352,6 +5384,10 @@ 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 df824ac2c..5fd3d64fd 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 + # 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 new file mode 100644 index 000000000..f7a1770a8 --- /dev/null +++ b/test/registered/unit/mem_cache/test_cp_hicache_symmetric_attach.py @@ -0,0 +1,185 @@ +"""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()