Model Authoring

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