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

[Hicache][1/2]Support Mamba branching in Unified Radix Cache with HiCache (#31181)

Browse Source 
Co-authored-by: Ke Bao <ispobaoke@gmail.com>
Co-authored-by: Zhangheng <hzh0425@apache.org>
This commit is contained in:
Jincong Chen
2026-07-25 19:44:40 +08:00
committed by GitHub
co-authored by Ke Bao Zhangheng
parent 1054060ef1
commit 7c4b22fae5
4 changed files with 104 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/base_prefix_cache.py
+4
View File
@@ -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,
)
python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py
+12 -11
View File
@@ -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:
python/sglang/srt/mem_cache/unified_radix_cache.py
+9 -1
View File
@@ -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:
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+79
View File
@@ -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")