了解 Core AI

一句话判断

Apple 终于把驱动 Apple Intelligence 的底层推理引擎（Inference framework）交出来了，Core AI 不是 Core ML 的简单升级，而是专为 Transformer 和端侧大模型重写的新一代范式。

这场 Session 讲了什么

过去几年大家在端侧跑大模型，用 Core ML 各种受罪：转换慢、不支持动态 Shape、没有原生的 KV Cache（Key-Value cache）支持。Apple 内部为了搞 Apple Intelligence，自己造了个新轮子叫 Core AI。现在他们把这个框架开放了。

它直接抛弃了以前那套繁琐的 CoreML Tools，改用 Python 的 coreai-torch 包直接对接 PyTorch 的 torch.export。在 Swift 端，新的 CoreAI framework 提供了极具表达力的 API，甚至支持通过 MutableView 直接操作底层内存来更新 KV Cache。

这对所有想在 App 里塞进本地 LLM（Large Language Model）或复杂视觉模型的开发者是降维打击。你不再需要自己手写 Metal shader 或者忍受 Core ML 的龟速转换，直接享受 Apple Silicon 全火力（CPU/GPU/ANE）加速。

值得深挖的点

拥抱 Stateful 模型：KV Cache 的原生支持

在 Core ML 时代，跑 Transformer 最痛苦的就是处理 KV Cache。因为 Core ML 偏向 Stateless（无状态）设计，每次推理都要把巨大的 Cache 作为 Input 传进去，再作为 Output 拿出来，导致严重的内存拷贝和带宽浪费。

Core AI 彻底改变了这一点。它允许你在 PyTorch 里用 register_buffer 定义 Cache，转换后，在 Swift 端可以通过 NDArray 的 mutableView() 直接获取这块内存的指针。推理时，模型直接原地（in-place）读写这块内存。这不仅省去了海量的数据搬运，还让自回归（Autoregressive）生成的延迟大幅降低。Apple 这个设计非常务实，完全对齐了当前 LLM 推理的主流工程实践，把端侧大模型的可用性拉到了及格线以上。

Specialization 与 AOT 编译：解决“首次加载卡顿”

端侧部署大模型有个致命坑：模型首次加载时，框架需要根据当前设备的硬件（比如 iPhone 15 Pro 和 iPad Pro 的 GPU 核心数不同）进行图优化和 Kernel 编译。这会导致用户点开 App 后卡好几秒，甚至触发 watchdog 崩溃。

Core AI 引入了 Specialization（特化）机制。你可以在 Xcode 里配置目标设备，利用 AOT（Ahead-of-Time，提前编译）在构建阶段就把模型编译成特定 Apple Silicon 架构的优化二进制。虽然这会让 App 包体积变大，但换来的是用户端真正的“即插即用”和零延迟启动。这是 Apple 利用自家软硬一体优势打出的王炸，第三方跨平台框架很难做到这种级别的深度优化。

代码片段

1. 用 coreai-torch 转换 PyTorch 模型

场景：把带有动态 Sequence Length 的 PyTorch 模型转成 Core AI 格式。

import torch
import coreai_torch

# 加载训练好的模型
pt_model = SnakeTransformer().load_checkpoint("snake.pt")
example_input = torch.randn(1, 5, 16)

# 定义动态维度，支持变长序列
seq_len = torch.export.Dim("seq_len", min=1, max=256)
exported = torch.export.export(
    pt_model, args=(example_input,), 
    dynamic_shapes={"features": {1: seq_len}},
)

# 运行分解并转换为 Core AI 图
exported = exported.run_decompositions(coreai_torch.get_decomp_table())
ai_program = coreai_torch.TorchConverter().add_exported_program(exported)
ai_program.save("SnakeTransformer.aimodel")

坑：dynamic_shapes 的 min 和 max 必须严格覆盖你运行时的所有可能长度，超范围直接报错。

2. Swift 端加载与基础推理

场景：在 App 里加载 .aimodel 并传入 NDArray 进行推理。

import CoreAI

// 异步加载 .aimodel 文件
let model = try await AIModel(contentsOf: modelURL)
// 加载主推理函数
let mainFunction = try model.loadFunction(named: "main")!

// 构造 N 维输入数据
var inputFeatures = NDArray(shape: [game.stepCount, hiddenDim], scalarType: .float32)
// 将游戏特征写入底层内存
writeFeatures(of: game, into: inputFeatures.mutableView())

// 执行推理并提取输出
var outputs = try await mainFunction.run(inputs: ["features": inputFeatures])
guard let logits = outputs.remove("logits")?.ndArray else {
    throw ModelError.missingOutput
}

坑：AIModel 的加载是 async 的，千万别在主线程同步阻塞等待，老老实实用 Task 包装。

3. 原地更新 KV Cache (Stateful 推理)

场景：利用 MutableView 避免大模型推理时的内存拷贝。

// 假设 model 已经加载，且包含 k_cache 和 v_cache 状态
let kCache = try model.loadState(named: "k_cache")!
let vCache = try model.loadState(named: "v_cache")!

// 获取底层内存的可变视图，直接传给底层操作
let kCacheView = kCache.mutableView()
let vCacheView = vCache.mutableView()

// 在推理时，将 state 作为参数传入，模型会原地修改这块内存
let inputs: [String: Any] = [
    "features": inputFeatures,
    "position_ids": positionIds
]
// Core AI 会自动识别 state_names 并原地更新 cache
let _ = try await mainFunction.run(inputs: inputs) 

坑：使用 mutableView() 时，确保你的 Swift 对象生命周期覆盖了整个推理过程，提前释放会导致野指针崩溃。

最佳实践

• 已有项目的迁移策略：如果你的项目还在用 Core ML 跑传统的 CV（计算机视觉）小模型，且性能已经达标，不要急着换。Core ML 依然稳健。但如果你在跑 Transformer、LLM 或需要处理长序列的音频模型，立刻评估 Core AI，KV Cache 的原生支持能帮你省掉至少 30% 的推理耗时。
• 新项目的采用建议：全面拥抱 PyTorch + coreai-torch 工作流。放弃以前那种“先转 ONNX 再转 Core ML”的曲线救国路线。直接在 Python 端用 torch.export 锁定计算图，用 coreai-torch 验证数值对齐（Numerical correctness）。
• 实战避坑：大模型上端，内存比算力更致命。务必在 Xcode 的新 Core AI Instrument 里盯紧内存峰值。利用 Specialization 提前编译模型，虽然会增加构建时间，但能避免用户在 App 启动时经历漫长的“模型预热”黑屏。

还有什么值得关注

• Core AI Debugger：可以直接把 Swift 端的推理 Tensor 值映射回 Python 源码行进行断点调试，找精度 loss 的神器。
• Core AI Models Repository：Apple 官方开源的模型库，里面有一堆调优好的热门架构，直接拿来用能省去几周调参时间。
• Debug Gauge：Xcode 里新增的实时仪表盘，能在 App 运行时流式监控 ANE/GPU 的利用率和推理延迟，不用再苦哈哈地看 Instruments 离线日志了。