From ac4382263451a9e3111fc9cba928fd34e0601d74 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sat, 22 Nov 2025 20:40:42 +0800 Subject: [PATCH] Refactor eagle bigram key matching (#13714) --- python/sglang/srt/managers/schedule_batch.py | 6 +- python/sglang/srt/managers/schedule_policy.py | 2 +- python/sglang/srt/mem_cache/hiradix_cache.py | 4 +- python/sglang/srt/mem_cache/radix_cache.py | 170 +++++++----------- .../sglang/srt/mem_cache/swa_radix_cache.py | 10 +- python/sglang/srt/mem_cache/utils.py | 12 +- test/srt/test_radix_cache_unit.py | 66 ------- 7 files changed, 89 insertions(+), 181 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index e1b986961..64c3eb1de 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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 - self.last_matched_prefix_len: int = 0 + # The prefix length that is inserted into the tree cache + 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 diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index e6cd8f708..d9af0c883 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/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) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 1feee6f51..87ed378ec 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 754a0f004..e528cf116 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/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__(self, token_ids: List[int], extra_key: Optional[str] = None): + def __init__( + 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: - # Make sure the value len equal to the EAGLE bigram key len - value = value[: len(key)] + key, value = self.maybe_bigram_convert(key, value) 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] - # 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 + token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_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: - page_aligned_len = actual_kv_len // self.page_size * self.page_size - page_aligned_kv_indices = kv_indices[:page_aligned_len].to( - dtype=torch.int64, copy=True - ) - 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 + # Maybe convert to bigram keys for EAGLE + keys = convert_to_bigram_key(req.fill_ids) if self.is_eagle else req.fill_ids + keys = self._page_align_keys(keys) + values = kv_indices[: len(keys)].to(dtype=torch.int64, copy=True) + radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle) # Radix Cache takes one ref in memory pool if is_insert: priority = getattr(req, "priority", 0) or 0 - new_prefix_len = self.insert( - RadixKey(token_ids[:page_aligned_token_len], req.extra_key), - page_aligned_kv_indices, - priority=priority, - ) + new_prefix_len = self.insert(radix_key, values, 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: - page_aligned_len = actual_kv_len // self.page_size * self.page_size - page_aligned_kv_indices = kv_indices[:page_aligned_len].to( - dtype=torch.int64, copy=True - ) - 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 + # Maybe convert to bigram keys for EAGLE + keys = convert_to_bigram_key(req.fill_ids) if self.is_eagle else req.fill_ids + keys = self._page_align_keys(keys) + values = kv_indices[: len(keys)].to(dtype=torch.int64, copy=True) + radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle) # 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), - page_aligned_kv_indices, + radix_key, + 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( - RadixKey(token_ids=page_aligned_token_ids, extra_key=req.extra_key) - ) + new_indices, new_last_node, _, _ = self.match_prefix(radix_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))), - new_indices[old_prefix_len:], + (req.req_pool_idx, slice(req.cache_protected_len, len(new_indices))), + 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. - req.last_matched_prefix_len = len(new_indices) + # So we introduce this `cache_protected_len` field to make sure the partial part can be freed correctly. + 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: - # Handle partial page, the partial part should be freed in the next cache_unfinished_req and final cache_finished_req. + # - page_size != 1: there is a partial page at the end, keep the full kv_indices + # - 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.prefix_indices = new_indices + req.last_node = new_last_node def pretty_print(self): diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index e19f9320c..2e97aca51 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/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) diff --git a/python/sglang/srt/mem_cache/utils.py b/python/sglang/srt/mem_cache/utils.py index 8f0e6d37d..d0abf02fe 100644 --- a/python/sglang/srt/mem_cache/utils.py +++ b/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)] diff --git a/test/srt/test_radix_cache_unit.py b/test/srt/test_radix_cache_unit.py index f8708eaf3..8cb75fb0b 100644 --- a/test/srt/test_radix_cache_unit.py +++ b/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(