chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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-18 09:30:24 +08:00
committed by GitHub
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

27
python/sglang/srt/models/deepseek_v2.py
View File

@@ -568,17 +568,25 @@ class DeepseekV2MoE(nn.Module):
use_reduce_scatter: bool = False,
gemm_output_zero_allocator: BumpAllocator = None,
) -> torch.Tensor:
current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream)
shared_output = self._forward_shared_experts(
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):
# router_logits: (num_tokens, n_experts)
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)
if not _is_cuda or isinstance(self.experts.quant_method, KTEPWrapperMethod):
final_hidden_states *= self.routed_scaling_factor
@@ -605,7 +613,12 @@ class DeepseekV2MoE(nn.Module):
self.shared_experts.gate_up_proj
):
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 (
not self._fuse_shared_experts_inside_sbo
@@ -615,7 +628,11 @@ class DeepseekV2MoE(nn.Module):
)
# router_logits: (num_tokens, n_experts)
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:
shared_output = None
topk_output = self.topk.empty_topk_output(hidden_states.device)