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 10, 04:00 UTC
AMD CI Test - unskip moe_sum test and moe_align_block_size tests (#32039)
Signed-off-by: Hongxia Yang <hongxia.yang@amd.com>
Hongxia Yang committed
Jan 14, 2026 at 02:25 UTC
048bb597286cc7cb1e30e03b366aef1e2b393b39
3 files changed
+11
-4
tests/kernels/moe/test_moe.py
-1
@@ -1161,7 +1161,6 @@ def test_batched_moe_align_block_size_opcheck():
1161
@pytest.mark.parametrize("topk", TOP_KS)
1162
@pytest.mark.parametrize("k", [128, 511, 1024])
1163
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
1164
-@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
1164
def test_moe_sum(m: int, topk: int, k: int, dtype: torch.dtype):
1165
input = torch.randn((m, topk, k), device="cuda", dtype=dtype)
1166
actual = torch.empty((m, k), device="cuda", dtype=dtype)
tests/kernels/moe/test_moe_align_block_size.py
-3
@@ -12,7 +12,6 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
12
batched_moe_align_block_size,
13
moe_align_block_size,
14
)
15
-from vllm.platforms import current_platform
15
from vllm.utils.math_utils import round_up
16
from vllm.utils.torch_utils import set_random_seed
17
@@ -185,7 +184,6 @@ def torch_moe_align_block_size(
184
@pytest.mark.parametrize("num_experts", NUM_EXPERTS)
185
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
186
@pytest.mark.parametrize("pad_sorted_ids", [False, True])
188
-@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
187
def test_moe_align_block_size(
188
m: int, topk: int, num_experts: int, block_size: int, pad_sorted_ids: bool
189
):
@@ -245,7 +243,6 @@ def test_moe_align_block_size(
243
@pytest.mark.parametrize("topk", [2, 4])
244
@pytest.mark.parametrize("num_experts", [8, 64])
245
@pytest.mark.parametrize("block_size", [64])
248
-@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
246
def test_moe_align_block_size_with_expert_map(
247
m: int, topk: int, num_experts: int, block_size: int
248
):
vllm/platforms/rocm.py
+11
@@ -187,6 +187,17 @@ class RocmPlatform(Platform):
187
if not on_gfx9():
188
supported_quantization += ["bitsandbytes"]
189
190
+ @classmethod
191
+ def import_kernels(cls) -> None:
192
+ """Import ROCm-specific kernels."""
193
+ super().import_kernels()
194
+
195
+ import contextlib
196
+
197
+ # Import ROCm-specific extension
198
+ with contextlib.suppress(ImportError):
199
+ import vllm._rocm_C # noqa: F401
200
+
201
@classmethod
202
def get_attn_backend_cls(
203
cls,