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

[UnifiedTree]: Fix UnifiedRadixCache device match semantics with HiCache (#25277)

Browse Source
This commit is contained in:
Zhangheng
2026-05-16 00:40:55 +08:00
committed by GitHub
parent 3f7e538b2f
commit 21b3ac52b4
7 changed files with 398 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/unified_cache_components/full_component.py
+9 -2
View File
@@ -41,8 +41,15 @@ class FullComponent(TreeComponent):
# HiCache state: set to host KV pool when HiCache enabled
self._full_kv_pool_host = None
def create_match_validator(self) -> Callable[[UnifiedTreeNode], bool]:
# HiCache: evicted + backuped nodes are valid match boundaries
def create_match_validator(
self, match_device_only: bool = False
) -> Callable[[UnifiedTreeNode], bool]:
if match_device_only:
return (
lambda node: node.component_data[self.component_type].value is not None
)
# HiCache: evicted + backuped nodes are valid match boundaries.
return lambda node: (
node.component_data[self.component_type].value is not None or node.backuped
)
python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py
+10 -2
View File
@@ -50,8 +50,13 @@ class MambaComponent(TreeComponent):
# HiCache state
self._mamba_pool_host = None # set to host mamba pool when HiCache enabled
def create_match_validator(self) -> Callable[[UnifiedTreeNode], bool]:
def create_match_validator(
self, match_device_only: bool = False
) -> Callable[[UnifiedTreeNode], bool]:
ct = self.component_type
if match_device_only:
return lambda node: node.component_data[ct].value is not None
# HiCache: evicted + backuped (host_value present) is also a valid match
return lambda node: (
node.component_data[ct].value is not None
@@ -69,7 +74,10 @@ class MambaComponent(TreeComponent):
req = params.req
last_node = result.best_match_node
if len(value_chunks) > best_value_len:
# HiCache can still use prefix matches and load back host-backed Mamba
# states. We temporarily skip branching-state fill in that mode and can
# add a HiCache-aware branching policy later.
if self.cache.cache_controller is None and len(value_chunks) > best_value_len:
chunk_size = get_global_server_args().mamba_cache_chunk_size
aligned_seqlen = (
sum(len(v) for v in value_chunks) // chunk_size
python/sglang/srt/mem_cache/unified_cache_components/swa_component.py
+4 -2
View File
@@ -69,7 +69,9 @@ class SWAComponent(TreeComponent):
self.cache.lru_lists[ct].insert_mru(node)
self.cache.component_evictable_size_[ct] += len(value)
def create_match_validator(self) -> Callable[[UnifiedTreeNode], bool]:
def create_match_validator(
self, match_device_only: bool = False
) -> Callable[[UnifiedTreeNode], bool]:
sliding_window_size = self.sliding_window_size
ct = self.component_type
state = {"len": float("inf")}
@@ -78,7 +80,7 @@ class SWAComponent(TreeComponent):
cd = node.component_data[ct]
# HiCache: a host-only tombstone is a valid match boundary too
# — load_back will restore SWA from host before use.
if cd.value is None and cd.host_value is None:
if cd.value is None and (match_device_only or cd.host_value is None):
state["len"] = 0
return False
state["len"] += len(node.key)
python/sglang/srt/mem_cache/unified_cache_components/tree_component.py
+6 -2
View File
@@ -116,11 +116,15 @@ class TreeComponent(ABC):
return len(value) if value is not None else 0
@abstractmethod
def create_match_validator(self) -> Callable[[UnifiedTreeNode], bool]:
def create_match_validator(
self, match_device_only: bool = False
) -> Callable[[UnifiedTreeNode], bool]:
"""Return a per-match stateful predicate that decides whether a node
is a valid match boundary for this component.
Called once per match_prefix; the returned closure may carry state.
- Full: always True (every node is valid).
When match_device_only is true, host-backed nodes must not be accepted
as valid match boundaries.
- Full: returns True if the node has full component data.
- SWA: tracks accumulated length since last gap; returns True only
when the contiguous window reaches swa_sliding_window_size.
- Mamba: returns True iff the node has mamba component data."""
python/sglang/srt/mem_cache/unified_radix_cache.py
+100 -84
View File
@@ -351,9 +351,18 @@ class UnifiedRadixCache(BasePrefixCache):
if len(key) == 0:
return self._empty_match_result
value, best_match_node, best_value_len = self._match_prefix_helper(key)
(
value,
best_match_node,
best_match_device_node,
best_match_device_value_len,
) = self._match_prefix_helper(key)
return self._match_post_processor(
params, value, best_match_node, best_value_len
params,
value,
best_match_node,
best_match_device_node,
best_match_device_value_len,
)
def insert(self, params: InsertParams) -> InsertResult:
@@ -585,67 +594,53 @@ class UnifiedRadixCache(BasePrefixCache):
# ---- Internal Helpers ----
def _match_prefix_helper_readonly(
self, key: RadixKey
) -> tuple[list[torch.Tensor], UnifiedTreeNode, int]:
"""Read-only version of _match_prefix_helper that does not split nodes.
Only considers fully matched nodes, ignores partial matches.
Not used yet; reserved for future read-only match operations."""
node = self.root_node
child_key = key.child_key(self.page_size)
value: list[torch.Tensor] = []
best_value_len = 0
best_match_node = node
validators = tuple(
comp.create_match_validator() for comp in self._components_tuple
)
def _update_best_if_valid(node):
nonlocal best_value_len, best_match_node
if all(v(node) for v in validators):
best_value_len = len(value)
best_match_node = node
while len(key) > 0 and child_key in node.children:
child = node.children[child_key]
# HiCache: dead node (evicted + not backuped) — stop traversal
if child.evicted and not child.backuped:
break
prefix_len = child.key.match(key, page_size=self.page_size)
if prefix_len < len(child.key):
# Read-only: do not split, ignore partial match and stop
break
if not child.evicted:
value.append(child.component_data[BASE_COMPONENT_TYPE].value)
node = child
_update_best_if_valid(node)
key = key[prefix_len:]
if len(key):
child_key = key.child_key(self.page_size)
return value, best_match_node, best_value_len
def _match_prefix_helper(
self, key: RadixKey
) -> tuple[list[torch.Tensor], UnifiedTreeNode, int]:
) -> tuple[list[torch.Tensor], UnifiedTreeNode, UnifiedTreeNode, int]:
# Non-HiCache mode has only device-resident matches, so the scheduler
# device anchor follows the best match. In HiCache mode, host-backed
# nodes can also match, so we separately track the best device-resident
# match for scheduler prefix indices and locking.
node = self.root_node
child_key = key.child_key(self.page_size)
value: list[torch.Tensor] = []
best_value_len = 0
best_match_node = node
validators = tuple(
comp.create_match_validator() for comp in self._components_tuple
)
best_match_device_node = node
best_match_device_value_len = 0
separate_device_match = self.cache_controller is not None
if separate_device_match:
validators = tuple(
comp.create_match_validator() for comp in self._components_tuple
)
device_validators = tuple(
comp.create_match_validator(match_device_only=True)
for comp in self._components_tuple
)
else:
validators = tuple(
comp.create_match_validator(match_device_only=True)
for comp in self._components_tuple
)
def _all_valid(validators, node):
return all([v(node) for v in validators])
def _update_best_if_valid(node):
nonlocal best_value_len, best_match_node
if all(v(node) for v in validators):
best_value_len = len(value)
nonlocal best_match_node
nonlocal best_match_device_value_len, best_match_device_node
matched = _all_valid(validators, node)
if matched:
best_match_node = node
if not separate_device_match:
if matched:
best_match_device_value_len = len(value)
best_match_device_node = node
return
if _all_valid(device_validators, node):
best_match_device_value_len = len(value)
best_match_device_node = node
while len(key) > 0 and child_key in node.children:
child = node.children[child_key]
@@ -668,14 +663,21 @@ class UnifiedRadixCache(BasePrefixCache):
key = key[prefix_len:]
if len(key):
child_key = key.child_key(self.page_size)
return value, best_match_node, best_value_len
return (
value,
best_match_node,
best_match_device_node,
best_match_device_value_len,
)
def _match_post_processor(
self,
params: MatchPrefixParams,
value: list[torch.Tensor],
best_match_node: UnifiedTreeNode,
best_value_len: int,
best_match_device_node: UnifiedTreeNode,
best_match_device_value_len: int,
) -> MatchResult:
node_update = best_match_node
for comp in self._components_tuple:
@@ -691,23 +693,21 @@ class UnifiedRadixCache(BasePrefixCache):
cur_time -= 0.00001
node_update = node_update.parent
# Walk up to find last_device_node
last_device_node = best_match_node
while last_device_node is not self.root_node and last_device_node.evicted:
last_device_node = last_device_node.parent
# Walk up to find last_host_node for full component.
if self.cache_controller is None:
last_host_node = best_match_device_node
else:
last_host_node = best_match_node
while last_host_node is not self.root_node and not last_host_node.backuped:
last_host_node = last_host_node.parent
# Walk up to find last_host_node
last_host_node = best_match_node
while last_host_node is not self.root_node and not last_host_node.backuped:
last_host_node = last_host_node.parent
if best_value_len > 0:
device_indices = torch.cat(value[:best_value_len])
if best_match_device_value_len > 0:
device_indices = torch.cat(value[:best_match_device_value_len])
else:
device_indices = self._empty_match_result.device_indices
result = MatchResult(
device_indices=device_indices,
last_device_node=last_device_node,
last_device_node=best_match_device_node,
last_host_node=last_host_node,
best_match_node=best_match_node,
host_hit_length=0,
@@ -718,7 +718,7 @@ class UnifiedRadixCache(BasePrefixCache):
result=result,
params=params,
value_chunks=value,
best_value_len=best_value_len,
best_value_len=best_match_device_value_len,
)
return result
@@ -1219,10 +1219,10 @@ class UnifiedRadixCache(BasePrefixCache):
best_match_node: UnifiedTreeNode,
mem_quota: Optional[int] = None,
req=None,
) -> Optional[torch.Tensor]:
) -> bool:
"""Load evicted KV data from host back to device (H→D)."""
if self.cache_controller is None:
return None
return False
# Build KV transfer
kv_xfer = self.components[BASE_COMPONENT_TYPE].build_hicache_transfers(
@@ -1255,7 +1255,7 @@ class UnifiedRadixCache(BasePrefixCache):
mem_quota is not None and kv_tokens > mem_quota + result.delta
):
self.dec_lock_ref(best_match_node, ancestor_lock_params)
return None
return False
avail = self.token_to_kv_pool_allocator.available_size()
if avail < kv_tokens:
@@ -1263,7 +1263,7 @@ class UnifiedRadixCache(BasePrefixCache):
result = self.evict(EvictParams(num_tokens=needed))
if result.num_tokens_evicted < needed:
self.dec_lock_ref(best_match_node, ancestor_lock_params)
return None
return False
# Load H→D
aux_xfers = [x for xfers in comp_xfers.values() for x in xfers]
@@ -1276,7 +1276,7 @@ class UnifiedRadixCache(BasePrefixCache):
self.dec_lock_ref(best_match_node, ancestor_lock_params)
if device_indices is None:
return None
return False
# Commit: each component gets only its own transfers
kv_xfer.device_indices = device_indices
@@ -1297,7 +1297,7 @@ class UnifiedRadixCache(BasePrefixCache):
best_match_node,
self.inc_lock_ref(best_match_node).to_dec_params(),
)
return device_indices
return True
def _build_sidecar_transfers(
self,
@@ -1432,25 +1432,41 @@ class UnifiedRadixCache(BasePrefixCache):
best_match_node = params.best_match_node
mem_quota = params.mem_quota
req = params.req
assert req is not None
last_best_match_device_node = req.last_node
def _collect_new_prefix_indices() -> torch.Tensor:
prefix_chunks: list[torch.Tensor] = []
node = best_match_node
while node is not last_best_match_device_node:
value = node.component_data[BASE_COMPONENT_TYPE].value
assert value is not None
prefix_chunks.append(value)
node = node.parent
if not prefix_chunks:
return self._empty_match_result.device_indices
prefix_chunks.reverse()
return torch.cat(prefix_chunks)
if best_match_node.evicted or params.host_hit_length > 0:
loading_values = self.load_back(best_match_node, mem_quota, req=req)
if loading_values is not None:
if self.load_back(best_match_node, mem_quota, req=req):
new_indices = _collect_new_prefix_indices()
if new_indices.numel() == 0:
return (
self._empty_match_result.device_indices,
last_best_match_device_node,
)
logger.debug(
"init_load_back success: loaded %d tokens for node %d",
len(loading_values),
len(new_indices),
best_match_node.id,
)
return loading_values, best_match_node
# Fallback: walk up to non-evicted ancestor
# TODO(ispobock): The fallback path is not correct. The last_device_node should consider all the components.
while best_match_node is not self.root_node and best_match_node.evicted:
best_match_node = best_match_node.parent
return new_indices, best_match_node
return (
self._empty_match_result.device_indices,
best_match_node,
last_best_match_device_node,
)
def check_hicache_events(self) -> None:
test/registered/radix_cache/test_unified_radix_hicache_kl.py
+1 -1
View File
@@ -155,7 +155,7 @@ class TestUnifiedDeepSeekV4FlashHiCache(UnifiedRadixTreeTestMixin, CustomTestCas
"--max-total-tokens",
"20000",
"--max-running-requests",
"4",
"2",
],
env={
"SGLANG_DSV4_FP4_EXPERTS": "0",
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+268 -17
View File
@@ -15,6 +15,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
DecLockRefParams,
EvictParams,
EvictResult,
InitLoadBackParams,
InsertParams,
MatchPrefixParams,
MatchResult,
@@ -908,6 +909,8 @@ class UnifiedRadixCacheSuite:
"""Verify readonly match does not modify tree structure (no split)."""
if self.cfg.page_size > 1 or self.cfg.has_mamba or self.cfg.has_swa:
self.skipTest("Full-only page_size=1 only")
if not hasattr(UnifiedRadixCache, "_match_prefix_helper_readonly"):
self.skipTest("_match_prefix_helper_readonly is not available")
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
self._insert(tree, allocator, req_to_token_pool, [1, 2, 3, 4, 5])
@@ -922,19 +925,27 @@ class UnifiedRadixCacheSuite:
self.assertEqual(node_count_before, 2)
tree._match_prefix_helper(RadixKey([1, 2]))
value, best_match_node, best_value_len = tree._match_prefix_helper(
RadixKey([1, 2, 3, 4])
)
(
value,
best_match_node,
best_match_device_node,
best_value_len,
) = tree._match_prefix_helper(RadixKey([1, 2, 3, 4]))
self.assertEqual(best_value_len, 2)
self.assertEqual(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)
value, best_match_node, best_value_len = tree._match_prefix_helper_readonly(
RadixKey([1, 2, 3])
)
(
value,
best_match_node,
best_match_device_node,
best_value_len,
) = tree._match_prefix_helper_readonly(RadixKey([1, 2, 3]))
self.assertEqual(best_value_len, 1)
self.assertEqual(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)
@@ -1258,8 +1269,8 @@ class UnifiedRadixCacheSuite:
# ================================================================
def _skip_unsupported_hicache_test(self):
if self.cfg.has_swa:
self.skipTest("HiCache tests do not run on SWA stacks")
if self.cfg.has_swa and self.cfg.has_mamba:
self.skipTest("HiCache unit fixture does not support SWA + Mamba stacks")
return False
def _simulate_backup(self, tree, node):
@@ -1343,14 +1354,14 @@ class UnifiedRadixCacheSuite:
self._backup_node(tree, node)
def _load_back_node(self, tree, node):
device_indices = tree.load_back(node)
self.assertIsNotNone(device_indices)
loaded = tree.load_back(node)
self.assertTrue(loaded)
producer_id = tree.ready_to_load_host_cache()
self.assertNotEqual(producer_id, -1)
for _, finish_event, _ in list(tree.cache_controller.ack_load_queue):
finish_event.synchronize()
tree.loading_check()
return device_indices
return node.component_data[ComponentType.FULL].value
def _get_full_kv_pool(self, allocator):
kv_pool = allocator.get_kvcache()
@@ -1474,7 +1485,9 @@ class UnifiedRadixCacheSuite:
def test_hicache_partial_match_splits_evicted_backed_up_node(self):
"""Partial matches on host-only nodes must keep the host prefix usable."""
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
if self._skip_unsupported_hicache_test():
return
tree, allocator, req_to_token_pool = self._build_hicache_fixture()
ps = self.cfg.page_size
seq = self._make_seq(1, 4)
expected_prefix = seq[: 2 * ps]
@@ -1484,7 +1497,7 @@ class UnifiedRadixCacheSuite:
self._insert(tree, allocator, req_to_token_pool, seq)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq)))
node = m.last_device_node
self._simulate_backup(tree, node)
self._backup_node(tree, node)
tree.evict(EvictParams(num_tokens=len(seq)))
self.assertTrue(node.evicted)
@@ -1684,6 +1697,244 @@ class UnifiedRadixCacheSuite:
chain.reverse()
return chain
def _release_ongoing_load_back_locks(self, tree):
for node, lock_params in list(tree.ongoing_load_back.values()):
tree.dec_lock_ref(node, lock_params)
tree.ongoing_load_back.clear()
def _finish_pending_loads(self, tree):
producer_id = tree.ready_to_load_host_cache()
self.assertNotEqual(producer_id, -1)
for _, finish_event, _ in list(tree.cache_controller.ack_load_queue):
finish_event.synchronize()
tree.loading_check()
def _match_tokens_for_chain(self, chain):
tokens: list[int] = []
for node in chain:
tokens.extend(node.key.token_ids)
return tokens
def _set_aux_host_tombstone(self, tree, node, component_type):
cd = node.component_data[component_type]
self.assertIsNotNone(cd.value)
if cd.host_value is None:
cd.host_value = cd.value.clone()
old_value = cd.value
cd.value = None
if component_type in tree.lru_lists and tree.lru_lists[component_type].in_list(
node
):
tree.lru_lists[component_type].remove_node(node)
tree.host_lru_lists[component_type].insert_mru(node)
tree.component_evictable_size_[component_type] -= len(old_value)
def test_match_prefix_best_and_device_node_without_hicache(self):
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
ps = self.cfg.page_size
min_tokens = 2 * ps
if self.cfg.has_swa:
min_tokens = max(min_tokens, self.cfg.sliding_window_size + ps)
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)))
self.assertEqual(len(result.device_indices), len(seq))
self.assertIs(result.best_match_node, result.last_device_node)
self.assertIs(result.last_host_node, result.last_device_node)
self.assertEqual(result.host_hit_length, 0)
def test_hicache_mamba_host_best_match_keeps_device_anchor(self):
if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1:
self.skipTest("requires page_size=1 Full+Mamba")
tree, allocator, req_to_token_pool = self._build_hicache_fixture()
chain = self._build_chain_pages(tree, allocator, req_to_token_pool, 3)
if len(chain) < 3:
self.skipTest("chain too short")
leaf = chain[-1]
parent = chain[-2]
tokens = self._match_tokens_for_chain(chain)
self._backup_node(tree, leaf)
tree.evict(EvictParams(num_tokens=len(leaf.key)))
self.assertTrue(leaf.evicted)
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens)))
self.assertIs(result.best_match_node, leaf)
self.assertIs(result.last_device_node, parent)
self.assertEqual(len(result.device_indices), len(tokens) - len(leaf.key))
self.assertEqual(result.host_hit_length, len(leaf.key))
def test_hicache_swa_host_best_match_keeps_device_anchor(self):
if not self.cfg.has_swa or self.cfg.has_mamba or self.cfg.page_size != 1:
self.skipTest("requires page_size=1 Full+SWA")
tree, allocator, req_to_token_pool = self._build_hicache_fixture()
chain = self._build_chain_pages(tree, allocator, req_to_token_pool, 3)
if len(chain) < 3:
self.skipTest("chain too short")
leaf = chain[-1]
parent = chain[-2]
tokens = self._match_tokens_for_chain(chain)
self._backup_node(tree, leaf)
tree.evict(EvictParams(num_tokens=len(leaf.key)))
self.assertTrue(leaf.evicted)
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens)))
self.assertIs(result.best_match_node, leaf)
self.assertIs(result.last_device_node, parent)
self.assertEqual(len(result.device_indices), len(tokens) - len(leaf.key))
self.assertEqual(result.host_hit_length, 1)
def test_mamba_branching_seqlen_disabled_under_hicache(self):
if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1:
self.skipTest("requires page_size=1 Full+Mamba")
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
chunk_size = get_global_server_args().mamba_cache_chunk_size
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))
).last_device_node
mamba_cd = leaf.component_data[ComponentType.MAMBA]
mamba_cd.value = None
no_hicache = tree.match_prefix(MatchPrefixParams(key=RadixKey(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)
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))
).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)))
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)
def test_scheduler_hicache_full_mamba_init_load_back_appends_new_indices(self):
if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1:
self.skipTest("requires page_size=1 Full+Mamba")
tree, allocator, req_to_token_pool = self._build_hicache_fixture()
chain = self._build_chain_pages(tree, allocator, req_to_token_pool, 3)
if len(chain) < 3:
self.skipTest("chain too short")
leaf = chain[-1]
tokens = self._match_tokens_for_chain(chain)
self._backup_node(tree, leaf)
tree.evict(EvictParams(num_tokens=len(leaf.key)))
self.assertTrue(leaf.evicted)
req = self._make_req(req_to_token_pool)
match = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens), req=req))
req.prefix_indices = match.device_indices
req.last_node = match.last_device_node
req.best_match_node = match.best_match_node
req.host_hit_length = match.host_hit_length
new_indices, new_node = tree.init_load_back(
InitLoadBackParams(
best_match_node=req.best_match_node,
host_hit_length=req.host_hit_length,
req=req,
)
)
self.assertIs(new_node, leaf)
self.assertEqual(len(torch.cat([req.prefix_indices, new_indices])), len(tokens))
self.assertIsNotNone(leaf.component_data[ComponentType.MAMBA].value)
self._finish_pending_loads(tree)
self._release_ongoing_load_back_locks(tree)
def test_scheduler_hicache_aux_only_load_back_appends_full_device_indices(self):
if self.cfg.page_size != 1:
self.skipTest("page_size=1 keeps the expected suffix precise")
aux = None
if self.cfg.has_swa and not self.cfg.has_mamba:
aux = ComponentType.SWA
elif self.cfg.has_mamba and not self.cfg.has_swa:
aux = ComponentType.MAMBA
if aux is None:
self.skipTest("requires exactly one aux component")
tree, allocator, req_to_token_pool = self._build_hicache_fixture()
chain = self._build_chain_pages(tree, allocator, req_to_token_pool, 3)
if len(chain) < 3:
self.skipTest("chain too short")
leaf = chain[-1]
tokens = self._match_tokens_for_chain(chain)
leaf_full = leaf.component_data[ComponentType.FULL].value.clone()
self._backup_node(tree, leaf)
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))
req.prefix_indices = match.device_indices
req.last_node = match.last_device_node
req.best_match_node = match.best_match_node
req.host_hit_length = match.host_hit_length
new_indices, new_node = tree.init_load_back(
InitLoadBackParams(
best_match_node=req.best_match_node,
host_hit_length=req.host_hit_length,
req=req,
)
)
self.assertIs(new_node, leaf)
self.assertEqual(new_indices.tolist(), leaf_full.tolist())
self.assertEqual(len(torch.cat([req.prefix_indices, new_indices])), len(tokens))
self.assertEqual(
leaf.component_data[ComponentType.FULL].value.tolist(),
leaf_full.tolist(),
)
self.assertIsNotNone(leaf.component_data[aux].value)
self._finish_pending_loads(tree)
self._release_ongoing_load_back_locks(tree)
def test_scheduler_hicache_load_back_fallback_keeps_old_anchor(self):
if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1:
self.skipTest("requires page_size=1 Full+Mamba")
tree, allocator, req_to_token_pool = self._build_hicache_fixture()
chain = self._build_chain_pages(tree, allocator, req_to_token_pool, 3)
if len(chain) < 3:
self.skipTest("chain too short")
leaf = chain[-1]
tokens = self._match_tokens_for_chain(chain)
self._backup_node(tree, leaf)
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))
req.prefix_indices = match.device_indices
req.last_node = match.last_device_node
req.best_match_node = match.best_match_node
req.host_hit_length = match.host_hit_length
new_indices, new_node = tree.init_load_back(
InitLoadBackParams(
best_match_node=req.best_match_node,
host_hit_length=req.host_hit_length,
req=req,
mem_quota=-1_000_000,
)
)
self.assertEqual(len(new_indices), 0)
self.assertIs(new_node, match.last_device_node)
self.assertIsNone(leaf.component_data[ComponentType.FULL].value)
self.assertIsNone(leaf.component_data[ComponentType.MAMBA].value)
def test_hicache_swa_load_back_min_suffix(self):
"""LOAD_BACK collects only the suffix nodes needed to cover sliding_window_size."""
if not self.cfg.has_swa:
@@ -1940,11 +2191,11 @@ class UnifiedRadixCacheSuite:
if chain_pages * ps > self.cfg.kv_size // 2:
self.skipTest("kv_size too small for the desired chain")
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
tree, allocator, req_to_token_pool = self._build_hicache_fixture()
chain = self._build_chain_pages(tree, allocator, req_to_token_pool, chain_pages)
if len(chain) < chain_pages:
self.skipTest("chain too short")
self._simulate_backup_tree(tree)
self._backup_tree(tree)
x = chain[-1]
y = chain[-window_pages]
@@ -1962,10 +2213,10 @@ class UnifiedRadixCacheSuite:
return tree, chain, n, y, x, tokens
def test_hicache_swa_match_prefix_picks_best_match_node_above_last_host(self):
tree, _, _, y, x, tokens = self._swa_anchor_setup()
tree, _, n, y, x, tokens = self._swa_anchor_setup()
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens)))
self.assertIs(result.best_match_node, x)
self.assertIs(result.last_device_node, x)
self.assertIs(result.last_device_node, n.parent)
self.assertIs(result.last_host_node, y)
def test_hicache_swa_load_back_anchored_on_best_match_node(self):