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

[Bugfix] Fix model output corruption caused by EPLB rebalance (Eager and CUDA Graph modes) (#18213)

Browse Source 
Co-authored-by: FortPercent <49947620+FortPercent@users.noreply.github.com>
This commit is contained in:
lviy
2026-03-17 18:30:24 -07:00
committed by GitHub
co-authored by FortPercent
parent 41c8651c0d
commit 944355c66f
1 changed files with 22 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/deepseek_v2.py
+22 -5
View File
@@ -568,17 +568,25 @@ class DeepseekV2MoE(nn.Module):
use_reduce_scatter: bool = False, use_reduce_scatter: bool = False,
gemm_output_zero_allocator: BumpAllocator = None, gemm_output_zero_allocator: BumpAllocator = None,
) -> torch.Tensor: ) -> torch.Tensor:
current_stream = torch.cuda.current_stream() current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream) self.alt_stream.wait_stream(current_stream)
shared_output = self._forward_shared_experts( shared_output = self._forward_shared_experts(
hidden_states, gemm_output_zero_allocator hidden_states, gemm_output_zero_allocator
) )
server_args = get_global_server_args()
dispatch_info = (
ExpertLocationDispatchInfo.init_new(layer_id=self.layer_id)
if server_args.enable_eplb
else None
)
with torch.cuda.stream(self.alt_stream): with torch.cuda.stream(self.alt_stream):
# router_logits: (num_tokens, n_experts) # router_logits: (num_tokens, n_experts)
router_logits = self.gate(hidden_states, gemm_output_zero_allocator) router_logits = self.gate(hidden_states, gemm_output_zero_allocator)
topk_output = self.topk(hidden_states, router_logits) topk_output = self.topk(
hidden_states,
router_logits,
expert_location_dispatch_info=dispatch_info,
)
final_hidden_states = self.experts(hidden_states, topk_output) final_hidden_states = self.experts(hidden_states, topk_output)
if not _is_cuda or isinstance(self.experts.quant_method, KTEPWrapperMethod): if not _is_cuda or isinstance(self.experts.quant_method, KTEPWrapperMethod):
final_hidden_states *= self.routed_scaling_factor final_hidden_states *= self.routed_scaling_factor
@@ -605,7 +613,12 @@ class DeepseekV2MoE(nn.Module):
self.shared_experts.gate_up_proj self.shared_experts.gate_up_proj
): ):
return self.forward_cpu(hidden_states, should_allreduce_fusion) return self.forward_cpu(hidden_states, should_allreduce_fusion)
server_args = get_global_server_args()
dispatch_info = (
ExpertLocationDispatchInfo.init_new(layer_id=self.layer_id)
if server_args.enable_eplb
else None
)
if hidden_states.shape[0] > 0: if hidden_states.shape[0] > 0:
if ( if (
not self._fuse_shared_experts_inside_sbo not self._fuse_shared_experts_inside_sbo
@@ -615,7 +628,11 @@ class DeepseekV2MoE(nn.Module):
) )
# router_logits: (num_tokens, n_experts) # router_logits: (num_tokens, n_experts)
router_logits = self.gate(hidden_states, gemm_output_zero_allocator) router_logits = self.gate(hidden_states, gemm_output_zero_allocator)
topk_output = self.topk(hidden_states, router_logits) topk_output = self.topk(
hidden_states,
router_logits,
expert_location_dispatch_info=dispatch_info,
)
else: else:
shared_output = None shared_output = None
topk_output = self.topk.empty_topk_output(hidden_states.device) topk_output = self.topk.empty_topk_output(hidden_states.device)