[CI/Build] Restrict ERNIE fused RoPE tests to NVIDIA CUDA (#7771)
Signed-off-by: andyluo7 <andy.luo@amd.com>
This commit is contained in:
1 parent
4c7a98c26f
commit
232dbc82f8
1 file changed
+14
-6
@@ -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
|
||||
|
||||
Reference in new issue
Block a user