minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Shuwen Wang
2026-09-18 15:29:29 +00:00
committed by GitHub
parent 45938a24ae
commit 191172fa74
3 changed files with 30 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/sglang-radix-tree/src/tests/unified_tree_core.rs
+11 -27
View File
@@ -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
+12 -7
View File
@@ -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={}",
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<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={}",
host_lru.len()
"{ct:?} host LRU: tree={host_only_count} != lru={host_listed_count}"
));
}
// Linked-list integrity