minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 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
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")