Block sparse MM MoEs

[angeloskath]

Implements all MoEs with the new mx.block_sparse_mm and mx.block_sparse_qmm (the latter is at ml-explore/mlx#1124 ).

One notable change is in the quantization predicate which accepts modules that define to_quantized() and delegates the check for the existence of quantized weights to the module using is_quantized(). This allows us to define SwitchMLP for each model and a quantized equivalent and auto convert to the quantized one.

Performance comparison at 4bits

Before

Qwen1.5-MoE-A2.7B-Chat     : 56.3 tps
Mixtral-8x7B-Instruct-v0.1 : 40.3 tps
phixtral-4x2_8             : 47.2 tps

After

Qwen1.5-MoE-A2.7B-Chat     : 121.7 tps
Mixtral-8x7B-Instruct-v0.1 :  59.3 tps
phixtral-4x2_8             :  92.8 tps

[angeloskath]

@awni I think this is ready for a review. I didn't implement the DoRASwitchLinear but it can easily be added in a later PR that also adds quantization to DoRA layers.

[awni]

Should we add use_dora to the condition here? Or throw if trying to DoRA-fy a switch linear layer (until we support Dora)?

[awni]

🚀 thanks for fixing the quantized models!