Normalize biases before encoding in gather_qmm_rhs

[dudududukim]

Proposed changes

Fixes #4055.

GatherQMM::eval_gpu normalizes w, scales and biases with
ensure_row_contiguous_matrix, which validates only the last two axes. A row
slice of an [E, 2R, D] expert weight passes that check while keeping the
leading stride of the original array, so the sorted path reads every expert
after the first from the wrong offset — silently, with no error.

The gather kernels index the leading axes themselves, so this switches those
three inputs to ensure_row_contiguous, which already exists in the file and
is already used for indices. x is left as is; the failure I can demonstrate
is in the weight path.

Verified on M5 Pro / macOS 26.5.1: the new test fails before the change and
passes after, and Ran 794 tests ... OK (skipped=46) for the full suite.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[zcbenz]

Thanks for the fix, I can verify the bug and the fix, but it is strange that gather_qmm_rhs actually uses ensure_row_contiguous inside it:

mlx/mlx/backend/metal/quantized.cpp

Lines 1612 to 1613 in f599c02

	array w = ensure_row_contiguous(w_, d, s);
	array scales = ensure_row_contiguous(scales_, d, s);

Do you have an idea what might go wrong?

[dudududukim]

Thank you for pointing that out — it led to the actual root cause.

The biases ensure_row_contiguous runs mid-encode: after
set_compute_pipeline_state and the first three set_input_array calls. For
a non-contiguous biases it encodes a copy kernel on the same encoder right
there, clobbering the state of the kernel being set up. w and scales are
normalized at the top of the function, so they were never affected — isolating
the inputs on 0.32.0 confirms only strided biases reproduces the bug
(w only: 0.0, scales only: 0.0, biases only: 1.41).

Updated the PR: the biases normalization is hoisted next to w/scales in
both gather_qmm_rhs and gather_qmm_rhs_nax, and the eval_gpu change is
reverted. Test unchanged, red before / green after, full test_quantized
passes on M5 Pro.

[zcbenz]

Nice fix!