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) <noreply@anthropic.com>
(cherry picked from commit 876a98177cd749b5a63623ba2ae8b868289b8e59)
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user