From fe6d38d2fa17228f1bd2e21d07b8099f21a96aa7 Mon Sep 17 00:00:00 2001 From: b8zhong Date: Thu, 11 Dec 2025 22:35:25 -0800 Subject: [PATCH] fix: trtllm mha attention auto-selection on sm120 (#14842) Co-authored-by: Brayden Zhong --- python/sglang/srt/server_args.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index e8e893da1..a7378b3b9 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -42,7 +42,6 @@ from sglang.srt.utils.common import ( get_device, get_device_memory_capacity, get_device_sm, - is_blackwell, is_blackwell_supported, is_cuda, is_fa3_default_architecture, @@ -1373,7 +1372,8 @@ class ServerArgs: 1. Models with MHA Architecture (e.g: Llama, QWen) 1.1 We will turn on FA3 on hopper unless user use spec decode with topk > 1 or page_size > 1. - 1.2 Use trtllm_mha for Blackwell excluding spec with topk > 1. + 1.2 Use trtllm_mha for SM100/SM103 (Blackwell B200/GB200/B300) excluding spec with topk > 1. + Note: trtllm_mha does not support SM120, which will fall back to flashinfer. 1.3 In other cases, we will use flashinfer if available, otherwise use triton. 2. Models with MLA Architecture and using FA3 2.1 We will use FA3 backend on hopper. @@ -1389,7 +1389,7 @@ class ServerArgs: and is_fa3_default_architecture(self.model_config.hf_config) ): self.attention_backend = "fa3" - elif is_blackwell() and is_no_spec_infer_or_topk_one(self): + elif is_sm100_supported() and is_no_spec_infer_or_topk_one(self): self.attention_backend = "trtllm_mha" elif is_hip(): self.attention_backend = "aiter"