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 12, 23:30 UTC
[Model][Mamba] Add selector for mamba attention backend and make it pluggable for other device (#26487)
Signed-off-by: shen-shanshan <467638484@qq.com>
Shanshan Shen committed
Nov 20, 2025 at 00:24 UTC
d44e9df7d49a9bb3400b002c38c06fae2dd7d1e8
12 files changed
+144
-85
docs/contributing/model/basic.md
+1
@@ -146,6 +146,7 @@ We use "mamba-like" to refer to layers that posses a state that is updated in-pl
146
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`.
147
It is also necessary to implement the "attention meta-data" class which handles the meta-data that is common across all layers.
148
Please see [`LinearAttentionMetadata`](../../../vllm/v1/attention/backends/linear_attn.py) or [`ShortConvAttentionMetadata`](../../../vllm/v1/attention/backends/short_conv_attn.py) for examples of this.
149
+It is also worth noting that we should update `MAMBA_TYPE_TO_BACKEND_MAP` and `MambaAttentionBackendEnum` in [`registry.py`](../../../vllm/attention/backends/registry.py) when adding a new mamba backend.
150
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.
151
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.
152
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.
vllm/attention/__init__.py
+2
-1
@@ -7,7 +7,7 @@ from vllm.attention.backends.abstract import (
7
AttentionType,
8
)
9
from vllm.attention.layer import Attention
10
-from vllm.attention.selector import get_attn_backend
10
+from vllm.attention.selector import get_attn_backend, get_mamba_attn_backend
11
12
__all__ = [
13
"Attention",
@@ -15,4 +15,5 @@ __all__ = [
15
"AttentionMetadata",
16
"AttentionType",
17
"get_attn_backend",
18
+ "get_mamba_attn_backend",
19
]
vllm/attention/backends/registry.py
+100
-14
@@ -2,8 +2,8 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Attention backend registry"""
4
5
-import enum
5
from collections.abc import Callable
6
+from enum import Enum, EnumMeta
7
from typing import TYPE_CHECKING, cast
8
9
from vllm.logger import init_logger
@@ -15,7 +15,7 @@ if TYPE_CHECKING:
15
logger = init_logger(__name__)
16
17
18
-class _AttentionBackendEnumMeta(enum.EnumMeta):
18
+class _AttentionBackendEnumMeta(EnumMeta):
19
"""Metaclass for AttentionBackendEnum to provide better error messages."""
20
21
def __getitem__(cls, name: str):
@@ -23,15 +23,15 @@ class _AttentionBackendEnumMeta(enum.EnumMeta):
23
try:
24
return super().__getitem__(name)
25
except KeyError:
26
- members = cast("dict[str, AttentionBackendEnum]", cls.__members__).values()
27
- valid_backends = ", ".join(m.name for m in members)
26
+ members = cast("dict[str, Enum]", cls.__members__).keys()
27
+ valid_backends = ", ".join(members)
28
raise ValueError(
29
f"Unknown attention backend: '{name}'. "
30
f"Valid options are: {valid_backends}"
31
) from None
32
33
34
-class AttentionBackendEnum(enum.Enum, metaclass=_AttentionBackendEnumMeta):
34
+class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
35
"""Enumeration of all supported attention backends.
36
37
The enum value is the default class path, but this can be overridden
@@ -83,7 +83,7 @@ class AttentionBackendEnum(enum.Enum, metaclass=_AttentionBackendEnumMeta):
83
Raises:
84
ValueError: If Backend.CUSTOM is used without being registered
85
"""
86
- path = _OVERRIDES.get(self, self.value)
86
+ path = _ATTN_OVERRIDES.get(self, self.value)
87
if not path:
88
raise ValueError(
89
f"Backend {self.name} must be registered before use. "
@@ -111,18 +111,93 @@ class AttentionBackendEnum(enum.Enum, metaclass=_AttentionBackendEnumMeta):
111
Returns:
112
True if the backend has a registered override
113
"""
114
- return self in _OVERRIDES
114
+ return self in _ATTN_OVERRIDES
115
116
def clear_override(self) -> None:
117
"""Clear any override for this backend, reverting to the default."""
118
- _OVERRIDES.pop(self, None)
118
+ _ATTN_OVERRIDES.pop(self, None)
119
120
121
-_OVERRIDES: dict[AttentionBackendEnum, str] = {}
121
+class MambaAttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
122
+ """Enumeration of all supported mamba attention backends.
123
+
124
+ The enum value is the default class path, but this can be overridden
125
+ at runtime using register_backend().
126
+
127
+ To get the actual backend class (respecting overrides), use:
128
+ backend.get_class()
129
+ """
130
+
131
+ MAMBA1 = "vllm.v1.attention.backends.mamba1_attn.Mamba1AttentionBackend"
132
+ MAMBA2 = "vllm.v1.attention.backends.mamba2_attn.Mamba2AttentionBackend"
133
+ SHORT_CONV = "vllm.v1.attention.backends.short_conv_attn.ShortConvAttentionBackend"
134
+ LINEAR = "vllm.v1.attention.backends.linear_attn.LinearAttentionBackend"
135
+ GDN_ATTN = "vllm.v1.attention.backends.gdn_attn.GDNAttentionBackend"
136
+ # Placeholder for third-party/custom backends - must be registered before use
137
+ CUSTOM = ""
138
+
139
+ def get_path(self, include_classname: bool = True) -> str:
140
+ """Get the class path for this backend (respects overrides).
141
+
142
+ Returns:
143
+ The fully qualified class path string
144
+
145
+ Raises:
146
+ ValueError: If Backend.CUSTOM is used without being registered
147
+ """
148
+ path = _MAMBA_ATTN_OVERRIDES.get(self, self.value)
149
+ if not path:
150
+ raise ValueError(
151
+ f"Backend {self.name} must be registered before use. "
152
+ f"Use register_backend(Backend.{self.name}, 'your.module.YourClass')"
153
+ )
154
+ if not include_classname:
155
+ path = path.rsplit(".", 1)[0]
156
+ return path
157
+
158
+ def get_class(self) -> "type[AttentionBackend]":
159
+ """Get the backend class (respects overrides).
160
+
161
+ Returns:
162
+ The backend class
163
+
164
+ Raises:
165
+ ImportError: If the backend class cannot be imported
166
+ ValueError: If Backend.CUSTOM is used without being registered
167
+ """
168
+ return resolve_obj_by_qualname(self.get_path())
169
+
170
+ def is_overridden(self) -> bool:
171
+ """Check if this backend has been overridden.
172
+
173
+ Returns:
174
+ True if the backend has a registered override
175
+ """
176
+ return self in _MAMBA_ATTN_OVERRIDES
177
+
178
+ def clear_override(self) -> None:
179
+ """Clear any override for this backend, reverting to the default."""
180
+ _MAMBA_ATTN_OVERRIDES.pop(self, None)
181
+
182
+
183
+MAMBA_TYPE_TO_BACKEND_MAP = {
184
+ "mamba1": MambaAttentionBackendEnum.MAMBA1.name,
185
+ "mamba2": MambaAttentionBackendEnum.MAMBA2.name,
186
+ "short_conv": MambaAttentionBackendEnum.SHORT_CONV.name,
187
+ "linear_attention": MambaAttentionBackendEnum.LINEAR.name,
188
+ "gdn_attention": MambaAttentionBackendEnum.GDN_ATTN.name,
189
+ "custom": MambaAttentionBackendEnum.CUSTOM.name,
190
+}
191
+
192
+
193
+_ATTN_OVERRIDES: dict[AttentionBackendEnum, str] = {}
194
+_MAMBA_ATTN_OVERRIDES: dict[MambaAttentionBackendEnum, str] = {}
195
196
197
def register_backend(
125
- backend: AttentionBackendEnum, class_path: str | None = None
198
+ backend: AttentionBackendEnum | MambaAttentionBackendEnum,
199
+ is_mamba: bool = False,
200
+ class_path: str | None = None,
201
) -> Callable[[type], type]:
202
"""Register or override a backend implementation.
203
@@ -135,12 +210,17 @@ def register_backend(
210
Decorator function if class_path is None, otherwise a no-op
211
212
Examples:
138
- # Override an existing backend
213
+ # Override an existing attention backend
214
@register_backend(AttentionBackendEnum.FLASH_ATTN)
215
class MyCustomFlashAttn:
216
...
217
143
- # Register a custom third-party backend
218
+ # Override an existing mamba attention backend
219
+ @register_backend(MambaAttentionBackendEnum.LINEAR, is_mamba=True)
220
+ class MyCustomMambaAttn:
221
+ ...
222
+
223
+ # Register a custom third-party attention backend
224
@register_backend(AttentionBackendEnum.CUSTOM)
225
class MyCustomBackend:
226
...
@@ -153,11 +233,17 @@ def register_backend(
233
"""
234
235
def decorator(cls: type) -> type:
156
- _OVERRIDES[backend] = f"{cls.__module__}.{cls.__qualname__}"
236
+ if is_mamba:
237
+ _MAMBA_ATTN_OVERRIDES[backend] = f"{cls.__module__}.{cls.__qualname__}" # type: ignore[index]
238
+ else:
239
+ _ATTN_OVERRIDES[backend] = f"{cls.__module__}.{cls.__qualname__}" # type: ignore[index]
240
return cls
241
242
if class_path is not None:
160
- _OVERRIDES[backend] = class_path
243
+ if is_mamba:
244
+ _MAMBA_ATTN_OVERRIDES[backend] = class_path # type: ignore[index]
245
+ else:
246
+ _ATTN_OVERRIDES[backend] = class_path # type: ignore[index]
247
return lambda x: x
248
249
return decorator
vllm/attention/selector.py
+32
-1
@@ -12,7 +12,11 @@ import torch
12
13
import vllm.envs as envs
14
from vllm.attention.backends.abstract import AttentionBackend
15
-from vllm.attention.backends.registry import AttentionBackendEnum
15
+from vllm.attention.backends.registry import (
16
+ MAMBA_TYPE_TO_BACKEND_MAP,
17
+ AttentionBackendEnum,
18
+ MambaAttentionBackendEnum,
19
+)
20
from vllm.config.cache import CacheDType
21
from vllm.logger import init_logger
22
from vllm.utils import STR_BACKEND_ENV_VAR
@@ -197,6 +201,33 @@ def _cached_get_attn_backend(
201
return backend
202
203
204
+def get_mamba_attn_backend(
205
+ mamba_type: str,
206
+) -> type[AttentionBackend]:
207
+ """Select which mamba attention backend to use and lazily import it."""
208
+ return _cached_get_mamba_attn_backend(mamba_type)
209
+
210
+
211
+@cache
212
+def _cached_get_mamba_attn_backend(
213
+ mamba_type: str,
214
+) -> type[AttentionBackend]:
215
+ assert mamba_type and isinstance(mamba_type, str)
216
+
217
+ selected_backend = None
218
+ try:
219
+ backend_name = MAMBA_TYPE_TO_BACKEND_MAP[mamba_type]
220
+ selected_backend = MambaAttentionBackendEnum[backend_name]
221
+ except KeyError as e:
222
+ raise ValueError(
223
+ f"Invalid mamba attention backend type: '{backend_name}'. Valid "
224
+ f"backends are: {list(MambaAttentionBackendEnum.__members__.keys())}"
225
+ ) from e
226
+
227
+ mamba_attn_backend = selected_backend.get_class()
228
+ return mamba_attn_backend
229
+
230
+
231
@contextmanager
232
def global_force_attn_backend_context_manager(
233
attn_backend: AttentionBackendEnum,
vllm/model_executor/layers/kda.py
+1
-7
@@ -5,7 +5,6 @@ import torch
5
from einops import rearrange
6
from torch import nn
7
8
-from vllm.attention import AttentionBackend
8
from vllm.attention.backends.abstract import AttentionMetadata
9
from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config
10
from vllm.distributed import (
@@ -83,12 +82,7 @@ direct_register_custom_op(
82
class KimiDeltaAttention(nn.Module, MambaBase):
83
@property
84
def mamba_type(self) -> str:
86
- return "linear_attention"
87
-
88
- def get_attn_backend(self) -> type["AttentionBackend"]:
89
- from vllm.v1.attention.backends.gdn_attn import GDNAttentionBackend
90
-
91
- return GDNAttentionBackend
85
+ return "gdn_attention"
86
87
def get_state_dtype(
88
self,
vllm/model_executor/layers/mamba/abstract.py
+5
-5
@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING
6
7
import torch
8
9
+from vllm.attention.selector import get_mamba_attn_backend
10
from vllm.config import VllmConfig
11
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
12
from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec
@@ -38,11 +39,6 @@ class MambaBase(AttentionLayerBase):
39
def mamba_type(self) -> str:
40
pass
41
41
- @abstractmethod
42
- def get_attn_backend(self) -> type["AttentionBackend"]:
43
- """Get the attention backend class for this Mamba layer."""
44
- pass
45
-
42
@abstractmethod
43
def get_state_dtype(self) -> tuple[torch.dtype, ...]:
44
pass
@@ -69,3 +65,7 @@ class MambaBase(AttentionLayerBase):
65
else 0
66
),
67
)
68
+
69
+ def get_attn_backend(self) -> type["AttentionBackend"]:
70
+ """Get the attention backend class for this Mamba layer."""
71
+ return get_mamba_attn_backend(self.mamba_type)
vllm/model_executor/layers/mamba/linear_attn.py
-14
@@ -2,12 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import math
5
-from typing import TYPE_CHECKING
6
-
7
-if TYPE_CHECKING:
8
- from vllm.attention.backends.abstract import AttentionBackend
9
-
10
-from typing import TYPE_CHECKING
5
6
import torch
7
import torch.nn.functional as F
@@ -37,9 +31,6 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
31
from vllm.utils.torch_utils import direct_register_custom_op
32
from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata
33
40
-if TYPE_CHECKING:
41
- from vllm.attention.backends.abstract import AttentionBackend
42
-
34
35
class MiniMaxText01RMSNormTP(CustomOp):
36
name = "MiniMaxText01RMSNormTP"
@@ -123,11 +114,6 @@ class MiniMaxText01LinearAttention(nn.Module, MambaBase):
114
def mamba_type(self) -> str:
115
return "linear_attention"
116
126
- def get_attn_backend(self) -> type["AttentionBackend"]:
127
- from vllm.v1.attention.backends.linear_attn import LinearAttentionBackend
128
-
129
- return LinearAttentionBackend
130
-
117
def get_state_dtype(self) -> tuple[torch.dtype]:
118
assert self.model_config is not None
119
assert self.cache_config is not None
vllm/model_executor/layers/mamba/mamba_mixer.py
+1
-9
@@ -1,10 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from typing import TYPE_CHECKING, NamedTuple
5
-
6
-if TYPE_CHECKING:
7
- from vllm.attention.backends.abstract import AttentionBackend
4
+from typing import NamedTuple
5
6
import torch
7
from torch import nn
@@ -452,11 +449,6 @@ class MambaMixer(MambaBase, CustomOp):
449
def mamba_type(self) -> str:
450
return "mamba1"
451
455
- def get_attn_backend(self) -> type["AttentionBackend"]:
456
- from vllm.v1.attention.backends.mamba1_attn import Mamba1AttentionBackend
457
-
458
- return Mamba1AttentionBackend
459
-
452
def _time_proj_bias(self) -> torch.Tensor | None:
453
if hasattr(self.dt_proj, "bias") and self.dt_proj.bias is not None:
454
return self.dt_proj.bias.float()
vllm/model_executor/layers/mamba/mamba_mixer2.py
-9
@@ -1,10 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from typing import TYPE_CHECKING
5
-
6
-if TYPE_CHECKING:
7
- from vllm.attention.backends.abstract import AttentionBackend
4
5
import torch
6
from torch import nn
@@ -908,11 +904,6 @@ class MambaMixer2(MambaBase, CustomOp):
904
def mamba_type(self) -> str:
905
return "mamba2"
906
911
- def get_attn_backend(self) -> type["AttentionBackend"]:
912
- from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionBackend
913
-
914
- return Mamba2AttentionBackend
915
-
907
908
def mamba_mixer2(
909
projected_states: torch.Tensor,
vllm/model_executor/layers/mamba/short_conv.py
-9
@@ -1,10 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from typing import TYPE_CHECKING
5
-
6
-if TYPE_CHECKING:
7
- from vllm.attention.backends.abstract import AttentionBackend
4
5
import torch
6
@@ -232,11 +228,6 @@ class ShortConv(MambaBase, CustomOp):
228
def mamba_type(self) -> str:
229
return "short_conv"
230
235
- def get_attn_backend(self) -> type["AttentionBackend"]:
236
- from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionBackend
237
-
238
- return ShortConvAttentionBackend
239
-
231
232
def short_conv(
233
hidden_states: torch.Tensor,
vllm/model_executor/models/plamo2.py
-9
@@ -4,10 +4,6 @@
4
5
from collections.abc import Iterable
6
from itertools import islice
7
-from typing import TYPE_CHECKING
8
-
9
-if TYPE_CHECKING:
10
- from vllm.attention.backends.abstract import AttentionBackend
7
8
import torch
9
from torch import nn
@@ -467,11 +463,6 @@ class Plamo2MambaMixer(MambaBase, CustomOp):
463
def mamba_type(self) -> str:
464
return "mamba2"
465
470
- def get_attn_backend(self) -> type["AttentionBackend"]:
471
- from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionBackend
472
-
473
- return Mamba2AttentionBackend
474
-
466
467
def plamo2_mamba_mixer(
468
hidden_states: torch.Tensor,
vllm/model_executor/models/qwen3_next.py
+2
-7
@@ -10,7 +10,7 @@ from einops import rearrange
10
from torch import nn
11
from transformers.activations import ACT2FN
12
13
-from vllm.attention import Attention, AttentionBackend, AttentionMetadata
13
+from vllm.attention import Attention, AttentionMetadata
14
from vllm.compilation.decorators import support_torch_compile
15
from vllm.config import (
16
CacheConfig,
@@ -216,12 +216,7 @@ class Qwen3NextSparseMoeBlock(nn.Module):
216
class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
217
@property
218
def mamba_type(self) -> str:
219
- return "linear_attention"
220
-
221
- def get_attn_backend(self) -> type["AttentionBackend"]:
222
- from vllm.v1.attention.backends.gdn_attn import GDNAttentionBackend
223
-
224
- return GDNAttentionBackend
219
+ return "gdn_attention"
220
221
def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:
222
return MambaStateDtypeCalculator.gated_delta_net_state_dtype(