retrieve chunk size from the model config instead of the mamba layers

Signed-off-by: nopperl <54780682+nopperl@users.noreply.github.com>

Author: nopperl

vllm/config.py

@@ -1330,6 +1330,17 @@ def get_num_layers_by_block_type(
 
             return sum(t == 1 for t in attn_type_list[start:end])
 
+    def get_mamba_chunk_size(self) -> Optional[int]:
+        """
+        Returns the mamba chunk size if it exists
+        """
+        # used by e.g. Bamba, FalconH1, Granite, PLaMo2
+        chunk_size = getattr(self.hf_text_config, "mamba_chunk_size", None)
+        if chunk_size is None:
+            # used by e.g. Mamba2, NemotronH, Zamba
+            chunk_size = getattr(self.hf_text_config, "chunk_size", None)
+        return chunk_size
+
     def get_multimodal_config(self) -> "MultiModalConfig":
         """
         Get the multimodal configuration of the model.

vllm/model_executor/layers/mamba/mamba_mixer2.py

@@ -224,8 +224,6 @@ class Mamba2Layer(ABC):
     Inherit from this class if you implement a custom Mamba2 layer.
     """
 
-    chunk_size: int
-
     # Contains the KV cache (mamba state) for the layer
     # in the shape specified by `self.get_state_shape`.
     # The outer list is for v0 PP virtual engine. Though this code path
@@ -257,22 +255,21 @@ class MambaMixer2(Mamba2Layer, CustomOp):
     """
 
     def __init__(
-            self,
-            hidden_size: int,
-            ssm_state_size: int,
-            conv_kernel_size: int,
-            intermediate_size: int,
-            use_conv_bias: bool,
-            use_bias: bool,
-            n_groups: int = 1,
-            num_heads: int = 128,
-            head_dim: int = 64,
-            rms_norm_eps: float = 1e-5,
-            activation: str = "silu",
-            use_rms_norm: bool = True,
-            quant_config: Optional[QuantizationConfig] = None,
-            prefix: str = "",
-            chunk_size: int = -1,  # the chunk size used by v1
+        self,
+        hidden_size: int,
+        ssm_state_size: int,
+        conv_kernel_size: int,
+        intermediate_size: int,
+        use_conv_bias: bool,
+        use_bias: bool,
+        n_groups: int = 1,
+        num_heads: int = 128,
+        head_dim: int = 64,
+        rms_norm_eps: float = 1e-5,
+        activation: str = "silu",
+        use_rms_norm: bool = True,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ):
         super().__init__()
 
@@ -454,10 +451,7 @@ def __init__(
             # of Attention + v0 PP.
             # The inner tuple is (conv_state, ssm_state)
             self.kv_cache = [(torch.tensor([]), torch.tensor([]))]
-            assert chunk_size != -1, "chunk_size must be set for v1"
 
-        # NOTE: chunk_size may be -1 for models without v1 support
-        self.chunk_size = chunk_size
         self.prefix = prefix
 
     def forward_native(

vllm/model_executor/models/bamba.py

@@ -99,8 +99,7 @@ def __init__(self,
                                 rms_norm_eps=config.rms_norm_eps,
                                 activation=config.hidden_act,
                                 quant_config=quant_config,
-                                prefix=f"{prefix}.mixer",
-                                chunk_size=config.mamba_chunk_size)
+                                prefix=f"{prefix}.mixer")
 
         self.feed_forward = BambaMLP(config, quant_config=quant_config)
         self.input_layernorm = RMSNorm(config.hidden_size,

vllm/model_executor/models/falcon_h1.py

@@ -109,7 +109,6 @@ def __init__(
             quant_config=quant_config,
             use_rms_norm=config.mamba_rms_norm,
             prefix=f"{prefix}.mixer",
-            chunk_size=config.mamba_chunk_size,
         )
         # n_groups is overridden later by `MambaMixer2`
         self.groups_time_state_size = self.mamba.n_groups * config.mamba_d_state

vllm/model_executor/models/granitemoehybrid.py

@@ -69,8 +69,7 @@ def __init__(self,
                                 rms_norm_eps=config.rms_norm_eps,
                                 activation=config.hidden_act,
                                 quant_config=quant_config,
-                                prefix=f"{prefix}.mixer",
-                                chunk_size=config.mamba_chunk_size)
+                                prefix=f"{prefix}.mixer")
 
         self.block_sparse_moe = None
         if getattr(config, "num_local_experts", 0) > 0:

vllm/model_executor/models/mamba2.py

@@ -62,8 +62,7 @@ def __init__(self,
                                  rms_norm_eps=config.layer_norm_epsilon,
                                  activation=config.hidden_act,
                                  quant_config=quant_config,
-                                 prefix=f"{prefix}.mixer",
-                                 chunk_size=config.chunk_size)
+                                 prefix=f"{prefix}.mixer")
 
         self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
 

vllm/model_executor/models/nemotron_h.py

@@ -154,7 +154,6 @@ def __init__(
             activation=config.mamba_hidden_act,
             quant_config=quant_config,
             prefix=f"{prefix}.mixer",
-            chunk_size=config.chunk_size,
         )
 
         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

vllm/model_executor/models/zamba2.py

@@ -501,8 +501,7 @@ def __init__(self,
                                  rms_norm_eps=config.rms_norm_eps,
                                  activation="silu",
                                  quant_config=quant_config,
-                                 prefix=f"{prefix}.mixer",
-                                 chunk_size=config.chunk_size)
+                                 prefix=f"{prefix}.mixer")
 
         # Input normalization
         self.input_layernorm = RMSNorm(config.hidden_size,

vllm/v1/attention/backends/mamba_attn.py

@@ -6,7 +6,6 @@
 import torch
 
 from vllm.attention.backends.abstract import AttentionBackend
-from vllm.config import VllmConfig, get_layers_from_vllm_config
 from vllm.model_executor.layers.mamba.mamba2_metadata import (
     _query_start_loc_to_chunk_indices_offsets)
 from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
@@ -20,15 +19,6 @@
     from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
 
-def get_mamba2_chunk_size(vllm_config: VllmConfig) -> int:
-    from vllm.model_executor.layers.mamba.mamba_mixer2 import Mamba2Layer
-    layers = get_layers_from_vllm_config(vllm_config, Mamba2Layer)
-    chunk_sizes = set(layer.chunk_size for layer in layers.values())
-    assert len(
-        chunk_sizes) == 1, "All Mamba2 layers must have the same chunk size"
-    return chunk_sizes.pop()
-
-
 class Mamba2AttentionBackend(AttentionBackend):
 
     @staticmethod
@@ -63,7 +53,10 @@ def __init__(self, runner: "GPUModelRunner", kv_cache_spec: MambaSpec,
         self.runner = runner
         self.kv_cache_spec = kv_cache_spec
         self.block_table = block_table
-        self.chunk_size = get_mamba2_chunk_size(runner.vllm_config)
+        self.chunk_size = runner.vllm_config.model_config.get_mamba_chunk_size(
+        )
+        assert self.chunk_size is not None, (
+            "chunk_size needs to be set in the model config for Mamba2 models")
 
     def reorder_batch(self, input_batch: "InputBatch",
                       scheduler_output: "SchedulerOutput") -> bool: