Update Int8DynActOnlyConfig

Author: lisjin

torchao/quantization/__init__.py

@@ -54,6 +54,7 @@
     GemliteUIntXWeightOnlyConfig,
     Int4DynamicActivationInt4WeightConfig,
     Int4WeightOnlyConfig,
+    Int8DynActOnlyConfig,
     Int8DynamicActivationInt4WeightConfig,
     Int8DynamicActivationInt8WeightConfig,
     Int8DynamicActivationIntxWeightConfig,
@@ -144,6 +145,7 @@
     "Int8DynamicActivationIntxWeightConfig",
     "Int4WeightOnlyConfig",
     "Float8DynamicActivationInt4WeightConfig",
+    "Int8DynActOnlyConfig",
     "Int8WeightOnlyConfig",
     "Float8WeightOnlyConfig",
     "Float8DynamicActivationFloat8WeightConfig",

torchao/quantization/qat/fake_quantize_config.py

@@ -292,6 +292,7 @@ def _infer_fake_quantize_configs(
     # avoid circular imports
     from torchao.quantization import (
         Int4WeightOnlyConfig,
+        Int8DynActOnlyConfig,
         Int8DynamicActivationInt4WeightConfig,
     )
 
@@ -315,5 +316,12 @@ def _infer_fake_quantize_configs(
             zero_point_domain=base_config.zero_point_domain,
         )
         return (None, weight_config)
+    elif isinstance(base_config, Int8DynActOnlyConfig):
+        act_config = IntxFakeQuantizeConfig(
+            dtype=torch.int8,
+            granularity="per_token",
+            is_symmetric=base_config.is_symmetric,
+        )
+        return (act_config, None)
     else:
         raise ValueError("Unexpected base config: %s" % base_config)

torchao/quantization/quant_api.py

@@ -148,6 +148,7 @@
     "gemlite_uintx_weight_only",
     "float8_dynamic_activation_float8_weight",
     "float8_static_activation_float8_weight",
+    "Int8DynActOnlyConfig",
     "Int8DynActInt4WeightQuantizer",
     "Float8DynamicActivationFloat8SemiSparseWeightConfig",
     "ModuleFqnToConfig",
@@ -1312,6 +1313,31 @@ def _float8_cutlass_quant_sparse(
     )
 
 
+@dataclass
+class Int8DynActOnlyConfig(AOBaseConfig):
+    """
+    Configuration for applying int8 dynamic symmetric per-token activation quantization to linear layers.
+    Args:
+        is_symmetric: bool = False - Whether to use symmetric quantization for activations.
+    """
+
+    is_symmetric: bool = False
+
+
+@register_quantize_module_handler(Int8DynActOnlyConfig)
+def _int8_dynamic_activation_transform(
+    module: torch.nn.Module, config: Int8DynActOnlyConfig
+) -> torch.nn.Module:
+    weight = module.weight
+    if config.is_symmetric == MappingType.SYMMETRIC:
+        input_quant_func = _int8_symm_per_token_reduced_range_quant
+    else:
+        input_quant_func = _int8_asymm_per_token_quant
+    weight = to_linear_activation_quantized(weight, input_quant_func)
+    module.weight = torch.nn.Parameter(weight, requires_grad=False)
+    return module
+
+
 @dataclass
 class Int8DynamicActivationInt8WeightConfig(AOBaseConfig):
     """