[Bugfix][ROCm][MoE] Fall back instead of crashing when AITER MoE is requested for a non-gated (is_act_and_mul=False) model (#56590)
Signed-off-by: Shantipriya Parida <shantipriya.parida@amd.com>
Shantipriya Parida committed
Sep 17, 2026 at 16:39 UTC
0eae9acd4d01574e12d4ecf6a0229813f7fdb799
2 files changed
+87
-1
tests/kernels/moe/test_unquantized_backend_selection.py
+60
@@ -6,6 +6,7 @@ import pytest
6
import torch
7
8
from tests.kernels.moe.utils import make_dummy_moe_config
9
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
10
from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
11
from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
12
UnquantizedMoeBackend,
@@ -228,6 +229,65 @@ def test_select_rocm_aiter_backend(mock_aiter_enabled, mock_has_flashinfer):
229
assert expert_cls is not None
230
231
232
+@patch(
233
+ "vllm.utils.flashinfer.has_flashinfer",
234
+ return_value=False,
235
+)
236
+@patch(
237
+ "vllm.model_executor.layers.fused_moe.oracle.unquantized.rocm_aiter_ops."
238
+ "is_fused_moe_enabled",
239
+ return_value=True,
240
+)
241
+@patch(
242
+ "vllm.model_executor.layers.fused_moe.oracle.unquantized.rocm_aiter_ops."
243
+ "is_rdna_aiter_enabled",
244
+ return_value=False,
245
+)
246
+@pytest.mark.skipif(
247
+ not current_platform.is_rocm(), reason="ROCm-specific backend selection test"
248
+)
249
+def test_select_rocm_aiter_backend_non_gated_activation_falls_back(
250
+ mock_rdna_disabled, mock_aiter_enabled, mock_has_flashinfer, monkeypatch
251
+):
252
+ """Test ROCm backend selection falls back (not raises) for non-gated MoE."""
253
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
254
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1")
255
+
256
+ with patch(
257
+ "vllm.model_executor.layers.fused_moe.oracle.unquantized.current_platform"
258
+ ) as mock_platform:
259
+ mock_platform.is_cuda.return_value = False
260
+ mock_platform.is_rocm.return_value = True
261
+ mock_platform.is_cpu.return_value = False
262
+ mock_platform.is_xpu.return_value = False
263
+ mock_platform.is_tpu.return_value = False
264
+ mock_platform.is_out_of_tree.return_value = False
265
+
266
+ moe_config = make_dummy_moe_config(activation=MoEActivation.SILU_NO_MUL)
267
+ assert moe_config.is_act_and_mul is False
268
+
269
+ selected_backend, expert_cls = select_unquantized_moe_backend(
270
+ moe_config=moe_config,
271
+ )
272
+
273
+ assert selected_backend != UnquantizedMoeBackend.AITER
274
+ assert expert_cls is not None
275
+
276
+
277
+@pytest.mark.skipif(
278
+ not current_platform.is_rocm(), reason="ROCm-specific backend selection test"
279
+)
280
+def test_explicit_aiter_backend_non_gated_activation_still_raises():
281
+ """Explicit `--moe-backend aiter` still raises for non-gated MoE;
282
+ only the env-var opt-in path falls back."""
283
+ moe_config = make_dummy_moe_config(activation=MoEActivation.SILU_NO_MUL)
284
+ moe_config.moe_backend = "aiter" # explicit pin, not "auto"
285
+ assert moe_config.is_act_and_mul is False
286
+
287
+ with pytest.raises(ValueError):
288
+ select_unquantized_moe_backend(moe_config=moe_config)
289
+
290
+
291
@patch(
292
"vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsMonolithic.is_supported_config",
293
return_value=(True, None),
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+27
-1
@@ -305,7 +305,33 @@ def select_unquantized_moe_backend(
305
AVAILABLE_BACKENDS.remove(UnquantizedMoeBackend.AITER)
306
else:
307
backend = UnquantizedMoeBackend.AITER
308
- return _return_or_raise(backend, moe_config, activation_format)
308
+ reason = None
309
+ for k_cls in backend_to_kernel_cls(backend):
310
+ supported, reason = k_cls.is_supported_config(
311
+ k_cls, moe_config, None, None, activation_format
312
+ )
313
+ if supported:
314
+ logger.info_once(_make_log_backend(backend))
315
+ return backend, k_cls
316
+ # AITER was explicitly requested but does not support this
317
+ # deployment configuration (e.g. a non-gated MoE activation,
318
+ # which AiterExperts._supports_no_act_and_mul() rejects
319
+ # unconditionally). Rather than hard-crashing at engine init,
320
+ # warn — explicitly stating that this is falling back, and
321
+ # which backends will be tried next — and fall back to the
322
+ # remaining backends in priority order below.
323
+ logger.warning_once(
324
+ "VLLM_ROCM_USE_AITER_MOE=1 was requested, but %s "
325
+ "Falling back to try the remaining available MoE backends: %s.",
326
+ _make_log_unsupported(backend, reason),
327
+ ", ".join(
328
+ b.value
329
+ for b in AVAILABLE_BACKENDS
330
+ if b != UnquantizedMoeBackend.AITER
331
+ ),
332
+ )
333
+ if UnquantizedMoeBackend.AITER in AVAILABLE_BACKENDS:
334
+ AVAILABLE_BACKENDS.remove(UnquantizedMoeBackend.AITER)
335
336
for backend in AVAILABLE_BACKENDS:
337
for k_cls in backend_to_kernel_cls(backend):