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

[RadixTree][1/N Refactor]: Support unified match_prefix params (#17142)

Browse Source 
Co-authored-by: yizhang2077 <1109276519@qq.com>
Co-authored-by: pansicheng <sicheng.pan.chn@gmail.com>
This commit is contained in:
zhangheng
2026-01-19 22:39:40 +08:00
committed by GitHub
co-authored by yizhang2077 pansicheng
parent ce8a6ac690
commit 20b0523eca
13 changed files with 117 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/radix_cache/test_mamba_unittest.py
+6 -5
View File
@@ -6,6 +6,7 @@ import torch
from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import MatchPrefixParams
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, HybridReqToTokenPool
@@ -289,7 +290,7 @@ class TestMamba(unittest.TestCase):
tree.pretty_print()
req5_token_ids = [1, 2, 3, 4, 5]
result = tree.match_prefix(RadixKey(req5_token_ids))
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req5_token_ids)))
kv_indices, last_node = result.device_indices, result.last_device_node
print(
f"req5: token_ids: {req5_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
@@ -297,7 +298,7 @@ class TestMamba(unittest.TestCase):
assert len(kv_indices) == 0
req6_token_ids = [1, 2, 3, 4, 5, 60, 70]
result = tree.match_prefix(RadixKey(req6_token_ids))
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req6_token_ids)))
kv_indices, last_node = result.device_indices, result.last_device_node
print(
f"req6: token_ids: {req6_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
@@ -306,7 +307,7 @@ class TestMamba(unittest.TestCase):
assert len(last_node.key) == 2
req7_token_ids = [1, 2, 3, 4, 5, 6, 7]
result = tree.match_prefix(RadixKey(req7_token_ids))
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req7_token_ids)))
kv_indices, last_node = result.device_indices, result.last_device_node
print(
f"req7: token_ids: {req7_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
@@ -320,7 +321,7 @@ class TestMamba(unittest.TestCase):
tree.pretty_print()
req8_token_ids = [1, 2, 3, 4, 5, 60, 70]
result = tree.match_prefix(RadixKey(req8_token_ids))
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req8_token_ids)))
kv_indices, last_node = result.device_indices, result.last_device_node
print(
f"req8: token_ids: {req8_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
@@ -331,7 +332,7 @@ class TestMamba(unittest.TestCase):
req9_token_ids = [1, 2, 3, 4, 5, 6, 7]
req9 = make_dummy_req()
result = tree.match_prefix(
RadixKey(req9_token_ids), **({"req": req9, "cow_mamba": True})
MatchPrefixParams(key=RadixKey(req9_token_ids), req=req9, cow_mamba=True)
)
kv_indices, last_node = result.device_indices, result.last_device_node
assert req9.mamba_pool_idx is not None
test/registered/radix_cache/test_radix_cache_unit.py
+17 -12
View File
@@ -31,6 +31,7 @@ import unittest.mock
import torch
from sglang.srt.disaggregation.kv_events import BlockRemoved, BlockStored
from sglang.srt.mem_cache.base_prefix_cache import MatchPrefixParams
from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
# Test constants
@@ -294,12 +295,12 @@ class TestRadixCache(unittest.TestCase):
self.assertEqual(cache.evictable_size(), 3)
# Test match_prefix
result = cache.match_prefix(RadixKey([1, 2, 3]))
result = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3])))
self.assertEqual(len(result.device_indices), 3)
torch.testing.assert_close(result.device_indices, value)
# Test partial match
result = cache.match_prefix(RadixKey([1, 2]))
result = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2])))
self.assertEqual(len(result.device_indices), 2)
torch.testing.assert_close(
result.device_indices, torch.tensor([10, 20], dtype=torch.int64)
@@ -402,10 +403,12 @@ class TestRadixCache(unittest.TestCase):
)
# Keys with different extra_key should not match each other
result1 = cache.match_prefix(RadixKey([1, 2, 3], "key1"))
result2 = cache.match_prefix(RadixKey([1, 2, 3], "key2"))
result3 = cache.match_prefix(RadixKey([1, 2, 3], None))
result4 = cache.match_prefix(RadixKey([1, 2, 3], "nonexistent"))
result1 = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3], "key1")))
result2 = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3], "key2")))
result3 = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3], None)))
result4 = cache.match_prefix(
MatchPrefixParams(key=RadixKey([1, 2, 3], "nonexistent"))
)
# Each should match only its own data
self.assertEqual(len(result1.device_indices), 3)
@@ -434,7 +437,7 @@ class TestRadixCache(unittest.TestCase):
cache.insert(RadixKey([1, 2, 3]), torch.tensor([10, 20, 30], dtype=torch.int64))
# Get node
result = cache.match_prefix(RadixKey([1, 2, 3]))
result = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3])))
node = result.last_device_node
initial_evictable = cache.evictable_size()
@@ -485,7 +488,7 @@ class TestRadixCache(unittest.TestCase):
tokens = list(range(sequence_length))
cache.insert(RadixKey(tokens), torch.tensor(tokens, dtype=torch.int64))
result = cache.match_prefix(RadixKey(tokens))
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(tokens)))
self.assertGreater(len(result.device_indices), 0)
# Match length should be page-aligned
@@ -541,23 +544,25 @@ class TestRadixCache(unittest.TestCase):
# Match that causes a split inside an existing node:
# take first 4 tokens of seq1, then diverge.
query1 = [1, 2, 3, 4, 999, 1000]
result1 = cache.match_prefix(RadixKey(query1))
result1 = cache.match_prefix(MatchPrefixParams(key=RadixKey(query1)))
torch.testing.assert_close(result1.device_indices, val1[:4])
# No data change after structural split during matching.
self.assertEqual(cache.total_size(), baseline_total)
# Full match of the long sequence still returns the full indices.
result_full = cache.match_prefix(RadixKey(seq1))
result_full = cache.match_prefix(MatchPrefixParams(key=RadixKey(seq1)))
torch.testing.assert_close(result_full.device_indices, val1)
# Another split deeper on the path (after matching 6 tokens, then diverge).
query2 = [1, 2, 3, 4, 5, 6, 777, 888]
result2 = cache.match_prefix(RadixKey(query2))
result2 = cache.match_prefix(MatchPrefixParams(key=RadixKey(query2)))
torch.testing.assert_close(result2.device_indices, val1[:6])
self.assertEqual(cache.total_size(), baseline_total)
# Matching the short diverging branch should return exactly its indices.
result_branch = cache.match_prefix(RadixKey(seq2))
result_branch = cache.match_prefix(
MatchPrefixParams(key=RadixKey(seq2))
)
torch.testing.assert_close(result_branch.device_indices, val2)
def test_hash_value_storage(self):
test/registered/radix_cache/test_swa_unittest.py
+5 -4
View File
@@ -2,6 +2,7 @@ import unittest
import torch
from sglang.srt.mem_cache.base_prefix_cache import MatchPrefixParams
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.mem_cache.radix_cache import RadixKey
@@ -188,7 +189,7 @@ class TestSWA(unittest.TestCase):
tree.pretty_print()
req5_token_ids = [1, 2, 3, 4, 5]
result = tree.match_prefix(RadixKey(req5_token_ids))
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req5_token_ids)))
kv_indices, last_node = result.device_indices, result.last_device_node
print(
f"req5: token_ids: {req5_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
@@ -196,7 +197,7 @@ class TestSWA(unittest.TestCase):
self.assertEqual(len(kv_indices), 0)
req6_token_ids = [1, 2, 3, 4, 5, 60, 70]
result = tree.match_prefix(RadixKey(req6_token_ids))
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req6_token_ids)))
kv_indices, last_node = result.device_indices, result.last_device_node
print(
f"req6: token_ids: {req6_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
@@ -329,7 +330,7 @@ class TestSWA(unittest.TestCase):
tree.pretty_print()
req5_token_ids = [1, 2, 3, 4, 5]
result = tree.match_prefix(RadixKey(req5_token_ids))
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req5_token_ids)))
kv_indices, last_node = result.device_indices, result.last_device_node
print(
f"req5: token_ids: {req5_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
@@ -337,7 +338,7 @@ class TestSWA(unittest.TestCase):
self.assertEqual(len(kv_indices), 0) # no swa prefix matched
req6_token_ids = [1, 2, 3, 4, 5, 60, 70]
result = tree.match_prefix(RadixKey(req6_token_ids))
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req6_token_ids)))
kv_indices, last_node = result.device_indices, result.last_device_node
print(
f"req6: token_ids: {req6_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"