WIP NVfp4

stack-info: PR: #2408, branch: drisspg/stack/78

Author: drisspg

test/prototype/mx_formats/test_mx_linear.py

@@ -25,7 +25,10 @@
     MXInferenceLinear,
     MXLinear,
 )
-from torchao.prototype.mx_formats.mx_subclass import MXFPInferenceConfig
+from torchao.prototype.mx_formats.mx_subclass import (
+    MXFPInferenceConfig,
+    NVFP4InferenceConfig,
+)
 from torchao.quantization import quantize_
 from torchao.quantization.utils import compute_error
 from torchao.testing.utils import skip_if_rocm
@@ -441,3 +444,37 @@ def test_inference_subclass(elem_dtype, bias: bool, compile: bool):
     assert sqnr >= SQNR_THRESHOLD, (
         f"Got a sqnr of {sqnr} for {elem_dtype} and bias={bias}"
     )
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+@pytest.mark.skipif(
+    not TORCH_VERSION_AT_LEAST_2_8, reason="torch.compile requires PyTorch 2.8+"
+)
+@pytest.mark.skipif(
+    not is_sm_at_least_100(), reason="CUDA capability >= 10.0 required for float4 gemm"
+)
+@pytest.mark.parametrize("bias", [True, False])
+# @pytest.mark.parametrize("compile", [True, False])
+@pytest.mark.parametrize("compile", [False])
+@torch.no_grad()
+@skip_if_rocm("ROCm float4 gemm require gfx950")
+def test_inference_subclass_nvfp4(bias: bool, compile: bool):
+    """
+    Test NVFP4 recipe with scale_dtype=float8_e4m3fn and block_size=16
+    """
+    m = nn.Linear(64, 256, bias=bias, dtype=torch.bfloat16, device="cuda")
+    m_mx = copy.deepcopy(m)
+
+    config = NVFP4InferenceConfig()
+    quantize_(m_mx, config=config)
+    if compile:
+        m_mx = torch.compile(m_mx, fullgraph=True)
+
+    x = torch.randn(128, 64, device="cuda", dtype=torch.bfloat16)
+    y_ref = m(x)
+    y_mx = m_mx(x)
+    sqnr = compute_error(y_ref, y_mx)
+    SQNR_THRESHOLD = 15.0  # Float4 threshold
+    assert sqnr >= SQNR_THRESHOLD, (
+        f"Got a sqnr of {sqnr} for NVFP4 recipe with bias={bias}"
+    )

test/prototype/mx_formats/test_nvfp4_reconstruction.py

@@ -0,0 +1,59 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+
+#!/usr/bin/env python3
+"""Test NVFP4 quantization and dequantization reconstruction loss."""
+
+import torch
+
+from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor
+
+
+def test_nvfp4_reconstruction():
+    # Test with different shapes and dtypes
+    test_configs = [
+        (torch.bfloat16, (32, 64)),
+        (torch.float32, (64, 128)),
+        (torch.bfloat16, (128, 256)),
+    ]
+
+    for dtype, shape in test_configs:
+        print(f"\nTesting shape={shape}, dtype={dtype}")
+
+        # Create a random tensor
+        x = torch.randn(shape, dtype=dtype)
+
+        # Quantize to NVFP4
+        x_nvfp4 = NVFP4Tensor.to_nvfp4(x)
+        print(f"Original tensor: shape={x.shape}, dtype={x.dtype}")
+        print(f"NVFP4 tensor: {x_nvfp4}")
+
+        # Dequantize back to original dtype
+        x_reconstructed = x_nvfp4.to_dtype(dtype)
+        print(
+            f"Reconstructed tensor: shape={x_reconstructed.shape}, dtype={x_reconstructed.dtype}"
+        )
+
+        # Calculate reconstruction error
+        abs_error = torch.abs(x - x_reconstructed)
+        rel_error = abs_error / (torch.abs(x) + 1e-8)
+
+        print(f"Max absolute error: {abs_error.max().item():.6f}")
+        print(f"Mean absolute error: {abs_error.mean().item():.6f}")
+        print(f"Max relative error: {rel_error.max().item():.6f}")
+        print(f"Mean relative error: {rel_error.mean().item():.6f}")
+
+        # Verify shapes match
+        assert x.shape == x_reconstructed.shape, (
+            f"Shape mismatch: {x.shape} vs {x_reconstructed.shape}"
+        )
+        assert x.dtype == x_reconstructed.dtype, (
+            f"Dtype mismatch: {x.dtype} vs {x_reconstructed.dtype}"
+        )
+
+
+if __name__ == "__main__":
+    test_nvfp4_reconstruction()

torchao/prototype/mx_formats/config.py

@@ -57,10 +57,10 @@ def _validate_gemm_kernel_choice(gemm_kernel_choice, block_size, elem_dtype):
             f"elem_dtype must be one of {valid_dtypes} to use the CUTLASS MX gemm kernels, got {elem_dtype}"
         )
     elif gemm_kernel_choice == MXGemmKernelChoice.CUBLAS:
-        assert block_size == 32, (
-            f"block_size must be 32 to use the cuBLAS MX gemm kernels, got {block_size}"
+        assert block_size in [16, 32], (
+            f"block_size must be in [16, 32] to use the cuBLAS MX gemm kernels, got {block_size}"
         )
-        valid_dtypes = [torch.float8_e4m3fn]
+        valid_dtypes = [torch.float8_e4m3fn, torch.float4_e2m1fn_x2]
         assert elem_dtype in valid_dtypes, (
             f"elem_dtype must be one of {valid_dtypes} to use the CUTLASS MX gemm kernels, got {elem_dtype}"
         )

torchao/prototype/mx_formats/mx_linear.py

@@ -104,6 +104,7 @@ def backward(ctx, grad_output_hp: torch.Tensor):
                 w_elem_dtype,
                 block_size,
                 weight_hp.dtype,
+                None,  # scale_dtype
                 False,
                 gemm_kernel_choice,
                 False,
@@ -133,6 +134,7 @@ def backward(ctx, grad_output_hp: torch.Tensor):
                 grad_elem_dtype,
                 block_size,
                 grad_output_hp_r.dtype,
+                None,  # scale_dtype
                 False,
                 gemm_kernel_choice,
                 False,
@@ -155,6 +157,7 @@ def backward(ctx, grad_output_hp: torch.Tensor):
                 in_elem_dtype,
                 block_size,
                 input_hp_r.dtype,
+                None,  # scale_dtype
                 False,
                 gemm_kernel_choice,
                 False,

torchao/prototype/mx_formats/mx_subclass.py

@@ -20,6 +20,7 @@
     _validate_gemm_kernel_choice,
 )
 from torchao.prototype.mx_formats.mx_tensor import MXTensor
+from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor
 from torchao.quantization.quant_api import to_linear_activation_quantized
 from torchao.quantization.transform_module import (
     register_quantize_module_handler,
@@ -63,7 +64,7 @@ class MXFPInferenceConfig(AOBaseConfig):
 
     block_size: int = 32
 
-    # Dtypes for Input and Weights
+    # Dtypes for Input and Weights, supports Fp8 and Fp4 formats
     activation_dtype: torch.dtype = torch.float8_e4m3fn
     weight_dtype: torch.dtype = torch.float8_e4m3fn
 
@@ -151,7 +152,106 @@ def _mx_inference_linear_transform(
     return module
 
 
+@dataclass
+class NVFP4InferenceConfig(AOBaseConfig):
+    """
+    NVIDIA FP4 (NVFP4) Inference Quantization Configuration
+
+    This is a specialized configuration for NVIDIA's FP4 format with UE4M3 scales.
+    It provides defaults optimized for NVFP4:
+    - Data: float4_e2m1fn_x2
+    - Scales: float8_e4m3fn (UE4M3)
+    - Block size: 16 (required for NVFP4)
+    - CUBLAS kernel (optimized for VEC16_UE4M3)
+    """
+
+    block_size: int = 16  # NVFP4 requires block size 16
+
+    # NVFP4 uses FP4 data
+    activation_dtype: torch.dtype = torch.float4_e2m1fn_x2
+    weight_dtype: torch.dtype = torch.float4_e2m1fn_x2
+
+    # NVFP4 uses E4M3 scales
+    scale_dtype: torch.dtype = torch.float8_e4m3fn
+
+    # CUBLAS is preferred for NVFP4 with VEC16_UE4M3 support
+    gemm_kernel_choice: MXGemmKernelChoice = MXGemmKernelChoice.CUBLAS
+
+    # Set some magic perf settings
+    set_inductor_config: bool = False
+
+    def __post_init__(self):
+        # Validate NVFP4 constraints
+        assert self.activation_dtype == torch.float4_e2m1fn_x2, (
+            f"NVFP4 requires activation_dtype=float4_e2m1fn_x2, got {self.activation_dtype}"
+        )
+        assert self.weight_dtype == torch.float4_e2m1fn_x2, (
+            f"NVFP4 requires weight_dtype=float4_e2m1fn_x2, got {self.weight_dtype}"
+        )
+        assert self.scale_dtype == torch.float8_e4m3fn, (
+            f"NVFP4 requires scale_dtype=float8_e4m3fn, got {self.scale_dtype}"
+        )
+        assert self.block_size == 16, (
+            f"NVFP4 requires block_size=16, got {self.block_size}"
+        )
+
+
+def _input_activation_quant_func_nvfp4(
+    x: torch.Tensor,
+    block_size: int = 16,
+    scale: Optional[torch.Tensor] = None,
+):
+    """NVFP4-specific activation quantization function"""
+    # TODO: scale for static quant
+    activation = NVFP4Tensor.to_nvfp4(
+        x,
+        block_size=block_size,
+    )
+    return activation
+
+
+@register_quantize_module_handler(NVFP4InferenceConfig)
+def _nvfp4_inference_linear_transform(
+    module: torch.nn.Module, config: NVFP4InferenceConfig
+):
+    """Quantization handler for NVFP4InferenceConfig"""
+    assert is_sm_at_least_100(), "NVFP4 is only supported on sm100+ machines"
+    if config.set_inductor_config:
+        torchao.quantization.utils.recommended_inductor_config_setter()
+
+    weight = module.weight
+    assert weight.dtype == torch.bfloat16, (
+        f"Only supporting bf16 out dtype for now, got {weight.dtype}"
+    )
+
+    # Convert weight to NVFP4 Tensor
+    quantized_weight = NVFP4Tensor.to_nvfp4(
+        weight,
+        block_size=config.block_size,
+    )
+
+    input_quant_func = _input_activation_quant_func_nvfp4
+    input_quant_kwargs = {
+        "block_size": config.block_size,
+        "scale": None,
+    }
+
+    quantized_weight = to_linear_activation_quantized(
+        quantized_weight, input_quant_func, quant_kwargs=input_quant_kwargs
+    )
+
+    module.weight = torch.nn.Parameter(quantized_weight, requires_grad=False)
+    module.extra_repr = types.MethodType(_linear_extra_repr, module)
+    return module
+
+
 if TORCH_VERSION_AT_LEAST_2_5:
     torch.serialization.add_safe_globals(
-        [MXTensor, MXGemmKernelChoice, _input_activation_quant_func_mxfp]
+        [
+            MXTensor,
+            NVFP4Tensor,
+            MXGemmKernelChoice,
+            _input_activation_quant_func_mxfp,
+            _input_activation_quant_func_nvfp4,
+        ]
     )