[Performance] PyTorch (MPS) is faster than MLX in backward of convolution layer

[arnold-yan]

Describe the bug
Recently I profiled the neural network layer performance from MLX and compared with PyTorch. I found that although MLX forwarding is consistently faster than PyTorch, in some chips (M1 Pro, M1 Max), PyTorch is much faster (3x~6x) for convolution forward + backward. While in some chips such as M3 Max, MLX is faster than PyTorch.

To Reproduce
To reproduce this, I have two minimal examples. The networks just have several convolution layers. You may try these two scripts to verify the performance.

time_pytorch_mlx.zip

[awni]

Same benchmark on an M2 Ultra

average time of Pytorch: 7.20261025428772
average time of MLX: 2.34059739112854

[abdussamettrkr]

On M2Pro

• average time of MLX: 4.472527980804443
• average time of Pytorch: 10.073836088180542

[alwint3r]

On M3 Max

• average time of MLX: 2.4813356399536133
• average time of PyTorch: 7.1081931591033936

[awni]

Thanks for the benchmarks everyone! There is clearly an unexpected performance cliff on M1 machines here as MLX is substantially faster on M2+. We'll need to take a deeper look at that to figure out where it's coming from.

[pyvadev]

On M1

• average time of MLX: 30.113215446472168
• average time of Pytorch: 15.948616743087769

[jrp2014]

M3 Max:
average time of MLX: 2.939736843109131
average time of Pytorch: 5.9829957485198975