PyTorch 2.15 MPS: FP8, Faster Decoding, and Metal Kernels

PyTorch 2.15, presented by Hugging Face‘s PyTorch team, pushes the MPS backend on Apple silicon closer to a first-class place to run and train PyTorch models. The release adds FP8 tensors, faster dense and MoE decoding, broader linear algebra coverage, and a long list of operations moved from MPSGraph to dedicated Metal kernels. The work is driven by running real models from the Hugging Face Hub across language, image and video and fixing gaps upstream, often with Metal kernels instead of MPSGraph. For Mac users working with transformers, diffusers or ComfyUI, the practical changes are faster local LLMs, video generation that fits in memory, corrected image editing with Qwen-Image-2.1, and FP8 weights. Requirements change too: MPS now needs macOS 15 or later, PyTorch 2.15 needs Python 3.11 or later, and the MPS availability check now explains why when those conditions are not met.

FP8 arrives as torch.float8_e4m3fn: one byte per value, half the storage of FP16 or BF16. These tensors can be multiplied directly on the GPU with F.scaled_mm, which supports tensorwise and rowwise scaling. The post’s example quantizes with abs-max scaling, then calls scaled_mm with scale_a, scale_b, ScalingType.TensorWise and output_dtype=torch.bfloat16.

Language-model decoding gets several wins. Dense linear layers benefit because F.linear now uses the GEMV kernels that previously only served torch.mm, with eligible bias additions fused; across 96 BF16 layer shapes from real models that is a 1.23x geometric-mean speedup, largest on small projections like MoE gates. Softmax and log_softmax run on dedicated Metal kernels and are about 8x faster, with an FP32 softmax over a [1, 128256] row of logits dropping from 57 µs to 7 µs. Grouped matmul now has a native Metal implementation that reads group offsets on the GPU rather than looping over groups through the CPU, including a separate kernel for the small groups common during decoding. The post reports MoE models decoding 5–6.5× faster than in 2.14 thanks to that native grouped matmul, a dedicated decoding path, fewer GPU-CPU syncs and faster top-k routing: on an M3 Ultra with torch.compile, Qwen3.6-35B-A3B goes from 9.4 to 61.3 tokens/s and LFM2.5-8B-A1B from 23.6 to 123.6 tokens/s. The dedicated decoding kernel shows up directly in throughput, with MoE models generating about 1.3× more tokens per second.

Linear algebra expands substantially. complex64 LU factorization, LU solve and triangular solve now run on MPS, unlocking complex determinants and other operations built on them. Batches of FP32 matrices up to 4×4 are inverted in a single Metal kernel that combines factorization and solving, useful for camera geometry and image transforms. QR uses blocked updates; triangular solves turn much of their work into matrix multiplication. Gated DeltaNet gets a dedicated fast path for batches of 16×16, 32×32 and 64×64 FP32 triangular systems with at least eight right-hand sides. SVD is also more robust at numerical extremes: it now correctly handles very small singular values, very large or very small input magnitudes and basis completion with full_matrices=True, which makes pseudoinverse and least-squares results more reliable.

Audio models, video decoders and training loops get several new Metal paths. Conv1d forward gets specialized kernels and an im2col-plus-matmul path for short outputs, useful for audio and streaming shapes. Constant padding moves to Metal, reflection and replication padding are faster and fix backward correctness, out-of-place dropout uses the fused kernel, upsampling backward uses a gather kernel, and channels-last conv2d backward stays on an efficient channels-last path. The constant-padding change has a dramatic side effect: a full 28-tile Wan2.2 VAE decode, with 8,764 padding calls, sees total driver-memory peak drop from about 270 GiB to 8 GiB. The same fix corrects Qwen-Image-2.1 VAE encoding on MPS, raising encode-decode PSNR from 10.06 dB to 29.95 dB against CPU. PyTorch 2.15 also exposes public MPS streams from Python via torch.mps.Stream and the torch.mps.stream context manager, adds Tensor.record_stream and more reliable event timing, and expands op coverage with frexp, logspace, take and Tensor.put_, heaviside, special.laguerre_polynomial_l and unsigned-integer comparisons; logspace does not yet support integer or boolean outputs with more than one element. SDPA’s math implementation now handles batch and head broadcasting correctly.

The post closes with a preview of PyTorch 2.16, noting that some changes have landed and others may change before release. Attention is dropping MPSGraph completely and will run entirely on Metal, with dedicated vector kernels for decoding and Metal Performance Primitives for half-precision prefill on newer macOS versions; on an M5 Pro, grouped-query attention decoding with a 4K context drops from 465 µs to 65 µs, and a 4K x 4K prefill with head dim 40 goes from 42 ms to 6 ms. Command buffers will be committed every 16 kernels so the GPU starts working while the CPU is still encoding; for compiled BF16 decoding at a 2048-token context, LFM2.5-8B-A1B improves by 16%, Qwen3.6-35B-A3B by 19%, and Qwen3.8-27B by 4%. More MPSGraph-to-Metal migrations have already landed: MSE loss is up to 12x faster on small inputs, std_mean and var_mean up to 4x faster, and bincount drops from 6.8 seconds to 0.5 milliseconds when counting 4 million integers. Bundled Metal shaders will build for Metal 3.1 by default, cutting libtorch_cpu.dylib from 384 MiB to 287 MiB and preparing for Metal 4.1 shaders with faster matmuls and native FP8 matrix multiplication.

MPS in PyTorch 2.15

View Original