diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 8351f519b..bc47648ad 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -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)) 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 3289aa3cc..a59f6ffd7 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 @@ -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, ):