feat: support shared expert fusion.

vllm/_aiter_ops.py

@@ -921,12 +921,12 @@ class rocm_aiter_ops:
         return cls._AITER_ENABLED and cls._RMSNORM_ENABLED
 
     @classmethod
-    @if_aiter_supported
+    # @if_aiter_supported
     def is_fused_moe_enabled(cls) -> bool:
         return cls._AITER_ENABLED and cls._FMOE_ENABLED
 
     @classmethod
-    @if_aiter_supported
+    # @if_aiter_supported
     def is_fusion_moe_shared_experts_enabled(cls) -> bool:
         return cls.is_fused_moe_enabled() and cls._MOE_SHARED_EXPERTS_ENABLED
 

vllm/envs.py

@@ -1055,7 +1055,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # Whether to use aiter triton fp4 bmm kernel
     # By default is enabled.
     "VLLM_ROCM_USE_AITER_FP4BMM": lambda: (
-        os.getenv("VLLM_ROCM_USE_AITER_FP4BMM", "True").lower() in ("true", "1")
+        os.getenv("VLLM_ROCM_USE_AITER_FP4BMM", "False").lower() in ("true", "1")
     ),
     # Use AITER triton unified attention for V1 attention
     "VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION": lambda: (

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

@@ -335,8 +335,10 @@ class GroupedTopKRouter(BaseRouter):
                 rocm_aiter_grouped_topk,
                 num_fused_shared_experts=self.num_fused_shared_experts,
             )
+            enable_shared_experts_fusion = True
         else:
             grouped_topk_impl = grouped_topk
+            enable_shared_experts_fusion = False
 
         if self.use_fused_gate:
             if envs.VLLM_USE_LIGHTOP:
@@ -347,7 +349,7 @@ class GroupedTopKRouter(BaseRouter):
                     self.num_expert_group,
                     self.topk_group,
                     self.top_k,
-                    0,
+                    self.num_fused_shared_experts if enable_shared_experts_fusion else 0,
                     self.routed_scaling_factor,
                 )
             else:
@@ -358,7 +360,7 @@ class GroupedTopKRouter(BaseRouter):
                     self.topk_group,
                     self.top_k,
                     routed_scaling_factor=self.routed_scaling_factor,
-                    n_share_experts_fusion=0,
+                    n_share_experts_fusion = (self.num_fused_shared_experts if enable_shared_experts_fusion else 0),
                 )
         else:
             topk_weights, topk_ids = grouped_topk_impl(

vllm/model_executor/models/deepseek_v2.py

@@ -324,8 +324,12 @@ class DeepseekV2MoE(nn.Module):
         self.experts = SharedFusedMoE(
             shared_experts=self.shared_experts,
             gate=self.gate,
-            num_experts=config.n_routed_experts,
-            top_k=config.num_experts_per_tok,
+            # num_experts=config.n_routed_experts,
+            # top_k=config.num_experts_per_tok,
+            num_experts=config.n_routed_experts 
+                    + (config.n_shared_experts if self.is_fusion_moe_shared_experts_enabled else 0),
+            top_k = config.num_experts_per_tok
+                    + (config.n_shared_experts if self.is_fusion_moe_shared_experts_enabled else 0),
             hidden_size=config.hidden_size,
             intermediate_size=config.moe_intermediate_size,
             reduce_results=False,