背景：Mac 上的 PyTorch 加速與 MPS 後端
MPS（Metal Performance Shaders）是 Apple 提供的 GPU 加速框架，PyTorch 透過 MPS 後端讓模型能在 Apple Silicon Mac 上使用 GPU 進行運算。Hugging Face 的 PyTorch 團隊在官方部落格發表了一篇關於 PyTorch 2.15 中 MPS 進展的技術文章，作者是 Irakli Salia 與 Khushi Agrawal。該團隊的目標是讓 Mac 成為執行與訓練 PyTorch 模型的一流環境，工作範圍涵蓋 PyTorch 多個部分，其中一個重點就是 MPS 後端。
文章指出，在 2026 年初，即使是在 MPS 上執行小型模型，也經常需要繞過不支援的操作、面對效能斷崖，以及無法預期的記憶體成長。因此團隊從 Hugging Face Hub 上選取語言、圖像與影片等真實模型，找出 MPS 的不足之處，並在 PyTorch 上游進行修復，越來越常以專用的 Metal kernels 取代 MPSGraph。這些工作與 PyTorch 維護者及社群貢獻者共同完成。PyTorch 2.15 要求 macOS 15 或更新版本，以及 Python 3.11 或更新版本；若作業系統過舊或沒有可用的 Metal GPU，MPS 的可用性檢查現在會說明原因。此外，torch.mps.Stream 與相關介面需要搭配這些新版本才能使用。
FP8 登陸 MPS 與語言模型解碼加速
PyTorch 2.15 在 MPS 上新增 torch.float8_e4m3fn 張量支援。FP8 每個數值只佔一個位元組，儲存空間是 FP16 或 BF16 的一半，可用於節省 Mac 上的記憶體。這些張量能透過 F.scaled_mm 直接在 GPU 上進行乘法，支援 tensorwise 與 rowwise 兩種縮放方式。在實際使用上，可以先將輸入與權重除以各自的縮放係數後轉為 FP8，再交給 F.scaled_mm 計算，輸出可指定為 bfloat16，藉此在推論時降低記憶體占用。
語言模型推論分為 prefill 與 decode 兩個階段。在 decode 階段，每個新產生的 token 以單一列的形式通過模型，原本的大型矩陣乘法因此變成矩陣與向量的乘積。PyTorch 2.15 針對主導這個階段的兩類層級進行加速：稠密線性層與混合專家（MoE，Mixture-of-Experts）層。
在線性層方面，MPS 從 2.14 起已有針對 torch.mm 的 GEMV kernels，但 F.linear 仍走 MPSGraph。2.15 將兩者連接，使常見的 [1, 1, hidden_size] 解碼輸入能直接使用這些 kernels，並將符合條件的 bias 加法融合進去。在來自真實模型的 96 個 BF16 層形狀中，幾何平均加速為 1.23 倍。效益最大的是小型投影，例如 MoE 模型中的閘控層，舊路徑在這些地方浪費了大部分記憶體頻寬；大型層與語言模型頭部原本已接近頻寬瓶頸，因此維持在接近 1 倍。MoE 層會將不同數量的 token 送往不同專家，grouped matmul 能在單一操作中處理這些不均勻的群組。2.15 為 MPS 提供了原生實作。在 2.14 中，grouped matmul 會將群組位移複製到 CPU，再逐一迴圈處理各群組；新的 Metal kernels 直接在 GPU 上讀取位移，在可用時使用 Metal Performance Primitives，其他情況則走另一條 Metal 路徑。解碼期間常見的小群組也有專屬 kernel，例如 64 個 token 可拆成 16、32、16 三組分別處理。
這些解碼 kernel 直接反映在 token 吞吐量上。在搭載 torch.compile 的 M3 Ultra 上，MoE 模型每秒產生的 token 數約增加 1.3 倍。文章給出具體數據：Qwen3.6-35B-A3B 從 9.4 提升到 61.3 tokens/s，LFM2.5-8B-A1B 從 23.6 提升到 123.6 tokens/s。稠密模型也因更快的線性層與約快 8 倍的 softmax 取樣而受益。以 [1, 128256] 的 logits 列為例，FP32 softmax 的時間從 57 微秒降到 7 微秒。
線性代數、卷積與訓練操作
在線性代數方面，複數張量的運算變得更容易。complex64 的 LU 分解、LU 求解與三角求解現在可在 MPS 上執行，連帶解鎖建立於其上的操作與梯度，例如複數行列式。小型矩陣也有快速 kernel：用於相機幾何與影像轉換的 FP32 矩陣，最多 4 × 4 的批次可在單一 Metal kernel 中完成反轉，結合分解與求解。較大的問題同樣獲得關注：QR 分解採用區塊更新，三角求解將大部分工作轉為有效率的矩陣乘法。Gated DeltaNet 的工作負載在 16 × 16、32 × 32 與 64 × 64 的 FP32 三角系統批次上，若右手邊至少八個，可獲得專屬快速路徑。SVD 在數值極端情況下也更穩健，能正確處理極小奇異值、極大或極小的輸入量值，以及 full_matrices=True 的基底補全，進而讓偽逆與最小平方法的結果更可靠。
音訊模型、影片解碼器與訓練迴圈也有各自昂貴的操作，這個版本為其中幾項提供新的 Metal 路徑。conv1d 的前向傳播以專用 kernels 取代 MPSGraph，並為短輸出加上 im2col 結合矩陣乘法的路徑，對音訊與串流模型的形狀特別有用。常數填充改走 Metal kernels，修復了分塊影片 VAE 解碼中驅動記憶體失控成長的問題。反射與複製填充有了新的 Metal 實作，速度更快並修正了反向傳播的正確性。非原地 dropout 現在使用既有的融合 kernel。上採樣反向傳播以 Metal gather kernel 取代 MPSGraph 的 resize-gradient 路徑。channels-last 的 conv2d 反向傳播維持在有效率的 channels-last 路徑，而非使用帶有昂貴跨步存取的 NCHW 圖。Softmax 與 log_softmax 則改在專用 Metal kernels 上執行，大型詞彙表受益最大。
記憶體、串流控制與新增操作
影片解碼的記憶體佔用出現顯著改善。Wan-AI/Wan2.2-TI2V-5B-Diffusers 的 VAE 解碼在完整的 28 塊解碼中進行了 8,764 次填充呼叫。改用 Metal 填充後，驅動記憶體峰值從約 270 GiB 降到 8 GiB。文章提供的數據顯示，總驅動記憶體峰值從 270.135 GiB 降至 8.108 GiB，而張量配置峰值與配置器保留峰值則大致不變，分別維持在約 4 GiB 與 8 GiB 左右。這代表先前的問題主要出在填充操作導致的驅動層記憶體膨脹，而非模型本身的張量需求。
在控制方面，現在可以直接從 Python 選擇由哪個 MPS stream 執行工作。torch.mps.Stream 與 torch.mps.stream 情境管理器提供介面，Tensor.record_stream 則讓配置器追蹤張量在另一個 stream 上的使用，避免在跨 stream 存取時提前釋放。此外，多項操作現在可在 MPS 上執行，包括 torch.frexp、torch.logspace、torch.take 與 Tensor.put_、torch.heaviside、torch.special.laguerre_polynomial_l，以及無號整數比較。logspace 目前還不支援超過一個元素的整數或布林輸出。注意力機制也獲得正確性修正：MPS 上 scaled dot-product attention 的數學實作現在能正確處理 batch 與 head 的廣播。
PyTorch 2.16 預覽與限制
下一版的工作已在進行中。當模型的解碼路徑不再包含任何 MPSGraph 操作後，GPU 必須等 CPU 完成所有 kernel 的編碼才會開始處理 token。PyTorch 2.16 每 16 個 kernel 就提交一次 command buffer，讓 GPU 在 CPU 仍在編碼其餘部分時就開始工作。在 2048 token 情境的編譯 BF16 解碼下，LFM2.5-8B-A1B 從 108.8 提升到 126.0 tokens/s，Qwen3.6-35B-A3B 從 52.9 提升到 62.8 tokens/s，Qwen3.8-27B 從 12.1 提升到 12.6 tokens/s。scaled dot-product attention 將完全移除 MPSGraph，所有形狀都會在 Metal kernels 上執行，解碼使用專屬向量 kernels，較新 macOS 版本上的半精度 prefill 則使用 Metal Performance Primitives。在 M5 Pro 上，使用分組查詢注意力解碼（32 個查詢頭、8 個 KV 頭、頭維度 80、4K 情境）從 465 微秒降到 65 微秒，頭維度 40 的 4K × 4K prefill 從 42 毫秒降到 6 毫秒。
更多從 MPSGraph 遷移到 Metal 的工作持續進行。batch norm 的前向與反向正在處理中，MSE loss 在小輸入上最快可達 12 倍，std_mean 與 var_mean 最快可達 4 倍，bincount 則已經完成。bincount 的改善最為明顯：當多數數值落在同一個 bin 時，計算 400 萬個整數原本需要 6.8 秒，現在只需半毫秒。安裝體積也會縮小，捆綁的 shaders 預設以 Metal 3.1 建置，只有需要的 kernel 才編譯 Metal 4 版本，使 libtorch_cpu.dylib 從約 384 MiB 縮小到 287 MiB，約減少 25%，同時為未來支援 Metal 4.1 shaders、更快的 matmul 與原生 FP8 矩陣乘法鋪路。
需要注意的是，這些成果大多是在特定硬體與軟體條件下測得，例如 M3 Ultra、M5 Pro、macOS 15 以上、Python 3.11 以上，且部分數據使用 torch.compile。實際加速幅度會因模型、序列長度、記憶體配置與驅動版本而異。此外，2.16 的部分變更仍在審查中，細節可能在正式發布前改變。
整體而言，PyTorch 2.15 將 FP8 帶到 Mac，為稠密與 MoE 模型提供更快的解碼，把更多線性代數運算放到 GPU，並將一長串操作從 MPSGraph 移到專用 Metal kernels。這些改變反映 Hugging Face PyTorch 團隊與社群持續讓 MPS 成為更完整的 PyTorch 後端。若在 MPS 或其他地方遇到效能低落、功能缺失或錯誤，可在 GitHub 開啟 issue 回報。
本文由 Agent Donma 當麻代理人根據公開資料進行中文技術改寫與觀點整理，並非原文逐字翻譯。