Repository navigation
Functionalize coreai-models export + new functions for authored models to define KV Cache updates and forward method - #337
Merged
Conversation
guru-desh
requested review from
Lewis300,
carinapeng,
srjoglekar246 and
tjia1818
October 10, 2026 00:12
tjia1818
approved these changes
Oct 10, 2026
…d build_model
_async_export_model's resolution, loading and compression move, unchanged, into two
functions: resolve_model reads the checkpoint's config and resolves it against the
registry, and build_model loads the weights and compresses them. _async_export_model
calls both.
The moved lines are the original lines, re-indented where they leave the `with` block.
The one exception is an assert that ruff re-wraps once it is less indented. Review with:
git show --color-moved=dimmed-zebra --color-moved-ws=allow-indentation-change
On a tiny Qwen3, the model, config and externalized model handed to the exporter are
identical before and after, for macOS uncompressed, macOS 4-bit and iOS palettized.
The if/elif chain in _resolve_export_config moves, unchanged, into resolve_compression, so another entry point can resolve a compression exactly as the export CLI does. The only edit is `args.compression_config` becoming the `compression_config` parameter.
The "Resolve compression preset" block moves, unchanged, from build_model into resolve_model, and ResolvedModel carries the two configs it resolves. A caller can then see which compression a model will get before any weight is read. The one behaviour change: a config naming both a quantization and a palettization now fails before the weights load rather than after.
quantize_for_export already takes the backend coreai-opt finalizes quantized weights for; build_model now passes one through, defaulting to Core AI as before. With another backend nothing is exported, so graph-mode quantization skips marking composite ops for externalization.
macOS still loads through from_hf_memory_efficient, with mmap_path=None, and the quantizer skips its disk checkpoints. The export always passes a temp_dir on macOS, so it is unaffected.
The config that exports a registry preset, with any value passed in winning over the preset's, for callers that start from a registry preset rather than from CLI arguments. The export CLI keeps its own resolution, unchanged.
new_eager_state, eager_forward_args and eager_step call forward on real inputs, one sequence at a time. macOS takes positions over the whole history and the KV pair; iOS takes only the query's positions, the write offset, and an additive mask over the context, which causal_mask builds. Export's palettization keeps its own trace inputs, so what it traces is unchanged.
guru-desh
force-pushed
the
functionalize-model-export-code
branch
from
October 10, 2026 00:35
4eb43b4 to
ce7cec7
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
2 participants
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.
This PR does the following:
ExportConfigto have afrom_presetmethodasync_export_modelBaseForCausalLMto define new_eager_state and eager_forward_args which define how KV Cache is stored and how a forward pass works. This is overriden inBaseForCausalLMForiOS