diff --git a/python/sglang/srt/mem_cache/unified_cache_components/full_component.py b/python/sglang/srt/mem_cache/unified_cache_components/full_component.py index 3263cfbd1..bccb866a7 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/full_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/full_component.py @@ -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 ) diff --git a/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py b/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py index c1ef99b88..7142bd9b3 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py @@ -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 diff --git a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py index 8a0d6bac2..63223625e 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py @@ -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) diff --git a/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py b/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py index 317ee3012..ae6f71167 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py @@ -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.""" diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index f275c9925..8c5870281 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -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: diff --git a/test/registered/radix_cache/test_unified_radix_hicache_kl.py b/test/registered/radix_cache/test_unified_radix_hicache_kl.py index c6e5cdb16..841af8733 100644 --- a/test/registered/radix_cache/test_unified_radix_hicache_kl.py +++ b/test/registered/radix_cache/test_unified_radix_hicache_kl.py @@ -155,7 +155,7 @@ class TestUnifiedDeepSeekV4FlashHiCache(UnifiedRadixTreeTestMixin, CustomTestCas "--max-total-tokens", "20000", "--max-running-requests", - "4", + "2", ], env={ "SGLANG_DSV4_FP4_EXPERTS": "0", diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index a96c8a82a..7074084d9 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -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):