From b717546fab883d32c06c014cae5b13de9dc8acdc Mon Sep 17 00:00:00 2001 From: ICENacl <40969745+ICENacl@users.noreply.github.com> Date: Fri, 10 Jul 2026 02:29:07 +0800 Subject: [PATCH] fix(mtp): avoid mtp perf regression in deepseek when enable eplb (#28982) Co-authored-by: chunjiang.hc --- python/sglang/srt/models/deepseek_v2.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index ce6ecdc6d..253c30afe 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -939,7 +939,7 @@ class DeepseekV2MoE(nn.Module): server_args = get_server_args() dispatch_info = ( ExpertLocationDispatchInfo.init_new(layer_id=self.layer_id) - if server_args.enable_eplb + if server_args.enable_eplb and not self.is_nextn else None ) with torch.cuda.stream(self.alt_stream): @@ -1036,7 +1036,7 @@ class DeepseekV2MoE(nn.Module): server_args = get_server_args() dispatch_info = ( ExpertLocationDispatchInfo.init_new(layer_id=self.layer_id) - if server_args.enable_eplb + if server_args.enable_eplb and not self.is_nextn else None ) defer_shared = not self.experts.moe_runner_config.inplace @@ -1237,8 +1237,12 @@ class DeepseekV2MoE(nn.Module): hidden_states, router_logits, num_token_non_padded=forward_batch.num_token_non_padded, - expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new( - layer_id=self.layer_id, + expert_location_dispatch_info=( + ExpertLocationDispatchInfo.init_new( + layer_id=self.layer_id, + ) + if not self.is_nextn + else None ), **topk_kwargs, ) @@ -1471,8 +1475,12 @@ class DeepseekV2MoE(nn.Module): hidden_states=hidden_states, router_logits=router_logits, num_token_non_padded=state.forward_batch.num_token_non_padded, - expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new( - layer_id=self.layer_id, + expert_location_dispatch_info=( + ExpertLocationDispatchInfo.init_new( + layer_id=self.layer_id, + ) + if not self.is_nextn + else None ), **topk_kwargs, )