Block activation support (#56)

* Block activation support

* Update logic for blocking activations

* add logic to preserve per-block granularity for shared observers

* address review comments

* fix failing tests

Author: vineet-g

src/coreai_opt/quantization/_eager/_prepare_for_export.py

@@ -28,6 +28,7 @@
     validate_fp4_export,
 )
 from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase
+from coreai_opt.quantization.spec.granularity import PerBlockGranularity
 
 logger = logging.getLogger(__name__)
 
@@ -198,7 +199,11 @@ def _import_coreai_torch_modules():
             CompressionTargetTensor.ACTIVATION,
         ):
             if is_float4_dtype(module.dtype):
-                raise ValueError("FP4 activation quantization is not supported for MLIR export.")
+                raise ValueError("Core AI export does not support FP4 activation quantization.")
+            if isinstance(module.granularity, PerBlockGranularity):
+                raise ValueError(
+                    "Core AI export does not support PerBlockGranularity on activations."
+                )
             modules_to_replace.append((name, module))
 
     # Replace each FakeQuantizeImplBase module

src/coreai_opt/quantization/_graph/_prepare_for_export.py

@@ -38,6 +38,7 @@
     resolve_attr,
 )
 from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase
+from coreai_opt.quantization.spec.granularity import PerBlockGranularity
 
 logger = logging.getLogger(__name__)
 
@@ -332,7 +333,10 @@ def _process_mlir_activation_quantization(
         fake_quant_mod: The fake quantization module
     """
     if is_float4_dtype(fake_quant_mod.dtype):
-        raise ValueError("FP4 activation quantization is not supported for MLIR export.")
+        raise ValueError("Core AI export does not support FP4 activation quantization.")
+
+    if isinstance(fake_quant_mod.granularity, PerBlockGranularity):
+        raise ValueError("Core AI export does not support PerBlockGranularity on activations.")
 
     def _import_coreai_custom_ops():
         import coreai_torch._compression.custom_layers  # noqa: PLC0415, F401

src/coreai_opt/quantization/_graph/_utils.py

@@ -11,6 +11,7 @@
 
 from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase
 from coreai_opt.quantization.spec.granularity import (
+    PerBlockGranularity,
     PerChannelGranularity,
     PerTensorGranularity,
     QuantizationGranularity,
@@ -267,8 +268,8 @@ def _shared_granularity_axis_is_safe(
     """Return True only if it's proven safe to keep ``fake_quant``'s
     granularity shared across ``op_node``; unproven cases default to unsafe.
 
-    Per-channel activation quantization on a shared observer is assumed
-    unsafe by default. Each op category below has exactly one condition
+    Per-channel and per-block activation quantization on a shared observer are
+    assumed unsafe by default. Each op category below has exactly one condition
     under which it stops being safe, checked directly against that
     category rather than composing generic checks that apply to every op:
 
@@ -284,10 +285,6 @@ def _shared_granularity_axis_is_safe(
     # No axis to violate — trivially safe.
     if isinstance(granularity, PerTensorGranularity):
         return True
-    # Anything other than PerChannelGranularity (e.g. PerBlockGranularity)
-    # has no condition checked below, so it's not safe.
-    if not isinstance(granularity, PerChannelGranularity):
-        return False
 
     output_shape = op_node.meta["val"].shape
     input_shape = input_fq_node.all_input_nodes[0].meta["val"].shape
@@ -296,14 +293,19 @@ def _shared_granularity_axis_is_safe(
     # index, so there's no condition to check here.
     if len(input_shape) != len(output_shape):
         return False
-    axis = QuantizationGranularity._resolve_axis(granularity, len(input_shape))
-    if axis is None:
-        return False
+
+    if isinstance(granularity, PerBlockGranularity):
+        axes: tuple[int, ...] = tuple(range(len(input_shape)))
+    else:
+        axis = QuantizationGranularity._resolve_axis(granularity, len(input_shape))
+        if axis is None:
+            return False
+        axes = (axis,)
 
     if op_node.target in _AXIS_RESIZING_ATEN_OPS:
-        return input_shape[axis] == output_shape[axis]
+        return all(input_shape[a] == output_shape[a] for a in axes)
     if op_node.target in _AXIS_REORDERING_ATEN_OPS:
-        return _op_preserves_axis_identity(op_node, axis)
+        return all(_op_preserves_axis_identity(op_node, a) for a in axes)
     # flatten/reshape/view/unsqueeze, or an unrecognized future op: no known
     # single condition to prove safety, so default to unsafe.
     return False

src/coreai_opt/quantization/spec/factory.py

@@ -35,19 +35,25 @@ class QuantizationComponentFactory(CompressionComponentFactoryBase):
     """
 
     @classmethod
-    def create_range_calculator(cls, spec: QuantizationSpec) -> RangeCalculatorBase:
+    def create_range_calculator(
+        cls,
+        spec: QuantizationSpec,
+        quantization_target: CompressionTargetTensor = CompressionTargetTensor.WEIGHT,
+    ) -> RangeCalculatorBase:
         """
         Create a RangeCalculatorBase instance from a QuantizationSpec.
 
         Args:
             spec: QuantizationSpec instance containing configuration
+            quantization_target: The target tensor for quantization (weight/activation).
 
         Returns:
             RangeCalculatorBase instance configured from the spec
         """
         # Standard arguments for range calculator
         common_args = {
             "granularity": spec.granularity,
+            "quantization_target": quantization_target,
         }
 
         # Automatically detect and include any extra arguments
@@ -96,7 +102,7 @@ def create_qparams_calculator(
             )
 
         # Create range calculator first
-        range_calculator = cls.create_range_calculator(spec)
+        range_calculator = cls.create_range_calculator(spec, quantization_target)
 
         # Standard arguments for qparams calculator
         common_args = {
@@ -202,7 +208,6 @@ def create_fake_quantizer(
             "quant_min": spec.quant_min,
             "quant_max": spec.quant_max,
             "qparams_calculator": qparams_calculator,
-            "quantization_target": quantization_target,
             "n_bits": spec.n_bits,
         }
 
@@ -283,7 +288,6 @@ def create_fake_quantizer_partial(
             "target_dtype": spec.target_dtype,
             "quant_min": spec.quant_min,
             "quant_max": spec.quant_max,
-            "quantization_target": quantization_target,
             "n_bits": spec.n_bits,
         }
 

src/coreai_opt/quantization/spec/fake_quantize.py

@@ -53,7 +53,6 @@ def __init__(
         quant_min: int | float,
         quant_max: int | float,
         qparams_calculator: QParamsCalculatorBase,
-        quantization_target: CompressionTargetTensor,
         n_bits: int | None = None,
         **kwargs,
     ):
@@ -65,7 +64,6 @@ def __init__(
         self.quant_min = quant_min
         self.quant_max = quant_max
         self.qparams_calculator = qparams_calculator
-        self.quantization_target = quantization_target
         self.register_buffer("_disabled", torch.tensor(False))
 
         # Infer n_bits from dtype if not provided
@@ -78,6 +76,11 @@ def qscheme(self) -> QuantizationScheme:
         """The quantization scheme, delegated to the qparams_calculator."""
         return self.qparams_calculator.qscheme
 
+    @property
+    def quantization_target(self) -> CompressionTargetTensor:
+        """Getter for quantization target."""
+        return self.qparams_calculator.quantization_target
+
     @property
     def granularity(self) -> QuantizationGranularity:
         """Getter for granularity."""
@@ -359,7 +362,7 @@ def _quantize_int(
 
         This function quantizes the values in tensor but keeps the quantized tensor dtype in FP.
         """
-        block_size = self.granularity.get_block_size(tensor.shape)
+        block_size = self.granularity.get_block_size(tensor.shape, self.quantization_target)
         original_shape, blockwise_shape, reduced_shape = _get_quantization_shapes(
             tensor, block_size
         )
@@ -384,7 +387,7 @@ def _dequantize_int(
         output_dtype: torch.dtype,
     ) -> torch.Tensor:
         """Integer dequantization. See :func:`_dequantize_int` for the math."""
-        block_size = self.granularity.get_block_size(tensor.shape)
+        block_size = self.granularity.get_block_size(tensor.shape, self.quantization_target)
         original_shape, blockwise_shape, reduced_shape = _get_quantization_shapes(
             tensor, block_size
         )
@@ -406,7 +409,7 @@ def _quantize_float(
         """
         Floating-point quantization: cast_to_low_precision(clamp(input / scale, min, max))
         """
-        block_size = self.granularity.get_block_size(tensor.shape)
+        block_size = self.granularity.get_block_size(tensor.shape, self.quantization_target)
         original_shape, blockwise_shape, reduced_shape = _get_quantization_shapes(
             tensor, block_size
         )
@@ -427,7 +430,7 @@ def _dequantize_float(
         output_dtype: torch.dtype,
     ) -> torch.Tensor:
         """Floating-point dequantization: input * scale"""
-        block_size = self.granularity.get_block_size(tensor.shape)
+        block_size = self.granularity.get_block_size(tensor.shape, self.quantization_target)
         original_shape, blockwise_shape, reduced_shape = _get_quantization_shapes(
             tensor, block_size
         )
@@ -449,7 +452,7 @@ def _fused_fake_quant_dequant(
 
         Dispatches to the int or float fused STE class based on self.dtype.
         """
-        block_size = self.granularity.get_block_size(tensor.shape)
+        block_size = self.granularity.get_block_size(tensor.shape, self.quantization_target)
         original_shape, blockwise_shape, reduced_shape = _get_quantization_shapes(
             tensor, block_size
         )