Refactor eagle bigram key matching (#13714)

python/sglang/srt/managers/schedule_batch.py

@@ -562,8 +562,8 @@ class Req:
         self.host_hit_length = 0
         # The node to lock until for swa radix tree lock ref
         self.swa_uuid_for_lock: Optional[int] = None
-        # The prefix length of the last prefix matching
+        # The prefix length that is inserted into the tree cache
-        self.last_matched_prefix_len: int = 0
+        self.cache_protected_len: int = 0
 
         # Whether or not if it is chunked. It increments whenever
         # it is chunked, and decrement whenever chunked request is
@@ -775,7 +775,7 @@ class Req:
                     else {}
                 ),
             )
-            self.last_matched_prefix_len = len(self.prefix_indices)
+            self.cache_protected_len = len(self.prefix_indices)
         self.extend_input_len = len(self.fill_ids) - len(self.prefix_indices)
 
     # Based on https://github.com/vllm-project/vllm/blob/7a64d24aad69e4d2548aa0bf528d9fe63428ab01/vllm/transformers_utils/detokenizer.py#L194-L313

python/sglang/srt/managers/schedule_policy.py

@@ -603,7 +603,7 @@ class PrefillAdder:
                 req.prefix_indices = torch.cat([req.prefix_indices, new_indices])
                 req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices)
                 prefix_len = len(req.prefix_indices)
-                req.last_matched_prefix_len = prefix_len
+                req.cache_protected_len = prefix_len
 
             input_tokens = self.ceil_paged_tokens(req.extend_input_len)
 

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -685,7 +685,7 @@ class HiRadixCache(RadixCache):
 
     def match_prefix(self, key: RadixKey, **kwargs):
         empty_value = torch.empty((0,), dtype=torch.int64, device=self.device)
-        key.token_ids = self.key_convert_fn(key.token_ids)
+        key, _ = self.maybe_bigram_convert(key)
         if self.disable or len(key) == 0:
             return MatchResult(
                 device_indices=empty_value,
@@ -857,7 +857,7 @@ class HiRadixCache(RadixCache):
     ):
         if priority is None:
             priority = 0
-        key.token_ids = self.key_convert_fn(key.token_ids)
+        key, value = self.maybe_bigram_convert(key, value)
 
         if len(key) == 0:
             return 0

python/sglang/srt/mem_cache/radix_cache.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from sglang.srt.mem_cache.utils import convert_to_bigram_key
+
 """
 Copyright 2023-2024 SGLang Team
 Licensed under the Apache License, Version 2.0 (the "License");
@@ -51,12 +53,18 @@ if TYPE_CHECKING:
 
 
 class RadixKey:
-
+    def __init__(
-    def __init__(self, token_ids: List[int], extra_key: Optional[str] = None):
+        self,
+        token_ids: List[int],
+        extra_key: Optional[str] = None,
+        is_bigram: bool = False,
+    ):
         # token ids sequence
         self.token_ids = token_ids
         # extra key (e.g. lora_id, cache_salt)
         self.extra_key = extra_key
+        # is bigram key
+        self.is_bigram = is_bigram
 
     def __len__(self) -> int:
         return len(self.token_ids)
@@ -178,16 +186,6 @@ def get_child_key(key: RadixKey, page_size: int = 1):
         return (key.extra_key, plain_key)
 
 
-def _convert_to_bigram_key(tokens: List[int]) -> List[Tuple[int, int]]:
-    # EAGLE uses bigram keys in the radix tree since draft sequence is the one-token-shifted version of target
-    # [1, 2, 3, 4] -> [(1,2), (2,3), (3,4)]
-    if len(tokens) < 2:
-        return []
-    if isinstance(tokens[0], tuple):
-        return tokens
-    return [(tokens[i], tokens[i + 1]) for i in range(len(tokens) - 1)]
-
-
 class RadixCache(BasePrefixCache):
     def __init__(
         self,
@@ -225,11 +223,6 @@ class RadixCache(BasePrefixCache):
             self.key_match_fn = partial(_key_match_paged, page_size=page_size)
             self.get_child_key_fn = partial(get_child_key, page_size=page_size)
 
-        if is_eagle:
-            self.key_convert_fn = _convert_to_bigram_key
-        else:
-            self.key_convert_fn = lambda key: key
-
         if eviction_policy.lower() == "lru":
             self.eviction_strategy: EvictionStrategy = LRUStrategy()
         elif eviction_policy.lower() == "lfu":
@@ -261,6 +254,16 @@ class RadixCache(BasePrefixCache):
         self.protected_size_ = 0
         self._record_all_cleared_event()
 
+    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:
+            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, key: RadixKey, **kwargs) -> MatchResult:
         """Find the longest cached prefix of ``key`` in the radix tree.
 
@@ -299,7 +302,7 @@ class RadixCache(BasePrefixCache):
                 to expose a precise boundary; this structural refinement improves
                 subsequent match efficiency and does not duplicate data.
         """
-        key.token_ids = self.key_convert_fn(key.token_ids)
+        key, _ = self.maybe_bigram_convert(key)
 
         def empty_match_result():
             return MatchResult(
@@ -337,78 +340,60 @@ class RadixCache(BasePrefixCache):
         if self.disable:
             return 0
 
-        key.token_ids = self.key_convert_fn(key.token_ids)
-
         if value is None:
             value = torch.tensor(key.token_ids, dtype=torch.int64)
 
-        if self.is_eagle:
+        key, value = self.maybe_bigram_convert(key, value)
-            # Make sure the value len equal to the EAGLE bigram key len
-            value = value[: len(key)]
 
         return self._insert_helper(self.root_node, key, value, priority)
 
+    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
         if self.disable_finished_insert:
             is_insert = False
 
-        committed_kv_len = req.pop_committed_kv_cache()
+        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, :committed_kv_len
+                req.req_pool_idx, :kv_committed_len
             ]
             self.token_to_kv_pool_allocator.free(kv_indices)
             self.req_to_token_pool.free(req.req_pool_idx)
             return
 
-        token_ids = (req.origin_input_ids + req.output_ids)[:committed_kv_len]
+        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 = committed_kv_len - 1 if self.is_eagle else committed_kv_len
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.req_pool_idx, :committed_kv_len
+            req.req_pool_idx, : len(token_ids)
         ]
 
-        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 = convert_to_bigram_key(req.fill_ids) if self.is_eagle else req.fill_ids
-            page_aligned_kv_indices = kv_indices[:page_aligned_len].to(
+        keys = self._page_align_keys(keys)
-                dtype=torch.int64, copy=True
+        values = kv_indices[: len(keys)].to(dtype=torch.int64, copy=True)
-            )
+        radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle)
-        else:
-            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.last_matched_prefix_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
         if is_insert:
             priority = getattr(req, "priority", 0) or 0
-            new_prefix_len = self.insert(
+            new_prefix_len = self.insert(radix_key, values, priority=priority)
-                RadixKey(token_ids[:page_aligned_token_len], req.extra_key),
-                page_aligned_kv_indices,
-                priority=priority,
-            )
             # Free the duplicates that were already in the tree
             self.token_to_kv_pool_allocator.free(
-                kv_indices[old_prefix_len:new_prefix_len]
+                kv_indices[req.cache_protected_len : new_prefix_len]
             )
         else:
             self.token_to_kv_pool_allocator.free(
-                kv_indices[old_prefix_len:page_aligned_len]
+                kv_indices[req.cache_protected_len : len(keys)]
             )
 
         # free the unaligned tail
-        self.token_to_kv_pool_allocator.free(kv_indices[page_aligned_len:])
+        self.token_to_kv_pool_allocator.free(kv_indices[len(keys) :])
 
         # Remove req slot release the cache lock
         self.req_to_token_pool.free(req.req_pool_idx)
@@ -420,77 +405,56 @@ class RadixCache(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:
+        # Maybe convert to bigram keys for EAGLE
-            page_aligned_len = actual_kv_len // self.page_size * self.page_size
+        keys = convert_to_bigram_key(req.fill_ids) if self.is_eagle else req.fill_ids
-            page_aligned_kv_indices = kv_indices[:page_aligned_len].to(
+        keys = self._page_align_keys(keys)
-                dtype=torch.int64, copy=True
+        values = kv_indices[: len(keys)].to(dtype=torch.int64, copy=True)
-            )
+        radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle)
-        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.last_matched_prefix_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
-        priority = getattr(req, "priority", 0) or 0
         new_prefix_len = self.insert(
-            RadixKey(page_aligned_token_ids, req.extra_key),
+            radix_key,
-            page_aligned_kv_indices,
+            values,
             chunked=chunked,
-            priority=priority,
+            priority=getattr(req, "priority", 0) or 0,
+        )
+
+        self.token_to_kv_pool_allocator.free(
+            kv_indices[req.cache_protected_len : new_prefix_len]
         )
-        self.token_to_kv_pool_allocator.free(kv_indices[old_prefix_len:new_prefix_len])
 
         # The prefix indices could be updated, reuse it
-        new_indices, new_last_node, _, _ = self.match_prefix(
+        new_indices, new_last_node, _, _ = self.match_prefix(radix_key)
-            RadixKey(token_ids=page_aligned_token_ids, extra_key=req.extra_key)
+        assert len(new_indices) == len(keys), f"{len(new_indices)=}, {len(keys)=}"
-        )
+
         self.req_to_token_pool.write(
-            (req.req_pool_idx, slice(old_prefix_len, len(new_indices))),
+            (req.req_pool_idx, slice(req.cache_protected_len, len(new_indices))),
-            new_indices[old_prefix_len:],
+            new_indices[req.cache_protected_len :],
         )
 
-        # The last_matched_prefix_len is not always equal to len(req.prefix_indices)
+        # The cache_protected_len is not always equal to len(req.prefix_indices)
         # since for page_size > 1, the partial part is added to req.prefix_indices, but that part of kv indices is not added to the tree.
         # It should be freed in the next cache_unfinished_req and final cache_finished_req to avoid memory leak.
-        # So we introduce this `last_matched_prefix_len` field to make sure the partial part can be freed correctly.
+        # So we introduce this `cache_protected_len` field to make sure the partial part can be freed correctly.
-        req.last_matched_prefix_len = len(new_indices)
+        req.cache_protected_len = len(new_indices)
 
         self.dec_lock_ref(req.last_node)
         self.inc_lock_ref(new_last_node)
 
         # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later
-        if self.page_size != 1:
+        # - page_size != 1: there is a partial page at the end, keep the full kv_indices
-            # Handle partial page, the partial part should be freed in the next cache_unfinished_req and final cache_finished_req.
+        # - eagle case: bigram keys will only cache len - 1 kv indices
+        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
 
     def pretty_print(self):

python/sglang/srt/mem_cache/swa_radix_cache.py

@@ -33,11 +33,11 @@ from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, MatchResult
 from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
 from sglang.srt.mem_cache.radix_cache import (
     RadixKey,
-    _convert_to_bigram_key,
     _key_match_page_size1,
     _key_match_paged,
     get_child_key,
 )
+from sglang.srt.mem_cache.utils import convert_to_bigram_key
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req
@@ -359,7 +359,7 @@ class SWARadixCache(BasePrefixCache):
             self.get_child_key_fn = partial(get_child_key, page_size=page_size)
 
         if is_eagle:
-            self.key_convert_fn = _convert_to_bigram_key
+            self.key_convert_fn = convert_to_bigram_key
         else:
             self.key_convert_fn = lambda key: key
 
@@ -474,7 +474,7 @@ class SWARadixCache(BasePrefixCache):
         )
 
         old_prefix_len = len(req.prefix_indices)
-        if self.is_eagle and old_prefix_len > req.last_matched_prefix_len:
+        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
@@ -536,7 +536,7 @@ class SWARadixCache(BasePrefixCache):
         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.last_matched_prefix_len:
+        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
@@ -562,7 +562,7 @@ class SWARadixCache(BasePrefixCache):
             new_indices[old_prefix_len:],
         )
 
-        req.last_matched_prefix_len = len(new_indices)
+        req.cache_protected_len = len(new_indices)
 
         self.dec_lock_ref(req.last_node, req.swa_uuid_for_lock)
         swa_uuid_for_lock = self.inc_lock_ref(new_last_node)

python/sglang/srt/mem_cache/utils.py

@@ -13,7 +13,7 @@
 # ==============================================================================
 """Common utilities."""
 
-from typing import Any, Optional, Tuple
+from typing import Any, List, Optional, Tuple
 
 import torch
 import triton
@@ -241,3 +241,13 @@ def maybe_init_custom_mem_pool(
         return init_mooncake_custom_mem_pool(device)
     else:
         return False, None, None
+
+
+def convert_to_bigram_key(tokens: List[int]) -> List[Tuple[int, int]]:
+    # EAGLE uses bigram keys in the radix tree since draft sequence is the one-token-shifted version of target
+    # [1, 2, 3, 4] -> [(1,2), (2,3), (3,4)]
+    if len(tokens) and isinstance(tokens[0], tuple):
+        return tokens
+    if len(tokens) < 2:
+        return []
+    return [(tokens[i], tokens[i + 1]) for i in range(len(tokens) - 1)]

test/srt/test_radix_cache_unit.py

@@ -307,72 +307,6 @@ class TestRadixCache(unittest.TestCase):
                     result.device_indices, torch.tensor([10, 20], dtype=torch.int64)
                 )
 
-    def test_insert_and_match_eagle(self):
-        """Test insert and match operations for EAGLE."""
-        cache = RadixCache(
-            req_to_token_pool=None,
-            token_to_kv_pool_allocator=None,
-            page_size=1,
-            disable=False,
-            is_eagle=True,
-        )
-
-        key = RadixKey([1, 2, 3, 4])
-        value = torch.tensor([10, 20, 30, 40], dtype=torch.int64)
-        prefix_len = cache.insert(key, value)
-
-        self.assertEqual(prefix_len, 0)  # No existing prefix
-        self.assertEqual(
-            cache.total_size(), 3
-        )  # The last token is ignored in bigram key
-        self.assertEqual(cache.evictable_size(), 3)
-
-        # Test match_prefix
-        result = cache.match_prefix(RadixKey([1, 2, 3, 4]))
-        self.assertEqual(len(result.device_indices), 3)
-        torch.testing.assert_close(
-            result.device_indices, torch.tensor([10, 20, 30], dtype=torch.int64)
-        )
-
-        # Test partial match
-        result = cache.match_prefix(RadixKey([1, 2]))
-        self.assertEqual(len(result.device_indices), 1)
-        torch.testing.assert_close(
-            result.device_indices, torch.tensor([10], dtype=torch.int64)
-        )
-
-    def test_insert_and_match_eagle_page_size(self):
-        """Test insert and match operations for EAGLE and page_size > 1."""
-        cache = RadixCache(
-            req_to_token_pool=None,
-            token_to_kv_pool_allocator=None,
-            page_size=2,
-            disable=False,
-            is_eagle=True,
-        )
-
-        key = RadixKey([1, 2, 3])
-        value = torch.tensor([10, 20, 30], dtype=torch.int64)
-        prefix_len = cache.insert(key, value)
-
-        self.assertEqual(prefix_len, 0)  # No existing prefix
-        self.assertEqual(cache.total_size(), 2)  # only one page is inserted
-        self.assertEqual(cache.evictable_size(), 2)
-
-        # Test match_prefix
-        result = cache.match_prefix(RadixKey([1, 2, 3, 4]))
-        self.assertEqual(len(result.device_indices), 2)
-        torch.testing.assert_close(
-            result.device_indices, torch.tensor([10, 20], dtype=torch.int64)
-        )
-
-        # Test unmatched
-        result = cache.match_prefix(RadixKey([1, 2]))
-        self.assertEqual(len(result.device_indices), 0)
-        torch.testing.assert_close(
-            result.device_indices, torch.tensor([], dtype=torch.int64)
-        )
-
     def test_insert_with_none_value(self):
         """Test insert with None value (should use token_ids as list)."""
         cache = RadixCache(