WWDC Quick Look 💓 By SwiftGGTeam
中文 EN 日本語
Deploy machine learning and AI models on-device with Core ML

Deploy machine learning and AI models on-device with Core ML

WWDC 2024
观看原视频

Highlight

Core ML 新增 MLTensor 类型用于高效张量计算，State API 优化有状态模型的推理效率，多函数模型支持单个 MLPackage 承载多种功能。

核心内容

端侧运行机器学习模型时，开发者面临一个现实问题：模型推理本身只是工作流的一部分。以大语言模型为例，生成文本需要解码器从词汇表分数中采样下一个 token，这部分”胶水代码”——包括概率分布计算、温度调整、Top-K 采样——之前要么手写实现，要么调用底层 API，代码量大且易出错。

Core ML 今年引入的 MLTensor 正是为解决这个问题。它是一个多维数组类型，提供类似 NumPy/PyTorch 的数学运算和变换操作，直接利用 Apple silicon 的 CPU、GPU 和 Neural Engine 加速。Session 中的演示表明，同样的解码逻辑，用 MLTensor 实现比手写底层 API 代码量少得多。

另一个痛点是 KV Cache 管理的有状态推理。之前手动管理缓存需要每次迭代时把缓存作为输入传入、从输出中读取更新值，数据在内存和计算单元之间反复搬运。Core ML 新的 State API 将状态管理交由框架内部处理，减少数据搬运开销。Session 中实测 Mistral 7B 模型在 M3 Max MacBook Pro 上，使用 State 后约 5 秒完成，不使用约 8 秒，提速约 1.6 倍（12:11）。

部署多个相关模型时也存在冗余——例如同一扩散模型配合不同风格适配器，之前需要部署多个独立模型。多函数模型允许单个 MLPackage 包含多个函数入口，共享基础权重，减少部署体积和内存占用。

此外，iOS 18 的推理栈做了底层优化，很多模型无需重新编译或修改代码即可获得性能提升（02:27）。

详细内容

MLTensor：端侧张量计算

MLTensor 是 Core ML 新增的多维数组类型，由 shape 和 scalar type 定义，支持从 MLShapedArray 或嵌套标量集合创建（04:03）：

// 从 MLShapedArray 创建
let tensor1 = MLTensor(shapedArray: shapedArray)

// 从嵌套标量集合创建
let tensor2 = MLTensor([[1.0, 2.0], [3.0, 4.0]])

关键点：

  • MLTensor 可从 MLShapedArray 或嵌套字面量构造
  • 与 MLShapedArray 一样由 shape 和 scalar type 描述
  • 嵌套字面量的层级决定维度数

张量支持逐元素数学运算、自动广播、比较和掩码操作（04:33）：

// 逐元素加法与乘法，字面量自动广播
let result = (tensor1 + 10) * tensor2

// 计算均值
let mean = result.mean()

// 比较运算创建布尔掩码
let mask = result .> mean

// 掩码相乘实现过滤
let filtered = mask * result

关键点：

  • 字面量与张量运算时自动广播兼容形状
  • .mean() 求所有元素的均值
  • .> 比较返回布尔型张量掩码
  • 布尔掩码与张量相乘，false 位置置零

索引和形状变换（04:57）：

// 切片：取矩阵第一行
let firstRow = matrix[0]

// 重塑形状
let reshaped = firstRow.reshape([1, -1])

关键点：

  • 索引语法按维度切片，与 Python 数值库类似
  • .reshape() 改变形状但不改变数据

所有张量操作异步执行，需显式 materialize 才能访问数据（05:12）：

let shapedArray = try await result.materialize()

关键点：

  • 异步执行确保上游操作完成后数据才可用
  • materialize 将 MLTensor 转回 MLShapedArray

用 MLTensor 简化 LLM 解码

Session 用 HuggingFace Swift Transformer + Mistral 7B 模型演示了解码器对比（06:33）。语言模型输出词汇表所有词的分数（logits），解码器负责从中选取下一个 token——策略包括贪心解码（取最高分）和 Top-K 采样（在概率最高的 K 个词中随机采样），还能通过 temperature 调整概率分布的平坦程度（06:57）。

对比结果：MLTensor 实现同样功能的代码量远少于手写底层 API 版本。MLTensor 适合这类常见 ML 运算场景，但底层 API 在需要精细控制时仍有其价值（07:57）。

State API：优化有状态推理

传统方式手动管理 KV Cache（11:04）：

// 手动预分配缓存
var keyValueCache = createEmptyCache()

for token in inputTokens {
    // 缓存作为输入传入
    let prediction = try model.prediction(from: input, using: keyValueCache)
    // 从输出中取出更新后的缓存
    keyValueCache = prediction.cache
}

使用 State API 后（11:26）：

// 模型实例创建并预分配状态
let state = model.makeState()

for token in inputTokens {
    // 传入状态，更新原位完成
    try model.prediction(from: input, state: state)
}

关键点：

  • model.makeState() 由 Core ML 预分配状态缓冲区，返回状态句柄
  • 通过句柄访问缓冲区并控制状态生命周期
  • 状态更新原位完成，无需从输出中读取再赋值
  • 模型需在准备阶段（coremltools 转换时）显式启用 State 支持

可通过 Xcode 模型预览的 Predictions 标签页确认模型是否有 State（10:48）。

多函数模型

Session 以 Stable Diffusion XL + 两个风格适配器（sticker 和 storybook）为例（14:25）。两个适配器共享同一个 UNet 基础模型，合并为单个 MLPackage 后各自暴露为一个函数：

// 加载时指定函数名
let model = try MLModel(
    contentsOf: url,
    configuration: MLModelConfiguration(functionName: "sticker")
)

关键点：

  • MLModelConfiguration(functionName:) 指定加载哪个函数
  • 未指定时使用默认函数
  • 不同函数可以有不同输入输出签名
  • Xcode 预览可查看模型所有可用函数

同一管线可以复用于两个函数，因为 sticker 和 storybook 的输入输出签名相同。Pipeline 中多个模型通过 MLTensor 无缝拼接（14:35）。

性能工具增强

Core ML 性能报告新增两项信息（16:05）：

  • 操作级耗时估算：每个操作的 estimated time = 中位预测时间 x 该操作估计相对成本，可排序找出瓶颈
  • 不支持原因提示：鼠标悬停在不支持的操作上，会提示原因（如数据类型不支持），方便回溯到准备阶段修改

报告现在支持导出和跨版本对比，评估模型变更效果无需写代码（16:56）。

新增 MLComputePlan API 提供编程式访问：模型结构、每个操作的支持/首选计算设备、支持状态、估计相对成本（17:13）。

核心启发

  • 用 MLTensor 替换手写解码逻辑：端侧 LLM 的 token 解码（概率计算、采样、温度调整）是最典型的适用场景。先用 MLTensor 重写概率计算和过滤部分，验证功能一致后再逐步替换整个解码器。代码量减少的同时，异步执行也避免了不必要的同步等待。

  • 启用 State API 减少 KV Cache 开销：如果你的模型涉及多轮迭代推理，用 coremltools 重新导出模型并启用 State。Mistral 7B 实测提速 1.6 倍，模型越大缓存越多，收益越明显。先在 Xcode 预览确认 State 可用，再用 makeState() 替换手动缓存管理代码。

  • 用多函数模型合并共享权重的模型：部署多个风格适配器或任务变体时，将它们合并为单个 MLPackage，每个变体对应一个函数。减少安装体积和运行时内存占用，加载时指定 functionName 即可切换功能。

关联 Session

评论

GitHub Issues · utterances