perf(dsv4): overlap shared expert computation with DeepEP combine

lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py

@@ -1,6 +1,6 @@
 import torch
 import threading
-from typing import Dict, Any, Optional, Tuple, List
+from typing import Dict, Any, Optional, Tuple, List, Callable
 from lightllm.common.basemodel.layer_weights.meta_weights.base_weight import BaseWeightTpl
 from lightllm.common.quantization.quantize_method import WeightPack
 from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_slicer import (
@@ -164,6 +164,11 @@ def experts(
             clamp_up_add_one=clamp_up_add_one,
         )
 
+    @property
+    def supports_decode_shared_expert_overlap(self) -> bool:
+        """Backend capability, independent of the model's overlap enable switch."""
+        return self.fuse_moe_impl.supports_decode_shared_expert_overlap
+
     def experts_with_topk(
         self,
         input_tensor: torch.Tensor,
@@ -174,7 +179,17 @@ def experts_with_topk(
         alpha: Optional[float] = None,
         limit: Optional[float] = None,
         clamp_up_add_one: bool = True,
+        shared_expert_func: Optional[Callable[[], torch.Tensor]] = None,
     ) -> torch.Tensor:
+        """Run routed experts, optionally including the shared expert in the result.
+
+        The callback requires decode overlap capability. It must treat input_tensor
+        as read-only and return a matching shared output without modifying routed
+        state. On successful eager/capture calls it runs exactly once on the current
+        compute stream, after combine SEND and before RECV. Graph replay executes
+        the captured GPU work rather than calling Python again. A callback result
+        is already added to the returned tensor; without one the result is routed-only.
+        """
         moe_capture_callback = get_moe_capture_callback(infer_state, self.layer_num_)
         if moe_capture_callback is not None:
             moe_capture_callback(topk_ids)
@@ -188,6 +203,7 @@ def experts_with_topk(
             alpha=alpha,
             limit=limit,
             clamp_up_add_one=clamp_up_add_one,
+            shared_expert_func=shared_expert_func,
         )
 
     def low_latency_dispatch(

lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py

@@ -12,6 +12,8 @@
 
 
 class FuseMoeBaseImpl:
+    supports_decode_shared_expert_overlap = False
+
     def __init__(
         self,
         n_routed_experts: int,

lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py

@@ -1,5 +1,5 @@
 import torch
-from typing import Optional, Tuple, Any
+from typing import Optional, Tuple, Any, Callable
 from .triton_impl import FuseMoeTriton
 from lightllm.distributed import dist_group_manager
 from lightllm.common.quantization.quantize_method import WeightPack
@@ -14,13 +14,18 @@
     masked_group_gemm,
     chunked_expanded_moe_forward,
     quantize_fused_experts_input,
+    use_sm100_mega_moe,
 )
 from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd
 from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType
 from lightllm.common.basemodel.triton_kernel.redundancy_topk_ids_repair import redundancy_topk_ids_repair
 
 
 class FuseMoeDeepGEMM(FuseMoeTriton):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.supports_decode_shared_expert_overlap = not use_sm100_mega_moe(self.quant_method)
+
     def _select_experts(
         self,
         input_tensor: torch.Tensor,
@@ -80,6 +85,7 @@ def _fused_experts(
         alpha: Optional[float] = None,
         limit: Optional[float] = None,
         clamp_up_add_one: bool = True,
+        shared_expert_func: Optional[Callable[[], torch.Tensor]] = None,
     ):
         output = fused_experts(
             hidden_states=input_tensor,
@@ -94,6 +100,7 @@ def _fused_experts(
             alpha=alpha,
             limit=limit,
             clamp_up_add_one=clamp_up_add_one,
+            shared_expert_func=shared_expert_func,
         )
         return output
 

lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py

@@ -101,7 +101,12 @@ def fused_experts_with_topk(
         alpha: Optional[float] = None,
         limit: Optional[float] = None,
         clamp_up_add_one: bool = True,
+        shared_expert_func: Optional[Callable[[], torch.Tensor]] = None,
     ):
+        extra_args = {}
+        if shared_expert_func is not None:
+            assert self.supports_decode_shared_expert_overlap
+            extra_args["shared_expert_func"] = shared_expert_func
         return self._fused_experts(
             input_tensor=input_tensor,
             w13=w13,
@@ -112,6 +117,7 @@ def fused_experts_with_topk(
             alpha=alpha,
             limit=limit,
             clamp_up_add_one=clamp_up_add_one,
+            **extra_args,
         )
 
     def __call__(

lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py

@@ -221,10 +221,13 @@ def fused_experts(
     alpha: Optional[float] = None,
     limit: Optional[float] = None,
     clamp_up_add_one: bool = True,
+    shared_expert_func: Optional[Callable[[], torch.Tensor]] = None,
 ):
     assert (limit is None and alpha is None) or (limit is not None and alpha is not None)
+    assert shared_expert_func is None or is_prefill is False
     check_ep_expert_dtype(quant_method)
     if use_sm100_mega_moe(quant_method):
+        assert shared_expert_func is None, "Mega MoE does not support shared expert overlap"
         if limit is not None:
             raise NotImplementedError("FP4 Mega MoE does not support clamped SwiGLU")
         return mega_moe_impl(hidden_states, w13, w2, topk_weights, topk_idx, quant_method)
@@ -249,6 +252,7 @@ def fused_experts(
         alpha=alpha,
         limit=limit,
         clamp_up_add_one=clamp_up_add_one,
+        shared_expert_func=shared_expert_func,
     )
 
 
@@ -271,6 +275,7 @@ def fused_experts_impl(
     limit: Optional[float] = None,
     clamp_up_add_one: bool = True,
     block_size_k: int = 128,
+    shared_expert_func: Optional[Callable[[], torch.Tensor]] = None,
 ):
     # Check constraints.
     assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch"
@@ -403,8 +408,19 @@ def fused_experts_impl(
         )
         # low latency combine
         combined_x, event_overlap, hook = buffer.low_latency_combine(
-            gemm_out_b, topk_idx, topk_weights, handle, async_finish=False, return_recv_hook=False
+            gemm_out_b,
+            topk_idx,
+            topk_weights,
+            handle,
+            async_finish=False,
+            return_recv_hook=shared_expert_func is not None,
         )
+        if shared_expert_func is not None:
+            # SEND, shared work, RECV and the final add use the current compute
+            # stream. The callback is read-only on hidden_states and runs once.
+            shared = shared_expert_func()
+            hook()
+            return combined_x + shared
     return combined_x
 
 

lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py

@@ -1,3 +1,4 @@
+import os
 import torch
 import triton
 from lightllm.common.basemodel import TransformerLayerInferTpl
@@ -52,6 +53,11 @@ def __init__(self, layer_num, network_config):
         self.flashmla_q_head_num_ = self.tp_q_head_num_
         self.tp_groups = self.o_groups // self.tp_world_size_
         self.enable_ep_moe = get_env_start_args().enable_ep_moe
+        self.enable_shared_expert_overlap = os.getenv("LIGHTLLM_DSV4_SHARED_EXPERT_OVERLAP", "1").upper() in (
+            "1",
+            "TRUE",
+            "ON",
+        )
         self.compressor = CompressorInfer(
             layer_idx=self.layer_num_, network_config=self.network_config_, tp_world_size=self.tp_world_size_
         )
@@ -550,9 +556,17 @@ def _ffn(self, x, infer_state: DeepseekV4InferStateInfo, layer_weight: DeepseekV
         x = self._tpsp_allgather(input=x, infer_state=infer_state)
 
         topk_weights, topk_indices = self._select_experts(x, infer_state, layer_weight)
-        # Routed experts can overwrite x in place, so compute the shared expert first.
-        shared = self._ffn_tp(input=x, infer_state=infer_state, layer_weight=layer_weight)
-        routed = layer_weight.experts.experts_with_topk(
+        if (
+            self.enable_shared_expert_overlap
+            and not infer_state.is_prefill
+            and layer_weight.experts.supports_decode_shared_expert_overlap
+        ):
+            shared_expert_func = lambda: self._ffn_tp(x, infer_state, layer_weight)
+        else:
+            # Routed experts can overwrite x in place, so compute shared first.
+            shared = self._ffn_tp(input=x, infer_state=infer_state, layer_weight=layer_weight)
+            shared_expert_func = None
+        out = layer_weight.experts.experts_with_topk(
             input_tensor=x,
             topk_weights=topk_weights,
             topk_ids=topk_indices,
@@ -561,8 +575,10 @@ def _ffn(self, x, infer_state: DeepseekV4InferStateInfo, layer_weight: DeepseekV
             alpha=1.0,
             limit=self.swiglu_limit,
             clamp_up_add_one=False,
+            shared_expert_func=shared_expert_func,
         )
-        out = routed + shared
+        if shared_expert_func is None:
+            out = out + shared
         if self.enable_ep_moe:
             return out
         return self._tpsp_reduce(input=out, infer_state=infer_state)