PyTorch 2.15のMPSバックエンド、FP8対応とMoE推論5〜6.5倍高速化

PyTorch 2.15で何が変わったか

Hugging FaceのPyTorchチームが、PyTorch 2.15におけるMPSバックエンドの改善内容をブログで報告しました。MPSはApple Silicon上のGPU（Metal）でPyTorchを動かすためのバックエンドです。チームは2026年初頭の時点では、小さなモデルでも未対応演算や性能の急な落ち込み、想定外のメモリ増加への回避策が必要だったと説明しています。そこでHugging Face Hub上の実モデルを言語・画像・動画にわたり実行し、不足点をPyTorch本体（upstream）で修正してきたという位置づけです。

2.15の柱は大きく4つに整理できます。Mac上でのFP8（float8_e4m3fn）テンソル対応、Mixture-of-Experts（MoE）を含むLLMデコードの高速化、複素LU分解やQR分解といった線形代数のGPU実行、そしてMPSGraphから専用Metalカーネルへの置き換えです。加えて動画VAEデコードのメモリ削減や、PythonからのMPSストリーム制御も含まれます。作業はPyTorchのメンテナやコミュニティの貢献者と共同で進められたとされています。

FP8は1値あたり1バイトで、FP16やBF16の半分の記憶域になります。F.scaled_mmはtensorwiseとrowwiseの両方のスケーリングに対応し、FP8テンソル同士の行列積をGPU上で直接実行できます。原文では1テンソル1スケールの量子化例と、出力dtypeにbfloat16を指定するコード例が示されています。

LLMデコードはどこが速くなったか

LLMの推論は、プロンプトをまとめて処理するprefillと、1トークンずつ生成するdecodeに分かれます。decodeでは入力が1行になるため、大きな行列積が行列ベクトル積（GEMV）に変わります。2.15では、これまでMPSGraph経由だったF.linearを、2.14から存在したGEMVカーネルに接続しました。実モデル由来のBF16の96種類のレイヤ形状で、幾何平均1.23倍の高速化が得られたと報告されています。効果が大きいのはMoEのゲートのような小さな投影で、大きな層やLMヘッドは元々メモリ帯域に近く、ほぼ1倍にとどまるということです。

MoE層では、エクスパートごとにトークン数が異なるため、grouped matmul（グループ化された行列積）で不均一なグループをまとめて処理します。2.14はグループのオフセットをCPUにコピーし、グループごとにループしていました。2.15ではネイティブのMetalカーネルがGPU上でオフセットを直接読み、デコード時に多い小さいグループ専用のカーネルも用意しています。この結果、M3 Ultraとtorch.compileの構成でQwen3.6-35B-A3Bが9.4から61.3トークン/秒、LFM2.5-8B-A1Bが23.6から123.6トークン/秒になったとされています。

トークンサンプリングも軽くなりました。softmaxとlog_softmaxが専用Metalカーネルで動き、[1, 128256]のFP32ロジットに対するsoftmaxは57マイクロ秒から7マイクロ秒へ、約8倍速くなっています。語彙が大きいモデルほど効果が出やすい部分です。

線形代数とメモリ削減

線形代数では、complex64のLU分解・LU求解・三角求解がMPS上で動くようになり、複素行列式など、それらに依存する演算と勾配も利用できるようになりました。カメラ幾何や画像変換で使う4×4以下のFP32行列のバッチ逆行列は、分解と求解をまとめた単一のMetalカーネルで処理します。QR分解はブロック更新、三角求解は行列積への変換で高速化され、SVDも非常に小さい特異値や極端な入力スケール、full_matrices=Trueの基底補完を正しく扱えるようになったとしています。

メモリ面で目立つのは動画VAEのデコードです。Wan2.2のタイル型VAEデコードは、28タイルの処理で8,764回のpadding呼び出しを行います。paddingがMetal上で動くようになった結果、ピークのドライバメモリが約270 GiBから8 GiBに下がりました。画像編集では、Qwen-Image-2.1のVAEエンコーダが2.14で画像を静かに破損させていた問題（ComfyUIで報告）が同じpaddingの変更で修正され、エンコード・デコードの往復でCPUと比較して29.95 dB PSNR（従来10.06 dB）になったと報告されています。

要件と2.16の予告

利用要件は見直されています。MPSはmacOS 15以降、PyTorch 2.15はPython 3.11以降が必須です。macOSが古い場合やMetal GPUが見つからない場合は、MPSの可用性チェックが理由を返すようになりました。またtorch.frexp、torch.logspace、torch.take、Tensor.put_、torch.heaviside、torch.special.laguerre_polynomial_l、符号なし整数の比較などがMPSで動くようになっていますが、logspaceは要素数が2以上の整数・真偽値出力には未対応です。scaled dot-product attentionのMPS実装は、バッチとヘッドのブロードキャストを正しく扱うよう修正されています。

2.16は開発中で、一部はmainに入りnightlyで利用可能、残りはレビュー中で詳細が変わり得ると明記されています。注目点はattentionの完全なMetal化、16カーネルごとのコマンドバッファcommit、MPSGraphからMetalへの移行継続です。M5 Proでの計測例として、grouped-query attention（クエリ32ヘッド、KV 8ヘッド、ヘッド次元80、4Kコンテキスト）のデコードが465マイクロ秒から65マイクロ秒、ヘッド次元40の4K×4K prefillが42ミリ秒から6ミリ秒になるとされています。bincountでは、400万個の整数カウントが6.8秒から0.5ミリ秒になったという例も示されています。