Refactor macOS export logic | Retrieve reference inputs / names / dynamic shapes from model class. - #166
Open
Lewis300 wants to merge 3 commits into
Open
Refactor macOS export logic | Retrieve reference inputs / names / dynamic shapes from model class.#166Lewis300 wants to merge 3 commits into
Lewis300 wants to merge 3 commits into
Conversation
tjia1818
approved these changes
Aug 11, 2026
gokulkrishna98
approved these changes
Aug 11, 2026
carinapeng
reviewed
Aug 12, 2026
carinapeng
reviewed
Aug 12, 2026
carinapeng
approved these changes
Aug 12, 2026
carinapeng
left a comment
Contributor
There was a problem hiding this comment.
Comments about error handling, everything else looks good. Thanks!
Lewis300
force-pushed
the
export-contract-hooks
branch
from
August 12, 2026 03:25
7b5984e to
163a9d4
Compare
`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
force-pushed
the
export-contract-hooks
branch
from
August 12, 2026 03:26
163a9d4 to
aa84d2d
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
export/macos.pyand onto the model class as overridable hooks.torch.export, and activation calibration read the cache length off a guessed input position.The hooks
Keyed by graph name — a macOS model has one graph,
main. The defaults are whatmacos.pyhardcoded, soexport_macos_modeland the pipeline's quantization step now just call them, andvalidate_export_contractcross-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.