minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

fix: trtllm mha attention auto-selection on sm120 (#14842)

Browse Source 
Co-authored-by: Brayden Zhong <b8zhong@users.noreply.github.com>
This commit is contained in:
b8zhong
2025-12-11 22:35:25 -08:00
committed by GitHub
co-authored by Brayden Zhong
parent dab31e4cc7
commit fe6d38d2fa
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/server_args.py
+3 -3
View File
@@ -42,7 +42,6 @@ from sglang.srt.utils.common import (
get_device, get_device,
get_device_memory_capacity, get_device_memory_capacity,
get_device_sm, get_device_sm,
is_blackwell,
is_blackwell_supported, is_blackwell_supported,
is_cuda, is_cuda,
is_fa3_default_architecture, is_fa3_default_architecture,
@@ -1373,7 +1372,8 @@ class ServerArgs:
1. Models with MHA Architecture (e.g: Llama, QWen) 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.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. 1.3 In other cases, we will use flashinfer if available, otherwise use triton.
2. Models with MLA Architecture and using FA3 2. Models with MLA Architecture and using FA3
2.1 We will use FA3 backend on hopper. 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) and is_fa3_default_architecture(self.model_config.hf_config)
): ):
self.attention_backend = "fa3" 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" self.attention_backend = "trtllm_mha"
elif is_hip(): elif is_hip():
self.attention_backend = "aiter" self.attention_backend = "aiter"