PyTorch 2.15 为 MPS 后端带来 FP8 与更快的 LLM 解码
MPS in PyTorch 2.15
Hugging Face 的 PyTorch 团队介绍 PyTorch 2.15 在 MPS 后端的更新,新增 torch.float8_e4m3fn 张量与 F.scaled_mm,FP8 张量占用为 BF16 的一半。
原文给出 MPS 后端在 FP8、MoE 解码和显存占用上的具体数字,可据此判断 Mac 本地跑模型的实际变化。
译文尚不完整,完整内容请切换到原文。
If you run models on your Mac or use MPS as your primary acceleration tool for computing, there's plenty to look forward to in PyTorch 2.15! FP8 arrives on MPS, language-model decoding gets faster paths through linear layers and softmax and new Metal kernels tackle some surprisingly expensive parts of audio, video and training workloads.
A quick word about who we are: we're the PyTorch team at Hugging Face. We contribute across PyTorch, and one of our focus areas is the MPS backend, where the goal is to make your Mac a first-class place to run and train PyTorch models. At the start of 2026, running even small models on MPS often meant working around unsupported operations, performance cliffs and unexpected memory growth. So we run real models from the Hugging Face Hub across language, image and video, find where MPS falls short and fix it upstream in PyTorch, more and more often by replacing MPSGraph with dedicated Metal kernels. Much of this post is the result of that work, done together with the PyTorch maintainers and community contributors.
TL;DR
- FP8 comes to the Mac: store
float8_e4m3fnvalues in half the space of BF16 and use scaled matrix multiplication and all other ops which are supported on CUDA/CPU! - MoE models decode 5–6.5× faster than in 2.14: Qwen3.6-35B-A3B goes from 9.4 to 61.3 tokens/s on an M3 Ultra, thanks to a native grouped matmul with a dedicated decoding path, fewer GPU-CPU syncs and faster top-k routing.
- Less time per token for every LLM: single-row linear layers use specialized matrix-vector kernels, and softmax over the vocabulary is about 8× faster.
- More linear algebra on the GPU: complex LU and triangular solves, faster small-matrix inverses, and faster QR decomposition.
- Audio, video, and training get attention too: faster paths for convolutions, padding, dropout, interpolation backward and reductions. Wan2.2's video VAE decode now peaks at 8 GiB of memory instead of 270 GiB.
- More control from Python: public MPS streams and more reliable event timing.
- Updated requirements: MPS now requires macOS 15 or later; PyTorch 2.15 requires Python 3.11 or later.
- A peek at 2.16: attention fully on Metal, more MPSGraph-to-Metal migrations and faster LLM decoding.
What this means for you
If you run models on a Mac with transformers, diffusers or tools like ComfyUI, here's what 2.15 changes in practice:
- Much faster local LLMs. Mixture-of-experts models decode 5–6.5× faster than in 2.14: on an M3 Ultra with
torch.compile, Qwen3.6-35B-A3B goes from 9.4 to 61.3 tokens/s and LFM2.5-8B-A1B from 23.6 to 123.6 tokens/s. Dense models benefit too, from faster linear layers and an 8× faster softmax for sampling. (More) - Video generation that fits in memory. Decoding a video with Wan2.2's tiled VAE used to push driver memory to 270 GiB; it now peaks at about 8 GiB. (More)
- Correct image editing with Qwen-Image-2.1. On 2.14, the Qwen-Image-2.1 VAE encoder silently corrupted images on MPS, breaking image-edit workflows (reported in ComfyUI). The same padding change fixes it: an encode-decode round trip now matches CPU at 29.95 dB PSNR, up from 10.06 dB.
- 您Mac上的FP8权重。 FP8张量占用BF16一半的内存,
F.scaled_mm直接在GPU上乘以它们。(更多)
内容
- 入门
- FP8来到MPS
- 更快的语言模型解码
- 线性代数获得提升
- 卷积、填充和训练
- Softmax和归约
- 视频解码的内存占用更小
- 使用MPS流获得更多控制权
- 更多操作可在MPS上运行
- PyTorch 2.16的预览
- 总结
入门
您需要带有macOS 15或更高版本和Python 3.11或更高版本的Apple Silicon Mac(请参阅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都有一个预填充阶段和解码阶段。在解码期间,每个新标记都作为单行通过模型,因此大的矩阵乘法变成了矩阵向量乘法。PyTorch 2.15加快了主导此步骤的两种类型的层:密集线性层和专家混合层。
线性层
MPS自2.14以来就有了用于torch.mm的GEMV内核,但F.linear仍然通过MPSGraph。PyTorch 2.15将它们连接到F.linear,因此常见的[1, 1, hidden_size]解码输入直接使用它们,符合条件的偏差添加被融合。在来自真实模型的96种BF16层形状中,这是1.23倍的几何平均加速。最大的收益落在小投影上,例如MoE模型中的门,其中旧路径留下了大部分的内存带宽未使用。大层和LM头已经接近带宽限制,因此它们仍然接近1×。
专家混合层
一个专家混合层将不同数量的标记发送给每个专家。分组矩阵乘法一次性处理这些不均匀的组,而2.15为MPS提供了原生实现。
在2.14中,分组矩阵乘法将组偏移复制到CPU,并循环遍历组。新的Metal内核直接在GPU上读取偏移量,在可用的情况下使用Metal Performance Primitives,并在其他地方使用单独的Metal路径。小组,解码时的常见情况,也获得了自己的内核。
例如,64个标记可以分成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])
解码内核直接显示在标记吞吐量中。在M3 Ultra上的torch.compile,MoE模型每秒生成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收集内核替换了MPSGraph的resize梯度路径。 |
通道最后conv2d反向 |
保持在高效的通道最后路径上,而不是使用昂贵的步进访问的NCHW图。 |
Softmax和归约
2.15 版本中,Token 采样速度也得到提升。 softmax 和 log_softmax 现在在专用的 Metal 内核上运行,大型词汇表从中受益最大:一个 FP32 softmax 在 [1, 128256] 行的 logits 上从 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 目前还不支持具有多个元素的整数或布尔输出。
注意力机制也获得了正确性修复:MPS 数学实现的缩放点积注意力现在 正确处理批处理和头广播。
PyTorch 2.16 的前瞻
2.16 的工作已经进展顺利。下面的一些更改已经 landed 在 main 上,并在夜更版本中可用;其他更改仍在审核中,因此在发布之前详细信息可能会有所变化。
更快的 LLM 解码
一旦模型的解码路径不再包含任何 MPSGraph 操作,GPU 只会在 CPU 完成对其所有内核的编码后才开始处理一个标记。PyTorch 2.16 每 16 个内核提交一次命令缓冲区,因此 GPU 可以在 CPU 仍在编码其余内容时开始工作。对于编译的 BF16 解码在 2048 个标记的上下文中:
| 模型 | 之前 | 之后 | 变化 |
|---|---|---|---|
| LFM2.5-8B-A1B | 108.8 标记/秒 | 126.0 标记/秒 | +16% |
| Qwen3.6-35B-A3B | 52.9 标记/秒 | 62.8 标记/秒 | +19% |
| Qwen3.8-27B | 12.1 标记/秒 | 12.6 标记/秒 | +4% |
注意力完全在 Metal 上运行
缩放点积注意力完全 放弃了 MPSGraph:每个形状都将运行在 Metal 内核上,解码将使用专用的向量内核,较新的 macOS 版本将使用 Metal Performance Primitives 进行半精度预填充。以前回退到 MPSGraph 的形状受益最大。在 M5 Pro 上,使用分组查询注意力(32 个查询头,8 个 KV 头,头维度 80,4K 上下文)的解码从 465 µs 降低到 65 µs,而 4K × 4K 的预填充(头维度 40)从 42 ms 降低到 6 ms。
更多 MPSGraph 到 Metal 的迁移
远离 MPSGraph 的迁移仍在继续。批归一化 前向 和 后向 正在进行中,而 MSE 损失(在小型输入上最多快 12 倍)、std_mean 和 var_mean(最多快 4 倍)和 bincount 已经 landed。bincount 最具戏剧性:当大多数值落入单个 bin 时,计数 400 万个整数以前需要 6.8 秒,现在只需要半毫秒。
更小的安装包
捆绑的着色器将 默认为 Metal 3.1 构建,Metal 4 版本只编译需要它们的内核。这使 libtorch_cpu.dylib 减少了大约 25%,从 384 MiB 减少到 287 MiB,并为具有更快 matmuls 和原生 FP8 矩阵乘法的 Metal 4.1 着色器奠定了基础。
总结
PyTorch 2.15 将 FP8 带到了 Mac,为密集和 MoE 模型提供了更快的 LLM 解码,在 GPU 上进行了更多线性代数,以及大量操作从 MPSGraph 迁移到专用的 Metal 内核。试试这个候选版本!
我们接下来要做什么
- 更多 MPS: 2.16 将完全在 Metal 上运行注意力,我们将继续将操作从 MPSGraph 移除。
torch.compile: 我们正在努力消除图断点,以便更多模型可以端到端编译。- 您的想法: 我们应该在 PyTorch 上做些什么?如果有什么慢、缺失或损坏的,在 MPS 或其他任何地方,在 GitHub 上打开一个问题 并告诉我们。
感谢 Nikita Shulga、Kurt Mohler 和 Aaron Gokaslan 审阅了大部分工作,以及感谢所有为本次 MPS 发布做出贡献的人。
来源: Hugging Face Blog · huggingface.co


