chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
zhangheng
2026-02-28 20:01:39 +08:00
committed by GitHub
parent 5705e02d28
commit a0d8a7ae6d
3 changed files with 190 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
python/sglang/srt/mem_cache/radix_cache.py
View File

@@ -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)

98
python/sglang/srt/mem_cache/swa_radix_cache.py
View File

@@ -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

139
test/registered/radix_cache/test_swa_unittest.py
View File

@@ -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()