Meet Core AI

Watch on Apple Developer

TL;DR

What Core AI provides

Core AI is the inference framework used for on-device Apple Intelligence, now exposed for developers building AI features in their own apps. It is designed for modern on-device workloads across Apple platforms and can execute across Apple Silicon CPU, GPU, and Neural Engine.

The framework is not just a runtime API. It covers the deployment lifecycle: model authoring, optimization, conversion, debugging, app integration, profiling, model specialization, ahead-of-time compilation, and lower-level customization when needed.

• Swift framework for loading .aimodel assets and running inference with memory-safe APIs.
• Python tooling built around familiar PyTorch workflows for conversion, authoring, optimization, and validation.
• Xcode integration for model inspection, Core AI Instruments profiling, a Core AI debug gauge, and a visual Core AI Debugger.
• Support for advanced needs such as stateful inference, custom GPU kernels with Metal 4, model specialization control, and optimized tight inference loops.

Convert and validate a PyTorch model

The demo converts a PyTorch snake-action transformer into Core AI format using the Core AI Torch Python package. The model is exported with torch.export, including a dynamic sequence-length dimension so the converted model is not fixed to the sample input length.

After conversion, the session validates numerics by running the same input through the original PyTorch model and the Core AI model using the Python runtime bindings, then checking that logits are close enough for the app's use case.

• Use torch.export.Dim and dynamic_shapes for dynamic dimensions.
• Run Core AI's decomposition table before conversion.
• Use TorchConverter().add_exported_program(...) to name model inputs and outputs.
• Save the result as a .aimodel asset for Xcode and the Swift runtime.

Convert a PyTorch model to Core AI

Exports a PyTorch model with a dynamic sequence dimension, converts it to a Core AI graph, and writes a .aimodel asset.

import torch
import coreai_torch

pt_model = SnakeTransformer().load_checkpoint("snake.pt")
example = torch.randn(1, 5, 16)

seq_len = torch.export.Dim("seq_len", min=1, max=256)
exported = torch.export.export(
    pt_model,
    args=(example,),
    dynamic_shapes={"features": {1: seq_len}},
)
exported = exported.run_decompositions(coreai_torch.get_decomp_table())

ai_program = coreai_torch.TorchConverter().add_exported_program(
    exported,
    input_names=["features"],
    output_names=["logits"],
).to_coreai()

ai_program.save_asset("SnakeTransformer.aimodel")

Validate converted model numerics

Compares Core AI output against PyTorch output before leaving the Python environment.

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

ai_model = await AIModel.load("SnakeTransformer.aimodel")
function = ai_model.load_function("main")

features = make_snake_features(game).astype(np.float32)[np.newaxis]

with torch.no_grad():
    pytorch_logits = pt_model(torch.from_numpy(features)).numpy()[0, -1]

result = await function({"features": NDArray(data=features)})
coreai_logits = result["logits"].numpy()[0, -1]

assert np.max(np.abs(pytorch_logits - coreai_logits)) < 0.01

Integrate a .aimodel in a Swift app

Xcode can open a .aimodel file and show metadata such as model size, operation distribution, and function signatures. In the demo, the model has one function that takes game-board features and outputs logits for the next movement direction.

The CoreAI Swift framework centers on AIModel, InferenceFunction, and NDArray. Apps typically create the AIModel and load the main function when preparing the feature, then construct NDArray inputs, run inference, and process output arrays.

• AIModel loads and inspects a .aimodel asset.
• InferenceFunction represents a runnable compute graph, commonly the model's main function.
• NDArray holds multidimensional input and output data.
• NDArray.MutableView and NDArray.View provide efficient, memory-safe access to array storage.

Load and run an inference function

Shows the core Swift runtime flow: load a model, load a function, pass NDArray inputs, and read NDArray outputs.

import CoreAI

let model = try await AIModel(contentsOf: modelURL)
let mainFunction = try model.loadFunction(named: "main")!

let inputNDArray: NDArray = nextInput()
var outputs = try await mainFunction.run(inputs: ["input": inputNDArray])

guard let outputNDArray = outputs.remove("output")?.ndArray else {
    throw ModelError.missingOutput
}

Prepare dynamic game features and run inference

The demo's model player builds a float32 NDArray, writes game features into it, runs inference, and chooses the direction with the highest logit.

extension ModelPlayer: SnakePlayer {
    mutating func chooseAction(game: SnakeGame) async throws -> Direction {
        var inputFeatures = NDArray(
            shape: [game.stepCount, hiddenDim],
            scalarType: .float32
        )
        writeFeatures(of: game, into: inputFeatures.mutableView())

        var outputs = try await nextActionFunction.run(
            inputs: ["features": inputFeatures]
        )
        guard let logits = outputs.remove("logits")?.ndArray else {
            throw ModelError.missingOutput
        }
        return predictedDirection(from: logits.view())
    }
}

Profile and optimize with states

The first implementation feeds the full game history into a transformer on every move. Instruments shows inference intervals growing over time, which matches transformer quadratic complexity as sequence length increases.

The optimization uses Core AI states for transformer key/value caches. States are model inputs that are read and updated in-place during inference. After the first step, the app can pass only the newest features while the cached state carries the relevant history.

• In PyTorch, key and value caches are registered as buffers so the exported program treats them as mutable buffers.
• During conversion, state_names maps those mutable buffers to Core AI state arguments.
• In Swift, the app owns NDArray cache buffers and passes mutable views through InferenceFunction.MutableViews.
• The updated implementation keeps game speed steady and reduces the growth rate of inference latency.

Convert a stateful model with key/value cache states

Adds state names during conversion so mutable key/value cache buffers become Core AI states.

exported = torch.export.export(
    stateful_model,
    args=(example_features, example_position_ids),
    dynamic_shapes={"position_ids": {1: seq_len}},
)
exported = exported.run_decompositions(coreai_torch.get_decomp_table())

ai_program = coreai_torch.TorchConverter().add_exported_program(
    exported,
    input_names=["features", "position_ids"],
    state_names=["keyCache", "valueCache"],
    output_names=["logits"],
).to_coreai()

ai_program.save_asset("SnakeTransformer.aimodel")

Pass mutable state views from Swift

Passes cache buffers as mutable state views so inference reads and updates them in-place.

var stateViews = InferenceFunction.MutableViews()
stateViews.insert(&keyCache, for: "keyCache")
stateViews.insert(&valueCache, for: "valueCache")

var outputs = try await nextActionFunction.run(
    inputs: ["features": inputFeatures],
    states: stateViews
)

Specialization, caching, and ahead-of-time compilation

A shipped .aimodel is a source representation that can run on Apple devices, but it must be specialized for the user's specific device before loading and inference. First specialization can take significant time for very large models, while later loads can be fast from cache.

Core AI exposes APIs to check and manage the default model cache, request specialization explicitly, configure specialization behavior, delete unused entries, control persistence, and share a cache between apps in the same app group. The session recommends avoiding specialization inside user-interactive flows.

Ahead-of-time compilation can move part of the compilation work to the development machine. The compiled model still needs device-specific specialization, but there is less work left on the user's device.

• Use cache checks to gate features or inform users that AI features are being prepared.
• Request specialization after downloading assets or when a user opts into a feature.
• Consult the Core AI documentation for SpecializationOptions, cache policies, app-group cache sharing, and ahead-of-time compilation.
• For tight inference loops, lower-level APIs can query optimal NDArray memory layouts, preallocate outputs, and pipeline functions with asynchronous values.

Check the model cache before loading

Checks whether a specialized model is already available before entering a user-visible flow.

let cache = AIModelCache.default
guard let model = try cache.model(for: modelURL, options: .default) else {
    Task { @MainActor in
        informUser("Preparing AI features. This may take a while...")
    }
    return
}

Request specialization explicitly

Starts specialization independently of model loading so the model can be prepared ahead of use.

try await AIModel.specialize(contentsOf: modelURL)

Resources

• Core AI PyTorch Extensions
• Core AI Python
• Core AI Optimization
• Core AI
• Compiling Core AI models ahead of time
• Managing model specialization and caching

Related Sessions

• Dive into Core AI model authoring and optimization
• Integrate on-device AI models into your app using Core AI
• Optimize custom machine learning operations with Metal tensors

Unofficial, not associated with Apple. By