[core/model] Use explicit model arch for Llama4 attention backend auto-selection (#24232)
This commit is contained in:
@@ -81,6 +81,10 @@ MIMO_V2_MODEL_ARCHS = (
|
|||||||
"MiMoV2ForCausalLM",
|
"MiMoV2ForCausalLM",
|
||||||
"MiMoV2FlashForCausalLM",
|
"MiMoV2FlashForCausalLM",
|
||||||
)
|
)
|
||||||
|
LLAMA4_MODEL_ARCHS = (
|
||||||
|
"Llama4ForConditionalGeneration",
|
||||||
|
"Llama4ForCausalLM",
|
||||||
|
)
|
||||||
|
|
||||||
SAMPLING_BACKEND_CHOICES = {"flashinfer", "pytorch", "ascend"}
|
SAMPLING_BACKEND_CHOICES = {"flashinfer", "pytorch", "ascend"}
|
||||||
|
|
||||||
@@ -2034,7 +2038,7 @@ class ServerArgs:
|
|||||||
logger.warning(
|
logger.warning(
|
||||||
"Disable hybrid SWA memory for Step3p5ForCausalLM model with hierarchical cache"
|
"Disable hybrid SWA memory for Step3p5ForCausalLM model with hierarchical cache"
|
||||||
)
|
)
|
||||||
elif "Llama4" in model_arch and self.device != "cpu":
|
elif model_arch in LLAMA4_MODEL_ARCHS and self.device != "cpu":
|
||||||
# Auto-select attention backend for Llama4 if not specified
|
# Auto-select attention backend for Llama4 if not specified
|
||||||
if self.attention_backend is None:
|
if self.attention_backend is None:
|
||||||
if is_sm100_supported():
|
if is_sm100_supported():
|
||||||
|
|||||||
Reference in New Issue
Block a user