PyTorch 2.15 为 Mac MPS 后端带来 FP8 与更快的 LLM 解码

如果你在 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 维护者和社区贡献者共同完成。

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，这得益于带有专用解码路径的原生分组矩阵乘法、更少的 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 倍：在配备 torch.compile 的 M3 Ultra 上，Qwen3.6-35B-A3B 从 9.4 提升到 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 中报告）。同样的填充改动修复了它：编码-解码往返现在与 CPU 匹配，PSNR 为 29.95 dB，高于之前的 10.06 dB。
• 在你的 Mac 上使用 FP8 权重。 FP8 张量占用的内存是 BF16 的一半，而 F.scaled_mm 直接在 GPU 上对它们进行乘法运算。（更多）

目录

• 快速开始
• FP8 登陆 MPS
• 更快的语言模型解码
• 线性代数得到提升
• 卷积、填充和训练
• Softmax 与归约
• 视频解码的内存占用大幅降低
• 通过 MPS stream 获得更多控制
• 更多操作可在 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 在 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 起就为 torch.mm 提供了 GEMV kernel，但 F.linear 仍然要经过 MPSGraph。PyTorch 2.15 将它们连接到 F.linear，因此常见的 [1, 1, hidden_size] 解码输入会直接使用它们，并融合符合条件的偏置加法。在来自真实模型的 96 种 BF16 层形状上，这带来了 1.23 倍的几何平均加速。收益最大的是小型投影，例如 MoE 模型中的门控，旧的路径让大部分内存带宽处于闲置状态。大型层和 LM head 本就已接近带宽瓶颈，因此它们基本保持在 1 倍左右。

混合专家层

混合专家层会向每个专家发送不同数量的 token。分组 matmul 在单次操作中处理这些不均匀的分组，而 2.15 为 MPS 提供了原生实现。

在 2.14 中，分组 matmul 会将分组偏移量复制到 CPU，并遍历各个分组。新的 Metal kernel 直接在 GPU 上读取偏移量，在可用时使用 Metal Performance Primitives，其他地方则使用单独的 Metal 路径。小组（解码期间的常见情况）也有自己的 kernel。

例如，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])

解码 kernel 直接体现在 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 中也变得更快。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 尚不支持元素数大于一的整数或布尔输出。

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

PyTorch 2.16 抢先看

2.16 的工作已经全面展开。下面的一些改动已经合入 main 并可在 nightly 版本中使用；另一些仍在审查中，因此细节可能会在正式发布前发生变化。

更快的 LLM 解码

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

注意力完全在 Metal 上运行

缩放点积注意力正在彻底弃用 MPSGraph：所有形状都将在 Metal kernel 上运行，解码使用专用向量 kernel，在较新的 macOS 版本上，半精度 prefill 使用 Metal Performance Primitives。过去回退到 MPSGraph 的那些形状受益最大。在 M5 Pro 上，使用分组查询注意力（32 个查询头、8 个 KV 头、head 维度 80、4K 上下文）的解码从 465 µs 降到 65 µs，而 head 维度为 40 的 4K × 4K prefill 从 42 ms 降到 6 ms。

更多从 MPSGraph 到 Metal 的迁移

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

更小的安装包

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

总结

PyTorch 2.15 将 FP8 带到 Mac，为稠密和 MoE 模型带来更快的 LLM 解码，在 GPU 上支持更多线性代数运算，并把一长串操作从 MPSGraph 迁移到专用 Metal kernel。来试试候选发布版吧！

我们的下一步计划

• 更多 MPS：2.16 将让注意力完全在 Metal 上运行，我们会继续将操作迁出 MPSGraph。
• torch.compile：我们正在努力消除 graph break，以便更多模型能端到端编译。
• 你的想法：我们还应该在 PyTorch 中做些什么？如果有东西在 MPS 或任何其他地方运行缓慢、缺失或损坏，在 GitHub 上提一个 issue 告诉我们。

非常感谢 Nikita Shulga、Kurt Mohler 和 Aaron Gokaslan 审阅了这项工作的大部分内容，也感谢本次发布中为 MPS 做出贡献的每一个人。