Meet Core AI

Highlight

Apple 推出了 Core AI，一个覆盖模型转换、优化、编译到端侧推理全链路的 AI 框架，让开发者用熟悉的 PyTorch 工作流就能把模型部署到 Apple Silicon 上本地运行，无需服务器、按 token 零成本。

核心内容

端侧 AI 的痛点

想在 App 里跑一个 AI 模型，过去的路径很长。你先要在 Python 里训练好模型，然后找工具把它转成设备能跑的格式，中间可能还要经过 ONNX 这样的中间层。转换完之后，你还得担心数值对不对、性能够不够好、首次加载会不会卡死用户。每一步都是摩擦。

Apple 自己跑 Apple Intelligence 时也遇到了同样的问题。Core ML 虽然能用，但面对 Transformer 这类现代 workload 时显得力不从心。Apple 没有打补丁，而是从头写了一套新框架，现在把它开放给所有开发者。

Core AI 是什么

Core AI 不只是个推理框架，它覆盖了模型部署的完整生命周期（00:59）：

• Python 端用 coreai-torch 做模型转换和优化
• Swift 端用 CoreAI framework 做推理
• Xcode 提供 AOT 编译、专用 Instruments 和可视化 Debugger

这套工具链的目标很明确：让开发者保持熟悉的 PyTorch 工作流，同时享受 Apple Silicon 的全栈加速（CPU、GPU、Neural Engine）。

一个具体的例子：AI 玩贪吃蛇

演讲者用 Core AI 做了一个双人对战的贪吃蛇游戏，其中一条蛇由 AI 模型驱动（03:51）。这个例子虽小，但用到的工具和 API 与跑 70B 参数大模型是同一套基础。

模型是一个简单的 Transformer，输入是游戏板状态特征，输出是四个方向的动作 logits。训练数据用 naive simulation 生成，跑几局游戏记录状态和动作就行。

模型转换：从 PyTorch 到 .aimodel

转换流程很直接（05:08）：

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")

关键点：

• torch.export 把 PyTorch 模型导出为静态计算图
• dynamic_shapes 声明序列长度是动态的，范围 1 到 256
• coreai_torch.get_decomp_table() 提供 Core AI 所需的算子分解规则
• 最终产物是 .aimodel 文件，Xcode 可以直接打开查看结构

转换后建议做数值验证（05:44）：

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

pt_model = SnakeTransformer().load_checkpoint("snake.pt")
ai_model = await AIModel.load("SnakeTransformer.aimodel")
function = ai_model.load_function("main")

features = np.array([extract_features(game) for _ in range(10)],
                    dtype=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]

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

关键点：

• 用同一组输入分别跑 PyTorch 和 Core AI
• 比较输出 logits 的最大差值
• 阈值根据业务需求设定，这里用 0.01

Swift 端推理：基础用法

Core AI 的 Swift API 采用渐进式披露设计（07:41）。最简单的用法只需要三个类型：

import CoreAI

let model = try await AIModel(contentsOf: modelURL)
let mainFunction: InferenceFunction = 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 {
    // Handle unexpected missing output
}

关键点：

• AIModel 从 .aimodel 文件加载，负责解析模型结构
• InferenceFunction 是实际可运行的计算图，通常叫 “main”
• NDArray 承载多维输入输出数据
• run 是异步的，不会阻塞主线程

把它封装到游戏玩家类型里（08:33）：

struct ModelPlayer {
    let nextActionFunction: InferenceFunction

    init(modelURL: URL) async throws {
        let model = try await AIModel(contentsOf: modelURL)
        self.nextActionFunction = try model.loadFunction(named: "main")!
    }
}

实际的决策逻辑（08:49）：

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())
    }
}

关键点：

• NDArray.MutableView 是非逃逸类型，安全高效地访问底层存储
• mutableView() 获取可写视图，view() 获取只读视图
• 输入特征的 shape 是 [stepCount, hiddenDim]，第一维动态增长

输入特征包括（10:10）：到四面墙的距离、最近食物的相对位置、当前方向 one-hot 编码、对手蛇的距离和方向。

性能问题：Transformer 的二次复杂度

游戏跑起来后发现越玩越慢（10:42）。Instruments 显示推理间隔随时间明显增大。原因是 Transformer 对序列长度是 O(n²) 复杂度，而贪吃蛇每一步都在增加历史长度。

解决方法是 KV Cache（11:20）。每次推理时，Transformer 会重新计算整个序列的 key 和 value 嵌入。如果把之前算好的 key/value 存下来，下次只需要算新 token 的部分，复杂度就从 O(n²) 降到 O(n)。

用 State 实现 KV Cache

Core AI 的 State 机制让这件事变得简单。State 是既被读取、又在推理过程中原地更新的输入（11:48）。

先在 PyTorch 模型里加 cache buffer（12:18）：

class SnakeTransformerStateful(nn.Module):
    def __init__(self, ...):
        super().__init__()
        self.register_buffer(
            "k_cache", torch.zeros(N_LAYERS, 1, MAX_SEQ_LEN, D_MODEL))
        self.register_buffer(
            "v_cache", torch.zeros(N_LAYERS, 1, MAX_SEQ_LEN, D_MODEL))

在 forward 里读写 cache（12:50）：

def forward(self, features, position_ids):
    new_k, new_v = [], []
    for i, block in enumerate(self.blocks):
        k_prev = self.k_cache[i]
        v_prev = self.v_cache[i]
        # ... compute q/k/v for the new token ...
        new_k.append(k_updated)
        new_v.append(v_updated)

    self.k_cache.copy_(torch.stack(new_k))
    self.v_cache.copy_(torch.stack(new_v))
    return self.action_head(self.ln_final(x))

重新转换模型时声明 state_names（12:59）：

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

Swift 端存储和传递 cache（13:17）：

struct ModelPlayer {
    let nextActionFunction: InferenceFunction
    var keyCache: NDArray
    var valueCache: NDArray

    init(modelURL: URL) async throws {
        let model = try await AIModel(contentsOf: modelURL)
        self.nextActionFunction = try model.loadFunction(named: "main")!
        self.keyCache = NDArray(
            shape: [layers, maxContext, hiddenDim],
            scalarType: .float32
        )
        self.valueCache = NDArray(
            shape: [layers, maxContext, hiddenDim],
            scalarType: .float32
        )
    }
}

推理时传入 state views（13:45）：

mutating func chooseAction(game: SnakeGame, snakeID: Int) async throws -> Direction {
    // ... prepare inputFeatures ...

    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
    )
    // ...
}

关键点：

• register_buffer 让 PyTorch 把 cache 当作可变 buffer 导出
• state_names 告诉转换器哪些 buffer 要变成 Core AI 的 state
• Swift 端用 MutableViews 集合包装多个 state，推理时原地更新
• 加了 KV Cache 后，游戏速度保持稳定，不再越玩越慢

模型特化与缓存

.aimodel 文件是源格式，能在任何 Apple 设备上运行。但真正加载前需要经过 specialization（特化），针对具体设备的 CPU/GPU/ANE 生成优化后的可执行代码（15:41）。

特化对大模型可能很慢。Core AI 提供了缓存机制（16:22）：

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...")
    }
}

也可以主动触发特化（16:42）：

try await AIModel.specialize(contentsOf: modelURL)

建议在用户主动开启功能后、或下载完模型资源后立刻调用，避免在交互流程中触发。

特化过程分两步（17:31）：编译（图分割、计划、优化）和生成设备专属的可执行产物。编译占主要时间。Xcode 支持 AOT（Ahead-of-Time）编译，在开发机器上预编译模型，减少用户端的特化时间（17:57）。

进阶优化 API

对于需要极致性能的 tight loop，Core AI 还提供了低层 API（18:39）：

• 查询 NDArray 的最优内存布局并预分配，避免推理时的布局转换
• 预分配输出值，避免推理中动态分配内存
• 用异步值（async values）流水线化多个推理函数的执行

大部分场景用高层 API 就够了，这些低层 API 在优化关键路径时备用。

详细内容

PyTorch 模型转换完整流程

从训练好的 PyTorch 模型到 .aimodel 文件，核心步骤是导出静态计算图、应用 Core AI 算子分解、然后转换保存。以下是一个完整的转换示例：

import torch
import coreai_torch

# 1. 加载 PyTorch 模型并准备示例输入
pt_model = SnakeTransformer().load_checkpoint("snake.pt")
example = torch.randn(1, 5, 16)

# 2. 定义动态维度（序列长度可变）
seq_len = torch.export.Dim("seq_len", min=1, max=256)

# 3. 导出为静态计算图
exported = torch.export.export(
    pt_model, args=(example,),
    dynamic_shapes={"features": {1: seq_len}},
)

# 4. 应用 Core AI 所需的算子分解
exported = exported.run_decompositions(coreai_torch.get_decomp_table())

# 5. 转换为 Core AI 程序并保存
ai_program = coreai_torch.TorchConverter().add_exported_program(
    exported, input_names=["features"], output_names=["logits"],
).to_coreai()

ai_program.save_asset("SnakeTransformer.aimodel")

关键点：

• torch.export.export() 将 PyTorch 动态图捕获为静态计算图，这是转换的前提
• dynamic_shapes 声明哪些维度是动态的，避免为每个序列长度单独转换
• coreai_torch.get_decomp_table() 提供 Core AI 支持的算子分解规则，把复杂 PyTorch 算子拆成 Core AI 能识别的原子操作
• .aimodel 是跨设备的源格式，Xcode 可直接打开查看模型结构和 tensor 形状

转换后务必做数值验证，确保精度损失在可接受范围：

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

pt_model = SnakeTransformer().load_checkpoint("snake.pt")
ai_model = await AIModel.load("SnakeTransformer.aimodel")
function = ai_model.load_function("main")

features = np.array([extract_features(game) for _ in range(10)],
                    dtype=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]

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

关键点：

• 用同一组输入分别跑 PyTorch 和 Core AI，对比最后一层 logits
• 阈值根据业务需求设定，分类任务通常 0.01 以内足够
• 若差值过大，检查 dynamic_shapes 范围是否覆盖测试输入、算子分解是否有遗漏

Swift 端推理与模型封装

Core AI 的 Swift API 设计遵循渐进式披露原则。最基础的用法只需要三个类型：

import CoreAI

// 1. 从 .aimodel 文件加载模型
let model = try await AIModel(contentsOf: modelURL)

// 2. 获取主推理函数（通常是 "main"）
let mainFunction: InferenceFunction = try model.loadFunction(named: "main")!

// 3. 准备输入并执行推理
let inputNDArray: NDArray = nextInput()
var outputs = try await mainFunction.run(inputs: ["input": inputNDArray])

guard let outputNDArray = outputs.remove("output")?.ndArray else {
    // Handle unexpected missing output
}

关键点：

• AIModel 负责解析 .aimodel 文件结构，不实际执行计算
• InferenceFunction 是编译后的可运行计算图，一个模型可包含多个 function
• NDArray 是多维数组的 Swift 表示，支持 float32、int64 等标量类型
• run 是 async 方法，不会阻塞主线程，适合在 UI 交互中调用

实际项目中通常会把推理逻辑封装到业务类型里。以下是一个完整的游戏 AI 玩家实现，包含模型加载、特征写入和动作决策：

struct ModelPlayer {
    let nextActionFunction: InferenceFunction

    init(modelURL: URL) async throws {
        let model = try await AIModel(contentsOf: modelURL)
        self.nextActionFunction = try model.loadFunction(named: "main")!
    }
}

extension ModelPlayer: SnakePlayer {
    mutating func chooseAction(game: SnakeGame) async throws -> Direction {
        // 1. 创建输入 NDArray，shape 为 [当前步数, 特征维度]
        var inputFeatures = NDArray(
            shape: [game.stepCount, hiddenDim],
            scalarType: .float32
        )

        // 2. 通过 MutableView 安全写入特征数据
        writeFeatures(of: game, into: inputFeatures.mutableView())

        // 3. 执行推理
        var outputs = try await nextActionFunction.run(
            inputs: ["features": inputFeatures]
        )

        // 4. 提取输出 logits 并解码为动作方向
        guard let logits = outputs.remove("logits")?.ndArray else {
            throw ModelError.missingOutput
        }

        return predictedDirection(from: logits.view())
    }
}

关键点：

• NDArray.MutableView 是 Swift 的非逃逸类型（non-escapable type），保证内存安全的同时零开销
• mutableView() 获取可写视图用于输入准备，view() 获取只读视图用于输出读取
• 输入特征的 shape 第一维随游戏步数动态增长，这正是需要 KV Cache 优化的原因
• 特征包括：到四面墙的距离、最近食物相对位置、当前方向 one-hot、对手蛇的距离和方向

模型特化与缓存机制

.aimodel 是跨设备的源格式，真正运行前需要 specialization（特化）——针对具体设备的 CPU/GPU/ANE 生成优化后的可执行代码。特化对大模型可能很慢，Core AI 提供了缓存和预触发机制：

// 检查缓存中是否已有特化后的模型
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...")
    }
}

// 主动触发特化（建议在用户开启功能后或下载完模型后立即调用）
try await AIModel.specialize(contentsOf: modelURL)

关键点：

• AIModelCache.default 按设备型号缓存特化结果，同一设备下次加载直接命中
• 缓存未命中时应给用户明确反馈，避免交互流程中突然卡顿
• specialize() 是异步操作，建议在非关键路径调用（如设置页面、下载完成后）
• 特化分两步：编译（图分割、计划、优化）和生成设备专属可执行产物，编译占主要时间
• Xcode 支持 AOT（Ahead-of-Time）编译，可在开发机器上预编译模型，减少用户端等待

Core AI 工具链全景

Core AI 的技术栈可以分成三层：

Python 层：coreai-torch 包支持从 PyTorch 转换、直接 authoring、优化 Apple Silicon、以及用 Metal 4 写自定义 kernel。演讲中提到的高级用法详见 session 325（Dive into Core AI model authoring and optimization）。

Swift 层：CoreAI framework 提供类型安全的 API，利用 Swift 的非逃逸类型（non-escapable types）保证内存安全的同时不牺牲性能。

工具层：Xcode 集成包括 AOT 编译、Core AI 专用 Instruments、可视化 Debugger（能追踪 tensor 值回到 Python 源码行）、以及 Xcode Debug Gauge（实时显示 Core AI 活动）。

Core AI Models 仓库

Apple 维护了一个官方模型仓库（19:26），包含：

• 热门模型的一键转换命令
• 精通 Core AI 转换和优化的 AI skills
• Swift package，为特定模型家族提供高层 API，内置低层优化
• 创建 Core AI Language Model 的 API，可接入 Foundation Models framework

核心启发

1. 在 App 里跑本地图像分类

做什么：用 coreai-torch 把训练好的 ResNet/ViT 转成 .aimodel，Swift 端几行代码加载推理。

为什么值得做：不需要网络请求，隐私数据不出设备，没有 API 调用成本。对于照片分类、物体识别这类高频功能，端侧推理的延迟和成本优势很明显。

怎么开始：用 coreai-torch 转换模型，Swift 端用 AIModel(contentsOf:) 加载，loadFunction(named: "main") 获取推理函数，传入 NDArray 执行。入口 API：AIModel(contentsOf:) + InferenceFunction.run

2. 给游戏加 AI 对手

做什么：像演讲里的贪吃蛇一样，用 Transformer 学习游戏策略，通过 Core AI 在设备上实时推理。

为什么值得做：端侧推理没有网络延迟，AI 对手可以实时响应游戏状态。KV Cache 解决 Transformer 序列长度增长导致的性能下降，长对局不会越玩越卡。

怎么开始：在 PyTorch 模型里用 register_buffer 定义 k/v cache，转换时通过 state_names 声明，Swift 端用 InferenceFunction.MutableViews 传入和更新 cache。入口 API：state_names + InferenceFunction.MutableViews

3. 端侧文本生成

做什么：从 Core AI Models 仓库获取预转换的 LLM，用 Foundation Models framework 接入。

为什么值得做：用户输入直接本地处理，零 token 成本，隐私数据不上传。对于聊天助手、文本摘要这类功能，端侧运行消除了对网络和服务器的依赖。

怎么开始：从 Core AI Models 仓库下载预转换模型，用 CoreAILanguageModel 加载，创建 LanguageModelSession 执行生成。入口 API：CoreAILanguageModel(resourcesAt:) + FoundationModels framework

4. 实时音频处理

做什么：用小模型做语音活动检测或说话人分离，Core AI 同时调度 CPU/GPU/ANE 保证低延迟。

为什么值得做：音频处理对延迟极度敏感，端侧推理避免了网络往返。ANE（Neural Engine）对这类小模型的能效比远高于 CPU/GPU，可以持续后台运行而不显著耗电。

怎么开始：用 AIModel.specialize(contentsOf:) 在用户首次开启功能时提前编译模型，避免交互流程中触发特化导致的卡顿。配合 Core AI 专用 Instruments 调优。入口 API：AIModel.specialize + Core AI Instruments

5. 自定义视觉模型

做什么：用 Python 端直接 authoring Core AI 计算图，写 Metal 4 kernel 实现自定义算子，再集成到 App。

为什么值得做：标准算子库覆盖不了所有研究前沿的模型结构。自定义 kernel 让你能把论文里的新算子直接部署到 Apple Silicon 上，同时享受 GPU/ANE 的硬件加速。

怎么开始：用 coreai-torch 的 authoring API 构建计算图，手写 Metal Shading Language kernel 实现自定义算子，转换时通过 register_custom_kernels 注册。入口 API：coreai-torch authoring API + Metal 4 custom kernels

关联 Session

• Dive into Core AI model authoring and optimization — 深入讲解 Core AI Python 包的模型 authoring、优化和自定义 kernel
• Integrate Core AI into your app — Core AI 在 App 中的集成最佳实践
• Explore MLX Swift — Apple 的 MLX 框架 Swift 绑定，与 Core AI 互补的本地 AI 方案
• Bring your model to Apple silicon — Metal 张量运算和模型部署的底层机制
• Foundation Models on device — 设备端大语言模型的 Foundation Models framework