PyTorch 2.15发布：MPS后端支持FP8，MoE解码提速超6倍

PyTorch 2.15 发布：MPS 后端迎来 FP8 支持与显著性能提升

对于在 Mac 上使用 PyTorch 进行模型推理或训练的用户而言，PyTorch 2.15 版本带来了多项关键改进。此次更新重点优化了 MPS（Metal Performance Shaders）后端，引入了对 FP8 数据格式的支持，大幅提升了语言模型解码速度，并针对音频、视频及训练工作负载中的高开销操作开发了新的 Metal 内核。

Irakli Salia (Isalia20)

Khushi Agrawal (khushi-411)

核心亮点概览

• FP8 登陆 Mac：支持将 float8_e4m3fn 值存储在 BF16 空间的一半中，并提供缩放矩阵乘法及其他 CUDA/CPU 已支持的操作。
• MoE 模型解码提速：相比 2.14 版本，混合专家模型（MoE）的解码速度提升 5–6.5 倍。例如，Qwen3.6-35B-A3B 在 M3 Ultra 芯片上的生成速度从 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 解码速度。

实际应用场景影响

对于使用 Transformers、扩散模型或 ComfyUI 等工具在 Mac 上运行模型的用户，2.15 版本带来以下具体变化：

• 本地大语言模型性能飞跃：启用 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。稠密模型也因线性层加速和采样时 Softmax 提速 8 倍而受益。
• 视频生成内存占用大幅降低：在使用 Wan2.2 进行瓦片式 VAE 解码视频时，此前驱动内存峰值高达 270 GiB，现在已降至约 8 GiB。
• 图像编辑流程修复：在 2.14 版本中，Qwen-Image-2.1 VAE 编码器在 MPS 上存在静默错误，导致图像编辑工作流（如 ComfyUI 中报告的问题）失效。通过改进填充实现，编码-解码往返的 PSNR 指标从 10.06 dB 恢复至与 CPU 一致的 29.95 dB。
• FP8 权重支持：用户可在 Mac 上使用 FP8 权重，并通过 F.scaled_mm 直接在 GPU 上进行乘法运算。

安装与环境准备

要体验这些新功能，你需要一台配备 Apple Silicon 芯片且运行 macOS 15 或更高版本的 Mac，同时安装 Python 3.11 或更高版本。若操作系统版本过低或缺少金属 GPU，MPS 可用性检查将明确提示原因。

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

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

FP8 支持详解

PyTorch 2.15 在 MPS 后端添加了 torch.float8_e4m3fn 张量类型，每个值仅占一个字节，存储量仅为 FP16 或 BF16 的一半。用户可以通过 F.scaled_mm 对这些张量进行乘法运算，该函数支持张量级缩放和行级缩放。

以下示例展示了如何使用张量级缩放进行 FP8 矩阵乘法：

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 执行。2.15 版本将两者连接，使常见的 [1, 1, hidden_size] 解码输入能直接使用这些内核，并融合符合条件的偏置加法。在来自真实模型的 96 种 BF16 层形状测试中，这一改动带来了 1.23× 的几何平均加速比。小型投影层（如 MoE 模型中的门控机制）提升最为显著，因为旧路径未能充分利用内存带宽；而大型层和 LM head 原本已接近带宽瓶颈，加速比维持在 1× 左右。

混合专家层（MoE）优化

MoE 层向不同专家发送数量不均等的 token。2.15 版本为 MPS 提供了原生的分组矩阵乘法实现，能够在单个操作中处理这些不均等的组。

在 2.14 版本中，分组矩阵乘法需要将分组偏移复制到 CPU 并在群组上循环。新的 Metal 内核直接在 GPU 上读取偏移，利用 Metal Performance Primitives 及其他专用 Metal 路径。解码过程中常见的小组也获得了专门的内核优化。

示例代码如下：

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.compile 的 M3 Ultra 上，MoE 模型的每秒 token 产生率约为 1.3 倍。

线性代数功能增强

复数计算支持

现在可以在 MPS 上运行复数 64LU 因子化、LU 解和三角解，这也解锁了基于这些操作的梯度计算，例如复数行列式。以下代码演示了在 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 分解使用了封闭更新，三角解决方案将大部分工作转化为高效的矩阵乘法。

Closed DeltaNet 工作负载中，至少包含八个右侧项的 16 × 16、32 × 32 和 64 × 64 FP32 三角系统批次获得了专用快速路径。此外，SVD 在数值极端情况下更加稳健，能正确处理极小的奇异值、极大或极小的输入大小以及 full_matrices=True 的基础完成，从而提高了伪逆和最小二乘结果的可靠性。

卷积、填充及训练相关操作优化

音频模型、视频解码器和训练循环中存在一些高开销操作，本次发行版为其提供了新的 Metal 路径：

Softmax 与归约操作

Token 采样在 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 数学实现的缩放点积注意力现在能正确处理批次和头部的广播。

PyTorch 2.16 前瞻

2.16 版本的工作已在进行中。部分更改已合并到 main 分支并在夜间构建中可用，其他更改仍在审查中，细节可能在发布前发生变化。

更快的 LLM 解码

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