From 2ce121a1c34b057c3103bd51b4e86b4996b4d3b6 Mon Sep 17 00:00:00 2001 From: roikoren755 <26850796+roikoren755@users.noreply.github.com> Date: Fri, 5 Dec 2025 12:23:58 +0200 Subject: [PATCH] Enable RadixCache for Mamba2 models (#13584) --- python/sglang/srt/managers/schedule_policy.py | 6 +++--- python/sglang/srt/managers/scheduler.py | 7 +++++-- .../sglang/srt/managers/scheduler_metrics_mixin.py | 10 +++++----- .../srt/managers/scheduler_runtime_checker_mixin.py | 8 ++++---- python/sglang/srt/model_executor/model_runner.py | 5 ----- python/sglang/srt/server_args.py | 12 ++++++++++++ 6 files changed, 29 insertions(+), 19 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index a99b6b16c..d6d1172e2 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -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() diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 59ac30e5b..b801fd8f8 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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) diff --git a/python/sglang/srt/managers/scheduler_metrics_mixin.py b/python/sglang/srt/managers/scheduler_metrics_mixin.py index 853341d6e..f121ed3d4 100644 --- a/python/sglang/srt/managers/scheduler_metrics_mixin.py +++ b/python/sglang/srt/managers/scheduler_metrics_mixin.py @@ -112,7 +112,7 @@ class SchedulerMetricsMixin: f"full token usage: {full_token_usage:.2f}, " f"swa token usage: {swa_token_usage:.2f}, " ) - elif self.is_hybrid_gdn: + elif self.is_ssm_model: ( full_num_used, _, @@ -166,7 +166,7 @@ class SchedulerMetricsMixin: self.stats.token_usage = token_usage if self.is_hybrid_swa: self.stats.swa_token_usage = swa_token_usage - if self.is_hybrid_gdn: + if self.is_ssm_model: self.stats.mamba_usage = mamba_usage self.stats.num_queue_reqs = len(self.waiting_queue) self.stats.num_grammar_queue_reqs = len(self.grammar_queue) @@ -238,7 +238,7 @@ class SchedulerMetricsMixin: f"#swa token: {swa_num_used}, " f"swa token usage: {swa_token_usage:.2f}, " ) - elif self.is_hybrid_gdn: + elif self.is_ssm_model: ( full_num_used, mamba_used, @@ -315,7 +315,7 @@ class SchedulerMetricsMixin: self.stats.token_usage = token_usage if self.is_hybrid_swa: self.stats.swa_token_usage = swa_token_usage - if self.is_hybrid_gdn: + if self.is_ssm_model: self.stats.mamba_usage = mamba_usage self.stats.gen_throughput = self.last_gen_throughput self.stats.num_queue_reqs = len(self.waiting_queue) @@ -401,7 +401,7 @@ class SchedulerMetricsMixin: if self.is_hybrid_swa: full_num_used, swa_num_used, *_ = self._get_swa_token_info() num_tokens = max(full_num_used, swa_num_used) - elif self.is_hybrid_gdn: + elif self.is_ssm_model: num_tokens = self._get_mamba_token_info()[0] else: num_tokens = self._get_token_info()[0] diff --git a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py index 70ffc90bf..09f1ba548 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py @@ -204,7 +204,7 @@ class SchedulerRuntimeCheckerMixin: def check_memory(self: Scheduler): if self.is_hybrid_swa: memory_leak, token_msg = self._check_hybrid_memory() - elif self.is_hybrid_gdn and isinstance(self.tree_cache, MambaRadixCache): + elif self.is_ssm_model and isinstance(self.tree_cache, MambaRadixCache): memory_leak, token_msg = self._check_mamba_memory() else: memory_leak, token_msg = self._check_radix_cache_memory() @@ -239,7 +239,7 @@ class SchedulerRuntimeCheckerMixin: ) = self._get_swa_token_info() num_used = max(full_num_used, swa_num_used) token_usage = max(full_token_usage, swa_token_usage) - elif self.is_hybrid_gdn: + elif self.is_ssm_model: ( num_used, _, @@ -278,7 +278,7 @@ class SchedulerRuntimeCheckerMixin: def check_tree_cache(self: Scheduler): if (self.is_hybrid_swa and isinstance(self.tree_cache, SWARadixCache)) or ( - self.is_hybrid_gdn and isinstance(self.tree_cache, MambaRadixCache) + self.is_ssm_model and isinstance(self.tree_cache, MambaRadixCache) ): self.tree_cache.sanity_check() @@ -322,7 +322,7 @@ class SchedulerRuntimeCheckerMixin: # Print batch size and memory pool info to check whether there are de-sync issues. if self.is_hybrid_swa: _, info_msg = self._check_hybrid_memory() - elif self.is_hybrid_gdn and isinstance(self.tree_cache, MambaRadixCache): + elif self.is_ssm_model and isinstance(self.tree_cache, MambaRadixCache): _, info_msg = self._check_mamba_memory() else: _, info_msg = self._check_radix_cache_memory() diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 568e9b848..4d58278b7 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -441,11 +441,6 @@ class ModelRunner: if architectures and not any("Llama4" in arch for arch in architectures): self.is_hybrid_swa = self.model_config.is_hybrid_swa = True - if config := self.mamba2_config: - class_name = config.__class__.__name__ - logger.warning(f"{class_name} model detected, disable radix cache") - self.server_args.disable_radix_cache = True - # For MTP models like DeepSeek-V3 or GLM-4.5, the MTP layer(s) are used separately as draft # models for speculative decoding. In those cases, `num_nextn_predict_layers` is used to # determine the number of layers. diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 0391d142a..f3ca2833e 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1278,6 +1278,18 @@ class ServerArgs: logger.info( "Use flashinfer_trtllm as MoE runner backend on sm100 for Qwen3NextForCausalLM" ) + elif model_arch in [ + "NemotronHForCausalLM", + "FalconH1ForCausalLM", + "JetNemotronForCausalLM", + "JetVLMForConditionalGeneration", + ]: + if not self.disable_radix_cache: + logger.warning( + "Disabling overlap schedule since MambaRadixCache is not compatible with " + "overlap schedule currently, try to use --disable-radix-cache if overlap schedule is necessary" + ) + self.disable_overlap_schedule = True def _handle_sampling_backend(self): if self.sampling_backend is None: