MPS in PyTorch 2.15

Source: https://huggingface.co/blog/mps-pytorch-2-15

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

https://huggingface.co/blog/mps-pytorch-2-15#tldrTL;DR

• FP8 登陆 Mac：float8_e4m3fn值占用 BF16 一半的存储空间，并支持缩放矩阵乘法及 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 解码峰值内存从 270 GiB 降至 8 GiB。
• **更多 Python 层面的控制：**公开 MPS 流以及更可靠的事件计时。
• 更新的系统要求：MPS 现在要求macOS 15 及以上版本；PyTorch 2.15 要求Python 3.11 及以上版本。
• **2.16 预览：**注意力机制完全迁移到 Metal、更多 MPSGraph 向 Metal 的迁移、更快的 LLM 解码。

https://huggingface.co/blog/mps-pytorch-2-15#what-this-means-for-you这对您意味着什么

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

• **本地 LLM 快得多。**混合专家（MoE）模型的解码速度比 2.14 快 5–6.5 倍：在 M3 Ultra 上配合torch.compile，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。（更多内容 (https://huggingface.co/blog/mps-pytorch-2-15#faster-language-model-decoding)）
• **视频生成不再爆内存。**此前使用 Wan2.2 的瓦片化 VAE 解码视频会将驱动内存推到 270 GiB；现在峰值仅约 8 GiB。（更多内容 (https://huggingface.co/blog/mps-pytorch-2-15#a-much-smaller-memory-footprint-for-video-decoding)）
• **Qwen-Image-2.1 图像编辑恢复正常。**在 2.14 上，Qwen-Image-2.1 的 VAE 编码器在 MPS 上会静默损坏图像，破坏图像编辑流程（在 ComfyUI 中报告 (https://github.com/Comfy-Org/ComfyUI/issues/16433)）。同一个填充修复 (https://github.com/pytorch/pytorch/pull/195368) 解决了这个问题：现在编码-解码往返结果与 CPU 的 PSNR 达到 29.95 dB，此前仅为 10.06 dB。
• **Mac 上的 FP8 权重。**FP8 张量占用 BF16 一半的内存，F.scaled_mm可直接在 GPU 上进行矩阵乘法。（更多内容 (https://huggingface.co/blog/mps-pytorch-2-15#fp8-comes-to-mps)）

https://huggingface.co/blog/mps-pytorch-2-15#contents目录

• 快速开始 (https://huggingface.co/blog/mps-pytorch-2-15#getting-started)
• FP8 登陆 MPS (https://huggingface.co/blog/mps-pytorch-2-15#fp8-comes-to-mps)
• 更快的语言模型解码 (https://huggingface.co/blog/mps-pytorch-2-15#faster-language-model-decoding)
• 线性代数性能提升 (https://huggingface.co/blog/mps-pytorch-2-15#linear-algebra-gets-a-boost)
• 卷积、填充与训练 (https://huggingface.co/blog/mps-pytorch-2-15#convolutions-padding-and-training)
• Softmax 与归约 (https://huggingface.co/blog/mps-pytorch-2-15#softmax-and-reductions)
• 视频解码的内存占用大幅降低 (https://huggingface.co/blog/mps-pytorch-2-15#a-much-smaller-memory-footprint-for-video-decoding)
• 通过 MPS 流获得更多控制 (https://huggingface.co/blog/mps-pytorch-2-15#more-control-with-mps-streams)
• 更多算子在 MPS 上可用 (https://huggingface.co/blog/mps-pytorch-2-15#more-operations-work-on-mps)
• PyTorch 2.16 预览 (https://huggingface.co/blog/mps-pytorch-2-15#a-sneak-peek-at-pytorch-216)
• 总结 (https://huggingface.co/blog/mps-pytorch-2-15#wrapping-up)

https://huggingface.co/blog/mps-pytorch-2-15#getting-started快速开始

你需要一台 Apple 芯片的 Mac，运行macOS 15 及以上版本，并安装Python 3.11 及以上版本（参见 2.15 兼容性矩阵 (https://github.com/pytorch/pytorch/blob/release/2.15/RELEASE.md#release-compatibility-matrix)）。如果你的操作系统版本过低或没有可用的 Metal GPU，MPS 可用性检查现在会告诉你原因。

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

uv pip install torch==2.15.0 --index-url https://download.pytorch.org/whl/test/cpu

https://huggingface.co/blog/mps-pytorch-2-15#fp8-comes-to-mpsFP8 登陆 MPS

FP8 为你在 Mac 上节省内存提供了另一种方式。PyTorch 2.15 在 MPS 上添加了torch.float8_e4m3fn张量 (https://github.com/pytorch/pytorch/pull/194074)：每个值占用 1 字节，仅为 FP16 或 BF16 存储量的一半。

你可以使用F.scaled_mm (https://github.com/pytorch/pytorch/pull/196916) 对这些张量进行矩阵乘法，它同时支持张量级（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) ``

https://huggingface.co/blog/mps-pytorch-2-15#faster-language-model-decoding更快的语言模型解码

每个 LLM 都有预填充（prefill）阶段和解码（decode）阶段。解码时，每个新 token 作为单行数据通过模型，大矩阵乘法变成了矩阵-向量乘。PyTorch 2.15 加速了主导这一步骤的两类层：稠密线性层和混合专家层。

https://huggingface.co/blog/mps-pytorch-2-15#linear-layers线性层

MPS 从 2.14 起就有用于torch.mm的 GEMV 内核，但F.linear仍然走 MPSGraph 路径。PyTorch 2.15 将它们接入了F.linear (https://github.com/pytorch/pytorch/pull/198721)，因此常见的\[1, 1, hidden_size\]解码输入直接使用这些内核，并将符合条件的偏置加法融合进来。在来自真实模型的 96 种 BF16 层形状上，这实现了1.23 倍几何平均加速。最大的收益出现在小型投影上，例如 MoE 模型中的门控层，旧路径在这些层上浪费了大部分内存带宽。大型层和 LM 头本身已接近带宽瓶颈，因此加速比维持在接近 1 倍。

https://huggingface.co/blog/mps-pytorch-2-15#mixture-of-experts-layers混合专家（Mixture-of-Experts）层

混合专家层会将不同数量的 token 发送给各个专家。分组矩阵乘法可以在一次操作中处理这些不均匀的分组，2.15 为 MPS 提供了原生实现 (https://github.com/pytorch/pytorch/pull/194933)。

在 2.14 中，分组矩阵乘法将分组偏移量复制到 CPU 并逐组循环。新的 Metal 内核直接在 GPU 上读取偏移量，可用时使用 Metal Performance Primitives，其他情况走独立的 Metal 路径。小分组（解码时的常见场景）也拥有专用内核 (https://github.com/pytorch/pytorch/pull/198725)。

例如，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 吞吐量上。在M3 Ultra 上配合torch.compile，MoE 模型生成速度提升约1.3 倍：

M3 Ultra 上配合 torch.compile 的 MoE 解码吞吐量柱状图。在上下文长度 512 时，Qwen3.6-35B-A3B 从 38.7 提升到 51.7 tokens/s，上下文长度 2048 时从 36.8 提升到 48.7；LFM2.5-8B-A1B 分别从 77.1 提升到 103.2 和从 74.9 提升到 98.7 tokens/s。(https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/mps-pytorch-2-15/moe-decode-throughput.png)

https://huggingface.co/blog/mps-pytorch-2-15#linear-algebra-gets-a-boost线性代数性能提升

https://huggingface.co/blog/mps-pytorch-2-15#complex-valued-computations复数运算

复数张量的操作变得更加容易。complex64 LU 分解 (https://github.com/pytorch/pytorch/pull/188473)、LU 求解 (https://github.com/pytorch/pytorch/pull/195369) 和三角求解 (https://github.com/pytorch/pytorch/pull/195861) 现在可以在 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()) ``

https://huggingface.co/blog/mps-pytorch-2-15#small-matrices-and-larger-solves小矩阵与大规模求解

小矩阵同样获得了快速内核。对于相机几何和图像变换场景，最大4 × 4的 FP32 矩阵批会在单个 Metal 内核中完成求逆 (https://github.com/pytorch/pytorch/pull/195944)，该内核将分解与求解合并在一起。更大规模的问题也得到了关注：QR 分解使用分块更新 (https://github.com/pytorch/pytorch/pull/198732)，三角求解 (https://github.com/pytorch/pytorch/pull/196241) 将大部分工作转化为高效的矩阵乘法。

Gated DeltaNet 工作负载获得专用快速路径 (https://github.com/pytorch/pytorch/pull/196900)，适用于16 × 16、32 × 32 和 64 × 64 FP32 三角系统的批量求解（至少 8 个右端项）。

SVD 在数值极端情况下也更加健壮。现在它能正确处理非常小的奇异值 (https://github.com/pytorch/pytorch/pull/200920)、非常大或非常小的输入幅度 (https://github.com/pytorch/pytorch/pull/201037)，以及使用full_matrices=True的基完备化 (https://github.com/pytorch/pytorch/pull/201198)，从而使得伪逆和最小二乘结果更加可靠。

https://huggingface.co/blog/mps-pytorch-2-15#convolutions-padding-and-training卷积、填充与训练

音频模型、视频解码器和训练循环各有自己的开销算子。本版本为其中多个算子提供了 Metal 新路径：

https://huggingface.co/blog/mps-pytorch-2-15#softmax-and-reductionsSoftmax 与归约

Token 采样在 2.15 中也变得更快。softmax (https://github.com/pytorch/pytorch/pull/196071) 和log_softmax (https://github.com/pytorch/pytorch/pull/195441) 现在运行在专用 Metal 内核上，大词表受益最明显：对\[1, 128256\]行 logits 的 FP32 softmax 从57 μs 降至 7 μs，比 MPSGraph 版本快约8 倍。这一加速在其他常见词表大小下同样成立：

FP32 softmax 单行 logits 在五个词表大小（从 128,256 到 201,088）下的耗时柱状图。MPSGraph 耗时 50 到 74 微秒；Metal 内核耗时 7 到 10 微秒，加速比 7.1 倍到 8.0 倍。(https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/mps-pytorch-2-15/softmax-sampling.png)

https://huggingface.co/blog/mps-pytorch-2-15#a-much-smaller-memory-footprint-for-video-decoding视频解码的内存占用大幅降低

本次发布中最引人注目的成果之一来自常数填充的改动。Wan-AI/Wan2.2-TI2V-5B-Diffusers中的 VAE 解码在完整的 28 瓦片解码过程中进行了8,764 次填充调用。使用 Metal 填充后，峰值驱动内存从约270 GiB 降至 8 GiB：

Wan2.2 瓦片化 VAE 解码期间峰值驱动内存的堆叠柱状图。使用 MPSGraph 填充时，分配器预留内存 8.07 GiB 加上分配器之外的 262 GiB，总计 270.1 GiB；使用 Metal 填充时总计 8.1 GiB。(https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/mps-pytorch-2-15/wan-vae-memory.png)

https://huggingface.co/blog/mps-pytorch-2-15#more-control-with-mps-streams通过 MPS 流获得更多控制

你现在可以直接从 Python 中选择由哪个 MPS 流来运行你的工作。torch.mps.Stream和torch.mps.stream上下文管理器 (https://github.com/pytorch/pytorch/pull/191415) 提供了接口，Tensor.record_stream (https://github.com/pytorch/pytorch/pull/193464) 让分配器可以跟踪张量在其他流上的使用情况。

下面是一个将工作发送到独立流并在读取结果前进行同步的简单示例：

`` 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) ``

https://huggingface.co/blog/mps-pytorch-2-15#more-operations-work-on-mps更多算子在 MPS 上可用

更多算子现在可以在 MPS 上运行：torch.frexp (https://github.com/pytorch/pytorch/pull/194110)、torch.logspace (https://github.com/pytorch/pytorch/pull/190495)、torch.take和Tensor.put