diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 86d93cdd2..e1b77562c 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -826,10 +826,9 @@ class DeepseekV2MoE(nn.Module): if server_args.enable_eplb else None ) + defer_shared = not self.experts.moe_runner_config.inplace if hidden_states.shape[0] > 0: - if ( - not self._fuse_shared_experts_inside_sbo - ): # TODO: check if it supports mtp + if not defer_shared and not self._fuse_shared_experts_inside_sbo: shared_output = self._forward_shared_experts( hidden_states, gemm_output_zero_allocator ) @@ -894,6 +893,15 @@ class DeepseekV2MoE(nn.Module): # fused in biased_grouped_topk so we can skip here final_hidden_states *= self.routed_scaling_factor + if ( + defer_shared + and hidden_states.shape[0] > 0 + and not self._fuse_shared_experts_inside_sbo + ): + shared_output = self._forward_shared_experts( + hidden_states, gemm_output_zero_allocator + ) + final_hidden_states = maybe_fuse_routed_scale_and_shared_add( self.experts, final_hidden_states,