minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

fix(mtp): avoid mtp perf regression in deepseek when enable eplb (#28982)

Browse Source 
Co-authored-by: chunjiang.hc <chunjiang.hc@alibaba-inc.com>
This commit is contained in:
ICENacl
2026-07-09 11:29:07 -07:00
committed by GitHub
co-authored by chunjiang.hc
parent 462b6171bd
commit b717546fab
1 changed files with 14 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/deepseek_v2.py
+14 -6
View File
@@ -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,
)