mlx-quantization-metal-kernels: fix batched inputs in host dispatch

[fenilsonani]

Related issue

huggingface/transformers#47331 (batched generate with MetalConfig silently corrupts every sequence after the first). @SunMarc asked there for the fix to happen kernel side instead of a transformers workaround so this is that.

Closes # huggingface/transformers#47331

What does this PR do?

fixes batched inputs in the mlx-quantization-metal-kernels host dispatch. the set_strides helper passes the row stride of scales/biases where the kernel expects the batch stride so batch element z reads scales shifted by z rows and past the end of the tensor for high rows that's why row 0 is always correct and the rest is garbage. i verified this before changing any code by predicting the garbage output exactly with row shifted scales and it matched the actual kernel output to 0.4%.

Motivation

any batched generate through transformers MetalConfig on mps is silently wrong today and the fix belongs in the kernel package so every consumer gets it not just transformers. testing also surfaced that affine_qmv batched was separately broken (host passed ndim-1 batch dims but the kernel expects ndim-2 and reads M from x_shape[batch_ndim] which was out of bounds and rows are covered by tid.x which the grid never set) and that the batched grid.z path pays a full 32 row tile per batch element so [32,1,K] decode on a lm_head sized weight was 60ms vs 2.2ms as a flattened gemm.

Changes

• set_strides sends batch strides (zeros when scales/biases are shared) instead of row strides, in all three .mm files
• set_batch_dims sends the full shape/strides so qmv's M = x_shape[batch_ndim] read is in bounds
• qmv dispatches use batch dims = ndim-2 with M on grid.x and the batch on grid.z
• public ops collapse dense batched input into one 2d gemm (correct and ~Bx faster at M=1) and keep the strided path as fallback
• contiguity guard: kernels assume row contiguous matrix dims and nothing enforced it so strided M gave silent garbage before

Testing

• built quantized.metal with the same metal compiler version as the shipped binaries (32023.883) and ran all shapes 1d to 5d including sliced and transposed batch dims with bits 4/8 and fp16/bf16 against a dequantized reference on torch 2.10 and 2.13 all pass
• batch 1 output is bitwise identical to the shipped builds on both torch versions so no regression
• batch 4 generate end to end through transformers MetalConfig with no transformers side change every row matches batch 1
• perf median of 50 on M4 Max bits=4 bf16: mlp up [8,1,K] 0.83ms -> 0.24ms, lm_head [8,1,K] 15.2ms -> 2.1ms, lm_head [32,1,K] 60.1ms -> 2.2ms, prefill [1,256,K] unchanged
• added the batched cases to tests/test_quantization.py
• note: the nax and mxfp4 variants get the same set_strides fix but i only exercised the affine quantized.mm path locally the nax metallib needs your build setup

Checklist

• [X ] This PR is linked to an issue that was discussed and approved
• I have tested these changes locally
• New/changed functionality has test coverage
• LLM disclosure:
  • I did not use an LLM to create this PR.
  • I used and LLM for assistance while creating this PR.
  • This PR was mostly or completely generated by an LLM.

[github-actions]

Hi @fenilsonani, thanks for your interest in contributing!

This project requires that pull request authors are vouched, and you are not in the list of vouched users.

This PR will be closed automatically. See https://github.com/huggingface/kernels-community/blob/main/CONTRIBUTING.md for more details.