chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[RadixTree][5/N Refactor]: Introduce pre and post-processing methods for key matching (#18147)

Browse Source
This commit is contained in:
zhangheng
2026-02-04 17:10:46 +08:00
committed by GitHub
parent d279520ba5
commit be557cbc5f
2 changed files with 101 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files

101
python/sglang/srt/mem_cache/mamba_radix_cache.py
View File

@@ -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

61
python/sglang/srt/mem_cache/swa_radix_cache.py
View File

@@ -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