From 191172fa742f8837423ca4e4be4919ebb52f1b51 Mon Sep 17 00:00:00 2001 From: Shuwen Wang <47200617+alphabetc1@users.noreply.github.com> Date: Fri, 18 Sep 2026 23:29:29 +0800 Subject: [PATCH] [Unified Tree] fix: exempt host-locked aux nodes from the sanity_check host-LRU check (#39980) --- .../unified_cache/unified_tree_core.py | 10 +++-- .../src/tests/unified_tree_core.rs | 38 ++++++------------- .../src/unified_tree_core.rs | 19 ++++++---- 3 files changed, 30 insertions(+), 37 deletions(-) diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index b105e4f57..0250b07ab 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py @@ -2863,16 +2863,20 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): f"{ct} device LRU: " f"+tree={tree_ids - lru_ids}, +lru={lru_ids - tree_ids}" ) - # Aux host-only states must match the host LRU. + # Aux host-only states must match the host LRU. A host lock + # delists its node, so locked nodes are exempt on both sides. host_lru = self.host_lru_lists[ct] + host_locked_ids = { + n.id for n in all_nodes if n.component_data[ct].host_lock_ref > 0 + } s3_ids = { n.id for n in all_nodes if n is not self.root_node and n.component_data[ct].value is None and n.component_data[ct].host_value is not None - } - host_lru_ids = set(host_lru.cache.keys()) + } - host_locked_ids + host_lru_ids = set(host_lru.cache.keys()) - host_locked_ids if s3_ids != host_lru_ids: E( f"{ct} host LRU: " diff --git a/rust/sglang-radix-tree/src/tests/unified_tree_core.rs b/rust/sglang-radix-tree/src/tests/unified_tree_core.rs index 5e81b9d0a..8bd7700ff 100644 --- a/rust/sglang-radix-tree/src/tests/unified_tree_core.rs +++ b/rust/sglang-radix-tree/src/tests/unified_tree_core.rs @@ -7820,33 +7820,6 @@ fn sanity_check_reports_a_cyclic_child_map_without_hanging() { tc.sanity_check(&[], &[]); } -#[test] -#[should_panic(expected = "host LRU mismatch")] -fn sanity_check_detects_a_host_locked_value_missing_from_the_lru() { - let mut tc = sane_tree(); - let leaf = tc - .match_prefix(&match_params(&vec![1, 2, 9])) - .best_match_node_id; - let parent = tc - .arena - .node(tc.arena.resolve(leaf).expect("live test node")) - .parent(); - tc.register_component_(Arc::new(SwaComponentForTest)); - // The arena was built Full-only; give the root the stub's lock too. - tc.arena.node_mut(tc.arena.root()).values[SWA.idx()].lock_ref = 1; - tc.arena - .node_mut(parent) - .state_mut_(ValueSlotIdx::host(FULL)) - .value = Some(Tensor::from_slice(&[10i64, 11])); - let leaf_node = tc - .arena - .node_mut(tc.arena.resolve(leaf).expect("live test node")); - leaf_node.state_mut_(ValueSlotIdx::host(FULL)).value = Some(Tensor::from_slice(&[30i64])); - leaf_node.state_mut_(ValueSlotIdx::host(SWA)).value = Some(Tensor::from_slice(&[30i64])); - leaf_node.state_mut_(ValueSlotIdx::host(SWA)).lock_ref = 1; - tc.sanity_check(&[], &[]); -} - // A backed-up leaf whose unlocked Swa value is host-only (no device value). fn host_only_aux_leaf(tc: &mut UnifiedTreeCore>) -> NodeIdx_ { let leaf = tc @@ -7890,6 +7863,17 @@ fn sanity_check_detects_a_host_only_value_missing_from_the_lru() { tc.sanity_check(&[], &[]); } +// A host lock delists its node: missing from the LRU is the in-flight state. +#[test] +fn sanity_check_accepts_a_host_locked_value_missing_from_the_lru() { + let mut tc = sane_tree(); + let leaf = host_only_aux_leaf(&mut tc); + tc.arena + .node_mut(leaf) + .set_lock_ref_(ValueSlotIdx::host(SWA), 1); + tc.sanity_check(&[], &[]); +} + #[test] #[should_panic(expected = "EvictLayer::All is not a single layer")] fn for_each_component_lru_rejects_the_all_layer() { diff --git a/rust/sglang-radix-tree/src/unified_tree_core.rs b/rust/sglang-radix-tree/src/unified_tree_core.rs index 2dadbcc4e..742f9c190 100644 --- a/rust/sglang-radix-tree/src/unified_tree_core.rs +++ b/rust/sglang-radix-tree/src/unified_tree_core.rs @@ -4364,6 +4364,8 @@ impl UnifiedTreeCore { // states must match the host LRU; never both at once. let mut device_count = 0; let mut host_only_count = 0; + // A host lock delists its node, so locked nodes are exempt. + let mut host_locked_listed = 0; for &node_id in &all_nodes { if self.arena.node(node_id).is_root() { continue; @@ -4377,17 +4379,20 @@ impl UnifiedTreeCore { )); } let host_only = !has_device && node.has_host_value(ct); - if host_only != host_lru.in_list(Some(node_id)) { + let host_listed = host_lru.in_list(Some(node_id)); + let host_locked = node.host_lock_ref(ct) > 0; + if host_locked { + host_locked_listed += host_listed as usize; + } else if host_only != host_listed { errors.push(format!( - "{ct:?} host LRU mismatch at node {node_id}: host_only={host_only} in_lru={}", - host_lru.in_list(Some(node_id)) + "{ct:?} host LRU mismatch at node {node_id}: host_only={host_only} in_lru={host_listed}" )); } if lru.in_list(Some(node_id)) && host_lru.in_list(Some(node_id)) { errors.push(format!("{ct:?} node {node_id} in both device and host LRU")); } device_count += has_device as usize; - host_only_count += host_only as usize; + host_only_count += (host_only && !host_locked) as usize; } if device_count != lru.len() { errors.push(format!( @@ -4395,10 +4400,10 @@ impl UnifiedTreeCore { lru.len() )); } - if host_only_count != host_lru.len() { + let host_listed_count = host_lru.len().saturating_sub(host_locked_listed); + if host_only_count != host_listed_count { errors.push(format!( - "{ct:?} host LRU: tree={host_only_count} != lru={}", - host_lru.len() + "{ct:?} host LRU: tree={host_only_count} != lru={host_listed_count}" )); } // Linked-list integrity