diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index 5b8dfab25b..2d3e48d435 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -125,7 +125,6 @@ def test_mxfp4_loading_and_execution_moe(vllm_runner, model_case: ModelCase): # if model_case.model_id == "fxmarty/qwen_1.5-moe-a2.7b-mxfp4": # llm.apply_model(check_model) - output = llm.generate_greedy("Today I am in the French Alps and", max_tokens=20) assert output @@ -187,6 +186,16 @@ def mxfp8_dequantize(x, scale): return x_float * scale +def mxfp4_quant_dequant(x: torch.Tensor) -> torch.Tensor: + shape = x.shape + quantized, scale = dynamic_mxfp4_quant(x.to(torch.bfloat16).flatten(0, -2)) + return ( + upcast_from_mxfp(quantized.view(torch.uint8), scale, torch.bfloat16, axis=-1) + .reshape(shape) + .float() + ) + + def reference_moe( roouting_logits, topk, @@ -219,6 +228,8 @@ def reference_moe( expert_weights = torch.nn.functional.softmax(experts.values, dim=1) expert_indices = experts.indices t = hidden_states.clone() + if act_type == "mxfp4": + t = mxfp4_quant_dequant(t) # MLP #1 mlp1_weight = w13[expert_indices, ...] mlp1_bias = bias13[expert_indices, ...] @@ -244,6 +255,8 @@ def reference_moe( t.to(torch.bfloat16), is_sf_swizzled_layout=False ) t = mxfp8_dequantize(t_quantized, t_scale) + elif act_type == "mxfp4": + t = mxfp4_quant_dequant(t) # MLP #2 mlp2_weight = w2[expert_indices, ...] mlp2_bias = bias2[expert_indices, ...] @@ -1392,6 +1405,14 @@ ROCM_BACKEND_CONFIGS = { "requires_aiter": True, "requires_gfx950": True, }, + "AITER_MXFP4_MXFP4": { + "activation": "SILU", + "act_type": "mxfp4", + "rtol": 1.0, + "percent": 0.8, + "requires_aiter": True, + "requires_gfx950": True, + }, } @@ -1662,7 +1683,7 @@ def test_rocm_mxfp4_moe_oracle( alpha=1.702 if activation == MoEActivation.SWIGLUOAI else 1.0, beta=1.0 if activation == MoEActivation.SWIGLUOAI else 0.0, limit=7.0 if activation == MoEActivation.SWIGLUOAI else None, - act_type="bf16", + act_type=str(config.get("act_type", "bf16")), activation=act_name, use_interleaved_layout=use_interleaved, ) diff --git a/tests/quantization/test_gfx950_moe.py b/tests/quantization/test_gfx950_moe.py index c8d34bb0ab..a4dac1c163 100644 --- a/tests/quantization/test_gfx950_moe.py +++ b/tests/quantization/test_gfx950_moe.py @@ -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."""