diff --git a/tests/diffusion/models/ernie_image/test_ernie_image_fused_rope.py b/tests/diffusion/models/ernie_image/test_ernie_image_fused_rope.py index bf3dce4f3..b73804433 100644 --- a/tests/diffusion/models/ernie_image/test_ernie_image_fused_rope.py +++ b/tests/diffusion/models/ernie_image/test_ernie_image_fused_rope.py @@ -5,9 +5,11 @@ import pytest import torch from vllm.triton_utils import HAS_TRITON +from tests.helpers.mark import hardware_test from vllm_omni.diffusion.models.ernie_image.ernie_image_transformer import _apply_rotary_emb +from vllm_omni.platforms import current_omni_platform -pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.diffusion] +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] @pytest.fixture(autouse=True) @@ -35,7 +37,8 @@ def _inputs(shape: tuple[int, int, int, int]): return query, key, freqs_cos, freqs_sin -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="NVIDIA CUDA required") @pytest.mark.skipif(not HAS_TRITON, reason="Triton required") @pytest.mark.parametrize( "shape", @@ -60,7 +63,8 @@ def test_fused_qk_rope_is_bit_exact(shape): assert torch.equal(actual_key, expected_key) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="NVIDIA CUDA required") @pytest.mark.skipif(not HAS_TRITON, reason="Triton required") def test_launch_failure_disables_runtime_key(monkeypatch): from vllm_omni.diffusion.models.ernie_image import fused_rope @@ -85,6 +89,7 @@ def test_launch_failure_disables_runtime_key(monkeypatch): assert len(fused_rope._FAILED_KEYS) == 1 +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) def test_failed_runtime_key_cache_is_bounded(): from vllm_omni.diffusion.models.ernie_image import fused_rope @@ -96,7 +101,8 @@ def test_failed_runtime_key_cache_is_bounded(): assert (fused_rope._FAILED_KEYS_MAX_SIZE,) in fused_rope._FAILED_KEYS -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="NVIDIA CUDA required") @pytest.mark.skipif(not HAS_TRITON, reason="Triton required") def test_compile_path_does_not_launch_fused_kernel(monkeypatch): from vllm_omni.diffusion.models.ernie_image import fused_rope @@ -116,7 +122,8 @@ def test_compile_path_does_not_launch_fused_kernel(monkeypatch): assert fused_rope.try_fused_qk_rotary_emb(query, key, freqs_cos, freqs_sin) is None -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="NVIDIA CUDA required") @pytest.mark.skipif(not HAS_TRITON, reason="Triton required") @pytest.mark.parametrize("requires_grad_input", ["query", "key", "cos", "sin"]) def test_gradient_inputs_do_not_launch_fused_kernel(monkeypatch, requires_grad_input): @@ -142,7 +149,8 @@ def test_gradient_inputs_do_not_launch_fused_kernel(monkeypatch, requires_grad_i ) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="NVIDIA CUDA required") @pytest.mark.skipif(not HAS_TRITON, reason="Triton required") def test_mismatched_key_dtype_or_device_does_not_launch(monkeypatch): from vllm_omni.diffusion.models.ernie_image import fused_rope