[ROCm][CI] Speed up `test_rocm_aiter_qk_norm_rope_kvcache_fusion` (#53004)
Signed-off-by: Micah Williamson <micah.williamson@amd.com> Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Micah Williamson committed
Aug 19, 2026 at 21:20 UTC
fbb4c04db42ff0c6a98af909321bf841798a20c2
1 file changed
+30
-20
tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py
+30
-20
@@ -2,9 +2,11 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import os
5
+from importlib.metadata import version
6
7
import pytest
8
import torch
9
+from packaging.version import Version
10
11
import vllm.config
12
from tests.compile.backend import TestBackend
@@ -40,10 +42,27 @@ from vllm.v1.attention.backend import (
42
from vllm.v1.attention.backends.registry import AttentionBackendEnum
43
from vllm.v1.kv_cache_interface import AttentionSpec
44
45
+pytestmark = pytest.mark.skip_global_cleanup
46
+
47
+if not is_aiter_found_and_supported():
48
+ pytest.skip(
49
+ "ROCm with supported AITER is required",
50
+ allow_module_level=True,
51
+ )
52
+
53
INDEX_SELECT_OP = torch.ops.aten.index.Tensor
54
FP8_DTYPE = current_platform.fp8_dtype()
55
56
57
+@pytest.fixture(scope="module", autouse=True)
58
+def module_global_cleanup():
59
+ # Cleanup once at the end of the module
60
+ from vllm.distributed import cleanup_dist_env_and_memory
61
+
62
+ yield
63
+ cleanup_dist_env_and_memory()
64
+
65
+
66
class QKNormRoPEKVCacheTestModel(torch.nn.Module):
67
"""Minimal model that reproduces the QK-norm + RoPE + KV cache update
68
pattern matched by QkNormRopeKvCacheFusionPass:
@@ -416,8 +435,7 @@ _FUSION_CONFIGS = [
435
AttentionBackendEnum.ROCM_AITER_FA,
436
],
437
)
419
-@pytest.mark.parametrize("num_tokens", [5, 16, 2048])
420
-@pytest.mark.parametrize("use_shuffle_kv_layout", ["1", "0"])
438
+@pytest.mark.parametrize("num_tokens", [5, 2048])
439
@pytest.mark.parametrize(
440
"kv_stride_order",
441
# FA/unified use the 4D packed cache (num_blocks, num_kv_heads, block_size,
@@ -428,11 +446,17 @@ _FUSION_CONFIGS = [
446
@pytest.mark.parametrize("block_size", [16, 32, 64])
447
@pytest.mark.parametrize("dtype", [torch.bfloat16])
448
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
431
-@pytest.mark.parametrize("rms_norm_eps", [1e-5, 1e-6])
449
+@pytest.mark.parametrize("rms_norm_eps", [1e-6])
450
@pytest.mark.parametrize("custom_op", ["+rotary_embedding", "+rms_norm"])
451
+# The fused_qk_norm_rope_cache kernel used by this test aborts on AITER < 0.1.20
452
+# (fused_qk_norm_rope_cache_quant.cu: "k_cache/v_cache must be contiguous within
453
+# a block")
454
@pytest.mark.skipif(
434
- not is_aiter_found_and_supported(),
435
- reason="Only test on ROCm with AITER installed and supported",
455
+ not Version(version("amd_aiter")) >= Version("0.1.20"),
456
+ reason=(
457
+ "Requires AITER >= 0.1.20; older kernels abort on "
458
+ "fused_qk_norm_rope_cache (k_cache/v_cache contiguity)"
459
+ ),
460
)
461
def test_qk_norm_rope_kvcache_fusion(
462
num_tokens: int,
@@ -443,7 +467,6 @@ def test_qk_norm_rope_kvcache_fusion(
467
is_neox: bool,
468
attn_backend: AttentionBackendEnum,
469
enable_aiter_triton_rope: bool,
446
- use_shuffle_kv_layout: str,
470
kv_stride_order: tuple[int, ...],
471
block_size: int,
472
dtype: torch.dtype,
@@ -452,19 +475,6 @@ def test_qk_norm_rope_kvcache_fusion(
475
custom_op: str,
476
monkeypatch: pytest.MonkeyPatch,
477
):
455
- if (
456
- attn_backend == AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN
457
- and use_shuffle_kv_layout == "1"
458
- ):
459
- pytest.skip("ROCM_AITER_UNIFIED_ATTN is NHD-only; shuffle env is ignored")
460
- if (
461
- attn_backend == AttentionBackendEnum.ROCM_AITER_FA
462
- and use_shuffle_kv_layout == "1"
463
- ):
464
- pytest.skip(
465
- "ROCM_AITER_FA gates the qk_norm+rope+kvcache fusion off under shuffle "
466
- "layout (defers to reshape_and_cache_shuffle_triton), so nothing fuses"
467
- )
478
_run_qk_norm_rope_kvcache_fusion_test(
479
attn_backend=attn_backend,
480
enable_aiter_triton_rope=enable_aiter_triton_rope,
@@ -475,7 +485,7 @@ def test_qk_norm_rope_kvcache_fusion(
485
rotary_dim=rotary_dim,
486
block_size=block_size,
487
is_neox=is_neox,
478
- use_shuffle_kv_layout=use_shuffle_kv_layout,
488
+ use_shuffle_kv_layout="0",
489
kv_stride_order=kv_stride_order,
490
dtype=dtype,
491
kv_cache_dtype=kv_cache_dtype,