Core AI 端侧推理栈：从 PyTorch 到 Neural Engine 的协同优化实践

Apple 在 WWDC26 正式向开发者开放了 Core AI 框架 —— 这套原本支撑 Apple Intelligence 端侧推理的基础设施。与 Siri 和云端大模型交互的产品层不同，Core AI 定位在底层执行栈，让开发者能够将自研或开源的 Transformer、CNN 等模型直接部署到 iPhone、iPad、Mac 乃至 Vision Pro 上运行，无需网络连接，也无需按 token 付费。

框架架构与执行模型

Core AI 的设计遵循「一次转换，多端运行」的思路。框架核心由三部分构成：Python 工具链（模型转换与优化）、运行时引擎（跨 Apple Silicon 异构计算调度）、以及 Swift API（应用层集成）。

在硬件抽象层，Core AI 自动调度 CPU、GPU 和 Neural Engine（ANE），根据模型算子特性、输入形状和当前设备负载动态选择执行后端。对于 Vision Pro 或 iPhone 17 Pro 等配备更强 ANE 的设备，框架会优先将卷积和矩阵运算下沉到专用神经网络引擎，显著降低功耗。

模型转换：从 PyTorch 到 .aimodel

Core AI 的模型转换流程基于 PyTorch 的 torch.export 机制。开发者通过 coreai-torch 库将训练好的模型导出为 .aimodel 格式 —— 这是一种针对 Apple Silicon 优化的中间表示。

转换代码结构清晰：首先使用 torch.export.export() 捕获计算图，支持通过 dynamic_shapes 参数声明可变维度（如序列长度）；随后调用分解表处理复杂算子；最终通过 TorchConverter 生成可在 Core AI 运行时加载的资产文件。关键代码模式如下：

seq_len = torch.export.Dim("seq_len", min=1, max=256)
exported = torch.export.export(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("Model.aimodel")

转换后的模型可通过数值验证确保一致性 —— 加载原始 PyTorch 模型和 Core AI 模型，对相同输入比较输出 logits 的差异，通常要求最大绝对误差小于 0.01。

Swift API 与推理执行

在应用侧，Core AI 提供了类型安全的 Swift API。开发者使用 AIModel 类加载 .aimodel 文件，通过 loadFunction(named:) 获取入口函数，然后构造 NDArray 输入并执行推理。

let model = try await AIModel(contentsOf: modelURL)
let function = try model.loadFunction(named: "main")!
let input = NDArray(shape: [1, seqLen, hiddenDim], scalarType: .float32)
// 写入特征数据...
let outputs = try await function.run(inputs: ["features": input])

NDArray 是 Core AI 的张量抽象，支持通过 mutableView() 获取可写视图，在预处理阶段直接填充特征数据，避免不必要的内存拷贝。

状态管理与 KV Cache 优化

对于自回归语言模型，Core AI 支持有状态推理（stateful inference）。开发者可以在 PyTorch 模型中通过 register_buffer 声明 KV Cache，转换时通过 state_names 参数将其标记为可变状态。这样，推理时只需传入新 token，框架会自动读写缓存，避免重复计算历史注意力，将时间复杂度从 O (n²) 降到 O (n)。

Swift 侧的集成需要显式管理状态缓冲区：

struct ModelPlayer {
    let nextActionFunction: InferenceFunction
    var keyCache: NDArray
    var valueCache: NDArray
    
    mutating func infer(...) async throws {
        var stateViews = InferenceFunction.MutableViews()
        stateViews.insert(&keyCache, for: "keyCache")
        stateViews.insert(&valueCache, for: "valueCache")
        let outputs = try await function.run(
            inputs: ["features": input],
            states: stateViews
        )
    }
}

AOT 编译与模型特化

Core AI 引入了 Ahead-of-Time（AOT）编译机制，允许开发者在构建阶段预先将模型编译为设备特定的二进制格式。通过命令行工具 coreai-build compile，可以为 iOS、macOS 等平台生成优化后的模型资产：

xcrun coreai-build compile MyModel.aimodel --platform iOS

AOT 编译将原本在首次运行时进行的「模型特化」（specialization）工作前置到开发机器上执行。特化过程包括算子融合、内存布局优化、以及针对目标设备 ANE 版本的指令生成。未经 AOT 编译的模型在首次加载时需要现场特化，可能导致数秒到数十秒的延迟，影响用户体验。

框架提供了 AIModelCache API 用于管理特化后的模型缓存：

let cache = AIModelCache.default
if let model = try cache.model(for: modelURL, options: .default) {
    // 使用缓存的特化模型
} else {
    // 首次加载，触发特化
    try await AIModel.specialize(contentsOf: modelURL)
}

性能分析与调试工具

Xcode 为 Core AI 提供了专门的 Instruments 模板，可实时采集模型推理延迟、内存占用和硬件调度信息。开发者能够识别性能瓶颈 —— 例如，若发现推理时间随序列长度呈二次增长，则表明 KV Cache 未正确配置。

此外，Core AI Debugger 支持数值级调试，允许逐层比对 PyTorch 参考输出与 Core AI 实际输出，精确定位转换过程中的数值偏差来源。

部署策略与最佳实践

在实际部署中，建议采用以下策略平衡包体大小与首次加载体验：

1. 模型分离：将大模型资产从主 App Bundle 中剥离，通过 Background Assets 按需下载，避免更新时全量下载模型文件
2. 渐进式体验：在首次启用 AI 功能时展示「准备中」状态，后台完成模型特化并缓存
3. 多版本管理：针对不同设备能力准备多个模型版本（如 0.6B 轻量版与 8B 高质量版），运行时动态选择
4. 批处理优化：在 macOS 上利用更大的内存和算力，对批量输入进行并行推理

结语

Core AI 的开放标志着 Apple 端侧 AI 生态从「系统独占」走向「开发者可用」。通过统一的 Python-to-Swift 工具链、自动异构调度、以及 AOT 编译等工程优化，开发者能够在保证隐私的前提下，将 Transformer、SAM 等先进模型部署到数亿台 Apple 设备上。对于追求零网络延迟、零推理成本的应用场景，Core AI 提供了目前移动端最完整的端侧推理解决方案。

资料来源

• Apple Developer, "Meet Core AI", WWDC26 Session 324
• Apple Developer, "Integrate on-device AI models into your app using Core AI", WWDC26 Session 326
• Apple Developer Documentation, Core AI Framework

ai-systems