From fcccaf9001ab88fabdc4ed457e2514936850855a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20Bernl=C3=B6hr?= <jan@bernloehrs.de>
Date: Tue, 25 Nov 2025 20:54:21 +0100
Subject: [PATCH] Add Llama4 attention backend auto-selection (#13421)

Signed-off-by: jbernloehr <jbernloehr@nvidia.com>
---
 docs/basic_usage/llama4.md       |  9 +++++++++
 python/sglang/srt/server_args.py | 20 +++++++++++++++-----
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/docs/basic_usage/llama4.md b/docs/basic_usage/llama4.md
index e663f9da6..b91720818 100644
--- a/docs/basic_usage/llama4.md
+++ b/docs/basic_usage/llama4.md
@@ -21,6 +21,15 @@ python3 -m sglang.launch_server \
 
 - **OOM Mitigation**: Adjust `--context-length` to avoid a GPU out-of-memory issue. For the Scout model, we recommend setting this value up to 1M on 8\*H100 and up to 2.5M on 8\*H200. For the Maverick model, we don't need to set context length on 8\*H200. When hybrid kv cache is enabled, `--context-length` can be set up to 5M on 8\*H100 and up to 10M on 8\*H200 for the Scout model.
 
+- **Attention Backend Auto-Selection**: SGLang automatically selects the optimal attention backend for Llama 4 based on your hardware. You typically don't need to specify `--attention-backend` manually:
+  - **Blackwell GPUs (B200/GB200)**: `trtllm_mha`
+  - **Hopper GPUs (H100/H200)**: `fa3`
+  - **AMD GPUs**: `aiter`
+  - **Intel XPU**: `intel_xpu`
+  - **Other platforms**: `triton` (fallback)
+
+  To override the auto-selection, explicitly specify `--attention-backend` with one of the supported backends: `fa3`, `aiter`, `triton`, `trtllm_mha`, or `intel_xpu`.
+
 - **Chat Template**: Add `--chat-template llama-4` for chat completion tasks.
 - **Enable Multi-Modal**: Add `--enable-multimodal` for multi-modal capabilities.
 - **Enable Hybrid-KVCache**: Add `--hybrid-kvcache-ratio` for hybrid kv cache. Details can be seen in [this PR](https://github.com/sgl-project/sglang/pull/6563)
diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py
index 4a06fdc80..b99eb4762 100644
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -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"