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: