[ROCM][CI] Attention test speedup (#52763)
Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>
stefankoncarevic committed
Aug 18, 2026 at 23:40 UTC
6066bb3d502107d27383a64b306d36148421d4b9
9 files changed
+24
-6
tests/kernels/attention/test_attention.py
+2
@@ -14,6 +14,8 @@ from vllm.platforms import current_platform
14
from vllm.utils.mem_utils import get_max_shared_memory_bytes
15
from vllm.utils.torch_utils import set_random_seed
16
17
+pytestmark = pytest.mark.skip_global_cleanup
18
+
19
FLOAT32_BYTES = torch.finfo(torch.float).bits // 8
20
# This will change depending on the compute capability.
21
# - 512 as a buffer
tests/kernels/attention/test_cache.py
+2
@@ -12,6 +12,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import scaled_deq
12
from vllm.platforms import current_platform
13
from vllm.utils.torch_utils import nvfp4_split_data_scale, set_random_seed
14
15
+pytestmark = pytest.mark.skip_global_cleanup
16
+
17
COPYING_DIRECTION = [("cuda", "cpu"), ("cuda", "cuda"), ("cpu", "cuda")]
18
DTYPES = [torch.bfloat16, torch.float]
19
NUM_TOKENS = [42] # Arbitrary values for testing
tests/kernels/attention/test_cutlass_mla_decode.py
+2
@@ -11,6 +11,8 @@ from vllm.platforms import current_platform
11
from vllm.triton_utils import triton
12
from vllm.utils.platform_utils import num_compute_units
13
14
+pytestmark = pytest.mark.skip_global_cleanup
15
+
16
17
def cal_diff(
18
x: torch.Tensor,
tests/kernels/attention/test_flashmla.py
+2
@@ -16,6 +16,8 @@ from vllm.v1.attention.ops.flashmla import (
16
is_flashmla_dense_supported,
17
)
18
19
+pytestmark = pytest.mark.skip_global_cleanup
20
+
21
22
def cal_diff(
23
x: torch.Tensor, y: torch.Tensor, name: str, use_fp8: bool = False
tests/kernels/attention/test_merge_attn_states.py
+8
-6
@@ -15,6 +15,14 @@ from vllm.v1.attention.ops.triton_merge_attn_states import (
15
merge_attn_states as merge_attn_states_triton,
16
)
17
18
+pytestmark = [
19
+ pytest.mark.skip_global_cleanup,
20
+ pytest.mark.skipif(
21
+ not current_platform.is_cuda_alike(),
22
+ reason="merge_attn_states kernels require CUDA or ROCm.",
23
+ ),
24
+]
25
+
26
27
# Naive PyTorch Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
28
# can be used to combine partial attention results (in the split-KV case)
@@ -153,12 +161,6 @@ def test_merge_attn_states(
161
input_dtype: torch.dtype,
162
use_fp8: bool,
163
):
156
- if not current_platform.is_cuda():
157
- pytest.skip(
158
- "Currently only support compare triton merge_attn_states "
159
- "with custom cuda merge_attn_states kernel"
160
- )
161
-
164
NUM_TOKENS = num_tokens
165
NUM_HEADS = num_query_heads
166
HEAD_SIZE = head_size
tests/kernels/attention/test_prefix_prefill.py
+2
@@ -17,6 +17,8 @@ from vllm.v1.attention.ops.chunked_prefill_paged_decode import (
17
)
18
from vllm.v1.attention.ops.prefix_prefill import context_attention_fwd
19
20
+pytestmark = pytest.mark.skip_global_cleanup
21
+
22
NUM_HEADS = [64]
23
NUM_QUERIES_PER_KV = [1, 64]
24
HEAD_SIZES = [24, 128]
tests/kernels/attention/test_triton_decode_attention.py
+2
@@ -8,6 +8,8 @@ from vllm.platforms import current_platform
8
from vllm.utils.math_utils import cdiv
9
from vllm.v1.attention.ops.triton_decode_attention import decode_attention_fwd
10
11
+pytestmark = pytest.mark.skip_global_cleanup
12
+
13
DEVICE_TYPE = current_platform.device_type
14
15
tests/kernels/attention/test_triton_unified_attention.py
+2
@@ -11,6 +11,8 @@ from vllm.utils.torch_utils import set_random_seed
11
from vllm.v1.attention.ops.triton_unified_attention import unified_attention
12
from vllm.v1.kv_cache_interface import KVQuantMode
13
14
+pytestmark = pytest.mark.skip_global_cleanup
15
+
16
DEVICE_TYPE = current_platform.device_type
17
18
NUM_HEADS = [(4, 4), (8, 2), (5, 1)]
tests/kernels/attention/test_triton_unified_attention_diffkv.py
+2
@@ -21,6 +21,8 @@ from vllm.v1.attention.ops.triton_unified_attention_diffkv import (
21
unified_attention_diffkv,
22
)
23
24
+pytestmark = pytest.mark.skip_global_cleanup
25
+
26
DEVICE_TYPE = current_platform.device_type
27
28
# (num_query_heads, num_kv_heads): MHA, GQA, and the num_kv_heads==1