[ROCm]Enable AITER MoE backend for MiniMax-M3-MXFP4 (#46419)
Signed-off-by: Qiang Li <qiang.li2@amd.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
qli88 committed
Jun 26, 2026 at 08:35 UTC
8e394244a59afc67a37bf47dab0ab76bf5ce5885
4 files changed
+27
-8
vllm/model_executor/layers/fused_moe/config.py
+2
@@ -1279,6 +1279,8 @@ class FusedMoEConfig:
1279
hidden_dim_unpadded: int | None = None
1280
# Defaults to intermediate_size_per_partition if not specified.
1281
intermediate_size_per_partition_unpadded: int | None = None
1282
+ # Model specific override
1283
+ intermediate_pad: int | None = None
1284
1285
moe_backend: MoEBackend = "auto"
1286
max_num_tokens: int = SchedulerConfig.DEFAULT_MAX_NUM_BATCHED_TOKENS_FOR_BATCHED_DP
vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py
+22
-8
@@ -251,12 +251,17 @@ def rocm_aiter_fused_experts(
251
if quant_config is None:
252
quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
253
254
+ # Gate/up interleave hint; only the SWIGLUOAI activations override it.
255
+ activation_interleave = None
256
if activation == MoEActivation.SILU:
257
activation_method = ActivationMethod.SILU
258
elif activation == MoEActivation.GELU:
259
activation_method = ActivationMethod.GELU
260
elif activation == MoEActivation.SWIGLUOAI:
261
activation_method = rocm_aiter_ops.get_aiter_activation_type("swiglu")
262
+ elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
263
+ activation_method = rocm_aiter_ops.get_aiter_activation_type("swiglu")
264
+ activation_interleave = False
265
else:
266
raise ValueError(f"Unsupported activation: {activation}")
267
@@ -337,9 +342,14 @@ def rocm_aiter_fused_experts(
342
assert moe_config.intermediate_size_per_partition_unpadded is not None
343
hidden_pad = hidden_states.shape[1] - moe_config.hidden_dim_unpadded
344
intermediate_pad = (
340
- moe_config.intermediate_size_per_partition
341
- - moe_config.intermediate_size_per_partition_unpadded
345
+ (
346
+ moe_config.intermediate_size_per_partition
347
+ - moe_config.intermediate_size_per_partition_unpadded
348
+ )
349
+ if moe_config.intermediate_pad is None
350
+ else moe_config.intermediate_pad
351
)
352
+
353
# Round hidden_pad/intermediate_pad to match AITER's CK/FlyDSL MoE
354
# dispatch (currently pinned to v0.1.13.post1):
355
# https://github.com/ROCm/aiter/blob/v0.1.13.post1/aiter/fused_moe.py#L1073
@@ -357,14 +367,17 @@ def rocm_aiter_fused_experts(
367
# `rocm_aiter_ops.shuffle_weight_a16w4` in `oracle/mxfp4.py`,
368
# which always sets `is_guinterleave=True`.
369
# Hence, we pass in GateMode.INTERLEAVE to match the weight shuffling.
370
+ from aiter.ops.flydsl.moe_common import GateMode
371
+
372
gate_mode = ""
373
if quant_config.use_mxfp4_w4a16:
362
- try:
363
- from aiter.ops.flydsl.moe_common import GateMode
364
-
365
- gate_mode = GateMode.INTERLEAVE.value
366
- except ImportError:
367
- pass
374
+ gate_mode = GateMode.INTERLEAVE.value
375
+ elif activation_interleave is not None:
376
+ gate_mode = (
377
+ GateMode.INTERLEAVE.value
378
+ if activation_interleave
379
+ else GateMode.SEPARATED.value
380
+ )
381
382
return rocm_aiter_ops.fused_moe(
383
hidden_states,
@@ -458,6 +471,7 @@ class AiterExperts(mk.FusedMoEExpertsModular):
471
MoEActivation.SILU,
472
MoEActivation.GELU,
473
MoEActivation.SWIGLUOAI,
474
+ MoEActivation.SWIGLUOAI_UNINTERLEAVE,
475
]
476
477
@staticmethod
vllm/model_executor/layers/fused_moe/layer.py
+2
@@ -105,6 +105,7 @@ def FusedMoE(
105
top_k: int,
106
hidden_size: int,
107
intermediate_size: int,
108
+ intermediate_pad: int | None = None,
109
params_dtype: torch.dtype | None = None,
110
renormalize: bool = True,
111
use_grouped_topk: bool = False,
@@ -311,6 +312,7 @@ def FusedMoE(
312
experts_per_token=top_k,
313
hidden_dim=hidden_size,
314
intermediate_size=intermediate_size,
315
+ intermediate_pad=intermediate_pad,
316
num_local_experts=expert_map_manager.local_num_experts,
317
num_logical_experts=logical_num_experts,
318
moe_parallel_config=moe_parallel_config,
vllm/models/minimax_m3/amd/model.py
+1
@@ -309,6 +309,7 @@ class MiniMaxM3MoE(nn.Module):
309
top_k=config.num_experts_per_tok,
310
hidden_size=config.hidden_size,
311
intermediate_size=config.intermediate_size,
312
+ intermediate_pad=0,
313
scoring_func=config.scoring_func,
314
e_score_correction_bias=self.e_score_correction_bias,
315
renormalize=True,