[Kernel][Performance] Add FlashInfer cutedsl NVFP4 GEMM backend (#42235)
Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
Mohammad Miadh Angkad committed
Jun 23, 2026 at 04:17 UTC
d1a38c276202b23b5dc8a7bfc7f0b3b83a1ac913
6 files changed
+93
-1
tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py
+3
-1
@@ -75,7 +75,7 @@ def get_ref_results(
75
@pytest.mark.parametrize("shape", SHAPES)
76
@pytest.mark.parametrize("seed", SEEDS)
77
@pytest.mark.parametrize("device", CUDA_DEVICES)
78
-@pytest.mark.parametrize("backend", ["cutlass", "cudnn", "trtllm", "b12x"])
78
+@pytest.mark.parametrize("backend", ["cute-dsl", "cutlass", "cudnn", "trtllm", "b12x"])
79
@pytest.mark.parametrize("autotune", [False, True])
80
@torch.inference_mode()
81
def test_flashinfer_nvfp4_gemm(
@@ -88,6 +88,8 @@ def test_flashinfer_nvfp4_gemm(
88
) -> None:
89
if "trtllm" in backend and dtype == torch.float16:
90
pytest.skip("Only torch.bfloat16 is supported for TRTLLM FP4 GEMM operations")
91
+ if backend == "cute-dsl" and not current_platform.is_device_capability_family(100):
92
+ pytest.skip("FlashInfer cutedsl backend is only supported on SM10x")
93
if backend == "b12x" and not current_platform.has_device_capability(120):
94
pytest.skip("b12x FP4 GEMM requires SM120+ (CC 12.0+)")
95
if backend == "b12x" and not has_flashinfer_b12x_gemm():
tests/models/quantization/test_nvfp4.py
+7
@@ -90,6 +90,7 @@ def test_models(example_prompts, model_name) -> None:
90
EAGER = [True, False]
91
92
SM_100_NVFP4_BACKENDS = [
93
+ "flashinfer_cutedsl",
94
"flashinfer_cudnn",
95
"flashinfer_trtllm",
96
"flashinfer_cutlass",
@@ -102,12 +103,18 @@ SM_100_NVFP4_BACKENDS = [
103
"backend",
104
[
105
"emulation",
106
+ "flashinfer_cutedsl",
107
"flashinfer_cudnn",
108
"flashinfer_trtllm", # the small seq_len ensures trtllm_8x4_layout backend is used
109
"flashinfer_cutlass",
110
],
111
)
112
def test_nvfp4(vllm_runner, model, eager, backend):
113
+ if backend == "flashinfer_cutedsl" and not (
114
+ current_platform.is_device_capability_family(100)
115
+ ):
116
+ pytest.skip("The flashinfer_cutedsl backend is only supported on SM10x")
117
+
118
if (
119
not current_platform.has_device_capability(100)
120
and backend in SM_100_NVFP4_BACKENDS
vllm/compilation/passes/fusion/collective_fusion.py
+9
@@ -956,6 +956,15 @@ class AsyncTPPass(VllmFusionPatternMatcherPass):
956
a_scale_view=a_scale_view,
957
)
958
)
959
+ self.register(
960
+ FlashInferAllGatherFP4Pattern(
961
+ self.model_dtype,
962
+ self.device,
963
+ "cute-dsl",
964
+ use_8x4_sf_layout=False,
965
+ a_scale_view="float8",
966
+ )
967
+ )
968
# NVFP4 reduce-scatter does not need scale communication: FP4
969
# scales are consumed by the local GEMM and only BF16 partial
970
# outputs are reduced. Keep this PR scoped to the all-gather
vllm/config/kernel.py
+2
@@ -141,6 +141,7 @@ LinearBackend = Literal[
141
"auto",
142
"cutlass",
143
"flashinfer_cutlass",
144
+ "flashinfer_cutedsl",
145
"flashinfer_trtllm",
146
"flashinfer_cudnn",
147
"flashinfer_b12x",
@@ -198,6 +199,7 @@ class KernelConfig:
199
- "auto": Automatically select the best backend based on model and hardware
200
- "cutlass": Use CUTLASS-based kernels
201
- "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels
202
+ - "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels
203
- "flashinfer_trtllm": Use FlashInfer with TensorRT-LLM kernels
204
- "flashinfer_cudnn": Use FlashInfer with cuDNN kernels
205
- "flashinfer_b12x": Use FlashInfer b12x CuteDSL NVFP4 GEMM (SM120+)
vllm/model_executor/kernels/linear/__init__.py
+6
@@ -115,6 +115,7 @@ from vllm.model_executor.kernels.linear.nvfp4.fbgemm import (
115
from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
116
FlashInferB12xNvFp4LinearKernel,
117
FlashInferCudnnNvFp4LinearKernel,
118
+ FlashInferCuteDslNvFp4LinearKernel,
119
FlashInferCutlassNvFp4LinearKernel,
120
FlashInferTrtllmNvFp4LinearKernel,
121
)
@@ -209,6 +210,9 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = {
210
FlashInferCutlassNvFp4LinearKernel,
211
FlashInferMxFp4LinearKernel,
212
},
213
+ "flashinfer_cutedsl": {
214
+ FlashInferCuteDslNvFp4LinearKernel,
215
+ },
216
"flashinfer_trtllm": {
217
FlashInferTrtllmNvFp4LinearKernel,
218
},
@@ -399,6 +403,7 @@ _POSSIBLE_MXFP8_KERNELS: dict[PlatformEnum, list[type[Mxfp8LinearKernel]]] = {
403
404
_POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = {
405
PlatformEnum.CUDA: [
406
+ FlashInferCuteDslNvFp4LinearKernel,
407
# FlashInferB12xNvFp4LinearKernel excluded from auto-selection until
408
# upstream CUTLASS SM121 MMA op guard is resolved; use
409
# --linear-backend flashinfer_b12x to opt in explicitly.
@@ -1038,6 +1043,7 @@ __all__ = [
1043
"CutlassNvFp4LinearKernel",
1044
"EmulationNvFp4LinearKernel",
1045
"FbgemmNvFp4LinearKernel",
1046
+ "FlashInferCuteDslNvFp4LinearKernel",
1047
"FlashInferB12xNvFp4LinearKernel",
1048
"FlashInferCutlassNvFp4LinearKernel",
1049
"FlashInferTrtllmNvFp4LinearKernel",
vllm/model_executor/kernels/linear/nvfp4/flashinfer.py
+66
@@ -28,6 +28,72 @@ from vllm.utils.flashinfer import (
28
from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
29
30
31
+class FlashInferCuteDslNvFp4LinearKernel(NvFp4LinearKernel):
32
+ """NVFP4 GEMM via FlashInfer's cutedsl backend."""
33
+
34
+ @classmethod
35
+ def is_supported(
36
+ cls, compute_capability: int | None = None
37
+ ) -> tuple[bool, str | None]:
38
+ if not current_platform.is_device_capability_family(100):
39
+ return False, "FlashInfer cutedsl requires sm_10x"
40
+ if not has_flashinfer():
41
+ return False, "FlashInfer required"
42
+ return True, None
43
+
44
+ @classmethod
45
+ def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
46
+ return True, None
47
+
48
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
49
+ # cutedsl uses the same swizzled + padded layout as cutlass.
50
+ layer.weight_scale = torch.nn.Parameter(
51
+ swizzle_blockscale(layer.weight_scale.data), requires_grad=False
52
+ )
53
+ padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(
54
+ layer.weight.data
55
+ )
56
+ layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False)
57
+ layer.weights_padding_cols = weights_padding_cols
58
+
59
+ def apply_weights(
60
+ self,
61
+ layer: torch.nn.Module,
62
+ x: torch.Tensor,
63
+ bias: torch.Tensor | None = None,
64
+ ) -> torch.Tensor:
65
+ output_size = layer.output_size_per_partition
66
+ output_dtype = x.dtype
67
+ output_shape = [*x.shape[:-1], output_size]
68
+
69
+ x_fp4, x_blockscale = scaled_fp4_quant(
70
+ x,
71
+ layer.input_global_scale_inv,
72
+ is_sf_swizzled_layout=True,
73
+ backend="flashinfer-cutedsl",
74
+ )
75
+
76
+ x_fp4 = pad_nvfp4_activation_for_cutlass(
77
+ x_fp4, getattr(layer, "weights_padding_cols", 0)
78
+ )
79
+
80
+ out = flashinfer_scaled_fp4_mm(
81
+ x_fp4,
82
+ layer.weight,
83
+ x_blockscale,
84
+ layer.weight_scale,
85
+ layer.alpha,
86
+ output_dtype,
87
+ backend="cute-dsl",
88
+ )
89
+
90
+ out = slice_nvfp4_output(out, output_size)
91
+
92
+ if bias is not None:
93
+ out = out + bias
94
+ return out.view(*output_shape)
95
+
96
+
97
class FlashInferCutlassNvFp4LinearKernel(NvFp4LinearKernel):
98
"""NVFP4 GEMM via FlashInfer's CUTLASS wrapper."""
99