add flashinfer cute-dsl backend for mxfp8 gemm (#34042)

Co-authored-by: Brayden Zhong <brayden@radixark.ai>

docs/docs/advanced_features/server_arguments.mdx

@@ -1436,9 +1436,9 @@ Please consult the documentation below and [server_args.py](https://github.com/s
     </tr>
         <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--fp8-gemm-backend`</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the runner backend for Blockwise FP8 GEMM operations. Options: 'auto' (default, auto-selects based on hardware), 'deep_gemm' (JIT-compiled; enabled by default on NVIDIA Hopper (SM90) and Blackwell (SM100) when DeepGEMM is installed), 'flashinfer_trtllm' (FlashInfer TRTLLM backend; SM100/SM103 only), 'flashinfer_cutlass' (FlashInfer CUTLASS backend, SM120 only), 'flashinfer_deepgemm' (Hopper SM90 only, uses swapAB optimization for small M dimensions in decoding), 'cutlass' (optimal for SM120 GPUs), 'triton' (fallback, widely compatible), 'aiter' (ROCm only).</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the runner backend for Blockwise FP8 GEMM operations. For MXFP8 dense GEMM, <code>auto</code> selects <code>flashinfer_cutedsl</code> when FlashInfer reports support (currently SM100/SM103), and otherwise selects <code>flashinfer_cutlass</code> on supported Blackwell GPUs. Options also include 'deep_gemm' (JIT-compiled), 'flashinfer_trtllm' (FlashInfer TRTLLM backend; SM100/SM103 only), 'flashinfer_cutlass' (FlashInfer CUTLASS backend), 'flashinfer_cutedsl' (FlashInfer CuTe DSL MXFP8 backend; SM100/SM103 only), 'flashinfer_deepgemm' (Hopper SM90 only, uses swapAB optimization for small M dimensions in decoding), 'cutlass' (optimal for SM120 GPUs), 'triton' (fallback, widely compatible), and 'aiter' (ROCm only).</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`auto`</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>auto</code>, <code>deep_gemm</code>, <code>flashinfer_trtllm</code>, <code>flashinfer_cutlass</code>, <code>flashinfer_deepgemm</code>, <code>cutlass</code>, <code>triton</code>, <code>aiter</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>auto</code>, <code>deep_gemm</code>, <code>flashinfer_trtllm</code>, <code>flashinfer_cutlass</code>, <code>flashinfer_cutedsl</code>, <code>flashinfer_deepgemm</code>, <code>cutlass</code>, <code>triton</code>, <code>aiter</code></td>
     </tr>
         <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--fp4-gemm-backend`</td>

python/sglang/srt/layers/quantization/fp8.py

@@ -773,7 +773,7 @@ class Fp8LinearMethod(LinearMethodBase):
                 .reshape_as(scale_u8)
                 .contiguous(),
             )
-        elif backend.is_flashinfer_cutlass():
+        elif backend.is_flashinfer_cutlass() or backend.is_flashinfer_cutedsl():
             from flashinfer import block_scale_interleave
 
             scale_u8 = layer.weight_scale_inv.data
@@ -974,7 +974,7 @@ class Fp8LinearMethod(LinearMethodBase):
         if self.use_mxfp8:
             backend = self.mxfp8_dense_backend
             extra_kwargs = {}
-            if backend.is_flashinfer_cutlass():
+            if backend.is_flashinfer_cutlass() or backend.is_flashinfer_cutedsl():
                 weight_scale = layer.weight_scale_inv_swizzled
             elif backend.is_flashinfer_trtllm():
                 weight_scale = layer.weight_scale_inv_shuffled

python/sglang/srt/layers/quantization/fp8_utils.py

@@ -38,6 +38,7 @@ from sglang.srt.utils import (
     get_bool_env_var,
     get_cuda_version,
     get_device_capability,
+    get_device_sm,
     get_hip_version,
     is_blackwell_supported,
     is_cuda,
@@ -276,6 +277,7 @@ class Fp8GemmRunnerBackend(Enum):
     AUTO = "auto"
     FLASHINFER_TRTLLM = "flashinfer_trtllm"
     FLASHINFER_CUTLASS = "flashinfer_cutlass"
+    FLASHINFER_CUTEDSL = "flashinfer_cutedsl"
     FLASHINFER_DEEPGEMM = "flashinfer_deepgemm"
     CUTLASS = "cutlass"
     DEEP_GEMM = "deep_gemm"
@@ -291,6 +293,9 @@ class Fp8GemmRunnerBackend(Enum):
     def is_flashinfer_cutlass(self) -> bool:
         return self == Fp8GemmRunnerBackend.FLASHINFER_CUTLASS
 
+    def is_flashinfer_cutedsl(self) -> bool:
+        return self == Fp8GemmRunnerBackend.FLASHINFER_CUTEDSL
+
     def is_flashinfer_deepgemm(self) -> bool:
         return self == Fp8GemmRunnerBackend.FLASHINFER_DEEPGEMM
 
@@ -312,6 +317,7 @@ class Mxfp8DenseGemmBackend(Enum):
     `Fp8GemmRunnerBackend`."""
 
     FLASHINFER_CUTLASS = "flashinfer_cutlass"
+    FLASHINFER_CUTEDSL = "flashinfer_cutedsl"
     FLASHINFER_TRTLLM = "flashinfer_trtllm"
     DEEP_GEMM = "deep_gemm"
     GFX95_DOT_SCALED = "gfx95_dot_scaled"
@@ -320,9 +326,15 @@ class Mxfp8DenseGemmBackend(Enum):
     def is_flashinfer_cutlass(self) -> bool:
         return self == Mxfp8DenseGemmBackend.FLASHINFER_CUTLASS
 
+    def is_flashinfer_cutedsl(self) -> bool:
+        return self == Mxfp8DenseGemmBackend.FLASHINFER_CUTEDSL
+
     def is_flashinfer_trtllm(self) -> bool:
         return self == Mxfp8DenseGemmBackend.FLASHINFER_TRTLLM
 
+    def is_flashinfer(self) -> bool:
+        return self.value.startswith("flashinfer_")
+
     def is_deep_gemm(self) -> bool:
         return self == Mxfp8DenseGemmBackend.DEEP_GEMM
 
@@ -533,6 +545,28 @@ def resolve_mxfp8_dense_gemm_backend() -> Mxfp8DenseGemmBackend:
             )
         return Mxfp8DenseGemmBackend.FLASHINFER_TRTLLM
 
+    if backend.is_flashinfer_cutedsl():
+        if not (
+            is_blackwell_supported()
+            and is_flashinfer_available()
+            and _raw_flashinfer_mm_mxfp8.is_backend_supported(
+                "cute-dsl", get_device_sm()
+            )
+        ):
+            raise RuntimeError(
+                "MXFP8 dense GEMM requested via --fp8-gemm-backend=flashinfer_cutedsl, "
+                "but that kernel requires an SM100/SM103 GPU and FlashInfer."
+            )
+        return Mxfp8DenseGemmBackend.FLASHINFER_CUTEDSL
+
+    if backend.is_flashinfer_cutlass():
+        if not (is_blackwell_supported() and is_flashinfer_available()):
+            raise RuntimeError(
+                "MXFP8 dense GEMM requested via --fp8-gemm-backend=flashinfer_cutlass, "
+                "but that kernel requires Blackwell GPUs and FlashInfer."
+            )
+        return Mxfp8DenseGemmBackend.FLASHINFER_CUTLASS
+
     if backend.is_deep_gemm():
         if not deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM:
             raise RuntimeError(
@@ -546,6 +580,8 @@ def resolve_mxfp8_dense_gemm_backend() -> Mxfp8DenseGemmBackend:
         return Mxfp8DenseGemmBackend.GFX95_DOT_SCALED
 
     if is_blackwell_supported() and is_flashinfer_available():
+        if _raw_flashinfer_mm_mxfp8.is_backend_supported("cute-dsl", get_device_sm()):
+            return Mxfp8DenseGemmBackend.FLASHINFER_CUTEDSL
         return Mxfp8DenseGemmBackend.FLASHINFER_CUTLASS
 
     if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM:
@@ -570,6 +606,8 @@ def dispatch_w8a8_mxfp8_linear() -> Callable:
         return partial(flashinfer_mxfp8_blockscaled_linear, backend="trtllm")
     elif backend.is_flashinfer_cutlass():
         return partial(flashinfer_mxfp8_blockscaled_linear, backend="cutlass")
+    elif backend.is_flashinfer_cutedsl():
+        return partial(flashinfer_mxfp8_blockscaled_linear, backend="cute-dsl")
     elif backend.is_unsupported():
         return _unsupported_mxfp8_linear
 
@@ -729,14 +767,6 @@ def initialize_fp8_gemm_config(server_args: ServerArgs) -> None:
 
     backend = Fp8GemmRunnerBackend(backend)
 
-    if (
-        backend.is_auto()
-        and server_args.quantization == "mxfp8"
-        and _is_sm100_supported
-        and is_flashinfer_available()
-    ):
-        backend = Fp8GemmRunnerBackend.FLASHINFER_CUTLASS
-
     FP8_GEMM_RUNNER_BACKEND = backend
 
 
@@ -1247,13 +1277,6 @@ def flashinfer_mxfp8_blockscaled_linear(
         else:
             output_dtype = torch.bfloat16
 
-    # At small M the persistent CUTLASS kernel is 2-5x slower than the
-    # CuTe-DSL swap-AB/split-K kernels (both consume the same swizzled
-    # 1D scales).
-    # CuTe-DSL has no mm_mxfp8 kernel on SM120, so the swap is SM100-only there.
-    if backend == "cutlass" and q_input.shape[0] <= 64 and _is_sm100_supported:
-        backend = "cute-dsl"
-
     if backend == "trtllm":
         weight_scale_t = weight_scale.view(-1)
     else:

python/sglang/srt/model_executor/runner/flashinfer_autotune.py

@@ -99,21 +99,16 @@ def should_run_flashinfer_autotune(
     )
 
     from sglang.srt.layers.quantization.fp8_utils import (
-        get_fp8_gemm_runner_backend,
+        flashinfer_per_tensor_fp8_supported,
+        resolve_mxfp8_dense_gemm_backend,
     )
-    from sglang.srt.utils import is_sm100_supported, is_sm120_supported
 
-    model_uses_modelopt_fp8 = model_quantization in (
+    if model_quantization == "mxfp8":
-        "modelopt",
+        fp8_gemm_needs_autotune = resolve_mxfp8_dense_gemm_backend().is_flashinfer()
-        "modelopt_fp8",
+    elif model_quantization in ("modelopt", "modelopt_fp8", "modelopt_mixed"):
-        "modelopt_mixed",
+        fp8_gemm_needs_autotune = flashinfer_per_tensor_fp8_supported()
-    )
+    else:
-    # SM120 satisfies is_blackwell_supported(), so resolve_mxfp8_dense_gemm_backend
+        fp8_gemm_needs_autotune = False
-    # sends it to the same tunable FlashInfer CUTLASS MXFP8 dense GEMM as SM100;
-    # without this the kernel always runs at tactic=-1.
-    fp8_gemm_needs_autotune = get_fp8_gemm_runner_backend().is_flashinfer_cutlass() or (
-        model_uses_modelopt_fp8 and (is_sm100_supported() or is_sm120_supported())
-    )
 
     if not (moe_needs_autotune or fp4_gemm_needs_autotune or fp8_gemm_needs_autotune):
         return False

python/sglang/srt/server_args.py

@@ -295,6 +295,7 @@ FP8_GEMM_RUNNER_BACKEND_CHOICES = [
     "flashinfer_trtllm",
     "flashinfer_cutlass",
     "flashinfer_deepgemm",
+    "flashinfer_cutedsl",
     "cutlass",
     "triton",
     "aiter",
@@ -1734,7 +1735,7 @@ class ServerArgs:
     fp8_gemm_runner_backend: A[
         str,
         Arg(
-            help="Choose the runner backend for Blockwise FP8 GEMM operations. Options: 'auto' (default, auto-selects based on hardware), 'deep_gemm' (JIT-compiled; enabled by default on NVIDIA Hopper (SM90) and Blackwell (SM100) when DeepGEMM is installed), 'flashinfer_trtllm' (optimal for Blackwell and low-latency), 'flashinfer_cutlass' (FlashInfer CUTLASS groupwise FP8 GEMM), 'flashinfer_deepgemm' (Hopper SM90 only; uses swapAB optimization for small M dimensions in decoding), 'cutlass' (optimal for SM120 GPUs), 'triton' (fallback, widely compatible), 'aiter' (ROCm only). ",
+            help="Choose the runner backend for Blockwise FP8 GEMM operations. Options: 'auto' (default, auto-selects based on hardware; MXFP8 dense picks flashinfer_cutedsl on SM100/SM103 and FlashInfer CUTLASS on other supported Blackwell GPUs), 'deep_gemm' (JIT-compiled; enabled by default on NVIDIA Hopper (SM90) and Blackwell (SM100) when DeepGEMM is installed), 'flashinfer_trtllm' (optimal for Blackwell and low-latency), 'flashinfer_cutlass' (FlashInfer CUTLASS groupwise FP8 GEMM), 'flashinfer_cutedsl' (FlashInfer CuTe DSL MXFP8 GEMM on SM100/SM103), 'flashinfer_deepgemm' (Hopper SM90 only; uses swapAB optimization for small M dimensions in decoding), 'cutlass' (optimal for SM120 GPUs), 'triton' (fallback, widely compatible), 'aiter' (ROCm only). ",
             cli_name="--fp8-gemm-backend",
             choices=FP8_GEMM_RUNNER_BACKEND_CHOICES,
             resolvable=True,

test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py

@@ -66,8 +66,13 @@ def _fp8_block_backends():
 
 def _mxfp8_backends():
     # MXFP8 linear is validated on SM100/103 only.
-    if 100 <= get_device_sm() < 110:
+    if get_device_sm() in (100, 103):
-        return ["triton", "flashinfer_trtllm", "flashinfer_cutlass"]
+        return [
+            "auto",
+            "flashinfer_trtllm",
+            "flashinfer_cutlass",
+            "flashinfer_cutedsl",
+        ]
     return []
 
 
@@ -188,15 +193,50 @@ class TestMxfp8LinearBackends(_LinearBackendCheck):
     def _run(self, backend: str):
         self._check_backend(backend, _mxfp8_backends(), MXFP8_SHAPES, self._build_layer)
 
-    def test_triton(self):
-        self._run("triton")
-
     def test_flashinfer_trtllm(self):
         self._run("flashinfer_trtllm")
 
     def test_flashinfer_cutlass(self):
         self._run("flashinfer_cutlass")
 
+    def test_flashinfer_cutedsl(self):
+        self._run("flashinfer_cutedsl")
+
+    def test_auto(self):
+        if "auto" not in _mxfp8_backends():
+            self.skipTest(f"auto not in SM{get_device_sm()} MXFP8 backend set")
+        with mock.patch.object(
+            fp8_utils,
+            "FP8_GEMM_RUNNER_BACKEND",
+            Fp8GemmRunnerBackend.AUTO,
+        ):
+            self.assertEqual(
+                fp8_utils.resolve_mxfp8_dense_gemm_backend(),
+                fp8_utils.Mxfp8DenseGemmBackend.FLASHINFER_CUTEDSL,
+            )
+        self._run("auto")
+
+    @unittest.skipUnless(get_device_sm() >= 100, "Requires Blackwell FlashInfer")
+    def test_auto_falls_back_when_cutedsl_is_unsupported(self):
+        with (
+            mock.patch.object(
+                fp8_utils,
+                "FP8_GEMM_RUNNER_BACKEND",
+                Fp8GemmRunnerBackend.AUTO,
+            ),
+            mock.patch.object(fp8_utils, "get_device_sm", return_value=107),
+            mock.patch.object(
+                fp8_utils._raw_flashinfer_mm_mxfp8,
+                "is_backend_supported",
+                return_value=False,
+            ) as is_backend_supported,
+        ):
+            self.assertEqual(
+                fp8_utils.resolve_mxfp8_dense_gemm_backend(),
+                fp8_utils.Mxfp8DenseGemmBackend.FLASHINFER_CUTLASS,
+            )
+            is_backend_supported.assert_called_once_with("cute-dsl", 107)
+
 
 @unittest.skipIf(get_device_sm() < 90, "FP8 GEMM backends require SM90+")
 class TestModeloptFp8PerTensorLinear(_LinearBackendCheck):