Offer to help improve the MPS backend fallbacks

[egeozkoc]

I'd like to help push the MPS backend off the 🐢 status in the README, and I'm set up to work on it directly (Apple Silicon, can build and test locally). I've contributed on the MPS side before (#1994) and on optimizers (#1998), so I'm somewhat familiar with the layout.

From reading the current backend (#1983, #2004), my understanding of the design is: try the kernels-community/bitsandbytes-mps Hub kernels when available, and fall back to pure PyTorch otherwise, with those fallbacks being what the 🐢 reflects. I don't want to cut across your in-flight work, so rather than guess, I wanted to ask where contributions would be most useful.

A few directions I could take on, if any are helpful:

• Profiling and speeding up specific PyTorch fallbacks in backends/mps/ops.py, in the spirit of perf: improve default/fallback backend implementation for blockwise quantization ops #1960.
• Filling in MPS-specific fallbacks for ops that currently only hit the generic slow path.
• Extending MPS test coverage and correctness for the fallback paths.

If the 8-bit optimizer path (the 🚧) is something you'd want help with, I'm glad to look at that too, though I'd want to check it fits your plans first. Let me know what would be most useful, or if you'd rather I stay in the smaller correctness and coverage lane for now.

[matthewdouglas]

Hi @egeozkoc,

Thanks for asking!

The README's 🐢 is a little out of date. With the newer Hub kernels the 4bit quantize/dequantize/gemv ops are in a much better place. I plan to release what we have in v0.50.0 very soon, so will likely be updating this in the README even if best performance needs macOS 26 + kernels.

I do have a plan already for the optimizers, which will be part of a future v0.51.0 release. I would ask for now that you hold off on that. Correctness is going to come first, then further perf if needed. I'll be implementing it similar to the CPU backend, which will rely on the separate quantize_blockwise, dequantize_blockwise ops, and do the rest in PyTorch. I think this will be good enough to start before jumping further to e.g. a fully fused kernel like we have on the CUDA side.

I did spend some time benchmarking as part of #1960 and #1983. But if you want to share any benchmarks, profile data, etc, then that could be useful. Especially if you've got different hardware than I do (M4, 10-core).

One obvious area that underperforms is the LLM.int8() ops. It's not exactly a priority, but there's opportunity there I'm sure. Another op that I know has room for improvement is quantize_blockwise.

I have not quite figured out exactly how I want to bring additional Metal kernels to the project in the future. I'm leaning on using torch.mps.compile_shader for this and doing JIT. But I want to avoid doing things like linking to libtorch without a stable ABI (makes build/dist more complex) or taking performance hits from not being able to access PyTorch's command buffer. The compile_shader approach avoids all of that. The other direction of course is to publish to the Hub separately, but right now it's kind of a stop-gap; we don't actually participate in maintaining anything in kernels-community at the moment.

With that said, if there's correctness fixes or similar well-isolated improvements to make, I'd be welcoming of that.