Apple M3 Ultra上实时扩散模型推理的系统优化

# 实时扩散模型在 Apple M3 Ultra 上的系统化优化推理
来源：https://arxiv.org/html/2605.16259
Yoichi Ochiai 筑波大学 图书馆、信息与媒体科学学院 wizard@slis\.tsukuba\.ac\.jp

###### 摘要

尽管使用扩散模型的实时图像生成在 NVIDIA GPU 上取得了快速进展，但在非 CUDA 平台（如 Apple Silicon）上的系统优化研究仍然极其有限。本研究针对 Apple M3 Ultra（60 核 GPU，512 GB 统一内存），以实现实时摄像头 img2img 转换为目标，在 10 个阶段进行了全面的优化实验。我们探索了多种技术，包括 CoreML 转换、量化、Token Merging、神经引擎利用、紧凑模型探索、帧插值、基于 kNN 搜索的合成、pix2pix-turbo、光流跳帧以及知识蒸馏，并定量评估了每种方法的有效性。最终，通过将蒸馏专用模型 SDXS-512 的 CoreML 转换与 3 线程摄像头流水线相结合，我们在 512×512 分辨率下实现了 22.7 FPS 的实时摄像头 img2img 转换。本研究的主要贡献是系统性地证明，针对 CUDA 建立的优化见解在 Apple Silicon 的统一内存架构上并不一定有效。我们揭示了一个与 NVIDIA GPU 根本不同的优化图景——包括量化无法加速、并行推理无效以及神经引擎不适合大规模模型——并为 Apple Silicon 上的扩散模型推理提供了实用指南。

## 1 引言

扩散模型\[1 (https://arxiv.org/html/2605.16259#bib.bib1),2 (https://arxiv.org/html/2605.16259#bib.bib2)\]已成为文本到图像生成和图像转换的主流范式。然而，其迭代去噪过程计算成本高昂，实时推理仍然是一个具有挑战性的问题。近年来，加速技术发展迅速，包括诸如 SD-Turbo\[3 (https://arxiv.org/html/2605.16259#bib.bib3)\] 和 SDXS\[4 (https://arxiv.org/html/2605.16259#bib.bib4)\] 等单步蒸馏方法，以及诸如潜在一致性模型\[5 (https://arxiv.org/html/2605.16259#bib.bib5)\] 等少步推理方法。StreamDiffusion\[8 (https://arxiv.org/html/2605.16259#bib.bib8)\] 通过流水线级优化在 NVIDIA GPU 上实现了超过 100 FPS 的实时推理。

然而，这些加速研究几乎都假设使用 NVIDIA GPU 和 CUDA 生态系统。CUDA 受益于数十年积累的资源，包括广泛的核函数库、成熟的性能分析工具以及通过 TensorRT 进行的推理优化。相比之下，针对非 CUDA 平台（如 Apple Silicon、Qualcomm Snapdragon 和 Intel Arc GPU）的扩散模型优化系统研究几乎不存在。

Apple M3 Ultra 是一款片上系统，具有多达 192 个 GPU 核心、192 GB 统一内存和 800 GB/s 内存带宽，采用独特的统一内存架构，CPU、GPU 和神经引擎共享同一内存空间。虽然这种设计消除了 CPU-GPU 数据传输的需要，但其内存访问模式和计算特性与 CUDA 根本不同，这意味着在 NVIDIA GPU 上有效的优化技术不一定能直接应用。

在本研究中，我们从 StreamDiffusion\[8 (https://arxiv.org/html/2605.16259#bib.bib8)\] 出发，在 M3 Ultra（60 核 GPU，512 GB 统一内存配置）上进行了跨 10 个阶段的系统优化实验，目标是实现实时摄像头 img2img 转换。本工作的贡献如下：

- • 涵盖超过 10 种技术在 Apple Silicon 上进行扩散模型推理的全面基准测试
- • 证明 CoreML 转换是唯一有效的 UNet 加速技术，并分析其根本原因
- • 系统阐明量化、Token Merging、并行推理等适用于 CUDA 环境的技术为何在 M3 Ultra 上无效
- • 展示利用 512 GB 内存替换扩散模型为 kNN 搜索的根本局限性
- • 使用 SDXS-512 CoreML 流水线实现 22.7 FPS 的实时 img2img 转换
- • 为统一内存架构上的扩散模型优化提供实用指南

## 2 相关工作

### 2.1 快速扩散模型

针对实时推理的 Stable Diffusion\[2 (https://arxiv.org/html/2605.16259#bib.bib2)\] 加速主要沿着两个方向进行：模型蒸馏和推理步骤减少。

SD-Turbo\[3 (https://arxiv.org/html/2605.16259#bib.bib3)\] 是第一个通过对抗性扩散蒸馏实现单步推理的实用模型，结合了对抗训练和分数蒸馏。SDXS\[4 (https://arxiv.org/html/2605.16259#bib.bib4)\] 采用了更先进的蒸馏方法，使用轻量级架构实现了高质量的单步图像生成，完全移除了 UNet 中间块，并将下/上块从 4 个减少到 3 个。参数量减少到标准 SD-Turbo 865.9M 的大约 38%（328.2M）。

潜在一致性模型\[5 (https://arxiv.org/html/2605.16259#bib.bib5)\] 通过一致性蒸馏能够在 2-4 步内实现高质量生成。LCM-LoRA\[6 (https://arxiv.org/html/2605.16259#bib.bib6)\] 使得这种蒸馏可以作为 LoRA 适配器应用，从而能够改造现有模型。Hyper-SD\[7 (https://arxiv.org/html/2605.16259#bib.bib7)\] 将分数蒸馏与基于人类反馈的强化学习相结合，在 1-4 步配置下实现了高质量生成。这些方法主要在 NVIDIA GPU 上评估，没有在 Apple Silicon 上的性能报告。

### 2.2 实时推理流水线

StreamDiffusion\[8 (https://arxiv.org/html/2605.16259#bib.bib8)\] 是一个在流水线层面优化实时扩散模型推理的框架。其核心创新是 Stream Batch 机制，将连续帧的不同去噪步骤打包到单个批次中，以最大化 GPU 计算并行性。例如，在 4 步推理的情况下，帧 t 的第 1 步、帧 t-1 的第 2 步、帧 t-2 的第 3 步和帧 t-3 的第 4 步作为单个批次处理。此外，Residual CFG 重用前一帧的分类器自由引导结果，将每帧的 UNet 调用次数减半。这些优化使得在 NVIDIA RTX 4090 上实现了超过 100 FPS 的实时推理。

然而，StreamDiffusion 的加速严重依赖于 NVIDIA 特定技术，如 CUDA Streams、TensorRT 和 xformers，移植到 Apple Silicon 需要替代方案。此外，当使用单步推理（SD-Turbo、SDXS 等）时，Stream Batch 机制不适用（只有一步可以批处理），因此优化其他流水线组件变得至关重要。

### 2.3 Apple Silicon 上的机器学习推理

Apple 发布了 ml-stable-diffusion\[9 (https://arxiv.org/html/2605.16259#bib.bib9)\]，一个基于 CoreML 的 Stable Diffusion 推理流水线。该框架通过称为 SPLIT_EINSUM_V2 的技术实现神经引擎上的推理，该技术对注意力操作进行拆分执行。然而，SPLIT_EINSUM_V2 是为 M1/M2 代神经引擎设计的，并未充分利用 M3 代 GPU 的高计算性能。

CoreML 是 Apple 的模型推理框架，它将模型从 PyTorch 或 TensorFlow 转换，并使用 Metal GPU 后端上的优化核函数执行推理。通过 coremltools\[9 (https://arxiv.org/html/2605.16259#bib.bib9)\] 进行的模型转换会自动应用计算图的静态优化、算子融合和内存布局优化。相比之下，PyTorch 的 Metal Performance Shaders (MPS) 后端是一个通用实现，模型特定的优化有限。这两个后端之间的性能差距对 Apple Silicon 上的扩散模型推理具有重要意义。

### 2.4 图像转换模型

pix2pix-turbo\[10 (https://arxiv.org/html/2605.16259#bib.bib10)\] 是一个基于 SD-Turbo 的端到端 img2img 转换模型，它在 VAE 编码器和解码器之间引入了跳接，将输入图像的结构信息直接传递到解码器。这些结构跳接在从边缘检测图像到逼真图像的转换中实现了高结构保持。然而，编码器和解码器共享中间特征张量的设计具有副作用，使得模型作为单个子图进行 CoreML 转换变得困难。

### 2.5 基于检索的图像生成

利用大规模内存的基于检索的图像生成\[11 (https://arxiv.org/html/2605.16259#bib.bib11)\] 是一种方法，它从预先计算的图像特征数据库中检索最近邻样本，并将其用作条件来辅助扩散模型生成。FAISS\[12 (https://arxiv.org/html/2605.16259#bib.bib12)\] 是一个高速近似最近邻搜索库，通过 IVF-PQ 索引，即使在十亿向量的数据库上也能实现亚毫秒级搜索。配备 512 GB 统一内存的 M3 Ultra 可以将典型 GPU（24 GB）无法容纳的大规模向量数据库保存在内存中，这为基于检索的方法提供了新的可能性。

## 3 实验设置

本研究使用的实验环境如表 1 (https://arxiv.org/html/2605.16259#S3.T1) 所示。Apple M3 Ultra 是 Apple 于 2024 年发布的旗舰芯片，由两个通过 UltraFusion 互联的 M3 Max 芯片组成。我们实验使用的配置具有 60 核 GPU（从最大 76 核配置中禁用了 16 核）和 512 GB 统一内存。理论 FP16 计算性能约为 22 TFLOPS，大约是 NVIDIA RTX 4090 约 330 TFLOPS 的 1/15。然而，CPU、GPU 和神经引擎通过统一内存共享 800 GB/s 内存带宽的能力是独立 GPU 配置所不具备的优势。

表 1：实验环境

在软件栈方面，我们使用 PyTorch 2.6.0 及 MPS 后端作为基准，并使用 CoreML Tools 9.0 进行模型转换。使用 Hugging Face diffusers 0.36.0 加载和预处理各种扩散模型，使用 OpenCV 进行摄像头输入/输出。

## 4 阶段 1：移植到 MPS 后端

作为第一阶段，我们将最初专为 CUDA 实现的 StreamDiffusion 移植到 Apple Metal Performance Shaders (MPS) 后端。三个主要修改是：(1) 将使用 torch.cuda.Event 的时间测量替换为 time.perf_counter，(2) 将设备规范从 cuda 更改为 mps，以及 (3) 在 MPS 上为了可复现性在 CPU 上初始化随机数生成器，然后传输到 GPU。

在我们的单步推理（SD-Turbo）实验中，StreamDiffusion 的 Stream Batch 机制不适用。这是因为 Stream Batch 旨在并行化多个步骤，而单步推理没有步骤可以批处理。因此，本研究中的优化重点转向加速整个单帧推理流水线。

移植后的基线性能如表 2 (https://arxiv.org/html/2605.16259#S4.T2) 所示。在 MPS 上，SD-Turbo（865.9M 参数）在 512×512 分辨率下的推理需要 95.8 毫秒/帧（10.4 FPS）。即使在降低的分辨率 256×256 下，也只达到了 79.5 毫秒（12.6 FPS），表明速度提升相对于计算量的减少是有限的，简单的分辨率降低是不够的。

表 2：MPS 移植后的基线性能（SD-Turbo，1 步）

## 5 阶段 2：加速技术的全面评估

在此阶段，我们系统评估了在 NVIDIA GPU 上报告有效的多种优化技术在 Apple M3 Ultra 上的效果。预先说明结论：只有 CoreML 转换被证明是有效的；所有其他技术要么无效，要么适得其反。

### 5.1 CoreML 转换

我们将 PyTorch 模型转换为 CoreML 格式（.mlpackage），并在 Metal GPU 后端上执行推理。使用 ct.convert 通过基于trace的转换进行，并指定 compute_units=CPU_AND_GPU。在转换过程中，会自动应用计算图的静态分析以进行算子融合、直接映射到 Metal 核函数以及内存布局优化。

如表 3 (https://arxiv.org/html/2605.16259#S5.T3) 所示，CoreML 转换将 UNet 推理时间从 87.6 毫秒减少到 53.4 毫秒，降低了 39%。这种改进归因于 CoreML 优化的 Metal 核函数（考虑了模型结构）相比 PyTorch 的 MPS 后端的通用 Metal Shader 实现的效果。在本研究中，CoreML 转换被证明是 Apple Silicon 上 UNet 推理唯一有效的加速技术。

表 3：通过 CoreML 转换加速 UNet 推理

### 5.2 量化

我们全面评估了 CoreML 的训后量化选项（INT8 线性、6 位调色板、4 位调色板和 2 位调色板）。在 NVIDIA GPU 上，通过使用 TensorRT INT8 量化和 AWQ/GPTQ 等技术降低内存带宽以实现推理加速已被广泛报道。

然而，如表 4 (https://arxiv.org/html/2605.16259#S5.T4) 所示，在 M3 Ultra 上，任何量化级别都没有观察到推理速度的变化。即使是 2 位调色板（理论上内存带宽为 1/16）也没有产生加速，这强烈表明 M3 Ultra GPU 上的 UNet 推理是计算受限的，而不是内存带宽受限。这可能是因为统一内存的 800 GB/s 带宽为传输 865.9M 参数模型的权重提供了足够的余量。

表 4：CoreML 量化的效果（UNet，512×512）

### 5.3 Token Merging

Token Merging (ToMe)\[13 (https://arxiv.org/html/2605.16259#bib.bib13)\] 是一种根据自注意力中 Key token 之间的相似性合并 token 的技术，从而降低注意力操作的计算成本。在 NVIDIA GPU 上，当注意力操作是瓶颈时，其有效性已有报道。

然而，在 MPS 上，token 相似性计算和合并操作的开销超过了注意力计算的减少，导致大约 10% 的速度下降。这表明 MPS 上的注意力实现与 NVIDIA GPU 上的 xformers 和 Flash Attention 具有不同的瓶颈特征。也就是说，在 MPS 上，注意力操作本身不是主要瓶颈，而 token 操作的额外成本相对较大。

### 5.4 CoreML 并行推理

为了最大化利用 M3 Ultra 的 60 核 GPU，我们尝试使用多个 CoreML 模型实例进行并行推理。在 NVIDIA GPU 上，通过 CUDA Streams 可以并行执行多个核函数，并广泛用于批处理推理和流水线并行化。

然而，使用 CoreML 时，1-4 个实例的并行执行没有观察到吞吐量提升。CoreML 将推理序列化。