如果你在 Mac 上运行模型,或使用 MPS 作为主要的计算加速工具,那么 PyTorch 2.15 有很多值得期待的内容!FP8 登陆 MPS,语言模型解码在 linear 层和 softmax 中获得了更快的路径,新的 Metal 内核则解决了音频、视频和训练工作负载中一些开销出乎意料地高的部分。
简要介绍一下我们是谁:我们是 Hugging Face 的 PyTorch 团队。我们全面参与 PyTorch 的贡献,其中一个重点领域是 MPS 后端,目标是让您的 Mac 成为一等公民的场所,用于运行和训练 PyTorch 模型。在2026年初,在 MPS 上运行哪怕很小的模型,也常常意味着要绕过不支持的操作、性能断崖和意外的内存增长。因此,我们将来自 Hugging Face Hub 的语言、图像和视频领域的真实模型实际运行起来,找出 MPS 的不足之处,并在 PyTorch 上游进行修复——越来越多的情况下,是用专用的 Metal 内核替换 MPSGraph。本文的大部分内容就是这项工作的成果,它是与 PyTorch 维护者和社区贡献者共同完成的。
简而言之
- FP8 来到 Mac: 商店
float8_e4m3fn以 BF16 一半的空间表示数值,并使用缩放矩阵乘法以及 CUDA/CPU 上支持的所有其他操作! - MoE 模型的解码速度比 2.14 快 5–6.5 倍: Qwen3.6-35B-A3B 在 M3 Ultra 上的速度从 9.4 提升到 61.3 tokens/s,这得益于原生的分组矩阵乘法(grouped matmul)配有一条专用的解码路径、更少的 GPU-CPU 同步以及更快的 top-k 路由。
- 每个 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 倍:在配备 M3 Ultra 的
torch.compile, 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 中报告)。 same padding change 这一修复解决了该问题:现在一次编码-解码往返在 29.95 dB PSNR 下与 CPU 结果一致,而此前仅为 10.06 dB。
- 在你的 Mac 上使用 FP8 权重。 FP8 张量占用的内存是 BF16 的一半,并且
F.scaled_mm可以直接在 GPU 上对它们进行乘法运算。(更多)
目录
- 入门指南
- FP8 登陆 MPS
- 更快的语言模型解码
- 线性代数获得提升
- 卷积、填充与训练
- Softmax 与归约操作
- 视频解码的内存占用大幅缩小
- 通过 MPS 流实现更多控制
- 更多操作支持 MPS
- PyTorch 2.16抢先看
- 总结
入门指南
你需要一台搭载Apple芯片的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 添加了 torch.float8_e4m3fn MPS 上的张量:每个值一个字节,存储量是 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 都有一个预填充阶段和一个解码阶段。在解码期间,每个新 token 以单行的形式通过模型,因此大矩阵乘法变成了矩阵-向量乘积。PyTorch 2.15 加速了主导这一步骤的两类层:稠密线性层和混合专家层。
线性层
MPS 自 2.14 起就已为 torch.mm 自 2.14 起,但 F.linear 仍然要经过 MPSGraph。PyTorch 2.15 将它们直接连接到 F.linear,因此常见的 [1, 1, hidden_size] 解码输入会直接使用它们,并将符合条件的偏置加法融合进来。在来自真实模型的 96 个 BF16 层形状上,这带来了 1.23× 的几何平均加速。最大的收益落在小型投影上,例如 MoE 模型中的门控,旧路径在这些地方浪费了大部分内存带宽。大型层和 LM 头本已接近带宽瓶颈,因此它们保持在接近 1× 的水平。
混合专家层
混合专家层会向每个专家发送不同数量的 token。分组 matmul 可以在单个操作中处理这些不均匀的分组,2.15 为 MPS 提供了 原生实现.
在 2.14 中,分组 matmul 会将分组偏移量复制到 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())
小型矩阵与更大规模的求解
小型矩阵也拥有快速内核。对于相机几何和图像变换,最多可达 32 个的 FP32 矩阵批次 4 × 4 的 FP32 矩阵批次可以 在单个 Metal 内核中完成求逆 ,该内核将分解和求解结合在一起。更大规模的问题也得到了关注: QR 使用分块更新,而 三角求解 则将其大部分工作转化为高效的矩阵乘法。
Gated DeltaNet 工作负载获得了一条 专用快速路径 适用于以下规模的批次: 16 × 16、32 × 32 和 64 × 64 的 FP32 三角系统 且至少有八个右侧项。
SVD 在数值极端情况下也更加稳健。它现在能正确处理 非常小的奇异值, 非常大或非常小的输入幅值 以及 基补全与 full_matrices=True,这反过来使伪逆和最小二乘结果更加可靠。
卷积、填充与训练
音频模型、视频解码器和训练循环都有各自开销高昂的操作。本次发布为其中多项操作提供了经由 Metal 的新路径:
| 操作 | 2.15 中的变化 |
|---|---|
conv1d forward |
专用内核取代了 MPSGraph,并针对短输出增加了 im2col 加 matmul 的路径。适用于音频和流式模型的形状。 |
| 常数填充 | 迁移到 Metal 内核,修复了分块视频 VAE 解码中驱动程序内存失控增长的问题(下文详述)。 |
| 反射填充和复制填充 | 全新的 Metal 实现,速度更快并修复了反向传播的正确性问题。 |
| 非就地 dropout | nn.Dropout 和 F.dropout 现在使用现有的融合内核。 |
| 上采样反向传播 | Metal gather 内核取代了 MPSGraph 的缩放梯度路径。 |
Channels-last conv2d 反向传播 |
保留在高效的 channels-last 路径上,而不是使用代价高昂的跨步访问的 NCHW 图。 |
Softmax 与归约
在 2.15 中,token 采样也变得更快。 softmax 和 log_softmax 现在运行在专用的 Metal 内核上,大型词表受益最大:对一个 FP32 softmax 而言, [1, 128256] 一行 logits 从 57 µs 降至 7 µs,大约 快 8 倍 于 MPSGraph 版本。这一加速在其他常见词表大小下同样成立:
视频解码的内存占用大幅缩小
本次发布中最引人注目的结果之一来自常数填充(constant-padding)的改动。 Wan-AI/Wan2.2-TI2V-5B-Diffusers 中的 VAE 解码在完整的 28 瓦片解码过程中会发起 8,764 次填充调用 超过一次完整的 28 瓦片解码。在 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 尚不支持包含多个元素的整数或布尔输出。
注意力机制也得到了一个正确性修复:scaled dot-product attention 的 MPS math 实现现在 能够正确处理 batch 和 head 的广播.
PyTorch 2.16 抢先看
2.16 的工作已经在顺利进行中。下面的一些更改已经合入 main 并可在 nightly 版本中使用;其他更改仍在审查中,因此在发布前细节可能会有所变化。
更快的 LLM 解码
一旦模型的解码路径不再包含任何 MPSGraph 操作,GPU 只有在 CPU 完成对一个 token 的所有内核编码之后才会开始处理该 token。PyTorch 2.16 每 16 个内核提交一次命令缓冲区,因此 GPU 可以在 CPU 仍在编码其余部分时开始工作。对于 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 上运行。
缩放点积注意力(Scaled dot-product attention)是 完全放弃 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 是最引人注目的:当大多数数值落入单个区间时,对 4 百万个整数进行计数过去需要 6.8 秒,现在只需半毫秒。
更小的安装体积
捆绑的着色器将 默认针对 Metal 3.1 构建,Metal 4 版本仅针对需要它们的内核编译。这将 libtorch_cpu.dylib 缩减了约 25%,从 384 MiB 降至 287 MiB,并为 Metal 4.1 着色器奠定了基础,后者将带来更快的矩阵乘法和原生 FP8 矩阵乘法。
总结
PyTorch 2.15 为 Mac 带来了 FP8,为稠密模型和 MoE 模型带来更快的 LLM 解码,增加了更多 GPU 上的线性代数功能,并将一长串操作从 MPSGraph 迁移到专用的 Metal 内核。快来试试这个候选版本吧!
我们的下一步计划
- 更多 MPS: 2.16 将把注意力机制完全放到 Metal 上运行,我们会继续将操作从 MPSGraph 迁出。
torch.compile: 我们正在努力消除 graph break,以便更多模型能够端到端编译。- 你的想法: 在 PyTorch 中我们还应该做些什么?如果有什么地方在 MPS 或其他任何地方运行缓慢、缺失或出错, 请在 GitHub 上提交一个 issue 并告诉我们。
特别感谢 Nikita Shulga, Kurt Mohler 以及 Aaron Gokaslan 审阅了这项工作的大部分内容,也感谢本版本中为 MPS 做出贡献的所有人。