[Kernel] add Flashinfer cutedsl w4a16 linear (#53014)
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>
Lain committed
Aug 31, 2026 at 19:49 UTC
882ca8d6964458ace0379e0eaa8aac9039b1bffc
8 files changed
+230
-10
docs/features/quantization/modelopt.md
+7
-3
@@ -17,6 +17,8 @@ following `quantization.quant_algo` values:
17
- `FP8_PER_CHANNEL_PER_TOKEN`: per-channel weight scale and dynamic per-token activation quantization.
18
- `FP8_PB_WO` (ModelOpt may emit `fp8_pb_wo`): block-scaled FP8 weight-only (typically 128×128 blocks).
19
- `NVFP4`: ModelOpt NVFP4 checkpoints (use `quantization="modelopt_fp4"`).
20
+- `W4A16_NVFP4`: weight-only ModelOpt NVFP4 checkpoints with 16-bit
21
+ activations (use `quantization="modelopt_fp4"`).
22
- `MXFP8`: ModelOpt MXFP8 checkpoints (use `quantization="modelopt_mxfp8"`).
23
24
!!! note
@@ -28,10 +30,12 @@ following `quantization.quant_algo` values:
30
workloads. Use `--linear-backend` to override the automatic selection
31
(this replaces the deprecated `VLLM_NVFP4_GEMM_BACKEND` environment
32
variable). Values relevant to NVFP4 include `cutlass`,
31
- `flashinfer_cutlass`, `flashinfer_trtllm`, `flashinfer_cudnn`, and
32
- `marlin`; the full list is documented under `KernelConfig` on the
33
+ `flashinfer_cutlass`, `flashinfer_cutedsl`, `flashinfer_trtllm`,
34
+ `flashinfer_cudnn`, and `marlin`; see `KernelConfig` on the
35
[Engine Arguments](../../configuration/engine_args.md) page and shown by
34
- `vllm serve --help=KernelConfig`.
36
+ `vllm serve --help=KernelConfig`. For `W4A16_NVFP4`, `auto` currently
37
+ selects Marlin. Models with BF16 activations can explicitly select the
38
+ FlashInfer CuTe-DSL backend with `--linear-backend flashinfer_cutedsl`.
39
40
!!! note
41
For models quantized to MXFP8 with BF16 activations on SM100-family GPUs,
tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py
+60
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import pytest
4
import torch
5
+import torch.nn.functional as F
6
from nvfp4_utils import (
7
FLOAT4_E2M1_MAX,
8
FLOAT8_E4M3_MAX,
@@ -10,6 +11,13 @@ from nvfp4_utils import (
11
)
12
13
from vllm import _custom_ops as ops
14
+from vllm.model_executor.kernels.linear.nvfp4 import NvFp4LinearLayerConfig
15
+from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
16
+ FlashInferCuteDslNvFp4W4A16LinearKernel,
17
+)
18
+from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
19
+ dequantize_to_dtype,
20
+)
21
from vllm.platforms import current_platform
22
from vllm.utils.flashinfer import (
23
flashinfer_scaled_fp4_mm,
@@ -166,3 +174,55 @@ def test_flashinfer_nvfp4_gemm(
174
)
175
176
torch.testing.assert_close(out, expected_out.to(dtype=dtype), atol=1e-1, rtol=1e-1)
177
+
178
+
179
+@pytest.mark.parametrize("shape", SHAPES)
180
+@pytest.mark.parametrize("seed", SEEDS)
181
+@pytest.mark.parametrize("device", CUDA_DEVICES)
182
+@torch.inference_mode()
183
+def test_flashinfer_cutedsl_nvfp4_w4a16_linear(
184
+ shape: tuple[int, int, int],
185
+ seed: int,
186
+ device: str,
187
+) -> None:
188
+ supported, reason = FlashInferCuteDslNvFp4W4A16LinearKernel.is_supported()
189
+ if not supported:
190
+ pytest.skip(reason)
191
+
192
+ set_random_seed(seed)
193
+ m, n, packed_k = shape
194
+ k = packed_k * 2
195
+ x = torch.randn((m, k), dtype=torch.bfloat16, device=device)
196
+ weight = torch.randn((n, k), dtype=torch.bfloat16, device=device)
197
+ weight_quant_scale = (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / weight.abs().max()).to(
198
+ torch.float32
199
+ )
200
+ weight_global_scale = weight_quant_scale.reciprocal()
201
+ weight_fp4, weight_scale = ops.scaled_fp4_quant(
202
+ weight,
203
+ weight_quant_scale,
204
+ is_sf_swizzled_layout=False,
205
+ )
206
+ weight_ref = dequantize_to_dtype(
207
+ weight_fp4,
208
+ weight_scale,
209
+ weight_global_scale,
210
+ dtype=torch.bfloat16,
211
+ swizzle=False,
212
+ )
213
+
214
+ layer = torch.nn.Module()
215
+ layer.output_size_per_partition = n
216
+ layer.weight = torch.nn.Parameter(weight_fp4, requires_grad=False)
217
+ layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
218
+ layer.weight_global_scale = torch.nn.Parameter(
219
+ weight_global_scale, requires_grad=False
220
+ )
221
+ kernel = FlashInferCuteDslNvFp4W4A16LinearKernel(NvFp4LinearLayerConfig())
222
+
223
+ kernel.process_weights_after_loading(layer)
224
+ output = kernel.apply_weights(layer, x)
225
+
226
+ expected = F.linear(x, weight_ref)
227
+ assert output.shape == (m, n)
228
+ torch.testing.assert_close(output, expected, atol=1e-1, rtol=1e-1)
tests/quantization/test_modelopt.py
+11
-1
@@ -16,6 +16,7 @@ from tests.quantization.utils import is_quant_method_supported
16
from vllm.config import VllmConfig, set_current_vllm_config
17
from vllm.config.model import ModelConfig
18
from vllm.model_executor.kernels.linear import (
19
+ FlashInferCuteDslNvFp4W4A16LinearKernel,
20
HummingNvFp4LinearKernel,
21
MarlinNvFp4LinearKernel,
22
)
@@ -584,10 +585,19 @@ def test_modelopt_nvfp4_config_dispatches_w4a16_method():
585
586
@pytest.mark.parametrize(
587
("linear_backend", "kernel_cls"),
587
- [("auto", MarlinNvFp4LinearKernel), ("humming", HummingNvFp4LinearKernel)],
588
+ [
589
+ ("auto", MarlinNvFp4LinearKernel),
590
+ ("humming", HummingNvFp4LinearKernel),
591
+ ("flashinfer_cutedsl", FlashInferCuteDslNvFp4W4A16LinearKernel),
592
+ ],
593
)
594
@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
595
def test_modelopt_w4a16_respects_linear_backend(linear_backend, kernel_cls):
596
+ if linear_backend != "auto":
597
+ is_supported, reason = kernel_cls.is_supported()
598
+ if not is_supported:
599
+ pytest.skip(reason)
600
+
601
vllm_config = VllmConfig()
602
vllm_config.kernel_config.linear_backend = linear_backend
603
with set_current_vllm_config(vllm_config):
vllm/config/kernel.py
+1
-1
@@ -276,7 +276,7 @@ class KernelConfig:
276
- "cutlass": Use CUTLASS-based kernels
277
- "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels
278
- "flashinfer_cutedsl": Use FlashInfer with CuTe-DSL kernels
279
- (BF16, NVFP4, MXFP8)
279
+ (BF16, NVFP4, MXFP8, W4A16_NVFP4)
280
- "flashinfer_trtllm": Use FlashInfer with TensorRT-LLM kernels
281
- "flashinfer_cudnn": Use FlashInfer with cuDNN kernels
282
- "flashinfer_b12x": Use FlashInfer b12x CuteDSL NVFP4 GEMM (SM120+)
vllm/model_executor/kernels/linear/__init__.py
+20
-3
@@ -149,6 +149,7 @@ from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
149
FlashInferB12xNvFp4LinearKernel,
150
FlashInferCudnnNvFp4LinearKernel,
151
FlashInferCuteDslNvFp4LinearKernel,
152
+ FlashInferCuteDslNvFp4W4A16LinearKernel,
153
FlashInferCutlassNvFp4LinearKernel,
154
FlashInferTrtllmNvFp4LinearKernel,
155
)
@@ -262,6 +263,7 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = {
263
},
264
"flashinfer_cutedsl": {
265
FlashInferCuteDslNvFp4LinearKernel,
266
+ FlashInferCuteDslNvFp4W4A16LinearKernel,
267
FlashInferCutedslMxfp8LinearKernel,
268
},
269
"flashinfer_trtllm": {
@@ -531,6 +533,7 @@ _POSSIBLE_MXFP8_KERNELS: dict[PlatformEnum, list[type[Mxfp8LinearKernel]]] = {
533
_POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = {
534
PlatformEnum.CUDA: [
535
FlashInferCuteDslNvFp4LinearKernel,
536
+ FlashInferCuteDslNvFp4W4A16LinearKernel,
537
FlashInferCutlassNvFp4LinearKernel,
538
FlashInferB12xNvFp4LinearKernel,
539
CutlassNvFp4LinearKernel,
@@ -1014,7 +1017,11 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel:
1017
"""Select and instantiate the best NVFP4 linear kernel for the
1018
current platform."""
1019
config = NvFp4LinearLayerConfig()
1017
- a16_kernels = (MarlinNvFp4LinearKernel, HummingNvFp4LinearKernel)
1020
+ a16_kernels = (
1021
+ FlashInferCuteDslNvFp4W4A16LinearKernel,
1022
+ MarlinNvFp4LinearKernel,
1023
+ HummingNvFp4LinearKernel,
1024
+ )
1025
1026
# VLLM_BATCH_INVARIANT forces deterministic execution. Prefer the
1027
# batch-invariant CUTLASS implementation when available, otherwise fall
@@ -1051,8 +1058,17 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel:
1058
)
1059
force_kernel = EmulationNvFp4LinearKernel
1060
elif linear_backend == "auto" and use_a16:
1054
- # Force a16 (Marlin) when running weight-only quantization.
1055
- force_kernel = MarlinNvFp4LinearKernel
1061
+ _cc = current_platform.get_device_capability()
1062
+ compute_capability = _cc.to_int() if _cc is not None else None
1063
+ # Weight-only: prefer FlashInfer CuTe-DSL W4A16 on SM100/103,
1064
+ # otherwise Marlin.
1065
+ cutedsl_ok, _ = FlashInferCuteDslNvFp4W4A16LinearKernel.is_supported(
1066
+ compute_capability
1067
+ )
1068
+ if compute_capability in (100, 103) and cutedsl_ok:
1069
+ force_kernel = FlashInferCuteDslNvFp4W4A16LinearKernel
1070
+ else:
1071
+ force_kernel = MarlinNvFp4LinearKernel
1072
1073
if force_kernel is not None:
1074
if use_a16 and force_kernel not in a16_kernels:
@@ -1232,6 +1248,7 @@ __all__ = [
1248
"EmulationNvFp4LinearKernel",
1249
"FbgemmNvFp4LinearKernel",
1250
"FlashInferCuteDslNvFp4LinearKernel",
1251
+ "FlashInferCuteDslNvFp4W4A16LinearKernel",
1252
"FlashInferB12xNvFp4LinearKernel",
1253
"FlashInferCutlassNvFp4LinearKernel",
1254
"FlashInferTrtllmNvFp4LinearKernel",
vllm/model_executor/kernels/linear/nvfp4/flashinfer.py
+80
@@ -20,14 +20,94 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
20
)
21
from vllm.platforms import current_platform
22
from vllm.utils.flashinfer import (
23
+ flashinfer_prepare_bf16_fp4_weights,
24
flashinfer_scaled_fp4_mm,
25
has_flashinfer,
26
has_flashinfer_b12x_gemm,
27
+ has_flashinfer_bf16_fp4,
28
)
29
30
from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
31
32
33
+class FlashInferCuteDslNvFp4W4A16LinearKernel(NvFp4LinearKernel):
34
+ """BF16 x NVFP4 GEMM via FlashInfer's CuTe-DSL backend."""
35
+
36
+ @classmethod
37
+ def is_supported(
38
+ cls, compute_capability: int | None = None
39
+ ) -> tuple[bool, str | None]:
40
+ if compute_capability is None:
41
+ if not current_platform.is_cuda():
42
+ return False, "FlashInfer CuTe-DSL W4A16 requires CUDA"
43
+ capability = current_platform.get_device_capability()
44
+ if capability is None:
45
+ return False, "CUDA compute capability is unavailable"
46
+ compute_capability = capability.to_int()
47
+
48
+ if compute_capability not in (100, 103) and not (
49
+ compute_capability >= 120 and compute_capability < 130
50
+ ):
51
+ return False, "FlashInfer CuTe-DSL W4A16 requires sm_100 or sm_12x"
52
+ if not has_flashinfer_bf16_fp4():
53
+ return False, "FlashInfer CuTe-DSL BF16 x FP4 GEMM is unavailable"
54
+ return True, None
55
+
56
+ @classmethod
57
+ def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
58
+ return True, None
59
+
60
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
61
+ padded_weight, weights_padding_bytes = pad_nvfp4_weight_for_cutlass(
62
+ layer.weight.data, alignment=64
63
+ )
64
+ swizzled_scale = swizzle_blockscale(layer.weight_scale.data)
65
+ weight, weight_scale, global_scale = flashinfer_prepare_bf16_fp4_weights(
66
+ padded_weight,
67
+ swizzled_scale,
68
+ layer.weight_global_scale.data.reshape(1),
69
+ backend="cute-dsl",
70
+ )
71
+ assert global_scale is not None
72
+
73
+ layer.weight = torch.nn.Parameter(weight, requires_grad=False)
74
+ layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
75
+ layer.weight_global_scale = torch.nn.Parameter(
76
+ global_scale, requires_grad=False
77
+ )
78
+ layer.weights_padding_cols = weights_padding_bytes
79
+
80
+ def apply_weights(
81
+ self,
82
+ layer: torch.nn.Module,
83
+ x: torch.Tensor,
84
+ bias: torch.Tensor | None = None,
85
+ ) -> torch.Tensor:
86
+ if x.dtype != torch.bfloat16:
87
+ raise ValueError(
88
+ f"FlashInfer CuTe-DSL W4A16 requires BF16 input, got {x.dtype}"
89
+ )
90
+
91
+ output_size = layer.output_size_per_partition
92
+ output_shape = [*x.shape[:-1], output_size]
93
+ x_2d = x.reshape(-1, x.shape[-1])
94
+ weights_padding_bytes = getattr(layer, "weights_padding_cols", 0)
95
+ if weights_padding_bytes:
96
+ x_2d = torch.nn.functional.pad(x_2d, (0, weights_padding_bytes * 2))
97
+ x_2d = x_2d.contiguous()
98
+
99
+ out = torch.ops.vllm.flashinfer_mm_bf16_fp4(
100
+ x_2d,
101
+ layer.weight,
102
+ layer.weight_scale,
103
+ layer.weight_global_scale,
104
+ )
105
+ out = slice_nvfp4_output(out, output_size)
106
+ if bias is not None:
107
+ out = out + bias
108
+ return out.view(*output_shape)
109
+
110
+
111
class FlashInferCuteDslNvFp4LinearKernel(NvFp4LinearKernel):
112
"""NVFP4 GEMM via FlashInfer's cutedsl backend."""
113
vllm/model_executor/layers/quantization/modelopt.py
+2
-2
@@ -1038,7 +1038,7 @@ class ModelOptNvFp4Config(ModelOptQuantConfigBase):
1038
1039
# Select LinearMethod implementation based on quant_algo (FP8 pattern).
1040
# NVFP4 -> W4A4: cutlass NVFP4 GEMM with input quantization
1041
- # W4A16_NVFP4 -> W4A16: FP4 Marlin GEMM with bf16/fp16 activations
1041
+ # W4A16_NVFP4 -> W4A16: flashinfer cute-dsl or Marlin with 16-bit inputs
1042
if quant_method == "NVFP4":
1043
self.LinearMethodCls = ModelOptNvFp4LinearMethod
1044
elif quant_method == "W4A16_NVFP4":
@@ -1283,7 +1283,7 @@ class ModelOptNvFp4W4A16LinearMethod(LinearMethodBase):
1283
self.quant_config = quant_config
1284
self.marlin_input_dtype = None
1285
# `init_nvfp4_linear_kernel(use_a16=True)` is best of both worlds:
1286
- # 1. `use_a16=True` forces `Marlin`: https://github.com/vllm-project/vllm/commit/e68988a#diff-7135ab92aa94dfacb1ad3c77fc13f9c4ffe0b977f8eac5d86c2afe243e5f92a6R842-R889
1286
+ # 1. `use_a16=True` forces flashinfer cutedsl for sm100 and `Marlin` for others.
1287
# for `--linear-backend=auto`, avoiding a W4A4 kernel that requires input_scale.
1288
# 2. Specifying e.g. `--linear-backend=humming` will override.
1289
self.kernel = init_nvfp4_linear_kernel(use_a16=True)
vllm/utils/flashinfer.py
+49
@@ -197,6 +197,9 @@ flashinfer_cutlass_fused_moe = _lazy_import_wrapper(
197
flashinfer_cutedsl_grouped_gemm_nt_masked = _lazy_import_wrapper(
198
"flashinfer.cute_dsl.blockscaled_gemm", "grouped_gemm_nt_masked"
199
)
200
+flashinfer_prepare_bf16_fp4_weights = _lazy_import_wrapper(
201
+ "flashinfer.gemm", "prepare_bf16_fp4_weights"
202
+)
203
flashinfer_fp4_quantize = _lazy_import_wrapper("flashinfer", "fp4_quantize")
204
flashinfer_mxfp4_quantize = _lazy_import_wrapper("flashinfer", "mxfp4_quantize")
205
nvfp4_batched_quantize = _lazy_import_wrapper("flashinfer", "nvfp4_batched_quantize")
@@ -333,6 +336,17 @@ def has_flashinfer_cutedsl() -> bool:
336
)
337
338
339
+@functools.cache
340
+def has_flashinfer_bf16_fp4() -> bool:
341
+ """Return ``True`` if FlashInfer's CuTe-DSL W4A16 GEMM is available."""
342
+ if not has_flashinfer_cutedsl():
343
+ return False
344
+ mod = _get_submodule("flashinfer.gemm")
345
+ return mod is not None and all(
346
+ hasattr(mod, name) for name in ("mm_bf16_fp4", "prepare_bf16_fp4_weights")
347
+ )
348
+
349
+
350
@functools.cache
351
def has_flashinfer_trtllm_fused_moe() -> bool:
352
"""Return `True` if FlashInfer TRTLLM fused MoE is available."""
@@ -706,6 +720,39 @@ if has_flashinfer():
720
) -> torch.Tensor:
721
return torch.empty(A.shape[0], B.shape[1], dtype=dtype, device=A.device)
722
723
+ @torch.library.custom_op(
724
+ "vllm::flashinfer_mm_bf16_fp4",
725
+ mutates_args=[],
726
+ device_types="cuda",
727
+ )
728
+ def flashinfer_mm_bf16_fp4(
729
+ A: torch.Tensor,
730
+ B: torch.Tensor,
731
+ B_scale: torch.Tensor,
732
+ global_scale: torch.Tensor,
733
+ ) -> torch.Tensor:
734
+ from flashinfer.gemm import mm_bf16_fp4
735
+
736
+ return mm_bf16_fp4(
737
+ A,
738
+ B,
739
+ B_scale,
740
+ global_scale,
741
+ backend="cute-dsl",
742
+ )
743
+
744
+ @torch.library.register_fake(
745
+ "vllm::flashinfer_mm_bf16_fp4",
746
+ )
747
+ def flashinfer_mm_bf16_fp4_fake(
748
+ A: torch.Tensor,
749
+ B: torch.Tensor,
750
+ B_scale: torch.Tensor,
751
+ global_scale: torch.Tensor,
752
+ ) -> torch.Tensor:
753
+ output_size = B.shape[0] if B.dtype == torch.uint8 else B.shape[1] // 2
754
+ return torch.empty(A.shape[0], output_size, dtype=A.dtype, device=A.device)
755
+
756
@torch.library.custom_op(
757
"vllm::flashinfer_mxfp4_quantize",
758
mutates_args=[],
@@ -1158,6 +1205,7 @@ __all__ = [
1205
"flashinfer_trtllm_fp8_block_scale_moe",
1206
"flashinfer_cutlass_fused_moe",
1207
"flashinfer_cutedsl_grouped_gemm_nt_masked",
1208
+ "flashinfer_prepare_bf16_fp4_weights",
1209
"flashinfer_fp4_quantize",
1210
"silu_and_mul_scaled_nvfp4_experts_quantize",
1211
"scaled_fp4_grouped_quantize",
@@ -1177,6 +1225,7 @@ __all__ = [
1225
"has_flashinfer_cutlass_fused_moe",
1226
"has_flashinfer_cutedsl_grouped_gemm_nt_masked",
1227
"has_flashinfer_cutedsl_moe_nvfp4",
1228
+ "has_flashinfer_bf16_fp4",
1229
"has_flashinfer_b12x_moe",
1230
"has_flashinfer_b12x_gemm",
1231
"has_flashinfer_fp8_blockscale_gemm",