PyTorch 2.15 为 MPS 带来 FP8 支持，MoE 模型解码提速 5–6.5 倍

OmniTools 10月20日消息，Hugging Face 的 PyTorch 团队于 10 月 20 日发布博文，介绍 PyTorch 2.15 中 MPS 后端的更新内容。该版本在 Apple silicon Mac 上新增 FP8 支持：torch.float8_e4m3fn 张量每个值占一个字节，存储空间为 FP16 或 BF16 的一半，并可通过 F.scaled_mm 在 GPU 上直接相乘，支持按张量和按行两种缩放方式。

在语言模型解码方面，2.15 将已有的 GEMV 内核接入 F.linear，在 96 个来自真实模型的 BF16 层形状上取得 1.23 倍几何平均加速；MoE 层新增原生 grouped matmul 实现，解码阶段不再把 group offsets 复制到 CPU。官方给出的数据是，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。此外，softmax 与 log_softmax 改由专用 Metal 内核执行，对 [1, 128256] 的 FP32 logits 做 softmax 的耗时从 57 微秒降至 7 微秒，约快 8 倍。

其他改动涉及线性代数、卷积与内存占用：complex64 的 LU 分解、LU 求解和三角求解可在 MPS 上运行，小矩阵求逆以及 QR 分解得到加速；conv1d 前向、常量填充和 reflection/replication 填充改用 Metal 内核。Wan2.2 的 tiled VAE 解码峰值驱动内存从约 270 GiB 降至约 8 GiB；Qwen-Image-2.1 的 VAE 编码器在 2.14 上静默损坏图像的问题也被修复，编解码往返 PSNR 从 10.06 dB 升至 29.95 dB。该版本要求 macOS 15 或更高版本以及 Python 3.11 或更高版本，团队还预告 2.16 将把 attention 完全放到 Metal 上。