From 43226af8121f20de0e217bc5c344fe04e746397f Mon Sep 17 00:00:00 2001 From: ziang663 <119752791+ziang663@users.noreply.github.com> Date: Mon, 17 Aug 2026 14:23:54 +0800 Subject: [PATCH] fix(hicache): limit load-back pending to write-back (#34519) Co-authored-by: Zhangheng --- .../unified_cache/unified_tree_core.py | 44 ++++++----- .../test_unified_radix_cache_unittest.py | 77 +++++++++++++++++++ 2 files changed, 103 insertions(+), 18 deletions(-) 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 0646734f3..e39a9e01c 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 @@ -1958,19 +1958,20 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): rebuild is deferred to the orchestration layer.""" node = self.node_by_id(node_id) cache_actions: list[CacheAction | ComponentAction] = [] - # Pin every node whose host slots the in-flight DMA reads (including - # aux-only nodes) against reclaim until the ack. - for xfers in ([kv_xfer], *comp_xfers.values()): - for xfer in xfers: - for nid in xfer.nodes_to_load or (): - pinned = self.node_by_id(nid) - # One live load-back per node; only the same anchor may - # re-pin (a node can sit in Full and aux transfer lists). - assert pinned.load_back_pending_id in (None, node_id), ( - f"node {nid} pinned by load-back " - f"{pinned.load_back_pending_id}, new anchor {node_id}" - ) - pinned.load_back_pending_id = node_id + if self.is_write_back: + # Write-back may reclaim a duplicate host copy while H->D DMA is + # still reading it, so pin every source node until the ack. + for xfers in ([kv_xfer], *comp_xfers.values()): + for xfer in xfers: + for nid in xfer.nodes_to_load or (): + pinned = self.node_by_id(nid) + # One live load-back per node; only the same anchor may + # re-pin (a node can sit in Full and aux transfer lists). + assert pinned.load_back_pending_id in (None, node_id), ( + f"node {nid} pinned by load-back " + f"{pinned.load_back_pending_id}, new anchor {node_id}" + ) + pinned.load_back_pending_id = node_id kv_xfer.device_indices = device_indices self.components_by_type[BASE_COMPONENT_TYPE].commit_hicache_transfer( node, @@ -1991,14 +1992,21 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): return cache_actions def finish_load_back(self, anchor_node_id: NodeId) -> None: - """Clear the in-flight H->D marks along the anchor's root path at ack - time; split fragments stay on the path, so the walk covers them.""" + """Finalize H->D load-back state along the anchor's root path. + + Write-back clears source-node pins at ack time. Write-through does not + use those pins, but still refreshes duplicate tracking after the device + copies become visible. Split fragments stay on the path, so the walk + covers them. + """ node = self.node_by_id(anchor_node_id) while node is not None and node is not self.root_node: - if node.load_back_pending_id == anchor_node_id: + if self.is_write_back: + if node.load_back_pending_id != anchor_node_id: + node = node.parent + continue node.load_back_pending_id = None - # The loaded copies become tracked duplicates only now. - self._update_duplicate_tracking(node) + self._update_duplicate_tracking(node) node = node.parent def mark_write_through_pending(self, node_id: NodeId) -> None: 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 8a7a344f5..932eae344 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 @@ -65,6 +65,7 @@ from sglang.srt.mem_cache.unified_cache.components.tree_component import ( EvictLayer, TreeComponent, ) +from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import ( DecSwaLockOnlyResult, DemoteResult, @@ -254,6 +255,82 @@ class TestUnifiedTreeNodeGetPrefixHashValues(CustomTestCase): self.assertEqual(n4.get_prefix_hash_values(n3), ["h1", "h2", "h3"]) +class TestUnifiedTreeCoreLoadBackPending(CustomTestCase): + def _build_core(self, *, is_write_back: bool): + component_types = (ComponentType.FULL,) + root = UnifiedTreeNode(component_types) + shared = UnifiedTreeNode(component_types) + anchor_a = UnifiedTreeNode(component_types) + anchor_b = UnifiedTreeNode(component_types) + shared.parent = root + anchor_a.parent = shared + anchor_b.parent = shared + + nodes = {node.id: node for node in (root, shared, anchor_a, anchor_b)} + core = mock.Mock() + core.is_write_back = is_write_back + core.root_node = root + core.node_by_id.side_effect = nodes.__getitem__ + core.components_by_type = {ComponentType.FULL: mock.Mock()} + core.full_host_duplicates = {} + core._is_settled_full_host_duplicate.side_effect = ( + lambda node: UnifiedTreeCore._is_settled_full_host_duplicate(core, node) + ) + core._update_duplicate_tracking.side_effect = ( + lambda node: UnifiedTreeCore._update_duplicate_tracking(core, node) + ) + return core, shared, anchor_a, anchor_b + + def _commit_load_back(self, core, anchor, source): + transfer = PoolTransfer( + name=PoolName.KV, + host_indices=torch.tensor([1], dtype=torch.int64), + nodes_to_load=[source.id], + ) + return UnifiedTreeCore.commit_load_back( + core, + anchor.id, + torch.tensor([2], dtype=torch.int64), + transfer, + {}, + ) + + def test_write_through_different_anchors_track_duplicate_without_pending(self): + core, shared, anchor_a, anchor_b = self._build_core(is_write_back=False) + full = shared.component_data[ComponentType.FULL] + full.value = torch.tensor([1], dtype=torch.int64) + full.host_value = torch.tensor([2], dtype=torch.int64) + + self._commit_load_back(core, anchor_a, shared) + self._commit_load_back(core, anchor_b, shared) + UnifiedTreeCore.finish_load_back(core, anchor_a.id) + + self.assertIsNone(shared.load_back_pending_id) + self.assertIn(shared.id, core.full_host_duplicates) + core._update_duplicate_tracking.assert_has_calls( + [mock.call(anchor_a), mock.call(shared)] + ) + + def test_write_back_pending_blocks_reclaim_until_ack(self): + core, shared, anchor_a, anchor_b = self._build_core(is_write_back=True) + full = shared.component_data[ComponentType.FULL] + full.value = torch.tensor([1], dtype=torch.int64) + full.host_value = torch.tensor([2], dtype=torch.int64) + + self._commit_load_back(core, anchor_a, shared) + + self.assertEqual(shared.load_back_pending_id, anchor_a.id) + self.assertFalse(UnifiedTreeCore._can_reclaim_full_host_duplicate(core, shared)) + with self.assertRaisesRegex(AssertionError, "new anchor"): + self._commit_load_back(core, anchor_b, shared) + + UnifiedTreeCore.finish_load_back(core, anchor_a.id) + + self.assertIsNone(shared.load_back_pending_id) + self.assertTrue(UnifiedTreeCore._can_reclaim_full_host_duplicate(core, shared)) + core._update_duplicate_tracking.assert_called_once_with(shared) + + def _write_backup(cache, node, write_back: bool = False) -> int: """Back up one node's KV D->H via the tree's build+execute primitives.""" return cache._execute_and_commit_kv_backup(