PyTorch 2.15 為 Mac MPS 加入 FP8 支援並加快模型解碼

如果你在 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 維護者及社羣貢獻者共同完成。

重點摘要

• FP8 登陸 Mac：儲存 float8_e4m3fn 數值所需空間只有 BF16 的一半，並可使用縮放矩陣乘法，以及 CUDA/CPU 支援的所有其他運算！
• MoE 模型的解碼速度較 2.14 快 5–6.5×：在 M3 Ultra 上，Qwen3.6-35B-A3B 的速度由 9.4 提升至 61.3 tokens/s；這要歸功於原生分組矩陣乘法及專用解碼路徑、更少的 GPU-CPU 同步，以及更快速的 top-k 路由。
• 所有 LLM 每個 token 所需時間更短：單列線性層採用專用矩陣向量核心，而對詞彙表進行 softmax 的速度則快約 8 倍。
• 更多線性代數運算交由 GPU 處理：複數 LU 分解和三角方程組求解、小型矩陣反矩陣速度更快，QR 分解也更快。
• 音訊、影片和訓練同樣獲得改進：卷積、填充、dropout、插值反向傳播和歸約運算的路徑更快。Wan2.2 的影片 VAE 解碼，記憶體峯值現為 8 GiB，而非 270 GiB。
• 透過 Python 進行更多控制：提供公開的 MPS 串流，事件計時亦更可靠。
• 更新後的要求：MPS 現時需要 macOS 15 或更新版本；PyTorch 2.15 需要 Python 3.11 或更新版本。
• 2.16 前瞻：注意力運算全面交由 Metal 處理，更多從 MPSGraph 遷移至 Metal 的工作，以及更快速的 LLM 解碼。

這對你有甚麼影響

如果你使用 transformers、diffusers 或 ComfyUI 等工具在 Mac 上執行模型，以下是 2.15 帶來的實際改變：

• 本機 LLM 速度大幅提升。MoE 模型的解碼速度較 2.14 快 5–6.5×：在配備 torch.compile 的 M3 Ultra 上，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 中回報此問題）。相同的填充變更修正了這個問題：現在，編碼及解碼一輪後的結果與 CPU 相符，PSNR 為 29.95 dB，高於原來的 10.06 dB。
• Mac 上的 FP8 權重。FP8 張量所需記憶體只有 BF16 的一半，而 F.scaled_mm 可直接在 GPU 上將它們相乘。（更多資訊）

目錄

• 開始使用
• MPS 支援 FP8
• 加快語言模型解碼
• 線性代數效能提升
• 卷積、填充及訓練
• Softmax 和歸約運算
• 大幅減少影片解碼所需的記憶體
• 透過 MPS 串流取得更多控制
• 更多運算支援 MPS
• 搶先一覽 PyTorch 2.16
• 總結

開始使用

你需要使用搭載 Apple silicon 的 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

MPS 支援 FP8

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 都有預填充階段和解碼階段。在解碼期間，每個新 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 Performance Primitives，其他情況則採用另一條 Metal 路徑。解碼期間常見的小型羣組也有專用核心。

例如，64 個 token 可以分成 16、32 和 16 個 token 的羣組：

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 吞吐量上。在配備 torch.compile 的 M3 Ultra 上，MoE 模型每秒生成的 token 數量約增加 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 路徑：

Softmax 與歸約

2.15 也加快了 token 抽樣。softmax和log_softmax現在會使用專用 Metal 核心；大型詞彙表受益最大：對一行 [1, 128256] logits 執行 FP32 softmax，耗時由 57 µs 降至 7 µs，速度約比 MPSGraph 版本快 8 倍。其他常見詞彙表大小也有類似的加速效果：

大幅減少影片解碼的記憶體佔用

這個版本其中一項最顯著的成果來自常數填充的改動。Wan-AI/Wan2.2-TI2V-5B-Diffusers 中的 VAE 解碼在完整解碼 28 個分塊時會進行 8,764 次填充呼叫。改用 Metal 填充後，驅動程式記憶體峯值由約 270 GiB 降至 8 GiB：

使用 MPS 串流，掌控更多

現在你可以直接在 Python 中選擇用哪個 MPS 串流執行工作。torch.mps.Stream 和 torch.mps.stream context manager提供相關介面，而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，並已包含在 nightly 版本中；其他變更仍在審查，因此正式發布前細節可能有所變動。

更快速的 LLM 解碼

當模型的解碼路徑不再包含任何 MPSGraph 運算時，GPU 必須等到 CPU 完成所有核心的編碼後，才會開始處理 token。PyTorch 2.16 每 16 個核心便提交一次命令緩衝區，讓 GPU 在 CPU 仍在編碼其餘核心時開始工作。以下是上下文長度為 2048 個 token 時，已編譯 BF16 解碼的效能：

注意力運算完全在 Metal 上執行

縮放點積注意力將完全不再使用 MPSGraph：所有形狀都會使用 Metal 核心執行，解碼會使用專用向量核心，而較新版本的 macOS 則會使用 Metal Performance Primitives 進行半精度預填。過去會退回使用 MPSGraph 的形狀，效能提升最為顯著。在 M5 Pro 上，使用分組查詢注意力解碼（32 個查詢頭、8 個 KV 頭、頭部維度為 80、上下文長度為 4K）的耗時由 465 µs 降至 65 µs；頭部維度為 40 的 4K × 4K 預填，耗時則由 42 ms 降至 6 ms。

更多從 MPSGraph 遷移至 Metal 的工作

遷離 MPSGraph 的工作仍在繼續。批次正規化的前向傳播和反向傳播仍在進行中；MSE 損失（小型輸入的速度最高提升 12 倍）、std_mean及var_mean（速度最高提升 4 倍），以及bincount則已經合併。bincount的效能提升最為驚人：當大部分數值落在同一個分箱時，計算 400 萬個整數過去需要 6.8 秒，現在只需半毫秒。

縮小安裝體積

隨附的著色器將預設以 Metal 3.1 為目標進行建置，只有需要 Metal 4 的核心才會編譯 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 提交 issue告訴我們。

非常感謝 Nikita Shulga、Kurt Mohler 和 Aaron Gokaslan 審閲這項工作的許多內容，亦感謝所有在這次發佈中為 MPS 作出貢獻的人。