Enable Llama 4 + TRTLLM MHA (#12003)

2025-10-23 18:22:58 -07:00
parent 823b442945
commit 47e12e082e
1 changed files with 7 additions and 1 deletions
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -966,7 +966,13 @@ class ServerArgs:
                "fa3",
                "aiter",
                "triton",
-            }, "fa3, aiter, or triton is required for Llama4 model"
+                "trtllm_mha",
+            }, "fa3, aiter, triton, or trtllm_mha is required for Llama4 model"
+            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"
+                )
        elif model_arch in [
            "Gemma2ForCausalLM",
            "Gemma3ForCausalLM",