chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[RadixTree][4/N Refactor]: Move available_and_evictable_str to individual radix cache classes (#17852)

Browse Source
This commit is contained in:
pansicheng
2026-02-19 17:03:15 +08:00
committed by GitHub
parent 82a0bafc1c
commit 48642d5384
7 changed files with 50 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
python/sglang/srt/mem_cache/base_prefix_cache.py
View File

@@ -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"

19
python/sglang/srt/mem_cache/common.py
View File

@@ -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()

20
python/sglang/srt/mem_cache/mamba_radix_cache.py
View File

@@ -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(

24
python/sglang/srt/mem_cache/swa_radix_cache.py
View File

@@ -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(