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

avoid cudaStreamSynchronize in DeepSeekV2AttentionMLA (#4577)

Browse Source 
Co-authored-by: Zhang Kaihong <zhangkaihong.zkh@alibaba-inc.com>
This commit is contained in:
strgrb
2025-03-20 01:02:26 +08:00
committed by GitHub
parent 4942074174
commit df7014a8d2
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -658,7 +658,7 @@ class DeepseekV2AttentionMLA(nn.Module):
and forward_batch.forward_mode.is_extend()
and not forward_batch.forward_mode.is_target_verify()
and not forward_batch.forward_mode.is_draft_extend()
and forward_batch.extend_prefix_lens.sum() == 0
and sum(forward_batch.extend_prefix_lens_cpu) == 0
)
else:
# Triton: Use normal computation for prefill and use weight absorption for extend/decode
@@ -666,7 +666,7 @@ class DeepseekV2AttentionMLA(nn.Module):
forward_batch.forward_mode.is_extend()
and not forward_batch.forward_mode.is_target_verify()
and not forward_batch.forward_mode.is_draft_extend()
and forward_batch.extend_prefix_lens.sum() == 0
and sum(forward_batch.extend_prefix_lens_cpu) == 0
)
def forward(