From 5844cb2fd82c43b2ab886f996dd08dd35f872228 Mon Sep 17 00:00:00 2001 From: Yi Zhang <1109276519@qq.com> Date: Mon, 26 Jan 2026 17:02:49 +0800 Subject: [PATCH] refactor mamba radix cache logic in server_args (#17645) --- python/sglang/srt/server_args.py | 207 ++++++++++++++----------------- 1 file changed, 91 insertions(+), 116 deletions(-) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 6370e9b93..1a049ced4 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1449,14 +1449,11 @@ class ServerArgs: f"Using {self.attention_backend} as attention backend for {model_arch}." ) elif model_arch in ["KimiLinearForCausalLM"]: - logger.warning( - f"Disabling Radix Cache for {model_arch} as it is not yet supported." + self._handle_mamba_radix_cache( + model_arch=model_arch, + support_mamba_cache=False, ) - self.disable_radix_cache = True elif model_arch in ["NemotronHForCausalLM"]: - assert ( - not self.enable_mamba_extra_buffer() - ), f"mamba extra_buffer is not supported for {model_arch} model" model_config = self.get_model_config() if model_config.quantization in [ "modelopt", @@ -1475,30 +1472,11 @@ class ServerArgs: self.quantization = model_config.quantization self.moe_runner_backend = "flashinfer_cutlass" - if not self.disable_radix_cache and self.speculative_algorithm is not None: - logger.warning( - "Disabling radix cache since speculative decoding for NemotronHForCausalLM is not supported with radix cache yet." - ) - self.disable_radix_cache = True - elif 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 - if is_sm100_supported(): - if self.attention_backend is None: - self.attention_backend = "flashinfer" - logger.info( - "Use flashinfer as attention backend on sm100 for NemotronHForCausalLM" - ) - if self.attention_backend == "trtllm_mha": - logger.warning( - "Disabling radix cache since trtllm_mha does not support page_size = 1, which is required by MambaRadixCache. " - "Try to use --attention-backend triton if radix cache is necessary." - ) - self.disable_radix_cache = True - self.disable_overlap_schedule = False + self._handle_mamba_radix_cache( + model_arch=model_arch, + support_mamba_cache_extra_buffer=False, + sm100_default_attention_backend="flashinfer", + ) assert self.attention_backend != "triton", ( "NemotronHForCausalLM does not support triton attention backend," "as the first layer might not be an attention layer" @@ -1541,21 +1519,12 @@ class ServerArgs: logger.info( "Use flashinfer_trtllm as MoE runner backend on sm100 for Qwen3NextForCausalLM" ) - if self.attention_backend is None: - self.attention_backend = "triton" - logger.info( - "Use triton as attention backend on sm100 for Qwen3NextForCausalLM" - ) - if ( - not self.disable_radix_cache - and self.attention_backend == "trtllm_mha" - ): - logger.warning( - "Disabling radix cache since trtllm_mha does not support page_size = 1, which is required by MambaRadixCache. " - "Try to use --attention-backend triton if radix cache is necessary." - ) - self.disable_radix_cache = True - self.disable_overlap_schedule = False + self._handle_mamba_radix_cache( + model_arch=model_arch, + support_mamba_cache_extra_buffer=True, + sm100_default_attention_backend="triton", + ) + elif model_arch in ["Glm4MoeForCausalLM"]: if is_sm100_supported(): quantization_config = getattr(hf_config, "quantization_config", None) @@ -1578,83 +1547,23 @@ class ServerArgs: "Use flashinfer_trtllm as MoE runner backend on sm100 for Glm4MoeForCausalLM" ) - # Mamba radix cache v2 - if self.enable_mamba_extra_buffer(): - assert ( - is_cuda() - ), "Mamba extra_buffer is only supported on CUDA devices with FLA backend" - if self.speculative_num_draft_tokens is not None: - assert ( - self.mamba_track_interval >= self.speculative_num_draft_tokens - ), f"mamba_track_interval {self.mamba_track_interval} must be greater than or equal to speculative_num_draft_tokens {self.speculative_num_draft_tokens}" - - if self.page_size is not None: - assert ( - self.mamba_track_interval % self.page_size == 0 - ), f"mamba_track_interval {self.mamba_track_interval} must be divisible by page_size {self.page_size}" - assert ( - max(FLA_CHUNK_SIZE, self.page_size) - % min(FLA_CHUNK_SIZE, self.page_size) - == 0 - ), f"For SSM models with extra buffer, either FLA_CHUNK_SIZE or page_size must be divisible by the other, got {FLA_CHUNK_SIZE=}, {self.page_size=}" - - elif not self.disable_radix_cache: - logger.warning( - "Disabling overlap schedule since MambaRadixCache no_buffer is not compatible with " - "overlap schedule currently, try to use --mamba-scheduler-strategy extra_buffer to enable overlap schedule" - ) - self.disable_overlap_schedule = True - elif model_arch in [ "FalconH1ForCausalLM", "JetNemotronForCausalLM", "JetVLMForConditionalGeneration", ]: - assert ( - not self.enable_mamba_extra_buffer() - ), f"mamba extra_buffer is not supported for {model_arch} model" - if not self.disable_radix_cache: - logger.warning( - "Disabling overlap schedule since mamba no_buffer is not compatible with " - "overlap schedule, try to use --disable-radix-cache if overlap schedule is necessary" - ) - self.disable_overlap_schedule = True - if is_sm100_supported(): - if self.attention_backend is None: - self.attention_backend = "triton" - logger.info( - f"Use triton as attention backend on sm100 for {model_arch}" - ) - if self.attention_backend == "trtllm_mha": - logger.warning( - "Disabling radix cache since trtllm_mha does not support page_size = 1, which is required by MambaRadixCache. " - "Try to use --attention-backend triton if radix cache is necessary." - ) - self.disable_radix_cache = True - self.disable_overlap_schedule = False + self._handle_mamba_radix_cache( + model_arch=model_arch, + support_mamba_cache_extra_buffer=False, + sm100_default_attention_backend="triton", + ) + elif model_arch in ["Lfm2ForCausalLM"]: - assert ( - not self.enable_mamba_extra_buffer() - ), f"mamba extra_buffer is not supported for {model_arch} model" - if not self.disable_radix_cache: - logger.warning( - "Disabling overlap schedule since mamba no_buffer is not compatible with " - "overlap schedule, try to use --disable-radix-cache if overlap schedule is necessary" - ) - self.disable_overlap_schedule = True - if is_sm100_supported(): - if self.attention_backend is None: - self.attention_backend = "flashinfer" - logger.info( - f"Use flashinfer as attention backend on sm100 for {model_arch}" - ) - if self.attention_backend == "trtllm_mha": - logger.warning( - "Disabling radix cache since trtllm_mha does not support page_size = 1, which is required by MambaRadixCache. " - "Try to use --attention-backend flashinfer if radix cache is necessary." - ) - self.disable_radix_cache = True - self.disable_overlap_schedule = False + self._handle_mamba_radix_cache( + model_arch=model_arch, + support_mamba_cache_extra_buffer=False, + sm100_default_attention_backend="flashinfer", + ) assert self.attention_backend != "triton", ( f"{model_arch} does not support triton attention backend, " "as the first layer might not be an attention layer" @@ -1692,6 +1601,72 @@ class ServerArgs: ): self.enable_flashinfer_allreduce_fusion = True + def _handle_mamba_radix_cache( + self, + model_arch: str, + support_mamba_cache: bool = True, + support_mamba_cache_extra_buffer: bool = True, + sm100_default_attention_backend: str = None, + ): + if ( + is_sm100_supported() + and self.attention_backend is None + and sm100_default_attention_backend is not None + ): + self.attention_backend = sm100_default_attention_backend + logger.info( + f"Use {sm100_default_attention_backend} as attention backend on sm100 for {model_arch}" + ) + + if not support_mamba_cache: + logger.warning( + f"Disabling Radix Cache for {model_arch} as it is not yet supported." + ) + self.disable_radix_cache = True + return + + if not support_mamba_cache_extra_buffer: + assert ( + not self.enable_mamba_extra_buffer() + ), f"mamba extra_buffer is not supported for {model_arch} model" + elif self.enable_mamba_extra_buffer(): # extra_buffer + assert ( + is_cuda() + ), "Mamba extra_buffer is only supported on CUDA devices with FLA backend" + if self.speculative_num_draft_tokens is not None: + assert ( + self.mamba_track_interval >= self.speculative_num_draft_tokens + ), f"mamba_track_interval {self.mamba_track_interval} must be greater than or equal to speculative_num_draft_tokens {self.speculative_num_draft_tokens}" + + if self.page_size is not None: + assert ( + self.mamba_track_interval % self.page_size == 0 + ), f"mamba_track_interval {self.mamba_track_interval} must be divisible by page_size {self.page_size}" + assert ( + max(FLA_CHUNK_SIZE, self.page_size) + % min(FLA_CHUNK_SIZE, self.page_size) + == 0 + ), f"For SSM models with extra buffer, either FLA_CHUNK_SIZE or page_size must be divisible by the other, got {FLA_CHUNK_SIZE=}, {self.page_size=}" + elif not self.disable_radix_cache: # no_buffer + if self.speculative_algorithm is None: + logger.warning( + "Disabling overlap schedule since mamba no_buffer is not compatible with " + "overlap schedule, try to use --disable-radix-cache if overlap schedule is necessary" + ) + self.disable_overlap_schedule = True + if self.attention_backend == "trtllm_mha": + logger.warning( + "Disabling radix cache since trtllm_mha does not support page_size = 1, which is required by MambaRadixCache. " + "Try to use --attention-backend triton if radix cache is necessary." + ) + self.disable_radix_cache = True + self.disable_overlap_schedule = False + else: + logger.warning( + f"Disabling radix cache since speculative decoding for {model_arch} is not supported with radix cache yet." + ) + self.disable_radix_cache = True + def _handle_sampling_backend(self): if self.sampling_backend is None: self.sampling_backend = (