From bc727bc4ee851d159f7d5a2a19db2bb279a61c55 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Sat, 5 Sep 2026 09:07:13 +0800 Subject: [PATCH] [FP8] SM120: route FP8 linear to per-tensor (cudnn/nvjet) instead of channelwise cutlass (#38006) Co-authored-by: BBuf --- python/sglang/srt/layers/quantization/fp8.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index b76bd2c3a..a4b0d9683 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -87,6 +87,7 @@ from sglang.srt.runtime_context import ( from sglang.srt.utils import ( cpu_has_amx_support, get_bool_env_var, + get_device_capability, is_cpu, is_cuda, is_flashinfer_available, @@ -911,9 +912,20 @@ class Fp8LinearMethod(LinearMethodBase): layer.input_scale.data, requires_grad=False ) + # On SM120 (Blackwell RTX 50) the per-tensor FP8 path dispatches + # to the fast cudnn/nvjet SM120 kernel, which beats the + # channelwise cutlass GemmUniversal by ~1.2-2.7x across the + # decode/prefill M range (matches vLLM's per-tensor SM120 + # choice). Requantize per-channel -> per-tensor (max scale) + # there instead. + use_sm120_fp8_pertensor = ( + self.cutlass_fp8_supported + and not self.use_marlin + and get_device_capability()[0] == 12 + ) # cutlass sgl-kernel and marlin only support per-channel scale; aiter supports per-channel scale if ( - self.cutlass_fp8_supported + (self.cutlass_fp8_supported and not use_sm120_fp8_pertensor) or self.use_marlin or (_use_aiter and self.use_aiter_fp8_per_token) ):