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,110 @@ 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])
+@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
+    assert sqnr >= SQNR_THRESHOLD, (
+        f"Got a sqnr of {sqnr} for NVFP4 recipe with 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("use_gelu", [True, False])
+@pytest.mark.parametrize("emulate", [True, False])
+@pytest.mark.parametrize("compile", [False])
+@pytest.mark.parametrize("bias", [True, False])
+@torch.no_grad()
+@skip_if_rocm("ROCm float4 gemm require gfx950")
+def test_nvfp4_matmul_with_amax(
+    use_gelu: bool, emulate: bool, compile: bool, bias: bool
+):
+    from torchao.prototype.mx_formats.nvfp4_tensor import (
+        NVFP4Tensor,
+        per_tensor_amax_to_scale,
+    )
+
+    m, k, n = 64, 256, 128
+
+    # Create activation tensor
+    if use_gelu:
+        x = torch.randn(m, k, dtype=torch.bfloat16, device="cuda")
+        A = torch.nn.functional.gelu(x)
+    else:
+        A = torch.randn(m, k, dtype=torch.bfloat16, device="cuda")
+
+    B = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
+    bias_tensor = torch.randn(n, dtype=torch.bfloat16, device="cuda") if bias else None
+
+    # Compute reference
+    C_ref = torch.matmul(A, B.t())
+    if bias:
+        C_ref = C_ref + bias_tensor
+
+    a_scale = per_tensor_amax_to_scale(torch.amax(torch.abs(A)))
+    b_scale = per_tensor_amax_to_scale(torch.amax(torch.abs(B)))
+    A_nvfp4 = NVFP4Tensor.to_nvfp4(A, per_tensor_scale=a_scale)
+    B_nvfp4 = NVFP4Tensor.to_nvfp4(B, per_tensor_scale=b_scale)
+
+    if emulate:
+        # Cast back to original dtype and compute
+        A_emulated = A_nvfp4.to_dtype(A.dtype)
+        B_emulated = B_nvfp4.to_dtype(B.dtype)
+        mm = torch.compile(torch.matmul, fullgraph=True) if compile else torch.matmul
+        C_emulated = mm(A_emulated, B_emulated.t())
+        if bias:
+            C_emulated = C_emulated + bias_tensor
+        sqnr = compute_error(C_ref, C_emulated)
+    else:
+        if bias:
+            linear_fn = (
+                torch.compile(torch.nn.functional.linear, fullgraph=True)
+                if compile
+                else torch.nn.functional.linear
+            )
+            C_nvfp4 = linear_fn(A_nvfp4, B_nvfp4, bias_tensor)
+        else:
+            mm = (
+                torch.compile(torch.matmul, fullgraph=True) if compile else torch.matmul
+            )
+            C_nvfp4 = mm(A_nvfp4, B_nvfp4.t())
+        sqnr = compute_error(C_ref, C_nvfp4)
+
+    # Check quality threshold
+    SQNR_THRESHOLD = 16.0
+    assert sqnr >= SQNR_THRESHOLD, (
+        f"SQNR {sqnr:.2f} < {SQNR_THRESHOLD}, use_gelu={use_gelu}, emulate={emulate}, compile={compile}, bias={bias}"
+    )

test/prototype/mx_formats/test_mx_tensor.py

@@ -14,6 +14,7 @@
 from torchao.prototype.mx_formats.constants import (
     DTYPE_FP6_E2M3,
     DTYPE_FP6_E3M2,
+    F4_E2M1_MAX,
     SUPPORTED_ELEM_DTYPES,
 )
 from torchao.prototype.mx_formats.kernels import pack_uint4, pack_uint6
@@ -591,3 +592,57 @@ def to_f8(x):
     torch.testing.assert_close(
         data_in_range_f8_c, data_out_of_range_f8_c, atol=0, rtol=0
     )
+
+
+@pytest.mark.parametrize(
+    "dtype,shape,use_per_tensor_scale",
+    [
+        (torch.bfloat16, (32, 64), False),
+        (torch.float32, (64, 128), False),
+        (torch.bfloat16, (128, 256), False),
+        (torch.bfloat16, (64, 128), True),
+    ],
+)
+@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+"
+)
+def test_nvfp4_reconstruction(dtype, shape, use_per_tensor_scale):
+    from torchao.prototype.mx_formats.nvfp4_tensor import (
+        NVFP4Tensor,
+        per_tensor_amax_to_scale,
+    )
+
+    x = torch.randn(shape, dtype=dtype, device="cuda")
+    if use_per_tensor_scale:
+        tensor_amax = torch.max(torch.abs(x))
+        scale = per_tensor_amax_to_scale(tensor_amax)
+    else:
+        scale = None
+
+    x_nvfp4 = NVFP4Tensor.to_nvfp4(x, per_tensor_scale=scale)
+    x_reconstructed = x_nvfp4.to_dtype(dtype)
+
+    def assert_sqnr_gt_threshold(orig, new, threshold):
+        sqnr = compute_error(orig, new)
+        if torch.all(torch.isnan(sqnr)):
+            # if both operands are full of zeroes, sqnr is nan and this is ok
+            # test for this explicitly
+            assert torch.all(orig == 0) and torch.all(new == 0)
+        else:
+            assert sqnr >= threshold
+
+    reconstructed_amax = x_nvfp4.get_scales().view(shape[0], -1, 1) * F4_E2M1_MAX
+    max_abs = torch.amax(
+        torch.abs(x.reshape(shape[0], -1, x_nvfp4._block_size)), dim=-1
+    ).unsqueeze(-1)
+
+    assert_sqnr_gt_threshold(max_abs, reconstructed_amax, 30.0)
+    assert_sqnr_gt_threshold(x, x_reconstructed, 8.0)
+
+    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}"
+    )

torchao/prototype/mx_formats/__init__.py

@@ -6,7 +6,10 @@
 )
 
 # Note: Prototype and subject to change
-from torchao.prototype.mx_formats.mx_subclass import MXFPInferenceConfig
+from torchao.prototype.mx_formats.mx_subclass import (
+    MXFPInferenceConfig,
+    NVFP4InferenceConfig,
+)
 
 # import mx_linear here to register the quantize_ transform logic
 # ruff: noqa: I001
@@ -18,4 +21,5 @@
     "MXLinearConfig",
     "MXLinearRecipeName",
     "MXFPInferenceConfig",
+    "NVFP4InferenceConfig",
 ]

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_subclass.py

@@ -5,12 +5,11 @@
 # LICENSE file in the root directory of this source tree.
 
 import types
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Optional
 
 import torch
 
-import torchao
 from torchao.core.config import AOBaseConfig
 from torchao.prototype.mx_formats import (
     MXGemmKernelChoice,
@@ -20,11 +19,16 @@
     _validate_gemm_kernel_choice,
 )
 from torchao.prototype.mx_formats.mx_tensor import MXTensor
+from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4MMConfig, NVFP4Tensor
 from torchao.quantization.quant_api import to_linear_activation_quantized
 from torchao.quantization.transform_module import (
     register_quantize_module_handler,
 )
-from torchao.utils import TORCH_VERSION_AT_LEAST_2_5, is_sm_at_least_100
+from torchao.utils import (
+    TORCH_VERSION_AT_LEAST_2_5,
+    TORCH_VERSION_AT_LEAST_2_8,
+    is_sm_at_least_100,
+)
 
 
 # Note: This API is extra prototype and will change in the future
@@ -63,16 +67,13 @@ 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
 
     # Which kernel to run for mm
     gemm_kernel_choice: MXGemmKernelChoice = MXGemmKernelChoice.CUBLAS
 
-    # Set some magic perf settings
-    set_inductor_config: bool = False
-
     def __post_init__(self):
         assert self.activation_dtype == self.weight_dtype, (
             "For now - we only support matching input/weight dtypes."
@@ -115,8 +116,6 @@ def _mx_inference_linear_transform(
     # TODO Sm120 has slightly more restrictive reqs
     # TODO handle AMD
     assert is_sm_at_least_100(), "MXFP is only supported on sm100 machiens for now"
-    if config.set_inductor_config:
-        torchao.quantization.utils.recommended_inductor_config_setter()
 
     activation_dtype = config.activation_dtype
     weight_dtype = config.weight_dtype
@@ -151,7 +150,90 @@ def _mx_inference_linear_transform(
     return module
 
 
+def _get_nvfp4_dtype():
+    """Factory function for NVFP4 dtype defaults."""
+    if not TORCH_VERSION_AT_LEAST_2_8:
+        raise RuntimeError("NVFP4InferenceConfig requires PyTorch 2.8 or later")
+    return torch.float4_e2m1fn_x2
+
+
+@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 = field(default_factory=_get_nvfp4_dtype)
+    weight_dtype: torch.dtype = field(default_factory=_get_nvfp4_dtype)
+
+    # 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
+
+    # Matrix multiplication configuration
+    mm_config: NVFP4MMConfig = NVFP4MMConfig.DYNAMIC
+
+    def __post_init__(self):
+        # Validate NVFP4 constraints
+        if not TORCH_VERSION_AT_LEAST_2_8:
+            raise RuntimeError("NVFP4InferenceConfig requires PyTorch 2.8 or later")
+
+        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}"
+        )
+
+
+@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"
+
+    weight = module.weight
+    assert weight.dtype == torch.bfloat16, (
+        f"Only supporting bf16 out dtype for now, got {weight.dtype}"
+    )
+
+    quantized_weight = NVFP4Tensor.to_nvfp4(
+        weight,
+        block_size=config.block_size,
+        mm_config=config.mm_config,
+    )
+
+    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,
+            NVFP4MMConfig,
+            MXGemmKernelChoice,
+            _input_activation_quant_func_mxfp,
+        ]
     )