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 9797081ba..3055924db 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 @@ -1927,6 +1927,20 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): ) if t: comp_xfers[comp.component_type] = t + # Reject transfers that would claim a node pinned by another load-back + # anchor; the empty spec makes the caller back off and recompute. + if any( + self.node_by_id(nid).load_back_pending_id not in (None, node_id) + for xfers in ([kv_xfer], *comp_xfers.values()) + for xfer in xfers + for nid in xfer.nodes_to_load or () + ): + empty_kv = PoolTransfer( + name=PoolName.KV, + host_indices=torch.empty((0,), dtype=torch.int64, device="cpu"), + nodes_to_load=[], + ) + return empty_kv, {} return kv_xfer, comp_xfers def prefetch_anchor_info( diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 657ba1c8c..9a749956b 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -1465,8 +1465,10 @@ class UnifiedRadixCache(BasePrefixCache): # Skip if there is nothing to load, or if the Full-KV transfer is too # small / exceeds memory quota. Aux transfers should still run even - # when the Full-KV load is skipped by thresholding. - if (kv_tokens < self.load_back_threshold and not comp_xfers) or ( + # when the Full-KV load is skipped by thresholding. max(1, ...): an + # entirely empty spec (e.g. foreign-pin rejection) must never report + # success, even at load_back_threshold <= 0. + if (kv_tokens < max(1, self.load_back_threshold) and not comp_xfers) or ( mem_quota is not None and kv_tokens > mem_quota + result.delta ): self.dec_lock_ref(node_id, ancestor_lock_params) 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 d0a85179c..f37f480cb 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 @@ -5698,6 +5698,62 @@ class UnifiedRadixCacheSuite: self.assertGreaterEqual(int(xfer.host_indices.numel()), sw) self.assertEqual(xfer.nodes_to_load, chain[-expected_pages:]) + def test_hicache_swa_load_back_rejects_foreign_pinned_window(self): + """A second load-back must not claim nodes pinned by an in-flight one. + + commit_load_back republishes Full device values before the DMA acks, so + a later anchor can claim just the still-host-only SWA window of a + pinned node and hit the commit pin assert. The spec build must degrade + to an empty spec instead (the caller recomputes). + """ + if not self.cfg.has_swa: + self.skipTest("requires SWA") + if self.cfg.has_mamba: + self.skipTest("SWA-only path keeps the chain construction simple") + if self.cfg.sliding_window_size <= self.cfg.page_size: + # A window within one page never reaches the pinned ancestor. + self.skipTest("window must span past the leaf to reach the pin") + + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + # commit_load_back pins its source nodes only under write-back; run the + # scenario in that mode so the foreign-pin path is reachable. + cache.is_write_back = True + chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3) + if len(chain) < 3: + self.skipTest("chain collapsed below the two-node suffix being tested") + self._simulate_backup_tree(cache) + + # Host-only suffix a -> b under a device-resident ancestor. + a, b = chain[-2], chain[-1] + for n in (a, b): + cache.tree_core.set_component_device_value_raw(n, ComponentType.FULL, None) + cache.tree_core.set_component_device_value_raw(n, ComponentType.SWA, None) + if cache.tree_core.is_node_in_device_lru(n, ComponentType.SWA): + cache.tree_core.remove_node_from_device_lru(n, ComponentType.SWA) + cache.tree_core.insert_node_into_host_lru(n, ComponentType.SWA) + + # Anchor `a`: a Full-only load whose SWA slice stays host-only. + kv_xfer, _comp_xfers = cache.tree_core.build_load_back_spec(a) + self.assertEqual(kv_xfer.nodes_to_load, [a]) + device_indices = torch.arange( + int(kv_xfer.host_indices.numel()), dtype=torch.int64, device=cache.device + ) + cache.tree_core.commit_load_back(a, device_indices, kv_xfer, {}) + self.assertEqual(cache.tree_core.node_by_id(a).load_back_pending_id, a) + + # Anchor `b` rejects its whole spec: its SWA window claims pinned `a`. + kv_xfer, comp_xfers = cache.tree_core.build_load_back_spec(b) + self.assertEqual(int(kv_xfer.host_indices.numel()), 0) + self.assertEqual(kv_xfer.nodes_to_load, []) + self.assertEqual(comp_xfers, {}) + + # After the ack unpins, the same spec builds fully. + cache.tree_core.finish_load_back(a) + self.assertIsNone(cache.tree_core.node_by_id(a).load_back_pending_id) + kv_xfer, comp_xfers = cache.tree_core.build_load_back_spec(b) + self.assertEqual(kv_xfer.nodes_to_load, [b]) + self.assertEqual(comp_xfers[ComponentType.SWA][0].nodes_to_load, [a, b]) + def _swa_finalize_setup(self): """Build a SWA chain long enough to fill at least the window plus one extra page, and host-back every node so we can flip