@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 2afa3f7e95

[Perf] Minimax M3 - Support cross-layer allreduce-norm fusion (#47631)

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>

Wei Zhao committed Jul 8, 2026 at 00:16 UTC 2afa3f7e950264bb179d030c23a1ed1f46558fd9

7 files changed +65 -17

vllm/model_executor/layers/fused_moe/layer.py

+13

@@ -127,6 +127,7 @@ def FusedMoE(

127 num_redundant_experts: int = 0,

128 has_bias: bool = False,

129 is_sequence_parallel: bool = False,

130 + reduce_results: bool = True,

131 ckpt_names: tuple[str, str, str] = ("gate_proj", "down_proj", "up_proj"),

132 n_shared_experts: int | None = None,

133 router_logits_dtype: torch.dtype | None = None,

@@ -184,6 +185,9 @@ def FusedMoE(

185 num_redundant_experts: Number of redundant experts for EPLB

186 has_bias: Whether expert layers have bias terms

187 is_sequence_parallel: Whether sequence parallelism is enabled

188 + reduce_results: Whether to all-reduce the final output. Setting this

189 + to False (to fuse the all-reduce downstream) is only honored on the

190 + late-AR path.

191 expert_mapping: Expert parameter mapping for weight loading

192 n_shared_experts: Number of shared experts to fuse into the routed

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

@@ -220,6 +224,14 @@ def FusedMoE(

224 parallel_config=vllm_config.parallel_config,

225 )

226

227 + # Resolve the deferred all-reduce request against the parallel config.

228 + skip_final_all_reduce = (

229 + not reduce_results

230 + and not moe_parallel_config.use_all2all_kernels

231 + and not moe_parallel_config.is_sequence_parallel

232 + and zero_expert_type is None

233 + )

234 +

235 global_num_experts, logical_num_experts, num_fused_shared_experts = (

236 determine_expert_counts(

237 num_experts,

@@ -340,6 +352,7 @@ def FusedMoE(

352 swiglu_alpha=swiglu_alpha,

353 swiglu_beta=swiglu_beta,

354 max_capture_size=vllm_config.compilation_config.max_cudagraph_capture_size,

355 + skip_final_all_reduce=skip_final_all_reduce,

356 )

357

358 logger.debug("FusedMoEConfig = %s", moe_config)

vllm/model_executor/layers/fused_moe/runner/moe_runner.py

+7 -1

@@ -428,7 +428,6 @@ class MoERunner(MoERunnerInterface):

428 if (

429 shared_output is not None

430 and not self.moe_config.is_sequence_parallel

431 - and not self.moe_config.skip_final_all_reduce

431 and self._fused_output_is_reduced

432 ):

433 shared_output = tensor_model_parallel_all_reduce(shared_output)

@@ -446,6 +445,13 @@ class MoERunner(MoERunnerInterface):

445 here. Skipped when sequence-parallel is active (SP handles its

446 own reduction) or when the early path already reduced both outputs.

447 """

448 + # skip_final_all_reduce must not coexist with a pre-reduced fused

449 + # output. This should be enforced by MoE config initialization.

450 + if self.moe_config.skip_final_all_reduce:

451 + assert not self._fused_output_is_reduced, (

452 + "skip_final_all_reduce requires an un-reduced fused output"

453 + )

454 +

455 # We don't need to reduce the final output if:

456 # - We are not running with TP or DP

457 # - The MK already reduced the fused output itself.

vllm/model_executor/models/deepseek_v2.py

@@ -279,6 +279,7 @@ class DeepseekV2MoE(nn.Module):

279 config: DeepseekV2Config | DeepseekV3Config,

280 parallel_config: ParallelConfig,

281 quant_config: QuantizationConfig | None = None,

282 + reduce_results: bool = True,

283 prefix: str = "",

284 apply_routed_scale_to_output: bool = False,

285 ):

@@ -379,6 +380,7 @@ class DeepseekV2MoE(nn.Module):

380 enable_eplb=self.enable_eplb,

381 num_redundant_experts=self.n_redundant_experts,

382 is_sequence_parallel=self.is_sequence_parallel,

383 + reduce_results=reduce_results,

384 n_shared_experts=config.n_shared_experts

385 if self.is_fusion_moe_shared_experts_enabled

386 else None,

vllm/models/deepseek_v32/nvidia/model.py

+4 -4

@@ -72,17 +72,17 @@ class DeepseekV32DecoderLayer(torch.nn.Module):

72 and layer_idx >= config.first_k_dense_replace

73 and layer_idx % moe_layer_freq == 0

74 ):

75 + # Defer the MoE cross-rank all-reduce; it is fused into the next

76 + # layer's input_layernorm (or the final norm) via

77 + # fused_allreduce_rms_norm.

78 self.mlp = DeepseekV2MoE(

79 config=config,

80 parallel_config=parallel_config,

81 quant_config=quant_config,

82 + reduce_results=False,

83 prefix=f"{prefix}.mlp",

84 apply_routed_scale_to_output=False,

85 )

82 - # Defer the MoE cross-rank all-reduce; it is fused into the next

83 - # layer's input_layernorm (or the final norm) via

84 - # fused_allreduce_rms_norm. self.mlp.experts is the MoERunner.

85 - self.mlp.experts.moe_config.skip_final_all_reduce = True

86 else:

87 self.mlp = DeepseekV2MLP(

88 hidden_size=config.hidden_size,

vllm/models/deepseek_v32/nvidia/mtp.py

+1 -1

@@ -89,7 +89,7 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):

89 hidden_states, residual = self.mtp_block(

90 positions=positions, hidden_states=hidden_states, residual=None

91 )

92 - # mtp_block's MoE output is left un-reduced (skip_final_all_reduce); the

92 + # mtp_block's MoE output is left un-reduced (reduce_results=False); the

93 # main model fuses that all-reduce into the next norm, but here the

94 # recycle hidden is consumed directly, so reduce it now.

95 hidden_states = tensor_model_parallel_all_reduce(hidden_states)

vllm/models/minimax_m3/nvidia/model.py

+37 -10

@@ -196,6 +196,7 @@ class MiniMaxM3MoE(nn.Module):

196 config: PretrainedConfig,

197 layer_id: int,

198 quant_config: QuantizationConfig | None = None,

199 + reduce_results: bool = True,

200 prefix: str = "",

201 ) -> None:

202 super().__init__()

@@ -261,6 +262,7 @@ class MiniMaxM3MoE(nn.Module):

262 router_logits_dtype=self.gate.out_dtype,

263 shared_experts=self.shared_experts,

264 quant_config=quant_config,

265 + reduce_results=reduce_results,

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

267 )

268

@@ -658,14 +660,10 @@ class MiniMaxM3DecoderLayer(nn.Module):

660 layer_id = int(prefix.split(sep=".")[-1])

661 self.layer_id = layer_id

662

661 - # Complete the preceding dense MLP's deferred all-reduce

662 - # (reduce_results=False), fused into this layer's input_layernorm.

663 - # Disable this fusion when PP is set

664 - self.fuse_input_allreduce = (

665 - layer_id > 0

666 - and not _is_moe_layer(config, layer_id - 1)

667 - and vllm_config.parallel_config.pipeline_parallel_size == 1

668 - )

663 + # When set, complete the preceding layer's deferred FFN all-reduce

664 + # fused into this layer's input_layernorm.

665 + # Configured by MiniMaxM3Model.__init__

666 + self.fuse_input_allreduce = False

667

668 is_sparse_attention_layer = (

669 force_sparse_attn or layer_id in _sparse_attention_layer_ids(config)

@@ -691,12 +689,19 @@ class MiniMaxM3DecoderLayer(nn.Module):

689

690 # Dense layers store the FFN under `mlp`; MoE layers under

691 # `block_sparse_moe` -- matching the checkpoint's naming.

692 + # Leave the FFN output un-reduced so its all-reduce fuses into the

693 + # next RMSNorm. MTP blocks add the residual directly and PP sends

694 + # hidden states across stages, so both must reduce.

695 + reduce_results = (

696 + is_mtp_block or vllm_config.parallel_config.pipeline_parallel_size > 1

697 + )

698 self.is_moe_layer = force_moe or _is_moe_layer(config, layer_id)

699 if self.is_moe_layer:

700 self.block_sparse_moe = MiniMaxM3MoE(

701 config=config,

702 layer_id=layer_id,

703 quant_config=quant_config,

704 + reduce_results=reduce_results,

705 prefix=f"{prefix}.block_sparse_moe",

706 )

707 else:

@@ -705,7 +710,7 @@ class MiniMaxM3DecoderLayer(nn.Module):

710 intermediate_size=config.dense_intermediate_size,

711 quant_config=quant_config,

712 prefix=f"{prefix}.mlp",

708 - reduce_results=vllm_config.parallel_config.pipeline_parallel_size > 1,

713 + reduce_results=reduce_results,

714 )

715

716 # config.use_gemma_norm is True for M3 -> Gemma-style RMSNorm.

@@ -744,6 +749,14 @@ class MiniMaxM3DecoderLayer(nn.Module):

749 hidden_states = ffn(hidden_states)

750 return hidden_states, residual

751

752 + @property

753 + def ffn_all_reduce_deferred(self) -> bool:

754 + """This layer's FFN output is left un-reduced; the caller fuses the

755 + all-reduce into the next RMSNorm."""

756 + if self.is_moe_layer:

757 + return self.block_sparse_moe.experts.moe_config.skip_final_all_reduce

758 + return not self.mlp.down_proj.reduce_results

759 +

760

761 class MiniMaxM3Model(nn.Module, EagleModelMixin):

762 fall_back_to_pt_during_load = False

@@ -807,6 +820,15 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):

820 ["hidden_states", "residual"], config.hidden_size

821 )

822

823 + # Configure cross-layer all-reduce/RMSNorm fusion: a layer whose FFN output

824 + # is left un-reduced has that all-reduce fused into the next layer's

825 + # input_layernorm (or the final norm).

826 + prev_defers = False

827 + for idx, layer in enumerate(self.layers[self.start_layer : self.end_layer]):

828 + layer.fuse_input_allreduce = idx > 0 and prev_defers

829 + prev_defers = layer.ffn_all_reduce_deferred

830 + self.fuse_final_norm_allreduce = prev_defers

831 +

832 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

833 return self.embed_tokens(input_ids)

834

@@ -841,7 +863,12 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):

863 {"hidden_states": hidden_states, "residual": residual}

864 )

865

844 - hidden_states, _ = self.norm(hidden_states, residual)

866 + if self.fuse_final_norm_allreduce:

867 + hidden_states, _ = fused_allreduce_gemma_rms_norm(

868 + hidden_states, residual, self.norm

869 + )

870 + else:

871 + hidden_states, _ = self.norm(hidden_states, residual)

872

873 if len(aux_hidden_states) > 0:

874 return hidden_states, aux_hidden_states

vllm/v1/attention/backends/mla/prefill/selector.py

+1 -1

@@ -134,7 +134,7 @@ def get_mla_prefill_backend(

134 f"Reason: {invalid_reasons}"

135 )

136 assert backend_cls is not None

137 - logger.info("Using %s MLA prefill backend.", selected_backend.name)

137 + logger.info_once("Using %s MLA prefill backend.", selected_backend.name)

138 return backend_cls

139

140 return _auto_select_mla_prefill_backend(