diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index cbc6a93f1..d99e9adfa 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -177,6 +177,8 @@ class MatchResult(NamedTuple): mamba_branching_seqlen: The mamba radix cache branching point, which is the longest page-aligned position that could've been cache hit if there exists a mamba state. + full_kv_hit_length: Longest Full-KV prefix available on either device or + host, independent of other components. """ device_indices: torch.Tensor @@ -188,6 +190,7 @@ class MatchResult(NamedTuple): mamba_host_hit_length: int = 0 mamba_branching_seqlen: Optional[int] = None cache_protected_len: Optional[int] = None + full_kv_hit_length: int = 0 def zero_match_result(tree_cache, match_result: MatchResult) -> MatchResult: @@ -205,6 +208,7 @@ def zero_match_result(tree_cache, match_result: MatchResult) -> MatchResult: host_hit_length=0, swa_host_hit_length=0, mamba_host_hit_length=0, + full_kv_hit_length=0, ) 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 58af3f9f3..ea7fa3369 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 @@ -55,6 +55,7 @@ class MambaComponent(TreeComponent): super().__init__(cache, params) self.enable_mamba_extra_buffer = params.enable_mamba_extra_buffer self.enable_mamba_extra_buffer_lazy = params.enable_mamba_extra_buffer_lazy + self.mamba_cache_chunk_size = get_server_args().mamba_cache_chunk_size self.mamba_max_states_per_path = get_server_args().mamba_max_states_per_path # HiCache state self._mamba_pool_host = None # set to host mamba pool when HiCache enabled @@ -107,17 +108,17 @@ class MambaComponent(TreeComponent): req = params.req last_node = result.best_match_node - # 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_server_args().mamba_cache_chunk_size - aligned_seqlen = ( - sum(len(v) for v in value_chunks) // chunk_size - ) * chunk_size - branching_seqlen = aligned_seqlen if aligned_seqlen > 0 else None - else: - branching_seqlen = None + mamba_boundary_len = len(result.device_indices) + result.host_hit_length + + # Full KV may extend beyond the latest reusable Mamba state. The branching + # point is the last Mamba-cache-chunk-aligned position within the Full-KV hit + # that lies beyond the current Mamba boundary. + aligned_seqlen = ( + result.full_kv_hit_length // self.mamba_cache_chunk_size + ) * self.mamba_cache_chunk_size + branching_seqlen = ( + aligned_seqlen if aligned_seqlen > mamba_boundary_len else None + ) mamba_value = last_node.component_data[self.component_type].value if cow_mamba and mamba_value is not None: diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 7caf33b45..a0c67820a 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -592,6 +592,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): best_match_node, best_match_device_node, best_match_device_value_len, + full_kv_hit_length, ) = self._match_prefix_helper(key) return self._match_post_processor( params, @@ -599,6 +600,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): best_match_node, best_match_device_node, best_match_device_value_len, + full_kv_hit_length, ) def insert(self, params: InsertParams) -> InsertResult: @@ -909,7 +911,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): def _match_prefix_helper( self, key: RadixKey - ) -> tuple[list[torch.Tensor], UnifiedTreeNode, UnifiedTreeNode, int]: + ) -> tuple[list[torch.Tensor], UnifiedTreeNode, UnifiedTreeNode, int, 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 @@ -920,6 +922,8 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): best_match_node = node best_match_device_node = node best_match_device_value_len = 0 + full_kv_hit_length = 0 + separate_device_match = self.cache_controller is not None if separate_device_match: validators = tuple( @@ -962,6 +966,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): break prefix_len = child.key.match(key, page_size=self.page_size) + full_kv_hit_length += prefix_len if prefix_len < len(child.key): node = self._split_node(child.key, child, prefix_len) if not node.evicted: @@ -982,6 +987,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): best_match_node, best_match_device_node, best_match_device_value_len, + full_kv_hit_length, ) def _match_post_processor( @@ -991,6 +997,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): best_match_node: UnifiedTreeNode, best_match_device_node: UnifiedTreeNode, best_match_device_value_len: int, + full_kv_hit_length: int, ) -> MatchResult: node_update = best_match_node for comp in self._components_tuple: @@ -1024,6 +1031,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): last_host_node=last_host_node, best_match_node=best_match_node, host_hit_length=0, + full_kv_hit_length=full_kv_hit_length, ) for component in self._components_tuple: 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 9c9247153..1794ae6f4 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 @@ -3344,6 +3344,85 @@ class UnifiedRadixCacheSuite: self.assertIs(with_hicache.last_device_node, tree_h.root_node) self.assertIsNone(with_hicache.mamba_branching_seqlen) + def test_mamba_branching_seqlen_uses_device_full_hit_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") + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + chunk_size = get_server_args().mamba_cache_chunk_size + prefix = self._make_seq(1, chunk_size) + tokens = prefix + self._make_seq(1000, chunk_size + 1) + self._insert(cache, allocator, req_to_token_pool, prefix) + self._insert(cache, allocator, req_to_token_pool, tokens) + + leaf = cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", tokens))) + ).last_device_node + parent = leaf.parent + leaf.component_data[ComponentType.MAMBA].value = None + + result = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens)))) + + self.assertIs(result.best_match_node, parent) + self.assertIs(result.last_device_node, parent) + self.assertEqual(len(result.device_indices), chunk_size) + self.assertEqual(result.host_hit_length, 0) + self.assertEqual(result.full_kv_hit_length, len(tokens)) + self.assertEqual(result.mamba_branching_seqlen, 2 * chunk_size) + + def test_mamba_branching_from_host_full_is_reusable_after_insert(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") + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + chunk_size = get_server_args().mamba_cache_chunk_size + prefix = self._make_seq(1, chunk_size) + tokens = prefix + self._make_seq(1000, chunk_size + 1) + self._insert(cache, allocator, req_to_token_pool, prefix) + self._insert(cache, allocator, req_to_token_pool, tokens) + + leaf = cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", tokens))) + ).last_device_node + parent = leaf.parent + self._backup_node(cache, leaf) + lock_result = cache.inc_lock_ref(parent) + try: + cache.evict(EvictParams(num_tokens=len(leaf.key))) + finally: + cache.dec_lock_ref( + parent, + DecLockRefParams( + swa_uuid_for_lock=getattr(lock_result, "swa_uuid_for_lock", None) + ), + ) + self.assertTrue(leaf.evicted) + self.assertTrue(leaf.backuped) + cache.components[ComponentType.MAMBA].evict_component( + leaf, target=EvictLayer.HOST + ) + + result = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens)))) + + self.assertIs(result.best_match_node, parent) + self.assertIs(result.last_device_node, parent) + self.assertEqual(len(result.device_indices), chunk_size) + self.assertEqual(result.host_hit_length, 0) + self.assertEqual(result.full_kv_hit_length, len(tokens)) + branching_seqlen = 2 * chunk_size + self.assertEqual(result.mamba_branching_seqlen, branching_seqlen) + + self._insert( + cache, + allocator, + req_to_token_pool, + tokens[:branching_seqlen], + ) + second_match = cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", tokens))) + ) + + self.assertEqual(len(second_match.device_indices), branching_seqlen) + self.assertIsNone(second_match.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")