From 48642d5384431cff6ea3eb697b141e98fe6aa106 Mon Sep 17 00:00:00 2001 From: pansicheng Date: Thu, 19 Feb 2026 17:03:15 +0800 Subject: [PATCH] [RadixTree][4/N Refactor]: Move available_and_evictable_str to individual radix cache classes (#17852) --- .../sglang/srt/mem_cache/base_prefix_cache.py | 5 ++++ python/sglang/srt/mem_cache/common.py | 19 ++------------- .../sglang/srt/mem_cache/mamba_radix_cache.py | 20 ++++++++-------- .../sglang/srt/mem_cache/swa_radix_cache.py | 24 +++++++++++-------- .../radix_cache/test_mamba_unittest.py | 5 ++++ .../radix_cache/test_radix_cache_unit.py | 9 +++++++ .../radix_cache/test_swa_unittest.py | 5 ++++ 7 files changed, 50 insertions(+), 37 deletions(-) diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 98680d314..a28adc751 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -222,3 +222,8 @@ class BasePrefixCache(ABC, PrefixCacheTrait): def is_tree_cache(self) -> bool: return not self.is_chunk_cache() + + def available_and_evictable_str(self) -> str: + available_size = self.token_to_kv_pool_allocator.available_size() + evictable_size = self.evictable_size() + return f"Available tokens: {available_size + evictable_size} ({available_size=} + {evictable_size=})\n" diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 7c7e82988..76de4912b 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -508,20 +508,5 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr tree_cache.req_to_token_pool.free(req) -def available_and_evictable_str(tree_cache) -> str: - token_to_kv_pool_allocator = tree_cache.token_to_kv_pool_allocator - if isinstance(token_to_kv_pool_allocator, SWATokenToKVPoolAllocator): - full_available_size = token_to_kv_pool_allocator.full_available_size() - swa_available_size = token_to_kv_pool_allocator.swa_available_size() - full_evictable_size = tree_cache.full_evictable_size() - swa_evictable_size = tree_cache.swa_evictable_size() - return ( - f"Available full tokens: {full_available_size + full_evictable_size} ({full_available_size=} + {full_evictable_size=})\n" - f"Available swa tokens: {swa_available_size + swa_evictable_size} ({swa_available_size=} + {swa_evictable_size=})\n" - f"Full LRU list evictable size: {tree_cache.full_lru_list_evictable_size()}\n" - f"SWA LRU list evictable size: {tree_cache.swa_lru_list_evictable_size()}\n" - ) - else: - available_size = token_to_kv_pool_allocator.available_size() - evictable_size = tree_cache.evictable_size() - return f"Available tokens: {available_size + evictable_size} ({available_size=} + {evictable_size=})\n" +def available_and_evictable_str(tree_cache: BasePrefixCache) -> str: + return tree_cache.available_and_evictable_str() diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index ebd5f6e82..07f5278d9 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -350,10 +350,10 @@ class LRUList: if self.mamba: evictable_size = tree_cache.mamba_evictable_size() - lru_list_evictable_size = tree_cache.mamba_lru_list_evictable_size() + lru_list_evictable_size = self.sanity_check_evictable_size() else: evictable_size = tree_cache.full_evictable_size() - lru_list_evictable_size = tree_cache.full_lru_list_evictable_size() + lru_list_evictable_size = self.sanity_check_evictable_size() assert ( evictable_size == lru_list_evictable_size @@ -857,14 +857,6 @@ class MambaRadixCache(BasePrefixCache): def mamba_evictable_size(self) -> int: return self.mamba_evictable_size_ - # Note: this is expensive, only use for debug - def full_lru_list_evictable_size(self) -> int: - return self.full_lru_list.sanity_check_evictable_size() - - # Note: this is expensive, only use for debug - def mamba_lru_list_evictable_size(self) -> int: - return self.mamba_lru_list.sanity_check_evictable_size() - def protected_size(self) -> Tuple[int, int]: # Note: use full_protected_size() and mamba_protected_size() instead. raise NotImplementedError @@ -900,6 +892,14 @@ class MambaRadixCache(BasePrefixCache): _dfs_helper(self.root_node) return torch.cat(values) if len(values) > 0 else torch.tensor([]) + def available_and_evictable_str(self) -> str: + full_available_size = self.token_to_kv_pool_allocator.available_size() + full_evictable_size = self.full_evictable_size() + return ( + f"Available full tokens: {full_available_size + full_evictable_size} ({full_available_size=} + {full_evictable_size=})\n" + f"Full LRU list evictable size: {self.full_lru_list.sanity_check_evictable_size()}\n" + ) + ##### Internal Helper Functions ##### def _match_prefix_helper( diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 1813a18ee..b0e548c17 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -322,10 +322,10 @@ class LRUList: if self.is_swa_list: evictable_size = tree_cache.swa_evictable_size() - lru_list_evictable_size = tree_cache.swa_lru_list_evictable_size() + lru_list_evictable_size = self.sanity_check_evictable_size() else: evictable_size = tree_cache.full_evictable_size() - lru_list_evictable_size = tree_cache.full_lru_list_evictable_size() + lru_list_evictable_size = self.sanity_check_evictable_size() assert ( evictable_size == lru_list_evictable_size @@ -781,14 +781,6 @@ class SWARadixCache(BasePrefixCache): def swa_evictable_size(self) -> int: return self.swa_evictable_size_ - # Note: this is expensive, only use for debug - def full_lru_list_evictable_size(self) -> int: - return self.full_lru_list.sanity_check_evictable_size() - - # Note: this is expensive, only use for debug - def swa_lru_list_evictable_size(self) -> int: - return self.swa_lru_list.sanity_check_evictable_size() - def protected_size(self) -> Tuple[int, int]: # Note: use full_protected_size() and swa_protected_size() instead. raise NotImplementedError @@ -812,6 +804,18 @@ class SWARadixCache(BasePrefixCache): _dfs_helper(self.root_node) return torch.cat(values) + def available_and_evictable_str(self) -> str: + full_available_size = self.token_to_kv_pool_allocator.full_available_size() + swa_available_size = self.token_to_kv_pool_allocator.swa_available_size() + full_evictable_size = self.full_evictable_size() + swa_evictable_size = self.swa_evictable_size() + return ( + f"Available full tokens: {full_available_size + full_evictable_size} ({full_available_size=} + {full_evictable_size=})\n" + f"Available swa tokens: {swa_available_size + swa_evictable_size} ({swa_available_size=} + {swa_evictable_size=})\n" + f"Full LRU list evictable size: {self.full_lru_list.sanity_check_evictable_size()}\n" + f"SWA LRU list evictable size: {self.swa_lru_list.sanity_check_evictable_size()}\n" + ) + ##### Internal Helper Functions ##### def _match_prefix_helper( diff --git a/test/registered/radix_cache/test_mamba_unittest.py b/test/registered/radix_cache/test_mamba_unittest.py index dd05ea1eb..0cebc9ff2 100644 --- a/test/registered/radix_cache/test_mamba_unittest.py +++ b/test/registered/radix_cache/test_mamba_unittest.py @@ -12,6 +12,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, ) from sglang.srt.mem_cache.cache_init_params import CacheInitParams +from sglang.srt.mem_cache.common import available_and_evictable_str from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, HybridReqToTokenPool from sglang.srt.mem_cache.radix_cache import RadixKey @@ -381,6 +382,10 @@ class TestMamba(unittest.TestCase): == mamba_pool.mamba_cache.temporal[:, last_node.mamba_value] ) + print(tree.available_and_evictable_str()) + print(available_and_evictable_str(tree)) + tree.sanity_check() + if __name__ == "__main__": unittest.main() diff --git a/test/registered/radix_cache/test_radix_cache_unit.py b/test/registered/radix_cache/test_radix_cache_unit.py index 73aa37b9f..41d202ad6 100644 --- a/test/registered/radix_cache/test_radix_cache_unit.py +++ b/test/registered/radix_cache/test_radix_cache_unit.py @@ -17,6 +17,7 @@ Usage: python -m pytest test_radix_cache_unit.py::TestRadixCache::test_insert_basic """ +from sglang.srt.mem_cache.common import available_and_evictable_str from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci # CPU-based unit test, runs quickly on any GPU runner @@ -764,6 +765,14 @@ class TestRadixCache(unittest.TestCase): # The cache size should be within reasonable bounds of the actual allocated memory. self.assertLess(torch_allocated, cache_size_bytes * 2) + def test_available_and_evictable_str(self): + mock_allocator = unittest.mock.Mock() + mock_allocator.available_size.return_value = 10 + cache: RadixCache = RadixCache.create_simulated(mock_allocator=mock_allocator) + + print(cache.available_and_evictable_str()) + print(available_and_evictable_str(cache)) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/radix_cache/test_swa_unittest.py b/test/registered/radix_cache/test_swa_unittest.py index 626fda769..3b1a57de8 100644 --- a/test/registered/radix_cache/test_swa_unittest.py +++ b/test/registered/radix_cache/test_swa_unittest.py @@ -9,6 +9,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, ) from sglang.srt.mem_cache.cache_init_params import CacheInitParams +from sglang.srt.mem_cache.common import available_and_evictable_str from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator @@ -231,6 +232,10 @@ class TestSWA(unittest.TestCase): self.assertEqual(last_node.key.token_ids[0], 60) self.assertEqual(last_node.key.token_ids[1], 70) + print(tree.available_and_evictable_str()) + print(available_and_evictable_str(tree)) + tree.sanity_check() + def test_swa_radix_cache_eagle(self): # args req_size = 10