[CI/Build][ROCm] Keep the CUDA-only kernel tests out of the ROCm run (#53113)
Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com>
stefankoncarevic committed
Aug 20, 2026 at 20:40 UTC
3b829cf176aa4ef357b7edd61eb9352419a8adbd
2 files changed
+10
-3
tests/kernels/test_fp32_router_gemm.py
+5
-2
@@ -16,6 +16,7 @@ import pytest
16
import torch
17
18
from vllm._custom_ops import fp32_router_gemm
19
+from vllm.platforms import current_platform
20
21
# (hidden_size, num_experts)
22
SHAPES = [(3072, 256), (6144, 128), (6144, 256)]
@@ -26,8 +27,10 @@ ATOL_BF16 = 2e-2 # bf16 activation has lower precision
27
28
29
def _requires_sm90():
29
- if not torch.cuda.is_available():
30
- pytest.skip("CUDA not available")
30
+ # ROCm reports a CUDA-like device capability (gfx950 -> (9, 5)), which would
31
+ # pass the SM90 check below for a kernel that is only built for CUDA.
32
+ if not current_platform.is_cuda():
33
+ pytest.skip("fp32_router_gemm is built for CUDA only")
34
major, minor = torch.cuda.get_device_capability()
35
if major * 10 + minor < 90:
36
pytest.skip(f"fp32_router_gemm requires SM90+, got SM{major}{minor}")
tests/kernels/test_kimi_k3_gemm_rs.py
+5
-1
@@ -12,7 +12,6 @@ from vllm.distributed.parallel_state import (
12
init_distributed_environment,
13
initialize_model_parallel,
14
)
15
-from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import GemmRS
15
from vllm.platforms import current_platform
16
from vllm.utils.network_utils import get_open_port
17
from vllm.utils.system_utils import update_environment_variables
@@ -70,6 +69,11 @@ def _assert_valid_rows_close(
69
70
71
def _worker(local_rank: int, world_size: int, master_port: int) -> None:
72
+ # This module pulls in cute_dsl, which is unavailable off CUDA, so importing
73
+ # it at module level would fail collection. Import it where it is used, as
74
+ # `kda.py` does.
75
+ from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import GemmRS
76
+
77
device = torch.device("cuda", local_rank)
78
torch.accelerator.set_device_index(device)
79
update_environment_variables(