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

python/sglang/srt/mem_cache/rust_tree_core/adapter.py

@@ -928,8 +928,10 @@ class RustUnifiedTreeCore(UnifiedTreeCoreInterface):
         result = DropSubtreeNoHostResult(is_dropped=binding_result.dropped)
         return _fill_evict_result(binding_result, result)
 
-    def mark_write_through_pending(self, node_id: NodeId) -> None:
-        self._binding.mark_write_through_pending(node_id)
+    def mark_write_through_pending(
+        self, node_ids: list[NodeId], ack_id: NodeId
+    ) -> list[NodeId]:
+        return self._binding.mark_write_through_pending(list(node_ids), ack_id)
 
     def finish_write_through(self, node_ids: list[NodeId], ack_id: int) -> None:
         self._binding.finish_write_through(list(node_ids), ack_id)

python/sglang/srt/mem_cache/unified_cache/unified_cache_linker.py

@@ -485,7 +485,7 @@ class UnifiedCacheLinkerWrapper:
             cache.dec_lock_ref(node_id, lock_params)
             return
 
-        cache.tree_core.mark_write_through_pending(node_id)
+        cache.tree_core.mark_write_through_pending([node_id], ack_id=node_id)
         node.external_cache_stored = True
         self.pending_offloads.append(_PendingOffload(node_id, lock_params, [node_id]))
 

python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py

@@ -2203,10 +2203,29 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
             self._update_duplicate_tracking(node)
             node = node.parent
 
-    def mark_write_through_pending(self, node_id: NodeId) -> None:
-        """Mark a node as having an in-flight write-through backup."""
-        node = self.node_by_id(node_id)
-        node.write_through_pending_id = node_id
+    def mark_write_through_pending(
+        self, node_ids: list[NodeId], ack_id: NodeId
+    ) -> list[NodeId]:
+        """Stamp ack_id on every covered node; returns them ancestors first."""
+        marked: list[tuple[int, NodeId]] = []
+        for node_id in node_ids:
+            node = self.node_by_id(node_id)
+            assert node.write_through_pending_id in (
+                None,
+                ack_id,
+            ), f"node {node.id} is already pending under a different write-through ack"
+            node.write_through_pending_id = ack_id
+            marked.append((self._depth_from_root(node), node_id))
+        marked.sort()
+        return [node_id for _, node_id in marked]
+
+    @staticmethod
+    def _depth_from_root(node: UnifiedTreeNode) -> int:
+        depth = 0
+        while node.parent is not None:
+            depth += 1
+            node = node.parent
+        return depth
 
     def finish_write_through(self, node_ids: list[NodeId], ack_id: int) -> None:
         """Clear the write-through-pending mark (when it matches ack_id) and record the

python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py

@@ -537,8 +537,11 @@ class UnifiedTreeCoreInterface(ABC):
     write_back_duplicate_reclaim_digest: int = 0
 
     @abstractmethod
-    def mark_write_through_pending(self, node_id: NodeId) -> None:
-        """Mark a node as having an in-flight write-through backup."""
+    def mark_write_through_pending(
+        self, node_ids: list[NodeId], ack_id: NodeId
+    ) -> list[NodeId]:
+        """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."""
         ...
 
     @abstractmethod

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -1370,10 +1370,26 @@ class UnifiedRadixCache(BasePrefixCache):
             lock_params = None
             if not write_back:
                 lock_params = self.inc_lock_ref(node_id).to_dec_params()
-            self._track_write_through_node(node_id, lock_params)
+            publish_node_ids = self._backup_publish_node_ids(node_id, comp_xfers)
+            self._track_write_through_node(
+                node_id, lock_params, publish_node_ids=publish_node_ids
+            )
             written = len(host_indices)
         return written
 
+    @staticmethod
+    def _backup_publish_node_ids(
+        node_id: NodeId, comp_xfers: dict[ComponentType, list[PoolTransfer]]
+    ) -> list[NodeId]:
+        """The acked node plus every node a component backup transfer covers."""
+        publish_node_ids: list[NodeId] = []
+        for transfers in comp_xfers.values():
+            for transfer in transfers:
+                publish_node_ids.extend(transfer.nodes_to_load or ())
+        if node_id not in publish_node_ids:
+            publish_node_ids.append(node_id)
+        return list(dict.fromkeys(publish_node_ids))
+
     def _build_backup_sidecar(self, device_value, comp_xfers):
         """Gather sidecar transfer spec."""
         kv_xfer = PoolTransfer(name=PoolName.KV, device_indices=device_value)
@@ -1399,10 +1415,13 @@ class UnifiedRadixCache(BasePrefixCache):
         self,
         node_id: NodeId,
         lock_params: Optional[DecLockRefParams],
+        publish_node_ids: list[NodeId],
     ) -> None:
-        self.tree_core.mark_write_through_pending(node_id)
+        publish_node_ids = self.tree_core.mark_write_through_pending(
+            publish_node_ids, ack_id=node_id
+        )
         self.ongoing_write_through[node_id] = _OngoingWriteThrough(
-            node_id, lock_params, [node_id]
+            node_id, lock_params, publish_node_ids
         )
 
     def _replace_pending_write_through_node(

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

test/registered/unit/mem_cache/test_rust_tree_core_integration.py

@@ -593,7 +593,7 @@ def test_hicache_write_through_and_load_back_round_trip():
     device_value, comp_xfers = core.build_backup_spec(leaf)
     assert device_value.tolist() == [10, 11]
     assert comp_xfers == {}
-    core.mark_write_through_pending(leaf)
+    core.mark_write_through_pending([leaf], ack_id=leaf)
     core.commit_backup(leaf, torch.tensor([100, 101], dtype=torch.int64), comp_xfers)
     core.finish_write_through([leaf], leaf)
     tracker = {ComponentType.FULL: 0}
@@ -617,6 +617,30 @@ def test_hicache_write_through_and_load_back_round_trip():
     core.sanity_check([], [])
 
 
+def test_cache_tracks_one_write_through_ack_across_rust_nodes():
+    from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
+
+    core = _tree_core()
+    _insert(core, [1], [10])
+    _insert(core, [1, 2], [10, 11])
+    parent = core.match_prefix(MatchPrefixParams(key=_key([1]))).best_match_node
+    leaf = core.match_prefix(MatchPrefixParams(key=_key([1, 2]))).best_match_node
+    cache = SimpleNamespace(tree_core=core, ongoing_write_through={})
+
+    # Child-first in, ancestors-first out: the publish side links every store
+    # event to its parent, and component transfer order is not tree order.
+    UnifiedRadixCache._track_write_through_node(
+        cache,
+        leaf,
+        lock_params=None,
+        publish_node_ids=[leaf, parent],
+    )
+
+    assert cache.ongoing_write_through[leaf].publish_node_ids == [parent, leaf]
+    core.finish_write_through([parent, leaf], ack_id=leaf)
+    core.sanity_check([], [])
+
+
 def test_invalid_demote_states_raise_assertion_error():
     core = _tree_core()
     core.set_hicache_enabled()
@@ -1588,7 +1612,7 @@ def test_split_of_a_write_through_pending_node_crosses_the_replace_action():
     core.set_hicache_enabled()
     _insert(core, [1, 2, 3, 4], [10, 11, 12, 13])
     leaf = core.match_prefix(MatchPrefixParams(key=_key([1, 2, 3, 4]))).best_match_node
-    core.mark_write_through_pending(leaf)
+    core.mark_write_through_pending([leaf], ack_id=leaf)
     # A divergent prefix splits the pending node; the publish list must follow.
     result = _insert(core, [1, 2], [10, 11])
     (replace,) = [

test/registered/unit/mem_cache/test_unified_cache_linker.py

@@ -150,8 +150,8 @@ def test_async_offload_pins_node_until_completion():
     cache = _cache_for_wrapper(
         tree_core=SimpleNamespace(
             enable_external_cache_linker=False,
-            mark_write_through_pending=lambda value: setattr(
-                node, "write_through_pending_id", value
+            mark_write_through_pending=lambda node_ids, ack_id: (
+                setattr(node, "write_through_pending_id", ack_id) or list(node_ids)
             ),
         ),
         _components_tuple=(_Component(),),
@@ -243,8 +243,10 @@ def test_failed_offload_rolls_back_split_fragments():
     )
     nodes = {child.id: child, parent.id: parent}
 
-    def mark_pending(node_id):
-        nodes[node_id].write_through_pending_id = node_id
+    def mark_pending(node_ids, ack_id):
+        for node_id in node_ids:
+            nodes[node_id].write_through_pending_id = ack_id
+        return list(node_ids)
 
     cache = _cache_for_wrapper(
         tree_core=SimpleNamespace(
@@ -317,8 +319,8 @@ def test_reset_quiesces_backend_before_releasing_pending_locks():
     cache = _cache_for_wrapper(
         tree_core=SimpleNamespace(
             enable_external_cache_linker=False,
-            mark_write_through_pending=lambda value: setattr(
-                node, "write_through_pending_id", value
+            mark_write_through_pending=lambda node_ids, ack_id: (
+                setattr(node, "write_through_pending_id", ack_id) or list(node_ids)
             ),
         ),
         _components_tuple=(_Component(),),

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -7524,6 +7524,19 @@ def _component_with_cache(component_type, cache):
 class TestUnifiedRadixCacheActionRouting(CustomTestCase):
     """CacheAction routing: each type forwards to the right Controller API."""
 
+    def test_backup_publish_node_ids_collects_component_nodes_once(self):
+        comp_xfers = {
+            ComponentType.SWA: [PoolTransfer(name=PoolName.SWA, nodes_to_load=[3, 4])],
+            ComponentType.MAMBA: [
+                PoolTransfer(name=PoolName.MAMBA, nodes_to_load=[4, 5])
+            ],
+        }
+
+        self.assertEqual(
+            UnifiedRadixCache._backup_publish_node_ids(7, comp_xfers),
+            [3, 4, 5, 7],
+        )
+
     def test_apply_cache_action_routes_replace_write_through(self):
         cache = mock.MagicMock()
         action = ReplaceWriteThroughOnNodeSplit(
@@ -7993,6 +8006,27 @@ class TestResumableInsertWalk(_InsertWalkSuite):
                 cache.evict(EvictParams(num_tokens=8))
         self.assertEqual(allocator.available_size(), available + 4)
 
+    def test_write_through_publish_list_is_ordered_ancestors_first(self):
+        """One ack spanning several nodes publishes a parent before its children,
+        whatever order the component transfers listed them in."""
+        cache, allocator, req_to_token_pool = self._build_hicache_fixture()
+        self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4])
+        (parent,) = _node_children(cache, cache.root_node_handle())
+        self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4, 5, 6])
+        (child,) = _node_children(cache, parent)
+
+        cache._track_write_through_node(
+            child, lock_params=None, publish_node_ids=[child, parent]
+        )
+
+        self.assertEqual(
+            cache.ongoing_write_through[child].publish_node_ids, [parent, child]
+        )
+        cache._finish_write_through_ack(child)
+        self.assertIsNone(cache.tree_core.get_write_through_pending_id(parent))
+        self.assertIsNone(cache.tree_core.get_write_through_pending_id(child))
+        cache.sanity_check()
+
     def test_match_split_relocation_survives_finalizer_failure(self):
         """A match-walk split's pending write-through relocation applies before
         the finalizers, so a finalizer failure cannot strand the stale record."""