[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>

.buildkite/hardware_tests/amd.yaml

@@ -6,6 +6,7 @@ steps:
   # differ ci_base is rebuilt and pushed automatically.
   - label: "AMD: :docker: ensure ci_base"
     key: ensure-ci-base-amd
+    soft_fail: false
     depends_on: []
     device: amd_cpu
     no_plugin: true
@@ -26,6 +27,7 @@ steps:
 
   - label: "AMD: :docker: build test image and artifacts"
     key: image-build-amd
+    soft_fail: false
     depends_on:
       - ensure-ci-base-amd
     device: amd_cpu

docs/design/attention_backends.md

@@ -168,7 +168,7 @@ Priority is **1 = highest** (tried first).
 | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
 | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
 | `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A |
-| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
+| `ROCM_AITER_UNIFIED_ATTN` | | bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
 | `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 |
 | `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 |
 | `TRITON_ATTN_DIFFKV` | | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |

tests/compile/passes/test_fusion_attn.py

@@ -306,6 +306,11 @@ def test_attention_quant_pattern(
     torch.manual_seed(42)
 
     backend_cls = backend.get_class()
+
+    # TODO: drop once AITER reenables fp16 unified attention.
+    if dtype not in backend_cls.supported_dtypes:
+        pytest.skip(f"{backend.name} does not support dtype {dtype}")
+
     block_size = backend_cls.get_preferred_block_size(16)
 
     model_config = ModelConfig(

tests/kernels/attention/test_rocm_aiter_unified_attn.py

@@ -30,7 +30,8 @@ NUM_Q_HEADS = 8
 NUM_KV_HEADS = 8
 HEAD_SIZES = [128, 256]
 BLOCK_SIZES = [16, 64]
-DTYPES = [torch.bfloat16, torch.float16]
+# TODO: re-add torch.float16 once AITER reenables fp16 unified attention.
+DTYPES = [torch.bfloat16]
 FP8_DTYPE = current_platform.fp8_dtype()
 
 # (query_len, kv_len) per sequence

tests/v1/attention/test_rocm_attention_backends_selection.py

@@ -136,9 +136,17 @@ def test_standard_attention_backend_selection(
     # Get the backend class path
     from vllm.platforms.rocm import RocmPlatform
 
+    # The AITER unified attention kernel only supports BF16/FP8 KV caches
+    # (its 3D kernel asserts on fp16), so it must be selected with bf16.
+    dtype = (
+        torch.bfloat16
+        if selected_backend == "ROCM_AITER_UNIFIED_ATTN"
+        else torch.float16
+    )
+
     attn_selector_config = AttentionSelectorConfig(
         head_size=128,
-        dtype=torch.float16,
+        dtype=dtype,
         kv_cache_dtype="auto",
         block_size=16,
         use_mla=False,

vllm/v1/attention/backends/rocm_aiter_unified_attn.py

@@ -2,10 +2,13 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Attention layer with PagedAttention and Triton prefix prefill."""
 
+from typing import ClassVar
+
 import torch
 
 from vllm import _custom_ops as ops
 from vllm._aiter_ops import rocm_aiter_ops
+from vllm.config.cache import CacheDType
 from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     QuantKey,
@@ -24,6 +27,14 @@ logger = init_logger(__name__)
 
 
 class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):
+    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
+    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
+        "auto",
+        "bfloat16",
+        "fp8",
+        "fp8_e4m3",
+    ]
+
     @staticmethod
     def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
         return [MultipleOf(16)]