@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c2507fb293

[ROCm] [MoE] [Perf] Shared-expert fusion for bias-routed MoE; enable on MiniMax-M3 mxfp8 model (#46545)

Signed-off-by: Hongxia Yang <hongxia.yang@amd.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Hongxia Yang committed Jun 26, 2026 at 10:05 UTC c2507fb2937aa8c8e74bea15719d04fb6090befe

5 files changed +110 -20

vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py

+7 -1

@@ -227,10 +227,16 @@ def fused_moe_mxfp8_native(

227

228 tiles = _mxfp8_moe_tiles(T)

229 block_m = tiles["block_m"]

230 + # Bin by the actual number of expert weight rows. With fused shared experts

231 + # the weight tensor has more rows than ``global_num_experts`` (the routed

232 + # count), and their ids fall outside [0, global_num_experts); binning by the

233 + # routed count would treat them as invalid. Under EP (expert_map set) the

234 + # tensor holds only local experts, so keep the global count for remapping.

235 + num_align_experts = w13.shape[0] if expert_map is None else global_num_experts

236 sorted_ids, expert_ids, num_post = moe_align_block_size(

237 topk_ids,

238 block_m,

233 - global_num_experts,

239 + num_align_experts,

240 expert_map,

241 ignore_invalid_experts=expert_map is not None,

242 )

vllm/model_executor/layers/fused_moe/layer.py

+27 -16

@@ -6,6 +6,7 @@ from typing import Any

7 import torch

9 +import vllm.envs as envs

10 from vllm._aiter_ops import rocm_aiter_ops

11 from vllm.config import ParallelConfig, get_current_vllm_config

12 from vllm.distributed import (

@@ -77,24 +78,20 @@ def determine_expert_counts(

78 ) -> tuple[int, int, int]:

79 global_num_experts = num_experts + num_redundant_experts

80 logical_num_experts = num_experts

80 - # ROCm aiter shared experts fusion

81 - # AITER only supports gated activations (silu/gelu), so disable it

82 - # for non-gated MoE (is_act_and_mul=False)

83 - # rocm_aiter_fmoe_enabled = rocm_aiter_ops.is_fused_moe_enabled() and is_act_and_mul

84 - aiter_fmoe_shared_expert_enabled = (

85 - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() and is_act_and_mul

86 - )

81 + # Shared-expert fusion: append the shared expert(s) as routed-expert slots

82 + # so they run in the same grouped GEMM. Gated by

83 + # VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: either the native aiter fused-MoE

84 + # path (env + master switch, via is_fusion_moe_shared_experts_enabled) or the

85 + # backend-neutral router-append path (env alone, independent of the master

86 + # switch; e.g. the MM3 triton/flydsl mxfp8 MoE). Gated activations only.

87 + fuse_shared_enabled = (

88 + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()

89 + or envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS

90 + ) and is_act_and_mul

92 num_fused_shared_experts = (

89 - n_shared_experts

90 - if n_shared_experts is not None and aiter_fmoe_shared_expert_enabled

91 - else 0

93 + n_shared_experts if n_shared_experts is not None and fuse_shared_enabled else 0

94 )

93 - if not aiter_fmoe_shared_expert_enabled and num_fused_shared_experts != 0:

94 - raise ValueError(

95 - "n_shared_experts is only supported on ROCm aiter when "

96 - "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled"

97 - )

96 return global_num_experts, logical_num_experts, num_fused_shared_experts

@@ -188,7 +185,8 @@ def FusedMoE(

185 has_bias: Whether expert layers have bias terms

186 is_sequence_parallel: Whether sequence parallelism is enabled

187 expert_mapping: Expert parameter mapping for weight loading

191 - n_shared_experts: Number of shared experts (ROCm aiter only)

188 + n_shared_experts: Number of shared experts to fuse into the routed

189 + grouped GEMM (ROCm; requires aiter FSE or the router-append path)

190 router_logits_dtype: Data type for router logits buffers

191 gate: Pre-configured gate module

192 shared_experts: Pre-configured shared experts module

@@ -289,6 +287,19 @@ def FusedMoE(

287 else 1.0,

288 e_score_correction_bias=e_score_correction_bias,

289 num_fused_shared_experts=num_fused_shared_experts,

290 + # Fused shared-expert slot weight. With apply_routed_scale_to_output

291 + # the runner scales the combined output by routed_scaling_factor, so

292 + # the shared slot weight must be 1/routed_scaling_factor for its net

293 + # contribution to be 1.0 (matching the un-scaled separate-MLP add).

294 + shared_expert_weight=(

295 + (1.0 / routed_scaling_factor)

296 + if (

297 + apply_routed_scale_to_output

298 + and num_fused_shared_experts > 0

299 + and routed_scaling_factor

300 + )

301 + else 1.0

302 + ),

303 zero_expert_type=zero_expert_type,

304 num_logical_experts=logical_num_experts,

305 hash_indices_table=hash_indices_table,

vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py

+26

@@ -334,6 +334,8 @@ class FusedTopKBiasRouter(BaseRouter):

334 *,

335 scoring_func: str = "sigmoid",

336 hash_indices_table: torch.Tensor | None = None,

337 + num_fused_shared_experts: int = 0,

338 + shared_expert_weight: float = 1.0,

339 ):

340 super().__init__(

341 top_k=top_k,

@@ -346,6 +348,11 @@ class FusedTopKBiasRouter(BaseRouter):

348 self.routed_scaling_factor = routed_scaling_factor

349 self.scoring_func = scoring_func

350 self._hash_indices_table = hash_indices_table

351 + # Fused shared experts: append constant slots (ids immediately after

352 + # the routed experts, [global, global+n)) routed to by every token at

353 + # ``shared_expert_weight``, AFTER the routed top-k is renormalized.

354 + self.num_fused_shared_experts = num_fused_shared_experts

355 + self.shared_expert_weight = shared_expert_weight

356

357 @property

358 def routing_method_type(self) -> RoutingMethodType:

@@ -382,4 +389,23 @@ class FusedTopKBiasRouter(BaseRouter):

389 routed_scaling_factor=self.routed_scaling_factor,

390 )

391

392 + if self.num_fused_shared_experts > 0:

393 + m = topk_ids.shape[0]

394 + n = self.num_fused_shared_experts

395 + # global_num_experts counts only the routed experts; the fused

396 + # shared experts occupy the slots immediately after them, i.e. ids

397 + # [global_num_experts, global_num_experts + n).

398 + base = self.global_num_experts

399 + shared_ids = torch.arange(

400 + base, base + n, dtype=topk_ids.dtype, device=topk_ids.device

401 + ).expand(m, n)

402 + shared_w = torch.full(

403 + (m, n),

404 + self.shared_expert_weight,

405 + dtype=topk_weights.dtype,

406 + device=topk_weights.device,

407 + )

408 + topk_ids = torch.cat([topk_ids, shared_ids], dim=-1)

409 + topk_weights = torch.cat([topk_weights, shared_w], dim=-1)

410 +

411 return topk_weights, topk_ids

vllm/model_executor/layers/fused_moe/router/router_factory.py

@@ -47,6 +47,7 @@ def create_fused_moe_router(

47 topk_group: int | None = None,

48 scoring_func: str = "softmax",

49 num_fused_shared_experts: int = 0,

50 + shared_expert_weight: float = 1.0,

51 # grouped topk + fused topk bias parameters

52 routed_scaling_factor: float = 1.0,

53 e_score_correction_bias: torch.Tensor | None = None,

@@ -188,6 +189,8 @@ def create_fused_moe_router(

189 routed_scaling_factor=routed_scaling_factor,

190 scoring_func=scoring_func,

191 hash_indices_table=hash_indices_table,

192 + num_fused_shared_experts=num_fused_shared_experts,

193 + shared_expert_weight=shared_expert_weight,

194 )

195

196 if (

vllm/models/minimax_m3/amd/model.py

+47 -3

@@ -23,6 +23,7 @@ import torch

23 from torch import nn

24 from transformers import PretrainedConfig

26 +import vllm.envs as envs

27 from vllm import _custom_ops as ops

28 from vllm.compilation.breakable_cudagraph import eager_break_during_capture

29 from vllm.config import (

@@ -95,6 +96,7 @@ from vllm.models.minimax_m3.common.sparse_attention import (

96 )

97 from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel

98 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 (

@@ -104,6 +106,23 @@ from vllm.v1.kv_cache_interface import (

106 )

107

108

109 +def _fuse_shared_experts_enabled(config: PretrainedConfig) -> bool:

110 + """Whether to fuse the shared expert into the routed grouped MoE.

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).

117 + """

118 + return bool(

119 + current_platform.is_rocm()

120 + and getattr(config, "n_shared_experts", None)

121 + and envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS

122 + and not get_current_vllm_config().parallel_config.enable_expert_parallel

123 + )

124 +

125 +

126 def _sparse_attention_layer_ids(config: PretrainedConfig) -> set[int]:

127 """Layer ids whose attention runs the extra sparse "index" branch."""

128 cfg = getattr(config, "sparse_attention_config", None)

@@ -294,8 +313,13 @@ class MiniMaxM3MoE(nn.Module):

313 prefix=f"{prefix}.gate",

314 )

315

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.

320 + self.fuse_shared_experts = _fuse_shared_experts_enabled(config)

321 self.shared_experts: MiniMaxM3MLP | None = None

298 - if self.n_shared_experts:

322 + if self.n_shared_experts and not self.fuse_shared_experts:

323 self.shared_experts = MiniMaxM3MLP(

324 config=config,

325 intermediate_size=config.intermediate_size * self.n_shared_experts,

@@ -321,6 +345,9 @@ class MiniMaxM3MoE(nn.Module):

345 apply_routed_scale_to_output=True,

346 router_logits_dtype=self.gate.out_dtype,

347 shared_experts=self.shared_experts,

348 + n_shared_experts=(

349 + self.n_shared_experts if self.fuse_shared_experts else None

350 + ),

351 quant_config=quant_config,

352 prefix=f"{prefix}.experts",

353 )

@@ -864,13 +891,18 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):

891 return hidden_states

892

893 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

867 - # Checkpoint experts use w1=gate, w2=down, w3=up.

894 + # Checkpoint experts use w1=gate, w2=down, w3=up. When fusing the shared

895 + # expert, include the appended slot (id == num_local_experts).

896 + n_shared = getattr(self.config, "n_shared_experts", 0) or 0

897 + num_experts = self.config.num_local_experts + (

898 + n_shared if _fuse_shared_experts_enabled(self.config) else 0

899 + )

900 return fused_moe_make_expert_params_mapping(

901 self,

902 ckpt_gate_proj_name="w1",

903 ckpt_down_proj_name="w2",

904 ckpt_up_proj_name="w3",

873 - num_experts=self.config.num_local_experts,

905 + num_experts=num_experts,

906 )

907

908 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

@@ -897,6 +929,7 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):

929

930 params_dict = dict(self.named_parameters())

931 loaded_params: set[str] = set()

932 + _fuse_shared = _fuse_shared_experts_enabled(self.config)

933 for name, loaded_weight in weights:

934 # The MTP module is not modeled yet.

935 if "mtp." in name:

@@ -907,6 +940,17 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):

940 if "weight_scale_inv" in name:

941 name = name.replace("weight_scale_inv", "weight_scale")

942

943 + # Shared-expert fusion: redirect the checkpoint shared expert into

944 + # routed-expert slot ``num_local_experts`` (gate->w1, up->w3,

945 + # down->w2) so it loads via the routed expert loader. Runs before the

946 + # stacked/dense mappings so shared_experts.gate_proj/up_proj are not

947 + # captured by the dense gate_up_proj mapping.

948 + if _fuse_shared and ".shared_experts." in name:

949 + sid = self.config.num_local_experts

950 + name = name.replace(".shared_experts.gate_proj.", f".experts.{sid}.w1.")

951 + name = name.replace(".shared_experts.up_proj.", f".experts.{sid}.w3.")

952 + name = name.replace(".shared_experts.down_proj.", f".experts.{sid}.w2.")

953 +

954 for param_name, weight_name, shard_id in stacked_params_mapping:

955 if weight_name not in name:

956 continue