From c7a4ebf3c88dc4a5230aec02ab4f42f2a0c0d5f0 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Mon, 20 Apr 2026 18:10:42 -0700 Subject: [PATCH] [Refactor] Replace `page_align_keys` helper with `RadixKey.page_aligned` method (#23107) --- python/sglang/srt/mem_cache/hiradix_cache.py | 12 ++-- python/sglang/srt/mem_cache/radix_cache.py | 56 ++++++++----------- .../sglang/srt/mem_cache/swa_radix_cache.py | 29 +++++----- .../srt/mem_cache/unified_radix_cache.py | 28 +++++----- .../unit/mem_cache/test_mamba_unittest.py | 40 +++++++------ .../test_radix_cache_slru_accuracy.py | 16 ++---- .../unit/mem_cache/test_radix_cache_unit.py | 5 +- .../unit/mem_cache/test_swa_unittest.py | 40 ++++++------- .../test_unified_radix_cache_bench.py | 13 +++-- .../test_unified_radix_cache_unittest.py | 22 +++++--- 10 files changed, 126 insertions(+), 135 deletions(-) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 878e530ac..d342bbef6 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -1216,10 +1216,8 @@ class HiRadixCache(RadixCache): host_hit_length=0, ) + key = key.page_aligned(self.page_size) page_aligned_len = len(key) - 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(self.root_node, key) if value: @@ -1394,15 +1392,15 @@ class HiRadixCache(RadixCache): if priority is None: priority = 0 + key, value = key.maybe_to_bigram_view(self.is_eagle, value) + key = key.page_aligned(self.page_size) + if value is not None: + value = value[: len(key)] if len(key) == 0: return InsertResult(prefix_len=0) - if self.is_eagle and value is not None: - # Make sure the value len equal to the EAGLE bigram key len - value = value[: len(key)] - node = self.root_node child_key = self.get_child_key_fn(key) total_prefix_length = 0 diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index ee098ac95..f4d54bdc3 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -122,6 +122,12 @@ class RadixKey: preview = self.token_ids[:10] return f"RadixKey(extra_key={self.extra_key!r}, token_ids={preview}{'...' if len(self.token_ids) > 10 else ''}, is_bigram={self.is_bigram})" + def page_aligned(self, page_size: int) -> "RadixKey": + if page_size == 1: + return self + aligned_len = len(self) // page_size * page_size + return self[:aligned_len] + def maybe_to_bigram_view( self, is_eagle: bool, @@ -136,24 +142,6 @@ class RadixKey: return self, value -def page_align_keys(key: list, page_size: int, is_bigram: bool = False) -> list: - """Truncate a raw token list so the resulting RadixKey length is page-aligned. - - In bigram mode, logical length = len(key) - 1, and we must keep one extra - boundary token so that bigram_count == aligned. - """ - if page_size == 1: - return key - if is_bigram: - logical_len = len(key) - 1 if len(key) > 0 else 0 - aligned = logical_len // page_size * page_size - if aligned == 0: - return [] - return key[: aligned + 1] - page_aligned_len = len(key) // page_size * page_size - return key[:page_aligned_len] - - class TreeNode: counter = 0 @@ -504,9 +492,7 @@ class RadixCache(BasePrefixCache): if self.disable or len(key) == 0: return empty_match_result() - if self.page_size != 1: - page_aligned_len = len(key) // self.page_size * self.page_size - key = key[:page_aligned_len] + key = key.page_aligned(self.page_size) if len(key) == 0: return empty_match_result() @@ -531,12 +517,13 @@ class RadixCache(BasePrefixCache): priority = params.priority chunked = params.chunked - if value is None: - # Debug/test fallback: use token ids themselves as values. Truncate - # to the logical key length so bigram mode gets len(key) entries. - value = torch.tensor(key.token_ids[: len(key)], dtype=torch.int64) - key, value = key.maybe_to_bigram_view(self.is_eagle, value) + key = key.page_aligned(self.page_size) + if value is not None: + value = value[: len(key)] + else: + # Debug/test fallback: use token ids themselves as values. + value = torch.tensor(key.token_ids[: len(key)], dtype=torch.int64) prefix_len = self._insert_helper(self.root_node, key, value, priority, chunked) return InsertResult(prefix_len=prefix_len) @@ -560,9 +547,11 @@ class RadixCache(BasePrefixCache): req.req_pool_idx, : len(token_ids) ] - keys = page_align_keys(token_ids, self.page_size, is_bigram=self.is_eagle) - radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle) - values = kv_indices[: len(radix_key)].to(dtype=torch.int64, copy=True) + radix_key = RadixKey( + token_ids, req.extra_key, is_bigram=self.is_eagle + ).page_aligned(self.page_size) + key_len = len(radix_key) + values = kv_indices[:key_len].to(dtype=torch.int64, copy=True) # Radix Cache takes one ref in memory pool if is_insert: @@ -577,11 +566,11 @@ class RadixCache(BasePrefixCache): ) else: self.token_to_kv_pool_allocator.free( - kv_indices[req.cache_protected_len : len(radix_key)] + kv_indices[req.cache_protected_len : key_len] ) # free the unaligned tail - self.token_to_kv_pool_allocator.free(kv_indices[len(radix_key) :]) + self.token_to_kv_pool_allocator.free(kv_indices[key_len:]) # Remove req slot release the cache lock self.dec_lock_ref(req.last_node) @@ -596,8 +585,9 @@ class RadixCache(BasePrefixCache): req.req_pool_idx, : len(token_ids) ] - keys = page_align_keys(token_ids, self.page_size, is_bigram=self.is_eagle) - radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle) + radix_key = RadixKey( + token_ids, req.extra_key, is_bigram=self.is_eagle + ).page_aligned(self.page_size) values = kv_indices[: len(radix_key)].to(dtype=torch.int64, copy=True) # Radix Cache takes one ref in memory pool diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 47682477d..b5450015e 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -46,7 +46,6 @@ from sglang.srt.mem_cache.radix_cache import ( _key_match_page_size1, _key_match_paged, get_child_key, - page_align_keys, ) from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.utils import convert_to_bigram_key @@ -430,10 +429,12 @@ class SWARadixCache(BasePrefixCache): prev_prefix_len = params.prev_prefix_len swa_evicted_seqlen = params.swa_evicted_seqlen - if value is None: - value = torch.tensor(key.token_ids[: len(key)], dtype=torch.int64) - key, value = key.maybe_to_bigram_view(self.is_eagle, value) + key = key.page_aligned(self.page_size) + if value is not None: + value = value[: len(key)] + else: + value = torch.tensor(key.token_ids[: len(key)], dtype=torch.int64) prefix_len = self._insert_helper( self.root_node, key, value, prev_prefix_len, swa_evicted_seqlen @@ -455,9 +456,9 @@ class SWARadixCache(BasePrefixCache): req.req_pool_idx, :kv_committed_len ] - # EAGLE: skip tuple materialization; is_bigram flag gives bigram semantics. - keys = page_align_keys(token_ids, self.page_size, is_bigram=self.is_eagle) - radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle) + radix_key = RadixKey( + token_ids, req.extra_key, is_bigram=self.is_eagle + ).page_aligned(self.page_size) page_aligned_len = len(radix_key) values = kv_indices[:page_aligned_len].to(dtype=torch.int64, copy=True) old_prefix_len = req.cache_protected_len @@ -502,8 +503,9 @@ class SWARadixCache(BasePrefixCache): req.req_pool_idx, : len(token_ids) ] - keys = page_align_keys(token_ids, self.page_size, is_bigram=self.is_eagle) - radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle) + radix_key = RadixKey( + token_ids, req.extra_key, is_bigram=self.is_eagle + ).page_aligned(self.page_size) values = kv_indices[: len(radix_key)].to(dtype=torch.int64, copy=True) old_prefix_len = req.cache_protected_len @@ -840,14 +842,11 @@ class SWARadixCache(BasePrefixCache): """Preprocess the key before matching.""" key = params.key key, _ = key.maybe_to_bigram_view(self.is_eagle) - 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] - + key = key.page_aligned(self.page_size) + if len(key) == 0: + return None return key def _match_post_processor( diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 3742905c8..0a1a83527 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -25,7 +25,6 @@ from sglang.srt.mem_cache.radix_cache import ( _key_match_page_size1, _key_match_paged, get_child_key, - page_align_keys, ) from sglang.srt.mem_cache.unified_cache_components import ( _NUM_COMPONENT_TYPES, @@ -249,9 +248,7 @@ class UnifiedRadixCache(BasePrefixCache): last_device_node=self.root_node, 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] + key = key.page_aligned(self.page_size) value, last_node, best_value_len = self._match_prefix_helper(key) return self._match_post_processor(params, value, last_node, best_value_len) @@ -262,10 +259,13 @@ class UnifiedRadixCache(BasePrefixCache): key = params.key value = params.value - if value is None: + key, value = key.maybe_to_bigram_view(self.is_eagle, value) + key = key.page_aligned(self.page_size) + if value is not None: + value = value[: len(key)] + else: value = torch.tensor(key.token_ids[: len(key)], dtype=torch.int64) - key, value = key.maybe_to_bigram_view(self.is_eagle, value) result = self._insert_helper(self.root_node, key, value, params) return result @@ -354,9 +354,9 @@ class UnifiedRadixCache(BasePrefixCache): token_ids = token_ids[:effective_cache_len] kv_indices = kv_indices[:effective_cache_len] - # Page align on raw tokens; bigram semantics via is_bigram flag. - keys = page_align_keys(token_ids, self.page_size, is_bigram=self.is_eagle) - radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle) + radix_key = RadixKey( + token_ids, req.extra_key, is_bigram=self.is_eagle + ).page_aligned(self.page_size) page_aligned_len = len(radix_key) values = kv_indices[:page_aligned_len].to(dtype=torch.int64, copy=True) @@ -420,11 +420,11 @@ class UnifiedRadixCache(BasePrefixCache): kv_indices = kv_indices_orig[:effective_cache_len] - # Page align on raw tokens; bigram semantics via is_bigram flag. - keys = page_align_keys( - token_ids[:effective_cache_len], self.page_size, is_bigram=self.is_eagle - ) - radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle) + radix_key = RadixKey( + token_ids[:effective_cache_len], + req.extra_key, + is_bigram=self.is_eagle, + ).page_aligned(self.page_size) page_aligned_len = len(radix_key) values = kv_indices[:page_aligned_len].to(dtype=torch.int64, copy=True) diff --git a/test/registered/unit/mem_cache/test_mamba_unittest.py b/test/registered/unit/mem_cache/test_mamba_unittest.py index adadc5f3c..b339aee02 100644 --- a/test/registered/unit/mem_cache/test_mamba_unittest.py +++ b/test/registered/unit/mem_cache/test_mamba_unittest.py @@ -156,10 +156,11 @@ class TestMamba(unittest.TestCase): print( f"req1: inserting, req1_token_ids: {req1_token_ids}, req1_kv_indices: {req1_kv_indices}" ) + key = RadixKey(req1_token_ids) result = tree.insert( InsertParams( - key=RadixKey(req1_token_ids), - value=req1_kv_indices, + key=key, + value=req1_kv_indices[: len(key)], mamba_value=req1.mamba_pool_idx.unsqueeze(0), ) ) @@ -173,10 +174,11 @@ class TestMamba(unittest.TestCase): print( f"req2: inserting, req2_token_ids: {req2_token_ids}, req2_kv_indices: {req2_kv_indices}" ) + key = RadixKey(req2_token_ids) result = tree.insert( InsertParams( - key=RadixKey(req2_token_ids), - value=req2_kv_indices, + key=key, + value=req2_kv_indices[: len(key)], mamba_value=req2.mamba_pool_idx.unsqueeze(0), ) ) @@ -191,10 +193,11 @@ class TestMamba(unittest.TestCase): print( f"req3: inserting, req3_token_ids: {req3_token_ids}, req3_kv_indices: {req3_kv_indices}" ) + key = RadixKey(req3_token_ids) result = tree.insert( InsertParams( - key=RadixKey(req3_token_ids), - value=req3_kv_indices, + key=key, + value=req3_kv_indices[: len(key)], mamba_value=req3.mamba_pool_idx.unsqueeze(0), ) ) @@ -208,10 +211,11 @@ class TestMamba(unittest.TestCase): print( f"req4: inserting, req4_token_ids: {req4_token_ids}, req4_kv_indices: {req4_kv_indices}" ) + key = RadixKey(req4_token_ids) result = tree.insert( InsertParams( - key=RadixKey(req4_token_ids), - value=req4_kv_indices, + key=key, + value=req4_kv_indices[: len(key)], mamba_value=req4.mamba_pool_idx.unsqueeze(0), ) ) @@ -400,10 +404,11 @@ class TestMamba(unittest.TestCase): # Step 1: Insert [1,2,3] to create first node req1 = make_dummy_req() + key1 = RadixKey([1, 2, 3]) tree.insert( InsertParams( - key=RadixKey([1, 2, 3]), - value=allocator.alloc(3), + key=key1, + value=allocator.alloc(3)[: len(key1)], mamba_value=req1.mamba_pool_idx.unsqueeze(0), ) ) @@ -412,10 +417,11 @@ class TestMamba(unittest.TestCase): # 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() + key2 = RadixKey([1, 2, 3, 4, 5, 6, 7]) result = tree.insert( InsertParams( - key=RadixKey([1, 2, 3, 4, 5, 6, 7]), - value=allocator.alloc(7), + key=key2, + value=allocator.alloc(7)[: len(key2)], mamba_value=req2.mamba_pool_idx.unsqueeze(0), prev_prefix_len=0, ) @@ -429,10 +435,11 @@ class TestMamba(unittest.TestCase): # 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() + key3 = RadixKey([1, 2, 3, 4, 5, 6, 7, 8]) result = tree.insert( InsertParams( - key=RadixKey([1, 2, 3, 4, 5, 6, 7, 8]), - value=allocator.alloc(8), + key=key3, + value=allocator.alloc(8)[: len(key3)], mamba_value=req3.mamba_pool_idx.unsqueeze(0), prev_prefix_len=2, ) @@ -445,10 +452,11 @@ class TestMamba(unittest.TestCase): # 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() + key4 = RadixKey([1, 2, 3, 4, 5, 6, 7, 8, 9]) result = tree.insert( InsertParams( - key=RadixKey([1, 2, 3, 4, 5, 6, 7, 8, 9]), - value=allocator.alloc(9), + key=key4, + value=allocator.alloc(9)[: len(key4)], mamba_value=req4.mamba_pool_idx.unsqueeze(0), prev_prefix_len=8, ) diff --git a/test/registered/unit/mem_cache/test_radix_cache_slru_accuracy.py b/test/registered/unit/mem_cache/test_radix_cache_slru_accuracy.py index e407c1741..bc28b40b7 100644 --- a/test/registered/unit/mem_cache/test_radix_cache_slru_accuracy.py +++ b/test/registered/unit/mem_cache/test_radix_cache_slru_accuracy.py @@ -62,9 +62,7 @@ class TestSLRUAccuracy(unittest.TestCase): """Test that SLRU eviction mechanism works correctly""" # Insert one key-value three times (high frequency access) - frequent_key = RadixKey( - token_ids=[1, 2], extra_key=None - ) # High hit rate, should be retained + frequent_key = RadixKey([1, 2]) # High hit rate, should be retained frequent_val = torch.tensor([10, 20], dtype=torch.int64) # Insert the frequent key multiple times to increase its hit count @@ -72,9 +70,7 @@ class TestSLRUAccuracy(unittest.TestCase): self.cache.insert(InsertParams(key=frequent_key, value=frequent_val)) # Insert first low-frequency key-value pair that should be evicted - first_low_freq_key = RadixKey( - token_ids=[5, 6], extra_key=None - ) # Low hit rate, should be evicted + first_low_freq_key = RadixKey([5, 6]) # Low hit rate, should be evicted first_low_freq_val = torch.tensor([50, 60], dtype=torch.int64) self.cache.insert( @@ -84,18 +80,14 @@ class TestSLRUAccuracy(unittest.TestCase): # Insert other key-values once each (low frequency access) - fill up the cache other_keys = [] for i in range(4): # Reduce the number to fit in our smaller cache - key = RadixKey( - token_ids=[i + 10], extra_key=None - ) # Unique keys for low-frequency items + key = RadixKey([i + 10]) # Unique keys for low-frequency items val = torch.tensor([i + 100], dtype=torch.int64) self.cache.insert(InsertParams(key=key, value=val)) other_keys.append(key) # Now insert more items to trigger evictions for i in range(6, 10): # Add more items to definitely exceed capacity - key = RadixKey( - token_ids=[i * 2], extra_key=None - ) # Different pattern to avoid conflicts + key = RadixKey([i * 2]) # Different pattern to avoid conflicts val = torch.tensor([i * 200], dtype=torch.int64) self.cache.insert(InsertParams(key=key, value=val)) diff --git a/test/registered/unit/mem_cache/test_radix_cache_unit.py b/test/registered/unit/mem_cache/test_radix_cache_unit.py index f01bc7736..c45e3cef8 100644 --- a/test/registered/unit/mem_cache/test_radix_cache_unit.py +++ b/test/registered/unit/mem_cache/test_radix_cache_unit.py @@ -547,10 +547,11 @@ class TestRadixCache(unittest.TestCase): cache = RadixCache.create_simulated(page_size=page_size) tokens = list(range(sequence_length)) + key = RadixKey(tokens) cache.insert( InsertParams( - key=RadixKey(tokens), - value=torch.tensor(tokens, dtype=torch.int64), + key=key, + value=torch.tensor(tokens, dtype=torch.int64)[: len(key)], ) ) diff --git a/test/registered/unit/mem_cache/test_swa_unittest.py b/test/registered/unit/mem_cache/test_swa_unittest.py index 0e2335d0f..e5b922166 100644 --- a/test/registered/unit/mem_cache/test_swa_unittest.py +++ b/test/registered/unit/mem_cache/test_swa_unittest.py @@ -215,9 +215,8 @@ class TestSWA(unittest.TestCase): print( f"req1: inserting, req1_token_ids: {req1_token_ids}, req1_kv_indices: {req1_kv_indices}" ) - result = tree.insert( - InsertParams(key=RadixKey(req1_token_ids), value=req1_kv_indices) - ) + key = RadixKey(req1_token_ids) + result = tree.insert(InsertParams(key=key, value=req1_kv_indices[: len(key)])) prefix_len = result.prefix_len print( f"req1: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}" @@ -227,9 +226,8 @@ class TestSWA(unittest.TestCase): print( f"req2: inserting, req2_token_ids: {req2_token_ids}, req2_kv_indices: {req2_kv_indices}" ) - result = tree.insert( - InsertParams(key=RadixKey(req2_token_ids), value=req2_kv_indices) - ) + key = RadixKey(req2_token_ids) + result = tree.insert(InsertParams(key=key, value=req2_kv_indices[: len(key)])) prefix_len = result.prefix_len print( f"req2: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}" @@ -239,9 +237,8 @@ class TestSWA(unittest.TestCase): print( f"req3: inserting, req3_token_ids: {req3_token_ids}, req3_kv_indices: {req3_kv_indices}" ) - result = tree.insert( - InsertParams(key=RadixKey(req3_token_ids), value=req3_kv_indices) - ) + key = RadixKey(req3_token_ids) + result = tree.insert(InsertParams(key=key, value=req3_kv_indices[: len(key)])) prefix_len = result.prefix_len print( f"req3: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}" @@ -251,9 +248,8 @@ class TestSWA(unittest.TestCase): print( f"req4: inserting, req4_token_ids: {req4_token_ids}, req4_kv_indices: {req4_kv_indices}" ) - result = tree.insert( - InsertParams(key=RadixKey(req4_token_ids), value=req4_kv_indices) - ) + key = RadixKey(req4_token_ids) + result = tree.insert(InsertParams(key=key, value=req4_kv_indices[: len(key)])) prefix_len = result.prefix_len print( f"req4: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}" @@ -374,9 +370,8 @@ class TestSWA(unittest.TestCase): print( f"req1: inserting, req1_token_ids: {req1_token_ids}, req1_kv_indices: {req1_kv_indices}" ) - result = tree.insert( - InsertParams(key=RadixKey(req1_token_ids), value=req1_kv_indices) - ) + key = RadixKey(req1_token_ids) + result = tree.insert(InsertParams(key=key, value=req1_kv_indices[: len(key)])) prefix_len = result.prefix_len self.assertEqual(prefix_len, 0) print( @@ -387,9 +382,8 @@ class TestSWA(unittest.TestCase): print( f"req2: inserting, req2_token_ids: {req2_token_ids}, req2_kv_indices: {req2_kv_indices}" ) - result = tree.insert( - InsertParams(key=RadixKey(req2_token_ids), value=req2_kv_indices) - ) + key = RadixKey(req2_token_ids) + result = tree.insert(InsertParams(key=key, value=req2_kv_indices[: len(key)])) prefix_len = result.prefix_len self.assertEqual(prefix_len, 2) print( @@ -400,9 +394,8 @@ class TestSWA(unittest.TestCase): print( f"req3: inserting, req3_token_ids: {req3_token_ids}, req3_kv_indices: {req3_kv_indices}" ) - result = tree.insert( - InsertParams(key=RadixKey(req3_token_ids), value=req3_kv_indices) - ) + key = RadixKey(req3_token_ids) + result = tree.insert(InsertParams(key=key, value=req3_kv_indices[: len(key)])) prefix_len = result.prefix_len self.assertEqual(prefix_len, 0) print( @@ -413,9 +406,8 @@ class TestSWA(unittest.TestCase): print( f"req4: inserting, req4_token_ids: {req4_token_ids}, req4_kv_indices: {req4_kv_indices}" ) - result = tree.insert( - InsertParams(key=RadixKey(req4_token_ids), value=req4_kv_indices) - ) + key = RadixKey(req4_token_ids) + result = tree.insert(InsertParams(key=key, value=req4_kv_indices[: len(key)])) prefix_len = result.prefix_len self.assertEqual(prefix_len, 4) print( diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py index 1a7389db3..5dce0295a 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py @@ -335,7 +335,8 @@ def _insert_seq(env, seq): if env.has_mamba: req = env.make_req() mamba_val = req.mamba_pool_idx.unsqueeze(0) - env.tree.insert(InsertParams(key=RadixKey(seq), value=v, mamba_value=mamba_val)) + key = RadixKey(seq) + env.tree.insert(InsertParams(key=key, value=v[: len(key)], mamba_value=mamba_val)) return True @@ -356,7 +357,10 @@ def _fill_no_evict(env): if env.has_mamba: req = env.make_req() mamba_val = req.mamba_pool_idx.unsqueeze(0) - env.tree.insert(InsertParams(key=RadixKey(seq), value=v, mamba_value=mamba_val)) + key = RadixKey(seq) + env.tree.insert( + InsertParams(key=key, value=v[: len(key)], mamba_value=mamba_val) + ) inserted += 1 return inserted @@ -501,8 +505,9 @@ def bench_match_prefix( queries.append([rng.randint(1, 32000)] * rng.randint(50, 300)) def verify_fn(q): - r1 = env.tree.match_prefix(MatchPrefixParams(key=RadixKey(q))) - r2 = env.tree.match_prefix(MatchPrefixParams(key=RadixKey(q))) + k = RadixKey(q) + r1 = env.tree.match_prefix(MatchPrefixParams(key=k)) + r2 = env.tree.match_prefix(MatchPrefixParams(key=k)) assert len(r1.device_indices) == len(r2.device_indices), "match not idempotent" warmup = min(20, len(queries) // 10) diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 5b1fa4aa9..7014290bd 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -258,10 +258,9 @@ class UnifiedRadixCacheSuite: def _insert(self, tree, allocator, req_to_token_pool, tokens): """Insert tokens, attaching mamba data when the config has mamba.""" - params = InsertParams( - key=RadixKey(tokens), - value=self._alloc(allocator, len(tokens)), - ) + key = RadixKey(tokens) + value = self._alloc(allocator, len(tokens)) + params = InsertParams(key=key, value=value[: len(key)]) if self.cfg.has_mamba: req = self._make_req(req_to_token_pool) params.mamba_value = req.mamba_pool_idx.unsqueeze(0) @@ -400,9 +399,11 @@ class UnifiedRadixCacheSuite: self.assertEqual(allocator.available_size(), initial_avail - len(seq_1p)) # Step 2: insert 2 pages with prev_prefix_len=0 → frees overlap of 1 page + key_2p = RadixKey(seq_2p) + value_2p = self._alloc(allocator, len(seq_2p)) params = InsertParams( - key=RadixKey(seq_2p), - value=self._alloc(allocator, len(seq_2p)), + key=key_2p, + value=value_2p[: len(key_2p)], prev_prefix_len=0, ) if self.cfg.has_mamba: @@ -417,9 +418,11 @@ class UnifiedRadixCacheSuite: # Step 3: insert 3 pages with prev_prefix_len=len(seq_2p) → nothing freed avail_before = allocator.available_size() + key_3p = RadixKey(seq_3p) + value_3p = self._alloc(allocator, len(seq_3p)) params = InsertParams( - key=RadixKey(seq_3p), - value=self._alloc(allocator, len(seq_3p)), + key=key_3p, + value=value_3p[: len(key_3p)], prev_prefix_len=len(seq_2p), ) if self.cfg.has_mamba: @@ -582,6 +585,7 @@ class UnifiedRadixCacheSuite: self.assertIsInstance(child_key, tuple) def test_paged_match_truncates_unaligned_key(self): + """match_prefix internally aligns keys to page boundary.""" if self.cfg.page_size == 1: self.skipTest("page_size > 1 only") ps = self.cfg.page_size @@ -589,10 +593,12 @@ class UnifiedRadixCacheSuite: seq = self._make_seq(1, 2) self._insert(tree, allocator, req_to_token_pool, seq) + # Tree truncates unaligned tail internally, so it matches the seq prefix. unaligned = seq + list(range(9000, 9000 + ps - 1)) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(unaligned))) self.assertEqual(len(m.device_indices), len(seq)) + # Below-page-size key aligns to 0 -> no match. m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq[: ps - 1]))) self.assertEqual(len(m.device_indices), 0)