diff --git a/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py b/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py index 900f90eb42..e385999db3 100644 --- a/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py +++ b/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py @@ -2,9 +2,11 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os +from importlib.metadata import version import pytest import torch +from packaging.version import Version import vllm.config from tests.compile.backend import TestBackend @@ -40,10 +42,27 @@ from vllm.v1.attention.backend import ( from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.kv_cache_interface import AttentionSpec +pytestmark = pytest.mark.skip_global_cleanup + +if not is_aiter_found_and_supported(): + pytest.skip( + "ROCm with supported AITER is required", + allow_module_level=True, + ) + INDEX_SELECT_OP = torch.ops.aten.index.Tensor FP8_DTYPE = current_platform.fp8_dtype() +@pytest.fixture(scope="module", autouse=True) +def module_global_cleanup(): + # Cleanup once at the end of the module + from vllm.distributed import cleanup_dist_env_and_memory + + yield + cleanup_dist_env_and_memory() + + class QKNormRoPEKVCacheTestModel(torch.nn.Module): """Minimal model that reproduces the QK-norm + RoPE + KV cache update pattern matched by QkNormRopeKvCacheFusionPass: @@ -416,8 +435,7 @@ _FUSION_CONFIGS = [ AttentionBackendEnum.ROCM_AITER_FA, ], ) -@pytest.mark.parametrize("num_tokens", [5, 16, 2048]) -@pytest.mark.parametrize("use_shuffle_kv_layout", ["1", "0"]) +@pytest.mark.parametrize("num_tokens", [5, 2048]) @pytest.mark.parametrize( "kv_stride_order", # FA/unified use the 4D packed cache (num_blocks, num_kv_heads, block_size, @@ -428,11 +446,17 @@ _FUSION_CONFIGS = [ @pytest.mark.parametrize("block_size", [16, 32, 64]) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) -@pytest.mark.parametrize("rms_norm_eps", [1e-5, 1e-6]) +@pytest.mark.parametrize("rms_norm_eps", [1e-6]) @pytest.mark.parametrize("custom_op", ["+rotary_embedding", "+rms_norm"]) +# The fused_qk_norm_rope_cache kernel used by this test aborts on AITER < 0.1.20 +# (fused_qk_norm_rope_cache_quant.cu: "k_cache/v_cache must be contiguous within +# a block") @pytest.mark.skipif( - not is_aiter_found_and_supported(), - reason="Only test on ROCm with AITER installed and supported", + not Version(version("amd_aiter")) >= Version("0.1.20"), + reason=( + "Requires AITER >= 0.1.20; older kernels abort on " + "fused_qk_norm_rope_cache (k_cache/v_cache contiguity)" + ), ) def test_qk_norm_rope_kvcache_fusion( num_tokens: int, @@ -443,7 +467,6 @@ def test_qk_norm_rope_kvcache_fusion( is_neox: bool, attn_backend: AttentionBackendEnum, enable_aiter_triton_rope: bool, - use_shuffle_kv_layout: str, kv_stride_order: tuple[int, ...], block_size: int, dtype: torch.dtype, @@ -452,19 +475,6 @@ def test_qk_norm_rope_kvcache_fusion( custom_op: str, monkeypatch: pytest.MonkeyPatch, ): - if ( - attn_backend == AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN - and use_shuffle_kv_layout == "1" - ): - pytest.skip("ROCM_AITER_UNIFIED_ATTN is NHD-only; shuffle env is ignored") - if ( - attn_backend == AttentionBackendEnum.ROCM_AITER_FA - and use_shuffle_kv_layout == "1" - ): - pytest.skip( - "ROCM_AITER_FA gates the qk_norm+rope+kvcache fusion off under shuffle " - "layout (defers to reshape_and_cache_shuffle_triton), so nothing fuses" - ) _run_qk_norm_rope_kvcache_fusion_test( attn_backend=attn_backend, enable_aiter_triton_rope=enable_aiter_triton_rope, @@ -475,7 +485,7 @@ def test_qk_norm_rope_kvcache_fusion( rotary_dim=rotary_dim, block_size=block_size, is_neox=is_neox, - use_shuffle_kv_layout=use_shuffle_kv_layout, + use_shuffle_kv_layout="0", kv_stride_order=kv_stride_order, dtype=dtype, kv_cache_dtype=kv_cache_dtype,