A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 02:15 UTC
[V1] Enable Mamba2 layers other than MambaMixer2 in the v1 engine (#20660)
Signed-off-by: nopperl <54780682+nopperl@users.noreply.github.com>
nopperl committed
Jul 11, 2025 at 14:53 UTC
5d09152ff13f7f2dabc45d5977e23dbf107ef5af
11 files changed
+68
-45
vllm/config.py
+11
@@ -1331,6 +1331,17 @@ class ModelConfig:
1331
1332
return sum(t == 1 for t in attn_type_list[start:end])
1333
1334
+ def get_mamba_chunk_size(self) -> Optional[int]:
1335
+ """
1336
+ Returns the mamba chunk size if it exists
1337
+ """
1338
+ # used by e.g. Bamba, FalconH1, Granite, PLaMo2
1339
+ chunk_size = getattr(self.hf_text_config, "mamba_chunk_size", None)
1340
+ if chunk_size is None:
1341
+ # used by e.g. Mamba2, NemotronH, Zamba
1342
+ chunk_size = getattr(self.hf_text_config, "chunk_size", None)
1343
+ return chunk_size
1344
+
1345
def get_multimodal_config(self) -> "MultiModalConfig":
1346
"""
1347
Get the multimodal configuration of the model.
vllm/model_executor/layers/mamba/abstract.py
new
+29
@@ -0,0 +1,29 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from abc import ABC, abstractmethod
4
+from collections.abc import Iterable
5
+
6
+import torch
7
+
8
+
9
+class MambaBase(ABC):
10
+ """
11
+ Base class for Mamba-like layers which support the v1 engine.
12
+ Inherit from this class if you implement a custom layer.
13
+ """
14
+
15
+ # Contains the KV cache (mamba state) for the layer
16
+ # in the shape specified by `self.get_state_shape`.
17
+ # The outer list is for v0 PP virtual engine. Though this code path
18
+ # only runs for v1, we have to do this to unify with the interface
19
+ # of Attention + v0 PP.
20
+ kv_cache: list[Iterable[torch.Tensor]]
21
+
22
+ @abstractmethod
23
+ def get_state_shape(self) -> Iterable[tuple[int, ...]]:
24
+ """
25
+ Defines the shape of the state.
26
+ For mamba layers this is usually a (conv_state, ssm_state) tuple.
27
+ In this case, returns (conv_state_shape, ssm_state_shape).
28
+ """
29
+ pass
vllm/model_executor/layers/mamba/mamba_mixer2.py
+17
-20
@@ -17,6 +17,7 @@ from vllm.forward_context import get_forward_context
17
from vllm.model_executor.custom_op import CustomOp
18
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
19
RowParallelLinear)
20
+from vllm.model_executor.layers.mamba.abstract import MambaBase
21
from vllm.model_executor.layers.mamba.mamba2_metadata import (Mamba2Metadata,
22
update_metadata)
23
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
@@ -219,7 +220,7 @@ def mamba_v2_sharded_weight_loader(
220
221
# Adapted from transformers.models.mamba.modeling_mamba.MambaMixer
222
@CustomOp.register("mamba_mixer2")
222
-class MambaMixer2(CustomOp):
223
+class MambaMixer2(MambaBase, CustomOp):
224
"""
225
Compute ∆, A, B, C, and D the state space parameters and compute
226
the `contextualized_states`. A, D are input independent
@@ -231,22 +232,21 @@ class MambaMixer2(CustomOp):
232
"""
233
234
def __init__(
234
- self,
235
- hidden_size: int,
236
- ssm_state_size: int,
237
- conv_kernel_size: int,
238
- intermediate_size: int,
239
- use_conv_bias: bool,
240
- use_bias: bool,
241
- n_groups: int = 1,
242
- num_heads: int = 128,
243
- head_dim: int = 64,
244
- rms_norm_eps: float = 1e-5,
245
- activation: str = "silu",
246
- use_rms_norm: bool = True,
247
- quant_config: Optional[QuantizationConfig] = None,
248
- prefix: str = "",
249
- chunk_size: int = -1, # the chunk size used by v1
235
+ self,
236
+ hidden_size: int,
237
+ ssm_state_size: int,
238
+ conv_kernel_size: int,
239
+ intermediate_size: int,
240
+ use_conv_bias: bool,
241
+ use_bias: bool,
242
+ n_groups: int = 1,
243
+ num_heads: int = 128,
244
+ head_dim: int = 64,
245
+ rms_norm_eps: float = 1e-5,
246
+ activation: str = "silu",
247
+ use_rms_norm: bool = True,
248
+ quant_config: Optional[QuantizationConfig] = None,
249
+ prefix: str = "",
250
):
251
super().__init__()
252
@@ -428,10 +428,7 @@ class MambaMixer2(CustomOp):
428
# of Attention + v0 PP.
429
# The inner tuple is (conv_state, ssm_state)
430
self.kv_cache = [(torch.tensor([]), torch.tensor([]))]
431
- assert chunk_size != -1, "chunk_size must be set for v1"
431
433
- # NOTE: chunk_size may be -1 for models without v1 support
434
- self.chunk_size = chunk_size
432
self.prefix = prefix
433
434
def forward_native(
vllm/model_executor/models/bamba.py
+1
-2
@@ -99,8 +99,7 @@ class BambaMixerDecoderLayer(nn.Module):
99
rms_norm_eps=config.rms_norm_eps,
100
activation=config.hidden_act,
101
quant_config=quant_config,
102
- prefix=f"{prefix}.mixer",
103
- chunk_size=config.mamba_chunk_size)
102
+ prefix=f"{prefix}.mixer")
103
104
self.feed_forward = BambaMLP(config, quant_config=quant_config)
105
self.input_layernorm = RMSNorm(config.hidden_size,
vllm/model_executor/models/falcon_h1.py
-1
@@ -109,7 +109,6 @@ class FalconH1SSMDecoderLayer(nn.Module):
109
quant_config=quant_config,
110
use_rms_norm=config.mamba_rms_norm,
111
prefix=f"{prefix}.mixer",
112
- chunk_size=config.mamba_chunk_size,
112
)
113
# n_groups is overridden later by `MambaMixer2`
114
self.groups_time_state_size = self.mamba.n_groups * config.mamba_d_state
vllm/model_executor/models/granitemoehybrid.py
+1
-2
@@ -69,8 +69,7 @@ class GraniteMoeHybridMambaDecoderLayer(nn.Module):
69
rms_norm_eps=config.rms_norm_eps,
70
activation=config.hidden_act,
71
quant_config=quant_config,
72
- prefix=f"{prefix}.mixer",
73
- chunk_size=config.mamba_chunk_size)
72
+ prefix=f"{prefix}.mixer")
73
74
self.block_sparse_moe = None
75
if getattr(config, "num_local_experts", 0) > 0:
vllm/model_executor/models/mamba2.py
+1
-2
@@ -62,8 +62,7 @@ class Mamba2DecoderLayer(nn.Module):
62
rms_norm_eps=config.layer_norm_epsilon,
63
activation=config.hidden_act,
64
quant_config=quant_config,
65
- prefix=f"{prefix}.mixer",
66
- chunk_size=config.chunk_size)
65
+ prefix=f"{prefix}.mixer")
66
67
self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
68
vllm/model_executor/models/nemotron_h.py
-1
@@ -154,7 +154,6 @@ class NemotronHMambaDecoderLayer(nn.Module):
154
activation=config.mamba_hidden_act,
155
quant_config=quant_config,
156
prefix=f"{prefix}.mixer",
157
- chunk_size=config.chunk_size,
157
)
158
159
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
vllm/model_executor/models/zamba2.py
+1
-2
@@ -501,8 +501,7 @@ class Zamba2MambaDecoderLayer(nn.Module):
501
rms_norm_eps=config.rms_norm_eps,
502
activation="silu",
503
quant_config=quant_config,
504
- prefix=f"{prefix}.mixer",
505
- chunk_size=config.chunk_size)
504
+ prefix=f"{prefix}.mixer")
505
506
# Input normalization
507
self.input_layernorm = RMSNorm(config.hidden_size,
vllm/v1/attention/backends/mamba_attn.py
+4
-11
@@ -7,7 +7,6 @@ from typing import TYPE_CHECKING, Optional
7
import torch
8
9
from vllm.attention.backends.abstract import AttentionBackend
10
-from vllm.config import VllmConfig, get_layers_from_vllm_config
10
from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
11
CommonAttentionMetadata)
12
from vllm.v1.kv_cache_interface import MambaSpec
@@ -19,15 +18,6 @@ if TYPE_CHECKING:
18
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
19
20
22
-def get_mamba2_chunk_size(vllm_config: VllmConfig) -> int:
23
- from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
24
- layers = get_layers_from_vllm_config(vllm_config, MambaMixer2)
25
- chunk_sizes = set(layer.chunk_size for layer in layers.values())
26
- assert len(
27
- chunk_sizes) == 1, "All Mamba2 layers must have the same chunk size"
28
- return chunk_sizes.pop()
29
-
30
-
21
def _query_start_loc_to_chunk_indices_offsets(query_start_loc: torch.Tensor,
22
chunk_size: int,
23
total_seqlens: int):
@@ -102,7 +92,10 @@ class Mamba2AttentionMetadataBuilder(
92
self.runner = runner
93
self.kv_cache_spec = kv_cache_spec
94
self.block_table = block_table
105
- self.chunk_size = get_mamba2_chunk_size(runner.vllm_config)
95
+ self.chunk_size = runner.vllm_config.model_config.get_mamba_chunk_size(
96
+ )
97
+ assert self.chunk_size is not None, (
98
+ "chunk_size needs to be set in the model config for Mamba2 models")
99
100
def reorder_batch(self, input_batch: "InputBatch",
101
scheduler_output: "SchedulerOutput") -> bool:
vllm/v1/worker/gpu_model_runner.py
+3
-4
@@ -30,7 +30,7 @@ from vllm.distributed.parallel_state import (
30
from vllm.forward_context import (DPMetadata, get_forward_context,
31
set_forward_context)
32
from vllm.logger import init_logger
33
-from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
33
+from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaBase
34
from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
35
from vllm.model_executor.model_loader import TensorizerLoader, get_model_loader
36
from vllm.model_executor.models.interfaces import (has_step_pooler,
@@ -2623,8 +2623,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2623
raise ValueError(
2624
f"Unknown attention type: {attn_module.attn_type}")
2625
2626
- mamba_layers = get_layers_from_vllm_config(self.vllm_config,
2627
- MambaMixer2)
2626
+ mamba_layers = get_layers_from_vllm_config(self.vllm_config, MambaBase)
2627
if len(mamba_layers) > 0:
2628
if self.vllm_config.speculative_config is not None:
2629
raise NotImplementedError(
@@ -2655,7 +2654,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2654
def _maybe_pad_mamba_page_size(
2655
self,
2656
attn_layers: dict[str, Attention],
2658
- mamba_layers: dict[str, MambaMixer2],
2657
+ mamba_layers: dict[str, MambaBase],
2658
kv_cache_spec: dict[str, KVCacheSpec],
2659
max_model_len: int,
2660
block_size: int,