[BE] Rename qparams for tinygemm (#2344)

Author: jainapurva

torchao/dtypes/affine_quantized_tensor.py

@@ -26,14 +26,14 @@
     choose_qparams_and_quantize_affine_hqq,
     dequantize_affine,
     dequantize_affine_float8,
-    dequantize_affine_float_zero_point,
     dequantize_affine_floatx,
     dequantize_affine_no_zero_point,
+    dequantize_affine_tinygemm,
     quantize_affine,
     quantize_affine_float8,
-    quantize_affine_float_zero_point,
     quantize_affine_floatx,
     quantize_affine_no_zero_point,
+    quantize_affine_tinygemm,
 )
 from torchao.utils import (
     TORCH_VERSION_AT_LEAST_2_5,
@@ -155,7 +155,7 @@ def dequantize(self, output_dtype: Optional[torch.dtype] = None) -> torch.Tensor
         else:
             data, scale, zero_point = self.tensor_impl.get_plain()
             if self.zero_point_domain == ZeroPointDomain.FLOAT:
-                dq = dequantize_affine_float_zero_point(
+                dq = dequantize_affine_tinygemm(
                     data,
                     self.block_size,
                     scale,
@@ -339,7 +339,7 @@ def from_hp_to_intx(
                     quant_max,
                 )
             elif zero_point_domain == ZeroPointDomain.FLOAT:
-                data = quantize_affine_float_zero_point(
+                data = quantize_affine_tinygemm(
                     input_float,
                     block_size,
                     scale,
@@ -410,7 +410,7 @@ def from_hp_to_intx_static(
                 quant_max,
             )
         elif zero_point_domain == ZeroPointDomain.FLOAT:
-            int_data = quantize_affine_float_zero_point(
+            int_data = quantize_affine_tinygemm(
                 input_float,
                 block_size,
                 scale,

torchao/dtypes/affine_quantized_tensor_ops.py

@@ -93,8 +93,8 @@
 from torchao.quantization.quant_primitives import (
     ZeroPointDomain,
     dequantize_affine,
-    dequantize_affine_float_zero_point,
     dequantize_affine_no_zero_point,
+    dequantize_affine_tinygemm,
 )
 from torchao.utils import (
     fill_defaults,
@@ -318,7 +318,7 @@ def _(func, types, args, kwargs):
     # we need to increase block size to correct dim
     new_blocks = idx.dim() - 1
     if args[1].zero_point_domain == ZeroPointDomain.FLOAT:
-        _dequantize_affine = dequantize_affine_float_zero_point
+        _dequantize_affine = dequantize_affine_tinygemm
     elif args[1].zero_point_domain == ZeroPointDomain.NONE:
         _dequantize_affine = dequantize_affine_no_zero_point
     else:

torchao/dtypes/uintx/int4_cpu_layout.py

@@ -19,7 +19,7 @@
 from torchao.dtypes.utils import AQTTensorImpl, Layout, is_device
 from torchao.quantization.quant_primitives import (
     ZeroPointDomain,
-    quantize_affine_float_zero_point,
+    quantize_affine_tinygemm,
 )
 from torchao.utils import (
     TORCH_VERSION_AT_LEAST_2_5,
@@ -266,7 +266,7 @@ def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         # TODO: move this to `unpack_tinygemm_scales_and_zeros`?
         scale = scale.reshape(scale.shape[:-1]).contiguous()
         zero = zero.reshape(zero.shape[:-1]).contiguous()
-        int_data = quantize_affine_float_zero_point(
+        int_data = quantize_affine_tinygemm(
             dequantized,
             block_size,
             scale,

torchao/dtypes/uintx/int4_xpu_layout.py

@@ -1,3 +1,9 @@
+# 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.
+
 from dataclasses import dataclass
 from typing import Optional, Tuple
 
@@ -372,7 +378,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs):
     def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         from torchao.quantization.quant_primitives import (
             quantize_affine,
-            quantize_affine_float_zero_point,
+            quantize_affine_tinygemm,
         )
         from torchao.quantization.utils import unpack_tinygemm_scales_and_zeros
 
@@ -423,7 +429,7 @@ def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
             # TODO: move this to `unpack_tinygemm_scales_and_zeros`?
             scale = scale.reshape(scale.shape[:-1]).contiguous()
             zero = zero.reshape(zero.shape[:-1]).contiguous()
-            int_data = quantize_affine_float_zero_point(
+            int_data = quantize_affine_tinygemm(
                 dequantized,
                 block_size,
                 scale,

torchao/dtypes/uintx/tensor_core_tiled_layout.py

@@ -21,7 +21,7 @@
 from torchao.quantization.quant_primitives import (
     ZeroPointDomain,
     _get_reduction_params,
-    quantize_affine_float_zero_point,
+    quantize_affine_tinygemm,
 )
 from torchao.utils import (
     TORCH_VERSION_AT_LEAST_2_5,
@@ -511,7 +511,7 @@ def dequant_4d(self):
         target_dtype = torch.int32
         quant_min = 0
         quant_max = 15
-        int_data = quantize_affine_float_zero_point(
+        int_data = quantize_affine_tinygemm(
             dequantized,
             self.block_size,
             scale,

torchao/prototype/parq/quant/uniform_torchao.py

@@ -18,11 +18,11 @@
     choose_qparams_affine_dont_preserve_zero,
     choose_qparams_affine_tinygemm,
     dequantize_affine,
-    dequantize_affine_float_zero_point,
     dequantize_affine_no_zero_point,
+    dequantize_affine_tinygemm,
     quantize_affine,
-    quantize_affine_float_zero_point,
     quantize_affine_no_zero_point,
+    quantize_affine_tinygemm,
 )
 
 from .quantizer import Quantizer
@@ -76,8 +76,8 @@ def quantize(
 
         if self.zero_point_domain == ZeroPointDomain.FLOAT and not self.preserve_zero:
             _choose_qparams_affine = choose_qparams_affine_tinygemm
-            _quantize_affine = quantize_affine_float_zero_point
-            _dequantize_affine = dequantize_affine_float_zero_point
+            _quantize_affine = quantize_affine_tinygemm
+            _dequantize_affine = dequantize_affine_tinygemm
         elif self.zero_point_domain == ZeroPointDomain.INT and not self.preserve_zero:
             _choose_qparams_affine = choose_qparams_affine_dont_preserve_zero
             _quantize_affine = quantize_affine

torchao/quantization/quant_primitives.py

@@ -30,10 +30,10 @@
     "choose_qparams_affine_floatx",
     "quantize_affine",
     "quantize_affine_no_zero_point",
-    "quantize_affine_float_zero_point",
+    "quantize_affine_tinygemm",
     "dequantize_affine",
     "dequantize_affine_no_zero_point",
-    "dequantize_affine_float_zero_point",
+    "dequantize_affine_tinygemm",
     "quantize_affine_floatx",
     "dequantize_affine_floatx",
     "fake_quantize_affine",
@@ -428,7 +428,7 @@ def _quantize_affine_no_dtype_cast(
     return quant
 
 
-def quantize_affine_float_zero_point(
+def quantize_affine_tinygemm(
     input: torch.Tensor,
     block_size: List[int],
     scale: torch.Tensor,
@@ -453,7 +453,7 @@ def quantize_affine_float_zero_point(
     # torch.uintx dtypes yet
     if output_dtype in _SUB_BYTE_UINT_BOUNDS:
         output_dtype = torch.uint8
-    return _quantize_affine_float_zero_point_no_dtype_cast(
+    return _quantize_affine_tinygemm_no_dtype_cast(
         input,
         block_size,
         scale,
@@ -463,7 +463,7 @@ def quantize_affine_float_zero_point(
     ).to(output_dtype)
 
 
-def _quantize_affine_float_zero_point_no_dtype_cast(
+def _quantize_affine_tinygemm_no_dtype_cast(
     input: torch.Tensor,
     block_size: Tuple[int, ...],
     scale: torch.Tensor,
@@ -803,7 +803,7 @@ def dequantize_affine_no_zero_point(
     )
 
 
-def _dequantize_affine_float_zero_point_no_dtype_check(
+def _dequantize_affine_tinygemm_no_dtype_check(
     input: torch.Tensor,
     block_size: List[int],
     scale: torch.Tensor,
@@ -848,7 +848,7 @@ def _dequantize_affine_float_zero_point_no_dtype_check(
     return dequant.view(original_shape).to(output_dtype)
 
 
-def dequantize_affine_float_zero_point(
+def dequantize_affine_tinygemm(
     input: torch.Tensor,
     block_size: Tuple[int, ...],
     scale: torch.Tensor,
@@ -887,7 +887,7 @@ def dequantize_affine_float_zero_point(
         torch.bfloat16,
     ], f"Unsupported output dtype: {output_dtype}"
     quant_min, quant_max = _get_and_check_qmin_qmax(input_dtype, quant_min, quant_max)
-    return _dequantize_affine_float_zero_point_no_dtype_check(
+    return _dequantize_affine_tinygemm_no_dtype_check(
         input,
         block_size,
         scale,
@@ -1013,8 +1013,8 @@ def _do_fake_quantize_affine(
         _quantize_affine = _quantize_affine_no_dtype_cast
         _dequantize_affine = _dequantize_affine_no_dtype_check
     elif zero_point_domain == ZeroPointDomain.FLOAT:
-        _quantize_affine = _quantize_affine_float_zero_point_no_dtype_cast
-        _dequantize_affine = _dequantize_affine_float_zero_point_no_dtype_check
+        _quantize_affine = _quantize_affine_tinygemm_no_dtype_cast
+        _dequantize_affine = _dequantize_affine_tinygemm_no_dtype_check
     elif ZeroPointDomain == ZeroPointDomain.NONE:
         _quantize_affine = _quantize_affine_no_zero_point_no_dtype_cast
         _dequantize_affine = _dequantize_affine_no_zero_point_no_dtype_check

torchao/quantization/utils.py

@@ -19,11 +19,11 @@
     choose_qparams_affine_dont_preserve_zero,
     choose_qparams_affine_tinygemm,
     dequantize_affine,
-    dequantize_affine_float_zero_point,
     dequantize_affine_no_zero_point,
+    dequantize_affine_tinygemm,
     quantize_affine,
-    quantize_affine_float_zero_point,
     quantize_affine_no_zero_point,
+    quantize_affine_tinygemm,
 )
 from torchao.utils import (
     TORCH_VERSION_AT_LEAST_2_5,
@@ -439,7 +439,7 @@ def groupwise_affine_quantize_tensor_from_qparams(
     if zero_point_domain == ZeroPointDomain.INT:
         _quantize_affine = quantize_affine
     elif zero_point_domain == ZeroPointDomain.FLOAT:
-        _quantize_affine = quantize_affine_float_zero_point
+        _quantize_affine = quantize_affine_tinygemm
     elif ZeroPointDomain == ZeroPointDomain.NONE:
         _quantize_affine = quantize_affine_no_zero_point
     else:
@@ -508,7 +508,7 @@ def groupwise_affine_dequantize_tensor_from_qparams(
     if zero_point_domain == ZeroPointDomain.INT:
         _dequantize_affine = dequantize_affine
     elif zero_point_domain == ZeroPointDomain.FLOAT:
-        _dequantize_affine = dequantize_affine_float_zero_point
+        _dequantize_affine = dequantize_affine_tinygemm
     else:
         _dequantize_affine = dequantize_affine_no_zero_point
     return _dequantize_affine(