[Hardware][AMD][CI] Fix Kernels Attention test groups

.buildkite/test-amd.yaml

@@ -1594,9 +1594,10 @@ steps:
 #----------------------------------------------------------  mi300 · kernels  ----------------------------------------------------------#
 
 - label: Kernels Attention Test %N # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 55
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_1
+  optional: true
   parallelism: 2
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
@@ -3040,7 +3041,7 @@ steps:
 #----------------------------------------------------------  mi355 · kernels  ----------------------------------------------------------#
 
 - label: Kernels (B200-MI355) # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 15
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   agent_pool: mi355_1
   working_dir: "/vllm-workspace/"
@@ -3064,11 +3065,10 @@ steps:
   - pytest -v -s tests/kernels/attention/test_attention_selector.py
 
 - label: Kernels Attention Test %N # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 60
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   agent_pool: mi355_1
   parallelism: 2
-  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - csrc/attention/

.buildkite/test_areas/kernels.yaml

@@ -74,6 +74,20 @@ steps:
   commands:
     - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
   parallelism: 2
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 55
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - csrc/attention/
+      - vllm/v1/attention
+      - vllm/model_executor/layers/attention
+      - tests/kernels/attention
+      - vllm/_aiter_ops.py
+      - vllm/envs.py
+      - vllm/platforms/rocm.py
 
 - label: Kernels Attention DiffKV Test (H100)
   key: kernels-attention-diffkv-test-h100

tests/kernels/attention/test_attention_selector.py

@@ -15,16 +15,14 @@
 from vllm.platforms import current_platform
 from vllm.platforms.cpu import CpuPlatform
 
-# CudaPlatform and RocmPlatform import their respective compiled C extensions
-# at module level, raising ModuleNotFoundError on incompatible builds.
-try:
+if current_platform.is_cuda():
     from vllm.platforms.cuda import CudaPlatform
-except (ImportError, ModuleNotFoundError):
+else:
     CudaPlatform = None
 
-try:
+if current_platform.is_rocm():
     from vllm.platforms.rocm import RocmPlatform
-except (ImportError, ModuleNotFoundError):
+else:
     RocmPlatform = None
 
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -434,9 +432,15 @@ def test_per_head_quant_scales_backend_selection(
     [
         ("FLASH_ATTN", True, True),  # FlashAttn supports non-causal
         ("FLASH_ATTN", False, True),  # FlashAttn also works with causal
-        ("FLASHINFER", True, False),  # FlashInfer does not support non-causal
-        ("FLASHINFER", False, True),  # FlashInfer works with causal
-    ],
+    ]
+    + (
+        [
+            ("FLASHINFER", True, False),  # FlashInfer does not support non-causal
+            ("FLASHINFER", False, True),  # FlashInfer works with causal
+        ]
+        if CudaPlatform is not None
+        else []
+    ),
 )
 def test_non_causal_backend_selection(
     backend_name: str, use_non_causal: bool, should_succeed: bool
@@ -459,11 +463,12 @@ def test_non_causal_backend_selection(
         attention_config=attention_config, cache_config=cache_config
     )
 
-    if CudaPlatform is None:
-        pytest.skip("CudaPlatform not available")
+    platform = CudaPlatform or RocmPlatform
+    if platform is None:
+        pytest.skip("CudaPlatform and RocmPlatform are not available")
     with (
         set_current_vllm_config(vllm_config),
-        patch("vllm.platforms.current_platform", CudaPlatform()),
+        patch("vllm.platforms.current_platform", platform()),
     ):
         if should_succeed:
             backend = get_attn_backend(

tests/kernels/attention/test_prefix_prefill.py

@@ -5,10 +5,12 @@
 import random
 import time
 from collections.abc import Callable
+from contextlib import nullcontext
 
 import pytest
 import torch
 import torch.nn.functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel
 
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
@@ -557,15 +559,21 @@ def _get_alibi_slopes(total_num_heads: int) -> torch.Tensor:
             query_len, seq_len, alibi_slopes, device, dtype
         )
 
-        # Compute attention
-        out = F.scaled_dot_product_attention(
-            q_sdpa,
-            k_sdpa,
-            v_sdpa,
-            attn_mask=alibi_mask,
-            dropout_p=0.0,
-            scale=scale,
-        )
+        # Compute attention. On ROCm we force use of the Math SDPA backend rather than
+        # the Flash or Mem-Efficient backends for increased numerical accuracy
+        if current_platform.is_rocm():
+            sdpa_context = sdpa_kernel(SDPBackend.MATH)
+        else:
+            sdpa_context = nullcontext()
+        with sdpa_context:
+            out = F.scaled_dot_product_attention(
+                q_sdpa,
+                k_sdpa,
+                v_sdpa,
+                attn_mask=alibi_mask,
+                dropout_p=0.0,
+                scale=scale,
+            )
 
         # Reshape output back to [query_len, num_heads, head_size]
         out = out.view(num_heads, query_len, head_size).permute(1, 0, 2)

tests/kernels/attention/test_rocm_triton_attn_dsv4.py

@@ -90,7 +90,9 @@ def _ref_sparse_prefill_ragged(
     return out.to(torch.bfloat16)
 
 
-def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor:
+def _pack_fp8_ds_mla_cache(
+    kv: torch.Tensor, block_size: int, is_extra: bool = False
+) -> torch.Tensor:
     assert kv.shape[-1] == HEAD_DIM
     num_tokens = kv.shape[0]
     num_blocks = (num_tokens + block_size - 1) // block_size
@@ -101,7 +103,9 @@ def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor:
     )
     cache_flat = cache.view(torch.uint8).flatten()
     kv_nope_fp8 = (
-        kv[:, :NOPE_HEAD_DIM].to(current_platform.fp8_dtype()).view(torch.uint8)
+        kv[:, :NOPE_HEAD_DIM]
+        .to(torch.float8_e4m3fn if is_extra else current_platform.fp8_dtype())
+        .view(torch.uint8)
     )
     kv_rope_u8 = kv[:, NOPE_HEAD_DIM:].contiguous().view(torch.uint8)
 
@@ -120,7 +124,7 @@ def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor:
 
 
 def _read_fp8_ds_mla_cache(
-    cache: torch.Tensor, slot: int, block_size: int
+    cache: torch.Tensor, slot: int, block_size: int, is_extra: bool = False
 ) -> torch.Tensor:
     cache_flat = cache.view(torch.uint8).flatten()
     block_idx = slot // block_size
@@ -129,7 +133,9 @@ def _read_fp8_ds_mla_cache(
     token_base = block_base + pos * 576
 
     nope_u8 = cache_flat[token_base : token_base + NOPE_HEAD_DIM]
-    nope = nope_u8.view(current_platform.fp8_dtype()).to(torch.float32)
+    nope = nope_u8.view(
+        torch.float8_e4m3fn if is_extra else current_platform.fp8_dtype()
+    ).to(torch.float32)
     rope_u8 = cache_flat[
         token_base + NOPE_HEAD_DIM : token_base + NOPE_HEAD_DIM + ROPE_HEAD_DIM * 2
     ]
@@ -157,7 +163,9 @@ def _ref_sparse_decode_ragged(
         ]
         if extra_cache is not None and extra_rows is not None:
             row_kv.extend(
-                _read_fp8_ds_mla_cache(extra_cache, int(slot), block_size)
+                _read_fp8_ds_mla_cache(
+                    extra_cache, int(slot), block_size, is_extra=True
+                )
                 for slot in extra_rows[query_idx]
             )
 
@@ -326,7 +334,7 @@ def test_sparse_attn_decode_ragged_kernel() -> None:
     main_kv = torch.randn(6, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
     extra_kv = torch.randn(5, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
     main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size)
-    extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size)
+    extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, is_extra=True)
     main_indices = torch.tensor([0, 2, 4, 1], dtype=torch.int32, device=device)
     main_indptr = torch.tensor([0, 2, 4], dtype=torch.int32, device=device)
     extra_indices = torch.tensor([1, 3, 0], dtype=torch.int32, device=device)
@@ -477,7 +485,7 @@ def test_sparse_attn_decode_split_k_kernel(
         rows = [[1, 3, 0, 5, 2, 4], [3, 0, 6]]
         extra_kv = torch.randn(7, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
         extra_rows = rows
-        extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size)
+        extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, is_extra=True)
         extra_indices, extra_indptr = _ragged_from_rows(rows, device)
 
     attn_sink = (

tests/kernels/attention/test_triton_unified_attention.py

@@ -18,11 +18,7 @@
 BLOCK_SIZES = [16]
 
 DTYPES = [torch.bfloat16]
-QDTYPES = (
-    [None, torch.float8_e4m3fn]
-    if not current_platform.is_rocm()
-    else [None, torch.float8_e4m3fnuz]
-)
+QDTYPES = [None, current_platform.fp8_dtype()]
 FP8_DTYPE = current_platform.fp8_dtype()
 
 # one value large enough to test overflow in index calculation.