nvidia-modelopt 0.35.0__py3-none-any.whl

modelopt/torch/quantization/qtensor/nf4_tensor.py

@@ -0,0 +1,199 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Implements NF4 quantization for efficient tensor storage and computation."""
+
+import numpy as np
+import torch
+
+from ..extensions import get_cuda_ext
+from ..qtensor.base_qtensor import BaseQuantizedTensor
+from ..utils import reduce_amax, reduce_block_padding
+
+nf4_table = torch.tensor(
+    [
+        -1.0000,
+        -0.6962,
+        -0.5251,
+        -0.3949,
+        -0.2844,
+        -0.1848,
+        -0.0911,
+        0.0000,
+        0.0796,
+        0.1609,
+        0.2461,
+        0.3379,
+        0.4407,
+        0.5626,
+        0.7230,
+        1.0000,
+    ],
+    dtype=torch.bfloat16,
+)
+
+__all__ = ["NF4QTensor"]
+
+
+def _dequantize_scalers(scales, double_scale, scale_zeros, dtype):
+    return (scales / double_scale.unsqueeze(-1)).to(dtype) + scale_zeros
+
+
+def _quantize_to_nearest_lut(flatten_tensor: torch.Tensor, lut: torch.Tensor) -> torch.Tensor:
+    """Quantize a float16 tensor to nearest value and return the indices of the look-up table."""
+    assert flatten_tensor.dim() == 1, (
+        f"Expect flatten tensor but got input with {flatten_tensor.dim()} dimensions."
+    )
+    diff = (flatten_tensor[:, None] - lut).abs()  # Shape: (numel, 16)
+    indices = diff.argmin(dim=-1)
+    return indices
+
+
+def _nf4_lookup(quantized_idx: torch.Tensor) -> torch.Tensor:
+    """Dequantize uint4 index to nf4 value."""
+    return nf4_table.to(quantized_idx.device)[quantized_idx]
+
+
+class NF4QTensor(BaseQuantizedTensor):
+    """Implements the NF4 quantization on tensors for more efficient storage or computation.
+
+    Attributes:
+        quantized_data (torch.Tensor): The quantized data stored as a packed uint8 tensor.
+    """
+
+    @classmethod
+    def quantize(
+        cls, input: torch.Tensor, block_size: int, scale_block_size: int | None = None
+    ) -> tuple:
+        """Converting a tensor to a quantized format based on NF4 double quantization.
+
+        Args:
+            input (torch.Tensor): The input tensor to be quantized.
+            block_size (int): The size of each block for quantization.
+            scale_block_size (int): The block size for scaling during quantization.
+
+        Returns:
+            tuple: Contains quantized data, input quantization config, and scale quantization config.
+        """
+        cuda_ext = get_cuda_ext()
+
+        # pad the input if needed
+        original_input = input
+        input = reduce_block_padding(input.view(-1), block_sizes={-1: block_size})
+
+        # get scales for each block
+        block_input = input.view(-1, block_size)
+        scales = reduce_amax(block_input, -1)
+
+        if cuda_ext and input.is_cuda:
+            packed_output_uint8 = cuda_ext.NF4_quantize(input, scales, block_size)
+        else:
+            # expand scalers to match shape of input
+            scales = scales.view(block_input.shape[0], -1)  # shape: (block_input.shape[0], 1)
+            scaled_blocks = block_input / scales
+
+            quantized_output = torch.empty(
+                block_input.numel(), dtype=torch.uint8, device=block_input.device
+            )
+            flattened = scaled_blocks.flatten()
+            quantized_output = _quantize_to_nearest_lut(
+                flattened, nf4_table.to(device=input.device, dtype=input.dtype)
+            )
+            quantized_output_uint8 = quantized_output.to(torch.uint8)
+
+            # pack the int4 weights into a uint8 tensor
+            # packing format: w0, w1, w2, w3, w4, w5, ...
+            #               | byte  | byte  | byte  |
+            packed_output_uint8 = quantized_output_uint8[::2] << 4 | quantized_output_uint8[1::2]
+
+        # pad the scales if needed
+        scales = reduce_block_padding(scales.view(-1), block_sizes={-1: scale_block_size})
+        return cls(original_input.shape, original_input.dtype, packed_output_uint8), scales
+
+    @classmethod
+    def double_quantization(cls, scales: torch.Tensor, scale_block_size: int, num_scale_bits: int):
+        """Perform double quantization on the scales.
+
+        Unlike the `quantize` method quantizing input data, this function quantizes float scales into
+        int8 to further reduce memory usage of scales.
+        """
+        # Double quantization for the scales, int8 per-block quantization
+        assert scales.numel() % scale_block_size == 0, (
+            "Number of scales elements is not divisible by the scale block size."
+        )
+        scale_quant_maxbound = 2 ** (num_scale_bits - 1) - 1
+        block_scales = scales.view(-1, scale_block_size)
+        num_scale_blocks = block_scales.shape[0]
+        scalers_zero_point = block_scales.mean()
+        block_scales = block_scales - scalers_zero_point
+
+        double_quant_scales = scale_quant_maxbound / reduce_amax(block_scales, -1)
+        quantized_scales = block_scales * double_quant_scales.expand(
+            num_scale_blocks, scale_block_size
+        )
+        quantized_scales = (
+            quantized_scales.round()
+            .clamp(-scale_quant_maxbound, scale_quant_maxbound)
+            .to(torch.int8)
+        )
+
+        return quantized_scales, double_quant_scales.flatten(), scalers_zero_point
+
+    def dequantize(self, dtype: torch.dtype = None, **kwarg):
+        """Dequantze NF4 packed tensor to a target dtype."""
+        if dtype is None:
+            dtype = self.metadata["dtype"]
+        cuda_ext = get_cuda_ext()
+
+        # get kwargs
+        scales = kwarg["scale"]
+        block_sizes = kwarg["block_sizes"]
+        double_scale = kwarg["double_scale"]
+        scale_zeros = kwarg["scale_zeros"]
+
+        # unpadd the scales if needed
+        scales = scales.view(-1)[: (self._quantized_data.numel() * 2) // block_sizes[-1]]
+
+        if cuda_ext and self._quantized_data.is_cuda:
+            # with a custom cuda kernel
+            scales = _dequantize_scalers(scales, double_scale, scale_zeros, dtype).flatten()
+            output = cuda_ext.NF4_dequantize(self._quantized_data, scales, block_sizes[-1])
+            return (
+                output.view(-1)[: np.prod(self.metadata["shape"])]  # handle padding
+                .reshape(self.metadata["shape"])
+                .to(dtype)
+            )
+        else:
+            # de-qauntize scales
+            scales = _dequantize_scalers(scales, double_scale, scale_zeros, dtype).flatten()
+            # indexing in torch required long dtype, we may need to optimize this with customized kernels
+            first_half_idx = (self._quantized_data >> 4).to(torch.long)
+            second_half_idx = (self._quantized_data & 0x0F).to(torch.long)
+            scaled_first_half = _nf4_lookup(first_half_idx)  # w0, w2, w4...
+            scaled_second_half = _nf4_lookup(second_half_idx)  # w1, w3, w5...
+
+            # de-quantize tensor
+            first_half = scaled_first_half.view(-1, block_sizes[-1] // 2) * scales.view(-1, 1)
+            second_half = scaled_second_half.view(-1, block_sizes[-1] // 2) * scales.view(-1, 1)
+
+            # merge the interleaving elements
+            first_half = first_half.flatten().unsqueeze(-1).transpose(0, 1)
+            second_half = second_half.flatten().unsqueeze(-1).transpose(0, 1)
+            return (
+                torch.stack([first_half, second_half], dim=-1)
+                .view(-1)[: np.prod(self.metadata["shape"])]  # handle padding
+                .reshape(self.metadata["shape"])
+                .to(dtype)
+            )

modelopt/torch/quantization/qtensor/nvfp4_tensor.py

@@ -0,0 +1,295 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Implements NVFP4 quantization for efficient tensor storage and computation."""
+
+import torch
+
+from ..backends.utils import fp4_compatible
+from ..qtensor.base_qtensor import BaseQuantizedTensor
+from ..utils import reduce_amax, reduce_block_amax, reduce_block_padding
+
+# Define conversion tables
+e2m1_bounds = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5])
+e2m1_values = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6, 0, -0.5, -1, -1.5, -2, -3, -4, -6])
+
+__all__ = ["NVFP4QTensor"]
+
+
+class NVFP4QTensor(BaseQuantizedTensor):
+    """Implements the INT4 quantization on tensors for more efficient storage or computation.
+
+    Attributes:
+        quantized_data (torch.Tensor): The quantized data stored as a packed uint8 tensor.
+    """
+
+    e2m1_values_on_device = {}
+    e2m1_bounds_on_device = {}
+
+    @classmethod
+    def get_e2m1_values(cls, device):
+        """Returns the e2m1 values on the device."""
+        if device not in cls.e2m1_values_on_device:
+            cls.e2m1_values_on_device[device] = e2m1_values.to(device)
+        return cls.e2m1_values_on_device[device]
+
+    @classmethod
+    def get_e2m1_bounds(cls, device):
+        """Returns the e2m1 values on the device."""
+        if device not in cls.e2m1_bounds_on_device:
+            cls.e2m1_bounds_on_device[device] = e2m1_bounds.to(device)
+        return cls.e2m1_bounds_on_device[device]
+
+    @classmethod
+    def get_weights_scaling_factor_2_from_quantizer(cls, weight_quantizer):
+        """Returns per tensor weight scaling factor from the weight_quantizer amax."""
+        # Assert that weight_quantizer has attribute amax
+        assert hasattr(weight_quantizer, "_amax"), "Weight quantizer does not have attribute amax"
+        return weight_quantizer._amax.float() / (6.0 * 448.0)
+
+    @classmethod
+    def get_weights_scaling_factor(
+        cls,
+        input: torch.Tensor,
+        block_size: int,
+        weights_scaling_factor_2: torch.Tensor | None = None,
+        keep_high_precision: bool = False,
+    ):
+        """Returns quantized per block weight scaling factor."""
+        if weights_scaling_factor_2 is None:
+            weights_scaling_factor_2 = cls.get_weights_scaling_factor_2(input)
+
+        # Get per_block amax
+        assert block_size != 0, "Block size is zero. Cannot return per_block amax for given input."
+
+        assert input.shape[-1] % block_size == 0, (
+            "Weight shape is not divisible for block size for block quantization."
+        )
+
+        # Get per block amax
+        per_block_amax = reduce_block_amax(input, block_sizes={-1: block_size}).float()
+        # Get per-block-scale
+        per_block_scale = per_block_amax / (6.0 * weights_scaling_factor_2)
+        # Set all zero values in scale to 1.0
+        per_block_scale[per_block_scale == 0] = 1.0
+        # Convert to torch.float8_e4m3fn
+        if not keep_high_precision:
+            per_block_scale = per_block_scale.to(torch.float8_e4m3fn)
+        return per_block_scale, weights_scaling_factor_2
+
+    @classmethod
+    def get_weights_scaling_factor_2(cls, input: torch.Tensor):
+        """Returns per tensor weight scaling factor."""
+        return reduce_amax(input).float() / (6.0 * 448.0)
+
+    @classmethod
+    def get_activation_scaling_factor(cls, quantizer):
+        """Returns the activation scaling factor for export."""
+        # TODO: Update to use module and not quantizer
+        if not quantizer.is_enabled:
+            return None
+
+        amax = quantizer.export_amax()
+
+        if amax is None:
+            return None
+
+        activation_scaling_factor = amax.float() / (quantizer.maxbound * 448.0)
+
+        assert torch.all(activation_scaling_factor > 0), (
+            f" activation scaling factor {activation_scaling_factor} not positive."
+        )
+
+        return activation_scaling_factor
+
+    @classmethod
+    def _cast_fp4(cls, weight: torch.Tensor):
+        """Converts tensor to uint4."""
+        device = weight.device
+
+        # Extract sign and compute absolute values in one pass
+        sign_bit = (weight < 0).to(torch.uint8)
+        weight_abs = weight.abs_()
+
+        # Get bounds and compute ordinal values
+        e2m1_bounds = cls.get_e2m1_bounds(device)
+        ord = torch.searchsorted(e2m1_bounds, weight_abs, out_int32=True).to(torch.uint8)
+
+        # Efficiently check for rounding at odd-indexed bounds [0.75, 1.75, 2.5]
+        # Only need to check bounds at indices 1, 3, 5
+        odd_bounds = e2m1_bounds[[1, 3, 5]]  # [0.75, 1.75, 2.5]
+        equals_odd_bounds = torch.any(weight_abs.unsqueeze(-1) == odd_bounds, dim=-1).to(
+            torch.uint8
+        )
+
+        # Combine sign, ordinal, and rounding adjustment
+        return (sign_bit << 3) + ord + equals_odd_bounds
+
+    @classmethod
+    def quantize(
+        cls,
+        input: torch.Tensor,
+        block_size: int,
+        weights_scaling_factor: torch.Tensor | None = None,
+        weights_scaling_factor_2: torch.Tensor | None = None,
+        keep_high_precision: bool = False,
+        try_tensorrt: bool = False,
+    ):
+        """Converting a tensor to a quantized format based on NVFP4 quantization.
+
+        Args:
+            input (torch.Tensor): The input tensor to be quantized.
+            block_size (int): The size of each block for quantization.
+            weights_scaling_factor (torch.Tensor): The scaling factor for the weights.
+            weights_scaling_factor_2 (torch.Tensor): The scaling factor for the weights.
+            keep_high_precision (bool): Whether to keep output scales at high precision.
+
+        Returns:
+        tuple: Contains quantized data, quantized per block scaling factor, and per tensor scaling factor.
+        """
+        # Get original input shape
+        input_shape = input.shape
+        input_dtype = input.dtype
+
+        # pad the input if needed
+        input = reduce_block_padding(input, block_sizes={-1: block_size})
+
+        if weights_scaling_factor_2 is None:
+            weights_scaling_factor_2 = cls.get_weights_scaling_factor_2(input)
+
+        # try call trtllm fp4 quantization if possible
+        if (
+            fp4_compatible()
+            and weights_scaling_factor is None
+            and try_tensorrt
+            and block_size == 16
+            and input.is_cuda
+            and input.dtype in [torch.half, torch.bfloat16]
+        ):
+            try:
+                import tensorrt_llm  # noqa: F401
+
+                # Make sure this utils is available for dequantize
+                from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import (
+                    cutlass_fp4_scale_to_modelopt_fp4_scale,  # noqa: F401
+                )
+
+                packed_weight, weights_scaling_factor = torch.ops.trtllm.fp4_quantize(
+                    input, 1.0 / weights_scaling_factor_2, block_size, False
+                )
+                # weights_scaling_factor is ready for nvfp4_gemm to use;
+                # however, it is different from the non trtllm version, so when dequantize,
+                # it will be converted.
+                return (
+                    cls(input_shape, input_dtype, packed_weight),
+                    weights_scaling_factor,
+                    weights_scaling_factor_2,
+                )
+            except ImportError:
+                pass
+
+        if weights_scaling_factor is None:
+            weights_scaling_factor, _ = cls.get_weights_scaling_factor(
+                input, block_size, weights_scaling_factor_2
+            )
+
+        # Reshape the weight and scale factors
+        original_shape = input.shape
+        input = input.view((*tuple(input.shape[:-1]), -1, block_size))
+
+        # Scale weights
+        scaled_weight = input / (
+            (weights_scaling_factor.to(torch.float32) * weights_scaling_factor_2).unsqueeze(-1)
+        )
+
+        # Reshape weights to original
+        scaled_weight = scaled_weight.view(original_shape)
+
+        if keep_high_precision:
+            return scaled_weight
+        # Cast weights to fp4
+        q_weight = cls._cast_fp4(scaled_weight)
+        # Pack weights
+        packed_weight = (q_weight[..., 1::2] << 4) | q_weight[..., 0::2]
+        return (
+            cls(input_shape, input_dtype, packed_weight),
+            weights_scaling_factor,
+            weights_scaling_factor_2,
+        )
+
+    def dequantize(self, dtype: torch.dtype = None, fast=False, **kwarg):
+        """Dequantze NVFP4 packed tensor to a target dtype."""
+        if dtype is None:
+            dtype = self.metadata["dtype"]
+
+        def _unpack_tensor(input: torch.Tensor):
+            # Initialize storage for unpacked tensor
+            unpacked_shape = list(input.shape)
+            unpacked_shape[-1] = unpacked_shape[-1] * 2
+            unpacked = torch.empty(unpacked_shape, dtype=dtype, device=input.device)
+
+            unpacked[..., 1::2] = input >> 4
+            unpacked[..., 0::2] = input & 0x0F
+
+            unpacked = unpacked.reshape(-1)
+            unpacked = self.get_e2m1_values(input.device)[unpacked.long()]
+
+            return unpacked.reshape(unpacked_shape)
+
+        # Get scales from kwargs
+        if kwarg["scale"].dtype == torch.uint8 and kwarg["scale"].ndim == 1:
+            # If quantization is done by trtllm, convert cutlass fp4 scale to modelopt fp4 scale
+            try:
+                from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import (
+                    cutlass_fp4_scale_to_modelopt_fp4_scale,
+                )
+
+                kwarg["scale"] = cutlass_fp4_scale_to_modelopt_fp4_scale(
+                    kwarg["scale"], self.metadata["shape"][-2:]
+                )
+            except ImportError as e:
+                raise ImportError(
+                    "This tensor is quantized by trtllm, but tensorrt_llm cannot be imported."
+                ) from e
+
+        if fast:
+            from ..triton.fp4_kernel import fp4_dequantize
+
+            return fp4_dequantize(
+                self._quantized_data,
+                kwarg["scale"],
+                kwarg["double_scale"],
+                block_size=kwarg["block_sizes"][-1],
+                dtype=dtype,
+            ).reshape(self.metadata["shape"])
+        else:
+            q_per_block_scale = (
+                kwarg["scale"].to(torch.float32)
+                if kwarg["scale"].dtype == torch.float8_e4m3fn
+                else kwarg["scale"]
+            )
+            block_size = kwarg["block_sizes"][-1]
+            per_block_quant_scale = kwarg["double_scale"]
+
+            # Dequantize scales
+            per_block_scale = q_per_block_scale * per_block_quant_scale
+
+            # Unpack and unscale weights
+            deq_data = _unpack_tensor(self._quantized_data)
+
+            deq_data = deq_data.view(
+                (*tuple(deq_data.shape[:-1]), -1, block_size)
+            ) * per_block_scale.unsqueeze(-1)
+            return deq_data.reshape(self.metadata["shape"]).to(dtype)

modelopt/torch/quantization/src/tensor_quant.cpp

@@ -0,0 +1,77 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "tensor_quant.h"
+
+at::Tensor NF4_dequantize(torch::Tensor quantized_data, torch::Tensor scales, int block_size) {
+  TORCH_CHECK(quantized_data.is_cuda());
+  return NF4_dequantize_cuda(quantized_data, scales, block_size);
+}
+at::Tensor NF4_quantize(torch::Tensor quantized_data, torch::Tensor scales, int block_size) {
+  TORCH_CHECK(quantized_data.is_cuda());
+  return NF4_quantize_cuda(quantized_data, scales, block_size);
+}
+at::Tensor INT4_dequantize(torch::Tensor quantized_data, torch::Tensor scales, int block_size) {
+  TORCH_CHECK(quantized_data.is_cuda());
+  return INT4_dequantize_cuda(quantized_data, scales, block_size);
+}
+at::Tensor INT4_quantize(torch::Tensor quantized_data, torch::Tensor scales, int block_size) {
+  TORCH_CHECK(quantized_data.is_cuda());
+  return INT4_quantize_cuda(quantized_data, scales, block_size);
+}
+
+void fake_tensor_quant_(at::Tensor inputs, at::Tensor amax, int num_bits = 8,
+                        bool is_unsigned = false, bool narrow_range = true) {
+  TORCH_CHECK(inputs.is_cuda());
+  TORCH_CHECK(inputs.is_contiguous()) // in-place on non-contiguous tensor is
+                                      // more difficult
+  TORCH_CHECK(amax.numel(), 1);
+  fake_tensor_quant_cuda_inplace(inputs, amax, num_bits, is_unsigned, narrow_range);
+}
+// TODO: Can we add support for CPU tensors here?
+at::Tensor fake_tensor_quant(at::Tensor inputs, at::Tensor amax, int num_bits = 8,
+                             bool is_unsigned = false, bool narrow_range = true) {
+  TORCH_CHECK(inputs.is_cuda());
+  TORCH_CHECK(amax.numel(), 1);
+  return fake_tensor_quant_cuda(inputs.contiguous(), amax.contiguous(), num_bits, is_unsigned,
+                                narrow_range);
+}
+
+at::Tensor fake_tensor_quant_with_axis(at::Tensor inputs, at::Tensor amax, int axis,
+                                       int num_bits = 8, bool is_unsigned = false,
+                                       bool narrow_range = true) {
+  TORCH_CHECK(inputs.is_cuda());
+  TORCH_CHECK(amax.numel(), inputs.size(axis));
+  return fake_tensor_quant_with_axis_cuda(inputs.contiguous(), amax.contiguous(), axis, num_bits,
+                                          is_unsigned, narrow_range);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("fake_tensor_quant_", &fake_tensor_quant_, "Fake Tensor Quant Inplace", py::arg("inputs"),
+        py::arg("amax"), py::arg("num_bits") = 8, py::arg("unsigned") = false,
+        py::arg("narrow_range") = true);
+  m.def("fake_tensor_quant", &fake_tensor_quant, "Fake Tensor Quant", py::arg("inputs"),
+        py::arg("amax"), py::arg("num_bits") = 8, py::arg("unsigned") = false,
+        py::arg("narrow_range") = true);
+  m.def("fake_tensor_quant_with_axis", &fake_tensor_quant_with_axis, "Fake Tensor Quant with axis",
+        py::arg("inputs"), py::arg("amax"), py::arg("axis"), py::arg("num_bits") = 8,
+        py::arg("unsigned") = false, py::arg("narrow_range") = true);
+  m.def("NF4_dequantize", &NF4_dequantize, "Dequantize NF4 weights.");
+  m.def("NF4_quantize", &NF4_quantize, "Quantize NF4 weights.");
+  m.def("INT4_dequantize", &INT4_dequantize, "Dequantize INT4 weights.");
+  m.def("INT4_quantize", &INT4_quantize, "Quantize INT4 weights.");
+}

modelopt/torch/quantization/src/tensor_quant.h

@@ -0,0 +1,69 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <torch/extension.h>
+
+// TODO: add descriptions for functions
+void fake_tensor_quant_cuda_inplace(at::Tensor, at::Tensor, int, bool, bool);
+at::Tensor fake_tensor_quant_cuda(at::Tensor, at::Tensor, int, bool, bool);
+at::Tensor fake_tensor_quant_with_axis_cuda(at::Tensor, at::Tensor, int, int, bool, bool);
+float bits_to_bound(int, int);
+at::Tensor fake_e4m3fy_cuda(at::Tensor inputs);
+
+// Dequantizes data using NF4 quantization scheme and per-block scaling factors.
+//
+// Args:
+//   quantized_data (torch::Tensor): Quantized data packed into uint8.
+//   scales (torch::Tensor): Scaling factors for each data block.
+//   block_size (int): Number of elements per data block.
+//
+// Returns:
+//   at::Tensor: Dequantized data as bf16 tensor.
+at::Tensor NF4_dequantize_cuda(torch::Tensor, torch::Tensor, int);
+
+// Quantizes data using NF4 quantization scheme and per-block scaling factors.
+//
+// Args:
+//   quantized_data (torch::Tensor): High precision data in fp32/fp16/bf16.
+//   scales (torch::Tensor): Scaling factors for each data block.
+//   block_size (int): Number of elements per data block.
+//
+// Returns:
+//   at::Tensor: Quantized data packed in uint8.
+at::Tensor NF4_quantize_cuda(torch::Tensor, torch::Tensor, int);
+
+// Dequantizes data using INT4 quantization scheme and per-block scaling factors.
+//
+// Args:
+//   quantized_data (torch::Tensor): Quantized data packed into uint8.
+//   scales (torch::Tensor): Scaling factors for each data block.
+//   block_size (int): Number of elements per data block.
+//
+// Returns:
+//   at::Tensor: Dequantized data as the same dtype of scales.
+at::Tensor INT4_dequantize_cuda(torch::Tensor, torch::Tensor, int);
+
+// Quantizes data using INT4 quantization scheme and per-block scaling factors.
+//
+// Args:
+//   quantized_data (torch::Tensor): High precision data in fp32/fp16/bf16.
+//   scales (torch::Tensor): Scaling factors for each data block.
+//   block_size (int): Number of elements per data block.
+//
+// Returns:
+//   at::Tensor: Quantized data packed in uint8.
+at::Tensor INT4_quantize_cuda(torch::Tensor, torch::Tensor, int);

modelopt/torch/quantization/src/tensor_quant_fp8.cpp

@@ -0,0 +1,48 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <ATen/ATen.h>
+#include <cuda_fp8.h>
+#include <torch/extension.h>
+
+at::Tensor fake_e4m3fy_cuda(at::Tensor inputs);
+at::Tensor fused_fake_e4m3fy_cuda(at::Tensor inputs, at::Tensor amax, const float zero_threshold);
+
+at::Tensor fake_e4m3fy(at::Tensor inputs) {
+  if (inputs.is_cuda()) {
+    return fake_e4m3fy_cuda(inputs.contiguous());
+  } else {
+    TORCH_CHECK(inputs.dtype() == at::ScalarType::Float);
+    TORCH_CHECK(inputs.is_contiguous());
+    auto out = at::zeros_like(inputs);
+    for (int i = 0; i < inputs.numel(); ++i) {
+      out.data_ptr<float>()[i] =
+          static_cast<float>(static_cast<__nv_fp8_e4m3>(inputs.data_ptr<float>()[i]));
+    }
+    return out;
+  }
+}
+
+at::Tensor fused_fake_e4m3fy(at::Tensor inputs, at::Tensor amax, const float zero_threshold) {
+  return fused_fake_e4m3fy_cuda(inputs.contiguous(), amax, zero_threshold);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("fake_e4m3fy", &fake_e4m3fy, "Reduce precision to E4M3", py::arg("inputs"));
+  m.def("fused_fake_e4m3fy", &fused_fake_e4m3fy, "Reduce precision to E4M3 (fused)",
+        py::arg("inputs"), py::arg("amax"), py::arg("zero_threshold"));
+}