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
+5 -2
View File
@@ -398,7 +398,10 @@ class Scheduler(
# Hybrid memory pool
self.is_hybrid_swa = self.tp_worker.is_hybrid_swa
self.is_hybrid_gdn = self.tp_worker.model_runner.hybrid_gdn_config is not None
self.is_ssm_model = (
self.tp_worker.model_runner.hybrid_gdn_config is not None
or self.tp_worker.model_runner.mamba2_config is not None
)
if self.is_hybrid_swa:
self.sliding_window_size = self.tp_worker.sliding_window_size
@@ -762,7 +765,7 @@ class Scheduler(
self.tree_cache = SWARadixCache(
params=params, sliding_window_size=self.sliding_window_size
)
elif self.is_hybrid_gdn:
elif self.is_ssm_model:
from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
self.tree_cache = MambaRadixCache(params)