feat(quantization): support int8 with float8_e8m0fnu scales

changelog.d/133.added

@@ -0,0 +1 @@
+Add MXINT8 support.

src/coreai_opt/_utils/torch_utils.py

@@ -24,6 +24,7 @@
 # Mapping from dtype to the largest power-of-2 component of its max value.
 # Used by e8m0 scale computation (OCP Microscaling spec, FLOOR mode).
 E8M0_TARGET_MAX_POW2: dict[torch.dtype, int] = {
+    torch.int8: 6,  # max = 127 = 1.98 * 2^6
     torch.float4_e2m1fn_x2: 2,
     torch.float8_e4m3fn: 8,  # max = 448.0 = 1.75 * 2^8
     torch.float8_e5m2: 15,  # max = 57344.0 = 1.75 * 2^15

src/coreai_opt/palettization/spec/spec.py

@@ -132,6 +132,13 @@ def validate_lut_qspec(self) -> "PalettizationSpec":
                 "Use lut_qspec.qformulation=ZP instead."
             )
 
+        if not self.lut_qspec.dtype.is_floating_point and self.lut_qspec.scale_dtype is not None:
+            msg = (
+                "lut_qspec.scale_dtype must be None for integer lut_qspec.dtype, "
+                f"got scale_dtype={self.lut_qspec.scale_dtype}"
+            )
+            raise ValueError(msg)
+
         return self
 
     def model_dump_preserve_objects(self) -> dict[str, Any]:

src/coreai_opt/quantization/_export_utils.py

@@ -371,20 +371,18 @@ def get_activation_export_handler(
     return _GRAPH_ACTIVATION_EXPORT_HANDLERS.get(type(granularity))
 
 
-def can_export_stateless_fake_quant(
+def has_activation_export_handler(
     fake_quant_mod: FakeQuantizeImplBase,
     backend: ExportBackend,
     execution_mode: ExecutionMode,
 ) -> bool:
-    """Return whether a recompute-every-forward calculator can still be exported.
+    """Return whether a registered activation handler will export this module.
 
-    The built-in export paths bake qparams into buffers, so a calculator that
-    recomputes them every forward has nothing to bake. A registered activation handler
-    emits its own ops and can express the recomputation, so it lifts the
-    restriction for the activation it covers.
+    A handler emits its own ops, so callers use this to lift restrictions of the
+    built-in export paths.
 
     Args:
-        fake_quant_mod (FakeQuantizeImplBase): The module whose calculator is stateless.
+        fake_quant_mod (FakeQuantizeImplBase): The module to export.
         backend (ExportBackend): The export target.
         execution_mode (ExecutionMode): The quantizer's execution mode.
 

src/coreai_opt/quantization/_graph/_provisional_qspec_generation.py

@@ -280,10 +280,11 @@ def _apply_op_intrinsic_override(
 
     Nothing else: the op has no opinion on the other fields.
 
-    The ``QSCHEME`` half is skipped for a floating-point dtype, which admits only
-    the symmetric scheme — see ``QuantizationSpec.validate_qscheme_for_fp_quant``.
+    The ``QSCHEME`` half is skipped for a floating-point dtype or an e8m0 scale,
+    which admit only the symmetric scheme — see
+    ``QuantizationSpec.validate_qscheme_for_fp_quant``.
     An op like ``relu`` proposes an asymmetric scheme to exploit its one-sided
-    range, but that is a representational choice the dtype forbids, so proposing it
+    range, but that is a representational choice the spec forbids, so proposing it
     would make the group unbuildable rather than merely lower quality. The
     ``FLOAT_RANGE`` half still applies: the bound is a fact about the data and
     holds whatever the dtype.
@@ -292,7 +293,7 @@ def _apply_op_intrinsic_override(
     assert intrinsic is not None, "caller must gate on _get_op_intrinsic"
     scheme, float_range = intrinsic
 
-    if not user_spec.dtype.is_floating_point:
+    if not user_spec.dtype.is_floating_point and user_spec.scale_dtype is None:
         _override_field(
             qspecs,
             output_slot,

src/coreai_opt/quantization/_graph/_qspec_resolution.py

@@ -26,10 +26,11 @@
 )
 from torchao.quantization.pt2e.quantizer.quantizer import Q_ANNOTATION_KEY
 
-from coreai_opt.quantization.spec import QuantizationSpec
+from coreai_opt.quantization.spec import QuantizationScheme, QuantizationSpec
 
 from ._annotation_config import AnnotationConfig
 from ._qspec_types import (
+    OP_INTRINSIC_PRIORITY,
     FieldName,
     NodeSlot,
     ProvisionalQSpec,
@@ -207,12 +208,15 @@ def _build_concrete_spec(qspec: ProvisionalQSpec) -> TorchAOQuantizationSpec | N
             f"{sorted(f.name for f in qspec.fields)}."
         )
 
-    spec = QuantizationSpec(
-        **{
-            kwarg: qspec.fields[field_name].value
-            for field_name, kwarg in _SPEC_KWARG_FROM_FIELD.items()
-        }
-    )
+    kwargs = {
+        kwarg: qspec.fields[field_name].value
+        for field_name, kwarg in _SPEC_KWARG_FROM_FIELD.items()
+    }
+    if qspec.fields[FieldName.QSCHEME].priority == OP_INTRINSIC_PRIORITY and (
+        kwargs["dtype"].is_floating_point or kwargs["scale_dtype"] is not None
+    ):
+        kwargs["qscheme"] = QuantizationScheme.SYMMETRIC
+    spec = QuantizationSpec(**kwargs)
     return _convert_to_pt2e_spec(spec, qspec.fields[FieldName.QUANTIZATION_TARGET].value)
 
 

src/coreai_opt/quantization/_graph/quantizer.py

@@ -1055,7 +1055,10 @@ def prepare(
         try:
             prepared_model = prepare_qat_pt2e(exported_model, quantizer)
         except Exception as e:
-            raise type(e)(f"prepare_qat_pt2e call failed, with error: {e}") from e
+            # Not type(e)(msg): exceptions such as pydantic's ValidationError take
+            # other constructor arguments, and rebuilding one raises a TypeError.
+            e.add_note("prepare_qat_pt2e call failed")
+            raise
         if torchao_requires_empty_kwargs:
             restore_kwargs(prepared_model.graph, saved_kwargs)
 

src/coreai_opt/quantization/quantizer.py

@@ -6,10 +6,12 @@
 from __future__ import annotations
 
 import warnings
+from collections.abc import Callable
 from contextlib import contextmanager
 from os import PathLike
 from typing import Any
 
+import torch
 import torch.nn as nn
 from torch import fx
 from torchao.quantization.pt2e import (
@@ -27,7 +29,7 @@
 from coreai_opt.common import ExportBackend
 from coreai_opt.quantization._eager import EagerQuantizer as _EagerQuantizer
 from coreai_opt.quantization._export_utils import (
-    can_export_stateless_fake_quant as _can_export_stateless_fake_quant,
+    has_activation_export_handler as _has_activation_export_handler,
 )
 from coreai_opt.quantization._graph import GraphQuantizer as _GraphQuantizer
 from coreai_opt.quantization.base_quantizer import _BaseQuantizer
@@ -417,16 +419,9 @@ def _validate_no_persistent_observer_calculators(
         ``EagerActivationExportHandler`` in eager mode.
 
         """
-        if backend == ExportBackend._TORCH:
-            return
-        model_to_check = model if model is not None else self._model
-        stateless_fq_names = [
-            name
-            for name, mod in model_to_check.named_modules()
-            if isinstance(mod, FakeQuantizeImplBase)
-            and mod.is_stateless
-            and not _can_export_stateless_fake_quant(mod, backend, self._execution_mode)
-        ]
+        stateless_fq_names = self._fake_quant_names_without_export_handler(
+            model, backend, lambda mod: mod.is_stateless
+        )
         if stateless_fq_names:
             raise NotImplementedError(
                 f"backend={backend} does not yet support qparams calculators that "
@@ -435,6 +430,54 @@ def _validate_no_persistent_observer_calculators(
                 f"backend=ExportBackend._TORCH for torch-only inference."
             )
 
+    def _validate_no_integer_e8m0_scales(
+        self,
+        model: nn.Module | fx.GraphModule | None,
+        backend: ExportBackend,
+    ) -> None:
+        """Reject CoreAI/CoreML export of integer quantization with e8m0 scales.
+
+        The built-in integer export ops cannot take an e8m0 scale; a registered
+        activation export handler lifts the restriction.
+        """
+        e8m0_int_fq_names = self._fake_quant_names_without_export_handler(
+            model,
+            backend,
+            lambda mod: (
+                not mod.dtype.is_floating_point
+                and mod.qparams_calculator.scale_dtype == torch.float8_e8m0fnu
+            ),
+        )
+        if e8m0_int_fq_names:
+            msg = (
+                f"backend={backend} does not support integer quantization with "
+                f"scale_dtype={torch.float8_e8m0fnu}. Affected FakeQuantize modules: "
+                f"{e8m0_int_fq_names}. Use backend=ExportBackend._TORCH for torch-only "
+                f"inference."
+            )
+            raise NotImplementedError(msg)
+
+    def _fake_quant_names_without_export_handler(
+        self,
+        model: nn.Module | fx.GraphModule | None,
+        backend: ExportBackend,
+        predicate: Callable[[FakeQuantizeImplBase], bool],
+    ) -> list[str]:
+        """Return FakeQuantize modules matching ``predicate`` with no activation export handler.
+
+        Always empty for ``ExportBackend._TORCH``.
+        """
+        if backend == ExportBackend._TORCH:
+            return []
+        model_to_check = model if model is not None else self._model
+        return [
+            name
+            for name, mod in model_to_check.named_modules()
+            if isinstance(mod, FakeQuantizeImplBase)
+            and predicate(mod)
+            and not _has_activation_export_handler(mod, backend, self._execution_mode)
+        ]
+
     def finalize(
         self,
         model: nn.Module | fx.GraphModule | None = None,
@@ -484,6 +527,7 @@ def finalize(
         """
         self._validate_mmap_dir_constraints(model, backend, mmap_dir)
         self._validate_no_persistent_observer_calculators(model, backend)
+        self._validate_no_integer_e8m0_scales(model, backend)
         return self._quantizer.finalize(model, backend, mmap_dir=mmap_dir)
 
     @contextmanager

src/coreai_opt/quantization/spec/qparams_calculator.py

@@ -221,12 +221,16 @@ def _compute_scale_zero_point_minval(
             scale = 2^(floor(log2(max_abs)) - target_max_pow2)
             where ``target_max_pow2`` comes from ``E8M0_TARGET_MAX_POW2``, which
             is the single source of truth for which dtypes support an e8m0 scale.
+            The spec pins this mode to symmetric ZP, so the zero point is zero.
         """
 
         # e8m0 path: power-of-2 scales
         if self.scale_dtype == torch.float8_e8m0fnu:
             max_abs = torch.maximum(torch.abs(min_val), torch.abs(max_val))
-            return self._compute_e8m0_scale(max_abs), None, None
+            scale = self._compute_e8m0_scale(max_abs)
+            if self.dtype.is_floating_point:
+                return scale, None, None
+            return scale, torch.zeros_like(scale, dtype=torch.int32), -max_abs
 
         # Default path: torchao handles integer and FP8 dtypes
         scale, zero_point = choose_qparams_affine_with_min_max(