From dd186a0dfb0e32a1745d709314e63c14a17c354e Mon Sep 17 00:00:00 2001 From: Lewis Panos Date: Tue, 11 Aug 2026 12:46:13 -0700 Subject: [PATCH 1/3] 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. --- .../coreai_models/{export => }/_constants.py | 14 +- .../src/coreai_models/export/compression.py | 112 ++++- python/src/coreai_models/export/macos.py | 94 ++-- python/src/coreai_models/export/pipeline.py | 39 +- python/src/coreai_models/model_registry.py | 2 +- python/src/coreai_models/models/base.py | 265 +++++++++++- .../coreai_models/primitives/macos/cache.py | 11 +- .../primitives/macos/cache_scatter.py | 3 +- python/tests/_runner_infra/testing_utils.py | 49 +-- .../tests/test_model_conversion/test_infra.py | 2 +- .../test_export/test_export_contract.py | 408 ++++++++++++++++++ 11 files changed, 829 insertions(+), 170 deletions(-) rename python/src/coreai_models/{export => }/_constants.py (64%) create mode 100644 python/tests/test_model_units/test_export/test_export_contract.py diff --git a/python/src/coreai_models/export/_constants.py b/python/src/coreai_models/_constants.py similarity index 64% rename from python/src/coreai_models/export/_constants.py rename to python/src/coreai_models/_constants.py index 94e1e588..7e790c5e 100644 --- a/python/src/coreai_models/export/_constants.py +++ b/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 diff --git a/python/src/coreai_models/export/compression.py b/python/src/coreai_models/export/compression.py index 4f5cfd41..30d2e5a4 100644 --- a/python/src/coreai_models/export/compression.py +++ b/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, diff --git a/python/src/coreai_models/export/macos.py b/python/src/coreai_models/export/macos.py index c8798301..da6d31b4 100644 --- a/python/src/coreai_models/export/macos.py +++ b/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...") diff --git a/python/src/coreai_models/export/pipeline.py b/python/src/coreai_models/export/pipeline.py index 7e2fd63d..abe49367 100644 --- a/python/src/coreai_models/export/pipeline.py +++ b/python/src/coreai_models/export/pipeline.py @@ -24,17 +24,15 @@ from coreai_opt.palettization.config.palettization_config import KMeansPalettizerConfig from transformers import AutoConfig, AutoTokenizer -from coreai_models.export._constants import ( +from coreai_models._constants import ( IOS_DEFAULT_MAX_CONTEXT_LENGTH, - QUANT_TRACE_OFFSET, - QUANT_TRACE_QUERY_LEN, TRACE_KV_CACHE_SEQ_LEN, ) from coreai_models.export.bundle import bundle_llm_asset from coreai_models.export.compression import ( get_c4, palettize_pytorch_model, - quantize_pytorch_model, + quantize_for_export, ) from coreai_models.export.ios import export_ios_model from coreai_models.export.macos import export_macos_model @@ -44,7 +42,6 @@ get_preset, ) from coreai_models.models.registry import get_model_entry -from coreai_models.primitives.macos.cache import KVCache logger = logging.getLogger(__name__) @@ -241,32 +238,6 @@ async def _async_export_model(config: ExportConfig) -> str: if torch_quantization_config is not None: logger.info(f"Applying pre-export torch quantization (preset={config.compression})") - 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) - ) - - saved_max_pos = hf_config.max_position_embeddings - hf_config.max_position_embeddings = TRACE_KV_CACHE_SEQ_LEN - k_cache, v_cache = KVCache.create_cache_tensors(hf_config, dtype=target_dtype) - hf_config.max_position_embeddings = saved_max_pos - - quantization_inputs = (input_ids, position_ids, k_cache, v_cache) - quantization_dynamic_shapes = { - "input_ids": {1: torch.export.Dim("seq_ids", max=TRACE_KV_CACHE_SEQ_LEN - 2)}, - "position_ids": { - 1: torch.export.Dim( - "seq_pos", min=QUANT_TRACE_QUERY_LEN, max=TRACE_KV_CACHE_SEQ_LEN - 1 - ) - }, - "k_cache": None, - "v_cache": None, - } - def get_calibration_data(): # type: ignore[no-untyped-def] tokenizer = AutoTokenizer.from_pretrained(config.hf_model_id) return get_c4(tokenizer) @@ -284,10 +255,10 @@ def get_calibration_data(): # type: ignore[no-untyped-def] if not isinstance(torch_quantization_config, dict) else dict(torch_quantization_config) ) - model = quantize_pytorch_model( + model = quantize_for_export( model, - quantization_inputs, - quantization_dynamic_shapes, + hf_config, + target_dtype, quant_cfg, calibration_data_fn=get_calibration_data, mmap_dir=quantizer_mmap_dir, diff --git a/python/src/coreai_models/model_registry.py b/python/src/coreai_models/model_registry.py index cc6a0bc9..d5b5717a 100644 --- a/python/src/coreai_models/model_registry.py +++ b/python/src/coreai_models/model_registry.py @@ -27,7 +27,7 @@ from dataclasses import asdict, dataclass from pathlib import Path -from coreai_models.export._constants import IOS_DEFAULT_MAX_CONTEXT_LENGTH +from coreai_models._constants import IOS_DEFAULT_MAX_CONTEXT_LENGTH # --------------------------------------------------------------------------- # Data model diff --git a/python/src/coreai_models/models/base.py b/python/src/coreai_models/models/base.py index 647e4043..b325abe2 100644 --- a/python/src/coreai_models/models/base.py +++ b/python/src/coreai_models/models/base.py @@ -7,13 +7,15 @@ import collections.abc import gc +import inspect import json import os import re from abc import abstractmethod from collections.abc import Callable +from dataclasses import dataclass from functools import wraps -from typing import TypeVar, cast +from typing import Any, NoReturn, TypeVar, cast import torch from huggingface_hub import snapshot_download @@ -21,14 +23,68 @@ from safetensors.torch import save_file from transformers import AutoConfig from transformers.modeling_utils import PreTrainedModel -from typing_extensions import Self - +from typing_extensions import Self, override + +from coreai_models._constants import ( + KEY_CACHE_NAME, + MAIN_GRAPH_NAME, + QUANT_TRACE_OFFSET, + QUANT_TRACE_QUERY_LEN, + TRACE_KV_CACHE_SEQ_LEN, + VALUE_CACHE_NAME, +) from coreai_models.primitives.ios.embedding import GatherEmbeddings, LoadEmbeddings from coreai_models.primitives.macos.cache import KVCache T = TypeVar("T", bound="BaseForCausalLM") +@dataclass(frozen=True) +class TraceSpec: + """Shapes for one ``torch.export`` trace of a causal LM forward. + + Passed to both ``build_reference_inputs`` and ``build_dynamic_shapes`` so the + tensors and their declared dims cannot disagree. + + Attributes: + max_context_length: Upper bound for the dynamic seq/cache dims. Must be at + least ``query_len + 2``. + cache_seq_len: Length the caches are *traced* at, to bound peak trace memory. + Unrelated to the inference cache size. Must not exceed + ``max_context_length``. + query_len: Trace-time ``input_ids`` length. + offset: Already-cached positions, so ``position_ids`` is + ``query_len + offset`` long. + """ + + max_context_length: int + cache_seq_len: int = TRACE_KV_CACHE_SEQ_LEN + query_len: int = QUANT_TRACE_QUERY_LEN + offset: int = QUANT_TRACE_OFFSET + + def __post_init__(self) -> None: + # Below this there is no legal `Dim(min=query_len, max=max_context_length - 1)`. + if self.max_context_length < self.query_len + 2: + raise ValueError( + f"max_context_length={self.max_context_length} is too small to trace: " + f"it must be at least query_len + 2 = {self.query_len + 2}." + ) + assert self.cache_seq_len <= self.max_context_length, ( + "cache_seq_len must not be greater than max_context_length. Received " + f"cache_seq_len = {self.cache_seq_len}, " + f"max_context_length = {self.max_context_length}" + ) + + @property + def caches_are_static(self) -> bool: + """Whether the cache dims must be pinned rather than declared dynamic. + + A cache traced at the full context has nowhere to grow, and ``Dim(min=max=n)`` + is illegal anyway. + """ + return self.cache_seq_len == self.max_context_length + + def _is_layer_key_beyond(key: str, num_layers: int) -> bool: """Return True if `key` refers to a transformer layer with index >= num_layers. @@ -300,6 +356,164 @@ def _mutate_state_dict(self: Self, state_dict: dict[str, torch.Tensor]) -> None: """ ... + # ------------------------------------------------------------------ + # Export contract + # + # Everything the exporters need to trace this model, keyed by graph name. A macOS + # model has one graph; iOS has several. These hooks supply only names and tensors; + # which callable each graph traces stays the exporter's business. + # + # Reference inputs bind to the traced signature, so they must be in its EXACT + # order. Names are looked up by name, so each list carries only the RELATIVE order + # of its own kind: for forward(input_ids, key_cache, position_ids, value_cache), + # input_names is (input_ids, position_ids) and state_names is (key_cache, + # value_cache). + # ------------------------------------------------------------------ + + @classmethod + def export_input_names(cls) -> dict[str, tuple[str, ...]]: + """Graph input names per graph, in relative order among the non-state args.""" + return {MAIN_GRAPH_NAME: ("input_ids", "position_ids")} + + @classmethod + def export_state_names(cls) -> dict[str, tuple[str, ...]]: + """Runner-visible state names per graph, in relative order among the state args. + + State args are mutated in place and surfaced through the runtime ``state=`` + kwarg rather than as ordinary inputs/outputs. + """ + return {MAIN_GRAPH_NAME: (KEY_CACHE_NAME, VALUE_CACHE_NAME)} + + @classmethod + def export_output_names(cls) -> dict[str, tuple[str, ...]]: + """Graph output names per graph, in return order.""" + return {MAIN_GRAPH_NAME: ("logits",)} + + def build_reference_inputs( + self, + config, + target_dtype: torch.dtype, + spec: TraceSpec, + ) -> dict[str, dict[str, Any]]: + """Reference tensors to trace with, per graph, keyed by parameter name. + + The inner dicts bind to the traced callable, so their keys must be its + parameters in *exact* signature order. Pass ``spec`` to + :meth:`build_dynamic_shapes` too, so the tensors and their dims cannot disagree. + """ + input_ids = torch.randint(1, config.vocab_size, (1, spec.query_len), dtype=torch.int32) + position_ids = ( + torch.arange(spec.query_len + spec.offset, dtype=torch.int32) + .unsqueeze(0) + .expand(1, spec.query_len + spec.offset) + ) + k_cache, v_cache = KVCache.create_cache_tensors( + config, dtype=target_dtype, seq_len=spec.cache_seq_len + ) + return { + MAIN_GRAPH_NAME: { + "input_ids": input_ids, + "position_ids": position_ids, + "k_cache": k_cache, + "v_cache": v_cache, + } + } + + def build_dynamic_shapes(self, config, spec: TraceSpec) -> dict[str, Any]: + """``dynamic_shapes`` per graph, matching :meth:`build_reference_inputs`. + + Keyed like the reference inputs; ``None`` pins that input to its traced shape. + """ + max_ctx = spec.max_context_length + shapes: dict[str, Any] = { + "input_ids": {1: torch.export.Dim("seq_ids", max=max_ctx - 2)}, + "position_ids": {1: torch.export.Dim("seq_pos", min=spec.query_len, max=max_ctx - 1)}, + } + seq_dim = KVCache.seq_len_dim() + if spec.caches_are_static: + shapes["k_cache"] = None + shapes["v_cache"] = None + else: + shapes["k_cache"] = { + seq_dim: torch.export.Dim("k_seq_len", min=spec.cache_seq_len, max=max_ctx) + } + shapes["v_cache"] = { + seq_dim: torch.export.Dim("v_seq_len", min=spec.cache_seq_len, max=max_ctx) + } + return {MAIN_GRAPH_NAME: shapes} + + def validate_export_contract( + self, + reference_inputs: dict[str, dict[str, Any]], + dynamic_shapes: dict[str, Any], + ) -> None: + """Check the five hooks agree with each other. + + The converter compares name *counts* only, so it would accept a contract whose + graphs disagree. Called by the exporters before tracing. + + Raises: + ValueError: On any disagreement between the hooks. + """ + cls_name = type(self).__name__ + inputs, states, outputs = ( + self.export_input_names(), + self.export_state_names(), + self.export_output_names(), + ) + named = { + "export_input_names": set(inputs), + "export_state_names": set(states), + "export_output_names": set(outputs), + "build_reference_inputs": set(reference_inputs), + "build_dynamic_shapes": set(dynamic_shapes), + } + graphs = named["export_input_names"] + for hook, keys in named.items(): + if keys != graphs: + raise ValueError( + f"{cls_name}: {hook} covers graphs {sorted(keys)} but " + f"export_input_names covers {sorted(graphs)}. Every hook must " + "describe the same graphs." + ) + + for graph in sorted(graphs): + refs = reference_inputs[graph] + declared = len(inputs[graph]) + len(states[graph]) + if declared != len(refs): + raise ValueError( + f"{cls_name}, graph {graph!r}: {len(inputs[graph])} input names + " + f"{len(states[graph])} state names = {declared}, but " + f"build_reference_inputs supplies {len(refs)} tensors " + f"{tuple(refs)}." + ) + shapes = dynamic_shapes[graph] + if shapes is not None and set(shapes) != set(refs): + raise ValueError( + f"{cls_name}, graph {graph!r}: dynamic_shapes keys " + f"{tuple(shapes)} do not match reference inputs {tuple(refs)}." + ) + + def reference_inputs_as_args(self, reference_inputs: dict[str, Any]) -> tuple[Any, ...]: + """One graph's reference inputs as positional args, validating the order. + + For the coreai-opt quantizer, which takes a tuple rather than kwargs; a dict + ordered differently from the signature would silently bind the wrong tensors. + + Raises: + ValueError: If the keys are not a contiguous in-order prefix of ``forward``. + """ + params = list(inspect.signature(self.forward).parameters) + keys = list(reference_inputs) + if keys != params[: len(keys)]: + raise ValueError( + f"{type(self).__name__}.build_reference_inputs returned keys {keys}, " + f"which are not a contiguous in-order prefix of forward's parameters " + f"{params}. Positional conversion would bind tensors to the wrong " + f"parameters; fix the dict's insertion order or the signature." + ) + return tuple(reference_inputs.values()) + @classmethod def _get_reauthored_config( cls, @@ -605,3 +819,48 @@ def __init__(self: Self, config, model_device: str, disable_embedding_quantizati def set_prefill_mode(self, prefill_mode: bool): self.extend.prefill_mode = prefill_mode + + # ------------------------------------------------------------------ + # Export contract -- not implemented for iOS + # + # The iOS graph matches the macOS defaults in no respect (different inputs, cache + # rank and output name -- see export/ios.py), and export_ios_model builds its own. + # Inheriting the defaults would report a plausible but wrong contract, and they are + # reachable: `--variant iOS --compression 4bit` resolves a macOS quantization preset + # and the quantize branch is not variant-gated. + # ------------------------------------------------------------------ + + @classmethod + def _export_contract_not_implemented(cls, hook: str) -> NoReturn: + raise NotImplementedError( + f"{cls.__name__}.{hook} is not yet implemented for iOS models: the iOS graph " + "has a different signature from the macOS one it would otherwise inherit." + ) + + @classmethod + @override + def export_input_names(cls) -> dict[str, tuple[str, ...]]: + cls._export_contract_not_implemented("export_input_names") + + @classmethod + @override + def export_state_names(cls) -> dict[str, tuple[str, ...]]: + cls._export_contract_not_implemented("export_state_names") + + @classmethod + @override + def export_output_names(cls) -> dict[str, tuple[str, ...]]: + cls._export_contract_not_implemented("export_output_names") + + @override + def build_reference_inputs( + self, + config, + target_dtype: torch.dtype, + spec: TraceSpec, + ) -> dict[str, dict[str, Any]]: + self._export_contract_not_implemented("build_reference_inputs") + + @override + def build_dynamic_shapes(self, config, spec: TraceSpec) -> dict[str, Any]: + self._export_contract_not_implemented("build_dynamic_shapes") diff --git a/python/src/coreai_models/primitives/macos/cache.py b/python/src/coreai_models/primitives/macos/cache.py index ece971fc..a3d51cde 100644 --- a/python/src/coreai_models/primitives/macos/cache.py +++ b/python/src/coreai_models/primitives/macos/cache.py @@ -36,15 +36,22 @@ def create_cache_tensors( cls, config, dtype: torch.dtype = torch.float32, + seq_len: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Create zero-initialized KV cache tensors from a model config. + Args: + config: Model config supplying the layer/head dimensions. + dtype: Cache dtype. + seq_len: Sequence-dim length; defaults to ``config.max_position_embeddings``. + Pass explicitly to build a trace-sized cache without mutating config. + Returns: - (k_cache, v_cache) tensors of shape (n_layers, 1, n_kv_heads, max_seq_len, head_dim). + (k_cache, v_cache) of shape (n_layers, 1, n_kv_heads, max_seq_len, head_dim). """ n_kv_heads = config.num_key_value_heads n_layers = config.num_hidden_layers - max_seq_len = config.max_position_embeddings + max_seq_len = config.max_position_embeddings if seq_len is None else seq_len if hasattr(config, "head_dim") and config.head_dim is not None: head_dim = config.head_dim else: diff --git a/python/src/coreai_models/primitives/macos/cache_scatter.py b/python/src/coreai_models/primitives/macos/cache_scatter.py index 6bba2214..a55c51c3 100644 --- a/python/src/coreai_models/primitives/macos/cache_scatter.py +++ b/python/src/coreai_models/primitives/macos/cache_scatter.py @@ -36,10 +36,11 @@ def create_cache_tensors( cls, config, dtype: torch.dtype = torch.float32, + seq_len: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: n_kv_heads = config.num_key_value_heads n_layers = config.num_hidden_layers - max_seq_len = config.max_position_embeddings + max_seq_len = config.max_position_embeddings if seq_len is None else seq_len if hasattr(config, "head_dim") and config.head_dim is not None: head_dim = config.head_dim else: diff --git a/python/tests/_runner_infra/testing_utils.py b/python/tests/_runner_infra/testing_utils.py index 8a5d7527..57439232 100644 --- a/python/tests/_runner_infra/testing_utils.py +++ b/python/tests/_runner_infra/testing_utils.py @@ -957,20 +957,14 @@ def test_weight_activation_quantization(self, activation_quantization) -> None: if activation_quantization: pytest.skip("Activation quantization temporarily disabled with eager mode quantization") - # We replicate the relevant parts of - # ``coreai_models.export.pipeline._async_export_model`` here: - # load HF -> apply torch quantization -> run macOS export. The output - # asset write is intentionally skipped; we only want to confirm + # Same building blocks as `_async_export_model`: load HF -> + # `quantize_for_export` -> `export_macos_model`. Calling those rather than + # reimplementing the calibration trace keeps this from drifting from + # production, as it had. The asset write is skipped; we only confirm # quantize + export produces a non-None AIProgram. - from coreai_models.export._constants import ( - QUANT_TRACE_OFFSET, - QUANT_TRACE_QUERY_LEN, - TRACE_KV_CACHE_SEQ_LEN, - ) - from coreai_models.export.compression import quantize_pytorch_model + from coreai_models.export.compression import quantize_for_export from coreai_models.export.macos import export_macos_model from coreai_models.export.pipeline import ExportConfig - from coreai_models.primitives.macos.cache import KVCache hf_config = transformers.AutoConfig.from_pretrained(self._toy_model_id) is_gemma = "gemma" in self._model_class.__name__.lower() @@ -1034,39 +1028,12 @@ def test_weight_activation_quantization(self, activation_quantization) -> None: hf_state_dict_prefix=hf_state_dict_prefix, ).eval() - # Build calibration / trace inputs for quantization - vocab_size = getattr(hf_config, "vocab_size", 32000) - input_ids = torch.randint(1, vocab_size, (1, QUANT_TRACE_QUERY_LEN), dtype=torch.int32) - position_ids = ( - torch.arange(QUANT_TRACE_QUERY_LEN + QUANT_TRACE_OFFSET, dtype=torch.int32) - .unsqueeze(0) - .expand(1, QUANT_TRACE_QUERY_LEN + QUANT_TRACE_OFFSET) - ) - saved_max_pos = hf_config.max_position_embeddings - hf_config.max_position_embeddings = TRACE_KV_CACHE_SEQ_LEN - k_cache, v_cache = KVCache.create_cache_tensors(hf_config, dtype=target_dtype) - hf_config.max_position_embeddings = saved_max_pos - - quantization_inputs = (input_ids, position_ids, k_cache, v_cache) - quantization_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": None, - "v_cache": None, - } - quantizer_mmap_dir = f"{tmpdir}/quantized" os.makedirs(quantizer_mmap_dir, exist_ok=True) - model = quantize_pytorch_model( + model = quantize_for_export( model, - quantization_inputs, - quantization_dynamic_shapes, + hf_config, + target_dtype, dict(torch_quantization_config), calibration_data_fn=None, mmap_dir=quantizer_mmap_dir, diff --git a/python/tests/test_model_conversion/test_infra.py b/python/tests/test_model_conversion/test_infra.py index cfe755ee..5cd714f0 100644 --- a/python/tests/test_model_conversion/test_infra.py +++ b/python/tests/test_model_conversion/test_infra.py @@ -17,8 +17,8 @@ Qwen3ForCausalLM as HFQwen3ForCausalLM, ) +from coreai_models._constants import IOS_DEFAULT_MAX_CONTEXT_LENGTH from coreai_models.export import pipeline as export_pipeline -from coreai_models.export._constants import IOS_DEFAULT_MAX_CONTEXT_LENGTH from coreai_models.export.ios import KEY_CACHE_INPUT_NAME, VALUE_CACHE_INPUT_NAME from coreai_models.export.pipeline import ExportConfig, _async_export_model diff --git a/python/tests/test_model_units/test_export/test_export_contract.py b/python/tests/test_model_units/test_export/test_export_contract.py new file mode 100644 index 00000000..ef8ddb6e --- /dev/null +++ b/python/tests/test_model_units/test_export/test_export_contract.py @@ -0,0 +1,408 @@ +# Copyright 2026 Apple Inc. +# +# 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 + +"""Tests for the export contract hooks on ``BaseForCausalLM``. + +Cover the graph contract independently of any real architecture, plus the +append-extras path. No hardware or HuggingFace weights required. +""" + +from types import SimpleNamespace + +import pytest +import torch +from typing_extensions import override + +from coreai_models._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 as MAIN, +) +from coreai_models.models.base import BaseForCausalLM, TraceSpec +from coreai_models.primitives.macos.cache import KVCache + +MAX_CONTEXT_LENGTH = 8192 + + +def _tiny_config() -> SimpleNamespace: + """The smallest config the contract hooks read.""" + return SimpleNamespace( + vocab_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + hidden_size=32, + head_dim=8, + max_position_embeddings=MAX_CONTEXT_LENGTH, + ) + + +class _StandardLM(BaseForCausalLM): + """A model with the default contract: (input_ids, position_ids, k_cache, v_cache).""" + + @override + def _init_model(self, config) -> None: + self.lm_head = torch.nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + @override + def _mutate_state_dict(self, state_dict) -> None: + pass + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.IntTensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + ) -> torch.Tensor: + raise NotImplementedError("shape contract only; never traced in these tests") + + +class _ExtraStateLM(_StandardLM): + """A model that appends state beyond the standard KV pair. + + The base KV pair first, then extra state args. + """ + + @override + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.IntTensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + extra_state: torch.Tensor = None, + ) -> torch.Tensor: + raise NotImplementedError("shape contract only; never traced in these tests") + + @classmethod + @override + def export_state_names(cls) -> dict[str, tuple[str, ...]]: + return {MAIN: (*super().export_state_names()[MAIN], "extraState")} + + @override + def build_reference_inputs(self, config, target_dtype, spec): + graphs = super().build_reference_inputs(config, target_dtype, spec) + graphs[MAIN]["extra_state"] = torch.zeros(2, spec.cache_seq_len, dtype=target_dtype) + return graphs + + @override + def build_dynamic_shapes(self, config, spec): + graphs = super().build_dynamic_shapes(config, spec) + graphs[MAIN]["extra_state"] = None + return graphs + + +@pytest.fixture +def config() -> SimpleNamespace: + return _tiny_config() + + +@pytest.fixture +def model(config) -> _StandardLM: + return _StandardLM(config) + + +class TestMacOSContract: + """The macOS model describes one graph, keyed ``main``.""" + + def _built(self, model, config, spec=None): + spec = spec or TraceSpec(max_context_length=MAX_CONTEXT_LENGTH) + return ( + model.build_reference_inputs(config, torch.float16, spec), + model.build_dynamic_shapes(config, spec), + ) + + def test_every_hook_describes_exactly_the_main_graph(self, model, config) -> None: + refs, shapes = self._built(model, config) + for hook in ( + model.export_input_names(), + model.export_state_names(), + model.export_output_names(), + refs, + shapes, + ): + assert list(hook) == [MAIN] + + def test_names(self, model) -> None: + assert model.export_input_names()[MAIN] == ("input_ids", "position_ids") + assert model.export_state_names()[MAIN] == (KEY_CACHE_NAME, VALUE_CACHE_NAME) + assert model.export_output_names()[MAIN] == ("logits",) + + def test_reference_inputs_are_in_exact_signature_order(self, model, config) -> None: + """They bind to the traced callable, so order is exact, not relative.""" + import inspect + + refs, _ = self._built(model, config) + params = list(inspect.signature(model.forward).parameters) + keys = list(refs[MAIN]) + assert keys == params[: len(keys)] + + def test_reference_input_shapes(self, model, config) -> None: + refs, _ = self._built(model, config) + graph = refs[MAIN] + assert graph["input_ids"].shape == (1, QUANT_TRACE_QUERY_LEN) + assert graph["input_ids"].dtype == torch.int32 + assert graph["position_ids"].shape == (1, QUANT_TRACE_QUERY_LEN + QUANT_TRACE_OFFSET) + expected = (2, 1, 2, TRACE_KV_CACHE_SEQ_LEN, 8) + for name in ("k_cache", "v_cache"): + assert graph[name].shape == expected + assert graph[name].dtype == torch.float16 + + def test_caches_traced_at_cache_seq_len(self, model, config) -> None: + spec = TraceSpec(max_context_length=MAX_CONTEXT_LENGTH, cache_seq_len=512) + refs, _ = self._built(model, config, spec) + assert refs[MAIN]["k_cache"].shape[KVCache.seq_len_dim()] == 512 + + def test_config_is_not_mutated(self, model, config) -> None: + """Regression: sizing the cache used to mutate and restore the config.""" + self._built( + model, config, TraceSpec(max_context_length=MAX_CONTEXT_LENGTH, cache_seq_len=512) + ) + assert config.max_position_embeddings == MAX_CONTEXT_LENGTH + + def test_dynamic_shape_bounds(self, model, config) -> None: + _, shapes = self._built(model, config) + graph = shapes[MAIN] + assert graph["input_ids"][1].max == MAX_CONTEXT_LENGTH - 2 + assert graph["position_ids"][1].min == QUANT_TRACE_QUERY_LEN + assert graph["position_ids"][1].max == MAX_CONTEXT_LENGTH - 1 + seq_dim = KVCache.seq_len_dim() + for name in ("k_cache", "v_cache"): + assert graph[name][seq_dim].min == TRACE_KV_CACHE_SEQ_LEN + assert graph[name][seq_dim].max == MAX_CONTEXT_LENGTH + + +class TestSmallContext: + """Contexts at or below the default trace cache length. + + Regression: the cache dim was built as ``Dim(min=TRACE_KV_CACHE_SEQ_LEN, + max=max_context_length)`` unconditionally, which raises from inside ``torch.export`` + whenever the context is <= the trace length. + """ + + def test_cache_seq_len_may_equal_the_context(self) -> None: + assert TraceSpec(max_context_length=512, cache_seq_len=512).cache_seq_len == 512 + + def test_cache_seq_len_above_the_context_is_rejected(self) -> None: + # A cache longer than the context it serves is meaningless, so the spec + # rejects it rather than quietly shrinking it. + with pytest.raises(AssertionError, match="must not be greater than"): + TraceSpec(max_context_length=512, cache_seq_len=513) + + def test_larger_context_leaves_trace_length_alone(self) -> None: + assert TraceSpec(max_context_length=8192).cache_seq_len == TRACE_KV_CACHE_SEQ_LEN + + def test_context_too_small_to_trace_is_rejected(self) -> None: + limit = TraceSpec(max_context_length=8192).query_len + 2 + TraceSpec(max_context_length=limit, cache_seq_len=limit) + for bad in (limit - 1, 2, 0, -5): + with pytest.raises(ValueError, match="too small to trace"): + TraceSpec(max_context_length=bad) + + @pytest.mark.parametrize("max_ctx", [512, TRACE_KV_CACHE_SEQ_LEN]) + def test_cache_dims_pin_instead_of_raising(self, model, max_ctx) -> None: + config = _tiny_config() + config.max_position_embeddings = max_ctx + spec = TraceSpec( + max_context_length=max_ctx, cache_seq_len=min(TRACE_KV_CACHE_SEQ_LEN, max_ctx) + ) + assert spec.caches_are_static + shapes = model.build_dynamic_shapes(config, spec)[MAIN] + assert shapes["k_cache"] is None + assert shapes["v_cache"] is None + + def test_dims_stay_dynamic_when_there_is_room(self, model, config) -> None: + spec = TraceSpec(max_context_length=MAX_CONTEXT_LENGTH) + assert not spec.caches_are_static + assert model.build_dynamic_shapes(config, spec)[MAIN]["k_cache"] is not None + + +class TestValidateExportContract: + """Cross-checks the five hooks against each other.""" + + def _built(self, model, config): + spec = TraceSpec(max_context_length=MAX_CONTEXT_LENGTH) + return ( + model.build_reference_inputs(config, torch.float16, spec), + model.build_dynamic_shapes(config, spec), + ) + + def test_accepts_the_default_contract(self, model, config) -> None: + model.validate_export_contract(*self._built(model, config)) + + def test_accepts_appended_state(self, config) -> None: + m = _ExtraStateLM(config) + m.validate_export_contract(*self._built(m, config)) + + def test_rejects_a_hook_covering_different_graphs(self, model, config) -> None: + refs, shapes = self._built(model, config) + refs["extra_graph"] = {} + with pytest.raises(ValueError, match="must describe the same graphs"): + model.validate_export_contract(refs, shapes) + + def test_rejects_a_name_count_mismatch(self, config) -> None: + class _TooFewNames(_StandardLM): + @classmethod + @override + def export_state_names(cls) -> dict[str, tuple[str, ...]]: + return {MAIN: (KEY_CACHE_NAME,)} + + m = _TooFewNames(config) + with pytest.raises(ValueError, match="build_reference_inputs supplies"): + m.validate_export_contract(*self._built(m, config)) + + def test_rejects_dynamic_shape_key_mismatch(self, model, config) -> None: + refs, shapes = self._built(model, config) + del shapes[MAIN]["position_ids"] + with pytest.raises(ValueError, match="do not match reference inputs"): + model.validate_export_contract(refs, shapes) + + +class TestReferenceInputsAsArgs: + """Positional conversion for the quantizer, which takes a tuple not kwargs.""" + + def _graph(self, model, config): + spec = TraceSpec(max_context_length=MAX_CONTEXT_LENGTH) + return model.build_reference_inputs(config, torch.float16, spec)[MAIN] + + def test_returns_signature_order(self, model, config) -> None: + graph = self._graph(model, config) + args = model.reference_inputs_as_args(graph) + assert len(args) == 4 + for arg, expected in zip(args, graph.values(), strict=True): + assert arg is expected + + def test_rejects_reordered_keys(self, model, config) -> None: + graph = self._graph(model, config) + swapped = {"position_ids": graph["position_ids"], "input_ids": graph["input_ids"]} + swapped.update({k: graph[k] for k in ("k_cache", "v_cache")}) + with pytest.raises(ValueError, match="not a contiguous in-order prefix"): + model.reference_inputs_as_args(swapped) + + def test_sees_through_the_logits_cast_decorator(self, config) -> None: + """Every real subclass decorates forward; introspection needs functools.wraps.""" + + class _Decorated(_StandardLM): + @BaseForCausalLM.cast_logits_bfloat16_to_float16 + @override + def forward(self, input_ids, position_ids, k_cache, v_cache): + raise NotImplementedError + + m = _Decorated(config) + assert len(m.reference_inputs_as_args(self._graph(m, config))) == 4 + + +class TestRegisteredModelsSatisfyTheContract: + """Pins the real registry, so renaming a forward parameter fails here.""" + + @staticmethod + def _macos_entries(): + from coreai_models.models.registry import _get_registry + + return sorted( + ((mt, e) for mt, e in _get_registry().items() if e.macos_class is not None), + key=lambda kv: kv[0], + ) + + def test_registry_is_not_empty(self) -> None: + assert self._macos_entries() + + def test_every_registered_model_validates(self) -> None: + import inspect + + from transformers import AutoConfig + + for model_type, entry in self._macos_entries(): + raw = AutoConfig.for_model(model_type) + cfg = ( + getattr(raw, entry.hf_config_attr) + if entry.hf_config_attr and hasattr(raw, entry.hf_config_attr) + else raw + ) + cfg.num_hidden_layers = 2 + cfg.max_position_embeddings = MAX_CONTEXT_LENGTH + m = entry.macos_class(cfg, model_device="meta") + + spec = TraceSpec(max_context_length=MAX_CONTEXT_LENGTH) + refs = m.build_reference_inputs(cfg, torch.float16, spec) + shapes = m.build_dynamic_shapes(cfg, spec) + m.validate_export_contract(refs, shapes) + + params = list(inspect.signature(m.forward).parameters) + keys = list(refs[MAIN]) + assert keys == params[: len(keys)], model_type + + +class TestIOSHooksRaise: + """iOS models must not inherit the macOS contract. + + ``export/ios.py`` never calls these hooks, but the defaults were reachable: + ``--variant iOS --compression 4bit`` resolves a macOS quantization preset and the + quantize branch is not variant-gated. + """ + + HOOKS = ( + "export_input_names", + "export_state_names", + "export_output_names", + "build_reference_inputs", + "build_dynamic_shapes", + ) + + @pytest.fixture + def ios_model(self): + from transformers import AutoConfig + + from coreai_models.models.ios.qwen3 import Qwen3ForCausalLMForiOS + + cfg = AutoConfig.for_model("qwen3") + cfg.num_hidden_layers = 2 + cfg.max_position_embeddings = MAX_CONTEXT_LENGTH + return Qwen3ForCausalLMForiOS(cfg, model_device="meta") + + def _call(self, model, hook): + spec = TraceSpec(max_context_length=MAX_CONTEXT_LENGTH) + args = { + "build_reference_inputs": (model.config, torch.float16, spec), + "build_dynamic_shapes": (model.config, spec), + }.get(hook, ()) + return getattr(model, hook)(*args) + + @pytest.mark.parametrize("hook", HOOKS) + def test_hook_raises_not_implemented(self, ios_model, hook) -> None: + with pytest.raises(NotImplementedError) as exc: + self._call(ios_model, hook) + message = str(exc.value) + assert "not yet implemented for iOS models" in message + assert hook in message, "the error should name the hook that was called" + assert type(ios_model).__name__ in message + + def test_every_base_hook_is_overridden(self) -> None: + """Guard against a hook being added to the base and silently inherited.""" + from coreai_models.models.base import BaseForCausalLMForiOS + + for hook in self.HOOKS: + assert hook in vars(BaseForCausalLMForiOS), ( + f"{hook} is not overridden on BaseForCausalLMForiOS; iOS models would " + "inherit the macOS default" + ) + + def test_the_reachable_quantization_path_raises_clearly(self, ios_model) -> None: + """``--variant iOS --compression 4bit`` reaches ``quantize_for_export``.""" + from coreai_models.export.compression import quantize_for_export + from coreai_models.export.presets import get_preset + + quantization_config = get_preset("4bit").get("torch_quantization_config") + assert quantization_config is not None + with pytest.raises(NotImplementedError, match="not yet implemented for iOS models"): + quantize_for_export( + ios_model, ios_model.config, torch.float16, dict(quantization_config) + ) From aa84d2d401766f3929d80fc0223e9fbdd7cab0eb Mon Sep 17 00:00:00 2001 From: Lewis Panos Date: Tue, 11 Aug 2026 20:23:51 -0700 Subject: [PATCH 2/3] Address comments --- python/src/coreai_models/export/compression.py | 17 ++++++++++------- python/src/coreai_models/models/base.py | 11 ++++++----- 2 files changed, 16 insertions(+), 12 deletions(-) diff --git a/python/src/coreai_models/export/compression.py b/python/src/coreai_models/export/compression.py index 30d2e5a4..bd83f513 100644 --- a/python/src/coreai_models/export/compression.py +++ b/python/src/coreai_models/export/compression.py @@ -166,13 +166,16 @@ def quantize_pytorch_model( 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)}" - ) + 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`). diff --git a/python/src/coreai_models/models/base.py b/python/src/coreai_models/models/base.py index b325abe2..783ed22f 100644 --- a/python/src/coreai_models/models/base.py +++ b/python/src/coreai_models/models/base.py @@ -69,11 +69,12 @@ def __post_init__(self) -> None: f"max_context_length={self.max_context_length} is too small to trace: " f"it must be at least query_len + 2 = {self.query_len + 2}." ) - assert self.cache_seq_len <= self.max_context_length, ( - "cache_seq_len must not be greater than max_context_length. Received " - f"cache_seq_len = {self.cache_seq_len}, " - f"max_context_length = {self.max_context_length}" - ) + if self.cache_seq_len > self.max_context_length: + raise ValueError( + "cache_seq_len must not be greater than max_context_length. Received " + f"cache_seq_len = {self.cache_seq_len}, " + f"max_context_length = {self.max_context_length}" + ) @property def caches_are_static(self) -> bool: From 15edf978bc39453f4d2f0f5fb2624f6ec9b9a174 Mon Sep 17 00:00:00 2001 From: Lewis Panos Date: Tue, 11 Aug 2026 20:31:02 -0700 Subject: [PATCH 3/3] Fix test --- .../tests/test_model_units/test_export/test_export_contract.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tests/test_model_units/test_export/test_export_contract.py b/python/tests/test_model_units/test_export/test_export_contract.py index ef8ddb6e..abebc774 100644 --- a/python/tests/test_model_units/test_export/test_export_contract.py +++ b/python/tests/test_model_units/test_export/test_export_contract.py @@ -194,7 +194,7 @@ def test_cache_seq_len_may_equal_the_context(self) -> None: def test_cache_seq_len_above_the_context_is_rejected(self) -> None: # A cache longer than the context it serves is meaningless, so the spec # rejects it rather than quietly shrinking it. - with pytest.raises(AssertionError, match="must not be greater than"): + with pytest.raises(ValueError, match="must not be greater than"): TraceSpec(max_context_length=512, cache_seq_len=513) def test_larger_context_leaves_trace_length_alone(self) -> None: