[Quantization] Route per-tensor FP8 checkpoints to FlashInfer on SM90 (#33148)

Co-authored-by: Brayden Zhong <b8zhong@uwaterloo.ca>

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)

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,