Speeding up PyTorch inference on Apple devices with AI-generated Metal kernels

They are comparing unoptimized PyTorch inference, something you would never deploy on a device, to a model with custom kernels.

Yes, of course the model with custom kernels is faster, whether it's written by a human or an AI.

Generally, PyTorch inference is meant to be used during the training process, and when running metrics, not when deploying. When deployed, you should export to ONNX, and then compile the ONNX to the native format of the device.

If you aren't familiar with the pipeline for ML deployment, this is the equivalent of comparing interpreted code to compiled code.

PyTorch is the baseline because that's what people prototype in, and the most common reference point. The aim here is to show that you can start from prototype code and automatically produce lower-level kernels (in this case Metal) that are more usable in real deployments, without additional work from the developer. Frontier models are capable at generating efficient Metal kernels automatically/immediately, and will only get better. We expect to see significant improvements as we refine the approach, but it's enough to show this seems to be a tractable problem for AI.

I have never really worked with pytorch professionally, but it feels to me a lot of the open source, especially generative oriented projects, just use pytorch like this. It makes hacking on the models a whole lot easier.

comfyui is a good example of a project like this.

> and then compile the ONNX to the native format of the device.

I'm assuming you are talking about https://github.com/onnx/onnx-mlir?

In your experience, how much faster is a "compiled" onnx model vs. using an onnx runtime?

For other people reading this:

Back in the day TensorFlow had tfdeploy which compiled TensorFlow terms into NumPy matrix operations. Our synthetic tests saw speedups of factor 50.

> Yes, of course the model with custom kernels is faster, whether it's written by a human or an AI.

But that’s the thing, I wouldn’t write a custom kernel before AI

I don't do that level of development or operate at that part of the stack but I’m very experienced in software development

AI significantly augments my skillsets in this area

The point is those kernels exist already, you can just use them off the shelf. In the case where you're trying to write a production grade kernel without operating at that part of the stack... well good luck with that.

vLLM is a LLM model serving framework written using raw PyTorch.

ONNX doesn’t support a bunch of operations that PyTorch does (it isn’t always possible to convert a PyTorch model to ONNX).

Torchserve runs raw PyTorch.

Generally speaking, PyTorch is pretty well optimized. For Mac it has been historically ignored, so the kernels for MPS were all missing or just bad, but on CUDA and Linux they are pretty good.