[RadixTree][6/N Refactor]: Refactor SWARadixTree to simplify the computation and alignment of bigram keys. (#19427)

python/sglang/srt/mem_cache/radix_cache.py

@@ -94,6 +94,26 @@ class RadixKey:
         return f"RadixKey(extra_key={self.extra_key!r}, token_ids={preview}{'...' if len(self.token_ids) > 10 else ''})"
 
 
+def maybe_bigram_convert(
+    is_eagle: bool,
+    key: RadixKey,
+    value: Optional[torch.Tensor] = None,
+) -> Tuple[RadixKey, Optional[torch.Tensor]]:
+    if is_eagle and not key.is_bigram:
+        key.token_ids = convert_to_bigram_key(key.token_ids)
+        key.is_bigram = True
+        if value is not None:
+            value = value[: len(key)]
+    return key, value
+
+
+def page_align_keys(key: list, page_size) -> list:
+    if page_size == 1:
+        return key
+    page_aligned_len = len(key) // page_size * page_size
+    return key[:page_aligned_len]
+
+
 class TreeNode:
 
     counter = 0
@@ -342,12 +362,7 @@ class RadixCache(BasePrefixCache):
     def maybe_bigram_convert(
         self, key: RadixKey, value: Optional[torch.Tensor] = None
     ) -> Tuple[RadixKey, Optional[torch.Tensor]]:
-        if self.is_eagle and not key.is_bigram:
+        return maybe_bigram_convert(self.is_eagle, key, value)
-            key.token_ids = convert_to_bigram_key(key.token_ids)
-            if value is not None:
-                value = value[: len(key)]
-
-        return key, value
 
     def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
         """Find the longest cached prefix of ``key`` in the radix tree.
@@ -437,12 +452,6 @@ class RadixCache(BasePrefixCache):
         prefix_len = self._insert_helper(self.root_node, key, value, priority)
         return InsertResult(prefix_len=prefix_len)
 
-    def _page_align_keys(self, key: list) -> list:
-        if self.page_size == 1:
-            return key
-        page_aligned_len = len(key) // self.page_size * self.page_size
-        return key[:page_aligned_len]
-
     def cache_finished_req(self, req: Req, is_insert: bool = True):
         """Cache request when it finishes."""
         # In deterministic mode, disable finished request insertion to radix cache
@@ -464,7 +473,7 @@ class RadixCache(BasePrefixCache):
 
         # Maybe convert to bigram keys for EAGLE
         keys = convert_to_bigram_key(token_ids) if self.is_eagle else token_ids
-        keys = self._page_align_keys(keys)
+        keys = page_align_keys(keys, self.page_size)
         values = kv_indices[: len(keys)].to(dtype=torch.int64, copy=True)
         radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle)
 
@@ -502,7 +511,7 @@ class RadixCache(BasePrefixCache):
 
         # Maybe convert to bigram keys for EAGLE
         keys = convert_to_bigram_key(token_ids) if self.is_eagle else token_ids
-        keys = self._page_align_keys(keys)
+        keys = page_align_keys(keys, self.page_size)
         values = kv_indices[: len(keys)].to(dtype=torch.int64, copy=True)
         radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle)
 

python/sglang/srt/mem_cache/swa_radix_cache.py

@@ -43,6 +43,8 @@ from sglang.srt.mem_cache.radix_cache import (
     _key_match_page_size1,
     _key_match_paged,
     get_child_key,
+    maybe_bigram_convert,
+    page_align_keys,
 )
 from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
 from sglang.srt.mem_cache.utils import convert_to_bigram_key
@@ -426,14 +428,10 @@ class SWARadixCache(BasePrefixCache):
         prev_prefix_len = params.prev_prefix_len
         swa_evicted_seqlen = params.swa_evicted_seqlen
 
-        key.token_ids = self.key_convert_fn(key.token_ids)
-
         if value is None:
             value = torch.tensor([x for x in key.token_ids], dtype=torch.int64)
 
-        if self.is_eagle:
+        key, value = maybe_bigram_convert(self.is_eagle, key, value)
-            # Make sure the value len equal to the EAGLE bigram key len
-            value = value[: len(key)]
 
         prefix_len = self._insert_helper(
             self.root_node, key, value, prev_prefix_len, swa_evicted_seqlen
@@ -451,40 +449,29 @@ class SWARadixCache(BasePrefixCache):
             return
 
         token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
-        # For EAGLE radix cache, we will convert the key to bigram key, e.g. [1,2,3,4] -> [(1,2), (2,3), (3,4)], the length will -1. ((len([(1,2), (2,3), (3,4)]) = len([1,2,3,4]) - 1))
-        # So for the corresponding kv length should also -1. Then we get the actual_kv_len, and use it to do later calculation and slicing.
-        actual_kv_len = kv_committed_len - 1 if self.is_eagle else kv_committed_len
         kv_indices = self.req_to_token_pool.req_to_token[
             req.req_pool_idx, :kv_committed_len
         ]
 
-        if self.page_size != 1:
+        # Maybe convert to bigram keys for EAGLE
-            page_aligned_len = actual_kv_len // self.page_size * self.page_size
+        keys = self.key_convert_fn(token_ids)
-            page_aligned_kv_indices = kv_indices[:page_aligned_len].to(
+        keys = page_align_keys(keys, self.page_size)
-                dtype=torch.int64, copy=True
+        page_aligned_len = len(keys)
+        values = kv_indices[:page_aligned_len].to(dtype=torch.int64, copy=True)
+        radix_key = RadixKey(
+            keys[:page_aligned_len],
+            req.extra_key,
+            is_bigram=self.is_eagle,
         )
-        else:
+        old_prefix_len = req.cache_protected_len
-            page_aligned_len = actual_kv_len
-            page_aligned_kv_indices = kv_indices.to(dtype=torch.int64, copy=True)
-
-        page_aligned_token_len = (
-            page_aligned_len + 1 if self.is_eagle else page_aligned_len
-        )
-
-        old_prefix_len = len(req.prefix_indices)
-        if self.is_eagle and old_prefix_len > req.cache_protected_len:
-            # In EAGLE chunked prefill case, the prefix_indices included one unmatched token (kv_indices[actual_kv_len:])
-            # Here we -1 to make sure the kv of the unmatched token can be freed correctly to avoid memory leak
-            old_prefix_len -= 1
 
         # Radix Cache takes one ref in memory pool
-        # insert the token_ids and kv_indices into the radix tree
         # Note: the insert function already frees the overlapped kv_indices
         if is_insert:
             self.insert(
                 InsertParams(
-                    key=RadixKey(token_ids[:page_aligned_token_len], req.extra_key),
+                    key=radix_key,
-                    value=page_aligned_kv_indices,
+                    value=values,
                     prev_prefix_len=old_prefix_len,
                     swa_evicted_seqlen=req.swa_evicted_seqlen,
                 )
@@ -512,58 +499,35 @@ class SWARadixCache(BasePrefixCache):
             return
 
         token_ids = req.fill_ids
-        all_token_len = len(token_ids)
-        # For EAGLE radix cache, we will convert the key to bigram key, e.g. [1,2,3,4] -> [(1,2), (2,3), (3,4)], the length will -1. ((len([(1,2), (2,3), (3,4)]) = len([1,2,3,4]) - 1))
-        # So for the corresponding kv length should also -1. Then we get the actual_kv_len, and use it to do later calculation and slicing.
-        actual_kv_len = all_token_len - 1 if self.is_eagle else all_token_len
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.req_pool_idx, :all_token_len
+            req.req_pool_idx, : len(token_ids)
         ]
 
-        if self.page_size != 1:
+        keys = self.key_convert_fn(token_ids)
-            page_aligned_len = actual_kv_len // self.page_size * self.page_size
+        keys = page_align_keys(keys, self.page_size)
-            page_aligned_kv_indices = kv_indices[:page_aligned_len].to(
+        values = kv_indices[: len(keys)].to(dtype=torch.int64, copy=True)
-                dtype=torch.int64, copy=True
+        radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle)
-            )
+        old_prefix_len = req.cache_protected_len
-        else:
-            page_aligned_len = actual_kv_len
-            page_aligned_kv_indices = kv_indices.to(dtype=torch.int64, copy=True)
-
-        # For EAGLE, the page_aligned_len is for the bigram key, the normal key len should +1
-        page_aligned_token_len = (
-            page_aligned_len + 1 if self.is_eagle else page_aligned_len
-        )
-        page_aligned_token_ids = token_ids[:page_aligned_token_len]
-
-        old_prefix_len = len(req.prefix_indices)
-        if self.is_eagle and old_prefix_len > req.cache_protected_len:
-            # In EAGLE chunked prefill case, the prefix_indices included one unmatched token (kv_indices[actual_kv_len:])
-            # Here we -1 to make sure the kv of the unmatched token can be freed correctly to avoid memory leak
-            old_prefix_len -= 1
 
         # Radix Cache takes one ref in memory pool
         # Note: the insert function already frees the overlapped kv_indices
         result = self.insert(
             InsertParams(
-                key=RadixKey(page_aligned_token_ids, req.extra_key),
+                key=radix_key,
-                value=page_aligned_kv_indices,
+                value=values,
                 prev_prefix_len=old_prefix_len,
             )
         )
         new_prefix_len = result.prefix_len
 
         # The prefix indices could be updated, reuse it
-        match_result = self.match_prefix(
+        match_result = self.match_prefix(MatchPrefixParams(key=radix_key))
-            MatchPrefixParams(key=RadixKey(page_aligned_token_ids, req.extra_key))
-        )
         new_indices, new_last_node = (
             match_result.device_indices,
             match_result.last_device_node,
         )
 
-        assert old_prefix_len <= len(
+        assert old_prefix_len <= len(new_indices), f"{old_prefix_len=}, {new_indices=}"
-            new_indices
-        ), f"{req.prefix_indices=}, {new_indices=}"
         assert new_prefix_len <= len(new_indices), f"{new_prefix_len=}, {new_indices=}"
         self.req_to_token_pool.write(
             (req.req_pool_idx, slice(old_prefix_len, len(new_indices))),
@@ -576,16 +540,10 @@ class SWARadixCache(BasePrefixCache):
         swa_uuid_for_lock = self.inc_lock_ref(new_last_node)
 
         # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later
-        if self.page_size != 1:
+        if len(new_indices) < len(kv_indices):
             req.prefix_indices = torch.cat(
                 [new_indices, kv_indices[len(new_indices) :]]
             )
-        else:
-            if self.is_eagle:
-                # Attach the kv index of the last token for EAGLE, it can be used in chunked prefill
-                req.prefix_indices = torch.cat(
-                    [new_indices, kv_indices[actual_kv_len:]]
-                )
         else:
             req.prefix_indices = new_indices
         req.last_node = new_last_node
@@ -874,7 +832,7 @@ class SWARadixCache(BasePrefixCache):
     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)
+        key, _ = maybe_bigram_convert(self.is_eagle, key)
 
         if self.disable or len(key) == 0:
             return None

test/registered/radix_cache/test_swa_unittest.py

@@ -22,6 +22,74 @@ register_amd_ci(est_time=10, suite="stage-b-test-small-1-gpu-amd")
 
 
 class TestSWA(unittest.TestCase):
+    class _DummyReq:
+        def __init__(self):
+            self._kv_committed_len = 0
+
+        def pop_committed_kv_cache(self):
+            return self._kv_committed_len
+
+    def _build_swa_tree(
+        self,
+        is_eagle: bool,
+        page_size: int = 1,
+        req_size: int = 8,
+        max_context_len: int = 64,
+        kv_size: int = 64,
+        kv_size_swa: int = 32,
+        sliding_window_size: int = 4,
+    ):
+        head_num = 8
+        head_dim = 128
+        num_layers = 24
+        global_interval = 4
+        dtype = torch.bfloat16
+        device = get_device()
+        full_attention_layer_ids = [i for i in range(0, num_layers, global_interval)]
+        full_attention_layer_ids_set = set(full_attention_layer_ids)
+        swa_attention_layer_ids = [
+            i for i in range(num_layers) if i not in full_attention_layer_ids_set
+        ]
+
+        req_to_token_pool = ReqToTokenPool(
+            size=req_size,
+            max_context_len=max_context_len,
+            device=device,
+            enable_memory_saver=False,
+        )
+        kv_pool = SWAKVPool(
+            size=kv_size,
+            size_swa=kv_size_swa,
+            page_size=page_size,
+            dtype=dtype,
+            head_num=head_num,
+            head_dim=head_dim,
+            swa_attention_layer_ids=swa_attention_layer_ids,
+            full_attention_layer_ids=full_attention_layer_ids,
+            enable_kvcache_transpose=False,
+            device=device,
+        )
+        allocator = SWATokenToKVPoolAllocator(
+            size=kv_size,
+            size_swa=kv_size_swa,
+            page_size=page_size,
+            dtype=dtype,
+            device=device,
+            kvcache=kv_pool,
+            need_sort=False,
+        )
+        tree = SWARadixCache(
+            params=CacheInitParams(
+                req_to_token_pool=req_to_token_pool,
+                token_to_kv_pool_allocator=allocator,
+                page_size=page_size,
+                disable=False,
+                is_eagle=is_eagle,
+                sliding_window_size=sliding_window_size,
+            ),
+        )
+        return tree, allocator, req_to_token_pool
+
     @classmethod
     def setUpClass(cls):
         pass
@@ -413,6 +481,77 @@ class TestSWA(unittest.TestCase):
         self.assertEqual(last_node.key.token_ids[0], (5, 60))
         self.assertEqual(last_node.key.token_ids[1], (60, 70))
 
+    def test_swa_cache_finished_req_eagle_uses_cache_protected_len_and_bigram_key(self):
+        tree, allocator, req_to_token_pool = self._build_swa_tree(is_eagle=True)
+
+        # Case 1: is_insert=True should pass bigram key and use cache_protected_len.
+        req = self._DummyReq()
+        req.req_pool_idx = 0
+        req.origin_input_ids = [1, 2, 3, 4, 5, 6]
+        req.output_ids = []
+        req._kv_committed_len = len(req.origin_input_ids)
+        kv_indices = allocator.alloc(req._kv_committed_len)
+        req_to_token_pool.write(
+            (req.req_pool_idx, slice(0, req._kv_committed_len)), kv_indices
+        )
+        req.extra_key = None
+        req.last_node = tree.root_node
+        req.swa_uuid_for_lock = None
+        req.swa_evicted_seqlen = 0
+        req.cache_protected_len = 1
+        # Intentionally mismatch to ensure code does not use len(prefix_indices).
+        req.prefix_indices = torch.tensor([7, 8, 9, 10, 11], device=tree.device)
+
+        captured = {}
+        original_insert = tree.insert
+
+        def wrapped_insert(params):
+            captured["prev_prefix_len"] = params.prev_prefix_len
+            captured["is_bigram"] = params.key.is_bigram
+            captured["key_len"] = len(params.key)
+            return original_insert(params)
+
+        tree.insert = wrapped_insert
+        tree.cache_finished_req(req, is_insert=True)
+
+        self.assertEqual(captured["prev_prefix_len"], req.cache_protected_len)
+        self.assertTrue(captured["is_bigram"])
+        self.assertEqual(captured["key_len"], len(req.origin_input_ids) - 1)
+
+        # Case 2: is_insert=False should free [cache_protected_len:page_aligned_len]
+        # even when len(prefix_indices) is intentionally larger.
+        req2 = self._DummyReq()
+        req2.req_pool_idx = 1
+        req2.origin_input_ids = [11, 12, 13, 14, 15, 16]
+        req2.output_ids = []
+        req2._kv_committed_len = len(req2.origin_input_ids)
+        kv_indices2 = allocator.alloc(req2._kv_committed_len)
+        req_to_token_pool.write(
+            (req2.req_pool_idx, slice(0, req2._kv_committed_len)), kv_indices2
+        )
+        req2.extra_key = None
+        req2.last_node = tree.root_node
+        req2.swa_uuid_for_lock = None
+        req2.swa_evicted_seqlen = 0
+        req2.cache_protected_len = 1
+        req2.prefix_indices = torch.tensor([21, 22, 23, 24, 25], device=tree.device)
+
+        freed_lens = []
+        original_free = allocator.free
+
+        def wrapped_free(indices):
+            freed_lens.append(int(indices.numel()))
+            return original_free(indices)
+
+        allocator.free = wrapped_free
+        tree.cache_finished_req(req2, is_insert=False)
+
+        # EAGLE + page_size=1 => page_aligned_len = committed_len - 1 = 5
+        # Expected frees:
+        #   overlap range [1:5] -> 4
+        #   tail range [5:]     -> 1
+        self.assertEqual(freed_lens, [4, 1])
+
 
 if __name__ == "__main__":
     unittest.main()