address nits

Author: jcaip

torchao/quantization/quant_api.py

@@ -1526,7 +1526,7 @@ class Int8DynamicActivationInt8WeightConfig(AOBaseConfig):
     layout: Optional[Layout] = PlainLayout()
     act_mapping_type: Optional[MappingType] = MappingType.SYMMETRIC
     weight_only_decode: bool = False
-    granularity: Union[PerRow, PerTensor] = PerRow()
+    granularity: Granularity = PerRow()
     set_inductor_config: bool = True
     version: int = 1
 

torchao/quantization/quantize_/common/quantize_tensor_kwargs.py

@@ -59,7 +59,7 @@ def _choose_quant_func_and_quantize_tensor(
         return Int8Tensor.from_hp(
             tensor,
             quant_kwargs.granularity,
-            quant_kwargs.mapping_type,
+            mapping_type=quant_kwargs.mapping_type,
         )
 
     raise NotImplementedError(f"Quant kwargs not supported: {quant_kwargs}")

torchao/quantization/quantize_/workflows/int8/int8_tensor.py

@@ -12,7 +12,7 @@
 
 from torchao.float8.inference import _slice_scale_for_dimension
 from torchao.kernel import int_scaled_matmul
-from torchao.quantization.granularity import Granularity, PerRow
+from torchao.quantization.granularity import Granularity
 from torchao.quantization.quant_primitives import (
     MappingType,
     choose_qparams_affine,
@@ -40,7 +40,7 @@ class QuantizeTensorToInt8Kwargs(QuantizeTensorKwargs):
         mapping_type: whether to use symmetric or asymmetric quant, only symmetric is supported currently
     """
 
-    granularity: Granularity = PerRow()
+    granularity: Granularity
     mapping_type: MappingType = MappingType.SYMMETRIC
 
 
@@ -113,7 +113,7 @@ def __repr__(self):
     def from_hp(
         cls,
         hp_tensor: torch.Tensor,
-        granularity: Granularity = PerRow(),
+        granularity: Granularity,
         act_quant_kwargs: Optional[QuantizeTensorToInt8Kwargs] = None,
         mapping_type=MappingType.SYMMETRIC,
     ):