From 48d6bea1ea4ad420054411a978c497844ee78d0a Mon Sep 17 00:00:00 2001 From: Hanming Lu <69857889+hanming-lu@users.noreply.github.com> Date: Mon, 3 Nov 2025 19:03:37 -0800 Subject: [PATCH] [GDN/SWA] mamba and swa radix cache edge case fix (#12111) Co-authored-by: yizhang2077 <1109276519@qq.com> --- python/sglang/srt/managers/scheduler.py | 6 ++ .../sglang/srt/mem_cache/mamba_radix_cache.py | 73 ++++++++++++------- .../sglang/srt/mem_cache/swa_radix_cache.py | 34 ++++++--- 3 files changed, 73 insertions(+), 40 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index bfa1c0fc3..cc9fab58b 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2386,6 +2386,12 @@ class Scheduler( - self.tree_cache.swa_evictable_size() ) num_tokens = max(num_tokens_full, num_tokens_swa) + elif self.is_hybrid_gdn: + num_tokens = ( + self.max_total_num_tokens + - self.token_to_kv_pool_allocator.available_size() + - self.tree_cache.full_evictable_size() + ) else: num_tokens = ( self.max_total_num_tokens diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 739b204ed..7d930f112 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -20,11 +20,11 @@ The radix tree data structure for managing the hybrid (full and Mamba) 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.mem_cache.allocator import TokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, MatchResult @@ -46,6 +46,7 @@ logger = logging.getLogger(__name__) class TreeNode: counter = 0 + last_access_time_counter_float = float64(1.0) def __init__(self, id: Optional[int] = None): self.children = defaultdict(TreeNode) @@ -61,7 +62,7 @@ class TreeNode: 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. - self.last_access_time = time.monotonic() + self.last_access_time = get_last_access_time() self.hit_count = 0 # store the host indices of KV cache @@ -90,6 +91,12 @@ class 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 @@ -382,8 +389,6 @@ class MambaRadixCache(BasePrefixCache): # copy mamba state to req local space if cow is true if cow_mamba and last_node.mamba_value is not None: - assert req.req_pool_idx is None # req_pool_idx is uninitialed - # for reqs without mamba cache if req.mamba_pool_idx is None: dst_index = self.req_to_token_pool.mamba_pool.alloc(1) @@ -421,7 +426,7 @@ class MambaRadixCache(BasePrefixCache): value = torch.tensor([x for x in key.token_ids], dtype=torch.int64) return self._insert_helper(self.root_node, key, value, mamba_value) - def cache_finished_req(self, req: Req) -> None: + def cache_finished_req(self, req: Req, is_insert=True) -> None: """Cache request when it finishes.""" if self.disable: kv_indices = self.req_to_token_pool.req_to_token[ @@ -449,15 +454,20 @@ class MambaRadixCache(BasePrefixCache): .clone() ) - new_prefix_len, mamba_exist = self.insert( - RadixKey(token_ids[:page_aligned_len], req.extra_key), - page_aligned_kv_indices, - mamba_value, - ) - - self.token_to_kv_pool_allocator.free( - kv_indices[len(req.prefix_indices) : new_prefix_len] - ) + if is_insert: + new_prefix_len, mamba_exist = self.insert( + RadixKey(token_ids[:page_aligned_len], req.extra_key), + page_aligned_kv_indices, + mamba_value, + ) + self.token_to_kv_pool_allocator.free( + kv_indices[len(req.prefix_indices) : new_prefix_len] + ) + else: + self.token_to_kv_pool_allocator.free( + kv_indices[len(req.prefix_indices) : page_aligned_len] + ) + mamba_exist = True self.req_to_token_pool.free(req.req_pool_idx, free_mamba_cache=mamba_exist) self.dec_lock_ref(req.last_node) @@ -767,15 +777,18 @@ class MambaRadixCache(BasePrefixCache): # update time for matched nodes, and make nodes closer to root to be least recently used # this allows mamba to evict nodes closer to root first - self.full_lru_list.reset_node_and_parents_mru(best_last_node, self.root_node) - self.mamba_lru_list.reset_node_and_parents_mru(best_last_node, self.root_node) + node_update = best_last_node + self.full_lru_list.reset_node_and_parents_mru(node_update, self.root_node) + self.mamba_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 = time.monotonic() - while node: - node.last_access_time = cur_time - cur_time -= 0.0001 - node = node.parent + 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 return value[:best_value_len], best_last_node @@ -791,7 +804,7 @@ class MambaRadixCache(BasePrefixCache): new_node.value = child.value[:split_len] # child time should be later than parent's time for mamba tombstone - child.last_access_time = time.monotonic() + child.last_access_time = get_last_access_time() self.full_lru_list.remove_node(child) if child.mamba_value is not None: @@ -819,7 +832,7 @@ class MambaRadixCache(BasePrefixCache): # Update the last access time from root to leaf, so that # mamba will tombstone the node closer to root first assert mamba_value is not None, "Mamba value should not be None here." - node.last_access_time = time.monotonic() + node.last_access_time = get_last_access_time() if node != self.root_node: self.full_lru_list.reset_node_mru(node) if node.mamba_value is not None: @@ -832,7 +845,7 @@ class MambaRadixCache(BasePrefixCache): total_prefix_length = 0 while len(key) > 0 and child_key in node.children.keys(): node = node.children[child_key] - node.last_access_time = time.monotonic() + node.last_access_time = get_last_access_time() self.full_lru_list.reset_node_mru(node) if node.mamba_value is not None: self.mamba_lru_list.reset_node_mru(node) @@ -856,17 +869,21 @@ class MambaRadixCache(BasePrefixCache): new_node.value = value new_node.mamba_value = mamba_value self.full_lru_list.insert_mru(new_node) - self.full_evictable_size_ += len(value) - self.mamba_evictable_size_ += len(mamba_value) 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) elif node.mamba_value is None: # add for mamba tombstone node.mamba_value = mamba_value - self.mamba_evictable_size_ += len(mamba_value) + self.full_lru_list.reset_node_mru(node) self.mamba_lru_list.insert_mru(node) - else: + 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) self.mamba_lru_list.reset_node_mru(node) + node.last_access_time = get_last_access_time() return total_prefix_length, mamba_value_exist diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 566c94784..43e816125 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -20,12 +20,12 @@ The radix tree data structure for managing the hybrid (full and SWA) KV cache. """ import heapq -import time from collections import defaultdict from functools import partial from typing import TYPE_CHECKING, List, Optional, Tuple import torch +from numpy import float64 from sglang.srt.mem_cache.allocator import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, MatchResult @@ -50,6 +50,7 @@ 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) @@ -64,7 +65,7 @@ class TreeNode: 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 = time.monotonic() + self.last_access_time = get_last_access_time() self.hit_count = 0 # store the host indices of KV cache @@ -99,6 +100,12 @@ def gen_swa_uuid() -> int: 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, swa: bool = False): self.swa = swa @@ -841,15 +848,18 @@ class SWARadixCache(BasePrefixCache): # 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(best_last_node, self.root_node) - self.swa_lru_list.reset_node_and_parents_mru(best_last_node, self.root_node) + node_update = best_last_node + 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 = time.monotonic() - while node: - node.last_access_time = cur_time - cur_time -= 0.0001 - node = node.parent + 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 return value[:best_value_len], best_last_node @@ -867,7 +877,7 @@ class SWARadixCache(BasePrefixCache): 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 = time.monotonic() + 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) @@ -892,7 +902,7 @@ class SWARadixCache(BasePrefixCache): ) -> 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 = time.monotonic() + 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: @@ -905,7 +915,7 @@ class SWARadixCache(BasePrefixCache): total_prefix_length = 0 while len(key) > 0 and child_key in node.children.keys(): node = node.children[child_key] - node.last_access_time = time.monotonic() + 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)