diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index a522a39153..5c8f7010f6 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -127,6 +127,7 @@ def FusedMoE( num_redundant_experts: int = 0, has_bias: bool = False, is_sequence_parallel: bool = False, + reduce_results: bool = True, ckpt_names: tuple[str, str, str] = ("gate_proj", "down_proj", "up_proj"), n_shared_experts: int | None = None, router_logits_dtype: torch.dtype | None = None, @@ -184,6 +185,9 @@ def FusedMoE( num_redundant_experts: Number of redundant experts for EPLB has_bias: Whether expert layers have bias terms is_sequence_parallel: Whether sequence parallelism is enabled + reduce_results: Whether to all-reduce the final output. Setting this + to False (to fuse the all-reduce downstream) is only honored on the + late-AR path. expert_mapping: Expert parameter mapping for weight loading n_shared_experts: Number of shared experts to fuse into the routed grouped GEMM (ROCm; requires aiter FSE or the router-append path) @@ -220,6 +224,14 @@ def FusedMoE( parallel_config=vllm_config.parallel_config, ) + # Resolve the deferred all-reduce request against the parallel config. + skip_final_all_reduce = ( + not reduce_results + and not moe_parallel_config.use_all2all_kernels + and not moe_parallel_config.is_sequence_parallel + and zero_expert_type is None + ) + global_num_experts, logical_num_experts, num_fused_shared_experts = ( determine_expert_counts( num_experts, @@ -340,6 +352,7 @@ def FusedMoE( swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, max_capture_size=vllm_config.compilation_config.max_cudagraph_capture_size, + skip_final_all_reduce=skip_final_all_reduce, ) logger.debug("FusedMoEConfig = %s", moe_config) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 1720e8b8a3..cf00fa452a 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -428,7 +428,6 @@ class MoERunner(MoERunnerInterface): if ( shared_output is not None and not self.moe_config.is_sequence_parallel - and not self.moe_config.skip_final_all_reduce and self._fused_output_is_reduced ): shared_output = tensor_model_parallel_all_reduce(shared_output) @@ -446,6 +445,13 @@ class MoERunner(MoERunnerInterface): here. Skipped when sequence-parallel is active (SP handles its own reduction) or when the early path already reduced both outputs. """ + # skip_final_all_reduce must not coexist with a pre-reduced fused + # output. This should be enforced by MoE config initialization. + if self.moe_config.skip_final_all_reduce: + assert not self._fused_output_is_reduced, ( + "skip_final_all_reduce requires an un-reduced fused output" + ) + # We don't need to reduce the final output if: # - We are not running with TP or DP # - The MK already reduced the fused output itself. diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 7cf119c7c2..208379403a 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -279,6 +279,7 @@ class DeepseekV2MoE(nn.Module): config: DeepseekV2Config | DeepseekV3Config, parallel_config: ParallelConfig, quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, prefix: str = "", apply_routed_scale_to_output: bool = False, ): @@ -379,6 +380,7 @@ class DeepseekV2MoE(nn.Module): enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, is_sequence_parallel=self.is_sequence_parallel, + reduce_results=reduce_results, n_shared_experts=config.n_shared_experts if self.is_fusion_moe_shared_experts_enabled else None, diff --git a/vllm/models/deepseek_v32/nvidia/model.py b/vllm/models/deepseek_v32/nvidia/model.py index db4398e4c0..22ced9fa94 100644 --- a/vllm/models/deepseek_v32/nvidia/model.py +++ b/vllm/models/deepseek_v32/nvidia/model.py @@ -72,17 +72,17 @@ class DeepseekV32DecoderLayer(torch.nn.Module): and layer_idx >= config.first_k_dense_replace and layer_idx % moe_layer_freq == 0 ): + # Defer the MoE cross-rank all-reduce; it is fused into the next + # layer's input_layernorm (or the final norm) via + # fused_allreduce_rms_norm. self.mlp = DeepseekV2MoE( config=config, parallel_config=parallel_config, quant_config=quant_config, + reduce_results=False, prefix=f"{prefix}.mlp", apply_routed_scale_to_output=False, ) - # Defer the MoE cross-rank all-reduce; it is fused into the next - # layer's input_layernorm (or the final norm) via - # fused_allreduce_rms_norm. self.mlp.experts is the MoERunner. - self.mlp.experts.moe_config.skip_final_all_reduce = True else: self.mlp = DeepseekV2MLP( hidden_size=config.hidden_size, diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index baeb8b25e0..060dfbf7bd 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -89,7 +89,7 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module): hidden_states, residual = self.mtp_block( positions=positions, hidden_states=hidden_states, residual=None ) - # mtp_block's MoE output is left un-reduced (skip_final_all_reduce); the + # mtp_block's MoE output is left un-reduced (reduce_results=False); the # main model fuses that all-reduce into the next norm, but here the # recycle hidden is consumed directly, so reduce it now. hidden_states = tensor_model_parallel_all_reduce(hidden_states) diff --git a/vllm/models/minimax_m3/nvidia/model.py b/vllm/models/minimax_m3/nvidia/model.py index ed6e751d43..66a0b67202 100644 --- a/vllm/models/minimax_m3/nvidia/model.py +++ b/vllm/models/minimax_m3/nvidia/model.py @@ -196,6 +196,7 @@ class MiniMaxM3MoE(nn.Module): config: PretrainedConfig, layer_id: int, quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, prefix: str = "", ) -> None: super().__init__() @@ -261,6 +262,7 @@ class MiniMaxM3MoE(nn.Module): router_logits_dtype=self.gate.out_dtype, shared_experts=self.shared_experts, quant_config=quant_config, + reduce_results=reduce_results, prefix=f"{prefix}.experts", ) @@ -658,14 +660,10 @@ class MiniMaxM3DecoderLayer(nn.Module): layer_id = int(prefix.split(sep=".")[-1]) self.layer_id = layer_id - # Complete the preceding dense MLP's deferred all-reduce - # (reduce_results=False), fused into this layer's input_layernorm. - # Disable this fusion when PP is set - self.fuse_input_allreduce = ( - layer_id > 0 - and not _is_moe_layer(config, layer_id - 1) - and vllm_config.parallel_config.pipeline_parallel_size == 1 - ) + # When set, complete the preceding layer's deferred FFN all-reduce + # fused into this layer's input_layernorm. + # Configured by MiniMaxM3Model.__init__ + self.fuse_input_allreduce = False is_sparse_attention_layer = ( force_sparse_attn or layer_id in _sparse_attention_layer_ids(config) @@ -691,12 +689,19 @@ class MiniMaxM3DecoderLayer(nn.Module): # Dense layers store the FFN under `mlp`; MoE layers under # `block_sparse_moe` -- matching the checkpoint's naming. + # Leave the FFN output un-reduced so its all-reduce fuses into the + # next RMSNorm. MTP blocks add the residual directly and PP sends + # hidden states across stages, so both must reduce. + reduce_results = ( + is_mtp_block or vllm_config.parallel_config.pipeline_parallel_size > 1 + ) self.is_moe_layer = force_moe or _is_moe_layer(config, layer_id) if self.is_moe_layer: self.block_sparse_moe = MiniMaxM3MoE( config=config, layer_id=layer_id, quant_config=quant_config, + reduce_results=reduce_results, prefix=f"{prefix}.block_sparse_moe", ) else: @@ -705,7 +710,7 @@ class MiniMaxM3DecoderLayer(nn.Module): intermediate_size=config.dense_intermediate_size, quant_config=quant_config, prefix=f"{prefix}.mlp", - reduce_results=vllm_config.parallel_config.pipeline_parallel_size > 1, + reduce_results=reduce_results, ) # config.use_gemma_norm is True for M3 -> Gemma-style RMSNorm. @@ -744,6 +749,14 @@ class MiniMaxM3DecoderLayer(nn.Module): hidden_states = ffn(hidden_states) return hidden_states, residual + @property + def ffn_all_reduce_deferred(self) -> bool: + """This layer's FFN output is left un-reduced; the caller fuses the + all-reduce into the next RMSNorm.""" + if self.is_moe_layer: + return self.block_sparse_moe.experts.moe_config.skip_final_all_reduce + return not self.mlp.down_proj.reduce_results + class MiniMaxM3Model(nn.Module, EagleModelMixin): fall_back_to_pt_during_load = False @@ -807,6 +820,15 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin): ["hidden_states", "residual"], config.hidden_size ) + # Configure cross-layer all-reduce/RMSNorm fusion: a layer whose FFN output + # is left un-reduced has that all-reduce fused into the next layer's + # input_layernorm (or the final norm). + prev_defers = False + for idx, layer in enumerate(self.layers[self.start_layer : self.end_layer]): + layer.fuse_input_allreduce = idx > 0 and prev_defers + prev_defers = layer.ffn_all_reduce_deferred + self.fuse_final_norm_allreduce = prev_defers + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) @@ -841,7 +863,12 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin): {"hidden_states": hidden_states, "residual": residual} ) - hidden_states, _ = self.norm(hidden_states, residual) + if self.fuse_final_norm_allreduce: + hidden_states, _ = fused_allreduce_gemma_rms_norm( + hidden_states, residual, self.norm + ) + else: + hidden_states, _ = self.norm(hidden_states, residual) if len(aux_hidden_states) > 0: return hidden_states, aux_hidden_states diff --git a/vllm/v1/attention/backends/mla/prefill/selector.py b/vllm/v1/attention/backends/mla/prefill/selector.py index a38b274dfc..e0d54eee10 100644 --- a/vllm/v1/attention/backends/mla/prefill/selector.py +++ b/vllm/v1/attention/backends/mla/prefill/selector.py @@ -134,7 +134,7 @@ def get_mla_prefill_backend( f"Reason: {invalid_reasons}" ) assert backend_cls is not None - logger.info("Using %s MLA prefill backend.", selected_backend.name) + logger.info_once("Using %s MLA prefill backend.", selected_backend.name) return backend_cls return _auto_select_mla_prefill_backend(