Enable RadixCache for Mamba2 models (#13584)
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user