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 8, 21:30 UTC
[CI][AMD][BugFix] Ensure VLLM_ROCM_USE_AITER is set so test_rocm_aiter_topk.py can run correctly (#33840)
Signed-off-by: Randall Smith <Randall.Smith@amd.com>
rasmith committed
Feb 4, 2026 at 23:05 UTC
c1395f72cd22d97eb39ecd67d9d22f2af3d20bda
1 file changed
+11
-5
tests/kernels/moe/test_rocm_aiter_topk.py
+11
-5
@@ -10,22 +10,28 @@
10
# and the platform is not ROCm.
11
12
import importlib.util
13
+import os
14
15
import pytest
16
import torch
17
18
+from vllm.platforms import current_platform
19
+
20
+if not current_platform.is_rocm():
21
+ pytest.skip("This test can only run on ROCm.", allow_module_level=True)
22
+
23
+# This environment variable must be set so ops will be registered.
24
+os.environ["VLLM_ROCM_USE_AITER"] = "1"
25
+
26
# this import statement is needed to ensure the ops are registered
27
import vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe # noqa: F401
19
-from vllm.platforms import current_platform
28
29
# need to import once to ensure the ops are registered
30
# Check if aiter package is installed
31
aiter_available = importlib.util.find_spec("aiter") is not None
32
25
-pytestmark = pytest.mark.skipif(
26
- not (current_platform.is_rocm() and aiter_available),
27
- reason="AITER ops are only available on ROCm with aiter package installed",
28
-)
33
+if not aiter_available:
34
+ pytest.skip("These tests require AITER to run.", allow_module_level=True)
35
36
37
def test_rocm_aiter_biased_grouped_topk_custom_op_registration():