Skip to content

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

Open
Lewis300 wants to merge 3 commits into
apple:mainfrom
Lewis300:export-contract-hooks
Open

Refactor macOS export logic | Retrieve reference inputs / names / dynamic shapes from model class.#166
Lewis300 wants to merge 3 commits into
apple:mainfrom
Lewis300:export-contract-hooks

Conversation

@Lewis300

Copy link
Copy Markdown
Contributor
  • Moves the macOS export graph contract — which inputs exist, which are state, what they're called, how they're shaped — out of export/macos.py and onto the model class as overridable hooks.
  • No behavior change for anything that exports today: all 8 registered macOS models produce bit-identical reference inputs, shape bounds, and graph names.
  • Fixes two real bugs found on the way: exports with a context ≤ 2048 crashed inside torch.export, and activation calibration read the cache length off a guessed input position.

The hooks

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}

Keyed by graph name — a macOS model has one graph, main. The defaults are what macos.py hardcoded, so export_macos_model and the pipeline's quantization step now just call them, and validate_export_contract cross-checks that all five agree.

Two ordering rules: reference inputs bind to the traced signature, so they must be in its exact order; name lists are looked up by name, so each carries only the relative order of its own kind.

iOS classes inherit the macOS-shaped defaults, which are wrong in every respect, so they raise until the iOS follow-up implements them.

@tjia1818 tjia1818 assigned Lewis300 and unassigned srjoglekar246 Aug 11, 2026
Comment thread python/src/coreai_models/models/base.py Outdated
Comment thread python/src/coreai_models/export/compression.py Outdated

@carinapeng carinapeng left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Comments about error handling, everything else looks good. Thanks!

@Lewis300
Lewis300 force-pushed the export-contract-hooks branch from 7b5984e to 163a9d4 Compare August 12, 2026 03:25
`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.
@Lewis300
Lewis300 force-pushed the export-contract-hooks branch from 163a9d4 to aa84d2d Compare August 12, 2026 03:26
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants