Update logic for blocking activations

Author: vineet-g

src/coreai_opt/quantization/_eager/_prepare_for_export.py

@@ -201,9 +201,7 @@ def _import_coreai_torch_modules():
             if is_float4_dtype(module.dtype):
                 raise ValueError("FP4 activation quantization is not supported for MLIR export.")
             if isinstance(module.granularity, PerBlockGranularity):
-                raise ValueError(
-                    "MLIR export does not support per-block granularity for activations."
-                )
+                raise ValueError("MLIR 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

@@ -340,7 +340,7 @@ def _process_mlir_activation_quantization(
         raise ValueError("FP4 activation quantization is not supported for MLIR export.")
 
     if isinstance(fake_quant_mod.granularity, PerBlockGranularity):
-        raise ValueError("MLIR export does not support per-block granularity for activations.")
+        raise ValueError("MLIR 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/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 = {

src/coreai_opt/quantization/spec/fake_quantize.py

@@ -359,7 +359,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 +384,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 +406,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 +427,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 +449,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
         )

src/coreai_opt/quantization/spec/granularity.py

@@ -12,6 +12,7 @@
 from pydantic import BaseModel, ConfigDict, Field, model_serializer
 
 from coreai_opt._utils.registry_utils import ConfigRegistryMixin as _ConfigRegistryMixin
+from coreai_opt.config.spec import CompressionTargetTensor as _CompressionTargetTensor
 from coreai_opt.quantization.spec.errors import _BlockSizeMismatchError
 
 
@@ -49,7 +50,11 @@ def _serialize_model(self) -> dict[str, Any]:
         return data
 
     @abstractmethod
-    def _get_block_size(self, block_sizes_list: list[int]) -> list[int]:
+    def _get_block_size(
+        self,
+        block_sizes_list: list[int],
+        quantization_target: _CompressionTargetTensor = _CompressionTargetTensor.WEIGHT,
+    ) -> list[int]:
         """
         Given an initial list of the tensor shape, return a list of block sizes
         corresponding to each axis:
@@ -61,6 +66,11 @@ def _get_block_size(self, block_sizes_list: list[int]) -> list[int]:
         - if per-block structuring is being done for a certain axis, set the block
           size for that specific axis
 
+        ``quantization_target`` distinguishes weight from activation tensors.
+        Only per-block granularity uses it (see
+        :meth:`PerBlockGranularity._handle_single_axis_block_size`); the other
+        granularities ignore it.
+
         Example:
             - ``[10, 5, 2]`` with per-channel structuring on axis 1 results in
               ``[10, 1, 2]``
@@ -70,11 +80,20 @@ def _get_block_size(self, block_sizes_list: list[int]) -> list[int]:
         """
         pass
 
-    def get_block_size(self, tensor_shape: torch.Size) -> tuple[int, ...]:
+    def get_block_size(
+        self,
+        tensor_shape: torch.Size,
+        quantization_target: _CompressionTargetTensor = _CompressionTargetTensor.WEIGHT,
+    ) -> tuple[int, ...]:
         """
         Get a list of block sizes based on the granularity.
+
+        Args:
+            tensor_shape: Shape of the tensor being quantized.
+            quantization_target: Whether the tensor is a weight or an activation.
+                Defaults to ``WEIGHT``, which preserves the historical behavior.
         """
-        return tuple(self._get_block_size(list(tensor_shape)))
+        return tuple(self._get_block_size(list(tensor_shape), quantization_target))
 
     # The axis resolution logic lives here because it is granularity-specific.
     # Currently only PerChannelGranularity has a meaningful axis to resolve, but
@@ -119,7 +138,11 @@ class PerTensorGranularity(QuantizationGranularity):
 
     axis: Literal[None] = None
 
-    def _get_block_size(self, block_sizes_list: list[int]) -> list[int]:
+    def _get_block_size(
+        self,
+        block_sizes_list: list[int],
+        quantization_target: _CompressionTargetTensor = _CompressionTargetTensor.WEIGHT,
+    ) -> list[int]:
         return block_sizes_list
 
 
@@ -138,7 +161,11 @@ class PerChannelGranularity(QuantizationGranularity):
 
     axis: int | None = None
 
-    def _get_block_size(self, block_sizes_list: list[int]) -> list[int]:
+    def _get_block_size(
+        self,
+        block_sizes_list: list[int],
+        quantization_target: _CompressionTargetTensor = _CompressionTargetTensor.WEIGHT,
+    ) -> list[int]:
 
         if self.axis is None:
             raise ValueError(
@@ -185,39 +212,63 @@ class PerBlockGranularity(QuantizationGranularity):
     ``Quantizer.prepare()`` automatically resolves the axis based on the module type
     for weight quantization.
 
+    Single-axis mode treats weights and activations differently. For weights only
+    the two leading channel axes participate in blocking, so trailing kernel
+    dimensions span a whole block. For activations every axis other than the block
+    axis gets its own scale.
+
     .. list-table::
        :header-rows: 1
 
-       * - Weight tensor shape (input)
+       * - Tensor shape (input)
+         - target
          - axis
          - block_size
-         - Weight shape of each block (output)
+         - Shape of each block (output)
        * - [C_out, C_in]
+         - weight
          - 1
          - 32
          - [1, 32]
        * - [C_out, C_in]
+         - weight
          - None
          - (4, 8)
          - [4, 8]
        * - [C_out, C_in, KH, KW]
+         - weight
          - 0
          - 16
          - [16, 1, KH, KW]
        * - [C_out, C_in, KH, KW]
+         - weight
          - None
          - (4, 16, 3, -1)
          - [4, 16, 3, KW]
+       * - [B, S, D]
+         - activation
+         - -1
+         - 16
+         - [1, 1, 16]
+       * - [B, C, H, W]
+         - activation
+         - 1
+         - 16
+         - [1, 16, 1, 1]
     """
 
     axis: int | None = None
     block_size: Annotated[int, Field(gt=0)] | tuple[Annotated[int, Field(gt=0)] | Literal[-1], ...]
 
-    def _get_block_size(self, block_sizes_list: list[int]) -> list[int]:
+    def _get_block_size(
+        self,
+        block_sizes_list: list[int],
+        quantization_target: _CompressionTargetTensor = _CompressionTargetTensor.WEIGHT,
+    ) -> list[int]:
         if isinstance(self.block_size, tuple):
             return self._handle_multi_axis_block_size(block_sizes_list)
         else:
-            return self._handle_single_axis_block_size(block_sizes_list)
+            return self._handle_single_axis_block_size(block_sizes_list, quantization_target)
 
     def _handle_multi_axis_block_size(self, block_sizes_list: list[int]) -> list[int]:
         """Handle blocking when self.block_size is a tuple"""
@@ -247,7 +298,11 @@ def _handle_multi_axis_block_size(self, block_sizes_list: list[int]) -> list[int
 
         return block_sizes_list
 
-    def _handle_single_axis_block_size(self, block_sizes_list: list[int]) -> list[int]:
+    def _handle_single_axis_block_size(
+        self,
+        block_sizes_list: list[int],
+        quantization_target: _CompressionTargetTensor = _CompressionTargetTensor.WEIGHT,
+    ) -> list[int]:
         """Handle blocking when self.block_size is an integer"""
         # TODO: Logic to be added where if self.axis is None,
         #  we can figure out the optimal axis for the user
@@ -272,12 +327,25 @@ def _handle_single_axis_block_size(self, block_sizes_list: list[int]) -> list[in
                 f"is not divisible by block size {self.block_size}"
             )
 
-        # Set the specified axis to block_size, and set the other channel axis
-        # (the one of {0, 1} that is not the block axis) to 1 so that it is
-        # quantized per-slice. Any remaining higher dimensions (index 2+) are
-        # left unchanged.
+        # How the non-block axes are treated depends on the quantization target,
+        #
+        # WEIGHT: only the two leading channel axes participate. The other
+        #   channel axis becomes 1 (one scale per slice) while any trailing
+        #   dimensions (index 2+, e.g. conv kernel dims) keep their full size so
+        #   each block spans the whole kernel.
+        #     [C_out, C_in, KH, KW], axis=0, block=16 -> [16, 1, KH, KW]
+        #
+        # ACTIVATION: blocking runs along a single axis and every other
+        #   dimension gets its own scale, so all non-block axes become 1.
+        #     [B, S, D],    axis=-1, block=16 -> [1, 1, 16]
+        #     [B, C, H, W], axis=1,  block=16 -> [1, 16, 1, 1]
+        if quantization_target == _CompressionTargetTensor.ACTIVATION:
+            collapse_upto = rank
+        else:
+            collapse_upto = 2
+
         block_sizes_list[axis] = self.block_size
-        for i, _ in enumerate(block_sizes_list[:2]):
+        for i, _ in enumerate(block_sizes_list[:collapse_upto]):
             if i != axis:
                 block_sizes_list[i] = 1
 

src/coreai_opt/quantization/spec/qparams_calculator.py

@@ -82,6 +82,11 @@ def granularity(self) -> QuantizationGranularity:
         """Getter for granularity."""
         return self._granularity
 
+    @property
+    def quantization_target(self):
+        """Whether the quantized tensor is a weight or an activation."""
+        return self.range_calculator.quantization_target
+
     @granularity.setter
     def granularity(self, granularity: QuantizationGranularity) -> None:
         """Update granularity for this calculator and its range calculator.
@@ -117,7 +122,9 @@ def _get_tensor_with_granularity_from_scalar(
         Return a tensor with dimensions equal to num blocks in each dimension, comprised
         of values equal to scalar.
         """
-        block_size_list = self.granularity.get_block_size(input_tensor.shape)
+        block_size_list = self.granularity.get_block_size(
+            input_tensor.shape, self.quantization_target
+        )
         num_blocks_list = [
             inp_size // block_size
             for inp_size, block_size in zip(input_tensor.shape, block_size_list, strict=True)
@@ -229,7 +236,7 @@ def _compute_scale_zero_point_minval(
             min_val=min_val,
             max_val=max_val,
             mapping_type=QuantizationScheme._to_mapping_type(self.qscheme),
-            block_size=self.granularity.get_block_size(tensor.shape),
+            block_size=self.granularity.get_block_size(tensor.shape, self.quantization_target),
             target_dtype=self.target_dtype,
             quant_min=self.quant_min,
             quant_max=self.quant_max,

src/coreai_opt/quantization/spec/range_calculator.py

@@ -10,6 +10,7 @@
 from torchao.quantization.quant_primitives import _get_reduction_params
 
 from coreai_opt._utils.registry_utils import ClassRegistryMixin as _ClassRegistryMixin
+from coreai_opt.config.spec import CompressionTargetTensor as _CompressionTargetTensor
 
 from .granularity import QuantizationGranularity
 
@@ -20,16 +21,22 @@ class RangeCalculatorBase(_ClassRegistryMixin, nn.Module):
     of a given tensor.
     """
 
-    def __init__(self, granularity: QuantizationGranularity, **kwargs):
+    def __init__(
+        self,
+        granularity: QuantizationGranularity,
+        quantization_target: _CompressionTargetTensor = _CompressionTargetTensor.WEIGHT,
+        **kwargs,
+    ):
         super().__init__()
         self.granularity = granularity
+        self.quantization_target = quantization_target
 
     def _reshape_min_max(self, range_tensor: torch.Tensor, input_shape: torch.Size):
         """
         Reshape range_tensor to have the same number of dimensions as input shape,
         taking block size into account.
         """
-        block_size_list = self.granularity.get_block_size(input_shape)
+        block_size_list = self.granularity.get_block_size(input_shape, self.quantization_target)
 
         # While reducing, each dimension with block size other than 1 or the original
         # dimension size will be split into 2 dimensions of num_blocks and block_size.
@@ -77,7 +84,7 @@ class MinMaxRangeCalculator(RangeCalculatorBase):
     """
 
     def _generate_min_max(self, tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        block_size_list = self.granularity.get_block_size(tensor.shape)
+        block_size_list = self.granularity.get_block_size(tensor.shape, self.quantization_target)
         shape_for_reduction, reduction_dims = _get_reduction_params(block_size_list, tensor.size())
 
         # If tensor is already the shape required, no minmaxing is needed.

tests/quantization/test_factory.py

@@ -38,7 +38,14 @@
 class TestQuantizationComponentFactory:
     """Test the QuantizationComponentFactory class"""
 
-    def test_create_range_calculator(self):
+    @pytest.mark.parametrize(
+        "quantization_target",
+        [
+            CompressionTargetTensor.WEIGHT,
+            CompressionTargetTensor.ACTIVATION,
+        ],
+    )
+    def test_create_range_calculator(self, quantization_target):
         """Test creating range calculator from spec"""
         spec = QuantizationSpec(
             dtype=torch.int8,
@@ -49,10 +56,11 @@ def test_create_range_calculator(self):
             range_calculator_cls=MinMaxRangeCalculator,
         )
 
-        range_calc = QuantizationComponentFactory.create_range_calculator(spec)
+        range_calc = QuantizationComponentFactory.create_range_calculator(spec, quantization_target)
 
         assert isinstance(range_calc, MinMaxRangeCalculator)
         assert range_calc.granularity == spec.granularity
+        assert range_calc.quantization_target == quantization_target
 
     @pytest.mark.parametrize(
         "range",