[Unified Tree] fix: exempt host-locked aux nodes from the sanity_check host-LRU check (#39980)

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_locked_ids
-                host_lru_ids = set(host_lru.cache.keys())
+                host_lru_ids = set(host_lru.cache.keys()) - host_locked_ids
                 if s3_ids != host_lru_ids:
                     E(
                         f"{ct} host LRU: "

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<Vec<i64>>) -> 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() {

rust/sglang-radix-tree/src/unified_tree_core.rs

@@ -4364,6 +4364,8 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
                 // 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<K: ChildKeyType> UnifiedTreeCore<K> {
                         ));
                     }
                     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={}",
+                            "{ct:?} host LRU mismatch at node {node_id}: host_only={host_only} in_lru={host_listed}"
-                            host_lru.in_list(Some(node_id))
                         ));
                     }
                     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<K: ChildKeyType> UnifiedTreeCore<K> {
                         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={}",
+                        "{ct:?} host LRU: tree={host_only_count} != lru={host_listed_count}"
-                        host_lru.len()
                     ));
                 }
                 // Linked-list integrity