diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index 13398b430..84ae861f6 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -730,7 +730,7 @@ class HiMambaRadixCache(MambaRadixCache): child_key = self.get_child_key_fn(key) new_node = TreeNode() new_node.parent = parent - new_node.key = key + new_node.key = key.compacted() new_node.value = value.clone() new_node.mamba_value = mamba_value self.full_lru_list.insert_mru(new_node) @@ -933,7 +933,7 @@ class HiMambaRadixCache(MambaRadixCache): new_node.mamba_value = None new_node.full_lock_ref = child.full_lock_ref new_node.mamba_lock_ref = 0 - new_node.key = child.key[:split_len] + new_node.key = child.key[:split_len].compacted() if child.backuped: new_node.host_value = child.host_value[:split_len].clone() @@ -947,7 +947,7 @@ class HiMambaRadixCache(MambaRadixCache): if child.mamba_value is not None: self.mamba_lru_list.remove_node(child) child.parent = new_node - child.key = child.key[split_len:] + child.key = child.key[split_len:].compacted() new_node.parent.children[self.get_child_key_fn(key)] = new_node if child.mamba_value is not None: self.mamba_lru_list.insert_mru(child) @@ -1649,7 +1649,7 @@ class HiMambaRadixCache(MambaRadixCache): if len(key): new_node = TreeNode() new_node.parent = node - new_node.key = key + new_node.key = key.compacted() new_node.value = None new_node.mamba_value = None new_node.host_value = host_value.clone() diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 0ca4a4cb4..f409ea415 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -4087,7 +4087,7 @@ class HiRadixCache(RadixCache): if len(key): new_node = TreeNode(priority=node.priority) new_node.parent = node - new_node.key = key + new_node.key = key.compacted() new_node.value = None new_node.host_value = host_value.clone() new_node.hash_value = hash_value @@ -4179,7 +4179,7 @@ class HiRadixCache(RadixCache): # If child is pinned, new parent inherits a host_ref_counter hold if child.pin_expiry > 0: new_node.host_ref_counter += 1 - new_node.key = child.key[:split_len] + new_node.key = child.key[:split_len].compacted() new_node.hit_count = child.hit_count # split value and host value if exists @@ -4203,7 +4203,7 @@ class HiRadixCache(RadixCache): child.hash_value, split_len, self.page_size ) child.parent = new_node - child.key = child.key[split_len:] + child.key = child.key[split_len:].compacted() new_node.parent.children[self.get_child_key_fn(key)] = new_node return new_node @@ -4330,7 +4330,7 @@ class HiRadixCache(RadixCache): priority=priority, ) new_node.parent = node - new_node.key = key + new_node.key = key.compacted() new_node.value = value.clone() node.children[child_key] = new_node self.evictable_size_ += self._node_device_resident_len(new_node) diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 9c66fd6be..5e3260230 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -1062,7 +1062,7 @@ class MambaRadixCache(BasePrefixCache): new_node.mamba_value = None # mamba cache can not be split new_node.full_lock_ref = child.full_lock_ref new_node.mamba_lock_ref = 0 - new_node.key = child.key[:split_len] + new_node.key = child.key[:split_len].compacted() new_node.value = child.value[:split_len].clone() # child time should be later than parent's time for mamba tombstone @@ -1072,7 +1072,7 @@ class MambaRadixCache(BasePrefixCache): if child.mamba_value is not None: self.mamba_lru_list.remove_node(child) child.parent = new_node - child.key = child.key[split_len:] + child.key = child.key[split_len:].compacted() child.value = child.value[split_len:].clone() new_node.parent.children[self.get_child_key_fn(key)] = new_node @@ -1134,7 +1134,7 @@ class MambaRadixCache(BasePrefixCache): if len(key): new_node = TreeNode() new_node.parent = node - new_node.key = key + new_node.key = key.compacted() new_node.value = value.clone() new_node.mamba_value = mamba_value self.full_lru_list.insert_mru(new_node) diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 35b6513ef..ca6cf0419 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -68,33 +68,87 @@ if TYPE_CHECKING: class RadixKey: + """A sequence of token ids used as a radix tree key. + + Slicing returns an O(1) view sharing the underlying token list instead of + copying it. The tree walks in ``match_prefix``/``insert`` re-slice the + remaining key at every node hop, so copying slices makes a single match + O(len * hops) — quadratic for long cached prefixes. Views keep it linear. + + Views keep the underlying list alive, so keys stored in long-lived tree + nodes must be compacted first (see ``compacted()``); only transient keys + inside a tree walk should remain views. + """ + + __slots__ = ("_tokens", "_start", "_stop", "extra_key", "is_bigram") + def __init__( self, token_ids: List[int], extra_key: Optional[str] = None, is_bigram: bool = False, ): - # token ids sequence - self.token_ids = token_ids + # token ids sequence (full backing list; this key covers + # [_start, _stop) of it) + self._tokens = token_ids + self._start = 0 + self._stop = len(token_ids) # extra key (e.g. lora_id, cache_salt) self.extra_key = extra_key # is bigram key self.is_bigram = is_bigram + @classmethod + def _view(cls, base: "RadixKey", start: int, stop: int) -> "RadixKey": + obj = cls.__new__(cls) + obj._tokens = base._tokens + obj._start = start + obj._stop = stop + obj.extra_key = base.extra_key + obj.is_bigram = base.is_bigram + return obj + + @property + def token_ids(self) -> List[int]: + if self._start == 0 and self._stop == len(self._tokens): + return self._tokens + return self._tokens[self._start : self._stop] + + @token_ids.setter + def token_ids(self, value: List[int]) -> None: + self._tokens = value + self._start = 0 + self._stop = len(value) + + def compacted(self) -> "RadixKey": + """Return an equivalent key that does not pin a larger backing list.""" + if self._start == 0 and self._stop == len(self._tokens): + return self + return RadixKey( + self._tokens[self._start : self._stop], self.extra_key, self.is_bigram + ) + def __len__(self) -> int: - return len(self.token_ids) + return self._stop - self._start def __iter__(self) -> Iterator[int]: - return iter(self.token_ids) + # NOTE: map over an index range instead of islice — islice skips the + # first `_start` elements one by one, which is O(_start) per call. + return map(self._tokens.__getitem__, range(self._start, self._stop)) def __getitem__(self, idx: Union[int, slice]) -> "RadixKey": if isinstance(idx, slice): - return RadixKey(self.token_ids[idx], self.extra_key) - return RadixKey([self.token_ids[idx]], self.extra_key) + start, stop, step = idx.indices(len(self)) + if step != 1: + raise ValueError(f"RadixKey does not support strided slicing: {idx}") + return RadixKey._view(self, self._start + start, self._start + stop) + if idx < 0: + idx += len(self) + return RadixKey([self._tokens[self._start + idx]], self.extra_key) def __repr__(self) -> str: preview = self.token_ids[:10] - return f"RadixKey(extra_key={self.extra_key!r}, token_ids={preview}{'...' if len(self.token_ids) > 10 else ''})" + return f"RadixKey(extra_key={self.extra_key!r}, token_ids={preview}{'...' if len(self) > 10 else ''})" def maybe_bigram_convert( @@ -207,22 +261,27 @@ def _check_extra_key(key0: RadixKey, key1: RadixKey): def _key_match_page_size1(key0: RadixKey, key1: RadixKey): _check_extra_key(key0, key1) + # Index the backing lists directly so that view keys (nonzero _start) are + # matched without materializing a copy of their suffix. + t0, s0 = key0._tokens, key0._start + t1, s1 = key1._tokens, key1._start + n = min(len(key0), len(key1)) i = 0 - for k0, k1 in zip(key0.token_ids, key1.token_ids): - if k0 != k1: - break + while i < n and t0[s0 + i] == t1[s1 + i]: i += 1 return i def _key_match_paged(key0: RadixKey, key1: RadixKey, page_size: int): _check_extra_key(key0, key1) + t0, s0 = key0._tokens, key0._start + t1, s1 = key1._tokens, key1._start min_len = min(len(key0), len(key1)) i = 0 while i < min_len: step = min(page_size, min_len - i) - if key0.token_ids[i : i + step] != key1.token_ids[i : i + step]: + if t0[s0 + i : s0 + i + step] != t1[s1 + i : s1 + i + step]: break i += step @@ -231,9 +290,9 @@ def _key_match_paged(key0: RadixKey, key1: RadixKey, page_size: int): def get_child_key(key: RadixKey, page_size: int = 1): if page_size == 1: - plain_key = key.token_ids[0] + plain_key = key._tokens[key._start] else: - plain_key = tuple(key.token_ids[:page_size]) + plain_key = tuple(key._tokens[key._start : key._start + page_size]) if key.extra_key is None: return plain_key else: @@ -892,10 +951,12 @@ class RadixCache(BasePrefixCache): new_node.children = {self.get_child_key_fn(key[split_len:]): child} new_node.parent = child.parent new_node.lock_ref = child.lock_ref - new_node.key = child.key[:split_len] + # Node keys must be compacted: a view would pin the backing list of the + # (possibly transient) key it was sliced from for the node's lifetime. + new_node.key = child.key[:split_len].compacted() new_node.value = child.value[:split_len].clone() child.parent = new_node - child.key = child.key[split_len:] + child.key = child.key[split_len:].compacted() child.value = child.value[split_len:].clone() new_node.parent.children[self.get_child_key_fn(key)] = new_node @@ -957,7 +1018,7 @@ class RadixCache(BasePrefixCache): if len(key): new_node = TreeNode(priority=priority) new_node.parent = node - new_node.key = key + new_node.key = key.compacted() new_node.value = value.clone() self._inc_hit_count(new_node, chunked) node.children[child_key] = new_node diff --git a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py index 9a82aa31f..989b77519 100644 --- a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py +++ b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py @@ -191,7 +191,7 @@ class LMCRadixCache(RadixCache): new_node = TreeNode(priority=last_node.priority) start = value.numel() end = start + fetched - new_node.key = key[start:end] + new_node.key = key[start:end].compacted() new_node.value = token_slots[:fetched] new_node.parent = last_node last_node.children[self.get_child_key_fn(new_node.key)] = new_node diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index dc0bae890..bee460d33 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -899,7 +899,7 @@ class SWARadixCache(BasePrefixCache): new_node.swa_tombstone = child.swa_tombstone new_node.full_lock_ref = child.full_lock_ref new_node.swa_lock_ref = child.swa_lock_ref - new_node.key = child.key[:split_len] + new_node.key = child.key[:split_len].compacted() assert len(new_node.key) > 0, f"new_node.key should not be empty" new_node.value = child.value[:split_len].clone() # parent inherits the swa_uuid from child for swa lock ref @@ -913,7 +913,7 @@ class SWARadixCache(BasePrefixCache): if not new_node.swa_tombstone: self.swa_lru_list.remove_node(child) child.parent = new_node - child.key = child.key[split_len:] + child.key = child.key[split_len:].compacted() assert len(child.key) > 0, f"child.key should not be empty" child.value = child.value[split_len:].clone() new_node.parent.children[self.get_child_key_fn(key)] = new_node @@ -1039,7 +1039,7 @@ class SWARadixCache(BasePrefixCache): assert len(key) > 0, f"key should not be empty" new_node = TreeNode() new_node.parent = parent - new_node.key = key + new_node.key = key.compacted() new_node.value = value.clone() new_node.swa_tombstone = swa_tombstone parent.children[self.get_child_key_fn(key)] = new_node diff --git a/test/registered/unit/mem_cache/test_radix_key_views.py b/test/registered/unit/mem_cache/test_radix_key_views.py new file mode 100644 index 000000000..3cd1d68c2 --- /dev/null +++ b/test/registered/unit/mem_cache/test_radix_key_views.py @@ -0,0 +1,219 @@ +""" +Unit tests for RadixKey zero-copy view semantics. + +RadixKey slicing returns O(1) views over a shared backing list instead of +copying. These tests pin down the semantics that the radix tree relies on: + +- Slicing/iteration/indexing behave exactly like list slicing. +- Views share the backing list; compacted() detaches them. +- Key matching and child-key derivation are view-position independent. +- Keys stored in tree nodes are always compact (never views), so lock-ref + walks, eviction, splits, and controller-thread reads never observe a key + that pins a transient backing list. + +Usage: + python -m pytest test_radix_key_views.py -v +""" + +import sys +import unittest.mock + +for _mod in ("sgl_kernel", "sgl_kernel.kvcacheio"): + if _mod not in sys.modules: + sys.modules[_mod] = unittest.mock.MagicMock() + +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci + +# CPU-based unit test, runs quickly on any GPU runner +register_cuda_ci(est_time=5, suite="stage-b-test-1-gpu-small") +register_amd_ci(est_time=5, suite="stage-b-test-1-gpu-small-amd") + +import random +import unittest + +import torch + +from sglang.srt.mem_cache.base_prefix_cache import InsertParams, MatchPrefixParams +from sglang.srt.mem_cache.radix_cache import ( + RadixCache, + RadixKey, + _key_match_page_size1, + _key_match_paged, + get_child_key, + maybe_bigram_convert, +) + + +class TestRadixKeyViewSemantics(unittest.TestCase): + def test_slice_is_view_sharing_backing(self): + backing = list(range(100)) + key = RadixKey(backing) + view = key[10:50] + self.assertIs(view._tokens, key._tokens) + self.assertEqual(view.token_ids, backing[10:50]) + self.assertEqual(len(view), 40) + + def test_nested_slicing_matches_list_semantics(self): + rng = random.Random(0) + tokens = [rng.randrange(1000) for _ in range(257)] + key = RadixKey(tokens) + ref = list(tokens) + for _ in range(200): + if len(ref) == 0: + break + a = rng.randrange(0, len(ref) + 1) + b = rng.randrange(a, len(ref) + 1) + key = key[a:b] + ref = ref[a:b] + self.assertEqual(key.token_ids, ref) + self.assertEqual(len(key), len(ref)) + self.assertEqual(list(key), ref) + + def test_open_ended_and_negative_slices(self): + tokens = list(range(20)) + key = RadixKey(tokens) + self.assertEqual(key[5:].token_ids, tokens[5:]) + self.assertEqual(key[:7].token_ids, tokens[:7]) + self.assertEqual(key[-4:].token_ids, tokens[-4:]) + self.assertEqual(key[3:1].token_ids, []) + + def test_int_indexing(self): + tokens = [7, 8, 9, 10] + key = RadixKey(tokens)[1:] + item = key[0] + self.assertEqual(item.token_ids, [8]) + self.assertEqual(key[-1].token_ids, [10]) + + def test_strided_slicing_rejected(self): + key = RadixKey(list(range(10))) + with self.assertRaises(ValueError): + key[::2] + + def test_compacted_detaches_views_and_is_identity_for_full_keys(self): + backing = list(range(50)) + key = RadixKey(backing, extra_key="lora1", is_bigram=False) + self.assertIs(key.compacted(), key) + + view = key[10:30] + compact = view.compacted() + self.assertIsNot(compact._tokens, backing) + self.assertEqual(compact.token_ids, backing[10:30]) + self.assertEqual(compact.extra_key, "lora1") + # Mutating the original backing list must not affect the compacted key. + backing[15] = -1 + self.assertNotIn(-1, compact.token_ids) + + def test_token_ids_setter_resets_view(self): + key = RadixKey(list(range(30)))[5:25] + key.token_ids = [1, 2, 3] + self.assertEqual(len(key), 3) + self.assertEqual(key.token_ids, [1, 2, 3]) + + def test_length_frozen_at_construction(self): + backing = [1, 2, 3] + key = RadixKey(backing) + backing.append(4) + self.assertEqual(len(key), 3) + self.assertEqual(key.token_ids, [1, 2, 3]) + + def test_extra_key_and_bigram_preserved_through_slices(self): + key = RadixKey([(1, 2), (2, 3), (3, 4)], extra_key="salt", is_bigram=True) + view = key[1:] + self.assertEqual(view.extra_key, "salt") + self.assertTrue(view.is_bigram) + self.assertTrue(view.compacted().is_bigram) + + def test_bigram_convert_on_view(self): + key = RadixKey(list(range(10)))[2:8] + converted, _ = maybe_bigram_convert(True, key) + self.assertTrue(converted.is_bigram) + self.assertEqual( + converted.token_ids[0][0] if converted.token_ids else None, + 2, + ) + + def test_key_match_is_view_position_independent(self): + rng = random.Random(1) + base = [rng.randrange(50) for _ in range(512)] + other = list(base[:300]) + [999] + base[301:] + for page_size in (1, 16, 64): + for off in (0, 1, 63, 128): + k0 = RadixKey([0] * off + base)[off:] + k1 = RadixKey(other) + k0_ref = RadixKey(list(base)) + if page_size == 1: + got = _key_match_page_size1(k0, k1) + want = _key_match_page_size1(k0_ref, k1) + else: + got = _key_match_paged(k0, k1, page_size) + want = _key_match_paged(k0_ref, k1, page_size) + self.assertEqual(got, want, f"{page_size=} {off=}") + + def test_get_child_key_on_views(self): + tokens = list(range(200)) + key = RadixKey(tokens, extra_key="e") + view = key[64:] + self.assertEqual(get_child_key(view, 64), get_child_key(view.compacted(), 64)) + self.assertEqual(get_child_key(view, 64), ("e", tuple(tokens[64:128]))) + self.assertEqual(get_child_key(view, 1), ("e", 64)) + + def test_tree_node_keys_are_always_compact(self): + """Node-stored keys must never be views: a view would pin the transient + request key's backing list, and lock/update paths (inc_lock_ref, + eviction, splits) would read keys whose backing outlives the match.""" + rng = random.Random(2) + cache = RadixCache.create_simulated(page_size=4) + seqs = [] + base = [rng.randrange(8) for _ in range(256)] + for _ in range(40): + # Shared prefixes of random length force splits at random points. + cut = rng.randrange(8, 256, 4) + seq = base[:cut] + [rng.randrange(8) for _ in range(rng.randrange(4, 64))] + seqs.append(seq) + cache.insert( + InsertParams( + key=RadixKey(seq), + value=torch.arange(len(seq), dtype=torch.int64), + ) + ) + cache.match_prefix(MatchPrefixParams(key=RadixKey(list(seq)))) + + stack = [cache.root_node] + checked = 0 + while stack: + node = stack.pop() + stack.extend(node.children.values()) + if node is cache.root_node or node.key is None: + continue + checked += 1 + self.assertEqual(node.key._start, 0, f"node {node.id} key is a view") + self.assertEqual( + node.key._stop, + len(node.key._tokens), + f"node {node.id} key is a view", + ) + self.assertGreater(checked, 10) + + def test_match_and_insert_equivalence_random(self): + """Differential test: tree behavior identical for fresh keys vs views.""" + rng = random.Random(3) + cache = RadixCache.create_simulated(page_size=2) + for _ in range(50): + seq = [rng.randrange(4) for _ in range(rng.randrange(2, 40, 2))] + cache.insert( + InsertParams( + key=RadixKey(seq), + value=torch.arange(len(seq), dtype=torch.int64), + ) + ) + probe = [rng.randrange(4) for _ in range(rng.randrange(2, 40, 2))] + res_fresh = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(probe)))) + padded = RadixKey([9] * 6 + list(probe))[6:] + res_view = cache.match_prefix(MatchPrefixParams(key=padded)) + self.assertTrue( + torch.equal(res_fresh.device_indices, res_view.device_indices) + ) + + +if __name__ == "__main__": + unittest.main()