[Unified Tree] Preserve aux LRU recency when splitting nodes (#38482)

python/sglang/srt/mem_cache/unified_cache/components/swa_component.py

@@ -615,6 +615,9 @@ class SWAComponent(TreeComponent):
                 new_parent.component_data[self.component_type].value is None
                 and parent_swa_data.host_lock_ref == 0
             ):
+                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

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
+            new_node,
-        )
+            lambda lru, node: lru.insert_after(child, node),
-        self._for_each_component_lru(
+            skip_existing=True,
-            child, UnifiedLRUList.insert_mru, skip_existing=True
         )
         child.last_access_time = get_and_increase_time_counter()
 

rust/sglang-radix-tree/src/components/swa.rs

@@ -599,8 +599,12 @@ impl<K: ChildKeyType> TreeComponent<K> 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 {
+                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);
             }

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]

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));

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

@@ -1790,14 +1790,6 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
         );
         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<K: ChildKeyType> UnifiedTreeCore<K> {
             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),
+            &mut |lru, node_id| lru.insert_after(child_id, node_id),
-            EvictLayer::Device,
-            /* skip_existing = */ true,
-        );
-        self.for_each_component_lru_(
-            child_id,
-            &mut |lru, node_id| lru.insert_mru(node_id),
             EvictLayer::Device,
             /* skip_existing = */ true,
         );

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")