MPS Improvements in PyTorch 2.15

MPS Improvements in PyTorch 2.15

PyTorch 2.15 adds several improvements for Apple silicon Macs using the Metal Performance Shaders (MPS) backend. The release introduces FP8 tensors on MPS, faster language-model decoding, additional Metal kernels, broader linear-algebra support, and better memory behavior for video workloads.

The release also updates its requirements: MPS now requires macOS 15 or later, while PyTorch 2.15 requires Python 3.11 or later.

Key changes

• FP8 on MPS: float8_e4m3fn values use half the storage of BF16, with scaled matrix multiplication and other supported operations available on the GPU.
• Faster mixture-of-experts decoding: Qwen3.6-35B-A3B increases from 9.4 to 61.3 tokens/s on an M3 Ultra, while LFM2.5-8B-A1B increases from 23.6 to 123.6 tokens/s in the reported benchmarks.
• Faster dense decoding: Specialized matrix-vector kernels improve single-row linear layers, while vocabulary softmax is approximately 8 times faster.
• Expanded linear algebra: MPS now supports complex LU factorization, LU solves, triangular solves, faster small-matrix inverses, and improved QR decomposition.
• Improved audio, video, and training operations: New paths cover convolutions, padding, dropout, interpolation backward passes, and reductions.
• Lower video-decoding memory use: Wan2.2's video VAE decode peaks at about 8 GiB instead of 270 GiB.
• Python-level stream control: Public MPS streams and more reliable event timing are available.

Getting started

An Apple silicon Mac with macOS 15 or later and Python 3.11 or later is required. The MPS availability check now reports why MPS cannot be used when the operating system is too old or a Metal GPU is unavailable.

A release candidate can be installed in a separate environment with:

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

FP8 comes to MPS

PyTorch 2.15 adds torch.float8_e4m3fn tensors on MPS. FP8 uses one byte per value, which is half the storage required by FP16 or BF16.

These tensors can be multiplied with torch.nn.functional.scaled_mm, which supports tensorwise and rowwise scaling. A tensorwise-scaling example is shown below:

import torch
import torch.nn.functional as F

def quantize_fp8(x):
    limit = torch.finfo(torch.float8_e4m3fn).max
    scale = x.abs().amax().clamp_min(1e-12) / limit
    return (x / scale).to(torch.float8_e4m3fn), scale

x = torch.randn(1, 1024, device="mps", dtype=torch.float32)
weight = torch.randn(1024, 1024, device="mps", dtype=torch.float32)
x_fp8, x_scale = quantize_fp8(x)
w_fp8, w_scale = quantize_fp8(weight)

y = F.scaled_mm(
    x_fp8,
    w_fp8.t(),
    scale_a=x_scale,
    scale_recipe_a=F.ScalingType.TensorWise,
    scale_b=w_scale,
    scale_recipe_b=F.ScalingType.TensorWise,
    output_dtype=torch.bfloat16,
)
print(y.shape, y.dtype)

Faster language-model decoding

Language-model inference has two main stages: prefill and decode. During decoding, each new token enters the model as a single row, turning large matrix multiplications into matrix-vector products. PyTorch 2.15 improves both dense linear layers and mixture-of-experts layers involved in this process.

Linear layers

MPS has supported GEMV kernels for torch.mm since PyTorch 2.14, but F.linear still used MPSGraph. PyTorch 2.15 connects F.linear to the GEMV kernels, allowing common decoding inputs shaped like [1, 1, hidden_size] to use them directly. Eligible bias additions are fused as well.

Across 96 BF16 layer shapes from real models, this produces a 1.23x geometric-mean speedup. The largest improvements occur in small projections, including MoE gates, where the previous path did not fully use available memory bandwidth. Larger layers and language-model heads were already close to bandwidth limits and remain near 1x.

Mixture-of-experts layers

A mixture-of-experts layer routes different numbers of tokens to different experts. Grouped matrix multiplication processes these uneven groups in one operation, and PyTorch 2.15 adds a native MPS implementation.

In PyTorch 2.14, grouped matrix multiplication copied group offsets to the CPU and looped over the groups. The new Metal kernels read offsets directly on the GPU, using Metal Performance Primitives where available and a separate Metal implementation elsewhere. Small groups, which are common during decoding, also use a dedicated kernel.

For example, 64 tokens can be divided into groups of 16, 32, and 16:

import torch
import torch.nn.functional as F

tokens = torch.randn(64, 64, device="mps", dtype=torch.bfloat16)
weights = torch.randn(3, 128, 64, device="mps", dtype=torch.bfloat16)
ends = torch.tensor([16, 48, 64], device="mps", dtype=torch.int32)

out = F.grouped_mm(tokens, weights.transpose(-2, -1), offs=ends)
print(out.shape)  # torch.Size([64, 128])

On an M3 Ultra with torch.compile, the reported MoE decoding benchmarks show approximately 1.3x higher token throughput:

• Qwen3.6-35B-A3B: 38.7 to 51.7 tokens/s at context 512, and 36.8 to 48.7 tokens/s at context 2048.
• LFM2.5-8B-A1B: 77.1 to 103.2 tokens/s at context 512, and 74.9 to 98.7 tokens/s at context 2048.

Linear algebra improvements

Complex-valued computations

Several complex tensor operations now run on MPS. These include complex64 LU factorization, LU solves, and triangular solves. Operations and gradients built on these functions, such as complex determinants, can use the new support as well.

import torch

a = torch.randn(8, 8, device="mps", dtype=torch.complex64)
a = a + 8 * torch.eye(8, device="mps", dtype=torch.complex64)
b = torch.randn(8, 2, device="mps", dtype=torch.complex64)

lu, pivots = torch.linalg.lu_factor(a)
x = torch.linalg.lu_solve(lu, pivots, b)
print((a @ x - b).abs().max().item())

Small matrices and larger solves

Batches of FP32 matrices up to 4 x 4 can be inverted in a single Metal kernel that combines factorization and solving. This supports workloads such as camera geometry and image transformations.

Larger problems also receive optimized paths. QR decomposition uses blocked updates, while triangular solves perform much of their work through matrix multiplication. Gated DeltaNet workloads have a dedicated path for batches of 16 x 16, 32 x 32, and 64 x 64 FP32 triangular systems with at least eight right-hand sides.

SVD handling is more robust at numerical extremes. PyTorch 2.15 improves behavior for very small singular values, very large or very small input magnitudes, and basis completion with full_matrices=True. These changes also improve pseudoinverse and least-squares results.

Convolutions, padding, and training

PyTorch 2.15 adds Metal paths for operations used in audio models, video decoders, and training loops:

Softmax and reductions

Token sampling also benefits from dedicated Metal kernels. softmax and log_softmax no longer use the previous MPSGraph path for these cases.

For an FP32 softmax over a [1, 128256] row of logits, execution time drops from 57 microseconds to 7 microseconds, approximately an 8x improvement. Other vocabulary sizes show speedups ranging from 7.1x to 8.0x.

Lower memory use for video decoding

The constant-padding change produces a significant memory reduction in Wan-AI/Wan2.2-TI2V-5B-Diffusers. A full 28-tile VAE decode makes 8,764 padding calls. Moving padding to Metal reduces peak driver memory from approximately 270 GiB to 8 GiB.

The same padding change also fixes a correctness issue affecting the Qwen-Image-2.1 VAE encoder on MPS. An encode-decode round trip now matches CPU at 29.95 dB PSNR, compared with 10.06 dB previously.

More control with MPS streams

PyTorch 2.15 exposes MPS streams from Python. torch.mps.Stream and the torch.mps.stream context manager select the stream used for work, while Tensor.record_stream lets the allocator track a tensor used on another stream.

import torch

x = torch.randn(1024, device="mps")
torch.mps.synchronize()
worker = torch.mps.Stream()

with torch.mps.stream(worker):
    y = x.square()
    x.record_stream(worker)

del x
worker.synchronize()
print(y.cpu().shape)

Additional MPS-supported operations

The release adds MPS support for several more operations:

• torch.frexp
• torch.logspace
• torch.take and Tensor.put_
• torch.heaviside
• torch.special.laguerre_polynomial_l
• Unsigned-integer comparisons

logspace does not yet support integer or boolean outputs with more than one element.

Scaled dot-product attention also receives a correctness fix. The MPS math implementation now handles batch and head broadcasting correctly.

Preview of PyTorch 2.16

Work on PyTorch 2.16 is underway. Some changes are available in nightly builds, while others remain under review and may change before release.

Faster LLM decoding

When a model's decode path contains no MPSGraph operations, the GPU may wait for the CPU to encode the complete set of kernels before starting work. PyTorch 2.16 commits the command buffer every 16 kernels, allowing the GPU to begin while the CPU continues encoding.

For compiled BF16 decoding at a 2048-token context, the reported results are:

Attention runs entirely on Metal

Scaled dot-product attention is being moved fully from MPSGraph to Metal kernels. The planned implementation includes dedicated vector kernels for decoding and Metal Performance Primitives for half-precision prefill on newer macOS versions.

On an M5 Pro, grouped-query attention with 32 query heads, 8 KV heads, head dimension 80, and a 4K context drops from 465 microseconds to 65 microseconds. A 4K x 4K prefill with head dimension 40 drops from 42 milliseconds to 6 milliseconds.

More MPSGraph-to-Metal migrations

The migration away from MPSGraph continues. Batch normalization forward and backward paths are in progress. MSE loss is already up to 12 times faster on small inputs, while std_mean and var_mean are up to 4 times faster.

bincount shows a larger change for a specific workload. Counting 4 million integers when most values fall into a single bin previously took 6.8 seconds and now takes approximately half a millisecond.

A smaller installation

Bundled shaders are planned to build for Metal 3.1 by default, with Metal 4 versions compiled only for kernels that require them. This reduces libtorch_cpu.dylib by approximately 25%, from 384 MiB to 287 MiB, and prepares for Metal 4.1 shaders with faster matrix multiplication and native FP8 matrix multiplication.

Summary

PyTorch 2.15 expands MPS support with FP8 tensors, faster dense and mixture-of-experts decoding, broader linear algebra capabilities, dedicated Metal kernels, improved stream control, and substantially lower memory use for Wan2.2 video VAE decoding. PyTorch 2.16 is expected to continue the move from MPSGraph to Metal, including fully Metal-based attention and additional decoding optimizations.