PyTorch 2.15 为 MPS 带来 FP8 与 MoE 解码提速

TL;DR

FP8 登陆 Mac：以 BF16 一半的空间存储 float8_e4m3fn 值，并使用缩放矩阵乘法以及 CUDA/CPU 上支持的所有其他运算！

MoE 模型解码速度比 2.14 快 5–6.5 倍：得益于原生分组矩阵乘法及专用解码路径、更少的 GPU-CPU 同步和更快的 top-k 路由，Qwen3.6-35B-A3B 在 M3 Ultra 上从 9.4 tokens/s 提升到 61.3 tokens/s。

每个 LLM 的每 token 耗时更少：单行线性层使用专用矩阵-向量内核，词表上的 softmax 约快 8 倍。

GPU 上更多线性代数运算：复数 LU 与三角求解、更快的小矩阵求逆，以及更快的 QR 分解。

音频、视频和训练也得到关注：卷积、填充、dropout、插值反向和归约的更快路径。Wan2.2 的视频 VAE 解码现在峰值内存为 8 GiB，而非 270 GiB。

从 Python 获得更多控制：公开的 MPS 流和更可靠的事件计时。

更新后的要求：MPS 现在要求 macOS 15 或更高版本；PyTorch 2.15 要求 Python 3.11 或更高版本。

2.16 前瞻：注意力完全在 Metal 上运行、更多 MPSGraph 到 Metal 的迁移，以及更快的 LLM 解码。

这对你意味着什么

如果你在 Mac 上使用 transformers、diffusers 或 ComfyUI 等工具运行模型，以下是 2.15 在实际使用中的变化：

本地 LLM 快得多。混合专家模型解码速度比 2.14 快 5–6.5 倍：在配备 torch.compile 的 M3 Ultra 上，Qwen3.6-35B-A3B 从 9.4 tokens/s 提升到 61.3 tokens/s，LFM2.5-8B-A1B 从 23.6 tokens/s 提升到 123.6 tokens/s。稠密模型也从中受益，得益于更快的线性层和用于采样的快 8 倍的 softmax。（更多）

视频生成不再超出内存。使用 Wan2.2 的分块 VAE 解码视频过去会将驱动内存推高到 270 GiB；现在峰值约为 8 GiB。（更多）

使用 Qwen-Image-2.1 正确进行图像编辑。在 2.14 上，Qwen-Image-2.1 的 VAE 编码器在 MPS 上会静默损坏图像，破坏图像编辑工作流（已在 ComfyUI 中报告）。同样的填充改动修复了该问题：编码-解码往返现在与 CPU 匹配，PSNR 为 29.95 dB，高于此前的 10.06 dB。

在你的 Mac 上使用 FP8 权重。FP8 张量占用的内存是 BF16 的一半，F.scaled_mm 直接在 GPU 上对它们进行乘法运算。（更多）

更快的语言模型解码

线性代数得到提升

卷积、填充与训练

Softmax 与归约

视频解码的内存占用大幅降低

通过 MPS 流获得更多控制

你需要一台搭载 macOS 15 或更高版本的 Apple 芯片 Mac，以及 Python 3.11 或更高版本（参见 2.15 兼容性矩阵）。如果你的操作系统版本过低或没有可用的 Metal GPU，MPS 可用性检查现在会告诉你原因。

要在单独的环境中试用候选版本：

FP8 登陆 MPS

FP8 为你在 Mac 上节省内存提供了另一种方式。PyTorch 2.15 在 MPS 上新增了 torch.float8_e4m3fn 张量：每个值一个字节，存储空间为 FP16 或 BF16 的一半。

你可以使用 F.scaled_mm 对这些张量进行乘法运算，它同时支持逐张量和逐行缩放。以下是一个每个张量一个缩放因子的小示例：

更快的语言模型解码

每个 LLM 都有预填充阶段和解码阶段。在解码期间，每个新 token 以单行形式通过模型，因此大型矩阵乘法转变为矩阵-向量乘积。PyTorch 2.15 加速了主导此步骤的两类层：稠密线性层和混合专家层。

线性层

自 2.14 起，MPS 已为 torch.mm 提供 GEMV 内核，但 F.linear 仍通过 MPSGraph 执行。PyTorch 2.15 将它们连接到 F.linear，因此常见的 [1, 1, hidden_size] 解码输入会直接使用这些内核，符合条件的偏置加法也会被融合进去。在来自真实模型的 96 个 BF16 层形状上，几何平均加速为 1.23 倍。最大的收益出现在小型投影上，例如 MoE 模型中的门控，旧路径在此处让大部分内存带宽闲置。大型层和 LM 头此前已接近带宽瓶颈，因此基本保持在 1 倍左右。

混合专家层

混合专家层会向每个专家发送不同数量的 token。分组矩阵乘法在单个操作中处理这些不均匀的分组，2.15 为 MPS 提供了原生实现。

在 2.14 中，分组矩阵乘法会将分组偏移量复制到 CPU，并循环遍历各分组。新的 Metal 内核直接在 GPU 上读取偏移量，在可用时使用 Metal Performance Primitives，在其他情况下使用单独的 Metal 路径。小分组（解码期间的常见情况）也有自己的内核。

例如，64 个 token 可以拆分为 16、32 和 16 的分组：

解码内核直接体现在 token 吞吐量上。在配备 torch.compile 的 M3 Ultra 上，MoE 模型每秒生成的 token 数约多 1.3 倍：

线性代数得到提升

复数计算

处理复数张量变得更加容易。complex64 LU 分解、LU 求解和三角求解现在可在 MPS 上运行，这也解锁了建立在其上的运算和梯度，例如复数行列式。以下是一个直接在 GPU 上求解的小型方程组：

小矩阵与更大规模的求解

小矩阵同样能获得快速内核。对于相机几何和图像变换，最多 4 × 4 的 FP32 矩阵批次会在单个 Metal 内核中求逆，该内核将分解与求解合并在一起。更大的问题同样受到关注：QR 使用分块更新，而三角求解将其大部分工作转化为高效的矩阵乘法。

Gated DeltaNet 工作负载获得了一条专用快速路径，适用于至少八个右端项的 16 × 16、32 × 32 和 64 × 64 FP32 三角系统批次。

SVD 在数值极端情况下也更加稳健。它现在能正确处理极小的奇异值、极大或极小的输入量级，以及 full_matrices=True 时的基补全，这反过来使伪逆和最小二乘结果更加可靠。

卷积、填充与训练

音频模型、视频解码器和训练循环都有各自开销高昂的运算。本次发布为其中若干运算提供了通过 Metal 的新路径：

Softmax 与归约

Token 采样在 2.15 中也变得更快。softmax 和 log_softmax 现在运行在专用 Metal 内核上，大词表受益最大：对一行 [1, 128256] 的 logits 执行 FP32 softmax，耗时从 57 µs 降至 7 µs，比 MPSGraph 版本快约 8 倍。这一加速在其他常见词表大小上同样成立：

视频解码的内存占用大幅降低

本次发布中最引人注目的结果之一来自常量填充的改动。Wan-AI/Wan2.2-TI2V-5B-Diffusers 中的 VAE 解码在完整的 28 块解码过程中发起了 8,764 次填充调用。在 Metal 上执行填充后，峰值驱动内存从约 270 GiB 降至 8 GiB：

借助 MPS 流获得更多控制

现在你可以直接从 Python 中选择由哪个 MPS 流来运行你的工作。torch.mps.Stream 和 torch.mps.stream 上下文管理器提供了这一接口，而 Tensor.record_stream 则让分配器能够跟踪张量在另一个流上的使用情况。

下面是一个小示例，它将工作发送到单独的流，并在读取结果前进行同步：

更多运算可在 MPS 上运行

现在有更多运算可以在 MPS 上运行：torch.frexp、torch.logspace、torch.take 和 Tensor.put_、torch.heaviside、torch.special.laguerre_polynomial_l 以及无符号整数比较。logspace 目前尚不支持元素数大于一的整数或布尔输出。

注意力机制也获得了一项正确性修复：缩放点积注意力的 MPS 数学实现现在能够正确处理批次和头的广播。

PyTorch 2.16 抢先看

2.16 的工作已经全面展开。下面的一些改动已经合入 main 并可在 nightly 版本中使用；另一些仍在审查中，因此细节可能会在正式发布前发生变化。

更快的 LLM 解码

一旦模型的解码路径不再包含任何 MPSGraph 运算，GPU 只有在 CPU 完成其所有内核的编码后才会开始处理一个 token。PyTorch 2.16 每 16 个内核提交一次命令缓冲区，因此 GPU 可以在 CPU 仍在编码其余部分时就开始工作。对于 2048 token 上下文下的编译式 BF16 解码：

注意力完全在 Metal 上运行

缩放点积注意力正在彻底弃用 MPSGraph：所有形状都将在 Metal 内核上运行，解码使用专用向量内核，在较新的 macOS 版本上，半精度预填充则使用 Metal Performance Primitives。过去会回退到 MPSGraph 的那些形状受益最大。在 M5 Pro 上，使用分组查询注意力进行解码（32 个查询头、8 个 KV 头、头维度 80、4K 上下文）从 465 µs 降至 65 µs，而头维度为 40 的 4K × 4K 预填充则从 42 ms 降至 6 ms。

更多从 MPSGraph 到 Metal 的迁移

离开 MPSGraph 的进程仍在继续。批归一化的前向和反向正在进行中，而 MSE 损失（在小输入上最高快 12 倍）、std_mean 和 var_mean（最高快 4 倍）以及 bincount 已经落地。bincount 最为显著：当大多数值落入同一个 bin 时，统计 400 万个整数过去需要 6.8 秒，现在只需半毫秒。

更小的安装体积

捆绑的着色器将默认针对 Metal 3.1 构建，仅为需要的内核编译 Metal 4 版本。这使 libtorch_cpu.dylib 缩小约 25%，从 384 MiB 降至 287 MiB，并为采用更快矩阵乘法和原生 FP8 矩阵乘法的 Metal 4.1 着色器奠定基础。

总结

PyTorch 2.15 将 FP8 带到 Mac，为稠密和 MoE 模型带来更快的 LLM 解码，在 GPU 上提供更多线性代数运算，并将一长串运算从 MPSGraph 迁移到专用 Metal 内核。快来试试候选发布版吧！

我们的下一步

更多 MPS：2.16 将完全在 Metal 上运行注意力，我们也会继续将运算从 MPSGraph 迁移出去。

torch.compile：我们正在努力消除图中断，让更多模型能够端到端编译。

你的想法：我们还能在 PyTorch 中做些什么？如果有哪里慢、缺失或损坏，无论是在 MPS 还是其他地方，请在 GitHub 上开一个 issue 告诉我们。

非常感谢 Nikita Shulga、Kurt Mohler 和 Aaron Gokaslan 审阅了本项工作的大部分内容，也感谢所有在本版本中为 MPS 做出贡献的人。

博客更多文章

tokenizers v1：编码、解码与扩展，实测

解锁连续批处理中的异步性