A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 07:00 UTC
[ROCm][CI] Guard CudaPlatform/RocmPlatform imports to fix test collection on cross-platform builds (#37617)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Mar 20, 2026 at 22:58 UTC
3ffa52009f35c4398a86b7cdd83d4031bf19651c
1 file changed
+35
-3
tests/kernels/attention/test_attention_selector.py
+35
-3
@@ -14,8 +14,19 @@ from vllm.config import (
14
)
15
from vllm.platforms import current_platform
16
from vllm.platforms.cpu import CpuPlatform
17
-from vllm.platforms.cuda import CudaPlatform
18
-from vllm.platforms.rocm import RocmPlatform
17
+
18
+# CudaPlatform and RocmPlatform import their respective compiled C extensions
19
+# at module level, raising ModuleNotFoundError on incompatible builds.
20
+try:
21
+ from vllm.platforms.cuda import CudaPlatform
22
+except (ImportError, ModuleNotFoundError):
23
+ CudaPlatform = None
24
+
25
+try:
26
+ from vllm.platforms.rocm import RocmPlatform
27
+except (ImportError, ModuleNotFoundError):
28
+ RocmPlatform = None
29
+
30
from vllm.v1.attention.backends.registry import AttentionBackendEnum
31
from vllm.v1.attention.selector import _cached_get_attn_backend, get_attn_backend
32
@@ -101,6 +112,8 @@ def test_backend_selection(
112
assert backend.get_name() == "CPU_ATTN"
113
114
elif device == "hip":
115
+ if RocmPlatform is None:
116
+ pytest.skip("RocmPlatform not available")
117
with patch("vllm.platforms.current_platform", RocmPlatform()):
118
if use_mla:
119
# ROCm MLA backend logic:
@@ -126,6 +139,8 @@ def test_backend_selection(
139
assert backend.get_name() == expected
140
141
elif device == "cuda":
142
+ if CudaPlatform is None:
143
+ pytest.skip("CudaPlatform not available")
144
with patch("vllm.platforms.current_platform", CudaPlatform()):
145
capability = torch.cuda.get_device_capability()
146
if use_mla:
@@ -214,7 +229,7 @@ def test_backend_selection(
229
assert backend.get_name() == expected
230
231
217
-@pytest.mark.parametrize("device", ["cpu", "cuda"])
232
+@pytest.mark.parametrize("device", ["cpu", "cuda", "hip"])
233
def test_fp32_fallback(device: str):
234
"""Test attention backend selection with fp32."""
235
# Use default config (no backend specified)
@@ -227,10 +242,25 @@ def test_fp32_fallback(device: str):
242
assert backend.get_name() == "CPU_ATTN"
243
244
elif device == "cuda":
245
+ if CudaPlatform is None:
246
+ pytest.skip("CudaPlatform not available")
247
with patch("vllm.platforms.current_platform", CudaPlatform()):
248
backend = get_attn_backend(16, torch.float32, None)
249
assert backend.get_name() == "FLEX_ATTENTION"
250
251
+ elif device == "hip":
252
+ if RocmPlatform is None:
253
+ pytest.skip("RocmPlatform not available")
254
+ # ROCm backends do not support head_size=16 (minimum is 32).
255
+ # No known HuggingFace transformer model uses head_size=16.
256
+ # Revisit if a real model with this head size is identified
257
+ # and accuracy-tested.
258
+ with (
259
+ patch("vllm.platforms.current_platform", RocmPlatform()),
260
+ pytest.raises(ValueError, match="No valid attention backend"),
261
+ ):
262
+ get_attn_backend(16, torch.float32, None)
263
+
264
265
def test_flash_attn(monkeypatch: pytest.MonkeyPatch):
266
"""Test FlashAttn validation."""
@@ -367,6 +397,8 @@ def test_per_head_quant_scales_backend_selection(
397
attention_config=attention_config, cache_config=cache_config
398
)
399
400
+ if CudaPlatform is None:
401
+ pytest.skip("CudaPlatform not available")
402
with (
403
set_current_vllm_config(vllm_config),
404
patch("vllm.platforms.current_platform", CudaPlatform()),