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

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

@@ -599,7 +599,11 @@ 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 {
-                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);

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