MPS in PyTorch 2.15

Introduction

PyTorch 2.15 is a major step forward for developers who run or train models on Apple Silicon Macs using the Metal Performance Shaders (MPS) backend. The release adds FP8 tensor support, sub‑8‑bit memory savings, and a suite of new Metal kernels that accelerate language‑model decoding, mixture‑of‑experts (MoE) routing, complex‑valued linear algebra, convolutions, padding, and softmax. In practice these changes translate into multi‑fold speedups for local LLMs, dramatically smaller video‑generation memory footprints, and finer‑grained control over MPS streams.

Faster language‑model decoding

• Dense linear layers now use dedicated GEMV kernels directly from `F.linear`, eliminating the previous MPSGraph fallback. Across a representative set of BF16 layer shapes this yields a ~1.2× geometric‑mean speedup, with the biggest gains on small projection layers.
• Mixture‑of‑experts layers receive a native grouped‑matmul implementation that runs entirely on the GPU. The new Metal kernels read token offsets on‑device, avoiding CPU‑side loops. Real‑world MoE models (e.g., Qwen‑3.6‑35B‑A3B) decode 5–6.5× faster than in 2.14 on an M3 Ultra when compiled with `torch.compile`.
• Softmax & log‑softmax are now Metal‑accelerated. A vocabulary of 128 k tokens drops from ~57 µs to ~7 µs, an 8× improvement that speeds token sampling for all LLMs.

Expanded linear‑algebra capabilities

• Complex tensors: LU factorization, LU solve, and triangular solve for `complex64` now run on MPS, unlocking gradients for complex‑valued networks.
• Small‑matrix kernels: Batches of 4×4‑8×8 FP32 matrices are inverted in a single Metal kernel; QR decomposition and triangular solves benefit from blocked updates that map efficiently to the GPU.
• SVD robustness: Edge‑case singular values and extreme magnitudes are now handled correctly, improving the reliability of pseudoinverse and least‑squares solutions.

Convolutions, padding, and training optimizations

• `conv1d` forward uses a specialized Metal kernel plus an im2col‑matmul path for short output sequences—ideal for audio and streaming models.
• Constant, reflection, and replication padding have been moved off MPSGraph onto Metal. This change alone reduces driver memory for tiled video VAE decoding from ~270 GiB to ~8 GiB.
• Dropout, upsampling backward, and channels‑last `conv2d` backward now follow fused or efficient Metal paths, cutting both latency and memory traffic.

More control with MPS streams

Developers can now create and select explicit MPS streams from Python via `torch.mps.Stream` and the `torch.mps.stream` context manager. Tensor lifetimes can be tracked across streams with `Tensor.record_stream`, enabling fine‑grained synchronization without sacrificing performance.

Practical takeaways

1. Upgrade requirements – macOS 15+ and Python 3.11+ are required. Install the release candidate with `pip install torch==2.15.0 --index-url https://download.pytorch.org/whl/test/cpu` in an isolated environment.
2. Enable FP8 – Use `torch.float8e4m3fn` tensors and `torch.nn.functional.scaledmm` to halve memory usage for weights and activations.
3. Boost LLM throughput – Compile your model with `torch.compile` and let the new MPS linear‑layer and MoE kernels do the heavy lifting; expect 5‑6× higher token rates on recent Apple Silicon.
4. Fit larger video pipelines – Switch to the Metal‑based padding kernels to keep driver memory under 10 GiB for tiled VAE decoders.
5. Leverage streams – For complex pipelines (e.g., simultaneous inference and preprocessing) create separate MPS streams to overlap work and reduce synchronization stalls.

Wrapping up

PyTorch 2.15 makes the Mac a far more competitive platform for both inference and training. By delivering FP8 support, dramatically faster LLM decoding, richer linear‑algebra primitives, and tighter memory usage, the release narrows the gap between Apple Silicon and traditional CUDA‑based systems. The upcoming 2.16 preview promises full‑Metal attention kernels and further MPSGraph‑to‑Metal migrations, so developers should start adopting 2.15 now to reap immediate performance gains and be ready for the next wave of improvements.