[Refactor] Replace page_align_keys helper with RadixKey.page_aligned method (#23107)

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

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:
+        key = key.page_aligned(self.page_size)
-            page_aligned_len = len(key) // self.page_size * self.page_size
-            key = key[:page_aligned_len]
 
         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(
-        radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle)
+            token_ids, req.extra_key, is_bigram=self.is_eagle
-        values = kv_indices[: len(radix_key)].to(dtype=torch.int64, copy=True)
+        ).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(
-        radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle)
+            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

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.
+        radix_key = RadixKey(
-        keys = page_align_keys(token_ids, self.page_size, is_bigram=self.is_eagle)
+            token_ids, req.extra_key, is_bigram=self.is_eagle
-        radix_key = RadixKey(keys, 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(
-        radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle)
+            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
-
+        key = key.page_aligned(self.page_size)
-        if self.page_size != 1:
+        if len(key) == 0:
-            page_aligned_len = len(key) // self.page_size * self.page_size
+            return None
-            key = key[:page_aligned_len]
-
         return key
 
     def _match_post_processor(

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:
+        key = key.page_aligned(self.page_size)
-            page_aligned_len = len(key) // self.page_size * self.page_size
-            key = key[:page_aligned_len]
 
         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.
+            radix_key = RadixKey(
-            keys = page_align_keys(token_ids, self.page_size, is_bigram=self.is_eagle)
+                token_ids, req.extra_key, is_bigram=self.is_eagle
-            radix_key = RadixKey(keys, 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.
+        radix_key = RadixKey(
-        keys = page_align_keys(
+            token_ids[:effective_cache_len],
-            token_ids[:effective_cache_len], self.page_size, is_bigram=self.is_eagle
+            req.extra_key,
-        )
+            is_bigram=self.is_eagle,
-        radix_key = RadixKey(keys, 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)
 

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),
+                key=key,
-                value=req1_kv_indices,
+                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),
+                key=key,
-                value=req2_kv_indices,
+                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),
+                key=key,
-                value=req3_kv_indices,
+                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),
+                key=key,
-                value=req4_kv_indices,
+                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]),
+                key=key1,
-                value=allocator.alloc(3),
+                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]),
+                key=key2,
-                value=allocator.alloc(7),
+                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]),
+                key=key3,
-                value=allocator.alloc(8),
+                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]),
+                key=key4,
-                value=allocator.alloc(9),
+                value=allocator.alloc(9)[: len(key4)],
                 mamba_value=req4.mamba_pool_idx.unsqueeze(0),
                 prev_prefix_len=8,
             )

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(
+        frequent_key = RadixKey([1, 2])  # High hit rate, should be retained
-            token_ids=[1, 2], extra_key=None
-        )  # 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(
+        first_low_freq_key = RadixKey([5, 6])  # Low hit rate, should be evicted
-            token_ids=[5, 6], extra_key=None
-        )  # 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(
+            key = RadixKey([i + 10])  # Unique keys for low-frequency items
-                token_ids=[i + 10], extra_key=None
-            )  # 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(
+            key = RadixKey([i * 2])  # Different pattern to avoid conflicts
-                token_ids=[i * 2], extra_key=None
-            )  # Different pattern to avoid conflicts
             val = torch.tensor([i * 200], dtype=torch.int64)
             self.cache.insert(InsertParams(key=key, value=val))
 

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),
+                        key=key,
-                        value=torch.tensor(tokens, dtype=torch.int64),
+                        value=torch.tensor(tokens, dtype=torch.int64)[: len(key)],
                     )
                 )
 

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(
+        key = RadixKey(req1_token_ids)
-            InsertParams(key=RadixKey(req1_token_ids), value=req1_kv_indices)
+        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(
+        key = RadixKey(req2_token_ids)
-            InsertParams(key=RadixKey(req2_token_ids), value=req2_kv_indices)
+        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(
+        key = RadixKey(req3_token_ids)
-            InsertParams(key=RadixKey(req3_token_ids), value=req3_kv_indices)
+        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(
+        key = RadixKey(req4_token_ids)
-            InsertParams(key=RadixKey(req4_token_ids), value=req4_kv_indices)
+        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(
+        key = RadixKey(req1_token_ids)
-            InsertParams(key=RadixKey(req1_token_ids), value=req1_kv_indices)
+        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(
+        key = RadixKey(req2_token_ids)
-            InsertParams(key=RadixKey(req2_token_ids), value=req2_kv_indices)
+        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(
+        key = RadixKey(req3_token_ids)
-            InsertParams(key=RadixKey(req3_token_ids), value=req3_kv_indices)
+        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(
+        key = RadixKey(req4_token_ids)
-            InsertParams(key=RadixKey(req4_token_ids), value=req4_kv_indices)
+        result = tree.insert(InsertParams(key=key, value=req4_kv_indices[: len(key)]))
-        )
         prefix_len = result.prefix_len
         self.assertEqual(prefix_len, 4)
         print(

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)))
+        k = RadixKey(q)
-        r2 = env.tree.match_prefix(MatchPrefixParams(key=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)

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)
-            key=RadixKey(tokens),
+        value = self._alloc(allocator, len(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),
+            key=key_2p,
-            value=self._alloc(allocator, len(seq_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),
+            key=key_3p,
-            value=self._alloc(allocator, len(seq_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)