Core AI (WWDC26): Apple's New On-Device Inference Stack

Part of the ANE Knowledge Base. Sources (fetched 2026-07-12): developer.apple.com/core-ai, the three WWDC26 sessions (324 — Meet Core AI, 325 — Dive into Core AI model authoring and optimization, 326 — Integrate on-device AI models into your app using Core AI), and the three repos pinned under references/: coreai-models (85e2f2d), coreai-torch (1b3cb3b, v0.4.1), coreai-optimization (5cdb1f1, v0.2.0).

At WWDC26 Apple introduced Core AI: the inference framework that powers on-device Apple Intelligence, opened to developers as "the next evolution of on-device AI execution across Apple platforms". It covers the full lifecycle — PyTorch authoring → optimization → conversion → AOT compilation → Swift integration → debugging/profiling — targeting everything from small vision models to 70B-parameter LLMs, entirely on device.

Positioning vs Core ML. WWDC26 has zero Core ML sessions (verified against the WWDC26 listing); Core AI got three, plus [330 — Optimize custom ML operations with Metal tensors]. Apple did not announce a Core ML deprecation; press coverage (InfoQ and others) describes Core AI as the successor for neural networks/transformers, with Core ML remaining for classic ML and MLX for research/custom weights. Requirements from the repos: iOS/macOS 27.0+, Xcode 27.

1. Architecture#

Piece	What it is
.aimodel	Source model asset, device-independent; inspectable in Xcode's model viewer (size, op distribution, function signatures, dynamic dims marked ?)
Specialization	On-device, per-hardware compilation at first load: segment/plan/optimize + generate device/OS-specific binaries; cached (AIModelCache, shareable across app groups)
AOT compilation	xcrun coreai-build compile MyModel.aimodel --platform iOS — moves most compilation to the dev machine; device only finishes specialization. Xcode-integrated
Runtime	Uses CPU, GPU, and ANE; fine-grained memory control (pre-allocation, zero-copy data paths), stateful execution
Swift API	AIModel → loadFunction(named:) → InferenceFunction.run(inputs:) with NDArray I/O; non-escapable MutableViews for memory-safe zero-copy access; states passed as InferenceFunction.MutableViews and updated in-place
Distribution	Models kept out of the app bundle; Background Assets for opt-in download

The "Meet Core AI" demo is directly on-theme for this KB: a transformer whose per-step latency grew with sequence length was fixed by adding KV-cache states (NDArray buffers passed as mutable state views) — constant-latency decode, the same pattern as doc 06 §4, now first-class in the API.

Foundation Models bridge#

CoreAILanguageModel (from the coreai-models Swift package) plugs a custom Core AI LLM into the Foundation Models LanguageModelSession API — same streaming and @Generable guided/structured generation as Apple's built-in model (session 326; see also WWDC26 session 339 "Bring an LLM provider to the Foundation Models framework").

2. The Python stack#

coreai-torch (conversion) — references/coreai-torch#

• Based on torch.export (not torch.jit.trace as in the Core ML era — doc 05): export with dynamic_shapes support → run_decompositions(coreai_torch.get_decomp_table()) → TorchConverter().add_exported_program(...).to_coreai() → save_asset("X.aimodel").
• Composite ops (coreai_torch/composite_ops/): _sdpa.py (incl. causal variants and sliding-window attention), _rms_norm.py, _rope.py, _gather_mm.py, _gated_delta_update.py — high-level ops preserved through the decomp table so the compiler maps them to pre-optimized hardware kernels (SDPA, LayerNorm/GroupNorm…). Hand-decomposing attention is no longer the author's job.
• Multi-entrypoint conversion: several exported functions → one .aimodel with multiple callable functions (e.g. SAM3 re-authored as image_encode / text_encode / detect; cached image embeddings gave 76% faster subsequent runs — session 325).
• Custom Metal kernels from Python: TorchMetalKernel DSL embeds MSL source (with a PyTorch reference implementation for tracing/verification) directly into the .aimodel.
• Numerical verification against PyTorch is a first-class workflow (coreai.runtime.AIModel in Python, compare logits, assert max-diff).

coreai-opt (compression) — references/coreai-optimization#

pip install coreai-opt. The productized descendant of the ct.optimize / argmaxtools-Palettizer lineage (doc 07):

• Schemes: INT4/INT8, FP4/FP8 quantization, k-means palettization, pruning; post-training and QAT; granularities per-tensor / per-channel / per-grouped-channel.
• Config-driven API: QuantizerConfig.presets.w4().without([nn.LayerNorm, "lm_head"]), presets.w8().only_for(nn.Linear, nn.Conv2d) (src/coreai_opt/config/compression_config.py) — sensitivity-aware selective compression as a one-liner.
• save_intermediates + the Core AI Debugger's comparison mode (PSNR/MSE/MAE, green/yellow/red per layer) = the per-layer divergence-profiling workflow of doc 07 §3, built into the toolchain.

coreai-models (model zoo + recipes) — references/coreai-models#

• models/: export recipes for whisper, qwen2/3(+MoE), gemma3, mistral/mixtral, gpt_oss, sam3, stable-diffusion, flux2, clip, clap, depth-anything, yolo, t5, roberta, wav2vec2, vlm… — single-command uv run scripts (PEP 723 inline deps; pinned coreai-core==1.0.0b2, coreai-torch==0.4.1).
• Declarative mixed-bit recipes: e.g. models/qwen3/qwen3_0_6b_mixed_4bit_8bit.yaml — global 4-bit per-grouped-channel (group 8), specific sensitive layers (regex-matched) at 8-bit per-tensor, embeddings skipped. The argmaxtools mixed-bit recipe (doc 07 §3), shipped as YAML.
• python/src/coreai_models/: reusable authoring primitives and per-model classes with an explicit platform registry (ios_class / macos_class in models/registry.py); export CLI takes --platform iOS.
• swift/: runtime packages (CoreAILanguageModels, segmentation, etc.) that hide tensor wrangling.
• skills/: agent skills (Claude Code / Codex / Gemini CLI) — model-authoring, model-compression-exploration, working-with-coreai.

3. Do the KB's ANE principles survive? — Yes, verbatim#

The most important finding for this KB. coreai_models/primitives/ has separate ios/ and macos/ implementations — the ANE-vs-GPU split made explicit at the framework level:

KB principle (2022)	Core AI iOS primitives (2026)
P1 — BC1S layout	ios/sdpa.py: query/key/value shaped (batch, n_heads*head_dim, 1, seq_len) — exactly BC1S
P1 — Conv2d 1×1, not Linear	ios/mlp.py: "uses Conv2d layers instead of Linear layers for better iOS performance" — gated-SiLU MLP from three 1×1 convs
P2 — per-head chunking	ios/sdpa.py docstring: "iOS requires each attention head to be computed individually to meet hardware constraints" — per-head split + per-head softmax over the key axis (softmax(1)), GQA-aware (kv_group_size)
P3 — minimize copies	keys transposed/permuted once in advance before the per-head loop
FP16 hygiene	scale folded into K before QKᵀ "for numerical stability"
KV cache w/ static shapes	ios/cache.py: 5D cache [n_layers, batch, n_kv_heads*head_dim, 1, max_seq_len], fixed max_seq_len, updated with the new mutable_slice_update in-place op. Note: "On iOS we must update on dim 4 (the last dim), whereas on macOS we use dim 3" — the last-axis constraint (doc 01 §2.2) still shapes the layout
Session guidance (325)	"static tensor shapes", "channels-first layouts", "convolutional projections instead of linear layers", "palettization over INT quantization for power efficiency on iOS"

Meanwhile macos/sdpa.py is a thin wrapper over the fused coreai_torch.composite_ops.SDPA — on GPU, none of the manual surgery is needed. The per-platform registry institutionalizes the lesson of doc 06 §3 (SDPA choice depends on workload/hardware).

What changed:

• In-place slice update replaces the one-hot mask blend for KV caches (mutable_slice_update with computed begin/end vs WhisperKit's cache*(1-m)+current*m).
• The toolchain absorbs the decomposition: you author with high-level composite ops (SDPA, RMSNorm, RoPE) and the compiler picks the hardware-optimal lowering per platform — with the option to drop to hand-written iOS primitives (as coreai-models does for LLMs) or custom Metal kernels.
• torch.export + dynamic shapes replaces trace-per-shape and multifunction workarounds.
• AOT compilation + explicit specialization/caching APIs replace "hide the first uncached load" folklore (doc 05).
• The Whisper story closes a loop: WhisperKit needed a hand-crafted re-implementation (doc 06); the Core AI recipe (models/whisper/export.py) exports the stock HF whisper-large-v3-turbo through the decomp table in ~200 lines.

3.5 Where the compiler's rules actually live#

"The toolchain owns the optimization" — but the rules are written in four layers of decreasing openness (verified against our pinned sources and the installed coreai-core 1.0.0b2 wheel):

Layer	What rules	Where	Open?
coreai-torch (frontend)	What survives decomposition: _decomp.py lists exactly 12 preserved ATen ops (scaled_dot_product_attention, silu, instance_norm, pads…). How ops translate: _aten_to_core.py, 3,760 lines of per-op lowerings (e.g. replace_sdpa() at line 3426). Your own rules: _torch_metal_kernel.py DSL	references/coreai-torch/coreai_torch/	✅ readable
coreai-core (middle-end)	Graph optimization passes behind .optimize(). The wheel reveals the substrate: coreai/_compiler/ is MLIR (LLVM compiler infra) with tblgen-generated dialects coreai, coreaix, udml, debuginfo — hence .aimodel's main.mlirb. Passes run in the MLIR PassManager but are compiled into _mlir.so	pip wheel (binary)	⚠️ IR inspectable, passes binary
OS Core AI framework (backend)	Specialization: compute-unit segmentation, layout assignment, per-chip kernel codegen (the ANECompiler lineage, private behind Core ML too)	macOS/iOS 27 system framework (+ _coreai_runtime_os.so bridge)	❌ closed, as ANE codegen always was
coreai-models primitives ("standard library")	What the compiler does not yet automate: primitives/ios/sdpa.py exists precisely because "iOS requires each attention head to be computed individually" — for ANE-bound LLMs Apple still writes the per-head form in the source graph (ios_class/macos_class registry, --platform iOS)	references/coreai-models/	✅ readable

Classic LLVM-style compiler architecture: open frontend, public IR with proprietary passes, closed hardware backend — plus an open standard library covering the gaps. The practical consequence: op-level optimization is automated only as far as the decomp table + lowerings reach; hardware-specific graph shaping still surfaces as authored primitives when the compiler falls short, which is where this KB's principles (docs 01–07) remain working knowledge rather than history.

4. Tooling: debugger, instruments, gauge#

• Core AI Instruments template: model load + specialization sub-events, inference latency over time (the "Meet" demo used it to spot quadratic KV-recompute growth).
• Core AI Debugger (macOS app): navigate the PyTorch module hierarchy ↔ converted graph ↔ original Python source lines; run on device and inspect intermediate tensors; comparison mode vs the PyTorch reference (PSNR default) — layer-sensitivity analysis before compression.
• Core AI debug gauge in Xcode: streaming activity view for a first look before Instruments.
• Deployment guidance (session 326): keep big models out of the bundle (Background Assets), never specialize in the hot path, AOT-compile for first-run UX; sizing example — SAM3 (623 MB) + Qwen3-0.6B on iPhone, Qwen3-8B on Mac.

5. What this means for the KB#

1. The hardware truths are permanent. 64-byte last-axis alignment, channels-first, static shapes, per-head computation, palettization-for-bandwidth — every 2022 principle reappears in Apple's 2026 first-party code. Understanding docs 01–07 is understanding why Core AI's iOS primitives look the way they do.
2. The division of labor moved up. Authors now write composite ops and YAML compression configs; the compiler owns lowering. Hand-optimization remains for the frontier (custom kernels, re-authoring like SAM3, platform-specific primitives).
3. Migration sketch (Core ML → Core AI): torch.jit.trace→torch.export; .mlpackage→.aimodel; MLModel/MLState→AIModel/InferenceFunction + state MutableViews; ct.optimize/argmaxtools recipes→coreai-opt presets/YAML; Xcode performance tab→Core AI Instruments + Debugger; multifunction models→multi-entrypoint conversion. Full comparison and decision guide: doc 10.
4. Open questions to research (not answered by the fetched material): exact ANE dispatch reporting in the new tooling (MLComputePlan equivalent), .aimodel op-level format, Core ML long-term support policy, whether WhisperKit-class projects migrate.