fix: propagate quantization mode in QuantizedAllToShardedLinear / QuantizedShardedToAllLinear

[vskiwi]

Summary

Fixes #3132

QuantizedAllToShardedLinear and QuantizedShardedToAllLinear in mlx/nn/layers/distributed.py do not accept, store, or pass the mode parameter to mx.quantized_matmul. When an MXFP8-quantized QuantizedLinear is converted via shard_linear(), the mode is silently lost. The resulting sharded layer calls quantized_matmul without mode=, which defaults to "affine" — interpreting FP8 packed weights as affine int8, producing garbage output with no error.

Additionally, MXFP8 does not use biases, but both classes unconditionally accessed self["biases"], which would raise ValueError once the mode fix is applied.

Changes

• Add mode parameter (default "affine") to both __init__ methods
• Store self.mode and pass it to mx.quantize and mx.quantized_matmul
• Use *biases unpacking to handle modes that don't produce biases (mxfp8, mxfp4)
• Use self.get("biases") instead of self["biases"] for safe access (consistent with QuantizedLinear)
• Propagate mode from source layer in from_quantized_linear
• Include mode in _extra_repr output
• Add distributed test for mxfp8 quantized shard_linear

Impact

This unblocks tensor parallel inference for all MXFP8-quantized models (and likely mxfp4). Confirmed working: GLM-5 754B (mlx-community/GLM-5-8bit-MXFP8, mode=mxfp8, group_size=32, bits=8) on 2× M3 Ultra 512GB at ~14 tok/s with tensor parallel.

No changes to the affine (default) code path — full backward compatibility.

Test plan

• Existing test_shard_linear test for affine quantization is unchanged and should still pass
• New mxfp8 test in test_shard_linear verifies mode propagation, biases=None, and output correctness
• Formatted with black

Made with Cursor

[angeloskath]

Thank you that looks great!

I'll merge after the tests pass.