From 197f80713435dee9320fe39b71995adb1c0a732d Mon Sep 17 00:00:00 2001 From: hzh0425 Date: Fri, 13 Mar 2026 12:28:32 +0800 Subject: [PATCH] [RadixTree][7/N Refactor]: Refactor mamba radix tree, release dup kvcache in insert func (#19429) --- .../sglang/srt/mem_cache/mamba_radix_cache.py | 21 +- .../radix_cache/test_mamba_unittest.py | 249 +++++++++++------- 2 files changed, 171 insertions(+), 99 deletions(-) diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 93ebd07a1..ac27d93c1 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -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:] diff --git a/test/registered/radix_cache/test_mamba_unittest.py b/test/registered/radix_cache/test_mamba_unittest.py index 0cebc9ff2..f067f141e 100644 --- a/test/registered/radix_cache/test_mamba_unittest.py +++ b/test/registered/radix_cache/test_mamba_unittest.py @@ -141,96 +141,9 @@ class TestMamba(unittest.TestCase): assert req_to_token_pool.mamba_pool.available_size() == mamba_cache_size - 1 def test_mamba_radix_cache_1(self): - set_global_server_args_for_scheduler( - ServerArgs(model_path="dummy", page_size=1) + tree, allocator, req_to_token_pool, make_dummy_req = ( + self._setup_tree_and_allocator() ) - # kv cache - size = 128 - dtype = torch.bfloat16 - head_num = 2 - head_dim = 256 - num_layers = 48 - global_interval = 4 - max_num_reqs = 10 - mamba_cache_size = 20 - max_context_len = 128 - device = get_device() - full_attention_layer_ids = [ - i for i in range(global_interval - 1, num_layers, global_interval) - ] - - # mamba - mamba_layers = [ - i for i in range(num_layers) if i not in full_attention_layer_ids - ] - with envs.SGLANG_MAMBA_SSM_DTYPE.override("bfloat16"): - shape = Mamba2StateShape.create( - tp_world_size=1, - intermediate_size=4096, - n_groups=16, - num_heads=32, - head_dim=128, - state_size=128, - conv_kernel=4, - ) - mamba2_cache_params = Mamba2CacheParams(shape=shape, layers=mamba_layers) - - req_to_token_pool = HybridReqToTokenPool( - size=max_num_reqs, - mamba_size=mamba_cache_size, - mamba_spec_state_size=max_num_reqs, - max_context_len=max_context_len, - device=device, - enable_memory_saver=False, - cache_params=mamba2_cache_params, - enable_mamba_extra_buffer=False, - speculative_num_draft_tokens=3, - ) - # setup kv pool - pool = HybridLinearKVPool( - size=size, - dtype=dtype, - page_size=1, - head_num=head_num, - head_dim=head_dim, - full_attention_layer_ids=full_attention_layer_ids, - enable_kvcache_transpose=False, - device=device, - enable_memory_saver=False, - mamba_pool=req_to_token_pool.mamba_pool, - ) - - # setup token to kv pool allocator - allocator = TokenToKVPoolAllocator( - size=size, - dtype=dtype, - device=device, - 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, - ) - # setup radix cache - tree = MambaRadixCache(params=params) - - def make_dummy_req(): - sampling_params = SamplingParams( - temperature=0, - max_new_tokens=1, - ) - req = Req( - rid=0, - origin_input_text="", - origin_input_ids=[], - sampling_params=sampling_params, - ) - req_to_token_pool.alloc([req]) - return req - mamba_pool = req_to_token_pool.mamba_pool # test print( @@ -386,6 +299,164 @@ class TestMamba(unittest.TestCase): print(available_and_evictable_str(tree)) tree.sanity_check() + def _setup_tree_and_allocator(self): + """Helper to create a MambaRadixCache with allocator for testing.""" + set_global_server_args_for_scheduler( + ServerArgs(model_path="dummy", page_size=1) + ) + size = 128 + dtype = torch.bfloat16 + head_num = 2 + head_dim = 256 + num_layers = 48 + global_interval = 4 + max_num_reqs = 10 + mamba_cache_size = 20 + max_context_len = 128 + device = get_device() + full_attention_layer_ids = [ + i for i in range(global_interval - 1, num_layers, global_interval) + ] + mamba_layers = [ + i for i in range(num_layers) if i not in full_attention_layer_ids + ] + with envs.SGLANG_MAMBA_SSM_DTYPE.override("bfloat16"): + shape = Mamba2StateShape.create( + tp_world_size=1, + intermediate_size=4096, + n_groups=16, + num_heads=32, + head_dim=128, + state_size=128, + conv_kernel=4, + ) + mamba2_cache_params = Mamba2CacheParams(shape=shape, layers=mamba_layers) + + req_to_token_pool = HybridReqToTokenPool( + size=max_num_reqs, + mamba_size=mamba_cache_size, + mamba_spec_state_size=max_num_reqs, + max_context_len=max_context_len, + device=device, + enable_memory_saver=False, + cache_params=mamba2_cache_params, + enable_mamba_extra_buffer=False, + speculative_num_draft_tokens=3, + ) + pool = HybridLinearKVPool( + size=size, + dtype=dtype, + page_size=1, + head_num=head_num, + head_dim=head_dim, + full_attention_layer_ids=full_attention_layer_ids, + enable_kvcache_transpose=False, + device=device, + enable_memory_saver=False, + mamba_pool=req_to_token_pool.mamba_pool, + ) + allocator = TokenToKVPoolAllocator( + size=size, + dtype=dtype, + device=device, + 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, + ) + tree = MambaRadixCache(params=params) + + def make_dummy_req(): + sampling_params = SamplingParams( + temperature=0, + max_new_tokens=1, + ) + req = Req( + rid=0, + origin_input_text="", + origin_input_ids=[], + sampling_params=sampling_params, + ) + req_to_token_pool.alloc([req]) + return req + + return tree, allocator, req_to_token_pool, make_dummy_req + + 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() + tree.insert( + InsertParams( + key=RadixKey([1, 2, 3]), + value=allocator.alloc(3), + mamba_value=req1.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() + result = tree.insert( + InsertParams( + key=RadixKey([1, 2, 3, 4, 5, 6, 7]), + value=allocator.alloc(7), + mamba_value=req2.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() + result = tree.insert( + InsertParams( + key=RadixKey([1, 2, 3, 4, 5, 6, 7, 8]), + value=allocator.alloc(8), + mamba_value=req3.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() + result = tree.insert( + InsertParams( + key=RadixKey([1, 2, 3, 4, 5, 6, 7, 8, 9]), + value=allocator.alloc(9), + mamba_value=req4.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()