[BUG] quantizedMM silently upcasts to float32 due to dtype mismatch between activations and scales/biases

[czkoko]

Describe the bug
Recently, I used mlx-swift to run the image model. The model needs bfloat16 accuracy to work normally. When using the unquantified model, it works normally and the speed is normal. However, when using the pre-quantified model, the running speed begins to decline. Finally, it is found that the calculation accuracy has become float32.

When using quantized models in mlx-swift, if there is a mismatch between the dtype of the input activations (e.g., bfloat16) and the quantization parameters scales/biases (e.g., float16), MLX silently promotes the calculation to float32.
Even if I turn the quantized weight to bfloat16, I can't solve the problem, and the speed will be more slower than float32.
If float16 accuracy is used throughout the whole process, the calculation will not be upgraded to float32. However, too much loss of accuracy of the image model causes a black screen.

To Reproduce

Include code snippet

import MLX

// 1. Setup inputs with mismatched dtypes
let x = MLXArray.zeros([1, 4096], dtype: .bfloat16)        // Input activations (bf16)
let w = MLXArray.zeros([512, 512], dtype: .uint32)         // Packed weights (4-bit, 4096/8=512 in_features, 512 out_features)

// Quantization parameters (fp16)
let scales = MLXArray.zeros([512, 64], dtype: .float16)    // group_size = 64 (4096/64 = 64)
let biases = MLXArray.zeros([512, 64], dtype: .float16)    // biases must be provided for affine mode

// 2. Perform quantized matrix multiplication
// This will silently upcast the execution to .float32 instead of running in bf16/fp16
let result = MLX.quantizedMM(
    x, 
    w, 
    scales: scales, 
    biases: biases, 
    transpose: true, 
    groupSize: 64, 
    bits: 4
)

print("Result dtype:", result.dtype) // Prints: float32