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

Fix FA3 DeepSeek prefill performance regression (#5624)

Browse Source 
Co-authored-by: ispobock <ispobaoke@gmail.com>
This commit is contained in:
JieXin Liang
2025-04-22 16:41:41 +08:00
committed by GitHub
parent 04f2abcb34
commit 4418f599a5
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -583,13 +583,17 @@ class DeepseekV2AttentionMLA(nn.Module):
return AttnForwardMethod.MLA return AttnForwardMethod.MLA
elif self.attention_backend == "fa3": elif self.attention_backend == "fa3":
# Flash Attention: Use MHA with chunked KV cache when prefilling on long sequences. # Flash Attention: Use MHA with chunked KV cache when prefilling on long sequences.
if forward_batch.extend_prefix_lens_cpu is not None:
sum_extend_prefix_lens = sum(forward_batch.extend_prefix_lens_cpu)
if ( if (
forward_batch.forward_mode.is_extend() forward_batch.forward_mode.is_extend()
and not self.disable_chunked_prefix_cache and not self.disable_chunked_prefix_cache
and not forward_batch.forward_mode.is_target_verify() and not forward_batch.forward_mode.is_target_verify()
and not forward_batch.forward_mode.is_draft_extend() and not forward_batch.forward_mode.is_draft_extend()
and sum(forward_batch.extend_prefix_lens_cpu) and (
>= self.chunked_prefix_cache_threshold sum_extend_prefix_lens >= self.chunked_prefix_cache_threshold
or sum_extend_prefix_lens == 0
)
): ):
return AttnForwardMethod.MHA_CHUNKED_KV return AttnForwardMethod.MHA_CHUNKED_KV
else: else: