[ROCm][Perf][MLA] Add AITER FlashAttention MLA prefill backend (`ROCM_AITER_FA`) (#45033)
Signed-off-by: Xavier Aguilar <xavier.aguilarfruto@amd.com> Signed-off-by: Xavier Aguilar <Xavier.AguilarFruto@amd.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
xaguilar-amd committed
Jun 28, 2026 at 16:52 UTC
5ecae3266cd5e6b814e14368d615957ffe85fdef
5 files changed
+269
tests/v1/attention/test_mla_prefill_registry.py
+17
@@ -133,3 +133,20 @@ def test_clear_override():
133
def test_unknown_backend_name_raises():
134
with pytest.raises(ValueError, match="Unknown MLA prefill backend"):
135
MLAPrefillBackendEnum["NONEXISTENT"]
136
+
137
+
138
+def test_rocm_aiter_fa_registered():
139
+ """ROCM_AITER_FA is a known backend pointing at the AITER FA class."""
140
+ assert "ROCM_AITER_FA" in MLAPrefillBackendEnum.__members__
141
+
142
+ path = MLAPrefillBackendEnum.ROCM_AITER_FA.get_path()
143
+ assert path == (
144
+ "vllm.v1.attention.backends.mla.prefill.aiter_flash_attn."
145
+ "AiterFlashAttnPrefillBackend"
146
+ )
147
+
148
+ backend_cls = MLAPrefillBackendEnum.ROCM_AITER_FA.get_class()
149
+ assert backend_cls.get_name() == "ROCM_AITER_FA"
150
+ # The AITER FA path is the fp16/bf16 generic-varlen prefill path.
151
+ assert backend_cls.supports_dtype(torch.bfloat16)
152
+ assert backend_cls.supports_dtype(torch.float16)
tests/v1/attention/test_mla_prefill_selector.py
+119
@@ -14,6 +14,7 @@ from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnu
14
from vllm.v1.attention.backends.mla.prefill.selector import (
15
MLAPrefillSelectorConfig,
16
_auto_select_mla_prefill_backend,
17
+ _get_mla_prefill_backend_priorities,
18
get_mla_prefill_backend,
19
)
20
@@ -166,6 +167,7 @@ class TestAutoSelectMLAPrefillBackend:
167
return
168
169
with (
170
+ patch("vllm.platforms.current_platform") as mock_platform,
171
patch.object(
172
MLAPrefillBackendEnum.FLASH_ATTN,
173
"get_class",
@@ -173,6 +175,8 @@ class TestAutoSelectMLAPrefillBackend:
175
),
176
patch.object(trtllm_cls, "validate_configuration", return_value=[]),
177
):
178
+ # Force the non-ROCm priority on the Blackwell.
179
+ mock_platform.is_rocm.return_value = False
180
backend = _auto_select_mla_prefill_backend(
181
capability,
182
selector_config,
@@ -272,6 +276,121 @@ class TestBackendValidation:
276
assert invalid_reasons == []
277
278
279
+class TestROCmAiterFAPrefillSelection:
280
+ """Tests for the ROCm AITER FlashAttention MLA prefill backend."""
281
+
282
+ def test_rocm_priorities_prefer_aiter_fa(self):
283
+ """On ROCm, ROCM_AITER_FA is tried first, FLASH_ATTN as fallback."""
284
+ with patch("vllm.platforms.current_platform") as mock_platform:
285
+ mock_platform.is_rocm.return_value = True
286
+ priorities = _get_mla_prefill_backend_priorities(
287
+ DeviceCapability(major=9, minor=5)
288
+ )
289
+
290
+ assert priorities == [
291
+ MLAPrefillBackendEnum.ROCM_AITER_FA,
292
+ MLAPrefillBackendEnum.FLASH_ATTN,
293
+ ]
294
+
295
+ def test_supported_dtypes_are_fp16_bf16_only(self):
296
+ from vllm.v1.attention.backends.mla.prefill.aiter_flash_attn import (
297
+ AiterFlashAttnPrefillBackend,
298
+ )
299
+
300
+ assert AiterFlashAttnPrefillBackend.supports_dtype(torch.bfloat16)
301
+ assert AiterFlashAttnPrefillBackend.supports_dtype(torch.float16)
302
+ # FP8 is served by the separate AITER ASM backend, not this one.
303
+ assert not AiterFlashAttnPrefillBackend.supports_dtype(torch.float8_e4m3fn)
304
+
305
+ def test_supports_compute_capability_on_rocm(self):
306
+ from vllm.v1.attention.backends.mla.prefill import aiter_flash_attn as mod
307
+
308
+ # Gating is decided by on_mi3xx(), not by capability
309
+ capability = MagicMock()
310
+
311
+ with patch.object(mod.current_platform, "is_rocm", return_value=False):
312
+ assert not mod.AiterFlashAttnPrefillBackend.supports_compute_capability(
313
+ capability
314
+ )
315
+
316
+ with (
317
+ patch.object(mod.current_platform, "is_rocm", return_value=True),
318
+ patch("vllm.platforms.rocm.on_mi3xx", return_value=False),
319
+ ):
320
+ assert not mod.AiterFlashAttnPrefillBackend.supports_compute_capability(
321
+ capability
322
+ )
323
+
324
+ with (
325
+ patch.object(mod.current_platform, "is_rocm", return_value=True),
326
+ patch("vllm.platforms.rocm.on_mi3xx", return_value=True),
327
+ ):
328
+ assert mod.AiterFlashAttnPrefillBackend.supports_compute_capability(
329
+ capability
330
+ )
331
+
332
+ def test_is_available_delegates_to_rocm_aiter_ops(self):
333
+ from vllm._aiter_ops import rocm_aiter_ops
334
+ from vllm.v1.attention.backends.mla.prefill import aiter_flash_attn as mod
335
+
336
+ with patch.object(rocm_aiter_ops, "is_enabled", return_value=False):
337
+ assert not mod.AiterFlashAttnPrefillBackend.is_available()
338
+
339
+ with patch.object(rocm_aiter_ops, "is_enabled", return_value=True):
340
+ assert mod.AiterFlashAttnPrefillBackend.is_available()
341
+
342
+ def test_auto_select_prefers_aiter_fa_on_rocm(self):
343
+ from vllm.v1.attention.backends.mla.prefill.aiter_flash_attn import (
344
+ AiterFlashAttnPrefillBackend,
345
+ )
346
+
347
+ # gfx gating is simulated via the mocked validate_configuration,
348
+ # not the capability.
349
+ capability = MagicMock()
350
+ selector_config = MLAPrefillSelectorConfig(dtype=torch.bfloat16)
351
+
352
+ with (
353
+ patch("vllm.platforms.current_platform") as mock_platform,
354
+ patch.object(
355
+ AiterFlashAttnPrefillBackend,
356
+ "validate_configuration",
357
+ return_value=[],
358
+ ),
359
+ ):
360
+ mock_platform.is_rocm.return_value = True
361
+ backend = _auto_select_mla_prefill_backend(capability, selector_config)
362
+ assert backend.get_name() == "ROCM_AITER_FA"
363
+
364
+ def test_auto_select_falls_back_to_flash_attn_when_aiter_invalid(self):
365
+ from vllm.v1.attention.backends.mla.prefill.aiter_flash_attn import (
366
+ AiterFlashAttnPrefillBackend,
367
+ )
368
+
369
+ try:
370
+ flash_attn_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class()
371
+ except ImportError:
372
+ pytest.skip("FLASH_ATTN backend not available")
373
+ return
374
+
375
+ # the fallback is forced by the mocked validate_configuration,
376
+ # not the capability.
377
+ capability = MagicMock()
378
+ selector_config = MLAPrefillSelectorConfig(dtype=torch.bfloat16)
379
+
380
+ with (
381
+ patch("vllm.platforms.current_platform") as mock_platform,
382
+ patch.object(
383
+ AiterFlashAttnPrefillBackend,
384
+ "validate_configuration",
385
+ return_value=["compute capability not supported"],
386
+ ),
387
+ patch.object(flash_attn_cls, "validate_configuration", return_value=[]),
388
+ ):
389
+ mock_platform.is_rocm.return_value = True
390
+ backend = _auto_select_mla_prefill_backend(capability, selector_config)
391
+ assert backend.get_name() == "FLASH_ATTN"
392
+
393
+
394
class TestMLAPrefillBackendParsing:
395
"""Tests for string-based mla_prefill_backend parsing from CLI args."""
396
vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py
new
+121
@@ -0,0 +1,121 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""AITER FlashAttention backend for MLA prefill (ROCm).
4
+
5
+This backend calls ``aiter.flash_attn_varlen_func`` directly, which natively
6
+supports different q/k and v head dims (qk headdim 192, v headdim 128) without
7
+padding V, and dispatches to the fast ``aiter::fmha_fwd_`` kernel on
8
+gfx942/gfx950 (fp16/bf16).
9
+"""
10
+
11
+from typing import TYPE_CHECKING
12
+
13
+import torch
14
+
15
+from vllm.platforms import current_platform
16
+from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
17
+
18
+if TYPE_CHECKING:
19
+ from vllm.config import VllmConfig
20
+ from vllm.platforms.interface import DeviceCapability
21
+
22
+
23
+class AiterFlashAttnPrefillBackend(MLAPrefillBackend):
24
+ """AITER FlashAttention backend for MLA prefill"""
25
+
26
+ @staticmethod
27
+ def get_name() -> str:
28
+ return "ROCM_AITER_FA"
29
+
30
+ @classmethod
31
+ def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool:
32
+ if not current_platform.is_rocm():
33
+ return False
34
+ from vllm.platforms.rocm import on_mi3xx
35
+
36
+ return on_mi3xx()
37
+
38
+ @classmethod
39
+ def is_available(cls) -> bool:
40
+ from vllm._aiter_ops import rocm_aiter_ops
41
+
42
+ return rocm_aiter_ops.is_enabled()
43
+
44
+ def __init__(
45
+ self,
46
+ num_heads: int,
47
+ scale: float,
48
+ kv_lora_rank: int,
49
+ qk_nope_head_dim: int,
50
+ qk_rope_head_dim: int,
51
+ v_head_dim: int,
52
+ vllm_config: "VllmConfig",
53
+ ) -> None:
54
+ super().__init__(
55
+ num_heads=num_heads,
56
+ scale=scale,
57
+ kv_lora_rank=kv_lora_rank,
58
+ qk_nope_head_dim=qk_nope_head_dim,
59
+ qk_rope_head_dim=qk_rope_head_dim,
60
+ v_head_dim=v_head_dim,
61
+ vllm_config=vllm_config,
62
+ )
63
+
64
+ from aiter import flash_attn_varlen_func
65
+
66
+ self.flash_attn_varlen_func = flash_attn_varlen_func
67
+
68
+ def run_prefill_new_tokens(
69
+ self,
70
+ q: torch.Tensor,
71
+ k: torch.Tensor,
72
+ v: torch.Tensor,
73
+ return_softmax_lse: bool,
74
+ out: torch.Tensor | None = None,
75
+ output_scale: torch.Tensor | None = None,
76
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
77
+ assert output_scale is None, (
78
+ "AiterFlashAttnPrefillBackend does not support fused quantized output."
79
+ )
80
+ result = self.flash_attn_varlen_func(
81
+ q=q,
82
+ k=k,
83
+ v=v,
84
+ cu_seqlens_q=self._prefill_metadata.query_start_loc,
85
+ cu_seqlens_k=self._prefill_metadata.query_start_loc,
86
+ max_seqlen_q=self._prefill_metadata.max_query_len,
87
+ max_seqlen_k=self._prefill_metadata.max_query_len,
88
+ softmax_scale=self.scale,
89
+ causal=True,
90
+ return_lse=return_softmax_lse,
91
+ out=out,
92
+ )
93
+
94
+ # aiter returns the bare output tensor when return_lse is False, and
95
+ # (out, softmax_lse) when it is True.
96
+ if return_softmax_lse:
97
+ return result[0], result[1]
98
+ return result
99
+
100
+ def run_prefill_context_chunk(
101
+ self,
102
+ chunk_idx: int,
103
+ q: torch.Tensor,
104
+ k: torch.Tensor,
105
+ v: torch.Tensor,
106
+ ) -> tuple[torch.Tensor, torch.Tensor]:
107
+ assert self._prefill_metadata.chunked_context is not None
108
+ chunked = self._prefill_metadata.chunked_context
109
+ out, lse = self.flash_attn_varlen_func(
110
+ q=q,
111
+ k=k,
112
+ v=v,
113
+ cu_seqlens_q=self._prefill_metadata.query_start_loc,
114
+ cu_seqlens_k=chunked.cu_seq_lens[chunk_idx],
115
+ max_seqlen_q=self._prefill_metadata.max_query_len,
116
+ max_seqlen_k=chunked.max_seq_lens[chunk_idx],
117
+ softmax_scale=self.scale,
118
+ causal=False,
119
+ return_lse=True,
120
+ )
121
+ return out, lse
vllm/v1/attention/backends/mla/prefill/registry.py
+4
@@ -48,6 +48,10 @@ class MLAPrefillBackendEnum(Enum, metaclass=_MLAPrefillBackendEnumMeta):
48
"vllm.v1.attention.backends.mla.prefill.tokenspeed_mla."
49
"TokenspeedMLAPrefillBackend"
50
)
51
+ ROCM_AITER_FA = (
52
+ "vllm.v1.attention.backends.mla.prefill.aiter_flash_attn."
53
+ "AiterFlashAttnPrefillBackend"
54
+ )
55
# Placeholder for third-party/custom backends - must be registered before use
56
# set to None to avoid alias with other backend, whose value is an empty string
57
CUSTOM = None
vllm/v1/attention/backends/mla/prefill/selector.py
+8
@@ -56,6 +56,14 @@ def _get_mla_prefill_backend_priorities(
56
Returns:
57
List of backends in priority order (highest priority first).
58
"""
59
+ from vllm.platforms import current_platform
60
+
61
+ if current_platform.is_rocm():
62
+ return [
63
+ MLAPrefillBackendEnum.ROCM_AITER_FA,
64
+ MLAPrefillBackendEnum.FLASH_ATTN,
65
+ ]
66
+
67
if device_capability.major == 10: # Blackwell
68
return [
69
MLAPrefillBackendEnum.FLASH_ATTN,