diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 19fd63f34..862488e5f 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -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( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 4adf91b53..d91ced805 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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]: diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 0e274b78b..f65f201ad 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -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)