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, 05:30 UTC
[Attention] Unify mamba and attention backend selection (#23171)
Signed-off-by: Ayush Satyam <ayushsatyam146@gmail.com>
Ayush Satyam committed
Aug 25, 2025 at 14:39 UTC
5c4b6e66fec0a59f702e32ec5c476cdd8a800566
11 files changed
+186
-72
tests/v1/attention/test_attention_backends_selection.py
new
+104
@@ -0,0 +1,104 @@
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
+from types import SimpleNamespace
6
+
7
+import pytest
8
+
9
+from vllm.model_executor.layers.mamba.mamba_mixer import MambaMixer
10
+from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
11
+from vllm.model_executor.layers.mamba.short_conv import ShortConv
12
+from vllm.model_executor.models.minimax_text_01 import (
13
+ MiniMaxText01LinearAttention)
14
+from vllm.v1.attention.backends.linear_attn import LinearAttentionBackend
15
+from vllm.v1.attention.backends.mamba1_attn import Mamba1AttentionBackend
16
+from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionBackend
17
+from vllm.v1.attention.backends.short_conv_attn import (
18
+ ShortConvAttentionBackend)
19
+
20
+
21
+@pytest.mark.parametrize(
22
+ "layer_class, init_kwargs, expected_backend, expected_mamba_type", [
23
+ (
24
+ MambaMixer,
25
+ dict(
26
+ hidden_size=128,
27
+ ssm_state_size=16,
28
+ conv_kernel_size=4,
29
+ intermediate_size=256,
30
+ time_step_rank=8,
31
+ use_conv_bias=True,
32
+ use_bias=False,
33
+ use_rms_norm=True,
34
+ ),
35
+ Mamba1AttentionBackend,
36
+ "mamba1",
37
+ ),
38
+ (
39
+ MambaMixer2,
40
+ dict(
41
+ hidden_size=128,
42
+ ssm_state_size=16,
43
+ conv_kernel_size=4,
44
+ intermediate_size=256,
45
+ use_conv_bias=True,
46
+ use_bias=False,
47
+ n_groups=1,
48
+ num_heads=8,
49
+ head_dim=32,
50
+ ),
51
+ Mamba2AttentionBackend,
52
+ "mamba2",
53
+ ),
54
+ (
55
+ MiniMaxText01LinearAttention,
56
+ dict(
57
+ hidden_size=128,
58
+ hidden_inner_size=256,
59
+ num_heads=8,
60
+ head_dim=32,
61
+ max_position=2048,
62
+ block_size=64,
63
+ num_hidden_layer=12,
64
+ layer_idx=0,
65
+ linear_layer_idx=0,
66
+ ),
67
+ LinearAttentionBackend,
68
+ "linear_attention",
69
+ ),
70
+ (
71
+ ShortConv,
72
+ dict(
73
+ config=SimpleNamespace(conv_L_cache=32, conv_bias=True),
74
+ dim=128,
75
+ layer_idx=0,
76
+ ),
77
+ ShortConvAttentionBackend,
78
+ "short_conv",
79
+ ),
80
+ ])
81
+def test_mamba_layers_get_attn_backend(dist_init, layer_class, init_kwargs,
82
+ expected_backend, expected_mamba_type):
83
+ """Test that Mamba-like layers return the correct attention backend."""
84
+ layer = layer_class(**init_kwargs)
85
+
86
+ backend_class = layer.get_attn_backend()
87
+ assert backend_class is expected_backend
88
+ assert layer.mamba_type == expected_mamba_type
89
+
90
+
91
+@pytest.mark.parametrize("layer_class,expected_backend,expected_mamba_type", [
92
+ (MambaMixer, Mamba1AttentionBackend, "mamba1"),
93
+ (MambaMixer2, Mamba2AttentionBackend, "mamba2"),
94
+ (MiniMaxText01LinearAttention, LinearAttentionBackend, "linear_attention"),
95
+ (ShortConv, ShortConvAttentionBackend, "short_conv"),
96
+])
97
+def test_mamba_layers_have_unified_interface(layer_class, expected_backend,
98
+ expected_mamba_type):
99
+ """Test that all Mamba layers have the unified get_attn_backend
100
+ interface."""
101
+ assert hasattr(layer_class, 'get_attn_backend'), (
102
+ f"{layer_class.__name__} should have get_attn_backend method")
103
+ assert hasattr(layer_class, 'mamba_type'), (
104
+ f"{layer_class.__name__} should have mamba_type property")
tests/v1/attention/test_mamba_selectors.py
deleted
-25
@@ -1,25 +0,0 @@
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.mamba2_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/attention/layer.py
+2
-1
@@ -18,6 +18,7 @@ from vllm.distributed.kv_transfer import (get_kv_transfer_group,
18
is_v1_kv_transfer_group)
19
from vllm.forward_context import ForwardContext, get_forward_context
20
from vllm.logger import init_logger
21
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
22
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
23
from vllm.model_executor.layers.quantization.base_config import (
24
QuantizationConfig)
@@ -54,7 +55,7 @@ def check_xformers_availability():
55
return USE_XFORMERS_OPS
56
57
57
-class Attention(nn.Module):
58
+class Attention(nn.Module, AttentionLayerBase):
59
"""Attention layer.
60
61
This class takes query, key, and value tensors as input. The input tensors
vllm/model_executor/layers/attention_layer_base.py
new
+23
@@ -0,0 +1,23 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Base class for attention-like layers."""
4
+from abc import ABC, abstractmethod
5
+from typing import TYPE_CHECKING
6
+
7
+if TYPE_CHECKING:
8
+ from vllm.attention.backends.abstract import AttentionBackend
9
+
10
+
11
+class AttentionLayerBase(ABC):
12
+ """
13
+ Base class for attention-like layers (Attention, Mamba, etc.)
14
+ that support the v1 engine.
15
+
16
+ This provides a common interface for getting attention backends
17
+ from different layer types.
18
+ """
19
+
20
+ @abstractmethod
21
+ def get_attn_backend(self) -> type["AttentionBackend"]:
22
+ """Get the attention backend class for this layer."""
23
+ pass
vllm/model_executor/layers/mamba/abstract.py
+13
-2
@@ -1,12 +1,18 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from abc import ABC, abstractmethod
3
+from abc import abstractmethod
4
from collections.abc import Iterable
5
+from typing import TYPE_CHECKING
6
7
import torch
8
9
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
10
9
-class MambaBase(ABC):
11
+if TYPE_CHECKING:
12
+ from vllm.attention.backends.abstract import AttentionBackend
13
+
14
+
15
+class MambaBase(AttentionLayerBase):
16
"""
17
Base class for Mamba-like layers which support the v1 engine.
18
Inherit from this class if you implement a custom layer.
@@ -32,3 +38,8 @@ class MambaBase(ABC):
38
@abstractmethod
39
def mamba_type(self) -> str:
40
pass
41
+
42
+ @abstractmethod
43
+ def get_attn_backend(self) -> type["AttentionBackend"]:
44
+ """Get the attention backend class for this Mamba layer."""
45
+ pass
vllm/model_executor/layers/mamba/mamba_mixer.py
+9
-1
@@ -1,7 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from typing import NamedTuple, Optional
4
+from typing import TYPE_CHECKING, NamedTuple, Optional
5
+
6
+if TYPE_CHECKING:
7
+ from vllm.attention.backends.abstract import AttentionBackend
8
9
import torch
10
from torch import nn
@@ -404,6 +407,11 @@ class MambaMixer(MambaBase, CustomOp):
407
def mamba_type(self) -> str:
408
return "mamba1"
409
410
+ def get_attn_backend(self) -> type["AttentionBackend"]:
411
+ from vllm.v1.attention.backends.mamba1_attn import (
412
+ Mamba1AttentionBackend)
413
+ return Mamba1AttentionBackend
414
+
415
def _time_proj_bias(self) -> Optional[torch.Tensor]:
416
if hasattr(self.dt_proj, "bias") and self.dt_proj.bias is not None:
417
return self.dt_proj.bias.float()
vllm/model_executor/layers/mamba/mamba_mixer2.py
+9
-1
@@ -1,7 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from typing import Optional, Union
4
+from typing import TYPE_CHECKING, Optional, Union
5
+
6
+if TYPE_CHECKING:
7
+ from vllm.attention.backends.abstract import AttentionBackend
8
9
import torch
10
from torch import nn
@@ -758,6 +761,11 @@ class MambaMixer2(MambaBase, CustomOp):
761
def mamba_type(self) -> str:
762
return "mamba2"
763
764
+ def get_attn_backend(self) -> type["AttentionBackend"]:
765
+ from vllm.v1.attention.backends.mamba2_attn import (
766
+ Mamba2AttentionBackend)
767
+ return Mamba2AttentionBackend
768
+
769
770
def mamba_mixer2(
771
hidden_states: torch.Tensor,
vllm/model_executor/layers/mamba/short_conv.py
+9
-1
@@ -1,7 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from typing import Optional
4
+from typing import TYPE_CHECKING, Optional
5
+
6
+if TYPE_CHECKING:
7
+ from vllm.attention.backends.abstract import AttentionBackend
8
9
import torch
10
@@ -232,6 +235,11 @@ class ShortConv(MambaBase, CustomOp):
235
def mamba_type(self) -> str:
236
return "short_conv"
237
238
+ def get_attn_backend(self) -> type["AttentionBackend"]:
239
+ from vllm.v1.attention.backends.short_conv_attn import (
240
+ ShortConvAttentionBackend)
241
+ return ShortConvAttentionBackend
242
+
243
244
def short_conv(
245
hidden_states: torch.Tensor,
vllm/model_executor/models/minimax_text_01.py
+9
-1
@@ -4,7 +4,10 @@
4
import copy
5
import math
6
from collections.abc import Iterable
7
-from typing import Optional, Union
7
+from typing import TYPE_CHECKING, Optional, Union
8
+
9
+if TYPE_CHECKING:
10
+ from vllm.attention.backends.abstract import AttentionBackend
11
12
import regex as re
13
import torch
@@ -339,6 +342,11 @@ class MiniMaxText01LinearAttention(nn.Module, MambaBase):
342
def mamba_type(self) -> str:
343
return "linear_attention"
344
345
+ def get_attn_backend(self) -> type["AttentionBackend"]:
346
+ from vllm.v1.attention.backends.linear_attn import (
347
+ LinearAttentionBackend)
348
+ return LinearAttentionBackend
349
+
350
def get_state_dtype(self) -> tuple[torch.dtype]:
351
return MambaStateDtypeCalculator.linear_attention_state_dtype(
352
self.model_config.dtype,
vllm/v1/attention/backends/mamba_selectors.py
deleted
-22
@@ -1,22 +0,0 @@
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.linear_attn import LinearAttentionBackend
5
-from vllm.v1.attention.backends.mamba1_attn import Mamba1AttentionBackend
6
-from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionBackend
7
-from vllm.v1.attention.backends.short_conv_attn import (
8
- ShortConvAttentionBackend)
9
-
10
-
11
-def get_mamba_attn_backend(mamba_type: str) -> type[AttentionBackend]:
12
- if mamba_type == "mamba1":
13
- return Mamba1AttentionBackend
14
- if mamba_type == "mamba2":
15
- return Mamba2AttentionBackend
16
- if mamba_type == "linear_attention":
17
- return LinearAttentionBackend
18
- if mamba_type == "short_conv":
19
- return ShortConvAttentionBackend
20
-
21
- raise NotImplementedError(f"Mamba Attention type {mamba_type} is not "
22
- "supported yet.")
vllm/v1/worker/gpu_model_runner.py
+8
-18
@@ -35,7 +35,8 @@ from vllm.distributed.parallel_state import (
35
from vllm.forward_context import (BatchDescriptor, DPMetadata,
36
set_forward_context)
37
from vllm.logger import init_logger
38
-from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaBase
38
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
39
+from vllm.model_executor.layers.mamba.abstract import MambaBase
40
from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
41
from vllm.model_executor.model_loader import TensorizerLoader, get_model_loader
42
from vllm.model_executor.models.interfaces import (is_mixture_of_experts,
@@ -55,7 +56,6 @@ from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
56
GiB_bytes, LazyLoader, cdiv, check_use_alibi,
57
get_dtype_size, is_pin_memory_available, round_up,
58
supports_dynamo)
58
-from vllm.v1.attention.backends.mamba_selectors import get_mamba_attn_backend
59
from vllm.v1.attention.backends.utils import (
60
AttentionCGSupport, AttentionMetadataBuilder, CommonAttentionMetadata,
61
make_kv_sharing_fast_prefill_attention_metadata,
@@ -2747,11 +2747,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2747
"""
2748
assert len(self.attn_groups) == 0, \
2749
"Attention backends are already initialized"
2750
- attn_layers = get_layers_from_vllm_config(self.vllm_config, Attention)
2750
2751
def get_attn_backends_for_layers(
2752
layer_names: list[str]
2753
) -> dict[type[AttentionBackend], list[str]]:
2754
+ layers = get_layers_from_vllm_config(self.vllm_config,
2755
+ AttentionLayerBase,
2756
+ layer_names)
2757
attn_backends = {}
2758
attn_backend_layers = defaultdict(list)
2759
# Dedupe based on full class name; this is a bit safer than using
@@ -2760,7 +2762,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2762
# they are cached correctly, there will be different objects per
2763
# layer.
2764
for layer_name in layer_names:
2763
- attn_backend = attn_layers[layer_name].get_attn_backend()
2765
+ attn_backend = layers[layer_name].get_attn_backend()
2766
key = attn_backend.full_cls_name()
2767
attn_backends[key] = attn_backend
2768
attn_backend_layers[key].append(layer_name)
@@ -2789,20 +2791,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2791
2792
for kv_cache_group_spec in kv_cache_config.kv_cache_groups:
2793
kv_cache_spec = kv_cache_group_spec.kv_cache_spec
2792
- if isinstance(kv_cache_spec, AttentionSpec):
2793
- attn_backends = get_attn_backends_for_layers(
2794
- kv_cache_group_spec.layer_names)
2795
- # TODO(lucas): move `get_mamba_attn_backend` into the mamba
2796
- # layers like above
2797
- elif isinstance(kv_cache_spec, MambaSpec):
2798
- attn_backends = {
2799
- get_mamba_attn_backend(kv_cache_spec.mamba_type):
2800
- kv_cache_group_spec.layer_names
2801
- }
2802
- else:
2803
- raise ValueError(
2804
- f"Unknown KV cache spec type: {type(kv_cache_spec)}")
2805
-
2794
+ attn_backends = get_attn_backends_for_layers(
2795
+ kv_cache_group_spec.layer_names)
2796
self.attn_groups.append(
2797
create_attn_groups(attn_backends, kv_cache_spec))
2798