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,92 @@ 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.parametrize("use_gelu", [True, False])
+@pytest.mark.parametrize("emulate", [True, False])
+@pytest.mark.parametrize("compile", [False])
+@torch.no_grad()
+def test_nvfp4_matmul_with_amax_emulate(use_gelu: bool, emulate: bool, compile: 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")
+
+    # Create weight tensor (always Gaussian)
+    B = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
+
+    # Compute reference
+    C_ref = torch.matmul(A, B.t())
+
+    a_scale = per_tensor_amax_to_scale(torch.amax(torch.abs(A)))
+    b_scale = per_tensor_amax_to_scale(torch.amax(torch.abs(B)))
+
+    # Quantize with per-tensor amax
+    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())
+        sqnr = compute_error(C_ref, C_emulated)
+    else:
+        # Use native nvfp4 matmul
+        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}"
+    )

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,53 @@ 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),
+    ],
+)
+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,7 +5,7 @@
 # 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
@@ -20,11 +20,16 @@
     _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,
 )
-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,7 +68,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 +156,113 @@ 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
+
+    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}"
+        )
+
+
+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,
+        ]
     )