[RadixTree][4/N Refactor]: Move available_and_evictable_str to individual radix cache classes (#17852)
This commit is contained in:
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user