diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index 01664d6a932..85619a91005 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -236,7 +236,6 @@ class MoETestConfig: use_gate: bool use_routed_input_transform: bool enable_eplb: bool = False - reduce_results: bool = False backend: str | None = None ep_size: int = 1 dp_size: int = 1 @@ -295,7 +294,6 @@ def generate_valid_test_configs( use_shared_experts, use_gate, use_routed_input_transform, - reduce_results, ) in product( SHAPE_COMBOS, NUM_EXPERTS, @@ -304,7 +302,6 @@ def generate_valid_test_configs( [False, True], # shared [False, True], # gate [False, True], # routed input exform - [False, True], # reduce results ): config = MoETestConfig( shape[0], # m @@ -318,7 +315,6 @@ def generate_valid_test_configs( use_gate, use_routed_input_transform, enable_eplb, - reduce_results, backend, ep_size, dp_size, @@ -395,18 +391,7 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]: and config.backend.startswith("flashinfer_nvlink") and not current_platform.has_device_capability(90) ): - return False, "flashinfer_nvlink needs an H100+ GPUs" - - # reduce_results incompatibilities - if config.reduce_results and config.use_shared_experts: - return False, "reduce_results=True is not compatible with shared_experts=True" - - if config.reduce_results and config.quantization is not None: - return ( - False, - "reduce_results=True only tested with unquantized data types in " - "order to limit number of tests run", - ) + return False, "flashinfer_nvlink needs H100+ GPUs" # Backend-specific checks if config.backend is not None: @@ -448,10 +433,6 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]: if config.enable_eplb and config.backend not in EPLB_SUPPORTED_BACKENDS: return False, f"EPLB not supported with {config.backend}." - world_size = config.tp_size * config.dp_size - if config.reduce_results and world_size == 1: - return False, "reduce_results=True only makes sense for multi-GPU tests" - if ( config.backend is not None and config.backend.startswith("flashinfer_nvlink") @@ -846,7 +827,6 @@ def make_fused_moe_layer( tp_size: int, ep_size: int, dp_size: int, - reduce_results: bool, w1: torch.Tensor, w2: torch.Tensor, top_k: int, @@ -874,7 +854,7 @@ def make_fused_moe_layer( routed_input_transform: torch.nn.Module | None = None, routed_output_transform: torch.nn.Module | None = None, pcp_size: int | None = 1, -) -> tuple[Callable, FusedMoE]: +) -> FusedMoE: quant_config, qw = make_quant_config(quantization, w1, w2, global_num_experts) kwargs = dict() @@ -887,8 +867,10 @@ def make_fused_moe_layer( # Add gate and routed_input_transform if provided if gate is not None: kwargs["gate"] = gate + if routed_input_transform is not None: kwargs["routed_input_transform"] = routed_input_transform + kwargs["routed_output_transform"] = routed_output_transform layer = builder( num_experts=global_num_experts, @@ -896,7 +878,6 @@ def make_fused_moe_layer( hidden_size=hidden_size, intermediate_size=intermediate_size, params_dtype=in_dtype, - reduce_results=reduce_results, renormalize=renormalize, use_grouped_topk=use_grouped_topk, num_expert_group=num_expert_group, @@ -936,36 +917,7 @@ def make_fused_moe_layer( layer.quant_method.process_weights_after_loading(layer) - def _moe( - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - ) -> torch.Tensor: - if shared_experts is None: - final_shared_states = None - final_hidden_states = layer(hidden_states, router_logits) - else: - final_shared_states, final_hidden_states = layer( - hidden_states, router_logits - ) - - # Apply routed output transform if provided - # (e.g., latent space -> original space) - if routed_output_transform is not None: - final_hidden_states = routed_output_transform(final_hidden_states) - - if shared_experts is not None: - assert not reduce_results - assert final_shared_states is not None - final_hidden_states += final_shared_states - - if not reduce_results and layer.tp_size > 1: - final_hidden_states = layer.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) - - return final_hidden_states - - return _moe, layer + return layer def make_fake_moe_layer( @@ -999,7 +951,6 @@ def make_fake_moe_layer( tp_size: int = 1, dp_size: int = 1, ep_size: int = 1, - reduce_results: bool = False, ) -> Callable: activation = MoEActivation.from_str(activation) @@ -1101,7 +1052,7 @@ def make_fake_moe_layer( def _test_body_regular( - moe_fn: Callable, + moe_layer: Callable, hidden_states: torch.Tensor, router_logits: torch.Tensor, vllm_config: VllmConfig, @@ -1118,13 +1069,12 @@ def _test_body_regular( num_tokens=num_tokens, num_tokens_across_dp=num_tokens_across_dp, ): - output = moe_fn(hidden_states, router_logits) + output = moe_layer(hidden_states, router_logits) return baseline_output, output def _test_body_eplb( - moe_fn: Callable, moe_layer: FusedMoE, hidden_states: torch.Tensor, router_logits: torch.Tensor, @@ -1145,7 +1095,6 @@ def _test_body_eplb( n: int, top_k: int, shared_experts, - reduce_results: bool, gate: torch.nn.Module | None, routed_input_transform: torch.nn.Module | None, routed_output_transform: torch.nn.Module | None, @@ -1161,7 +1110,7 @@ def _test_body_eplb( num_tokens=num_tokens, num_tokens_across_dp=num_tokens_across_dp, ): - output_before = moe_fn(hidden_states, router_logits) + output_before = moe_layer(hidden_states, router_logits) # Create a fresh FusedMoE layer with enable_eplb=True # Delete the original layer's registration so the constructor can @@ -1174,7 +1123,7 @@ def _test_body_eplb( # When using routed_input_transform, experts operate in latent space hidden_size_for_layer = k // 2 if routed_input_transform is not None else k - moe_fn, moe_layer = make_fused_moe_layer( + eplb_moe_layer = make_fused_moe_layer( quantization=quantization, use_ep=use_ep, hidden_size=hidden_size_for_layer, @@ -1183,7 +1132,6 @@ def _test_body_eplb( tp_size=tp_size, ep_size=ep_size, dp_size=dp_size, - reduce_results=reduce_results, w1=w1, w2=w2, top_k=top_k, @@ -1196,14 +1144,14 @@ def _test_body_eplb( ) # Necessary? - if moe_layer._expert_map is not None: - moe_layer._expert_map = moe_layer._expert_map.to(device) + if eplb_moe_layer._expert_map is not None: + eplb_moe_layer._expert_map = eplb_moe_layer._expert_map.to(device) # All ranks must generate the same permutation initial_indices = torch.arange(num_experts, dtype=torch.long) shuffled_indices = initial_indices[torch.randperm(num_experts)] - expert_weights = [list(moe_layer.get_expert_weights())] + expert_weights = [list(eplb_moe_layer.get_expert_weights())] communicator = create_eplb_communicator( group_coordinator=get_eplb_group(), @@ -1227,7 +1175,7 @@ def _test_body_eplb( num_experts, dtype=torch.int32, device=device ) - moe_layer.set_eplb_state( + eplb_moe_layer.set_eplb_state( moe_layer_idx=0, expert_load_view=torch.zeros( (1, num_experts), @@ -1244,7 +1192,7 @@ def _test_body_eplb( ), ) - moe_layer.eplb_state.should_record_tensor = torch.ones( + eplb_moe_layer.eplb_state.should_record_tensor = torch.ones( (), dtype=torch.bool, device=device ) @@ -1255,7 +1203,7 @@ def _test_body_eplb( num_tokens=num_tokens, num_tokens_across_dp=num_tokens_across_dp, ): - output_after = moe_fn(hidden_states, router_logits) + output_after = eplb_moe_layer(hidden_states, router_logits) return output_before, output_after @@ -1274,7 +1222,6 @@ def _run_one_config( num_experts: int, top_k: int, quantization: str | None, - reduce_results: bool, backend: str | None, test_body_fn: Callable, use_shared_experts: bool, @@ -1341,7 +1288,6 @@ def _run_one_config( tp_size=tp_size, ep_size=ep_size, dp_size=dp_size, - reduce_results=reduce_results, ) baseline_output = baseline_layer(hidden_states, router_logits) @@ -1369,7 +1315,7 @@ def _run_one_config( hidden_size_for_layer = k // 2 if routed_input_transform is not None else k # Create initial MoE layer - moe_fn, moe_layer = make_fused_moe_layer( + moe_layer = make_fused_moe_layer( quantization=quantization, use_ep=use_ep, hidden_size=hidden_size_for_layer, @@ -1378,7 +1324,6 @@ def _run_one_config( tp_size=tp_size, ep_size=ep_size, dp_size=dp_size, - reduce_results=reduce_results, w1=w1, w2=w2, top_k=top_k, @@ -1402,7 +1347,6 @@ def _run_one_config( # Call the test body function with all necessary context expected, actual = test_body_fn( - moe_fn=moe_fn, moe_layer=moe_layer, hidden_states=hidden_states, router_logits=router_logits, @@ -1423,7 +1367,6 @@ def _run_one_config( m=m, top_k=top_k, shared_experts=shared_experts, - reduce_results=reduce_results, gate=gate, routed_input_transform=routed_input_transform, routed_output_transform=routed_output_transform, @@ -1520,7 +1463,6 @@ def test_moe_layer_no_parallel( test_config.num_experts, test_config.top_k, test_config.quantization, - test_config.reduce_results, test_config.backend, _test_body_regular, use_shared_experts=test_config.use_shared_experts, @@ -1578,7 +1520,6 @@ def _parallel_worker( test_config.num_experts, test_config.top_k, test_config.quantization, - test_config.reduce_results, test_config.backend, functools.partial( _test_body_config, test_config=test_config, cpu_group=cpu_group @@ -1597,7 +1538,7 @@ def _parallel_worker( failed = failed + 1 if verbosity > 0: traceback.print_exc() - print(f"\n{str(ex)}\nFAILED {ex.__class__}") + print(f"\n{str(ex)}\nFAILED") else: print("F", end="") finally: diff --git a/tests/kernels/moe/test_shared_fused_moe_routed_transform.py b/tests/kernels/moe/test_shared_fused_moe_routed_transform.py index 89e23eb0d74..464754c9f1b 100644 --- a/tests/kernels/moe/test_shared_fused_moe_routed_transform.py +++ b/tests/kernels/moe/test_shared_fused_moe_routed_transform.py @@ -165,7 +165,6 @@ def test_routed_input_transform_inside_vs_outside( top_k=top_k, hidden_size=latent_size, intermediate_size=intermediate_size, - reduce_results=False, renormalize=True, params_dtype=dtype, tp_size=1, @@ -183,7 +182,6 @@ def test_routed_input_transform_inside_vs_outside( top_k=top_k, hidden_size=latent_size, intermediate_size=intermediate_size, - reduce_results=False, renormalize=True, params_dtype=dtype, tp_size=1, @@ -212,34 +210,20 @@ def test_routed_input_transform_inside_vs_outside( hidden_states = torch.randn(num_tokens, hidden_size, device="cuda", dtype=dtype) router_logits = torch.randn(num_tokens, num_experts, device="cuda", dtype=dtype) - # Clone inputs so any in-place modification by Method A - # cannot affect Method B's computation. - hidden_states_A = hidden_states.clone() - router_logits_A = router_logits.clone() - with set_forward_context(None, vllm_config, num_tokens=num_tokens): - shared_out_A, routed_out_A = moe_with_transform( - hidden_states_A, router_logits_A - ) + # Method A: combined output (shared + routed) + combined_A = moe_with_transform(hidden_states, router_logits) + # Method B: manually transform, get routed output, add shared transformed_hidden = routed_transform(hidden_states) - shared_out_B, routed_out_B = moe_without_transform( - transformed_hidden, router_logits - ) + routed_out_B = moe_without_transform(transformed_hidden, router_logits) + shared_out_B = shared_experts(hidden_states) + combined_B = shared_out_B + routed_out_B - expected_shared_out = shared_experts(hidden_states) - - _assert_close( - routed_out_A, - routed_out_B, + torch.testing.assert_close( + combined_A, + combined_B, atol=1e-3, rtol=1e-3, - label="Routed output: transform inside vs outside", - ) - _assert_close( - shared_out_A, - expected_shared_out, - atol=1e-3, - rtol=1e-3, - label="Shared expert output", + msg="Combined output should match: transform inside vs outside", ) diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index 835bffe58ca..d6eec675c6d 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -592,9 +592,6 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA): def forward(self, *args, **kwargs): return self.base_layer.forward(*args, **kwargs) - def maybe_all_reduce_tensor_model_parallel(self, *args, **kwargs): - return self.base_layer.maybe_all_reduce_tensor_model_parallel(*args, **kwargs) - @property def quant_method(self): return self.base_layer.quant_method diff --git a/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py b/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py index 6c916cf3cb6..daae5b6bd16 100644 --- a/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py @@ -716,7 +716,7 @@ class MarlinExperts(MarlinExpertsBase): ): assert self.w1_scale is not None assert self.w2_scale is not None - return fused_marlin_moe( + fused_marlin_moe( hidden_states=hidden_states, w1=w1, w2=w2, diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 190a9cc3b5d..bf10bc9d5c4 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -230,11 +230,18 @@ class FusedMoE(PluggableLayer): hidden_size: Input hidden state size of the transformer intermediate_size: Intermediate size of the experts params_dtype: Data type for the parameters. - reduce_results: Whether to all_reduce on the output of the layer renormalize: Whether to renormalize the logits in the fused_moe kernel quant_config: Quantization configure. enable_eplb: Whether to enable expert parallelism load balancer. router_logits_dtype: Data type for router logits buffers. + routed_scaling_factor: A scaling factor that is applied to the topk_weights + by the router or the output of the layer depending + on the value of `apply_routed_scale_to_output` + apply_routed_scale_to_output: Determine whether or not `routed_scaling_factor` + is applied to the topk_weights or to the experts + output. It is applied to the experts output + instead of the topk_weights when this feature is + not supported by the router (or the experts). """ # --8<-- [end:fused_moe] @@ -246,7 +253,6 @@ class FusedMoE(PluggableLayer): hidden_size: int, intermediate_size: int, params_dtype: torch.dtype | None = None, - reduce_results: bool = False, renormalize: bool = True, use_grouped_topk: bool = False, num_expert_group: int | None = None, @@ -274,12 +280,12 @@ class FusedMoE(PluggableLayer): gate: torch.nn.Module | None = None, shared_experts: torch.nn.Module | None = None, routed_input_transform: torch.nn.Module | None = None, + routed_output_transform: torch.nn.Module | None = None, + apply_routed_scale_to_output: bool = False, zero_expert_type: str | None = None, ): super().__init__() - self._routed_input_transform = routed_input_transform - if params_dtype is None: params_dtype = torch.get_default_dtype() self.params_dtype = params_dtype @@ -425,7 +431,6 @@ class FusedMoE(PluggableLayer): assert intermediate_size % self.tp_size == 0 intermediate_size_per_partition = intermediate_size // self.tp_size - self.reduce_results = reduce_results self.renormalize = renormalize # TODO(bnell): these attributes are only used by monolithic kernels. @@ -437,7 +442,14 @@ class FusedMoE(PluggableLayer): self.topk_group = topk_group self.custom_routing_function = custom_routing_function self.scoring_func = scoring_func - self.routed_scaling_factor = routed_scaling_factor + # When apply_routed_scale_to_output is True, we set the scaling factor + # to 1.0 so it ends up being a nop. Applying the scale will be handled + # by the runner in this case. + # The member variable must be set in the same way as the router since + # some quantization methods can access it. + self.routed_scaling_factor = ( + routed_scaling_factor if not apply_routed_scale_to_output else 1.0 + ) self.e_score_correction_bias = e_score_correction_bias # TODO(bnell): end attributes @@ -456,7 +468,7 @@ class FusedMoE(PluggableLayer): topk_group=topk_group, custom_routing_function=custom_routing_function, scoring_func=scoring_func, - routed_scaling_factor=routed_scaling_factor, + routed_scaling_factor=self.routed_scaling_factor, e_score_correction_bias=e_score_correction_bias, num_fused_shared_experts=self.num_fused_shared_experts, enable_eplb=enable_eplb, @@ -578,12 +590,18 @@ class FusedMoE(PluggableLayer): layer_name=self.layer_name, moe_config=self.moe_config, router=self.router, - routed_input_transform=self._routed_input_transform, gate=gate, shared_experts=shared_experts, quant_method=self.quant_method, - reduce_results=self.reduce_results, enable_dbo=self.vllm_config.parallel_config.enable_dbo, + routed_input_transform=routed_input_transform, + routed_output_transform=routed_output_transform, + # When apply_routed_scale_to_output is True, we allow + # the scaling factor to be passed to the runner, otherwise + # we pass 1.0 so it ends up being a nop. + routed_scaling_factor=routed_scaling_factor + if apply_routed_scale_to_output + else 1.0, ) # TODO(bnell): This method is provided as a hook so vllm/lora/layers/fused_moe.py @@ -1514,32 +1532,11 @@ class FusedMoE(PluggableLayer): self.ensure_moe_quant_config_init() return self.quant_method.moe_quant_config - def must_reduce_shared_expert_outputs(self) -> bool: - """ - The shared_experts are typically computed using the RowParallelLinear - layer. The result of this function is typically used as - the reduce_results argument to the module. - When just tensor-parallel is used, it is not required to reduce - the shared_experts results immediately. Instead we reduce at the - once at the end of the MoE op. (Refer to DeepSeekV2MoE module) - With EP and all2all kernels - this is no longer viable as all - GPU ranks in DP, produce the complete set of hidden_states. - Therefore it is required that we reduce the shared_experts output - early. - """ - return self.runner.must_reduce_shared_expert_outputs() - - def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor): - """ - Some combine kernels reduce across GPU ranks by default. - """ - return self.runner.maybe_all_reduce_tensor_model_parallel(final_hidden_states) - def forward( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: return self.runner.forward( hidden_states, router_logits, @@ -1613,7 +1610,6 @@ class FusedMoE(PluggableLayer): f"intermediate_size_per_partition={self.intermediate_size_per_partition}, " # noqa: E501 f"tp_size={self.tp_size},\n" f"ep_size={self.ep_size}, " - f"reduce_results={self.reduce_results}, " ) return s diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index f2e6e2560e7..376fcdf5b65 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1261,7 +1261,7 @@ class FusedMoEKernelModularImpl: topk_ids: torch.Tensor, apply_router_weight_on_input: bool, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: """ The _finalize method is a wrapper around self.prepare_finalize.finalize that handles DBO, async and shared expert overlap. diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 85c6563c084..8cd2fc65704 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -37,10 +37,6 @@ class DefaultMoERunner(MoERunnerBase): for different configurations (e.g., with/without shared experts, gates, etc.). """ - @property - def reduce_results(self) -> bool: - return self._reduce_results - @property def do_naive_dispatch_combine(self) -> bool: return ( 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 9ffbf3108f8..199ceab0659 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -26,18 +26,7 @@ class MoERunner(ABC): self, hidden_states: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - raise NotImplementedError - - @abstractmethod - def must_reduce_shared_expert_outputs(self) -> bool: - raise NotImplementedError - - @abstractmethod - def maybe_all_reduce_tensor_model_parallel( - self, - final_hidden_states: torch.Tensor, - ): + ) -> torch.Tensor: raise NotImplementedError @abstractmethod diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index 692d45d3460..2d2d0bb4ebe 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -81,7 +81,9 @@ def _resolve_layer_name(layer_name: str | LayerName) -> str: # Note: _moe_forward and _moe_forward_shared should not contain any # implementation details, They should merely pass along control to -# the runner's 'forward_dispatch' method. +# the runner's '_forward_dispatch' method. +# These functions should never be called directly since they do not +# include all the functionality of the MoE layer. def _moe_forward( hidden_states: torch.Tensor, router_logits: torch.Tensor, @@ -89,7 +91,7 @@ def _moe_forward( layer_name: _layer_name_type, ) -> torch.Tensor: layer = get_layer_from_name(_resolve_layer_name(layer_name)) - return layer.runner.forward_dispatch( + return layer.runner._forward_dispatch( layer, hidden_states, router_logits, @@ -113,7 +115,7 @@ def _moe_forward_shared( layer_name: _layer_name_type, ) -> tuple[torch.Tensor, torch.Tensor]: layer = get_layer_from_name(_resolve_layer_name(layer_name)) - return layer.runner.forward_dispatch( + return layer.runner._forward_dispatch( layer, hidden_states, router_logits, @@ -143,7 +145,7 @@ def _moe_forward_shared_fake( direct_register_custom_op( op_name="moe_forward", op_func=_moe_forward, - mutates_args=["hidden_states"], # is this still true? + mutates_args=["hidden_states"], fake_impl=_moe_forward_fake, tags=(torch.Tag.needs_fixed_stride_order,), ) @@ -157,6 +159,15 @@ direct_register_custom_op( ) +def _unpack( + result: torch.Tensor | tuple[torch.Tensor, torch.Tensor], +) -> tuple[torch.Tensor | None, torch.Tensor]: + if isinstance(result, tuple): + return result + else: + return (None, result) + + class MoERunnerBase(MoERunner): """ Abstract base class providing common functionality for MoE runner implementations. @@ -174,7 +185,6 @@ class MoERunnerBase(MoERunner): allowing flexibility in the actual MoE computation implementation. Key abstract methods that subclasses must implement: - - reduce_results: Determines whether results should be reduced across ranks - _forward_impl: The core MoE computation logic specific to each runner type """ @@ -187,17 +197,23 @@ class MoERunnerBase(MoERunner): gate: torch.nn.Module | None, shared_experts: torch.nn.Module | None, quant_method: FusedMoEMethodBase, - reduce_results: bool, enable_dbo: bool, + routed_output_transform: torch.nn.Module | None = None, + routed_scaling_factor: float = 1.0, ): super().__init__() self.moe_config = moe_config self.router = router self.routed_input_transform = routed_input_transform + self.routed_output_transform = routed_output_transform + self.routed_scaling_factor = routed_scaling_factor self.gate = gate self.quant_method = quant_method - self._reduce_results = reduce_results self.enable_dbo = enable_dbo + self._fused_output_is_reduced = ( + self.quant_method.moe_kernel is not None + and self.quant_method.moe_kernel.output_is_reduced() + ) self._shared_experts: SharedExperts | None = None if shared_experts is not None: @@ -209,7 +225,6 @@ class MoERunnerBase(MoERunner): # called, i.e. by a MK or by the MoERunner. # Once the MK can be created upfront, we can just pass in the proper # flags derived from the quant_method's MK. - reduce_results=reduce_results, quant_method=quant_method, enable_dbo=enable_dbo, ) @@ -217,7 +232,7 @@ class MoERunnerBase(MoERunner): # Needed for string -> FusedMoE layer lookup in custom ops. self.layer_name = layer_name - self.forward_entry = self._select_forward() + self._forward_entry = self._select_forward() def _select_forward(self) -> Callable: if current_platform.is_tpu() or current_platform.is_cpu(): @@ -245,38 +260,6 @@ class MoERunnerBase(MoERunner): def is_internal_router(self) -> bool: return self.gate is not None - @property - @abstractmethod - def reduce_results(self) -> bool: - raise NotImplementedError - - def must_reduce_shared_expert_outputs(self) -> bool: - """ - The shared_experts are typically computed using the RowParallelLinear - layer. The result of this function is typically used as - the reduce_results argument to the module. - When just tensor-parallel is used, it is not required to reduce - the shared_experts results immediately. Instead we reduce at the - once at the end of the MoE op. (Refer to DeepSeekV2MoE module) - With EP and all2all kernels - this is no longer viable as all - GPU ranks in DP, produce the complete set of hidden_states. - Therefore it is required that we reduce the shared_experts output - early. - """ - return ( - self.quant_method.moe_kernel is not None - and self.quant_method.moe_kernel.output_is_reduced() - ) - - def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor): - """ - Some combine kernels reduce across GPU ranks by default. - """ - if self.must_reduce_shared_expert_outputs(): - return final_hidden_states - else: - return tensor_model_parallel_all_reduce(final_hidden_states) - def apply_routed_input_transform( self, hidden_states: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor | None]: @@ -286,10 +269,6 @@ class MoERunnerBase(MoERunner): is saved separately so shared experts get [S, hidden_size] while routed experts get the transformed [S, moe_latent_size]. - TODO: For latent MoE bandwidth optimization, fc2_latent_proj could be - moved inside SharedFusedMoE to all-reduce on the smaller latent - dimension. - Returns (possibly transformed) hidden states and the input for shared experts (or None if there are no shared experts). """ @@ -306,33 +285,79 @@ class MoERunnerBase(MoERunner): hidden_states if self._shared_experts is not None else None, ) - def _maybe_reduce_output( + def apply_routed_output_transform( self, - states: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - trunc_sizes: list[int], - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - def trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor: - return x[..., :trunc_size] + fused_output: torch.Tensor, + ) -> torch.Tensor: + """Apply transform to routed expert output (e.g., latent to full dim). - def reduce_and_trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor: - return trunc(self.maybe_all_reduce_tensor_model_parallel(x), trunc_size) + Used by latent MoE models (e.g., NemotronH) where routed experts + operate in a compressed latent space and need projection back to + the full hidden dimension before combining with shared expert output. + """ + if self.routed_output_transform is not None: + r = self.routed_output_transform(fused_output) + fused_output = r[0] if isinstance(r, tuple) else r + return fused_output + def _maybe_apply_routed_scale_to_output( + self, + shared_output: torch.Tensor | None, + fused_output: torch.Tensor, + ) -> tuple[torch.Tensor | None, torch.Tensor]: + """Apply routed_scaling_factor to the output with FP16 overflow + protection. + + Scale the fused expert output by routed_scaling_factor. For FP16, + avoid overflow by dividing shared_output by the scale instead + (the decoder layer compensates with matching divisions). + """ + if self.routed_scaling_factor != 1.0: + if fused_output.dtype != torch.float16: + fused_output *= self.routed_scaling_factor + elif shared_output is not None: + shared_output *= 1.0 / self.routed_scaling_factor + return shared_output, fused_output + + def _maybe_reduce_shared_expert_output( + self, + shared_output: torch.Tensor | None, + ) -> torch.Tensor | None: + """All-reduce shared expert output when the combine kernel already + reduced fused output. + + This is the "early" all-reduce path. When the combine kernel produces + already-reduced fused output, shared output must be reduced separately + to match. + """ + if self._fused_output_is_reduced: + assert shared_output is not None + shared_output = tensor_model_parallel_all_reduce(shared_output) + return shared_output + + def _maybe_reduce_final_output( + self, + states: torch.Tensor, + trunc_size: int, + ) -> torch.Tensor: + """Truncate padded dimensions and all-reduce the combined output. + + This is the "late" all-reduce path. When neither fused nor shared + output was individually reduced, the combined sum is all-reduced + here. Skipped when sequence-parallel is active (SP handles its + own reduction) or when the early path already reduced both outputs. + """ + # 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. if ( not self.moe_config.is_sequence_parallel - and self.reduce_results and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) + and not self._fused_output_is_reduced ): - func = reduce_and_trunc - else: - func = trunc + states = tensor_model_parallel_all_reduce(states) - if isinstance(states, tuple): - return tuple( - [func(s, trunc_size) for s, trunc_size in zip(states, trunc_sizes)] - ) - else: - assert len(trunc_sizes) == 1 - return func(states, trunc_sizes[0]) + return states[..., :trunc_size] def _encode_layer_name(self) -> str | LayerName: if _USE_LAYERNAME: @@ -349,7 +374,15 @@ class MoERunnerBase(MoERunner): self, shared_experts_input: torch.Tensor | None, hidden_states: torch.Tensor, - ) -> tuple[torch.Tensor, list[int]]: + ) -> tuple[torch.Tensor, int]: + """Pad hidden_states to moe_config.hidden_dim and compute the + original dimension for later truncation. + + For latent MoE, the routed hidden_states may be smaller than + hidden_dim. Padding ensures uniform tensor sizes through the + fused MoE kernel. The returned trunc_size is used by + _maybe_reduce_final_output to strip the padding from the result. + """ shared_experts_hidden_dim = ( shared_experts_input.shape[-1] if shared_experts_input is not None else 0 ) @@ -365,10 +398,10 @@ class MoERunnerBase(MoERunner): value=0.0, ) - if self._shared_experts is not None: - orig_hidden_dims = [shared_experts_hidden_dim, transformed_hidden_dim] + if self.routed_output_transform is not None and shared_experts_hidden_dim > 0: + orig_hidden_dims = shared_experts_hidden_dim else: - orig_hidden_dims = [transformed_hidden_dim] + orig_hidden_dims = transformed_hidden_dim return hidden_states, orig_hidden_dims @@ -388,6 +421,12 @@ class MoERunnerBase(MoERunner): router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, ) -> tuple[torch.Tensor | None, torch.Tensor]: + """Run expert routing and the fused MoE kernel via the quant method. + + Orchestrates shared expert execution (before/after), expert selection + via the router, and the actual fused MoE computation. Returns + (shared_expert_output, fused_expert_output). + """ # Run this before quant_method to avoid inplace issues. # TODO(bnell): probably not needed anymore since inplace is # disabled when shared experts are present. @@ -428,6 +467,13 @@ class MoERunnerBase(MoERunner): ) def _sequence_parallel_context(self): + """Return a context manager for sequence-parallel token + redistribution. + + When sequence parallelism is active, returns a context that handles + local size tracking for proper token scatter/gather. Otherwise + returns a no-op context. + """ ctx = get_forward_context() return ( ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size) @@ -448,22 +494,25 @@ class MoERunnerBase(MoERunner): def _maybe_add_zero_expert_output( self, - result: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + result: torch.Tensor, + ) -> torch.Tensor: + """Add the zero expert's contribution to the final result. + + When a ZeroExpertRouter is used, it computes a bias-like output + from the "zero expert" that is added to the combined routed+shared + expert output. + """ if isinstance(self.router, ZeroExpertRouter): zero_expert_output = self.router.zero_expert_output assert zero_expert_output is not None - if isinstance(result, tuple): - result = (result[0], result[1] + zero_expert_output) - else: - result = result + zero_expert_output + result = result + zero_expert_output return result def forward( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: """Invoke the fused moe layer. Input: @@ -472,55 +521,88 @@ class MoERunnerBase(MoERunner): Output: - The new hidden_states. - or - - A tuple of (shared experts output, new hidden_states). Calling sequence - forward - - self.forward_entry (_moe_forward or _moe_forward_shared custom op) - - forward_dispatch + - self._forward_entry (_moe_forward or _moe_forward_shared custom op) + - _forward_dispatch - _forward_impl Note: The existence of _moe_forward and _moe_forward_shared custom ops are due to the following reasons: 1. the chunking loop in ChunkingMoERunner._forward_impl cannot be compiled by torch.compile - 2. pytorch cannot handle union types in custom op signatures so _moe_forward - and _moe_forward_shared must be split. + 2. pytorch cannot handle union types in custom op signatures so + _moe_forward and _moe_forward_shared must be split. If ChunkingMoERunner._forward_impl can be implemented via torch.scan we can potentially get rid of _moe_forward and _moe_forward_shared and collapse the whole sequence into the 'forward' method. """ - # Apply transform for routed experts (e.g., latent projection for latent MoE) + # Apply transform for routed experts (e.g., latent projection + # for latent MoE) hidden_states, shared_experts_input = self.apply_routed_input_transform( hidden_states ) - hidden_states, og_hidden_dims = self._maybe_pad_hidden_states( + hidden_states, og_hidden_dim = self._maybe_pad_hidden_states( shared_experts_input, hidden_states, ) - fused_output = self.forward_entry( + result = self._forward_entry( hidden_states, router_logits, shared_experts_input, self._encode_layer_name(), ) - result = self._maybe_reduce_output(fused_output, og_hidden_dims) + # + # Note: there are two all-reduce points below. They are mutually + # exclusive, controlled by _fused_output_is_reduced + # - When True: the combine kernel already reduced fused_output, + # so we reduce shared_output here to match, then skip the + # all-reduce in _maybe_reduce_final_output. + # - When False: neither output is reduced yet, so we combine + # them first and all-reduce the sum in _maybe_reduce_final_output. + + # Extract outputs from result + shared_output, fused_output = _unpack(result) + + # If combine kernel already reduced fused, reduce shared to match. + # See note above re: the two all-reduce points. + shared_output = self._maybe_reduce_shared_expert_output(shared_output) + + shared_output, fused_output = self._maybe_apply_routed_scale_to_output( + shared_output, fused_output + ) + + # Apply output transform (e.g. latent -> full dim) + fused_output = self.apply_routed_output_transform(fused_output) + + if shared_output is not None: + result = shared_output + fused_output + else: + result = fused_output + + result = self._maybe_reduce_final_output(result, og_hidden_dim) return self._maybe_add_zero_expert_output(result) - def forward_dispatch( + def _forward_dispatch( self, layer: torch.nn.Module, hidden_states: torch.Tensor, router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Entry point called by the custom op to run the MoE computation. + + Handles pre-dispatch setup (gate application, external shared expert + triggering, quant config init) then delegates to _forward_impl within + the sequence-parallel context. + """ # TODO(bnell): this can be removed after MK migration is complete. layer.ensure_moe_quant_config_init() @@ -549,4 +631,11 @@ class MoERunnerBase(MoERunner): router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Core MoE computation to be implemented by subclasses. + + Performs expert routing, fused MoE kernel execution, and shared + expert computation. Returns a single tensor (fused output only) + or a tuple of (shared_output, fused_output) when shared experts + are present. + """ raise NotImplementedError diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py index 2143fa3ce08..feb4614d837 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py @@ -29,8 +29,9 @@ def create_moe_runner( gate: torch.nn.Module | None, shared_experts: SharedExperts | None, quant_method: FusedMoEMethodBase, - reduce_results: bool, enable_dbo: bool, + routed_output_transform: torch.nn.Module | None = None, + routed_scaling_factor: float = 1.0, ) -> MoERunner: return DefaultMoERunner( layer_name, @@ -40,6 +41,7 @@ def create_moe_runner( gate, shared_experts, quant_method, - reduce_results, enable_dbo, + routed_output_transform=routed_output_transform, + routed_scaling_factor=routed_scaling_factor, ) diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index 827a6e6bd3e..3a08d5d0fc2 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -5,10 +5,6 @@ from enum import IntEnum import torch import vllm.envs as envs -from vllm.distributed import ( - get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, -) from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, @@ -48,7 +44,6 @@ class SharedExperts: layer: torch.nn.Module, moe_config: FusedMoEConfig, quant_method: QuantizeMethodBase, - reduce_results: bool, enable_dbo: bool, ): from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( @@ -68,7 +63,6 @@ class SharedExperts: self._layer = layer self._moe_config = moe_config self._quant_method = quant_method - self._reduce_results = reduce_results # Allow disabling of the separate shared experts stream for # debug purposes. @@ -139,18 +133,6 @@ class SharedExperts: return output - def _maybe_reduce_shared_out(self, shared_out: torch.Tensor) -> torch.Tensor: - # Reduce shared expert outputs if necessary, since the MLP - # should have been created with reduce_results=False. - if ( - self._reduce_results - and self._quant_method.moe_kernel is not None - and self._quant_method.moe_kernel.output_is_reduced() - and get_tensor_model_parallel_world_size() > 1 - ): - shared_out = tensor_model_parallel_all_reduce(shared_out) - return shared_out - @property def _output_idx(self) -> int: return dbo_current_ubatch_id() if self.enable_dbo else 0 diff --git a/vllm/model_executor/layers/fused_moe/shared_fused_moe.py b/vllm/model_executor/layers/fused_moe/shared_fused_moe.py index ed243e99263..9cfcb1baa9b 100644 --- a/vllm/model_executor/layers/fused_moe/shared_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/shared_fused_moe.py @@ -18,12 +18,8 @@ class SharedFusedMoE(FusedMoE): self, hidden_states: torch.Tensor, router_logits: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - result = super().forward( + ) -> torch.Tensor: + return super().forward( hidden_states=hidden_states, router_logits=router_logits, ) - if self.shared_experts is None: - return None, result - else: - return result diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py index f5ed4400fb6..d42fbed42ae 100644 --- a/vllm/model_executor/models/AXK1.py +++ b/vllm/model_executor/models/AXK1.py @@ -100,7 +100,7 @@ class AXK1MoE(nn.Module): self.tp_size = get_tensor_model_parallel_world_size() self.tp_rank = get_tensor_model_parallel_rank() - self.routed_scaling_factor = config.routed_scaling_factor + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) self.ep_group = get_ep_group().device_group self.ep_rank = get_ep_group().rank_in_group @@ -170,7 +170,6 @@ class AXK1MoE(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, @@ -180,9 +179,8 @@ class AXK1MoE(nn.Module): scoring_func=config.scoring_func, # we do scaling outside, set factor to 1.0 to avoid double mul # aiter applies routed_scaling_factor internally - routed_scaling_factor=1.0 - if not self.is_rocm_aiter_moe_enabled - else self.routed_scaling_factor, + routed_scaling_factor=self.routed_scaling_factor, + apply_routed_scale_to_output=not self.is_rocm_aiter_moe_enabled, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, @@ -204,43 +202,20 @@ class AXK1MoE(nn.Module): hidden_states = sequence_parallel_chunk(hidden_states) if self.experts.is_internal_router: - # In this case, the gate/router runs inside the FusedMoE class - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=hidden_states ) else: - # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - shared_output, final_hidden_states = fused_moe_out - if self.shared_experts is None: - assert shared_output is None - - # Fix FP16 overflow - # See AXK1DecoderLayer for more details. - if hidden_states.dtype != torch.float16: - if not self.is_rocm_aiter_moe_enabled: - final_hidden_states *= self.routed_scaling_factor - elif self.shared_experts is not None: - assert shared_output is not None - shared_output *= 1.0 / self.routed_scaling_factor - - if self.shared_experts is not None: - assert shared_output is not None - final_hidden_states += shared_output - if self.is_sequence_parallel: final_hidden_states = tensor_model_parallel_all_gather( final_hidden_states, 0 ) final_hidden_states = final_hidden_states[:num_tokens] - elif self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/afmoe.py b/vllm/model_executor/models/afmoe.py index 22037336411..5bad52a0c49 100644 --- a/vllm/model_executor/models/afmoe.py +++ b/vllm/model_executor/models/afmoe.py @@ -131,7 +131,6 @@ class AfmoeMoE(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=self.route_norm if self.score_func == "sigmoid" else False, quant_config=quant_config, use_grouped_topk=True, @@ -152,20 +151,10 @@ class AfmoeMoE(nn.Module): router_logits = self.gate(hidden_states.to(dtype=torch.float32)) - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - if self.shared_experts is not None: - shared_output, final_hidden_states = fused_moe_out - final_hidden_states = final_hidden_states + shared_output - else: - final_hidden_states = fused_moe_out - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) - return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/aria.py b/vllm/model_executor/models/aria.py index 7b891f8ee42..7a079c56540 100644 --- a/vllm/model_executor/models/aria.py +++ b/vllm/model_executor/models/aria.py @@ -283,7 +283,6 @@ class AriaTextMoELayer(nn.Module): hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, quant_config=quant_config, - reduce_results=True, prefix=f"{prefix}.experts", ) @@ -301,12 +300,7 @@ class AriaTextMoELayer(nn.Module): router_output = torch.nn.functional.linear(hidden_states, self.router_weight) - sparse_expert_output = self.experts(hidden_states, router_output) - - if self.shared_experts is not None: - return sparse_expert_output[0] + sparse_expert_output[1] - else: - return sparse_expert_output + return self.experts(hidden_states, router_output) class AriaTextDecoderLayer(LlamaDecoderLayer): diff --git a/vllm/model_executor/models/bailing_moe.py b/vllm/model_executor/models/bailing_moe.py index 7725dfa2a88..510d605f804 100644 --- a/vllm/model_executor/models/bailing_moe.py +++ b/vllm/model_executor/models/bailing_moe.py @@ -291,7 +291,6 @@ class BailingMoE(nn.Module): top_k=self.top_k, hidden_size=self.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=self.norm_expert_prob, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -301,6 +300,7 @@ class BailingMoE(nn.Module): topk_group=self.topk_group, use_grouped_topk=self.use_grouped_topk, router_logits_dtype=self.router_dtype, + routed_scaling_factor=self.routed_scaling_factor, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -314,21 +314,6 @@ class BailingMoE(nn.Module): final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - - if self.shared_experts is not None: - shared_output, final_hidden_states = final_hidden_states - else: - shared_output = None - - final_hidden_states *= self.routed_scaling_factor - - if shared_output is not None: - final_hidden_states = final_hidden_states + shared_output - - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) return final_hidden_states.view(num_tokens, hidden_size) diff --git a/vllm/model_executor/models/bailing_moe_linear.py b/vllm/model_executor/models/bailing_moe_linear.py index ecc5d63ced7..a63ad83f45b 100644 --- a/vllm/model_executor/models/bailing_moe_linear.py +++ b/vllm/model_executor/models/bailing_moe_linear.py @@ -358,7 +358,6 @@ class BailingMoeV25(nn.Module): top_k=self.top_k, hidden_size=self.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=self.norm_expert_prob, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -368,6 +367,8 @@ class BailingMoeV25(nn.Module): topk_group=self.topk_group, use_grouped_topk=self.use_grouped_topk, router_logits_dtype=self.router_dtype, + routed_scaling_factor=self.routed_scaling_factor, + apply_routed_scale_to_output=True, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -383,22 +384,6 @@ class BailingMoeV25(nn.Module): hidden_states=hidden_states, router_logits=router_logits ) - # Handle tuple return from SharedFusedMoE - if self.shared_experts is not None: - shared_output, final_hidden_states = final_hidden_states - else: - shared_output = None - - final_hidden_states *= self.routed_scaling_factor - - if shared_output is not None: - final_hidden_states = final_hidden_states + shared_output - - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) - return final_hidden_states.view(num_tokens, hidden_size) diff --git a/vllm/model_executor/models/dbrx.py b/vllm/model_executor/models/dbrx.py index ca6e6a49a98..a72f4e48716 100644 --- a/vllm/model_executor/models/dbrx.py +++ b/vllm/model_executor/models/dbrx.py @@ -85,7 +85,6 @@ class DbrxExperts(FusedMoE): hidden_size=config.d_model, intermediate_size=config.ffn_config.ffn_hidden_size, params_dtype=params_dtype, - reduce_results=True, renormalize=True, quant_config=quant_config, tp_size=get_tensor_model_parallel_world_size(), diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index cd28fb0192f..1b01caded94 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -318,7 +318,6 @@ class DeepseekV2MoE(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, @@ -326,11 +325,9 @@ class DeepseekV2MoE(nn.Module): topk_group=getattr(config, "topk_group", 1), prefix=f"{prefix}.experts", scoring_func=getattr(config, "scoring_func", "softmax"), - # we do scaling outside, set factor to 1.0 to avoid double mul # aiter applies routed_scaling_factor internally - routed_scaling_factor=1.0 - if not self.is_rocm_aiter_moe_enabled - else self.routed_scaling_factor, + routed_scaling_factor=self.routed_scaling_factor, + apply_routed_scale_to_output=not self.is_rocm_aiter_moe_enabled, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, @@ -363,43 +360,20 @@ class DeepseekV2MoE(nn.Module): hidden_states = sequence_parallel_chunk(hidden_states) if self.experts.is_internal_router: - # In this case, the gate/router runs inside the FusedMoE class - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=hidden_states ) else: - # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - shared_output, final_hidden_states = fused_moe_out - if self.shared_experts is None: - assert shared_output is None - - # Fix FP16 overflow - # See DeepseekV2DecoderLayer for more details. - if hidden_states.dtype != torch.float16: - if not self.is_rocm_aiter_moe_enabled: - final_hidden_states *= self.routed_scaling_factor - elif self.shared_experts is not None: - assert shared_output is not None - shared_output *= 1.0 / self.routed_scaling_factor - - if self.shared_experts is not None: - assert shared_output is not None - final_hidden_states += shared_output - if self.is_sequence_parallel: final_hidden_states = tensor_model_parallel_all_gather( final_hidden_states, 0 ) final_hidden_states = final_hidden_states[:num_tokens] - elif self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/dots1.py b/vllm/model_executor/models/dots1.py index 4e393145462..c176b736568 100644 --- a/vllm/model_executor/models/dots1.py +++ b/vllm/model_executor/models/dots1.py @@ -37,7 +37,6 @@ from vllm.config import CacheConfig, ModelConfig, VllmConfig from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, ) from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention @@ -120,7 +119,6 @@ class Dots1MoE(nn.Module): prefix: str = "", ): super().__init__() - self.tp_size = get_tensor_model_parallel_world_size() self.routed_scaling_factor = config.routed_scaling_factor self.n_shared_experts = config.n_shared_experts @@ -163,7 +161,6 @@ class Dots1MoE(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, @@ -171,9 +168,9 @@ class Dots1MoE(nn.Module): topk_group=config.topk_group, prefix=f"{prefix}.experts", scoring_func=config.scoring_func, - # we do scaling outside, set factor to 1.0 to avoid double mul - routed_scaling_factor=1.0, e_score_correction_bias=self.gate.e_score_correction_bias, + routed_scaling_factor=self.routed_scaling_factor, + apply_routed_scale_to_output=True, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -182,16 +179,9 @@ class Dots1MoE(nn.Module): router_logits, _ = self.gate(hidden_states) - shared_out, routed_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - if self.shared_experts is not None: - final_hidden_states = (routed_out + shared_out) * self.routed_scaling_factor - else: - final_hidden_states = routed_out * self.routed_scaling_factor - - if self.tp_size > 1: - final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/ernie45_moe.py b/vllm/model_executor/models/ernie45_moe.py index f038cfb21f2..c92e230bcd2 100644 --- a/vllm/model_executor/models/ernie45_moe.py +++ b/vllm/model_executor/models/ernie45_moe.py @@ -194,7 +194,6 @@ class Ernie4_5_MoeMoE(nn.Module): top_k=config.moe_k, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=True, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -215,16 +214,6 @@ class Ernie4_5_MoeMoE(nn.Module): hidden_states=hidden_states, router_logits=router_logits ) - if self.has_shared_experts: - final_hidden_states = final_hidden_states[0] + final_hidden_states[1] - else: - final_hidden_states = final_hidden_states[1] - - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) - return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/ernie45_vl_moe.py b/vllm/model_executor/models/ernie45_vl_moe.py index 418fdcfa072..e4b7ac6fb00 100644 --- a/vllm/model_executor/models/ernie45_vl_moe.py +++ b/vllm/model_executor/models/ernie45_vl_moe.py @@ -263,7 +263,6 @@ class Ernie4_5_VLMoeMoE(nn.Module): top_k=config.moe_k, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size[0], - reduce_results=False, renormalize=True, quant_config=quant_config, e_score_correction_bias=self.e_score_correction_bias[0], @@ -301,7 +300,6 @@ class Ernie4_5_VLMoeMoE(nn.Module): top_k=config.moe_k, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size[1], - reduce_results=False, renormalize=True, quant_config=quant_config, e_score_correction_bias=self.e_score_correction_bias[1], @@ -342,9 +340,6 @@ class Ernie4_5_VLMoeMoE(nn.Module): visual_token_mask = visual_token_mask.repeat(1, self.hidden_size).bool() text_token_mask = ~visual_token_mask final_experts_hidden_states = torch.zeros_like(hidden_states) - final_shared_output = ( - torch.zeros_like(hidden_states) if self.has_shared_experts else None - ) text_hidden_states = hidden_states[text_token_mask].reshape( -1, self.hidden_size @@ -356,26 +351,20 @@ class Ernie4_5_VLMoeMoE(nn.Module): text_router_logits, _ = self.text_experts_gate( text_hidden_states.to(dtype=torch.float32) ) - text_shared_output, text_experts_output = self.text_experts( + text_output = self.text_experts( hidden_states=text_hidden_states, router_logits=text_router_logits ) - final_experts_hidden_states[text_token_mask] = text_experts_output.flatten() - if self.has_shared_experts: - final_shared_output[text_token_mask] = text_shared_output.flatten() + final_experts_hidden_states[text_token_mask] = text_output.flatten() vision_router_logits, _ = self.vision_experts_gate( vision_hidden_states.to(dtype=torch.float32) ) - vision_shared_output, vision_experts_output = self.vision_experts( + vision_output = self.vision_experts( hidden_states=vision_hidden_states, router_logits=vision_router_logits ) - final_experts_hidden_states[visual_token_mask] = ( - vision_experts_output.flatten() - ) - if self.has_shared_experts: - final_shared_output[visual_token_mask] = vision_shared_output.flatten() + final_experts_hidden_states[visual_token_mask] = vision_output.flatten() - final_hidden_states = (final_shared_output, final_experts_hidden_states) + final_hidden_states = final_experts_hidden_states else: # only text modal input text_router_logits, _ = self.text_experts_gate( @@ -386,20 +375,6 @@ class Ernie4_5_VLMoeMoE(nn.Module): hidden_states=hidden_states, router_logits=text_router_logits ) - if self.has_shared_experts: - # for shared_experts model - final_hidden_states = final_hidden_states[0] + final_hidden_states[1] - else: - # for not shared_experts model - final_hidden_states = final_hidden_states[1] - - if self.tp_size > 1: - final_hidden_states = ( - self.text_experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) - ) - return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/exaone_moe.py b/vllm/model_executor/models/exaone_moe.py index d7282edcf4f..a46cadf007e 100644 --- a/vllm/model_executor/models/exaone_moe.py +++ b/vllm/model_executor/models/exaone_moe.py @@ -31,6 +31,7 @@ from vllm.distributed import ( get_tensor_model_parallel_world_size, ) from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -116,12 +117,26 @@ class ExaoneMoe(nn.Module): self.physical_expert_start + self.n_local_physical_experts ) - self.experts = FusedMoE( + if getattr(config, "num_shared_experts", 0) > 0: + intermediate_size = config.moe_intermediate_size * config.num_shared_experts + self.shared_experts = ExaoneMoeGatedMLP( + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.shared_experts", + ) + else: + self.shared_experts = None + + self.experts = SharedFusedMoE( + shared_experts=self.shared_experts, + gate=self.gate, num_experts=self.n_routed_experts, top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, @@ -135,41 +150,16 @@ class ExaoneMoe(nn.Module): num_redundant_experts=self.n_redundant_experts, ) - if getattr(config, "num_shared_experts", 0) > 0: - intermediate_size = config.moe_intermediate_size * config.num_shared_experts - self.shared_experts = ExaoneMoeGatedMLP( - hidden_size=config.hidden_size, - intermediate_size=intermediate_size, - hidden_act=config.hidden_act, - quant_config=quant_config, - reduce_results=self.experts.must_reduce_shared_expert_outputs(), - prefix=f"{prefix}.shared_experts", - ) - else: - self.shared_experts = None - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # NOTE: hidden_states can have either 1D or 2D shape. orig_shape = hidden_states.shape hidden_dim = hidden_states.shape[-1] hidden_states = hidden_states.view(-1, hidden_dim) - # router_logits: (num_tokens, n_experts) - router_logits, _ = self.gate(hidden_states) - final_hidden_states = self.experts( - hidden_states=hidden_states, router_logits=router_logits + hidden_states=hidden_states, router_logits=hidden_states ) - if self.shared_experts is not None: - shared_output = self.shared_experts(hidden_states) - final_hidden_states = final_hidden_states + shared_output - - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501 - final_hidden_states - ) - return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/flex_olmo.py b/vllm/model_executor/models/flex_olmo.py index 67be99a879f..1b2047eb231 100644 --- a/vllm/model_executor/models/flex_olmo.py +++ b/vllm/model_executor/models/flex_olmo.py @@ -76,7 +76,6 @@ class FlexOlmoMoE(nn.Module): top_k=hf_config.num_experts_per_tok, hidden_size=hf_config.hidden_size, intermediate_size=hf_config.intermediate_size, - reduce_results=True, renormalize=False, quant_config=None, tp_size=tp_size, diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py index d166a9df38a..42762e36f81 100644 --- a/vllm/model_executor/models/gemma4.py +++ b/vllm/model_executor/models/gemma4.py @@ -349,7 +349,6 @@ class Gemma4MoE(nn.Module): "moe_intermediate_size", getattr(config, "expert_intermediate_size", None), ), - reduce_results=True, renormalize=True, quant_config=quant_config, prefix=f"{prefix}.experts", diff --git a/vllm/model_executor/models/glm4_moe.py b/vllm/model_executor/models/glm4_moe.py index d0e6cb6ada8..671e868da0a 100644 --- a/vllm/model_executor/models/glm4_moe.py +++ b/vllm/model_executor/models/glm4_moe.py @@ -184,7 +184,6 @@ class Glm4MoE(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, @@ -192,8 +191,8 @@ class Glm4MoE(nn.Module): topk_group=config.topk_group, prefix=f"{prefix}.experts", scoring_func="sigmoid", - # we do scaling outside, set factor to 1.0 to avoid double mul - routed_scaling_factor=1.0, + routed_scaling_factor=self.routed_scaling_factor, + apply_routed_scale_to_output=True, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, @@ -207,23 +206,9 @@ class Glm4MoE(nn.Module): # router_logits: (num_tokens, n_experts) router_logits = self.gate(hidden_states.to(dtype=torch.float32)) - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - - if self.shared_experts is not None: - shared_output, final_hidden_states = fused_moe_out - assert shared_output is not None - final_hidden_states = ( - final_hidden_states * self.routed_scaling_factor + shared_output - ) - else: - final_hidden_states = fused_moe_out * self.routed_scaling_factor - - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/gpt_oss.py b/vllm/model_executor/models/gpt_oss.py index 4e4eb581842..b6edc344302 100644 --- a/vllm/model_executor/models/gpt_oss.py +++ b/vllm/model_executor/models/gpt_oss.py @@ -189,7 +189,6 @@ class MLPBlock(torch.nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, - reduce_results=True, renormalize=True, quant_config=quant_config, prefix=f"{prefix}.experts", diff --git a/vllm/model_executor/models/granitemoe.py b/vllm/model_executor/models/granitemoe.py index 171b2e0ec5a..f57a8c942bb 100644 --- a/vllm/model_executor/models/granitemoe.py +++ b/vllm/model_executor/models/granitemoe.py @@ -104,7 +104,6 @@ class GraniteMoeMoE(nn.Module): hidden_size=hidden_size, intermediate_size=intermediate_size, params_dtype=params_dtype, - reduce_results=True, renormalize=True, quant_config=quant_config, tp_size=tp_size, diff --git a/vllm/model_executor/models/grok1.py b/vllm/model_executor/models/grok1.py index 0bd6a8f3d60..c9aa3d2068f 100644 --- a/vllm/model_executor/models/grok1.py +++ b/vllm/model_executor/models/grok1.py @@ -209,7 +209,6 @@ class Grok1MoE(nn.Module): hidden_size=hidden_size, intermediate_size=intermediate_size, params_dtype=params_dtype, - reduce_results=True, renormalize=renormalize, quant_config=quant_config, tp_size=tp_size, diff --git a/vllm/model_executor/models/hunyuan_v1.py b/vllm/model_executor/models/hunyuan_v1.py index a0130402c66..35d30006a66 100644 --- a/vllm/model_executor/models/hunyuan_v1.py +++ b/vllm/model_executor/models/hunyuan_v1.py @@ -39,7 +39,6 @@ from vllm.distributed import ( get_ep_group, get_pp_group, get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, ) from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention @@ -445,7 +444,6 @@ class HunYuanSparseMoeBlock(nn.Module): top_k=top_k, hidden_size=config.hidden_size, intermediate_size=intermediate_size, - reduce_results=False, renormalize=top_k > 1, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -464,11 +462,6 @@ class HunYuanSparseMoeBlock(nn.Module): final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - if self.shared_mlp is not None: - final_hidden_states = final_hidden_states[0] + final_hidden_states[1] - - if self.tp_size > 1: - final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index 28331b8ef3e..9612ea57b2c 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -176,7 +176,6 @@ class InternS1ProMoeSparseMoeBlock(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=True, renormalize=config.norm_topk_prob, quant_config=quant_config, prefix=f"{prefix}.experts", diff --git a/vllm/model_executor/models/jamba.py b/vllm/model_executor/models/jamba.py index 980bcffb5f9..b4b3b6873db 100644 --- a/vllm/model_executor/models/jamba.py +++ b/vllm/model_executor/models/jamba.py @@ -90,7 +90,6 @@ class JambaMoE(nn.Module): self.intermediate_size, tp_size=tp_size, params_dtype=params_dtype, - reduce_results=True, renormalize=False, use_grouped_topk=False, quant_config=quant_config, diff --git a/vllm/model_executor/models/kimi_linear.py b/vllm/model_executor/models/kimi_linear.py index 4cd7b63c147..e586a3ac346 100644 --- a/vllm/model_executor/models/kimi_linear.py +++ b/vllm/model_executor/models/kimi_linear.py @@ -11,11 +11,10 @@ from vllm.config import CacheConfig, ModelConfig, ParallelConfig, VllmConfig from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, ) from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe import SharedFusedMoE from vllm.model_executor.layers.kda import KimiDeltaAttention from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( @@ -132,22 +131,6 @@ class KimiMoE(nn.Module): self.gate.e_score_correction_bias = nn.Parameter(torch.empty(num_experts)) - self.experts = FusedMoE( - num_experts=num_experts, - top_k=config.num_experts_per_token, - hidden_size=hidden_size, - intermediate_size=moe_intermediate_size, - reduce_results=False, - renormalize=moe_renormalize, - quant_config=quant_config, - use_grouped_topk=config.use_grouped_topk, - num_expert_group=config.num_expert_group, - topk_group=config.topk_group, - prefix=f"{prefix}.experts", - scoring_func=config.moe_router_activation_func, - e_score_correction_bias=self.gate.e_score_correction_bias, - ) - if self.num_shared_experts is not None: intermediate_size = moe_intermediate_size * self.num_shared_experts self.shared_experts = KimiMLP( @@ -158,22 +141,33 @@ class KimiMoE(nn.Module): reduce_results=False, prefix=f"{prefix}.shared_experts", ) + else: + self.shared_experts = None + + self.experts = SharedFusedMoE( + shared_experts=self.shared_experts, + num_experts=num_experts, + top_k=config.num_experts_per_token, + hidden_size=hidden_size, + intermediate_size=moe_intermediate_size, + renormalize=moe_renormalize, + quant_config=quant_config, + use_grouped_topk=config.use_grouped_topk, + num_expert_group=config.num_expert_group, + topk_group=config.topk_group, + prefix=f"{prefix}.experts", + scoring_func=config.moe_router_activation_func, + e_score_correction_bias=self.gate.e_score_correction_bias, + routed_scaling_factor=self.routed_scaling_factor, + ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: num_tokens, hidden_size = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_size) - if self.num_shared_experts is not None: - shared_output = self.shared_experts(hidden_states) router_logits, _ = self.gate(hidden_states) - final_hidden_states = ( - self.experts(hidden_states=hidden_states, router_logits=router_logits) - * self.routed_scaling_factor + final_hidden_states = self.experts( + hidden_states=hidden_states, router_logits=router_logits ) - if shared_output is not None: - final_hidden_states = final_hidden_states + shared_output - - if self.tp_size > 1: - final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(num_tokens, hidden_size) @@ -482,7 +476,7 @@ class KimiLinearModel(nn.Module): if self.config.is_moe: # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) - expert_params_mapping = FusedMoE.make_expert_params_mapping( + expert_params_mapping = SharedFusedMoE.make_expert_params_mapping( self, ckpt_gate_proj_name="w1", ckpt_down_proj_name="w2", diff --git a/vllm/model_executor/models/lfm2_moe.py b/vllm/model_executor/models/lfm2_moe.py index d955b7127ad..4b49430c1fa 100644 --- a/vllm/model_executor/models/lfm2_moe.py +++ b/vllm/model_executor/models/lfm2_moe.py @@ -150,7 +150,6 @@ class Lfm2MoeSparseMoeBlock(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, # needed for softmax score func @@ -161,6 +160,7 @@ class Lfm2MoeSparseMoeBlock(nn.Module): num_redundant_experts=self.n_redundant_experts, scoring_func="sigmoid", e_score_correction_bias=self.gate.e_score_correction_bias, + routed_scaling_factor=self.routed_scaling_factor, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -170,16 +170,10 @@ class Lfm2MoeSparseMoeBlock(nn.Module): # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) - final_hidden_states = ( - self.experts(hidden_states=hidden_states, router_logits=router_logits) - * self.routed_scaling_factor + final_hidden_states = self.experts( + hidden_states=hidden_states, router_logits=router_logits ) - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501 - final_hidden_states - ) - return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/llama4.py b/vllm/model_executor/models/llama4.py index b84b4e2ae51..a1c0ac89605 100644 --- a/vllm/model_executor/models/llama4.py +++ b/vllm/model_executor/models/llama4.py @@ -135,7 +135,6 @@ class Llama4MoE(nn.Module): custom_routing_function=Llama4MoE.custom_routing_function, intermediate_size=intermediate_size_moe, apply_router_weight_on_input=True, - reduce_results=False, renormalize=False, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -151,19 +150,14 @@ class Llama4MoE(nn.Module): router_logits, _ = self.router(hidden_states) - shared_out, routed_out = self.experts( + experts_out = self.experts( hidden_states=hidden_states, router_logits=router_logits, ) - experts_out = routed_out + shared_out if self.is_sequence_parallel: experts_out = tensor_model_parallel_all_gather(experts_out, 0) experts_out = experts_out[:num_tokens] - elif self.tp_size > 1: - experts_out = self.experts.maybe_all_reduce_tensor_model_parallel( - experts_out - ) return experts_out diff --git a/vllm/model_executor/models/longcat_flash.py b/vllm/model_executor/models/longcat_flash.py index 375b0b69b1f..945fcb61509 100644 --- a/vllm/model_executor/models/longcat_flash.py +++ b/vllm/model_executor/models/longcat_flash.py @@ -300,7 +300,6 @@ class LongcatMoe(nn.Module): top_k=top_k, hidden_size=hidden_size, intermediate_size=intermediate_size, - reduce_results=True, params_dtype=params_dtype, renormalize=False, quant_config=quant_config, diff --git a/vllm/model_executor/models/mimo_v2_flash.py b/vllm/model_executor/models/mimo_v2_flash.py index 43475ed690c..0b466f16601 100644 --- a/vllm/model_executor/models/mimo_v2_flash.py +++ b/vllm/model_executor/models/mimo_v2_flash.py @@ -162,7 +162,6 @@ class MiMoV2MoE(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=True, renormalize=config.norm_topk_prob, quant_config=quant_config, prefix=f"{prefix}.experts", diff --git a/vllm/model_executor/models/minimax_m2.py b/vllm/model_executor/models/minimax_m2.py index 2a9e5f4e08a..84d8dda533f 100644 --- a/vllm/model_executor/models/minimax_m2.py +++ b/vllm/model_executor/models/minimax_m2.py @@ -36,7 +36,6 @@ from vllm.config import CacheConfig, ModelConfig, VllmConfig from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, ) from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import FusedMoE @@ -104,7 +103,6 @@ class MiniMaxM2MoE(nn.Module): e_score_correction_bias=self.e_score_correction_bias, hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, - reduce_results=False, renormalize=True, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -134,9 +132,6 @@ class MiniMaxM2MoE(nn.Module): final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - final_hidden_states = final_hidden_states - if self.tp_size > 1: - final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/minimax_text_01.py b/vllm/model_executor/models/minimax_text_01.py index 21d74d8b058..67d7cb2d8bc 100644 --- a/vllm/model_executor/models/minimax_text_01.py +++ b/vllm/model_executor/models/minimax_text_01.py @@ -162,7 +162,6 @@ class MiniMaxText01MoE(nn.Module): hidden_size=self.hidden_size, intermediate_size=self.intermediate_size * self.tp_size, params_dtype=self.params_dtype, - reduce_results=True, renormalize=True, quant_config=self.quant_config, tp_size=self.tp_size, diff --git a/vllm/model_executor/models/mixtral.py b/vllm/model_executor/models/mixtral.py index 376fd7a1709..c182444f667 100644 --- a/vllm/model_executor/models/mixtral.py +++ b/vllm/model_executor/models/mixtral.py @@ -132,7 +132,6 @@ class MixtralMoE(nn.Module): hidden_size=hidden_size, intermediate_size=intermediate_size, params_dtype=params_dtype, - reduce_results=True, renormalize=True, quant_config=quant_config, tp_size=tp_size, diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index 8abbc808cce..fa068639648 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -216,7 +216,6 @@ class NemotronHMoE(nn.Module): top_k=config.num_experts_per_tok, hidden_size=self.moe_hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, @@ -231,6 +230,9 @@ class NemotronHMoE(nn.Module): num_redundant_experts=self.n_redundant_experts, is_sequence_parallel=self.is_sequence_parallel, routed_input_transform=self.fc1_latent_proj, + routed_output_transform=self.fc2_latent_proj, + routed_scaling_factor=self.routed_scaling_factor, + apply_routed_scale_to_output=True, router_logits_dtype=self.gate.out_dtype, ) @@ -244,38 +246,15 @@ class NemotronHMoE(nn.Module): # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) - # SharedFusedMoE handles: - # - shared experts (with original hidden_states) - # - routed_input_transform (fc1_latent_proj) for latent MoE - # - multistream parallelism between shared and routed experts - shared_output, final_hidden_states = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - # Fix FP16 overflow - # See DeepseekV2DecoderLayer for more details. - if hidden_states.dtype != torch.float16: - final_hidden_states *= self.routed_scaling_factor - elif self.shared_experts is not None: - shared_output *= 1.0 / self.routed_scaling_factor - - # TODO: See SharedFusedMoE.apply_routed_input_transform - # for bandwidth optimization - if self.use_latent_moe: - final_hidden_states, _ = self.fc2_latent_proj(final_hidden_states) - - if self.shared_experts is not None: - final_hidden_states += shared_output - if self.is_sequence_parallel: final_hidden_states = tensor_model_parallel_all_gather( final_hidden_states, 0 ) final_hidden_states = final_hidden_states[:num_tokens] - elif self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/olmoe.py b/vllm/model_executor/models/olmoe.py index f0afe0e997c..fcde2e41afb 100644 --- a/vllm/model_executor/models/olmoe.py +++ b/vllm/model_executor/models/olmoe.py @@ -98,7 +98,6 @@ class OlmoeMoE(nn.Module): top_k=top_k, hidden_size=hidden_size, intermediate_size=intermediate_size, - reduce_results=True, renormalize=False, quant_config=quant_config, tp_size=tp_size, diff --git a/vllm/model_executor/models/openpangu.py b/vllm/model_executor/models/openpangu.py index 994ae82529a..7de84da5193 100644 --- a/vllm/model_executor/models/openpangu.py +++ b/vllm/model_executor/models/openpangu.py @@ -206,7 +206,6 @@ class OpenPanguMoE(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, @@ -214,8 +213,8 @@ class OpenPanguMoE(nn.Module): topk_group=1, prefix=f"{prefix}.experts", scoring_func="sigmoid", - # we do scaling outside, set factor to 1.0 to avoid double mul - routed_scaling_factor=1.0, + routed_scaling_factor=self.routed_scaling_factor, + apply_routed_scale_to_output=True, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, @@ -234,33 +233,15 @@ class OpenPanguMoE(nn.Module): router_logits, _ = self.gate(hidden_states) - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - shared_output, final_hidden_states = fused_moe_out - if self.shared_experts is None: - assert shared_output is None - - if hidden_states.dtype != torch.float16: - final_hidden_states *= self.routed_scaling_factor - elif self.shared_experts is not None: - assert shared_output is not None - shared_output *= 1.0 / self.routed_scaling_factor - - if self.shared_experts is not None: - assert shared_output is not None - final_hidden_states += shared_output - if self.is_sequence_parallel: final_hidden_states = tensor_model_parallel_all_gather( final_hidden_states, 0 ) final_hidden_states = final_hidden_states[:num_tokens] - elif self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/param2moe.py b/vllm/model_executor/models/param2moe.py index 925f94e0600..fddd1a8f173 100644 --- a/vllm/model_executor/models/param2moe.py +++ b/vllm/model_executor/models/param2moe.py @@ -359,7 +359,6 @@ class Param2MoEMoEBlock(nn.Module): top_k=self.top_k, hidden_size=self.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=self.norm_expert_prob, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -388,24 +387,11 @@ class Param2MoEMoEBlock(nn.Module): self.gate.weight.float(), ).to(hidden_states.dtype) - final_hidden = self.experts( + expert_output = self.experts( hidden_states=hidden_states, router_logits=router_logits, ) - if self.shared_experts is not None: - shared_output, expert_output = final_hidden - else: - shared_output, expert_output = None, final_hidden - - if shared_output is not None: - expert_output = expert_output + shared_output - - if self.tp_size > 1: - expert_output = self.experts.maybe_all_reduce_tensor_model_parallel( - expert_output - ) - return expert_output.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/phimoe.py b/vllm/model_executor/models/phimoe.py index 0b55b7ec839..7d6083f202e 100644 --- a/vllm/model_executor/models/phimoe.py +++ b/vllm/model_executor/models/phimoe.py @@ -281,7 +281,6 @@ class PhiMoE(nn.Module): hidden_size=hidden_size, intermediate_size=intermediate_size, params_dtype=params_dtype, - reduce_results=True, renormalize=False, quant_config=quant_config, tp_size=tp_size, diff --git a/vllm/model_executor/models/qwen2_moe.py b/vllm/model_executor/models/qwen2_moe.py index 750835cccd9..b5d13e926d7 100644 --- a/vllm/model_executor/models/qwen2_moe.py +++ b/vllm/model_executor/models/qwen2_moe.py @@ -170,7 +170,6 @@ class Qwen2MoeSparseMoeBlock(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -187,12 +186,6 @@ class Qwen2MoeSparseMoeBlock(nn.Module): final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - if self.shared_expert is not None: - final_hidden_states = final_hidden_states[0] + final_hidden_states[1] - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501 - final_hidden_states - ) return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/qwen3_moe.py b/vllm/model_executor/models/qwen3_moe.py index f2ce070be8b..f0f69d43537 100644 --- a/vllm/model_executor/models/qwen3_moe.py +++ b/vllm/model_executor/models/qwen3_moe.py @@ -212,7 +212,6 @@ class Qwen3MoeSparseMoeBlock(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -234,22 +233,15 @@ class Qwen3MoeSparseMoeBlock(nn.Module): # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) - shared_out, fused_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - final_hidden_states = ( - shared_out + fused_out if shared_out is not None else fused_out - ) if self.is_sequence_parallel: final_hidden_states = tensor_model_parallel_all_gather( final_hidden_states, 0 ) final_hidden_states = final_hidden_states[:num_tokens] - elif self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501 - final_hidden_states - ) # return to 1d if input is 1d return final_hidden_states.squeeze(0) if is_input_1d else final_hidden_states diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index 3bd026d9dc9..50d44dbbf63 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -153,7 +153,6 @@ class Qwen3NextSparseMoeBlock(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=getattr(config, "norm_topk_prob", True), quant_config=quant_config, prefix=f"{prefix}.experts", @@ -183,18 +182,11 @@ class Qwen3NextSparseMoeBlock(nn.Module): hidden_states=hidden_states, router_logits=router_logits ) - if self.shared_expert is not None: - final_hidden_states = final_hidden_states[0] + final_hidden_states[1] - if self.is_sequence_parallel: final_hidden_states = tensor_model_parallel_all_gather( final_hidden_states, 0 ) final_hidden_states = final_hidden_states[:num_tokens] - elif self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501 - final_hidden_states - ) return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/sarvam.py b/vllm/model_executor/models/sarvam.py index fa5ec44d7e7..3656fc921b2 100644 --- a/vllm/model_executor/models/sarvam.py +++ b/vllm/model_executor/models/sarvam.py @@ -341,7 +341,6 @@ class SarvamMLAMoE(nn.Module): top_k=self.top_k, hidden_size=self.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=self.norm_expert_prob, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -370,20 +369,7 @@ class SarvamMLAMoE(nn.Module): router_logits=router_logits, ) - if self.shared_experts is not None: - shared_output, expert_output = final_hidden - else: - shared_output, expert_output = None, final_hidden - - if shared_output is not None: - expert_output = expert_output + shared_output - - if self.tp_size > 1: - expert_output = self.experts.maybe_all_reduce_tensor_model_parallel( - expert_output - ) - - return expert_output.view(num_tokens, hidden_dim) + return final_hidden.view(num_tokens, hidden_dim) class SarvamMLABlock(nn.Module): diff --git a/vllm/model_executor/models/step3_text.py b/vllm/model_executor/models/step3_text.py index 18b689166a5..636a121c590 100644 --- a/vllm/model_executor/models/step3_text.py +++ b/vllm/model_executor/models/step3_text.py @@ -14,7 +14,6 @@ from vllm.config import CacheConfig, ModelConfig, VllmConfig from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, ) from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul @@ -71,7 +70,6 @@ class FusedMoEBlock(nn.Module): top_k=config.moe_top_k, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_expert_weight, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -94,8 +92,6 @@ class FusedMoEBlock(nn.Module): final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - if self.tp_size > 1: - final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/step3p5.py b/vllm/model_executor/models/step3p5.py index bb4bf14a963..018f7895602 100644 --- a/vllm/model_executor/models/step3p5.py +++ b/vllm/model_executor/models/step3p5.py @@ -379,7 +379,6 @@ class FusedMoEBlock(nn.Module): top_k=config.moe_top_k, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_expert_weight, quant_config=quant_config, activation=activation, @@ -397,30 +396,16 @@ class FusedMoEBlock(nn.Module): hidden_states = hidden_states.view(-1, hidden_dim) if self.experts.is_internal_router: - # In this case, the gate/router runs inside the FusedMoE class - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=hidden_states ) else: - # router_logits: (num_tokens, n_experts) + # TODO(bnell): this gate could be moved into the FusedMoE? router_logits, _ = self.gate(hidden_states) - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - shared_output, final_hidden_states = fused_moe_out - if self.share_expert is None: - assert shared_output is None - - if self.share_expert is not None: - assert shared_output is not None - final_hidden_states += shared_output - - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) - return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py index 81d21abbd06..cf13958ef76 100644 --- a/vllm/model_executor/models/transformers/moe.py +++ b/vllm/model_executor/models/transformers/moe.py @@ -204,8 +204,6 @@ class MoEMixin(MixtureOfExperts): ) assert intermediate_size is not None - # If there are shared experts, the results are - # reduced after mlp.forward() not inside FusedMoE num_shared_experts = getattr_iter( text_config, [ @@ -214,17 +212,6 @@ class MoEMixin(MixtureOfExperts): ], 0, ) - reduce_results = num_shared_experts == 0 - - def add_all_reduce(mlp: nn.Module): - """Adds an all-reduce to the output of `mlp.forward()`.""" - - class MLPWithAllReduce(mlp.__class__): - def forward(self, *args, **kwargs): - output = super().forward(*args, **kwargs) - return self.experts.maybe_all_reduce_tensor_model_parallel(output) - - mlp.__class__ = MLPWithAllReduce # Unused kwargs since we use custom_routing_function: # - `scoring_func` and `e_score_correction_bias` only used for grouped @@ -289,14 +276,11 @@ class MoEMixin(MixtureOfExperts): if "bias" in experts_param_name: has_bias = True break - # Double check there are no shared experts - nonlocal reduce_results - if reduce_results: + # If the config does not specify num_shared_experts, but + # the model has shared experts, we assume there is one. + if self.num_shared_experts == 0: for mlp_param_name, _ in mlp.named_parameters(): if "shared_expert" in mlp_param_name: - reduce_results = False - # If the config does not specify num_shared_experts, but - # the model has shared experts, we assume there is one. self.num_shared_experts = 1 break # Replace experts module with FusedMoE @@ -305,7 +289,6 @@ class MoEMixin(MixtureOfExperts): top_k=top_k, hidden_size=hidden_size, intermediate_size=intermediate_size, - reduce_results=reduce_results, renormalize=renormalize, # Hard coded because topk happens in Transformers use_grouped_topk=False, @@ -326,13 +309,6 @@ class MoEMixin(MixtureOfExperts): self.moe_layers.append(fused_experts) self.expert_weights.append(fused_experts.get_expert_weights()) self.num_moe_layers += 1 - # If results are not all-reduced in FusedMoE, ensure they - # are all-reduced at the end of mlp.forward() if tensor - # parallel or expert parallel is enabled - if not reduce_results and ( - fused_experts.tp_size > 1 or fused_experts.ep_size > 1 - ): - add_all_reduce(mlp) else: _recursive_replace(child_module, prefix=qual_name)