EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

[perf] slightly imporve DeepSeek-R1-FP4 TP8 (#7481)

Browse Source
This commit is contained in:
JieXin Liang
2025-06-24 03:05:30 +08:00
committed by GitHub
parent 34c3f9b2d3
commit bdbb8d009a
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -362,12 +362,14 @@ class DeepseekV2MoE(nn.Module):
return self.forward_deepep(hidden_states, forward_batch)
def forward_normal_dual_stream(self, hidden_states: torch.Tensor) -> torch.Tensor:
# router_logits: (num_tokens, n_experts)
router_logits = self.gate(hidden_states)
current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream)
shared_output = self._forward_shared_experts(hidden_states)
with torch.cuda.stream(self.alt_stream):
# router_logits: (num_tokens, n_experts)
router_logits = self.gate(hidden_states)
final_hidden_states = self.experts(
hidden_states=hidden_states, router_logits=router_logits
)