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

[feat] Support different attention backends for prefill and decode (#6338)

Browse Source 
Co-authored-by: tianqilin.99 <tianqilin.99@bytedance.com>
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
This commit is contained in:
Qiaolin Yu
2025-07-27 20:42:29 -07:00
committed by GitHub
parent fe6a445d1e
commit 2810338401
9 changed files with 350 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -925,7 +925,10 @@ class DeepseekV2AttentionMLA(nn.Module):
self.disable_chunked_prefix_cache = global_server_args_dict[
"disable_chunked_prefix_cache"
]
self.attention_backend = global_server_args_dict["attention_backend"]
self.current_attention_backend = (
None # Attention backend used by current forward batch
)
self.rocm_fused_decode_mla = get_bool_env_var(
"SGLANG_ROCM_FUSED_DECODE_MLA", "false"
)
@@ -1009,9 +1012,16 @@ class DeepseekV2AttentionMLA(nn.Module):
else:
return AttnForwardMethod.MLA
if self.attention_backend == "ascend":
# Determine attention backend used by current forward batch
if forward_batch.forward_mode.is_decode_or_idle():
attention_backend = global_server_args_dict["decode_attention_backend"]
else:
attention_backend = global_server_args_dict["prefill_attention_backend"]
self.current_attention_backend = attention_backend
if attention_backend == "ascend":
return AttnForwardMethod.MLA
elif self.attention_backend == "flashinfer":
elif attention_backend == "flashinfer":
# Flashinfer MLA: Do not absorb when enabling ragged prefill
if (
not self.flashinfer_mla_disable_ragged
@@ -1023,7 +1033,7 @@ class DeepseekV2AttentionMLA(nn.Module):
return AttnForwardMethod.MHA
else:
return _dispatch_mla_subtype()
elif self.attention_backend == "fa3":
elif attention_backend == "fa3":
# 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)
@@ -1040,7 +1050,7 @@ class DeepseekV2AttentionMLA(nn.Module):
return AttnForwardMethod.MHA_CHUNKED_KV
else:
return _dispatch_mla_subtype()
elif self.attention_backend == "aiter":
elif attention_backend == "aiter":
if (
forward_batch.forward_mode.is_extend()
and not forward_batch.forward_mode.is_target_verify()
@@ -1288,9 +1298,9 @@ class DeepseekV2AttentionMLA(nn.Module):
self, q_pe, k_pe, q_nope_out, k_nope, forward_batch, zero_allocator
):
if (
self.attention_backend == "fa3"
or self.attention_backend == "flashinfer"
or self.attention_backend == "cutlass_mla"
self.current_attention_backend == "fa3"
or self.current_attention_backend == "flashinfer"
or self.current_attention_backend == "cutlass_mla"
):
attn_output = self.attn_mqa(
q_nope_out, k_nope, k_nope, forward_batch, q_rope=q_pe, k_rope=k_pe