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
+23 -41
View File
@@ -1768,49 +1768,31 @@ class NativeSparseAttnBackend(
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)
# (G)B300 (SM103) hangs with TRTLLM attention at high concurrency.
# Use TRTLLm ragged attention for SM100 (Blackwell/B200) to avoid FA4 accuracy issues
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,
)
import flashinfer
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
seq_lens = metadata.cache_seqlens_int32
return flashinfer.prefill.trtllm_ragged_attention_deepseek(
query=q,
key=k,
value=v,
workspace_buffer=self.workspace_buffer,
seq_lens=seq_lens,
max_q_len=metadata.max_seq_len_q,
max_kv_len=max_seqlen_k,
bmm1_scale=layer.scaling,
bmm2_scale=1.0,
o_sf_scale=1.0,
batch_size=forward_batch.batch_size,
window_left=-1,
cum_seq_lens_q=cu_seqlens_q,
cum_seq_lens_kv=cu_seqlens_k,
enable_pdl=False,
is_causal=causal,
return_lse=False,
skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_PREFILL_THRESHOLD_SCALE_FACTOR.get(),
)
seq_lens = metadata.cache_seqlens_int32
return flashinfer.prefill.trtllm_ragged_attention_deepseek(
query=q,
key=k,
value=v,
workspace_buffer=self.workspace_buffer,
seq_lens=seq_lens,
max_q_len=metadata.max_seq_len_q,
max_kv_len=max_seqlen_k,
bmm1_scale=layer.scaling,
bmm2_scale=1.0,
o_sf_scale=1.0,
batch_size=forward_batch.batch_size,
window_left=-1,
cum_seq_lens_q=cu_seqlens_q,
cum_seq_lens_kv=cu_seqlens_k,
enable_pdl=False,
is_causal=causal,
return_lse=False,
skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_PREFILL_THRESHOLD_SCALE_FACTOR.get(),
)
# Use FA3 for SM90 (Hopper/H200)
return flash_attn_varlen_func(
python/sglang/srt/server_args.py
+11 -27
View File
@@ -57,7 +57,6 @@ from sglang.srt.utils.common import (
is_remote_url,
is_sm90_supported,
is_sm100_supported,
is_sm103_supported,
is_sm120_supported,
is_triton_kernels_available,
is_xpu,
@@ -1464,13 +1463,8 @@ class ServerArgs:
self.nsa_decode_backend = "tilelang"
elif kv_cache_dtype == "fp8_e4m3":
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_decode_backend = "trtllm"
self.nsa_prefill_backend = "trtllm"
self.nsa_decode_backend = "trtllm"
else:
# flashmla_auto dispatches to flashmla_sparse/flashmla_kv based on hardware and heuristics
if not user_set_prefill:
@@ -1605,7 +1599,7 @@ class ServerArgs:
if not self.disable_piecewise_cuda_graph:
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 (
self.attention_backend is None
and self.prefill_attention_backend is None
@@ -1703,7 +1697,7 @@ class ServerArgs:
elif model_arch in ["GptOssForCausalLM"]:
# Set attention backend for GPT-OSS
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"
elif is_sm90_supported():
self.attention_backend = "fa3"
@@ -1841,7 +1835,7 @@ class ServerArgs:
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() and not is_sm103_supported():
if is_sm100_supported():
self.attention_backend, platform = "trtllm_mha", "sm100"
elif is_sm90_supported():
self.attention_backend, platform = "fa3", "sm90"
@@ -1900,7 +1894,7 @@ class ServerArgs:
self.disable_hybrid_swa_memory = True
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"
elif is_cuda() and get_device_sm() >= 80:
self.attention_backend = "fa3"
@@ -1986,7 +1980,7 @@ class ServerArgs:
"Qwen3_5ForConditionalGeneration",
]:
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.
# _get_default_attn_backend handles the eagle_topk check.
# There is only one case where page_size=1 is required,
@@ -2135,7 +2129,6 @@ class ServerArgs:
):
if (
is_sm100_supported()
and not is_sm103_supported()
and self.attention_backend is None
and sm100_default_attention_backend is not None
):
@@ -2250,7 +2243,6 @@ class ServerArgs:
return "fa3"
elif (
is_sm100_supported()
and not is_sm103_supported()
and is_no_spec_infer_or_topk_one(self)
and (
self.speculative_algorithm is None
@@ -2378,13 +2370,9 @@ class ServerArgs:
if self.prefill_attention_backend is not None
else self.attention_backend
)
if prefill_backend == "trtllm_mha" and (
not is_sm100_supported() or is_sm103_supported()
):
if prefill_backend == "trtllm_mha" and not is_sm100_supported():
raise ValueError(
"TRTLLM MHA backend for prefill is only supported on SM100. "
"(G)B300 (SM103) is temporarily disabled due to hangs at high concurrency. "
"Please use a different prefill backend."
"TRTLLM MHA backend for prefill is only supported on Blackwell GPUs (SM100). Please use a different prefill backend."
)
# Check decode backend
@@ -2394,14 +2382,10 @@ class ServerArgs:
else self.attention_backend
)
if decode_backend == "trtllm_mha" and not (
is_sm90_supported()
or (is_sm100_supported() and not is_sm103_supported())
or is_sm120_supported()
is_sm90_supported() or is_sm100_supported() or is_sm120_supported()
):
raise ValueError(
"TRTLLM MHA backend for decode is only supported on Hopper (SM90), SM100, and SM120 GPUs. "
"(G)B300 (SM103) is temporarily disabled due to hangs at high concurrency. "
"Please use a different decode backend."
"TRTLLM MHA backend for decode is only supported on Hopper (SM90), Blackwell (SM100) and (SM120) GPUs. Please use a different decode backend."
)
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)(
partial(
_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)
)
)
# 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)(
partial(
_check_cuda_device_version, device_capability_majors=[9], cuda_version=(12, 3)