| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for MXFP4 MoE oracle backend selection on mi355x (GFX950). |
| 4 | |
| 5 | These tests run on real hardware — no mocks. Skipped on non-GFX950 platforms. |
| 6 | """ |
| 7 | |
| 8 | import pytest |
| 9 | import torch |
| 10 | |
| 11 | from vllm.config.kernel import MoEBackend |
| 12 | from vllm.model_executor.layers.fused_moe.config import ( |
| 13 | FusedMoEConfig, |
| 14 | FusedMoEParallelConfig, |
| 15 | RoutingMethodType, |
| 16 | ) |
| 17 | from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( |
| 18 | Mxfp4MoeBackend, |
| 19 | _requires_qwen38_tep8_emulation, |
| 20 | select_mxfp4_moe_backend, |
| 21 | ) |
| 22 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 23 | kMxfp4Dynamic, |
| 24 | ) |
| 25 | from vllm.platforms import current_platform |
| 26 | |
| 27 | ROCM_AVAILABLE = current_platform.is_rocm() |
| 28 | ROCM_GFX950 = False |
| 29 | ROCM_AITER_SUPPORTED = False |
| 30 | |
| 31 | if ROCM_AVAILABLE: |
| 32 | from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops |
| 33 | from vllm.platforms.rocm import on_gfx950 |
| 34 | |
| 35 | ROCM_GFX950 = on_gfx950() |
| 36 | ROCM_AITER_SUPPORTED = is_aiter_found_and_supported() |
| 37 | |
| 38 | |
| 39 | def set_rocm_aiter(monkeypatch: pytest.MonkeyPatch, enabled: bool) -> None: |
| 40 | value = "1" if enabled else "0" |
| 41 | monkeypatch.setenv("VLLM_ROCM_USE_AITER", value) |
| 42 | monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", value) |
| 43 | monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", enabled) |
| 44 | monkeypatch.setattr(rocm_aiter_ops, "_FMOE_ENABLED", enabled) |
| 45 | |
| 46 | |
| 47 | @pytest.fixture |
| 48 | def enable_rocm_aiter(monkeypatch: pytest.MonkeyPatch): |
| 49 | set_rocm_aiter(monkeypatch, True) |
| 50 | |
| 51 | |
| 52 | @pytest.fixture |
| 53 | def disable_rocm_aiter(monkeypatch: pytest.MonkeyPatch): |
| 54 | set_rocm_aiter(monkeypatch, False) |
| 55 | |
| 56 | |
| 57 | def _make_w4a4_moe_config(moe_backend: MoEBackend = "auto") -> FusedMoEConfig: |
| 58 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 59 | |
| 60 | return FusedMoEConfig( |
| 61 | num_experts=8, |
| 62 | experts_per_token=2, |
| 63 | hidden_dim=256, |
| 64 | intermediate_size=256, |
| 65 | num_local_experts=8, |
| 66 | num_logical_experts=8, |
| 67 | moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), |
| 68 | activation=MoEActivation.SILU, |
| 69 | in_dtype=torch.bfloat16, |
| 70 | device="cuda", |
| 71 | routing_method=RoutingMethodType.Renormalize, |
| 72 | moe_backend=moe_backend, |
| 73 | ) |
| 74 | |
| 75 | |
| 76 | def _make_qwen38_tep8_moe_config() -> FusedMoEConfig: |
| 77 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 78 | |
| 79 | return FusedMoEConfig( |
| 80 | num_experts=512, |
| 81 | experts_per_token=9, |
| 82 | hidden_dim=2560, |
| 83 | intermediate_size=640, |
| 84 | num_local_experts=64, |
| 85 | num_logical_experts=512, |
| 86 | moe_parallel_config=FusedMoEParallelConfig( |
| 87 | tp_size=1, |
| 88 | tp_rank=0, |
| 89 | pcp_size=1, |
| 90 | pcp_rank=0, |
| 91 | dp_size=1, |
| 92 | dp_rank=0, |
| 93 | ep_size=8, |
| 94 | ep_rank=0, |
| 95 | sp_size=1, |
| 96 | use_ep=True, |
| 97 | all2all_backend="allgather_reducescatter", |
| 98 | enable_eplb=False, |
| 99 | ), |
| 100 | activation=MoEActivation.SILU, |
| 101 | in_dtype=torch.bfloat16, |
| 102 | device="cuda", |
| 103 | routing_method=RoutingMethodType.Renormalize, |
| 104 | ) |
| 105 | |
| 106 | |
| 107 | def test_qwen38_tep8_requires_emulation_only_on_gfx950(monkeypatch): |
| 108 | import vllm.model_executor.layers.fused_moe.oracle.mxfp4 as mxfp4_oracle |
| 109 | |
| 110 | config = _make_qwen38_tep8_moe_config() |
| 111 | monkeypatch.setattr(current_platform, "is_rocm", lambda: True) |
| 112 | monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: True) |
| 113 | |
| 114 | assert _requires_qwen38_tep8_emulation(config, kMxfp4Dynamic) |
| 115 | |
| 116 | config.moe_parallel_config.ep_size = 4 |
| 117 | config.num_local_experts = 128 |
| 118 | assert not _requires_qwen38_tep8_emulation(config, kMxfp4Dynamic) |
| 119 | |
| 120 | config = _make_qwen38_tep8_moe_config() |
| 121 | monkeypatch.setattr(mxfp4_oracle.current_platform, "is_rocm", lambda: False) |
| 122 | assert not _requires_qwen38_tep8_emulation(config, kMxfp4Dynamic) |
| 123 | |
| 124 | |
| 125 | @pytest.mark.parametrize( |
| 126 | "requested_backend,expected_backend", |
| 127 | [ |
| 128 | ("auto", Mxfp4MoeBackend.EMULATION), |
| 129 | ("aiter", Mxfp4MoeBackend.AITER_MXFP4_MXFP4), |
| 130 | ], |
| 131 | ) |
| 132 | def test_qwen38_tep8_auto_fallback_respects_explicit_backend( |
| 133 | requested_backend, |
| 134 | expected_backend, |
| 135 | monkeypatch, |
| 136 | ): |
| 137 | import vllm.model_executor.layers.fused_moe.oracle.mxfp4 as mxfp4_oracle |
| 138 | |
| 139 | class SupportedExperts: |
| 140 | @staticmethod |
| 141 | def is_supported_config(*args, **kwargs): |
| 142 | return True, None |
| 143 | |
| 144 | config = _make_qwen38_tep8_moe_config() |
| 145 | config.moe_backend = requested_backend |
| 146 | monkeypatch.setattr(current_platform, "is_rocm", lambda: True) |
| 147 | monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: True) |
| 148 | monkeypatch.setattr(mxfp4_oracle, "_user_moe_activation_override", lambda: None) |
| 149 | monkeypatch.setattr( |
| 150 | mxfp4_oracle, "backend_to_kernel_cls", lambda backend: [SupportedExperts] |
| 151 | ) |
| 152 | |
| 153 | backend, experts_cls = select_mxfp4_moe_backend( |
| 154 | config, activation_key=kMxfp4Dynamic |
| 155 | ) |
| 156 | |
| 157 | assert backend == expected_backend |
| 158 | assert experts_cls is SupportedExperts |
| 159 | |
| 160 | |
| 161 | @pytest.fixture |
| 162 | def mxfp4_oracle_config(): |
| 163 | """Stub the config the oracle reads (``model_config.quantization_config``) |
| 164 | so backend dispatch resolves without a real model / user override.""" |
| 165 | from unittest.mock import patch |
| 166 | |
| 167 | with patch( |
| 168 | "vllm.model_executor.layers.fused_moe.oracle.mxfp4.get_current_vllm_config" |
| 169 | ) as mock_get_config: |
| 170 | mock_get_config.return_value.model_config.quantization_config = None |
| 171 | yield |
| 172 | |
| 173 | |
| 174 | @pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)") |
| 175 | @pytest.mark.skipif(not ROCM_AITER_SUPPORTED, reason="Requires supported AITER") |
| 176 | def test_w4a4_dispatches_to_aiter(mxfp4_oracle_config, enable_rocm_aiter): |
| 177 | """With AITER enabled + GFX950, W4A4 selects AITER_MXFP4_MXFP4.""" |
| 178 | config = _make_w4a4_moe_config() |
| 179 | backend, experts_cls = select_mxfp4_moe_backend( |
| 180 | config, activation_key=kMxfp4Dynamic |
| 181 | ) |
| 182 | assert backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4 |
| 183 | assert experts_cls is not None |
| 184 | |
| 185 | |
| 186 | @pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)") |
| 187 | def test_w4a4_falls_back_without_aiter( |
| 188 | mxfp4_oracle_config, |
| 189 | disable_rocm_aiter, |
| 190 | ): |
| 191 | config = _make_w4a4_moe_config() |
| 192 | backend, experts_cls = select_mxfp4_moe_backend( |
| 193 | config, activation_key=kMxfp4Dynamic |
| 194 | ) |
| 195 | assert backend == Mxfp4MoeBackend.EMULATION |
| 196 | assert experts_cls is not None |
| 197 | |
| 198 | |
| 199 | @pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)") |
| 200 | def test_w4a4_dispatches_to_emulation_with_moe_backend(mxfp4_oracle_config): |
| 201 | """With --moe-backend emulation, W4A4 selects EMULATION.""" |
| 202 | config = _make_w4a4_moe_config(moe_backend="emulation") |
| 203 | backend, experts_cls = select_mxfp4_moe_backend( |
| 204 | config, activation_key=kMxfp4Dynamic |
| 205 | ) |
| 206 | assert backend == Mxfp4MoeBackend.EMULATION |
| 207 | assert experts_cls is not None |