[Hardware][AMD][CI] Fix Kernels Attention test groups (#46080)
Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>
Matt committed
Jun 21, 2026 at 17:10 UTC
4f0d0049a0a2a188fcfaa2d07317d629b79b81d9
6 files changed
+68
-37
.buildkite/test-amd.yaml
+4
-4
@@ -1594,9 +1594,10 @@ steps:
1594
#---------------------------------------------------------- mi300 · kernels ----------------------------------------------------------#
1595
1596
- label: Kernels Attention Test %N # TBD
1597
- timeout_in_minutes: 180
1597
+ timeout_in_minutes: 55
1598
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
1599
agent_pool: mi300_1
1600
+ optional: true
1601
parallelism: 2
1602
working_dir: "/vllm-workspace/tests"
1603
source_file_dependencies:
@@ -3041,7 +3042,7 @@ steps:
3042
#---------------------------------------------------------- mi355 · kernels ----------------------------------------------------------#
3043
3044
- label: Kernels (B200-MI355) # TBD
3044
- timeout_in_minutes: 180
3045
+ timeout_in_minutes: 15
3046
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
3047
agent_pool: mi355_1
3048
working_dir: "/vllm-workspace/"
@@ -3065,11 +3066,10 @@ steps:
3066
- pytest -v -s tests/kernels/attention/test_attention_selector.py
3067
3068
- label: Kernels Attention Test %N # TBD
3068
- timeout_in_minutes: 180
3069
+ timeout_in_minutes: 60
3070
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
3071
agent_pool: mi355_1
3072
parallelism: 2
3072
- optional: true
3073
working_dir: "/vllm-workspace/tests"
3074
source_file_dependencies:
3075
- csrc/attention/
.buildkite/test_areas/kernels.yaml
+14
@@ -74,6 +74,20 @@ steps:
74
commands:
75
- pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
76
parallelism: 2
77
+ mirror:
78
+ amd:
79
+ device: mi325_1
80
+ timeout_in_minutes: 55
81
+ depends_on:
82
+ - image-build-amd
83
+ source_file_dependencies:
84
+ - csrc/attention/
85
+ - vllm/v1/attention
86
+ - vllm/model_executor/layers/attention
87
+ - tests/kernels/attention
88
+ - vllm/_aiter_ops.py
89
+ - vllm/envs.py
90
+ - vllm/platforms/rocm.py
91
92
- label: Kernels Attention DiffKV Test (H100)
93
key: kernels-attention-diffkv-test-h100
tests/kernels/attention/test_attention_selector.py
+17
-12
@@ -15,16 +15,14 @@ from vllm.config import (
15
from vllm.platforms import current_platform
16
from vllm.platforms.cpu import CpuPlatform
17
18
-# CudaPlatform and RocmPlatform import their respective compiled C extensions
19
-# at module level, raising ModuleNotFoundError on incompatible builds.
20
-try:
18
+if current_platform.is_cuda():
19
from vllm.platforms.cuda import CudaPlatform
22
-except (ImportError, ModuleNotFoundError):
20
+else:
21
CudaPlatform = None
22
25
-try:
23
+if current_platform.is_rocm():
24
from vllm.platforms.rocm import RocmPlatform
27
-except (ImportError, ModuleNotFoundError):
25
+else:
26
RocmPlatform = None
27
28
from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -434,9 +432,15 @@ def test_per_head_quant_scales_backend_selection(
432
[
433
("FLASH_ATTN", True, True), # FlashAttn supports non-causal
434
("FLASH_ATTN", False, True), # FlashAttn also works with causal
437
- ("FLASHINFER", True, False), # FlashInfer does not support non-causal
438
- ("FLASHINFER", False, True), # FlashInfer works with causal
439
- ],
435
+ ]
436
+ + (
437
+ [
438
+ ("FLASHINFER", True, False), # FlashInfer does not support non-causal
439
+ ("FLASHINFER", False, True), # FlashInfer works with causal
440
+ ]
441
+ if CudaPlatform is not None
442
+ else []
443
+ ),
444
)
445
def test_non_causal_backend_selection(
446
backend_name: str, use_non_causal: bool, should_succeed: bool
@@ -459,11 +463,12 @@ def test_non_causal_backend_selection(
463
attention_config=attention_config, cache_config=cache_config
464
)
465
462
- if CudaPlatform is None:
463
- pytest.skip("CudaPlatform not available")
466
+ platform = CudaPlatform or RocmPlatform
467
+ if platform is None:
468
+ pytest.skip("CudaPlatform and RocmPlatform are not available")
469
with (
470
set_current_vllm_config(vllm_config),
466
- patch("vllm.platforms.current_platform", CudaPlatform()),
471
+ patch("vllm.platforms.current_platform", platform()),
472
):
473
if should_succeed:
474
backend = get_attn_backend(
tests/kernels/attention/test_prefix_prefill.py
+17
-9
@@ -5,10 +5,12 @@ import math
5
import random
6
import time
7
from collections.abc import Callable
8
+from contextlib import nullcontext
9
10
import pytest
11
import torch
12
import torch.nn.functional as F
13
+from torch.nn.attention import SDPBackend, sdpa_kernel
14
15
from vllm.platforms import current_platform
16
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
@@ -557,15 +559,21 @@ def test_contexted_kv_attention_alibi(
559
query_len, seq_len, alibi_slopes, device, dtype
560
)
561
560
- # Compute attention
561
- out = F.scaled_dot_product_attention(
562
- q_sdpa,
563
- k_sdpa,
564
- v_sdpa,
565
- attn_mask=alibi_mask,
566
- dropout_p=0.0,
567
- scale=scale,
568
- )
562
+ # Compute attention. On ROCm we force use of the Math SDPA backend rather than
563
+ # the Flash or Mem-Efficient backends for increased numerical accuracy
564
+ if current_platform.is_rocm():
565
+ sdpa_context = sdpa_kernel(SDPBackend.MATH)
566
+ else:
567
+ sdpa_context = nullcontext()
568
+ with sdpa_context:
569
+ out = F.scaled_dot_product_attention(
570
+ q_sdpa,
571
+ k_sdpa,
572
+ v_sdpa,
573
+ attn_mask=alibi_mask,
574
+ dropout_p=0.0,
575
+ scale=scale,
576
+ )
577
578
# Reshape output back to [query_len, num_heads, head_size]
579
out = out.view(num_heads, query_len, head_size).permute(1, 0, 2)
tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+15
-7
@@ -90,7 +90,9 @@ def _ref_sparse_prefill_ragged(
90
return out.to(torch.bfloat16)
91
92
93
-def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor:
93
+def _pack_fp8_ds_mla_cache(
94
+ kv: torch.Tensor, block_size: int, is_extra: bool = False
95
+) -> torch.Tensor:
96
assert kv.shape[-1] == HEAD_DIM
97
num_tokens = kv.shape[0]
98
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:
103
)
104
cache_flat = cache.view(torch.uint8).flatten()
105
kv_nope_fp8 = (
104
- kv[:, :NOPE_HEAD_DIM].to(current_platform.fp8_dtype()).view(torch.uint8)
106
+ kv[:, :NOPE_HEAD_DIM]
107
+ .to(torch.float8_e4m3fn if is_extra else current_platform.fp8_dtype())
108
+ .view(torch.uint8)
109
)
110
kv_rope_u8 = kv[:, NOPE_HEAD_DIM:].contiguous().view(torch.uint8)
111
@@ -120,7 +124,7 @@ def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor:
124
125
126
def _read_fp8_ds_mla_cache(
123
- cache: torch.Tensor, slot: int, block_size: int
127
+ cache: torch.Tensor, slot: int, block_size: int, is_extra: bool = False
128
) -> torch.Tensor:
129
cache_flat = cache.view(torch.uint8).flatten()
130
block_idx = slot // block_size
@@ -129,7 +133,9 @@ def _read_fp8_ds_mla_cache(
133
token_base = block_base + pos * 576
134
135
nope_u8 = cache_flat[token_base : token_base + NOPE_HEAD_DIM]
132
- nope = nope_u8.view(current_platform.fp8_dtype()).to(torch.float32)
136
+ nope = nope_u8.view(
137
+ torch.float8_e4m3fn if is_extra else current_platform.fp8_dtype()
138
+ ).to(torch.float32)
139
rope_u8 = cache_flat[
140
token_base + NOPE_HEAD_DIM : token_base + NOPE_HEAD_DIM + ROPE_HEAD_DIM * 2
141
]
@@ -157,7 +163,9 @@ def _ref_sparse_decode_ragged(
163
]
164
if extra_cache is not None and extra_rows is not None:
165
row_kv.extend(
160
- _read_fp8_ds_mla_cache(extra_cache, int(slot), block_size)
166
+ _read_fp8_ds_mla_cache(
167
+ extra_cache, int(slot), block_size, is_extra=True
168
+ )
169
for slot in extra_rows[query_idx]
170
)
171
@@ -326,7 +334,7 @@ def test_sparse_attn_decode_ragged_kernel() -> None:
334
main_kv = torch.randn(6, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
335
extra_kv = torch.randn(5, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
336
main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size)
329
- extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size)
337
+ extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, is_extra=True)
338
main_indices = torch.tensor([0, 2, 4, 1], dtype=torch.int32, device=device)
339
main_indptr = torch.tensor([0, 2, 4], dtype=torch.int32, device=device)
340
extra_indices = torch.tensor([1, 3, 0], dtype=torch.int32, device=device)
@@ -477,7 +485,7 @@ def test_sparse_attn_decode_split_k_kernel(
485
rows = [[1, 3, 0, 5, 2, 4], [3, 0, 6]]
486
extra_kv = torch.randn(7, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
487
extra_rows = rows
480
- extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size)
488
+ extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, is_extra=True)
489
extra_indices, extra_indptr = _ragged_from_rows(rows, device)
490
491
attn_sink = (
tests/kernels/attention/test_triton_unified_attention.py
+1
-5
@@ -18,11 +18,7 @@ HEAD_SIZES = [128, 256]
18
BLOCK_SIZES = [16]
19
20
DTYPES = [torch.bfloat16]
21
-QDTYPES = (
22
- [None, torch.float8_e4m3fn]
23
- if not current_platform.is_rocm()
24
- else [None, torch.float8_e4m3fnuz]
25
-)
21
+QDTYPES = [None, current_platform.fp8_dtype()]
22
FP8_DTYPE = current_platform.fp8_dtype()
23
24
# one value large enough to test overflow in index calculation.