[awni]

We need a more sophisticated policy to set ops per buffer based on the device. This is a start to that.

For inference on B200 it helps a lot for inference to increase it at very little memory cost.

mlx_lm.benchmark --model meta-llama/Meta-Llama-3.1-8B --p 128 -g 128 -b 1 -n 4

Pre: generation_tps=244.456, peak_memory=16.166
Post: generation_tps=283.073, peak_memory=16.224

For training 0.6B it's a double win, faster and less RAM 💪

	Toks/sec	Mem (GB)
Pre	60631	54.51
Post	64078	51.63

[angeloskath]

Beautiful 🚀

[awni]

This will probably need more tuning in the future especially for devices that I didn't add yet. But for now I think it's good to merge.