minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Brayden Zhong <brayden@radixark.ai>
This commit is contained in:
Carrie Chen
2026-08-13 08:50:01 +08:00
committed by GitHub
co-authored by Brayden Zhong
parent 40eaf34428
commit 6a5a9eccaa
6 changed files with 97 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/fp8.py
+2 -2
View File
@@ -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 -15
View File
@@ -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
+8 -13
View File
@@ -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
python/sglang/srt/server_args.py
+2 -1
View File
@@ -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,