[Kernel] Add NVFP4 support to the torch linear backend (#53319)
Signed-off-by: Michael Lazos <mlazos@meta.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Michael Lazos committed
Sep 8, 2026 at 17:30 UTC
28a73ccfba9eb42df3bdf9b1f11c6ca5fac528c4
3 files changed
+102
tests/models/quantization/test_nvfp4.py
+2
@@ -94,6 +94,7 @@ SM_100_NVFP4_BACKENDS = [
94
"flashinfer_cudnn",
95
"flashinfer_trtllm",
96
"flashinfer_cutlass",
97
+ "torch",
98
]
99
100
@@ -107,6 +108,7 @@ SM_100_NVFP4_BACKENDS = [
108
"flashinfer_cudnn",
109
"flashinfer_trtllm", # the small seq_len ensures trtllm_8x4_layout backend is used
110
"flashinfer_cutlass",
111
+ "torch",
112
],
113
)
114
def test_nvfp4(vllm_runner, model, eager, backend):
vllm/model_executor/kernels/linear/__init__.py
+6
@@ -159,6 +159,9 @@ from vllm.model_executor.kernels.linear.nvfp4.humming import (
159
from vllm.model_executor.kernels.linear.nvfp4.marlin import (
160
MarlinNvFp4LinearKernel,
161
)
162
+from vllm.model_executor.kernels.linear.nvfp4.pytorch import (
163
+ TorchNvFp4LinearKernel,
164
+)
165
from vllm.model_executor.kernels.linear.scaled_mm import (
166
Fp8BlockScaledMMLinearKernel,
167
FP8ScaledMMLinearKernel,
@@ -304,6 +307,7 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = {
307
ChannelWiseTorchFP8ScaledMMLinearKernel,
308
RowWiseTorchFP8ScaledMMLinearKernel,
309
BlockWiseTorchFP8ScaledMMLinearKernel,
310
+ TorchNvFp4LinearKernel,
311
},
312
"aiter": {
313
AiterInt8ScaledMMLinearKernel,
@@ -543,6 +547,7 @@ _POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = {
547
FlashInferCudnnNvFp4LinearKernel,
548
FbgemmNvFp4LinearKernel,
549
B12xNvFp4LinearKernel,
550
+ TorchNvFp4LinearKernel,
551
EmulationNvFp4LinearKernel,
552
HummingNvFp4LinearKernel,
553
],
@@ -1255,6 +1260,7 @@ __all__ = [
1260
"FlashInferTrtllmNvFp4LinearKernel",
1261
"FlashInferCudnnNvFp4LinearKernel",
1262
"MarlinNvFp4LinearKernel",
1263
+ "TorchNvFp4LinearKernel",
1264
"_KernelT",
1265
"DeepGemmFp8BlockScaledMMKernel",
1266
"FlashInferFp8DeepGEMMDynamicBlockScaledKernel",
vllm/model_executor/kernels/linear/nvfp4/pytorch.py
new
+94
@@ -0,0 +1,94 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+from vllm._custom_ops import scaled_fp4_quant
7
+from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
8
+ pad_nvfp4_weight_for_cutlass,
9
+ slice_nvfp4_output,
10
+ swizzle_blockscale,
11
+)
12
+from vllm.platforms import current_platform
13
+
14
+from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
15
+
16
+
17
+class TorchNvFp4LinearKernel(NvFp4LinearKernel):
18
+ """NVFP4 GEMM implemented with PyTorch's native ``torch._scaled_mm``.
19
+
20
+ Eager execution uses PyTorch's native scaled-matmul implementation. Under
21
+ ``torch.compile``, TorchInductor may select any enabled scaled-GEMM backend,
22
+ including NVGEMM on supported Blackwell systems.
23
+ """
24
+
25
+ @classmethod
26
+ def is_supported(
27
+ cls, compute_capability: int | None = None
28
+ ) -> tuple[bool, str | None]:
29
+ if not current_platform.is_device_capability_family(100):
30
+ return False, "Torch NVFP4 requires sm_10x (Blackwell)"
31
+ if not hasattr(torch, "_scaled_mm"):
32
+ return False, "torch._scaled_mm not available"
33
+ if not hasattr(torch, "float4_e2m1fn_x2"):
34
+ return False, "torch.float4_e2m1fn_x2 not available"
35
+ return True, None
36
+
37
+ @classmethod
38
+ def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
39
+ return True, None
40
+
41
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
42
+ layer.weight_scale = torch.nn.Parameter(
43
+ swizzle_blockscale(layer.weight_scale.data), requires_grad=False
44
+ )
45
+ padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(
46
+ layer.weight.data
47
+ )
48
+ layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False)
49
+ layer.weights_padding_cols = weights_padding_cols
50
+
51
+ def apply_weights(
52
+ self,
53
+ layer: torch.nn.Module,
54
+ x: torch.Tensor,
55
+ bias: torch.Tensor | None = None,
56
+ ) -> torch.Tensor:
57
+ output_size = layer.output_size_per_partition
58
+ output_dtype = x.dtype
59
+ output_shape = [*x.shape[:-1], output_size]
60
+ weights_padding_bytes = getattr(layer, "weights_padding_cols", 0)
61
+
62
+ x_fp4, x_blockscale = scaled_fp4_quant(
63
+ x,
64
+ layer.input_global_scale_inv,
65
+ is_sf_swizzled_layout=True,
66
+ backend="cutlass",
67
+ padded_n=x.shape[-1] + weights_padding_bytes * 2,
68
+ )
69
+
70
+ # scaled_fp4_quant returns packed bytes, while torch._scaled_mm uses
71
+ # the corresponding shell dtype to describe the two FP4 values stored
72
+ # in each byte.
73
+ x_fp4 = x_fp4.view(torch.float4_e2m1fn_x2)
74
+ weight = layer.weight
75
+ if weight.dtype != torch.float4_e2m1fn_x2:
76
+ weight = weight.view(torch.float4_e2m1fn_x2)
77
+
78
+ out = torch._scaled_mm(
79
+ x_fp4,
80
+ weight.t(),
81
+ scale_a=x_blockscale.reshape(-1),
82
+ scale_b=layer.weight_scale.reshape(-1),
83
+ out_dtype=output_dtype,
84
+ )
85
+
86
+ # ModelOpt stores the product of the activation and weight global
87
+ # dequantization scales as alpha. TorchInductor can fold this multiply
88
+ # into an NVGEMM output-scale argument.
89
+ out = out * layer.alpha
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)