Enable trtllm_mha as gemma4 default attn backend. (#25006)

Co-authored-by: Khoa Pham <khoa.pham@radixark.ai>
This commit is contained in:
Shu Wang
2026-05-17 14:58:12 -07:00
committed by GitHub
co-authored by Khoa Pham
parent 7158a255eb
commit c67b287056
+6 -2
View File
@@ -2192,9 +2192,13 @@ class ServerArgs:
) )
self.disable_hybrid_swa_memory = True self.disable_hybrid_swa_memory = True
elif model_arch == "Gemma4ForConditionalGeneration": elif model_arch == "Gemma4ForConditionalGeneration":
if self.is_attention_backend_not_set(): if is_sm100_supported():
self.attention_backend = "trtllm_mha"
else:
self.attention_backend = "triton" self.attention_backend = "triton"
logger.info("Use triton as default attention backend for Gemma4") logger.info(
f"Use {self.attention_backend} as default attention backend for Gemma4"
)
elif model_arch == "MossVLForConditionalGeneration": elif model_arch == "MossVLForConditionalGeneration":
if self.is_attention_backend_not_set(): if self.is_attention_backend_not_set():
self.prefill_attention_backend = "flashinfer" self.prefill_attention_backend = "flashinfer"