苹果神经引擎可直接训练运行Transformer

背景

长期以来，在 Apple Silicon 设备上运行 Transformer 模型（如 GPT、Llama 等）主要依赖 CoreML 或 Metal 后端。CoreML 虽然提供了统一的推理接口，但针对 Apple Neural Engine（ANE，神经引擎）的优化并不彻底——模型需要通过 MIL（Model Intermediate Language）编译，且每次解码步骤都可能涉及重新编译或数据搬运，导致性能瓶颈。与此同时，llama.cpp 等流行方案并未提供 ANE 后端，只能回退到 GPU（Metal）或 CPU，无法充分利用 ANE 的高能效计算能力。

在此背景下，开发者 Christopher Karani 发布了名为 Espresso 的开源框架，直接通过逆向工程发现的私有 API（_ANEClient、_ANEInMemoryModel）与 ANE 硬件通信，完全跳过 CoreML，实现 Transformer 模型的训练与推理。该项目发布在 GitHub 上（christopherkarani/Espresso），采用 MIT 许可协议，面向内部工具、研究及企业分发场景。

核心内容

Espresso 是一个纯 Swift 6.2 实现的无依赖库，允许用户在 Apple Silicon（M 系列芯片）上直接编译、训练和运行 Transformer 模型。其核心架构如下：

• 直接 ANE 推理：通过逆向的私有 API 将 MIL 程序直接编译到 ANE 硅片上，生成 E5 二进制缓存。编译过程使用 _ANEClient 接口，推理时通过 _ANEInMemoryModel 执行。
• 性能数据：在相同 6 层 Transformer 模型（dim=768，12 头，32k 词表，seqLen=256）测试中，Espresso 的解码速度为 1.08 ms/token，而 CoreML 为 5.09 ms/token，提升约 4.76 倍。
• 融合多层内核：将 3 个 Transformer 层融合为一个 ANE 调度（dispatch），6 层模型仅需 2 次 ANE eval 调用，而非逐层 6 次。
• 零拷贝 I/O：使用 IOSurface 缓冲区管理输入/输出和 KV 缓存，通过 NEON 向量化读取和 vDSP argmax 操作，无需数据序列化/反序列化（marshaling）。
• 完整训练能力：支持在 ANE 上进行前向传播、反向传播、梯度累积和 Adam 优化器。
• 纯 Swift 6.2 实现：采用 ~Copyable move-only 张量、严格并发、typed throws，零外部依赖。

快速使用示例：

1. 克隆仓库并运行 ./espresso，自动构建、下载演示权重、启动 TUI。
2. 在 Swift Package 中添加依赖，然后创建 ANEKernel：传入 MIL 文本、权重 blob、输入/输出尺寸，调用 eval() 即可在 ANE 上运行推理，返回零拷贝的 IOSurface 输出。
3. 命令行工具：espresso "Hello" 生成文本；espresso doctor 检查主机准备状态；espresso compare --no-power "Hello" 与 CoreML 对比；espresso install 安装到 ~/.local/bin。

模型打包与运行时选择：

• .esp 是规范的便携式模型包格式。
• espc 将原生模型目录打包为 .esp，并生成编译缓存层。
• esprun 检查、解析并运行包中的 artifact。
• espresso-generate --bundle <path> 使用与运行时相同的包边界。

技术栈分层：

• ANEInterop（ObjC/C 私有 API 桥接）
• ANETypes（~Copyable 值类型，IOSurface I/O）
• MILGenerator（28+ 内核变种）
• ANERuntime（编译、eval、表面管理）
• Espresso（训练、生成、解码）
• CPUOps（基于 Accelerate/vDSP 的内核）

解码循环：程序仅编译一次，所有步骤重用。KV 缓存直接位于 IOSurface 缓冲区中，不经过 CoreML。每个解码步骤生成两个精确 token，并经过验证。融合的三层内核每次 dispatch 处理 3 层，6 层只需 2 次 eval 调用。

硬件/OS 限制：需要 macOS 15+；不支持 iOS/tvOS（私有 API entitlement 因平台不同）。Apple App Store 会拒绝使用私有 ANE API（_ANEClient、_ANEInMemoryModel）的应用，但在内部工具、研究、侧载应用和企业分发中均可使用。

基准测试再现：提供 scripts/reproduce_local_real_artifact_claim.sh 脚本，可通过 RESULTS_DIR、REPEATS、WARMUP、ITERATIONS 环境变量自定义。机器可读结果输出到 artifacts/benchmarks/。

测试套件：7 个测试套件涵盖 MIL 生成、张量运算、CPU 内核、ANE 编译、硬件 eval、交叉验证和端到端生成。运行 swift test 执行单元测试（无需 ANE），ANE_HARDWARE_TESTS=1 swift test 执行硬件测试，OBJC_CROSS_VALIDATION=1 ANE_HARDWARE_TESTS=1 执行奇偶校验。

关键要点

• 性能飞跃：解码速度比 CoreML 快 4.76 倍（1.08 ms/token vs 5.09 ms/token），主要得益于融合内核、零拷贝 I/O 和避免逐 token 重新编译。
• 完全绕过 CoreML：通过逆向的私有 API 直接与 ANE 通信，无需 CoreML 介入，从而消除中间层开销。
• 训练与推理一体化：不仅支持推理，还支持完整的 ANE 训练（前向+反向+梯度累积+Adam），这在现有方案中罕见。
• 纯 Swift 6.2 实现，零依赖：使用现代 Swift 特性（~Copyable、严格并发、typed throws），无需任何外部库，只依赖系统框架。
• 便携式模型包 .esp：定义了标准化打包格式和编译缓存层，便于分发和运行时选择。
• 适用范围明确：不能上架 App Store，但适用于内部工具、研究、企业分发和侧载应用。需要 macOS 15+。
• 开源与可验证：MIT 许可，提供完整的基准测试再现脚本和测试套件，确保结果可重复。
• 非苹果官方项目：使用通过运行时内省发现的未文档化私有 API，不隶属于苹果公司；性能结果依赖硬件和 OS 版本。

意义与影响

Espresso 项目展示了通过逆向工程挖掘 Apple 私有 API 的潜力，为在 Apple Silicon 设备上高效运行 Transformer 模型开辟了新途径。其意义体现在：

1. 释放 ANE 潜能：ANE 等神经引擎设计初衷是加速特定推理任务，但官方工具（CoreML）并未充分暴露其底层能力。Espresso 证明，通过直接访问私有 API，可以获得远超 CoreML 的性能（4.76 倍），同时支持训练，使得 ANE 成为真正可用的通用加速器。

2. 对开发者生态的影响：为 macOS 研究者和内部工具开发者提供了高性能、低延迟的本地 Transformer 推理方案，尤其适用于需要低功耗、低延迟的应用（如实时文本生成、本地 AI 助手）。强调零依赖和 Swift 现代特性，降低了集成门槛。

3. 挑战与风险：使用私有 API 意味着不稳定性——未来的 macOS 更新可能改变或移除这些接口，且无法在 iOS/tvOS 上使用。这一做法不适合面向大众的 App Store 应用，但为实验性项目和垂直领域（如企业内部 AI 工具）提供了可行路径。

4. 推动竞争与改进：可能促使苹果在未来的 CoreML 或新框架中引入类似的高效接口，或者招募相关人才。同时，开源社区的贡献（如通过 GitHub Issues 报告 bug 和功能请求）有望进一步完善项目。