🚀 The feature, motivation and pitch

MPS currently has no equivalent of CUDA Graphs. Every PyTorch op on MPS goes through Python dispatch, MPSGraph compilation/validation, cache lookup, and command buffer encoding on each call. For dispatch-heavy models (deep transformers, small LLMs, activation chains), this CPU overhead is significant relative to actual GPU compute time.

I propose adding a graph capture/replay API to the MPS backend exposed through torch.mps.graph_capture() and torch.mps.graph_replay() that records MPS operations on the first pass and replays them in a single dispatch on subsequent passes, skipping all per-op CPU overhead.

The proposed API:

model.eval()
x = torch.randn(batch, seq, d_model, device="mps")

# Capture pass: ops run normally and get recorded
with torch.no_grad():
    with torch.mps.graph_capture():
        out = model(x)

# Replay: re-encode all ops in a single dispatch, no Python/compilation overhead
for batch_data in loader:
    x.copy_(batch_data)        # update input in-place
    torch.mps.graph_replay()   # replay recorded ops
    results.append(out.cpu())

Also available as a class-based API mirroring torch.cuda.CUDAGraph:

g = torch.mps.MPSGraph()
with torch.mps.graph(g):
    out = model(x)

x.copy_(new_input)
g.replay()

The implementation captures both MPSGraph ops (matmul, linear, etc.) and raw Metal kernel dispatches (elementwise, activations, etc.) through a centralized recording mechanism at MPSStream::commandEncoder(). A wrapper object MPSRecordingEncoder intercepts the 5 key Metal encoder methods (setComputePipelineState:, setBuffer:offset:atIndex:, setBytes:length:atIndex:, dispatchThreads:threadsPerThreadgroup:, dispatchThreadgroups:threadsPerThreadgroup:) and records kernel state while forwarding calls to the underlying encoder. This ensures all Metal kernels are captured automatically with no per-site recording code.

Additionally includes:

• Executable caching: compiled MPSGraphExecutable objects cached per graph to skip re-validation
• Buffer size validation on replay: asserts buffer sizes match what was recorded, catching tensor reallocation with a clear error instead of silent incorrectness

Constraints:

• Tensor shapes must not change between capture and replays
• Input data must be updated in-place via .copy_() before each replay
• MPS profiling must be disabled during capture

Using a 100-config benchmark suite (LLMs, diffusion, dispatch-heavy models) on Apple M4 Pro, 10 trials, 5 warmup, 10% trimmed mean:

Zero regressions across all 100 configs. Gains are concentrated on dispatch-heavy models where CPU overhead dominates GPU compute time.

Alternatives

Relying on eager execution without capture. For large GPU-bound models the dispatch overhead is negligible, but for dispatch-heavy workloads (small models, deep stacks, inference serving with many small batches) the CPU overhead adds up significantly.

Additional context

No response

cc @jhavukainen @jerryzh168 @kulinseth @malfet @DenisVieriu97 @aditvenk