MPS in PyTorch 2.15

MPS is PyTorch's backend for the GPU in Apple silicon Macs. Many of its operations used to run through MPSGraph, Apple's graph framework, rather than through kernels written specifically for each job. In a Hugging Face blog post, the company's PyTorch team explains that it has been running real models from the Hub on Macs, finding where MPS falls short and fixing those problems upstream, increasingly by replacing MPSGraph with dedicated Metal kernels. Much of what the post describes in PyTorch 2.15 comes from that work. The post presents 2.15 as a release candidate to try. MPS now requires macOS 15 or later, and PyTorch 2.15 requires Python 3.11 or later.

The largest gain is in language-model decoding, where the model generates one token at a time and each step feeds a single row through the network. Mixture-of-experts (MoE) models route each token to a few specialist sub-networks, so a layer has to multiply unevenly sized groups of tokens. In 2.14, that grouped multiplication copied the group boundaries to the CPU and looped over the groups. In 2.15, new Metal kernels read the boundaries directly on the GPU, and small groups, the common case during decoding, get a kernel of their own. Combined with fewer GPU-CPU synchronizations and faster top-k routing, the team reports MoE decoding 5 to 6.5 times faster than in 2.14. On an M3 Ultra with torch.compile, Qwen3.6-35B-A3B goes from 9.4 to 61.3 tokens per second, and LFM2.5-8B-A1B goes from 23.6 to 123.6. The post separately credits the decoding kernel with about 1.3 times more tokens per second for MoE models. For dense models it reports smaller, per-operation gains. Routing single-row linear layers to matrix-vector kernels gives a 1.23 times geometric-mean speedup across 96 real layer shapes, mostly on small projections, and softmax over a 128,256-token vocabulary drops from 57 to 7 microseconds.

The most dramatic memory result comes from moving constant padding to Metal. The video VAE decoder in Wan2.2-TI2V-5B makes 8,764 padding calls over a 28-tile decode. With the old MPSGraph padding, peak driver memory reached about 270 GiB, while the tensors themselves peaked near 4.4 GiB. With Metal padding, the driver peak falls to about 8.1 GiB. So the waste was in the driver rather than the model's data. The same padding change fixes the Qwen-Image-2.1 VAE encoder, which had silently corrupted images on MPS. An encode-decode round trip now matches the CPU result at 29.95 dB PSNR, up from 10.06 dB.

The release also adds float8_e4m3fn tensors on MPS, which use one byte per value, half the storage of BF16. These tensors can be multiplied directly on the GPU with F.scaled_mm, which supports both tensorwise and rowwise scaling. The post previews more of the same direction for 2.16, including attention running entirely on Metal, and notes that some of that work is still in review and may change.