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 11, 03:45 UTC
[Bugfix] fix the alias bug of AttentionBackendEnum when register CUSTOM attention backend to vllm (#30869)
Signed-off-by: zejunchen-zejun <zejun.chen@amd.com>
zejunchen-zejun committed
Dec 20, 2025 at 09:03 UTC
d52c5096d7305abc7f266026cb042121ff5bccda
2 files changed
+173
-2
tests/test_attention_backend_registry.py
new
+169
@@ -0,0 +1,169 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from vllm.attention.backends.abstract import (
4
+ AttentionBackend,
5
+ AttentionImpl,
6
+)
7
+from vllm.attention.backends.registry import (
8
+ AttentionBackendEnum,
9
+ MambaAttentionBackendEnum,
10
+ register_backend,
11
+)
12
+
13
+
14
+class CustomAttentionImpl(AttentionImpl):
15
+ """Mock custom attention implementation for testing."""
16
+
17
+ def __init__(self, *args, **kwargs):
18
+ super().__init__()
19
+
20
+ def forward(self, *args, **kwargs):
21
+ """Mock forward pass."""
22
+ pass
23
+
24
+
25
+class CustomAttentionBackend(AttentionBackend):
26
+ """Mock custom attention backend for testing."""
27
+
28
+ @staticmethod
29
+ def get_name():
30
+ return "CUSTOM"
31
+
32
+ @staticmethod
33
+ def get_impl_cls():
34
+ return CustomAttentionImpl
35
+
36
+ @staticmethod
37
+ def get_builder_cls():
38
+ """Mock builder class."""
39
+ return None
40
+
41
+ @staticmethod
42
+ def get_required_kv_cache_layout():
43
+ """Mock KV cache layout."""
44
+ return None
45
+
46
+
47
+class CustomMambaAttentionImpl(AttentionImpl):
48
+ """Mock custom mamba attention implementation for testing."""
49
+
50
+ def __init__(self, *args, **kwargs):
51
+ super().__init__()
52
+
53
+ def forward(self, *args, **kwargs):
54
+ """Mock forward pass."""
55
+ pass
56
+
57
+
58
+class CustomMambaAttentionBackend(AttentionBackend):
59
+ """Mock custom mamba attention backend for testing."""
60
+
61
+ @staticmethod
62
+ def get_name():
63
+ return "CUSTOM_MAMBA"
64
+
65
+ @staticmethod
66
+ def get_impl_cls():
67
+ return CustomMambaAttentionImpl
68
+
69
+ @staticmethod
70
+ def get_builder_cls():
71
+ """Mock builder class."""
72
+ return None
73
+
74
+ @staticmethod
75
+ def get_required_kv_cache_layout():
76
+ """Mock KV cache layout."""
77
+ return None
78
+
79
+
80
+def test_custom_is_not_alias_of_any_backend():
81
+ # Get all members of AttentionBackendEnum
82
+ all_backends = list(AttentionBackendEnum)
83
+
84
+ # Find any aliases of CUSTOM
85
+ aliases = []
86
+ for backend in all_backends:
87
+ if backend.name != "CUSTOM" and backend is AttentionBackendEnum.CUSTOM:
88
+ aliases.append(backend.name)
89
+
90
+ # CUSTOM should not be an alias of any other backend
91
+ assert len(aliases) == 0, (
92
+ f"BUG! CUSTOM is an alias of: {', '.join(aliases)}!\n"
93
+ f"CUSTOM.value = {repr(AttentionBackendEnum.CUSTOM.value)}\n"
94
+ f"This happens when CUSTOM has the same value as another backend.\n"
95
+ f"When you register to CUSTOM, you're actually registering to {aliases[0]}!\n"
96
+ f"All backend values:\n"
97
+ + "\n".join(f" {b.name}: {repr(b.value)}" for b in all_backends)
98
+ )
99
+
100
+ # Verify CUSTOM has its own unique identity
101
+ assert AttentionBackendEnum.CUSTOM.name == "CUSTOM", (
102
+ f"CUSTOM.name should be 'CUSTOM', but got '{AttentionBackendEnum.CUSTOM.name}'"
103
+ )
104
+
105
+
106
+def test_register_custom_backend_with_class_path():
107
+ # Register with explicit class path
108
+ register_backend(
109
+ backend=AttentionBackendEnum.CUSTOM,
110
+ class_path="tests.test_attention_backend_registry.CustomAttentionBackend",
111
+ is_mamba=False,
112
+ )
113
+
114
+ # Check that CUSTOM backend is registered
115
+ assert AttentionBackendEnum.CUSTOM.is_overridden(), (
116
+ "CUSTOM should be overridden after registration"
117
+ )
118
+
119
+ # Get the registered class path
120
+ class_path = AttentionBackendEnum.CUSTOM.get_path()
121
+ assert class_path == "tests.test_attention_backend_registry.CustomAttentionBackend"
122
+
123
+ # Get the backend class
124
+ backend_cls = AttentionBackendEnum.CUSTOM.get_class()
125
+ assert backend_cls.get_name() == "CUSTOM"
126
+ assert backend_cls.get_impl_cls() == CustomAttentionImpl
127
+
128
+
129
+def test_mamba_custom_is_not_alias_of_any_backend():
130
+ # Get all mamba backends
131
+ all_backends = list(MambaAttentionBackendEnum)
132
+
133
+ # Find any aliases of CUSTOM
134
+ aliases = []
135
+ for backend in all_backends:
136
+ if backend.name != "CUSTOM" and backend is MambaAttentionBackendEnum.CUSTOM:
137
+ aliases.append(backend.name)
138
+
139
+ # CUSTOM should not be an alias of any other backend
140
+ assert len(aliases) == 0, (
141
+ f"BUG! MambaAttentionBackendEnum.CUSTOM is an alias of: {', '.join(aliases)}!\n"
142
+ f"CUSTOM.value = {repr(MambaAttentionBackendEnum.CUSTOM.value)}\n"
143
+ f"All mamba backend values:\n"
144
+ + "\n".join(f" {b.name}: {repr(b.value)}" for b in all_backends)
145
+ )
146
+
147
+
148
+def test_register_custom_mamba_backend_with_class_path():
149
+ # Register with explicit class path
150
+ register_backend(
151
+ backend=MambaAttentionBackendEnum.CUSTOM,
152
+ class_path="tests.test_attention_backend_registry.CustomMambaAttentionBackend",
153
+ is_mamba=True,
154
+ )
155
+
156
+ # Check that the backend is registered
157
+ assert MambaAttentionBackendEnum.CUSTOM.is_overridden()
158
+
159
+ # Get the registered class path
160
+ class_path = MambaAttentionBackendEnum.CUSTOM.get_path()
161
+ assert (
162
+ class_path
163
+ == "tests.test_attention_backend_registry.CustomMambaAttentionBackend"
164
+ )
165
+
166
+ # Get the backend class
167
+ backend_cls = MambaAttentionBackendEnum.CUSTOM.get_class()
168
+ assert backend_cls.get_name() == "CUSTOM_MAMBA"
169
+ assert backend_cls.get_impl_cls() == CustomMambaAttentionImpl
vllm/attention/backends/registry.py
+4
-2
@@ -77,7 +77,8 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
77
)
78
CPU_ATTN = "vllm.v1.attention.backends.cpu_attn.CPUAttentionBackend"
79
# Placeholder for third-party/custom backends - must be registered before use
80
- CUSTOM = ""
80
+ # set to None to avoid alias with other backend, whose value is an empty string
81
+ CUSTOM = None
82
83
def get_path(self, include_classname: bool = True) -> str:
84
"""Get the class path for this backend (respects overrides).
@@ -139,7 +140,8 @@ class MambaAttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
140
LINEAR = "vllm.v1.attention.backends.linear_attn.LinearAttentionBackend"
141
GDN_ATTN = "vllm.v1.attention.backends.gdn_attn.GDNAttentionBackend"
142
# Placeholder for third-party/custom backends - must be registered before use
142
- CUSTOM = ""
143
+ # set to None to avoid alias with other backend, whose value is an empty string
144
+ CUSTOM = None
145
146
def get_path(self, include_classname: bool = True) -> str:
147
"""Get the class path for this backend (respects overrides).