[Models] Apply SharedFusedMoE to all models with shared experts

Signed-off-by: Bill Nell <bnell@redhat.com>

Author: bnellnm

vllm/model_executor/layers/quantization/fp8.py

@@ -630,6 +630,8 @@ def create_weights(self, layer: Module, num_experts: int, hidden_size: int,
             layer.w13_input_scale = None
             layer.w2_input_scale = None
 
+        self.rocm_aiter_moe_enabled = False
+
     def process_weights_after_loading(self, layer: Module) -> None:
         # Lazy import to avoid importing triton too early.
         from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (

vllm/model_executor/layers/shared_fused_moe/shared_fused_moe.py

@@ -18,13 +18,15 @@ class SharedFusedMoE(FusedMoE):
 
     def __init__(
         self,
-        shared_experts: torch.nn.Module,
+        shared_experts: Optional[torch.nn.Module],
         use_overlapped: bool = True,
         **kwargs,
     ):
         super().__init__(**kwargs)
         self._shared_experts = shared_experts
-        self.use_overlapped = use_overlapped
+        self.use_overlapped = use_overlapped and not (
+            self.use_ep or self.use_flashinfer_cutlass_kernels
+        ) and self.shared_experts is not None
 
     @property
     def shared_experts(self) -> Optional[torch.nn.Module]:
@@ -36,13 +38,16 @@ def forward(
         router_logits: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         if not self.use_overlapped:
-            shared_out = self._shared_experts(hidden_states)
-
-            # Reduce outputs if necessary, since the MLP should
-            # have been created with reduce_results=False.
-            if (self.reduce_results and self.tp_size > 1
-                    and self.must_reduce_shared_expert_outputs()):
-                shared_out = tensor_model_parallel_all_reduce(shared_out)
+            if self._shared_experts is not None:
+                shared_out = self._shared_experts(hidden_states)
+
+                # Reduce shared expert outputs if necessary, since the MLP
+                # should have been created with reduce_results=False.
+                if (self.reduce_results and self.tp_size > 1
+                        and self.must_reduce_shared_expert_outputs()):
+                    shared_out = tensor_model_parallel_all_reduce(shared_out)
+            else:
+                shared_out = None
 
             fused_out = super().forward(
                 hidden_states=hidden_states,

vllm/model_executor/models/aria.py

@@ -13,6 +13,7 @@
 from vllm.distributed import get_tensor_model_parallel_rank
 from vllm.model_executor.layers.activation import get_act_fn
 from vllm.model_executor.layers.fused_moe import FusedMoE
+from vllm.model_executor.layers.shared_fused_moe import SharedFusedMoE
 from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                                RowParallelLinear)
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -199,7 +200,7 @@ def forward(
         return out
 
 
-class AriaFusedMoE(FusedMoE):
+class AriaFusedMoE(SharedFusedMoE):
 
     def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor,
                       shard_id: str) -> None:
@@ -253,22 +254,23 @@ def __init__(
             torch.empty(
                 (self.config.moe_num_experts, self.config.hidden_size)))
 
+        self.shared_experts = LlamaMLP(
+            config.hidden_size,
+            config.intermediate_size * config.moe_num_shared_experts,
+            "silu",
+            quant_config=quant_config,
+            bias=config.mlp_bias,
+        )
+
         self.experts = AriaFusedMoE(
+            shared_experts=self.shared_experts,
             num_experts=config.moe_num_experts,
             top_k=config.moe_topk,
             hidden_size=config.hidden_size,
             intermediate_size=config.intermediate_size,
             quant_config=quant_config,
-            reduce_results=True,
             prefix=f"{prefix}.experts",
         )
-        self.shared_experts = LlamaMLP(
-            config.hidden_size,
-            config.intermediate_size * config.moe_num_shared_experts,
-            "silu",
-            quant_config=quant_config,
-            bias=config.mlp_bias,
-        )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         """
@@ -285,12 +287,8 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         router_output = torch.nn.functional.linear(hidden_states,
                                                    self.router_weight)
 
-        hidden_states_copy = hidden_states.clone()
-        # NOTE: hidden_states will be modified inplace by `FusedMoE`
-        sparse_expert_output = self.experts(hidden_states, router_output)
-        shared_expert_output = self.shared_experts(hidden_states_copy)
-
-        return sparse_expert_output + shared_expert_output
+        # NOTE: hidden_states will be modified inplace by `SharedFusedMoE`
+        return self.experts(hidden_states, router_output)
 
 
 class AriaTextDecoderLayer(LlamaDecoderLayer):

vllm/model_executor/models/bailing_moe.py

@@ -36,8 +36,7 @@
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
 from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
-                              get_tensor_model_parallel_world_size,
-                              tensor_model_parallel_all_reduce)
+                              get_tensor_model_parallel_world_size)
 from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.fused_moe import FusedMoE
 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -47,6 +46,7 @@
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.shared_fused_moe import SharedFusedMoE
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead, VocabParallelEmbedding)
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
@@ -266,22 +266,6 @@ def __init__(
             # default value for scoring_func
             self.score_function = "softmax"
 
-        self.experts = FusedMoE(
-            num_experts=self.num_experts,
-            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",
-            scoring_func=self.score_function,
-            e_score_correction_bias=self.gate.expert_bias,
-            num_expert_group=self.n_group,
-            topk_group=self.topk_group,
-            use_grouped_topk=self.use_grouped_topk,
-        )
-
         if self.num_shared_experts > 0:
             if hasattr(config, "moe_shared_expert_intermediate_size"):
                 intermediate_size = config.moe_shared_expert_intermediate_size
@@ -294,29 +278,52 @@ def __init__(
                 quant_config=quant_config,
                 reduce_results=False,
                 prefix=f"{prefix}.shared_experts")
+
+            self.experts = SharedFusedMoE(
+                shared_experts=self.shared_experts,
+                fused_output_scaling_factor=self.routed_scaling_factor,
+                shared_output_scaling_factor=1.0,
+                num_experts=self.num_experts,
+                top_k=self.top_k,
+                hidden_size=self.hidden_size,
+                intermediate_size=config.moe_intermediate_size,
+                renormalize=self.norm_expert_prob,
+                quant_config=quant_config,
+                prefix=f"{prefix}.experts",
+                scoring_func=self.score_function,
+                e_score_correction_bias=self.gate.expert_bias,
+                num_expert_group=self.n_group,
+                topk_group=self.topk_group,
+                use_grouped_topk=self.use_grouped_topk,
+            )
         else:
+            self.experts = FusedMoE(
+                num_experts=self.num_experts,
+                top_k=self.top_k,
+                hidden_size=self.hidden_size,
+                intermediate_size=config.moe_intermediate_size,
+                renormalize=self.norm_expert_prob,
+                quant_config=quant_config,
+                prefix=f"{prefix}.experts",
+                scoring_func=self.score_function,
+                e_score_correction_bias=self.gate.expert_bias,
+                num_expert_group=self.n_group,
+                topk_group=self.topk_group,
+                use_grouped_topk=self.use_grouped_topk,
+            )
             self.shared_experts = None
 
     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.shared_experts:
-            shared_output = self.shared_experts(hidden_states)
+
         # router_logits: (num_tokens, n_experts)
         router_logits = self.gate(hidden_states.to(self.router_dtype))
         router_logits = router_logits.to(hidden_states.dtype)
 
         final_hidden_states = self.experts(hidden_states=hidden_states,
                                            router_logits=router_logits)
 
-        final_hidden_states *= self.routed_scaling_factor
-
-        if self.shared_experts:
-            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)
 
 

vllm/model_executor/models/deepseek_v2.py

@@ -36,7 +36,7 @@
 from vllm.attention.backends.abstract import AttentionBackend
 from vllm.attention.ops.common import pack_seq_triton, unpack_seq_triton
 from vllm.compilation.decorators import support_torch_compile
-from vllm.config import (CacheConfig, ParallelConfig, VllmConfig,
+from vllm.config import (CacheConfig, ModelConfig, ParallelConfig, VllmConfig,
                          get_current_vllm_config)
 from vllm.distributed import (get_ep_group, get_pp_group,
                               get_tensor_model_parallel_rank,
@@ -133,6 +133,7 @@ class DeepseekV2MoE(nn.Module):
     def __init__(
         self,
         config: Union[DeepseekV2Config, DeepseekV3Config],
+        model_config: ModelConfig,
         parallel_config: ParallelConfig,
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
@@ -182,27 +183,9 @@ def __init__(
                                     self.n_local_physical_experts)
 
         if config.n_shared_experts is None:
-            self.experts = FusedMoE(
-                num_experts=config.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,
-                num_expert_group=config.n_group,
-                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,
-                enable_eplb=self.enable_eplb,
-                num_redundant_experts=self.n_redundant_experts,
-                is_sequence_parallel=self.is_sequence_parallel,
-            )
             self.shared_experts = None
+            fused_output_scaling_factor = 1.0
+            shared_output_scaling_factor = 1.0
         else:
             intermediate_size = (config.moe_intermediate_size *
                                  config.n_shared_experts)
@@ -213,31 +196,42 @@ def __init__(
                 hidden_act=config.hidden_act,
                 quant_config=quant_config,
                 is_sequence_parallel=self.is_sequence_parallel,
-                reduce_results=False,
+                reduce_results=False,  # XXXXX
                 prefix=f"{prefix}.shared_experts",
             )
 
-            self.experts = SharedFusedMoE(
-                shared_experts=self.shared_experts,
-                num_experts=config.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,
-                num_expert_group=config.n_group,
-                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,
-                enable_eplb=self.enable_eplb,
-                num_redundant_experts=self.n_redundant_experts,
-                is_sequence_parallel=self.is_sequence_parallel,
-            )
+            # Fix FP16 overflow
+            # See DeepseekV2DecoderLayer for more details.
+            if model_config.dtype != torch.float16:
+                fused_output_scaling_factor = self.routed_scaling_factor
+                shared_output_scaling_factor = 1.0
+            else:
+                fused_output_scaling_factor = 1.0
+                shared_output_scaling_factor = (1. /
+                                                self.routed_scaling_factor)
+
+        self.experts = SharedFusedMoE(
+            shared_experts=self.shared_experts,
+            fused_output_scaling_factor=fused_output_scaling_factor,
+            shared_output_scaling_factor=shared_output_scaling_factor,
+            num_experts=config.n_routed_experts,
+            top_k=config.num_experts_per_tok,
+            hidden_size=config.hidden_size,
+            intermediate_size=config.moe_intermediate_size,
+            renormalize=config.norm_topk_prob,
+            quant_config=quant_config,
+            use_grouped_topk=True,
+            num_expert_group=config.n_group,
+            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,
+            enable_eplb=self.enable_eplb,
+            num_redundant_experts=self.n_redundant_experts,
+            is_sequence_parallel=self.is_sequence_parallel,
+        )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         num_tokens, hidden_dim = hidden_states.shape
@@ -253,36 +247,15 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         # router_logits: (num_tokens, n_experts)
         router_logits, _ = self.gate(hidden_states)
 
-        fused_moe_out = 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
-        else:
-            shared_output = None
-            final_hidden_states = fused_moe_out
-
-        # 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:
-            assert shared_output is not None
-            shared_output *= (1. / self.routed_scaling_factor)
-
-        if self.shared_experts is not None:
-            assert shared_output is not None
-            final_hidden_states += shared_output
+        final_hidden_states = self.experts(hidden_states=hidden_states,
+                                           router_logits=router_logits)
 
         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))
 
+        # TODO(bnell): why is this view needed?
         return final_hidden_states.view(num_tokens, hidden_dim)
 
 
@@ -1036,6 +1009,7 @@ def __init__(self,
                 and layer_idx % config.moe_layer_freq == 0):
             self.mlp = DeepseekV2MoE(
                 config=config,
+                model_config=model_config,
                 parallel_config=parallel_config,
                 quant_config=quant_config,
                 prefix=f"{prefix}.mlp",