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

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