diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 42b169728..9292facb0 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/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: - key.token_ids = convert_to_bigram_key(key.token_ids) - if value is not None: - value = value[: len(key)] - - return key, value + return maybe_bigram_convert(self.is_eagle, 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) diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index b0e548c17..50b24a313 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/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: - # Make sure the value len equal to the EAGLE bigram key len - value = value[: len(key)] + key, value = maybe_bigram_convert(self.is_eagle, key, value) 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: - 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 + # Maybe convert to bigram keys for EAGLE + keys = self.key_convert_fn(token_ids) + keys = page_align_keys(keys, self.page_size) + 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, ) - - 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 + old_prefix_len = req.cache_protected_len # 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), - value=page_aligned_kv_indices, + key=radix_key, + 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: - 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.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 + keys = self.key_convert_fn(token_ids) + 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) + old_prefix_len = req.cache_protected_len # 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), - value=page_aligned_kv_indices, + key=radix_key, + 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( - MatchPrefixParams(key=RadixKey(page_aligned_token_ids, req.extra_key)) - ) + match_result = self.match_prefix(MatchPrefixParams(key=radix_key)) new_indices, new_last_node = ( match_result.device_indices, match_result.last_device_node, ) - assert old_prefix_len <= len( - new_indices - ), f"{req.prefix_indices=}, {new_indices=}" + assert old_prefix_len <= len(new_indices), f"{old_prefix_len=}, {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,18 +540,12 @@ 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.prefix_indices = new_indices req.last_node = new_last_node req.swa_uuid_for_lock = swa_uuid_for_lock @@ -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 diff --git a/test/registered/radix_cache/test_swa_unittest.py b/test/registered/radix_cache/test_swa_unittest.py index 3b1a57de8..f027c1fae 100644 --- a/test/registered/radix_cache/test_swa_unittest.py +++ b/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()