From 5263568bcbf9da7b463e25c136f6aa92eedc3c08 Mon Sep 17 00:00:00 2001
From: weireweire <weiliangl@nvidia.com>
Date: Thu, 27 Aug 2026 01:30:38 +0800
Subject: [PATCH] [HiCache] Keep auxiliary load-back out of Full KV pending
 ownership (#36317)

Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>
---
 .../unified_cache/unified_tree_core.py        | 23 +++++++--------
 .../test_unified_radix_cache_unittest.py      | 28 +++++++++++++++++++
 2 files changed, 38 insertions(+), 13 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 a50d4fd27..c86e1d807 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
@@ -1983,19 +1983,16 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
         node = self.node_by_id(node_id)
         cache_actions: list[CacheAction | ComponentAction] = []
         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
+            # Pin Full KV host slots against duplicate reclaim until the ack.
+            # Auxiliary pools have independent host locks and may legitimately
+            # load the same radix node under a different anchor.
+            for nid in kv_xfer.nodes_to_load or ():
+                pinned = self.node_by_id(nid)
+                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,
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 bb3fe457e..ec0bd04b7 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
@@ -331,6 +331,34 @@ class TestUnifiedTreeCoreLoadBackPending(CustomTestCase):
         self.assertTrue(UnifiedTreeCore._can_reclaim_full_host_duplicate(core, shared))
         core._update_duplicate_tracking.assert_called_once_with(shared)
 
+    def test_auxiliary_load_does_not_reuse_full_pending_pin(self):
+        core, shared, anchor_a, anchor_b = self._build_core(is_write_back=True)
+        core.components_by_type[ComponentType.SWA] = mock.Mock()
+
+        self._commit_load_back(core, anchor_a, shared)
+        self.assertEqual(shared.load_back_pending_id, anchor_a.id)
+
+        kv_transfer = PoolTransfer(
+            name=PoolName.KV,
+            host_indices=torch.tensor([1], dtype=torch.int64),
+            nodes_to_load=[anchor_b.id],
+        )
+        swa_transfer = PoolTransfer(
+            name=PoolName.SWA,
+            host_indices=torch.tensor([2], dtype=torch.int64),
+            nodes_to_load=[shared.id],
+        )
+        UnifiedTreeCore.commit_load_back(
+            core,
+            anchor_b.id,
+            torch.tensor([3], dtype=torch.int64),
+            kv_transfer,
+            {ComponentType.SWA: [swa_transfer]},
+        )
+
+        self.assertEqual(shared.load_back_pending_id, anchor_a.id)
+        self.assertEqual(anchor_b.load_back_pending_id, anchor_b.id)
+
 
 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."""