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, )