[RadixTree][3/N Refactor]：Support unified insert/evict params (#17401)

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),
+                        InsertParams(
-                        torch.empty(len(prefix_ids), dtype=torch.bool),
+                            key=RadixKey(token_ids=prefix_ids, extra_key=extra_key),
+                            value=torch.empty(len(prefix_ids), dtype=torch.bool),
+                        )
                     )
         return temporary_deprioritized
 

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

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):
+    def evict(self, params: EvictParams) -> EvictResult:
-        pass
+        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):
+    def evict(self, params: EvictParams) -> EvictResult:
-        pass
+        return EvictResult()

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)

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(
+    def insert(self, params: InsertParams) -> InsertResult:
-        self,
+        key = params.key
-        key: RadixKey,
+        value = params.value
-        value=None,
+        chunked = params.chunked
-        chunked: bool = False,
+        priority = params.priority
-        priority: int | None = None,
+
-    ):
         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):

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(
+            result = self.insert(
-                RadixKey(token_ids[:page_aligned_len], req.extra_key),
+                InsertParams(
-                page_aligned_kv_indices,
+                    key=RadixKey(token_ids[:page_aligned_len], req.extra_key),
-                mamba_value,
+                    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(
+        result = self.insert(
-            RadixKey(page_aligned_token_ids, req.extra_key),
+            InsertParams(
-            page_aligned_kv_indices,
+                key=RadixKey(page_aligned_token_ids, req.extra_key),
-            mamba_value_forked,
+                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.

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(
+        result = self.insert(
-            radix_key,
+            InsertParams(
-            values,
+                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(InsertParams(key=RadixKey(token_ids=[1, 2, 3], extra_key=None)))
-    tree.insert(RadixKey(token_ids=[1, 2, 3], extra_key=None))
+    tree.insert(InsertParams(key=RadixKey(token_ids=[1, 2, 3], extra_key=None)))
-    tree.insert(RadixKey(token_ids=[1, 2, 4, 5], extra_key=None))
+    tree.insert(InsertParams(key=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(
-    tree.insert(RadixKey(token_ids=[8, 9, 10, 11, 12], extra_key=None))
+        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(

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_evicted, start_time)
-        self.update_eviction_metrics(num_tokens, start_time)
+        return EvictResult(num_tokens_evicted=num_evicted)
 
     def evictable_size(self):
         return self.tree.evictable_size()

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

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(
+    def insert(self, params: InsertParams) -> InsertResult:
-        self,
-        key: RadixKey,
-        value=None,
-        prev_prefix_len: int = 0,
-        swa_evicted_seqlen: int = 0,
-    ) -> int:
         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),
+                InsertParams(
-                page_aligned_kv_indices,
+                    key=RadixKey(token_ids[:page_aligned_token_len], req.extra_key),
-                old_prefix_len,
+                    value=page_aligned_kv_indices,
-                req.swa_evicted_seqlen,
+                    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(
+        result = self.insert(
-            RadixKey(page_aligned_token_ids, req.extra_key),
+            InsertParams(
-            page_aligned_kv_indices,
+                key=RadixKey(page_aligned_token_ids, req.extra_key),
-            old_prefix_len,
+                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]:
         """

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(
+        result = tree.insert(
-            RadixKey(req1_token_ids), req1_kv_indices, req1.mamba_pool_idx.unsqueeze(0)
+            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(
+        result = tree.insert(
-            RadixKey(req2_token_ids), req2_kv_indices, req2.mamba_pool_idx.unsqueeze(0)
+            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(
+        result = tree.insert(
-            RadixKey(req3_token_ids), req3_kv_indices, req3.mamba_pool_idx.unsqueeze(0)
+            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(
+        result = tree.insert(
-            RadixKey(req4_token_ids), req4_kv_indices, req4.mamba_pool_idx.unsqueeze(0)
+            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]

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.insert(
-        cache.evict(3)
+            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(
-        cache.insert(RadixKey([3, 4]), torch.tensor([30, 40], dtype=torch.int64))
+            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(
-        cache.insert(RadixKey([3, 4]), torch.tensor([30, 40], dtype=torch.int64))
+            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
 

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]