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

[DeepSeek V3] Run routed experts on main stream in dual-stream MoE (#29142)

Browse Source 
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Khoa Pham
2026-06-26 00:48:59 -07:00
committed by GitHub
co-authored by Cursor Claude Opus 4.8
parent aeb4e98108
commit 999199f9ff
1 changed files with 47 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/deepseek_v2.py
+10 -8
View File
@@ -936,16 +936,12 @@ class DeepseekV2MoE(nn.Module):
) -> 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(
hidden_states, gemm_output_zero_allocator
)
server_args = get_global_server_args() server_args = get_global_server_args()
dispatch_info = ( dispatch_info = (
ExpertLocationDispatchInfo.init_new(layer_id=self.layer_id) ExpertLocationDispatchInfo.init_new(layer_id=self.layer_id)
if server_args.enable_eplb if server_args.enable_eplb
else None else None
) )
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)
if use_flashinfer_trtllm_bypass: if use_flashinfer_trtllm_bypass:
@@ -966,8 +962,11 @@ class DeepseekV2MoE(nn.Module):
expert_location_dispatch_info=dispatch_info, expert_location_dispatch_info=dispatch_info,
**topk_kwargs, **topk_kwargs,
) )
has_shared_output = (
hidden_states.shape[0] > 0 and self.num_fused_shared_experts == 0
)
deferred_finalize = ( deferred_finalize = (
shared_output is not None has_shared_output
and not self._shared_expert_tp1 and not self._shared_expert_tp1
and topk_output.format == TopKOutputFormat.BYPASSED and topk_output.format == TopKOutputFormat.BYPASSED
and self.experts.supports_deferred_finalize and self.experts.supports_deferred_finalize
@@ -977,9 +976,7 @@ class DeepseekV2MoE(nn.Module):
hidden_states, topk_output hidden_states, topk_output
) )
elif use_flashinfer_trtllm_bypass: elif use_flashinfer_trtllm_bypass:
final_hidden_states = self.experts.forward_impl( final_hidden_states = self.experts.forward_impl(hidden_states, topk_output)
hidden_states, topk_output
)
else: else:
final_hidden_states = self.experts(hidden_states, topk_output) final_hidden_states = self.experts(hidden_states, topk_output)
if ( if (
@@ -990,6 +987,11 @@ class DeepseekV2MoE(nn.Module):
): ):
final_hidden_states *= self.routed_scaling_factor final_hidden_states *= self.routed_scaling_factor
with torch.cuda.stream(self.alt_stream):
shared_output = self._forward_shared_experts(
hidden_states, gemm_output_zero_allocator
)
current_stream.wait_stream(self.alt_stream) current_stream.wait_stream(self.alt_stream)
if deferred_finalize: if deferred_finalize: