diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index b89c3b15c..d4ee55a2f 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -433,11 +433,8 @@ class MambaRadixCache(BasePrefixCache): The last node create a new child if the prefix is shorter than the last node's value. """ - key = params.key - cow_mamba = params.cow_mamba - req = params.req - - if self.disable or len(key) == 0: + key = self._match_pre_processor(params) + if key is None: return MatchResult( device_indices=torch.empty( (0,), @@ -448,39 +445,8 @@ class MambaRadixCache(BasePrefixCache): last_host_node=self.root_node, ) - value, last_node, mamba_branching_seqlen = self._match_prefix_helper(key) - - # copy mamba state to req local space if cow is true - if cow_mamba and last_node.mamba_value is not None: - # for reqs without mamba cache - if req.mamba_pool_idx is None: - dst_index = self.req_to_token_pool.mamba_pool.alloc(1) - # try to alloc again, protect last_node from eviction - 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_pool.alloc(1) - self.dec_lock_ref(last_node) - assert dst_index is not None, "Can not alloc mamba cache" - src_index = last_node.mamba_value - self.req_to_token_pool.mamba_pool.copy_from(src_index, dst_index) - req.mamba_pool_idx = dst_index[0] - else: - src_index = last_node.mamba_value - dst_index = req.mamba_pool_idx.unsqueeze(0) - self.req_to_token_pool.mamba_pool.copy_from(src_index, dst_index) - - 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, - mamba_branching_seqlen=mamba_branching_seqlen, - ) + 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: @@ -940,7 +906,7 @@ class MambaRadixCache(BasePrefixCache): def _match_prefix_helper( self, key: RadixKey - ) -> Tuple[List[torch.Tensor], TreeNode, Optional[int]]: + ) -> 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, @@ -978,9 +944,31 @@ class MambaRadixCache(BasePrefixCache): 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 + # 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 - node_update = best_last_node + node_update = 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) @@ -1008,7 +996,38 @@ class MambaRadixCache(BasePrefixCache): else: mamba_branching_seqlen = None - return value[:best_value_len], best_last_node, mamba_branching_seqlen + # Copy mamba state to req local space if cow is true + if cow_mamba and last_node.mamba_value is not None: + # for reqs without mamba cache + if req.mamba_pool_idx is None: + dst_index = self.req_to_token_pool.mamba_pool.alloc(1) + # try to alloc again, protect last_node from eviction + 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_pool.alloc(1) + self.dec_lock_ref(last_node) + assert dst_index is not None, "Can not alloc mamba cache" + src_index = last_node.mamba_value + self.req_to_token_pool.mamba_pool.copy_from(src_index, dst_index) + req.mamba_pool_idx = dst_index[0] + else: + src_index = last_node.mamba_value + dst_index = req.mamba_pool_idx.unsqueeze(0) + self.req_to_token_pool.mamba_pool.copy_from(src_index, dst_index) + + 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, + mamba_branching_seqlen=mamba_branching_seqlen, + ) def _split_node(self, key: RadixKey, child: TreeNode, split_len: int) -> TreeNode: # new_node -> child diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 1feaca670..8deadf251 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -401,10 +401,9 @@ class SWARadixCache(BasePrefixCache): The last node create a new child if the prefix is shorter than the last node's value. """ - key = params.key - key.token_ids = self.key_convert_fn(key.token_ids) - if self.disable or len(key) == 0: + key = self._match_pre_processor(params) + if key is None: return MatchResult( device_indices=torch.empty( (0,), @@ -415,20 +414,8 @@ class SWARadixCache(BasePrefixCache): last_host_node=self.root_node, ) - if self.page_size != 1: - page_aligned_len = len(key) // self.page_size * self.page_size - key = key[:page_aligned_len] - - value, last_node = self._match_prefix_helper(key) - 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, - ) + 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: @@ -829,7 +816,7 @@ class SWARadixCache(BasePrefixCache): def _match_prefix_helper( self, key: RadixKey - ) -> Tuple[List[torch.Tensor], TreeNode]: + ) -> 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, @@ -878,9 +865,33 @@ class SWARadixCache(BasePrefixCache): 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.token_ids = self.key_convert_fn(key.token_ids) + + if self.disable or len(key) == 0: + return None + + if self.page_size != 1: + page_aligned_len = len(key) // self.page_size * self.page_size + key = key[:page_aligned_len] + + 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 - 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) @@ -893,7 +904,17 @@ class SWARadixCache(BasePrefixCache): ) node_update = node_update.parent - return value[:best_value_len], best_last_node + 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, + ) def _split_node(self, key: RadixKey, child: TreeNode, split_len: int) -> TreeNode: # new_node -> child