From c9506d023f33ed1161714dce1da859157033ea73 Mon Sep 17 00:00:00 2001 From: Aditya Kamat <81307590+adityakamat24@users.noreply.github.com> Date: Wed, 5 Aug 2026 17:54:10 -0700 Subject: [PATCH] [Quantization] Route per-tensor FP8 checkpoints to FlashInfer on SM90 (#33148) Co-authored-by: Brayden Zhong --- .../srt/layers/quantization/fp8_utils.py | 24 ++++++++++++------- .../srt/layers/quantization/modelopt_quant.py | 18 ++++++++------ 2 files changed, 27 insertions(+), 15 deletions(-) diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 38a524bea..da675c84b 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -311,14 +311,15 @@ class Fp8GemmRunnerBackend(Enum): FP8_GEMM_RUNNER_BACKEND: Fp8GemmRunnerBackend | None = None -if is_blackwell_supported() and is_flashinfer_available(): - from flashinfer import SfLayout - from flashinfer import bmm_fp8 as _raw_flashinfer_bmm_fp8 - from flashinfer import mm_mxfp8 as _raw_flashinfer_mm_mxfp8 - from flashinfer import mxfp8_quantize as _raw_flashinfer_mxfp8_quantize - from flashinfer.gemm import gemm_fp8_nt_groupwise as _raw_gemm_fp8_nt_groupwise +@lru_cache(maxsize=1) +def flashinfer_per_tensor_fp8_supported() -> bool: + return is_flashinfer_available() and ( + is_sm90_supported() or is_sm100_supported() or is_sm120_supported() + ) - from sglang.srt.utils.custom_op import register_custom_op + +if flashinfer_per_tensor_fp8_supported(): + from flashinfer import bmm_fp8 as _raw_flashinfer_bmm_fp8 @register_custom_op( op_name="flashinfer_bmm_fp8", @@ -344,6 +345,13 @@ if is_blackwell_supported() and is_flashinfer_available(): backend="cublas", ).view(m, n) + +if is_blackwell_supported() and is_flashinfer_available(): + from flashinfer import SfLayout + from flashinfer import mm_mxfp8 as _raw_flashinfer_mm_mxfp8 + from flashinfer import mxfp8_quantize as _raw_flashinfer_mxfp8_quantize + from flashinfer.gemm import gemm_fp8_nt_groupwise as _raw_gemm_fp8_nt_groupwise + @lru_cache(maxsize=1) def _get_flashinfer_groupwise_backend() -> str: if get_fp8_gemm_runner_backend().is_flashinfer_cutlass(): @@ -1820,7 +1828,7 @@ def apply_fp8_linear_bmm_flashinfer( input_scale: torch.Tensor, bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: - """Per-tensor static fp8 linear via flashinfer bmm_fp8 (SM100/SM120 Blackwell).""" + """Per-tensor static fp8 linear via flashinfer bmm_fp8 (SM90 and newer).""" output_shape = [*input.shape[:-1], weight.shape[1]] input_2d = input.view(-1, input.shape[-1]) qinput, x_scale = static_quant_fp8(input_2d, input_scale, repeat_scale=False) diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index cd2e6e67a..c99dad370 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -42,6 +42,7 @@ from sglang.srt.layers.quantization.fp8_utils import ( apply_fp8_linear_bmm_flashinfer, can_auto_enable_marlin_fp8, cutlass_fp8_supported, + flashinfer_per_tensor_fp8_supported, is_blackwell_supported, ) from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod @@ -66,8 +67,6 @@ from sglang.srt.layers.utils import alias_or_bind_derived_param, copy_or_rebind_ from sglang.srt.utils.common import ( get_device_capability, is_cuda, - is_flashinfer_available, - is_sm100_supported, is_sm120_supported, round_up, set_weight_attrs, @@ -505,9 +504,7 @@ class ModelOptFp8LinearMethod(LinearMethodBase): super().__init__() self.quant_config = quant_config self.cutlass_fp8_supported = cutlass_fp8_supported() - self.enable_flashinfer_bmm = ( - is_sm100_supported() or is_sm120_supported() - ) and is_flashinfer_available() + self.enable_flashinfer_bmm = flashinfer_per_tensor_fp8_supported() self.use_marlin = False if is_cuda(): self.use_marlin = ( @@ -565,13 +562,20 @@ class ModelOptFp8LinearMethod(LinearMethodBase): ) layer.register_parameter(scale_name, scale) + def _can_use_flashinfer_bmm(self, layer: torch.nn.Module) -> bool: + if not self.enable_flashinfer_bmm or layer.input_scale is None: + return False + k, n = layer.weight.shape + return k % 16 == 0 and n % 16 == 0 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: """Requantizes weights after loading using the maximum scale.""" max_w_scale, quantized_weight = requantize_with_max_scale( layer.weight, layer.weight_scale, layer.logical_widths ) layer.weight = Parameter(quantized_weight.t(), requires_grad=False) - if self.cutlass_fp8_supported and not self.enable_flashinfer_bmm: + layer.use_flashinfer_bmm = self._can_use_flashinfer_bmm(layer) + if self.cutlass_fp8_supported and not layer.use_flashinfer_bmm: max_w_scale = convert_to_channelwise(max_w_scale, layer.logical_widths) layer.weight_scale = Parameter(max_w_scale, requires_grad=False) layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False) @@ -597,7 +601,7 @@ class ModelOptFp8LinearMethod(LinearMethodBase): size_k=layer.input_size_per_partition, bias=bias, ) - if self.enable_flashinfer_bmm and layer.input_scale is not None: + if layer.use_flashinfer_bmm: return apply_fp8_linear_bmm_flashinfer( input=x, weight=layer.weight,