| 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) |