[ROCm][CI] Expand AITER W4A4 MoE Coverage (#52647)
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
This commit is contained in:
@@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user