PyTorch MPS backend replaces MSE forward MPSGraph with Metal kernels
PyTorch's Apple Silicon MPS backend now uses custom Metal kernels instead of MPSGraph for MSE loss forward, bringing up to 12x speedup for small tensors though some large tensor shapes regress.
Summary
A PyTorch commit migrates the MSE (mean squared error) loss forward pass on the MPS backend from Apple's MPSGraph framework to hand-written Metal kernels. The change aligns the implementation with CPU and CUDA backends by using an elementwise kernel followed by a separate reduction. Benchmarks show substantial speedups for small and medium tensors, with the best case of 12.35x for a 1×16×8×8 channels-last bf16 tensor. However, several larger tensor configurations show slowdowns, including a 0.20x regression for a 32×512×768 target-expanded bf16 mean reduction. torch.compile performance is unaffected.
What changed in MSE forward for MPS
- Implementation shift: MPSGraph-based MSE forward replaced with two-step Metal kernel (elementwise op + separate reduction)
- CPU/CUDA parity: New structure matches CPU and CUDA implementations, improving cross-backend consistency
- torch.compile unaffected: Compiled graphs continue using existing optimized paths; only eager mode sees changes
- Performance profile: Best gains on small tensors (1×16×8×8 shapes, up to 12.35×); mixed results on 8×64×64×64 shapes (1.25× to 0.74×); regressions on very large or sliced tensors (0.20× worst case for 32×512×768 target-expanded bf16 mean)
- Data type coverage: Both fp32 and bf16 layouts supported across contiguous, channels-last, strided, and sliced tensor configurations
Why it matters
Apple Silicon Mac users running PyTorch training or inference in eager mode will see faster MSE loss computation for typical small-to-medium tensor shapes common in model development and debugging. The regressions on large tensors suggest the new kernels may need further tuning for production-scale workloads.
Original reference: github.com →