PyTorch 2.15 MPS 后端发布：FP8 上 Mac，MoE 解码提速 5–6.5 倍

如果你在 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 解码峰值内存现在从 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 倍：在配备 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 中有报告）。同样的 padding 修改解决了这个问题：一次编码-解码往返现在与 CPU 一致，达到 29.95 dB PSNR，而此前仅为 10.06 dB。
• 在你的 Mac 上使用 FP8 权重。 FP8 张量占用的内存只有 BF16 的一半，而 F.scaled_mm 可以直接在 GPU 上对它们做乘法。（更多）

目录

• 快速上手
• MPS 迎来 FP8
• 更快的语言模型解码
• 线性代数提速
• 卷积、填充与训练
• 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

MPS 迎来 FP8

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)

更快的语言模型解码

每个大语言模型都有预填充（prefill）阶段和解码（decode）阶段。在解码阶段，每个新 token 以单行的形式通过模型，因此大型矩阵乘法变成了矩阵-向量乘积。PyTorch 2.15 加速了在这一步骤中占主导的两类层：稠密线性层和混合专家层。

线性层

MPS 自 2.14 起就拥有针对 torch.mm 的 GEMV 内核，但 F.linear 仍然要经过 MPSGraph。PyTorch 2.15 将它们直接接入 F.linear，因此常见的 [1, 1, hidden_size] 解码输入可以直接使用这些内核，并把符合条件的偏置加法融合进来。在来自真实模型的 96 个 BF16 层形状上，这带来了 1.23 倍的几何平均加速。最大的收益出现在小型投影层，例如 MoE 模型中的门控层——旧路径让大部分内存带宽闲置了。大型层和 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 吞吐量上。在配备 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 的新路径：

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：

通过 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 尚不支持包含多个元素的整数或布尔输出。

Attention 也获得了正确性修复：缩放点积注意力的 MPS 数学实现现在可以正确处理 batch 和 head 的广播。

PyTorch 2.16 抢先看

2.16 的开发工作已经进展顺利。下面的一些改动已经合入 main 并可在 nightly 版本中使用；其他改动仍在评审中，因此在发布前细节可能会有所变化。

更快的 LLM 解码

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

Attention 完全在 Metal 上运行

缩放点积注意力将完全弃用 MPSGraph：每种形状都将在 Metal kernel 上运行，解码使用专用的向量 kernel，在较新的 macOS 版本上半精度 prefill 则使用 Metal Performance Primitives。以前回退到 MPSGraph 的形状受益最大。在 M5 Pro 上，使用分组查询注意力（32 个 query head、8 个 KV head、head dim 80、4K 上下文）的解码从 465 µs 降至 65 µs，head dim 40 的 4K × 4K prefill 则从 42 ms 降至 6 ms。

更多从 MPSGraph 到 Metal 的迁移

摆脱 MPSGraph 的迁移仍在继续。Batch norm 的 前向 和 反向 正在进行中，而 MSE 损失（小输入上最高快 12 倍）、std_mean 和 var_mean（最高快 4 倍）以及 bincount 已经合入。bincount 的效果最惊人：当大多数值落入同一个 bin 时，统计 400 万个整数过去需要 6.8 秒，现在只需半毫秒。

更小的安装体积

打包的 shader 将默认为 Metal 3.1 构建，Metal 4 版本仅为需要的 kernel 编译。这使得 libtorch_cpu.dylib 缩小了约 25%，从 384 MiB 降至 287 MiB，并为具有更快 matmul 和原生 FP8 矩阵乘法的 Metal 4.1 shader 奠定了基础。

结语

PyTorch 2.15 为 Mac 带来了 FP8、更快的稠密和 MoE 模型 LLM 解码、更多 GPU 上的线性代数，以及一长串从 MPSGraph 迁移到专用 Metal kernel 的操作。来试试这个候选版本吧！

我们接下来的计划

• 更多 MPS：2.16 将让 attention 完全在 Metal 上运行，我们会继续把操作从 MPSGraph 上迁移走。
• torch.compile：我们正在努力消除 graph break，让更多模型能够端到端编译。
• 你的想法：我们还应该在 PyTorch 上做什么？如果有东西慢、缺失或损坏，无论是在 MPS 还是其他地方，请在 GitHub 上提交 issue 告诉我们。

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