[Revision] Replace enable_flashinfer_mla argument with attention_backend (#5052)

This commit is contained in:
Baizhou Zhang
2025-04-05 01:23:02 -07:00
committed by GitHub
parent ca8d02abd5
commit efbae697b3
9 changed files with 92 additions and 82 deletions

View File

@@ -686,7 +686,6 @@ class DeepseekV2AttentionMLA(nn.Module):
self.w_vc = None
self.w_scale = None
self.enable_flashinfer_mla = global_server_args_dict["enable_flashinfer_mla"]
self.flashinfer_mla_disable_ragged = global_server_args_dict[
"flashinfer_mla_disable_ragged"
]
@@ -694,7 +693,7 @@ class DeepseekV2AttentionMLA(nn.Module):
self.rocm_fused_decode_mla = os.getenv("SGLANG_ROCM_FUSED_DECODE_MLA") == "1"
def no_absorb(self, forward_batch: ForwardBatch) -> bool:
if self.enable_flashinfer_mla:
if self.attention_backend == "flashinfer":
# Flashinfer MLA: Do not absorb when enabling ragged prefill
return (
not self.flashinfer_mla_disable_ragged