Fix Gemma4 NVFP4 MoE default attention backend (#26791)

This commit is contained in:
Mohammad Miadh Angkad
2026-06-09 14:33:08 +08:00
committed by GitHub
parent d145a6127a
commit a287ab83c0
+13 -3
View File
@@ -1809,7 +1809,8 @@ class ServerArgs:
if parse_connector_type(self.model_path) == ConnectorType.INSTANCE:
return
hf_config = self.get_model_config().hf_config
model_config = self.get_model_config()
hf_config = model_config.hf_config
model_arch = hf_config.architectures[0]
_hybrid_spec = get_linear_attn_spec_by_arch(model_arch)
@@ -2366,8 +2367,17 @@ class ServerArgs:
"Gemma4ForCausalLM",
"Gemma4UnifiedForConditionalGeneration",
):
is_gemma4_modelopt_fp4 = model_config.quantization == "modelopt_fp4"
is_gemma4_moe = getattr(
model_config.hf_text_config, "enable_moe_block", False
)
is_gemma4_modelopt_fp4_moe = is_gemma4_modelopt_fp4 and is_gemma4_moe
# TODO: switch Gemma4 modelopt_fp4 MoE back to trtllm_mha by default
# after the SM10X trtllm_mha accuracy issue is fixed.
default_attention_backend = (
"trtllm_mha" if is_sm100_supported() else "triton"
"trtllm_mha"
if is_sm100_supported() and not is_gemma4_modelopt_fp4_moe
else "triton"
)
if self.is_attention_backend_not_set():
self.attention_backend = default_attention_backend
@@ -2392,7 +2402,7 @@ class ServerArgs:
)
if is_sm100_supported() and self.moe_runner_backend == "auto":
if self.get_model_config().quantization == "modelopt_fp4":
if is_gemma4_modelopt_fp4:
self.quantization = "modelopt_fp4"
self.moe_runner_backend = "flashinfer_trtllm"
logger.info(