From 28a73ccfba9eb42df3bdf9b1f11c6ca5fac528c4 Mon Sep 17 00:00:00 2001 From: Michael Lazos Date: Tue, 8 Sep 2026 17:30:05 -0700 Subject: [PATCH] [Kernel] Add NVFP4 support to the torch linear backend (#53319) Signed-off-by: Michael Lazos Co-authored-by: OpenAI Codex --- tests/models/quantization/test_nvfp4.py | 2 + .../model_executor/kernels/linear/__init__.py | 6 ++ .../kernels/linear/nvfp4/pytorch.py | 94 +++++++++++++++++++ 3 files changed, 102 insertions(+) create mode 100644 vllm/model_executor/kernels/linear/nvfp4/pytorch.py diff --git a/tests/models/quantization/test_nvfp4.py b/tests/models/quantization/test_nvfp4.py index f5022a461b..00f211b271 100644 --- a/tests/models/quantization/test_nvfp4.py +++ b/tests/models/quantization/test_nvfp4.py @@ -94,6 +94,7 @@ SM_100_NVFP4_BACKENDS = [ "flashinfer_cudnn", "flashinfer_trtllm", "flashinfer_cutlass", + "torch", ] @@ -107,6 +108,7 @@ SM_100_NVFP4_BACKENDS = [ "flashinfer_cudnn", "flashinfer_trtllm", # the small seq_len ensures trtllm_8x4_layout backend is used "flashinfer_cutlass", + "torch", ], ) def test_nvfp4(vllm_runner, model, eager, backend): diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 1c03dcd468..f8e4f5ecb4 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -159,6 +159,9 @@ from vllm.model_executor.kernels.linear.nvfp4.humming import ( from vllm.model_executor.kernels.linear.nvfp4.marlin import ( MarlinNvFp4LinearKernel, ) +from vllm.model_executor.kernels.linear.nvfp4.pytorch import ( + TorchNvFp4LinearKernel, +) from vllm.model_executor.kernels.linear.scaled_mm import ( Fp8BlockScaledMMLinearKernel, FP8ScaledMMLinearKernel, @@ -304,6 +307,7 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = { ChannelWiseTorchFP8ScaledMMLinearKernel, RowWiseTorchFP8ScaledMMLinearKernel, BlockWiseTorchFP8ScaledMMLinearKernel, + TorchNvFp4LinearKernel, }, "aiter": { AiterInt8ScaledMMLinearKernel, @@ -543,6 +547,7 @@ _POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = { FlashInferCudnnNvFp4LinearKernel, FbgemmNvFp4LinearKernel, B12xNvFp4LinearKernel, + TorchNvFp4LinearKernel, EmulationNvFp4LinearKernel, HummingNvFp4LinearKernel, ], @@ -1255,6 +1260,7 @@ __all__ = [ "FlashInferTrtllmNvFp4LinearKernel", "FlashInferCudnnNvFp4LinearKernel", "MarlinNvFp4LinearKernel", + "TorchNvFp4LinearKernel", "_KernelT", "DeepGemmFp8BlockScaledMMKernel", "FlashInferFp8DeepGEMMDynamicBlockScaledKernel", diff --git a/vllm/model_executor/kernels/linear/nvfp4/pytorch.py b/vllm/model_executor/kernels/linear/nvfp4/pytorch.py new file mode 100644 index 0000000000..571f96d2d3 --- /dev/null +++ b/vllm/model_executor/kernels/linear/nvfp4/pytorch.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm._custom_ops import scaled_fp4_quant +from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( + pad_nvfp4_weight_for_cutlass, + slice_nvfp4_output, + swizzle_blockscale, +) +from vllm.platforms import current_platform + +from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig + + +class TorchNvFp4LinearKernel(NvFp4LinearKernel): + """NVFP4 GEMM implemented with PyTorch's native ``torch._scaled_mm``. + + Eager execution uses PyTorch's native scaled-matmul implementation. Under + ``torch.compile``, TorchInductor may select any enabled scaled-GEMM backend, + including NVGEMM on supported Blackwell systems. + """ + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_device_capability_family(100): + return False, "Torch NVFP4 requires sm_10x (Blackwell)" + if not hasattr(torch, "_scaled_mm"): + return False, "torch._scaled_mm not available" + if not hasattr(torch, "float4_e2m1fn_x2"): + return False, "torch.float4_e2m1fn_x2 not available" + return True, None + + @classmethod + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.weight_scale = torch.nn.Parameter( + swizzle_blockscale(layer.weight_scale.data), requires_grad=False + ) + padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass( + layer.weight.data + ) + layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False) + layer.weights_padding_cols = weights_padding_cols + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + output_size = layer.output_size_per_partition + output_dtype = x.dtype + output_shape = [*x.shape[:-1], output_size] + weights_padding_bytes = getattr(layer, "weights_padding_cols", 0) + + x_fp4, x_blockscale = scaled_fp4_quant( + x, + layer.input_global_scale_inv, + is_sf_swizzled_layout=True, + backend="cutlass", + padded_n=x.shape[-1] + weights_padding_bytes * 2, + ) + + # scaled_fp4_quant returns packed bytes, while torch._scaled_mm uses + # the corresponding shell dtype to describe the two FP4 values stored + # in each byte. + x_fp4 = x_fp4.view(torch.float4_e2m1fn_x2) + weight = layer.weight + if weight.dtype != torch.float4_e2m1fn_x2: + weight = weight.view(torch.float4_e2m1fn_x2) + + out = torch._scaled_mm( + x_fp4, + weight.t(), + scale_a=x_blockscale.reshape(-1), + scale_b=layer.weight_scale.reshape(-1), + out_dtype=output_dtype, + ) + + # ModelOpt stores the product of the activation and weight global + # dequantization scales as alpha. TorchInductor can fold this multiply + # into an NVGEMM output-scale argument. + out = out * layer.alpha + out = slice_nvfp4_output(out, output_size) + + if bias is not None: + out = out + bias + return out.view(*output_shape)