composite ops boundary quantization coverage

Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>

Author: pkmandke

tests/export/test_composite_op_externalize.py

@@ -20,18 +20,19 @@
 
 from coreai_opt import ExportBackend
 from coreai_opt.quantization import (
-    ModuleQuantizerConfig,
     Quantizer,
     QuantizerConfig,
 )
 from coreai_opt.quantization.spec import (
     PerTensorGranularity,
     QuantizationScheme,
     QuantizationSpec,
-    default_activation_quantization_spec,
-    default_weight_quantization_spec,
 )
-from tests.fixtures.quantization import make_graph_mode_ptq_config
+from tests.fixtures.quantization import (
+    COMPOSITE_BOUNDARY_ACT_DTYPE,
+    make_graph_mode_composite_boundary_config,
+    make_graph_mode_ptq_config,
+)
 from tests.models.composite import (
     CompositeRMSNormModel,
     CompositeRMSNormOnlyModel,
@@ -104,15 +105,10 @@ class TestCompositeOpIOQuantization:
     while every other quantized edge carries the global dtype.
     """
 
-    _COMPOSITE_ACT_DTYPE = torch.uint8
+    _COMPOSITE_ACT_DTYPE = COMPOSITE_BOUNDARY_ACT_DTYPE
 
     @classmethod
     def _composite_act_spec(cls) -> QuantizationSpec:
-        # The composite config must use a dtype DISTINCT from the global
-        # (default) activation dtype: a matching dtype collapses via observer
-        # sharing into a vacuous no-op, so the composite's effect at the
-        # boundary would not be observable.
-        assert cls._COMPOSITE_ACT_DTYPE != default_activation_quantization_spec().dtype
         return QuantizationSpec(
             dtype=cls._COMPOSITE_ACT_DTYPE,
             qscheme=QuantizationScheme.SYMMETRIC,
@@ -127,21 +123,11 @@ def _config(
         target_by: str,
         module_input_spec: dict | None = None,
     ) -> QuantizerConfig:
-        composite_act = cls._composite_act_spec()
-        global_config = ModuleQuantizerConfig(
-            op_state_spec={"weight": default_weight_quantization_spec()},
-            op_input_spec={"*": default_activation_quantization_spec()},
-            op_output_spec={"*": default_activation_quantization_spec()},
-        )
-        composite_config = ModuleQuantizerConfig(
-            module_input_spec=module_input_spec or {"*": composite_act},
-            module_output_spec={"*": composite_act},
+        return make_graph_mode_composite_boundary_config(
+            module_name=module_name if target_by == "name" else None,
+            module_type=spec.target_class if target_by == "type" else None,
+            module_input_spec=module_input_spec,
         )
-        if target_by == "name":
-            scope = {"module_name_configs": {module_name: composite_config}}
-        else:
-            scope = {"module_type_configs": {spec.target_class: composite_config}}
-        return QuantizerConfig(global_config=global_config, execution_mode="graph", **scope)
 
     def _finalize(
         self,

tests/export/test_graph_mode_mlir_export.py

@@ -31,6 +31,7 @@
 from tests.fixtures.fp8 import ParametrizedFP8Configs
 from tests.fixtures.quantization import (
     ParametrizedQuantConfigs,
+    make_graph_mode_composite_boundary_config,
     make_graph_mode_ptq_config,
 )
 from tests.models.composite import (
@@ -452,47 +453,66 @@ def test_integer_quant_minval_export(
 
 # Composite-op externalize export coverage
 
+# (model, externalize spec, composite submodule path, expected coreai.quantize count
+# per config kind)
+_EXTERNALIZE_EXPORT_CASES = [
+    pytest.param(
+        CompositeRMSNormModel,
+        rmsnorm_externalize_spec(),
+        "norm",
+        {"w8": 0, "w8a8": 4, "w8a8-boundary": 6},
+        id="rmsnorm",
+    ),
+    pytest.param(
+        CompositeSDPAModel,
+        sdpa_externalize_spec(),
+        "composite",
+        {"w8": 0, "w8a8": 4, "w8a8-boundary": 8},
+        id="sdpa",
+    ),
+]
 
+
+@pytest.mark.parametrize("config_kind", ["w8", "w8a8", "w8a8-boundary"])
 @pytest.mark.parametrize(
-    "quantize_activations, expected_quantize_count",
-    [
-        pytest.param(False, 0, id="w8-weight-only"),
-        pytest.param(True, 4, id="w8a8"),
-    ],
-)
-@pytest.mark.parametrize(
-    "model_cls, externalize_spec",
-    [
-        pytest.param(CompositeRMSNormModel, rmsnorm_externalize_spec(), id="rmsnorm"),
-        pytest.param(CompositeSDPAModel, sdpa_externalize_spec(), id="sdpa"),
-    ],
+    "model_cls, externalize_spec, composite_module, expected_quantize_counts",
+    _EXTERNALIZE_EXPORT_CASES,
 )
 def test_composite_externalize_export(
     model_cls: type[torch.nn.Module],
     externalize_spec: ExternalizeSpec,
-    quantize_activations: bool,
-    expected_quantize_count: int,
+    composite_module: str,
+    expected_quantize_counts: Mapping[str, int],
+    config_kind: str,
 ) -> None:
     """End-to-end CoreAI export of a model with an externalized composite op.
 
-    Marks the composite op for externalization, runs graph-mode PTQ
-    (w8 weight-only or w8a8), then lowers the finalized graph to a
-    .aimodel and runs it. ``convert_and_verify`` handles SNR / PSNR
-    on the runtime output and op-count verification on the exported
-    program (``constexpr_blockwise_shift_scale`` for weight quantizers
-    and ``quantize`` / ``dequantize`` for activation quantizers).
+    Marks the composite op for externalization, runs graph-mode PTQ, then lowers the
+    finalized graph to a .aimodel and runs it. ``convert_and_verify`` handles SNR /
+    PSNR on the runtime output and op-count verification on the exported program
+    (``constexpr_blockwise_shift_scale`` for weight quantizers and ``quantize`` /
+    ``dequantize`` for activation quantizers).
+
+    Three configs:
+
+    - ``w8`` / ``w8a8``: global config only.
+    - ``w8a8-boundary``: adds a module-scoped ``module_input_spec`` /
+      ``module_output_spec`` so the composite op's i/o edges are quantized
     """
     model = model_cls().eval().half()
     input_data = torch.randn(2, 4, 32, dtype=torch.float16)
 
     _patch_model_for_externalization(model, [externalize_spec])
 
-    quantizer = Quantizer(
-        model, make_graph_mode_ptq_config(quantize_activations=quantize_activations)
-    )
+    if config_kind == "w8a8-boundary":
+        config = make_graph_mode_composite_boundary_config(module_name=composite_module)
+    else:
+        config = make_graph_mode_ptq_config(quantize_activations=config_kind == "w8a8")
+
+    quantizer = Quantizer(model, config)
     prepared_model = quantizer.prepare((input_data,))
 
-    if quantize_activations:
+    if config_kind != "w8":
         with quantizer.calibration_mode(), torch.no_grad():
             prepared_model(input_data)
 
@@ -501,6 +521,7 @@ def test_composite_externalize_export(
 
     finalized_model = quantizer.finalize(backend=ExportBackend.CoreAI)
 
+    expected_quantize_count = expected_quantize_counts[config_kind]
     export_utils.convert_and_verify(
         finalized_model=finalized_model,
         input_data=input_data,

tests/fixtures/quantization.py

@@ -104,6 +104,61 @@ def make_graph_mode_ptq_config(*, quantize_activations: bool) -> QuantizerConfig
     )
 
 
+COMPOSITE_BOUNDARY_ACT_DTYPE = torch.uint8
+
+
+def make_graph_mode_composite_boundary_config(
+    *,
+    module_name: str | None = None,
+    module_type: type | None = None,
+    module_input_spec: dict | None = None,
+) -> QuantizerConfig:
+    """Build a graph-mode w8a8 config that also quantizes the composite op module's
+    own i/o boundary.
+
+    Args:
+        module_name: Target the module at this path (``module_name_configs``).
+        module_type: Target modules of this type (``module_type_configs``).
+            Exactly one of module_name / module_type must be given.
+        module_input_spec: Override the boundary input spec, e.g.
+            ``{0: spec, 2: spec}`` to select individual positional args.
+            Defaults to the ``"*"`` wildcard over every boundary input.
+
+    Returns:
+        QuantizerConfig: global w8a8 plus a module-scoped boundary spec.
+    """
+    if (module_name is None) == (module_type is None):
+        msg = "pass exactly one of module_name / module_type"
+        raise ValueError(msg)
+    assert COMPOSITE_BOUNDARY_ACT_DTYPE != default_activation_quantization_spec().dtype
+
+    def _boundary_spec() -> QuantizationSpec:
+        return QuantizationSpec(
+            dtype=COMPOSITE_BOUNDARY_ACT_DTYPE,
+            qscheme=QuantizationScheme.SYMMETRIC,
+            granularity=PerTensorGranularity(),
+        )
+
+    boundary_config = ModuleQuantizerConfig(
+        module_input_spec=module_input_spec or {"*": _boundary_spec()},
+        module_output_spec={"*": _boundary_spec()},
+    )
+    scope = (
+        {"module_name_configs": {module_name: boundary_config}}
+        if module_name is not None
+        else {"module_type_configs": {module_type: boundary_config}}
+    )
+    return QuantizerConfig(
+        global_config=ModuleQuantizerConfig(
+            op_state_spec={"weight": default_weight_quantization_spec()},
+            op_input_spec={"*": default_activation_quantization_spec()},
+            op_output_spec={"*": default_activation_quantization_spec()},
+        ),
+        execution_mode="graph",
+        **scope,
+    )
+
+
 @dataclass
 class ParametrizedQuantConfigs:
     """Container for parametrized Eager and PT2E quantization configs.