fix(hicache): limit load-back pending to write-back (#34519)

Co-authored-by: Zhangheng <hzh0425@apache.org>

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
+        if self.is_write_back:
-        # aux-only nodes) against reclaim until the ack.
+            # Write-back may reclaim a duplicate host copy while H->D DMA is
-        for xfers in ([kv_xfer], *comp_xfers.values()):
+            # still reading it, so pin every source node until the ack.
-            for xfer in xfers:
+            for xfers in ([kv_xfer], *comp_xfers.values()):
-                for nid in xfer.nodes_to_load or ():
+                for xfer in xfers:
-                    pinned = self.node_by_id(nid)
+                    for nid in xfer.nodes_to_load or ():
-                    # One live load-back per node; only the same anchor may
+                        pinned = self.node_by_id(nid)
-                    # re-pin (a node can sit in Full and aux transfer lists).
+                        # One live load-back per node; only the same anchor may
-                    assert pinned.load_back_pending_id in (None, node_id), (
+                        # re-pin (a node can sit in Full and aux transfer lists).
-                        f"node {nid} pinned by load-back "
+                        assert pinned.load_back_pending_id in (None, node_id), (
-                        f"{pinned.load_back_pending_id}, new anchor {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
+                        )
+                        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
+        """Finalize H->D load-back state along the anchor's root path.
-        time; split fragments stay on the path, so the walk covers them."""
+
+        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:

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(