Enable RadixCache for Mamba2 models (#13584)

This commit is contained in:
roikoren755
2025-12-05 18:23:58 +08:00
committed by GitHub
parent 35ba6fe19e
commit 2ce121a1c3
6 changed files with 29 additions and 19 deletions
@@ -362,7 +362,7 @@ class PrefillAdder:
self.is_hybrid_swa = isinstance(
self.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator
)
self.is_hybrid_gdn_cache = isinstance(self.tree_cache, MambaRadixCache)
self.is_ssm_radix_cache = isinstance(self.tree_cache, MambaRadixCache)
self.priority_scheduling_preemption_threshold = (
priority_scheduling_preemption_threshold
@@ -387,7 +387,7 @@ class PrefillAdder:
self.token_to_kv_pool_allocator.swa_available_size()
+ self.tree_cache.swa_evictable_size(),
)
elif self.is_hybrid_gdn_cache:
elif self.is_ssm_radix_cache:
available_and_evictable = (
self.token_to_kv_pool_allocator.available_size()
+ self.tree_cache.full_evictable_size()
@@ -409,7 +409,7 @@ class PrefillAdder:
self.token_to_kv_pool_allocator.swa_available_size()
+ self.tree_cache.swa_evictable_size(),
)
elif self.is_hybrid_gdn_cache:
elif self.is_ssm_radix_cache:
available_and_evictable = (
self.token_to_kv_pool_allocator.available_size()
+ self.tree_cache.full_evictable_size()