Meet Core AI

Hi everyone, my name is Ben and I'm an engineer on the Core AI team.

Today I'll be giving an introduction to Core AI, and showing how you can use it to add intelligent features into your apps.

AI is advancing faster than ever. New models and capabilities that previously seemed out of reach are emerging constantly.

Core AI is built to help you harness that momentum and build on top of it. Core AI marks the next evolution of on-device AI execution across Apple platforms. It's built from the ground up for modern workloads, and delivers the high-performance inference you need to build advanced AI features. Core AI is the inference framework powering on-device Apple Intelligence. And now, it's available for you to use, bringing that same power to your app's own intelligence. Core AI is more than just a framework. It's a complete set of technologies, covering the model deployment lifecycle, from model optimization and conversion to debugging and integration into your app. All designed to support the fast, iterative cycle that building great AI features requires.

Core AI allows you to leverage all of Apple Silicon. It provides blazing fast inference across the CPU, GPU, and Neural Engine.

The framework comes with a modern Swift API. It's an expressive API that delivers the performance your app demands without compromising on memory safety.

The broader set of technologies fit naturally into common ML engineering workflows, reusing familiar Python and PyTorch foundations for model authoring, optimization and conversion.

Core AI also supports extensive customization from fine-grained inference management and model specialization to custom GPU kernels.

And all of this is tightly integrated into a new developer toolchain, with ahead-of-time compilation, dedicated Core AI Instruments, and a powerful visual Debugger to trace tensor values directly back to your original Python source code. Core AI is designed to scale to your needs and available compute.

Whether you want your app to identify who's talking in a live meeting with a small speaker diarization model, your users to point their camera at anything, ask a question, and instantly get an answer with a larger vision language model, or let them hand off complex, multi-step tasks to a powerful agentic assistant powered by a 70 billion parameter LLM. Core AI has you covered. With all of it running locally on Apple devices, with no server and no cost per token. In this talk, I'll start by showing how to get your model into the Core AI format.

Then I'll go over how to integrate the converted model into your app.

I'll then dive a little deeper into optimizing the performance of your model and app. And lastly I'll highlight some additional features of Core AI and its associated tools that you may find useful.

Let's get started. Every great app experience starts with an idea. Maybe you want to build something that feels a little magical, something that responds intelligently, or makes a decision that would otherwise require a human or hard coded rules.

Machine learning and AI are what make those kinds of experiences possible. Once you have that idea, the next step is finding or building a model that can power it. Just like your idea itself will evolve over time, finding the right model is an iterative process. You'll try things, evaluate them against your requirements and refine. Core AI is designed to support that iteration and make it as fast and frictionless as possible.

So to make this concrete, I'll implement a fun game idea I had. It's an app that lets you play a two player snake game where one snake is powered by an AI model run through Core AI. The app will follow traditional snake rules where snakes can grow by eating food and must avoid hitting walls, themselves and the other snake. The last snake standing wins. At each time step, the AI model will see a set of features describing the current board state, and those features will be accumulated into the full game history that gets fed to the model. It will then predict the best direction to move. While snake is a simple game, the tools and APIs used to create this experience are the same foundation that scale all the way up to the larger, more complex use cases.

I was curious to see what I could put together with PyTorch for this project. With a little help from an AI coding assistant I was able to sketch out a simple snake action prediction model pretty quickly. To train it, I used a naive simulation to generate training data, just running the game and recording states and actions. The idea was to start simple and get the model working in my app.

So the next step is taking this PyTorch model and converting it to Core AI.

I'll use the new Core AI Torch Python package to easily perform the conversion.

First I'll load the trained checkpoint of the SnakeTransformer module, and prepare a sample input.

Then I'll export the torch program using torch.export and also make sure to use the dynamic_shapes argument to specify that the sequence length of the features is dynamic, that way it doesn't get traced with the static sample length of 5. Also I'll run decompositions on the converted program using Core AI's decomposition table.

Next I'll run Core AI's TorchConverter, specify the names of the inputs and outputs, and finally save the converted Core AI model to disk. Before leaving the Python environment, one more thing I'll do is run a test to verify that the converted Core AI model matches the numerics of my original PyTorch model. This can be done easily with the Core AI framework Python bindings. First I'll load the PyTorch and Core AI models. Then prepare a sample snake game input.

Then run that same input through both the PyTorch module and the Core AI inference function. And finally assert a sufficiently small delta for my use case between the PyTorch and Core AI outputs.

Now that I have the converted AI model, the next step is to hop into Xcode and integrate the model into my app. First I'll open the AI model file with Xcode, which shows information about the model.

It includes the model size, the distribution of operations and other helpful metadata. Also in the Functions tab it shows you the exact function signature of each unique function in the model.

In this case the model just has one function, which takes the features of the game board as an input and produces logits as an output which indicate which direction the model thinks would be best to move. Also note that the question mark in the NDArray values denotes that the dimension has a dynamic shape, which matches how I converted the model with a dynamic sequence length. Now that I've included the AI model file in my Xcode project and have examined its structure, the next step is to use the Core AI framework to run the model.

The Core AI framework is a new Swift API surface for loading and running Core AI models.

It offers a progressively disclosing set of APIs, which makes it simple to get things up and running, while also having deeper layers of flexibility for supporting performance critical applications.

Also, it uses modern Swift language features like non-escapable types, to offer memory-safe APIs while not sacrificing performance.

Let's begin by discussing the core types within the framework.

An AIModel is initialized from a URL to a .aimodel file and is used primarily to inspect and load one or more inference functions.

An InferenceFunction is the runnable object which represents a single loaded compute graph. In the common case, your AIModel will only have a single main InferenceFunction, though you can convert a single model with multiple functions. The AIModel and InferenceFunction are typically objects you'll construct when preparing your app's AI feature. For example this could be on app initialization.

NDArray is the type which holds your multi-dimensional input and output data and you use the run method on an InferenceFunction to run inference with that data.

Finally you can read and process the outputs of the inference.

So for implementing the snake game, I'll start by making the ModelPlayer type. At app initialization time, it'll be initialized with the URL to the AI model file that it should use. Then it will initialize the AIModel, and load the main inference function from it.

Next is the logic for the model player to make decisions. It'll conform to the SnakePlayer protocol that I've defined in my app.

The main protocol requirement is the chooseAction function which is passed in the game's history, and returns the next action that the snake should take.

The first thing to do is create an NDArray to populate with the input features.

For this inference function, the expected structure of the NDArray is 2 dimensional with float32 data, where the first dimension of the shape is the current sequence length, and the second is the fixed hidden dimension size.

Then it'll write the features into that NDArray using this writeFeatures helper function which takes the game and a mutable view of the NDArray. The NDArray.MutableView type is a non-escapable type which provides safe and efficient access to the backing storage of the NDArray.

After preparing the inputs, it'll run inference with them, and extract the expected output logits ndarray.

The last step is to sample the output logits to pick the next direction that the snake will move, by passing an ndarray view into the helper function which will read the values and choose the direction with the largest corresponding logit. The writeFeatures function is what's populating the input features. Let's briefly go over what these features include.

They have the normalized distance of the AI snake's head to all the walls. The normalized relative X and Y distance to the nearest food.

Four elements encoding it's current direction.

The normalized distance to the other snake.

And finally the opponent's direction.

Now with this put together I'm going to try a test run with both snakes powered by the AI model to see how it does.

Running it shows that the model is working. However, I see that the game is getting slower as it goes on.

Alongside the Core AI framework, there's a new instrument in Xcode to help you profile the Core AI models running in your app. In this case I've ran the app with Instruments and I can see the inference intervals getting notably larger over time, which means the inference calls are increasing in latency. This makes sense because transformer models have quadratic time complexity with respect to the sequence length. And in our game the sequence length is increasing with every move the model makes. The next step in this case is to optimize the performance of the model usage.

Each time the input sequence is increased, the transformer model recomputes a set of internal key and value embeddings for every element in the sequence. A common strategy used to improve the performance of decoding loops like this when using transformers is to cache keys and values that are computed for each element in the sequence, as opposed to re-computing them all from scratch with each inference. This can be achieved through Core AI by using states.

States are inputs to the model which are both read, and updated in-place during inference. By introducing the key and value caches as states on the model, we both avoid recomputing them on each inference, and also remove the need to provide the full history of the game as an input since the data needed from older steps are stored in the states.

So after the first input, each subsequent step uses the cache for history and only takes the new features of the latest board state.

To implement the key/value caching, I'll go back to the original authoring code and make a few changes to add in the key and value caches. First I'll update the torch module by adding key and value cache tensors as buffers within the transformer module, by using the torch register_buffer API. This will later result in these tensors being mutable buffers in the exported torch program which Core AI will convert to states. Then in the forward function of the module, I'll add the logic to actually use the caches. This involves reading previous features keys and values out of the cache. Then writing the computed keys and values for the new features back into the cache. Lastly, I'll rerun the same code from before to re-convert the model, but now adding in the state_names argument to the convert call to specify the names of the new state arguments. Now that I've re-converted the model with the new function signature, I'll update the app code to handle it. To start, I'll update the ModelPlayer to store the key and value cache NDArrays which will be the state arguments passed to each inference. I'll initialize them with the expected shape for the transformer. In this case I converted the model such that it expects the key and value caches to always be a fixed size for a maximum possible context length.

Then when it's time to run inference, I'll construct a collection of MutableViews containing both views of the key and value caches. Then provide those as the states argument of the InferenceFunction.run method. Now the caches will be both read and updated in-place during each inference. Now with the updated model, I'll re-run the app. This time I can see it maintains a steady speed, no longer slowing down overtime.

When tracing the updated app in Instruments, I can confirm that the inference latency is growing at a much slower rate.

Before wrapping up, I'll show some features that I didn't use while making the snake game, but that you may find useful when developing your own apps. When converting the snake game models, I used the coreai-torch package to directly convert the PyTorch module. This flow is simple and works great for many use cases, but sometimes you may need more control over how your model is authored, and potentially even how the operations within the model are run.

We've only touched the surface of what the Core AI Python package has to offer. It also has support for directly authoring your model with Core AI APIs, optimizing the model for Apple Silicon, and defining custom kernel implementations with Metal 4. To learn more about these advanced model authoring flows, see the talk "Dive into Core AI model authoring and optimization". In addition to debugging performance, it's also crucial to be able to debug the numerics of your converted model. For this you can use the Core AI Debugger which allows you to visualize your converted model, easily inspect intermediate tensor values, and trace back operations in the converted model to the Python source code which introduced them.

There is also a convenient Core AI debug gauge which shows you streaming Core AI activity while your app is running in Xcode. This is a great place to spot performance issues before jumping into instruments.

One thing that was glossed over in the snake game implementation is the process of model specialization.

When you ship an AI model with your app, that is a source representation of the model, which can be run on any Apple device. However, to actually load and run the model within your app, it must be specialized for the device that the app is running on. When your model is loaded it is checked to see if it has already been specialized and cached. The specialization process can take a significant amount of time for very large models.

While future loads are from the cache and fast, that first time is something you may need to plan for. It is recommended you avoid having model specialization occur within user interactive flows.

Core AI can help you with that. First, Core AI gives you programmatic access to the default model cache for your app. You can request to load models directly from it. If nil is returned, it is not present and requires specialization. You can use this to gate features or inform the users that they may need to wait a bit while your app prepares the model.

Second, you can request model specialization explicitly in your app independent of it being loaded.

You can do this after downloading assets or when the user opts in to a feature so the model is ready to go ahead of time. And there is a lot more control available. SpecializationOptions help configure how you want your model to be optimized for inference.

With the AIModelCache you can also delete entries you no longer need, and control the policy on how long entries persist.

You can even share a cache between multiple apps in the same app group.

Check out the "Managing model specialization and caching" article on developer.apple.com to learn more.

Independent of when specialization occurs, it still takes time. Lets take a quick peak inside. During specialization, the model goes through two main transformations. First, it goes through a core set of compilation steps which segment, plan and optimize compute. Second, executable artifacts are generated for the compute units used. These artifacts are tied to the device and OS version they were generated on. Of these two steps, compilation is the one which incurs most of the latency.

The Core AI toolchain can help you reduce that time by allowing some compilation to occur ahead of time on your development machine, producing a compiled version of the model.

While that compiled model still needs to be specialized for the specific users device, there is now much less work to do and finishes significantly faster. To learn more about this option, check out the "Compiling Core AI models ahead of time" article on developer.apple.com. Controlling when, where, and how specialization happens is one way to help you optimize your users experience. Another area you may want to optimize is removing any overheads in tight inference loops using your model. The Core AI Framework has several APIs to help you here.

You can dynamically check the optimal memory layout of NDArray arguments and allocate them with that structure to avoid layout conversions at inference time.

You can also pre-allocate output values for the framework to write into, to avoid allocating new output values during inference.

And you can also use asynchronous values to efficiently pipeline execution of multiple inference functions together. For most use cases, the higher-level inference APIs will get you exactly where you need to be.

But when you're optimizing a tight inference loop or integrating a model into a complex compute pipeline, these lower-level APIs are there when you need them.

Whether you're just getting started or diving deep, the Core AI Models repository is a great place to find what you need.

It has a collection of popular models, each just a single command away from being converted and optimized for your app.

AI skills that are experts in Core AI model authoring, optimization, and conversion.

And a Swift package with libraries for specific families of models that give you higher-level APIs that already have many of those low-level inference optimizations built in.

It also provides an API for creating a Core AI Language model, which plugs right in to the Foundation Models framework, letting you bring your own custom models and token sampling strategies.

To wrap things up: Core AI is available on all Apple Silicon to help you build cutting edge AI experiences on all Apple platforms.

It has tight integration with the existing Python tools that you're already familiar with, a modern Swift framework for running your models efficiently within your app, and state of the art debugging tools to help you understand how your models are running on Apple devices.

We can't wait to see what sorts of experiences you build.