PyTorch 2.15 为 Mac 的 MPS 加入 FP8 并加快模型解码

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

太长不看

• 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 解码。

这对你意味着什么

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

• 本地 LLM 快得多。混合专家模型的解码速度比 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 上对它们做乘法。（更多）

目录

• 入门
• FP8 来到 MPS
• 更快的语言模型解码
• 线性代数获得加速
• 卷积、填充与训练
• Softmax 与归约
• 视频解码的内存占用大幅降低
• 通过 MPS streams 获得更多控制
• 更多运算可在 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 对这些张量做乘法，它同时支持 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 加速了主导这一步骤的两类层：稠密线性层和混合专家层。

线性层

自 2.14 起，MPS 就为 torch.mm 提供了 GEMV 内核，但 F.linear 仍要经过 MPSGraph。PyTorch 2.15 将它们连接到 F.linear，因此常见的 [1, 1, hidden_size] 解码输入会直接使用它们，并融合符合条件的偏置加法。在来自真实模型的 96 种 BF16 层形状上，几何平均加速比为 1.23×。最大的收益出现在小型投影上，例如 MoE 模型中的门控，旧路径让大部分内存带宽处于闲置。大型层和 LM head 已经接近带宽受限，因此仍保持在接近 1×。

混合专家层

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

在 2.14 中，grouped 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 吞吐量上。在 配备 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 中也同样更快。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 数学实现现在能正确处理 batch 与 head 的广播。

抢先了解 PyTorch 2.16

2.16 的开发已经全面展开。下面的部分改动已经合入 main，可在 nightly 版本中使用；其余仍在审查中，因此发布前细节可能会有变化。

更快的 LLM 解码

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

注意力完全在 Metal 上运行

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

更多从 MPSGraph 到 Metal 的迁移

脱离 MPSGraph 的迁移仍在继续。Batch norm 的前向与反向正在进行中，而 MSE loss（小输入上最高快 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 为 Mac 带来了 FP8、面向 dense 与 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 做出贡献的人。