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

perf: migrate Req token-id storage to array.array('q') in Scheduler (#25098)

Browse Source 
Co-authored-by: jialino <jialino@fb.com>
This commit is contained in:
Jialin Ouyang
2026-05-22 10:51:07 -07:00
committed by GitHub
co-authored by jialino
parent 5e9bd21979
commit 06c23d55b5
34 changed files with 833 additions and 299 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/managers/test_scheduler_chunked_req_gate.py
+4 -3
View File
@@ -1,6 +1,7 @@
"""Regression tests for the SWA chunked-req stash gate (#24252)."""
import unittest
from array import array
from types import SimpleNamespace
from unittest.mock import MagicMock
@@ -27,9 +28,9 @@ def _make_req(
) -> Req:
req = Req.__new__(Req)
req.rid = "test-req"
req.origin_input_ids = list(fill_ids)
req.output_ids = []
req.fill_ids = list(fill_ids)
req.origin_input_ids = array("q", fill_ids)
req.output_ids = array("q")
req.fill_ids = array("q", fill_ids)
req.prefix_indices = prefix_indices
req.req_pool_idx = req_pool_idx
req.extend_input_len = extend_input_len
test/registered/unit/mem_cache/test_decode_radix_lock_ref.py
+17 -10
View File
@@ -26,6 +26,7 @@ register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd")
import unittest
from array import array
from unittest.mock import MagicMock
import torch
@@ -65,11 +66,11 @@ class MockReq:
"""Minimal mock Req with fields needed by cache_unfinished/finished_req."""
def __init__(self, fill_ids, req_pool_idx=0, cache_protected_len=0, last_node=None):
self.fill_ids = list(fill_ids)
self.origin_input_ids = (
list(fill_ids[:-1]) if len(fill_ids) > 1 else list(fill_ids)
self.fill_ids = array("q", fill_ids)
self.origin_input_ids = array(
"q", fill_ids[:-1] if len(fill_ids) > 1 else fill_ids
)
self.output_ids = [fill_ids[-1]] if len(fill_ids) > 1 else []
self.output_ids = array("q", [fill_ids[-1]] if len(fill_ids) > 1 else [])
self.req_pool_idx = req_pool_idx
self.cache_protected_len = cache_protected_len
self.last_node = last_node
@@ -99,7 +100,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
"""Insert a prefix into the tree so future requests can match it."""
cache.insert(
InsertParams(
key=RadixKey(prefix_ids),
key=RadixKey(array("q", prefix_ids)),
value=torch.tensor(prefix_values, dtype=torch.int64),
)
)
@@ -119,7 +120,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
self._populate_prefix(cache, prefix, prefix_vals)
# Match prefix (simulates _match_prefix_and_lock in pop_preallocated)
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(prefix)))
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", prefix))))
matched_node = result.last_device_node
prefix_len = len(result.device_indices)
self.assertEqual(prefix_len, 3)
@@ -164,7 +165,9 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
# No prefix in tree -- match returns root
full_ids = [10, 20, 30]
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(full_ids)))
result = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", full_ids)))
)
matched_node = result.last_device_node
self.assertEqual(len(result.device_indices), 0) # no match
# matched_node is root
@@ -212,7 +215,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
self._populate_prefix(cache, prefix, prefix_vals)
# Match and lock
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(prefix)))
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", prefix))))
matched_node = result.last_device_node
prefix_len = len(result.device_indices)
@@ -256,7 +259,9 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
# No prefix in tree -- match returns root (simulates _match_prefix_and_lock)
full_ids = [10, 20, 30]
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(full_ids)))
result = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", full_ids)))
)
matched_node = result.last_device_node
self.assertIs(matched_node, cache.root_node)
@@ -356,7 +361,9 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
self._populate_prefix(cache, prefix, prefix_vals)
for iteration in range(5):
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(prefix)))
result = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", prefix)))
)
matched_node = result.last_device_node
prefix_len = len(result.device_indices)
test/registered/unit/mem_cache/test_mamba_unittest.py
+34 -23
View File
@@ -1,4 +1,5 @@
import unittest
from array import array
import torch
@@ -116,7 +117,7 @@ class TestMamba(unittest.TestCase):
req = Req(
rid=0,
origin_input_text="",
origin_input_ids=[],
origin_input_ids=array("q"),
sampling_params=sampling_params,
)
@@ -158,7 +159,7 @@ class TestMamba(unittest.TestCase):
print(
f"req1: inserting, req1_token_ids: {req1_token_ids}, req1_kv_indices: {req1_kv_indices}"
)
key = RadixKey(req1_token_ids)
key = RadixKey(array("q", req1_token_ids))
result = tree.insert(
InsertParams(
key=key,
@@ -176,7 +177,7 @@ class TestMamba(unittest.TestCase):
print(
f"req2: inserting, req2_token_ids: {req2_token_ids}, req2_kv_indices: {req2_kv_indices}"
)
key = RadixKey(req2_token_ids)
key = RadixKey(array("q", req2_token_ids))
result = tree.insert(
InsertParams(
key=key,
@@ -195,7 +196,7 @@ class TestMamba(unittest.TestCase):
print(
f"req3: inserting, req3_token_ids: {req3_token_ids}, req3_kv_indices: {req3_kv_indices}"
)
key = RadixKey(req3_token_ids)
key = RadixKey(array("q", req3_token_ids))
result = tree.insert(
InsertParams(
key=key,
@@ -213,7 +214,7 @@ class TestMamba(unittest.TestCase):
print(
f"req4: inserting, req4_token_ids: {req4_token_ids}, req4_kv_indices: {req4_kv_indices}"
)
key = RadixKey(req4_token_ids)
key = RadixKey(array("q", req4_token_ids))
result = tree.insert(
InsertParams(
key=key,
@@ -244,7 +245,9 @@ class TestMamba(unittest.TestCase):
tree.pretty_print()
req5_token_ids = [1, 2, 3, 4, 5]
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req5_token_ids)))
result = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", 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}"
@@ -252,7 +255,9 @@ class TestMamba(unittest.TestCase):
assert len(kv_indices) == 0
req6_token_ids = [1, 2, 3, 4, 5, 60, 70]
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req6_token_ids)))
result = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", 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}"
@@ -261,7 +266,9 @@ class TestMamba(unittest.TestCase):
assert len(last_node.key) == 2
req7_token_ids = [1, 2, 3, 4, 5, 6, 7]
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req7_token_ids)))
result = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", 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}"
@@ -278,7 +285,9 @@ class TestMamba(unittest.TestCase):
tree.pretty_print()
req8_token_ids = [1, 2, 3, 4, 5, 60, 70]
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req8_token_ids)))
result = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", 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}"
@@ -289,7 +298,9 @@ class TestMamba(unittest.TestCase):
req9_token_ids = [1, 2, 3, 4, 5, 6, 7]
req9 = make_dummy_req()
result = tree.match_prefix(
MatchPrefixParams(key=RadixKey(req9_token_ids), req=req9, cow_mamba=True)
MatchPrefixParams(
key=RadixKey(array("q", 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
@@ -315,7 +326,7 @@ class TestMamba(unittest.TestCase):
stored_hashes = []
req1 = make_dummy_req()
key1 = RadixKey([1, 2, 3])
key1 = RadixKey(array("q", [1, 2, 3]))
tree.insert(
InsertParams(
key=key1,
@@ -330,7 +341,7 @@ class TestMamba(unittest.TestCase):
stored_hashes.extend(e.block_hashes[0] for e in stored_events)
req2 = make_dummy_req()
key2 = RadixKey([1, 2, 3, 4, 5])
key2 = RadixKey(array("q", [1, 2, 3, 4, 5]))
tree.insert(
InsertParams(
key=key2,
@@ -367,7 +378,7 @@ class TestMamba(unittest.TestCase):
tree.take_events() # Clear the reset event.
req1 = make_dummy_req()
key1 = RadixKey([1, 2, 3, 4])
key1 = RadixKey(array("q", [1, 2, 3, 4]))
tree.insert(
InsertParams(
key=key1,
@@ -382,7 +393,7 @@ class TestMamba(unittest.TestCase):
split_parent_hash = first_insert_events[1].block_hashes[0]
req2 = make_dummy_req()
key2 = RadixKey([1, 2, 5, 6])
key2 = RadixKey(array("q", [1, 2, 5, 6]))
tree.insert(
InsertParams(
key=key2,
@@ -394,7 +405,7 @@ class TestMamba(unittest.TestCase):
e for e in tree.take_events() if isinstance(e, BlockStored)
]
self.assertEqual(len(second_insert_events), 2)
self.assertEqual(second_insert_events[0].token_ids, [5])
self.assertEqual(list(second_insert_events[0].token_ids), [5])
self.assertEqual(second_insert_events[0].parent_block_hash, split_parent_hash)
def _setup_tree_and_allocator(self, enable_kv_cache_events=False):
@@ -478,7 +489,7 @@ class TestMamba(unittest.TestCase):
req = Req(
rid=0,
origin_input_text="",
origin_input_ids=[],
origin_input_ids=array("q"),
sampling_params=sampling_params,
)
req_to_token_pool.alloc([req])
@@ -492,9 +503,9 @@ class TestMamba(unittest.TestCase):
parent = TreeNode()
deleted = TreeNode()
root.key = RadixKey([])
parent.key = RadixKey([1])
deleted.key = RadixKey([2])
root.key = RadixKey(array("q", []))
parent.key = RadixKey(array("q", [1]))
deleted.key = RadixKey(array("q", [2]))
parent.parent = root
deleted.parent = parent
parent.value = torch.tensor([1], dtype=torch.int64)
@@ -668,7 +679,7 @@ class TestMamba(unittest.TestCase):
# Step 1: Insert [1,2,3] to create first node
req1 = make_dummy_req()
key1 = RadixKey([1, 2, 3])
key1 = RadixKey(array("q", [1, 2, 3]))
tree.insert(
InsertParams(
key=key1,
@@ -681,7 +692,7 @@ class TestMamba(unittest.TestCase):
# Step 2: Insert [1,2,3,4,5,6,7] with prev_prefix_len=0 (free all matched)
# Creates tree: [1,2,3] -> [4,5,6,7]
req2 = make_dummy_req()
key2 = RadixKey([1, 2, 3, 4, 5, 6, 7])
key2 = RadixKey(array("q", [1, 2, 3, 4, 5, 6, 7]))
result = tree.insert(
InsertParams(
key=key2,
@@ -699,7 +710,7 @@ class TestMamba(unittest.TestCase):
# Matched prefix = 7 (across two nodes: [1,2,3] len=3, [4,5,6,7] len=4)
# Protected [0..1], freed [2..6] = 5 slots, new [7] = 1 slot stored
req3 = make_dummy_req()
key3 = RadixKey([1, 2, 3, 4, 5, 6, 7, 8])
key3 = RadixKey(array("q", [1, 2, 3, 4, 5, 6, 7, 8]))
result = tree.insert(
InsertParams(
key=key3,
@@ -716,7 +727,7 @@ class TestMamba(unittest.TestCase):
# Step 4: Insert [1,2,3,4,5,6,7,8,9] with prev_prefix_len=8 (covers all matched)
# Matched prefix = 8, prev_prefix_len=8 => nothing freed
req4 = make_dummy_req()
key4 = RadixKey([1, 2, 3, 4, 5, 6, 7, 8, 9])
key4 = RadixKey(array("q", [1, 2, 3, 4, 5, 6, 7, 8, 9]))
result = tree.insert(
InsertParams(
key=key4,
test/registered/unit/mem_cache/test_radix_cache_slru_accuracy.py
+7 -4
View File
@@ -1,4 +1,5 @@
import unittest
from array import array
import torch
@@ -63,7 +64,7 @@ class TestSLRUAccuracy(unittest.TestCase):
"""Test that SLRU eviction mechanism works correctly"""
# Insert one key-value three times (high frequency access)
frequent_key = RadixKey([1, 2]) # High hit rate, should be retained
frequent_key = RadixKey(array("q", [1, 2])) # High hit rate, should be retained
frequent_val = torch.tensor([10, 20], dtype=torch.int64)
# Insert the frequent key multiple times to increase its hit count
@@ -71,7 +72,9 @@ class TestSLRUAccuracy(unittest.TestCase):
self.cache.insert(InsertParams(key=frequent_key, value=frequent_val))
# Insert first low-frequency key-value pair that should be evicted
first_low_freq_key = RadixKey([5, 6]) # Low hit rate, should be evicted
first_low_freq_key = RadixKey(
array("q", [5, 6])
) # Low hit rate, should be evicted
first_low_freq_val = torch.tensor([50, 60], dtype=torch.int64)
self.cache.insert(
@@ -81,14 +84,14 @@ class TestSLRUAccuracy(unittest.TestCase):
# Insert other key-values once each (low frequency access) - fill up the cache
other_keys = []
for i in range(4): # Reduce the number to fit in our smaller cache
key = RadixKey([i + 10]) # Unique keys for low-frequency items
key = RadixKey(array("q", [i + 10])) # Unique keys for low-frequency items
val = torch.tensor([i + 100], dtype=torch.int64)
self.cache.insert(InsertParams(key=key, value=val))
other_keys.append(key)
# Now insert more items to trigger evictions
for i in range(6, 10): # Add more items to definitely exceed capacity
key = RadixKey([i * 2]) # Different pattern to avoid conflicts
key = RadixKey(array("q", [i * 2])) # Different pattern to avoid conflicts
val = torch.tensor([i * 200], dtype=torch.int64)
self.cache.insert(InsertParams(key=key, value=val))
test/registered/unit/mem_cache/test_radix_cache_unit.py
+85 -53
View File
@@ -28,6 +28,7 @@ import random
import time
import unittest
import unittest.mock
from array import array
import torch
@@ -50,30 +51,30 @@ class TestRadixKey(unittest.TestCase):
def test_init_basic(self):
"""Test basic initialization of RadixKey."""
token_ids = [1, 2, 3, 4]
key = RadixKey(token_ids)
self.assertEqual(key.token_ids, token_ids)
key = RadixKey(array("q", token_ids))
self.assertEqual(list(key.token_ids), token_ids)
self.assertIsNone(key.extra_key)
def test_init_with_extra_key(self):
"""Test initialization with extra_key."""
token_ids = [1, 2, 3]
extra_key = "test_key"
key = RadixKey(token_ids, extra_key)
self.assertEqual(key.token_ids, token_ids)
key = RadixKey(array("q", token_ids), extra_key)
self.assertEqual(list(key.token_ids), token_ids)
self.assertEqual(key.extra_key, extra_key)
def test_len(self):
"""Test __len__ method."""
key = RadixKey([1, 2, 3])
key = RadixKey(array("q", [1, 2, 3]))
self.assertEqual(len(key), 3)
empty_key = RadixKey([])
empty_key = RadixKey(array("q", []))
self.assertEqual(len(empty_key), 0)
def test_iter(self):
"""Test __iter__ method."""
token_ids = [1, 2, 3, 4]
key = RadixKey(token_ids)
key = RadixKey(array("q", token_ids))
self.assertEqual(list(key), token_ids)
def test_len_and_iter(self):
@@ -86,7 +87,7 @@ class TestRadixKey(unittest.TestCase):
for tokens, expected in test_cases:
with self.subTest(tokens=tokens):
key = RadixKey(tokens)
key = RadixKey(array("q", tokens))
self.assertEqual(len(key), expected)
self.assertEqual(list(key), tokens)
@@ -100,34 +101,34 @@ class TestRadixKey(unittest.TestCase):
for tokens, index, expected in test_cases:
with self.subTest(tokens=tokens, index=index):
key = RadixKey(tokens)
key = RadixKey(array("q", tokens))
result = key[index]
self.assertIsInstance(result, RadixKey)
self.assertEqual(result.token_ids, expected)
self.assertEqual(list(result.token_ids), expected)
def test_getitem_slice(self):
"""Test __getitem__ with slice and edge cases."""
key = RadixKey([1, 2, 3, 4, 5], "extra")
key = RadixKey(array("q", [1, 2, 3, 4, 5]), "extra")
# Basic slice
sliced = key[1:4]
self.assertIsInstance(sliced, RadixKey)
self.assertEqual(sliced.token_ids, [2, 3, 4])
self.assertEqual(list(sliced.token_ids), [2, 3, 4])
self.assertEqual(sliced.extra_key, "extra")
# Edge cases
self.assertEqual(key[2:2].token_ids, []) # Empty slice
self.assertEqual(key[:].token_ids, [1, 2, 3, 4, 5]) # Full slice
self.assertEqual(list(key[2:2].token_ids), []) # Empty slice
self.assertEqual(list(key[:].token_ids), [1, 2, 3, 4, 5]) # Full slice
def test_getitem_invalid_index(self):
"""Test __getitem__ with invalid indices."""
key = RadixKey([1, 2, 3])
key = RadixKey(array("q", [1, 2, 3]))
with self.assertRaises(IndexError):
_ = key[10] # Out of bounds
def test_repr(self):
"""Test __repr__ method."""
key = RadixKey([1, 2, 3], "test")
key = RadixKey(array("q", [1, 2, 3]), "test")
repr_str = repr(key)
self.assertIn("RadixKey", repr_str)
self.assertIn("extra_key='test'", repr_str)
@@ -136,7 +137,7 @@ class TestRadixKey(unittest.TestCase):
def test_repr_long_token_ids(self):
"""Test __repr__ with long token_ids."""
long_tokens = list(range(15))
key = RadixKey(long_tokens)
key = RadixKey(array("q", long_tokens))
repr_str = repr(key)
self.assertIn("...", repr_str) # Should be truncated
@@ -274,7 +275,7 @@ class TestRadixCache(unittest.TestCase):
# Insert some data
cache.insert(
InsertParams(
key=RadixKey([1, 2, 3]),
key=RadixKey(array("q", [1, 2, 3])),
value=torch.tensor([10, 20, 30], dtype=torch.int64),
)
)
@@ -292,7 +293,7 @@ class TestRadixCache(unittest.TestCase):
with self.subTest(disable_cache=disable_cache):
cache = RadixCache.create_simulated(disable=disable_cache)
key = RadixKey([1, 2, 3])
key = RadixKey(array("q", [1, 2, 3]))
value = torch.tensor([10, 20, 30], dtype=torch.int64)
result = cache.insert(InsertParams(key=key, value=value))
prefix_len = result.prefix_len
@@ -307,12 +308,16 @@ class TestRadixCache(unittest.TestCase):
self.assertEqual(cache.evictable_size(), 3)
# Test match_prefix
result = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3])))
result = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", [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(MatchPrefixParams(key=RadixKey([1, 2])))
result = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", [1, 2])))
)
self.assertEqual(len(result.device_indices), 2)
torch.testing.assert_close(
result.device_indices, torch.tensor([10, 20], dtype=torch.int64)
@@ -322,7 +327,7 @@ class TestRadixCache(unittest.TestCase):
"""Test insert with None value (should use token_ids as list)."""
cache = RadixCache.create_simulated()
key = RadixKey([1, 2, 3])
key = RadixKey(array("q", [1, 2, 3]))
result = cache.insert(InsertParams(key=key, value=None))
prefix_len = result.prefix_len
@@ -338,7 +343,7 @@ class TestRadixCache(unittest.TestCase):
cache.insert(
InsertParams(
key=RadixKey([1, 2, 3]),
key=RadixKey(array("q", [1, 2, 3])),
value=torch.tensor([10, 20, 30], dtype=torch.int64),
)
)
@@ -346,7 +351,8 @@ class TestRadixCache(unittest.TestCase):
cache.insert(
InsertParams(
key=RadixKey([4, 5]), value=torch.tensor([40, 50], dtype=torch.int64)
key=RadixKey(array("q", [4, 5])),
value=torch.tensor([40, 50], dtype=torch.int64),
)
)
self.assertEqual(cache.total_size(), 5)
@@ -366,7 +372,9 @@ class TestRadixCache(unittest.TestCase):
)
# Insert data
cache.insert(InsertParams(key=RadixKey([1, 2, 3, 4, 5]), value=None))
cache.insert(
InsertParams(key=RadixKey(array("q", [1, 2, 3, 4, 5])), value=None)
)
# Take events
events = cache.take_events()
@@ -395,7 +403,7 @@ class TestRadixCache(unittest.TestCase):
# Insert and then evict data
cache.insert(
InsertParams(
key=RadixKey([1, 2, 3]),
key=RadixKey(array("q", [1, 2, 3])),
value=torch.tensor([10, 20, 30], dtype=torch.int64),
)
)
@@ -427,29 +435,35 @@ class TestRadixCache(unittest.TestCase):
# Insert same token sequence with different extra keys
cache.insert(
InsertParams(
key=RadixKey([1, 2, 3], "key1"),
key=RadixKey(array("q", [1, 2, 3]), "key1"),
value=torch.tensor([10, 20, 30], dtype=torch.int64),
)
)
cache.insert(
InsertParams(
key=RadixKey([1, 2, 3], "key2"),
key=RadixKey(array("q", [1, 2, 3]), "key2"),
value=torch.tensor([40, 50, 60], dtype=torch.int64),
)
)
cache.insert(
InsertParams(
key=RadixKey([1, 2, 3], None),
key=RadixKey(array("q", [1, 2, 3]), None),
value=torch.tensor([70, 80, 90], dtype=torch.int64),
)
)
# Keys with different extra_key should not match each other
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)))
result1 = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", [1, 2, 3]), "key1"))
)
result2 = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", [1, 2, 3]), "key2"))
)
result3 = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", [1, 2, 3]), None))
)
result4 = cache.match_prefix(
MatchPrefixParams(key=RadixKey([1, 2, 3], "nonexistent"))
MatchPrefixParams(key=RadixKey(array("q", [1, 2, 3]), "nonexistent"))
)
# Each should match only its own data
@@ -478,13 +492,15 @@ class TestRadixCache(unittest.TestCase):
# Insert sequence
cache.insert(
InsertParams(
key=RadixKey([1, 2, 3]),
key=RadixKey(array("q", [1, 2, 3])),
value=torch.tensor([10, 20, 30], dtype=torch.int64),
)
)
# Get node
result = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3])))
result = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", [1, 2, 3])))
)
node = result.last_device_node
initial_evictable = cache.evictable_size()
@@ -510,12 +526,14 @@ class TestRadixCache(unittest.TestCase):
# Insert sequences
cache.insert(
InsertParams(
key=RadixKey([1, 2]), value=torch.tensor([10, 20], dtype=torch.int64)
key=RadixKey(array("q", [1, 2])),
value=torch.tensor([10, 20], dtype=torch.int64),
)
)
cache.insert(
InsertParams(
key=RadixKey([3, 4]), value=torch.tensor([30, 40], dtype=torch.int64)
key=RadixKey(array("q", [3, 4])),
value=torch.tensor([30, 40], dtype=torch.int64),
)
)
@@ -547,7 +565,7 @@ class TestRadixCache(unittest.TestCase):
cache = RadixCache.create_simulated(page_size=page_size)
tokens = list(range(sequence_length))
key = RadixKey(tokens)
key = RadixKey(array("q", tokens))
cache.insert(
InsertParams(
key=key,
@@ -555,7 +573,9 @@ class TestRadixCache(unittest.TestCase):
)
)
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(tokens)))
result = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", tokens)))
)
self.assertGreater(len(result.device_indices), 0)
# Match length should be page-aligned
@@ -568,7 +588,7 @@ class TestRadixCache(unittest.TestCase):
cache.insert(
InsertParams(
key=RadixKey([1, 2, 3]),
key=RadixKey(array("q", [1, 2, 3])),
value=torch.tensor([10, 20, 30], dtype=torch.int64),
)
)
@@ -585,12 +605,14 @@ class TestRadixCache(unittest.TestCase):
cache.insert(
InsertParams(
key=RadixKey([1, 2]), value=torch.tensor([10, 20], dtype=torch.int64)
key=RadixKey(array("q", [1, 2])),
value=torch.tensor([10, 20], dtype=torch.int64),
)
)
cache.insert(
InsertParams(
key=RadixKey([3, 4]), value=torch.tensor([30, 40], dtype=torch.int64)
key=RadixKey(array("q", [3, 4])),
value=torch.tensor([30, 40], dtype=torch.int64),
)
)
@@ -609,12 +631,12 @@ class TestRadixCache(unittest.TestCase):
# Insert a long sequence that will be split later.
seq1 = [1, 2, 3, 4, 5, 6, 7, 8]
val1 = torch.tensor([x * 10 for x in seq1], dtype=torch.int64)
cache.insert(InsertParams(key=RadixKey(seq1), value=val1))
cache.insert(InsertParams(key=RadixKey(array("q", seq1)), value=val1))
# Insert a diverging branch to create an internal node on the path.
seq2 = [1, 2, 9, 10]
val2 = torch.tensor([x * 10 for x in seq2], dtype=torch.int64)
cache.insert(InsertParams(key=RadixKey(seq2), value=val2))
cache.insert(InsertParams(key=RadixKey(array("q", seq2)), value=val2))
print(cache.pretty_print())
baseline_total = cache.total_size()
@@ -624,24 +646,30 @@ 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(MatchPrefixParams(key=RadixKey(query1)))
result1 = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", 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(MatchPrefixParams(key=RadixKey(seq1)))
result_full = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", 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(MatchPrefixParams(key=RadixKey(query2)))
result2 = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", 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(
MatchPrefixParams(key=RadixKey(seq2))
MatchPrefixParams(key=RadixKey(array("q", seq2)))
)
torch.testing.assert_close(result_branch.device_indices, val2)
@@ -653,7 +681,9 @@ class TestRadixCache(unittest.TestCase):
)
# Insert a sequence
cache.insert(InsertParams(key=RadixKey([1, 2, 3, 4, 5, 6, 7, 8]), value=None))
cache.insert(
InsertParams(key=RadixKey(array("q", [1, 2, 3, 4, 5, 6, 7, 8])), value=None)
)
# Trigger event emission to compute hash_value lazily
cache.take_events()
@@ -679,7 +709,9 @@ class TestRadixCache(unittest.TestCase):
)
# Insert a sequence with repeating token pattern: [1,2,3,4, 1,2,3,4]
cache.insert(InsertParams(key=RadixKey([1, 2, 3, 4, 1, 2, 3, 4]), value=None))
cache.insert(
InsertParams(key=RadixKey(array("q", [1, 2, 3, 4, 1, 2, 3, 4])), value=None)
)
events = cache.take_events()
block_stored_events = [e for e in events if isinstance(e, BlockStored)]
@@ -713,11 +745,11 @@ class TestRadixCache(unittest.TestCase):
)
# Insert a sequence that will cause a split
cache.insert(InsertParams(key=RadixKey([1, 2, 3, 4]), value=None))
cache.insert(InsertParams(key=RadixKey(array("q", [1, 2, 3, 4])), value=None))
cache.take_events() # Clear events and compute hash_value for first node
# Insert a diverging sequence that will cause a split at page boundary
cache.insert(InsertParams(key=RadixKey([1, 2, 5, 6]), value=None))
cache.insert(InsertParams(key=RadixKey(array("q", [1, 2, 5, 6])), value=None))
cache.take_events() # Trigger event emission to compute hash_value
# Find the split node
@@ -754,7 +786,7 @@ class TestRadixCache(unittest.TestCase):
cache: RadixCache = RadixCache.create_simulated()
for key, value in zip(keys, values):
cache.insert(InsertParams(key=RadixKey(key), value=value))
cache.insert(InsertParams(key=RadixKey(array("q", key)), value=value))
del values
test/registered/unit/mem_cache/test_radix_force_miss.py
+8 -5
View File
@@ -11,6 +11,7 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu")
import unittest
import unittest.mock
from array import array
import torch
@@ -27,8 +28,8 @@ from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey
class _StubReq:
def __init__(self, token_ids):
self.origin_input_ids = list(token_ids)
self.output_ids = []
self.origin_input_ids = array("q", token_ids)
self.output_ids = array("q")
self.extra_key = None
self.prefix_indices = None
self.last_node = None
@@ -42,9 +43,9 @@ class _StubReq:
class TestZeroMatchResult(unittest.TestCase):
def test_zero_replaces_indices_and_nodes(self):
tree = RadixCache.create_simulated()
tree.insert(InsertParams(key=RadixKey(token_ids=[1, 2, 3, 4, 5])))
tree.insert(InsertParams(key=RadixKey(token_ids=array("q", [1, 2, 3, 4, 5]))))
match = tree.match_prefix(
MatchPrefixParams(key=RadixKey(token_ids=[1, 2, 3, 9]))
MatchPrefixParams(key=RadixKey(token_ids=array("q", [1, 2, 3, 9])))
)
self.assertGreater(len(match.device_indices), 0)
zeroed = zero_match_result(tree, match)
@@ -76,7 +77,9 @@ class TestMatchPrefixForReqForceMiss(unittest.TestCase):
def test_force_miss_zeros_req_prefix(self):
tree = RadixCache.create_simulated()
tree.insert(
InsertParams(key=RadixKey(token_ids=[10, 11, 12, 13, 14, 15, 16, 17]))
InsertParams(
key=RadixKey(token_ids=array("q", [10, 11, 12, 13, 14, 15, 16, 17]))
)
)
# Sanity: without the flag, the same lookup hits.
test/registered/unit/mem_cache/test_swa_lock_release_lifecycle.py
+2
View File
@@ -12,6 +12,7 @@ Covers:
"""
import unittest
from array import array
import torch
@@ -110,6 +111,7 @@ def _swa_alloc(allocator, need_size):
def _insert_chain(tree, allocator, token_ids):
token_ids = array("q", token_ids)
indices = _swa_alloc(allocator, len(token_ids))
assert indices is not None
tree.insert(InsertParams(key=RadixKey(token_ids), value=indices))
test/registered/unit/mem_cache/test_swa_unittest.py
+31 -20
View File
@@ -1,4 +1,5 @@
import unittest
from array import array
import torch
@@ -113,12 +114,12 @@ def _swa_alloc(allocator, need_size):
def _insert(tree, allocator, token_ids):
indices = _swa_alloc(allocator, len(token_ids))
assert indices is not None
tree.insert(InsertParams(key=RadixKey(token_ids), value=indices))
tree.insert(InsertParams(key=RadixKey(array("q", token_ids)), value=indices))
def _insert_chain(tree, allocator, token_ids):
_insert(tree, allocator, token_ids)
match = tree.match_prefix(MatchPrefixParams(key=RadixKey(token_ids)))
match = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", token_ids))))
return match.last_device_node
@@ -193,7 +194,7 @@ class TestSWA(unittest.TestCase):
e for e in tree.take_events() if isinstance(e, BlockStored)
]
self.assertEqual(len(second_insert_events), 2)
self.assertEqual(second_insert_events[0].token_ids, [5])
self.assertEqual(list(second_insert_events[0].token_ids), [5])
self.assertEqual(second_insert_events[0].parent_block_hash, split_parent_hash)
def test_swa_memory_pool(self):
@@ -313,7 +314,7 @@ class TestSWA(unittest.TestCase):
print(
f"req1: inserting, req1_token_ids: {req1_token_ids}, req1_kv_indices: {req1_kv_indices}"
)
key = RadixKey(req1_token_ids)
key = RadixKey(array("q", req1_token_ids))
result = tree.insert(InsertParams(key=key, value=req1_kv_indices[: len(key)]))
prefix_len = result.prefix_len
print(
@@ -324,7 +325,7 @@ class TestSWA(unittest.TestCase):
print(
f"req2: inserting, req2_token_ids: {req2_token_ids}, req2_kv_indices: {req2_kv_indices}"
)
key = RadixKey(req2_token_ids)
key = RadixKey(array("q", req2_token_ids))
result = tree.insert(InsertParams(key=key, value=req2_kv_indices[: len(key)]))
prefix_len = result.prefix_len
print(
@@ -335,7 +336,7 @@ class TestSWA(unittest.TestCase):
print(
f"req3: inserting, req3_token_ids: {req3_token_ids}, req3_kv_indices: {req3_kv_indices}"
)
key = RadixKey(req3_token_ids)
key = RadixKey(array("q", req3_token_ids))
result = tree.insert(InsertParams(key=key, value=req3_kv_indices[: len(key)]))
prefix_len = result.prefix_len
print(
@@ -346,7 +347,7 @@ class TestSWA(unittest.TestCase):
print(
f"req4: inserting, req4_token_ids: {req4_token_ids}, req4_kv_indices: {req4_kv_indices}"
)
key = RadixKey(req4_token_ids)
key = RadixKey(array("q", req4_token_ids))
result = tree.insert(InsertParams(key=key, value=req4_kv_indices[: len(key)]))
prefix_len = result.prefix_len
print(
@@ -376,7 +377,9 @@ class TestSWA(unittest.TestCase):
tree.pretty_print()
req5_token_ids = [1, 2, 3, 4, 5]
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req5_token_ids)))
result = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", 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}"
@@ -384,7 +387,9 @@ class TestSWA(unittest.TestCase):
self.assertEqual(len(kv_indices), 0)
req6_token_ids = [1, 2, 3, 4, 5, 60, 70]
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req6_token_ids)))
result = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", 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}"
@@ -468,7 +473,7 @@ class TestSWA(unittest.TestCase):
print(
f"req1: inserting, req1_token_ids: {req1_token_ids}, req1_kv_indices: {req1_kv_indices}"
)
key = RadixKey(req1_token_ids)
key = RadixKey(array("q", req1_token_ids))
result = tree.insert(InsertParams(key=key, value=req1_kv_indices[: len(key)]))
prefix_len = result.prefix_len
self.assertEqual(prefix_len, 0)
@@ -480,7 +485,7 @@ class TestSWA(unittest.TestCase):
print(
f"req2: inserting, req2_token_ids: {req2_token_ids}, req2_kv_indices: {req2_kv_indices}"
)
key = RadixKey(req2_token_ids)
key = RadixKey(array("q", req2_token_ids))
result = tree.insert(InsertParams(key=key, value=req2_kv_indices[: len(key)]))
prefix_len = result.prefix_len
self.assertEqual(prefix_len, 2)
@@ -492,7 +497,7 @@ class TestSWA(unittest.TestCase):
print(
f"req3: inserting, req3_token_ids: {req3_token_ids}, req3_kv_indices: {req3_kv_indices}"
)
key = RadixKey(req3_token_ids)
key = RadixKey(array("q", req3_token_ids))
result = tree.insert(InsertParams(key=key, value=req3_kv_indices[: len(key)]))
prefix_len = result.prefix_len
self.assertEqual(prefix_len, 0)
@@ -504,7 +509,7 @@ class TestSWA(unittest.TestCase):
print(
f"req4: inserting, req4_token_ids: {req4_token_ids}, req4_kv_indices: {req4_kv_indices}"
)
key = RadixKey(req4_token_ids)
key = RadixKey(array("q", req4_token_ids))
result = tree.insert(InsertParams(key=key, value=req4_kv_indices[: len(key)]))
prefix_len = result.prefix_len
self.assertEqual(prefix_len, 4)
@@ -553,7 +558,9 @@ class TestSWA(unittest.TestCase):
tree.pretty_print()
req5_token_ids = [1, 2, 3, 4, 5]
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req5_token_ids)))
result = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", 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}"
@@ -561,7 +568,9 @@ 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(MatchPrefixParams(key=RadixKey(req6_token_ids)))
result = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", 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}"
@@ -578,8 +587,8 @@ class TestSWA(unittest.TestCase):
# Case 1: is_insert=True should pass bigram key and use cache_protected_len.
req = _DummyReq()
req.req_pool_idx = 0
req.origin_input_ids = [1, 2, 3, 4, 5, 6]
req.output_ids = []
req.origin_input_ids = array("q", [1, 2, 3, 4, 5, 6])
req.output_ids = array("q")
req._kv_committed_len = len(req.origin_input_ids)
kv_indices = allocator.alloc(req._kv_committed_len)
req_to_token_pool.write(
@@ -613,8 +622,8 @@ class TestSWA(unittest.TestCase):
# even when len(prefix_indices) is intentionally larger.
req2 = _DummyReq()
req2.req_pool_idx = 1
req2.origin_input_ids = [11, 12, 13, 14, 15, 16]
req2.output_ids = []
req2.origin_input_ids = array("q", [11, 12, 13, 14, 15, 16])
req2.output_ids = array("q")
req2._kv_committed_len = len(req2.origin_input_ids)
kv_indices2 = allocator.alloc(req2._kv_committed_len)
req_to_token_pool.write(
@@ -730,7 +739,9 @@ class TestSWASplitLeafOnInsert(CustomTestCase):
with envs.SGLANG_OPT_SWA_SPLIT_LEAF_ON_INSERT.override(True):
inserted_leaf = _insert_chain(tree, allocator, token_ids)
self.assertEqual(len(inserted_leaf.value), 4)
match = tree.match_prefix(MatchPrefixParams(key=RadixKey(token_ids)))
match = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", token_ids)))
)
self.assertEqual(match.device_indices.shape[0], 12)
self.assertIs(match.last_device_node, inserted_leaf)
test/registered/unit/mem_cache/test_unified_radix_cache_bench.py
+10 -9
View File
@@ -12,6 +12,7 @@ import random
import statistics
import time
import unittest
from array import array
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Callable
@@ -335,7 +336,7 @@ def _insert_seq(env, seq):
if env.has_mamba:
req = env.make_req()
mamba_val = req.mamba_pool_idx.unsqueeze(0)
key = RadixKey(seq)
key = RadixKey(array("q", seq))
env.tree.insert(InsertParams(key=key, value=v[: len(key)], mamba_value=mamba_val))
return True
@@ -357,7 +358,7 @@ def _fill_no_evict(env):
if env.has_mamba:
req = env.make_req()
mamba_val = req.mamba_pool_idx.unsqueeze(0)
key = RadixKey(seq)
key = RadixKey(array("q", seq))
env.tree.insert(
InsertParams(key=key, value=v[: len(key)], mamba_value=mamba_val)
)
@@ -505,7 +506,7 @@ def bench_match_prefix(
queries.append([rng.randint(1, 32000)] * rng.randint(50, 300))
def verify_fn(q):
k = RadixKey(q)
k = RadixKey(array("q", q))
r1 = env.tree.match_prefix(MatchPrefixParams(key=k))
r2 = env.tree.match_prefix(MatchPrefixParams(key=k))
assert len(r1.device_indices) == len(r2.device_indices), "match not idempotent"
@@ -514,7 +515,7 @@ def bench_match_prefix(
return bench_api(
"match_prefix",
lambda: queries,
lambda q: env.tree.match_prefix(MatchPrefixParams(key=RadixKey(q))),
lambda q: env.tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", q)))),
min(len(queries) - warmup, num_seqs),
env.avg_tokens,
warmup,
@@ -566,7 +567,7 @@ def bench_lock_unlock(
nodes = []
for seq in env.seqs[: num_seqs // 2]:
r = env.tree.match_prefix(MatchPrefixParams(key=RadixKey(seq)))
r = env.tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
if r.last_device_node != env.tree.root_node:
nodes.append(r.last_device_node)
if not nodes:
@@ -613,7 +614,7 @@ def bench_cache_finished(
# Pre-build Req objects with token IDs filled into req_to_token
req_items: list = []
for seq in env.seqs:
key = RadixKey(seq)
key = RadixKey(array("q", seq))
mr = env.tree.match_prefix(MatchPrefixParams(key=key))
matched_len = len(mr.device_indices)
node = mr.last_device_node
@@ -635,9 +636,9 @@ def bench_cache_finished(
kv_indices = mr.device_indices
req = env.make_req()
req.origin_input_ids = list(seq)
req.output_ids = []
req.fill_ids = list(seq)
req.origin_input_ids = array("q", seq)
req.output_ids = array("q")
req.fill_ids = array("q", seq)
req.last_node = node
req.cache_protected_len = matched_len
req.kv_committed_len = len(seq)
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+113 -92
View File
@@ -1,6 +1,7 @@
"""Unit tests for UnifiedRadixCache"""
import unittest
from array import array
from dataclasses import dataclass
from typing import Optional
from unittest import mock
@@ -272,7 +273,7 @@ class UnifiedRadixCacheSuite:
def _insert(self, tree, allocator, req_to_token_pool, tokens):
"""Insert tokens, attaching mamba data when the config has mamba."""
key = RadixKey(tokens)
key = RadixKey(array("q", tokens))
value = self._alloc(allocator, len(tokens))
params = InsertParams(key=key, value=value[: len(key)])
if self.cfg.has_mamba:
@@ -290,15 +291,17 @@ class UnifiedRadixCacheSuite:
result = self._insert(tree, allocator, req_to_token_pool, seq_b)
self.assertEqual(result.prefix_len, len(seq_a))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_b)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_b))))
self.assertEqual(len(m.device_indices), len(seq_b))
m = tree.match_prefix(
MatchPrefixParams(key=RadixKey(seq_a + self._make_seq(9000, 1)))
MatchPrefixParams(key=RadixKey(array("q", seq_a + self._make_seq(9000, 1))))
)
self.assertEqual(len(m.device_indices), len(seq_a))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(self._make_seq(5000, 2))))
m = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", self._make_seq(5000, 2))))
)
self.assertEqual(len(m.device_indices), 0)
tree.sanity_check()
@@ -317,11 +320,11 @@ class UnifiedRadixCacheSuite:
self.assertEqual(result_b.prefix_len, len(base))
for seq in (branch_a, branch_b):
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
self.assertEqual(len(m.device_indices), len(seq))
m = tree.match_prefix(
MatchPrefixParams(key=RadixKey(base + self._make_seq(999, 1)))
MatchPrefixParams(key=RadixKey(array("q", base + self._make_seq(999, 1))))
)
self.assertEqual(len(m.device_indices), len(base))
tree.sanity_check()
@@ -350,13 +353,13 @@ class UnifiedRadixCacheSuite:
self._insert(tree, allocator, req_to_token_pool, seq_a)
self._insert(tree, allocator, req_to_token_pool, seq_b)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_a)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_a))))
lock_result = tree.inc_lock_ref(m.last_device_node)
result = tree.evict(EvictParams(num_tokens=len(seq_a) + len(seq_b)))
self.assertGreaterEqual(result.num_tokens_evicted, len(seq_b))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_a)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_a))))
self.assertEqual(len(m.device_indices), len(seq_a))
# Unlock -> should now be evictable
@@ -395,7 +398,7 @@ class UnifiedRadixCacheSuite:
if self.cfg.has_mamba:
self.assertEqual(tree.mamba_evictable_size(), 0)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seqs[0])))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seqs[0]))))
self.assertEqual(len(m.device_indices), 0)
tree.sanity_check()
@@ -413,7 +416,7 @@ class UnifiedRadixCacheSuite:
self.assertEqual(allocator.available_size(), initial_avail - len(seq_1p))
# Step 2: insert 2 pages with prev_prefix_len=0 → frees overlap of 1 page
key_2p = RadixKey(seq_2p)
key_2p = RadixKey(array("q", seq_2p))
value_2p = self._alloc(allocator, len(seq_2p))
params = InsertParams(
key=key_2p,
@@ -432,7 +435,7 @@ class UnifiedRadixCacheSuite:
# Step 3: insert 3 pages with prev_prefix_len=len(seq_2p) → nothing freed
avail_before = allocator.available_size()
key_3p = RadixKey(seq_3p)
key_3p = RadixKey(array("q", seq_3p))
value_3p = self._alloc(allocator, len(seq_3p))
params = InsertParams(
key=key_3p,
@@ -461,11 +464,11 @@ class UnifiedRadixCacheSuite:
self.assertEqual(result.prefix_len, len(base))
for seq in (fork_a, fork_b):
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
self.assertEqual(len(m.device_indices), len(seq))
m = tree.match_prefix(
MatchPrefixParams(key=RadixKey(base + self._make_seq(999, 1)))
MatchPrefixParams(key=RadixKey(array("q", base + self._make_seq(999, 1))))
)
self.assertEqual(len(m.device_indices), len(base))
tree.sanity_check()
@@ -477,8 +480,8 @@ class UnifiedRadixCacheSuite:
req = self._make_req(req_to_token_pool)
input_ids = self._make_seq(1, 3)
output_ids = self._make_seq(2000, 1)
req.origin_input_ids = input_ids
req.output_ids = output_ids
req.origin_input_ids = array("q", input_ids)
req.output_ids = array("q", output_ids)
kv_len = len(input_ids) + len(output_ids)
kv_indices = self._alloc(allocator, kv_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
@@ -487,7 +490,7 @@ class UnifiedRadixCacheSuite:
req.cache_protected_len = 0
req.swa_uuid_for_lock = None
req.extra_key = None
req.fill_ids = input_ids + output_ids
req.fill_ids = array("q", input_ids + output_ids)
if self.cfg.has_mamba:
req.mamba_last_track_seqlen = kv_len
@@ -495,7 +498,9 @@ class UnifiedRadixCacheSuite:
all_ids = input_ids + output_ids
aligned_len = (len(all_ids) // ps) * ps
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(all_ids[:aligned_len])))
m = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", all_ids[:aligned_len])))
)
self.assertEqual(len(m.device_indices), aligned_len)
tree.sanity_check()
@@ -506,9 +511,9 @@ class UnifiedRadixCacheSuite:
req = self._make_req(req_to_token_pool)
prompt_ids = self._make_seq(1, 3)
output_ids = self._make_seq(2000, 7)
req.origin_input_ids = prompt_ids
req.output_ids = output_ids
req.fill_ids = prompt_ids + output_ids
req.origin_input_ids = array("q", prompt_ids)
req.output_ids = array("q", output_ids)
req.fill_ids = array("q", prompt_ids + output_ids)
kv_len = len(req.fill_ids)
kv_indices = self._alloc(allocator, kv_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
@@ -538,7 +543,9 @@ class UnifiedRadixCacheSuite:
prompt_aligned = (len(prompt_ids) // ps) * ps
# Thinking+answer must not be reachable past the prompt.
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(prompt_ids + output_ids)))
m = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", prompt_ids + output_ids)))
)
self.assertEqual(len(m.device_indices), prompt_aligned)
# Only prompt-aligned pages remain owned by the tree.
self.assertEqual(
@@ -550,8 +557,8 @@ class UnifiedRadixCacheSuite:
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
req = self._make_req(req_to_token_pool)
tokens = self._make_seq(1, 2)
req.origin_input_ids = tokens
req.output_ids = []
req.origin_input_ids = array("q", tokens)
req.output_ids = array("q")
kv_len = len(tokens)
kv_indices = self._alloc(allocator, kv_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
@@ -560,13 +567,13 @@ class UnifiedRadixCacheSuite:
req.cache_protected_len = 0
req.swa_uuid_for_lock = None
req.extra_key = None
req.fill_ids = tokens
req.fill_ids = array("q", tokens)
avail_before = allocator.available_size()
tree.cache_finished_req(req, is_insert=False)
self.assertEqual(allocator.available_size(), avail_before + kv_len)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens))))
self.assertEqual(len(m.device_indices), 0)
tree.sanity_check()
@@ -575,9 +582,9 @@ class UnifiedRadixCacheSuite:
req = self._make_req(req_to_token_pool)
tokens = self._make_seq(1, 3)
req.origin_input_ids = tokens
req.output_ids = []
req.fill_ids = tokens[:]
req.origin_input_ids = array("q", tokens)
req.output_ids = array("q")
req.fill_ids = array("q", tokens)
kv_len = len(tokens)
kv_indices = self._alloc(allocator, kv_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
@@ -628,11 +635,11 @@ class UnifiedRadixCacheSuite:
for suffix_start in [100, 200, 300]:
seq = base + self._make_seq(suffix_start, 2)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
self.assertEqual(len(m.device_indices), len(seq))
m = tree.match_prefix(
MatchPrefixParams(key=RadixKey(base + self._make_seq(999, 1)))
MatchPrefixParams(key=RadixKey(array("q", base + self._make_seq(999, 1))))
)
self.assertEqual(len(m.device_indices), len(base))
tree.sanity_check()
@@ -641,7 +648,7 @@ class UnifiedRadixCacheSuite:
if self.cfg.page_size == 1:
self.skipTest("page_size > 1 only")
tree, _, _ = build_fixture(self.cfg)
key = RadixKey(self._make_seq(1, 1))
key = RadixKey(array("q", self._make_seq(1, 1)))
child_key = key.child_key(tree.page_size)
self.assertIsInstance(child_key, tuple)
@@ -656,11 +663,13 @@ class UnifiedRadixCacheSuite:
# Tree truncates unaligned tail internally, so it matches the seq prefix.
unaligned = seq + list(range(9000, 9000 + ps - 1))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(unaligned)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", unaligned))))
self.assertEqual(len(m.device_indices), len(seq))
# Below-page-size key aligns to 0 -> no match.
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq[: ps - 1])))
m = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", seq[: ps - 1])))
)
self.assertEqual(len(m.device_indices), 0)
tree.sanity_check()
@@ -679,12 +688,12 @@ class UnifiedRadixCacheSuite:
# Mismatch in second page → only first page matches
bad_page2 = seq[:ps] + [9999] * ps
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(bad_page2)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", bad_page2))))
self.assertEqual(len(m.device_indices), ps)
# Mismatch in first page → 0 match
bad_page1 = [9999] + seq[1:]
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(bad_page1)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", bad_page1))))
self.assertEqual(len(m.device_indices), 0)
tree.sanity_check()
@@ -699,8 +708,8 @@ class UnifiedRadixCacheSuite:
tail_extra = ps // 2
input_ids = self._make_seq(1, 1) + list(range(8000, 8000 + tail_extra))
req = self._make_req(req_to_token_pool)
req.origin_input_ids = input_ids
req.output_ids = []
req.origin_input_ids = array("q", input_ids)
req.output_ids = array("q")
kv_len = len(input_ids)
kv_indices = self._alloc(allocator, kv_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
@@ -709,7 +718,7 @@ class UnifiedRadixCacheSuite:
req.cache_protected_len = 0
req.swa_uuid_for_lock = None
req.extra_key = None
req.fill_ids = input_ids
req.fill_ids = array("q", input_ids)
if self.cfg.has_mamba:
req.mamba_last_track_seqlen = kv_len
@@ -718,7 +727,7 @@ class UnifiedRadixCacheSuite:
self.assertEqual(allocator.available_size(), avail_before + tail_extra)
aligned = input_ids[: (len(input_ids) // ps) * ps]
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(aligned)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", aligned))))
self.assertEqual(len(m.device_indices), len(aligned))
tree.sanity_check()
@@ -749,7 +758,7 @@ class UnifiedRadixCacheSuite:
tree.evict(EvictParams(num_tokens=0, mamba_num=10))
self.assertEqual(tree.mamba_evictable_size(), 0)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_long)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_long))))
self.assertEqual(len(m.device_indices), 0)
tree.sanity_check()
@@ -788,7 +797,7 @@ class UnifiedRadixCacheSuite:
req2 = self._make_req(req_to_token_pool)
m = tree.match_prefix(
MatchPrefixParams(key=RadixKey(seq), cow_mamba=True, req=req2)
MatchPrefixParams(key=RadixKey(array("q", seq)), cow_mamba=True, req=req2)
)
self.assertEqual(len(m.device_indices), len(seq))
self.assertIsNotNone(req2.mamba_pool_idx)
@@ -809,7 +818,7 @@ class UnifiedRadixCacheSuite:
seq = self._make_seq(1, 3)
self._insert(tree, allocator, req_to_token_pool, seq)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
self.assertEqual(len(m.device_indices), len(seq))
tree.sanity_check()
@@ -866,13 +875,13 @@ class UnifiedRadixCacheSuite:
self._insert(tree, allocator, req_to_token_pool, seq_a)
self._insert(tree, allocator, req_to_token_pool, seq_b)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_a)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_a))))
lock_result = tree.inc_lock_ref(m.last_device_node)
result = tree.evict(EvictParams(num_tokens=len(seq_a) + len(seq_b)))
self.assertGreaterEqual(result.num_tokens_evicted, len(seq_b))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_a)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_a))))
self.assertEqual(len(m.device_indices), len(seq_a))
tree.dec_lock_ref(
@@ -886,8 +895,8 @@ class UnifiedRadixCacheSuite:
parent = UnifiedTreeNode(self.cfg.components)
deleted = UnifiedTreeNode(self.cfg.components)
parent.key = RadixKey(self._make_seq(1, 1))
deleted.key = RadixKey(self._make_seq(1000, 1))
parent.key = RadixKey(array("q", self._make_seq(1, 1)))
deleted.key = RadixKey(array("q", self._make_seq(1000, 1)))
parent.parent = tree.root_node
deleted.parent = parent
parent.component_data[ComponentType.FULL].value = torch.arange(
@@ -924,15 +933,15 @@ class UnifiedRadixCacheSuite:
node_count_before = count_nodes(tree.root_node)
self.assertEqual(node_count_before, 2)
tree._match_prefix_helper(RadixKey([1, 2]))
tree._match_prefix_helper(RadixKey(array("q", [1, 2])))
(
value,
best_match_node,
best_match_device_node,
best_value_len,
) = tree._match_prefix_helper(RadixKey([1, 2, 3, 4]))
) = tree._match_prefix_helper(RadixKey(array("q", [1, 2, 3, 4])))
self.assertEqual(best_value_len, 2)
self.assertEqual(best_match_node.key.token_ids, [3, 4])
self.assertEqual(list(best_match_node.key.token_ids), [3, 4])
self.assertIs(best_match_device_node, best_match_node)
node_count_after_regular = count_nodes(tree.root_node)
self.assertEqual(node_count_after_regular, node_count_before + 2)
@@ -942,9 +951,9 @@ class UnifiedRadixCacheSuite:
best_match_node,
best_match_device_node,
best_value_len,
) = tree._match_prefix_helper_readonly(RadixKey([1, 2, 3]))
) = tree._match_prefix_helper_readonly(RadixKey(array("q", [1, 2, 3])))
self.assertEqual(best_value_len, 1)
self.assertEqual(best_match_node.key.token_ids, [1, 2])
self.assertEqual(list(best_match_node.key.token_ids), [1, 2])
self.assertIs(best_match_device_node, best_match_node)
node_count_after_readonly = count_nodes(tree.root_node)
self.assertEqual(node_count_after_readonly, node_count_after_regular)
@@ -971,7 +980,7 @@ class UnifiedRadixCacheSuite:
seq = self._make_seq(1, 2)
self._insert(tree, allocator, req_to_token_pool, seq)
match = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq)))
match = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
node = match.last_device_node
full_cd = node.component_data[ComponentType.FULL]
aux_cd = node.component_data[aux]
@@ -1040,7 +1049,7 @@ class UnifiedRadixCacheSuite:
self._insert(tree, allocator, req_to_token_pool, leaf)
# Lock the base node to prevent it from being evicted
m_base = tree.match_prefix(MatchPrefixParams(key=RadixKey(base)))
m_base = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", base))))
lock_result = tree.inc_lock_ref(m_base.last_device_node)
# Evict the leaf — parent (base) should become D-leaf after unlock
@@ -1089,14 +1098,14 @@ class UnifiedRadixCacheSuite:
self._insert(tree, allocator, req_to_token_pool, seq_new)
# Touch seq_new to make it MRU
tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_new)))
tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_new))))
# Evict just enough for one sequence
tree.evict(EvictParams(num_tokens=len(seq_old)))
# seq_old should be gone (LRU), seq_new should remain
m_old = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_old)))
m_new = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_new)))
m_old = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_old))))
m_new = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_new))))
self.assertEqual(len(m_old.device_indices), 0)
self.assertEqual(len(m_new.device_indices), len(seq_new))
tree.sanity_check()
@@ -1136,13 +1145,13 @@ class UnifiedRadixCacheSuite:
self._insert(tree, allocator, req_to_token_pool, branch_b)
# Lock branch_b
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(branch_b)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", branch_b))))
lr = tree.inc_lock_ref(m.last_device_node)
# Evict — branch_a should go, base + branch_b stay
tree.evict(EvictParams(num_tokens=len(branch_a)))
m_b = tree.match_prefix(MatchPrefixParams(key=RadixKey(branch_b)))
m_b = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", branch_b))))
self.assertEqual(len(m_b.device_indices), len(branch_b))
tree.dec_lock_ref(
@@ -1173,7 +1182,7 @@ class UnifiedRadixCacheSuite:
self._insert(tree, allocator, req_to_token_pool, seq_b)
# Lock seq_a
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_a)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_a))))
lr = tree.inc_lock_ref(m.last_device_node)
# Try to evict everything
@@ -1181,7 +1190,7 @@ class UnifiedRadixCacheSuite:
result = tree.evict(EvictParams(num_tokens=total))
# seq_a should still be matchable (protected)
m2 = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_a)))
m2 = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_a))))
self.assertEqual(len(m2.device_indices), len(seq_a))
tree.dec_lock_ref(
@@ -1220,7 +1229,7 @@ class UnifiedRadixCacheSuite:
# Re-insert
seq_b = self._make_seq(500, 2)
self._insert(tree, allocator, req_to_token_pool, seq_b)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_b)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_b))))
self.assertEqual(len(m.device_indices), len(seq_b))
tree.sanity_check()
@@ -1247,7 +1256,7 @@ class UnifiedRadixCacheSuite:
self._insert(tree, allocator, req_to_token_pool, s)
# Lock some, evict some, unlock
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seqs[0])))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seqs[0]))))
lr = tree.inc_lock_ref(m.last_device_node)
tree.evict(EvictParams(num_tokens=len(seqs[1])))
@@ -1409,7 +1418,7 @@ class UnifiedRadixCacheSuite:
self._insert(tree, allocator, req_to_token_pool, seq)
# Find the leaf node
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
node = m.last_device_node
self.assertIsNot(node, tree.root_node)
@@ -1435,7 +1444,7 @@ class UnifiedRadixCacheSuite:
seq = self._make_seq(1, 2)
self._insert(tree, allocator, req_to_token_pool, seq)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
node = m.last_device_node
self._backup_node(tree, node)
@@ -1469,7 +1478,7 @@ class UnifiedRadixCacheSuite:
self._backup_tree(tree)
# Lock leaf so only base can be evicted
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(leaf)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", leaf))))
lr = tree.inc_lock_ref(m.last_device_node)
# Evict base (inner node won't be evicted while child is locked)
@@ -1479,7 +1488,7 @@ class UnifiedRadixCacheSuite:
m.last_device_node,
DecLockRefParams(swa_uuid_for_lock=getattr(lr, "swa_uuid_for_lock", None)),
)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(leaf)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", leaf))))
self.assertGreaterEqual(len(m.device_indices), len(base))
tree.sanity_check()
@@ -1495,7 +1504,7 @@ class UnifiedRadixCacheSuite:
query = expected_prefix + self._make_seq(9000, 1)
self._insert(tree, allocator, req_to_token_pool, seq)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
node = m.last_device_node
self._backup_node(tree, node)
@@ -1503,7 +1512,7 @@ class UnifiedRadixCacheSuite:
self.assertTrue(node.evicted)
self.assertTrue(node.backuped)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(query)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", query))))
self.assertEqual(len(m.device_indices), 0)
self.assertIs(m.last_device_node, tree.root_node)
@@ -1512,8 +1521,8 @@ class UnifiedRadixCacheSuite:
self.assertIsNot(split_parent, tree.root_node)
self.assertTrue(split_parent.evicted)
self.assertTrue(split_parent.backuped)
self.assertEqual(split_parent.key.token_ids, expected_prefix)
self.assertEqual(node.key.token_ids, expected_suffix)
self.assertEqual(list(split_parent.key.token_ids), expected_prefix)
self.assertEqual(list(node.key.token_ids), expected_suffix)
if self.cfg.has_mamba:
self.assertEqual(m.host_hit_length, 0)
@@ -1536,7 +1545,7 @@ class UnifiedRadixCacheSuite:
self._insert(tree, allocator, req_to_token_pool, s)
for i in range(2):
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seqs[i])))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seqs[i]))))
self._backup_node(tree, m.last_device_node)
# Evict one backed-up node
@@ -1555,7 +1564,7 @@ class UnifiedRadixCacheSuite:
seq = self._make_seq(1, 2)
self._insert(tree, allocator, req_to_token_pool, seq)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
node = m.last_device_node
self._backup_node(tree, node)
@@ -1580,7 +1589,7 @@ class UnifiedRadixCacheSuite:
base = self._make_seq(1, 2)
self._insert(tree, allocator, req_to_token_pool, base)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(base)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", base))))
node = m.last_device_node
original_device_indices = m.device_indices.clone()
self._fill_full_kv(allocator, original_device_indices, marker=3)
@@ -1662,7 +1671,7 @@ class UnifiedRadixCacheSuite:
seq = self._make_seq(1, 2)
self._insert(tree, allocator, req_to_token_pool, seq)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
node = m.last_device_node
for aux in aux_types:
@@ -1688,7 +1697,7 @@ class UnifiedRadixCacheSuite:
for i in range(num_pages):
seq = seq + self._make_seq(1000 * (i + 1), 1)
self._insert(tree, allocator, req_to_token_pool, seq)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
chain: list = []
cur = m.last_device_node
while cur is not tree.root_node:
@@ -1738,7 +1747,7 @@ class UnifiedRadixCacheSuite:
seq = self._make_seq(1, (min_tokens + ps - 1) // ps)
self._insert(tree, allocator, req_to_token_pool, seq)
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq)))
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
self.assertEqual(len(result.device_indices), len(seq))
self.assertIs(result.best_match_node, result.last_device_node)
@@ -1760,7 +1769,7 @@ class UnifiedRadixCacheSuite:
tree.evict(EvictParams(num_tokens=len(leaf.key)))
self.assertTrue(leaf.evicted)
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens)))
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens))))
self.assertIs(result.best_match_node, leaf)
self.assertIs(result.last_device_node, parent)
@@ -1782,7 +1791,7 @@ class UnifiedRadixCacheSuite:
tree.evict(EvictParams(num_tokens=len(leaf.key)))
self.assertTrue(leaf.evicted)
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens)))
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens))))
self.assertIs(result.best_match_node, leaf)
self.assertIs(result.last_device_node, parent)
@@ -1797,12 +1806,14 @@ class UnifiedRadixCacheSuite:
tokens = self._make_seq(1, chunk_size + 1)
self._insert(tree, allocator, req_to_token_pool, tokens)
leaf = tree.match_prefix(
MatchPrefixParams(key=RadixKey(tokens))
MatchPrefixParams(key=RadixKey(array("q", tokens)))
).last_device_node
mamba_cd = leaf.component_data[ComponentType.MAMBA]
mamba_cd.value = None
no_hicache = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens)))
no_hicache = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", tokens)))
)
self.assertIs(no_hicache.best_match_node, tree.root_node)
self.assertIs(no_hicache.last_device_node, tree.root_node)
self.assertEqual(no_hicache.mamba_branching_seqlen, chunk_size)
@@ -1810,11 +1821,13 @@ class UnifiedRadixCacheSuite:
tree_h, allocator_h, req_to_token_pool_h = self._build_hicache_fixture()
self._insert(tree_h, allocator_h, req_to_token_pool_h, tokens)
leaf_h = tree_h.match_prefix(
MatchPrefixParams(key=RadixKey(tokens))
MatchPrefixParams(key=RadixKey(array("q", tokens)))
).last_device_node
self._backup_node(tree_h, leaf_h)
tree_h.evict(EvictParams(num_tokens=len(tokens)))
with_hicache = tree_h.match_prefix(MatchPrefixParams(key=RadixKey(tokens)))
with_hicache = tree_h.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", tokens)))
)
self.assertIs(with_hicache.best_match_node, leaf_h)
self.assertIs(with_hicache.last_device_node, tree_h.root_node)
self.assertIsNone(with_hicache.mamba_branching_seqlen)
@@ -1834,7 +1847,9 @@ class UnifiedRadixCacheSuite:
self.assertTrue(leaf.evicted)
req = self._make_req(req_to_token_pool)
match = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens), req=req))
match = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", tokens)), req=req)
)
req.prefix_indices = match.device_indices
req.last_node = match.last_device_node
req.best_match_node = match.best_match_node
@@ -1876,7 +1891,9 @@ class UnifiedRadixCacheSuite:
self._set_aux_host_tombstone(tree, leaf, aux)
req = self._make_req(req_to_token_pool)
match = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens), req=req))
match = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", tokens)), req=req)
)
req.prefix_indices = match.device_indices
req.last_node = match.last_device_node
req.best_match_node = match.best_match_node
@@ -1915,7 +1932,9 @@ class UnifiedRadixCacheSuite:
tree.evict(EvictParams(num_tokens=len(leaf.key)))
req = self._make_req(req_to_token_pool)
match = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens), req=req))
match = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", tokens)), req=req)
)
req.prefix_indices = match.device_indices
req.last_node = match.last_device_node
req.best_match_node = match.best_match_node
@@ -1986,7 +2005,7 @@ class UnifiedRadixCacheSuite:
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
seq = self._make_seq(1, 2)
self._insert(tree, allocator, req_to_token_pool, seq)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
node = m.last_device_node
self._simulate_backup(tree, node)
@@ -2077,7 +2096,9 @@ class UnifiedRadixCacheSuite:
)
result = swa_comp.finalize_match_result(
result=result,
params=MatchPrefixParams(key=RadixKey(self._make_seq(1, 1))),
params=MatchPrefixParams(
key=RadixKey(array("q", self._make_seq(1, 1)))
),
value_chunks=[],
best_value_len=0,
)
@@ -2214,7 +2235,7 @@ class UnifiedRadixCacheSuite:
def test_hicache_swa_match_prefix_picks_best_match_node_above_last_host(self):
tree, _, n, y, x, tokens = self._swa_anchor_setup()
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens)))
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens))))
self.assertIs(result.best_match_node, x)
self.assertIs(result.last_device_node, n.parent)
self.assertIs(result.last_host_node, y)
@@ -2247,7 +2268,7 @@ class UnifiedRadixCacheSuite:
)
result = swa_comp.finalize_match_result(
result=base,
params=MatchPrefixParams(key=RadixKey(self._make_seq(1, 1))),
params=MatchPrefixParams(key=RadixKey(array("q", self._make_seq(1, 1)))),
value_chunks=[],
best_value_len=0,
)
@@ -2365,7 +2386,7 @@ class UnifiedRadixCacheSuite:
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
seq = self._make_seq(1, 2)
self._insert(tree, allocator, req_to_token_pool, seq)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
node = m.last_device_node
cd = node.component_data[ComponentType.MAMBA]
old_mamba = cd.value
@@ -2414,7 +2435,7 @@ class UnifiedRadixCacheSuite:
tree.sanity_check()
for i in range(3):
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seqs[i])))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seqs[i]))))
self._backup_node(tree, m.last_device_node)
# Evict to free some tokens
@@ -2443,7 +2464,7 @@ class UnifiedRadixCacheSuite:
self._insert(tree, allocator, req_to_token_pool, base)
self._insert(tree, allocator, req_to_token_pool, leaf_seq)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(leaf_seq)))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", leaf_seq))))
leaf = m.last_device_node
parent = leaf.parent
self.assertIsNot(parent, tree.root_node)
test/registered/unit/utils/test_common.py
+50
View File
@@ -0,0 +1,50 @@
import unittest
from array import array
import torch
from sglang.srt.utils.common import flatten_arrays_to_int64_tensor
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=5, stage="base-b", runner_config="1-gpu-small")
@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
class TestFlattenArraysToInt64Tensor(CustomTestCase):
"""`flatten_arrays_to_int64_tensor` is invoked by `prepare_for_extend`
to build the per-batch input_ids tensor (pinned, async H2D) from a
list of array.array('q') per-req fill_ids slices. Tests the full
matrix of (device, pin) the production code paths through.
"""
DEVICES = ("cpu", "cuda")
PIN_OPTIONS = (False, True)
def _check(self, parts: list, expected: list[int]) -> None:
for device in self.DEVICES:
for pin in self.PIN_OPTIONS:
with self.subTest(device=device, pin=pin):
out = flatten_arrays_to_int64_tensor(parts, device, pin)
if device == "cuda":
torch.cuda.synchronize()
self.assertEqual(out.dtype, torch.int64)
self.assertEqual(out.device.type, device)
self.assertEqual(out.shape, (len(expected),))
self.assertEqual(out.cpu().tolist(), expected)
def test_single_part(self):
parts = [array("q", [1, 2, 3, 4, 5])]
self._check(parts, [1, 2, 3, 4, 5])
def test_multiple_parts(self):
parts = [
array("q", [10, 20, 30]),
array("q", [100, 200]),
array("q", [1000]),
]
self._check(parts, [10, 20, 30, 100, 200, 1000])
if __name__ == "__main__":
unittest.main()