From 3d1b591aa15b2b0d98493705d8e6b2ad763cd74a Mon Sep 17 00:00:00 2001
From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com>
Date: Sat, 6 Dec 2025 06:26:03 +0800
Subject: [PATCH] Tiny use trtllm_mha as default when possible (#14291)

---
 python/sglang/srt/server_args.py   | 6 +++++-
 test/srt/test_flash_attention_4.py | 2 ++
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py
index 69440d16a..1f6531b24 100644
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -40,6 +40,7 @@ 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,
@@ -1313,7 +1314,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 In other cases, we will use flashinfer if available, otherwise use triton.
+                1.2 Use trtllm_mha for Blackwell excluding spec with topk > 1.
+                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.
                 2.2 We will use Flashinfer backend on blackwell.
@@ -1328,6 +1330,8 @@ 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):
+                    self.attention_backend = "trtllm_mha"
                 elif is_hip():
                     self.attention_backend = "aiter"
                 else:
diff --git a/test/srt/test_flash_attention_4.py b/test/srt/test_flash_attention_4.py
index 4322263c4..44623a132 100644
--- a/test/srt/test_flash_attention_4.py
+++ b/test/srt/test_flash_attention_4.py
@@ -22,6 +22,8 @@ class TestFlashAttention4(unittest.TestCase):
             "0.8",
             "--prefill-attention-backend",
             "fa4",
+            "--decode-attention-backend",
+            "flashinfer",
         ]
         cls.process = popen_launch_server(
             cls.model,