PyTorch 2.15 发布 MPS 后端更新：Mac 支持 FP8 与模型解码加速
MPS in PyTorch 2.15
PyTorch 2.15 为 Mac 的 MPS 后端带来 FP8 支持与专用 Metal 内核，MoE 模型解码较 2.14 提速 5–6.5 倍，Wan2.2 视频 VAE 解码显存从 270 GiB 降至 8 GiB。softmax 与线性代数运算也显著加快，并新增 Python 层 MPS 流控制。
PyTorch 2.15 在 Mac 的 MPS 后端加入 FP8 与 Metal 内核，使 MoE 模型解码提速 5 至 6.5 倍，视频生成显存由 270 GiB 降至 8 GiB。
如果你在 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 维护者和社区贡献者合作完成的成果。
摘要
- 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 提升到 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 速度大幅提升。 混合专家（Mixture-of-experts）模型的解码速度比 2.14 快 5–6.5 倍：在配备
torch.compile的 M3 Ultra 上，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 匹配，PSNR 为 29.95 dB，从 10.06 dB 提升。
- 在你的 Mac 上使用 FP8 权重。 FP8 张量占用的内存是 BF16 的一半，并且
F.scaled_mm可以直接在 GPU 上对它们进行乘法运算。（更多）
目录
- 入门指南
- FP8 登陆 MPS
- 更快的语言模型解码
- 线性代数性能提升
- 卷积、填充与训练
- Softmax 与归约运算
- 视频解码的内存占用大幅降低
- 通过 MPS 流获得更多控制
- MPS 支持更多操作
- PyTorch 2.16 抢先看
- 总结
入门指南
你需要一台配备 macOS 15 或更高版本 和 Python 3.11 或更高版本 的 Apple 芯片 Mac（参见 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 对这些张量进行乘法运算，它支持逐张量（tensorwise）和逐行（rowwise）缩放。下面是一个每个张量使用一个缩放因子的小示例：
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)
更快的语言模型解码
每个 LLM 都有预填充（prefill）阶段和解码（decode）阶段。在解码期间，每个新 token 作为单行通过模型，因此大型矩阵乘法变成了矩阵-向量乘积。PyTorch 2.15 加速了主导此步骤的两种层类型：稠密线性层和混合专家（mixture-of-experts）层。
线性层
MPS 自 2.14 版本起就拥有针对 torch.mm 的 GEMV 内核，但 F.linear 仍然通过 MPSGraph 处理。PyTorch 2.15 将它们连接到 F.linear，因此常见的 [1, 1, hidden_size] 解码输入直接使用它们，并将符合条件的偏置加法融合在内。在来自真实模型的 96 种 BF16 层形状中，这带来了 1.23 倍的几何平均加速。最大的提升出现在小型投影上，例如 MoE 模型中的门控（gates），旧路径在那里未充分利用大部分内存带宽。大型层和 LM 头已经接近带宽瓶颈，因此它们保持在 1 倍左右。
混合专家层
混合专家层将不同数量的 token 发送给每个专家。分组矩阵乘法（Grouped matmul）在单次操作中处理这些不均匀的组，而 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 吞吐量上。在配备 M3 Ultra 和 torch.compile 的设备上，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 加矩阵乘法路径。适用于音频和流式模型形状。 |
| 常量填充 | 迁移到 Metal 内核，修复了分块视频 VAE 解码中驱动程序内存失控增长的问题（详见下文）。 |
| 反射和复制填充 | 新的 Metal 实现，速度更快并修复了反向传播的正确性。 |
| 非原地 dropout | nn.Dropout 和 F.dropout 现在使用现有的融合内核。 |
| 上采样反向传播 | Metal 收集内核取代了 MPSGraph 的 resize 梯度路径。 |
通道在后 conv2d 反向传播 |
保持在高效的通道在后路径上，而不是使用具有昂贵跨步访问的 NCHW 图。 |
Softmax 与归约
2.15 版本中，Token 采样也变快了。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：
| 内存指标 | MPSGraph 填充 | Metal 填充 |
|---|---|---|
| 当前张量分配峰值 | 4.422 GiB | 4.095 GiB |
| 分配器保留峰值 | 8.070 GiB | 8.070 GiB |
| 驱动程序内存总峰值 | 270.135 GiB | 8.108 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 尚不支持包含多个元素的整数或布尔输出。
注意力机制也获得了正确性修复：缩放点积注意力的 MPS 数学实现现在可以正确处理批次和头部的广播。
PyTorch 2.16 抢先看
2.16 的开发工作已全面展开。以下部分更改已合并到 main 并在每日构建版本中可用；其他更改仍在审查中，因此发布前细节可能会发生变化。
更快的 LLM 解码
一旦模型的解码路径不再包含任何 MPSGraph 操作，GPU 只有在 CPU 完成所有内核编码后才会开始处理一个 token。PyTorch 2.16 每 16 个内核提交一次命令缓冲区，因此当 CPU 仍在编码其余部分时，GPU 就可以开始工作。对于 2048 token 上下文下的编译 BF16 解码：
| 模型 | 之前 | 之后 | 变化 |
|---|---|---|---|
| LFM2.5-8B-A1B | 108.8 tokens/s | 126.0 tokens/s | +16% |
| Qwen3.6-35B-A3B | 52.9 tokens/s | 62.8 tokens/s | +19% |
| Qwen3.8-27B | 12.1 tokens/s | 12.6 tokens/s | +4% |
注意力机制完全在 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 的迁移仍在继续。批量归一化（Batch norm）的前向和反向操作正在进行中，而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 为 Mac 带来了 FP8 支持、针对稠密模型和 MoE 模型更快的 LLM 解码、更多在 GPU 上运行的线性代数操作，以及一长串从 MPSGraph 迁移到专用 Metal 内核的操作。请试用候选版本！
我们的下一步计划
- 更多 MPS： 2.16 将完全在 Metal 上运行注意力机制，我们将继续将操作从 MPSGraph 迁移出去。
torch.compile： 我们正在努力消除图中断（graph breaks），以便更多模型可以端到端编译。- 您的想法： 我们还应该在 PyTorch 中处理什么？如果某些功能在 MPS 或其他地方运行缓慢、缺失或损坏，请在 GitHub 上提交 issue 并告诉我们。
非常感谢 Nikita Shulga、Kurt Mohler 和 Aaron Gokaslan 审阅了这项工作的很大一部分，以及感谢所有为本次 MPS 版本做出贡献的人。
来源：Hugging Face Blog · huggingface.co