PyTorch 2.15 中的 MPS 更新

如果你在 Mac 上运行模型，或者把 MPS 作为主要的计算加速工具，那么 PyTorch 2.15 有很多值得期待的内容！FP8 正式登陆 MPS，语言模型解码通过线性层和 softmax 获得更快路径，全新的 Metal 内核则解决了音频、视频和训练工作负载中一些出人意料的昂贵部分。

简单介绍一下我们：我们是 Hugging Face 的 PyTorch 团队。我们在整个 PyTorch 中贡献代码，其中一个重点领域是 MPS 后端，目标是让你的 Mac 成为运行和训练 PyTorch 模型的一流平台。到 2026 年初，即使在 MPS 上运行小型模型，也常常需要绕过不支持的操作、性能悬崖和意外的内存增长。因此我们从 Hugging Face Hub 上选取语言、图像和视频领域的真实模型，找出 MPS 的不足，并在上游的 PyTorch 中修复，越来越多地通过用专用 Metal 内核替换 MPSGraph 来实现。这篇文章的大部分内容正是这些工作的成果，由我们与 PyTorch 维护者和社区贡献者共同完成。

TL;DR

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

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

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

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

音频、视频和训练同样得到关注：卷积、填充、dropout、插值反向传播和归约都有更快路径。Wan2.2 的视频 VAE 解码峰值内存从 270 GiB 降至 8 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 倍：在 M3 Ultra 上配合 torch.compile，Qwen3.6-35B-A3B 从 9.4 提升到 61.3 tokens/s，LFM2.5-8B-A1B 从 23.6 提升到 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 匹配，达到 29.95 dB PSNR，从 10.06 dB 提升。

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

目录

入门指南

FP8 登陆 MPS

更快的语言模型解码

线性代数获得提升

卷积、填充与训练

Softmax 与归约

视频解码的内存占用大幅减小

通过 MPS 流获得更多控制

更多操作可在 MPS 上运行

PyTorch 2.16 抢先看

总结

入门指南

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

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

uv pip install torch==2.15.0 --index-url https://download.pytorch.org/whl/test/cpu

FP8 登陆 MPS

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

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

import torch
import torch.nn.functional as F

def quantize_fp8(x):
 limit = torch.finfo(torch.float8_e4m3fn).max
 scale = x.abs().amax().clamp_min(1e-12) / limit
 return (x / scale).to(torch.float8_e4m3fn), scale

x = torch.randn(1, 1024, device="mps", dtype=torch.float32)
weight = torch.randn(1024, 1024, device="mps", dtype=torch.float32)
x_fp8, x_scale = quantize_fp8(x)
w_fp8, w_scale = quantize_fp8(weight)

y = F.scaled_mm(
 x_fp8,
 w_fp8.t(),
 scale_a=x_scale,
 scale_recipe_a=F.ScalingType.TensorWise,
 scale_b=w_scale,
 scale_recipe_b=F.ScalingType.TensorWise,
 output_dtype=torch.bfloat16,
)
print(y.shape, y.dtype)

更快的语言模型解码

每个大语言模型都有预填充阶段和解码阶段。在解码阶段，每个新 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 模型中的门控，旧路径会留下大部分内存带宽未使用。大型层和语言模型头已经接近带宽受限，因此仍接近 1×。

混合专家层

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

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

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

import torch
import torch.nn.functional as F

tokens = torch.randn(64, 64, device="mps", dtype=torch.bfloat16)
weights = torch.randn(3, 128, 64, device="mps", dtype=torch.bfloat16)
ends = torch.tensor([16, 48, 64], device="mps", dtype=torch.int32)

out = F.grouped_mm(tokens, weights.transpose(-2, -1), offs=ends)
print(out.shape) # torch.Size([64, 128])

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

线性代数获得提升

复数值计算

处理复数张量变得更容易。complex64 的 LU 分解、LU 求解和三角求解现在可在 MPS 上运行，这也解锁了基于它们的操作和梯度，例如复数行列式。下面是一个直接在 GPU 上求解的小系统：

import torch

a = torch.randn(8, 8, device="mps", dtype=torch.complex64)
a = a + 8 * torch.eye(8, device="mps", dtype=torch.complex64)
b = torch.randn(8, 2, device="mps", dtype=torch.complex64)

lu, pivots = torch.linalg.lu_factor(a) x = torch.linalg.lu_solve(lu, pivots, b) print((a @ x - b).abs().max().item())

小矩阵与更大求解

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

Gated DeltaNet 工作负载为至少包含八个右端项的 16 × 16、32 × 32 和 64 × 64 FP32 三角系统批次提供了专用快速路径。

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

卷积、填充与训练

音频模型、视频解码器和训练循环有各自昂贵的操作。本次发布为其中多个操作提供了通过 Metal 的新路径：

操作 2.15 中的变化

conv1d 前向 专用内核替代 MPSGraph，并为短输出增加了 im2col-plus-matmul 路径。适用于音频和流式模型形状。

常数填充 迁移到 Metal 内核，修复了分块视频 VAE 解码中驱动器内存失控增长的问题（详见下文）。

反射与复制填充 新的 Metal 实现，更快并修复了反向正确性。

非原地 dropout nn.Dropout 和 F.dropout 现在使用现有的融合内核。

上采样反向 Metal gather 内核替代 MPSGraph 的 resize-gradient 路径。

Channels-last conv2d 反向 保持高效的 channels-last 路径，而不是带有代价高昂跨步访问的 NCHW 图。

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-tile 解码过程中进行了 8,764 次填充调用。使用 Metal 上的填充后，峰值驱动内存从约 270 GiB 降至 8 GiB：

通过 MPS 流获得更多控制

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

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

import torch

x = torch.randn(1024, device="mps")
torch.mps.synchronize()
worker = torch.mps.Stream()

with torch.mps.stream(worker):
 y = x.square()
 x.record_stream(worker)

del x worker.synchronize() print(y.cpu().shape)

更多操作可在 MPS 上运行

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

注意力机制也得到了正确性修复：scaled dot-product attention 的 MPS 数学实现现在正确处理了批量和头广播。

PyTorch 2.16 预览

2.16 的工作已在顺利进行中。以下部分更改已合入 main 并在 nightly 版本中可用；其他仍在审查中，因此发布前细节可能会有变化。

更快的 LLM 解码

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

注意力完全在 Metal 上运行

Scaled dot-product attention 将完全放弃 MPSGraph：每种形状都将在 Metal 内核上运行，解码使用专用向量内核，在较新的 macOS 版本上半精度 prefill 使用 Metal Performance Primitives。以前回退到 MPSGraph 的形状受益最大。在 M5 Pro 上，使用分组查询注意力（32 个查询头、8 个 KV 头、头维度 80、4K 上下文）的解码从 465 µs 降至 65 µs，而头维度 40 的 4K × 4K prefill 从 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，并为具有更快 matmul 和原生 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 做出贡献的人。