Skip to content

Refactor iOS export logic - #168

Open
Lewis300 wants to merge 1 commit into
apple:mainfrom
Lewis300:lpanos/ios-export-contract
Open

Refactor iOS export logic#168
Lewis300 wants to merge 1 commit into
apple:mainfrom
Lewis300:lpanos/ios-export-contract

Conversation

@Lewis300

Copy link
Copy Markdown
Contributor

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.py and onto the model class, replacing the hooks that were
left raising NotImplementedError.

  • iOS models now implement the five contract hooks instead of raising.
    export/ios.py drops from ~320 lines of inline construction to tracing and conversion
    only.
  • Adds three iOS-only hooks for things macOS has no equivalent of: the static shape
    ladder, the buffer layout constraints, and the state output names.
  • Moves the iOS graph I/O names into _constants.py, so the contract, the exporter, and
    the tests read one definition instead of importing them from the exporter.
  • No behavior change intended — the hooks return the names, shapes, and constraints
    ios.py hardcoded.

Four entrypoints, three contract entries

iOS emits four entrypoints over three callables:

entrypoint callable contract entry
load_embeddings model.load_embeddings load_embeddings
gather_embeddings model.gather_embeddings gather_embeddings
extend model.extend (prefill off) extend
prompt_opt model.extend (prefill on) extend

The 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 extend entry. Which callable each entrypoint traces, and the prefill
toggle between them, stays in export/ios.py: that's export-side, not contract.

This is also why model.forward can't describe the graph. It composes the three
callables 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
NotImplementedError was standing in for.

The new hooks

export_state_output_names()                        -> {graph: (name, ...)}
export_static_shape_configs(config, max_ctx)       -> {graph: {label: {name: shape}}}
export_hardware_constraints(max_ctx)               -> {graph: {name: HardwareConstraints}}
  • export_state_output_names — iOS names its state outputs because the layout
    constraints 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 is
    built 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_constraintsIOSurface allocation and interleave/alignment for
    the 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 a
model with different requirements overrides rather than forks the method.

export_ios_model now also runs validate_export_contract, which iOS was previously
unable to reach.

Ordering

Unchanged from the macOS PR, and the iOS graphs exercise both rules in a way macOS
didn't:

  • Reference inputs bind to the traced signature — exact order. extend declares
    embedding_table after both caches, so the dict interleaves inputs and state.
  • Name lists are looked up by name — each carries only the relative order of its own
    kind. gather_embeddings takes a parameter called input_ids but declares it as
    in_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 actually
differ. These build a tiny Qwen3 iOS model and stop before the Core AI converter, so they
need no runtime.

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.
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.

2 participants