[ROCm][CI] Expand AITER W4A4 MoE Coverage (#52647)

Signed-off-by: Micah Williamson <micah.williamson@amd.com>
This commit is contained in:
Micah Williamson
2026-08-17 15:46:22 -05:00
committed by GitHub
parent f08a95f8d8
commit 8878ebd8fd
2 changed files with 59 additions and 7 deletions
+36 -5
View File
@@ -25,14 +25,32 @@ from vllm.platforms import current_platform
ROCM_AVAILABLE = current_platform.is_rocm()
ROCM_GFX950 = False
ROCM_AITER_AVAILABLE = False
ROCM_AITER_SUPPORTED = False
if ROCM_AVAILABLE:
from vllm._aiter_ops import rocm_aiter_ops
from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
from vllm.platforms.rocm import on_gfx950
ROCM_GFX950 = on_gfx950()
ROCM_AITER_AVAILABLE = rocm_aiter_ops.is_fused_moe_enabled()
ROCM_AITER_SUPPORTED = is_aiter_found_and_supported()
def set_rocm_aiter(monkeypatch: pytest.MonkeyPatch, enabled: bool) -> None:
value = "1" if enabled else "0"
monkeypatch.setenv("VLLM_ROCM_USE_AITER", value)
monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", value)
monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", enabled)
monkeypatch.setattr(rocm_aiter_ops, "_FMOE_ENABLED", enabled)
@pytest.fixture
def enable_rocm_aiter(monkeypatch: pytest.MonkeyPatch):
set_rocm_aiter(monkeypatch, True)
@pytest.fixture
def disable_rocm_aiter(monkeypatch: pytest.MonkeyPatch):
set_rocm_aiter(monkeypatch, False)
def _make_w4a4_moe_config(moe_backend: str = "auto") -> FusedMoEConfig:
@@ -68,8 +86,8 @@ def mxfp4_oracle_config():
@pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)")
@pytest.mark.skipif(not ROCM_AITER_AVAILABLE, reason="Requires AITER enabled")
def test_w4a4_dispatches_to_aiter(mxfp4_oracle_config):
@pytest.mark.skipif(not ROCM_AITER_SUPPORTED, reason="Requires supported AITER")
def test_w4a4_dispatches_to_aiter(mxfp4_oracle_config, enable_rocm_aiter):
"""With AITER enabled + GFX950, W4A4 selects AITER_MXFP4_MXFP4."""
config = _make_w4a4_moe_config()
backend, experts_cls = select_mxfp4_moe_backend(
@@ -79,6 +97,19 @@ def test_w4a4_dispatches_to_aiter(mxfp4_oracle_config):
assert experts_cls is not None
@pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)")
def test_w4a4_falls_back_without_aiter(
mxfp4_oracle_config,
disable_rocm_aiter,
):
config = _make_w4a4_moe_config()
backend, experts_cls = select_mxfp4_moe_backend(
config, activation_key=kMxfp4Dynamic
)
assert backend == Mxfp4MoeBackend.EMULATION
assert experts_cls is not None
@pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)")
def test_w4a4_dispatches_to_emulation_with_moe_backend(mxfp4_oracle_config):
"""With --moe-backend emulation, W4A4 selects EMULATION."""