[Performance] Run shared_experts on a separate cuda stream (in parallel with the FusedMoE)

Signed-off-by: Alexander Matveev <amatveev@redhat.com>

Author: alexm-redhat

vllm/model_executor/layers/fused_moe/layer.py

@@ -1042,6 +1042,9 @@ def __init__(
         expert_mapping: list[tuple[str, str, int, str]] | None = None,
     ):
         super().__init__()
+
+        self.shared_experts_stream = torch.cuda.Stream()
+
         if params_dtype is None:
             params_dtype = torch.get_default_dtype()
         self.params_dtype = params_dtype
@@ -1265,6 +1268,10 @@ def __init__(
     def shared_experts(self) -> torch.nn.Module | None:
         return None
 
+    @property
+    def gate(self) -> Optional[torch.nn.Module]:
+        return None
+
     @property
     def tp_size(self):
         return self.moe_parallel_config.tp_size
@@ -2054,6 +2061,7 @@ def forward_impl_chunked(
         self,
         full_hidden_states: torch.Tensor,
         full_router_logits: torch.Tensor,
+        has_separate_shared_experts: bool,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         assert self.batched_hidden_states is not None
         assert self.batched_router_logits is not None
@@ -2102,11 +2110,19 @@ def process_chunk(chunk_start, chunk_end, skip_result_store=False):
 
             # 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)
+            if has_separate_shared_experts:
+                assert self.shared_experts is not None
+
+                # For chunked, we start the shared experts stream here
+                # (Note that no concurrency with the router/gate)
+                current_stream = torch.cuda.current_stream()
+                self.shared_experts_stream.wait_stream(current_stream)
+
+                with torch.cuda.stream(self.shared_experts_stream):
+                    # Note that staged_hidden_states clone() is necessary
+                    # here to avoid conflict with the main stream
+                    shared_output = self.shared_experts(staged_hidden_states.clone())
+
             else:
                 shared_output = None
 
@@ -2133,9 +2149,13 @@ def process_chunk(chunk_start, chunk_end, skip_result_store=False):
                 logical_replica_count=self.logical_replica_count,
             )
 
-            if shared_output is not None:
+            if has_separate_shared_experts:
                 assert not isinstance(final_hidden_states, tuple)
                 assert self.shared_experts is not None
+
+                # Here we finish the shared experts stream
+                current_stream.wait_stream(self.shared_experts_stream)
+
                 final_hidden_states = (
                     shared_output,
                     final_hidden_states,
@@ -2205,20 +2225,42 @@ def forward_impl(
 
         self.ensure_moe_quant_config()
 
-        if self.use_dp_chunking:
-            return self.forward_impl_chunked(hidden_states, router_logits)
+        has_separate_shared_experts = (
+            not isinstance(self.quant_method.fused_experts, FusedMoEModularKernel)
+            and self.shared_experts is not None
+        )
+
+        use_chunked_impl = self.use_dp_chunking
+
+        if has_separate_shared_experts and not use_chunked_impl:
+            # Start the separate shared experts stream here since we want
+            # to run in parallel with the router/gate (next op below)
+            current_stream = torch.cuda.current_stream()
+            self.shared_experts_stream.wait_stream(current_stream)
+
+        # If router/gate provided, then apply it here
+        if self.gate is not None:
+            router_logits, _ = self.gate(hidden_states)
+
+        if use_chunked_impl:
+            return self.forward_impl_chunked(
+                hidden_states, router_logits, has_separate_shared_experts
+            )
 
         do_naive_dispatch_combine: bool = (
             self.dp_size > 1 and not self.quant_method.using_modular_kernel
         )
 
         # 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(hidden_states)
+        if has_separate_shared_experts:
+            assert self.shared_experts is not None
+
+            # Run shared experts in parallel on a separate stream
+            with torch.cuda.stream(self.shared_experts_stream):
+                # Note that hidden_states clone() is necessary here to avoid
+                # conflict with the main stream
+                shared_output = self.shared_experts(hidden_states.clone())
         else:
             shared_output = None
 
@@ -2259,9 +2301,13 @@ def forward_impl(
                 logical_replica_count=self.logical_replica_count,
             )
 
-            if shared_output is not None:
+            if has_separate_shared_experts:
                 assert not isinstance(final_hidden_states, tuple)
                 assert self.shared_experts is not None
+
+                # Wait for the parallel shared experts stream to finish here
+                current_stream.wait_stream(self.shared_experts_stream)
+
                 final_hidden_states = (
                     shared_output,
                     final_hidden_states,

vllm/model_executor/layers/fused_moe/shared_fused_moe.py

@@ -18,25 +18,36 @@ class SharedFusedMoE(FusedMoE):
     def __init__(
         self,
         shared_experts: torch.nn.Module | None,
+        gate: torch.nn.Module | None = None,
         use_overlapped: bool = True,
         **kwargs,
     ):
         super().__init__(**kwargs)
         self._shared_experts = shared_experts
+
         # Disable shared expert overlap if EP is disabled or we are not using
         # flashinfer + DP since there is nothing to be gained in this case.
         # Disabling the overlap optimization also prevents the shared experts
         # from being hidden from torch.compile.
         self.use_overlapped = (
             use_overlapped
-            and not (self.use_ep or self.use_flashinfer_cutlass_kernels)
+            and not (
+                self.use_ep
+                or (self.use_flashinfer_cutlass_kernels and self.dp_size > 1)
+            )
             and self._shared_experts is not None
         )
 
+        self._gate = gate
+
     @property
     def shared_experts(self) -> torch.nn.Module | None:
         return self._shared_experts if self.use_overlapped else None
 
+    @property
+    def gate(self) -> torch.nn.Module | None:
+        return self._gate if self.use_overlapped else None
+
     def forward(
         self,
         hidden_states: torch.Tensor,

vllm/model_executor/models/deepseek_v2.py

@@ -220,6 +220,7 @@ def __init__(
 
         self.experts = SharedFusedMoE(
             shared_experts=self.shared_experts,
+            gate=self.gate,
             num_experts=config.n_routed_experts,
             top_k=config.num_experts_per_tok,
             hidden_size=config.hidden_size,
@@ -251,12 +252,16 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         if self.is_sequence_parallel:
             hidden_states = sequence_parallel_chunk(hidden_states)
 
-        # 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 isinstance(self.experts, SharedFusedMoE):
+            fused_moe_out = 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(
+                hidden_states=hidden_states, router_logits=router_logits
+            )
 
         if self.shared_experts is not None:
             shared_output, final_hidden_states = fused_moe_out