Skip to content

[CI/Build] Restrict ERNIE fused RoPE tests to NVIDIA CUDA #7771

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yenuo26 merged 2 commits into from
Sep 19, 2026
Merged

[CI/Build] Restrict ERNIE fused RoPE tests to NVIDIA CUDA #7771

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
19ea2fc
[CI/Build] Restrict ERNIE fused RoPE tests to NVIDIA CUDA
andyluo7 Sep 18, 2026
9482260
[CI/Build] Mark ERNIE fused RoPE hardware requirements
andyluo7 Sep 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 14 additions & 6 deletions tests/diffusion/models/ernie_image/test_ernie_image_fused_rope.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +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)
Expand Down Expand Up @@ -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",
Expand All @@ -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
Expand All @@ -85,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

Expand All @@ -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
Expand All @@ -116,7 +122,8 @@ 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")
@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):
Expand All @@ -142,7 +149,8 @@ def unexpected_launch(*args):
)


@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
Expand Down
Loading