[V1] Enable Mamba2 layers other than MambaMixer2 in the v1 engine

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

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from abc import ABC, abstractmethod
+from collections.abc import Iterable
 from typing import Optional, Union
 
 import torch
@@ -216,9 +218,32 @@ def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
     return loader
 
 
+class Mamba2Layer(ABC):
+    """
+    Base class for all Mamba2 layers which support the v1 engine.
+    Inherit from this class if you implement a custom Mamba2 layer.
+    """
+
+    # 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
+    # only runs for v1, we have to do this to unify with the interface
+    # of Attention + v0 PP.
+    kv_cache: list[Iterable[torch.Tensor]]
+
+    @abstractmethod
+    def get_state_shape(self) -> Iterable[tuple[int, ...]]:
+        """
+        Defines the shape of the mamba state.
+        Usually, the mamba state is a (conv_state, ssm_state) tuple.
+        In this case, returns (conv_state_shape, ssm_state_shape).
+        """
+        pass
+
+
 # Adapted from transformers.models.mamba.modeling_mamba.MambaMixer
 @CustomOp.register("mamba_mixer2")
-class MambaMixer2(CustomOp):
+class MambaMixer2(Mamba2Layer, CustomOp):
     """
     Compute ∆, A, B, C, and D the state space parameters and compute
     the `contextualized_states`. A, D are input independent
@@ -230,22 +255,21 @@ class MambaMixer2(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__()
 
@@ -427,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 MambaMixer2
-    layers = get_layers_from_vllm_config(vllm_config, MambaMixer2)
-    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:

vllm/v1/worker/gpu_model_runner.py

@@ -31,7 +31,7 @@
 from vllm.forward_context import (DPMetadata, get_forward_context,
                                   set_forward_context)
 from vllm.logger import init_logger
-from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
+from vllm.model_executor.layers.mamba.mamba_mixer2 import Mamba2Layer
 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
 from vllm.model_executor.model_loader import TensorizerLoader, get_model_loader
 from vllm.model_executor.models.interfaces import (has_step_pooler,
@@ -2660,7 +2660,7 @@ def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
                     f"Unknown attention type: {attn_module.attn_type}")
 
         mamba_layers = get_layers_from_vllm_config(self.vllm_config,
-                                                   MambaMixer2)
+                                                   Mamba2Layer)
         if len(mamba_layers) > 0:
             if self.vllm_config.speculative_config is not None:
                 raise NotImplementedError(
@@ -2691,7 +2691,7 @@ def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
     def _maybe_pad_mamba_page_size(
         self,
         attn_layers: dict[str, Attention],
-        mamba_layers: dict[str, MambaMixer2],
+        mamba_layers: dict[str, Mamba2Layer],
         kv_cache_spec: dict[str, KVCacheSpec],
         max_model_len: int,
         block_size: int,