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

Revert "[Bugfix] Temporarily skip TRTLLM attention on (G)B300 (SM103) to avoid high-concurrency hang" (#22098)

Browse Source
This commit is contained in:
Baizhou Zhang
2026-04-04 02:17:19 -07:00
committed by GitHub
parent 46bf19cdab
commit bf984ae65d
3 changed files with 34 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/nsa_backend.py
+1 -19
View File
@@ -1768,26 +1768,8 @@ class NativeSparseAttnBackend(
f"cu_seqlens_k has {len(cu_seqlens_k)-1} requests" f"cu_seqlens_k has {len(cu_seqlens_k)-1} requests"
) )
# TODO: Revert FA4 back to trtllm once FlashInfer fixes TRTLLM attention on SM103 (#21904) # Use TRTLLm ragged attention for SM100 (Blackwell/B200) to avoid FA4 accuracy issues
# (G)B300 (SM103) hangs with TRTLLM attention at high concurrency.
if self.device_sm_major >= 10: if self.device_sm_major >= 10:
if self.device_capability == (10, 3):
from sglang.jit_kernel.flash_attention_v4 import (
flash_attn_varlen_func as flash_attn_varlen_func_v4,
)
return flash_attn_varlen_func_v4(
q=q,
k=k,
v=v,
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_k=cu_seqlens_k,
max_seqlen_q=metadata.max_seq_len_q,
max_seqlen_k=max_seqlen_k,
softmax_scale=layer.scaling,
causal=causal,
)
else:
import flashinfer import flashinfer
seq_lens = metadata.cache_seqlens_int32 seq_lens = metadata.cache_seqlens_int32
python/sglang/srt/server_args.py
+9 -25
View File
@@ -57,7 +57,6 @@ from sglang.srt.utils.common import (
is_remote_url, is_remote_url,
is_sm90_supported, is_sm90_supported,
is_sm100_supported, is_sm100_supported,
is_sm103_supported,
is_sm120_supported, is_sm120_supported,
is_triton_kernels_available, is_triton_kernels_available,
is_xpu, is_xpu,
@@ -1464,11 +1463,6 @@ class ServerArgs:
self.nsa_decode_backend = "tilelang" self.nsa_decode_backend = "tilelang"
elif kv_cache_dtype == "fp8_e4m3": elif kv_cache_dtype == "fp8_e4m3":
if major >= 10: if major >= 10:
# TODO: Set sm103 default to trtllm after the hanging bug is fixed (#21904)
if is_sm103_supported():
self.nsa_prefill_backend = "flashmla_sparse"
self.nsa_decode_backend = "flashmla_kv"
else:
self.nsa_prefill_backend = "trtllm" self.nsa_prefill_backend = "trtllm"
self.nsa_decode_backend = "trtllm" self.nsa_decode_backend = "trtllm"
else: else:
@@ -1605,7 +1599,7 @@ class ServerArgs:
if not self.disable_piecewise_cuda_graph: if not self.disable_piecewise_cuda_graph:
logger.info("Piecewise CUDA graph is enabled, use MLA for prefill.") logger.info("Piecewise CUDA graph is enabled, use MLA for prefill.")
if is_sm100_supported() and not is_sm103_supported(): if is_sm100_supported():
if ( if (
self.attention_backend is None self.attention_backend is None
and self.prefill_attention_backend is None and self.prefill_attention_backend is None
@@ -1703,7 +1697,7 @@ class ServerArgs:
elif model_arch in ["GptOssForCausalLM"]: elif model_arch in ["GptOssForCausalLM"]:
# Set attention backend for GPT-OSS # Set attention backend for GPT-OSS
if self.is_attention_backend_not_set(): if self.is_attention_backend_not_set():
if is_sm100_supported() and not is_sm103_supported(): if is_sm100_supported():
self.attention_backend = "trtllm_mha" self.attention_backend = "trtllm_mha"
elif is_sm90_supported(): elif is_sm90_supported():
self.attention_backend = "fa3" self.attention_backend = "fa3"
@@ -1841,7 +1835,7 @@ class ServerArgs:
elif "Llama4" in model_arch and self.device != "cpu": elif "Llama4" in model_arch and self.device != "cpu":
# Auto-select attention backend for Llama4 if not specified # Auto-select attention backend for Llama4 if not specified
if self.attention_backend is None: if self.attention_backend is None:
if is_sm100_supported() and not is_sm103_supported(): if is_sm100_supported():
self.attention_backend, platform = "trtllm_mha", "sm100" self.attention_backend, platform = "trtllm_mha", "sm100"
elif is_sm90_supported(): elif is_sm90_supported():
self.attention_backend, platform = "fa3", "sm90" self.attention_backend, platform = "fa3", "sm90"
@@ -1900,7 +1894,7 @@ class ServerArgs:
self.disable_hybrid_swa_memory = True self.disable_hybrid_swa_memory = True
if self.attention_backend is None: if self.attention_backend is None:
if is_cuda() and is_sm100_supported() and not is_sm103_supported(): if is_cuda() and is_sm100_supported():
self.attention_backend = "trtllm_mha" self.attention_backend = "trtllm_mha"
elif is_cuda() and get_device_sm() >= 80: elif is_cuda() and get_device_sm() >= 80:
self.attention_backend = "fa3" self.attention_backend = "fa3"
@@ -1986,7 +1980,7 @@ class ServerArgs:
"Qwen3_5ForConditionalGeneration", "Qwen3_5ForConditionalGeneration",
]: ]:
sm100_default_attn_backend = "triton" sm100_default_attn_backend = "triton"
if is_sm100_supported() and not is_sm103_supported(): if is_sm100_supported():
# trtllm_mha requires speculative_eagle_topk == 1 and page_size > 1. # trtllm_mha requires speculative_eagle_topk == 1 and page_size > 1.
# _get_default_attn_backend handles the eagle_topk check. # _get_default_attn_backend handles the eagle_topk check.
# There is only one case where page_size=1 is required, # There is only one case where page_size=1 is required,
@@ -2135,7 +2129,6 @@ class ServerArgs:
): ):
if ( if (
is_sm100_supported() is_sm100_supported()
and not is_sm103_supported()
and self.attention_backend is None and self.attention_backend is None
and sm100_default_attention_backend is not None and sm100_default_attention_backend is not None
): ):
@@ -2250,7 +2243,6 @@ class ServerArgs:
return "fa3" return "fa3"
elif ( elif (
is_sm100_supported() is_sm100_supported()
and not is_sm103_supported()
and is_no_spec_infer_or_topk_one(self) and is_no_spec_infer_or_topk_one(self)
and ( and (
self.speculative_algorithm is None self.speculative_algorithm is None
@@ -2378,13 +2370,9 @@ class ServerArgs:
if self.prefill_attention_backend is not None if self.prefill_attention_backend is not None
else self.attention_backend else self.attention_backend
) )
if prefill_backend == "trtllm_mha" and ( if prefill_backend == "trtllm_mha" and not is_sm100_supported():
not is_sm100_supported() or is_sm103_supported()
):
raise ValueError( raise ValueError(
"TRTLLM MHA backend for prefill is only supported on SM100. " "TRTLLM MHA backend for prefill is only supported on Blackwell GPUs (SM100). Please use a different prefill backend."
"(G)B300 (SM103) is temporarily disabled due to hangs at high concurrency. "
"Please use a different prefill backend."
) )
# Check decode backend # Check decode backend
@@ -2394,14 +2382,10 @@ class ServerArgs:
else self.attention_backend else self.attention_backend
) )
if decode_backend == "trtllm_mha" and not ( if decode_backend == "trtllm_mha" and not (
is_sm90_supported() is_sm90_supported() or is_sm100_supported() or is_sm120_supported()
or (is_sm100_supported() and not is_sm103_supported())
or is_sm120_supported()
): ):
raise ValueError( raise ValueError(
"TRTLLM MHA backend for decode is only supported on Hopper (SM90), SM100, and SM120 GPUs. " "TRTLLM MHA backend for decode is only supported on Hopper (SM90), Blackwell (SM100) and (SM120) GPUs. Please use a different decode backend."
"(G)B300 (SM103) is temporarily disabled due to hangs at high concurrency. "
"Please use a different decode backend."
) )
if self.page_size not in [16, 32, 64]: if self.page_size not in [16, 32, 64]:
python/sglang/srt/utils/common.py
-17
View File
@@ -234,18 +234,6 @@ def _check_cuda_device_version(
) )
def _check_cuda_device_exact(
device_capability: Tuple[int, int], cuda_version: Tuple[int, int]
):
"""Check for an exact compute capability (major, minor) match."""
if not is_cuda():
return False
return (
torch.cuda.get_device_capability() == device_capability
and tuple(map(int, torch.version.cuda.split(".")[:2])) >= cuda_version
)
is_ampere_with_cuda_12_3 = lru_cache(maxsize=1)( is_ampere_with_cuda_12_3 = lru_cache(maxsize=1)(
partial( partial(
_check_cuda_device_version, device_capability_majors=[8], cuda_version=(12, 3) _check_cuda_device_version, device_capability_majors=[8], cuda_version=(12, 3)
@@ -273,11 +261,6 @@ is_sm100_supported = lru_cache(maxsize=1)(
_check_cuda_device_version, device_capability_majors=[10], cuda_version=(12, 8) _check_cuda_device_version, device_capability_majors=[10], cuda_version=(12, 8)
) )
) )
# TODO(mmangkad): Remove the TRTLLM attention skips for SM103 once FlashInfer
# ships a fix. Tracking: https://github.com/flashinfer-ai/flashinfer/issues/2939
is_sm103_supported = lru_cache(maxsize=1)(
partial(_check_cuda_device_exact, device_capability=(10, 3), cuda_version=(13, 0))
)
is_sm90_supported = lru_cache(maxsize=1)( is_sm90_supported = lru_cache(maxsize=1)(
partial( partial(
_check_cuda_device_version, device_capability_majors=[9], cuda_version=(12, 3) _check_cuda_device_version, device_capability_majors=[9], cuda_version=(12, 3)