nvidia-modelopt 0.35.0__py3-none-any.whl

modelopt/torch/quantization/tensor_quant.py

@@ -0,0 +1,816 @@
+# 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.
+
+"""Basic tensor quantization functions."""
+
+import warnings
+
+import torch
+from torch.autograd import Function
+from torch.onnx import symbolic_helper
+
+import modelopt.torch.quantization.triton as triton_kernel
+
+from .config import QuantizerAttributeConfig
+from .extensions import get_cuda_ext, get_cuda_ext_fp8, get_cuda_ext_mx
+
+mx_format_map = {
+    (4, 3): "E4M3",
+    (5, 2): "E5M2",
+    (3, 2): "E3M2",
+    (2, 3): "E2M3",
+    8: "INT8",
+    (8, 0): "E8M0",
+    (2, 1): "E2M1",
+    (1, 2): "E1M2",
+    (0, 3): "E0M3",
+    (3, 0): "E3M0",
+}
+
+DISABLE_TRITON_KERNEL = False
+
+
+def scaled_e4m3_impl(
+    inputs: torch.Tensor,  # TODO: check support for multiple inputs
+    amax: torch.Tensor,
+    disable_fused_kernel=True,
+) -> torch.Tensor:
+    """Implementation of fake quantizing input to FP8.
+
+    Args:
+        inputs: Torch tensor.
+        amax: Absolute max range of the input tensor.
+
+    Returns:
+        Input tensors faked quantized to FP8.
+    """
+    cuda_ext_fp8 = get_cuda_ext_fp8(raise_if_failed=True)
+
+    def is_fusable():
+        # ignore no scaling and shape([]) cases
+        if amax is None or len(amax.shape) == 0:
+            return False
+        else:
+            # can't have amax.shape = [1, 1, 4, 1] and the like
+            amax_last_dim_only = amax.numel() == amax.shape[-1]
+            # must be cuda
+            all_cuda = inputs.is_cuda and amax.is_cuda
+
+            # also check explicit disable.
+            return amax_last_dim_only and all_cuda and (not disable_fused_kernel)
+
+    with torch.cuda.device(
+        None if inputs.device.index == torch.cuda.current_device() else inputs.device.index
+    ):
+        # differentiate between fused & unfused cases
+        if is_fusable():
+            zero_threshold = 1.0 / (1 << 24)
+            outputs = cuda_ext_fp8.fused_fake_e4m3fy(inputs, amax.float(), zero_threshold)
+        else:
+            zero_mask = inputs.abs() < 1.0 / (1 << 24)
+
+            if amax is None:
+                outputs = cuda_ext_fp8.fake_e4m3fy(inputs)
+            else:
+                scale = 448.0 / amax
+                outputs = cuda_ext_fp8.fake_e4m3fy(inputs * scale) / scale
+
+            # Zero out values that are tiny.
+            # Tiny values could lead to tiny amax and then large scale which cause overflow/saturation
+            # and won't go back to normal value after dividing by scale. The right behavior is to mark them
+            # as zero which also get rid of inf/nan
+            outputs[zero_mask] = 0.0
+
+        return outputs
+
+
+def fake_quant_impl(
+    inputs: torch.Tensor,
+    amax: torch.Tensor,
+    num_bits=8,
+    unsigned=False,
+    narrow_range=True,
+):
+    """Implementation of fake quantizing input according to number of bits."""
+    cuda_ext = get_cuda_ext()
+
+    with torch.cuda.device(
+        None if inputs.device.index == torch.cuda.current_device() else inputs.device.index
+    ):
+        if amax.numel() == 1:
+            outputs = cuda_ext.fake_tensor_quant(inputs, amax, num_bits, unsigned, narrow_range)
+        else:
+            axis = amax.shape.index(amax.numel())
+            outputs = cuda_ext.fake_tensor_quant_with_axis(
+                inputs, amax.squeeze(), axis, num_bits, unsigned, narrow_range
+            )
+        return outputs
+
+
+def _quantize_impl(
+    inputs: torch.Tensor,
+    amax: torch.Tensor,
+    num_bits: int = 8,
+    exponent_bits: int = 0,
+    unsigned: bool = False,
+    narrow_range: bool = True,
+):
+    if num_bits == 8 and exponent_bits == 4:
+        return scaled_e4m3_impl(inputs=inputs, amax=amax)
+    elif isinstance(num_bits, int):
+        return fake_quant_impl(
+            inputs=inputs,
+            amax=amax,
+            num_bits=num_bits,
+            unsigned=unsigned,
+            narrow_range=narrow_range,
+        )
+    else:
+        raise ValueError(
+            f"Invalid combination of (num_bits, exponent_bits): ({num_bits}, {exponent_bits})."
+        )
+
+
+def _quantize_impl_abstract(
+    input: torch.Tensor,
+    amax: torch.Tensor,
+    num_bits: int = 8,
+    exponent_bits: int = 0,
+    unsigned: bool = False,
+    narrow_range: bool = True,
+) -> torch.Tensor:
+    """Register an abstract implementation for quantizing tensor.
+
+    This abstract function returns an empty tensor with the same shape and dtype.
+    """
+    output = torch.empty_like(input)
+
+    return output
+
+
+# Argument types: Tensor, int, NoneType, int, int, int, int,
+def _dynamic_block_quantize_impl(
+    inputs: torch.Tensor,
+    block_size: int,
+    amax: torch.Tensor,
+    num_bits: int,
+    exponent_bits: int,
+    scale_num_bits: int,
+    scale_exponent_bits: int,
+):
+    scale_bits = (scale_exponent_bits, scale_num_bits - scale_exponent_bits - 1)
+    if exponent_bits != 0:
+        num_bits = (exponent_bits, num_bits - exponent_bits - 1)
+    if num_bits in mx_format_map:
+        assert scale_bits in mx_format_map, f"Scale bits should be in {mx_format_map.keys()}"
+        if scale_bits != (8, 0):
+            assert amax.is_cuda, "amax must be a CUDA tensor for dynamic block quantization."
+            if amax.numel() != 1:
+                amax = amax.amax()
+        with torch.cuda.device(
+            None if inputs.device.index == torch.cuda.current_device() else inputs.device.index
+        ):
+            if (
+                num_bits == (2, 1)  # type: ignore[comparison-overlap]
+                and scale_bits == (4, 3)
+                and triton_kernel.IS_AVAILABLE
+                and not DISABLE_TRITON_KERNEL
+                and amax is not None
+            ):
+                return triton_kernel.fp4_fake_quant_block(inputs, amax)
+            cuda_ext_mx = get_cuda_ext_mx(raise_if_failed=True)
+            return cuda_ext_mx.fused_amax_convert(
+                inputs,
+                block_size,
+                getattr(cuda_ext_mx.Types, mx_format_map[num_bits]),
+                getattr(cuda_ext_mx.Types, mx_format_map[scale_bits]),
+                amax,
+            )
+    else:
+        raise NotImplementedError(
+            f"Unsupported num_bits: {num_bits}, scale_bits: {scale_bits} for dynamic block quantization."
+        )
+
+
+def _dynamic_block_quantize_impl_abstract(
+    inputs: torch.Tensor,
+    block_size: int,
+    amax: torch.Tensor,
+    num_bits: int,
+    exponent_bits: int,
+    scale_num_bits: int,
+    scale_exponent_bits: int,
+):
+    """Register an abstract implementation for dynamic block quantization.
+
+    This abstract function returns an empty tensor with the same shape and dtype.
+    """
+    output = torch.empty_like(inputs)
+
+    return output
+
+
+quantize_op = _quantize_impl
+dynamic_block_quantize_op = _dynamic_block_quantize_impl
+# Define custom operators via torch.library if supported:
+# 1. quantize_op: Applies static quantization to the input tensor using a specified amax.
+# 2. dynamic_block_quantize_op: Performs blockwise double quantization with dynamically
+#    determined scales, governed by the given quantization format (scale_num_bits, scale_exponent_bits).
+try:
+    torch.library.define(
+        "tensorrt::quantize_op",
+        "(Tensor input, Tensor amax, int num_bits, int exponent_bits, "
+        "bool unsigned, bool narrow_range) -> Tensor",
+    )
+    torch.library.define(
+        "tensorrt::dynamic_block_quantize_op",
+        "(Tensor input, int block_size, Tensor amax, int num_bits, int exponent_bits, "
+        "int scale_num_bits, int scale_exponent_bits) -> Tensor",
+    )
+    torch.library.define(
+        "tensorrt::dynamic_block_quantize_op.overload",
+        "(Tensor input, int block_size, None amax, int num_bits, int exponent_bits, "
+        "int scale_num_bits, int scale_exponent_bits) -> Tensor",
+    )
+
+    # Implement the None amax case
+    def _dynamic_block_quantize_impl_none_amax(
+        inputs: torch.Tensor,
+        block_size: int,
+        amax: None,
+        num_bits: int,
+        exponent_bits: int,
+        scale_num_bits: int,
+        scale_exponent_bits: int,
+    ):
+        return torch.empty_like(inputs)
+
+    # Register the implementation for both CPU and CUDA
+    torch.library.impl("tensorrt::quantize_op", ["cpu", "cuda"])(_quantize_impl)
+    torch.library.impl("tensorrt::dynamic_block_quantize_op", ["cpu", "cuda"])(
+        _dynamic_block_quantize_impl
+    )
+    torch.library.impl("tensorrt::dynamic_block_quantize_op.overload", ["cpu", "cuda"])(
+        _dynamic_block_quantize_impl_none_amax
+    )
+
+    # Register the fake implementation
+    torch.library.register_fake("tensorrt::quantize_op")(_quantize_impl_abstract)
+    torch.library.register_fake("tensorrt::dynamic_block_quantize_op")(
+        _dynamic_block_quantize_impl_abstract
+    )
+    torch.library.register_fake("tensorrt::dynamic_block_quantize_op.overload")(
+        _dynamic_block_quantize_impl_abstract
+    )
+
+    quantize_op = torch.ops.tensorrt.quantize_op
+    dynamic_block_quantize_op = torch.ops.tensorrt.dynamic_block_quantize_op
+except (AttributeError, RuntimeError) as e:
+    # torch.library is an experimental feature, the function signatures may change overtime.
+    warnings.warn(
+        "Unable to register operators with torch.library. Exporting quantized models with"
+        f" torch.export will not be supported.\n{e}"
+    )
+
+# Predefined descriptors
+QUANT_DESC_8BIT_PER_TENSOR = QuantizerAttributeConfig(num_bits=8)
+QUANT_DESC_UNSIGNED_8BIT_PER_TENSOR = QuantizerAttributeConfig(num_bits=8, unsigned=True)
+QUANT_DESC_8BIT_CONV1D_WEIGHT_PER_CHANNEL = QuantizerAttributeConfig(num_bits=8, axis=(0))
+QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL = QuantizerAttributeConfig(num_bits=8, axis=(0))
+QUANT_DESC_8BIT_CONV3D_WEIGHT_PER_CHANNEL = QuantizerAttributeConfig(num_bits=8, axis=(0))
+QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW = QuantizerAttributeConfig(num_bits=8, axis=(0))
+QUANT_DESC_8BIT_CONVTRANSPOSE1D_WEIGHT_PER_CHANNEL = QuantizerAttributeConfig(num_bits=8, axis=(0))
+QUANT_DESC_8BIT_CONVTRANSPOSE2D_WEIGHT_PER_CHANNEL = QuantizerAttributeConfig(num_bits=8, axis=(0))
+QUANT_DESC_8BIT_CONVTRANSPOSE3D_WEIGHT_PER_CHANNEL = QuantizerAttributeConfig(num_bits=8, axis=(0))
+
+
+@torch.jit.script
+def _fake_tensor_quant_backward(inputs, amax: torch.Tensor | None, grad_outputs):
+    # Skip clip for MX formats
+    if amax is None:
+        return grad_outputs
+
+    zero = grad_outputs.new_zeros(1)
+    grad_inputs = torch.where(inputs.abs() <= amax, grad_outputs, zero)
+    return grad_inputs
+
+
+def _fake_quant_backward_function(ctx, grad_outputs, num_args=1):
+    saved_tensors = ctx.saved_tensors
+    if len(saved_tensors) == 0:
+        return (grad_outputs,) + (None,) * (num_args - 1)
+    inputs, amax = saved_tensors
+    return (_fake_tensor_quant_backward(inputs, amax, grad_outputs),) + (None,) * (num_args - 1)
+
+
+def _save_for_backward_if_needed(ctx, pass_through_bwd, inputs, amax):
+    if not pass_through_bwd and amax is not None:
+        amax = (
+            amax
+            if isinstance(amax, torch.Tensor)
+            else torch.tensor(amax, device=inputs.device, dtype=inputs.dtype)
+        )
+        ctx.save_for_backward(inputs, amax)
+
+
+class FakeTensorQuantFunction(Function):
+    """Fake version of TensorQuantFunction use CUDA extension."""
+
+    @staticmethod
+    @symbolic_helper.parse_args("v", "t", "t", "i", "b", "b", "s", "b", "i", "i")
+    def symbolic(
+        g,
+        inputs,
+        amax,
+        bias=None,
+        num_bits=8,
+        unsigned=False,
+        narrow_range=True,
+        trt_high_precision_dtype=None,
+        pass_through_bwd=False,
+        block_size=None,
+        axis=None,
+    ):
+        """ONNX symbolic function."""
+        from .export_onnx import export_int4, export_int8
+
+        if num_bits == 4:
+            return export_int4(
+                g, inputs, amax, num_bits, trt_high_precision_dtype, block_size, axis
+            )
+
+        return export_int8(
+            g, inputs, amax, num_bits, unsigned, narrow_range, trt_high_precision_dtype
+        )
+
+    @staticmethod
+    def forward(
+        ctx,
+        inputs,
+        amax,
+        bias=None,
+        num_bits=8,
+        unsigned=False,
+        narrow_range=True,
+        trt_high_precision_dtype=None,
+        pass_through_bwd=False,
+        block_size=None,
+        axis=None,
+    ):
+        """Forward method."""
+        if bias is not None:
+            inputs = inputs - bias
+
+        _save_for_backward_if_needed(ctx, pass_through_bwd, inputs, amax)
+
+        def legacy_quant_func():
+            # The LegacyFakeTensorQuantFunction support cpu and amax with any shape that can be broadcasted to inputs.
+            outputs, scale = _tensor_quant(inputs, amax, num_bits, unsigned, narrow_range)
+            return outputs / scale.to(inputs.dtype)
+
+        if not inputs.is_cuda:
+            outputs = legacy_quant_func()
+        else:
+            try:
+                outputs = quantize_op(
+                    inputs,
+                    amax,
+                    num_bits=num_bits,
+                    exponent_bits=0,
+                    unsigned=unsigned,
+                    narrow_range=narrow_range,
+                )
+            except (AttributeError, ValueError):
+                # AttributeError: cuda_ext is not imported, possibly due to CPU only installation
+                # ValueError: cuda_ext is installed, but trying to perform multidimensional quantization (amax dim > 1)
+                outputs = legacy_quant_func()
+
+        if bias is not None:
+            outputs = outputs + bias
+
+        return outputs
+
+    @staticmethod
+    def backward(ctx, grad_outputs):
+        """Implements straight through estimation with clipping."""
+        return _fake_quant_backward_function(ctx, grad_outputs, num_args=10)
+
+
+class ScaledE4M3Function(Function):
+    """E4M3fy input with scale."""
+
+    @staticmethod
+    @symbolic_helper.parse_args("v", "t", "t", "i", "i", "s", "b")
+    def symbolic(
+        g,
+        inputs,
+        amax=None,
+        bias=None,
+        E=4,  # noqa: N803
+        M=3,  # noqa: N803
+        trt_high_precision_dtype=None,
+        pass_through_bwd=False,
+    ):
+        """ONNX symbolic function."""
+        from .export_onnx import export_fp8
+
+        return export_fp8(g, inputs, amax, trt_high_precision_dtype)
+
+    @staticmethod
+    # Default values could cause errors from TorchDynamo during torch.export
+    def forward(
+        ctx,
+        inputs,
+        amax,
+        bias,
+        E,  # noqa: N803
+        M,  # noqa: N803
+        trt_high_precision_dtype=None,
+        pass_through_bwd=False,
+    ):
+        """Forward method."""
+        if E != 4 or M != 3:
+            raise NotImplementedError("Only support E=4 & M=3 for now.")
+
+        if bias is not None:
+            inputs = inputs - bias
+
+        _save_for_backward_if_needed(ctx, pass_through_bwd, inputs, amax)
+
+        outputs = quantize_op(
+            inputs,
+            amax,
+            num_bits=8,
+            exponent_bits=4,
+            unsigned=False,
+            narrow_range=False,
+        )
+
+        if bias is not None:
+            outputs = outputs + bias
+
+        return outputs
+
+    @staticmethod
+    def backward(ctx, grad_outputs):
+        """Implements straight through estimation with clipping."""
+        return _fake_quant_backward_function(ctx, grad_outputs, num_args=7)
+
+
+def _dynamic_block_quantize_forward(
+    ctx,
+    inputs,
+    block_size,
+    amax,
+    num_bits,
+    scale_bits,
+    trt_high_precision_dtype=None,
+    onnx_quantizer_type="dynamic",
+    pass_through_bwd=True,
+):
+    """Forward method."""
+    if isinstance(num_bits, int):
+        # special case for INT dynamic block quantization, e.g. MXINT8
+        exponent_bits = 0
+    else:
+        assert isinstance(num_bits, tuple) and len(num_bits) == 2
+        exponent_bits = num_bits[0]
+        num_bits = num_bits[0] + num_bits[1] + 1
+    assert isinstance(scale_bits, tuple) and len(scale_bits) == 2
+    scale_exponent_bits = scale_bits[0]
+    scale_num_bits = scale_bits[0] + scale_bits[1] + 1
+    outputs = dynamic_block_quantize_op(
+        inputs,
+        block_size,
+        amax,
+        num_bits,
+        exponent_bits,
+        scale_num_bits,
+        scale_exponent_bits,
+    )
+    return outputs
+
+
+class DynamicBlockQuantizationFunction(Function):
+    """Dynamic block quantization functional."""
+
+    @staticmethod
+    @symbolic_helper.parse_args("v", "i", "v", "t", "is", "is", "s", "s", "b")
+    def symbolic(
+        g,
+        inputs,
+        block_size,
+        amax,
+        bias,
+        num_bits,
+        scale_bits,
+        trt_high_precision_dtype=None,
+        onnx_quantizer_type="dynamic",
+        pass_through_bwd=True,
+    ):
+        """ONNX symbolic function."""
+        from .export_onnx import export_fp4, export_mxfp8
+
+        if num_bits == (2, 1) and scale_bits == (4, 3):
+            return export_fp4(
+                g,
+                inputs,
+                block_size,
+                amax,
+                num_bits,
+                trt_high_precision_dtype,
+                onnx_quantizer_type,
+            )
+        if num_bits == (4, 3) and scale_bits == (8, 0):
+            return export_mxfp8(
+                g,
+                inputs,
+                onnx_quantizer_type,
+                block_size,
+            )
+        raise NotImplementedError(
+            f"Unsupported num_bits: {num_bits} and scale_bits: {scale_bits} for ONNX export."
+        )
+
+    @staticmethod
+    def forward(
+        ctx,
+        inputs,
+        block_size,
+        amax,
+        bias,
+        num_bits,
+        scale_bits,
+        trt_high_precision_dtype=None,
+        onnx_quantizer_type="dynamic",
+        pass_through_bwd=True,
+    ):
+        """Forward method."""
+        _save_for_backward_if_needed(ctx, pass_through_bwd, inputs, amax)
+        return _dynamic_block_quantize_forward(
+            ctx,
+            inputs,
+            block_size,
+            amax,
+            num_bits,
+            scale_bits,
+            trt_high_precision_dtype,
+            onnx_quantizer_type,
+            pass_through_bwd,
+        )
+
+    @staticmethod
+    def backward(ctx, grad_outputs):
+        """Implements straight through estimation with clipping."""
+        return _fake_quant_backward_function(ctx, grad_outputs, num_args=9)
+
+
+class TensorQuantFunction(Function):
+    """A universal tensor quantization function.
+
+    Take an input tensor, output an quantized tensor. The granularity of scale can be interpreted from the
+    shape of amax.
+    output_dtype indicates whether the quantized value will be stored in integer or float. The reason we want to store
+    it in float is the pytorch function takes the quantized value may not accept integer input, e.g. Conv2D.
+
+    It uses 2^num_bits -1 values instead of 2^num_bits. e.g., for num_bits=8, it uses [-127, 127] instead of [-128, 127]
+    """
+
+    @staticmethod
+    @symbolic_helper.parse_args("v", "t", "t", "i", "b", "b", "s")
+    def symbolic(
+        g,
+        inputs,
+        amax,
+        bias=None,
+        num_bits=8,
+        unsigned=False,
+        narrow_range=True,
+        trt_high_precision_dtype=None,
+    ):
+        """ONNX symbolic function."""
+        from .export_onnx import export_int8
+
+        return export_int8(
+            g, inputs, amax, num_bits, unsigned, narrow_range, trt_high_precision_dtype
+        )
+
+    @staticmethod
+    def forward(
+        ctx,
+        inputs,
+        amax,
+        bias=None,
+        num_bits=8,
+        unsigned=False,
+        narrow_range=True,
+        trt_high_precision_dtype=None,
+    ):
+        """Forward method.
+
+        Follow tensorflow convention, max value is passed in and used to decide scale, instead of inputting scale
+        directly. Though inputting scale directly may be more natural to use.
+
+        Args:
+            ctx: A Context object to store tensors for backward.
+            inputs: A Tensor of type float32.
+            amax: A Tensor of type float32. Inputs will be quantized within range [-amax, amax]
+                amax will be broadcasted to inputs tensor.
+            num_bits: A integer used to calculate scaling factor, scale = (2^(num_bits-1) - 1) / max
+                Effectively, it indicates how many integer bits is used to represent the value. Default 8.
+            output_dtype: A type of Tensor. torch.int32 or torch.float32.
+            unsigned: A boolean. Use unsigned integer range. E.g. [0, 255] for num_bits=8. Default False.
+            narrow_range: A boolean. Use symmetric integer range for signed quantization
+                E.g. [-127,127] instead of [-128,127] for num_bits=8. Default True.
+
+        Returns:
+            outputs: A Tensor of type output_dtype.
+            scale: A Tensor of type float32. outputs / scale will dequantize outputs tensor.
+
+        Raises:
+            ValueError:
+        """
+        if bias is not None:
+            inputs = inputs - bias
+
+        ctx.save_for_backward(inputs, amax)
+
+        outputs, scale = _tensor_quant(inputs, amax, num_bits, unsigned, narrow_range)
+        # Check if scale overflows FP16
+        if outputs.dtype == torch.half and scale.max() > 65504:
+            raise ValueError(f"scale is too large for FP16 with amax={amax}")
+
+        if bias is not None:
+            outputs = outputs + bias
+
+        return outputs, scale.to(inputs.dtype)
+
+    @staticmethod
+    def backward(ctx, grad_outputs, grad_scale):
+        """Implements straight through estimation with clipping.
+
+        For -amax <= input <= amax the gradient passes straight through, otherwise the gradient is zero.
+
+        Args:
+            ctx: A Context object with saved tensors from forward.
+            grad_outputs: A tensor of gradient of outputs.
+            grad_scale: A tensor of gradient of scale.
+
+        Returns:
+            grad_inputs: A tensor of gradient.
+        """
+        inputs, amax = ctx.saved_tensors
+        zero = grad_outputs.new_zeros(1)  # create a zero tensor with the same type and device
+        grad_inputs = torch.where(inputs.abs() <= amax, grad_outputs, zero)
+        return grad_inputs, None, None, None, None, None, None
+
+
+class LegacyFakeTensorQuantFunction(Function):
+    """Fake version of TensorQuantFunction.
+
+    See comments of TensorQuantFunction, arguments are the same.
+    """
+
+    @staticmethod
+    def forward(ctx, inputs, amax, bias, num_bits=8, unsigned=False, narrow_range=True):
+        """Forward method."""
+        if bias is not None:
+            inputs = inputs - bias
+
+        ctx.save_for_backward(inputs, amax)
+
+        outputs, scale = _tensor_quant(inputs, amax, num_bits, unsigned, narrow_range)
+
+        if bias is not None:
+            outputs = outputs + bias
+
+        return outputs / scale.to(inputs.dtype)
+
+    @staticmethod
+    def backward(ctx, grad_outputs):
+        """Implements straight through estimation."""
+        inputs, amax = ctx.saved_tensors
+        zero = grad_outputs.new_zeros(1)
+        grad_inputs = torch.where(inputs.abs() <= amax, grad_outputs, zero)
+        return grad_inputs, None, None, None, None, None
+
+
+def _tensor_quant(inputs, amax, num_bits=8, unsigned=False, narrow_range=True):
+    """Shared function body between TensorQuantFunction and FakeTensorQuantFunction."""
+    # Fine scale, per channel scale will be handled by broadcasting, which could be tricky. Pop a warning.
+    if unsigned and inputs.min() < 0.0:
+        raise TypeError("Negative values encountered in unsigned quantization.")
+
+    # Computation can be done in FP32 to prevent potential over flow.
+    input_dtype = inputs.dtype
+    if inputs.dtype == torch.half:
+        inputs = inputs.float()
+    if amax.dtype == torch.half:
+        amax = amax.float()
+
+    min_amax = amax.min()
+    if min_amax < 0:
+        raise ValueError("Negative values in amax")
+
+    max_bound = torch.tensor((2.0 ** (num_bits - 1 + int(unsigned))) - 1.0, device=amax.device)
+    if unsigned:
+        min_bound = 0
+    elif narrow_range:
+        min_bound = -max_bound
+    else:
+        min_bound = -max_bound - 1
+    scale = max_bound / amax
+
+    epsilon = 1.0 / (1 << 24)
+    if min_amax <= epsilon:  # Treat amax smaller than minimum representable of fp16 0
+        zero_amax_mask = amax <= epsilon
+        scale[zero_amax_mask] = 0  # Value quantized with amax=0 should all be 0
+
+    outputs = torch.clamp((inputs * scale).round_(), min_bound, max_bound)
+
+    if min_amax <= epsilon:
+        scale[zero_amax_mask] = (
+            1.0  # Return 1 makes more sense for values quantized to 0 with amax=0
+        )
+
+    if input_dtype == torch.half:
+        outputs = outputs.half()
+
+    return outputs, scale
+
+
+class FakeAffineTensorQuantFunction(Function):
+    """Fake version of affine quantization.
+
+    gemmlowp style scale+shift quantization. See more details in
+    https://github.com/google/gemmlowp/blob/master/doc/quantization.md.
+
+    We DO NOT recommend affine quantization on weights for performance reason. There might be value to affine quantize
+    activation as it can be cancelled by bias and comes with no performance penalty. This functionality is only added
+    for experimental purpose.
+    """
+
+    @staticmethod
+    def forward(ctx, inputs, min_range, max_range, num_bits=8):
+        """As it will be only applied on activation with per tensor granularity, broadcast is not needed.
+
+        Args:
+            ctx: Pytorch convention.
+            inputs: A Tensor of type float32.
+            min_range: A float.
+            max_range: A float.
+            num_bits: An integer
+
+        Returns:
+            outputs: A Tensor of type output_dtype
+        """
+        ctx.save_for_backward(inputs, min_range, max_range)
+
+        step_size = (max_range - min_range) / (2.0**num_bits - 1)
+
+        min_bound = -(2.0 ** (num_bits - 1))
+        max_bound = 2.0 ** (num_bits - 1) - 1
+
+        quant_zero = torch.round(min_range / step_size) - min_bound
+        quantized = torch.round(inputs / step_size) - quant_zero
+        quantized = torch.clamp(quantized, min_bound, max_bound)
+
+        outputs = (quantized + quant_zero) * step_size
+
+        return outputs
+
+    @staticmethod
+    def backward(ctx, grad_outputs):
+        """Implements straight through estimation with clipping.
+
+        Args:
+            ctx: Pytorch convention.
+            grad_output: A tensor of gradient of outputs.
+
+        Returns:
+            grad_inputs: A tensor of gradient
+        """
+        inputs, min_range, max_range = ctx.saved_tensors
+        zero = grad_outputs.new_zeros(1)
+        grad_inputs = torch.where((inputs <= max_range) * (inputs >= min_range), grad_outputs, zero)
+        return grad_inputs, None, None, None
+
+
+tensor_quant = TensorQuantFunction.apply
+legacy_fake_tensor_quant = LegacyFakeTensorQuantFunction.apply
+fake_tensor_quant = FakeTensorQuantFunction.apply
+fake_affine_tensor_quant = FakeAffineTensorQuantFunction.apply
+scaled_e4m3 = ScaledE4M3Function.apply
+dynamic_block_quant = DynamicBlockQuantizationFunction.apply