fix: align write-through pending across tree cores (#37278)

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

@@ -1820,9 +1820,14 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
         })
     }
 
-    /// Mark a node as having an in-flight write-through backup.
-    fn mark_write_through_pending(&self, py: Python<'_>, node_id: NodeId) {
-        py.allow_threads(|| self.core().mark_write_through_pending(node_id));
+    /// Mark the nodes one write-through backup covers; returns them ancestors first.
+    fn mark_write_through_pending(
+        &self,
+        py: Python<'_>,
+        node_ids: Vec<NodeId>,
+        ack_id: NodeId,
+    ) -> Vec<NodeId> {
+        py.allow_threads(|| self.core().mark_write_through_pending(node_ids, ack_id))
     }
 
     /// Clear the write-through-pending mark on the acked nodes.
@@ -2799,9 +2804,14 @@ macro_rules! tree_core_binding {
                 self.inner.drop_subtree_no_host(py, node_id)
             }
 
-            /// Mark a node as having an in-flight write-through backup.
-            fn mark_write_through_pending(&self, py: Python<'_>, node_id: NodeId) {
-                self.inner.mark_write_through_pending(py, node_id)
+            /// Mark the nodes one write-through backup covers; returns them ancestors first.
+            fn mark_write_through_pending(
+                &self,
+                py: Python<'_>,
+                node_ids: Vec<NodeId>,
+                ack_id: NodeId,
+            ) -> Vec<NodeId> {
+                self.inner.mark_write_through_pending(py, node_ids, ack_id)
             }
 
             /// Clear the write-through-pending mark on the acked nodes.

rust/sglang-radix-tree/src/tests/components/full.rs

@@ -425,7 +425,7 @@ fn host_drive_spares_coexisting_host_values_under_an_in_flight_transfer() {
         .match_prefix(&match_params(&vec![1, 2]))
         .best_match_node_id;
     tc.commit_backup(handle, Tensor::from_slice(&[20i64, 21]), HashMap::new());
-    tc.mark_write_through_pending(handle);
+    tc.mark_write_through_pending(vec![handle], /* ack_id = */ handle);
 
     let (mut tr, mut df, mut hf) = (tracker(), frees(), frees());
     accumulate_step(

rust/sglang-radix-tree/src/tests/components/mamba.rs

@@ -548,7 +548,7 @@ fn reinsert_full_backed_target_schedules_mamba_only_backup() {
             .equal(&Tensor::from_slice(&[7i64]))
     );
 
-    tc.mark_write_through_pending(leaf);
+    tc.mark_write_through_pending(vec![leaf], /* ack_id = */ leaf);
     let pending = tc.insert(&insert_params_mamba(&key, &[30, 31], Some(9)));
     assert!(
         !pending

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

@@ -2132,11 +2132,11 @@ fn insert_threshold_crossing_emits_the_backup_kv_action() {
 }
 
 #[test]
-fn mark_write_through_pending_stamps_the_node_id_as_the_ack() {
+fn mark_write_through_pending_stamps_the_supplied_ack() {
     let mut tc = core();
     tc.insert(&insert_params(&vec![1], &[10]));
     let leaf = tc.match_prefix(&match_params(&vec![1])).best_match_node_id;
-    tc.mark_write_through_pending(leaf);
+    tc.mark_write_through_pending(vec![leaf], /* ack_id = */ leaf);
     assert_eq!(
         tc.arena
             .node(tc.arena.resolve(leaf))
@@ -2145,12 +2145,60 @@ fn mark_write_through_pending_stamps_the_node_id_as_the_ack() {
     );
 }
 
+#[test]
+fn mark_write_through_pending_stamps_one_ack_on_every_published_node() {
+    let mut tc = core();
+    tc.insert(&insert_params(&vec![1], &[10]));
+    tc.insert(&insert_params(&vec![1, 2], &[10, 11]));
+    let parent = tc.match_prefix(&match_params(&vec![1])).best_match_node_id;
+    let leaf = tc
+        .match_prefix(&match_params(&vec![1, 2]))
+        .best_match_node_id;
+
+    let published = tc.mark_write_through_pending(vec![parent, leaf], /* ack_id = */ leaf);
+
+    assert_eq!(published, vec![parent, leaf]);
+    for node_id in [parent, leaf] {
+        assert_eq!(
+            tc.arena
+                .node(tc.arena.resolve(node_id))
+                .write_through_pending_id,
+            Some(leaf)
+        );
+    }
+    tc.finish_write_through(vec![parent, leaf], /* ack_id = */ leaf);
+    for node_id in [parent, leaf] {
+        assert_eq!(
+            tc.arena
+                .node(tc.arena.resolve(node_id))
+                .write_through_pending_id,
+            None
+        );
+    }
+}
+
+#[test]
+fn mark_write_through_pending_returns_the_published_nodes_ancestors_first() {
+    let mut tc = core();
+    tc.insert(&insert_params(&vec![1], &[10]));
+    tc.insert(&insert_params(&vec![1, 2], &[10, 11]));
+    let parent = tc.match_prefix(&match_params(&vec![1])).best_match_node_id;
+    let leaf = tc
+        .match_prefix(&match_params(&vec![1, 2]))
+        .best_match_node_id;
+
+    // The caller merges per-component transfers, whose order is not tree order.
+    let published = tc.mark_write_through_pending(vec![leaf, parent], /* ack_id = */ leaf);
+
+    assert_eq!(published, vec![parent, leaf]);
+}
+
 #[test]
 fn finish_write_through_clears_only_the_matching_ack() {
     let mut tc = core();
     tc.insert(&insert_params(&vec![1], &[10]));
     let leaf = tc.match_prefix(&match_params(&vec![1])).best_match_node_id;
-    tc.mark_write_through_pending(leaf);
+    tc.mark_write_through_pending(vec![leaf], /* ack_id = */ leaf);
     tc.finish_write_through(vec![leaf], /* ack_id = */ 999_999);
     assert_eq!(
         tc.arena
@@ -2202,7 +2250,7 @@ fn split_of_a_pending_node_transfers_the_ack_and_emits_the_replace_action() {
     let node = tc
         .match_prefix(&match_params(&vec![1, 2, 3]))
         .best_match_node_id;
-    tc.mark_write_through_pending(node);
+    tc.mark_write_through_pending(vec![node], /* ack_id = */ node);
     let (new_node, action) = tc.split_node_(tc.arena.resolve(node), /* split_len = */ 1);
     assert_eq!(tc.arena.node(new_node).write_through_pending_id, Some(node));
     assert_eq!(
@@ -2806,7 +2854,7 @@ fn finish_write_through_after_a_split_publishes_both_fragments() {
     let leaf = tc
         .match_prefix(&match_params(&vec![1, 2, 3, 4]))
         .best_match_node_id;
-    tc.mark_write_through_pending(leaf);
+    tc.mark_write_through_pending(vec![leaf], /* ack_id = */ leaf);
     let _ = tc.take_events();
     let result = tc.insert(&insert_params(&vec![1, 2, 5, 6], &[20, 21, 22, 23]));
     let new_node_id = result
@@ -3179,7 +3227,7 @@ fn insert_host_drop_preserves_split_actions_and_lengths() {
     let leaf = tc
         .match_prefix(&match_params(&vec![1, 2, 3]))
         .best_match_node_id;
-    tc.mark_write_through_pending(leaf);
+    tc.mark_write_through_pending(vec![leaf], /* ack_id = */ leaf);
     let root = tc.arena.root();
     let result = tc.insert_host(
         tc.arena.node(root).id,

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

@@ -3668,10 +3668,39 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
         }
     }
 
-    /// Mark a node as having an in-flight write-through backup.
-    pub fn mark_write_through_pending(&mut self, node_id: NodeId) {
-        let node_idx = self.arena.resolve(node_id);
-        self.arena.node_mut(node_idx).write_through_pending_id = Some(node_id);
+    /// Mark every node covered by one in-flight write-through backup, and return
+    /// them ancestors first: publish links each host store event to its parent.
+    pub fn mark_write_through_pending(
+        &mut self,
+        node_ids: Vec<NodeId>,
+        ack_id: NodeId,
+    ) -> Vec<NodeId> {
+        let mut marked: Vec<(usize, NodeId)> = Vec::with_capacity(node_ids.len());
+        for node_id in node_ids {
+            let node_idx = self.arena.resolve(node_id);
+            let depth = self.depth_from_root_(node_idx);
+            let node = self.arena.node_mut(node_idx);
+            assert!(
+                node.write_through_pending_id.is_none()
+                    || node.write_through_pending_id == Some(ack_id),
+                "node {} is already pending under a different write-through ack",
+                node.id
+            );
+            node.write_through_pending_id = Some(ack_id);
+            marked.push((depth, node_id));
+        }
+        marked.sort_unstable();
+        marked.into_iter().map(|(_, node_id)| node_id).collect()
+    }
+
+    fn depth_from_root_(&self, node_idx: NodeIdx_) -> usize {
+        let mut depth = 0;
+        let mut node = self.arena.node(node_idx);
+        while !node.is_root() {
+            depth += 1;
+            node = self.arena.node(node.parent());
+        }
+        depth
     }
 
     /// Clear the write-through-pending mark (when it matches ack_id) and record the