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

Signed-off-by: Alexander Matveev <[email protected]>

Author: alexm-redhat

examples/offline_inference/basic/basic.py

@@ -16,7 +16,13 @@
 
 def main():
     # Create an LLM.
-    llm = LLM(model="facebook/opt-125m")
+    llm = LLM(model="deepseek-ai/DeepSeek-R1-0528", tensor_parallel_size=8)
+    # llm = LLM(
+    #     model="nvidia/DeepSeek-R1-FP4",
+    #     tensor_parallel_size=8,
+    #     quantization="modelopt_fp4",
+    # )
+
     # Generate texts from the prompts.
     # The output is a list of RequestOutput objects
     # that contain the prompt, generated text, and other information.

vllm/model_executor/layers/fused_moe/layer.py

@@ -1038,6 +1038,9 @@ def __init__(
         expert_mapping: Optional[list[tuple[str, str, int, str]]] = 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
@@ -1275,6 +1278,10 @@ def __init__(
     def shared_experts(self) -> Optional[torch.nn.Module]:
         return None
 
+    @property
+    def gate(self) -> Optional[torch.nn.Module]:
+        return None
+
     @property
     def tp_size(self):
         return self.moe_parallel_config.tp_size
@@ -2058,6 +2065,7 @@ def forward_impl_chunked(
         self,
         full_hidden_states: torch.Tensor,
         full_router_logits: torch.Tensor,
+        has_separate_shared_experts: bool,
     ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
         assert self.batched_hidden_states is not None
         assert self.batched_router_logits is not None
@@ -2106,11 +2114,16 @@ 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
+
+                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
 
@@ -2137,9 +2150,12 @@ 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
+
+                current_stream.wait_stream(self.shared_experts_stream)
+
                 final_hidden_states = (
                     shared_output,
                     final_hidden_states,
@@ -2215,12 +2231,31 @@ def forward_impl(
             self.dp_size > 1 and self.use_flashinfer_cutlass_kernels
         )
 
-        if (
+        has_separate_shared_experts = (
+            not isinstance(self.quant_method.fused_experts, FusedMoEModularKernel)
+            and self.shared_experts is not None
+        )
+
+        use_chunked_impl = (
             self.moe_parallel_config.use_pplx_kernels
             or self.moe_parallel_config.use_deepep_ll_kernels
             or _use_flashinfer_cutlass_kernels
-        ):
-            return self.forward_impl_chunked(hidden_states, router_logits)
+        )
+
+        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
@@ -2230,11 +2265,14 @@ def forward_impl(
 
         # 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
 
@@ -2275,9 +2313,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/shared_fused_moe/shared_fused_moe.py

@@ -19,23 +19,31 @@ class SharedFusedMoE(FusedMoE):
     def __init__(
         self,
         shared_experts: torch.nn.Module,
+        gate: Optional[torch.nn.Module] = None,
         use_overlapped: bool = True,
         **kwargs,
     ):
         super().__init__(**kwargs)
         self._shared_experts = shared_experts
+        self._gate = gate
         self.use_overlapped = use_overlapped
 
     @property
     def shared_experts(self) -> Optional[torch.nn.Module]:
         return self._shared_experts if self.use_overlapped else None
 
+    @property
+    def gate(self) -> Optional[torch.nn.Module]:
+        return self._gate if self.use_overlapped else None
+
     def forward(
         self,
         hidden_states: torch.Tensor,
         router_logits: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         if not self.use_overlapped:
+            assert self.gate is None
+
             shared_out = self._shared_experts(hidden_states)
 
             # Reduce outputs if necessary, since the MLP should

vllm/model_executor/models/deepseek_v2.py

@@ -241,6 +241,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,
@@ -272,12 +273,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