Model Authoring

This skill contains the hard-won empirical knowledge for making PyTorch models compile and run correctly on Apple hardware via Core AI. The rules here are stable across Core AI releases — they reflect hardware behavior, not API shapes.

Reference material

Use these resources on-demand — do not read all files upfront. Consult the relevant reference when the user's task requires specific patterns for a target platform, or when debugging.

coreai-models: working reference implementations

For complex models (LLMs, MoE, multimodal, diffusion), explore the coreai-models repo before writing primitives from scratch. It has complete authoring primitives for both GPU and Neural Engine, including advanced patterns like iOS embedding quantization, MoE routing, and memory-efficient weight loading for large models. If the user has a local clone, explore it directly. If not, suggest cloning it.

Online docs: coreai-torch composite ops | externalization | composite ops API

Model optimization — use working-with-coreai

Model optimization decisions (precision, compression, device compatibility) are resolved by the working-with-coreai skill.

• If the active plan contains deployment decisions (platform, compression approach), follow those. The plan uses "optimize for energy efficiency" (BC1S, Conv2d, static shapes, fp16) and "optimize for scalable performance" (standard layout, nn.Linear, dynamic shapes supported).
• If no deployment context exists and the user's intent is ambiguous, invoke Skill("coreai-skills:working-with-coreai") before authoring.

If the user explicitly names an accelerator (Neural Engine, GPU, CPU), use their choice. Otherwise, infer from context and use outcome-oriented language in your responses — say "optimized for energy-efficient inference on iPhone" rather than "targets Neural Engine". Mirror the user's vocabulary: if they say Neural Engine, match them.

Compute unit characteristics

Neural Engine and GPU at a glance

Quick reference for the key authoring differences. Consult neural_engine_rules.md or gpu_rules.md for full details.

Authoring workflow

Phase 1: Architecture discovery

Run code, don't read code. Running gives ground truth instantly.

1. Print model structure and state dict keys with shapes
2. Trace forward pass with register_forward_hook — capture intermediates
3. Document target hardware, IO boundary, module hierarchy, activation type, KV cache layout

Phase 2: Primitive implementation (bottom-up)

Author in this order — each depends on the previous:

1. Norm — layout and weight shape depend on target
2. Linear projections — Conv2d(in, out, 1) for Neural Engine; nn.Linear for GPU
3. Attention — layout, K@Q convention, causal mask depend on target
4. MLP / FFN — activation must match source exactly
5. Full decoder block — compose primitives with KV cache wiring

Verification gates

Verify each primitive individually before composing the full model. Also compare the full re-authored model's outputs against a baseline export (direct from HuggingFace without re-authoring) — both in Python and after compilation on device — to confirm end-to-end parity.

The from_source_model classmethod

Every re-authored model gets a factory classmethod — no hardcoded constants:

@classmethod
def from_source_model(cls, source_model) -> "MyDecoder":
    cfg = source_model.config
    model = cls(
        n_layers=cfg.num_hidden_layers,
        hidden=cfg.hidden_size,
        n_heads=cfg.num_attention_heads,
        # ...
    )
    model.load_weights_from(source_model.state_dict())
    return model

KV cache conventions

Both Neural Engine and GPU require explicit KV cache management, but the patterns differ:

Key rule: Do not use stateful transforms for token generation — state resets between inference calls. Use the readonly KV I/O pattern (Neural Engine) or the stateful export wrapper (GPU) instead.

Palettization (weight compression)

Apply after authoring float16 model passes verification, before Core AI export.

For compression exploration and configuration, use Skill("coreai-skills:model-compression-exploration") which covers coreai-opt quantization and palettization sweeps.

Key facts for authoring:

• 4-bit palettization: ~4x size reduction, PSNR ~40 dB vs float16
• Palettize Conv2d / Linear only — skip embeddings, norms, bias
• State dict keys gain ._data, ._lut, ._indices suffixes after compression
• Size reduction is realized in the compiled .aimodel, not in the PyTorch checkpoint

Working with Core AI

Deploy PyTorch models on Apple silicon: export with coreai-torch, compile with coreai-build, run with the Core AI runtime (Swift or Python).

Related skills: Skill("coreai-skills:model-authoring") (Neural Engine and GPU authoring patterns, use when re-structuring model architecture) | Skill("coreai-skills:model-compression-exploration") (quantization/palettization sweeps — use when exploring compression tradeoffs)

Documentation and reference material

The Core AI toolchain has extensive documentation. Use these as reference — do not read all pages upfront. Instead, consult the relevant docs when you need specifics about a particular step.

coreai-models: the reference implementation

The coreai-models repo is the canonical source for how to export and run models with Core AI. Before writing export code from scratch, always explore this repo — it has working export recipes for many model families, Swift and Python runtime utilities, and reusable primitives. If the user has a local clone, explore it. If not, suggest cloning it.

Explore these directories to find relevant patterns:

• models/ — Per-model export recipes with READMEs and CLI commands for many popular model families (LLMs, vision, audio, diffusion).
• python/src/coreai_models/export/ — Export pipeline code covering macOS and iOS export paths, compression presets, and custom MLIR lowerings.
• swift/Sources/ — Runtime utilities for LLMs (engines, text generation, KV cache, sampling, decode loops), diffusion pipelines, object detection, image segmentation, and constrained decoding.

Pipeline overview

The Core AI pipeline transforms a PyTorch model into an optimized on-device asset:

1. AUTHOR        Re-structure model for target platform
                  → Skill("coreai-skills:model-authoring")

2. COMPRESS      Explore quantization/palettization tradeoffs
                  → Skill("coreai-skills:model-compression-exploration")

3. EXPORT        Convert PyTorch → AIProgram via TorchConverter
                  → coreai-torch docs

4. COMPILE       Ahead-of-time compilation for target platform
                  → coreai-build CLI

5. RUN           Load and run on device (Swift or Python)
                  → CoreAI framework / coreai Python API

Steps 1 and 2 are optional — many models export directly without re-authoring or compression. Start with export, then add authoring or compression if needed (poor accuracy, poor performance, too large).

For models already in coreai-models, the export recipes handle all steps. Check the models/ directory first — if the user's model family is there, point them to the recipe.

Export (Python — coreai-torch)

import torch
from coreai_torch import TorchConverter, get_decomp_table

model = MyModel().eval()
ep = torch.export.export(model, args=(torch.randn(1, 3, 224, 224),))
ep = ep.run_decompositions(get_decomp_table())

program = (
    TorchConverter()
    .add_exported_program(ep, input_names=["image"], output_names=["logits"])
    .to_coreai()
)
program.optimize()
program.save_asset("model.aimodel")

This is the simplest export pattern. Real models often need more — consult the coreai-torch docs and explore the export code in the coreai-models repo for patterns around:

• Externalization of composite ops via add_pytorch_module() with externalize_modules
• Mutable state (e.g. KV cache) via state_names
• Custom Metal kernels via TorchMetalKernel and register_torch_lowering()
• iOS static shape specialization via set_static_shape_config()
• Compression presets for macOS vs iOS (different default strategies per platform)

Compile (coreai-build CLI)

Ahead of time (AOT) compilation of models can optionally be performed with:

xcrun coreai-build compile model.aimodel --platform iOS

Docs: Ahead-of-time compilation

Run (Swift)

import CoreAI

let model = try await AIModel(contentsOf: modelURL)
guard let fn = try model.loadFunction(named: "main") else { return }

var input = NDArray(shape: [1, 3, 224, 224], scalarType: .float32)
var view = input.mutableView(as: Float32.self)
// fill view with data...

var outputs = try await fn.run(inputs: ["image": input])
let result = outputs.remove("logits")?.ndArray

For LLMs, diffusion, and other complex models, explore the Swift runtime utilities in the coreai-models repo — they provide complete inference engines, decode loops, sampling, and KV cache management that handle the complexity beyond basic AIModel usage.

Docs: CoreAI framework

Run (Python)

from coreai.runtime import AIModel, NDArray
import numpy as np

model = await AIModel.load("model.aimodel")
fn = model.load_function("main")
outputs = await fn(
    {"image": NDArray(np.random.randn(1, 3, 224, 224).astype(np.float32))}
)
logits = outputs["logits"].numpy()

Docs: coreai Python API

Verifying outputs (PSNR)

Compare on-device outputs against PyTorch reference:

Onboarding new models

When helping a user deploy a model they haven't deployed before, follow this protocol. If the user already knows what they want (specific export question, compilation flag, etc.), skip directly to the relevant section above.

First response: conversation, not code

Your first response is always a conversation. Even if the deployment path seems obvious, ask first.

Rules:

• Summarize what you understood (2-3 sentences)
• Ask 1-2 clarifying questions about what's unclear
• Do not recommend or dive into compute unit specific optimizations unless you have already reviewed the general and platform level guidance in guidance.md
• Do not produce compilation commands, code, or technical recommendations yet
• Only discuss what's on the user's path — don't mention irrelevant options
• Always check coreai-models for existing recipes before suggesting a from-scratch approach

What to learn:

• What are they building? (real-time camera, batch processing, background service, LLM chat)
• What platform? (iOS, macOS, multi-platform)
• What does the model look like? (architecture, size, input shapes, dynamic vs static)
• What matters most? (speed, energy efficiency, quality, broad device reach)

Example interaction:

User: I want to deploy my image classifier on iPhone.

Assistant: Got it — you're looking to run an image classification model on iPhone. A couple of questions to help me give you the best guidance:

1. Roughly how large is the model (parameter count or file size)?
2. Does it need to run in real-time (e.g., processing camera frames), or is batch/offline processing fine?

User: It's a ResNet-50, about 25M params. Needs to run at 30fps on camera frames.

Assistant: Perfect — a 25M param ResNet-50 for real-time camera on iPhone is a great fit for efficient on-device inference. Here's the plan...

After the conversation

1. Check if the model family exists in coreai-models — explore the models/ directory for a matching recipe
2. If the user's needs involve platform targeting, model sizing, or compression strategy, read guidance.md to resolve the deployment path — present the outcome, not the reasoning
3. Walk through the pipeline steps relevant to their situation, consulting the docs above as needed
4. If the model needs architectural changes, invoke Skill("coreai-skills:model-authoring")
5. If compression tradeoffs need exploration, invoke Skill("coreai-skills:model-compression-exploration")

Model Compression Exploration

Systematically explore weight-only compression configurations for a PyTorch model using coreai_opt. The goal is to present the user with a clear overview of accuracy-vs-size tradeoff options across quantization and palettization, organized into three experiment groups.

Supporting files

Bundled scripts

The deterministic helpers are unit-tested and importable. Prefer them over hand-rolled equivalents — they encode formulas and edge cases that have already been debugged.

CoreAI Opt

CoreAI Opt (coreai-opt) is a package that helps with model compression and model optimization in a hardware-aware manner.

For the full coreai-opt documentation, fetch: https://apple.github.io/coreai-optimization/llms-full.txt

Check to see that the package is installed in the current python scope (venv, conda env). The package is called coreai-opt and is imported as coreai_opt. If it is not installed, prompt the user to install it.

For API verification at runtime, use help(coreai_opt) or inspect to confirm current signatures.

Setup

Step 1: Gathering input from user

The user has to provide information on how to load the model, how to perform a forward pass, what are the inputs to be used and how to check the quality of the outputs. This information is very important to acquire from the user, since every model is different and making assumptions can lead us to meaningless results. For example, choosing to use random inputs instead of a valid input can result in the output quality being meaningless.

1. Model: get_model() -> nn.Module - How to load/create the model (imports, weights, model class).
2. Data: get_reference_data() -> tuple[torch.Tensor, ...] | dict[str, torch.Tensor] - A representative batch of real inputs. Even 1-3 real samples suffice — random inputs produce meaningless PSNR because they don't exercise learned weight structure.
3. Forward pass: verify that get_model()(*get_reference_data()) (or the dict-spread equivalent get_model()(**get_reference_data())) actually runs. If neither works, ask the user how to invoke the model end-to-end.
4. Quality Metric: get_quality_metric(model_out) -> list[str] - we need to understand what metric to use for checking the quality of each output against the uncompressed output. Ask the user for each output produced by the model, should we use one of {"psnr", "snr", "iou"}. For example, if we have a mask as an output, PSNR isn't the right metric. IoU is a right metric.

This information is required to proceed to the next step.

Step 2: Check the setup

• Compute the uncompressed baseline output (store as detached tensor) and the uncompressed baseline size (total parameters × 2 bytes, assuming fp16 storage). The bundled _call_model helper in references/experiment_runner.md handles both tuple- and dict-shaped reference data; reuse it everywhere you call the model.

Step 3: Estimate the time it takes for doing our compression exploration

1. Take the default global weight quantization preset QuantizerConfig.presets.w8() (graph mode is the default). Apply it to a fresh model and time a single forward pass through the prepared model. If Quantizer.prepare(...) errors — e.g., torch.export guard failure, dynamic control flow — fall back to QuantizerConfig.presets.w8(execution_mode=ExecutionMode.EAGER) and time again. The mode that succeeded here is the mode you should use for the entire sweep, so the timing reflects real wall-clock cost. Record this mode and reuse it.

   The single elapsed time becomes avg_quant_time. Pass quantizer=quantizer to extract_layer_specs(...) so it can read graph-mode FQ metadata via Quantizer._get_fake_quantize_modules(); otherwise the walker won't see graph-mode quantization and would misreport every layer as fp16.

2. Take the default palettization preset KMeansPalettizerConfig.presets.w6() (6-bit, per-grouped-channel, group_size=16). Apply the palettization config and run a forward pass while calculating the time it takes to compute a palettized model pass. This will be the average time it takes to run a single palettization pass: avg_palett_time. Palettization is eager-only — there's no graph/eager fallback to do here.

3. Below, we enumerate 3 groups of config options, totaling around ~15 quant configs and ~15 palett configs. Estimate the time required as avg_quant_time * 15 + avg_palett_time * 15.

4. Ask the user if this time estimate is in-line with their expectation before proceeding. Use the AskUserQuestion tool here to provide the estimate and ask if it is okay to proceed, or if they want to cut short the time.

Step 4: Experimentation

How to run each experiment

Use coreai_opt.quantization.Quantizer for Groups 1-2 and coreai_opt.palettization.KMeansPalettizer for Group 3. Run the loop in references/experiment_runner.md (memory-safe, plus the canonical extract_layer_specs(prepared, quantizer=compressor) pattern that works in both modes). The execution mode was decided once in Step 3 — use that mode for every config in the sweep.

For each config:

1. Re-create a fresh model
2. Apply compression via prepare() by loading the config
3. Compute theoretical size, average bitwidth and compression ratio using references/size_estimation.md
4. Run a forward pass on the prepared model in eval() mode with no_grad
5. Compute the per-output quality metrics chosen by the user
6. Append the record to results.jsonl (Output Report section)
7. Free memory (template handles this)
8. Do not call finalize(). Calibration is not needed for weight-only compression.

What are the experiments to run

Build configs through QuantizerConfig.presets / KMeansPalettizerConfig.presets where the shape matches; for the variations they don't cover (asymmetric, symmetric_with_clipping, alternative block sizes, enable_per_channel_scale=True), see references/experiment_runner.md for the spec-construction patterns. Verify the preset namespace at runtime with dir(QuantizerConfig.presets) and dir(KMeansPalettizerConfig.presets) — new presets are added over time.

1a: Channel-structured quantization — 6 configs

Cross-product of {int8, int4} × {symmetric, asymmetric, symmetric_with_clipping}, all per-channel. The two symmetric corners match QuantizerConfig.presets.w8() and .w4() directly; the other four are variations that swap qscheme=.

1b: Block-structured quantization — 9 configs

Cross-product of {block_size: 16, 32, 128} × {symmetric, asymmetric, symmetric_with_clipping}, all int4 per-block. The block_size=32, symmetric corner matches QuantizerConfig.presets.w4_per_block(block_size=32); the rest swap block_size= and qscheme=.

Scale overhead reminder: per-block stores one fp16 scale per block. At block_size=16 with int4, effective bitwidth is ~5 bits/weight — account for this in compute_average_bitwidth (it already does).

2: Palettization — 15 configs

Cross-product of {(8-bit, per-tensor), (6-bit, per-tensor), (6-bit, gs=4|8|16), (4-bit, gs=4|8|16)} × {enable_per_channel_scale: True, False} minus the one undefined entry (8-bit per-tensor with enable_per_channel_scale=True is sometimes folded into the 8-bit per-tensor row — keep both for completeness, totaling 15). The (8-bit, per-tensor, False) corner matches KMeansPalettizerConfig.presets.w8(); (6-bit, gs=16, False) matches presets.w6(); (4-bit, gs=16, False) matches presets.w4().

Per-Group Refinement

After the main sweep within a group:

1. Filter — drop configs that errored or scored below the floor (PSNR < 10 dB or IoU < 0.1). These are too far gone for layer-skipping to rescue.
2. Pick two refinement seeds per group:
   • 95th-percentile config — best surviving quality, modest size win. Refining this tells us how much smaller we can go without losing quality.
   • 75th-percentile config — mid-quality, larger size win. Refining this tells us how much quality we can recover at an aggressive compression target.
3. Run 5 layer-skip variants per seed (10 extra runs/group):
   • Skip first layer
   • Skip last layer
   • Skip first and last layer
   • Skip all layers of the smallest-aggregate-parameter type, breaking ties by Embedding > Linear > Conv. Compute parameter counts per type and pick the type with the smallest sum so compression ratio barely moves.
   • Skip first/last and the smallest-parameter type — the safest combination.
4. Apply skips via set_module_name overrides on top of the seed's preset. Refinement runs inherit any divisibility overrides from the seed — don't rebuild the config from scratch.
5. Sub-models — for multi-modal architectures (ViT backbone + text encoder feeding an encoder-decoder), use model.named_children() to enumerate top-level submodules. Boundary layers exist within each submodule; the "first/last layer" skip should consider entry/exit projections of each major child, not only the outermost first/last of the whole model.

Output Report

Use a JSON structure to track all the details of the experiment. We want to track the following:

{
  "group": "2",
  "config": {
    "name": "palette_grouped_gs4_6bit_pcs0_skip-Embedding",
    "path": "path/to/config",
  },
  "time_taken": 1000,
  "output_quality_metrics": [
    {"name": "bbox", "metric": "iou", "value": 0.7},
    {"name": "logits", "metric": "psnr", "value": 16}
  ],
  "compression_metrics": {
    "average_bitwidth": 5,
    "compression_ratio": 1.7,
    "theoretical_model_size": 402
  }
}

After all sweeps complete, the JSONL holds 40-50 records — too many to surface to the user verbatim. For each group, pick exactly 5 configs that span the accuracy-vs-size tradeoff and put only those in the report. Concretely, after filtering out configs that errored or fell below the floor (PSNR < 10 dB / IoU < 0.1):

1. Highest quality — best primary-metric value in the group.
2. Highest compression — best compression ratio in the group (after the floor filter).
3. Three points on the frontier between (1) and (2) — pick configs that maximize spread, not similarity. A simple rule: sort survivors by compression ratio, then take the configs whose (quality, ratio) points are furthest from the line connecting (1) and (2). If two configs have nearly identical (quality, ratio), prefer the one with the simpler config (fewer overrides, larger block/group size).

The goal is that a reader scanning the table can see the shape of the tradeoff in one glance: not 30 indistinguishable rows, but 5 anchors covering the frontier from "barely compressed, near-perfect quality" to "maximum compression, quality at the floor". If a group has fewer than 5 survivors, surface them all and note the count.

Produce one report per group with these columns:

Refer to references/output_report.md for more details on output formatting. Following this format consistently keeps the qualitative picture comparable across runs; consumers grep these tables to compare model variants.

Generate a PSNR-vs-compression-ratio scatter plot (matplotlib) with annotated config names. Save as compression_exploration.png and include the link in the report.

Parallelization

The full sweep is ~30 main-sweep configs + ~30 refinement configs × per-config-time. Launch one subagent per group (1a, 1b, 2) so they run in parallel:

• Agent "group-1a runner" → channel-structured quant (6 configs) + its 10 refinement runs
• Agent "group-1b runner" → block-structured quant (9 configs) + its 10 refinement runs
• Agent "group-2 runner" → palettization (15 configs) + its 10 refinement runs

Each agent appends to a shared results.jsonl file (one JSON record per line). JSONL append is safe in practice when each agent writes one complete line at a time — use a flush after each write. The main agent uses /loop 5m (a slash command in this repo's plugin set that re-runs a prompt on a schedule) to read results.jsonl and report per-group completed/total to the user. Long sweeps look hung without progress signals; surface counts so the user can ctrl-C if something is clearly broken.

See references/experiment_runner.md for the append_record() and status_snapshot() helpers.

Common Pitfalls

Read references/compression_patterns.md for the full list. Critical ones:

1. Silent skip on indivisible block/group size: Per-block and per-grouped-channel silently skip layers where the weight dimension isn't divisible. Pre-check with check_divisibility() and override with per-channel.
2. Graph mode failures: covered by the dry-run probe in Step 3 — fall back to execution_mode=ExecutionMode.EAGER for the whole sweep. Always pass quantizer= to extract_layer_specs(...) so it works in either mode.
3. Axis defaults: coreai-opt picks the correct default axis per module type. Only override for non-standard behavior.
4. Scale/ZP overhead: At 2-4 bit with fine granularity, overhead can be 5-15% of total size.
5. LUT overhead: 8-bit per-channel stores 256 × fp16 entries per output channel — significant for wide layers.