MPS in PyTorch 2.15

What it claims

PyTorch 2.15 changes Apple silicon (MPS) workloads in several ways beyond FP8 and MoE speed. It fixes a silent image corruption bug in the Qwen-Image-2.1 VAE encoder, routes F.linear decode inputs to GEMV kernels, adds public MPS streams, and moves more linear algebra onto the GPU. It also raises the minimum requirements: MPS needs macOS 15 or later, and PyTorch 2.15 needs Python 3.11 or later. The Hugging Face PyTorch team wrote the post. Most gains come from replacing MPSGraph with dedicated Metal kernels.

How it works

• MoE decoding. In 2.14, grouped matmul copied group offsets to the CPU and looped over the groups. The 2.15 Metal kernels read offsets on the GPU. They use Metal Performance Primitives where available and a separate Metal path elsewhere, with their own kernel for small decode groups.
• Dense decoding. The [1, 1, hidden_size] decode input to F.linear now uses GEMV kernels, with eligible bias additions fused in.
• Wan2.2 memory. The memory drop comes from moving constant padding to Metal. The VAE makes 8,764 padding calls over a 28-tile decode.
• Linear algebra. complex64 LU, LU solve and triangular solve now run on MPS. FP32 matrices up to 4 × 4 are inverted in one fused kernel. QR uses blocked updates. Gated DeltaNet gets a fast path for 16, 32 and 64 square FP32 triangular systems with at least eight right-hand sides.

Evidence

• F.linear GEMV: 1.23× geometric-mean speedup across 96 BF16 layer shapes from real models. Small projections such as MoE gates gain most. Large layers and LM heads stay near 1×.
• MoE on an M3 Ultra with torch.compile: LFM2.5-8B-A1B goes from 23.6 to 123.6 tokens/s. The dedicated decoding kernel alone accounts for about 1.3× more tokens/s.
• Softmax: an FP32 softmax over a [1, 128256] row drops from 57 µs to 7 µs.
• Qwen-Image-2.1 encode-decode round trip: 29.95 dB PSNR against CPU, up from 10.06 dB in 2.14. The bug was reported in ComfyUI.

How to use it

Install the release candidate in a separate environment:

uv pip install torch==2.15.0 --index-url https://download.pytorch.org/whl/test/cpu

The new APIs are:

• FP8: F.scaled_mm takes torch.float8_e4m3fn inputs with scale_recipe_a/scale_recipe_b set to F.ScalingType.TensorWise or rowwise, plus an output_dtype such as torch.bfloat16.
• Grouped matmul: F.grouped_mm(tokens, weights.transpose(-2, -1), offs=ends), where offs is an int32 tensor of cumulative group ends.
• Streams: torch.mps.Stream(), the torch.mps.stream(...) context manager and Tensor.record_stream.
• Newly supported ops: torch.frexp, torch.logspace, torch.take, Tensor.put_, torch.heaviside, torch.special.laguerre_polynomial_l and unsigned-integer comparisons.

The MPS availability check now reports why MPS is unavailable when the OS is too old or no Metal GPU is present.

Limits

The post covers a release candidate. logspace does not yet support integer or boolean outputs with more than one element. Large linear layers see no meaningful speedup.

The 2.16 previews are partly unmerged, and the authors say details may change. These include:

• SDPA fully on Metal: 465 µs to 65 µs for GQA decode on an M5 Pro.
• Command-buffer commits every 16 kernels.
• bincount going from 6.8 s to 0.5 ms on skewed input.
• A 25% smaller libtorch_cpu.dylib.

Native FP8 matmul is described only as future groundwork for Metal 4.1 shaders. The post does not say whether 2.15's FP8 path is native.