[feat] Support extra_buffer in Mamba2-based models (#15829)

Signed-off-by: Roi Koren <roik@nvidia.com>

python/sglang/srt/arg_groups/nemotron_h_hook.py

@@ -41,9 +41,8 @@ def apply_nemotron_h_defaults(server_args: "ServerArgs", model_arch: str) -> Non
 
     server_args._handle_mamba_radix_cache(
         model_arch=model_arch,
-        support_mamba_cache=True,
-        support_mamba_cache_extra_buffer=False,
         sm100_default_attention_backend="flashinfer",
+        fallback_attention_backend="flashinfer",
     )
     assert server_args.attention_backend != "triton", (
         "NemotronHForCausalLM does not support triton attention backend,"

python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py

@@ -23,12 +23,6 @@ from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
 from sglang.srt.speculative.spec_info import SpecInput
-from sglang.srt.utils import is_cpu
-
-if not is_cpu():
-    from sglang.srt.layers.attention.fla.chunk_delta_h import (
-        CHUNK_SIZE as FLA_CHUNK_SIZE,
-    )
 
 logger = logging.getLogger(__name__)
 
@@ -278,9 +272,9 @@ class MambaAttnBackendBase(AttentionBackend):
         After processing a prefill chunk, we need to save the last `conv_state_len` tokens
         of the processed region for prefix caching.
 
-        The key insight is that FLA (Flash Linear Attention) processes sequences in chunks
-        of FLA_CHUNK_SIZE. We only track the conv state up to the last complete chunk boundary
-        (aligned_len).
+        The key insight is that FLA (Flash Linear Attention) and Mamba2 processes sequences in chunks
+        of the chunk size (FLA_CHUNK_SIZE=64 for FLA, mamba_chunk_size for Mamba2).
+        We only track the conv state up to the last complete chunk boundary (aligned_len).
 
         start_indices is the starting token index of the conv state to track in this extend batch.
         indices include all pos to track in this extend batch, conv_state_len for each req that
@@ -317,28 +311,33 @@ class MambaAttnBackendBase(AttentionBackend):
         Compute source and destination indices for tracking SSM states for prefix caching.
 
         After processing a prefill, we need to save the SSM recurrent state for prefix caching.
-        The FLA kernel outputs intermediate hidden states `h` at each chunk boundary,
+        The kernel outputs intermediate hidden states `h` at each chunk boundary,
         plus a `last_recurrent_state` at the end of the chunked prefill size.
 
+        The chunk size varies by model type:
+        - FLA models: FLA_CHUNK_SIZE (64)
+        - Mamba2 models: mamba_chunk_size (256)
+
         The challenge is that sequences may or may not end on a chunk boundary:
-          - Aligned case (len % FLA_CHUNK_SIZE == 0): In this case, FLA will store the to-cache
-            state in the last_recurrent_state.
-          - Unaligned case (len % FLA_CHUNK_SIZE != 0): The last_recurrent_state includes the
+          - Aligned case (len % chunk_size == 0): The to-cache state is stored in
+            the last_recurrent_state.
+          - Unaligned case (len % chunk_size != 0): The last_recurrent_state includes the
             unaligned position, but we only want state up to the last chunk boundary.
             We must extract from the intermediate `h` tensor at the appropriate chunk index.
 
         We compute the src and dst indices for all requests that need to be cached
         (i.e. mamba_track_mask is True) based on the rule above.
 
-        For example:
-        1. If chunked prefill length is < 64, then only final state has value. In this case we
-           cache `final` state.
-        2. if chunked prefill length == 64, then only final state has value. In this case we
-           cache pos 64, from `final` state
-        3. if chunked prefill length >64 and < 128, then both h and final state have value.
-           We cache pos 64 from `h` state
-        4. if chunked prefill length ==128, then both h and final state have value. We cache
-           pos 128 from `final` state. Note `h` doesn't include the pos 128.
+        For example (assuming chunk_size=64):
+        1. If chunked prefill length is < chunk_size, then only final state has value.
+           In this case we cache `final` state.
+        2. If chunked prefill length == chunk_size, then only final state has value.
+           In this case we cache pos chunk_size, from `final` state.
+        3. If chunked prefill length > chunk_size and < 2 * chunk_size, then both h and
+           final state have value. We cache pos chunk_size from `h` state.
+        4. If chunked prefill length == 2 * chunk_size, then both h and final state have
+           value. We cache pos 2 * chunk_size from `final` state. Note `h` doesn't include
+           the final position.
 
         Returns:
             track_ssm_h_src: Source indices into the packed `h` tensor (for unaligned seqs)
@@ -346,6 +345,7 @@ class MambaAttnBackendBase(AttentionBackend):
             track_ssm_final_src: Source indices into last_recurrent_state buffer (for aligned seqs)
             track_ssm_final_dst: Destination cache slot indices (for aligned seqs)
         """
+        mamba_cache_chunk_size = get_global_server_args().mamba_cache_chunk_size
         # Move to CPU to avoid kernel launches for masking operations
         mamba_track_mask = forward_batch.mamba_track_mask.cpu()
         extend_seq_lens = forward_batch.extend_seq_lens.cpu()
@@ -355,7 +355,10 @@ class MambaAttnBackendBase(AttentionBackend):
         prefix_lens = forward_batch.extend_prefix_lens.cpu()
 
         # Calculate the number of hidden states per request
-        num_h_states = (extend_seq_lens - 1) // FLA_CHUNK_SIZE + 1
+        if isinstance(self, Mamba2AttnBackend):
+            num_h_states = extend_seq_lens // mamba_cache_chunk_size
+        else:
+            num_h_states = (extend_seq_lens - 1) // mamba_cache_chunk_size + 1
 
         # Calculate the starting offset for each sequence in the packed batch
         track_ssm_src_offset = torch.zeros_like(num_h_states)
@@ -368,17 +371,17 @@ class MambaAttnBackendBase(AttentionBackend):
         dst_masked = mamba_track_indices[mamba_track_mask]
 
         # Determine if the sequence ends at a chunk boundary
-        is_aligned = (lens_masked % FLA_CHUNK_SIZE) == 0
+        is_aligned = (lens_masked % mamba_cache_chunk_size) == 0
 
         # Case 1: Aligned. Use last_recurrent_state from ssm_states.
         track_ssm_final_src = mamba_cache_indices[mamba_track_mask][is_aligned]
         track_ssm_final_dst = dst_masked[is_aligned]
 
         # Case 2: Unaligned. Use intermediate state from h.
-        # TODO: if support FLA_CHUNK_SIZE % page size != 0, then need to modify this
+        # TODO: if support mamba_cache_chunk_size % page size != 0, then need to modify this
         not_aligned = ~is_aligned
         track_ssm_h_src = offset_masked[not_aligned] + (
-            lens_masked[not_aligned] // FLA_CHUNK_SIZE
+            lens_masked[not_aligned] // mamba_cache_chunk_size
         )
         track_ssm_h_dst = dst_masked[not_aligned]
 
@@ -639,10 +642,10 @@ class MambaAttnBackendBase(AttentionBackend):
         """
         Track and copy SSM states during extend for prefix caching.
 
-        After the FLA chunked prefill kernel runs, we need to save the SSM recurrent
+        After the chunked prefill kernel runs, we need to save the SSM recurrent
         state at the last chunk boundary so it can be reused for prefix caching.
         The source of the state depends on whether the sequence length is aligned
-        to FLA_CHUNK_SIZE. See `_init_track_ssm_indices` for more details on how
+        to the chunk size. See `_init_track_ssm_indices` for more details on how
         the source and destination indices are computed.
 
         Note: Conv state tracking for extend is handled separately via gather operations
@@ -669,6 +672,17 @@ class Mamba2AttnBackend(MambaAttnBackendBase):
         config = model_runner.mamba2_config
         assert config is not None
         self.mamba_chunk_size = config.mamba_chunk_size
+        self.conv_states_shape = (
+            model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape
+        )
+
+        if model_runner.server_args.enable_mamba_extra_buffer():
+            assert (
+                self.conv_states_shape[-1] < self.mamba_chunk_size
+            ), f"{self.conv_states_shape[-1]=} should be less than {self.mamba_chunk_size}"
+            assert (
+                model_runner.server_args.mamba_track_interval >= self.mamba_chunk_size
+            ), f"mamba_track_interval ({model_runner.server_args.mamba_track_interval}) must be >= mamba_chunk_size ({self.mamba_chunk_size})"
 
     def init_forward_metadata(self, forward_batch: ForwardBatch):
         self._execute_deferred_mamba_cow_and_clear(forward_batch)
@@ -726,20 +740,46 @@ class Mamba2AttnBackend(MambaAttnBackendBase):
         hidden_states: torch.Tensor,
         output: torch.Tensor,
         layer_id: int,
+        forward_batch: ForwardBatch,
         mup_vector: Optional[torch.Tensor] = None,
         use_triton_causal_conv: bool = False,
     ):
         assert isinstance(self.forward_metadata, Mamba2Metadata)
         layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer_id)
-        return mixer.forward(
+        intermediate_states = mixer.forward(
             hidden_states=hidden_states,
             output=output,
             layer_cache=layer_cache,
             metadata=self.forward_metadata,
+            forward_batch=forward_batch,
             mup_vector=mup_vector,
             use_triton_causal_conv=use_triton_causal_conv,
         )
 
+        if forward_batch.mamba_track_mask is not None:
+            if (
+                intermediate_states is not None
+                and forward_batch.mamba_track_mask is not None
+                and forward_batch.mamba_track_mask.any()
+            ):
+                self._track_mamba_state_extend(
+                    forward_batch,
+                    intermediate_states,
+                    layer_cache.temporal,
+                    self.forward_metadata,
+                )
+
+            if self.forward_metadata.num_decodes > 0:
+                num_decodes = self.forward_metadata.num_decodes
+                track_mamba_states_if_needed(
+                    layer_cache.conv[0],
+                    layer_cache.temporal,
+                    self.forward_metadata.mamba_cache_indices[-num_decodes:],
+                    forward_batch.mamba_track_mask[-num_decodes:],
+                    forward_batch.mamba_track_indices[-num_decodes:],
+                    num_decodes,
+                )
+
     def forward_decode(self, *args, **kwargs):
         raise NotImplementedError(
             "Mamba2AttnBackend's forward is called directly instead of through HybridLinearAttnBackend, as it supports mixed prefill and decode"

python/sglang/srt/layers/attention/mamba/mamba.py

@@ -26,6 +26,7 @@ from sglang.srt.layers.linear import (
 )
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.mem_cache.memory_pool import MambaPool
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import (
     composed_weight_loader,
     sharded_weight_loader,
@@ -410,6 +411,7 @@ class MambaMixer2(torch.nn.Module):
         output: torch.Tensor,
         layer_cache: MambaPool.State,
         metadata: Mamba2Metadata,
+        forward_batch: ForwardBatch,
         mup_vector: Optional[torch.Tensor] = None,
         use_triton_causal_conv: bool = False,
     ):
@@ -420,6 +422,7 @@ class MambaMixer2(torch.nn.Module):
         state_indices_tensor = metadata.mamba_cache_indices
         conv_state = layer_cache.conv[0]
         ssm_state = layer_cache.temporal
+        intermediate_states = None
 
         query_start_loc = metadata.query_start_loc
 
@@ -517,6 +520,14 @@ class MambaMixer2(torch.nn.Module):
             x = hidden_states_B_C_p.transpose(
                 0, 1
             )  # this is the form that causal-conv see
+            if (
+                forward_batch.mamba_track_mask is not None
+                and forward_batch.mamba_track_mask.any()
+                and metadata.track_conv_indices is not None
+            ):
+                x_to_track = x[:, metadata.track_conv_indices].transpose(0, 1)
+                mask_indices = forward_batch.mamba_track_mask.nonzero(as_tuple=True)[0]
+                conv_state[forward_batch.mamba_track_indices[mask_indices]] = x_to_track
             ccfn = (
                 causal_conv1d_fn
                 if not use_triton_causal_conv
@@ -546,7 +557,7 @@ class MambaMixer2(torch.nn.Module):
                 )
 
             # NOTE: final output is an in-place update of out tensor
-            varlen_state = mamba_chunk_scan_combined(
+            intermediate_states, varlen_state = mamba_chunk_scan_combined(
                 hidden_states_p.view(
                     1, num_prefill_tokens, self.num_heads // self.tp_size, self.head_dim
                 ),
@@ -565,6 +576,7 @@ class MambaMixer2(torch.nn.Module):
                 initial_states=initial_states,
                 return_varlen_states=True,
                 return_final_states=False,
+                return_intermediate_states=True,
                 dt_softplus=True,
                 dt_limit=(0.0, float("inf")),
                 out=preallocated_ssm_out_p.view(
@@ -708,6 +720,8 @@ class MambaMixer2(torch.nn.Module):
         # 5. Final linear projection
         output[:num_actual_tokens], _ = self.out_proj(hidden_states)
 
+        return intermediate_states
+
     @property
     def mamba_type(self) -> str:
         return "mamba2"

python/sglang/srt/layers/attention/mamba/mamba2_metadata.py

@@ -171,6 +171,11 @@ class Mamba2Metadata(ForwardMetadata):
             retrieve_next_token=forward_metadata.retrieve_next_token,
             retrieve_next_sibling=forward_metadata.retrieve_next_sibling,
             retrieve_parent_token=forward_metadata.retrieve_parent_token,
+            track_conv_indices=forward_metadata.track_conv_indices,
+            track_ssm_h_src=forward_metadata.track_ssm_h_src,
+            track_ssm_h_dst=forward_metadata.track_ssm_h_dst,
+            track_ssm_final_src=forward_metadata.track_ssm_final_src,
+            track_ssm_final_dst=forward_metadata.track_ssm_final_dst,
             num_decodes=len(seq_lens),
             num_prefills=0,
             num_prefill_tokens=0,
@@ -239,6 +244,11 @@ class Mamba2Metadata(ForwardMetadata):
             retrieve_next_token=forward_metadata.retrieve_next_token,
             retrieve_next_sibling=forward_metadata.retrieve_next_sibling,
             retrieve_parent_token=forward_metadata.retrieve_parent_token,
+            track_conv_indices=forward_metadata.track_conv_indices,
+            track_ssm_h_src=forward_metadata.track_ssm_h_src,
+            track_ssm_h_dst=forward_metadata.track_ssm_h_dst,
+            track_ssm_final_src=forward_metadata.track_ssm_final_src,
+            track_ssm_final_dst=forward_metadata.track_ssm_final_dst,
             num_prefills=num_prefills,
             num_prefill_tokens=num_prefill_tokens,
             num_decodes=num_decodes,

python/sglang/srt/managers/schedule_batch.py

@@ -68,7 +68,6 @@ from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST, DisaggregationM
 from sglang.srt.distributed.parallel_state import get_tensor_model_parallel_rank
 from sglang.srt.dllm.mixin.req import ReqDllmMixin
 from sglang.srt.environ import envs
-from sglang.srt.layers.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE
 from sglang.srt.managers.embed_types import PositionalEmbeds
 from sglang.srt.managers.scheduler_components.new_token_ratio_tracker import (
     NewTokenRatioTracker,
@@ -2090,18 +2089,19 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
         self,
         req: Req,
     ) -> "_MambaRadixCacheV2TrackEntry":
+        mamba_cache_chunk_size = get_global_server_args().mamba_cache_chunk_size
+
         def _force_track_h(i: int) -> int:
-            assert i % FLA_CHUNK_SIZE == 0
+            assert i % mamba_cache_chunk_size == 0
             # There are 3 cases for mamba_track_seqlen passed to mamba_track_seqlens_cpu:
-            # 1) aligned with FLA_CHUNK_SIZE-> retrieve from last_recurrent_state
+            # 1) aligned with mamba_cache_chunk_size-> retrieve from last_recurrent_state
             #    a) is the last position -> retrieve from last_recurrent_state
             #    b) is NOT the last position -> retrieve from h
-            # 2) unaligned with FLA_CHUNK_SIZE -> retrieve from h
+            # 2) unaligned with mamba_cache_chunk_size -> retrieve from h
             # Currently, the math calculation only supports case 1a and 2. So for 1b, we need to add 1
             # to force the math calculation to retrieve the correct mamba state from h.
             return i + 1
 
-        mamba_cache_chunk_size = get_global_server_args().mamba_cache_chunk_size
         mask = req.extend_input_len >= mamba_cache_chunk_size
         track_index = req.mamba_ping_pong_track_buffer[req.mamba_next_track_idx].item()
         mamba_track_seqlen = -1
@@ -2123,13 +2123,14 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
                 * mamba_cache_chunk_size
             )
 
-            # mamba_track_fla_chunk_aligned is the aligned seqlen based on FLA_CHUNK_SIZE
+            # mamba_track_fla_chunk_aligned is the aligned seqlen based on mamba_cache_chunk_size
             # If mamba_track_fla_chunk_aligned != mamba_track_seqlen_aligned, which can be true when
-            # page_size > FLA_CHUNK_SIZE, we need to force the math calculation to retrieve the correct mamba state from h
+            # page_size > mamba_cache_chunk_size, we need to force the math calculation to retrieve the correct mamba state from h
             # by _force_track_h()
             mamba_track_fla_chunk_aligned = (
                 len(req.prefix_indices)
-                + (req.extend_input_len // FLA_CHUNK_SIZE) * FLA_CHUNK_SIZE
+                + (req.extend_input_len // mamba_cache_chunk_size)
+                * mamba_cache_chunk_size
             )
             if mamba_track_fla_chunk_aligned != mamba_track_seqlen_aligned:
                 # We want to track mamba_track_seqlen_aligned, and it's not the last position,

python/sglang/srt/mem_cache/mamba_radix_cache.py

@@ -28,7 +28,6 @@ import torch
 from numpy import float64
 
 from sglang.srt.distributed import get_tensor_model_parallel_rank
-from sglang.srt.layers.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE
 from sglang.srt.mem_cache.allocator import (
     PagedTokenToKVPoolAllocator,
     TokenToKVPoolAllocator,
@@ -426,6 +425,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
         ) or isinstance(params.token_to_kv_pool_allocator, PagedTokenToKVPoolAllocator)
         self.req_to_token_pool: HybridReqToTokenPool = params.req_to_token_pool
         self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator
+        self.mamba_cache_chunk_size = get_global_server_args().mamba_cache_chunk_size
 
         self.page_size = params.page_size
         self.disable = params.disable
@@ -637,7 +637,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
 
         assert page_aligned_len == len(
             kv_indices
-        ), f"page_aligned_len != len(kv_indices), {page_aligned_len=}, {len(kv_indices)=}, {cache_len=}, {self.page_size=}, {FLA_CHUNK_SIZE=}"
+        ), f"page_aligned_len != len(kv_indices), {page_aligned_len=}, {len(kv_indices)=}, {cache_len=}, {self.page_size=}, {self.mamba_cache_chunk_size=}"
 
         page_aligned_token_ids = token_ids[:page_aligned_len]
 
@@ -1041,14 +1041,11 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
         # Calculate the branching point. It is defined as the last aligned position that
         # does not have a mamba value.
         if len(value) > best_value_len:
-            mamba_cache_chunk_size = get_global_server_args().mamba_cache_chunk_size
-            mamba_cache_chunk_aligned_seqlen = (
-                sum(len(v) for v in value) // mamba_cache_chunk_size
-            ) * mamba_cache_chunk_size
+            chunk_aligned_seqlen = (
+                sum(len(v) for v in value) // self.mamba_cache_chunk_size
+            ) * self.mamba_cache_chunk_size
             mamba_branching_seqlen = (
-                mamba_cache_chunk_aligned_seqlen
-                if mamba_cache_chunk_aligned_seqlen > 0
-                else None
+                chunk_aligned_seqlen if chunk_aligned_seqlen > 0 else None
             )
         else:
             mamba_branching_seqlen = None

python/sglang/srt/models/falcon_h1.py

@@ -349,6 +349,7 @@ class FalconH1HybridAttentionDecoderLayer(nn.Module):
                 hidden_states * self.ssm_in_multiplier,
                 mamba_hidden_states,
                 layer_id=self.layer_id,
+                forward_batch=forward_batch,
                 mup_vector=self.mup_vector,
             )
             mamba_hidden_states = mamba_hidden_states * self.ssm_out_multiplier

python/sglang/srt/models/granitemoehybrid.py

@@ -148,6 +148,7 @@ class GraniteMoeHybridMambaDecoderLayer(nn.Module):
             layer_id=self.layer_idx,
             hidden_states=hidden_states,
             output=output,
+            forward_batch=forward_batch,
             use_triton_causal_conv=True,
         )
 

python/sglang/srt/models/nemotron_h.py

@@ -423,6 +423,7 @@ class NemotronHMambaDecoderLayer(nn.Module):
             layer_id=self.layer_id,
             hidden_states=hidden_states,
             output=output,
+            forward_batch=forward_batch,
             use_triton_causal_conv=True,
         )
         return output

python/sglang/srt/server_args.py

@@ -2456,8 +2456,6 @@ class ServerArgs:
         ]:
             self._handle_mamba_radix_cache(
                 model_arch=model_arch,
-                support_mamba_cache=True,
-                support_mamba_cache_extra_buffer=False,
                 sm100_default_attention_backend="triton",
             )
 
@@ -2470,8 +2468,7 @@ class ServerArgs:
             if has_mamba:
                 self._handle_mamba_radix_cache(
                     model_arch=model_arch,
-                    support_mamba_cache_extra_buffer=False,
-                    sm100_default_attention_backend="triton",
+                    sm100_default_attention_backend="flashinfer",
                 )
 
         elif model_arch in ["Lfm2ForCausalLM"]:
@@ -2556,6 +2553,7 @@ class ServerArgs:
         support_mamba_cache: bool = True,
         support_mamba_cache_extra_buffer: bool = True,
         sm100_default_attention_backend: str = None,
+        fallback_attention_backend: str = "triton",
     ):
         if (
             is_sm100_supported()
@@ -2582,7 +2580,7 @@ class ServerArgs:
         if self.enable_mamba_extra_buffer():  # extra_buffer
             if self.disable_radix_cache:
                 raise ValueError(
-                    "mamba extra_buffer is not compatible with --disable-radix-cache "
+                    "mamba extra_buffer is not compatible with --disable-radix-cache. "
                     "Overlap scheduling is already supported with no_buffer + disable_radix_cache. "
                     "Please use --mamba-scheduler-strategy no_buffer instead."
                 )
@@ -2599,11 +2597,7 @@ class ServerArgs:
                 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=}"
+                assert self.mamba_cache_chunk_size is not None
         elif not self.disable_radix_cache:  # no_buffer
             if self.page_size is not None and self.page_size != 1:
                 logger.warning(
@@ -2621,7 +2615,7 @@ class ServerArgs:
                 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."
+                        f"Try to use --attention-backend {fallback_attention_backend} if radix cache is necessary."
                     )
                     self.disable_radix_cache = True
                     self.disable_overlap_schedule = False
@@ -7047,9 +7041,17 @@ class ServerArgs:
 
     @property
     def mamba_cache_chunk_size(self) -> int:
-        # For mamba cache with extra buffer, the chunk size is the max of FLA_CHUNK_SIZE and page_size.
+        # For mamba cache with extra buffer, the chunk size is the max of FLA_CHUNK_SIZE
+        # (or mamba_chunk_size if it is defined in the model's config) and page_size.
         # It is used to determine the caching point in a sequence during prefill.
-        return max(FLA_CHUNK_SIZE, self.page_size)
+        if not hasattr(self, "_mamba_cache_chunk_size"):
+            hf_config = self.get_model_config().hf_config
+            chunk_size = getattr(hf_config, "mamba_chunk_size", FLA_CHUNK_SIZE)
+            assert (
+                max(chunk_size, self.page_size) % min(chunk_size, self.page_size) == 0
+            ), f"For SSM models, either chunk_size or page_size must be divisible by the other, got {chunk_size=}, {self.page_size=}"
+            self._mamba_cache_chunk_size = max(chunk_size, self.page_size)
+        return self._mamba_cache_chunk_size
 
     def check_server_args(self):
         # Check parallel size constraints