PyTorch 2.15 speeds up MoE models on Mac by 5–6.5x and adds FP8

PyTorch 2.15 turns the Mac from a tolerable platform for experiments into a fully workable machine for local LLMs and video generation. The PyTorch team at Hugging Face, which works on the MPS backend, added FP8 support and replaced some MPSGraph operations with dedicated Metal kernels. On an M3 Ultra with torch.compile, Qwen3.6-35B-A3B now generates 61.3 tokens per second instead of 9.4, and video decoding in Wan2.2 no longer inflates driver memory to 270 GiB.

FP8 and MoE decoding: where 61 tokens per second came from

MPS now has the torch.float8_e4m3fn type. It uses one byte per value, half the size of FP16 or BF16. Such tensors can be multiplied with F.scaled_mm, which supports both tensorwise and rowwise scaling. The other operations available for this type on CUDA and CPU work as well.

Most of the speedup comes from the decode phase. There, each new token goes through the model as a single row, so large matrix-matrix multiplications turn into matrix-vector ones. In 2.14, only torch.mm had GEMV kernels, while F.linear went through MPSGraph. Now the typical [1, 1, hidden_size] input uses them directly, and suitable bias additions are fused into the kernel. Across 96 BF16 layer shapes taken from real models, this gives a geometric-mean speedup of 1.23×. Small projections, such as MoE gates, benefit most. Large layers and the LM head were already limited by memory bandwidth, so they gain roughly 1×.

For MoE, the team built a native grouped matmul. Previously, group offsets were copied to the CPU and the operation looped over groups. The new kernels read the offsets directly on the GPU. They use Metal Performance Primitives where available and a separate Metal path elsewhere. Small groups, typical of decoding, get their own kernel. The team also cut GPU-CPU synchronizations and sped up top-k routing.

Small fixes that were eating memory and time

The most striking result came from a fix to constant padding. When fully decoding 28 tiles, the VAE decoder of Wan-AI/Wan2.2-TI2V-5B-Diffusers makes 8 764 padding calls. Through MPSGraph, peak driver memory reached 270.135 GiB; with Metal padding it is 8.108 GiB. Peak allocated tensors barely changed: 4.422 versus 4.095 GiB. What had been ballooning was the memory the driver held for graphs.

The same fix repaired the Qwen-Image-2.1 VAE encoder. In 2.14 it silently corrupted images on MPS, which broke editing workflows in ComfyUI. Now an encode-decode round trip matches the CPU at 29.95 dB PSNR, up from 10.06 dB.

The rest of the 2.15 changes:

• softmax and log_softmax: dedicated Metal kernels; FP32 softmax over a [1, 128256] row dropped from 57 to 7 µs, roughly 8x.
• conv1d forward: specialized kernels instead of MPSGraph, plus an im2col-plus-matmul path for short outputs, useful for audio and streaming models.
• Reflection and replication padding: new Metal implementations, faster and with a fixed backward pass.
• Dropout: nn.Dropout and F.dropout use the existing fused kernel.
• Upsampling backward: a Metal gather kernel instead of the resize gradient graph.
• Channels-last conv2d backward: stays on the efficient channels-last path, without an NCHW graph and expensive strided access.
• Complex linear algebra: LU, LU solve and triangular solve for complex64, along with complex determinants and gradients.
• Small matrices: FP32 batches up to 4×4 are inverted by a single Metal kernel that combines factorization and solve.
• QR and triangular solve: blocked updates that reduce much of the work to matrix multiplication. For Gated DeltaNet there is a fast path for 16×16, 32×32 and 64×64 FP32 systems with eight or more right-hand sides.
• SVD: correctly handles very small singular values, extreme input scales and full_matrices=True, which improves pseudoinverses and least squares.
• MPS streams: public torch.mps.Stream, the torch.mps.stream context and Tensor.record_stream, plus more reliable event-based timing.
• New operations: torch.frexp, torch.logspace (no integer or boolean output yet for several elements), torch.take, Tensor.put_, torch.heaviside, torch.special.laguerre_polynomial_l and unsigned integer comparisons. Batch and head broadcasting is fixed in the math implementation of SDPA.

What's coming in 2.16 and how to try 2.15

Some of the plans for 2.16 are already in the nightly builds. The command buffer will be committed every 16 kernels, so the GPU starts working while the CPU encodes the rest.

In compressed BF16 decoding at a 2048-token context, LFM2.5-8B-A1B rises from 108.8 to 126.

0 tokens/s (+16%), Qwen3.6-35B-A3B from 52.9 to 62.

8 (+19%), and Qwen3.8-27B from 12.1 to 12.

6 (+4%). Attention will move entirely to Metal. On an M5 Pro, GQA decoding (32 query heads, 8 KV heads, head dim 80, 4K context) will drop from 465 to 65 µs, and a 4K×4K prefill with head dim 40 from 42 to 6 ms.

Among other changes, bincount over 4 million integers with nearly identical values will take half a millisecond instead of 6.8 s, and libtorch_cpu.dylib will shrink from 384 to 287 MiB because shaders will be built for Metal 3.

1 by default.

Here is the practical sequence. You need an Apple silicon Mac with macOS 15 or newer and Python 3.11+; otherwise the MPS availability check will now explain why it refused. Install the release candidate in a separate environment with uv pip install torch==2.15.0 --index-url https://download.pytorch.org/whl/test/cpu. If you run MoE models through transformers, it makes sense to enable torch.compile and compare tokens per second against 2.14. ComfyUI users with Qwen-Image-2.1 should update because of the VAE fix, and owners of Macs with a lot of memory who generate video in Wan2.2 will avoid the unexpected driver-memory bloat. If something is still slow or doesn't work, the team asks users to open an issue on GitHub.