[ROCm] Fix AITER_UNIFIED_ATTN Dispatching After AITER Bump (#46780)
Signed-off-by: Micah Williamson <micah.williamson@amd.com> Signed-off-by: Rohan138 <rohanpotdar138@gmail.com> Co-authored-by: Rohan138 <rohanpotdar138@gmail.com>
Micah Williamson committed
Jun 26, 2026 at 04:09 UTC
d980a3cc6ed9fc83386894211170f1ca85ac9735
6 files changed
+30
-3
.buildkite/hardware_tests/amd.yaml
+2
@@ -6,6 +6,7 @@ steps:
6
# differ ci_base is rebuilt and pushed automatically.
7
- label: "AMD: :docker: ensure ci_base"
8
key: ensure-ci-base-amd
9
+ soft_fail: false
10
depends_on: []
11
device: amd_cpu
12
no_plugin: true
@@ -26,6 +27,7 @@ steps:
27
28
- label: "AMD: :docker: build test image and artifacts"
29
key: image-build-amd
30
+ soft_fail: false
31
depends_on:
32
- ensure-ci-base-amd
33
device: amd_cpu
docs/design/attention_backends.md
+1
-1
@@ -168,7 +168,7 @@ Priority is **1 = highest** (tried first).
168
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
169
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
170
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A |
171
-| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
171
+| `ROCM_AITER_UNIFIED_ATTN` | | bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
172
| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
173
| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int4_per_token_head`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ✅ | ❌ | All | Any |
174
| `TRITON_ATTN_DIFFKV` | | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
tests/compile/passes/test_fusion_attn.py
+5
@@ -306,6 +306,11 @@ def test_attention_quant_pattern(
306
torch.manual_seed(42)
307
308
backend_cls = backend.get_class()
309
+
310
+ # TODO: drop once AITER reenables fp16 unified attention.
311
+ if dtype not in backend_cls.supported_dtypes:
312
+ pytest.skip(f"{backend.name} does not support dtype {dtype}")
313
+
314
block_size = backend_cls.get_preferred_block_size(16)
315
316
model_config = ModelConfig(
tests/kernels/attention/test_rocm_aiter_unified_attn.py
+2
-1
@@ -30,7 +30,8 @@ NUM_Q_HEADS = 8
30
NUM_KV_HEADS = 8
31
HEAD_SIZES = [128, 256]
32
BLOCK_SIZES = [16, 64]
33
-DTYPES = [torch.bfloat16, torch.float16]
33
+# TODO: re-add torch.float16 once AITER reenables fp16 unified attention.
34
+DTYPES = [torch.bfloat16]
35
FP8_DTYPE = current_platform.fp8_dtype()
36
37
# (query_len, kv_len) per sequence
tests/v1/attention/test_rocm_attention_backends_selection.py
+9
-1
@@ -136,9 +136,17 @@ def test_standard_attention_backend_selection(
136
# Get the backend class path
137
from vllm.platforms.rocm import RocmPlatform
138
139
+ # The AITER unified attention kernel only supports BF16/FP8 KV caches
140
+ # (its 3D kernel asserts on fp16), so it must be selected with bf16.
141
+ dtype = (
142
+ torch.bfloat16
143
+ if selected_backend == "ROCM_AITER_UNIFIED_ATTN"
144
+ else torch.float16
145
+ )
146
+
147
attn_selector_config = AttentionSelectorConfig(
148
head_size=128,
141
- dtype=torch.float16,
149
+ dtype=dtype,
150
kv_cache_dtype="auto",
151
block_size=16,
152
use_mla=False,
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+11
@@ -2,10 +2,13 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Attention layer with PagedAttention and Triton prefix prefill."""
4
5
+from typing import ClassVar
6
+
7
import torch
8
9
from vllm import _custom_ops as ops
10
from vllm._aiter_ops import rocm_aiter_ops
11
+from vllm.config.cache import CacheDType
12
from vllm.logger import init_logger
13
from vllm.model_executor.layers.quantization.utils.quant_utils import (
14
QuantKey,
@@ -24,6 +27,14 @@ logger = init_logger(__name__)
27
28
29
class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):
30
+ supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
31
+ supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
32
+ "auto",
33
+ "bfloat16",
34
+ "fp8",
35
+ "fp8_e4m3",
36
+ ]
37
+
38
@staticmethod
39
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
40
return [MultipleOf(16)]