From 19ea2fc36dc69ce1823dc3a5614bce6cfefc617d Mon Sep 17 00:00:00 2001 From: andyluo7 Date: Thu, 17 Sep 2026 21:31:19 -0700 Subject: [PATCH 1/2] [CI/Build] Restrict ERNIE fused RoPE tests to NVIDIA CUDA Signed-off-by: andyluo7 --- .../models/ernie_image/test_ernie_image_fused_rope.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) 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 bf3dce4f326..e6b171ad0dd 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 @@ -6,6 +6,7 @@ from vllm.triton_utils import HAS_TRITON 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] @@ -35,7 +36,7 @@ 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") +@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 +61,7 @@ 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") +@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 @@ -96,7 +97,7 @@ 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") +@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 +117,7 @@ def unexpected_launch(*args): 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") +@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 +143,7 @@ def unexpected_launch(*args): ) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@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 From 94822606be6d77dcdc2e1b432cb8eb122b1a5821 Mon Sep 17 00:00:00 2001 From: andyluo7 Date: Fri, 18 Sep 2026 08:02:48 -0700 Subject: [PATCH 2/2] [CI/Build] Mark ERNIE fused RoPE hardware requirements Signed-off-by: andyluo7 --- .../models/ernie_image/test_ernie_image_fused_rope.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) 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 e6b171ad0dd..b73804433d1 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,10 +5,11 @@ 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) @@ -36,6 +37,7 @@ def _inputs(shape: tuple[int, int, int, int]): return query, key, freqs_cos, freqs_sin +@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( @@ -61,6 +63,7 @@ def test_fused_qk_rope_is_bit_exact(shape): assert torch.equal(actual_key, expected_key) +@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): @@ -86,6 +89,7 @@ def failing_launch(*args): 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 @@ -97,6 +101,7 @@ def test_failed_runtime_key_cache_is_bounded(): assert (fused_rope._FAILED_KEYS_MAX_SIZE,) in fused_rope._FAILED_KEYS +@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): @@ -117,6 +122,7 @@ def unexpected_launch(*args): assert fused_rope.try_fused_qk_rotary_emb(query, key, freqs_cos, freqs_sin) is None +@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"]) @@ -143,6 +149,7 @@ def unexpected_launch(*args): ) +@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):