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

[Refactor] Move radix-cache utils onto RadixKey as methods (#23209)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-04-20 23:11:58 -07:00
committed by GitHub
parent a490632416
commit 2b2cad70d6
9 changed files with 156 additions and 239 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/hi_mamba_radix_cache.py
+15 -15
View File
@@ -548,7 +548,7 @@ class HiMambaRadixCache(MambaRadixCache):
self._discard_from_leaf_sets(node) self._discard_from_leaf_sets(node)
parent = node.parent parent = node.parent
key = self.get_child_key_fn(node.key) key = node.key.child_key(self.page_size)
v = parent.children.pop(key, None) v = parent.children.pop(key, None)
assert v == node, f"parent does not have child key, {key}" assert v == node, f"parent does not have child key, {key}"
@@ -584,7 +584,7 @@ class HiMambaRadixCache(MambaRadixCache):
self._discard_from_leaf_sets(node) self._discard_from_leaf_sets(node)
parent = node.parent parent = node.parent
key = self.get_child_key_fn(node.key) key = node.key.child_key(self.page_size)
v = parent.children.pop(key, None) v = parent.children.pop(key, None)
assert v == node, f"parent does not have child key, {key}" assert v == node, f"parent does not have child key, {key}"
@@ -598,7 +598,7 @@ class HiMambaRadixCache(MambaRadixCache):
assert node.mamba_host_value is None, f"has mamba host value, {node.id=}" assert node.mamba_host_value is None, f"has mamba host value, {node.id=}"
assert len(node.children) == 0, f"leaf node has children, {node.id=}" assert len(node.children) == 0, f"leaf node has children, {node.id=}"
parent = node.parent parent = node.parent
key = self.get_child_key_fn(node.key) key = node.key.child_key(self.page_size)
v = parent.children.pop(key, None) v = parent.children.pop(key, None)
assert v == node, f"parent does not have child key, {key}" assert v == node, f"parent does not have child key, {key}"
@@ -824,7 +824,7 @@ class HiMambaRadixCache(MambaRadixCache):
if len(key) == 0: if len(key) == 0:
return 0, True return 0, True
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
total_prefix_length = 0 total_prefix_length = 0
while len(key) > 0 and child_key in node.children.keys(): while len(key) > 0 and child_key in node.children.keys():
@@ -836,7 +836,7 @@ class HiMambaRadixCache(MambaRadixCache):
if node.mamba_value is not None: if node.mamba_value is not None:
self.mamba_lru_list.reset_node_mru(node) self.mamba_lru_list.reset_node_mru(node)
prefix_len = self.key_match_fn(node.key, key) prefix_len = node.key.match(key, page_size=self.page_size)
if prefix_len < len(node.key): if prefix_len < len(node.key):
new_node = self._split_node(node.key, node, prefix_len) new_node = self._split_node(node.key, node, prefix_len)
@@ -855,7 +855,7 @@ class HiMambaRadixCache(MambaRadixCache):
value = value[prefix_len:] value = value[prefix_len:]
if len(key): if len(key):
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
mamba_value_exist = False mamba_value_exist = False
if len(key): if len(key):
@@ -884,7 +884,7 @@ class HiMambaRadixCache(MambaRadixCache):
value: torch.Tensor, value: torch.Tensor,
mamba_value: torch.Tensor, mamba_value: torch.Tensor,
) -> TreeNode: ) -> TreeNode:
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
new_node = TreeNode() new_node = TreeNode()
new_node.parent = parent new_node.parent = parent
new_node.key = key new_node.key = key
@@ -924,7 +924,7 @@ class HiMambaRadixCache(MambaRadixCache):
) -> Tuple[List[torch.Tensor], TreeNode, int]: ) -> Tuple[List[torch.Tensor], TreeNode, int]:
"""Walk tree to find best_last_node (mamba boundary).""" """Walk tree to find best_last_node (mamba boundary)."""
node = self.root_node node = self.root_node
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
value: List[torch.Tensor] = [] value: List[torch.Tensor] = []
best_value_len = 0 best_value_len = 0
@@ -940,7 +940,7 @@ class HiMambaRadixCache(MambaRadixCache):
best_value_len = len(value) best_value_len = len(value)
best_last_node = node best_last_node = node
prefix_len = self.key_match_fn(child.key, key) prefix_len = child.key.match(key, page_size=self.page_size)
if prefix_len < len(child.key): if prefix_len < len(child.key):
new_node = self._split_node(child.key, child, prefix_len) new_node = self._split_node(child.key, child, prefix_len)
if not new_node.evicted: if not new_node.evicted:
@@ -953,7 +953,7 @@ class HiMambaRadixCache(MambaRadixCache):
node = child node = child
key = key[prefix_len:] key = key[prefix_len:]
if len(key): if len(key):
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
if node.mamba_value is not None or node.mamba_backuped: if node.mamba_value is not None or node.mamba_backuped:
best_value_len = len(value) best_value_len = len(value)
@@ -1074,7 +1074,7 @@ class HiMambaRadixCache(MambaRadixCache):
self.evictable_full_host_leaves.discard(child) self.evictable_full_host_leaves.discard(child)
new_node = TreeNode() new_node = TreeNode()
new_node.children = {self.get_child_key_fn(key[split_len:]): child} new_node.children = {key[split_len:].child_key(self.page_size): child}
new_node.parent = child.parent new_node.parent = child.parent
new_node.value = None new_node.value = None
new_node.mamba_value = None new_node.mamba_value = None
@@ -1095,7 +1095,7 @@ class HiMambaRadixCache(MambaRadixCache):
self.mamba_lru_list.remove_node(child) self.mamba_lru_list.remove_node(child)
child.parent = new_node child.parent = new_node
child.key = child.key[split_len:] child.key = child.key[split_len:]
new_node.parent.children[self.get_child_key_fn(key)] = new_node new_node.parent.children[key.child_key(self.page_size)] = new_node
if child.mamba_value is not None: if child.mamba_value is not None:
self.mamba_lru_list.insert_mru(child) self.mamba_lru_list.insert_mru(child)
@@ -1836,7 +1836,7 @@ class HiMambaRadixCache(MambaRadixCache):
if len(key) == 0: if len(key) == 0:
return 0 return 0
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
matched_length = 0 matched_length = 0
while len(key) > 0 and child_key in node.children.keys(): while len(key) > 0 and child_key in node.children.keys():
@@ -1844,7 +1844,7 @@ class HiMambaRadixCache(MambaRadixCache):
node.last_access_time = get_last_access_time() node.last_access_time = get_last_access_time()
if node != self.root_node and node.mamba_value is not None: if node != self.root_node and node.mamba_value is not None:
self.mamba_lru_list.reset_node_mru(node) self.mamba_lru_list.reset_node_mru(node)
prefix_len = self.key_match_fn(node.key, key) prefix_len = node.key.match(key, page_size=self.page_size)
key = key[prefix_len:] key = key[prefix_len:]
host_value = host_value[prefix_len:] host_value = host_value[prefix_len:]
@@ -1856,7 +1856,7 @@ class HiMambaRadixCache(MambaRadixCache):
node = new_node node = new_node
if len(key): if len(key):
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
leaf_node: Optional[TreeNode] = None leaf_node: Optional[TreeNode] = None
if len(key): if len(key):
python/sglang/srt/mem_cache/hiradix_cache.py
+13 -17
View File
@@ -781,11 +781,7 @@ class HiRadixCache(RadixCache):
return self.evictable_size_ return self.evictable_size_
def _to_radix_key(self, token_ids: List[int]) -> RadixKey: def _to_radix_key(self, token_ids: List[int]) -> RadixKey:
"""Convert raw token_ids to a RadixKey for tree walking. """Convert raw token_ids to a RadixKey; must be list (not tuple) for paged match."""
Must use list (not tuple) to match scheduler's RadixKey format,
since _key_match_paged compares slices directly and list != tuple.
"""
return RadixKey(token_ids=list(token_ids)) return RadixKey(token_ids=list(token_ids))
def inc_lock_ref(self, node: TreeNode) -> IncLockRefResult: def inc_lock_ref(self, node: TreeNode) -> IncLockRefResult:
@@ -938,7 +934,7 @@ class HiRadixCache(RadixCache):
self._record_remove_event(x, medium=StorageMedium.CPU) self._record_remove_event(x, medium=StorageMedium.CPU)
num_evicted += self.cache_controller.evict_host(x.host_value) num_evicted += self.cache_controller.evict_host(x.host_value)
key = self.get_child_key_fn(x.key) key = x.key.child_key(self.page_size)
v = x.parent.children.pop(key, None) v = x.parent.children.pop(key, None)
assert v == x, f"parent does not have child key, {key}" assert v == x, f"parent does not have child key, {key}"
if x in self.evictable_host_leaves: if x in self.evictable_host_leaves:
@@ -1312,13 +1308,13 @@ class HiRadixCache(RadixCache):
if len(key) == 0: if len(key) == 0:
return 0 return 0
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
matched_length = 0 matched_length = 0
while len(key) > 0 and child_key in node.children.keys(): while len(key) > 0 and child_key in node.children.keys():
node = node.children[child_key] node = node.children[child_key]
node.last_access_time = time.monotonic() node.last_access_time = time.monotonic()
prefix_len = self.key_match_fn(node.key, key) prefix_len = node.key.match(key, page_size=self.page_size)
key = key[prefix_len:] key = key[prefix_len:]
host_value = host_value[prefix_len:] host_value = host_value[prefix_len:]
hash_value = hash_value[prefix_len // self.page_size :] hash_value = hash_value[prefix_len // self.page_size :]
@@ -1329,7 +1325,7 @@ class HiRadixCache(RadixCache):
node = new_node node = new_node
if len(key): if len(key):
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
if len(key): if len(key):
new_node = TreeNode(priority=node.priority) new_node = TreeNode(priority=node.priority)
@@ -1350,13 +1346,13 @@ class HiRadixCache(RadixCache):
def _match_prefix_helper(self, node: TreeNode, key: RadixKey): def _match_prefix_helper(self, node: TreeNode, key: RadixKey):
node.last_access_time = time.monotonic() node.last_access_time = time.monotonic()
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
value = [] value = []
while len(key) > 0 and child_key in node.children.keys(): while len(key) > 0 and child_key in node.children.keys():
child = node.children[child_key] child = node.children[child_key]
child.last_access_time = time.monotonic() child.last_access_time = time.monotonic()
prefix_len = self.key_match_fn(child.key, key) prefix_len = child.key.match(key, page_size=self.page_size)
if prefix_len < len(child.key): if prefix_len < len(child.key):
new_node = self._split_node(child.key, child, prefix_len) new_node = self._split_node(child.key, child, prefix_len)
if not new_node.evicted: if not new_node.evicted:
@@ -1370,14 +1366,14 @@ class HiRadixCache(RadixCache):
key = key[prefix_len:] key = key[prefix_len:]
if len(key): if len(key):
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
return value, node return value, node
def _split_node(self, key: RadixKey, child: TreeNode, split_len: int): def _split_node(self, key: RadixKey, child: TreeNode, split_len: int):
# child node split into new_node -> child # child node split into new_node -> child
new_node = TreeNode(priority=child.priority) new_node = TreeNode(priority=child.priority)
new_node.children = {self.get_child_key_fn(key[split_len:]): child} new_node.children = {key[split_len:].child_key(self.page_size): child}
new_node.parent = child.parent new_node.parent = child.parent
new_node.lock_ref = child.lock_ref new_node.lock_ref = child.lock_ref
new_node.key = child.key[:split_len] new_node.key = child.key[:split_len]
@@ -1398,7 +1394,7 @@ class HiRadixCache(RadixCache):
) )
child.parent = new_node child.parent = new_node
child.key = child.key[split_len:] child.key = child.key[split_len:]
new_node.parent.children[self.get_child_key_fn(key)] = new_node new_node.parent.children[key.child_key(self.page_size)] = new_node
return new_node return new_node
@@ -1420,14 +1416,14 @@ class HiRadixCache(RadixCache):
return InsertResult(prefix_len=0) return InsertResult(prefix_len=0)
node = self.root_node node = self.root_node
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
total_prefix_length = 0 total_prefix_length = 0
while len(key) > 0 and child_key in node.children.keys(): while len(key) > 0 and child_key in node.children.keys():
node = node.children[child_key] node = node.children[child_key]
node.last_access_time = time.monotonic() node.last_access_time = time.monotonic()
node.priority = max(node.priority, priority) node.priority = max(node.priority, priority)
prefix_len = self.key_match_fn(node.key, key) prefix_len = node.key.match(key, page_size=self.page_size)
if prefix_len == len(node.key): if prefix_len == len(node.key):
if node.evicted: if node.evicted:
@@ -1463,7 +1459,7 @@ class HiRadixCache(RadixCache):
value = value[prefix_len:] value = value[prefix_len:]
if len(key): if len(key):
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
if len(key): if len(key):
new_node = TreeNode(priority=priority) new_node = TreeNode(priority=priority)
python/sglang/srt/mem_cache/mamba_radix_cache.py
+15 -26
View File
@@ -21,7 +21,7 @@ The radix tree data structure for managing the hybrid (full and Mamba) KV cache.
import heapq import heapq
from collections import defaultdict from collections import defaultdict
from functools import lru_cache, partial from functools import lru_cache
from typing import TYPE_CHECKING, List, Optional, Tuple from typing import TYPE_CHECKING, List, Optional, Tuple
import torch import torch
@@ -46,12 +46,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
MatchResult, MatchResult,
) )
from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool
from sglang.srt.mem_cache.radix_cache import ( from sglang.srt.mem_cache.radix_cache import RadixKey
RadixKey,
_key_match_page_size1,
_key_match_paged,
get_child_key,
)
from sglang.srt.server_args import get_global_server_args from sglang.srt.server_args import get_global_server_args
if TYPE_CHECKING: if TYPE_CHECKING:
@@ -445,12 +440,6 @@ class MambaRadixCache(BasePrefixCache):
if params.enable_metrics: if params.enable_metrics:
self.init_metrics_collector() self.init_metrics_collector()
if self.page_size == 1:
self.key_match_fn = _key_match_page_size1
self.get_child_key_fn = get_child_key
else:
self.key_match_fn = partial(_key_match_paged, page_size=self.page_size)
self.get_child_key_fn = partial(get_child_key, page_size=self.page_size)
self.reset() self.reset()
##### Public API ##### ##### Public API #####
@@ -963,7 +952,7 @@ class MambaRadixCache(BasePrefixCache):
node is greater than or equal to the sliding window size. node is greater than or equal to the sliding window size.
""" """
node = self.root_node node = self.root_node
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
value: List[torch.Tensor] = [] value: List[torch.Tensor] = []
best_value_len = 0 best_value_len = 0
@@ -975,7 +964,7 @@ class MambaRadixCache(BasePrefixCache):
best_value_len = len(value) best_value_len = len(value)
best_last_node = node best_last_node = node
prefix_len = self.key_match_fn(child.key, key) prefix_len = child.key.match(key, page_size=self.page_size)
if prefix_len < len(child.key): if prefix_len < len(child.key):
new_node = self._split_node(child.key, child, prefix_len) new_node = self._split_node(child.key, child, prefix_len)
value.append(new_node.value) value.append(new_node.value)
@@ -987,7 +976,7 @@ class MambaRadixCache(BasePrefixCache):
key = key[prefix_len:] key = key[prefix_len:]
if len(key): if len(key):
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
# handle best_value_len and best_last_node, for the case that last node is fully matched # handle best_value_len and best_last_node, for the case that last node is fully matched
if node.mamba_value is not None: if node.mamba_value is not None:
best_value_len = len(value) best_value_len = len(value)
@@ -1081,7 +1070,7 @@ class MambaRadixCache(BasePrefixCache):
def _split_node(self, key: RadixKey, child: TreeNode, split_len: int) -> TreeNode: def _split_node(self, key: RadixKey, child: TreeNode, split_len: int) -> TreeNode:
# new_node -> child # new_node -> child
new_node = TreeNode() new_node = TreeNode()
new_node.children = {self.get_child_key_fn(key[split_len:]): child} new_node.children = {key[split_len:].child_key(self.page_size): child}
new_node.parent = child.parent new_node.parent = child.parent
new_node.mamba_value = None # mamba cache can not be split new_node.mamba_value = None # mamba cache can not be split
new_node.full_lock_ref = child.full_lock_ref new_node.full_lock_ref = child.full_lock_ref
@@ -1098,7 +1087,7 @@ class MambaRadixCache(BasePrefixCache):
child.parent = new_node child.parent = new_node
child.key = child.key[split_len:] child.key = child.key[split_len:]
child.value = child.value[split_len:].clone() child.value = child.value[split_len:].clone()
new_node.parent.children[self.get_child_key_fn(key)] = new_node new_node.parent.children[key.child_key(self.page_size)] = new_node
# insert the new node and child into the lru lists, insert # insert the new node and child into the lru lists, insert
# parent first so that parent is after child in the lru list # parent first so that parent is after child in the lru list
@@ -1128,7 +1117,7 @@ class MambaRadixCache(BasePrefixCache):
if len(key) == 0: if len(key) == 0:
return 0, True return 0, True
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
total_prefix_length = 0 total_prefix_length = 0
while len(key) > 0 and child_key in node.children.keys(): while len(key) > 0 and child_key in node.children.keys():
@@ -1137,7 +1126,7 @@ class MambaRadixCache(BasePrefixCache):
self.full_lru_list.reset_node_mru(node) self.full_lru_list.reset_node_mru(node)
if node.mamba_value is not None: if node.mamba_value is not None:
self.mamba_lru_list.reset_node_mru(node) self.mamba_lru_list.reset_node_mru(node)
prefix_len = self.key_match_fn(node.key, key) prefix_len = node.key.match(key, page_size=self.page_size)
if prev_prefix_len < total_prefix_length + prefix_len: if prev_prefix_len < total_prefix_length + prefix_len:
start = max(0, prev_prefix_len - total_prefix_length) start = max(0, prev_prefix_len - total_prefix_length)
@@ -1152,7 +1141,7 @@ class MambaRadixCache(BasePrefixCache):
node = new_node node = new_node
if len(key): if len(key):
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
mamba_value_exist = False mamba_value_exist = False
if len(key): if len(key):
@@ -1208,7 +1197,7 @@ class MambaRadixCache(BasePrefixCache):
node.mamba_value is not None node.mamba_value is not None
), f"Invariant violated: leaf node is a tombstone, {node.id=}" ), f"Invariant violated: leaf node is a tombstone, {node.id=}"
assert len(node.children) == 0, f"leaf node has children, {node.id=}" assert len(node.children) == 0, f"leaf node has children, {node.id=}"
key = self.get_child_key_fn(node.key) key = node.key.child_key(self.page_size)
v = node.parent.children.pop(key, None) v = node.parent.children.pop(key, None)
assert v == node, f"parent does not have child key, {key}" assert v == node, f"parent does not have child key, {key}"
@@ -1225,7 +1214,7 @@ class MambaRadixCache(BasePrefixCache):
node.mamba_value is None node.mamba_value is None
), f"Deleting a unexpected non-tombstone leaf node, {node.id=}" ), f"Deleting a unexpected non-tombstone leaf node, {node.id=}"
assert len(node.children) == 0, f"leaf node has children, {node.id=}" assert len(node.children) == 0, f"leaf node has children, {node.id=}"
key = self.get_child_key_fn(node.key) key = node.key.child_key(self.page_size)
v = node.parent.children.pop(key, None) v = node.parent.children.pop(key, None)
assert v == node, f"parent does not have child key, {key}" assert v == node, f"parent does not have child key, {key}"
@@ -1270,9 +1259,9 @@ class MambaRadixCache(BasePrefixCache):
for key, child in current_node.children.items(): for key, child in current_node.children.items():
stack.append((child, current_indent + 2)) stack.append((child, current_indent + 2))
assert key == self.get_child_key_fn( assert key == child.key.child_key(
child.key self.page_size
), f"{key=}, {self.get_child_key_fn(child.key)=}" ), f"{key=}, {child.key.child_key(self.page_size)=}"
def _total_size_helper(self) -> Tuple[int, int]: def _total_size_helper(self) -> Tuple[int, int]:
total_size = 0 total_size = 0
python/sglang/srt/mem_cache/radix_cache.py
+81 -124
View File
@@ -21,12 +21,13 @@ limitations under the License.
The radix tree data structure for managing the KV cache. The radix tree data structure for managing the KV cache.
""" """
import hashlib
import heapq import heapq
import logging import logging
import sys import sys
import time import time
from collections import defaultdict from collections import defaultdict
from functools import lru_cache, partial from functools import lru_cache
from typing import TYPE_CHECKING, Any, Iterator, List, Optional, Tuple, Union from typing import TYPE_CHECKING, Any, Iterator, List, Optional, Tuple, Union
import torch import torch
@@ -141,6 +142,71 @@ class RadixKey:
value = value[: len(self)] value = value[: len(self)]
return self, value return self, value
def _check_compatible(self, other: "RadixKey") -> None:
if self.extra_key != other.extra_key:
raise ValueError(
f"RadixKey operations require matching extra_key, but got "
f"{self.extra_key=} != {other.extra_key=}"
)
def match(self, other: "RadixKey", page_size: int = 1) -> int:
"""Logical-unit prefix length shared with ``other``. Result is rounded down to ``page_size``."""
self._check_compatible(other)
t0, t1 = self.token_ids, other.token_ids
if self.is_bigram:
# Walk raw tokens; L matching tokens imply L-1 matching bigrams.
i = 0
for a, b in zip(t0, t1):
if a != b:
break
i += 1
matched = max(0, min(i - 1, len(self), len(other)))
return (matched // page_size) * page_size if page_size > 1 else matched
if page_size == 1:
i = 0
for a, b in zip(t0, t1):
if a != b:
break
i += 1
return i
min_len = min(len(self), len(other))
i = 0
while i < min_len:
if t0[i : i + page_size] != t1[i : i + page_size]:
break
i += page_size
return i
def child_key(self, page_size: int = 1):
"""Hashable dict-key for the first ``page_size`` logical units, namespaced by ``extra_key``."""
t = self.token_ids
if self.is_bigram:
if page_size == 1:
plain = (t[0], t[1])
else:
plain = tuple((t[j], t[j + 1]) for j in range(page_size))
else:
plain = t[0] if page_size == 1 else tuple(t[:page_size])
return plain if self.extra_key is None else (self.extra_key, plain)
def hash_page(self, start: int, end: int, prior_hash: Optional[str] = None) -> str:
"""SHA256 for logical units [start, end); bigram mode feeds overlapping (t_i, t_{i+1}) byte pairs."""
hasher = hashlib.sha256()
if prior_hash:
hasher.update(bytes.fromhex(prior_hash))
t = self.token_ids
if self.is_bigram:
for j in range(start, end):
hasher.update(t[j].to_bytes(4, byteorder="little", signed=False))
hasher.update(t[j + 1].to_bytes(4, byteorder="little", signed=False))
else:
for j in range(start, end):
hasher.update(t[j].to_bytes(4, byteorder="little", signed=False))
return hasher.hexdigest()
class TreeNode: class TreeNode:
@@ -205,82 +271,8 @@ class TreeNode:
return self.last_access_time < other.last_access_time return self.last_access_time < other.last_access_time
def _check_extra_key(key0: RadixKey, key1: RadixKey):
if key0.extra_key != key1.extra_key:
raise ValueError(
f"_key_match should be run on the same extra key, but got key0.extra_key={key0.extra_key} != key1.extra_key={key1.extra_key}"
)
def _key_match_page_size1(key0: RadixKey, key1: RadixKey):
_check_extra_key(key0, key1)
# In bigram mode we compare raw tokens position-by-position; matching L
# consecutive tokens implies L-1 matching bigrams. In plain mode, matching
# tokens == matching units directly.
t0 = key0.token_ids
t1 = key1.token_ids
i = 0
for a, b in zip(t0, t1):
if a != b:
break
i += 1
if key0.is_bigram:
# Clamp by logical bigram length of each side (guards short tails).
return max(0, min(i - 1, len(key0), len(key1)))
return i
def _key_match_paged(key0: RadixKey, key1: RadixKey, page_size: int):
_check_extra_key(key0, key1)
if key0.is_bigram:
# Walk raw tokens, convert to bigram count, then round to page boundary.
t0 = key0.token_ids
t1 = key1.token_ids
i = 0
for a, b in zip(t0, t1):
if a != b:
break
i += 1
bigram_matched = max(0, i - 1)
bigram_matched = min(bigram_matched, len(key0), len(key1))
return (bigram_matched // page_size) * page_size
min_len = min(len(key0), len(key1))
i = 0
while i < min_len:
if key0.token_ids[i : i + page_size] != key1.token_ids[i : i + page_size]:
break
i += page_size
return i
def get_child_key(key: RadixKey, page_size: int = 1):
if key.is_bigram:
t = key.token_ids
if page_size == 1:
# first bigram -> (tokens[0], tokens[1])
plain_key = (t[0], t[1])
else:
# first page_size bigrams spanning tokens[0 : page_size + 1]
plain_key = tuple((t[j], t[j + 1]) for j in range(page_size))
else:
if page_size == 1:
plain_key = key.token_ids[0]
else:
plain_key = tuple(key.token_ids[:page_size])
if key.extra_key is None:
return plain_key
else:
return (key.extra_key, plain_key)
def compute_node_hash_values(node: "TreeNode", page_size: int) -> List[str]: def compute_node_hash_values(node: "TreeNode", page_size: int) -> List[str]:
"""Compute SHA256-based hash values for position-aware identification. """Compute SHA256-based hash values for position-aware identification."""
In bigram mode, each page logically covers `page_size` bigrams over
`page_size + 1` raw tokens; we feed overlapping (t_i, t_{i+1}) byte pairs
to the hasher so the output matches the pre-optimization tuple-based hash.
"""
hash_values = [] hash_values = []
parent_hash = None parent_hash = None
@@ -288,45 +280,17 @@ def compute_node_hash_values(node: "TreeNode", page_size: int) -> List[str]:
if len(node.parent.key) > 0 and len(node.parent.hash_value) > 0: if len(node.parent.key) > 0 and len(node.parent.hash_value) > 0:
parent_hash = node.parent.hash_value[-1] parent_hash = node.parent.hash_value[-1]
raw = node.key.token_ids
is_bigram = node.key.is_bigram
logical_len = len(node.key) logical_len = len(node.key)
for start in range(0, logical_len, page_size): for start in range(0, logical_len, page_size):
end = min(start + page_size, logical_len) end = min(start + page_size, logical_len)
if end <= start: if end <= start:
continue continue
hash_val = _hash_page(raw, start, end, is_bigram, parent_hash) hash_val = node.key.hash_page(start, end, parent_hash)
hash_values.append(hash_val) hash_values.append(hash_val)
parent_hash = hash_val parent_hash = hash_val
return hash_values return hash_values
def _hash_page(
raw_tokens: List[int],
start: int,
end: int,
is_bigram: bool,
prior_hash: Optional[str],
) -> str:
import hashlib
hasher = hashlib.sha256()
if prior_hash:
hasher.update(bytes.fromhex(prior_hash))
if is_bigram:
for j in range(start, end):
hasher.update(raw_tokens[j].to_bytes(4, byteorder="little", signed=False))
hasher.update(
raw_tokens[j + 1].to_bytes(4, byteorder="little", signed=False)
)
else:
for j in range(start, end):
hasher.update(raw_tokens[j].to_bytes(4, byteorder="little", signed=False))
return hasher.hexdigest()
def split_node_hash_value( def split_node_hash_value(
child_hash_value: Optional[List[str]], split_len: int, page_size: int child_hash_value: Optional[List[str]], split_len: int, page_size: int
) -> tuple[Optional[List[str]], Optional[List[str]]]: ) -> tuple[Optional[List[str]], Optional[List[str]]]:
@@ -375,13 +339,6 @@ class RadixCache(BasePrefixCache):
else: else:
self.device = torch.device("cpu") self.device = torch.device("cpu")
if self.page_size == 1:
self.key_match_fn = _key_match_page_size1
self.get_child_key_fn = get_child_key
else:
self.key_match_fn = partial(_key_match_paged, page_size=self.page_size)
self.get_child_key_fn = partial(get_child_key, page_size=self.page_size)
if self.eviction_policy == "lru": if self.eviction_policy == "lru":
self.eviction_strategy: EvictionStrategy = LRUStrategy() self.eviction_strategy: EvictionStrategy = LRUStrategy()
elif self.eviction_policy == "lfu": elif self.eviction_policy == "lfu":
@@ -737,13 +694,13 @@ class RadixCache(BasePrefixCache):
access_time = time.monotonic() access_time = time.monotonic()
node.last_access_time = access_time node.last_access_time = access_time
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
value = [] value = []
while len(key) > 0 and child_key in node.children.keys(): while len(key) > 0 and child_key in node.children.keys():
child = node.children[child_key] child = node.children[child_key]
child.last_access_time = access_time child.last_access_time = access_time
prefix_len = self.key_match_fn(child.key, key) prefix_len = child.key.match(key, page_size=self.page_size)
if prefix_len < len(child.key): if prefix_len < len(child.key):
new_node = self._split_node(child.key, child, prefix_len) new_node = self._split_node(child.key, child, prefix_len)
value.append(new_node.value) value.append(new_node.value)
@@ -755,7 +712,7 @@ class RadixCache(BasePrefixCache):
key = key[prefix_len:] key = key[prefix_len:]
if len(key): if len(key):
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
return value, node return value, node
@@ -764,7 +721,7 @@ class RadixCache(BasePrefixCache):
# New node inherits child's priority (represents shared prefix) # New node inherits child's priority (represents shared prefix)
new_node = TreeNode(priority=child.priority) new_node = TreeNode(priority=child.priority)
new_node.hit_count = child.hit_count new_node.hit_count = child.hit_count
new_node.children = {self.get_child_key_fn(key[split_len:]): child} new_node.children = {key[split_len:].child_key(self.page_size): child}
new_node.parent = child.parent new_node.parent = child.parent
new_node.lock_ref = child.lock_ref new_node.lock_ref = child.lock_ref
new_node.key = child.key[:split_len] new_node.key = child.key[:split_len]
@@ -772,7 +729,7 @@ class RadixCache(BasePrefixCache):
child.parent = new_node child.parent = new_node
child.key = child.key[split_len:] child.key = child.key[split_len:]
child.value = child.value[split_len:].clone() child.value = child.value[split_len:].clone()
new_node.parent.children[self.get_child_key_fn(key)] = new_node new_node.parent.children[key.child_key(self.page_size)] = new_node
# Split hash_value if it was already computed, otherwise leave as None # Split hash_value if it was already computed, otherwise leave as None
new_node.hash_value, child.hash_value = split_node_hash_value( new_node.hash_value, child.hash_value = split_node_hash_value(
@@ -807,13 +764,13 @@ class RadixCache(BasePrefixCache):
if len(key) == 0: if len(key) == 0:
return 0 return 0
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
total_prefix_length = 0 total_prefix_length = 0
while len(key) > 0 and child_key in node.children.keys(): while len(key) > 0 and child_key in node.children.keys():
node = node.children[child_key] node = node.children[child_key]
node.last_access_time = access_time node.last_access_time = access_time
prefix_len = self.key_match_fn(node.key, key) prefix_len = node.key.match(key, page_size=self.page_size)
total_prefix_length += prefix_len total_prefix_length += prefix_len
key = key[prefix_len:] key = key[prefix_len:]
value = value[prefix_len:] value = value[prefix_len:]
@@ -827,7 +784,7 @@ class RadixCache(BasePrefixCache):
node.priority = max(node.priority, priority) node.priority = max(node.priority, priority)
self._inc_hit_count(node, chunked) self._inc_hit_count(node, chunked)
if len(key): if len(key):
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
if len(key): if len(key):
new_node = TreeNode(priority=priority) new_node = TreeNode(priority=priority)
@@ -857,12 +814,12 @@ class RadixCache(BasePrefixCache):
for key, child in current_node.children.items(): for key, child in current_node.children.items():
stack.append((child, current_indent + 2)) stack.append((child, current_indent + 2))
assert key == self.get_child_key_fn( assert key == child.key.child_key(
child.key self.page_size
), f"{key=}, {self.get_child_key_fn(child.key)=}" ), f"{key=}, {child.key.child_key(self.page_size)=}"
def _delete_leaf(self, node): def _delete_leaf(self, node):
key = self.get_child_key_fn(node.key) key = node.key.child_key(self.page_size)
v = node.parent.children.pop(key, None) v = node.parent.children.pop(key, None)
assert v == node, f"parent does not have child key, {key}" assert v == node, f"parent does not have child key, {key}"
python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py
+1 -1
View File
@@ -194,7 +194,7 @@ class LMCRadixCache(RadixCache):
new_node.key = key[start:end] new_node.key = key[start:end]
new_node.value = token_slots[:fetched] new_node.value = token_slots[:fetched]
new_node.parent = last_node new_node.parent = last_node
last_node.children[self.get_child_key_fn(new_node.key)] = new_node last_node.children[new_node.key.child_key(self.page_size)] = new_node
last_node = new_node last_node = new_node
value = torch.cat([value, token_slots[:fetched]]) value = torch.cat([value, token_slots[:fetched]])
python/sglang/srt/mem_cache/swa_radix_cache.py
+15 -28
View File
@@ -22,7 +22,6 @@ The radix tree data structure for managing the hybrid (full and SWA) KV cache.
import heapq import heapq
import time import time
from collections import defaultdict from collections import defaultdict
from functools import partial
from typing import TYPE_CHECKING, List, Optional, Tuple from typing import TYPE_CHECKING, List, Optional, Tuple
import torch import torch
@@ -41,12 +40,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
MatchResult, MatchResult,
) )
from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.radix_cache import ( from sglang.srt.mem_cache.radix_cache import RadixKey
RadixKey,
_key_match_page_size1,
_key_match_paged,
get_child_key,
)
from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.utils import convert_to_bigram_key from sglang.srt.mem_cache.utils import convert_to_bigram_key
@@ -353,13 +347,6 @@ class SWARadixCache(BasePrefixCache):
else: else:
self.device = torch.device("cpu") self.device = torch.device("cpu")
if self.page_size == 1:
self.key_match_fn = _key_match_page_size1
self.get_child_key_fn = get_child_key
else:
self.key_match_fn = partial(_key_match_paged, page_size=self.page_size)
self.get_child_key_fn = partial(get_child_key, page_size=self.page_size)
if self.is_eagle: if self.is_eagle:
self.key_convert_fn = convert_to_bigram_key self.key_convert_fn = convert_to_bigram_key
else: else:
@@ -795,7 +782,7 @@ class SWARadixCache(BasePrefixCache):
node is greater than or equal to the sliding window size. node is greater than or equal to the sliding window size.
""" """
node = self.root_node node = self.root_node
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
value = [] value = []
# for path connected to root without tombstone, always match, so set to inf # for path connected to root without tombstone, always match, so set to inf
@@ -813,7 +800,7 @@ class SWARadixCache(BasePrefixCache):
# reset match_len_since_tombstone if we hit a tombstone node # reset match_len_since_tombstone if we hit a tombstone node
match_len_since_tombstone = 0 match_len_since_tombstone = 0
prefix_len = self.key_match_fn(child.key, key) prefix_len = child.key.match(key, page_size=self.page_size)
if prefix_len < len(child.key): if prefix_len < len(child.key):
new_node = self._split_node(child.key, child, prefix_len) new_node = self._split_node(child.key, child, prefix_len)
value.append(new_node.value) value.append(new_node.value)
@@ -829,7 +816,7 @@ class SWARadixCache(BasePrefixCache):
key = key[prefix_len:] key = key[prefix_len:]
if len(key): if len(key):
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
# handle best_value_len and best_last_node, for the case that last node is fully matched # handle best_value_len and best_last_node, for the case that last node is fully matched
if match_len_since_tombstone >= self.sliding_window_size: if match_len_since_tombstone >= self.sliding_window_size:
@@ -887,7 +874,7 @@ class SWARadixCache(BasePrefixCache):
def _split_node(self, key: RadixKey, child: TreeNode, split_len: int) -> TreeNode: def _split_node(self, key: RadixKey, child: TreeNode, split_len: int) -> TreeNode:
# new_node -> child # new_node -> child
new_node = TreeNode() new_node = TreeNode()
new_node.children = {self.get_child_key_fn(key[split_len:]): child} new_node.children = {key[split_len:].child_key(self.page_size): child}
new_node.parent = child.parent new_node.parent = child.parent
new_node.swa_tombstone = child.swa_tombstone new_node.swa_tombstone = child.swa_tombstone
new_node.full_lock_ref = child.full_lock_ref new_node.full_lock_ref = child.full_lock_ref
@@ -909,7 +896,7 @@ class SWARadixCache(BasePrefixCache):
child.key = child.key[split_len:] child.key = child.key[split_len:]
assert len(child.key) > 0, f"child.key should not be empty" assert len(child.key) > 0, f"child.key should not be empty"
child.value = child.value[split_len:].clone() child.value = child.value[split_len:].clone()
new_node.parent.children[self.get_child_key_fn(key)] = new_node new_node.parent.children[key.child_key(self.page_size)] = new_node
# insert the new node and child into the lru lists, insert # insert the new node and child into the lru lists, insert
# parent first so that parent is after child in the lru list # parent first so that parent is after child in the lru list
@@ -938,7 +925,7 @@ class SWARadixCache(BasePrefixCache):
if len(key) == 0: if len(key) == 0:
return 0 return 0
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
total_prefix_length = 0 total_prefix_length = 0
while len(key) > 0 and child_key in node.children.keys(): while len(key) > 0 and child_key in node.children.keys():
@@ -947,7 +934,7 @@ class SWARadixCache(BasePrefixCache):
self.full_lru_list.reset_node_mru(node) self.full_lru_list.reset_node_mru(node)
if not node.swa_tombstone: if not node.swa_tombstone:
self.swa_lru_list.reset_node_mru(node) self.swa_lru_list.reset_node_mru(node)
prefix_len = self.key_match_fn(node.key, key) prefix_len = node.key.match(key, page_size=self.page_size)
if prefix_len < len(node.key): if prefix_len < len(node.key):
new_node = self._split_node(node.key, node, prefix_len) new_node = self._split_node(node.key, node, prefix_len)
@@ -1002,7 +989,7 @@ class SWARadixCache(BasePrefixCache):
value = value[prefix_len:] value = value[prefix_len:]
if len(key): if len(key):
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
if len(key): if len(key):
# Layout: |--- total_prefix_length ---|--- len(key) ---| # Layout: |--- total_prefix_length ---|--- len(key) ---|
@@ -1055,7 +1042,7 @@ class SWARadixCache(BasePrefixCache):
new_node.key = key new_node.key = key
new_node.value = value.clone() new_node.value = value.clone()
new_node.swa_tombstone = swa_tombstone new_node.swa_tombstone = swa_tombstone
parent.children[self.get_child_key_fn(key)] = new_node parent.children[key.child_key(self.page_size)] = new_node
self.full_lru_list.insert_mru(new_node) self.full_lru_list.insert_mru(new_node)
self.full_evictable_size_ += len(value) self.full_evictable_size_ += len(value)
if not swa_tombstone: if not swa_tombstone:
@@ -1091,7 +1078,7 @@ class SWARadixCache(BasePrefixCache):
not node.swa_tombstone not node.swa_tombstone
), f"Invariant violated: leaf node is a tombstone, {node.id=}" ), f"Invariant violated: leaf node is a tombstone, {node.id=}"
assert len(node.children) == 0, f"leaf node has children, {node.id=}" assert len(node.children) == 0, f"leaf node has children, {node.id=}"
key = self.get_child_key_fn(node.key) key = node.key.child_key(self.page_size)
v = node.parent.children.pop(key, None) v = node.parent.children.pop(key, None)
assert v == node, f"parent does not have child key, {key}" assert v == node, f"parent does not have child key, {key}"
self.full_evictable_size_ -= len(node.key) self.full_evictable_size_ -= len(node.key)
@@ -1107,7 +1094,7 @@ class SWARadixCache(BasePrefixCache):
node.swa_tombstone node.swa_tombstone
), f"Deleting a unexpected non-tombstone leaf node, {node.id=}" ), f"Deleting a unexpected non-tombstone leaf node, {node.id=}"
assert len(node.children) == 0, f"leaf node has children, {node.id=}" assert len(node.children) == 0, f"leaf node has children, {node.id=}"
key = self.get_child_key_fn(node.key) key = node.key.child_key(self.page_size)
v = node.parent.children.pop(key, None) v = node.parent.children.pop(key, None)
assert v == node, f"parent does not have child key, {key}" assert v == node, f"parent does not have child key, {key}"
@@ -1152,9 +1139,9 @@ class SWARadixCache(BasePrefixCache):
for key, child in current_node.children.items(): for key, child in current_node.children.items():
stack.append((child, current_indent + 2)) stack.append((child, current_indent + 2))
assert key == self.get_child_key_fn( assert key == child.key.child_key(
child.key self.page_size
), f"{key=}, {self.get_child_key_fn(child.key)=}" ), f"{key=}, {child.key.child_key(self.page_size)=}"
def _total_size_helper(self) -> Tuple[int, int]: def _total_size_helper(self) -> Tuple[int, int]:
total_size = 0 total_size = 0
python/sglang/srt/mem_cache/unified_cache_components/README.md
+1 -1
View File
@@ -93,7 +93,7 @@ Find the longest cached prefix for a token sequence.
**Algorithm detail:** **Algorithm detail:**
1. Calls `create_match_validator()` once per component — returns a stateful closure (e.g., SWA tracks accumulated window length) 1. Calls `create_match_validator()` once per component — returns a stateful closure (e.g., SWA tracks accumulated window length)
2. Walks tree edges via `key_match_fn`; at each node, calls all validator closures — the match boundary is only advanced when **all** validators return `True` 2. Walks tree edges via `RadixKey.match()`; at each node, calls all validator closures — the match boundary is only advanced when **all** validators return `True`
3. If match ends mid-node, calls `_split_node` → triggers `redistribute_on_node_split()` per component 3. If match ends mid-node, calls `_split_node` → triggers `redistribute_on_node_split()` per component
4. Post-match (`_match_post_processor`): 4. Post-match (`_match_post_processor`):
- Promotes matched path to MRU in each component's LRU via `node_has_component_data()` as filter - Promotes matched path to MRU in each component's LRU via `node_has_component_data()` as filter
python/sglang/srt/mem_cache/unified_radix_cache.py
+14 -26
View File
@@ -20,12 +20,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
MatchPrefixParams, MatchPrefixParams,
MatchResult, MatchResult,
) )
from sglang.srt.mem_cache.radix_cache import ( from sglang.srt.mem_cache.radix_cache import RadixKey
RadixKey,
_key_match_page_size1,
_key_match_paged,
get_child_key,
)
from sglang.srt.mem_cache.unified_cache_components import ( from sglang.srt.mem_cache.unified_cache_components import (
_NUM_COMPONENT_TYPES, _NUM_COMPONENT_TYPES,
BASE_COMPONENT_TYPE, BASE_COMPONENT_TYPE,
@@ -188,13 +183,6 @@ class UnifiedRadixCache(BasePrefixCache):
if params.enable_metrics: if params.enable_metrics:
self.init_metrics_collector() self.init_metrics_collector()
if self.page_size == 1:
self.key_match_fn = _key_match_page_size1
self.get_child_key_fn = get_child_key
else:
self.key_match_fn = partial(_key_match_paged, page_size=self.page_size)
self.get_child_key_fn = partial(get_child_key, page_size=self.page_size)
assert params.tree_components is not None assert params.tree_components is not None
self.tree_components = tuple(params.tree_components) self.tree_components = tuple(params.tree_components)
self.components: dict[ComponentType, TreeComponent] = { self.components: dict[ComponentType, TreeComponent] = {
@@ -484,7 +472,7 @@ class UnifiedRadixCache(BasePrefixCache):
Not used yet; reserved for future read-only match operations.""" Not used yet; reserved for future read-only match operations."""
node = self.root_node node = self.root_node
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
value: list[torch.Tensor] = [] value: list[torch.Tensor] = []
best_value_len = 0 best_value_len = 0
best_node = node best_node = node
@@ -500,7 +488,7 @@ class UnifiedRadixCache(BasePrefixCache):
while len(key) > 0 and child_key in node.children: while len(key) > 0 and child_key in node.children:
child = node.children[child_key] child = node.children[child_key]
prefix_len = self.key_match_fn(child.key, key) prefix_len = child.key.match(key, page_size=self.page_size)
if prefix_len < len(child.key): if prefix_len < len(child.key):
# Read-only: do not split, ignore partial match and stop # Read-only: do not split, ignore partial match and stop
break break
@@ -509,14 +497,14 @@ class UnifiedRadixCache(BasePrefixCache):
_update_best_if_valid(node) _update_best_if_valid(node)
key = key[prefix_len:] key = key[prefix_len:]
if len(key): if len(key):
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
return value, best_node, best_value_len return value, best_node, best_value_len
def _match_prefix_helper( def _match_prefix_helper(
self, key: RadixKey self, key: RadixKey
) -> tuple[list[torch.Tensor], UnifiedTreeNode, int]: ) -> tuple[list[torch.Tensor], UnifiedTreeNode, int]:
node = self.root_node node = self.root_node
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
value: list[torch.Tensor] = [] value: list[torch.Tensor] = []
best_value_len = 0 best_value_len = 0
best_node = node best_node = node
@@ -532,7 +520,7 @@ class UnifiedRadixCache(BasePrefixCache):
while len(key) > 0 and child_key in node.children: while len(key) > 0 and child_key in node.children:
child = node.children[child_key] child = node.children[child_key]
prefix_len = self.key_match_fn(child.key, key) prefix_len = child.key.match(key, page_size=self.page_size)
if prefix_len < len(child.key): if prefix_len < len(child.key):
node = self._split_node(child.key, child, prefix_len) node = self._split_node(child.key, child, prefix_len)
value.append(node.component_data[BASE_COMPONENT_TYPE].value) value.append(node.component_data[BASE_COMPONENT_TYPE].value)
@@ -543,7 +531,7 @@ class UnifiedRadixCache(BasePrefixCache):
_update_best_if_valid(node) _update_best_if_valid(node)
key = key[prefix_len:] key = key[prefix_len:]
if len(key): if len(key):
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
return value, best_node, best_value_len return value, best_node, best_value_len
def _match_post_processor( def _match_post_processor(
@@ -587,7 +575,7 @@ class UnifiedRadixCache(BasePrefixCache):
self, key: RadixKey, child: UnifiedTreeNode, split_len: int self, key: RadixKey, child: UnifiedTreeNode, split_len: int
) -> UnifiedTreeNode: ) -> UnifiedTreeNode:
new_node = UnifiedTreeNode(self.tree_components) new_node = UnifiedTreeNode(self.tree_components)
new_node.children = {self.get_child_key_fn(key[split_len:]): child} new_node.children = {key[split_len:].child_key(self.page_size): child}
new_node.parent = child.parent new_node.parent = child.parent
new_node.key = child.key[:split_len] new_node.key = child.key[:split_len]
new_node.component_data[BASE_COMPONENT_TYPE].value = ( new_node.component_data[BASE_COMPONENT_TYPE].value = (
@@ -604,7 +592,7 @@ class UnifiedRadixCache(BasePrefixCache):
for component in self._components_tuple: for component in self._components_tuple:
component.redistribute_on_node_split(new_parent=new_node, child=child) component.redistribute_on_node_split(new_parent=new_node, child=child)
new_node.parent.children[self.get_child_key_fn(key)] = new_node new_node.parent.children[key.child_key(self.page_size)] = new_node
self._for_each_component_lru(new_node, UnifiedLRUList.insert_mru) self._for_each_component_lru(new_node, UnifiedLRUList.insert_mru)
self._for_each_component_lru(child, UnifiedLRUList.insert_mru) self._for_each_component_lru(child, UnifiedLRUList.insert_mru)
@@ -626,7 +614,7 @@ class UnifiedRadixCache(BasePrefixCache):
new_node.parent = parent new_node.parent = parent
new_node.key = key new_node.key = key
new_node.component_data[BASE_COMPONENT_TYPE].value = value.clone() new_node.component_data[BASE_COMPONENT_TYPE].value = value.clone()
parent.children[self.get_child_key_fn(key)] = new_node parent.children[key.child_key(self.page_size)] = new_node
self.lru_lists[BASE_COMPONENT_TYPE].insert_mru(new_node) self.lru_lists[BASE_COMPONENT_TYPE].insert_mru(new_node)
self.component_evictable_size_[BASE_COMPONENT_TYPE] += len(value) self.component_evictable_size_[BASE_COMPONENT_TYPE] += len(value)
return new_node return new_node
@@ -642,12 +630,12 @@ class UnifiedRadixCache(BasePrefixCache):
if len(key) == 0: if len(key) == 0:
return InsertResult(prefix_len=0, mamba_exist=True) return InsertResult(prefix_len=0, mamba_exist=True)
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
total_prefix_length = 0 total_prefix_length = 0
while len(key) > 0 and child_key in node.children: while len(key) > 0 and child_key in node.children:
node = node.children[child_key] node = node.children[child_key]
self._touch_node(node) self._touch_node(node)
prefix_len = self.key_match_fn(node.key, key) prefix_len = node.key.match(key, page_size=self.page_size)
if prefix_len < len(node.key): if prefix_len < len(node.key):
node = self._split_node(node.key, node, prefix_len) node = self._split_node(node.key, node, prefix_len)
@@ -674,7 +662,7 @@ class UnifiedRadixCache(BasePrefixCache):
key = key[prefix_len:] key = key[prefix_len:]
value = value[prefix_len:] value = value[prefix_len:]
if len(key): if len(key):
child_key = self.get_child_key_fn(key) child_key = key.child_key(self.page_size)
is_new_leaf = False is_new_leaf = False
# Create new leaf for remaining suffix # Create new leaf for remaining suffix
@@ -737,7 +725,7 @@ class UnifiedRadixCache(BasePrefixCache):
self._iteratively_delete_tombstone_leaf(node, tracker) self._iteratively_delete_tombstone_leaf(node, tracker)
def _remove_leaf_from_parent(self, node: UnifiedTreeNode): def _remove_leaf_from_parent(self, node: UnifiedTreeNode):
key = self.get_child_key_fn(node.key) key = node.key.child_key(self.page_size)
v = node.parent.children.pop(key, None) v = node.parent.children.pop(key, None)
assert v == node assert v == node
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+1 -1
View File
@@ -632,7 +632,7 @@ class UnifiedRadixCacheSuite:
self.skipTest("page_size > 1 only") self.skipTest("page_size > 1 only")
tree, _, _ = build_fixture(self.cfg) tree, _, _ = build_fixture(self.cfg)
key = RadixKey(self._make_seq(1, 1)) key = RadixKey(self._make_seq(1, 1))
child_key = tree.get_child_key_fn(key) child_key = key.child_key(tree.page_size)
self.assertIsInstance(child_key, tuple) self.assertIsInstance(child_key, tuple)
def test_paged_match_truncates_unaligned_key(self): def test_paged_match_truncates_unaligned_key(self):