From f33022d039c61ec139499b6651a43017ea4dd131 Mon Sep 17 00:00:00 2001 From: zhangheng Date: Thu, 22 Jan 2026 17:36:31 +0800 Subject: [PATCH] =?UTF-8?q?[RadixTree][3/N=20Refactor]=EF=BC=9ASupport=20u?= =?UTF-8?q?nified=20insert/evict=20params=20(#17401)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- python/sglang/srt/managers/schedule_policy.py | 12 +- .../sglang/srt/mem_cache/base_prefix_cache.py | 47 ++++++- python/sglang/srt/mem_cache/chunk_cache.py | 16 ++- python/sglang/srt/mem_cache/common.py | 10 +- python/sglang/srt/mem_cache/hiradix_cache.py | 32 +++-- .../sglang/srt/mem_cache/mamba_radix_cache.py | 73 +++++++--- python/sglang/srt/mem_cache/radix_cache.py | 53 ++++--- .../sglang/srt/mem_cache/radix_cache_cpp.py | 12 +- .../storage/lmcache/lmc_radix_cache.py | 15 +- .../sglang/srt/mem_cache/swa_radix_cache.py | 52 ++++--- .../radix_cache/test_mamba_unittest.py | 57 ++++++-- .../radix_cache/test_radix_cache_unit.py | 132 ++++++++++++++---- .../radix_cache/test_swa_unittest.py | 89 ++++++++++-- 13 files changed, 462 insertions(+), 138 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 57ed79f4e..f4e2746bf 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -35,7 +35,11 @@ import torch from sglang.srt.dllm.config import DllmConfig from sglang.srt.layers.attention.nsa.utils import is_nsa_prefill_cp_in_seq_split from sglang.srt.managers.schedule_batch import DllmStagingReqs, Req, ScheduleBatch -from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, MatchPrefixParams +from sglang.srt.mem_cache.base_prefix_cache import ( + BasePrefixCache, + InsertParams, + MatchPrefixParams, +) from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator from sglang.srt.server_args import ServerArgs @@ -228,8 +232,10 @@ class SchedulePolicy: else: # Insert with a dummy key self.waiting_queue_radix_tree.insert( - RadixKey(token_ids=prefix_ids, extra_key=extra_key), - torch.empty(len(prefix_ids), dtype=torch.bool), + InsertParams( + key=RadixKey(token_ids=prefix_ids, extra_key=extra_key), + value=torch.empty(len(prefix_ids), dtype=torch.bool), + ) ) return temporary_deprioritized diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index bdf8730a3..e01f98dcb 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -43,6 +43,51 @@ class MatchPrefixParams: req: Optional[Req] = None +@dataclasses.dataclass +class InsertParams: + """Unified parameters for insert across different cache types""" + + key: RadixKey + value: Optional[torch.Tensor] = None + + # Mamba specific + mamba_value: Optional[torch.Tensor] = None + + # SWA specific + prev_prefix_len: int = 0 + swa_evicted_seqlen: int = 0 + + # General + chunked: bool = False + priority: int = 0 + + +@dataclasses.dataclass +class InsertResult: + """Result of an insert operation""" + + prefix_len: int + mamba_exist: bool = False + + +@dataclasses.dataclass +class EvictParams: + """Unified parameters for evict across different cache types""" + + num_tokens: int + swa_num_tokens: int = 0 + mamba_num: int = 0 + + +@dataclasses.dataclass +class EvictResult: + """Result of an evict operation""" + + num_tokens_evicted: int = 0 + swa_num_tokens_evicted: int = 0 + mamba_num_evicted: int = 0 + + class MatchResult(NamedTuple): """Result of a prefix match operation. @@ -102,7 +147,7 @@ class BasePrefixCache(ABC, PrefixCacheTrait): pass @abstractmethod - def evict(self, num_tokens: int): + def evict(self, params: EvictParams) -> EvictResult: pass @abstractmethod diff --git a/python/sglang/srt/mem_cache/chunk_cache.py b/python/sglang/srt/mem_cache/chunk_cache.py index 87f88c7c6..a4377989b 100644 --- a/python/sglang/srt/mem_cache/chunk_cache.py +++ b/python/sglang/srt/mem_cache/chunk_cache.py @@ -9,6 +9,10 @@ import torch from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, + EvictParams, + EvictResult, + InsertParams, + InsertResult, MatchPrefixParams, MatchResult, ) @@ -54,6 +58,10 @@ class ChunkCache(BasePrefixCache): last_host_node=None, ) + def insert(self, params: InsertParams) -> InsertResult: + # ChunkCache does not support prefix caching, so insert is a no-op + return InsertResult(prefix_len=0) + def cache_finished_req(self, req: Req, is_insert: bool = True): kv_committed_len = req.pop_committed_kv_cache() # For decode server: if req.output_ids is empty, we want to free all req.origin_input_ids @@ -70,8 +78,8 @@ class ChunkCache(BasePrefixCache): # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True) - def evict(self, num_tokens: int): - pass + def evict(self, params: EvictParams) -> EvictResult: + return EvictResult() def inc_lock_ref(self, node: Any): return 0 @@ -103,5 +111,5 @@ class SWAChunkCache(ChunkCache): ), "sliding_window_size must be set for SWAChunkCache" return True - def evict(self, num_tokens: int): - pass + def evict(self, params: EvictParams) -> EvictResult: + return EvictResult() diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 374614be5..5f7589998 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -7,7 +7,7 @@ import torch import triton import triton.language as tl -from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache +from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator from sglang.srt.server_args import get_global_server_args @@ -244,11 +244,13 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): if full_available_size < num_tokens or swa_available_size < num_tokens: full_num_tokens = max(0, num_tokens - full_available_size) swa_num_tokens = max(0, num_tokens - swa_available_size) - tree_cache.evict(full_num_tokens, swa_num_tokens) + tree_cache.evict( + EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) + ) else: # Standard allocator if allocator.available_size() < num_tokens: - tree_cache.evict(num_tokens) + tree_cache.evict(EvictParams(num_tokens=num_tokens)) def alloc_paged_token_slots_extend( @@ -311,7 +313,7 @@ def alloc_req_slots( if mamba_available_size < mamba_state_needed: if tree_cache is not None and tree_cache.supports_mamba(): mamba_num = max(0, mamba_state_needed - mamba_available_size) - tree_cache.evict_mamba(mamba_num) + tree_cache.evict(EvictParams(num_tokens=0, mamba_num=mamba_num)) req_pool_indices = req_to_token_pool.alloc(num_reqs, reqs) else: req_pool_indices = req_to_token_pool.alloc(num_reqs) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index d4e3de146..443d654b9 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -10,7 +10,14 @@ from typing import TYPE_CHECKING, List, Optional import torch from sglang.srt.managers.cache_controller import HiCacheController, PrefetchOperation -from sglang.srt.mem_cache.base_prefix_cache import MatchPrefixParams, MatchResult +from sglang.srt.mem_cache.base_prefix_cache import ( + EvictParams, + EvictResult, + InsertParams, + InsertResult, + MatchPrefixParams, + MatchResult, +) from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, MLATokenToKVPool from sglang.srt.mem_cache.memory_pool_host import ( MHATokenToKVPoolHost, @@ -331,8 +338,9 @@ class HiRadixCache(RadixCache): def evictable_size(self): return self.evictable_size_ - def evict(self, num_tokens: int): + def evict(self, params: EvictParams) -> EvictResult: start_time = time.perf_counter() + num_tokens = params.num_tokens leaves = self._collect_leaves_device() eviction_heap = [ (self.eviction_strategy.get_priority(node), node) for node in leaves @@ -374,6 +382,7 @@ class HiRadixCache(RadixCache): self._evict_backuped(node) self.update_eviction_metrics(num_evicted, start_time) + return EvictResult(num_tokens_evicted=num_evicted) def _evict_backuped(self, node: TreeNode): # evict a node already written to host @@ -453,7 +462,7 @@ class HiRadixCache(RadixCache): host_indices=host_indices, node_id=last_hit_node.id ) if device_indices is None: - self.evict(len(host_indices)) + self.evict(EvictParams(num_tokens=len(host_indices))) device_indices = self.cache_controller.load( host_indices=host_indices, node_id=last_hit_node.id ) @@ -854,19 +863,18 @@ class HiRadixCache(RadixCache): new_node.parent.children[self.get_child_key_fn(key)] = new_node return new_node - def insert( - self, - key: RadixKey, - value=None, - chunked: bool = False, - priority: int | None = None, - ): + def insert(self, params: InsertParams) -> InsertResult: + key = params.key + value = params.value + chunked = params.chunked + priority = params.priority + if priority is None: priority = 0 key, value = self.maybe_bigram_convert(key, value) if len(key) == 0: - return 0 + return InsertResult(prefix_len=0) if self.is_eagle and value is not None: # Make sure the value len equal to the EAGLE bigram key len @@ -924,7 +932,7 @@ class HiRadixCache(RadixCache): if self.cache_controller.write_policy != "write_back": self._inc_hit_count(new_node, chunked) - return total_prefix_length + return InsertResult(prefix_len=total_prefix_length) def _collect_leaves_device(self): def is_leaf(node): diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index c8bae716b..3ade398ba 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -35,6 +35,10 @@ from sglang.srt.mem_cache.allocator import ( ) from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, + EvictParams, + EvictResult, + InsertParams, + InsertResult, MatchPrefixParams, MatchResult, ) @@ -454,7 +458,7 @@ class MambaRadixCache(BasePrefixCache): # try to alloc again, protect last_node from eviction if dst_index is None: self.inc_lock_ref(last_node) - self.evict_mamba(1) + self.evict(EvictParams(num_tokens=0, mamba_num=1)) dst_index = self.req_to_token_pool.mamba_pool.alloc(1) self.dec_lock_ref(last_node) assert dst_index is not None, "Can not alloc mamba cache" @@ -478,13 +482,20 @@ class MambaRadixCache(BasePrefixCache): mamba_branching_seqlen=mamba_branching_seqlen, ) - def insert(self, key: RadixKey, value=None, mamba_value=None) -> Tuple[int, bool]: + def insert(self, params: InsertParams) -> InsertResult: if self.disable: - return 0, False + return InsertResult(prefix_len=0, mamba_exist=False) + + key = params.key + value = params.value + mamba_value = params.mamba_value if value is None: value = torch.tensor([x for x in key.token_ids], dtype=torch.int64) - return self._insert_helper(self.root_node, key, value, mamba_value) + prefix_len, mamba_exist = self._insert_helper( + self.root_node, key, value, mamba_value + ) + return InsertResult(prefix_len=prefix_len, mamba_exist=mamba_exist) def cache_finished_req(self, req: Req, is_insert: bool = True) -> None: """Cache request when it finishes.""" @@ -556,11 +567,14 @@ class MambaRadixCache(BasePrefixCache): mamba_value = req.mamba_pool_idx.unsqueeze(-1).clone() mamba_ping_pong_track_buffer_to_keep = None - new_prefix_len, mamba_exist = self.insert( - RadixKey(token_ids[:page_aligned_len], req.extra_key), - page_aligned_kv_indices, - mamba_value, + result = self.insert( + InsertParams( + key=RadixKey(token_ids[:page_aligned_len], req.extra_key), + value=page_aligned_kv_indices, + mamba_value=mamba_value, + ) ) + new_prefix_len, mamba_exist = result.prefix_len, result.mamba_exist self.token_to_kv_pool_allocator.free( kv_indices[req.cache_protected_len : new_prefix_len] @@ -644,16 +658,19 @@ class MambaRadixCache(BasePrefixCache): # if alloc mamba cache failed, do evict and alloc again if mamba_value_forked is None: - self.evict_mamba(1) + self.evict(EvictParams(num_tokens=0, mamba_num=1)) mamba_value_forked = self.req_to_token_pool.mamba_pool.fork_from( mamba_value ) assert mamba_value_forked is not None, "Can not alloc mamba cache" - new_prefix_len, mamba_exist = self.insert( - RadixKey(page_aligned_token_ids, req.extra_key), - page_aligned_kv_indices, - mamba_value_forked, + result = self.insert( + InsertParams( + key=RadixKey(page_aligned_token_ids, req.extra_key), + value=page_aligned_kv_indices, + mamba_value=mamba_value_forked, + ) ) + new_prefix_len, mamba_exist = result.prefix_len, result.mamba_exist self.token_to_kv_pool_allocator.free( kv_indices[req.cache_protected_len : new_prefix_len] ) @@ -735,9 +752,26 @@ class MambaRadixCache(BasePrefixCache): full_num_evicted += leaf_full_num_evicted return full_num_evicted, mamba_num_evicted, x, x_next - def evict_mamba(self, mamba_num: int) -> None: + def evict(self, params: EvictParams) -> EvictResult: + if self.disable: + return EvictResult() + + full_num_evicted = 0 + mamba_num_evicted = 0 + + if params.num_tokens > 0: + full_num_evicted = self.evict_full(params.num_tokens) + if params.mamba_num > 0: + mamba_num_evicted = self.evict_mamba(params.mamba_num) + + return EvictResult( + num_tokens_evicted=full_num_evicted, mamba_num_evicted=mamba_num_evicted + ) + + def evict_mamba(self, mamba_num: int) -> int: + """Evict mamba states. Returns the number of mamba states evicted.""" if self.disable or mamba_num <= 0: - return + return 0 # get the least recently used node that is not locked, doesn't have to be a leaf x = self.mamba_lru_list.get_lru_no_lock() mamba_num_evicted = 0 @@ -767,9 +801,12 @@ class MambaRadixCache(BasePrefixCache): x = x_next - def evict(self, full_num_tokens: int) -> None: + return mamba_num_evicted + + def evict_full(self, full_num_tokens: int) -> int: + """Evict full KV cache. Returns the number of tokens evicted.""" if self.disable or full_num_tokens <= 0: - return + return 0 full_num_evicted = 0 # get the least recently used leaf node that is not locked @@ -789,6 +826,8 @@ class MambaRadixCache(BasePrefixCache): x = x_next + return full_num_evicted + def inc_lock_ref(self, node: TreeNode) -> Optional[int]: """ Increment the lock reference count for the node. diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 2bebf4bf4..4a7bb7229 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -41,6 +41,10 @@ from sglang.srt.disaggregation.kv_events import ( ) from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, + EvictParams, + EvictResult, + InsertParams, + InsertResult, MatchPrefixParams, MatchResult, ) @@ -413,16 +417,21 @@ class RadixCache(BasePrefixCache): last_host_node=last_node, ) - def insert(self, key: RadixKey, value=None, chunked=False, priority: int = 0): + def insert(self, params: InsertParams) -> InsertResult: if self.disable: - return 0 + return InsertResult(prefix_len=0) + + key = params.key + value = params.value + priority = params.priority if value is None: value = torch.tensor(key.token_ids, dtype=torch.int64) key, value = self.maybe_bigram_convert(key, value) - return self._insert_helper(self.root_node, key, value, priority) + prefix_len = self._insert_helper(self.root_node, key, value, priority) + return InsertResult(prefix_len=prefix_len) def _page_align_keys(self, key: list) -> list: if self.page_size == 1: @@ -459,7 +468,10 @@ class RadixCache(BasePrefixCache): # Radix Cache takes one ref in memory pool if is_insert: priority = getattr(req, "priority", 0) or 0 - new_prefix_len = self.insert(radix_key, values, priority=priority) + result = self.insert( + InsertParams(key=radix_key, value=values, priority=priority) + ) + new_prefix_len = result.prefix_len # Free the duplicates that were already in the tree self.token_to_kv_pool_allocator.free( kv_indices[req.cache_protected_len : new_prefix_len] @@ -493,12 +505,15 @@ class RadixCache(BasePrefixCache): radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle) # Radix Cache takes one ref in memory pool - new_prefix_len = self.insert( - radix_key, - values, - chunked=chunked, - priority=getattr(req, "priority", 0) or 0, + result = self.insert( + InsertParams( + key=radix_key, + value=values, + chunked=chunked, + priority=getattr(req, "priority", 0) or 0, + ) ) + new_prefix_len = result.prefix_len self.token_to_kv_pool_allocator.free( kv_indices[req.cache_protected_len : new_prefix_len] @@ -545,11 +560,12 @@ class RadixCache(BasePrefixCache): def total_size(self): return self._total_size_helper() - def evict(self, num_tokens: int): + def evict(self, params: EvictParams) -> EvictResult: if self.disable: - return + return EvictResult() start_time = time.perf_counter() + num_tokens = params.num_tokens leaves = self._collect_leaves() eviction_heap = [ (self.eviction_strategy.get_priority(node), node) for node in leaves @@ -571,6 +587,7 @@ class RadixCache(BasePrefixCache): self._record_remove_event(x) self.update_eviction_metrics(num_evicted, start_time) + return EvictResult(num_tokens_evicted=num_evicted) def inc_lock_ref(self, node: TreeNode): if self.disable: @@ -842,11 +859,15 @@ if __name__ == "__main__": tree = RadixCache.create_simulated() # Example token id sequences (as lists of ints) - tree.insert(RadixKey(token_ids=[1, 2, 3], extra_key=None)) - tree.insert(RadixKey(token_ids=[1, 2, 3], extra_key=None)) - tree.insert(RadixKey(token_ids=[1, 2, 4, 5], extra_key=None)) - tree.insert(RadixKey(token_ids=[1, 2, 4, 5, 6, 7], extra_key=None)) - tree.insert(RadixKey(token_ids=[8, 9, 10, 11, 12], extra_key=None)) + tree.insert(InsertParams(key=RadixKey(token_ids=[1, 2, 3], extra_key=None))) + tree.insert(InsertParams(key=RadixKey(token_ids=[1, 2, 3], extra_key=None))) + tree.insert(InsertParams(key=RadixKey(token_ids=[1, 2, 4, 5], extra_key=None))) + tree.insert( + InsertParams(key=RadixKey(token_ids=[1, 2, 4, 5, 6, 7], extra_key=None)) + ) + tree.insert( + InsertParams(key=RadixKey(token_ids=[8, 9, 10, 11, 12], extra_key=None)) + ) tree.pretty_print() print( diff --git a/python/sglang/srt/mem_cache/radix_cache_cpp.py b/python/sglang/srt/mem_cache/radix_cache_cpp.py index c1a55c929..c2b46f7b6 100644 --- a/python/sglang/srt/mem_cache/radix_cache_cpp.py +++ b/python/sglang/srt/mem_cache/radix_cache_cpp.py @@ -8,6 +8,8 @@ import torch from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, + EvictParams, + EvictResult, MatchPrefixParams, MatchResult, ) @@ -137,14 +139,18 @@ class RadixCacheCpp(BasePrefixCache): """ self.tree.lock_ref(node, True) - def evict(self, num_tokens: int): + def evict(self, params: EvictParams) -> EvictResult: start_time = time.perf_counter() + num_tokens = params.num_tokens evicted_device_indices = self.tree.evict(num_tokens) + + num_evicted = 0 for indice in evicted_device_indices: + num_evicted += len(indice) self.token_to_kv_pool_allocator.free(indice) - # FIXME: not sure about the real evict length here - self.update_eviction_metrics(num_tokens, start_time) + self.update_eviction_metrics(num_evicted, start_time) + return EvictResult(num_tokens_evicted=num_evicted) def evictable_size(self): return self.tree.evictable_size() diff --git a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py index 3b4702943..9a82aa31f 100644 --- a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py +++ b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py @@ -6,7 +6,12 @@ from typing import TYPE_CHECKING, Optional import torch -from sglang.srt.mem_cache.base_prefix_cache import MatchPrefixParams, MatchResult +from sglang.srt.mem_cache.base_prefix_cache import ( + EvictParams, + EvictResult, + MatchPrefixParams, + MatchResult, +) from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode try: @@ -151,7 +156,7 @@ class LMCRadixCache(RadixCache): prefix_pad = value.numel() % chunk_size if self.token_to_kv_pool_allocator.available_size() < uncached_len: - self.evict(uncached_len) + self.evict(EvictParams(num_tokens=uncached_len)) token_slots = self.token_to_kv_pool_allocator.alloc(uncached_len) if token_slots is None: @@ -248,10 +253,10 @@ class LMCRadixCache(RadixCache): with self._node_lock: self._in_flight_nodes.append(new_last_node) - def evict(self, num_tokens: int) -> None: # type: ignore[override] + def evict(self, params: EvictParams) -> EvictResult: """Before base eviction, wait for any outstanding stores and release locks.""" if self.disable: - return + return EvictResult() self.store_stream.synchronize() with self._node_lock: @@ -259,7 +264,7 @@ class LMCRadixCache(RadixCache): self.dec_lock_ref(node) self._in_flight_nodes.clear() - super().evict(num_tokens) + return super().evict(params) def pretty_print(self): # type: ignore[override] super().pretty_print() diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index eacbf634c..4b07b841f 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -30,6 +30,10 @@ from numpy import float64 from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, + EvictParams, + EvictResult, + InsertParams, + InsertResult, MatchPrefixParams, MatchResult, ) @@ -426,15 +430,14 @@ class SWARadixCache(BasePrefixCache): last_host_node=last_node, ) - def insert( - self, - key: RadixKey, - value=None, - prev_prefix_len: int = 0, - swa_evicted_seqlen: int = 0, - ) -> int: + def insert(self, params: InsertParams) -> InsertResult: if self.disable: - return 0 + return InsertResult(prefix_len=0) + + key = params.key + value = params.value + prev_prefix_len = params.prev_prefix_len + swa_evicted_seqlen = params.swa_evicted_seqlen key.token_ids = self.key_convert_fn(key.token_ids) @@ -445,9 +448,10 @@ class SWARadixCache(BasePrefixCache): # Make sure the value len equal to the EAGLE bigram key len value = value[: len(key)] - return self._insert_helper( + prefix_len = self._insert_helper( self.root_node, key, value, prev_prefix_len, swa_evicted_seqlen ) + return InsertResult(prefix_len=prefix_len) def cache_finished_req(self, req: Req, is_insert: bool = True) -> None: """Cache request when it finishes.""" @@ -492,10 +496,12 @@ class SWARadixCache(BasePrefixCache): # Note: the insert function already frees the overlapped kv_indices if is_insert: self.insert( - RadixKey(token_ids[:page_aligned_token_len], req.extra_key), - page_aligned_kv_indices, - old_prefix_len, - req.swa_evicted_seqlen, + InsertParams( + key=RadixKey(token_ids[:page_aligned_token_len], req.extra_key), + value=page_aligned_kv_indices, + prev_prefix_len=old_prefix_len, + swa_evicted_seqlen=req.swa_evicted_seqlen, + ) ) else: self.token_to_kv_pool_allocator.free( @@ -552,11 +558,14 @@ class SWARadixCache(BasePrefixCache): # Radix Cache takes one ref in memory pool # Note: the insert function already frees the overlapped kv_indices - new_prefix_len = self.insert( - RadixKey(page_aligned_token_ids, req.extra_key), - page_aligned_kv_indices, - old_prefix_len, + result = self.insert( + InsertParams( + key=RadixKey(page_aligned_token_ids, req.extra_key), + value=page_aligned_kv_indices, + prev_prefix_len=old_prefix_len, + ) ) + new_prefix_len = result.prefix_len # The prefix indices could be updated, reuse it match_result = self.match_prefix( @@ -605,10 +614,12 @@ class SWARadixCache(BasePrefixCache): def total_size(self) -> Tuple[int, int]: return self._total_size_helper() - def evict(self, full_num_tokens: int, swa_num_tokens: int = 0) -> None: + def evict(self, params: EvictParams) -> EvictResult: if self.disable: - return + return EvictResult() start_time = time.perf_counter() + full_num_tokens = params.num_tokens + swa_num_tokens = params.swa_num_tokens full_num_evicted = 0 swa_num_evicted = 0 if full_num_tokens > 0: @@ -689,6 +700,9 @@ class SWARadixCache(BasePrefixCache): x = x_next self.update_eviction_metrics(full_num_evicted + swa_num_evicted, start_time) + return EvictResult( + num_tokens_evicted=full_num_evicted, swa_num_tokens_evicted=swa_num_evicted + ) def inc_lock_ref(self, node: TreeNode) -> Optional[int]: """ diff --git a/test/registered/radix_cache/test_mamba_unittest.py b/test/registered/radix_cache/test_mamba_unittest.py index 0ad099b6a..7e66dadcf 100644 --- a/test/registered/radix_cache/test_mamba_unittest.py +++ b/test/registered/radix_cache/test_mamba_unittest.py @@ -6,7 +6,11 @@ import torch from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator -from sglang.srt.mem_cache.base_prefix_cache import MatchPrefixParams +from sglang.srt.mem_cache.base_prefix_cache import ( + EvictParams, + InsertParams, + MatchPrefixParams, +) from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, HybridReqToTokenPool @@ -234,9 +238,14 @@ class TestMamba(unittest.TestCase): print( f"req1: inserting, req1_token_ids: {req1_token_ids}, req1_kv_indices: {req1_kv_indices}" ) - prefix_len = tree.insert( - RadixKey(req1_token_ids), req1_kv_indices, req1.mamba_pool_idx.unsqueeze(0) + result = tree.insert( + InsertParams( + key=RadixKey(req1_token_ids), + value=req1_kv_indices, + mamba_value=req1.mamba_pool_idx.unsqueeze(0), + ) ) + prefix_len = result.prefix_len print( f"req1: prefix_len: {prefix_len}, allocator mamba available size: {mamba_pool.available_size()}, full available size: {allocator.available_size()}" ) @@ -246,9 +255,14 @@ class TestMamba(unittest.TestCase): print( f"req2: inserting, req2_token_ids: {req2_token_ids}, req2_kv_indices: {req2_kv_indices}" ) - prefix_len = tree.insert( - RadixKey(req2_token_ids), req2_kv_indices, req2.mamba_pool_idx.unsqueeze(0) + result = tree.insert( + InsertParams( + key=RadixKey(req2_token_ids), + value=req2_kv_indices, + mamba_value=req2.mamba_pool_idx.unsqueeze(0), + ) ) + prefix_len = result.prefix_len print( f"req2: prefix_len: {prefix_len}, allocator mamba available size: {mamba_pool.available_size()}, full available size: {allocator.available_size()}" ) @@ -259,9 +273,14 @@ class TestMamba(unittest.TestCase): print( f"req3: inserting, req3_token_ids: {req3_token_ids}, req3_kv_indices: {req3_kv_indices}" ) - prefix_len = tree.insert( - RadixKey(req3_token_ids), req3_kv_indices, req3.mamba_pool_idx.unsqueeze(0) + result = tree.insert( + InsertParams( + key=RadixKey(req3_token_ids), + value=req3_kv_indices, + mamba_value=req3.mamba_pool_idx.unsqueeze(0), + ) ) + prefix_len = result.prefix_len print( f"req3: prefix_len: {prefix_len}, allocator mamba available size: {mamba_pool.available_size()}, full available size: {allocator.available_size()}" ) @@ -271,9 +290,14 @@ class TestMamba(unittest.TestCase): print( f"req4: inserting, req4_token_ids: {req4_token_ids}, req4_kv_indices: {req4_kv_indices}" ) - prefix_len = tree.insert( - RadixKey(req4_token_ids), req4_kv_indices, req4.mamba_pool_idx.unsqueeze(0) + result = tree.insert( + InsertParams( + key=RadixKey(req4_token_ids), + value=req4_kv_indices, + mamba_value=req4.mamba_pool_idx.unsqueeze(0), + ) ) + prefix_len = result.prefix_len print( f"req4: prefix_len: {prefix_len}, allocator mamba available size: {mamba_pool.available_size()}, full available size: {allocator.available_size()}" ) @@ -281,12 +305,18 @@ class TestMamba(unittest.TestCase): tree.pretty_print() full_num_tokens = 1 print(f"evicting {full_num_tokens} full token") - tree.evict(full_num_tokens=full_num_tokens) + result = tree.evict(EvictParams(num_tokens=full_num_tokens)) + assert ( + result.num_tokens_evicted >= full_num_tokens + ), f"evicted {result.num_tokens_evicted} full tokens, expected {full_num_tokens}" tree.pretty_print() mamba_num = 1 print(f"evicting {mamba_num} mamba") - tree.evict_mamba(mamba_num=mamba_num) + result = tree.evict(EvictParams(num_tokens=0, mamba_num=mamba_num)) + assert ( + result.mamba_num_evicted >= mamba_num + ), f"evicted {result.mamba_num_evicted} mamba states, expected {mamba_num}" tree.pretty_print() req5_token_ids = [1, 2, 3, 4, 5] @@ -317,7 +347,10 @@ class TestMamba(unittest.TestCase): mamba_num = 1 print(f"evicting {mamba_num} mamba") - tree.evict_mamba(mamba_num=mamba_num) + result = tree.evict(EvictParams(num_tokens=0, mamba_num=mamba_num)) + assert ( + result.mamba_num_evicted >= mamba_num + ), f"evicted {result.mamba_num_evicted} mamba states, expected {mamba_num}" tree.pretty_print() req8_token_ids = [1, 2, 3, 4, 5, 60, 70] diff --git a/test/registered/radix_cache/test_radix_cache_unit.py b/test/registered/radix_cache/test_radix_cache_unit.py index d6cac0e9f..73aa37b9f 100644 --- a/test/registered/radix_cache/test_radix_cache_unit.py +++ b/test/registered/radix_cache/test_radix_cache_unit.py @@ -31,7 +31,12 @@ import unittest.mock import torch from sglang.srt.disaggregation.kv_events import BlockRemoved, BlockStored -from sglang.srt.mem_cache.base_prefix_cache import MatchPrefixParams +from sglang.srt.mem_cache.base_prefix_cache import ( + EvictParams, + EvictResult, + InsertParams, + MatchPrefixParams, +) from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode # Test constants @@ -266,7 +271,12 @@ class TestRadixCache(unittest.TestCase): cache = RadixCache.create_simulated() # Insert some data - cache.insert(RadixKey([1, 2, 3]), torch.tensor([10, 20, 30], dtype=torch.int64)) + cache.insert( + InsertParams( + key=RadixKey([1, 2, 3]), + value=torch.tensor([10, 20, 30], dtype=torch.int64), + ) + ) self.assertGreater(cache.total_size(), 0) # Reset @@ -283,7 +293,8 @@ class TestRadixCache(unittest.TestCase): key = RadixKey([1, 2, 3]) value = torch.tensor([10, 20, 30], dtype=torch.int64) - prefix_len = cache.insert(key, value) + result = cache.insert(InsertParams(key=key, value=value)) + prefix_len = result.prefix_len if disable_cache: self.assertEqual(prefix_len, 0) @@ -311,7 +322,8 @@ class TestRadixCache(unittest.TestCase): cache = RadixCache.create_simulated() key = RadixKey([1, 2, 3]) - prefix_len = cache.insert(key, None) + result = cache.insert(InsertParams(key=key, value=None)) + prefix_len = result.prefix_len # When None is passed, it should create value from token_ids self.assertEqual(prefix_len, 0) @@ -323,10 +335,19 @@ class TestRadixCache(unittest.TestCase): self.assertEqual(cache.total_size(), 0) - cache.insert(RadixKey([1, 2, 3]), torch.tensor([10, 20, 30], dtype=torch.int64)) + cache.insert( + InsertParams( + key=RadixKey([1, 2, 3]), + value=torch.tensor([10, 20, 30], dtype=torch.int64), + ) + ) self.assertEqual(cache.total_size(), 3) - cache.insert(RadixKey([4, 5]), torch.tensor([40, 50], dtype=torch.int64)) + cache.insert( + InsertParams( + key=RadixKey([4, 5]), value=torch.tensor([40, 50], dtype=torch.int64) + ) + ) self.assertEqual(cache.total_size(), 5) def test_kv_cache_events(self): @@ -344,7 +365,7 @@ class TestRadixCache(unittest.TestCase): ) # Insert data - cache.insert(RadixKey([1, 2, 3, 4, 5]), None) + cache.insert(InsertParams(key=RadixKey([1, 2, 3, 4, 5]), value=None)) # Take events events = cache.take_events() @@ -371,8 +392,19 @@ class TestRadixCache(unittest.TestCase): ) # Insert and then evict data - cache.insert(RadixKey([1, 2, 3]), torch.tensor([10, 20, 30], dtype=torch.int64)) - cache.evict(3) + cache.insert( + InsertParams( + key=RadixKey([1, 2, 3]), + value=torch.tensor([10, 20, 30], dtype=torch.int64), + ) + ) + result = cache.evict(EvictParams(num_tokens=3)) + self.assertIsInstance(result, EvictResult) + self.assertGreaterEqual( + result.num_tokens_evicted, + 3, + f"evicted {result.num_tokens_evicted} tokens, expected at least 3", + ) # Take events - should include both store and remove events events = cache.take_events() @@ -393,13 +425,22 @@ class TestRadixCache(unittest.TestCase): # Insert same token sequence with different extra keys cache.insert( - RadixKey([1, 2, 3], "key1"), torch.tensor([10, 20, 30], dtype=torch.int64) + InsertParams( + key=RadixKey([1, 2, 3], "key1"), + value=torch.tensor([10, 20, 30], dtype=torch.int64), + ) ) cache.insert( - RadixKey([1, 2, 3], "key2"), torch.tensor([40, 50, 60], dtype=torch.int64) + InsertParams( + key=RadixKey([1, 2, 3], "key2"), + value=torch.tensor([40, 50, 60], dtype=torch.int64), + ) ) cache.insert( - RadixKey([1, 2, 3], None), torch.tensor([70, 80, 90], dtype=torch.int64) + InsertParams( + key=RadixKey([1, 2, 3], None), + value=torch.tensor([70, 80, 90], dtype=torch.int64), + ) ) # Keys with different extra_key should not match each other @@ -434,7 +475,12 @@ class TestRadixCache(unittest.TestCase): cache = RadixCache.create_simulated() # Insert sequence - cache.insert(RadixKey([1, 2, 3]), torch.tensor([10, 20, 30], dtype=torch.int64)) + cache.insert( + InsertParams( + key=RadixKey([1, 2, 3]), + value=torch.tensor([10, 20, 30], dtype=torch.int64), + ) + ) # Get node result = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3]))) @@ -461,13 +507,27 @@ class TestRadixCache(unittest.TestCase): cache = RadixCache.create_simulated(mock_allocator=mock_allocator) # Insert sequences - cache.insert(RadixKey([1, 2]), torch.tensor([10, 20], dtype=torch.int64)) - cache.insert(RadixKey([3, 4]), torch.tensor([30, 40], dtype=torch.int64)) + cache.insert( + InsertParams( + key=RadixKey([1, 2]), value=torch.tensor([10, 20], dtype=torch.int64) + ) + ) + cache.insert( + InsertParams( + key=RadixKey([3, 4]), value=torch.tensor([30, 40], dtype=torch.int64) + ) + ) initial_size = cache.total_size() # Evict some tokens - cache.evict(2) + result = cache.evict(EvictParams(num_tokens=2)) + self.assertIsInstance(result, EvictResult) + self.assertGreaterEqual( + result.num_tokens_evicted, + 2, + f"evicted {result.num_tokens_evicted} tokens, expected at least 2", + ) # Should have called free and reduced size mock_allocator.free.assert_called() @@ -486,7 +546,12 @@ class TestRadixCache(unittest.TestCase): cache = RadixCache.create_simulated(page_size=page_size) tokens = list(range(sequence_length)) - cache.insert(RadixKey(tokens), torch.tensor(tokens, dtype=torch.int64)) + cache.insert( + InsertParams( + key=RadixKey(tokens), + value=torch.tensor(tokens, dtype=torch.int64), + ) + ) result = cache.match_prefix(MatchPrefixParams(key=RadixKey(tokens))) self.assertGreater(len(result.device_indices), 0) @@ -499,7 +564,12 @@ class TestRadixCache(unittest.TestCase): """Test pretty_print produces output.""" cache = RadixCache.create_simulated() - cache.insert(RadixKey([1, 2, 3]), torch.tensor([10, 20, 30], dtype=torch.int64)) + cache.insert( + InsertParams( + key=RadixKey([1, 2, 3]), + value=torch.tensor([10, 20, 30], dtype=torch.int64), + ) + ) # Just test that it doesn't crash try: @@ -511,8 +581,16 @@ class TestRadixCache(unittest.TestCase): """Test all_values_flatten method.""" cache = RadixCache.create_simulated() - cache.insert(RadixKey([1, 2]), torch.tensor([10, 20], dtype=torch.int64)) - cache.insert(RadixKey([3, 4]), torch.tensor([30, 40], dtype=torch.int64)) + cache.insert( + InsertParams( + key=RadixKey([1, 2]), value=torch.tensor([10, 20], dtype=torch.int64) + ) + ) + cache.insert( + InsertParams( + key=RadixKey([3, 4]), value=torch.tensor([30, 40], dtype=torch.int64) + ) + ) all_values = cache.all_values_flatten() self.assertEqual(len(all_values), 4) @@ -529,12 +607,12 @@ class TestRadixCache(unittest.TestCase): # Insert a long sequence that will be split later. seq1 = [1, 2, 3, 4, 5, 6, 7, 8] val1 = torch.tensor([x * 10 for x in seq1], dtype=torch.int64) - cache.insert(RadixKey(seq1), val1) + cache.insert(InsertParams(key=RadixKey(seq1), value=val1)) # Insert a diverging branch to create an internal node on the path. seq2 = [1, 2, 9, 10] val2 = torch.tensor([x * 10 for x in seq2], dtype=torch.int64) - cache.insert(RadixKey(seq2), val2) + cache.insert(InsertParams(key=RadixKey(seq2), value=val2)) print(cache.pretty_print()) baseline_total = cache.total_size() @@ -573,7 +651,7 @@ class TestRadixCache(unittest.TestCase): ) # Insert a sequence - cache.insert(RadixKey([1, 2, 3, 4, 5, 6, 7, 8]), None) + cache.insert(InsertParams(key=RadixKey([1, 2, 3, 4, 5, 6, 7, 8]), value=None)) # Trigger event emission to compute hash_value lazily cache.take_events() @@ -599,7 +677,7 @@ class TestRadixCache(unittest.TestCase): ) # Insert a sequence with repeating token pattern: [1,2,3,4, 1,2,3,4] - cache.insert(RadixKey([1, 2, 3, 4, 1, 2, 3, 4]), None) + cache.insert(InsertParams(key=RadixKey([1, 2, 3, 4, 1, 2, 3, 4]), value=None)) events = cache.take_events() block_stored_events = [e for e in events if isinstance(e, BlockStored)] @@ -633,11 +711,11 @@ class TestRadixCache(unittest.TestCase): ) # Insert a sequence that will cause a split - cache.insert(RadixKey([1, 2, 3, 4]), None) + cache.insert(InsertParams(key=RadixKey([1, 2, 3, 4]), value=None)) cache.take_events() # Clear events and compute hash_value for first node # Insert a diverging sequence that will cause a split at page boundary - cache.insert(RadixKey([1, 2, 5, 6]), None) + cache.insert(InsertParams(key=RadixKey([1, 2, 5, 6]), value=None)) cache.take_events() # Trigger event emission to compute hash_value # Find the split node @@ -674,7 +752,7 @@ class TestRadixCache(unittest.TestCase): cache: RadixCache = RadixCache.create_simulated() for key, value in zip(keys, values): - cache.insert(RadixKey(key), value) + cache.insert(InsertParams(key=RadixKey(key), value=value)) del values diff --git a/test/registered/radix_cache/test_swa_unittest.py b/test/registered/radix_cache/test_swa_unittest.py index deb0bd2c0..24c5615de 100644 --- a/test/registered/radix_cache/test_swa_unittest.py +++ b/test/registered/radix_cache/test_swa_unittest.py @@ -2,7 +2,12 @@ import unittest import torch -from sglang.srt.mem_cache.base_prefix_cache import MatchPrefixParams +from sglang.srt.mem_cache.base_prefix_cache import ( + EvictParams, + EvictResult, + InsertParams, + MatchPrefixParams, +) from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.mem_cache.radix_cache import RadixKey @@ -140,7 +145,10 @@ class TestSWA(unittest.TestCase): print( f"req1: inserting, req1_token_ids: {req1_token_ids}, req1_kv_indices: {req1_kv_indices}" ) - prefix_len = tree.insert(RadixKey(req1_token_ids), req1_kv_indices) + result = tree.insert( + InsertParams(key=RadixKey(req1_token_ids), value=req1_kv_indices) + ) + prefix_len = result.prefix_len print( f"req1: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}" ) @@ -149,7 +157,10 @@ class TestSWA(unittest.TestCase): print( f"req2: inserting, req2_token_ids: {req2_token_ids}, req2_kv_indices: {req2_kv_indices}" ) - prefix_len = tree.insert(RadixKey(req2_token_ids), req2_kv_indices) + result = tree.insert( + InsertParams(key=RadixKey(req2_token_ids), value=req2_kv_indices) + ) + prefix_len = result.prefix_len print( f"req2: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}" ) @@ -158,7 +169,10 @@ class TestSWA(unittest.TestCase): print( f"req3: inserting, req3_token_ids: {req3_token_ids}, req3_kv_indices: {req3_kv_indices}" ) - prefix_len = tree.insert(RadixKey(req3_token_ids), req3_kv_indices) + result = tree.insert( + InsertParams(key=RadixKey(req3_token_ids), value=req3_kv_indices) + ) + prefix_len = result.prefix_len print( f"req3: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}" ) @@ -167,7 +181,10 @@ class TestSWA(unittest.TestCase): print( f"req4: inserting, req4_token_ids: {req4_token_ids}, req4_kv_indices: {req4_kv_indices}" ) - prefix_len = tree.insert(RadixKey(req4_token_ids), req4_kv_indices) + result = tree.insert( + InsertParams(key=RadixKey(req4_token_ids), value=req4_kv_indices) + ) + prefix_len = result.prefix_len print( f"req4: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}" ) @@ -175,17 +192,23 @@ class TestSWA(unittest.TestCase): tree.pretty_print() full_num_tokens, swa_num_tokens = 1, 0 print(f"evicting {full_num_tokens} full token and {swa_num_tokens} swa token") - tree.evict(full_num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) + tree.evict( + EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) + ) tree.pretty_print() full_num_tokens, swa_num_tokens = 0, 1 print(f"evicting {full_num_tokens} full token and {swa_num_tokens} swa token") - tree.evict(full_num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) + tree.evict( + EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) + ) tree.pretty_print() full_num_tokens, swa_num_tokens = 1, 2 print(f"evicting {full_num_tokens} full token and {swa_num_tokens} swa token") - tree.evict(full_num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) + tree.evict( + EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) + ) tree.pretty_print() req5_token_ids = [1, 2, 3, 4, 5] @@ -277,7 +300,10 @@ class TestSWA(unittest.TestCase): print( f"req1: inserting, req1_token_ids: {req1_token_ids}, req1_kv_indices: {req1_kv_indices}" ) - prefix_len = tree.insert(RadixKey(req1_token_ids), req1_kv_indices) + result = tree.insert( + InsertParams(key=RadixKey(req1_token_ids), value=req1_kv_indices) + ) + prefix_len = result.prefix_len self.assertEqual(prefix_len, 0) print( f"req1: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}" @@ -287,7 +313,10 @@ class TestSWA(unittest.TestCase): print( f"req2: inserting, req2_token_ids: {req2_token_ids}, req2_kv_indices: {req2_kv_indices}" ) - prefix_len = tree.insert(RadixKey(req2_token_ids), req2_kv_indices) + result = tree.insert( + InsertParams(key=RadixKey(req2_token_ids), value=req2_kv_indices) + ) + prefix_len = result.prefix_len self.assertEqual(prefix_len, 2) print( f"req2: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}" @@ -297,7 +326,10 @@ class TestSWA(unittest.TestCase): print( f"req3: inserting, req3_token_ids: {req3_token_ids}, req3_kv_indices: {req3_kv_indices}" ) - prefix_len = tree.insert(RadixKey(req3_token_ids), req3_kv_indices) + result = tree.insert( + InsertParams(key=RadixKey(req3_token_ids), value=req3_kv_indices) + ) + prefix_len = result.prefix_len self.assertEqual(prefix_len, 0) print( f"req3: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}" @@ -307,7 +339,10 @@ class TestSWA(unittest.TestCase): print( f"req4: inserting, req4_token_ids: {req4_token_ids}, req4_kv_indices: {req4_kv_indices}" ) - prefix_len = tree.insert(RadixKey(req4_token_ids), req4_kv_indices) + result = tree.insert( + InsertParams(key=RadixKey(req4_token_ids), value=req4_kv_indices) + ) + prefix_len = result.prefix_len self.assertEqual(prefix_len, 4) print( f"req4: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}" @@ -316,17 +351,41 @@ class TestSWA(unittest.TestCase): tree.pretty_print() full_num_tokens, swa_num_tokens = 1, 0 print(f"evicting {full_num_tokens} full token and {swa_num_tokens} swa token") - tree.evict(full_num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) + evict_result = tree.evict( + EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) + ) + assert isinstance(evict_result, EvictResult) + assert ( + evict_result.num_tokens_evicted >= full_num_tokens + ) # May evict more due to node granularity + print( + f"evicted {evict_result.num_tokens_evicted} full tokens, {evict_result.swa_num_tokens_evicted} swa tokens" + ) tree.pretty_print() full_num_tokens, swa_num_tokens = 0, 1 print(f"evicting {full_num_tokens} full token and {swa_num_tokens} swa token") - tree.evict(full_num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) + evict_result = tree.evict( + EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) + ) + assert isinstance(evict_result, EvictResult) + assert ( + evict_result.swa_num_tokens_evicted >= swa_num_tokens + ), f"evicted {evict_result.swa_num_tokens_evicted} swa tokens, expected {swa_num_tokens}" tree.pretty_print() full_num_tokens, swa_num_tokens = 1, 2 print(f"evicting {full_num_tokens} full token and {swa_num_tokens} swa token") - tree.evict(full_num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) + evict_result = tree.evict( + EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) + ) + assert isinstance(evict_result, EvictResult) + assert ( + evict_result.num_tokens_evicted >= full_num_tokens + ), f"evicted {evict_result.num_tokens_evicted} full tokens, expected {full_num_tokens}" + assert ( + evict_result.swa_num_tokens_evicted >= swa_num_tokens + ), f"evicted {evict_result.swa_num_tokens_evicted} swa tokens, expected {swa_num_tokens}" tree.pretty_print() req5_token_ids = [1, 2, 3, 4, 5]