Model Authoring

---
name: model-authoring
description: Empirical rules for authoring PyTorch models for on-device execution on Apple platforms, covering energy-efficient inference, scalable compute, and correctness testing. Use this skill whenever the user is writing, debugging, or reviewing PyTorch model code intended for on-device execution — even if they don't explicitly mention Neural Engine or Core AI. Covers BC1S layout, op compatibility, KV cache patterns, precision rules, PSNR verification, activation functions, and common issues.
---

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

| Resource | When to consult |
| ------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
| [`neural_engine_rules.md`](references/neural_engine_rules.md) | Neural Engine patterns: BC1S layout, Conv2d projections, per-head attention, KV cache readonly pattern, stride/dilation/pooling rules, causal mask, RoPE, chunked prefill |
| [`gpu_rules.md`](references/gpu_rules.md) | GPU patterns: fused QKV, native SDPA, KV cache stateful pattern, MoE (GatherMM/SwitchLinear), memory-efficient loading, RMSNorm variants |
| [`common_issues.md`](references/common_issues.md) | Debugging: PSNR issues, compilation errors, runtime problems, stale flags |
| [coreai-models repo](https://github.com/apple/coreai-models) | Complete working reference implementations for LLMs, vision, audio, diffusion. Explore `primitives/` and `models/` directories. |

### 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](https://apple.github.io/coreai-torch/guides/composite-ops.html) | [externalization](https://apple.github.io/coreai-torch/guides/externalization.html) | [composite ops API](https://apple.github.io/coreai-torch/api/composite-ops.html)