[UnifiedTree] gate load back pre-evict on full-attn availability only (#26302)

Signed-off-by: Vladislav Nosivskoy <vladnosiv@gmail.com>
2026-05-29 00:30:03 +08:00
parent e33bbbb467
commit 34ea682a07
2 changed files with 91 additions and 1 deletions
@@ -1419,7 +1419,10 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
            self.dec_lock_ref(best_match_node, ancestor_lock_params)
            return False

-        avail = self.token_to_kv_pool_allocator.available_size()
+        if self.supports_swa():
+            avail = self.token_to_kv_pool_allocator.full_available_size()
+        else:
+            avail = self.token_to_kv_pool_allocator.available_size()
        if avail < kv_tokens:
            needed = kv_tokens - avail
            result = self.evict(EvictParams(num_tokens=needed))
@@ -2521,6 +2521,93 @@ class UnifiedRadixCacheSuite:
        tree.dec_lock_ref(leaf, request_lock.to_dec_params())
        self.assertEqual(cd.lock_ref, 0)

+    def test_hicache_swa_load_back_uses_full_pool_capacity(self):
+        """load_back should gate Full KV load on Full pool capacity only."""
+        if not self.cfg.has_swa:
+            self.skipTest("requires SWA")
+        if self.cfg.has_mamba:
+            self.skipTest("SWA-only path")
+        if self.cfg.page_size > 1:
+            self.skipTest("page_size==1 for direct swa_attn_allocator access")
+
+        tree, allocator, req_to_token_pool = self._build_hicache_fixture()
+
+        sw = self.cfg.sliding_window_size
+        kv_tokens = sw + 2
+        chain = self._build_chain_pages(tree, allocator, req_to_token_pool, kv_tokens)
+        if len(chain) < kv_tokens:
+            self.skipTest("chain too short")
+        leaf = chain[-1]
+
+        self._backup_tree(tree)
+        result = tree.evict(EvictParams(num_tokens=kv_tokens))
+        self.assertGreaterEqual(result.num_tokens_evicted, kv_tokens)
+        self.assertIsNone(leaf.component_data[ComponentType.FULL].value)
+
+        kv_xfer = tree.components[ComponentType.FULL].build_hicache_transfers(
+            leaf, CacheTransferPhase.LOAD_BACK
+        )[0]
+        self.assertEqual(int(kv_xfer.host_indices.numel()), kv_tokens)
+
+        swa_xfer = tree.components[ComponentType.SWA].build_hicache_transfers(
+            leaf, CacheTransferPhase.LOAD_BACK
+        )[0]
+        self.assertEqual(int(swa_xfer.host_indices.numel()), sw)
+
+        # Leave tree-owned SWA available for controller-side SWA eviction.
+        unrelated_seq = self._make_seq(100_000, sw)
+        self._insert(tree, allocator, req_to_token_pool, unrelated_seq)
+        self.assertGreaterEqual(tree.swa_evictable_size(), sw)
+
+        # Make raw SWA availability smaller than both load-back transfers.
+        target_swa_avail = sw - 1
+        swa_avail = allocator.swa_attn_allocator.available_size()
+        self.assertGreaterEqual(swa_avail, target_swa_avail)
+        if swa_avail > target_swa_avail:
+            external_swa = allocator.swa_attn_allocator.alloc(
+                swa_avail - target_swa_avail
+            )
+            self.assertIsNotNone(external_swa)
+
+        self.assertGreaterEqual(
+            allocator.full_attn_allocator.available_size(),
+            int(kv_xfer.host_indices.numel()),
+        )
+        self.assertLess(
+            allocator.swa_attn_allocator.available_size(),
+            int(kv_xfer.host_indices.numel()),
+        )
+        self.assertLess(
+            allocator.swa_attn_allocator.available_size(),
+            int(swa_xfer.host_indices.numel()),
+        )
+
+        with mock.patch.object(tree, "evict", wraps=tree.evict) as evict_mock:
+            self.assertTrue(tree.load_back(leaf))
+
+        # Full pre-eviction must not be triggered by SWA pool pressure.
+        full_pre_evict_calls = [
+            call
+            for call in evict_mock.call_args_list
+            if call.args and call.args[0].num_tokens > 0
+        ]
+        self.assertEqual(full_pre_evict_calls, [])
+
+        # SWA shortage is handled by the controller through SWA-only eviction.
+        self.assertTrue(
+            any(
+                call.args
+                and call.args[0].num_tokens == 0
+                and call.args[0].swa_num_tokens > 0
+                for call in evict_mock.call_args_list
+            )
+        )
+
+        self._finish_pending_loads(tree)
+        self.assertIsNotNone(leaf.component_data[ComponentType.FULL].value)
+        self._release_ongoing_load_back_locks(tree)
+        tree.sanity_check()
+
    def test_hicache_full_temp_lock_skips_evicted_anchor_and_mirrors_on_release(
        self,
    ):