[ROCm][Perf] Fused shared expert for Minimax M3 (#46474)
Signed-off-by: Fangzhou-Ai <fangzhouai@gmail.com> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com>
Fangzhou Ai committed
Jun 27, 2026 at 05:34 UTC
51a99565c398c8320de8131e07731c75c52eb87c
1 file changed
+46
-13
vllm/models/minimax_m3/amd/model.py
+46
-13
@@ -23,8 +23,9 @@ import torch
23
from torch import nn
24
from transformers import PretrainedConfig
25
26
-import vllm.envs as envs
26
from vllm import _custom_ops as ops
27
+from vllm import envs
28
+from vllm._aiter_ops import rocm_aiter_ops
29
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
30
from vllm.config import (
31
CacheConfig,
@@ -96,7 +97,6 @@ from vllm.models.minimax_m3.common.sparse_attention import (
97
)
98
from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel
99
from vllm.multimodal import MULTIMODAL_REGISTRY
99
-from vllm.platforms import current_platform
100
from vllm.sequence import IntermediateTensors
101
from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype
102
from vllm.v1.kv_cache_interface import (
@@ -107,14 +107,15 @@ from vllm.v1.kv_cache_interface import (
107
108
109
def _fuse_shared_experts_enabled(config: PretrainedConfig) -> bool:
110
- """Whether to fuse the shared expert into the routed grouped MoE.
110
+ """Whether to fuse the shared expert with routed experts.
111
112
ROCm only. Opt-in via ``VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS`` (the
113
- router-append fusion runs on the triton/flydsl mxfp8 MoE independent of the
114
- aiter master switch); requires a shared expert and is disabled under expert
115
- parallelism (the shared slot is appended to the routed top-k, which the EP
116
- expert-map path does not handle).
113
+ router-append fusion runs on both aiter and non-aiter MoE);
114
+ it is disabled under expert parallelism (the shared slot is appended to
115
+ the routed top-k, which the EP expert-mapping path does not handle).
116
"""
117
+ from vllm.platforms import current_platform
118
+
119
return bool(
120
current_platform.is_rocm()
121
and getattr(config, "n_shared_experts", None)
@@ -268,6 +269,24 @@ class MiniMaxM3MLP(nn.Module):
269
return x
270
271
272
+def _aiter_moe_fused_shared_experts_enabled(config: PretrainedConfig) -> bool:
273
+ """Whether the fused shared expert routes through aiter's grouped top-k MoE.
274
+
275
+ A strict sub-case of :func:`_fuse_shared_experts_enabled`: shared-expert
276
+ fusion must already be opted in (``VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS``)
277
+ and allowed (not under expert parallelism). When additionally on gfx950 with
278
+ an active aiter MoE backend, the shared expert is appended inside aiter's
279
+ biased grouped top-k kernel (``num_fused_shared_experts``) instead of the
280
+ vLLM router's torch concat. Otherwise FSE still runs via the vLLM top-k bias
281
+ router.
282
+ """
283
+ if not _fuse_shared_experts_enabled(config):
284
+ return False
285
+ from vllm.platforms.rocm import on_gfx950
286
+
287
+ return on_gfx950() and rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
288
+
289
+
290
class MiniMaxM3MoE(nn.Module):
291
"""Sigmoid-routed MoE block with a routing-bias correction and a shared
292
expert."""
@@ -313,11 +332,14 @@ class MiniMaxM3MoE(nn.Module):
332
prefix=f"{prefix}.gate",
333
)
334
316
- # Fuse the shared expert into the routed grouped GEMM when opted in via
317
- # VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: it becomes routed-expert slot
318
- # ``num_local_experts``, reached by every token, eliminating the
319
- # separate dense-MLP launches. Not supported under expert parallelism.
335
+ # Shared-expert fusion (opt-in via VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS,
336
+ # off under expert parallelism) folds the shared expert into the routed
337
+ # MoE call as the last expert slot, so we don't build a separate module.
338
+ # On gfx950 with aiter MoE the append is fused inside aiter's grouped
339
+ # top-k kernel; otherwise it goes through the vLLM top-k bias router.
340
self.fuse_shared_experts = _fuse_shared_experts_enabled(config)
341
+ self.use_aiter_moe_fse = _aiter_moe_fused_shared_experts_enabled(config)
342
+
343
self.shared_experts: MiniMaxM3MLP | None = None
344
if self.n_shared_experts and not self.fuse_shared_experts:
345
self.shared_experts = MiniMaxM3MLP(
@@ -328,6 +350,13 @@ class MiniMaxM3MoE(nn.Module):
350
prefix=f"{prefix}.shared_experts",
351
)
352
353
+ # The aiter MoE fused path goes through aiter's biased grouped top-k
354
+ # (GroupedTopKRouter, as in DeepSeek-V4): M3 is not group-routed, so a
355
+ # trivial single group (num_expert_group=topk_group=1) reduces to plain
356
+ # top-k while applying the sigmoid + bias correction and appending the
357
+ # always-on shared expert; aiter applies the routed scaling internally.
358
+ # Every other path (vLLM top-k bias router, or no fusion) applies the
359
+ # routed scaling to the MoE output here.
360
self.experts = FusedMoE(
361
num_experts=config.num_local_experts,
362
top_k=config.num_experts_per_tok,
@@ -337,12 +366,15 @@ class MiniMaxM3MoE(nn.Module):
366
scoring_func=config.scoring_func,
367
e_score_correction_bias=self.e_score_correction_bias,
368
renormalize=True,
369
+ use_grouped_topk=self.use_aiter_moe_fse,
370
+ num_expert_group=1 if self.use_aiter_moe_fse else None,
371
+ topk_group=1 if self.use_aiter_moe_fse else None,
372
activation="swigluoai_uninterleave",
373
swiglu_limit=config.swiglu_limit,
374
swiglu_alpha=config.swiglu_alpha,
375
swiglu_beta=config.swiglu_beta,
376
routed_scaling_factor=self.routed_scaling_factor,
345
- apply_routed_scale_to_output=True,
377
+ apply_routed_scale_to_output=not self.use_aiter_moe_fse,
378
router_logits_dtype=self.gate.out_dtype,
379
shared_experts=self.shared_experts,
380
n_shared_experts=(
@@ -927,9 +959,10 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
959
# (param_name, weight_name, expert_id, shard_id)
960
expert_params_mapping = self.get_expert_mapping()
961
962
+ _fuse_shared = _fuse_shared_experts_enabled(self.config)
963
+
964
params_dict = dict(self.named_parameters())
965
loaded_params: set[str] = set()
932
- _fuse_shared = _fuse_shared_experts_enabled(self.config)
966
for name, loaded_weight in weights:
967
# The MTP module is not modeled yet.
968
if "mtp." in name: