MPS in PyTorch 2.15

PyTorch 2.15 中的 MPS

返回文章列表

[-1

]

[0

发布于 2026 年 10 月 20 日

GitHub 上的更新

点赞

1

[

Irakli Salia

Isalia20

关注

Khushi Agrawal

khushi-411

如果你在 Mac 上运行模型，或使用 MPS 作为主要计算加速工具，PyTorch 2.15 有很多值得期待的改进！FP8 支持现已登陆 MPS，语言模型解码通过线性层和 softmax 的路径显著加速，新的 Metal 内核优化了音频、视频和训练任务中一些出人意料的高消耗部分。

关于我们的简要说明：我们是 Hugging Face 的 PyTorch 团队。我们在 PyTorch 生态中广泛参与，其中一个重要方向是 MPS 后端，目标是让 Mac 成为运行和训练 PyTorch 模型的一流平台。2026 年初，即使在 MPS 上运行小型模型也常常需要绕过不支持的操作、性能断崖和意外的内存增长。因此我们通过 Hugging Face Hub 在语言、图像和视频领域运行真实模型，定位 MPS 的不足之处，并在 PyTorch 中上游修复这些问题，越来越多地通过将 MPSGraph 替换为专用 Metal 内核实现。本文的许多内容正是这些工作与 PyTorch 维护者及社区贡献者共同完成的成果。

TL;DR

• FP8 登陆 Mac：以 half 的空间存储 float8_e4m3fn 值，支持与 CUDA/CPU 相同的缩放矩阵乘法和其他操作！

• MoE 模型解码速度提升 5–6.5 倍：得益于原生分组矩阵乘法、专用解码路径、更少的 GPU-CPU 同步和更快的 top-k 路由，Qwen3.6-35B-A3B 在 M3 Ultra 上从 9.4 提升至 61.3 tokens/s。

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

对你的意义

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

• 本地 LLM 速度大幅提升。专家混合模型解码速度比 2.14 提升 5–6.5 倍：在 M3 Ultra 上使用 torch.compile，Qwen3.6-35B-A3B 从 9.4 提升至 61.3 tokens/s，LFM2.5-8B-A1B 从 23.6 提升至 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 中报告）。相同的填充修复解决了问题：现在编码-解码往返的 PSNR 从 10.06 dB 提升至 29.95 dB。

• Mac 上的 FP8 权重。FP8 张量占用 BF16 的一半内存，F.scaled_mm 可直接在 GPU 上进行乘法运算。（更多）

内容

• 入门指南

/think

• FP8 支持 MPS

• 更快的语言模型解码

• 线性代数性能提升

• 卷积、填充和训练

• Softmax 与归约操作

• 视频解码内存占用大幅降低

• 更精细的 MPS 流控制

• 更多操作支持 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 张量：每个值仅占用 1 字节，存储空间仅为 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)

更快的语言模型解码

每个大语言模型都有预填充阶段和解码阶段。在解码阶段，每个新 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 头已接近带宽极限，因此提升接近 1 倍。

专家混合层

专家混合层会将不同数量的 token 分配给每个专家。分组矩阵乘法可在一个操作中处理这些不均匀分组，2.15 为 MPS 提供了原生实现。

在 2.14 版本中，分组矩阵乘法需将分组偏移量复制到 CPU 并逐个处理分组。新的 Metal 内核直接在 GPU 上读取偏移量，使用可用的 Metal 性能原语或其他 Metal 路径。解码时常见的小分组也获得了专属内核。

例如，64 个 token 可拆分为 16、32 和 16 的分组：

import
torch
import
torch.nn.functional
as
F

 /think

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

解码内核直接影响令牌吞吐量。在配备 torch.compile 的 M3 Ultra 上，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-plus-matmul 路径。适用于音频和流式模型形状。

常量填充

迁移到 Metal 内核，解决了分块视频 VAE 解码中的驱动程序内存增长问题（详见下文）。

反射和复制填充

新的 Metal 实现，速度更快且修复了反向正确性问题。

非原地 dropout

nn.Dropout

和

F.dropout

现在使用现有的融合内核。

上采样反向传播

Metal 收集内核取代 MPSGraph 的调整梯度路径。

通道优先

conv2d

反向传播

保持在高效的通道优先路径上，而不是使用代价高昂的 NCHW 图和跨步访问。

softmax 和归约操作

2.15 版本中令牌采样也变得更加快速。softmax 和 log_softmax 现在在专用的 Metal 内核上运行，大词汇表受益最大：对 [1, 128256] 行的 FP32 softmax 计算时间从 57 微秒降至 7 微秒，比 MPSGraph 版本快约 8 倍。这一加速效果在其他常见词汇量规模上同样存在：

视频解码的内存占用大幅减少

本版本中最具影响力的成果之一来自于常量填充（constant-padding）的改进。在 Wan-AI/Wan2.2-TI2V-5B-Diffusers 的 VAE 解码过程中，完整 28-tile 解码需要进行 8,764 次填充操作。使用 Metal 进行填充后，驱动程序内存峰值从约 270 GiB 降至 8 GiB：

内存指标

MPSGraph 填充

Metal 填充

当前张量分配峰值

4.422 GiB

4.095 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 版本的工作已全面展开。以下部分更改已合并到主分支并在 nightly 版本中可用，其余更改仍在审核中，发布前细节可能发生变化。

更快的 LLM 解码

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

模型

之前

之后

提升

LFM2.5-8B-A1B

108.8 token/秒

126.0 token/秒

+16%

Qwen3.6-35B-A3B

52.9 token/秒

62.8 token/秒

+19%

Qwen3.8-27B

12.1 token/秒

12.6 token/秒

+4%

注意力完全在 Metal 上运行

缩放点积注意力已完全移除 MPSGraph：所有形状都将使用 Metal 内核，解码使用专用向量内核，新 macOS 版本的半精度预填充使用 Metal Performance Primitives。原本需要回退到 MPSGraph 的形状受益最大。在 M5 Pro 上，使用分组查询注意力（32 个查询头，8 个 KV 头，头维度 80，4K 上下文）的解码时间从 465 微秒降至 65 微秒，4K×4K 预填充（头维度 40）时间从 42 毫秒降至 6 毫秒。

更多 MPSGraph 到 Metal 的迁移

继续减少对 MPSGraph 的依赖。批量归一化前向和反向传播正在推进中，而 MSE 损失（小输入最多快 12 倍）、std_mean 和 var_mean（最多快 4 倍）以及 bincount 已完成迁移。bincount 的改进最为显著：当大多数值落入单个 bin 时，统计 400 万个整数的时间从 6.8 秒降至 0.5 毫秒。

更小的安装包

捆绑的着色器默认会为 Metal 3.1 构建，仅对需要 Metal 4 的内核进行编译。这使 libtorch_cpu.dylib 的体积缩小了约 25%，从 384 MiB 减少到 287 MiB，并为 Metal 4.1 着色器奠定基础，后者具备更快的矩阵乘法和原生 FP8 矩阵乘法功能。

总结

PyTorch 2.15 为 Mac 带来 FP8 支持，提升密集型和 MoE 模型的 LLM 解码速度，在 GPU 上提供更多线性代数运算，并将 MPSGraph 中的大量操作迁移至专用 Metal 内核。请尝试体验这个候选版本！

我们的下一步计划

• 更多 MPS 功能：2.16 版本将完全在 Metal 上运行注意力机制，我们将继续将操作从 MPSGraph 迁移出去。

• torch.compile：我们正在解决图中断问题，使更多模型能够端到端编译。

• 您的建议：PyTorch 还需要在哪些方面改进？如果在 MPS 或其他任何地方遇到性能瓶颈、功能缺失或缺陷，请在 GitHub 上提交问题告诉我们。

衷心感谢 Nikita Shulga、Kurt Mohler 和 Aaron Gokaslan 对大量工作的审阅，以及所有为本次发布贡献 MPS 功能的开发者。

更多博客文章

tokenizers

tokenization

rust

tokenizers v1：编码、解码与扩展性，性能实测

91

2026年9月21日

transformers

pytorch

optimization

在连续批处理中解锁异步性

68

2026年5月14日

社区

编辑

预览

通过拖拽、粘贴或点击此处上传图片、音频和视频

.

轻点或粘贴此处上传图片

评论

· 注册或登录以发表评论