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:
|
if parse_connector_type(self.model_path) == ConnectorType.INSTANCE:
|
||||||
return
|
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]
|
model_arch = hf_config.architectures[0]
|
||||||
|
|
||||||
_hybrid_spec = get_linear_attn_spec_by_arch(model_arch)
|
_hybrid_spec = get_linear_attn_spec_by_arch(model_arch)
|
||||||
@@ -2366,8 +2367,17 @@ class ServerArgs:
|
|||||||
"Gemma4ForCausalLM",
|
"Gemma4ForCausalLM",
|
||||||
"Gemma4UnifiedForConditionalGeneration",
|
"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 = (
|
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():
|
if self.is_attention_backend_not_set():
|
||||||
self.attention_backend = default_attention_backend
|
self.attention_backend = default_attention_backend
|
||||||
@@ -2392,7 +2402,7 @@ class ServerArgs:
|
|||||||
)
|
)
|
||||||
|
|
||||||
if is_sm100_supported() and self.moe_runner_backend == "auto":
|
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.quantization = "modelopt_fp4"
|
||||||
self.moe_runner_backend = "flashinfer_trtllm"
|
self.moe_runner_backend = "flashinfer_trtllm"
|
||||||
logger.info(
|
logger.info(
|
||||||
|
|||||||
Reference in New Issue
Block a user