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 3, 22:45 UTC
[Misc] Replace mamba_type string literals with MambaAttentionBackendEnum (#41188)
Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>
wangxiyuan committed
May 11, 2026 at 11:59 UTC
5536fc0c019db9b1f42afcc18fc38f4666b65e23
18 files changed
+64
-58
docs/contributing/model/basic.md
+1
-1
@@ -142,7 +142,7 @@ We use "mamba-like" to refer to layers that possess a state that is updated in-p
142
For implementing new custom mamba-like layers, one should inherit from `MambaBase` and implement the methods `get_state_dtype`, `get_state_shape` to calculate the data types and state shapes at runtime, as well as `mamba_type` and `get_attn_backend`.
143
It is also necessary to implement the "attention meta-data" class which handles the meta-data that is common across all layers.
144
Please see [`LinearAttentionMetadata`](../../../vllm/v1/attention/backends/linear_attn.py) or [`ShortConvAttentionMetadata`](../../../vllm/v1/attention/backends/short_conv_attn.py) for examples of this.
145
-It is also worth noting that we should update `MAMBA_TYPE_TO_BACKEND_MAP` and `MambaAttentionBackendEnum` in [`registry.py`](../../../vllm/v1/attention/backends/registry.py) when adding a new mamba backend.
145
+It is also worth noting that we should update `MambaAttentionBackendEnum` in [`registry.py`](../../../vllm/v1/attention/backends/registry.py) when adding a new mamba backend.
146
Finally, if one wants to support torch compile and CUDA graphs, it necessary to wrap the call to the mamba-like layer inside a custom op and register it.
147
Please see the calls to `direct_register_custom_op` in [vllm/model_executor/models/minimax_text_01.py](../../../vllm/model_executor/models/minimax_text_01.py) or [vllm/model_executor/layers/mamba/short_conv.py](../../../vllm/model_executor/layers/mamba/short_conv.py) for examples of this.
148
The new custom op should then be added to the list `_attention_ops` in [vllm/config/compilation.py](../../../vllm/config/compilation.py) to ensure that piecewise CUDA graphs works as intended.
tests/kernels/mamba/test_ssu_dispatch.py
+10
-2
@@ -13,6 +13,7 @@ from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
13
selective_state_update,
14
)
15
from vllm.utils.torch_utils import set_random_seed
16
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
17
from vllm.v1.kv_cache_interface import (
18
KVCacheConfig,
19
KVCacheGroupSpec,
@@ -27,7 +28,9 @@ except ImportError:
28
HAS_FLASHINFER = False
29
30
30
-def _kv_cache_config_with_ssu(mamba_type: str = "mamba2") -> KVCacheConfig:
31
+def _kv_cache_config_with_ssu(
32
+ mamba_type: MambaAttentionBackendEnum = MambaAttentionBackendEnum.MAMBA2,
33
+) -> KVCacheConfig:
34
spec = MambaSpec(
35
block_size=16,
36
shapes=((16, 64),),
@@ -77,7 +80,12 @@ def test_uninitialized_backend_raises():
80
81
82
@pytest.mark.parametrize(
80
- "mamba_type", ["linear_attention", "gdn_attention", "short_conv"]
83
+ "mamba_type",
84
+ [
85
+ MambaAttentionBackendEnum.LINEAR,
86
+ MambaAttentionBackendEnum.GDN_ATTN,
87
+ MambaAttentionBackendEnum.SHORT_CONV,
88
+ ],
89
)
90
def test_init_is_noop_for_non_ssu_mamba_type(mamba_type):
91
import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod
tests/v1/attention/test_attention_backends_selection.py
+13
-8
@@ -13,6 +13,7 @@ from vllm.model_executor.models.minimax_text_01 import MiniMaxText01LinearAttent
13
from vllm.v1.attention.backends.linear_attn import LinearAttentionBackend
14
from vllm.v1.attention.backends.mamba1_attn import Mamba1AttentionBackend
15
from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionBackend
16
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
17
from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionBackend
18
19
@@ -32,7 +33,7 @@ from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionBackend
33
use_rms_norm=True,
34
),
35
Mamba1AttentionBackend,
35
- "mamba1",
36
+ MambaAttentionBackendEnum.MAMBA1,
37
),
38
(
39
MambaMixer2,
@@ -48,7 +49,7 @@ from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionBackend
49
head_dim=32,
50
),
51
Mamba2AttentionBackend,
51
- "mamba2",
52
+ MambaAttentionBackendEnum.MAMBA2,
53
),
54
(
55
MiniMaxText01LinearAttention,
@@ -64,7 +65,7 @@ from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionBackend
65
linear_layer_idx=0,
66
),
67
LinearAttentionBackend,
67
- "linear_attention",
68
+ MambaAttentionBackendEnum.LINEAR,
69
),
70
(
71
ShortConv,
@@ -74,7 +75,7 @@ from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionBackend
75
layer_idx=0,
76
),
77
ShortConvAttentionBackend,
77
- "short_conv",
78
+ MambaAttentionBackendEnum.SHORT_CONV,
79
),
80
],
81
)
@@ -97,10 +98,14 @@ def test_mamba_layers_get_attn_backend(
98
@pytest.mark.parametrize(
99
"layer_class,expected_backend,expected_mamba_type",
100
[
100
- (MambaMixer, Mamba1AttentionBackend, "mamba1"),
101
- (MambaMixer2, Mamba2AttentionBackend, "mamba2"),
102
- (MiniMaxText01LinearAttention, LinearAttentionBackend, "linear_attention"),
103
- (ShortConv, ShortConvAttentionBackend, "short_conv"),
101
+ (MambaMixer, Mamba1AttentionBackend, MambaAttentionBackendEnum.MAMBA1),
102
+ (MambaMixer2, Mamba2AttentionBackend, MambaAttentionBackendEnum.MAMBA2),
103
+ (
104
+ MiniMaxText01LinearAttention,
105
+ LinearAttentionBackend,
106
+ MambaAttentionBackendEnum.LINEAR,
107
+ ),
108
+ (ShortConv, ShortConvAttentionBackend, MambaAttentionBackendEnum.SHORT_CONV),
109
],
110
)
111
def test_mamba_layers_have_unified_interface(
vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py
+2
-1
@@ -13,6 +13,7 @@ from dataclasses import dataclass
13
import torch
14
15
from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
16
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
17
from vllm.v1.kv_cache_interface import MambaSpec
18
19
@@ -103,7 +104,7 @@ def derive_mamba_conv_split(
104
MambaConvSplitInfo with per-rank x_local, b_local, conv_rows,
105
conv_dtype_size, and ssm_sizes (conv_state_bytes, ssm_state_bytes).
106
"""
106
- if mamba_spec.mamba_type != "mamba2":
107
+ if mamba_spec.mamba_type != MambaAttentionBackendEnum.MAMBA2:
108
raise NotImplementedError(
109
f"3-read conv transfer only supports Mamba2 models, "
110
f"got mamba_type={mamba_spec.mamba_type!r}. "
vllm/model_executor/layers/kda.py
+3
-2
@@ -17,6 +17,7 @@ from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader
17
from vllm.model_executor.utils import set_weight_attrs
18
from vllm.utils.torch_utils import direct_register_custom_op
19
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
20
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
21
22
from .fla.ops.kda import (
23
FusedRMSNormGated,
@@ -84,8 +85,8 @@ direct_register_custom_op(
85
86
class KimiDeltaAttention(nn.Module, MambaBase):
87
@property
87
- def mamba_type(self) -> str:
88
- return "gdn_attention"
88
+ def mamba_type(self) -> MambaAttentionBackendEnum:
89
+ return MambaAttentionBackendEnum.GDN_ATTN
90
91
def get_state_dtype(
92
self,
vllm/model_executor/layers/mamba/abstract.py
+2
-1
@@ -8,6 +8,7 @@ import torch
8
from vllm.config import VllmConfig
9
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
10
from vllm.v1.attention.backend import AttentionBackend
11
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
12
from vllm.v1.attention.selector import get_mamba_attn_backend
13
from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec
14
@@ -33,7 +34,7 @@ class MambaBase(AttentionLayerBase):
34
35
@property
36
@abstractmethod
36
- def mamba_type(self) -> str:
37
+ def mamba_type(self) -> MambaAttentionBackendEnum:
38
pass
39
40
@abstractmethod
vllm/model_executor/layers/mamba/gdn_linear_attn.py
+3
-2
@@ -64,6 +64,7 @@ from vllm.utils.torch_utils import (
64
direct_register_custom_op,
65
)
66
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
67
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
68
69
# Optional ROCm AITER Triton kernels for the GDN decode fast-path.
70
# Availability is checked centrally via rocm_aiter_ops; the actual function
@@ -237,8 +238,8 @@ class ChunkGatedDeltaRule(CustomOp):
238
@PluggableLayer.register("gated_delta_net_attention")
239
class GatedDeltaNetAttention(PluggableLayer, MambaBase):
240
@property
240
- def mamba_type(self) -> str:
241
- return "gdn_attention"
241
+ def mamba_type(self) -> MambaAttentionBackendEnum:
242
+ return MambaAttentionBackendEnum.GDN_ATTN
243
244
def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:
245
return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
vllm/model_executor/layers/mamba/linear_attn.py
+3
-2
@@ -32,6 +32,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
32
from vllm.utils.torch_utils import direct_register_custom_op
33
from vllm.v1.attention.backend import AttentionMetadata
34
from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata
35
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
36
37
38
@CustomOp.register("minimax_text01_rmsnorm_tp")
@@ -246,8 +247,8 @@ class MiniMaxText01LinearKernel:
247
248
class MiniMaxText01LinearAttention(nn.Module, MambaBase):
249
@property
249
- def mamba_type(self) -> str:
250
- return "linear_attention"
250
+ def mamba_type(self) -> MambaAttentionBackendEnum:
251
+ return MambaAttentionBackendEnum.LINEAR
252
253
def get_state_dtype(self) -> tuple[torch.dtype]:
254
assert self.model_config is not None
vllm/model_executor/layers/mamba/mamba_mixer.py
+3
-2
@@ -42,6 +42,7 @@ from vllm.utils.torch_utils import (
42
)
43
from vllm.v1.attention.backend import AttentionMetadata
44
from vllm.v1.attention.backends.mamba1_attn import Mamba1AttentionMetadata
45
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
46
47
48
# Adapted from transformers.models.mamba.modeling_mamba.MambaMixer
@@ -476,8 +477,8 @@ class MambaMixer(MambaBase, PluggableLayer):
477
)
478
479
@property
479
- def mamba_type(self) -> str:
480
- return "mamba1"
480
+ def mamba_type(self) -> MambaAttentionBackendEnum:
481
+ return MambaAttentionBackendEnum.MAMBA1
482
483
def _time_proj_bias(self) -> torch.Tensor | None:
484
if hasattr(self.dt_proj, "bias") and self.dt_proj.bias is not None:
vllm/model_executor/layers/mamba/mamba_mixer2.py
+3
-2
@@ -52,6 +52,7 @@ from vllm.utils.torch_utils import (
52
)
53
from vllm.v1.attention.backend import AttentionMetadata
54
from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata
55
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
56
57
# Added by the IBM Team, 2024
58
@@ -935,8 +936,8 @@ class MambaMixer2(MambaBase, PluggableLayer):
936
)
937
938
@property
938
- def mamba_type(self) -> str:
939
- return "mamba2"
939
+ def mamba_type(self) -> MambaAttentionBackendEnum:
940
+ return MambaAttentionBackendEnum.MAMBA2
941
942
943
def mamba_mixer2(
vllm/model_executor/layers/mamba/ops/ssu_dispatch.py
+3
-1
@@ -14,6 +14,7 @@ import torch
14
15
from vllm.config.mamba import MambaBackendEnum, MambaConfig
16
from vllm.logger import init_logger
17
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
18
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
19
from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
20
@@ -200,7 +201,8 @@ def initialize_mamba_ssu_backend(
201
"""
202
if not any(
203
isinstance(g.kv_cache_spec, MambaSpec)
203
- and g.kv_cache_spec.mamba_type in ("mamba1", "mamba2")
204
+ and g.kv_cache_spec.mamba_type
205
+ in (MambaAttentionBackendEnum.MAMBA1, MambaAttentionBackendEnum.MAMBA2)
206
for g in kv_cache_config.kv_cache_groups
207
):
208
return
vllm/model_executor/layers/mamba/short_conv.py
+3
-2
@@ -25,6 +25,7 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
25
)
26
from vllm.utils.torch_utils import direct_register_custom_op
27
from vllm.v1.attention.backend import AttentionMetadata
28
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
29
from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionMetadata
30
31
@@ -223,8 +224,8 @@ class ShortConv(MambaBase, CustomOp):
224
)
225
226
@property
226
- def mamba_type(self) -> str:
227
- return "short_conv"
227
+ def mamba_type(self) -> MambaAttentionBackendEnum:
228
+ return MambaAttentionBackendEnum.SHORT_CONV
229
230
231
def short_conv(
vllm/model_executor/models/bailing_moe_linear.py
+3
-2
@@ -64,6 +64,7 @@ from vllm.model_executor.models.bailing_moe import BailingMLP
64
from vllm.sequence import IntermediateTensors
65
from vllm.v1.attention.backend import AttentionMetadata
66
from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata
67
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
68
69
from .interfaces import HasInnerState, IsHybrid, SupportsPP
70
from .utils import (
@@ -444,8 +445,8 @@ class BailingMoELinearAttention(PluggableLayer, MambaBase):
445
# --8<-- [end:bailing_moe_linear_attention]
446
447
@property
447
- def mamba_type(self) -> str:
448
- return "linear_attention"
448
+ def mamba_type(self) -> MambaAttentionBackendEnum:
449
+ return MambaAttentionBackendEnum.LINEAR
450
451
def get_state_shape(self) -> tuple[tuple[int, ...], ...]:
452
"""Return state shape for linear attention cache.
vllm/model_executor/models/olmo_hybrid.py
+3
-2
@@ -92,6 +92,7 @@ from vllm.triton_utils.allocation import set_triton_allocator
92
from vllm.utils.torch_utils import direct_register_custom_op
93
from vllm.v1.attention.backend import AttentionMetadata
94
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
95
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
96
97
from .interfaces import HasInnerState, IsHybrid, SupportsLoRA, SupportsPP
98
from .utils import (
@@ -136,8 +137,8 @@ class OlmoHybridGatedDeltaNet(nn.Module, MambaBase):
137
"""
138
139
@property
139
- def mamba_type(self) -> str:
140
- return "gdn_attention"
140
+ def mamba_type(self) -> MambaAttentionBackendEnum:
141
+ return MambaAttentionBackendEnum.GDN_ATTN
142
143
def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:
144
return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
vllm/model_executor/models/plamo2.py
+3
-2
@@ -72,6 +72,7 @@ from vllm.sequence import IntermediateTensors
72
from vllm.utils.torch_utils import direct_register_custom_op
73
from vllm.v1.attention.backend import AttentionMetadata
74
from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata
75
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
76
77
# Only used for type hinting.
78
if TYPE_CHECKING:
@@ -478,8 +479,8 @@ class Plamo2MambaMixer(MambaBase, PluggableLayer):
479
)
480
481
@property
481
- def mamba_type(self) -> str:
482
- return "mamba2"
482
+ def mamba_type(self) -> MambaAttentionBackendEnum:
483
+ return MambaAttentionBackendEnum.MAMBA2
484
485
486
def plamo2_mamba_mixer(
vllm/v1/attention/backends/registry.py
-10
@@ -193,16 +193,6 @@ class MambaAttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
193
_MAMBA_ATTN_OVERRIDES.pop(self, None)
194
195
196
-MAMBA_TYPE_TO_BACKEND_MAP = {
197
- "mamba1": MambaAttentionBackendEnum.MAMBA1.name,
198
- "mamba2": MambaAttentionBackendEnum.MAMBA2.name,
199
- "short_conv": MambaAttentionBackendEnum.SHORT_CONV.name,
200
- "linear_attention": MambaAttentionBackendEnum.LINEAR.name,
201
- "gdn_attention": MambaAttentionBackendEnum.GDN_ATTN.name,
202
- "custom": MambaAttentionBackendEnum.CUSTOM.name,
203
-}
204
-
205
-
196
_ATTN_OVERRIDES: dict[AttentionBackendEnum, str] = {}
197
_MAMBA_ATTN_OVERRIDES: dict[MambaAttentionBackendEnum, str] = {}
198
vllm/v1/attention/selector.py
+4
-15
@@ -12,7 +12,6 @@ from vllm.logger import init_logger
12
from vllm.utils.import_utils import resolve_obj_by_qualname
13
from vllm.v1.attention.backend import AttentionBackend, AttentionType
14
from vllm.v1.attention.backends.registry import (
15
- MAMBA_TYPE_TO_BACKEND_MAP,
15
MambaAttentionBackendEnum,
16
)
17
@@ -138,7 +137,7 @@ def _cached_get_attn_backend(
137
138
139
def get_mamba_attn_backend(
141
- mamba_type: str,
140
+ mamba_type: MambaAttentionBackendEnum,
141
) -> type[AttentionBackend]:
142
"""Select which mamba attention backend to use and lazily import it."""
143
return _cached_get_mamba_attn_backend(mamba_type)
@@ -146,21 +145,11 @@ def get_mamba_attn_backend(
145
146
@cache
147
def _cached_get_mamba_attn_backend(
149
- mamba_type: str,
148
+ mamba_type: MambaAttentionBackendEnum,
149
) -> type[AttentionBackend]:
151
- assert mamba_type and isinstance(mamba_type, str)
150
+ assert mamba_type and isinstance(mamba_type, MambaAttentionBackendEnum)
151
153
- selected_backend = None
154
- try:
155
- backend_name = MAMBA_TYPE_TO_BACKEND_MAP[mamba_type]
156
- selected_backend = MambaAttentionBackendEnum[backend_name]
157
- except KeyError as e:
158
- raise ValueError(
159
- f"Invalid mamba attention backend type: '{mamba_type}'. Valid "
160
- f"types are: {list(MAMBA_TYPE_TO_BACKEND_MAP.keys())}"
161
- ) from e
162
-
163
- mamba_attn_backend = selected_backend.get_class()
152
+ mamba_attn_backend = mamba_type.get_class()
153
if envs.VLLM_BATCH_INVARIANT and not mamba_attn_backend.supports_batch_invariance():
154
raise RuntimeError(
155
"VLLM batch_invariant mode is not supported for "
vllm/v1/kv_cache_interface.py
+2
-1
@@ -16,6 +16,7 @@ from typing_extensions import Self
16
from vllm.logger import init_logger
17
from vllm.utils.math_utils import cdiv, round_up
18
from vllm.utils.torch_utils import get_dtype_size, nvfp4_kv_cache_full_dim
19
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
20
21
if TYPE_CHECKING:
22
from vllm.config import VllmConfig
@@ -532,7 +533,7 @@ class MambaSpec(KVCacheSpec):
533
shapes: tuple[tuple[int, ...], ...]
534
dtypes: tuple[torch.dtype]
535
page_size_padded: int | None = None
535
- mamba_type: str = "mamba2"
536
+ mamba_type: MambaAttentionBackendEnum = MambaAttentionBackendEnum.MAMBA2
537
mamba_cache_mode: str = "none"
538
num_speculative_blocks: int = 0
539