Move the macOS export graph contract onto the model classes

`export/macos.py` hardcoded the standard text-LLM forward -- two inputs, two KV cache
states, one logits output -- so a model needing anything else cannot use the pipeline at
all. A model threading extra conv or recurrent state alongside the KV cache is stuck, and
one mixing sliding and full attention hits the same wall, because the two layer types need
separate cache tensors with different head counts and head dims.

Adds overridable hooks to `BaseForCausalLM`, keyed by graph name:

    export_input_names()  -> {graph: (name, ...)}
    export_state_names()  -> {graph: (name, ...)}
    export_output_names() -> {graph: (name, ...)}
    build_reference_inputs(config, target_dtype, spec) -> {graph: {param: tensor}}
    build_dynamic_shapes(config, spec)                 -> {graph: shapes}

A macOS model has one graph, `main`. The defaults are exactly what `macos.py` hardcoded,
so `export_macos_model` and the pipeline's quantization step now just call them.

Two ordering rules, because the two are consumed differently. Reference inputs bind to
the traced signature, so they must be in its exact order, interleaved where the signature
interleaves inputs and states. The name lists are looked up by name at runtime, so each
carries only the relative order of its own kind.

Also in here:

* `quantize_for_export` builds the calibration trace from the hooks, so the pipeline and
  any standalone recipe share it. `quantize_pytorch_model` now requires `cache_seq_len`
  and `state_indices` rather than guessing them from input positions.
* `KVCache.create_cache_tensors` takes an explicit `seq_len`, replacing a
  mutate-`config.max_position_embeddings`-and-restore hack. `cache_scatter`'s copy gets
  it too, since its docstring promises the two are drop-in interchangeable.
* `export/_constants.py` moves to `coreai_models/_constants.py`. These are graph/runner
  contract constants that `models/` now needs, and importing them from `export/` would
  reverse the package dependency direction.
* Fixes a crash for contexts at or below the trace cache length: the cache dim was
  declared `Dim(min=TRACE_KV_CACHE_SEQ_LEN, max=max_context_length)` unconditionally, so
  `--max-context-length 2048` raised "Cannot create Dim with inconsistent min/max" from
  inside torch.export. `TraceSpec` now requires `cache_seq_len <= max_context_length` --
  a cache longer than the context it serves is meaningless, so callers cap it at the
  context they are exporting -- and pins the cache dims when the two are equal, since a
  cache traced at the full context has nowhere to grow.

No behavior change for anything that exports today: all 8 registered macOS model classes
produce bit-identical reference inputs, dynamic-shape bounds, and graph names.

Author: Lewis300

…n/src/coreai_models/export/_constants.py python/src/coreai_models/_constants.pypython/src/coreai_models/export/_constants.py renamed to python/src/coreai_models/_constants.py python/src/coreai_models/export/_constants.py renamed to 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,18 @@ 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:
+            assert pos >= 2, (
+                "States cannot occupy the first two input positions. "
+                "Those must be reserved for input_ids and position_ids"
+            )
+            assert pos < len(inputs), (
+                f"State index out of bounds, got {pos}, while the number of inputs is {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 +184,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 +229,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...")