diff --git a/.buildkite/hardware_tests/amd.yaml b/.buildkite/hardware_tests/amd.yaml index c2510f38aa..a18241cf18 100644 --- a/.buildkite/hardware_tests/amd.yaml +++ b/.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 diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index d268d5b4db..9278ab6761 100644 --- a/docs/design/attention_backends.md +++ b/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 | diff --git a/tests/compile/passes/test_fusion_attn.py b/tests/compile/passes/test_fusion_attn.py index b776f6af98..531d26e008 100644 --- a/tests/compile/passes/test_fusion_attn.py +++ b/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( diff --git a/tests/kernels/attention/test_rocm_aiter_unified_attn.py b/tests/kernels/attention/test_rocm_aiter_unified_attn.py index 9e33f24ea2..c02a457c98 100644 --- a/tests/kernels/attention/test_rocm_aiter_unified_attn.py +++ b/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 diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index 8f9e8acac6..48c6de8f8b 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/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, diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index 984fc20eca..d8363169a8 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/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)]