Handle shared experts correctly

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

Author: MatthewBonanni

vllm/model_executor/layers/fused_moe/layer.py

@@ -1801,6 +1801,15 @@ def process_chunk(chunk_start, chunk_end, skip_result_store=False):
             staged_hidden_states.copy_(hidden_states, non_blocking=True)
             staged_router_logits.copy_(router_logits, non_blocking=True)
 
+            # If there are shared experts but we are not using a modular kernel,
+            # the shared experts must be called here
+            if (not isinstance(self.quant_method.fused_experts,
+                               FusedMoEModularKernel)
+                    and self.shared_experts is not None):
+                shared_output = self.shared_experts(staged_hidden_states)
+            else:
+                shared_output = None
+
             # Matrix multiply.
             final_hidden_states = self.quant_method.apply(
                 layer=self,
@@ -1824,8 +1833,13 @@ def process_chunk(chunk_start, chunk_end, skip_result_store=False):
                 logical_replica_count=self.logical_replica_count,
             )
 
-            assert self.shared_experts is None or isinstance(
-                final_hidden_states, tuple)
+            if shared_output is not None:
+                assert not isinstance(final_hidden_states, tuple)
+                assert self.shared_experts is not None
+                final_hidden_states = (
+                    shared_output,
+                    final_hidden_states,
+                )
 
             if not skip_result_store:
                 if self.shared_experts is None: