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 7, 05:15 UTC
[ROCm] Refactor ROCm attention backend selection logic (#35246)
Signed-off-by: Sage Moore <sage@neuralmagic.com>
Sage Moore committed
Mar 5, 2026 at 08:51 UTC
8c760b6ab6993c6a0d5f639747baefedb4612525
6 files changed
+170
-114
docs/design/attention_backends.md
+1
-1
@@ -211,6 +211,6 @@ configuration.
211
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
212
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
213
| `ROCM_AITER_MLA` | fp16, bf16 | `auto` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
214
-| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto` | Any | 576 | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
214
+| `ROCM_AITER_MLA_SPARSE` | bf16 | `auto` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
215
| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
216
| `TRITON_MLA` | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
tests/kernels/attention/test_attention_selector.py
+4
-5
@@ -103,21 +103,20 @@ def test_backend_selection(
103
104
if name == "TRITON_MLA" and block_size == 1:
105
# TRITON_MLA doesn't support block_size == 1
106
- with pytest.raises(ValueError) as exc_info:
106
+ with pytest.raises(ValueError):
107
get_attn_backend(
108
- 16, torch.float16, None, block_size, use_mla=use_mla
108
+ 576, torch.float16, None, block_size, use_mla=use_mla
109
)
110
- assert f"The selected backend, {name}" in str(exc_info.value)
110
else:
111
# Valid backend-block_size combination
112
backend = get_attn_backend(
114
- 16, torch.float16, None, block_size, use_mla=use_mla
113
+ 576, torch.float16, None, block_size, use_mla=use_mla
114
)
115
expected = name
116
assert backend.get_name() == expected
117
else:
118
backend = get_attn_backend(
120
- 16, torch.float16, None, block_size, use_mla=use_mla
119
+ 32, torch.float16, None, block_size, use_mla=use_mla
120
)
121
expected = "ROCM_ATTN"
122
assert backend.get_name() == expected
vllm/platforms/rocm.py
+136
-104
@@ -306,6 +306,52 @@ def flash_attn_triton_available() -> bool:
306
return False
307
308
309
+def _get_backend_priorities(
310
+ use_mla: bool,
311
+ use_sparse: bool,
312
+) -> list[AttentionBackendEnum]:
313
+ from vllm._aiter_ops import rocm_aiter_ops
314
+
315
+ if use_sparse:
316
+ return [AttentionBackendEnum.ROCM_AITER_MLA_SPARSE]
317
+
318
+ if use_mla:
319
+ if rocm_aiter_ops.is_mla_enabled():
320
+ return [
321
+ AttentionBackendEnum.ROCM_AITER_MLA,
322
+ AttentionBackendEnum.TRITON_MLA,
323
+ AttentionBackendEnum.ROCM_AITER_TRITON_MLA,
324
+ ]
325
+ else:
326
+ return [
327
+ AttentionBackendEnum.TRITON_MLA,
328
+ ]
329
+
330
+ backends = []
331
+
332
+ # Priority 1: Check for AITER Unified Attention (must check before MHA)
333
+ if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION:
334
+ backends.append(AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN)
335
+
336
+ # Priority 2: Check for AITER MHA (Flash Attention)
337
+ if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA:
338
+ backends.append(AttentionBackendEnum.ROCM_AITER_FA)
339
+
340
+ # Priority 3: Check for ROCM_ATTN (prefill-decode split)
341
+ from vllm.config import get_current_vllm_config_or_none
342
+
343
+ vllm_config = get_current_vllm_config_or_none()
344
+ if (
345
+ vllm_config is not None
346
+ and vllm_config.attention_config.use_prefill_decode_attention
347
+ ):
348
+ backends.append(AttentionBackendEnum.ROCM_ATTN)
349
+
350
+ # Default: Triton Unified Attention
351
+ backends.append(AttentionBackendEnum.TRITON_ATTN)
352
+ return backends
353
+
354
+
355
class RocmPlatform(Platform):
356
_enum = PlatformEnum.ROCM
357
device_name: str = "rocm"
@@ -349,6 +395,39 @@ class RocmPlatform(Platform):
395
with contextlib.suppress(ImportError):
396
import vllm._rocm_C # noqa: F401
397
398
+ @classmethod
399
+ def get_valid_backends(
400
+ cls,
401
+ device_capability: DeviceCapability,
402
+ attn_selector_config: "AttentionSelectorConfig",
403
+ num_heads: int | None = None,
404
+ ) -> tuple[
405
+ list[tuple["AttentionBackendEnum", int]],
406
+ dict["AttentionBackendEnum", list[str]],
407
+ ]:
408
+ valid_backends_priorities = []
409
+ invalid_reasons = {}
410
+
411
+ backend_priorities = _get_backend_priorities(
412
+ attn_selector_config.use_mla,
413
+ attn_selector_config.use_sparse,
414
+ )
415
+ for priority, backend in enumerate(backend_priorities):
416
+ try:
417
+ backend_class = backend.get_class()
418
+ invalid_reasons_i = backend_class.validate_configuration(
419
+ device_capability=device_capability,
420
+ **attn_selector_config._asdict(),
421
+ )
422
+ except ImportError:
423
+ invalid_reasons_i = ["ImportError"]
424
+ if invalid_reasons_i:
425
+ invalid_reasons[backend] = invalid_reasons_i
426
+ else:
427
+ valid_backends_priorities.append((backend, priority))
428
+
429
+ return valid_backends_priorities, invalid_reasons
430
+
431
@classmethod
432
def get_attn_backend_cls(
433
cls,
@@ -356,118 +435,71 @@ class RocmPlatform(Platform):
435
attn_selector_config: "AttentionSelectorConfig",
436
num_heads: int | None = None,
437
) -> str:
359
- from vllm._aiter_ops import rocm_aiter_ops
360
-
361
- block_size = attn_selector_config.block_size
362
- kv_cache_dtype = attn_selector_config.kv_cache_dtype
363
-
364
- if attn_selector_config.use_sparse:
365
- if kv_cache_dtype and kv_cache_dtype.startswith("fp8"):
366
- raise ValueError(
367
- "ROCMAiterMLASparseBackend doesn't support fp8 kv_cache_dtype."
438
+ device_capability = cls.get_device_capability()
439
+ assert device_capability is not None
440
+
441
+ # First try checking just the selected backend, if there is one.
442
+ if selected_backend is not None:
443
+ try:
444
+ backend_class = selected_backend.get_class()
445
+ invalid_reasons = backend_class.validate_configuration(
446
+ device_capability=device_capability,
447
+ **attn_selector_config._asdict(),
448
)
369
- assert block_size == 1, (
370
- "Sparse MLA backend on ROCm only supports block size 1 for now."
371
- )
372
- logger.info_once("Using Sparse MLA backend.")
373
- return AttentionBackendEnum.ROCM_AITER_MLA_SPARSE.get_path()
374
-
375
- if attn_selector_config.use_mla:
376
- if selected_backend is None:
377
- selected_backend = (
378
- AttentionBackendEnum.ROCM_AITER_MLA
379
- if rocm_aiter_ops.is_mla_enabled() or block_size == 1
380
- else AttentionBackendEnum.TRITON_MLA
381
- )
382
- if selected_backend == AttentionBackendEnum.TRITON_MLA:
383
- if block_size != 1:
384
- logger.info_once("Using Triton MLA backend.")
385
- return AttentionBackendEnum.TRITON_MLA.get_path()
449
+ except ImportError:
450
+ invalid_reasons = ["ImportError"]
451
+ if invalid_reasons:
452
raise ValueError(
387
- f" The selected backend, {selected_backend.name},"
388
- f"does not support block size {block_size}."
453
+ f"Selected backend {selected_backend} is not valid for "
454
+ f"this configuration. Reason: {invalid_reasons}"
455
)
390
- if selected_backend == AttentionBackendEnum.ROCM_AITER_MLA:
391
- logger.info("Using AITER MLA backend.")
392
- return AttentionBackendEnum.ROCM_AITER_MLA.get_path()
393
- if selected_backend == AttentionBackendEnum.ROCM_AITER_TRITON_MLA:
394
- logger.info("Using AITER TRITON MLA backend.")
395
- return AttentionBackendEnum.ROCM_AITER_TRITON_MLA.get_path()
396
-
456
+ else:
457
+ logger.info("Using %s backend.", selected_backend)
458
+ return selected_backend.get_path()
459
+
460
+ # No selected backend or the selected backend is invalid,
461
+ # so we try finding a valid backend.
462
+ valid_backends_priorities, invalid_reasons = cls.get_valid_backends(
463
+ device_capability=device_capability,
464
+ attn_selector_config=attn_selector_config,
465
+ num_heads=num_heads,
466
+ )
467
+ reasons_str = (
468
+ "{"
469
+ + ", ".join(
470
+ f"{backend.name}: [{', '.join(reasons)}]"
471
+ for backend, reasons in invalid_reasons.items()
472
+ )
473
+ + "}"
474
+ )
475
+ config_str = attn_selector_config.__repr__()
476
+ logger.debug_once(
477
+ f"Some attention backends are not valid for {cls.device_name} with "
478
+ f"{config_str}. Reasons: {reasons_str}."
479
+ )
480
+ if len(valid_backends_priorities) == 0:
481
raise ValueError(
398
- f" The selected backend, {selected_backend.name},"
399
- f"is not MLA type while requested for MLA backend."
482
+ f"No valid attention backend found for {cls.device_name} "
483
+ f"with {config_str}. Reasons: {reasons_str}."
484
)
485
402
- if selected_backend == AttentionBackendEnum.FLEX_ATTENTION:
403
- logger.info("Using FlexAttention backend.")
404
- return AttentionBackendEnum.FLEX_ATTENTION.get_path()
405
-
406
- if selected_backend == AttentionBackendEnum.TRITON_ATTN:
407
- logger.info("Using Triton Attention backend.")
408
- return AttentionBackendEnum.TRITON_ATTN.get_path()
409
-
410
- if selected_backend == AttentionBackendEnum.ROCM_ATTN:
411
- logger.info("Using Rocm Attention backend.")
412
- return AttentionBackendEnum.ROCM_ATTN.get_path()
413
-
414
- if selected_backend == AttentionBackendEnum.ROCM_AITER_FA:
415
- if on_gfx9():
416
- logger.info("Using Aiter Flash Attention backend.")
417
- return AttentionBackendEnum.ROCM_AITER_FA.get_path()
418
- else:
419
- raise ValueError(
420
- f"The selected backend, {selected_backend.name}, "
421
- "is only supported on gfx9 architectures."
422
- )
423
-
424
- if selected_backend == AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN:
425
- logger.info("Using Aiter Unified Attention backend.")
426
- return AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path()
427
-
428
- # Handle automatic backend selection based on environment variables
429
- if selected_backend is None:
430
- # Priority 1: Check for AITER Unified Attention (must check before MHA)
431
- if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION:
432
- logger.info("Using Aiter Unified Attention backend.")
433
- return AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path()
434
-
435
- # Priority 2: Check for AITER MHA (Flash Attention)
436
- # Only use if explicitly enabled (not just VLLM_ROCM_USE_AITER=1)
437
- if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA and on_gfx9():
438
- logger.info("Using Aiter Flash Attention backend.")
439
- return AttentionBackendEnum.ROCM_AITER_FA.get_path()
440
-
441
- # Priority 3: Check for ROCM_ATTN (prefill-decode split)
442
- from vllm.config import get_current_vllm_config_or_none
443
-
444
- vllm_config = get_current_vllm_config_or_none()
445
- if (
446
- vllm_config is not None
447
- and vllm_config.attention_config.use_prefill_decode_attention
448
- ):
449
- logger.info("Using Rocm Attention backend.")
450
- return AttentionBackendEnum.ROCM_ATTN.get_path()
451
-
452
- # Priority 4: Check for AITER enabled without specific flags
453
- # This defaults to AITER FA only if MHA is not explicitly disabled
454
- if (
455
- envs.VLLM_ROCM_USE_AITER
456
- and on_gfx9()
457
- and envs.VLLM_ROCM_USE_AITER_MHA is not False
458
- ):
459
- logger.info("Using Aiter Flash Attention backend.")
460
- return AttentionBackendEnum.ROCM_AITER_FA.get_path()
461
-
462
- # Default: Triton Unified Attention
463
- logger.info("Using Triton Attention backend.")
464
- return AttentionBackendEnum.TRITON_ATTN.get_path()
465
-
466
- raise RuntimeError(
467
- f"Attention backend {selected_backend.name} is not supported on "
468
- "ROCm. Note that V0 attention backends have been removed."
486
+ # We have found some valid backends. Select the one with the
487
+ # highest priority.
488
+ sorted_indices = sorted(
489
+ range(len(valid_backends_priorities)),
490
+ key=lambda i: valid_backends_priorities[i][1],
491
+ )
492
+ selected_index = sorted_indices[0]
493
+ selected_backend = valid_backends_priorities[selected_index][0]
494
+ logger.info_once(
495
+ "Using %s attention backend out of potential backends: %s.",
496
+ selected_backend.name,
497
+ "[" + ", ".join(f"'{b[0].name}'" for b in valid_backends_priorities) + "]",
498
+ scope="local",
499
)
500
501
+ return selected_backend.get_path()
502
+
503
@classmethod
504
def get_supported_vit_attn_backends(cls) -> list["AttentionBackendEnum"]:
505
return [
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+14
-4
@@ -77,6 +77,7 @@ def fetch_id_to_ragged_triton(
77
78
class ROCMAiterMLASparseBackend(AttentionBackend):
79
accept_output_buffer: bool = True
80
+ supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
81
82
@staticmethod
83
def get_name() -> str:
@@ -104,14 +105,23 @@ class ROCMAiterMLASparseBackend(AttentionBackend):
105
) -> tuple[int, ...]:
106
return (num_blocks, block_size, head_size)
107
107
- @classmethod
108
- def get_supported_dtypes(cls) -> list[torch.dtype]:
109
- return [torch.bfloat16]
110
-
108
@classmethod
109
def get_supported_head_sizes(cls) -> list[int]:
110
return [576]
111
112
+ @classmethod
113
+ def is_mla(cls) -> bool:
114
+ return True
115
+
116
+ @classmethod
117
+ def is_sparse(cls) -> bool:
118
+ return True
119
+
120
+ @classmethod
121
+ def supports_block_size(cls, block_size: int | None) -> bool:
122
+ # The only supported block_size is 1
123
+ return block_size is None or block_size == 1
124
+
125
126
@dataclass
127
class ROCMAiterMLASparseMetadata(AttentionMetadata):
vllm/v1/attention/backends/mla/triton_mla.py
+5
@@ -45,6 +45,11 @@ class TritonMLABackend(MLACommonBackend):
45
def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
46
return True
47
48
+ @classmethod
49
+ def supports_block_size(cls, block_size: int | None) -> bool:
50
+ # The only unsupported block_size is 1
51
+ return block_size is None or block_size != 1
52
+
53
54
class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
55
can_return_lse_for_decode: bool = True
vllm/v1/attention/backends/rocm_aiter_fa.py
+10
@@ -12,6 +12,7 @@ from vllm.config import VllmConfig, get_layers_from_vllm_config
12
from vllm.logger import init_logger
13
from vllm.model_executor.layers.attention import Attention
14
from vllm.platforms import current_platform
15
+from vllm.platforms.interface import DeviceCapability
16
from vllm.utils.math_utils import cdiv
17
from vllm.utils.platform_utils import num_compute_units
18
from vllm.v1.attention.backend import (
@@ -766,6 +767,15 @@ class AiterFlashAttentionBackend(AttentionBackend):
767
raise ValueError("Block size must be a multiple of 16.")
768
return (2, num_blocks, block_size, num_kv_heads, head_size)
769
770
+ @classmethod
771
+ def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
772
+ from vllm.platforms.rocm import on_mi3xx
773
+
774
+ # DeviceCapability is currently created using torch.cuda.get_device_capability()
775
+ # which is known to be buggy on rocm systems. on_mi3xx uses amd-smi which is
776
+ # more reliable.
777
+ return on_mi3xx()
778
+
779
780
class AiterFlashAttentionImpl(AttentionImpl):
781
def __init__(