diff --git a/docs/docs/advanced_features/server_arguments.mdx b/docs/docs/advanced_features/server_arguments.mdx index 06111fe8c..898cdd948 100644 --- a/docs/docs/advanced_features/server_arguments.mdx +++ b/docs/docs/advanced_features/server_arguments.mdx @@ -1436,9 +1436,9 @@ Please consult the documentation below and [server_args.py](https://github.com/s `--fp8-gemm-backend` - 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). + Choose the runner backend for Blockwise FP8 GEMM operations. For MXFP8 dense GEMM, auto selects flashinfer_cutedsl when FlashInfer reports support (currently SM100/SM103), and otherwise selects flashinfer_cutlass 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). `auto` - auto, deep_gemm, flashinfer_trtllm, flashinfer_cutlass, flashinfer_deepgemm, cutlass, triton, aiter + auto, deep_gemm, flashinfer_trtllm, flashinfer_cutlass, flashinfer_cutedsl, flashinfer_deepgemm, cutlass, triton, aiter `--fp4-gemm-backend` diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 4990d60bd..a77e1fbe4 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/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 diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index f5a044969..ec9987b10 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/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: diff --git a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py index 85818efa7..34ef49ea2 100644 --- a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py +++ b/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 ( - "modelopt", - "modelopt_fp8", - "modelopt_mixed", - ) - # SM120 satisfies is_blackwell_supported(), so resolve_mxfp8_dense_gemm_backend - # 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 model_quantization == "mxfp8": + fp8_gemm_needs_autotune = resolve_mxfp8_dense_gemm_backend().is_flashinfer() + elif model_quantization in ("modelopt", "modelopt_fp8", "modelopt_mixed"): + fp8_gemm_needs_autotune = flashinfer_per_tensor_fp8_supported() + else: + fp8_gemm_needs_autotune = False if not (moe_needs_autotune or fp4_gemm_needs_autotune or fp8_gemm_needs_autotune): return False diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 03a92501f..7deebe8b3 100644 --- a/python/sglang/srt/server_args.py +++ b/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, diff --git a/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py b/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py index ac1d628be..b25995b03 100644 --- a/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py +++ b/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: - return ["triton", "flashinfer_trtllm", "flashinfer_cutlass"] + if get_device_sm() in (100, 103): + 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):