chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Support HiCache for MambaRadixCache (#19663)

Browse Source 
Co-authored-by: hzh0425 <hzh0425@apache.org>
This commit is contained in:
Ke Bao
2026-03-08 00:36:25 +08:00
committed by GitHub
parent 17721b00fd
commit 5867c3fa80
6 changed files with 1739 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
python/sglang/srt/managers/cache_controller.py
View File

@@ -264,7 +264,12 @@ class HiCacheController:
):
self.tp_group = tp_group
self.mem_pool_device_allocator = token_to_kv_pool_allocator
self.mem_pool_device = token_to_kv_pool_allocator.get_kvcache()
mem_pool_device = token_to_kv_pool_allocator.get_kvcache()
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
if isinstance(mem_pool_device, HybridLinearKVPool):
mem_pool_device = mem_pool_device.full_kv_pool
self.mem_pool_device = mem_pool_device
self.mem_pool_host = mem_pool_host
self.write_policy = write_policy
self.page_size = page_size

13
python/sglang/srt/managers/schedule_batch.py
View File

@@ -639,6 +639,7 @@ class Req(ReqDllmMixin):
self.extend_logprob_start_len = 0
self.last_node: Any = None
self.last_host_node: Any = None
self.last_host_backup_node: Any = None
self.host_hit_length = 0
# Tokens loaded from storage backend (L3) during prefetch for this request
self.storage_hit_length = 0
@@ -857,7 +858,11 @@ class Req(ReqDllmMixin):
# Whether request reached finished condition
return self.finished_reason is not None
def init_next_round_input(self, tree_cache: Optional[BasePrefixCache] = None):
def init_next_round_input(
self,
tree_cache: Optional[BasePrefixCache] = None,
cow_mamba: Optional[bool] = None,
):
if self.is_dllm():
self._init_fill_ids_for_dllm()
self.determine_dllm_phase()
@@ -873,23 +878,27 @@ class Req(ReqDllmMixin):
token_ids = self.fill_ids[:max_prefix_len]
if tree_cache is not None:
if cow_mamba is None:
cow_mamba = tree_cache.supports_mamba()
match_result = tree_cache.match_prefix(
MatchPrefixParams(
key=RadixKey(token_ids=token_ids, extra_key=self.extra_key),
req=self,
cow_mamba=tree_cache.supports_mamba(),
cow_mamba=cow_mamba,
)
)
(
self.prefix_indices,
self.last_node,
self.last_host_node,
self.last_host_backup_node,
self.host_hit_length,
self.mamba_branching_seqlen,
) = (
match_result.device_indices,
match_result.last_device_node,
match_result.last_host_node,
match_result.last_host_backup_node,
match_result.host_hit_length,
match_result.mamba_branching_seqlen,
)

32
python/sglang/srt/managers/scheduler.py
View File

@@ -694,9 +694,20 @@ class Scheduler(
logger.info("Using experimental C++ radix tree implementation.")
self.tree_cache = RadixCacheCpp(params=params, server_args=server_args)
elif self.enable_hierarchical_cache:
from sglang.srt.mem_cache.hiradix_cache import HiRadixCache
if self.is_hybrid_ssm:
from sglang.srt.mem_cache.hi_mamba_radix_cache import (
HiMambaRadixCache,
)
self.tree_cache = HiRadixCache(params=params, server_args=server_args)
self.tree_cache = HiMambaRadixCache(
params=params, server_args=server_args
)
else:
from sglang.srt.mem_cache.hiradix_cache import HiRadixCache
self.tree_cache = HiRadixCache(
params=params, server_args=server_args
)
self.tp_worker.register_hicache_layer_transfer_counter(
self.tree_cache.cache_controller.layer_done_counter
)
@@ -1707,22 +1718,25 @@ class Scheduler(
def _prefetch_kvcache(self, req: Req):
if self.enable_hicache_storage:
req.init_next_round_input(self.tree_cache)
if req.last_node.backuped:
# only to initiate the prefetch if the last node is backuped
# otherwise, the allocated GPU memory must be locked for integrity
last_hash = req.last_host_node.get_last_hash_value()
req.init_next_round_input(self.tree_cache, cow_mamba=False)
last_host_node = (
req.last_host_backup_node
if req.last_host_backup_node is not None
else req.last_host_node
)
if last_host_node.backuped or last_host_node is self.tree_cache.root_node:
last_hash = last_host_node.get_last_hash_value()
matched_len = len(req.prefix_indices) + req.host_hit_length
new_input_tokens = req.fill_ids[matched_len:]
prefix_keys = (
req.last_node.get_prefix_hash_values(req.last_node.parent)
last_host_node.get_prefix_hash_values(last_host_node.parent)
if self.tree_cache.hicache_storage_pass_prefix_keys
else None
)
self.tree_cache.prefetch_from_storage(
req.rid,
req.last_host_node,
last_host_node,
new_input_tokens,
last_hash,
prefix_keys,

2
python/sglang/srt/mem_cache/base_prefix_cache.py
View File

@@ -97,6 +97,7 @@ class MatchResult(NamedTuple):
last_host_node : The last TreeNode on the host that was matched.
Note that if HiCache is not enabled,
this **must** be the same as `last_device_node`.
last_host_backup_node: The deepest backuped node for prefetch from storage.
host_hit_length : Length of the KV cache hit on the host, if applicable.
0 if HiCache is not enabled.
mamba_branching_seqlen: The mamba radix cache branching point, which is the longest
@@ -107,6 +108,7 @@ class MatchResult(NamedTuple):
device_indices: torch.Tensor
last_device_node: Any
last_host_node: Any
last_host_backup_node: Any = None
host_hit_length: int = 0
mamba_branching_seqlen: Optional[int] = None

1668
python/sglang/srt/mem_cache/hi_mamba_radix_cache.py Normal file
View File

File diff suppressed because it is too large Load Diff

31
python/sglang/srt/mem_cache/mamba_radix_cache.py
View File

@@ -21,7 +21,7 @@ The radix tree data structure for managing the hybrid (full and Mamba) KV cache.
import heapq
from collections import defaultdict
from functools import partial
from functools import lru_cache, partial
from typing import TYPE_CHECKING, List, Optional, Tuple
import torch
@@ -82,8 +82,11 @@ class TreeNode:
self.last_access_time = get_last_access_time()
self.hit_count = 0
self.host_ref_counter = 0
# store the host indices of KV cache
self.host_value = None
# store hash values of each pages
self.hash_value: Optional[List[str]] = None
# for lru list, invariant:
# 1. prev has greater last_access_time
@@ -104,6 +107,29 @@ class TreeNode:
def backuped(self):
return self.host_value is not None
def protect_host(self):
"""Protect the host value from eviction."""
self.host_ref_counter += 1
def release_host(self):
"""Release the host value, allowing it to be evicted."""
if self.host_ref_counter > 0:
self.host_ref_counter -= 1
else:
raise RuntimeError("Host reference counter is already zero.")
def get_last_hash_value(self) -> Optional[str]:
"""Returns the hash value of the last page in this node."""
if self.hash_value is None or len(self.hash_value) == 0:
return None
return self.hash_value[-1]
@lru_cache(maxsize=1)
def get_prefix_hash_values(self, node: "TreeNode") -> List[str]:
if node is None or node.hash_value is None:
return []
return node.get_prefix_hash_values(node.parent) + node.hash_value
def __lt__(self, other: "TreeNode"):
return self.last_access_time < other.last_access_time
@@ -457,7 +483,7 @@ class MambaRadixCache(BasePrefixCache):
if value is None:
value = torch.tensor([x for x in key.token_ids], dtype=torch.int64)
prefix_len, mamba_exist = self._insert_helper(
self.root_node, key, value, mamba_value
self.root_node, key, value, mamba_value, params.chunked
)
return InsertResult(prefix_len=prefix_len, mamba_exist=mamba_exist)
@@ -1063,6 +1089,7 @@ class MambaRadixCache(BasePrefixCache):
key: RadixKey,
value,
mamba_value,
chunked: bool = False,
) -> Tuple[int, bool]:
# Update the last access time from root to leaf, so that
# mamba will tombstone the node closer to root first