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

[RadixTree][7/N Refactor]: Refactor mamba radix tree, release dup kvcache in insert func (#19429)

Browse Source
This commit is contained in:
hzh0425
2026-03-13 12:28:32 +08:00
committed by GitHub
parent f605612b87
commit 197f807134
2 changed files with 171 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -479,18 +479,18 @@ class MambaRadixCache(BasePrefixCache):
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.token_ids], dtype=torch.int64)
prefix_len, mamba_exist = self._insert_helper(
self.root_node, key, value, mamba_value, params.chunked
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) -> None:
"""Cache request when it finishes."""
kv_committed_len = req.pop_committed_kv_cache()
if self.disable:
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :kv_committed_len
@@ -555,13 +555,10 @@ class MambaRadixCache(BasePrefixCache):
key=RadixKey(token_ids[:page_aligned_len], req.extra_key),
value=page_aligned_kv_indices,
mamba_value=mamba_value,
prev_prefix_len=req.cache_protected_len,
)
)
new_prefix_len, mamba_exist = result.prefix_len, result.mamba_exist
self.token_to_kv_pool_allocator.free(
kv_indices[req.cache_protected_len : new_prefix_len]
)
mamba_exist = result.mamba_exist
else:
self.token_to_kv_pool_allocator.free(kv_indices[req.cache_protected_len :])
mamba_exist = True
@@ -651,12 +648,10 @@ class MambaRadixCache(BasePrefixCache):
key=RadixKey(page_aligned_token_ids, req.extra_key),
value=page_aligned_kv_indices,
mamba_value=mamba_value_forked,
prev_prefix_len=req.cache_protected_len,
)
)
new_prefix_len, mamba_exist = result.prefix_len, result.mamba_exist
self.token_to_kv_pool_allocator.free(
kv_indices[req.cache_protected_len : new_prefix_len]
)
# there is a mamba cache in radix cache, release it
if mamba_exist:
self.req_to_token_pool.mamba_pool.free(mamba_value_forked)
@@ -1090,6 +1085,7 @@ class MambaRadixCache(BasePrefixCache):
value,
mamba_value,
chunked: bool = False,
prev_prefix_len: int = 0,
) -> Tuple[int, bool]:
# Update the last access time from root to leaf, so that
# mamba will tombstone the node closer to root first
@@ -1112,6 +1108,11 @@ class MambaRadixCache(BasePrefixCache):
if node.mamba_value is not None:
self.mamba_lru_list.reset_node_mru(node)
prefix_len = self.key_match_fn(node.key, key)
if prev_prefix_len < total_prefix_length + prefix_len:
start = max(0, prev_prefix_len - total_prefix_length)
self.token_to_kv_pool_allocator.free(value[start:prefix_len])
total_prefix_length += prefix_len
key = key[prefix_len:]
value = value[prefix_len:]