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

简单介绍一下我们是谁:我们是 HF Mirror 的 PyTorch 团队。我们在 PyTorch 的多个方面都有贡献,其中一个重点领域是 MPS 后端,目标是让你的 Mac 成为运行和训练 PyTorch 模型的一流平台。在 2026 年初,即使在 MPS 上运行小型模型也常常意味着要绕过不支持的操作、性能悬崖和意外的内存增长。因此,我们从 HF Mirror 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 tokens/s 提升到 61.3 tokens/s,这得益于带有专用解码路径的原生分组 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.compileQwen3.6-35B-A3B 从 9.4 tokens/s 提升到 61.3 tokens/s,LFM2.5-8B-A1B 从 23.6 提升到 123.6 tokens/s。稠密模型也受益,线性层更快,采样用的 softmax 快了 8 倍。(更多)
  • 视频生成不再爆内存。用 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 上相乘。(更多)

目录

快速开始

你需要一台 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)

更快的语言模型解码

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

线性层

MPS 自 2.14 起就已经拥有 GEMV 内核,torch.mm但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 的分组:

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×:

Bar chart of MoE decoding throughput on M3 Ultra with torch.compile. Qwen3.6-35B-A3B goes from 38.7 to 51.7 tokens/s at context 512 and 36.8 to 48.7 at context 2048; LFM2.5-8B-A1B goes from 77.1 to 103.2 and 74.9 to 98.7 tokens/s.

线性代数获得提升

复数计算

处理复数张量变得更简单。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 加 matmul 的路径用于短输出。适用于音频和流式模型的形状。
常量填充 迁移到 Metal 内核,修复了分块视频 VAE 解码中驱动内存失控增长的问题(详见下文)。
反射与复制填充 新的 Metal 实现,速度更快并修复了反向传播的正确性。
非原地 dropout nn.Dropout和F.dropout现在使用现有的融合内核。
上采样反向 一个 Metal gather 内核取代了 MPSGraph 的 resize 梯度路径。
Channels-lastconv2d反向 保持在高效的 channels-last 路径上,而不是使用带有昂贵跨步访问的 NCHW 图。

Softmax 与归约

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

Bar chart of FP32 softmax time for a single row of logits at five vocabulary sizes from 128,256 to 201,088. MPSGraph takes 50 to 74 microseconds; the Metal kernel takes 7 to 10 microseconds, a 7.1x to 8.0x speedup.

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

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

Stacked bar chart of peak driver memory during the Wan2.2 tiled VAE decode. With MPSGraph padding, 8.07 GiB of allocator-reserved memory plus 262 GiB outside the allocator gives 270.1 GiB; with Metal padding the total is 8.1 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尚不支持元素数大于 1 的整数或布尔输出。

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

PyTorch 2.16 抢先看

2.16 的工作已经在稳步推进。下面的一些改动已经合入main并可在 nightly 版本中使用;另一些仍在评审中,因此细节在发布前可能发生变化。

更快的 LLM 解码

一旦模型的解码路径不再包含任何 MPSGraph 算子,GPU 只有在 CPU 完成其所有内核的编码后才会开始处理一个 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 上运行

缩放点积注意力完全弃用 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是最引人注目的:当大多数值落入同一个分箱时,统计 400 万个整数过去需要 6.8 秒,现在只需半毫秒。

更小的安装体积

捆绑的着色器将默认针对 Metal 3.1 构建,仅对需要它们的 kernel 编译 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 kernel。试试候选发布版吧!

我们的下一步计划

  • 更多 MPS:2.16 将完全在 Metal 上运行注意力,我们还会继续把操作从 MPSGraph 迁出。
  • torch.compile:我们正在努力消除图断点,让更多模型能够端到端编译。
  • 你的想法:我们接下来还应该在 PyTorch 中做些什么?如果有哪里慢、缺失或损坏,无论是在 MPS 还是其他地方,在 GitHub 上开一个 issue并告诉我们。

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