chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Add Llama4 attention backend auto-selection (#13421)

Browse Source 
Signed-off-by: jbernloehr <jbernloehr@nvidia.com>
This commit is contained in:
Jan Bernlöhr
2025-11-25 20:54:21 +01:00
committed by GitHub
parent 215a97fa6c
commit fcccaf9001
2 changed files with 24 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
python/sglang/srt/server_args.py
View File

@@ -1110,6 +1110,21 @@ class ServerArgs:
self.disable_hybrid_swa_memory = True
elif "Llama4" in model_arch and self.device != "cpu":
# Auto-select attention backend for Llama4 if not specified
if self.attention_backend is None:
if is_sm100_supported():
self.attention_backend, platform = "trtllm_mha", "sm100"
elif is_sm90_supported():
self.attention_backend, platform = "fa3", "sm90"
elif is_hip():
self.attention_backend, platform = "aiter", "hip"
elif self.device == "xpu":
self.attention_backend, platform = "intel_xpu", "xpu"
else:
self.attention_backend, platform = "triton", "other platforms"
logger.warning(
f"Use {self.attention_backend} as attention backend on {platform} for Llama4 model"
)
assert self.attention_backend in {
"fa3",
"aiter",
@@ -1117,11 +1132,6 @@ class ServerArgs:
"trtllm_mha",
"intel_xpu",
}, f"fa3, aiter, triton, trtllm_mha or intel_xpu is required for Llama4 model but got {self.attention_backend}"
if is_sm100_supported() and self.attention_backend is None:
self.attention_backend = "trtllm_mha"
logger.warning(
"Use trtllm_mha as attention backend on sm100 for Llama4 model"
)
if is_sm100_supported() and self.moe_runner_backend == "auto":
if self.quantization in {"fp8", "modelopt_fp8"}:
self.moe_runner_backend = "flashinfer_trtllm"