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 14, 23:45 UTC
[v1][mamba] Added mamba_type into MambaSpec (#21715)
Signed-off-by: asafg <asafg@ai21.com> Co-authored-by: asafg <asafg@ai21.com>
Asaf Joseph Gardin committed
Jul 28, 2025 at 11:15 UTC
a6c050286a1a59ba35d133c4aa048a51c0b170d9
6 files changed
+52
-4
tests/v1/attention/test_mamba_selectors.py
new
+25
@@ -0,0 +1,25 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for mamba attention backend selectors."""
4
+
5
+import pytest
6
+
7
+from vllm.v1.attention.backends.mamba_attn import Mamba2AttentionBackend
8
+from vllm.v1.attention.backends.mamba_selectors import get_mamba_attn_backend
9
+
10
+
11
+@pytest.mark.parametrize(argnames=["mamba_type", "expected_backend"],
12
+ argvalues=[("mamba2", Mamba2AttentionBackend)])
13
+def test_get_mamba_attn_backend_mamba2(mamba_type, expected_backend):
14
+ backend_class = get_mamba_attn_backend(mamba_type)
15
+
16
+ assert backend_class is expected_backend
17
+
18
+
19
+def test_get_mamba_attn_backend_unsupported():
20
+ unsupported_types = ["mamba", ""]
21
+
22
+ for mamba_type in unsupported_types:
23
+ err_message = f"Mamba Attention type {mamba_type} is not supported yet."
24
+ with pytest.raises(NotImplementedError, match=err_message):
25
+ get_mamba_attn_backend(mamba_type)
vllm/model_executor/layers/mamba/abstract.py
+5
@@ -27,3 +27,8 @@ class MambaBase(ABC):
27
In this case, returns (conv_state_shape, ssm_state_shape).
28
"""
29
pass
30
+
31
+ @property
32
+ @abstractmethod
33
+ def mamba_type(self) -> str:
34
+ pass
vllm/model_executor/layers/mamba/mamba_mixer2.py
+4
@@ -732,6 +732,10 @@ class MambaMixer2(MambaBase, CustomOp):
732
conv_kernel=self.conv_kernel_size,
733
)
734
735
+ @property
736
+ def mamba_type(self) -> str:
737
+ return "mamba2"
738
+
739
740
def mamba_mixer2(
741
hidden_states: torch.Tensor,
vllm/v1/attention/backends/mamba_selectors.py
new
+12
@@ -0,0 +1,12 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from vllm.attention.backends.abstract import AttentionBackend
4
+from vllm.v1.attention.backends.mamba_attn import Mamba2AttentionBackend
5
+
6
+
7
+def get_mamba_attn_backend(mamba_type: str) -> type[AttentionBackend]:
8
+ if mamba_type == "mamba2":
9
+ return Mamba2AttentionBackend
10
+
11
+ raise NotImplementedError(f"Mamba Attention type {mamba_type} is not "
12
+ "supported yet.")
vllm/v1/kv_cache_interface.py
+2
-1
@@ -200,13 +200,14 @@ class MambaSpec(KVCacheSpec):
200
shapes: tuple[tuple[int, ...], ...]
201
dtype: torch.dtype
202
page_size_padded: Optional[int] = None
203
+ mamba_type: str = "mamba2"
204
205
def __post_init__(self):
206
self.num_elements = sum(prod(shape) for shape in self.shapes)
207
208
@property
209
def type_id(self) -> str:
209
- return f"mamba_{self.shapes}_{self.dtype}"
210
+ return f"mamba_{self.shapes}_{self.dtype}_{self.mamba_type}"
211
212
@property
213
def page_size_bytes(self) -> int:
vllm/v1/worker/gpu_model_runner.py
+4
-3
@@ -44,7 +44,7 @@ from vllm.tasks import GenerationTask, PoolingTask, SupportedTask
44
from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
45
GiB_bytes, LazyLoader, check_use_alibi, get_dtype_size,
46
is_pin_memory_available, round_up)
47
-from vllm.v1.attention.backends.mamba_attn import Mamba2AttentionBackend
47
+from vllm.v1.attention.backends.mamba_selectors import get_mamba_attn_backend
48
from vllm.v1.attention.backends.utils import (
49
AttentionMetadataBuilder, CommonAttentionMetadata,
50
make_local_attention_virtual_batches)
@@ -2539,7 +2539,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2539
"Non-Attention backend is not supported by V1 "
2540
"GPUModelRunner.")
2541
elif isinstance(kv_cache_spec, MambaSpec):
2542
- attn_backend_i = Mamba2AttentionBackend
2542
+ attn_backend_i = get_mamba_attn_backend(kv_cache_spec.mamba_type)
2543
else:
2544
raise ValueError(
2545
f"Unknown KV cache spec type: {type(kv_cache_spec)}")
@@ -2919,7 +2919,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2919
shapes=mamba_module.get_state_shape(),
2920
dtype=self.kv_cache_dtype,
2921
block_size=max_model_len,
2922
- page_size_padded=page_size_padded)
2922
+ page_size_padded=page_size_padded,
2923
+ mamba_type=mamba_module.mamba_type)
2924
2925
return kv_cache_spec
2926