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
Add FLASHINFER_MLA to backend selector test (#24753)
Signed-off-by: Matthew Bonanni <mbonanni001@gmail.com>
Matthew Bonanni committed
Sep 12, 2025 at 18:30 UTC
5fe643fc265633b6f38752321b7883899a6a8903
2 files changed
+43
-19
tests/kernels/attention/test_attention_selector.py
+41
-19
@@ -22,7 +22,10 @@ def clear_cache():
22
23
# Define MLA and non-MLA backends separately
24
DEVICE_MLA_BACKENDS = {
25
- "cuda": ["TRITON_MLA", "FLASHMLA", "FLASH_ATTN_MLA", "CUTLASS_MLA"],
25
+ "cuda": [
26
+ "TRITON_MLA", "FLASHMLA", "FLASHINFER_MLA", "FLASH_ATTN_MLA",
27
+ "CUTLASS_MLA"
28
+ ],
29
"hip": ["TRITON_MLA", "ROCM_AITER_MLA"],
30
"cpu": [],
31
}
@@ -90,8 +93,8 @@ def test_env(
93
94
with patch("vllm.attention.selector.current_platform",
95
CpuPlatform()):
93
- backend = get_attn_backend(16, torch.float16, torch.float16,
94
- block_size, False)
96
+ backend = get_attn_backend(16, torch.float16, None, block_size,
97
+ False)
98
assert backend.get_name() == "TORCH_SDPA_VLLM_V1"
99
100
elif device == "hip":
@@ -109,7 +112,7 @@ def test_env(
112
with pytest.raises(ValueError) as exc_info:
113
get_attn_backend(16,
114
torch.float16,
112
- torch.float16,
115
+ None,
116
block_size,
117
False,
118
use_mla=use_mla)
@@ -120,7 +123,7 @@ def test_env(
123
with pytest.raises(ValueError) as exc_info:
124
get_attn_backend(16,
125
torch.float16,
123
- torch.float16,
126
+ None,
127
block_size,
128
False,
129
use_mla=use_mla)
@@ -130,7 +133,7 @@ def test_env(
133
# Valid backend-block_size combination
134
backend = get_attn_backend(16,
135
torch.float16,
133
- torch.float16,
136
+ None,
137
block_size,
138
False,
139
use_mla=use_mla)
@@ -139,7 +142,7 @@ def test_env(
142
else:
143
backend = get_attn_backend(16,
144
torch.float16,
142
- torch.float16,
145
+ None,
146
block_size,
147
False,
148
use_mla=use_mla)
@@ -153,6 +156,8 @@ def test_env(
156
# CUDA MLA backend logic:
157
# - CUTLASS_MLA: only supported with block_size == 128
158
# and Blackwell GPUs (SM 10.0), V1 only
159
+ # - FLASHINFER_MLA: only supported on Blackwell GPUs
160
+ # (SM 10.0+), V1 only
161
# - FLASHMLA: only supported with block_size == 64
162
# - FLASH_ATTN_MLA: V1 only
163
# - TRITON_MLA: fallback for other cases
@@ -169,12 +174,31 @@ def test_env(
174
else:
175
backend = get_attn_backend(16,
176
torch.float16,
172
- torch.float16,
177
+ None,
178
block_size,
179
False,
180
use_mla=use_mla)
181
expected = "CUTLASS_MLA_VLLM_V1"
182
assert backend.get_name() == expected
183
+ elif name == "FLASHINFER_MLA":
184
+ if not use_v1:
185
+ # FlashInfer MLA only supported on V1 engine
186
+ pytest.skip(
187
+ "FlashInfer MLA only supported on V1 engine")
188
+ elif block_size not in [32, 64]:
189
+ # FlashInfer MLA only supports block_size 32 or 64
190
+ pytest.skip(
191
+ "FlashInfer MLA only supports block_size 32 "
192
+ "or 64")
193
+ else:
194
+ backend = get_attn_backend(16,
195
+ torch.float16,
196
+ None,
197
+ block_size,
198
+ False,
199
+ use_mla=use_mla)
200
+ expected = "FLASHINFER_MLA"
201
+ assert backend.get_name() == expected
202
elif name == "FLASHMLA":
203
if block_size != 64:
204
# FlashMLA only supports block_size == 64
@@ -189,7 +213,7 @@ def test_env(
213
else:
214
backend = get_attn_backend(16,
215
torch.float16,
192
- torch.float16,
216
+ None,
217
block_size,
218
False,
219
use_mla=use_mla)
@@ -204,7 +228,7 @@ def test_env(
228
else:
229
backend = get_attn_backend(16,
230
torch.float16,
207
- torch.float16,
231
+ None,
232
block_size,
233
False,
234
use_mla=use_mla)
@@ -214,7 +238,7 @@ def test_env(
238
# TRITON_MLA or other fallback
239
backend = get_attn_backend(16,
240
torch.float16,
217
- torch.float16,
241
+ None,
242
block_size,
243
False,
244
use_mla=use_mla)
@@ -224,7 +248,7 @@ def test_env(
248
elif name == "FLASHINFER":
249
backend = get_attn_backend(16,
250
torch.float16,
227
- torch.float16,
251
+ None,
252
block_size,
253
False,
254
use_mla=use_mla)
@@ -233,7 +257,7 @@ def test_env(
257
else:
258
backend = get_attn_backend(32,
259
torch.float16,
236
- torch.float16,
260
+ None,
261
block_size,
262
False,
263
use_mla=use_mla)
@@ -243,7 +267,7 @@ def test_env(
267
if use_v1:
268
backend = get_attn_backend(16,
269
torch.float16,
246
- torch.float16,
270
+ None,
271
block_size,
272
False,
273
use_mla=use_mla)
@@ -269,15 +293,13 @@ def test_fp32_fallback(
293
294
with patch("vllm.attention.selector.current_platform",
295
CpuPlatform()):
272
- backend = get_attn_backend(16, torch.float32, torch.float32,
273
- 16, False)
296
+ backend = get_attn_backend(16, torch.float32, None, 16, False)
297
assert backend.get_name() == "TORCH_SDPA_VLLM_V1"
298
299
elif device == "cuda":
300
with patch("vllm.attention.selector.current_platform",
301
CudaPlatform()):
279
- backend = get_attn_backend(16, torch.float32, torch.float32,
280
- 16, False)
302
+ backend = get_attn_backend(16, torch.float32, None, 16, False)
303
assert (backend.get_name() == "FLEX_ATTENTION"
304
if use_v1 else "XFORMERS")
305
@@ -331,7 +353,7 @@ def test_flash_attn(monkeypatch: pytest.MonkeyPatch):
353
assert backend.get_name() != STR_FLASH_ATTN_VAL
354
355
# Attention-free models should bypass env and use PlaceholderAttention
334
- backend = get_attn_backend(16, torch.float16, torch.float16, 16, True)
356
+ backend = get_attn_backend(16, torch.float16, None, 16, True)
357
assert backend.get_name() != STR_FLASH_ATTN_VAL
358
359
tests/v1/attention/utils.py
+2
@@ -141,6 +141,8 @@ def get_attention_backend(backend_name: _Backend):
141
"vllm.v1.attention.backends.mla.flashmla.FlashMLABackend",
142
_Backend.FLASH_ATTN_MLA:
143
"vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend",
144
+ _Backend.FLASHINFER_MLA:
145
+ "vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend",
146
_Backend.TRITON_MLA_VLLM_V1:
147
"vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend",
148
}