Refactor iOS export logic - #168
Open
Lewis300 wants to merge 1 commit into
Open
Conversation
iOS models previously raised NotImplementedError from every export-contract hook and export/ios.py built its own reference inputs, graph I/O names, static shapes and hardware constraints inline. Move all of that onto BaseForCausalLMForiOS so the contract describes the iOS graphs, and reduce export/ios.py to tracing and conversion. iOS emits four entrypoints over three callables: the embedding-table loader, the token gather, and the transformer traced twice with prefill off and on. The two transformer entrypoints have identical inputs, states and outputs, so the contract carries three entries and the exporter maps both onto the transformer entry. Also: - Move the iOS graph I/O names to _constants.py, shared by the contract, the exporter and the tests. - Add export_state_output_names, export_static_shape_configs and export_hardware_constraints hooks; iOS needs the static shape ladder and the buffer layout constraints, both per-graph. - Give quantize_for_export a clear error on multi-graph models, which is what `--variant iOS --compression 4bit` reaches now that the hooks no longer raise.
tjia1818
approved these changes
Aug 13, 2026
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.
Make the iOS export use the contract
Follow-up to the macOS contract refactor. Does the same thing for iOS: moves the graph
contract out of
export/ios.pyand onto the model class, replacing the hooks that wereleft raising
NotImplementedError.export/ios.pydrops from ~320 lines of inline construction to tracing and conversiononly.
ladder, the buffer layout constraints, and the state output names.
_constants.py, so the contract, the exporter, andthe tests read one definition instead of importing them from the exporter.
ios.pyhardcoded.Four entrypoints, three contract entries
iOS emits four entrypoints over three callables:
load_embeddingsmodel.load_embeddingsload_embeddingsgather_embeddingsmodel.gather_embeddingsgather_embeddingsextendmodel.extend(prefill off)extendprompt_optmodel.extend(prefill on)extendThe two transformer entrypoints have identical inputs, states, and outputs — they differ
only by module state — so the contract carries three entries and the exporter maps
both onto the
extendentry. Which callable each entrypoint traces, and the prefilltoggle between them, stays in
export/ios.py: that's export-side, not contract.This is also why
model.forwardcan't describe the graph. It composes the threecallables for eager use and is never exported, so the inherited single-graph defaults
described the wrong graph entirely — which is what the previous PR's
NotImplementedErrorwas standing in for.The new hooks
export_state_output_names— iOS names its state outputs because the layoutconstraints attach to the mutated output as well as the input. The macOS converter
surfaces state in/out implicitly, so it doesn't need this.
export_static_shape_configs— iOS compiles for fixed shapes, so each graph isbuilt once per shape it must serve: the transformer over (cache length, query length),
the gather over query length alone. Cache length starts at 256 and doubles up to the
context.
export_hardware_constraints—IOSurfaceallocation and interleave/alignment forthe buffers the runner shares with the compiled graphs; both sides must agree on
layout.
All three are keyed by graph name like the existing hooks, and an empty entry means the
graph needs none. The tuning knobs (
IOS_QUERY_LEN,IOS_STATIC_QUERY_LENS,IOS_STATIC_MIN_CACHE_LEN,KV_CACHE_INTERLEAVE_FACTOR) are class attributes, so amodel with different requirements overrides rather than forks the method.
export_ios_modelnow also runsvalidate_export_contract, which iOS was previouslyunable to reach.
Ordering
Unchanged from the macOS PR, and the iOS graphs exercise both rules in a way macOS
didn't:
extenddeclaresembedding_tableafter both caches, so the dict interleaves inputs and state.kind.
gather_embeddingstakes a parameter calledinput_idsbut declares it asin_new_token_ids; graph names are not parameter names.Tests
New
test_ios_export_contract.py(24 tests) covers the three-entries-not-four shape,both ordering rules, the shape ladder and its scaling with context, the constraints, and
tracing through
_export_programs— including that the two transformer programs actuallydiffer. These build a tiny Qwen3 iOS model and stop before the Core AI converter, so they
need no runtime.