From a8a4d86be9f483fabd097350b06b6ae6e6905874 Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Sun, 20 Sep 2026 16:16:27 +0800 Subject: [PATCH] Remove swa and mamba radix cache (#40313) --- python/sglang/srt/arg_groups/mamba_hook.py | 10 +- .../arg_groups/model_overrides/qwen4_exp.py | 2 +- python/sglang/srt/environ.py | 3 - .../srt/kv_canary/radix_cache_walker.py | 10 +- .../scheduler_components/invariant_checker.py | 4 +- python/sglang/srt/mem_cache/README.md | 2 +- python/sglang/srt/mem_cache/common.py | 4 +- .../sglang/srt/mem_cache/mamba_radix_cache.py | 1445 ---------------- .../srt/mem_cache/pure_swa_radix_cache.py | 4 +- python/sglang/srt/mem_cache/qsa_kv_pool.py | 4 +- .../sglang/srt/mem_cache/swa_radix_cache.py | 1462 ----------------- .../unified_cache/components/README.md | 2 +- .../unified_cache/components/mamba.py | 3 +- python/sglang/srt/utils/common.py | 2 +- .../test/scripted_runtime/context/radix.py | 3 - .../dsv4/test_dsv4_swa_radix_retract.py | 8 +- .../kv_canary/test_self_unit_radix_walker.py | 60 - .../unit/mem_cache/test_mamba_unittest.py | 435 +---- .../unit/mem_cache/test_radix_cache_unit.py | 3 +- .../mem_cache/test_swa_eviction_boundary.py | 550 ------- .../test_swa_lock_release_lifecycle.py | 440 ----- .../unit/mem_cache/test_swa_unittest.py | 743 +-------- .../test_unified_radix_cache_bench.py | 52 +- .../test_unified_radix_cache_unittest.py | 2 +- 24 files changed, 69 insertions(+), 5184 deletions(-) delete mode 100644 python/sglang/srt/mem_cache/mamba_radix_cache.py delete mode 100644 python/sglang/srt/mem_cache/swa_radix_cache.py delete mode 100644 test/registered/unit/mem_cache/test_swa_eviction_boundary.py delete mode 100644 test/registered/unit/mem_cache/test_swa_lock_release_lifecycle.py diff --git a/python/sglang/srt/arg_groups/mamba_hook.py b/python/sglang/srt/arg_groups/mamba_hook.py index 9107f432e..3310bf039 100644 --- a/python/sglang/srt/arg_groups/mamba_hook.py +++ b/python/sglang/srt/arg_groups/mamba_hook.py @@ -72,12 +72,10 @@ def handle_mamba_backend(server_args: Any): def handle_int8_mamba_checkpoint(server_args: Any): - # The int8 mamba checkpoint pool is only wired into the built-in - # MambaRadixCache. The host-offload path (enabled by - # --enable-hierarchical-cache) and custom radix-cache backends are NOT - # int8-aware: they would read int8 checkpoint slots as bf16 active slots - # (wrong pool / out-of-range). Reject the combination up front rather than - # silently corrupting state. + # The host-offload path (enabled by --enable-hierarchical-cache) and + # custom radix-cache backends are NOT int8-aware: they would read int8 + # checkpoint slots as bf16 active slots (wrong pool / out-of-range). + # Reject the combination up front rather than silently corrupting state. cfg = resolving_view(server_args) if not cfg.enable_int8_mamba_checkpoint: return diff --git a/python/sglang/srt/arg_groups/model_overrides/qwen4_exp.py b/python/sglang/srt/arg_groups/model_overrides/qwen4_exp.py index 2c96fc22f..a35b0770f 100644 --- a/python/sglang/srt/arg_groups/model_overrides/qwen4_exp.py +++ b/python/sglang/srt/arg_groups/model_overrides/qwen4_exp.py @@ -25,7 +25,7 @@ def _qwen4_exp_overrides(server_args: Any, hf_config: Any) -> dict: """Compressed QSA must own ``page_size`` here, so the qwen3_5 hybrid attention-shape policy is restated rather than shared. page_size=64 needs page-aligned full-KV allocation (slots are full_slot // ratio), - which MambaRadixCache allows only with mamba extra-buffer or --disable-radix-cache. + which in turn needs the mamba extra-buffer strategy or --disable-radix-cache. """ cfg = resolving_view(server_args) if ( diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 61e39fb36..19d91d52e 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -678,9 +678,6 @@ class Envs: SGLANG_ENABLE_UNIFIED_RADIX_TREE = EnvBool(False) # Registered TreeCore backend serving the unified radix cache. SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND = EnvStr("python") - # TODO(DSV4): @ispobock this has bug on main branch when retract - SGLANG_OPT_SWA_RADIX_CACHE_COMPACT = EnvBool(False) - SGLANG_OPT_SWA_SPLIT_LEAF_ON_INSERT = EnvBool(False) SGLANG_OPT_SWA_RELEASE_LEAF_LOCK_AFTER_WINDOW = EnvBool(False) # =================================================================== diff --git a/python/sglang/srt/kv_canary/radix_cache_walker.py b/python/sglang/srt/kv_canary/radix_cache_walker.py index 718e67e58..b65b4c3d1 100644 --- a/python/sglang/srt/kv_canary/radix_cache_walker.py +++ b/python/sglang/srt/kv_canary/radix_cache_walker.py @@ -5,7 +5,6 @@ from typing import TYPE_CHECKING, Any import torch from sglang.srt.mem_cache.radix_cache import RadixCache -from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import ( RadixCacheWalkResult, ) @@ -34,7 +33,7 @@ def walk_radix_cache_for_canary( return radix_cache.tree_core.walk_for_kv_canary( unlocked_only=unlocked_only, swa_resident_only=swa_resident_only ) - if cache_type is not RadixCache and cache_type is not SWARadixCache: + if cache_type is not RadixCache: raise NotImplementedError( f"walk_radix_cache_for_canary does not support {cache_type.__name__}" ) @@ -133,9 +132,6 @@ def _node_is_unlocked_for_canary( if type(radix_cache) is RadixCache: return node.lock_ref == 0 - if type(radix_cache) is SWARadixCache: - return node.full_lock_ref == 0 - raise NotImplementedError( f"walk_radix_cache_for_canary does not support {type(radix_cache).__name__}" ) @@ -146,7 +142,5 @@ def _node_is_swa_resident_for_canary( node: TreeNode, radix_cache: BasePrefixCache, ) -> bool: - if type(radix_cache) is SWARadixCache: - return not node.swa_tombstone - + # RadixCache has no SWA tier, so every node it holds is resident. return True diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index 16a7704a3..b10f18d5e 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -99,8 +99,8 @@ class SchedulerInvariantChecker: session_held = self.pool_stats_observer.session_held_full_tokens() total = ps.full_capacity elif self.is_hybrid_ssm: - # Branch on cache type for the protected accessor (MambaRadixCache - # splits full/mamba; ChunkCache only has the single protected_size). + # Branch on cache type for the protected accessor (a mamba-capable + # cache splits full/mamba; ChunkCache only has the single protected_size). # Use the allocator's `.size` for `total`: static max_total_num_tokens for # non-unified pools, the dynamic byte-coordinated cap (matching # `available_size`) for the unified pool. diff --git a/python/sglang/srt/mem_cache/README.md b/python/sglang/srt/mem_cache/README.md index 310a06dda..ea2ff10d4 100644 --- a/python/sglang/srt/mem_cache/README.md +++ b/python/sglang/srt/mem_cache/README.md @@ -38,7 +38,7 @@ to keep. The layout is specified in Two groups sit outside that stack: - **Radix cache** is its own axis. The per-model variants (`radix_cache.py`, - `swa_radix_cache.py`, `mamba_radix_cache.py`, `hiradix_cache.py`, `chunk_cache.py`) + `hiradix_cache.py`, `chunk_cache.py`) are converging onto the **Unified Radix Cache** (`unified_cache/`, [#20415](https://github.com/sgl-project/sglang/issues/20415)), whose Full/SWA/Mamba component model is documented in diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index fda235acf..9640071cc 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -271,10 +271,10 @@ def retraction_discard(req: Req, tree_cache: BasePrefixCache, backend: str) -> N def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = True): assert (not req.kv.holds_kv) == req.kv.is_kv_released - # MambaRadixCache may alloc mamba state before alloc KV cache + # A mamba-capable cache may alloc mamba state before alloc KV cache if not req.kv.holds_kv: assert tree_cache.supports_mamba(), ( - "Only MambaRadixCache allow freeing before alloc" + "Only a mamba-capable tree cache allows freeing before alloc" ) # TODO (csy, hanming): clean up this early allocation logic if req.kv.holds_mamba: diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py deleted file mode 100644 index e2faa6efe..000000000 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ /dev/null @@ -1,1445 +0,0 @@ -from __future__ import annotations - -""" -Copyright 2023-2024 SGLang Team -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -""" - -""" -The radix tree data structure for managing the hybrid (full and Mamba) KV cache. -""" - -import os -from array import array -from collections import defaultdict -from typing import TYPE_CHECKING, List, Optional, Tuple - -import torch -from numpy import float64 - -from sglang.srt.mem_cache.allocator import ( - PagedTokenToKVPoolAllocator, - TokenToKVPoolAllocator, -) -from sglang.srt.mem_cache.allocator.unified_mamba import ( - UnifiedMambaTokenToKVPoolAllocator, -) -from sglang.srt.mem_cache.base_prefix_cache import ( - BasePrefixCache, - DecLockRefParams, - DecLockRefResult, - EvictParams, - EvictResult, - IncLockRefResult, - InsertParams, - InsertResult, - MatchPrefixParams, - MatchResult, -) -from sglang.srt.mem_cache.events import KVCacheEventRecorder -from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool -from sglang.srt.mem_cache.radix_cache import RadixKey -from sglang.srt.mem_cache.utils import split_node_hash_value -from sglang.srt.runtime_context import ( - get_parallel, - mamba_cache_chunk_size, -) - -if TYPE_CHECKING: - from sglang.srt.managers.schedule_batch import Req - from sglang.srt.mem_cache.cache_init_params import CacheInitParams - -import logging - -logger = logging.getLogger(__name__) - -# Debug-only invariant checks in the Mamba slot-donation path call tensor.item(), -# which forces a per-request cudaStreamSynchronize on the scheduler thread. Under -# load this can serialize/stall the scheduler. Gate them off by default; set -# SGLANG_MAMBA_DEBUG_ASSERTS=1 to re-enable for debugging. -_MAMBA_DEBUG_ASSERTS = os.environ.get("SGLANG_MAMBA_DEBUG_ASSERTS", "0") == "1" - - -class TreeNode: - counter = 0 - last_access_time_counter_float = float64(1.0) - - def __init__(self, id: Optional[int] = None): - self.children = defaultdict(TreeNode) - self.parent: TreeNode = None - self.key: RadixKey = None - self.value: Optional[torch.Tensor] = None - self.mamba_value: Optional[torch.Tensor] = None - self.mamba_host_value: Optional[torch.Tensor] = None - # invariant: for any node, if mamba_lock_ref is locked, full_lock_ref must be locked; - # if full_lock_ref is locked, mamba_lock_ref doesn't need to be locked. So, - # full_lock_ref is always >= mamba_lock_ref. - # for full_lock, once it is locked, its parent must be locked as well - # for mamba_lock, it only need lock node itself - self.full_lock_ref = 0 - self.mamba_lock_ref = 0 - # last access time is only used for sanity check. LRU is maintained by the lru list. - # `last_access_time` tracks the full LRU (whole matched path is reused as prefix); - # `mamba_last_access_time` tracks the mamba LRU, which only touches the single state - # actually consumed per access, so the two orders diverge and need separate stamps. - self.last_access_time = get_last_access_time() - self.mamba_last_access_time = self.last_access_time - - self.hit_count = 0 - self.host_ref_counter = 0 - self.host_mamba_ref_counter = 0 - # store the host indices of KV cache - self.host_value = None - # store hash values of each pages - self.hash_value: Optional[List[str]] = None - # Namespace-aware hashes used only for external KV events. - self.event_hash_value: Optional[List[str]] = None - - # for lru list, invariant: - # 1. prev has greater last_access_time - # 2. next has smaller last_access_time - self.prev = None - self.next = None - self.mamba_prev = None - self.mamba_next = None - self.host_mamba_prev = None - self.host_mamba_next = None - - self.id = TreeNode.counter if id is None else id - TreeNode.counter += 1 - - @property - def evicted(self): - return self.value is None - - @property - def mamba_evicted(self): - return self.mamba_value is None - - @property - def backuped(self): - return self.host_value is not None - - @property - def mamba_backuped(self): - return self.mamba_host_value is not None - - def protect_host(self): - """Protect the host KV value from eviction.""" - self.host_ref_counter += 1 - - def release_host(self): - """Release the host KV value, allowing it to be evicted.""" - if self.host_ref_counter > 0: - self.host_ref_counter -= 1 - else: - raise RuntimeError("Host reference counter is already zero.") - - def protect_host_mamba(self): - """Protect the host mamba value from eviction.""" - self.host_mamba_ref_counter += 1 - - def release_host_mamba(self): - """Release the host mamba value, allowing it to be evicted.""" - if self.host_mamba_ref_counter > 0: - self.host_mamba_ref_counter -= 1 - else: - raise RuntimeError("Host mamba reference counter is already zero.") - - def get_last_hash_value(self) -> Optional[str]: - """Returns the hash value of the last page in this node.""" - if self.hash_value is None or len(self.hash_value) == 0: - return None - return self.hash_value[-1] - - def get_prefix_hash_values(self, node: TreeNode) -> List[str]: - chunks = [] - while node is not None and node.hash_value is not None: - chunks.append(node.hash_value) - node = node.parent - return [value for chunk in reversed(chunks) for value in chunk] - - def __lt__(self, other: TreeNode): - return self.last_access_time < other.last_access_time - - -def get_last_access_time() -> float64: - ret = TreeNode.last_access_time_counter_float - TreeNode.last_access_time_counter_float += 1.0 - return ret - - -class LRUList: - def __init__(self, mamba: bool = False): - self.mamba = mamba - if self.mamba: - self.prv = "mamba_prev" - self.nxt = "mamba_next" - self.lock_ref = "mamba_lock_ref" - self.time_attr = "mamba_last_access_time" - else: - self.prv = "prev" - self.nxt = "next" - self.lock_ref = "full_lock_ref" - self.time_attr = "last_access_time" - # Initialize dummy head and tail nodes - self.head = TreeNode() # Most recently used side - self.tail = TreeNode() # Least recently used side - setattr(self.head, self.nxt, self.tail) # self.head.next = self.tail - setattr(self.tail, self.prv, self.head) # self.tail.prev = self.head - self.cache = {} - - def _add_node(self, node): - """Helper to add node right after head (most recently used)""" - self._add_node_after(self.head, node) - - def _add_node_after(self, old_node, new_node): - """Helper to add node right after old_node""" - setattr(new_node, self.prv, old_node) # new_node.prev = old_node - setattr( - new_node, self.nxt, getattr(old_node, self.nxt) - ) # new_node.next = old_node.next - setattr( - getattr(old_node, self.nxt), self.prv, new_node - ) # old_node.next.prev = new_node - setattr(old_node, self.nxt, new_node) # old_node.next = new_node - - def _remove_node(self, node): - """Helper to remove node from linked list""" - setattr( - getattr(node, self.prv), self.nxt, getattr(node, self.nxt) - ) # node.prev.next = node.next - setattr( - getattr(node, self.nxt), self.prv, getattr(node, self.prv) - ) # node.next.prev = node.prev - # Clear self pointers to break reference cycles among evicted nodes. - setattr(node, self.prv, None) - setattr(node, self.nxt, None) - - def _get_lru(self) -> Optional[TreeNode]: - """ - Get the least recently used node - """ - if len(self.cache) == 0: - return None - return getattr(self.tail, self.prv) - - def reset_node_mru(self, node): - """ - Move a (existing) node to most recently used position - """ - assert node.id in self.cache, f"Resetting node {node.id=} not in lru list" - assert not self.mamba or node.mamba_value is not None, ( - f"Resetting mamba tombstone node in mamba lru list: {node.id=}" - ) - if self.mamba: - node.mamba_last_access_time = get_last_access_time() - self._remove_node(node) - self._add_node(node) - - def reset_node_and_parents_mru(self, node, root_node): - """ - Move an (existing) node and its parents to most recently used position. Child node is - more recently used than parent node. - """ - prev_node = self.head - while node != root_node: - if not self.mamba or node.mamba_value is not None: - assert node.id in self.cache, ( - f"Resetting node {node.id=} not in lru list when resetting node and parents mru" - ) - self._remove_node(node) - self._add_node_after(prev_node, node) - prev_node = node - node = node.parent - - def insert_mru(self, node): - """ - Insert a (new) node as most recently used - """ - assert not self.mamba or node.mamba_value is not None, ( - f"Inserting mamba tombstone node in mamba lru list: {node.id=}" - ) - assert node.id not in self.cache, ( - f"Inserting node {node.id=} already in lru list, existing node: {self.cache[node.id].id=}" - ) - if self.mamba: - node.mamba_last_access_time = get_last_access_time() - self.cache[node.id] = node - self._add_node(node) - - def remove_node(self, node: TreeNode): - """ - Remove node from lru list - """ - assert node.id in self.cache, f"Removing node {node.id=} not in lru list" - assert not self.mamba or node.mamba_value is not None, ( - f"Removing mamba tombstone node from mamba lru list: {node.id=}" - ) - del self.cache[node.id] - self._remove_node(node) - - def get_lru_no_lock(self) -> Optional[TreeNode]: - """ - Get the least recently used node that is not locked - """ - return self.get_prev_no_lock(self.tail, check_id=False) - - def get_leaf_lru_no_lock(self) -> Optional[TreeNode]: - """ - Get the least recently used leaf node that is not locked - """ - return self.get_prev_leaf_no_lock(self.tail, check_id=False) - - def get_prev_no_lock( - self, node: TreeNode, check_id: bool = True - ) -> Optional[TreeNode]: - """ - Get the previous (i.e. more recently used) node that is not locked - """ - if check_id: - assert node.id in self.cache, ( - f"Getting prev of node {node.id=} not in lru list" - ) - x = getattr(node, self.prv) # x = node.prev - while getattr(x, self.lock_ref) > 0: - x = getattr(x, self.prv) # x = x.prev - # if x is the head, it means there is no node in the lru list without lock - if x == self.head: - return None - return x - - def get_prev_leaf_no_lock(self, node: TreeNode, check_id: bool = True): - """ - Get the previous (i.e. more recently used) leaf node that is not locked - """ - if check_id: - assert node.id in self.cache, ( - f"Getting prev of node {node.id=} not in lru list" - ) - x = getattr(node, self.prv) # x = node.prev - while getattr(x, self.lock_ref) > 0 or len(x.children) > 0: - x = getattr(x, self.prv) # x = x.prev - # if x is the head, it means there is no leaf node in the lru list without lock - if x == self.head: - return None - return x - - def in_list(self, node: Optional[TreeNode]): - """ - Check if the node is in the lru list - """ - if not node: - return False - return node.id in self.cache - - def pretty_print(self, tree_cache: Optional[MambaRadixCache] = None): - """ - Pretty print the lru list - """ - msg = f"{self.mamba=} LRU list: " - x_lru = self._get_lru() - while x_lru is not None and x_lru.id in self.cache: - msg += f"[{x_lru.id}] {getattr(x_lru, self.time_attr):f} -> " - x_lru = getattr(x_lru, self.prv) - print(msg) - - if not tree_cache: - return - msg = f"{self.mamba=} Nodes (sorted by {self.time_attr}): " - if self.mamba: - nodes = tree_cache._collect_nontombstone_nodes() - else: - nodes = tree_cache._collect_all_nodes() - nodes.sort(key=lambda n: getattr(n, self.time_attr)) - for x in nodes: - msg += f"[{x.id}] {getattr(x, self.time_attr):f} -> " - print(msg) - - # Note: this is expensive, only use for debug - def sanity_check_evictable_size(self): - """ - Check the evictable size (i.e. the size of the nodes that are not locked) - """ - node = self.get_lru_no_lock() - evictable_size = 0 - while self.in_list(node): - evictable_size += ( - len(node.value) if not self.mamba else len(node.mamba_value) - ) - node = self.get_prev_no_lock(node) - return evictable_size - - # Note: this is expensive, only use for debug or idle check - def sanity_check(self, tree_cache: MambaRadixCache): - """ - Check the lru list is valid by rebuilding it from the tree, sorting by this list's - access-time stamp, and checking the order matches the linked list. - """ - try: - if self.mamba: - nodes = tree_cache._collect_nontombstone_nodes() - else: - nodes = tree_cache._collect_all_nodes() - total_nodes = len(nodes) - total_lru = len(self.cache) - # rebuild expected order from this list's own access-time stamp (full and mamba - # lists have independent recency, so they use different stamps) - nodes.sort(key=lambda n: getattr(n, self.time_attr)) - # the root node is not in the lru list - assert len(nodes) == (total_lru + (0 if self.mamba else 1)), ( - f"len(nodes): {len(nodes)}, total_lru: {total_lru}" - ) - - x_lru = self._get_lru() - for x in nodes: - if x == tree_cache.root_node: - # root node is not in the lru list - continue - assert x_lru is not None and x_lru.id in self.cache, ( - f"Incorrect LRU list, x_lru is None or not in cache: {x_lru=}, {x.id=}" - ) - - assert x == x_lru, ( - f"Incorrect LRU list, {self.mamba=}, x: {x.id=} != x_lru: {x_lru.id=}, {getattr(x, self.time_attr)=}, {getattr(x_lru, self.time_attr)=}" - ) - assert x_lru.full_lock_ref == 0, ( - f"x_lru should not be locked when idle, {x_lru.full_lock_ref=}, {x_lru.id=}" - ) - assert x_lru.mamba_lock_ref == 0, ( - f"x_lru should not be locked when idle, {x_lru.mamba_lock_ref=}, {x_lru.id=}" - ) - x_lru = getattr(x, self.prv) - - if self.mamba: - evictable_size = tree_cache.mamba_evictable_size() - lru_list_evictable_size = self.sanity_check_evictable_size() - else: - evictable_size = tree_cache.full_evictable_size() - lru_list_evictable_size = self.sanity_check_evictable_size() - - assert evictable_size == lru_list_evictable_size, ( - f"{self.mamba=}, total nodes: {total_nodes}, total lru: {total_lru}, evictable size: {evictable_size} != lru list evictable size: {lru_list_evictable_size}" - ) - except Exception as e: - if get_parallel().tp_rank == 0: - msg = f"Mamba Radix tree sanity check failed, ping @yizhang2077: {e}" - logger.error(msg) - tree_cache.pretty_print() - tree_cache.full_lru_list.pretty_print(tree_cache) - tree_cache.mamba_lru_list.pretty_print(tree_cache) - raise Exception(msg) - - -class MambaRadixCache(BasePrefixCache): - def __init__(self, params: CacheInitParams): - assert ( - isinstance(params.token_to_kv_pool_allocator, TokenToKVPoolAllocator) - or isinstance( - params.token_to_kv_pool_allocator, PagedTokenToKVPoolAllocator - ) - or isinstance( - params.token_to_kv_pool_allocator, UnifiedMambaTokenToKVPoolAllocator - ) - ) - self.req_to_token_pool: HybridReqToTokenPool = params.req_to_token_pool - self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator - self.mamba_cache_chunk_size = mamba_cache_chunk_size() - - self.page_size = params.page_size - self.disable = params.disable - self.enable_mamba_extra_buffer = params.enable_mamba_extra_buffer - self.enable_mamba_extra_buffer_lazy = params.enable_mamba_extra_buffer_lazy - self.kv_events = KVCacheEventRecorder( - enabled=params.enable_kv_cache_events, page_size=self.page_size - ) - - if not self.enable_mamba_extra_buffer: - assert self.page_size == 1, ( - f"Page size must be 1 for MambaRadixCache v1, got {self.page_size}" - ) - - if self.token_to_kv_pool_allocator: - self.device = self.token_to_kv_pool_allocator.device - else: - self.device = torch.device("cpu") - - if params.enable_metrics: - self.init_metrics_collector() - - self.reset() - - ##### Public API ##### - - def supports_mamba(self) -> bool: - return True - - def reset(self) -> None: - self.root_node = TreeNode() - self.root_node.key = RadixKey(array("q"), None) - self.root_node.value = [] - self.root_node.hash_value = [] - self.root_node.full_lock_ref = 1 - self.root_node.mamba_lock_ref = 1 - self.full_evictable_size_ = 0 - self.mamba_evictable_size_ = 0 - self.full_protected_size_ = 0 - self.mamba_protected_size_ = 0 - # LRU lists are used to maintain the order of eviction of the nodes in the tree - self.full_lru_list = LRUList(mamba=False) - self.mamba_lru_list = LRUList(mamba=True) - self.kv_events.record_all_cleared() - - def match_prefix(self, params: MatchPrefixParams) -> MatchResult: - """Find the matching prefix from the radix tree. - Args: - params: MatchPrefixParams containing key and optional Mamba-specific parameters. - Returns: - A tuple of a tensor of matching prefix token IDs and - the last node that contains the prefix values. Note that - this API can modify the internal state of the Radix tree. - The last node create a new child if the prefix is shorter - than the last node's value. - """ - key = self._match_pre_processor(params) - if key is None: - return MatchResult( - device_indices=torch.empty( - (0,), - dtype=torch.int64, - device=self.device, - ), - last_device_node=self.root_node, - last_host_node=self.root_node, - best_match_node=self.root_node, - ) - - value, last_node, best_value_len = self._match_prefix_helper(key) - return self._match_post_processor(params, value, last_node, best_value_len) - - def insert(self, params: InsertParams) -> InsertResult: - if self.disable: - return InsertResult(prefix_len=0, mamba_exist=False) - - key = params.key - value = params.value - mamba_value = params.mamba_value - prev_prefix_len = params.prev_prefix_len - - if value is None: - value = torch.tensor([x for x in key.raw_token_ids()], dtype=torch.int64) - prefix_len, mamba_exist = self._insert_helper( - self.root_node, key, value, mamba_value, params.chunked, prev_prefix_len - ) - return InsertResult(prefix_len=prefix_len, mamba_exist=mamba_exist) - - def cache_finished_req( - self, req: Req, is_insert: bool = True, *, owned_kv_len: int - ) -> None: - """Cache request when it finishes.""" - if self.disable: - kv_indices = self.req_to_token_pool.req_to_token[ - req.kv.req_pool_idx, :owned_kv_len - ] - self.token_to_kv_pool_allocator.free_segment(kv_indices, start_pos=0) - self.req_to_token_pool.free_mamba_cache(req) - return - - token_ids = (req.origin_input_ids + req.output_ids)[:owned_kv_len] - kv_indices = self.req_to_token_pool.req_to_token[ - req.kv.req_pool_idx, :owned_kv_len - ] - - if is_insert: - if self.enable_mamba_extra_buffer: - cache_len = req.kv.mamba_last_track_seqlen - else: - cache_len = len(token_ids) - # ReplaySSM (no_buffer): `temporal[slot]` lags the live state by - # the slot's unflushed ring depth (`write_pos`), so cap the - # donate to the last flush boundary (where temporal is current) - # and reset the cursor, keeping the donated checkpoint consistent - # with its key length. page_size is asserted == 1, so no realign. - mamba_pool = self.req_to_token_pool.mamba_pool - write_pos_buf = mamba_pool.replayssm_write_pos - cursor_idx = req.kv.mamba_pool_idx - if write_pos_buf is None: - write_pos_buf = getattr( - mamba_pool, "replayssm_spec_write_pos", None - ) - cursor_idx = req.kv.req_pool_idx - if write_pos_buf is not None: - cache_len -= int(write_pos_buf[cursor_idx].item()) - write_pos_buf[cursor_idx] = 0 - if ( - getattr(mamba_pool, "replayssm_spec_write_pos", None) - is not None - ): - mamba_pool.replayssm_cache_base[cursor_idx] = 0 - mamba_pool.replayssm_is_flush[cursor_idx] = 0 - if cache_len is None: - cache_len = 0 - if cache_len != len(token_ids): - cache_end_idx = max(cache_len, req.kv.cache_protected_len) - self.token_to_kv_pool_allocator.free_segment( - kv_indices[cache_end_idx:], start_pos=cache_end_idx - ) - token_ids = token_ids[:cache_len] - kv_indices = kv_indices[:cache_len] - - if self.page_size != 1: - page_aligned_len = len(kv_indices) // self.page_size * self.page_size - page_aligned_kv_indices = kv_indices[:page_aligned_len].to( - dtype=torch.int64, copy=True - ) - else: - page_aligned_len = len(kv_indices) - page_aligned_kv_indices = kv_indices.to(dtype=torch.int64, copy=True) - - assert cache_len == page_aligned_len, ( - f"It is required {cache_len=}, {page_aligned_len=}, {owned_kv_len=}, {len(req.origin_input_ids)=}, {len(req.output_ids)=} ping @yizhang2077 if you see this" - ) - - # Radix Cache takes one ref in memory pool - # insert the token_ids and kv_indices into the radix tree - if self.enable_mamba_extra_buffer: - mamba_ping_pong_track_buffer_to_keep = ( - self.req_to_token_pool.get_mamba_ping_pong_keep_idx(req) - ) - src_active = req.kv.mamba_ping_pong_track_buffer[ - mamba_ping_pong_track_buffer_to_keep - ].unsqueeze(-1) - if _MAMBA_DEBUG_ASSERTS: - # .item() forces a cudaStreamSynchronize; only pay it when debugging. - assert src_active.item() != -1, ( - f"Cached mamba slot is -1: keep_idx={mamba_ping_pong_track_buffer_to_keep}, " - f"buf={req.kv.mamba_ping_pong_track_buffer.tolist()}, " - f"next_track_idx={req.kv.mamba_next_track_idx}, " - f"last_track_seqlen={req.kv.mamba_last_track_seqlen}, " - f"rid={req.rid}" - ) - if self.int8_ckpt_pool is not None: - mamba_value = self._commit_int8_checkpoint(src_active) - # quantized -> no ping-pong slot needs keeping - mamba_ping_pong_track_buffer_to_keep = None - else: - mamba_value = src_active.clone() - else: - if self.int8_ckpt_pool is not None: - mamba_value = self._commit_int8_checkpoint( - req.kv.mamba_pool_idx.unsqueeze(-1) - ) - else: - mamba_value = req.kv.mamba_pool_idx.unsqueeze(-1).clone() - mamba_ping_pong_track_buffer_to_keep = None - - result = self.insert( - InsertParams( - key=RadixKey( - token_ids[:page_aligned_len], - req.extra_key, - cache_salt=req.cache_salt, - ), - value=page_aligned_kv_indices, - mamba_value=mamba_value, - prev_prefix_len=req.kv.cache_protected_len, - ) - ) - mamba_exist = result.mamba_exist - if mamba_exist and self.int8_ckpt_pool is not None: - # state already cached -> the int8 slot we just allocated is a duplicate - self.int8_ckpt_pool.free(mamba_value) - else: - self.token_to_kv_pool_allocator.free_segment( - kv_indices[req.kv.cache_protected_len :], - start_pos=req.kv.cache_protected_len, - ) - mamba_exist = True - - if mamba_exist: - mamba_ping_pong_track_buffer_to_keep = None - - # With int8 checkpoints the radix owns an int8 slot (not the request's active - # slot), so the active mamba slot must always be returned to the active pool. - free_mamba_cache = ( - True - if (self.enable_mamba_extra_buffer or self.int8_ckpt_pool is not None) - else mamba_exist - ) - - if free_mamba_cache: - self.req_to_token_pool.free_mamba_cache( - req, - mamba_ping_pong_track_buffer_to_keep=mamba_ping_pong_track_buffer_to_keep, - ) - - self.dec_lock_ref(req.last_node) - - def cache_unfinished_req(self, req: Req, chunked=False) -> None: - """Cache request when it is unfinished.""" - - def _skip_cache_unfinished_req(req: Req) -> None: - kv_indices = self.req_to_token_pool.req_to_token[ - req.kv.req_pool_idx, : req.extend_range.end - ] - - # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later - req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True) - return - - token_ids = req.get_fill_ids() - cache_len = ( - req.kv.mamba_last_track_seqlen - if self.enable_mamba_extra_buffer - else len(token_ids) - ) - spec_write_pos = getattr( - self.req_to_token_pool.mamba_pool, "replayssm_spec_write_pos", None - ) - if not self.enable_mamba_extra_buffer and spec_write_pos is not None: - cache_len -= int(spec_write_pos[req.kv.req_pool_idx].item()) - if self.disable or cache_len is None: - return _skip_cache_unfinished_req(req) - - kv_indices_orig = self.req_to_token_pool.req_to_token[ - req.kv.req_pool_idx, : len(token_ids) - ] - # kv_indices is the kv indices to be cached - kv_indices = kv_indices_orig[:cache_len] - if self.page_size != 1: - page_aligned_len = len(kv_indices) // self.page_size * self.page_size - page_aligned_kv_indices = kv_indices[:page_aligned_len].to( - dtype=torch.int64, copy=True - ) - else: - page_aligned_len = len(kv_indices) - page_aligned_kv_indices = kv_indices.to(dtype=torch.int64, copy=True) - - assert page_aligned_len == len(kv_indices), ( - f"page_aligned_len != len(kv_indices), {page_aligned_len=}, {len(kv_indices)=}, {cache_len=}, {self.page_size=}, {self.mamba_cache_chunk_size=}" - ) - - page_aligned_token_ids = token_ids[:page_aligned_len] - - # Donate the mamba index to the radix cache instead of copying. - # This avoids a data copy that would race with the forward stream. - if self.int8_ckpt_pool is not None: - # int8 path: quantize the to-be-cached active state into an int8 slot - # (strategy-agnostic donate hook). - if self.enable_mamba_extra_buffer: - new_slot = self._alloc_mamba_slot() - src_active = self.req_to_token_pool.donate_mamba_ping_pong_slot( - req, new_slot - ) - mamba_value_donated = self._commit_int8_checkpoint(src_active) - self.req_to_token_pool.mamba_allocator.free(src_active) - else: - mamba_value_donated = self._commit_int8_checkpoint( - req.kv.mamba_pool_idx.view(-1) - ) - elif self.enable_mamba_extra_buffer: - new_slot = self._alloc_mamba_slot() - mamba_value_donated = self.req_to_token_pool.donate_mamba_ping_pong_slot( - req, new_slot - ) - else: - mamba_value_donated = self._alloc_mamba_slot() - # mamba_pool is a pure PHYSICAL store; translate both slot ids - # virtual->physical (identity for the non-unified memory pool) before the copy. - translate = self.req_to_token_pool.translate_mamba_indices - self.req_to_token_pool.mamba_pool.copy_from( - translate(req.kv.mamba_pool_idx.unsqueeze(0)), - translate(mamba_value_donated), - ) - - result = self.insert( - InsertParams( - key=RadixKey( - page_aligned_token_ids, - req.extra_key, - cache_salt=req.cache_salt, - ), - value=page_aligned_kv_indices, - mamba_value=mamba_value_donated, - prev_prefix_len=req.kv.cache_protected_len, - chunked=chunked, - ) - ) - new_prefix_len, mamba_exist = result.prefix_len, result.mamba_exist - if mamba_exist: - self._free_mamba_value(mamba_value_donated) - - # The prefix indices could be updated, reuse it - match_result = self.match_prefix( - MatchPrefixParams( - key=RadixKey( - page_aligned_token_ids, - req.extra_key, - cache_salt=req.cache_salt, - ) - ) - ) - new_indices, new_last_node = ( - match_result.device_indices, - match_result.last_device_node, - ) - - if not mamba_exist: - assert torch.equal(new_last_node.mamba_value, mamba_value_donated) - - assert req.kv.cache_protected_len <= len(new_indices) + self.page_size - 1, ( - f"{req.kv.cache_protected_len=}, {len(new_indices)=}, {len(page_aligned_token_ids)=}, {mamba_exist=}" - ) - assert new_prefix_len <= len(new_indices), ( - f"{new_prefix_len=}, {len(new_indices)=}" - ) - - self.req_to_token_pool.write( - (req.kv.req_pool_idx, slice(req.kv.cache_protected_len, len(new_indices))), - new_indices[req.kv.cache_protected_len :], - ) - - self.dec_lock_ref(req.last_node) - self.inc_lock_ref(new_last_node) - - # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later - # NOTE: this is needed for both page_size == 1 and page_size > 1 - req.prefix_indices = torch.cat( - [new_indices, kv_indices_orig[len(new_indices) :]] - ) - req.kv.cache_protected_len = len(new_indices) - req.kv.mamba_last_track_seqlen = None - req.last_node = new_last_node - - def pretty_print(self) -> None: - self._print_helper(self.root_node, 0) - total_size, total_mamba_size = self._total_size_helper() - print(f"#full_tokens: {total_size}, #mamba_num: {total_mamba_size}") - - def total_size(self) -> Tuple[int, int]: - return self._total_size_helper() - - def _evict_leaf_node( - self, x: TreeNode, is_evict_mamba: bool - ) -> Tuple[int, int, TreeNode, TreeNode]: - assert x.full_lock_ref == 0 and x.mamba_lock_ref == 0, ( - f"evict leaf node invalid with {x.id=} {x.full_lock_ref=} {x.mamba_lock_ref=}" - ) - - assert x.mamba_value is not None, f"leaf node mamba value is not None, {x.id=}" - # 1. a leaf node, free full tokens and mamba - self.kv_events.record_remove(x) - # Tree values are page-aligned copies of a kv row: page-exact segment. - self.token_to_kv_pool_allocator.free_segment(x.value, start_pos=0) - full_num_evicted = len(x.value) - self._free_mamba_value(x.mamba_value) - mamba_num_evicted = len(x.mamba_value) - - # 2. get the next node, update the lru lists - if is_evict_mamba: - x_next = self.mamba_lru_list.get_prev_no_lock(x) - else: - x_next = self.full_lru_list.get_prev_leaf_no_lock(x) - self.full_lru_list.remove_node(x) - self.mamba_lru_list.remove_node(x) - - # 3. delete the leaf node - self._delete_leaf(x) - - # 4. Iteratively delete tombstone leaves to maintain invariant that leaf nodes are not tombstone - x, leaf_full_num_evicted = self._iteratively_delete_tombstone_leaf(x) - full_num_evicted += leaf_full_num_evicted - return full_num_evicted, mamba_num_evicted, x, x_next - - 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 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 - # evict lru leaf nodes until mamba_num_tokens is reached - while mamba_num_evicted < mamba_num and (self.mamba_lru_list.in_list(x)): - assert x.mamba_value is not None, f"node has no mamba value, {x.id=}" - assert len(x.mamba_value) == 1, ( - f"node has abnormal mamba length, {x.id=}, {len(x.mamba_value)=}" - ) - assert x != self.root_node, f"root node is not evictable, {x.id=}" - assert x.mamba_lock_ref == 0, f"node is in use by mamba kv indices, {x.id=}" - - if len(x.children) > 0: - # 1. an internal node, free mamba tokens. - self._free_mamba_value(x.mamba_value) - mamba_num_evicted += len(x.mamba_value) - - # 2. get the next node, update the lru lists - x_next = self.mamba_lru_list.get_prev_no_lock(x) - self.mamba_lru_list.remove_node(x) - - # 3. tombstone the node - self._tombstone_internal_node(x) - else: - _, mamba_evicted_delta, _, x_next = self._evict_leaf_node(x, True) - mamba_num_evicted += mamba_evicted_delta - - x = x_next - - 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 0 - - full_num_evicted = 0 - # get the least recently used leaf node that is not locked - x = self.full_lru_list.get_leaf_lru_no_lock() - - while full_num_evicted < full_num_tokens and self.full_lru_list.in_list(x): - assert x != self.root_node, ( - f"root node should not exist in full lru list, {x.id=}" - ) - full_num_evicted_delta, _, x, x_next = self._evict_leaf_node(x, False) - full_num_evicted += full_num_evicted_delta - - # if parent has no more children, it is a leaf. It is possible that this node is lru, so - # we need to get the first leaf node in the lru list - if len(x.parent.children) == 0: - x_next = self.full_lru_list.get_leaf_lru_no_lock() - - x = x_next - - return full_num_evicted - - def inc_lock_ref(self, node: TreeNode) -> IncLockRefResult: - """ - Increment the lock reference count for the node. - It locks the full_lock_ref for nodes between the [last node, root), exclusive. - It locks the mamba_lock_ref for current node if its mamba_value exists. - """ - if self.disable: - return IncLockRefResult() - - # protect mamba value in current node if it exists - if node.mamba_value is not None: - if node.mamba_lock_ref == 0: - self.mamba_evictable_size_ -= len(node.mamba_value) - self.mamba_protected_size_ += len(node.mamba_value) - node.mamba_lock_ref += 1 - - while node != self.root_node: - # lock full from node to root - assert node.full_lock_ref >= 0, ( - f"inc_lock_ref on node with {node.full_lock_ref=}, {node.id=}" - ) - if node.full_lock_ref == 0: - self.full_evictable_size_ -= len(node.value) - self.full_protected_size_ += len(node.value) - node.full_lock_ref += 1 - node = node.parent - return IncLockRefResult() - - def dec_lock_ref( - self, node: TreeNode, params: Optional[DecLockRefParams] = None - ) -> DecLockRefResult: - """ - Decrement the lock reference count for the node. - It unlocks the full_lock_ref for nodes between the [last node, root), exclusive. - It unlocks the mamba_lock_ref for current node if its mamba_value exists. - """ - if self.disable: - return DecLockRefResult() - - if node.mamba_value is not None: - assert node.mamba_lock_ref > 0, ( - f"dec_lock_ref on node with {node.mamba_lock_ref=}, {node.id=}" - ) - if node.mamba_lock_ref == 1: - self.mamba_evictable_size_ += len(node.mamba_value) - self.mamba_protected_size_ -= len(node.mamba_value) - node.mamba_lock_ref -= 1 - - while node != self.root_node: - assert node.full_lock_ref > 0, ( - f"dec_lock_ref on node with {node.full_lock_ref=}, {node.id=}" - ) - if node.full_lock_ref == 1: - self.full_evictable_size_ += len(node.value) - self.full_protected_size_ -= len(node.value) - node.full_lock_ref -= 1 - node = node.parent - - return DecLockRefResult() - - def sanity_check(self): - if self.disable: - return - self.full_lru_list.sanity_check(self) - self.mamba_lru_list.sanity_check(self) - - def evictable_size(self) -> Tuple[int, int]: - # Note: use full_evictable_size() and mamba_evictable_size() instead. - raise NotImplementedError - - def full_evictable_size(self) -> int: - return self.full_evictable_size_ - - def mamba_evictable_size(self) -> int: - return self.mamba_evictable_size_ - - def protected_size(self) -> Tuple[int, int]: - # Note: use full_protected_size() and mamba_protected_size() instead. - raise NotImplementedError - - def full_protected_size(self) -> int: - # protected size refers to the size of the full cache that is locked - return self.full_protected_size_ - - def mamba_protected_size(self) -> int: - # protected size refers to the size of the mamba cache that is locked - return self.mamba_protected_size_ - - def all_values_flatten(self) -> torch.Tensor: - values = [] - - def _dfs_helper(node: TreeNode): - for _, child in node.children.items(): - values.append(child.value) - _dfs_helper(child) - - _dfs_helper(self.root_node) - return torch.cat(values) if len(values) > 0 else torch.tensor([]) - - def all_mamba_values_flatten(self) -> torch.Tensor: - values = [] - - def _dfs_helper(node: TreeNode): - if node.mamba_value is not None: - values.append(node.mamba_value) - for _, child in node.children.items(): - _dfs_helper(child) - - _dfs_helper(self.root_node) - return torch.cat(values) if len(values) > 0 else torch.tensor([]) - - def available_and_evictable_str(self) -> str: - full_available_size = self.token_to_kv_pool_allocator.available_size() - full_evictable_size = self.full_evictable_size() - return ( - f"Available full tokens: {full_available_size + full_evictable_size} ({full_available_size=} + {full_evictable_size=})\n" - f"Full LRU list evictable size: {self.full_lru_list.sanity_check_evictable_size()}\n" - ) - - ##### Internal Helper Functions ##### - - def _alloc_mamba_slot(self) -> torch.Tensor: - """Allocate one mamba pool slot, evicting if necessary.""" - slot = self.req_to_token_pool.mamba_allocator.alloc(1) - if slot is None: - self.evict(EvictParams(num_tokens=0, mamba_num=1)) - slot = self.req_to_token_pool.mamba_allocator.alloc(1) - assert slot is not None, "Can not alloc mamba cache" - return slot - - @property - def int8_ckpt_pool(self): - """The int8 checkpoint pool, or None when --enable-int8-mamba-checkpoint is off. - When enabled, radix-cached mamba states live HERE (int8), not in the active - bf16 pool -> ~2x cached-prefix capacity at fixed memory.""" - return getattr(self.req_to_token_pool, "mamba_ckpt_pool", None) - - def _alloc_int8_ckpt_slot(self) -> torch.Tensor: - """Allocate one int8 checkpoint slot, evicting cached states if the pool is full.""" - slot = self.int8_ckpt_pool.alloc(1) - if slot is None: - self.evict(EvictParams(num_tokens=0, mamba_num=1)) - slot = self.int8_ckpt_pool.alloc(1) - assert slot is not None, "Can not alloc int8 mamba checkpoint slot" - return slot - - def _commit_int8_checkpoint(self, active_slots: torch.Tensor) -> torch.Tensor: - """Quantize the active-pool state at ``active_slots`` into a fresh int8 - checkpoint slot and return that slot. Strategy-agnostic donate hook: both - no_buffer (copy_from) and extra_buffer (ping-pong) converge here. The caller - frees ``active_slots`` separately.""" - ckpt_slot = self._alloc_int8_ckpt_slot() - self.int8_ckpt_pool.store_from_active( - self.req_to_token_pool.mamba_pool, active_slots, ckpt_slot - ) - return ckpt_slot - - def _free_mamba_value(self, mamba_value: torch.Tensor) -> None: - """Free a node's mamba_value to the right allocator (int8 ckpt pool or the - active mamba allocator).""" - if self.int8_ckpt_pool is not None: - self.int8_ckpt_pool.free(mamba_value) - else: - self.req_to_token_pool.mamba_allocator.free(mamba_value) - - def _match_prefix_helper( - self, key: RadixKey - ) -> Tuple[List[torch.Tensor], TreeNode, int]: - """ - Mamba prefix matching helper. It factors in the sliding window size such that - the matched node is guaranteed to either 1. connected to root without mamba tombstone, - or 2. the number of matching tokens from the matched node to the last mamba tombstone - node is greater than or equal to the sliding window size. - """ - node = self.root_node - child_key = key.child_key(self.page_size) - - value: List[torch.Tensor] = [] - best_value_len = 0 - best_last_node = node - while len(key) > 0 and child_key in node.children.keys(): - child = node.children[child_key] - # update best_value_len and best_last_node if needed - if node.mamba_value is not None: - best_value_len = len(value) - best_last_node = node - - prefix_len = child.key.match(key, page_size=self.page_size) - if prefix_len < len(child.key): - new_node = self._split_node(child.key, child, prefix_len) - value.append(new_node.value) - node = new_node - break - else: - value.append(child.value) - node = child - key = key[prefix_len:] - - if len(key): - child_key = key.child_key(self.page_size) - # handle best_value_len and best_last_node, for the case that last node is fully matched - if node.mamba_value is not None: - best_value_len = len(value) - best_last_node = node - - return value, best_last_node, best_value_len - - def _match_pre_processor(self, params: MatchPrefixParams) -> Optional[RadixKey]: - """Preprocess the key before matching.""" - key = params.key - - if self.disable or len(key) == 0: - return None - - return key - - def _match_post_processor( - self, - params: MatchPrefixParams, - value: List[torch.Tensor], - last_node: TreeNode, - best_value_len: int, - ) -> MatchResult: - """Post-process the matched result.""" - cow_mamba = params.cow_mamba - req = params.req - - # Full KV of the whole matched path is reused as prefix, so refresh the entire - # chain (nodes closer to root end up least recently used, evicted first). - node_update = last_node - self.full_lru_list.reset_node_and_parents_mru(node_update, self.root_node) - # Mamba only consumes last_node's state (cf. inc_lock_ref, which locks just this - # node's mamba_value). Refreshing ancestors would keep a whole session's states - # adjacent in the mamba LRU and evict cold sessions wholesale; touch only the used - # state so older leaves survive. - if last_node is not self.root_node and last_node.mamba_value is not None: - self.mamba_lru_list.reset_node_mru(last_node) - - # This last_access_time is for sanity check, can be deleted after validation in production - cur_time = get_last_access_time() - while node_update: - node_update.last_access_time = cur_time - cur_time -= ( - 0.00001 # assuming less than 100000 nodes in a branch of the tree - ) - node_update = node_update.parent - - # Calculate the branching point. It is defined as the last aligned position that - # does not have a mamba value. - if len(value) > best_value_len: - chunk_aligned_seqlen = ( - sum(len(v) for v in value) // self.mamba_cache_chunk_size - ) * self.mamba_cache_chunk_size - mamba_branching_seqlen = ( - chunk_aligned_seqlen if chunk_aligned_seqlen > 0 else None - ) - else: - mamba_branching_seqlen = None - - # Defer COW to forward stream: record source index, allocate destination - if cow_mamba and last_node.mamba_value is not None: - if not req.kv.holds_mamba: - dst_index = self.req_to_token_pool.mamba_allocator.alloc(1) - if dst_index is None: - self.inc_lock_ref(last_node) - self.evict(EvictParams(num_tokens=0, mamba_num=1)) - dst_index = self.req_to_token_pool.mamba_allocator.alloc(1) - self.dec_lock_ref(last_node) - assert dst_index is not None, "Can not alloc mamba cache" - req.kv.mamba_pool_idx = dst_index[0] - req.kv.mamba_cow_src_index = last_node.mamba_value - req.kv.mamba_needs_clear = False - - value = value[:best_value_len] - if value: - value = torch.cat(value) - else: - value = torch.empty((0,), dtype=torch.int64, device=self.device) - - return MatchResult( - device_indices=value, - last_device_node=last_node, - last_host_node=last_node, - best_match_node=last_node, - mamba_branching_seqlen=mamba_branching_seqlen, - ) - - def _split_node(self, key: RadixKey, child: TreeNode, split_len: int) -> TreeNode: - # new_node -> child - new_node = TreeNode() - new_node.children = {key[split_len:].child_key(self.page_size): child} - new_node.parent = child.parent - new_node.mamba_value = None # mamba cache can not be split - new_node.full_lock_ref = child.full_lock_ref - new_node.mamba_lock_ref = 0 - new_node.key = child.key[:split_len] - new_node.value = child.value[:split_len].clone() - - # child time should be later than the new parent's time in the full LRU - child.last_access_time = get_last_access_time() - - # A split does not change the set of live mamba states (child keeps its value, - # new_node is a mamba tombstone), so the mamba LRU is left untouched — only the - # full LRU reorders around the new intermediate node. - self.full_lru_list.remove_node(child) - child.parent = new_node - child.key = child.key[split_len:] - child.value = child.value[split_len:].clone() - new_node.parent.children[key.child_key(self.page_size)] = new_node - new_node.hash_value, child.hash_value = split_node_hash_value( - child.hash_value, split_len, self.page_size - ) - new_node.event_hash_value, child.event_hash_value = split_node_hash_value( - child.event_hash_value, split_len, self.page_size - ) - - # insert the new node and child into the full lru list, insert - # parent first so that parent is after child in the lru list - self.full_lru_list.insert_mru(new_node) - self.full_lru_list.insert_mru(child) - return new_node - - def _insert_helper( - self, - node: TreeNode, - key: RadixKey, - value, - mamba_value, - chunked: bool = False, - prev_prefix_len: int = 0, - ) -> Tuple[int, bool]: - # Refresh the full LRU from root to leaf (the whole path is reused as prefix). - # The mamba states of these existing nodes were not recomputed this insert, so - # the mamba LRU is left untouched here; only genuinely new mamba states (the new - # leaf / a revived tombstone below) are inserted. - assert mamba_value is not None, "Mamba value should not be None here." - node.last_access_time = get_last_access_time() - if node != self.root_node: - self.full_lru_list.reset_node_mru(node) - if len(key) == 0: - return 0, True - - child_key = key.child_key(self.page_size) - - total_prefix_length = 0 - while len(key) > 0 and child_key in node.children.keys(): - node = node.children[child_key] - node.last_access_time = get_last_access_time() - self.full_lru_list.reset_node_mru(node) - prefix_len = node.key.match(key, page_size=self.page_size) - - if prev_prefix_len < total_prefix_length + prefix_len: - start = max(0, prev_prefix_len - total_prefix_length) - # value sits at offset total_prefix_length of the kv row; match() - # rounds prefix_len to page multiples, so frees never share a page. - self.token_to_kv_pool_allocator.free_segment( - value[start:prefix_len], - start_pos=total_prefix_length + start, - ) - - total_prefix_length += prefix_len - key = key[prefix_len:] - value = value[prefix_len:] - - if prefix_len < len(node.key): - new_node = self._split_node(node.key, node, prefix_len) - node = new_node - - if len(key): - child_key = key.child_key(self.page_size) - - mamba_value_exist = False - if len(key): - new_node = TreeNode() - new_node.parent = node - new_node.key = key - new_node.value = value.clone() - new_node.mamba_value = mamba_value - self.full_lru_list.insert_mru(new_node) - self.mamba_lru_list.insert_mru(new_node) - node.children[child_key] = new_node - self.full_evictable_size_ += len(value) - self.mamba_evictable_size_ += len(mamba_value) - self.kv_events.record_store(new_node) - elif node.mamba_value is None: # add for mamba tombstone - node.mamba_value = mamba_value - self.full_lru_list.reset_node_mru(node) - self.mamba_lru_list.insert_mru(node) - self.mamba_evictable_size_ += len(mamba_value) - node.last_access_time = get_last_access_time() - else: # mamba value already exists - mamba_value_exist = True - self.full_lru_list.reset_node_mru(node) - node.last_access_time = get_last_access_time() - - return total_prefix_length, mamba_value_exist - - def _iteratively_delete_tombstone_leaf( - self, node: TreeNode - ) -> Tuple[TreeNode, int]: - full_num_evicted = 0 - while node.parent.mamba_value is None and len(node.parent.children) == 0: - # root node is not evictable - if node.parent == self.root_node: - break - # if locked, means node is in use, skip - if node.parent.full_lock_ref > 0: - break - assert node.parent.mamba_lock_ref == 0, ( - f"tombstone mamba_lock_ref should always be 0, {node.parent.full_lock_ref=}, {node.parent.mamba_lock_ref=}, {node.parent.id=}" - ) - # delete tombstone node evicts full tokens - self.kv_events.record_remove(node.parent) - self.token_to_kv_pool_allocator.free_segment(node.parent.value, start_pos=0) - full_num_evicted += len(node.parent.value) - self.full_lru_list.remove_node(node.parent) - self._delete_tombstone_leaf(node.parent) - node = node.parent - - return node, full_num_evicted - - def _delete_leaf(self, node: TreeNode) -> None: - assert node.mamba_value is not None, ( - f"Invariant violated: leaf node is a tombstone, {node.id=}" - ) - assert len(node.children) == 0, f"leaf node has children, {node.id=}" - key = node.key.child_key(self.page_size) - v = node.parent.children.pop(key, None) - assert v == node, f"parent does not have child key, {key}" - - self.full_evictable_size_ -= len(node.key) - self.mamba_evictable_size_ -= len(node.mamba_value) - - def _tombstone_internal_node(self, node: TreeNode) -> None: - assert len(node.children) != 0, f"Cannot tombstone a leaf node, {node.id=}" - self.mamba_evictable_size_ -= len(node.mamba_value) - node.mamba_value = None - - def _delete_tombstone_leaf(self, node: TreeNode) -> None: - assert node.mamba_value is None, ( - f"Deleting a unexpected non-tombstone leaf node, {node.id=}" - ) - assert len(node.children) == 0, f"leaf node has children, {node.id=}" - key = node.key.child_key(self.page_size) - v = node.parent.children.pop(key, None) - assert v == node, f"parent does not have child key, {key}" - - self.full_evictable_size_ -= len(node.key) - - def _collect_nontombstone_nodes(self) -> List[TreeNode]: - ret_list = [] - stack = [self.root_node] - - while stack: - cur_node = stack.pop() - if cur_node.mamba_value is not None: - ret_list.append(cur_node) - stack.extend(cur_node.children.values()) - - return ret_list - - def _collect_all_nodes(self) -> List[TreeNode]: - ret_list = [] - stack = [self.root_node] - while stack: - cur_node = stack.pop() - ret_list.append(cur_node) - stack.extend(cur_node.children.values()) - return ret_list - - def _print_helper(self, node: TreeNode, indent: int) -> None: - """Prints the radix tree in a human-readable format.""" - stack = [(node, indent)] - while stack: - current_node, current_indent = stack.pop() - print( - " " * current_indent, - f"[{current_node.id}]", - len(current_node.key), - f"fr={current_node.full_lock_ref}", - f"mr={current_node.mamba_lock_ref}", - f"fll={self.full_lru_list.in_list(current_node)}", - f"mll={self.mamba_lru_list.in_list(current_node)}", - f"mv={current_node.mamba_value}", - ) - for key, child in current_node.children.items(): - stack.append((child, current_indent + 2)) - - assert key == child.key.child_key(self.page_size), ( - f"{key=}, {child.key.child_key(self.page_size)=}" - ) - - def _total_size_helper(self) -> Tuple[int, int]: - total_size = 0 - total_mamba_size = 0 - stack = [self.root_node] - while stack: - current_node = stack.pop() - total_size += len(current_node.value) - if current_node.mamba_value is not None: - total_mamba_size += len(current_node.mamba_value) - for child in current_node.children.values(): - if child.evicted: - continue - stack.append(child) - return total_size, total_mamba_size diff --git a/python/sglang/srt/mem_cache/pure_swa_radix_cache.py b/python/sglang/srt/mem_cache/pure_swa_radix_cache.py index 2a652fc31..9926f5aa0 100644 --- a/python/sglang/srt/mem_cache/pure_swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/pure_swa_radix_cache.py @@ -52,8 +52,8 @@ class PureSWARadixCache(RadixCache): return 0 def sanity_check(self): - """No-op: PureSWARadixCache uses RadixCache's simple tree structure - which doesn't need the dual-LRU sanity checks of SWARadixCache.""" + """No-op: an all-SWA model has no full tier, so there is no full/SWA + split to cross-check.""" pass def evict(self, params: EvictParams) -> EvictResult: diff --git a/python/sglang/srt/mem_cache/qsa_kv_pool.py b/python/sglang/srt/mem_cache/qsa_kv_pool.py index 61308c2e3..b3e4f6010 100644 --- a/python/sglang/srt/mem_cache/qsa_kv_pool.py +++ b/python/sglang/srt/mem_cache/qsa_kv_pool.py @@ -74,8 +74,8 @@ class QSATokenToKVPool(HybridLinearKVPool): "compressed QSA requires a paged full-KV cache with the page " "a multiple of the compress ratio (compressed slots are " f"full_slot // ratio): page_size={page_size}, " - f"ratio={qsa_compress_ratio}. With MambaRadixCache this " - "needs the mamba extra-buffer strategy or " + f"ratio={qsa_compress_ratio}. This needs the mamba " + "extra-buffer strategy or " "--disable-radix-cache (see the Qwen4-Exp arg overrides)." ) # super().__init__ computes mem_usage via the overridden get_kv_size_bytes, diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py deleted file mode 100644 index ffce936f0..000000000 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ /dev/null @@ -1,1462 +0,0 @@ -from __future__ import annotations - -""" -Copyright 2023-2024 SGLang Team -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -""" - -""" -The radix tree data structure for managing the hybrid (full and SWA) KV cache. -""" - -import heapq -import time -from collections import defaultdict -from typing import TYPE_CHECKING, List, Optional, Tuple - -import torch -from numpy import float64 - -from sglang.srt.environ import envs -from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator -from sglang.srt.mem_cache.base_prefix_cache import ( - BasePrefixCache, - DecLockRefParams, - DecLockRefResult, - EvictParams, - EvictResult, - IncLockRefResult, - InsertParams, - InsertResult, - MatchPrefixParams, - MatchResult, -) -from sglang.srt.mem_cache.cache_init_params import CacheInitParams -from sglang.srt.mem_cache.common import free_kv_row_segments -from sglang.srt.mem_cache.events import KVCacheEventRecorder -from sglang.srt.mem_cache.radix_cache import RadixKey -from sglang.srt.mem_cache.utils import split_node_hash_value - -if TYPE_CHECKING: - from sglang.srt.managers.schedule_batch import Req - -import logging - -logger = logging.getLogger(__name__) - - -class TreeNode: - counter = 0 - swa_uuid_counter = 1 - last_access_time_counter_float = float64(1.0) - - def __init__(self, id: Optional[int] = None): - self.children = defaultdict(TreeNode) - self.parent: TreeNode = None - self.key: RadixKey = None - self.value: Optional[torch.Tensor] = None - # swa_tombstone is used to indicate the kv indices have been freed for swa layers - self.swa_tombstone = False - # invariant: for any node, if swa_lock_ref is locked, full_lock_ref must be locked; - # if full_lock_ref is locked, swa_lock_ref doesn't need to be locked. So, - # full_lock_ref is always >= swa_lock_ref. - self.full_lock_ref = 0 - self.swa_lock_ref = 0 - # last access time is only used for sanity check. LRU is maintained by the lru list. - self.last_access_time = get_last_access_time() - - self.hit_count = 0 - # store the host indices of KV cache - self.host_value = None - # store hash values of each page - self.hash_value: Optional[List[str]] = None - # Namespace-aware hashes used only for external KV events. - self.event_hash_value: Optional[List[str]] = None - - # for lru list, invariant: - # 1. prev has greater last_access_time - # 2. next has smaller last_access_time - self.prev = None - self.next = None - self.swa_prev = None - self.swa_next = None - - self.id = TreeNode.counter if id is None else id - TreeNode.counter += 1 - self.swa_uuid = None - - @property - def evicted(self): - return self.value is None - - @property - def backuped(self): - return self.host_value is not None - - def __lt__(self, other: TreeNode): - return self.last_access_time < other.last_access_time - - -def gen_swa_uuid() -> int: - TreeNode.swa_uuid_counter += 1 - return TreeNode.swa_uuid_counter - - -def get_last_access_time() -> float64: - ret = TreeNode.last_access_time_counter_float - TreeNode.last_access_time_counter_float += 1.0 - return ret - - -class LRUList: - def __init__(self, is_swa_list: bool = False): - self.is_swa_list = is_swa_list - if self.is_swa_list: - self.prv = "swa_prev" - self.nxt = "swa_next" - self.lock_ref = "swa_lock_ref" - else: - self.prv = "prev" - self.nxt = "next" - self.lock_ref = "full_lock_ref" - # Initialize dummy head and tail nodes - self.head = TreeNode() # Most recently used side - self.tail = TreeNode() # Least recently used side - setattr(self.head, self.nxt, self.tail) # self.head.next = self.tail - setattr(self.tail, self.prv, self.head) # self.tail.prev = self.head - self.cache = {} - - def _add_node(self, node): - """Helper to add node right after head (most recently used)""" - self._add_node_after(self.head, node) - - def _add_node_after(self, old_node, new_node): - """Helper to add node right after old_node""" - setattr(new_node, self.prv, old_node) # new_node.prev = old_node - setattr( - new_node, self.nxt, getattr(old_node, self.nxt) - ) # new_node.next = old_node.next - setattr( - getattr(old_node, self.nxt), self.prv, new_node - ) # old_node.next.prev = new_node - setattr(old_node, self.nxt, new_node) # old_node.next = new_node - - def _remove_node(self, node): - """Helper to remove node from linked list""" - setattr( - getattr(node, self.prv), self.nxt, getattr(node, self.nxt) - ) # node.prev.next = node.next - setattr( - getattr(node, self.nxt), self.prv, getattr(node, self.prv) - ) # node.next.prev = node.prev - # Clear self pointers to break reference cycles among evicted nodes. - setattr(node, self.prv, None) - setattr(node, self.nxt, None) - - def _get_lru(self) -> Optional[TreeNode]: - """ - Get the least recently used node - """ - if len(self.cache) == 0: - return None - return getattr(self.tail, self.prv) - - def reset_node_mru(self, node): - """ - Move a (existing) node to most recently used position - """ - assert node.id in self.cache, f"Resetting node {node.id=} not in lru list" - assert not self.is_swa_list or not node.swa_tombstone, ( - f"Resetting swa tombstone node in swa lru list: {node.id=}" - ) - self._remove_node(node) - self._add_node(node) - - def reset_node_and_parents_mru(self, node, root_node): - """ - Move an (existing) node and its parents to most recently used position. Child node is - more recently used than parent node. - """ - prev_node = self.head - while node != root_node: - # for swa lru list, only reset non-tombstone nodes - if not self.is_swa_list or not node.swa_tombstone: - assert node.id in self.cache, ( - f"Resetting node {node.id=} not in lru list when resetting node and parents mru" - ) - self._remove_node(node) - self._add_node_after(prev_node, node) - prev_node = node - node = node.parent - - def insert_mru(self, node): - """ - Insert a (new) node as most recently used - """ - assert not self.is_swa_list or not node.swa_tombstone, ( - f"Inserting swa tombstone node in swa lru list: {node.id=}" - ) - assert node.id not in self.cache, ( - f"Inserting node {node.id=} already in lru list, existing node: {self.cache[node.id].id=}" - ) - self.cache[node.id] = node - self._add_node(node) - - def remove_node(self, node: TreeNode): - """ - Remove node from lru list - """ - assert node.id in self.cache, f"Removing node {node.id=} not in lru list" - assert not self.is_swa_list or not node.swa_tombstone, ( - f"Removing swa tombstone node from swa lru list: {node.id=}" - ) - del self.cache[node.id] - self._remove_node(node) - - def get_lru_no_lock(self) -> Optional[TreeNode]: - """ - Get the least recently used node that is not locked - """ - return self.get_prev_no_lock(self.tail, check_id=False) - - def get_leaf_lru_no_lock(self) -> Optional[TreeNode]: - """ - Get the least recently used leaf node that is not locked - """ - return self.get_prev_leaf_no_lock(self.tail, check_id=False) - - def get_prev_no_lock( - self, node: TreeNode, check_id: bool = True - ) -> Optional[TreeNode]: - """ - Get the previous (i.e. more recently used) node that is not locked - """ - if check_id: - assert node.id in self.cache, ( - f"Getting prev of node {node.id=} not in lru list" - ) - x = getattr(node, self.prv) # x = node.prev - while getattr(x, self.lock_ref) > 0: - x = getattr(x, self.prv) # x = x.prev - # if x is the head, it means there is no node in the lru list without lock - if x == self.head: - return None - return x - - def get_prev_leaf_no_lock(self, node: TreeNode, check_id: bool = True): - """ - Get the previous (i.e. more recently used) leaf node that is not locked - """ - if check_id: - assert node.id in self.cache, ( - f"Getting prev of node {node.id=} not in lru list" - ) - x = getattr(node, self.prv) # x = node.prev - while getattr(x, self.lock_ref) > 0 or len(x.children) > 0: - x = getattr(x, self.prv) # x = x.prev - # if x is the head, it means there is no leaf node in the lru list without lock - if x == self.head: - return None - return x - - def in_list(self, node: Optional[TreeNode]): - """ - Check if the node is in the lru list - """ - if not node: - return False - return node.id in self.cache - - # Note: this is expensive, only use for debug - def sanity_check_evictable_size(self): - """ - Check the evictable size (i.e. the size of the nodes that are not locked) - """ - node = self.get_lru_no_lock() - evictable_size = 0 - while self.in_list(node): - evictable_size += len(node.value) - node = self.get_prev_no_lock(node) - return evictable_size - - # Note: this is expensive, only use for debug or idle check - def sanity_check(self, tree_cache: SWARadixCache): - """ - Check if the lru list is valid by rebuilding the lru list from the tree, heapifying it, and - checking if the lru list is valid. - """ - try: - if self.is_swa_list: - nodes = tree_cache._collect_nontombstone_nodes() - else: - nodes = tree_cache._collect_all_nodes() - total_nodes = len(nodes) - total_lru_plus_1 = len(self.cache) + 1 - # heapify based on last_access_time - heapq.heapify(nodes) - # the root node is not in the lru list - assert len(nodes) == len(self.cache) + 1, ( - f"len(nodes): {len(nodes)} != len(self.cache) + 1: {len(self.cache) + 1}" - ) - - x_lru = self._get_lru() - while len(nodes): - x = heapq.heappop(nodes) - if x == tree_cache.root_node: - # root node is not in the lru list - continue - assert x == x_lru, ( - f"Incorrect LRU list, {self.is_swa_list=}, x: {x.id=} != x_lru: {x_lru.id=}" - ) - assert x_lru.full_lock_ref == 0, ( - f"x_lru should not be locked when idle, {x_lru.full_lock_ref=}, {x_lru.swa_uuid=}, {x_lru.id=}" - ) - assert x_lru.swa_lock_ref == 0, ( - f"x_lru should not be locked when idle, {x_lru.swa_lock_ref=}, {x_lru.swa_uuid=}, {x_lru.id=}" - ) - x_lru = getattr(x, self.prv) - - if self.is_swa_list: - evictable_size = tree_cache.swa_evictable_size() - lru_list_evictable_size = self.sanity_check_evictable_size() - else: - evictable_size = tree_cache.full_evictable_size() - lru_list_evictable_size = self.sanity_check_evictable_size() - - assert evictable_size == lru_list_evictable_size, ( - f"{self.is_swa_list=}, total nodes: {total_nodes}, total lru plus 1: {total_lru_plus_1}, evictable size: {evictable_size} != lru list evictable size: {lru_list_evictable_size}" - ) - except Exception as e: - msg = f"SWA Radix tree sanity check failed, ping @hanming-lu: {e}" - logger.error(msg) - raise Exception(msg) - - -class SWARadixCache(BasePrefixCache): - def __init__(self, params: CacheInitParams): - assert isinstance(params.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator) - self.req_to_token_pool = params.req_to_token_pool - self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator - self.page_size = params.page_size - self.disable = params.disable - self.is_eagle = params.is_eagle - self.kv_events = KVCacheEventRecorder( - enabled=params.enable_kv_cache_events, page_size=self.page_size - ) - - if self.token_to_kv_pool_allocator: - self.device = self.token_to_kv_pool_allocator.device - else: - self.device = torch.device("cpu") - - if params.enable_metrics: - self.init_metrics_collector() - - self.sliding_window_size = params.sliding_window_size - self.reset() - - ##### Public API ##### - - def supports_swa(self) -> bool: - assert self.sliding_window_size is not None, ( - "sliding_window_size must be set for SWARadixCache" - ) - return True - - def swa_reprefill_tail_tokens(self) -> int: - """Tokens at the tail of a matched prefix that must NOT be reused. - - The DeepSeek-V4 unified_kv layout keeps SWA in a per-request ring - (addressed by ``req_pool_idx * window + pos % window``), which is NOT - content-stable and is never stored in the radix tree. A reused prefix - therefore carries another request's stale SWA in the ring. Hold back the - trailing sliding window from the match so it gets re-prefilled into THIS - request's ring, making the decode window read freshly-written data. - - No-op (0) for the index-addressed SWA pool, whose slots are - content-stable and safe to reuse. - """ - from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( - is_unified_kv_triton, - ) - - if self.sliding_window_size and is_unified_kv_triton(): - return self.sliding_window_size - return 0 - - def reset(self) -> None: - self.root_node = TreeNode() - self.root_node.key = [] - self.root_node.value = [] - self.root_node.hash_value = [] - self.root_node.full_lock_ref = 1 - self.root_node.swa_lock_ref = 1 - self.full_evictable_size_ = 0 - self.swa_evictable_size_ = 0 - self.full_protected_size_ = 0 - self.swa_protected_size_ = 0 - # LRU lists are used to maintain the order of eviction of the nodes in the tree - self.full_lru_list = LRUList(is_swa_list=False) - self.swa_lru_list = LRUList(is_swa_list=True) - self.kv_events.record_all_cleared() - - def match_prefix(self, params: MatchPrefixParams) -> MatchResult: - """Find the matching prefix from the radix tree. - Args: - params: MatchPrefixParams containing key. - Returns: - A tuple of a tensor of matching prefix token IDs and - the last node that contains the prefix values. Note that - this API can modify the internal state of the Radix tree. - The last node create a new child if the prefix is shorter - than the last node's value. - """ - - key = self._match_pre_processor(params) - if key is None: - return MatchResult( - device_indices=torch.empty( - (0,), - dtype=torch.int64, - device=self.device, - ), - last_device_node=self.root_node, - last_host_node=self.root_node, - best_match_node=self.root_node, - ) - - value, last_node, best_value_len = self._match_prefix_helper(key) - return self._match_post_processor(params, value, last_node, best_value_len) - - def insert(self, params: InsertParams) -> InsertResult: - if self.disable: - 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, value = key.maybe_to_bigram_view(self.is_eagle, value) - key = key.page_aligned(self.page_size) - if value is not None: - value = value[: len(key)] - else: - value = torch.tensor(key.token_ids[: len(key)], dtype=torch.int64) - - 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, *, owned_kv_len: int - ) -> None: - """Cache request when it finishes.""" - if self.disable: - self.free_kv_row(req.kv, [(0, owned_kv_len)]) - return - - token_ids = (req.origin_input_ids + req.output_ids)[:owned_kv_len] - kv_indices = self.req_to_token_pool.req_to_token[ - req.kv.req_pool_idx, :owned_kv_len - ] - - radix_key = RadixKey( - token_ids, - req.extra_key, - is_bigram=self.is_eagle, - cache_salt=req.cache_salt, - ).page_aligned(self.page_size) - page_aligned_len = len(radix_key) - values = kv_indices[:page_aligned_len].to(dtype=torch.int64, copy=True) - old_prefix_len = req.kv.cache_protected_len - - # Radix Cache takes one ref in memory pool - # Note: the insert function already frees the overlapped kv_indices - if is_insert: - self.insert( - InsertParams( - key=radix_key, - value=values, - prev_prefix_len=old_prefix_len, - swa_evicted_seqlen=req.kv.swa_evicted_seqlen, - ) - ) - else: - self.free_kv_row(req.kv, [(old_prefix_len, page_aligned_len)]) - - # free the unaligned tail - self.free_kv_row(req.kv, [(page_aligned_len, owned_kv_len)]) - - # Remove req slot release the cache lock - self.dec_lock_ref( - req.last_node, req.lock_receipt, skip_swa=req.swa_prefix_lock_released - ) - req.swa_prefix_lock_released = False - - def cache_unfinished_req(self, req: Req, chunked=False) -> None: - """Cache request when it is unfinished.""" - if self.disable: - kv_indices = self.req_to_token_pool.req_to_token[ - req.kv.req_pool_idx, : req.extend_range.end - ] - - # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later - req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True) - return - - token_ids = req.get_fill_ids() - kv_indices = self.req_to_token_pool.req_to_token[ - req.kv.req_pool_idx, : len(token_ids) - ] - - radix_key = RadixKey( - token_ids, - req.extra_key, - is_bigram=self.is_eagle, - cache_salt=req.cache_salt, - ).page_aligned(self.page_size) - values = kv_indices[: len(radix_key)].to(dtype=torch.int64, copy=True) - old_prefix_len = req.kv.cache_protected_len - - # Radix Cache takes one ref in memory pool - # Note: the insert function already frees the overlapped kv_indices - # The prefix below swa_evicted_seqlen has no SWA peers left; the insert - # tombstones it instead of claiming live SWA KV. - result = self.insert( - InsertParams( - key=radix_key, - value=values, - prev_prefix_len=old_prefix_len, - swa_evicted_seqlen=req.kv.swa_evicted_seqlen, - ) - ) - new_prefix_len = result.prefix_len - - # The prefix indices could be updated, reuse it - match_result = self.match_prefix(MatchPrefixParams(key=radix_key)) - new_indices, new_last_node = ( - match_result.device_indices, - match_result.last_device_node, - ) - - assert old_prefix_len <= len(new_indices), f"{old_prefix_len=}, {new_indices=}" - assert new_prefix_len <= len(new_indices), f"{new_prefix_len=}, {new_indices=}" - self.req_to_token_pool.write( - (req.kv.req_pool_idx, slice(old_prefix_len, len(new_indices))), - new_indices[old_prefix_len:], - ) - - req.kv.cache_protected_len = len(new_indices) - - self.dec_lock_ref( - req.last_node, req.lock_receipt, skip_swa=req.swa_prefix_lock_released - ) - req.swa_prefix_lock_released = False - lock_receipt = self.inc_lock_ref(new_last_node).to_dec_params() - - # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later - if len(new_indices) < len(kv_indices): - req.prefix_indices = torch.cat( - [new_indices, kv_indices[len(new_indices) :]] - ) - else: - req.prefix_indices = new_indices - req.last_node = new_last_node - req.lock_receipt = lock_receipt - - def pretty_print(self) -> None: - self._print_helper(self.root_node, 0) - total_size, total_swa_size = self._total_size_helper() - print(f"#full_tokens: {total_size}, #swa_tokens: {total_swa_size}") - - def total_size(self) -> Tuple[int, int]: - return self._total_size_helper() - - def _free_node_value( - self, node: TreeNode, value: Optional[torch.Tensor] = None - ) -> Tuple[int, int]: - if value is None: - value = node.value - num_tokens = len(value) - if node.swa_tombstone: - # SWA peers went back in `dec_swa_lock_only` or an SWA evict, so - # only the full side is still ours; `free` would hand the SWA pool - # mapping entries that read as the padding slot. - self.token_to_kv_pool_allocator.free_full_segment(value, start_pos=0) - return num_tokens, 0 - self.token_to_kv_pool_allocator.free_segment(value, start_pos=0) - return num_tokens, num_tokens - - def evict(self, params: EvictParams) -> EvictResult: - if self.disable: - 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: - # get the least recently used leaf node that is not locked - x = self.full_lru_list.get_leaf_lru_no_lock() - - while full_num_evicted < full_num_tokens and self.full_lru_list.in_list(x): - assert x != self.root_node, ( - f"root node should not exist in full lru list, {x.id=}" - ) - assert x.full_lock_ref == 0, f"node is in use, {x.id=}" - - # 1. free node kv indices, evict full and swa tokens - self.kv_events.record_remove(x) - node_full_evicted, node_swa_evicted = self._free_node_value(x) - full_num_evicted += node_full_evicted - swa_num_evicted += node_swa_evicted - - # 2. get the next leaf, update the lru lists - x_next = self.full_lru_list.get_prev_leaf_no_lock(x) - self.full_lru_list.remove_node(x) - if not x.swa_tombstone: - self.swa_lru_list.remove_node(x) - - # 3. delete the leaf node - self._delete_leaf(x) - - # 4. Iteratively delete tombstone leaves to maintain invariant that leaf nodes are not tombstone - x, leaf_full_num_evicted = self._iteratively_delete_tombstone_leaf(x) - full_num_evicted += leaf_full_num_evicted - - # 5. if parent has no more children, it is a leaf. It is possible that this node is lru, so - # we need to get the first leaf node in the lru list - if len(x.parent.children) == 0: - x_next = self.full_lru_list.get_leaf_lru_no_lock() - - x = x_next - - if swa_num_evicted < swa_num_tokens: - # get the least recently used node that is not locked, doesn't have to be a leaf - x = self.swa_lru_list.get_lru_no_lock() - - # evict lru leaf nodes until swa_num_tokens is reached - while swa_num_evicted < swa_num_tokens and (self.swa_lru_list.in_list(x)): - assert not x.swa_tombstone, f"duplicate swa tombstone node, {x.id=}" - assert x != self.root_node, f"root node is not evictable, {x.id=}" - assert x.swa_lock_ref == 0, f"node is in use by swa kv indices, {x.id=}" - - if len(x.children) > 0: - # 1. an internal node, free swa tokens. - self.token_to_kv_pool_allocator.free_swa_segment( - x.value, start_pos=0 - ) - swa_num_evicted += len(x.value) - - # 2. get the next node, update the lru lists - x_next = self.swa_lru_list.get_prev_no_lock(x) - self.swa_lru_list.remove_node(x) - - # 3. tombstone the node - self._tombstone_internal_node(x) - elif x.full_lock_ref > 0: - # Leaf still holds a full-side lock (can happen when the - # SWA leaf-lock early-release optimization revived a - # tombstoned leaf. Treat it like an internal tombstone. - self.token_to_kv_pool_allocator.free_swa_segment( - x.value, start_pos=0 - ) - swa_num_evicted += len(x.value) - - x_next = self.swa_lru_list.get_prev_no_lock(x) - self.swa_lru_list.remove_node(x) - - self.swa_evictable_size_ -= len(x.value) - x.swa_tombstone = True - else: - assert x.full_lock_ref == 0, ( - f"leaf node with full lock must also have swa lock, {x.id=}" - ) - # 1. a leaf node, free full and swa tokens - self.kv_events.record_remove(x) - node_full_evicted, node_swa_evicted = self._free_node_value(x) - full_num_evicted += node_full_evicted - swa_num_evicted += node_swa_evicted - - # 2. get the next node, update the lru lists - x_next = self.swa_lru_list.get_prev_no_lock(x) - self.full_lru_list.remove_node(x) - self.swa_lru_list.remove_node(x) - - # 3. delete the leaf node - self._delete_leaf(x) - - # 4. Iteratively delete tombstone leaves to maintain invariant that leaf nodes are not tombstone - self._iteratively_delete_tombstone_leaf(x) - - 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) -> IncLockRefResult: - """ - Increment the lock reference count for the node. Returns the swa_uuid_for_lock, which needs - to be passed to dec_lock_ref. - It locks the full_lock_ref for nodes between the [last node, root), exclusive. - It locks the swa_lock_ref for nodes between the [last node, swa_uuid_for_lock], inclusive. - """ - if self.disable: - return IncLockRefResult() - - swa_lock_size = 0 - swa_uuid_for_lock = None - while node != self.root_node: - # lock full from node to root - assert node.full_lock_ref >= 0, ( - f"inc_lock_ref on node with {node.full_lock_ref=}, {node.id=}" - ) - if node.full_lock_ref == 0: - self.full_evictable_size_ -= len(node.value) - self.full_protected_size_ += len(node.value) - node.full_lock_ref += 1 - - # lock swa if we have not reached the sliding window size. - # When we reach the sliding window size, we will set the swa_uuid_for_lock. - # caller needs to pass the swa_uuid_for_lock to dec_lock_ref - if swa_lock_size < self.sliding_window_size: - assert not node.swa_tombstone, ( - f"inc_lock_swa on swa_tombstone node, {node.id=}" - ) - if node.swa_lock_ref == 0: - self.swa_evictable_size_ -= len(node.value) - self.swa_protected_size_ += len(node.value) - node.swa_lock_ref += 1 - swa_lock_size += len(node.value) - if swa_lock_size >= self.sliding_window_size: - if node.swa_uuid is None: - node.swa_uuid = gen_swa_uuid() - swa_uuid_for_lock = node.swa_uuid - node = node.parent - return IncLockRefResult(swa_uuid_for_lock=swa_uuid_for_lock) - - def dec_lock_ref( - self, - node: TreeNode, - params: Optional[DecLockRefParams] = None, - skip_swa: bool = False, - ) -> DecLockRefResult: - """ - Decrement the lock reference count for the node. - It unlocks the full_lock_ref for nodes between the [last node, root), exclusive. - It unlocks the swa_lock_ref for nodes between the [last node, swa_uuid_for_lock], inclusive. - If swa_uuid_for_lock is None, it unlocks to the root, exclusive. - - If skip_swa is True, only the full_lock_ref is decremented; the SWA lock is - assumed to have been released already (e.g. via `dec_swa_lock_only`). - """ - swa_uuid_for_lock = params.swa_uuid_for_lock if params is not None else None - - if self.disable: - return DecLockRefResult() - - dec_lock_swa = not skip_swa - while node != self.root_node: - assert node.full_lock_ref > 0, ( - f"dec_lock_ref on node with {node.full_lock_ref=}, {node.id=}" - ) - if node.full_lock_ref == 1: - self.full_evictable_size_ += len(node.value) - self.full_protected_size_ -= len(node.value) - node.full_lock_ref -= 1 - - if dec_lock_swa: - assert not node.swa_tombstone, ( - f"dec_lock_ref on swa_tombstone node, {node.id=}" - ) - assert node.swa_lock_ref > 0, ( - f"dec_lock_ref on node with {node.swa_lock_ref=}, {node.id=}" - ) - - if node.swa_lock_ref == 1: - self.swa_evictable_size_ += len(node.value) - self.swa_protected_size_ -= len(node.value) - node.swa_lock_ref -= 1 - if swa_uuid_for_lock and node.swa_uuid == swa_uuid_for_lock: - dec_lock_swa = False - - node = node.parent - - return DecLockRefResult() - - def dec_swa_lock_only( - self, - node: TreeNode, - params: DecLockRefParams, - ): - """ - Decrement only the swa_lock_ref (and swa_protected_size_) along the chain - [node, receipt boundary uuid], inclusive. The full_lock_ref is left - untouched so the caller's full-cache protection is preserved. Of the - receipt this cache consumes only ``swa_uuid_for_lock``; it has no - lower-priority components to drop. - - Used to early-release the SWA portion of a request's tree lock once the - request's decode position has advanced past the sliding window, so the - protected window can be reclaimed. - - For internal nodes, the standard protected -> evictable transition is - applied (node stays in swa_lru_list and may be evicted by SWA LRU later). - For leaf nodes, since `swa_lru_list` cannot contain a leaf with - `full_lock_ref > 0` (SWA-eviction would also delete the still-referenced - leaf), we instead free the SWA pool slots immediately and mark the leaf - as `swa_tombstone=True`. The full kv stays alive until the full-side - lock drops; future prefix-matches stop before this tombstoned leaf. - - Caller must ensure this is invoked at most once per (node, boundary uuid) - pair (track via e.g. `Req.swa_prefix_lock_released`). When the request - finally releases its full lock via `dec_lock_ref`, pass `skip_swa=True` - to avoid touching SWA state again. - """ - if self.disable: - return - swa_uuid_for_lock = params.swa_uuid_for_lock - - while node != self.root_node: - assert not node.swa_tombstone, ( - f"dec_swa_lock_only on swa_tombstone node, {node.id=}" - ) - assert node.swa_lock_ref > 0, ( - f"dec_swa_lock_only on node with {node.swa_lock_ref=}, {node.id=}" - ) - - if node.swa_lock_ref == 1: - self.swa_protected_size_ -= len(node.value) - if len(node.children) == 0: - # Leaf: free SWA pool slots and tombstone, and remove from - # swa_lru_list so SWA-eviction won't pick this tombstoned - # leaf (which still holds full_lock_ref > 0). The full kv - # stays alive until the request releases its full lock. - self.token_to_kv_pool_allocator.free_swa_segment( - node.value, start_pos=0 - ) - self.swa_lru_list.remove_node(node) - node.swa_tombstone = True - else: - # Internal: standard protected -> evictable. - self.swa_evictable_size_ += len(node.value) - node.swa_lock_ref -= 1 - - if swa_uuid_for_lock and node.swa_uuid == swa_uuid_for_lock: - break - node = node.parent - - def sanity_check(self): - self.full_lru_list.sanity_check(self) - self.swa_lru_list.sanity_check(self) - - def evictable_size(self) -> Tuple[int, int]: - # Note: use full_evictable_size() and swa_evictable_size() instead. - raise NotImplementedError - - def full_evictable_size(self) -> int: - return self.full_evictable_size_ - - def swa_evictable_size(self) -> int: - return self.swa_evictable_size_ - - def protected_size(self) -> Tuple[int, int]: - # Note: use full_protected_size() and swa_protected_size() instead. - raise NotImplementedError - - def full_protected_size(self) -> int: - # protected size refers to the size of the full cache that is locked - return self.full_protected_size_ - - def swa_protected_size(self) -> int: - # protected size refers to the size of the swa cache that is locked - return self.swa_protected_size_ - - def all_values_flatten(self) -> torch.Tensor: - values = [] - - def _dfs_helper(node: TreeNode): - for _, child in node.children.items(): - values.append(child.value) - _dfs_helper(child) - - _dfs_helper(self.root_node) - return torch.cat(values) - - def available_and_evictable_str(self) -> str: - full_available_size = self.token_to_kv_pool_allocator.full_available_size() - swa_available_size = self.token_to_kv_pool_allocator.swa_available_size() - full_evictable_size = self.full_evictable_size() - swa_evictable_size = self.swa_evictable_size() - return ( - f"Available full tokens: {full_available_size + full_evictable_size} ({full_available_size=} + {full_evictable_size=})\n" - f"Available swa tokens: {swa_available_size + swa_evictable_size} ({swa_available_size=} + {swa_evictable_size=})\n" - f"Full LRU list evictable size: {self.full_lru_list.sanity_check_evictable_size()}\n" - f"SWA LRU list evictable size: {self.swa_lru_list.sanity_check_evictable_size()}\n" - ) - - ##### Internal Helper Functions ##### - - def _match_prefix_helper( - self, key: RadixKey - ) -> Tuple[List[torch.Tensor], TreeNode, int]: - """ - SWA prefix matching helper. It factors in the sliding window size such that - the matched node is guaranteed to either 1. connected to root without swa tombstone, - or 2. the number of matching tokens from the matched node to the last swa tombstone - node is greater than or equal to the sliding window size. - """ - node = self.root_node - child_key = key.child_key(self.page_size) - - value = [] - # for path connected to root without tombstone, always match, so set to inf - match_len_since_tombstone = float("inf") - best_value_len = 0 - best_last_node = node - enable_compact = envs.SGLANG_OPT_SWA_RADIX_CACHE_COMPACT.get() - while len(key) > 0 and child_key in node.children.keys(): - child = node.children[child_key] - - if enable_compact: - self._compact_single_child_chain(child) - - if child.swa_tombstone: - # update best_value_len and best_last_node if needed - if match_len_since_tombstone >= self.sliding_window_size: - best_value_len = len(value) - best_last_node = node - # reset match_len_since_tombstone if we hit a tombstone node - match_len_since_tombstone = 0 - - prefix_len = child.key.match(key, page_size=self.page_size) - if prefix_len < len(child.key): - new_node = self._split_node(child.key, child, prefix_len) - value.append(new_node.value) - if not new_node.swa_tombstone: - match_len_since_tombstone += len(new_node.value) - node = new_node - break - else: - value.append(child.value) - if not child.swa_tombstone: - match_len_since_tombstone += len(child.value) - node = child - key = key[prefix_len:] - - if len(key): - child_key = key.child_key(self.page_size) - - # handle best_value_len and best_last_node, for the case that last node is fully matched - if match_len_since_tombstone >= self.sliding_window_size: - best_value_len = len(value) - best_last_node = node - - return value, best_last_node, best_value_len - - def _match_pre_processor(self, params: MatchPrefixParams) -> Optional[RadixKey]: - """Preprocess the key before matching.""" - key = params.key - key, _ = key.maybe_to_bigram_view(self.is_eagle) - if self.disable or len(key) == 0: - return None - key = key.page_aligned(self.page_size) - if len(key) == 0: - return None - return key - - def _match_post_processor( - self, - params: MatchPrefixParams, - value: List[torch.Tensor], - last_node: TreeNode, - best_value_len: int, - ) -> MatchResult: - """Post-process the matched result.""" - node_update = last_node - # update time for matched nodes, and make nodes closer to root to be least recently used - # this allows swa to evict nodes closer to root first - self.full_lru_list.reset_node_and_parents_mru(node_update, self.root_node) - self.swa_lru_list.reset_node_and_parents_mru(node_update, self.root_node) - - # This last_access_time is for sanity check, can be deleted after validation in production - cur_time = get_last_access_time() - while node_update: - node_update.last_access_time = cur_time - cur_time -= ( - 0.00001 # assuming less than 100000 nodes in a branch of the tree - ) - node_update = node_update.parent - - value = value[:best_value_len] - if value: - value = torch.cat(value) - else: - value = torch.empty((0,), dtype=torch.int64, device=self.device) - - return MatchResult( - device_indices=value, - last_device_node=last_node, - last_host_node=last_node, - best_match_node=last_node, - ) - - def _compact_single_child_chain(self, node: TreeNode) -> None: - # FIXME(ispobock): drifts retract pool accounting (commit 6348cb506); - # also overwrites active swa_uuid when window > page_size. Off by - # default via SGLANG_OPT_SWA_RADIX_CACHE_COMPACT. - while len(node.children) == 1: - child = next(iter(node.children.values())) - if len(child.children) == 0: - break - sum_gc_full_lock_ref = sum( - gc.full_lock_ref for gc in child.children.values() - ) - if child.full_lock_ref > sum_gc_full_lock_ref: - break - if ( - child.swa_tombstone != node.swa_tombstone - or child.full_lock_ref != node.full_lock_ref - or child.swa_lock_ref != node.swa_lock_ref - ): - break - - # Preserve is_bigram: main #23106 made bigram an O(1) flag on RadixKey; - # the constructor defaults to False, so concat without explicit flag - # silently demotes EAGLE/MTP bigram keys → match() returns 0 → - # _split_node assert. - node.key = RadixKey( - node.key.token_ids + child.key.token_ids, - node.key.extra_key, - is_bigram=node.key.is_bigram, - cache_salt=node.key.cache_salt, - ) - node.value = torch.cat([node.value, child.value]) - node.children = child.children - for grandchild in node.children.values(): - grandchild.parent = node - - if child.swa_uuid is not None: - node.swa_uuid = child.swa_uuid - - if node.hash_value is not None and child.hash_value is not None: - node.hash_value = list(node.hash_value) + list(child.hash_value) - else: - node.hash_value = None - if node.event_hash_value is not None and child.event_hash_value is not None: - node.event_hash_value = list(node.event_hash_value) + list( - child.event_hash_value - ) - else: - node.event_hash_value = None - - self.full_lru_list.remove_node(child) - if not child.swa_tombstone: - self.swa_lru_list.remove_node(child) - - def _maybe_split_leaf_for_swa_lock(self, leaf: TreeNode) -> TreeNode: - """``inc_lock_ref`` protects ``len(leaf.value)`` SWA tokens for the - leaf even though SWA only actually needs the last - ``sliding_window_size`` tokens. With chunked prefill, leaves can be - thousands of tokens long, which inflates ``swa_protected_size_`` by - ~``chunked_prefill_size / sliding_window_size`` and causes premature - SWA pool exhaustion / retract thrashing. - """ - if ( - leaf is self.root_node - or leaf.swa_lock_ref > 0 - or leaf.swa_tombstone - or len(leaf.value) == 0 - ): - return leaf - - # Smallest page-aligned size that still covers the sliding window. - tail_size = ( - (self.sliding_window_size + self.page_size - 1) - // self.page_size - * self.page_size - ) - if len(leaf.value) <= tail_size: - return leaf - - split_at = len(leaf.value) - tail_size - - if split_at <= 0 or split_at >= len(leaf.value): - return leaf - if self.page_size > 1 and ( - split_at % self.page_size != 0 or len(leaf.value) % self.page_size != 0 - ): - return leaf - - self._split_node(leaf.key, leaf, split_at) - return leaf - - def _split_node(self, key: RadixKey, child: TreeNode, split_len: int) -> TreeNode: - # new_node -> child - new_node = TreeNode() - new_node.children = {key[split_len:].child_key(self.page_size): child} - new_node.parent = child.parent - new_node.swa_tombstone = child.swa_tombstone - new_node.full_lock_ref = child.full_lock_ref - new_node.swa_lock_ref = child.swa_lock_ref - new_node.key = child.key[:split_len] - assert len(new_node.key) > 0, f"new_node.key should not be empty" - new_node.value = child.value[:split_len].clone() - # parent inherits the swa_uuid from child for swa lock ref - new_node.swa_uuid = child.swa_uuid - child.swa_uuid = None - # child time should be later than parent's time for swa tombstone - child.last_access_time = get_last_access_time() - - # remove the child from the lru lists because it is being split - self.full_lru_list.remove_node(child) - if not new_node.swa_tombstone: - self.swa_lru_list.remove_node(child) - child.parent = new_node - child.key = child.key[split_len:] - assert len(child.key) > 0, f"child.key should not be empty" - child.value = child.value[split_len:].clone() - new_node.parent.children[key.child_key(self.page_size)] = new_node - new_node.hash_value, child.hash_value = split_node_hash_value( - child.hash_value, split_len, self.page_size - ) - new_node.event_hash_value, child.event_hash_value = split_node_hash_value( - child.event_hash_value, split_len, self.page_size - ) - - # insert the new node and child into the lru lists, insert - # parent first so that parent is after child in the lru list - self.full_lru_list.insert_mru(new_node) - self.full_lru_list.insert_mru(child) - if not new_node.swa_tombstone: - self.swa_lru_list.insert_mru(new_node) - self.swa_lru_list.insert_mru(child) - return new_node - - def _insert_helper( - self, - node: TreeNode, - key: RadixKey, - value, - update_kv_after_len: int, - swa_evicted_seqlen: int = 0, - ) -> int: - # Update the last access time from root to leaf, so that - # swa will tombstone the node closer to root first - node.last_access_time = get_last_access_time() - if node != self.root_node: - self.full_lru_list.reset_node_mru(node) - if not node.swa_tombstone: - self.swa_lru_list.reset_node_mru(node) - if len(key) == 0: - return 0 - - child_key = key.child_key(self.page_size) - - total_prefix_length = 0 - while len(key) > 0 and child_key in node.children.keys(): - node = node.children[child_key] - node.last_access_time = get_last_access_time() - self.full_lru_list.reset_node_mru(node) - if not node.swa_tombstone: - self.swa_lru_list.reset_node_mru(node) - prefix_len = node.key.match(key, page_size=self.page_size) - - if prefix_len < len(node.key): - new_node = self._split_node(node.key, node, prefix_len) - node = new_node - - # if tombstone after update_kv_after_len, update node.value to be the input value. - # This is needed because it is possible that the last sliding window size tokens - # contains tombstone. If this is the case and we don't update the kv value, then - # the prefill prefix matching will stuck. - if update_kv_after_len < total_prefix_length + prefix_len: - # For page_size > 1 and chunked prefill case, update_kv_after_len may be not page-aligned due to a trailing partial page - # (kept in the request but not inserted into the radix tree) appended to prefix_indices. - if node.swa_tombstone: - assert node.swa_lock_ref == 0, ( - f"tombstone swa_lock_ref should always be 0, {node.full_lock_ref=}, {node.swa_lock_ref=}, {node.id=}" - ) - assert swa_evicted_seqlen % self.page_size == 0, ( - f"swa_evicted_seqlen must be page aligned, {swa_evicted_seqlen=}, {self.page_size=}" - ) - if swa_evicted_seqlen <= total_prefix_length: - # Branch 1: all swa tokens of value[:prefix_len] are not evicted, so we can insert it to the tree directly. - if node.full_lock_ref > 0: - # Full KV is still locked by a running request. Keep it - # and adopt the incoming SWA instead of freeing in-flight - # Full slots. - self._recover_tombstone_keeping_locked_full( - node, value[:prefix_len] - ) - else: - # Free full tokens in the original tree node. - self.token_to_kv_pool_allocator.free_full_segment( - node.value[:prefix_len], start_pos=0 - ) - # Overwrite the new value in request to the tree node. - node.value = value[:prefix_len].clone() - node.swa_tombstone = False - self.swa_lru_list.insert_mru(node) - self.swa_evictable_size_ += len(node.value) - elif swa_evicted_seqlen < total_prefix_length + prefix_len: - # Branch 2: part of swa tokens of value[:prefix_len] are evicted, so we need to split the node and insert the value to new node. - start_update_idx = swa_evicted_seqlen - total_prefix_length - if node.full_lock_ref > 0: - # Split first so the recovered suffix keeps the locked - # Full slots, then adopt the incoming SWA for that suffix. - self._split_node(node.key, node, start_update_idx) - self._recover_tombstone_keeping_locked_full( - node, value[start_update_idx:prefix_len] - ) - self.token_to_kv_pool_allocator.free_full_segment( - value[:start_update_idx], start_pos=0 - ) - else: - self.token_to_kv_pool_allocator.free_full_segment( - node.value[start_update_idx:prefix_len], start_pos=0 - ) - self._split_node(node.key, node, start_update_idx) - # Here node is the new node after split, so we can overwrite the value to the new node. - # The old node is still swa tombstone and the full token is not freed. - node.value = value[start_update_idx:prefix_len].clone() - self.token_to_kv_pool_allocator.free_full_segment( - value[:start_update_idx], start_pos=0 - ) - node.swa_tombstone = False - self.swa_lru_list.insert_mru(node) - self.swa_evictable_size_ += len(node.value) - else: - # Branch 3: all swa tokens of value[:prefix_len] are evicted, so we don't need to update the node. - self.token_to_kv_pool_allocator.free_full_segment( - value[:prefix_len], start_pos=0 - ) - else: - # The node is not tombstone, so we don't need to update the node. - # The incoming slice can still straddle this request's own - # eviction floor, so split it there. - free_kv_row_segments( - self.token_to_kv_pool_allocator, - [(value[:prefix_len], total_prefix_length)], - swa_evicted_seqlen=swa_evicted_seqlen, - ) - - total_prefix_length += prefix_len - key = key[prefix_len:] - value = value[prefix_len:] - - if len(key): - child_key = key.child_key(self.page_size) - - if len(key): - # Layout: |--- total_prefix_length ---|--- len(key) ---| - # ^ ^ ^ - # 0 total_prefix_length total_length - # - # Cases based on swa_evicted_seqlen position: - # 1. swa_evicted_seqlen <= total_prefix_length: - # Already handled in the while loop above. All of len(key) is non-tombstone. - # 2. total_prefix_length < swa_evicted_seqlen < total_length: - # Split: [total_prefix_length, swa_evicted_seqlen) as tombstone, - # [swa_evicted_seqlen, total_length) as non-tombstone. - # 3. swa_evicted_seqlen == total_length: - # All remaining tokens are evicted. Free value and return without - # creating a node (leaf nodes must not be tombstone). - # Note: the -page_size fix in _evict_swa prevents this case from - # occurring in normal operation. This check is a defensive guard - # against unexpected eviction states from other code paths. - if swa_evicted_seqlen == total_prefix_length + len(key): - self.token_to_kv_pool_allocator.free_full_segment(value, start_pos=0) - return total_prefix_length - - if ( - swa_evicted_seqlen > total_prefix_length - and swa_evicted_seqlen < total_prefix_length + len(key) - ): - swa_tombstone_len = swa_evicted_seqlen - total_prefix_length - node = self._add_new_node( - node, - key[:swa_tombstone_len], - value[:swa_tombstone_len], - swa_tombstone=True, - ) - key = key[swa_tombstone_len:] - value = value[swa_tombstone_len:] - - new_leaf = self._add_new_node(node, key, value, swa_tombstone=False) - - if envs.SGLANG_OPT_SWA_SPLIT_LEAF_ON_INSERT.get(): - # Cap the leaf at one (page-aligned) sliding window so a future - # inc_lock_ref only protects `sliding_window_size` tokens of SWA pool. - self._maybe_split_leaf_for_swa_lock(new_leaf) - - return total_prefix_length - - def _recover_tombstone_keeping_locked_full( - self, node: TreeNode, incoming_full: torch.Tensor - ) -> None: - """Recover a tombstoned node whose Full KV is locked by a running request. - - Keep node.value, the locked Full slots, and re-point its full->SWA - mapping at the incoming request's fresh SWA. Free only the incoming - redundant Full slots, not their SWA slots. - """ - assert len(node.value) == len(incoming_full), ( - f"locked-full recover size mismatch: {len(node.value)=}, " - f"{len(incoming_full)=}" - ) - - allocator = self.token_to_kv_pool_allocator - swa_value = allocator.translate_loc_from_full_to_swa(incoming_full) - allocator.set_full_to_swa_mapping(node.value, swa_value) - allocator.clear_full_to_swa_mapping(incoming_full) - allocator.free_full_segment(incoming_full, start_pos=0) - - node.swa_tombstone = False - self.swa_lru_list.insert_mru(node) - self.swa_evictable_size_ += len(node.value) - - def _add_new_node( - self, - parent: TreeNode, - key: RadixKey, - value: torch.Tensor, - swa_tombstone: bool = False, - ) -> TreeNode: - assert len(key) > 0, f"key should not be empty" - new_node = TreeNode() - new_node.parent = parent - new_node.key = key - new_node.value = value.clone() - new_node.swa_tombstone = swa_tombstone - parent.children[key.child_key(self.page_size)] = new_node - self.full_lru_list.insert_mru(new_node) - self.full_evictable_size_ += len(value) - if not swa_tombstone: - self.swa_lru_list.insert_mru(new_node) - self.swa_evictable_size_ += len(value) - self.kv_events.record_store(new_node) - return new_node - - def _iteratively_delete_tombstone_leaf( - self, node: TreeNode - ) -> Tuple[TreeNode, int]: - full_num_evicted = 0 - while node.parent.swa_tombstone and len(node.parent.children) == 0: - # root node is not evictable - if node.parent == self.root_node: - break - # if locked, means node is in use, skip - if node.parent.full_lock_ref > 0: - break - assert node.parent.swa_lock_ref == 0, ( - f"tombstone swa_lock_ref should always be 0, {node.parent.full_lock_ref=}, {node.parent.swa_lock_ref=}, {node.parent.id=}" - ) - # delete tombstone node evicts full tokens - self.kv_events.record_remove(node.parent) - node_full_evicted, _ = self._free_node_value(node.parent) - full_num_evicted += node_full_evicted - self.full_lru_list.remove_node(node.parent) - self._delete_tombstone_leaf(node.parent) - node = node.parent - - return node, full_num_evicted - - def _delete_leaf(self, node: TreeNode) -> None: - assert len(node.children) == 0, f"leaf node has children, {node.id=}" - key = node.key.child_key(self.page_size) - v = node.parent.children.pop(key, None) - assert v == node, f"parent does not have child key, {key}" - self.full_evictable_size_ -= len(node.key) - # Tombstoned leaves were never (re-)added to swa_lru_list and were - # already removed from swa_evictable_size_ when they were tombstoned. - if not node.swa_tombstone: - self.swa_evictable_size_ -= len(node.key) - - def _tombstone_internal_node(self, node: TreeNode) -> None: - assert len(node.children) != 0, f"Cannot tombstone a leaf node, {node.id=}" - node.swa_tombstone = True - self.swa_evictable_size_ -= len(node.key) - - def _delete_tombstone_leaf(self, node: TreeNode) -> None: - assert node.swa_tombstone, ( - f"Deleting a unexpected non-tombstone leaf node, {node.id=}" - ) - assert len(node.children) == 0, f"leaf node has children, {node.id=}" - key = node.key.child_key(self.page_size) - v = node.parent.children.pop(key, None) - assert v == node, f"parent does not have child key, {key}" - - self.full_evictable_size_ -= len(node.key) - - def _collect_nontombstone_nodes(self) -> List[TreeNode]: - ret_list = [] - stack = [self.root_node] - - while stack: - cur_node = stack.pop() - if not cur_node.swa_tombstone: - ret_list.append(cur_node) - stack.extend(cur_node.children.values()) - - return ret_list - - def _collect_all_nodes(self) -> List[TreeNode]: - ret_list = [] - stack = [self.root_node] - while stack: - cur_node = stack.pop() - ret_list.append(cur_node) - stack.extend(cur_node.children.values()) - return ret_list - - def _print_helper(self, node: TreeNode, indent: int) -> None: - """Prints the radix tree in a human-readable format.""" - stack = [(node, indent)] - while stack: - current_node, current_indent = stack.pop() - print( - " " * current_indent, - current_node.id, - len(current_node.key), - f"fr={current_node.full_lock_ref}", - f"sr={current_node.swa_lock_ref}", - f"fll={self.full_lru_list.in_list(current_node)}", - f"sll={self.swa_lru_list.in_list(current_node)}", - f"ts={current_node.swa_tombstone}", - ) - for key, child in current_node.children.items(): - stack.append((child, current_indent + 2)) - - assert key == child.key.child_key(self.page_size), ( - f"{key=}, {child.key.child_key(self.page_size)=}" - ) - - def _total_size_helper(self) -> Tuple[int, int]: - total_size = 0 - total_swa_size = 0 - stack = [self.root_node] - while stack: - current_node = stack.pop() - total_size += len(current_node.value) - if not current_node.swa_tombstone: - total_swa_size += len(current_node.value) - for child in current_node.children.values(): - if child.evicted: - continue - stack.append(child) - return total_size, total_swa_size diff --git a/python/sglang/srt/mem_cache/unified_cache/components/README.md b/python/sglang/srt/mem_cache/unified_cache/components/README.md index c5d0e5cf1..491c5f78e 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/README.md +++ b/python/sglang/srt/mem_cache/unified_cache/components/README.md @@ -4,7 +4,7 @@ A component-based, pluggable prefix cache framework for SGLang that unifies Full ## Design Goals -1. **Unified tree structure** — One radix tree manages all KV cache types instead of separate specialized implementations (`SWARadixCache`, `MambaRadixCache`, etc.). +1. **Unified tree structure** — One radix tree manages all KV cache types, replacing the separate specialized implementations that preceded it. 2. **Pluggable components** — Each attention/state type (Full, SWA, Mamba) is a `TreeComponent` that implements hook interfaces. Adding a new cache type only requires adding a new component. 3. **Per-component resource isolation** — Each component has its own lock reference counting, evictable/protected size tracking, and eviction driver. Auxiliary components use per-component LRUs; Full uses device/host leaf sets. 4. **Cascade eviction with priority** — When a component evicts a node, lower-or-equal-priority components on the same node are evicted together, maintaining cross-component consistency. diff --git a/python/sglang/srt/mem_cache/unified_cache/components/mamba.py b/python/sglang/srt/mem_cache/unified_cache/components/mamba.py index 3c18cee30..1070e0b5c 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/mamba.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/mamba.py @@ -537,8 +537,7 @@ class MambaComponent(TreeComponent): # slot's unflushed ring depth (`write_pos`), so on request finish cap # the donate to the last flush boundary (where temporal is current) # and reset the cursor, keeping the donated checkpoint consistent with - # its key length. page_size is asserted == 1, so no realign. Mirrors - # MambaRadixCache.cache_finished_req. + # its key length. page_size is asserted == 1, so no realign. if is_finished: write_pos_buf = ( self.cache.req_to_token_pool.mamba_pool.replayssm_write_pos diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 22ed8e168..1db8f8e6d 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -4611,7 +4611,7 @@ def get_extend_input_len_swa_limit( sliding_window_size: int, chunked_prefill_size: int, page_size: int ) -> int: # 1. a factor of 2x is because each prefill contains chunked_prefill_size tokens, - # and between prefills, we run swa_radix_cache.cache_unfinished_req(), + # and between prefills, we run the tree cache's cache_unfinished_req(), # so we unlock the previously locked nodes. # 2. max is to handle the case that chunked_prefill_size is larger than sliding_window_size. # in that case, each prefill contains chunked_prefill_size tokens, diff --git a/python/sglang/test/scripted_runtime/context/radix.py b/python/sglang/test/scripted_runtime/context/radix.py index 9ca8c0f2a..91403b2c0 100644 --- a/python/sglang/test/scripted_runtime/context/radix.py +++ b/python/sglang/test/scripted_runtime/context/radix.py @@ -2,7 +2,6 @@ from __future__ import annotations from typing import TYPE_CHECKING, Any, Callable, Dict -from sglang.srt.mem_cache.swa_radix_cache import TreeNode as SWATreeNode from sglang.srt.mem_cache.unified_radix_cache import UnifiedTreeNode if TYPE_CHECKING: @@ -18,8 +17,6 @@ def get_all_node_lock_refs(ctx: ScriptedContext) -> Dict[int, int]: def _node_lock_ref(node: Any) -> int: - if isinstance(node, SWATreeNode): - return node.full_lock_ref + node.swa_lock_ref if isinstance(node, UnifiedTreeNode): return sum(cd.lock_ref for cd in node.component_data) return node.lock_ref diff --git a/test/manual/dsv4/test_dsv4_swa_radix_retract.py b/test/manual/dsv4/test_dsv4_swa_radix_retract.py index 95fe2a807..7db064d58 100644 --- a/test/manual/dsv4/test_dsv4_swa_radix_retract.py +++ b/test/manual/dsv4/test_dsv4_swa_radix_retract.py @@ -1,7 +1,10 @@ """DSV4 stress test for SWA radix cache + tombstone + retract interaction. -Reproduces the assert in `swa_radix_cache.cache_unfinished_req`: +Regression test for the former SWA `cache_unfinished_req` assertion: assert old_prefix_len <= len(new_indices) +The unified cache reads `req.kv.cache_protected_len` and tolerates page_size - 1 +of alignment slack, so this reproduces the historical trip conditions rather +than a line that still exists. Trip conditions (all required): 1. Fork-only SWA leaf early-release on (`SGLANG_OPT_SWA_RELEASE_LEAF_LOCK_AFTER_WINDOW=1`) @@ -88,7 +91,6 @@ class TestDSV4FlashSWARadixRetract(CustomTestCase): env = { "SGLANG_DSV4_FP4_EXPERTS": "0", "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024", - "SGLANG_OPT_SWA_RADIX_CACHE_COMPACT": "0", "SGLANG_TEST_RETRACT": "1", "SGLANG_TEST_RETRACT_INTERVAL": "3", } @@ -130,7 +132,7 @@ class TestDSV4FlashSWARadixRetract(CustomTestCase): """Stress: 64 concurrent long-prompt reqs with long generation force retract under SWA pool pressure. Reqs share a 30k+ token prefix so tombstoned leaves from retracted reqs are on the radix path of new - reqs. Scheduler must not crash on the swa_radix_cache assert.""" + reqs. Scheduler must not crash on the SWA insert assert.""" random.seed(0) concurrency = 64 diff --git a/test/registered/kv_canary/test_self_unit_radix_walker.py b/test/registered/kv_canary/test_self_unit_radix_walker.py index 7bc6c446c..3d94ca4ae 100644 --- a/test/registered/kv_canary/test_self_unit_radix_walker.py +++ b/test/registered/kv_canary/test_self_unit_radix_walker.py @@ -9,7 +9,6 @@ import torch from sglang.srt.kv_canary.radix_cache_walker import walk_radix_cache_for_canary from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.radix_cache import RadixKey -from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache, TreeNode from sglang.srt.mem_cache.unified_cache.components import ( BASE_COMPONENT_TYPE, ComponentType, @@ -75,65 +74,6 @@ class TestSelfUnitRadixWalker(CustomTestCase): result = walk_radix_cache_for_canary(radix_cache=cache, unlocked_only=True) self.assertEqual(result.slot_indices.tolist(), [3, 4]) - def test_walk_unlocked_only_uses_swa_full_lock_ref(self): - """Verify SWA radix walking honors full-pool lock references.""" - cache = SWARadixCache.__new__(SWARadixCache) - cache.device = self.device - cache.page_size = 1 - cache.disable = False - - root = TreeNode() - root.value = torch.tensor([], dtype=torch.int32, device=self.device) - cache.root_node = root - - locked_child = TreeNode() - locked_child.value = torch.tensor([1, 2], dtype=torch.int32, device=self.device) - locked_child.parent = root - locked_child.full_lock_ref = 1 - root.children[locked_child.id] = locked_child - - unlocked_child = TreeNode() - unlocked_child.value = torch.tensor( - [3, 4], dtype=torch.int32, device=self.device - ) - unlocked_child.parent = root - root.children[unlocked_child.id] = unlocked_child - - result = walk_radix_cache_for_canary(radix_cache=cache, unlocked_only=True) - self.assertEqual(result.slot_indices.tolist(), [3, 4]) - - def test_swa_resident_only_skips_tombstoned_nodes(self): - """Verify SWA radix walking skips nodes whose SWA storage was evicted.""" - cache = SWARadixCache.__new__(SWARadixCache) - cache.device = self.device - cache.page_size = 1 - cache.disable = False - - root = TreeNode() - root.value = torch.tensor([], dtype=torch.int32, device=self.device) - cache.root_node = root - - tombstoned_child = TreeNode() - tombstoned_child.value = torch.tensor( - [1, 2], dtype=torch.int32, device=self.device - ) - tombstoned_child.parent = root - tombstoned_child.swa_tombstone = True - root.children[tombstoned_child.id] = tombstoned_child - - resident_child = TreeNode() - resident_child.value = torch.tensor( - [3, 4], dtype=torch.int32, device=self.device - ) - resident_child.parent = root - root.children[resident_child.id] = resident_child - - result = walk_radix_cache_for_canary( - radix_cache=cache, - swa_resident_only=True, - ) - self.assertEqual(result.slot_indices.tolist(), [3, 4]) - def test_unified_swa_sweep_gates_on_swa_lock_not_full_lock(self): """With unlocked_only + swa_resident_only, the sweep filters on the SWA component lock: a FULL-locked node whose SWA lock was already released diff --git a/test/registered/unit/mem_cache/test_mamba_unittest.py b/test/registered/unit/mem_cache/test_mamba_unittest.py index 410fc3eea..e2d0c21e1 100755 --- a/test/registered/unit/mem_cache/test_mamba_unittest.py +++ b/test/registered/unit/mem_cache/test_mamba_unittest.py @@ -7,24 +7,14 @@ import torch from sglang.kernels.ops.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape -from sglang.srt.disaggregation.kv_events import BlockRemoved, BlockStored from sglang.srt.environ import envs 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 ( - EvictParams, - InsertParams, - MatchPrefixParams, -) -from sglang.srt.mem_cache.cache_init_params import CacheInitParams -from sglang.srt.mem_cache.common import available_and_evictable_str -from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache from sglang.srt.mem_cache.memory_pool import ( HybridLinearKVPool, HybridReqToTokenPool, MambaPool, ) -from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler from sglang.srt.utils import get_device @@ -241,338 +231,10 @@ class TestMamba(unittest.TestCase): view[0, 0, 0, 1, 0] = -1 self.assertEqual(view[0, 0, 1, 0, 0].item(), -1) - def test_mamba_radix_cache_1(self): - tree, allocator, req_to_token_pool, make_dummy_req = ( - self._setup_tree_and_allocator() - ) - mamba_allocator = req_to_token_pool.mamba_allocator - mamba_pool = req_to_token_pool.mamba_pool - # test - print( - f"[Start] allocator mamba available size: {mamba_allocator.available_size()}, full available size: {allocator.available_size()}" - ) - req1 = make_dummy_req() - req1_token_ids, req1_kv_indices = [1, 2, 3], allocator.alloc(3) - assert len(req1_token_ids) == len(req1_kv_indices) - print( - f"req1: inserting, req1_token_ids: {req1_token_ids}, req1_kv_indices: {req1_kv_indices}" - ) - key = RadixKey(array("q", req1_token_ids)) - result = tree.insert( - InsertParams( - key=key, - value=req1_kv_indices[: len(key)], - mamba_value=req1.kv.mamba_pool_idx.unsqueeze(0), - ) - ) - prefix_len = result.prefix_len - print( - f"req1: prefix_len: {prefix_len}, allocator mamba available size: {mamba_allocator.available_size()}, full available size: {allocator.available_size()}" - ) - req2 = make_dummy_req() - req2_token_ids, req2_kv_indices = [1, 2, 3, 4, 5, 6, 7], allocator.alloc(7) - assert len(req2_token_ids) == len(req2_kv_indices) - print( - f"req2: inserting, req2_token_ids: {req2_token_ids}, req2_kv_indices: {req2_kv_indices}" - ) - key = RadixKey(array("q", req2_token_ids)) - result = tree.insert( - InsertParams( - key=key, - value=req2_kv_indices[: len(key)], - mamba_value=req2.kv.mamba_pool_idx.unsqueeze(0), - ) - ) - prefix_len = result.prefix_len - print( - f"req2: prefix_len: {prefix_len}, allocator mamba available size: {mamba_allocator.available_size()}, full available size: {allocator.available_size()}" - ) - - req3 = make_dummy_req() - req3_token_ids, req3_kv_indices = [10, 11, 12], allocator.alloc(3) - assert len(req3_token_ids) == len(req3_kv_indices) - print( - f"req3: inserting, req3_token_ids: {req3_token_ids}, req3_kv_indices: {req3_kv_indices}" - ) - key = RadixKey(array("q", req3_token_ids)) - result = tree.insert( - InsertParams( - key=key, - value=req3_kv_indices[: len(key)], - mamba_value=req3.kv.mamba_pool_idx.unsqueeze(0), - ) - ) - prefix_len = result.prefix_len - print( - f"req3: prefix_len: {prefix_len}, allocator mamba available size: {mamba_allocator.available_size()}, full available size: {allocator.available_size()}" - ) - req4 = make_dummy_req() - req4_token_ids, req4_kv_indices = [1, 2, 3, 4, 5, 60, 70], allocator.alloc(7) - assert len(req4_token_ids) == len(req4_kv_indices) - print( - f"req4: inserting, req4_token_ids: {req4_token_ids}, req4_kv_indices: {req4_kv_indices}" - ) - key = RadixKey(array("q", req4_token_ids)) - result = tree.insert( - InsertParams( - key=key, - value=req4_kv_indices[: len(key)], - mamba_value=req4.kv.mamba_pool_idx.unsqueeze(0), - ) - ) - prefix_len = result.prefix_len - print( - f"req4: prefix_len: {prefix_len}, allocator mamba available size: {mamba_allocator.available_size()}, full available size: {allocator.available_size()}" - ) - - tree.pretty_print() - full_num_tokens = 1 - print(f"evicting {full_num_tokens} full token") - 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") - 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] - result = tree.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", req5_token_ids))) - ) - kv_indices, last_node = result.device_indices, result.last_device_node - print( - f"req5: token_ids: {req5_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}" - ) - assert len(kv_indices) == 0 - - req6_token_ids = [1, 2, 3, 4, 5, 60, 70] - result = tree.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", req6_token_ids))) - ) - kv_indices, last_node = result.device_indices, result.last_device_node - print( - f"req6: token_ids: {req6_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}" - ) - assert len(kv_indices) == 7 - assert len(last_node.key) == 2 - - req7_token_ids = [1, 2, 3, 4, 5, 6, 7] - result = tree.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", req7_token_ids))) - ) - kv_indices, last_node = result.device_indices, result.last_device_node - print( - f"req7: token_ids: {req7_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}" - ) - assert len(kv_indices) == 7 - assert len(last_node.key) == 2 - - mamba_num = 1 - print(f"evicting {mamba_num} mamba") - 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] - result = tree.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", req8_token_ids))) - ) - kv_indices, last_node = result.device_indices, result.last_device_node - print( - f"req8: token_ids: {req8_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}" - ) - assert len(kv_indices) == 0 - assert len(last_node.key) == 0 - - req9_token_ids = [1, 2, 3, 4, 5, 6, 7] - req9 = make_dummy_req() - result = tree.match_prefix( - MatchPrefixParams( - key=RadixKey(array("q", req9_token_ids)), req=req9, cow_mamba=True - ) - ) - kv_indices, last_node = result.device_indices, result.last_device_node - assert req9.kv.holds_mamba - assert torch.all( - mamba_pool.mamba_cache.conv[0][:, req9.kv.mamba_pool_idx] - == mamba_pool.mamba_cache.conv[0][:, last_node.mamba_value] - ) - assert torch.all( - mamba_pool.mamba_cache.temporal[:, req9.kv.mamba_pool_idx] - == mamba_pool.mamba_cache.temporal[:, last_node.mamba_value] - ) - - print(tree.available_and_evictable_str()) - print(available_and_evictable_str(tree)) - tree.sanity_check() - - def test_mamba_lru_match_refreshes_only_used_node(self): - """A prefix-cache hit must refresh only the matched leaf's mamba state in - the mamba LRU, not its ancestors. Whole-chain refresh clustered a session's - states adjacently, so under mamba-pool pressure eviction dropped whole cold - sessions instead of the intermediate states reuse never needs. Guards against - reverting the mamba list to reset_node_and_parents_mru. - """ - tree, allocator, req_to_token_pool, make_dummy_req = ( - self._setup_tree_and_allocator() - ) - - def insert(token_ids): - req = make_dummy_req() - kv = allocator.alloc(len(token_ids)) - tree.insert( - InsertParams( - key=RadixKey(array("q", token_ids)), - value=kv, - mamba_value=req.kv.mamba_pool_idx.unsqueeze(0), - ) - ) - - def match_leaf(token_ids): - return tree.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", token_ids))) - ).last_device_node - - def mamba_lru_mru_to_lru(): - lst = tree.mamba_lru_list - order, x = [], getattr(lst.head, lst.nxt) - while x is not None and x is not lst.tail and x.id in lst.cache: - order.append(x) - x = getattr(x, lst.nxt) - return order - - # Two independent sessions, each a 2-node mamba chain: - # root -> a1 -> b1 and root -> a2 -> b2 - insert([1, 2, 3]) - insert([1, 2, 3, 4, 5, 6]) - insert([7, 8, 9]) - insert([7, 8, 9, 10, 11, 12]) - - b1 = match_leaf([1, 2, 3, 4, 5, 6]) - a1 = b1.parent - b2 = match_leaf([7, 8, 9, 10, 11, 12]) - a2 = b2.parent - # Session 2 was matched last, so session 1's ancestor a1 is older than a2. - order = mamba_lru_mru_to_lru() - self.assertGreater(order.index(a1), order.index(a2)) - - # Re-access session 1. Only its consumed leaf (b1) moves to MRU; its ancestor - # a1 must stay put -- whole-chain reset would bump a1 right behind b1, making - # it newer than a2. - self.assertIs(match_leaf([1, 2, 3, 4, 5, 6]), b1) - order = mamba_lru_mru_to_lru() - self.assertIs(order[0], b1) - self.assertGreater(order.index(a1), order.index(a2)) - tree.sanity_check() - - def test_mamba_radix_cache_kv_events(self): - tree, allocator, _, make_dummy_req = self._setup_tree_and_allocator( - enable_kv_cache_events=True - ) - tree.take_events() # Clear the reset event. - - stored_hashes = [] - - req1 = make_dummy_req() - key1 = RadixKey(array("q", [1, 2, 3])) - tree.insert( - InsertParams( - key=key1, - value=allocator.alloc(3)[: len(key1)], - mamba_value=req1.kv.mamba_pool_idx.unsqueeze(0), - ) - ) - events = tree.take_events() - stored_events = [e for e in events if isinstance(e, BlockStored)] - self.assertEqual(len(stored_events), 1) - self.assertEqual(list(stored_events[0].token_ids), [1, 2, 3]) - stored_hashes.extend( - block_hash for event in stored_events for block_hash in event.block_hashes - ) - - req2 = make_dummy_req() - key2 = RadixKey(array("q", [1, 2, 3, 4, 5])) - tree.insert( - InsertParams( - key=key2, - value=allocator.alloc(5)[: len(key2)], - mamba_value=req2.kv.mamba_pool_idx.unsqueeze(0), - ) - ) - events = tree.take_events() - stored_events = [e for e in events if isinstance(e, BlockStored)] - self.assertEqual(len(stored_events), 1) - self.assertEqual(list(stored_events[0].token_ids), [4, 5]) - stored_hashes.extend( - block_hash for event in stored_events for block_hash in event.block_hashes - ) - - # Evicting an internal mamba state creates a tombstone but does not - # remove full-attention KV blocks, so it must not emit BlockRemoved. - result = tree.evict(EvictParams(num_tokens=0, mamba_num=1)) - self.assertEqual(result.num_tokens_evicted, 0) - self.assertEqual(result.mamba_num_evicted, 1) - events = tree.take_events() - self.assertEqual([e for e in events if isinstance(e, BlockRemoved)], []) - - result = tree.evict(EvictParams(num_tokens=1)) - self.assertGreaterEqual(result.num_tokens_evicted, 1) - events = tree.take_events() - removed_hashes = _event_hashes( - [e for e in events if isinstance(e, BlockRemoved)] - ) - self.assertCountEqual(removed_hashes, stored_hashes) - - def test_mamba_radix_cache_kv_events_split_hash(self): - tree, allocator, _, make_dummy_req = self._setup_tree_and_allocator( - enable_kv_cache_events=True - ) - tree.take_events() # Clear the reset event. - - req1 = make_dummy_req() - key1 = RadixKey(array("q", [1, 2, 3, 4])) - tree.insert( - InsertParams( - key=key1, - value=allocator.alloc(4)[: len(key1)], - mamba_value=req1.kv.mamba_pool_idx.unsqueeze(0), - ) - ) - first_insert_events = [ - e for e in tree.take_events() if isinstance(e, BlockStored) - ] - self.assertEqual(len(first_insert_events), 1) - split_parent_hash = first_insert_events[0].block_hashes[1] - - req2 = make_dummy_req() - key2 = RadixKey(array("q", [1, 2, 5, 6])) - tree.insert( - InsertParams( - key=key2, - value=allocator.alloc(4)[: len(key2)], - mamba_value=req2.kv.mamba_pool_idx.unsqueeze(0), - ) - ) - second_insert_events = [ - e for e in tree.take_events() if isinstance(e, BlockStored) - ] - self.assertEqual(len(second_insert_events), 1) - self.assertEqual(list(second_insert_events[0].token_ids), [5, 6]) - self.assertEqual(second_insert_events[0].parent_block_hash, split_parent_hash) - - def _setup_tree_and_allocator(self, enable_kv_cache_events=False): - """Helper to create a MambaRadixCache with allocator for testing.""" + def _setup_pools(self): + """Build the hybrid req/KV pools and an allocator for pool-level tests.""" server_args = ServerArgs(model_path="dummy", page_size=1) - # MambaRadixCache reads mamba_cache_chunk_size, whose property otherwise + # The mamba pool reads mamba_cache_chunk_size, whose property otherwise # loads the HF config for self.model_path — impossible for the dummy model. # Mirror the property's default for a dummy HF config: FLA_CHUNK_SIZE. server_args._mamba_cache_chunk_size = FLA_CHUNK_SIZE @@ -635,14 +297,6 @@ class TestMamba(unittest.TestCase): kvcache=pool, need_sort=False, ) - params = CacheInitParams( - req_to_token_pool=req_to_token_pool, - token_to_kv_pool_allocator=allocator, - page_size=1, - disable=False, - enable_kv_cache_events=enable_kv_cache_events, - ) - tree = MambaRadixCache(params=params) def make_dummy_req(): sampling_params = SamplingParams( @@ -658,7 +312,7 @@ class TestMamba(unittest.TestCase): req_to_token_pool.alloc([req]) return req - return tree, allocator, req_to_token_pool, make_dummy_req + return allocator, req_to_token_pool, make_dummy_req # Qwen4-Exp's PLE N-gram window is 2 wide (ngram_size=3) and its "no history" # sentinel is the eos id; pick a recognisable one for the tests. @@ -700,7 +354,7 @@ class TestMamba(unittest.TestCase): def test_slot_siblings_registered(self): """Enabled PLE side states register on the pool that owns the slots; disabled ones stay off so the host-offload payload keeps its legacy shape.""" - _, _, base_pool, _ = self._setup_tree_and_allocator() + _, base_pool, _ = self._setup_pools() # The default hybrid setup has no PLE config: no siblings ride along. self.assertEqual(len(base_pool.mamba_pool._slot_siblings), 0) pool = self._setup_pool_with_ngram() @@ -813,7 +467,7 @@ class TestMamba(unittest.TestCase): def test_mamba_pool_cpu_offload(self): """MambaPool.get_cpu_copy / load_cpu_copy round-trips conv and temporal state.""" - _, _, req_to_token_pool, _ = self._setup_tree_and_allocator() + _, req_to_token_pool, _ = self._setup_pools() mamba_pool = req_to_token_pool.mamba_pool n = 3 indices = req_to_token_pool.mamba_allocator.alloc(n) @@ -862,7 +516,7 @@ class TestMamba(unittest.TestCase): def test_hybrid_kv_pool_cpu_offload(self): """HybridLinearKVPool.get_cpu_copy / load_cpu_copy saves and restores both the full-attention KV cache and Mamba state in a single round-trip.""" - _, allocator, req_to_token_pool, _ = self._setup_tree_and_allocator() + allocator, req_to_token_pool, _ = self._setup_pools() mamba_pool = req_to_token_pool.mamba_pool hybrid_pool = allocator._kvcache # HybridLinearKVPool @@ -934,81 +588,6 @@ class TestMamba(unittest.TestCase): mamba_cpu_none, "mamba_cpu should be None when mamba_indices=None" ) - def test_insert_prev_prefix_len(self): - """Test that prev_prefix_len correctly controls which KV indices are freed - during insert, covering: full free, partial free across multi-node, and no free. - """ - tree, allocator, req_to_token_pool, make_dummy_req = ( - self._setup_tree_and_allocator() - ) - - initial_avail = allocator.available_size() - - # Step 1: Insert [1,2,3] to create first node - req1 = make_dummy_req() - key1 = RadixKey(array("q", [1, 2, 3])) - tree.insert( - InsertParams( - key=key1, - value=allocator.alloc(3)[: len(key1)], - mamba_value=req1.kv.mamba_pool_idx.unsqueeze(0), - ) - ) - assert allocator.available_size() == initial_avail - 3 - - # Step 2: Insert [1,2,3,4,5,6,7] with prev_prefix_len=0 (free all matched) - # Creates tree: [1,2,3] -> [4,5,6,7] - req2 = make_dummy_req() - key2 = RadixKey(array("q", [1, 2, 3, 4, 5, 6, 7])) - result = tree.insert( - InsertParams( - key=key2, - value=allocator.alloc(7)[: len(key2)], - mamba_value=req2.kv.mamba_pool_idx.unsqueeze(0), - prev_prefix_len=0, - ) - ) - assert result.prefix_len == 3 - # alloc 7, freed 3 (dup prefix [0..2]), stored 4 in new node => net -4 - assert allocator.available_size() == initial_avail - 3 - 4 - avail_after_step2 = allocator.available_size() - - # Step 3: Insert [1,2,3,4,5,6,7,8] with prev_prefix_len=2 - # Matched prefix = 7 (across two nodes: [1,2,3] len=3, [4,5,6,7] len=4) - # Protected [0..1], freed [2..6] = 5 slots, new [7] = 1 slot stored - req3 = make_dummy_req() - key3 = RadixKey(array("q", [1, 2, 3, 4, 5, 6, 7, 8])) - result = tree.insert( - InsertParams( - key=key3, - value=allocator.alloc(8)[: len(key3)], - mamba_value=req3.kv.mamba_pool_idx.unsqueeze(0), - prev_prefix_len=2, - ) - ) - assert result.prefix_len == 7 - # alloc 8, freed 5, stored 1 => net -3 - assert allocator.available_size() == avail_after_step2 - 3 - avail_after_step3 = allocator.available_size() - - # Step 4: Insert [1,2,3,4,5,6,7,8,9] with prev_prefix_len=8 (covers all matched) - # Matched prefix = 8, prev_prefix_len=8 => nothing freed - req4 = make_dummy_req() - key4 = RadixKey(array("q", [1, 2, 3, 4, 5, 6, 7, 8, 9])) - result = tree.insert( - InsertParams( - key=key4, - value=allocator.alloc(9)[: len(key4)], - mamba_value=req4.kv.mamba_pool_idx.unsqueeze(0), - prev_prefix_len=8, - ) - ) - assert result.prefix_len == 8 - # alloc 9, freed 0, stored 1 => net -9 - assert allocator.available_size() == avail_after_step3 - 9 - - tree.sanity_check() - if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_radix_cache_unit.py b/test/registered/unit/mem_cache/test_radix_cache_unit.py index 49319dfcf..31f165402 100644 --- a/test/registered/unit/mem_cache/test_radix_cache_unit.py +++ b/test/registered/unit/mem_cache/test_radix_cache_unit.py @@ -45,7 +45,6 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, ) from sglang.srt.mem_cache.events import KVCacheEventRecorder -from sglang.srt.mem_cache.mamba_radix_cache import TreeNode as MambaTreeNode from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode from sglang.srt.utils import get_device from sglang.test.test_utils import CustomTestCase @@ -338,7 +337,7 @@ class TestTreeNode(unittest.TestCase): def test_get_prefix_hash_values_not_shared_across_calls(self): """Regression guard for cached mutable prefix hash lists.""" - for node_cls in (TreeNode, MambaTreeNode): + for node_cls in (TreeNode,): with self.subTest(node_cls=node_cls.__module__): root = node_cls() n1 = node_cls() diff --git a/test/registered/unit/mem_cache/test_swa_eviction_boundary.py b/test/registered/unit/mem_cache/test_swa_eviction_boundary.py deleted file mode 100644 index e486db1fd..000000000 --- a/test/registered/unit/mem_cache/test_swa_eviction_boundary.py +++ /dev/null @@ -1,550 +0,0 @@ -"""Unit tests for SWA eviction boundary fixes. - -Bug: when page_size > sliding_window_size, _evict_swa could advance the -eviction frontier to exactly page_floor(seq_len), making all tokens being -inserted into the radix tree fully evicted (case 3). _insert_helper had no -handling for this, creating an incorrect non-tombstone node that caused -inflated swa_evictable_size_, negative usage, and potential double-free. - -Two-sided fix: -1. _evict_swa subtracts max(window, page) on the radix path (preventive). -2. _insert_helper early-returns on case 3 (defensive). - -Tests use real tree/allocator/pool with mock Req/ScheduleBatch wrappers. -""" - -import unittest -from types import SimpleNamespace - -import torch - -from sglang.srt.managers.schedule_batch import ReqKvInfo, ScheduleBatch -from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator -from sglang.srt.mem_cache.base_prefix_cache import DecLockRefParams -from sglang.srt.mem_cache.cache_init_params import CacheInitParams -from sglang.srt.mem_cache.common import free_swa_out_of_window_slots -from sglang.srt.mem_cache.memory_pool import ReqToTokenPool -from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool -from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache -from sglang.srt.utils import get_device -from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci - -register_cuda_ci(est_time=13, stage="base-b", runner_config="1-gpu-large") -register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd") - -# --------------------------------------------------------------------------- -# Infrastructure helpers (shared setup, not logic) -# --------------------------------------------------------------------------- - - -def _swa_alloc(allocator, need_size): - """Allocate from SWA allocator for any page_size. - - SWATokenToKVPoolAllocator.alloc() asserts page_size == 1. For page_size > 1, - allocate from the underlying paged allocators directly and set up the mapping. - """ - if allocator.page_size == 1: - return allocator.alloc(need_size) - - if need_size > allocator.full_attn_allocator.available_size(): - return None - if need_size > allocator.swa_attn_allocator.available_size(): - return None - - full_indices = allocator.full_attn_allocator.alloc(need_size) - swa_indices = allocator.swa_attn_allocator.alloc(need_size) - assert full_indices is not None and swa_indices is not None - allocator.full_to_swa_index_mapping[full_indices] = swa_indices - return full_indices - - -def _build_swa_tree(page_size, sliding_window_size, kv_size=1024, kv_size_swa=512): - head_num, head_dim, num_layers, global_interval = 8, 128, 24, 4 - dtype = torch.bfloat16 - device = get_device() - full_ids = list(range(0, num_layers, global_interval)) - swa_ids = [i for i in range(num_layers) if i not in set(full_ids)] - - pool = ReqToTokenPool( - size=8, max_context_len=2048, device=device, enable_memory_saver=False - ) - kv_pool = SWAKVPool( - size=kv_size, - size_swa=kv_size_swa, - page_size=page_size, - dtype=dtype, - head_num=head_num, - head_dim=head_dim, - swa_attention_layer_ids=swa_ids, - full_attention_layer_ids=full_ids, - device=device, - ) - allocator = SWATokenToKVPoolAllocator( - size=kv_size, - size_swa=kv_size_swa, - page_size=page_size, - dtype=dtype, - device=device, - kvcache=kv_pool, - need_sort=False, - ) - tree = SWARadixCache( - params=CacheInitParams( - req_to_token_pool=pool, - token_to_kv_pool_allocator=allocator, - page_size=page_size, - disable=False, - is_eagle=False, - sliding_window_size=sliding_window_size, - ), - ) - return tree, allocator, pool - - -def _make_req(req_pool_idx, token_ids, cache_protected_len, tree): - """Mock Req with fields needed by _evict_swa and cache_finished_req.""" - req = SimpleNamespace( - origin_input_ids=token_ids, - output_ids=[], - kv=ReqKvInfo( - req_pool_idx=req_pool_idx, cache_protected_len=cache_protected_len - ), - extra_key=None, - cache_salt=None, - last_node=tree.root_node, - lock_receipt=DecLockRefParams(), - swa_prefix_lock_released=False, - prefix_indices=torch.tensor([], dtype=torch.int64, device=tree.device), - _kv_committed_len=len(token_ids), - ) - return req - - -def _make_batch(tree, allocator, pool): - """Mock ScheduleBatch with fields needed by _evict_swa.""" - return SimpleNamespace( - tree_cache=tree, - req_to_token_pool=pool, - token_to_kv_pool_allocator=allocator, - ) - - -# --------------------------------------------------------------------------- -# Tests -# --------------------------------------------------------------------------- - - -class TestSWAEvictionBoundary(unittest.TestCase): - # -- Eviction formula: page_size > window -- - - def test_formula_page_gt_window_sweep(self): - """Sweep page_size > window combinations. The -page_size fix must - prevent eviction from reaching page_floor(seq_len).""" - for page_size in [4, 8, 16, 32, 64, 128, 256]: - for window in [1, 2, 4, 8]: - if page_size <= window: - continue - tree, allocator, pool = _build_swa_tree( - page_size=page_size, - sliding_window_size=window, - kv_size=max(4096, page_size * 20), - kv_size_swa=max(2048, page_size * 10), - ) - for seq_len in range(page_size + 1, page_size * 5): - alloc_size = (seq_len + page_size - 1) // page_size * page_size - kv = _swa_alloc(allocator, alloc_size) - if kv is None: - break - pool.write((0, slice(0, alloc_size)), kv) - - req = _make_req(0, list(range(seq_len)), 0, tree) - batch = _make_batch(tree, allocator, pool) - ScheduleBatch._evict_swa(batch, req, seq_len - 1) - - insert_len = seq_len // page_size * page_size - self.assertLess( - req.kv.swa_evicted_seqlen, - insert_len, - f"page={page_size}, win={window}, seq={seq_len}", - ) - allocator.free(kv) - - # -- Eviction formula: page_size <= window -- - - def test_formula_page_leq_window(self): - """page_size <= window: -page_size fix causes no regression.""" - page_size, window = 4, 8 - tree, allocator, pool = _build_swa_tree( - page_size=page_size, sliding_window_size=window - ) - - for seq_len in [13, 17, 25, 33]: - alloc_size = (seq_len + page_size - 1) // page_size * page_size - kv = _swa_alloc(allocator, alloc_size) - pool.write((0, slice(0, alloc_size)), kv) - - req = _make_req(0, list(range(seq_len)), 0, tree) - batch = _make_batch(tree, allocator, pool) - ScheduleBatch._evict_swa(batch, req, seq_len - 1) - - insert_len = seq_len // page_size * page_size - self.assertLess(req.kv.swa_evicted_seqlen, insert_len) - - tree.cache_finished_req( - req, is_insert=True, owned_kv_len=req._kv_committed_len - ) - tree.sanity_check() - - # -- Retention floor: never free past the last state checkpoint -- - - def test_retain_floor_clamps_eviction(self): - """A hybrid cache keeps SWA down to the last state checkpoint, not to the - window behind the tail, because that is where a prefix match lands. The - floor must clamp the frontier even though the tail has moved far past it.""" - page_size, window = 8, 16 - tree, allocator, pool = _build_swa_tree( - page_size=page_size, sliding_window_size=window - ) - seq_len = 200 - checkpoint = 96 - kv = _swa_alloc(allocator, seq_len) - pool.write((0, slice(0, seq_len)), kv) - req = _make_req(0, list(range(seq_len)), 0, tree) - batch = _make_batch(tree, allocator, pool) - - free_swa_out_of_window_slots( - req, - seq_len - 1, - sliding_window_size=window, - page_size=page_size, - req_to_token_pool=batch.req_to_token_pool, - token_to_kv_pool_allocator=batch.token_to_kv_pool_allocator, - retain_floor=checkpoint - window, - ) - - # Without the floor this would reach page_floor(199 - 16) = 176. - self.assertLessEqual(req.kv.swa_evicted_seqlen, checkpoint - window) - self.assertEqual(req.kv.swa_evicted_seqlen % page_size, 0) - - def test_retain_floor_ignored_for_chunk_cache(self): - """Chunk cache builds no tree, so a retained checkpoint could never be - matched. Holding it would cost SWA slots for nothing.""" - page_size, window = 8, 16 - seq_len = 200 - tree, allocator, pool = _build_swa_tree( - page_size=page_size, sliding_window_size=window - ) - kv = _swa_alloc(allocator, seq_len) - pool.write((0, slice(0, seq_len)), kv) - req = _make_req(0, list(range(seq_len)), 0, tree) - batch = _make_batch(tree, allocator, pool) - - free_swa_out_of_window_slots( - req, - seq_len - 1, - sliding_window_size=window, - page_size=page_size, - req_to_token_pool=batch.req_to_token_pool, - token_to_kv_pool_allocator=batch.token_to_kv_pool_allocator, - is_chunk_cache=True, - retain_floor=16, - ) - - expected = (seq_len - 1 - window) // page_size * page_size - self.assertEqual(req.kv.swa_evicted_seqlen, expected) - - def test_retain_floor_none_matches_old_behaviour(self): - """retain_floor=None must reproduce the pre-change frontier exactly, so a - cache without a second state stream is unaffected.""" - page_size, window = 8, 16 - seq_len = 200 - frontiers = [] - for floor in (None, "absent"): - tree, allocator, pool = _build_swa_tree( - page_size=page_size, sliding_window_size=window - ) - kv = _swa_alloc(allocator, seq_len) - pool.write((0, slice(0, seq_len)), kv) - req = _make_req(0, list(range(seq_len)), 0, tree) - batch = _make_batch(tree, allocator, pool) - kwargs = {} if floor == "absent" else {"retain_floor": None} - free_swa_out_of_window_slots( - req, - seq_len - 1, - sliding_window_size=window, - page_size=page_size, - req_to_token_pool=batch.req_to_token_pool, - token_to_kv_pool_allocator=batch.token_to_kv_pool_allocator, - **kwargs, - ) - frontiers.append(req.kv.swa_evicted_seqlen) - - expected = (seq_len - 1 - max(window, page_size)) // page_size * page_size - self.assertEqual(frontiers[0], expected) - self.assertEqual(frontiers[1], expected) - - def test_retain_floor_above_threshold_is_inert(self): - """The floor is a min(), so a checkpoint that is already inside the window - must not hold anything extra.""" - page_size, window = 8, 16 - seq_len = 200 - tree, allocator, pool = _build_swa_tree( - page_size=page_size, sliding_window_size=window - ) - kv = _swa_alloc(allocator, seq_len) - pool.write((0, slice(0, seq_len)), kv) - req = _make_req(0, list(range(seq_len)), 0, tree) - batch = _make_batch(tree, allocator, pool) - - free_swa_out_of_window_slots( - req, - seq_len - 1, - sliding_window_size=window, - page_size=page_size, - req_to_token_pool=batch.req_to_token_pool, - token_to_kv_pool_allocator=batch.token_to_kv_pool_allocator, - retain_floor=seq_len, - ) - - expected = (seq_len - 1 - max(window, page_size)) // page_size * page_size - self.assertEqual(req.kv.swa_evicted_seqlen, expected) - - def test_retain_floor_does_not_unfree(self): - """The frontier only advances. A floor arriving after slots were already - freed must not claim them back, which would double-free on the next pass.""" - page_size, window = 8, 16 - tree, allocator, pool = _build_swa_tree( - page_size=page_size, sliding_window_size=window - ) - seq_len = 200 - kv = _swa_alloc(allocator, seq_len) - pool.write((0, slice(0, seq_len)), kv) - req = _make_req(0, list(range(seq_len)), 0, tree) - batch = _make_batch(tree, allocator, pool) - common_kwargs = dict( - sliding_window_size=window, - page_size=page_size, - req_to_token_pool=batch.req_to_token_pool, - token_to_kv_pool_allocator=batch.token_to_kv_pool_allocator, - ) - - free_swa_out_of_window_slots(req, seq_len - 1, **common_kwargs) - advanced = req.kv.swa_evicted_seqlen - self.assertGreater(advanced, 0) - - free_swa_out_of_window_slots(req, seq_len - 1, retain_floor=0, **common_kwargs) - self.assertEqual(req.kv.swa_evicted_seqlen, advanced) - - # -- Eviction formula: page_size == 1 -- - - def test_formula_page_size_1(self): - """page_size=1: radix keeps max(window, page)=window, so the frontier is pre_len - window.""" - page_size, window = 1, 4 - tree, allocator, pool = _build_swa_tree( - page_size=page_size, sliding_window_size=window - ) - - for seq_len in range(window + 2, 30): - kv = _swa_alloc(allocator, seq_len) - pool.write((0, slice(0, seq_len)), kv) - - req = _make_req(0, list(range(seq_len)), 0, tree) - batch = _make_batch(tree, allocator, pool) - ScheduleBatch._evict_swa(batch, req, seq_len - 1) - - self.assertLess(req.kv.swa_evicted_seqlen, seq_len) - self.assertEqual( - req.kv.swa_evicted_seqlen, max(0, seq_len - 1 - max(window, page_size)) - ) - - tree.cache_finished_req( - req, is_insert=True, owned_kv_len=req._kv_committed_len - ) - tree.sanity_check() - - # -- Eviction formula: no-op when seq too short -- - - def test_formula_noop_short_sequence(self): - """pre_len - window - page_size < 0: eviction stays at 0.""" - page_size, window = 8, 4 - tree, allocator, pool = _build_swa_tree( - page_size=page_size, sliding_window_size=window - ) - - seq_len = page_size + window - 2 # = 10, formula gives 10-1-4-8 = -3 - alloc_size = (seq_len + page_size - 1) // page_size * page_size - kv = _swa_alloc(allocator, alloc_size) - pool.write((0, slice(0, alloc_size)), kv) - - req = _make_req(0, list(range(seq_len)), 0, tree) - batch = _make_batch(tree, allocator, pool) - ScheduleBatch._evict_swa(batch, req, seq_len - 1) - - self.assertEqual(req.kv.swa_evicted_seqlen, 0) - - # -- Insert case 1: swa_evicted <= total_prefix_length -- - - def test_insert_case1_evicted_within_matched(self): - """Eviction within matched region. New tokens all non-tombstone.""" - page_size, window = 8, 2 - tree, allocator, pool = _build_swa_tree( - page_size=page_size, sliding_window_size=window - ) - - # First request: populate tree with 16 tokens (2 pages) - first_len = page_size * 2 - kv1 = _swa_alloc(allocator, first_len) - pool.write((0, slice(0, first_len)), kv1) - req1 = _make_req(0, list(range(first_len)), 0, tree) - tree.cache_finished_req( - req1, is_insert=True, owned_kv_len=req1._kv_committed_len - ) - tree.sanity_check() - - # Second request: 24 tokens, first 16 overlap with tree - second_len = page_size * 3 - kv2 = _swa_alloc(allocator, second_len) - pool.write((1, slice(0, second_len)), kv2) - - req2 = _make_req(1, list(range(second_len)), 0, tree) - batch = _make_batch(tree, allocator, pool) - - # pre_len=15: 15-2-8=5, floor to 8 -> 0. Eviction stays within matched. - ScheduleBatch._evict_swa(batch, req2, first_len - 1) - self.assertLessEqual(req2.kv.swa_evicted_seqlen, first_len) - - swa_evictable_before = tree.swa_evictable_size_ - tree.cache_finished_req( - req2, is_insert=True, owned_kv_len=req2._kv_committed_len - ) - - # New tokens [16, 24) should all be non-tombstone - new_tokens = second_len // page_size * page_size - first_len - self.assertEqual(tree.swa_evictable_size_, swa_evictable_before + new_tokens) - tree.sanity_check() - - # -- Insert case 2: total_prefix_length < swa_evicted < total_length -- - - def test_insert_case2_partial_tombstone(self): - """Partial eviction: some tombstone, some non-tombstone.""" - page_size, window = 8, 2 - tree, allocator, pool = _build_swa_tree( - page_size=page_size, sliding_window_size=window - ) - - # seq_len=25: insert_length=24, evicted should be 8 (1 page) - seq_len = page_size * 3 + 1 - alloc_size = (seq_len + page_size - 1) // page_size * page_size - kv = _swa_alloc(allocator, alloc_size) - pool.write((0, slice(0, alloc_size)), kv) - - req = _make_req(0, list(range(seq_len)), 0, tree) - batch = _make_batch(tree, allocator, pool) - swa_evictable_before = tree.swa_evictable_size_ - - ScheduleBatch._evict_swa(batch, req, seq_len - 1) - insert_len = seq_len // page_size * page_size - self.assertGreater(req.kv.swa_evicted_seqlen, 0, "Should have some eviction") - self.assertLess(req.kv.swa_evicted_seqlen, insert_len, "Should be partial") - - tree.cache_finished_req(req, is_insert=True, owned_kv_len=req._kv_committed_len) - - non_tombstone = insert_len - req.kv.swa_evicted_seqlen - self.assertEqual(tree.swa_evictable_size_, swa_evictable_before + non_tombstone) - self.assertGreater(tree.full_evictable_size_, 0) - tree.sanity_check() - - # -- Insert case 3: swa_evicted == total_length (defensive) -- - - def test_insert_case3_defensive_early_return(self): - """Simulate OLD formula to trigger case 3. Defensive early return - must prevent non-tombstone node creation.""" - page_size, window = 8, 2 - tree, allocator, pool = _build_swa_tree( - page_size=page_size, sliding_window_size=window - ) - - # seq_len=11: OLD formula -> evicted=8 == insert_length=8 - seq_len = page_size + window + 1 - alloc_size = (seq_len + page_size - 1) // page_size * page_size - kv = _swa_alloc(allocator, alloc_size) - pool.write((0, slice(0, alloc_size)), kv) - - # OLD formula (without -page_size) - pre_len = seq_len - 1 - old_evicted = max(0, (pre_len - window) // page_size * page_size) - insert_len = seq_len // page_size * page_size - self.assertEqual( - old_evicted, insert_len, "Precondition: old formula hits boundary" - ) - - # Manually free SWA as _evict_swa would - allocator.free_swa(pool.req_to_token[0, :old_evicted]) - - req = _make_req(0, list(range(seq_len)), 0, tree) - req.kv.swa_evicted_seqlen = old_evicted - swa_evictable_before = tree.swa_evictable_size_ - - tree.cache_finished_req(req, is_insert=True, owned_kv_len=req._kv_committed_len) - - self.assertEqual(tree.swa_evictable_size_, swa_evictable_before) - - # -- Integration: multiple decode turns -- - - def test_multiple_decodes(self): - """Multiple decode turns with page_size > window. No over-eviction, - tree stays consistent throughout.""" - page_size, window = 8, 2 - tree, allocator, pool = _build_swa_tree( - page_size=page_size, sliding_window_size=window - ) - - for turn in range(4): - seq_len = page_size * (turn + 2) + 1 - idx = turn % pool.size - alloc_size = (seq_len + page_size - 1) // page_size * page_size - kv = _swa_alloc(allocator, alloc_size) - assert kv is not None, f"Alloc failed at turn {turn}" - pool.write((idx, slice(0, alloc_size)), kv) - - req = _make_req(idx, list(range(seq_len)), 0, tree) - batch = _make_batch(tree, allocator, pool) - ScheduleBatch._evict_swa(batch, req, seq_len - 1) - - insert_len = seq_len // page_size * page_size - self.assertLess(req.kv.swa_evicted_seqlen, insert_len, f"turn {turn}") - - tree.cache_finished_req( - req, is_insert=True, owned_kv_len=req._kv_committed_len - ) - tree.sanity_check() - - # -- Integration: page_size=1 full flow -- - - def test_page_size_1_full_flow(self): - """End-to-end with page_size=1. Fix is near no-op.""" - page_size, window = 1, 4 - tree, allocator, pool = _build_swa_tree( - page_size=page_size, sliding_window_size=window - ) - - for seq_len in [10, 20, 30]: - kv = _swa_alloc(allocator, seq_len) - pool.write((0, slice(0, seq_len)), kv) - - req = _make_req(0, list(range(seq_len)), 0, tree) - batch = _make_batch(tree, allocator, pool) - ScheduleBatch._evict_swa(batch, req, seq_len - 1) - - self.assertEqual( - req.kv.swa_evicted_seqlen, max(0, seq_len - 1 - max(window, page_size)) - ) - - tree.cache_finished_req( - req, is_insert=True, owned_kv_len=req._kv_committed_len - ) - tree.sanity_check() - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/unit/mem_cache/test_swa_lock_release_lifecycle.py b/test/registered/unit/mem_cache/test_swa_lock_release_lifecycle.py deleted file mode 100644 index 1d19ae00f..000000000 --- a/test/registered/unit/mem_cache/test_swa_lock_release_lifecycle.py +++ /dev/null @@ -1,440 +0,0 @@ -"""Regression for SWA lock release lifecycle. - -Hybrid-SWA early-release protocol: once a request's decode position passes -the sliding window, drop its prefill SWA lock without touching the full -lock, freeing SWA pages back to LRU. - -Covers: -- SWARadixCache.dec_swa_lock_only (leaf tombstone + free, internal protected->evictable) -- SWARadixCache.dec_lock_ref(skip_swa=True) -- SWARadixCache.evict swa branch for leaf with full_lock_ref > 0 -- SWARadixCache._delete_leaf skipping swa_evictable_size_ on tombstoned leaves -""" - -import unittest -from array import array - -import torch - -from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator -from sglang.srt.mem_cache.base_prefix_cache import ( - DecLockRefParams, - EvictParams, - 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 -from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool -from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache -from sglang.srt.utils import get_device -from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci -from sglang.test.test_utils import CustomTestCase - -register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small") -register_amd_ci(est_time=12, suite="stage-b-test-1-gpu-small-amd") - - -def _build_tree( - *, - sliding_window_size: int = 4, - page_size: int = 1, - kv_size: int = 128, - kv_size_swa: int = 64, -): - head_num, head_dim, num_layers, global_interval = 8, 128, 24, 4 - dtype = torch.bfloat16 - device = get_device() - full_ids = list(range(0, num_layers, global_interval)) - swa_ids = [i for i in range(num_layers) if i not in set(full_ids)] - - pool = ReqToTokenPool( - size=8, max_context_len=256, device=device, enable_memory_saver=False - ) - kv_pool = SWAKVPool( - size=kv_size, - size_swa=kv_size_swa, - page_size=page_size, - dtype=dtype, - head_num=head_num, - head_dim=head_dim, - swa_attention_layer_ids=swa_ids, - full_attention_layer_ids=full_ids, - device=device, - ) - allocator = SWATokenToKVPoolAllocator( - size=kv_size, - size_swa=kv_size_swa, - page_size=page_size, - dtype=dtype, - device=device, - kvcache=kv_pool, - need_sort=False, - ) - tree = SWARadixCache( - params=CacheInitParams( - req_to_token_pool=pool, - token_to_kv_pool_allocator=allocator, - page_size=page_size, - disable=False, - is_eagle=False, - sliding_window_size=sliding_window_size, - ), - ) - return tree, allocator, pool - - -def _swa_alloc(allocator, need_size): - """Allocate from SWA allocator for any page_size. - - SWATokenToKVPoolAllocator.alloc() asserts page_size == 1; for page_size > 1 - we drive the underlying paged allocators directly (mirrors the helper in - test_swa_eviction_boundary.py). Required: need_size is a multiple of - page_size when page_size > 1. - """ - if allocator.page_size == 1: - return allocator.alloc(need_size) - - assert need_size % allocator.page_size == 0, ( - f"page_size > 1 requires page-aligned alloc, got {need_size=} " - f"with {allocator.page_size=}" - ) - if need_size > allocator.full_attn_allocator.available_size(): - return None - if need_size > allocator.swa_attn_allocator.available_size(): - return None - full_indices = allocator.full_attn_allocator.alloc(need_size) - swa_indices = allocator.swa_attn_allocator.alloc(need_size) - assert full_indices is not None and swa_indices is not None - allocator.full_to_swa_index_mapping[full_indices] = swa_indices - return full_indices - - -def _insert_chain(tree, allocator, token_ids): - token_ids = array("q", token_ids) - indices = _swa_alloc(allocator, len(token_ids)) - assert indices is not None - tree.insert(InsertParams(key=RadixKey(token_ids), value=indices)) - match = tree.match_prefix(MatchPrefixParams(key=RadixKey(token_ids))) - return match.last_device_node - - -def _release_swa_lock_chain_in_place(tree, leaf, swa_uuid_for_lock): - # Mirrors dec_swa_lock_only's non-tombstone arm (protected->evictable on - # internal nodes) but skips the leaf-free + tombstone step, to construct - # the post-revival state where SWA was already early-released yet the - # leaf is back in swa_lru_list with full_lock_ref still > 0. - node = leaf - while node is not tree.root_node: - if node.swa_lock_ref > 0: - if node.swa_lock_ref == 1: - tree.swa_protected_size_ -= len(node.value) - tree.swa_evictable_size_ += len(node.value) - node.swa_lock_ref -= 1 - if swa_uuid_for_lock and node.swa_uuid == swa_uuid_for_lock: - break - node = node.parent - - -class TestSWALockReleaseLifecycle(CustomTestCase): - """Each test pins one component of the early-release fix; method names - are prefixed with the API surface they exercise so pytest output groups - them naturally.""" - - def test_dec_swa_lock_only_leaf_tombstones_and_frees(self): - tree, allocator, _ = _build_tree(sliding_window_size=4) - leaf = _insert_chain(tree, allocator, [1, 2, 3, 4, 5, 6, 7, 8]) - self.assertEqual(len(leaf.value), 8) - - inc_res = tree.inc_lock_ref(leaf) - swa_uuid = inc_res.swa_uuid_for_lock - self.assertIsNotNone(swa_uuid) - - swa_avail_before = allocator.swa_available_size() - full_avail_before = allocator.full_available_size() - self.assertEqual(leaf.swa_lock_ref, 1) - self.assertEqual(leaf.full_lock_ref, 1) - self.assertFalse(leaf.swa_tombstone) - self.assertTrue(tree.swa_lru_list.in_list(leaf)) - - tree.dec_swa_lock_only(leaf, DecLockRefParams(swa_uuid_for_lock=swa_uuid)) - - self.assertTrue(leaf.swa_tombstone) - self.assertFalse(tree.swa_lru_list.in_list(leaf)) - self.assertEqual(leaf.swa_lock_ref, 0) - self.assertEqual( - allocator.swa_available_size(), swa_avail_before + len(leaf.value) - ) - self.assertEqual(leaf.full_lock_ref, 1) - self.assertEqual(allocator.full_available_size(), full_avail_before) - - # sanity_check forbids live locks; release the full half before checking. - tree.dec_lock_ref( - leaf, DecLockRefParams(swa_uuid_for_lock=swa_uuid), skip_swa=True - ) - tree.sanity_check() - - def test_dec_swa_lock_only_internal_no_tombstone_no_free(self): - # Two siblings force an internal node at the shared prefix. - tree, allocator, _ = _build_tree(sliding_window_size=4) - leaf_a = _insert_chain(tree, allocator, [1, 2, 3, 4, 5, 6, 7, 8]) - _insert_chain(tree, allocator, [1, 2, 3, 4, 5, 6, 7, 9]) - - # Post-split: leaf_a now carries [8] only, parent holds the shared 7. - self.assertEqual(len(leaf_a.value), 1) - internal = leaf_a.parent - self.assertGreater(len(internal.children), 1) - self.assertEqual(len(internal.value), 7) - - inc_res = tree.inc_lock_ref(leaf_a) - swa_uuid = inc_res.swa_uuid_for_lock - # window=4, value 1 (leaf) + 7 (internal): swa lock chain ends at internal. - self.assertEqual(swa_uuid, internal.swa_uuid) - - swa_protected_before = tree.swa_protected_size_ - swa_evictable_before = tree.swa_evictable_size_ - swa_avail_before = allocator.swa_available_size() - - tree.dec_swa_lock_only(leaf_a, DecLockRefParams(swa_uuid_for_lock=swa_uuid)) - - self.assertFalse(internal.swa_tombstone) - self.assertTrue(tree.swa_lru_list.in_list(internal)) - self.assertEqual(internal.swa_lock_ref, 0) - self.assertEqual( - tree.swa_protected_size_, swa_protected_before - (len(leaf_a.value) + 7) - ) - self.assertEqual(tree.swa_evictable_size_, swa_evictable_before + 7) - self.assertEqual( - allocator.swa_available_size(), swa_avail_before + len(leaf_a.value) - ) - - tree.dec_lock_ref( - leaf_a, DecLockRefParams(swa_uuid_for_lock=swa_uuid), skip_swa=True - ) - tree.sanity_check() - - def test_dec_lock_ref_skip_swa_true_drops_full_only(self): - tree, allocator, _ = _build_tree(sliding_window_size=4) - leaf = _insert_chain(tree, allocator, [1, 2, 3, 4, 5, 6, 7, 8]) - - inc_res = tree.inc_lock_ref(leaf) - swa_uuid = inc_res.swa_uuid_for_lock - - tree.dec_swa_lock_only(leaf, DecLockRefParams(swa_uuid_for_lock=swa_uuid)) - self.assertTrue(leaf.swa_tombstone) - self.assertEqual(leaf.full_lock_ref, 1) - - swa_avail_after_release = allocator.swa_available_size() - swa_protected_after_release = tree.swa_protected_size_ - - # Without skip_swa, dec_lock_ref would assert on the swa_tombstone leaf. - tree.dec_lock_ref( - leaf, DecLockRefParams(swa_uuid_for_lock=swa_uuid), skip_swa=True - ) - - self.assertEqual(leaf.full_lock_ref, 0) - self.assertEqual(allocator.swa_available_size(), swa_avail_after_release) - self.assertEqual(tree.swa_protected_size_, swa_protected_after_release) - tree.sanity_check() - - def test_dec_lock_ref_skip_swa_false_drops_both(self): - # Default skip_swa=False must keep legacy behavior intact. - tree, allocator, _ = _build_tree(sliding_window_size=4) - leaf = _insert_chain(tree, allocator, [1, 2, 3, 4, 5, 6, 7, 8]) - - inc_res = tree.inc_lock_ref(leaf) - swa_uuid = inc_res.swa_uuid_for_lock - - full_avail_before = allocator.full_available_size() - swa_avail_before = allocator.swa_available_size() - - tree.dec_lock_ref(leaf, DecLockRefParams(swa_uuid_for_lock=swa_uuid)) - - self.assertEqual(leaf.full_lock_ref, 0) - self.assertEqual(leaf.swa_lock_ref, 0) - self.assertEqual(tree.full_protected_size_, 0) - self.assertEqual(tree.swa_protected_size_, 0) - # dec_lock_ref releases locks but doesn't free; eviction does. - self.assertEqual(allocator.full_available_size(), full_avail_before) - self.assertEqual(allocator.swa_available_size(), swa_avail_before) - tree.sanity_check() - - def test_evict_swa_leaf_with_full_lock_tombstones_in_place(self): - # Large window so inc_lock_ref locks the entire SWA chain. - tree, allocator, _ = _build_tree(sliding_window_size=64) - leaf = _insert_chain(tree, allocator, [1, 2, 3, 4]) - self.assertEqual(len(leaf.value), 4) - - inc_res = tree.inc_lock_ref(leaf) - _release_swa_lock_chain_in_place(tree, leaf, inc_res.swa_uuid_for_lock) - - self.assertEqual(leaf.full_lock_ref, 1) - self.assertEqual(leaf.swa_lock_ref, 0) - self.assertFalse(leaf.swa_tombstone) - self.assertTrue(tree.swa_lru_list.in_list(leaf)) - - swa_avail_before = allocator.swa_available_size() - swa_evictable_before = tree.swa_evictable_size_ - - # num_tokens=0 skips the full eviction loop; swa loop hits the new branch. - evict_res = tree.evict(EvictParams(num_tokens=0, swa_num_tokens=4)) - - self.assertGreaterEqual(evict_res.swa_num_tokens_evicted, 4) - self.assertTrue(leaf.swa_tombstone) - self.assertFalse(tree.swa_lru_list.in_list(leaf)) - self.assertEqual(leaf.full_lock_ref, 1) - self.assertEqual( - allocator.swa_available_size(), swa_avail_before + len(leaf.value) - ) - # Full lock prevents _delete_leaf, so the node stays attached. - self.assertIs(leaf.parent.children[leaf.key.child_key(tree.page_size)], leaf) - self.assertEqual( - tree.swa_evictable_size_, swa_evictable_before - len(leaf.value) - ) - - tree.dec_lock_ref( - leaf, - DecLockRefParams(swa_uuid_for_lock=inc_res.swa_uuid_for_lock), - skip_swa=True, - ) - tree.sanity_check() - - def test_delete_leaf_skips_swa_size_on_tombstone(self): - # Tombstone removes the count once; _delete_leaf must not subtract again. - tree, allocator, _ = _build_tree(sliding_window_size=4) - leaf = _insert_chain(tree, allocator, [1, 2, 3, 4, 5, 6, 7, 8]) - - inc_res = tree.inc_lock_ref(leaf) - swa_uuid = inc_res.swa_uuid_for_lock - - tree.dec_swa_lock_only(leaf, DecLockRefParams(swa_uuid_for_lock=swa_uuid)) - self.assertTrue(leaf.swa_tombstone) - - swa_evictable_before_delete = tree.swa_evictable_size_ - tree.full_lru_list.remove_node(leaf) - tree._delete_leaf(leaf) - - self.assertEqual(tree.swa_evictable_size_, swa_evictable_before_delete) - - def test_dec_swa_lock_only_leaf_page_size_variants(self): - """Single-leaf tombstone+free across all (page_size, window) regimes. - - Sweep covers: - - window multiple of page_size (page_size=2, window=4) - - page_size > window (page_size=8, window=4) - - window not multiple of page (page_size=4, window=6) - - With page_size > 1, _swa_alloc routes through the paged allocators; - free_swa(leaf.value) must release exactly len(leaf.value) tokens - (page-aligned) regardless of how page_size relates to the window. - """ - for page_size, window in [(2, 4), (8, 4), (4, 6)]: - with self.subTest(page_size=page_size, window=window): - tree, allocator, _ = _build_tree( - sliding_window_size=window, - page_size=page_size, - kv_size=max(128, 32 * page_size), - kv_size_swa=max(64, 16 * page_size), - ) - n_tokens = max(window, 2 * page_size) - n_tokens = (n_tokens + page_size - 1) // page_size * page_size - leaf = _insert_chain(tree, allocator, list(range(1, n_tokens + 1))) - self.assertEqual(len(leaf.value), n_tokens) - self.assertEqual(len(leaf.value) % page_size, 0) - - inc_res = tree.inc_lock_ref(leaf) - swa_uuid = inc_res.swa_uuid_for_lock - self.assertIsNotNone( - swa_uuid, - f"inc_lock_ref must reach the window with leaf.value=" - f"{len(leaf.value)} >= window={window}", - ) - - swa_avail_before = allocator.swa_available_size() - full_avail_before = allocator.full_available_size() - - tree.dec_swa_lock_only( - leaf, DecLockRefParams(swa_uuid_for_lock=swa_uuid) - ) - - self.assertTrue(leaf.swa_tombstone) - self.assertFalse(tree.swa_lru_list.in_list(leaf)) - self.assertEqual(leaf.swa_lock_ref, 0) - self.assertEqual( - allocator.swa_available_size(), - swa_avail_before + len(leaf.value), - "free_swa must release the leaf's full page-aligned slot count", - ) - self.assertEqual(leaf.full_lock_ref, 1) - self.assertEqual(allocator.full_available_size(), full_avail_before) - - tree.dec_lock_ref( - leaf, - DecLockRefParams(swa_uuid_for_lock=swa_uuid), - skip_swa=True, - ) - tree.sanity_check() - - def test_dec_swa_lock_only_internal_page_size_gt_1(self): - """Internal-node chain release with page_size > 1. - - Two siblings sharing a page-aligned prefix force a radix split on a - page boundary. The swa lock chain therefore spans leaf -> internal, - and dec_swa_lock_only must: - - tombstone the leaf and free len(leaf.value) SWA tokens - - flip the internal node from protected -> evictable (no free, - no tombstone) - """ - page_size, window = 2, 6 - tree, allocator, _ = _build_tree( - sliding_window_size=window, page_size=page_size - ) - # Shared prefix len 4 (2 pages); divergent suffix len 2 (1 page each). - leaf_a = _insert_chain(tree, allocator, [1, 2, 3, 4, 5, 6]) - _insert_chain(tree, allocator, [1, 2, 3, 4, 7, 8]) - - self.assertEqual(len(leaf_a.value), 2) - internal = leaf_a.parent - self.assertGreater(len(internal.children), 1) - self.assertEqual(len(internal.value), 4) - - inc_res = tree.inc_lock_ref(leaf_a) - swa_uuid = inc_res.swa_uuid_for_lock - # leaf_a (2) + internal (4) = 6 >= window=6, so uuid stops at internal. - self.assertEqual(swa_uuid, internal.swa_uuid) - - swa_protected_before = tree.swa_protected_size_ - swa_evictable_before = tree.swa_evictable_size_ - swa_avail_before = allocator.swa_available_size() - - tree.dec_swa_lock_only(leaf_a, DecLockRefParams(swa_uuid_for_lock=swa_uuid)) - - # Leaf side: tombstoned and pages freed. - self.assertTrue(leaf_a.swa_tombstone) - self.assertFalse(tree.swa_lru_list.in_list(leaf_a)) - self.assertEqual( - allocator.swa_available_size(), - swa_avail_before + len(leaf_a.value), - ) - # Internal side: protected -> evictable, still in lru, no free. - self.assertFalse(internal.swa_tombstone) - self.assertTrue(tree.swa_lru_list.in_list(internal)) - self.assertEqual(internal.swa_lock_ref, 0) - self.assertEqual( - tree.swa_protected_size_, - swa_protected_before - (len(leaf_a.value) + len(internal.value)), - ) - self.assertEqual( - tree.swa_evictable_size_, - swa_evictable_before + len(internal.value), - ) - - tree.dec_lock_ref( - leaf_a, DecLockRefParams(swa_uuid_for_lock=swa_uuid), skip_swa=True - ) - tree.sanity_check() - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/unit/mem_cache/test_swa_unittest.py b/test/registered/unit/mem_cache/test_swa_unittest.py index 0adadf515..cab272388 100644 --- a/test/registered/unit/mem_cache/test_swa_unittest.py +++ b/test/registered/unit/mem_cache/test_swa_unittest.py @@ -1,12 +1,10 @@ import unittest -from array import array from types import SimpleNamespace from unittest import mock from unittest.mock import patch import torch -from sglang.srt.disaggregation.kv_events import BlockRemoved, BlockStored from sglang.srt.environ import InvariantCheckLevel, envs from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.allocator.swa import ( @@ -15,21 +13,12 @@ from sglang.srt.mem_cache.allocator.swa import ( ) from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, - DecLockRefParams, - EvictParams, - EvictResult, - InsertParams, - MatchPrefixParams, ) -from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.common import ( - available_and_evictable_str, free_kv_row_segments, ) from sglang.srt.mem_cache.memory_pool import ReqToTokenPool -from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool -from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache from sglang.srt.utils import get_device from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.test_utils import CustomTestCase @@ -38,10 +27,6 @@ register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-large") register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd") -def _event_hashes(events): - return [block_hash for event in events for block_hash in event.block_hashes] - - class _DummyReq: def __init__(self): self._kv_committed_len = 0 @@ -100,18 +85,7 @@ def _build_swa_tree( need_sort=False, req_to_token_pool=req_to_token_pool, ) - tree = SWARadixCache( - params=CacheInitParams( - req_to_token_pool=req_to_token_pool, - token_to_kv_pool_allocator=allocator, - page_size=page_size, - disable=False, - is_eagle=is_eagle, - sliding_window_size=sliding_window_size, - enable_kv_cache_events=enable_kv_cache_events, - ), - ) - return tree, allocator, req_to_token_pool + return allocator, req_to_token_pool def _sync_error(fn): @@ -164,23 +138,6 @@ def _swa_alloc(allocator, need_size): return full_indices -def _insert(tree, allocator, token_ids): - indices = _swa_alloc(allocator, len(token_ids)) - assert indices is not None - tree.insert(InsertParams(key=RadixKey(array("q", token_ids)), value=indices)) - - -def _insert_chain(tree, allocator, token_ids): - _insert(tree, allocator, token_ids) - match = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", token_ids)))) - return match.last_device_node - - -def _expected_tail_size(window: int, page_size: int) -> int: - """Mirror of _maybe_split_leaf_for_swa_lock's tail_size formula.""" - return (window + page_size - 1) // page_size * page_size - - class TestSWA(unittest.TestCase): @classmethod def setUpClass(cls): @@ -190,71 +147,9 @@ class TestSWA(unittest.TestCase): def tearDownClass(cls): pass - def test_swa_radix_cache_kv_events(self): - tree, allocator, _ = _build_swa_tree( - is_eagle=False, enable_kv_cache_events=True - ) - tree.take_events() # Clear the reset event. - - _insert(tree, allocator, [1, 2, 3, 4]) - first_insert_events = [ - e for e in tree.take_events() if isinstance(e, BlockStored) - ] - self.assertEqual(len(first_insert_events), 1) - self.assertEqual(list(first_insert_events[0].token_ids), [1, 2, 3, 4]) - - _insert(tree, allocator, [1, 2, 3, 4, 5, 6]) - second_insert_events = [ - e for e in tree.take_events() if isinstance(e, BlockStored) - ] - self.assertEqual(len(second_insert_events), 1) - self.assertEqual(list(second_insert_events[0].token_ids), [5, 6]) - - stored_hashes = [ - block_hash - for event in first_insert_events + second_insert_events - for block_hash in event.block_hashes - ] - - # Evicting only SWA tokens tombstones nodes but keeps full KV blocks. - result = tree.evict(EvictParams(num_tokens=0, swa_num_tokens=1)) - self.assertEqual(result.num_tokens_evicted, 0) - self.assertGreaterEqual(result.swa_num_tokens_evicted, 1) - self.assertEqual( - [e for e in tree.take_events() if isinstance(e, BlockRemoved)], [] - ) - - result = tree.evict(EvictParams(num_tokens=1, swa_num_tokens=0)) - self.assertGreaterEqual(result.num_tokens_evicted, 1) - removed_hashes = _event_hashes( - [e for e in tree.take_events() if isinstance(e, BlockRemoved)] - ) - self.assertCountEqual(removed_hashes, stored_hashes) - - def test_swa_radix_cache_kv_events_split_hash(self): - tree, allocator, _ = _build_swa_tree( - is_eagle=False, enable_kv_cache_events=True - ) - tree.take_events() # Clear the reset event. - - _insert(tree, allocator, [1, 2, 3, 4]) - first_insert_events = [ - e for e in tree.take_events() if isinstance(e, BlockStored) - ] - self.assertEqual(len(first_insert_events), 1) - split_parent_hash = first_insert_events[0].block_hashes[1] - - _insert(tree, allocator, [1, 2, 5, 6]) - second_insert_events = [ - e for e in tree.take_events() if isinstance(e, BlockStored) - ] - self.assertEqual(len(second_insert_events), 1) - self.assertEqual(list(second_insert_events[0].token_ids), [5, 6]) - self.assertEqual(second_insert_events[0].parent_block_hash, split_parent_hash) - def test_swa_memory_pool_paged_free_clears_full_page_mapping(self): page_size = 4 - _, allocator, _ = _build_swa_tree( + allocator, _ = _build_swa_tree( is_eagle=False, page_size=page_size, kv_size=16, @@ -281,7 +176,7 @@ class TestSWA(unittest.TestCase): """Clearing the full-to-SWA mapping must not block the stream; writing a host-resident scalar into it does. """ - _, allocator, _ = _build_swa_tree(is_eagle=False) + allocator, _ = _build_swa_tree(is_eagle=False) full_indices = _swa_alloc(allocator, 4) mapping = allocator.full_to_swa_index_mapping @@ -307,7 +202,7 @@ class TestSWA(unittest.TestCase): self._free_swa_group_owns_deferred_indices(page_size) def _free_swa_group_owns_deferred_indices(self, page_size): - _, allocator, _ = _build_swa_tree( + allocator, _ = _build_swa_tree( is_eagle=False, page_size=page_size, kv_size=32 * page_size, @@ -344,7 +239,7 @@ class TestSWA(unittest.TestCase): ) def test_free_swa_group_owns_mapping_at_enqueue_time(self): - _, allocator, _ = _build_swa_tree( + allocator, _ = _build_swa_tree( is_eagle=False, kv_size=8, kv_size_swa=8, @@ -376,7 +271,7 @@ class TestSWA(unittest.TestCase): ) def _build_two_mapped_slots(self, page_size=1): - _, allocator, _ = _build_swa_tree( + allocator, _ = _build_swa_tree( is_eagle=False, page_size=page_size, kv_size=8 * page_size, @@ -447,524 +342,6 @@ class TestSWA(unittest.TestCase): allocator.full_to_swa_index_mapping[indices], indices ) - def test_swa_radix_cache_1(self): - # args - req_size = 10 - max_context_len = 128 - kv_size = 128 - kv_size_swa = 64 - page_size = 1 - sliding_window_size = 4 - head_num = 8 - head_dim = 128 - num_layers = 48 - global_interval = 4 - dtype = torch.bfloat16 - device = get_device() - full_attention_layer_ids = [i for i in range(0, num_layers, global_interval)] - full_attention_layer_ids_set = set(full_attention_layer_ids) - swa_attention_layer_ids = [ - i for i in range(num_layers) if i not in full_attention_layer_ids_set - ] - # setup req to token pool - req_to_token_pool = ReqToTokenPool( - size=req_size, - max_context_len=max_context_len, - device=device, - enable_memory_saver=False, - ) - # setup kv pool - kv_pool = SWAKVPool( - size=kv_size, - size_swa=kv_size_swa, - page_size=page_size, - dtype=dtype, - head_num=head_num, - head_dim=head_dim, - swa_attention_layer_ids=swa_attention_layer_ids, - full_attention_layer_ids=full_attention_layer_ids, - device=device, - ) - # setup token to kv pool allocator - allocator = SWATokenToKVPoolAllocator( - size=kv_size, - size_swa=kv_size_swa, - page_size=page_size, - dtype=dtype, - device=device, - kvcache=kv_pool, - need_sort=False, - ) - # setup radix cache - tree = SWARadixCache( - params=CacheInitParams( - req_to_token_pool=req_to_token_pool, - token_to_kv_pool_allocator=allocator, - disable=False, - page_size=page_size, - sliding_window_size=sliding_window_size, - ), - ) - - # test - print( - f"[Start] allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}" - ) - req1_token_ids, req1_kv_indices = [1, 2, 3], allocator.alloc(3) - self.assertEqual(len(req1_token_ids), len(req1_kv_indices)) - print( - f"req1: inserting, req1_token_ids: {req1_token_ids}, req1_kv_indices: {req1_kv_indices}" - ) - key = RadixKey(array("q", req1_token_ids)) - result = tree.insert(InsertParams(key=key, value=req1_kv_indices[: len(key)])) - 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()}" - ) - req2_token_ids, req2_kv_indices = [1, 2, 3, 4, 5, 6, 7], allocator.alloc(7) - self.assertEqual(len(req2_token_ids), len(req2_kv_indices)) - print( - f"req2: inserting, req2_token_ids: {req2_token_ids}, req2_kv_indices: {req2_kv_indices}" - ) - key = RadixKey(array("q", req2_token_ids)) - result = tree.insert(InsertParams(key=key, value=req2_kv_indices[: len(key)])) - 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()}" - ) - req3_token_ids, req3_kv_indices = [10, 11, 12], allocator.alloc(3) - self.assertEqual(len(req3_token_ids), len(req3_kv_indices)) - print( - f"req3: inserting, req3_token_ids: {req3_token_ids}, req3_kv_indices: {req3_kv_indices}" - ) - key = RadixKey(array("q", req3_token_ids)) - result = tree.insert(InsertParams(key=key, value=req3_kv_indices[: len(key)])) - 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()}" - ) - req4_token_ids, req4_kv_indices = [1, 2, 3, 4, 5, 60, 70], allocator.alloc(7) - self.assertEqual(len(req4_token_ids), len(req4_kv_indices)) - print( - f"req4: inserting, req4_token_ids: {req4_token_ids}, req4_kv_indices: {req4_kv_indices}" - ) - key = RadixKey(array("q", req4_token_ids)) - result = tree.insert(InsertParams(key=key, value=req4_kv_indices[: len(key)])) - 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()}" - ) - - 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( - 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( - 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( - EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) - ) - tree.pretty_print() - - req5_token_ids = [1, 2, 3, 4, 5] - result = tree.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", req5_token_ids))) - ) - kv_indices, last_node = result.device_indices, result.last_device_node - print( - f"req5: token_ids: {req5_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}" - ) - self.assertEqual(len(kv_indices), 0) - - req6_token_ids = [1, 2, 3, 4, 5, 60, 70] - result = tree.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", req6_token_ids))) - ) - kv_indices, last_node = result.device_indices, result.last_device_node - print( - f"req6: token_ids: {req6_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}" - ) - self.assertEqual(len(kv_indices), 7) - self.assertEqual(len(last_node.key), 2) - self.assertEqual(last_node.key.token_ids[0], 60) - self.assertEqual(last_node.key.token_ids[1], 70) - - print(tree.available_and_evictable_str()) - print(available_and_evictable_str(tree)) - tree.sanity_check() - - def test_swa_radix_cache_eagle(self): - # args - req_size = 10 - max_context_len = 128 - kv_size = 128 - kv_size_swa = 64 - page_size = 1 - sliding_window_size = 4 - head_num = 8 - head_dim = 128 - num_layers = 48 - global_interval = 4 - dtype = torch.bfloat16 - device = get_device() - full_attention_layer_ids = [i for i in range(0, num_layers, global_interval)] - full_attention_layer_ids_set = set(full_attention_layer_ids) - swa_attention_layer_ids = [ - i for i in range(num_layers) if i not in full_attention_layer_ids_set - ] - # setup req to token pool - req_to_token_pool = ReqToTokenPool( - size=req_size, - max_context_len=max_context_len, - device=device, - enable_memory_saver=False, - ) - # setup kv pool - kv_pool = SWAKVPool( - size=kv_size, - size_swa=kv_size_swa, - page_size=page_size, - dtype=dtype, - head_num=head_num, - head_dim=head_dim, - swa_attention_layer_ids=swa_attention_layer_ids, - full_attention_layer_ids=full_attention_layer_ids, - device=device, - ) - # setup token to kv pool allocator - allocator = SWATokenToKVPoolAllocator( - size=kv_size, - size_swa=kv_size_swa, - page_size=page_size, - dtype=dtype, - device=device, - kvcache=kv_pool, - need_sort=False, - ) - # setup radix cache - tree = SWARadixCache( - params=CacheInitParams( - req_to_token_pool=req_to_token_pool, - token_to_kv_pool_allocator=allocator, - page_size=page_size, - disable=False, - is_eagle=True, - sliding_window_size=sliding_window_size, - ), - ) - - # test - print( - f"[Start] allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}" - ) - req1_token_ids, req1_kv_indices = [1, 2, 3], allocator.alloc(3) - self.assertEqual(len(req1_token_ids), len(req1_kv_indices)) - print( - f"req1: inserting, req1_token_ids: {req1_token_ids}, req1_kv_indices: {req1_kv_indices}" - ) - key = RadixKey(array("q", req1_token_ids)) - result = tree.insert(InsertParams(key=key, value=req1_kv_indices[: len(key)])) - 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()}" - ) - req2_token_ids, req2_kv_indices = [1, 2, 3, 4, 5, 6, 7], allocator.alloc(7) - self.assertEqual(len(req2_token_ids), len(req2_kv_indices)) - print( - f"req2: inserting, req2_token_ids: {req2_token_ids}, req2_kv_indices: {req2_kv_indices}" - ) - key = RadixKey(array("q", req2_token_ids)) - result = tree.insert(InsertParams(key=key, value=req2_kv_indices[: len(key)])) - 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()}" - ) - req3_token_ids, req3_kv_indices = [10, 11, 12], allocator.alloc(3) - self.assertEqual(len(req3_token_ids), len(req3_kv_indices)) - print( - f"req3: inserting, req3_token_ids: {req3_token_ids}, req3_kv_indices: {req3_kv_indices}" - ) - key = RadixKey(array("q", req3_token_ids)) - result = tree.insert(InsertParams(key=key, value=req3_kv_indices[: len(key)])) - 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()}" - ) - req4_token_ids, req4_kv_indices = [1, 2, 3, 4, 5, 60, 70], allocator.alloc(7) - self.assertEqual(len(req4_token_ids), len(req4_kv_indices)) - print( - f"req4: inserting, req4_token_ids: {req4_token_ids}, req4_kv_indices: {req4_kv_indices}" - ) - key = RadixKey(array("q", req4_token_ids)) - result = tree.insert(InsertParams(key=key, value=req4_kv_indices[: len(key)])) - 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()}" - ) - - 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") - 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") - 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") - 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] - result = tree.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", req5_token_ids))) - ) - kv_indices, last_node = result.device_indices, result.last_device_node - print( - f"req5: token_ids: {req5_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}" - ) - self.assertEqual(len(kv_indices), 0) # no swa prefix matched - - req6_token_ids = [1, 2, 3, 4, 5, 60, 70] - result = tree.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", req6_token_ids))) - ) - kv_indices, last_node = result.device_indices, result.last_device_node - print( - f"req6: token_ids: {req6_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}" - ) - self.assertEqual(len(kv_indices), 6) - self.assertEqual(len(last_node.key), 2) - # Bigram view: token_ids holds raw tokens; iteration yields bigram tuples. - self.assertTrue(last_node.key.is_bigram) - self.assertEqual(list(last_node.key), [(5, 60), (60, 70)]) - - def test_swa_cache_finished_req_eagle_uses_cache_protected_len_and_bigram_key(self): - tree, allocator, req_to_token_pool = _build_swa_tree(is_eagle=True) - - # Case 1: is_insert=True should pass bigram key and use cache_protected_len. - req = _DummyReq() - req.kv.req_pool_idx = 0 - req.origin_input_ids = array("q", [1, 2, 3, 4, 5, 6]) - req.output_ids = array("q") - req._kv_committed_len = len(req.origin_input_ids) - kv_indices = allocator.alloc(req._kv_committed_len) - req_to_token_pool.write( - (req.kv.req_pool_idx, slice(0, req._kv_committed_len)), kv_indices - ) - req.extra_key = None - req.cache_salt = None - req.last_node = tree.root_node - req.lock_receipt = DecLockRefParams() - req.kv.swa_evicted_seqlen = 0 - req.kv.cache_protected_len = 1 - # Intentionally mismatch to ensure code does not use len(prefix_indices). - req.prefix_indices = torch.tensor([7, 8, 9, 10, 11], device=tree.device) - - captured = {} - original_insert = tree.insert - - def wrapped_insert(params): - captured["prev_prefix_len"] = params.prev_prefix_len - captured["is_bigram"] = params.key.is_bigram - captured["key_len"] = len(params.key) - return original_insert(params) - - tree.insert = wrapped_insert - tree.cache_finished_req(req, is_insert=True, owned_kv_len=req._kv_committed_len) - - self.assertEqual(captured["prev_prefix_len"], req.kv.cache_protected_len) - self.assertTrue(captured["is_bigram"]) - self.assertEqual(captured["key_len"], len(req.origin_input_ids) - 1) - - # Case 2: is_insert=False should free [cache_protected_len:page_aligned_len] - # even when len(prefix_indices) is intentionally larger. - req2 = _DummyReq() - req2.kv.req_pool_idx = 1 - req2.origin_input_ids = array("q", [11, 12, 13, 14, 15, 16]) - req2.output_ids = array("q") - req2._kv_committed_len = len(req2.origin_input_ids) - kv_indices2 = allocator.alloc(req2._kv_committed_len) - req_to_token_pool.write( - (req2.kv.req_pool_idx, slice(0, req2._kv_committed_len)), kv_indices2 - ) - req2.extra_key = None - req2.cache_salt = None - req2.last_node = tree.root_node - req2.lock_receipt = DecLockRefParams() - req2.kv.swa_evicted_seqlen = 0 - req2.kv.cache_protected_len = 1 - req2.prefix_indices = torch.tensor([21, 22, 23, 24, 25], device=tree.device) - - freed_lens = [] - original_free_segment = allocator.free_segment - - def wrapped_free_segment(indices, *, start_pos): - freed_lens.append(int(indices.numel())) - return original_free_segment(indices, start_pos=start_pos) - - allocator.free_segment = wrapped_free_segment - tree.cache_finished_req( - req2, is_insert=False, owned_kv_len=req2._kv_committed_len - ) - - # EAGLE + page_size=1 => page_aligned_len = committed_len - 1 = 5 - # Expected frees: - # overlap range [1:5] -> 4 - # tail range [5:] -> 1 - self.assertEqual(freed_lens, [4, 1]) - - -# Optimization: SGLANG_OPT_SWA_SPLIT_LEAF_ON_INSERT. -# Splits a freshly-inserted leaf at the (page-aligned) sliding-window -# boundary so a future inc_lock_ref protects only ~sliding_window_size SWA -# tokens instead of the whole chunked-prefill chain. -class TestSWASplitLeafOnInsert(CustomTestCase): - def _insert_and_lock(self, *, window, page_size, leaf_len, flag_on): - tree, allocator, _ = _build_swa_tree( - is_eagle=False, - kv_size=128, - kv_size_swa=64, - sliding_window_size=window, - page_size=page_size, - ) - token_ids = list(range(leaf_len)) - with envs.SGLANG_OPT_SWA_SPLIT_LEAF_ON_INSERT.override(flag_on): - leaf = _insert_chain(tree, allocator, token_ids) - result = tree.inc_lock_ref(leaf) - return tree, leaf, result - - def test_flag_off_protects_full_leaf(self): - tree, leaf, _ = self._insert_and_lock( - window=4, page_size=1, leaf_len=12, flag_on=False - ) - self.assertEqual(len(leaf.value), 12) - self.assertEqual(tree.swa_protected_size_, 12) - - def test_flag_on_caps_protection_at_window(self): - # (window, page_size, leaf_len, expected_tail_size); leaf_len picked - # > tail_size and page-aligned for page_size > 1. - cases = [ - (4, 1, 12, 4), - (4, 1, 5, 4), - (1, 1, 5, 1), - (4, 2, 12, 4), - (8, 2, 12, 8), - (4, 4, 12, 4), - # window NOT page-aligned -> tail rounds up to page boundary. - (3, 2, 12, 4), - (5, 4, 12, 8), - (3, 4, 12, 4), - ] - for window, page_size, leaf_len, expected_tail in cases: - with self.subTest(window=window, page_size=page_size, leaf_len=leaf_len): - self.assertEqual(_expected_tail_size(window, page_size), expected_tail) - tree, leaf, _ = self._insert_and_lock( - window=window, - page_size=page_size, - leaf_len=leaf_len, - flag_on=True, - ) - self.assertEqual(len(leaf.value), expected_tail) - self.assertEqual(tree.swa_protected_size_, expected_tail) - - def test_flag_on_no_split_when_leaf_within_window(self): - # leaf_len <= tail_size: split must no-op. - cases = [ - (4, 1, 4), - (4, 1, 3), - (4, 2, 4), - (3, 2, 4), - (8, 2, 4), - (4, 4, 4), - ] - for window, page_size, leaf_len in cases: - with self.subTest(window=window, page_size=page_size, leaf_len=leaf_len): - tree, leaf, _ = self._insert_and_lock( - window=window, - page_size=page_size, - leaf_len=leaf_len, - flag_on=True, - ) - self.assertEqual(len(leaf.value), leaf_len) - self.assertEqual(tree.swa_protected_size_, leaf_len) - - def test_match_prefix_returns_full_chain_after_split(self): - tree, allocator, _ = _build_swa_tree( - is_eagle=False, - kv_size=128, - kv_size_swa=64, - sliding_window_size=4, - page_size=1, - ) - token_ids = list(range(12)) - with envs.SGLANG_OPT_SWA_SPLIT_LEAF_ON_INSERT.override(True): - inserted_leaf = _insert_chain(tree, allocator, token_ids) - self.assertEqual(len(inserted_leaf.value), 4) - match = tree.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", token_ids))) - ) - self.assertEqual(match.device_indices.shape[0], 12) - self.assertIs(match.last_device_node, inserted_leaf) - - def test_dec_lock_ref_after_split_balances_to_zero(self): - tree, leaf, result = self._insert_and_lock( - window=4, page_size=1, leaf_len=12, flag_on=True - ) - self.assertEqual(tree.swa_protected_size_, 4) - self.assertEqual(tree.full_protected_size_, 12) - - tree.dec_lock_ref( - leaf, - params=DecLockRefParams(swa_uuid_for_lock=result.swa_uuid_for_lock), - ) - - self.assertEqual(tree.swa_protected_size_, 0) - self.assertEqual(tree.full_protected_size_, 0) - tree.sanity_check() - class _SinglePoolAllocator(BaseTokenToKVPoolAllocator): """Minimal single-pool allocator: no SWA peer, so the whole range dies @@ -995,7 +372,7 @@ class TestFreeFullPartition(CustomTestCase): """`free_full` releases only the full side of a hybrid SWA allocator.""" def setUp(self): - _, self.allocator, _ = _build_swa_tree(is_eagle=False) + self.allocator, _ = _build_swa_tree(is_eagle=False) self.full_baseline = self.allocator.full_available_size() self.swa_baseline = self.allocator.swa_available_size() @@ -1043,7 +420,7 @@ class TestFreeKvRow(CustomTestCase): whole, the SWA side only from the floor up.""" def setUp(self): - _, self.allocator, _ = _build_swa_tree(is_eagle=False) + self.allocator, _ = _build_swa_tree(is_eagle=False) self.full_baseline = self.allocator.full_available_size() self.swa_baseline = self.allocator.swa_available_size() @@ -1080,7 +457,7 @@ class TestFreeKvRow(CustomTestCase): self.assertEqual(self._sizes(), (self.full_baseline, self.swa_baseline)) def test_below_floor_pieces_go_back_through_the_full_side(self): - _, allocator, _ = _build_swa_tree(is_eagle=False, page_size=4) + allocator, _ = _build_swa_tree(is_eagle=False, page_size=4) indices = _swa_alloc(allocator, 8) allocator.free_swa(indices) after_alloc = allocator.full_available_size() @@ -1100,7 +477,7 @@ class TestFreeKvRow(CustomTestCase): self.assertEqual(allocator.full_available_size(), after_alloc + 8) def test_grouped_full_side_frees_defer_and_skip_the_unique_path(self): - _, allocator, _ = _build_swa_tree(is_eagle=False, page_size=4) + allocator, _ = _build_swa_tree(is_eagle=False, page_size=4) indices = _swa_alloc(allocator, 12) allocator.free_swa(indices[:8]) after_alloc = allocator.full_available_size() @@ -1163,7 +540,7 @@ class TestSWAPeerMappedContract(CustomTestCase): return bool(assert_async.call_args.args[0]) def test_segment_free_flags_a_page_whose_peer_is_already_gone(self): - _, allocator, _ = _build_swa_tree(is_eagle=False, page_size=4) + allocator, _ = _build_swa_tree(is_eagle=False, page_size=4) live = _swa_alloc(allocator, 8) stale = _swa_alloc(allocator, 8) allocator.clear_full_to_swa_mapping(stale) @@ -1176,7 +553,7 @@ class TestSWAPeerMappedContract(CustomTestCase): """page_size > 1: page reps by stride replace the page expansion's filter and the inner allocator's torch.unique, in and out of a group.""" ps = 4 - _, allocator, _ = _build_swa_tree(is_eagle=False, page_size=ps) + allocator, _ = _build_swa_tree(is_eagle=False, page_size=ps) def grouped(indices): allocator.free_group_begin() @@ -1204,7 +581,7 @@ class TestSWAPeerMappedContract(CustomTestCase): self.assertIsNone(_sync_error(lambda: grouped(second[: 2 * ps - 1]))) def test_free_swa_flags_a_slot_whose_peer_is_already_gone(self): - _, allocator, _ = _build_swa_tree(is_eagle=False) + allocator, _ = _build_swa_tree(is_eagle=False) live = _swa_alloc(allocator, 4) stale = _swa_alloc(allocator, 4) # Whoever released the peer left the mapping reading as the padding slot. @@ -1217,7 +594,7 @@ class TestSWAPeerMappedContract(CustomTestCase): def test_free_swa_does_not_synchronize(self): """The filter's output shape was data-dependent, so it read a count back to the host; the gather that replaced it has a fixed shape.""" - _, allocator, _ = _build_swa_tree(is_eagle=False) + allocator, _ = _build_swa_tree(is_eagle=False) mapping = allocator.full_to_swa_index_mapping # Warm up outside the window: a first-time cudaMalloc can synchronize on @@ -1241,7 +618,7 @@ class TestSWAReqRingFree(CustomTestCase): def _allocated_ring(self): ps = self.PS - _, allocator, req_pool = _build_swa_tree( + allocator, req_pool = _build_swa_tree( is_eagle=False, page_size=ps, req_size=2, @@ -1361,7 +738,7 @@ class TestSWAPageRepsFree(CustomTestCase): PS = 4 def _allocator(self): - _, allocator, _ = _build_swa_tree(is_eagle=False, page_size=self.PS) + allocator, _ = _build_swa_tree(is_eagle=False, page_size=self.PS) return allocator def _sizes(self, allocator): @@ -1371,7 +748,7 @@ class TestSWAPageRepsFree(CustomTestCase): def test_free_swa_segment_npu_uses_reference_path(self): for page_size in (1, 4): with self.subTest(page_size=page_size): - _, allocator, _ = _build_swa_tree( + allocator, _ = _build_swa_tree( is_eagle=False, page_size=page_size, kv_size=8 * page_size, @@ -1423,7 +800,7 @@ class TestSWAPageRepsFree(CustomTestCase): with self.subTest(name=name): num_allocated_tokens = max(2 * page_size, num_tokens) kv_size = max(8 * page_size, num_allocated_tokens) - _, allocator, _ = _build_swa_tree( + allocator, _ = _build_swa_tree( is_eagle=False, page_size=page_size, kv_size=kv_size, @@ -1467,90 +844,6 @@ class TestSWAPageRepsFree(CustomTestCase): self.assertTrue(torch.all(mapping[indices[:touched]] == 0)) self.assertTrue(torch.all(mapping[indices[touched:]] > 0)) - def test_node_frees_take_the_page_path_through_the_tree(self): - """Tree values are page-aligned copies of a kv row, so SWA eviction and - the full eviction of its tombstones both free by page reps.""" - ps = self.PS - tree, allocator, _ = _build_swa_tree( - is_eagle=False, page_size=ps, sliding_window_size=ps - ) - full_before, swa_before = self._sizes(allocator) - _insert(tree, allocator, list(range(1, 3 * ps + 1))) - - # Either inner `free` is the torch.unique path a caller falls back to - # when it hands no start position. - with ( - patch.object( - allocator.full_attn_allocator, - "free", - side_effect=AssertionError("full side took the unique path"), - ), - patch.object( - allocator.swa_attn_allocator, - "free", - side_effect=AssertionError("swa side took the unique path"), - ), - ): - tree.evict(EvictParams(num_tokens=0, swa_num_tokens=ps)) - tree.evict(EvictParams(num_tokens=3 * ps, swa_num_tokens=0)) - - self.assertEqual(self._sizes(allocator), (full_before, swa_before)) - - -class TestCacheUnfinishedReqEvictedPrefix(CustomTestCase): - """An unfinished request whose SWA prefix is already gone must insert that - prefix as a tombstone, not as live SWA KV.""" - - def test_evicted_prefix_inserts_as_tombstone(self): - page_size, window, num_tokens, evicted = 4, 4, 16, 8 - tree, allocator, req_to_token_pool = _build_swa_tree( - is_eagle=False, page_size=page_size, sliding_window_size=window - ) - kv_indices = _swa_alloc(allocator, num_tokens) - req_to_token_pool.write((0, slice(0, num_tokens)), kv_indices) - # Drop the prefix's SWA peers, as window eviction would. - allocator.free_swa(kv_indices[:evicted]) - swa_before = allocator.swa_available_size() - - token_ids = array("q", range(1, num_tokens + 1)) - req = _DummyReq() - req.kv.req_pool_idx = 0 - req.origin_input_ids = token_ids - req.output_ids = array("q") - req.get_fill_ids = lambda: token_ids - req.extra_key = None - req.cache_salt = None - req.kv.cache_protected_len = 0 - req.last_node = tree.root_node - req.lock_receipt = DecLockRefParams() - req.prefix_indices = torch.empty(0, dtype=torch.int64, device=tree.device) - req.kv.swa_evicted_seqlen = evicted - - tree.cache_unfinished_req(req) - - # The insert itself frees nothing. - self.assertEqual(allocator.swa_available_size(), swa_before) - # The live leaf holds a full window, so the whole key stays matchable. - self.assertEqual(req.kv.cache_protected_len, num_tokens) - # [0, evicted) is a tombstone; only [evicted, num_tokens) counts as SWA. - (first,) = tree.root_node.children.values() - self.assertTrue(first.swa_tombstone) - self.assertEqual(len(first.value), evicted) - self.assertEqual( - tree.swa_evictable_size_ + tree.swa_protected_size_, - num_tokens - evicted, - ) - - # Finishing drops the locks, which sanity_check needs; the accounting - # must survive the re-walk. - tree.cache_finished_req(req, owned_kv_len=num_tokens) - self.assertEqual(allocator.swa_available_size(), swa_before) - self.assertEqual( - tree.swa_evictable_size_ + tree.swa_protected_size_, - num_tokens - evicted, - ) - tree.sanity_check() - if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py index c88f6e9a2..d02c06ee5 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py @@ -1,7 +1,7 @@ """Large-scale benchmark + fuzz correctness tests for UnifiedRadixCache. Usage (standalone): - bench: python3 test/registered/unit/mem_cache/test_unified_radix_cache_bench.py --bench --num-seqs 5000 --verify --components mamba legacy-mamba swa legacy-swa + bench: python3 test/registered/unit/mem_cache/test_unified_radix_cache_bench.py --bench --num-seqs 5000 --verify --components mamba swa CI Test: python -m pytest test/registered/unit/mem_cache/test_unified_radix_cache_bench.py -v -s """ @@ -30,10 +30,8 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, ) from sglang.srt.mem_cache.cache_init_params import CacheInitParams -from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, HybridReqToTokenPool from sglang.srt.mem_cache.radix_cache import RadixKey -from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache from sglang.srt.mem_cache.unified_cache.components.base import ComponentType from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler @@ -131,7 +129,6 @@ def create_bench_cache( max_context_len, components, page_size=1, - tree_cls=None, sliding_window_size=_SWA_WINDOW_SIZE, ): """Create cache. Returns (tree, allocator, req_to_token_pool, make_req).""" @@ -225,21 +222,19 @@ def create_bench_cache( ) # --- tree --- - if tree_cls is None: - tree_cls = UnifiedRadixCache backend_override = ( envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.override(_TREE_CORE_TEST_BACKEND) - if _TREE_CORE_TEST_BACKEND is not None and tree_cls is UnifiedRadixCache + if _TREE_CORE_TEST_BACKEND is not None else nullcontext() ) with backend_override: - tree = tree_cls( + tree = UnifiedRadixCache( params=CacheInitParams( req_to_token_pool=req_to_token_pool, token_to_kv_pool_allocator=allocator, page_size=page_size, disable=False, - tree_components=components if tree_cls is UnifiedRadixCache else None, + tree_components=components, sliding_window_size=sliding_window_size if has_swa else None, ) ) @@ -279,7 +274,7 @@ class _Env: avg_tokens: int -def _make_env(num_seqs, chunk_len, kv_size, components, tree_cls=None, page_size=1): +def _make_env(num_seqs, chunk_len, kv_size, components, page_size=1): """Create sequences + cache, return shared _Env.""" if components is None: components = _DEFAULT_COMPONENTS @@ -293,7 +288,6 @@ def _make_env(num_seqs, chunk_len, kv_size, components, tree_cls=None, page_size max_context_len=max_seq_len + 10, components=components, page_size=page_size, - tree_cls=tree_cls, ) return _Env( tree, @@ -467,11 +461,10 @@ def bench_insert( kv_size=500_000, components=None, verify=False, - tree_cls=None, page_size=1, ): """Insert throughput (alloc + evict-fallback + insert).""" - env = _make_env(num_seqs, chunk_len, kv_size, components, tree_cls, page_size) + env = _make_env(num_seqs, chunk_len, kv_size, components, page_size) warmup = min(20, num_seqs // 10) return bench_api( @@ -491,11 +484,10 @@ def bench_match_prefix( kv_size=500_000, components=None, verify=False, - tree_cls=None, page_size=1, ): """Prefix matching throughput (hit / partial / miss mix).""" - env = _make_env(num_seqs, chunk_len, kv_size, components, tree_cls, page_size) + env = _make_env(num_seqs, chunk_len, kv_size, components, page_size) _populate(env, num_seqs // 2) rng = random.Random(123) @@ -535,11 +527,10 @@ def bench_evict( kv_size=500_000, components=None, verify=False, - tree_cls=None, page_size=1, ): """Eviction throughput — fill pool then repeatedly evict batches.""" - env = _make_env(num_seqs, chunk_len, kv_size, components, tree_cls, page_size) + env = _make_env(num_seqs, chunk_len, kv_size, components, page_size) inserted = _fill_no_evict(env) evict_batch = max(100, kv_size // 200) @@ -564,11 +555,10 @@ def bench_lock_unlock( kv_size=500_000, components=None, verify=False, - tree_cls=None, page_size=1, ): """Lock/unlock throughput — match nodes then cycle lock/unlock.""" - env = _make_env(num_seqs, chunk_len, kv_size, components, tree_cls, page_size) + env = _make_env(num_seqs, chunk_len, kv_size, components, page_size) _populate(env, num_seqs // 2) nodes = [] @@ -608,14 +598,13 @@ def bench_cache_finished( kv_size=500_000, components=None, verify=False, - tree_cls=None, page_size=1, ): """cache_finished_req throughput — full request lifecycle. Simulates: match_prefix → inc_lock_ref → alloc → fill req_to_token → cache_finished_req. """ - env = _make_env(num_seqs, chunk_len, kv_size, components, tree_cls, page_size) + env = _make_env(num_seqs, chunk_len, kv_size, components, page_size) # Pre-build Req objects with token IDs filled into req_to_token req_items: list = [] @@ -691,7 +680,6 @@ def run_all_benchmarks( components=None, verify=False, benchmarks=None, - tree_cls=None, page_size=1, ): if components is None: @@ -700,12 +688,12 @@ def run_all_benchmarks( benchmarks = list(ALL_BENCHMARKS.keys()) server_args = ServerArgs(model_path="dummy", page_size=page_size) - # MambaRadixCache reads mamba_cache_chunk_size, whose property otherwise + # The mamba component reads mamba_cache_chunk_size, whose property otherwise # loads the HF config for self.model_path — impossible for the dummy model. server_args._mamba_cache_chunk_size = max(FLA_CHUNK_SIZE, page_size) set_global_server_args_for_scheduler(server_args) - impl_name = (tree_cls or UnifiedRadixCache).__name__ + impl_name = UnifiedRadixCache.__name__ results = [] for name in benchmarks: if name not in ALL_BENCHMARKS: @@ -718,7 +706,6 @@ def run_all_benchmarks( kv_size=kv_size, components=components, verify=verify, - tree_cls=tree_cls, page_size=page_size, ) ) @@ -823,12 +810,10 @@ del _cfg, _name # CLI # =================================================================== _TREE_CONFIGS = { - "full": ((ComponentType.FULL,), None), - "mamba": ((ComponentType.FULL, ComponentType.MAMBA), None), - "swa": ((ComponentType.FULL, ComponentType.SWA), None), - "all": ((ComponentType.FULL, ComponentType.SWA, ComponentType.MAMBA), None), - "legacy-mamba": ((ComponentType.FULL, ComponentType.MAMBA), MambaRadixCache), - "legacy-swa": ((ComponentType.FULL, ComponentType.SWA), SWARadixCache), + "full": (ComponentType.FULL,), + "mamba": (ComponentType.FULL, ComponentType.MAMBA), + "swa": (ComponentType.FULL, ComponentType.SWA), + "all": (ComponentType.FULL, ComponentType.SWA, ComponentType.MAMBA), } @@ -841,7 +826,7 @@ def _run_bench_cli(): "--components", nargs="+", choices=list(_TREE_CONFIGS.keys()), - default=["mamba", "legacy-mamba"], + default=["mamba"], help="Component configs to benchmark", ) parser.add_argument("--page-size", type=int, default=1) @@ -857,7 +842,7 @@ def _run_bench_cli(): args, _ = parser.parse_known_args() for comp_name in args.components: - components, tree_cls = _TREE_CONFIGS[comp_name] + components = _TREE_CONFIGS[comp_name] run_all_benchmarks( num_seqs=args.num_seqs, chunk_len=args.chunk_len, @@ -865,7 +850,6 @@ def _run_bench_cli(): components=components, verify=args.verify, benchmarks=args.benchmarks, - tree_cls=tree_cls, page_size=args.page_size, ) diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 50ca52f1a..2ff2651ae 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -471,7 +471,7 @@ def build_fixture( page_size=cfg.page_size, enable_int8_mamba_checkpoint=cfg.enable_int8_mamba_checkpoint, ) - # MambaRadixCache reads mamba_cache_chunk_size, whose property otherwise + # The mamba component reads mamba_cache_chunk_size, whose property otherwise # loads the HF config for self.model_path — impossible for the dummy model. # Mirror the property's default for a dummy HF config: FLA_CHUNK_SIZE. server_args._mamba_cache_chunk_size = (