From 19b46863f3514236f299df75607d9c487099a1da Mon Sep 17 00:00:00 2001 From: Shuwen Wang <47200617+alphabetc1@users.noreply.github.com> Date: Fri, 4 Sep 2026 22:08:25 +0800 Subject: [PATCH] fix: align write-through pending across tree cores (#37278) --- .../srt/mem_cache/rust_tree_core/adapter.py | 6 +- .../unified_cache/unified_cache_linker.py | 2 +- .../unified_cache/unified_tree_core.py | 27 +++++++-- .../unified_tree_core_interface.py | 7 ++- .../srt/mem_cache/unified_radix_cache.py | 25 +++++++- rust/sglang-radix-tree/src/python_bindings.rs | 22 +++++-- .../src/tests/components/full.rs | 2 +- .../src/tests/components/mamba.rs | 2 +- .../src/tests/unified_tree_core.rs | 60 +++++++++++++++++-- .../src/unified_tree_core.rs | 37 ++++++++++-- .../test_rust_tree_core_integration.py | 28 ++++++++- .../mem_cache/test_unified_cache_linker.py | 14 +++-- .../test_unified_radix_cache_unittest.py | 34 +++++++++++ 13 files changed, 228 insertions(+), 38 deletions(-) diff --git a/python/sglang/srt/mem_cache/rust_tree_core/adapter.py b/python/sglang/srt/mem_cache/rust_tree_core/adapter.py index a553d00ad..a9459ec6a 100644 --- a/python/sglang/srt/mem_cache/rust_tree_core/adapter.py +++ b/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) diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_cache_linker.py b/python/sglang/srt/mem_cache/unified_cache/unified_cache_linker.py index 8e4dae0c9..1dfc6b999 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_cache_linker.py +++ b/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])) diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index 9ef9bbd33..37264a7d9 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py index ec91abc9c..db7125902 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 38b236af7..d100a416d 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/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( diff --git a/rust/sglang-radix-tree/src/python_bindings.rs b/rust/sglang-radix-tree/src/python_bindings.rs index 6b6bc3823..af470a613 100644 --- a/rust/sglang-radix-tree/src/python_bindings.rs +++ b/rust/sglang-radix-tree/src/python_bindings.rs @@ -1820,9 +1820,14 @@ impl TreeCoreBinding { }) } - /// 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, + ack_id: NodeId, + ) -> Vec { + 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, + ack_id: NodeId, + ) -> Vec { + self.inner.mark_write_through_pending(py, node_ids, ack_id) } /// Clear the write-through-pending mark on the acked nodes. diff --git a/rust/sglang-radix-tree/src/tests/components/full.rs b/rust/sglang-radix-tree/src/tests/components/full.rs index 0794cc636..e6017c0f7 100644 --- a/rust/sglang-radix-tree/src/tests/components/full.rs +++ b/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( diff --git a/rust/sglang-radix-tree/src/tests/components/mamba.rs b/rust/sglang-radix-tree/src/tests/components/mamba.rs index c749aa0d1..2fc8d8a9e 100644 --- a/rust/sglang-radix-tree/src/tests/components/mamba.rs +++ b/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 diff --git a/rust/sglang-radix-tree/src/tests/unified_tree_core.rs b/rust/sglang-radix-tree/src/tests/unified_tree_core.rs index 61fec2de9..105f70014 100644 --- a/rust/sglang-radix-tree/src/tests/unified_tree_core.rs +++ b/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, diff --git a/rust/sglang-radix-tree/src/unified_tree_core.rs b/rust/sglang-radix-tree/src/unified_tree_core.rs index 55989fd07..1e8426c1d 100644 --- a/rust/sglang-radix-tree/src/unified_tree_core.rs +++ b/rust/sglang-radix-tree/src/unified_tree_core.rs @@ -3668,10 +3668,39 @@ impl UnifiedTreeCore { } } - /// 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, + ack_id: NodeId, + ) -> Vec { + 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 diff --git a/test/registered/unit/mem_cache/test_rust_tree_core_integration.py b/test/registered/unit/mem_cache/test_rust_tree_core_integration.py index 945029445..008e04269 100644 --- a/test/registered/unit/mem_cache/test_rust_tree_core_integration.py +++ b/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,) = [ diff --git a/test/registered/unit/mem_cache/test_unified_cache_linker.py b/test/registered/unit/mem_cache/test_unified_cache_linker.py index c64e503cf..9fa239a0f 100644 --- a/test/registered/unit/mem_cache/test_unified_cache_linker.py +++ b/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(),), diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 35d92ded4..527c6a423 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/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."""