From dc7e92c0fa63739ee844a3924d6ef308d363c9d0 Mon Sep 17 00:00:00 2001 From: Matthew Wong Date: Tue, 16 Jun 2026 07:35:19 +0000 Subject: [PATCH 1/5] Fix platform detection in kernels/attention/test_attention_selector Signed-off-by: Matthew Wong --- .../attention/test_attention_selector.py | 29 +++++++++++-------- 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/tests/kernels/attention/test_attention_selector.py b/tests/kernels/attention/test_attention_selector.py index db4dcc8a636e..447b502293b2 100644 --- a/tests/kernels/attention/test_attention_selector.py +++ b/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( From 30337eb267c254f5840fe4b35a5ffef80e0072f1 Mon Sep 17 00:00:00 2001 From: Matthew Wong Date: Tue, 16 Jun 2026 09:16:11 +0000 Subject: [PATCH 2/5] Use right FP8 dtype in kernels/attention/test_triton_unified_attention Signed-off-by: Matthew Wong --- tests/kernels/attention/test_triton_unified_attention.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/tests/kernels/attention/test_triton_unified_attention.py b/tests/kernels/attention/test_triton_unified_attention.py index 6440ba3156e3..d3435ea665db 100644 --- a/tests/kernels/attention/test_triton_unified_attention.py +++ b/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. From fe1f22d72cd84cbbb15b3c5afb141eb99c6a520f Mon Sep 17 00:00:00 2001 From: Matthew Wong Date: Tue, 16 Jun 2026 09:15:09 +0000 Subject: [PATCH 3/5] Use Math SDPA backend for increased accuracy in kernels/attention/test_prefix_prefill Signed-off-by: Matthew Wong --- .../kernels/attention/test_prefix_prefill.py | 26 ++++++++++++------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/tests/kernels/attention/test_prefix_prefill.py b/tests/kernels/attention/test_prefix_prefill.py index de63b4548f2d..f1c591fb671b 100644 --- a/tests/kernels/attention/test_prefix_prefill.py +++ b/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) From 7c977e3a51d9bad22fd75173e233ab3ca5b16f56 Mon Sep 17 00:00:00 2001 From: Matthew Wong Date: Tue, 16 Jun 2026 19:41:47 +0000 Subject: [PATCH 4/5] Promote Kernels Attention (MI300) to gating and track MI355 more Signed-off-by: Matthew Wong --- .buildkite/test-amd.yaml | 8 ++++---- .buildkite/test_areas/kernels.yaml | 14 ++++++++++++++ 2 files changed, 18 insertions(+), 4 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 5550c0a0c185..e529045a8265 100644 --- a/.buildkite/test-amd.yaml +++ b/.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/ diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index ebcb95a9d821..7bad72d8ebce 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.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 From a29cb0ca7003fed9e69c10419675b58508f201a3 Mon Sep 17 00:00:00 2001 From: Matthew Wong Date: Sun, 21 Jun 2026 02:14:10 +0000 Subject: [PATCH 5/5] Fix reference sparse MLA decode ragged in tests Signed-off-by: Matthew Wong --- .../attention/test_rocm_triton_attn_dsv4.py | 22 +++++++++++++------ 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index daf73b82e614..e00726f64d80 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/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 = (