chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Revert "Fix CP HiCache writing_check deadlock at the source: rank-replicated symmetric attach"

Browse Source 
This reverts commit 85585fcf4b.
This commit is contained in:
孙远航
2026-06-22 21:35:27 +00:00
parent d08ee9f8d9
commit f353cf4702
3 changed files with 11 additions and 238 deletions
Download Patch File Download Diff File Expand all files Collapse all files

58
python/sglang/srt/mem_cache/hiradix_cache.py
View File

@@ -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:

6
python/sglang/srt/mem_cache/radix_cache.py
View File

@@ -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

185
test/registered/unit/mem_cache/test_cp_hicache_symmetric_attach.py
View File

@@ -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()