Fix Gemma4 NVFP4 MoE default attention backend (#26791)
This commit is contained in:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user