PyTorch 2.15 MPS 后端升级：支持 FP8 与专用 Metal 内核

返回文章列表

PyTorch 2.15 中的 MPS

Irakli Salia Isalia20 Khushi Agrawal khushi-411

如果你在 Mac 上运行模型，或者将 MPS 作为主要的计算加速工具，PyTorch 2.15 中有许多值得期待的新特性！FP8 正式登陆 MPS，语言模型的解码通过线性层和 softmax 获得了更快的路径，新的 Metal kernels 则解决了音频、视频和训练工作负载中一些出人意料的高开销部分。

简单介绍一下我们是谁：我们是 Hugging Face 的 PyTorch 团队。我们在整个 PyTorch 项目中都有贡献，而我们的重点方向之一是 MPS backend，目标是让你的 Mac 成为运行和训练 PyTorch 模型的一流平台。在 2026 年初，即使在 MPS 上运行小型模型也常常意味着要绕过不支持的操作、性能断崖以及意外的内存增长。因此，我们在 Hugging Face Hub 上跨语言、图像和视频领域运行真实模型，找出 MPS 的不足之处，并在 PyTorch 上游进行修复，越来越多地通过专用 Metal kernels 替换 MPSGraph。本文的大部分内容正是这项工作的成果，由我们与 PyTorch 维护者及社区贡献者共同完成。

TL;DR

• FP8 登陆 Mac： float8_e4m3fn 值的存储空间仅为 BF16 的一半，并支持缩放矩阵乘法（scaled matrix multiplication）以及在 CUDA/CPU 上支持的所有其他操作！
• MoE 模型解码速度比 2.14 版本快 5–6.5 倍： Qwen3.6-35B-A3B 在 M3 Ultra 上的生成速度从 9.4 tokens/s 提升至 61.3 tokens/s，这得益于带有专用解码路径的原生 grouped matmul、更少的 GPU-CPU 同步以及更快的 top-k routing。
• 每个 LLM token 的处理时间减少： 单行线性层使用专用的矩阵-向量 kernels，且针对词汇表的 softmax 速度提升了约 8 倍。
• GPU 上执行更多线性代数运算： 复数 LU 分解和三角求解、更快的小矩阵求逆，以及更快的 QR 分解。
• 音频、视频和训练同样受到关注： 卷积、padding、dropout、插值反向传播和 reductions 均获得了更快的路径。Wan2.2 的视频 VAE 解码峰值内存占用从 270 GiB 降至 8 GiB。
• Python 端提供更多控制： 公开的 MPS streams 和更可靠的事件计时。
• 更新的要求： MPS 现在要求 macOS 15 或更高版本；PyTorch 2.15 要求 Python 3.11 或更高版本。
• 2.16 版本前瞻： attention 完全迁移至 Metal，更多的 MPSGraph-to-Metal 迁移，以及更快的 LLM 解码。

这对您意味着什么

如果您在 Mac 上使用 transformers、diffusers 或 ComfyUI 等工具运行模型，以下是 PyTorch 2.15 在实际应用中的变化：

• 本地 LLM 速度大幅提升。 混合专家（Mixture-of-experts）模型的解码速度比 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。稠密（Dense）模型也从中受益，得益于更快的线性层以及采样时 softmax 速度提升 8 倍。(更多详情)
• 视频生成不再受内存限制。 使用 Wan2.2 的分块 VAE 解码视频时，驱动内存占用曾高达 270 GiB；现在峰值仅为约 8 GiB。(更多详情)
• Qwen-Image-2.1 图像编辑功能修复。 在 2.14 版本中，Qwen-Image-2.1 的 VAE 编码器在 MPS 后端会静默损坏图像，导致图像编辑工作流失效（已在 ComfyUI 中报告）。通过相同的 padding 更改已修复此问题：编码-解码往返测试的 PSNR 值现已与 CPU 匹配，达到 29.95 dB，较之前的 10.06 dB 显著提升。
• Mac 支持 FP8 权重。 FP8 张量的内存占用仅为 BF16 的一半，且 F.scaled_mm 可直接在 GPU 上进行乘法运算。(更多详情)

目录

• 快速上手
• FP8 登陆 MPS
• 更快的语言模型解码
• 线性代数性能提升
• 卷积、填充与训练
• Softmax 与归约操作
• 视频解码内存占用大幅降低
• 通过 MPS Streams 实现更多控制
• 更多算子支持在 MPS 上运行
• PyTorch 2.16 前瞻
• 总结

快速上手

你需要一台搭载 Apple Silicon 的 Mac，并安装 macOS 15 或更高版本以及 Python 3.11 或更高版本（请参阅 2.15 兼容性矩阵）。如果你的操作系统版本过低或没有可用的 Metal GPU，MPS 可用性检查现在会告知你具体原因。

要在独立环境中尝试该发布候选版本：

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, MoE）层。

线性层

自 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× 左右。

混合专家（MoE）层

混合专家层会将不同数量的 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)  

解码内核直接体现在 token 吞吐量上。在启用 torch.compile 的 M3 Ultra 上，MoE 模型的生成速度约为 每秒多 1.3× 的 token：

线性代数性能提升

复数运算

处理复数张量变得更加轻松。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-tile 解码过程中会发起 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 目前尚不支持包含多个元素的整数或布尔输出。

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

PyTorch 2.16 前瞻

针对 2.16 版本的工作已经在紧锣密鼓地进行中。以下列出的一些更改已经合并到 main 分支并可在夜间构建版（nightlies）中使用；另一些仍在审查中，因此在正式发布前细节可能会有所变动。

更快的 LLM 解码

当模型的解码路径中不再包含任何 MPSGraph 算子时，GPU 只有在 CPU 完成该 token 所有内核的编码后才会开始处理。PyTorch 2.16 每 16 个内核提交一次命令缓冲区，因此 GPU 可以在 CPU 继续编码剩余部分的同时开始工作。对于在 2048-token 上下文长度下的编译 BF16 解码：

Attention 完全运行在 Metal 上

缩放点积注意力（Scaled dot-product attention）正在彻底弃用 MPSGraph：所有形状都将运行在 Metal 内核上，其中解码使用专用的向量内核，而在新版 macOS 上的半精度预填充则使用 Metal Performance Primitives。那些以前回退到 MPSGraph 的形状受益最大。在 M5 Pro 上，使用分组查询注意力（grouped-query attention，32 个查询头，8 个 KV 头，头维度 80，4K 上下文）进行解码时，耗时从 465 µs 降至 65 µs；头维度为 40 的 4K × 4K 预填充耗时从 42 ms 降至 6 ms。

更多从 MPSGraph 迁移到 Metal 的工作

摆脱 MPSGraph 的进程仍在继续。Batch norm 的前向传播和反向传播正在进行中，而 MSE loss（在小输入上速度提升高达 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，并为支持更快矩阵乘法（matmul）和原生 FP8 矩阵乘法的 Metal 4.1 着色器奠定了基础。

总结

PyTorch 2.15 为 Mac 带来了 FP8 支持、更快的密集模型和 MoE 模型 LLM 解码、更多的 GPU 线性代数运算，以及大量从 MPSGraph 迁移到专用 Metal 内核的操作。不妨试试发布候选版本！

我们接下来的计划

• 更多 MPS 支持： PyTorch 2.16 将完全在 Metal 上运行注意力机制，并持续将操作从 MPSGraph 迁移出来。
• torch.compile： 我们正在努力消除图断点（graph breaks），以便更多模型能够端到端编译。
• 你的想法： PyTorch 还有哪些地方需要我们改进？如果 MPS 或其他任何部分存在性能慢、功能缺失或错误问题，请在 GitHub 上提交 issue 告诉我们。

衷心感谢 Nikita Shulga、Kurt Mohler 和 Aaron Gokaslan 对大部分工作的审阅，以及所有为本版本 MPS 贡献代码的开发者。