Dive into Core AI model authoring and optimization

Watch on Apple Developer

TL;DR

Core AI Python deployment workflow

Core AI is presented as a deployment lifecycle for Apple silicon: convert, optimize, debug, specialize, and integrate model assets. This session focuses on the Python ecosystem, especially PyTorch-based authoring and conversion.

The basic workflow uses a PyTorch ExportedProgram as the input to Core AI. coreai-torch converts that graph into a Core AI program, optimization produces a deployable .aimodel asset, and Python can load a specialized function and execute it with NumPy-backed tensors.

• Install the Python entry point with pip install coreai-torch; this installs both coreai and coreai-torch.
• Use torch.export.export to capture weights, operations, and shapes in a graph that Core AI can convert.
• Use TorchConverter with explicit input and output names, then call to_coreai(), optimize(), and save_asset().
• Core AI conversion can also assemble multiple models into one artifact, register custom lowerings, and inline Metal kernels.

Export a PyTorch model for Core AI conversion

torch.export produces the graph representation consumed by coreai-torch.

import torch
import torch.nn as nn

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(256, 512)
        self.fc2 = nn.Linear(512, 10)

    def forward(self, x):
        return self.fc2(torch.relu(self.fc1(x)))

model = MLP().eval()
example_input = (torch.randn(1, 256),)
exported_program = torch.export.export(model, example_input)

Convert, optimize, save, and run a Core AI model

The converted Core AI program becomes an .aimodel asset and can be specialized and executed from Python.

import coreai
import coreai_torch
from coreai.runtime import NDArray

converter = coreai_torch.TorchConverter()
converter.add_exported_program(
    exported_program,
    input_names=["features"],
    output_names=["logits"],
)
core_ai_program = converter.to_coreai()

core_ai_program.optimize()
asset = core_ai_program.save_asset("mlp.aimodel")

specialized_model = await AIModel.load("mlp.aimodel")
specialized_function = specialized_model.load_function("main")
result = await specialized_function({"features": NDArray(example[0].numpy())})

Model repositories, reusable components, and Core AI Skills

The coreai-models repository is described as both a source of ready-to-go model architectures and reusable components for bringing custom models to Core AI. It also includes a Swift package for running LLMs in apps.

Core AI Skills are agent skills that can be installed into a coding assistant. They are intended to help an agent translate high-level deployment goals into concrete Core AI plans, including PyTorch changes, conversion steps, optimization choices, and target-hardware considerations.

• Use coreai-models for examples engineered around different use cases and constraints.
• Agent skills may ask for model type, target hardware families, and application constraints before suggesting a workflow.
• Use these examples and skills as starting points for authoring PyTorch that maps well to Apple silicon.

Compression with coreai-opt

coreai-opt provides config-driven model compression. The session demonstrates it with SAM3, an 850-million-parameter prompt-based image segmentation model with image encoder, text encoder, and detector components.

The first compression pass uses the presets.w4 configuration for 4-bit per-channel symmetric quantization. This reduces the converted SAM3 asset from over 3 GB to about 430 MB, but applying the same aggressive compression to every layer causes a missed detection in the output.

The key guidance is to compress selectively based on the model's structure and quality sensitivity. In the SAM3 example, the image and text encoders contain most parameters, while the detector is only about 4% of parameters and is more sensitive to compression.

• coreai-opt supports int4, int8, FP4, and FP8 weight compression with flexible granularity.
• Quantization APIs can use calibration data or quantization-aware training.
• ExecutionMode.EAGER is called out as a good fit for weight compression; activations use graph mode.
• Preserving high-level semantics such as attention during export may require Core AI's custom decomposition table.
• Casting exported programs to 16-bit floating point is supported through a coreai-opt helper when needed.

Core AI Debugger for structure, runtime, and comparison analysis

Core AI Debugger is a standalone app for inspecting Core AI models on Apple platforms. It combines a PyTorch-module navigator, a structure graph, a source viewer tied to original Python code, and an inspector for operation inputs, outputs, and tensor details.

The debugger can run a model on a selected device, specialize it for that target, and expose intermediate output tensors without modifying the model. In the SAM3 example, this confirms that the quantized model's final mask is missing one flower.

For correctness analysis, the workflow saves PyTorch intermediate tensors using a new save-intermediates API, loads that file into Core AI Debugger, and compares specialized Core AI outputs against the PyTorch reference. The debugger creates sync points and reports similarity metrics such as PSNR, making it easier to find where results diverge.

• Navigator groups operations by PyTorch module hierarchy, which is useful for large models.
• Structure viewer shows operation connectivity, execution order, and data dependencies.
• Inspector can preview tensors after running on-device.
• Comparison mode pairs specialized-model operations with reference PyTorch operations as sync points.
• In the SAM3 case, low-PSNR sync points concentrated in the detector decoder, motivating a revised quantization scheme that skips the detector.

Custom Metal kernels and fused operations

Advanced Core AI authoring can tune the computational graph rather than converting it end-to-end as-is. One approach is operation fusion: replacing several graph operations with one kernel dispatch.

Core AI includes prepackaged fast kernels and primitives for heavy operations such as Scaled Dot Product Attention. For lower-level customization, Core AI can embed custom Metal Shading Language source directly into an .aimodel asset. The PyTorch model calls a Python-side reference function for tracing, while the Core AI converter bundles the matching Metal implementation into the asset.

• Define a PyTorch reference implementation that torch.export can trace.
• Define the MSL kernel source and bind parameters such as thread_position_in_grid.
• Create a TorchMetalKernel with input names, result names, Metal source, reference function, and dtype templates.
• Call the custom kernel from PyTorch like a function, specifying grid sizes and result shapes.
• Register custom kernels with TorchConverter before adding the exported program.

Define a SiLU custom Metal kernel with a PyTorch reference

The reference function is visible to PyTorch export; the MSL implementation is embedded in the Core AI asset.

import torch
from coreai_torch.dsl import TorchMetalKernel, MetalParameter

def silu_torch(x):
    return x * torch.sigmoid(x)

SILU_MSL = """
float val = float(x[gid]);
float sig = 1.0f / (1.0f + exp(-val));
y[gid] = TYPE(val * sig);
"""

silu_kernel = TorchMetalKernel(
    name="fused_silu",
    input_names=["x"],
    result_names=["y"],
    src=SILU_MSL,
    torch_defn=silu_torch,
    metal_params=[MetalParameter("gid", "uint", "thread_position_in_grid")],
    template_dtypes={"x": "TYPE"},
)

Register and use a custom kernel during conversion

Registering the kernel lets the converter integrate the Metal implementation into the resulting .aimodel.

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(256, 256)

    def forward(self, x):
        h = self.linear(x)
        n = h.numel()
        return silu_kernel(
            h,
            threads_per_grid_size=(n, 1, 1),
            threads_per_thread_group=(min(n, 256), 1, 1),
            result_shapes=[h.shape],
        )

exported_program = torch.export.export(MyModel(), (torch.randn(1, 256),))
converter = coreai_torch.TorchConverter()
converter.register_custom_kernels([silu_kernel])
converter.add_exported_program(exported_program, input_names=["x"], output_names=["y"])
deployable = converter.to_coreai()

Model re-authoring for target-specific performance

For more demanding constraints, especially on iOS, the session recommends model re-authoring: rewriting the PyTorch implementation with the target platform and usage pattern in mind. This can involve different operations, tensor layouts, static shapes, new interfaces, and reusable Core AI-friendly patterns.

The SAM3 re-authoring example splits a single end-to-end model into three Core AI functions: image_encode, text_encode, and detect. This lets an app run parts of the model at different cadences, cache intermediate results, and compress each function independently.

The re-authored SAM3 uses convolutional projections and channels-first/static-shape patterns for the encoders, applies 4-bit palettization with per-channel scales to the encoders, keeps the compression-sensitive detector uncompressed, casts exported programs to half precision, and packages all three functions into one asset.

• Use semantic PyTorch patterns that Core AI can map to optimized runtime implementations, such as in-place key-value cache updates for LLMs.
• Prefer static tensor shapes, channels-first layouts, and convolutional operation patterns when targeting iOS constraints.
• Test re-authored models at both module and model level with unit or integration tests.
• In the SAM3 example, changing the prompt from flowers to butterfly reuses cached image embeddings and only reruns the text encoder and detector, making the second inference 76% faster after warmup.

Resources

• Core AI PyTorch Extensions
• Core AI Python
• Core AI Optimization
• Inspecting, debugging, and profiling Core AI models
• Inspecting Core AI models with Core AI Debugger
• Core AI

Related Sessions

• Explore distributed inference and training with MLX
• Explore numerical computing in Swift with MLX
• Run local agentic AI on the Mac using MLX

Unofficial, not associated with Apple. By