refactor mamba radix cache logic in server_args (#17645)

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(
+            self._handle_mamba_radix_cache(
-                f"Disabling Radix Cache for {model_arch} as it is not yet supported."
+                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:
+            self._handle_mamba_radix_cache(
-                logger.warning(
+                model_arch=model_arch,
-                    "Disabling radix cache since speculative decoding for NemotronHForCausalLM is not supported with radix cache yet."
+                support_mamba_cache_extra_buffer=False,
-                )
+                sm100_default_attention_backend="flashinfer",
-                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
             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._handle_mamba_radix_cache(
-                    self.attention_backend = "triton"
+                model_arch=model_arch,
-                    logger.info(
+                support_mamba_cache_extra_buffer=True,
-                        "Use triton as attention backend on sm100 for Qwen3NextForCausalLM"
+                sm100_default_attention_backend="triton",
-                    )
+            )
-                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
         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 (
+            self._handle_mamba_radix_cache(
-                not self.enable_mamba_extra_buffer()
+                model_arch=model_arch,
-            ), f"mamba extra_buffer is not supported for {model_arch} model"
+                support_mamba_cache_extra_buffer=False,
-            if not self.disable_radix_cache:
+                sm100_default_attention_backend="triton",
-                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
         elif model_arch in ["Lfm2ForCausalLM"]:
-            assert (
+            self._handle_mamba_radix_cache(
-                not self.enable_mamba_extra_buffer()
+                model_arch=model_arch,
-            ), f"mamba extra_buffer is not supported for {model_arch} model"
+                support_mamba_cache_extra_buffer=False,
-            if not self.disable_radix_cache:
+                sm100_default_attention_backend="flashinfer",
-                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
             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 = (