PyTorch 2.15 为 Mac MPS 带来 FP8 支持与更快的模型解码
MPS in PyTorch 2.15
PyTorch 2.15 为 Mac 上的 MPS 后端加入 FP8 支持，并优化 LLM 解码、线性代数及音视频和训练相关操作。M3 Ultra 上，Qwen3.6-35B-A3B 的解码速度从 2.14 的 9.4 tokens/s 提升至 61.3 tokens/s；Wan2.2 视频 VAE 解码的峰值驱动内存则从约 270 GiB 降至约 8 GiB。
文章将 MPS 2.15 的性能提升与内存变化落到具体模型和工作负载，便于判断 Mac 本地运行 PyTorch 模型时的实际改进。
如果你在 Mac 上运行模型，或使用 MPS 作为主要计算加速工具，那么 PyTorch 2.15 有不少值得期待的更新！MPS 现已支持 FP8；语言模型解码在线性层和 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 倍：得益于原生分组矩阵乘法及专用解码路径、更少的 GPU-CPU 同步和更快的 top-k 路由，Qwen3.6-35B-A3B 在 M3 Ultra 上的速度从 9.4 tokens/s 提升至 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 运行速度大幅提升。MoE 模型的解码速度比 2.14 快 5–6.5 倍：在配备
torch.compile的 M3 Ultra 上，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 中报告）。同一项填充改动修复了这个问题：编码-解码往返结果现在与 CPU 上的结果一致，PSNR 达到 29.95 dB，高于原先的 10.06 dB。
- Mac 上的 FP8 权重。 FP8 张量占用的内存只有 BF16 的一半，而
F.scaled_mm可直接在 GPU 上对其执行乘法。（了解更多）
目录
- 入门
- MPS 支持 FP8
- 更快的语言模型解码
- 线性代数加速
- 卷积、填充与训练
- 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
MPS 支持 FP8
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 加速了主导这一阶段的两类层：稠密线性层和混合专家层。
线性层
自 2.14 起，MPS 就为 torch.mm 提供了 GEMV 内核，但 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 吞吐量上。在搭载 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 加矩阵乘法路径。适用于音频和流式模型的形状。 |
| 常量填充 | 改用 Metal 内核，修复分块视频 VAE 解码时驱动程序内存不断增长的问题（详见下文）。 |
| 反射填充与复制填充 | 新增 Metal 实现，速度更快，并修复反向计算的正确性问题。 |
| 非原地 dropout | nn.Dropout 和 F.dropout 现在使用现有的融合内核。 |
| 上采样反向计算 | Metal gather 内核取代了 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目前还不支持包含多个元素的整数或布尔型输出。
Attention 也修复了一项正确性问题：MPS 数学实现中的缩放点积注意力现在可以正确处理批次和注意力头的广播。
抢先一窥 PyTorch 2.16
2.16 版的开发已在如火如荼地进行。下面部分改动已合入 main，并已在 nightly 版本中提供；其他改动仍在审核中，因此在正式发布前细节可能会有变化。
更快的 LLM 解码
模型的解码路径一旦不再包含任何 MPSGraph 操作，GPU 就只能等 CPU 完成所有内核的编码后，才开始处理一个 token。PyTorch 2.16 会每编码 16 个内核就提交一次命令缓冲区，这样 CPU 还在编码其余内核时，GPU 就能开始工作。对于编译后的 BF16 解码，在 2048-token 上下文下：
| 模型 | 之前 | 之后 | 变化 |
|---|---|---|---|
| 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 的前向和反向操作正在开发中；而均方误差损失（小输入下快达 12 倍）、std_mean 和 var_mean（快达 4 倍）以及bincount已经合入。bincount的提升最为显著：当大多数值都落在同一个分箱中时，统计 400 万个整数过去需要 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：我们正努力消除计算图断点，让更多模型能够端到端编译。- 你的想法：你希望我们在 PyTorch 中继续开发哪些功能？如果 MPS 或其他任何地方存在速度慢、缺失或出错的问题，请在 GitHub 上提交 issue告诉我们。
衷心感谢 Nikita Shulga、Kurt Mohler 和 Aaron Gokaslan 审阅了本项工作的很大一部分，也感谢所有为本次发布中的 MPS 作出贡献的人。
来源：Hugging Face Blog · huggingface.co