MPS backend

Active contributors: malfet

Purpose

The MPS (Metal Performance Shaders) backend lets PyTorch use the Apple Silicon GPU. Like the CUDA backend, it spans low-level runtime helpers, op kernels, and a Python module (torch.mps).

The implementation mixes Objective-C++ (.mm files) for Metal API access with C++/Python for the rest of the stack.

Directory layout

Path	Contents
aten/src/ATen/mps/	Runtime helpers (allocator, streams, profiling)
aten/src/ATen/mps/MPSAllocator.mm	MPS caching allocator
aten/src/ATen/mps/MPSStream.mm	Stream / command queue
aten/src/ATen/native/mps/	MPS op kernels (.mm files using MPSGraph + custom Metal shaders)
aten/src/ATen/native/mps/operations/	Per-op kernels
c10/metal/	Headers for Metal kernels
torch/mps/	Python torch.mps package
torch/csrc/mps/	Python bindings for torch.mps
test/test_mps.py	The MPS-specific test file

Key abstractions

Type	File	Purpose
at::mps::MPSDevice	aten/src/ATen/mps/MPSDevice.h	Singleton Metal device
at::mps::MPSStream	aten/src/ATen/mps/MPSStream.h	Wraps MTLCommandQueue + MTLCommandBuffer
at::mps::MPSAllocator	aten/src/ATen/mps/MPSAllocator.mm	Caching allocator on top of MTLBuffer
at::mps::MPSEvent	aten/src/ATen/mps/MPSEvent.h	Synchronization event
at::mps::MPSProfiler	aten/src/ATen/mps/MPSProfiler.h	Capture for Instruments / Metal frame capture

How it works

Op implementation styles

There are two styles of MPS op:

1. MPSGraph-based — most ops use Apple's high-level MPSGraph API. The kernel author builds a graph of MPSGraphTensors, compiles it once (cached), and runs it on the input tensors. Cleaner code but slightly more overhead per call.
2. Custom Metal shaders — for performance-critical or unsupported ops, kernels are written as .metal shader files (under aten/src/ATen/native/mps/kernels/ and similar) and dispatched via MTLComputeCommandEncoder. The shader source is embedded at build time.

Allocator

MPSAllocator (~3K lines) implements a caching allocator over MTLBuffers with semantics analogous to the CUDA caching allocator: free-list buckets, stream-aware deferred frees, configurable via PYTORCH_MPS_HIGH_WATERMARK_RATIO and friends. Because Apple Silicon has unified memory, allocations come from a shared system pool.

Streams

The single MPSStream per device wraps a MTLCommandQueue. Multiple "streams" are emulated via separate MTLCommandBuffers; PyTorch tracks which buffers are in flight. Synchronization with CPU work uses MPSEvent (a MTLSharedEvent).

Quirks vs. CUDA

• Apple GPUs don't have float64. Op kernels for double inputs raise rather than silently downgrading.
• Some dtypes (bfloat16) only became supported in macOS 14.
• Inductor's MPS backend is newer than CUDA's; some torch.compile features (e.g., reduction tile sizes) have separate code paths.

Integration points

• Dispatcher. MPS uses DispatchKey::MPS and DispatchKey::AutogradMPS.
• Autocast. aten/src/ATen/autocast_mode.cpp registers AutocastMPS.
• Inductor. A dedicated GPU codegen path that emits Metal shaders rather than Triton.

Entry points for modification

• New MPS op → write an .mm under aten/src/ATen/native/mps/operations/, register in native_functions.yaml under MPS:. Use MPSGraph if Apple already provides the op; drop to a custom Metal kernel for control.
• Allocator behaviour → MPSAllocator.mm.
• For debugging, MPSProfiler produces a Metal frame capture that opens in Instruments.

Key source files

File	Purpose
aten/src/ATen/mps/MPSAllocator.mm	Caching allocator
aten/src/ATen/mps/MPSStream.mm	Stream / command queue
aten/src/ATen/native/mps/operations/	Per-op kernels
c10/metal/	Shared Metal headers
torch/mps/__init__.py	Python torch.mps
torch/csrc/mps/Module.cpp	C++ binding