From b9cb96496dcd2f0c6c4dce72fe2672dc56aaf1e8 Mon Sep 17 00:00:00 2001 From: Shuwen Wang <47200617+alphabetc1@users.noreply.github.com> Date: Sat, 12 Sep 2026 18:17:38 +0800 Subject: [PATCH] [Unified Tree] Preserve aux LRU recency when splitting nodes (#38482) --- .../unified_cache/components/swa_component.py | 5 +++- .../unified_cache/unified_tree_core.py | 17 ++++++++----- rust/sglang-radix-tree/src/components/swa.rs | 6 ++++- .../src/tests/unified_tree_core.rs | 8 +++---- .../sglang-radix-tree/src/unified_lru_list.rs | 5 ++++ .../src/unified_tree_core.rs | 17 ++----------- .../test_unified_radix_cache_unittest.py | 24 +++++++++++++++++++ 7 files changed, 55 insertions(+), 27 deletions(-) diff --git a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py index 6878bc03f..1c2c95156 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py @@ -615,7 +615,10 @@ class SWAComponent(TreeComponent): new_parent.component_data[self.component_type].value is None and parent_swa_data.host_lock_ref == 0 ): - host_lru.insert_mru(new_parent) + if host_lru.in_list(child): + host_lru.insert_after(child, new_parent) + else: + host_lru.insert_mru(new_parent) if ( child.component_data[self.component_type].value is None and child_swa_data.host_lock_ref == 0 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 84e77362e..8e0555f5f 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 @@ -239,6 +239,12 @@ class UnifiedLRUList: self.cache[node.id] = node self._add_node(node) + def insert_after(self, prev_node: UnifiedTreeNode, node: UnifiedTreeNode): + assert prev_node.id in self.cache + assert node.id not in self.cache + self.cache[node.id] = node + self._add_node_after(prev_node, node) + def remove_node(self, node: UnifiedTreeNode): assert node.id in self.cache del self.cache[node.id] @@ -1307,8 +1313,6 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): # owner (b + P) % N on both sides of the split). new_node.rotation_base = child.rotation_base - self._for_each_component_lru(child, UnifiedLRUList.remove_node) - child.parent = new_node child.key = child.key[split_len:] new_node.hash_value, child.hash_value = split_node_hash_value( @@ -1334,11 +1338,12 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): new_child_node_id=child.id, ) + # Splitting does not access the suffix; retain its recency and place + # the inherited prefix beside it, in the same session partition. self._for_each_component_lru( - new_node, UnifiedLRUList.insert_mru, skip_existing=True - ) - self._for_each_component_lru( - child, UnifiedLRUList.insert_mru, skip_existing=True + new_node, + lambda lru, node: lru.insert_after(child, node), + skip_existing=True, ) child.last_access_time = get_and_increase_time_counter() diff --git a/rust/sglang-radix-tree/src/components/swa.rs b/rust/sglang-radix-tree/src/components/swa.rs index 262853daf..84525821a 100644 --- a/rust/sglang-radix-tree/src/components/swa.rs +++ b/rust/sglang-radix-tree/src/components/swa.rs @@ -599,7 +599,11 @@ impl TreeComponent for SwaComponent { let child_parks = !child.has_device_value(SWA) && child.host_lock_ref(SWA) == 0; let host_lru = tree_core.host_lru_list_mut(SWA); if parent_parks { - host_lru.insert_mru(new_parent_id); + if host_lru.in_list(Some(child_id)) { + host_lru.insert_after(child_id, new_parent_id); + } else { + host_lru.insert_mru(new_parent_id); + } } if child_parks && !host_lru.in_list(Some(child_id)) { host_lru.insert_mru(child_id); 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 7ebc290a4..285eeec13 100644 --- a/rust/sglang-radix-tree/src/tests/unified_tree_core.rs +++ b/rust/sglang-radix-tree/src/tests/unified_tree_core.rs @@ -990,13 +990,13 @@ fn split_updates_the_leaf_sets() { } #[test] -fn split_readmits_aux_lru_cells() { +fn split_preserves_aux_lru_position() { let mut tc = core(); tc.register_component_(Arc::new(SwaComponentForTest)); let c = split_setup(&mut tc); tc.arena.node_mut(c).values[SWA.idx()].value = Some(Tensor::from_slice(&[0i64])); tc.device_lru_list_mut(SWA).insert_mru(c); - // A second listed node makes the child's detach-and-readmit observable. + // A newer node must stay ahead of the unmatched suffix after the split. let root = tc.arena.root(); let s = tc .arena @@ -1009,10 +1009,10 @@ fn split_readmits_aux_lru_cells() { .unwrap(); tc.device_lru_list_mut(SWA).insert_mru(s); let (new_node, _) = tc.split_node_(c, /* split_len = */ 2); - // The child re-enters the SWA LRU at MRU; the value-less prefix node does not. + // The child stays cold; the value-less prefix node does not enter the LRU. assert!(tc.device_lru_list(SWA).in_list(Some(c))); assert!(!tc.device_lru_list(SWA).in_list(Some(new_node))); - assert_eq!(tc.device_lru_list(SWA).get_lru_where(|_| true), Some(s)); + assert_eq!(tc.device_lru_list(SWA).get_lru_where(|_| true), Some(c)); } #[test] diff --git a/rust/sglang-radix-tree/src/unified_lru_list.rs b/rust/sglang-radix-tree/src/unified_lru_list.rs index 064937ad1..454a42661 100644 --- a/rust/sglang-radix-tree/src/unified_lru_list.rs +++ b/rust/sglang-radix-tree/src/unified_lru_list.rs @@ -149,6 +149,11 @@ impl UnifiedLRUList { self.add_node_(Self::cell_of_(node_id)); } + /// Insert beside an existing member without refreshing its recency. + pub fn insert_after(&mut self, prev_node_id: NodeIdx_, node_id: NodeIdx_) { + self.add_node_after_(Self::cell_of_(prev_node_id), Self::cell_of_(node_id)); + } + /// Remove a member node, resetting its cell; panics if not a member. pub fn remove_node(&mut self, node_id: NodeIdx_) { self.remove_node_(Self::cell_of_(node_id)); diff --git a/rust/sglang-radix-tree/src/unified_tree_core.rs b/rust/sglang-radix-tree/src/unified_tree_core.rs index c19713914..a08dd007b 100644 --- a/rust/sglang-radix-tree/src/unified_tree_core.rs +++ b/rust/sglang-radix-tree/src/unified_tree_core.rs @@ -1790,14 +1790,6 @@ impl UnifiedTreeCore { ); self.arena.node_mut(new_node_id).external_cache_stored = child_external_cache_stored; - // The child's aux LRU cells detach while it is re-linked. - self.for_each_component_lru_( - child_id, - &mut |lru, node_id| lru.remove_node(node_id), - EvictLayer::Device, - /* skip_existing = */ false, - ); - let child = self.arena.node_mut(child_id); child.parent = Some(new_node_id); child.key = key_tail; @@ -1846,15 +1838,10 @@ impl UnifiedTreeCore { None }; + // A split does not access the suffix; keep both fragments at its old position. self.for_each_component_lru_( new_node_id, - &mut |lru, node_id| lru.insert_mru(node_id), - EvictLayer::Device, - /* skip_existing = */ true, - ); - self.for_each_component_lru_( - child_id, - &mut |lru, node_id| lru.insert_mru(node_id), + &mut |lru, node_id| lru.insert_after(child_id, node_id), EvictLayer::Device, /* skip_existing = */ true, ); diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 61a597f05..cfa4be4a3 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -2518,6 +2518,30 @@ class UnifiedRadixCacheSuite: ) cache.sanity_check() + def test_partial_match_keeps_unmatched_suffix_lru_position(self): + if not self.cfg.has_swa and not self.cfg.has_mamba: + self.skipTest("requires an aux LRU") + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + cold_tokens = self._make_seq(1, 2) + hot_tokens = self._make_seq(100, 2) + cold = self._insert(cache, allocator, req_to_token_pool, cold_tokens) + hot = self._insert(cache, allocator, req_to_token_pool, hot_tokens) + + cache.match_prefix( + MatchPrefixParams( + key=RadixKey(array("q", cold_tokens[: self.cfg.page_size])) + ) + ) + + for ct in self.cfg.components: + if ct == ComponentType.FULL: + continue + order = cache.tree_core.get_component_device_lru_node_ids(ct) + self.assertLess( + order.index(hot.last_device_node), order.index(cold.last_device_node) + ) + cache.sanity_check() + def test_swa_lru_match_only_refreshes_window_cushion(self): if not self._swa_pinning_cfg_supported(): self.skipTest("requires SWA-only config with node size >= cushion")