Refactor macOS export logic | Retrieve reference inputs / names / dynamic shapes from model class.

python/src/coreai_models/export/_constants.py → python/src/coreai_models/_constants.py

@@ -3,15 +3,21 @@
 # Use of this source code is governed by a BSD-3-clause license that can
 # be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
 
-"""Constants for the export pipeline."""
+"""Graph and runner contract constants.
+
+A leaf module: imported by both ``models/`` and ``export/``, imports nothing from
+either.
+"""
+
+# Graph name for a single-graph (macOS) export. iOS uses its own entrypoint names.
+MAIN_GRAPH_NAME = "main"
 
 # KV cache names used by the Swift runner
 KEY_CACHE_NAME = "keyCache"
 VALUE_CACHE_NAME = "valueCache"
 
-# Trace-time KV cache sequence length. Used only for export/quantization tracing
-# to bound peak memory; at inference the actual cache size is determined
-# dynamically.
+# Trace-time KV cache sequence length, to bound peak trace memory. At inference the
+# cache size is dynamic.
 TRACE_KV_CACHE_SEQ_LEN = 2048
 
 # Trace-time `input_ids` length and `position_ids` offset for export/quantization

python/src/coreai_models/export/compression.py

@@ -11,15 +11,18 @@
 """
 
 import logging
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 
 import torch
 import torch.nn as nn
 
-from coreai_models.export._constants import (
+from coreai_models._constants import (
+    MAIN_GRAPH_NAME,
     QUANT_TRACE_OFFSET,
     QUANT_TRACE_QUERY_LEN,
+    TRACE_KV_CACHE_SEQ_LEN,
 )
+from coreai_models.models.base import BaseForCausalLM, TraceSpec
 
 logger = logging.getLogger(__name__)
 
@@ -100,6 +103,8 @@ def quantize_pytorch_model(
     inputs: tuple,
     dynamic_shapes: dict,
     quantization_config: dict,
+    cache_seq_len: int,
+    state_indices: Sequence[int],
     calibration_data_fn: Callable[[], list] | None = None,
     export_backend: object | None = None,
     mmap_dir: str | None = None,
@@ -119,6 +124,10 @@ def quantize_pytorch_model(
             coreai-opt expects under `quantization_config`. Includes a
             `calibrate_activations` key (popped here before constructing the
             coreai-opt config).
+        cache_seq_len: Sequence-dim length the caches in ``inputs`` were traced at,
+            used to bound the calibration query length.
+        state_indices: Positions in ``inputs`` that are state and must be reset
+            between calibration samples.
         calibration_data_fn: Optional function that returns calibration data samples.
             Required when calibrate_activations is enabled.
         export_backend: Backend for the finalized quantized model.
@@ -142,7 +151,11 @@ def quantize_pytorch_model(
 
     # When doing activation quantization, run real calibration data through the
     # prepared model so the activation observers see representative ranges.
-    # `inputs` follows the model forward contract: (input_ids, position_ids, k_cache, v_cache).
+    #
+    # `inputs[0]` must be input_ids and `inputs[1]` position_ids -- token-based
+    # calibration cannot do anything else. Which of the rest are state comes from
+    # `state_indices`; a non-state input keeps its traced tensor, which is only correct
+    # if its shape does not depend on query length.
     if run_calibration:
         if calibration_data_fn is None:
             raise ValueError(
@@ -151,9 +164,21 @@ def quantize_pytorch_model(
         calibration_data = calibration_data_fn()
         device = next(model.parameters()).device
 
-        cache_seq_len = inputs[2].shape[-2]
-        # Match the dynamic-shape upper bound declared by the caller:
-        #   position_ids.shape[1] <= cache_seq_len - 1   (see pipeline.py `seq_pos` Dim)
+        reset_positions = set(state_indices)
+        for pos in reset_positions:
+            if pos < 2:
+                raise ValueError(
+                    "States cannot occupy the first two input positions. "
+                    "Those must be reserved for input_ids and position_ids"
+                )
+            if pos >= len(inputs):
+                raise IndexError(
+                    f"State index out of bounds, got {pos}, while the number of inputs is "
+                    f"{len(inputs)}"
+                )
+
+        # Match the caller's declared bound: position_ids.shape[1] <= cache_seq_len - 1
+        # (the `seq_pos` Dim in `BaseForCausalLM.build_dynamic_shapes`).
         # position_ids has length QUANT_TRACE_OFFSET + query_len, so:
         #   query_len <= cache_seq_len - QUANT_TRACE_OFFSET - 1
         max_calib_query_len = cache_seq_len - QUANT_TRACE_OFFSET - 1
@@ -162,16 +187,20 @@ def quantize_pytorch_model(
         min_calib_query_len = QUANT_TRACE_QUERY_LEN - QUANT_TRACE_OFFSET
 
         def _prep_calib_inputs(sample: torch.Tensor) -> tuple:
-            sample = sample[:, :max_calib_query_len].to(device)
-            position_ids = (
-                torch.arange(QUANT_TRACE_OFFSET + sample.shape[1], dtype=torch.int32)
+            prepared = list(inputs)
+            prepared[0] = sample[:, :max_calib_query_len].to(device)
+            prepared[1] = (
+                torch.arange(QUANT_TRACE_OFFSET + prepared[0].shape[1], dtype=torch.int32)
                 .unsqueeze(0)
                 .to(device)
             )
-            zero_cache = tuple(
-                torch.zeros(inp.shape, dtype=inp.dtype, device=device) for inp in inputs[2:]
-            )
-            return (sample, position_ids, *zero_cache)
+            for i in range(2, len(prepared)):
+                inp = inputs[i]
+                if i in reset_positions:
+                    prepared[i] = torch.zeros(inp.shape, dtype=inp.dtype, device=device)
+                elif isinstance(inp, torch.Tensor):
+                    prepared[i] = inp.to(device)
+            return tuple(prepared)
 
         calibration_data = [s for s in calibration_data if s.shape[1] >= min_calib_query_len]
         if not calibration_data:
@@ -203,6 +232,66 @@ def _prep_calib_inputs(sample: torch.Tensor) -> tuple:
     return finalized_model
 
 
+def quantize_for_export(
+    model: BaseForCausalLM,
+    config,
+    target_dtype: torch.dtype,
+    quantization_config: dict,
+    calibration_data_fn: Callable[[], list] | None = None,
+    mmap_dir: str | None = None,
+) -> nn.Module:
+    """Apply pre-export torch quantization using the model's own graph contract.
+
+    Builds the calibration trace from the export hooks rather than hardcoding a forward
+    signature, so a model with extra inputs or states calibrates without the caller
+    knowing about them, and activation calibration resets exactly the states.
+
+    Args:
+        model: The loaded model, in eval mode.
+        config: The config the model was built from.
+        target_dtype: Dtype for the trace's cache tensors.
+        quantization_config: Inner coreai-opt ``quantization_config`` dict.
+        calibration_data_fn: Calibration samples; required when the recipe enables
+            ``calibrate_activations``.
+        mmap_dir: Directory for the quantizer's disk checkpointing.
+    """
+    spec = TraceSpec(max_context_length=TRACE_KV_CACHE_SEQ_LEN)
+    reference_inputs = model.build_reference_inputs(config, target_dtype, spec)
+    dynamic_shapes = model.build_dynamic_shapes(config, spec)
+    # Same check the export path runs, so a bad contract fails identically on both.
+    model.validate_export_contract(reference_inputs, dynamic_shapes)
+
+    graph_inputs = reference_inputs[MAIN_GRAPH_NAME]
+    keys = list(graph_inputs)
+    if quantization_config.get("calibrate_activations") and keys[:2] != [
+        "input_ids",
+        "position_ids",
+    ]:
+        raise ValueError(
+            f"{type(model).__name__}: activation calibration feeds tokenized samples as "
+            f"input_ids and rebuilds position_ids, so those must be the first two "
+            f"parameters of forward; got {tuple(keys[:2])}."
+        )
+
+    # Which *positions* are state cannot be read off the contract: the name lists carry
+    # only relative order. So this assumes the declared inputs precede the states, which
+    # holds for the macOS graph -- the only graph calibration runs on. A model that
+    # interleaved a non-state parameter after a cache would need them passed in.
+    n_inputs = len(model.export_input_names()[MAIN_GRAPH_NAME])
+    state_indices = tuple(range(n_inputs, len(keys)))
+
+    return quantize_pytorch_model(
+        model,
+        model.reference_inputs_as_args(graph_inputs),
+        dynamic_shapes[MAIN_GRAPH_NAME],
+        quantization_config,
+        calibration_data_fn=calibration_data_fn,
+        mmap_dir=mmap_dir,
+        cache_seq_len=spec.cache_seq_len,
+        state_indices=state_indices,
+    )
+
+
 def palettize_pytorch_model(
     model: nn.Module,
     example_inputs: tuple,

python/src/coreai_models/export/macos.py

@@ -11,24 +11,19 @@
 """
 
 import logging
+from typing import Any
 
 import coreai_torch
 import coreai_torch.composite_ops
 import torch
 from coreai.authoring import AIProgram
 
-from coreai_models.export._constants import (
-    KEY_CACHE_NAME,
-    QUANT_TRACE_OFFSET,
-    QUANT_TRACE_QUERY_LEN,
-    TRACE_KV_CACHE_SEQ_LEN,
-    VALUE_CACHE_NAME,
-)
+from coreai_models._constants import MAIN_GRAPH_NAME, TRACE_KV_CACHE_SEQ_LEN
 from coreai_models.export.mlir_ops import (
     register_custom_torch_lowering,
     remove_functionalization,
 )
-from coreai_models.primitives.macos.cache import KVCache
+from coreai_models.models.base import BaseForCausalLM, TraceSpec
 
 logger = logging.getLogger(__name__)
 
@@ -64,69 +59,31 @@
 
 
 def _build_reference_inputs(
-    model: torch.nn.Module,
+    model: BaseForCausalLM,
     config,
     target_dtype: torch.dtype,
     max_context_length: int,
-) -> tuple[dict[str, torch.Tensor], dict]:
-    """Build reference inputs and dynamic shapes for macOS model export.
-
-    Args:
-        model: The PyTorch model (used only to read config).
-        config: HuggingFace model config.
-        target_dtype: Data type for cache tensors.
-        max_context_length: Maximum context length for the model.
+) -> tuple[dict[str, Any], dict]:
+    """Reference inputs and dynamic shapes for macOS export.
 
-    Returns:
-        Tuple of (reference_inputs dict, dynamic_shapes dict).
+    Thin wrapper over the model's export-contract hooks, where the per-model variation
+    lives. Returns ``(reference_inputs, dynamic_shapes)``.
     """
-    batch_size = 1
-    vocab_size = config.vocab_size
-
-    input_ids = torch.randint(1, vocab_size, (batch_size, QUANT_TRACE_QUERY_LEN), dtype=torch.int32)
-    position_ids = (
-        torch.arange(QUANT_TRACE_QUERY_LEN + QUANT_TRACE_OFFSET, dtype=torch.int32)
-        .unsqueeze(0)
-        .expand(batch_size, QUANT_TRACE_QUERY_LEN + QUANT_TRACE_OFFSET)
+    # The trace cache length only bounds peak memory, so cap it at the context it serves.
+    spec = TraceSpec(
+        max_context_length=max_context_length,
+        cache_seq_len=min(TRACE_KV_CACHE_SEQ_LEN, max_context_length),
     )
-
-    # Clamp `max_position_embeddings` so KVCache.create_cache_tensors doesn't
-    # allocate a full-context cache for huge models
-    saved_max_pos = config.max_position_embeddings
-    config.max_position_embeddings = TRACE_KV_CACHE_SEQ_LEN
-    k_cache, v_cache = KVCache.create_cache_tensors(config, dtype=target_dtype)
-    config.max_position_embeddings = saved_max_pos
-
-    reference_inputs = {
-        "input_ids": input_ids,
-        "position_ids": position_ids,
-        "k_cache": k_cache,
-        "v_cache": v_cache,
-    }
-
-    dynamic_shapes = {
-        "input_ids": {1: torch.export.Dim("seq_ids", max=max_context_length - 2)},
-        "position_ids": {
-            1: torch.export.Dim("seq_pos", min=QUANT_TRACE_QUERY_LEN, max=max_context_length - 1)
-        },
-        "k_cache": {
-            KVCache.seq_len_dim(): torch.export.Dim(
-                "k_seq_len", min=TRACE_KV_CACHE_SEQ_LEN, max=max_context_length
-            )
-        },
-        "v_cache": {
-            KVCache.seq_len_dim(): torch.export.Dim(
-                "v_seq_len", min=TRACE_KV_CACHE_SEQ_LEN, max=max_context_length
-            )
-        },
-    }
-
-    return reference_inputs, dynamic_shapes
+    reference_inputs = model.build_reference_inputs(config, target_dtype, spec)
+    dynamic_shapes = model.build_dynamic_shapes(config, spec)
+    model.validate_export_contract(reference_inputs, dynamic_shapes)
+    # A macOS model has exactly one graph.
+    return reference_inputs[MAIN_GRAPH_NAME], dynamic_shapes[MAIN_GRAPH_NAME]
 
 
 def export_to_coreai(
     model: torch.nn.Module,
-    reference_inputs: dict[str, torch.Tensor],
+    reference_inputs: dict[str, Any],
     dynamic_shapes: dict | None = None,
     input_names: tuple[str, ...] | None = None,
     output_names: tuple[str, ...] | None = None,
@@ -197,7 +154,7 @@ def export_fn(module: torch.nn.Module) -> torch.export.ExportedProgram:
 
 
 def export_macos_model(
-    model: torch.nn.Module,
+    model: BaseForCausalLM,
     config,
     export_config,
 ) -> AIProgram:
@@ -209,7 +166,8 @@ def export_macos_model(
     3. Optimizes the resulting AIProgram
 
     Args:
-        model: A loaded PyTorch model (already in the correct dtype).
+        model: A loaded PyTorch model (already in the correct dtype). Its
+            export-contract hooks supply the graph's inputs, states, and names.
         config: HuggingFace model config (used for cache dimensions, vocab size, etc.).
         export_config: An ExportConfig instance (used for max_context_length, etc.).
 
@@ -231,18 +189,14 @@ def export_macos_model(
         model, config, target_dtype, max_context_length
     )
 
-    input_names = ("input_ids", "position_ids")
-    output_names = ("logits",)
-    state_names = (KEY_CACHE_NAME, VALUE_CACHE_NAME)
-
     logger.info("Exporting model to Core AI dialect...")
     coreai_program = export_to_coreai(
         model,
         reference_inputs,
         dynamic_shapes=dynamic_shapes,
-        input_names=input_names,
-        output_names=output_names,
-        state_names=state_names,
+        input_names=model.export_input_names()[MAIN_GRAPH_NAME],
+        output_names=model.export_output_names()[MAIN_GRAPH_NAME],
+        state_names=model.export_state_names()[MAIN_GRAPH_NAME],
     )
 
     logger.info("Optimizing AIProgram...")