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

Add fmha_v2 attention backend for SM90/120 (#23112)

Browse Source 
Co-authored-by: Yangmin Li <yangminl@nvidia.com>
This commit is contained in:
akhilg-nv
2026-08-18 18:41:03 -07:00
committed by GitHub
co-authored by Yangmin Li
parent baa2251847
commit 5d12280ae7
2 changed files with 92 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/server_args.py
+31 -2
View File
@@ -6013,9 +6013,28 @@ class ServerArgs:
prefill_backend, decode_backend = self._resolved_attention_backends()
if "trtllm_mha" in (prefill_backend, decode_backend):
if prefill_backend == "trtllm_mha" and not is_sm100_supported():
if prefill_backend == "trtllm_mha" and not (
is_sm90_supported() or is_sm100_supported() or is_sm120_supported()
):
raise ValueError(
"TRTLLM MHA backend for prefill is only supported on Blackwell GPUs (SM100). Please use a different prefill backend."
"TRTLLM MHA backend for prefill requires Hopper (SM90), Blackwell (SM100), or SM120 GPUs. "
"Please use a different prefill backend."
)
if (
prefill_backend == "trtllm_mha"
and is_sm120_supported()
and (
self.kv_cache_dtype == "fp8_e4m3"
or (
envs.SGLANG_SKIP_SOFTMAX_PREFILL_THRESHOLD_SCALE_FACTOR.get()
or 0.0
)
> 0
)
):
raise ValueError(
"TRTLLM FMHAv2 prefill on SM120 does not support "
"fp8_e4m3 KV cache or skip-softmax."
)
if decode_backend == "trtllm_mha" and not (
is_sm90_supported() or is_sm100_supported() or is_sm120_supported()
@@ -6023,6 +6042,16 @@ class ServerArgs:
raise ValueError(
"TRTLLM MHA backend for decode is only supported on Hopper (SM90), Blackwell (SM100) and (SM120) GPUs. Please use a different decode backend."
)
if (
prefill_backend == "trtllm_mha"
and not is_sm100_supported()
and (self.enable_prefill_context_parallel or self.attn_cp_size > 1)
):
raise ValueError(
"Prefill context parallelism with the TRTLLM MHA prefill backend "
"requires SM100 (trtllm-gen context kernel): the SM90/SM120 "
"fmha_v2 prefill path does not implement CP shard masking."
)
run_post_process_pass(self, _attention_backend_fa3_fp8_fallback)