MetalConfig quantization: batched generate silently corrupts all rows after the first — affine_qmm_t computes only batch element 0 of 3D inputs (verified, one-line fix)
Author: fenilsonani
Pull request: no

### System Info

```
- `transformers` version: 5.13.1
- Platform: macOS-26.5.2-arm64-arm-64bit
- Python version: 3.11.15
- Huggingface_hub version: 1.23.0
- Safetensors version: 0.8.0
- Accelerate version: 1.14.0
- Accelerate config: not found
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.13.0 (MPS)
- kernels version: 0.15.2
- Device: Apple M4 Max, 36 GB, MPS backend
- Using distributed or parallel set-up in script?: No
```

### Who can help?

@MekkCyber @vasqu

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder
- [x] My own task or dataset (give details below)

### Reproduction

Models quantized with `MetalConfig` (bits=4 and bits=8 identical) generate correctly at
batch size 1, but at batch size > 1 only the first row is correct; every later row's
argmax collapses to token 0 (decoding as `!!!!...`). Deterministic — repeated runs are
bit-identical.

**Model-level repro** (greedy, same prompt duplicated — every row should match batch 1):

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, MetalConfig

MODEL = "Qwen/Qwen3-0.6B"
tok = AutoTokenizer.from_pretrained(MODEL)

def generate(model, batch):
    ids = tok(["The capital of France is"] * batch, return_tensors="pt").to("mps")
    with torch.inference_mode():
        out = model.generate(**ids, max_new_tokens=20, do_sample=False)
    return [tok.decode(o[ids["input_ids"].shape[1]:], skip_special_tokens=True) for o in out]

q = AutoModelForCausalLM.from_pretrained(
    MODEL, quantization_config=MetalConfig(bits=8), dtype=torch.bfloat16
).eval().to("mps")
print(generate(q, 2))
# ['" Paris. The capital of Italy is Rome. ..."', '"!!!!!!!!!!!!!!!!!!!!"']
# row 0 == batch-1 output exactly; row 1 garbage. Unquantized bf16: all rows correct.
```

**Root cause (verified at kernel level, no model involved).** `MetalLinear.forward`
passes its input to `affine_qmm_t` as-is; during HF decode that input is 3D
`[batch, 1, hidden]`. Calling the kernel directly with a random quantized weight and
checking each output row against `x @ dequant(W).T`:

| input shape | result |
|---|---|
| `[B, K]` 2D, any B | all rows correct |
| `[1, S, K]` any S | all rows correct |
| `[B, 1, K]`, B > 1 | **batch element 0 correct, elements 1+ garbage** |
| `[2, 4, K]` | **rows 0–3 (batch elem 0) correct, rows 4–7 (elem 1) garbage** |

Identical pattern for bits=4/8 × bf16/fp16. The `[2, 4, K]` case pins it: the kernel
computes the **entire first batch element** and never touches the rest — it treats the
input as 2D `[shape[-2], K]`, ignoring the leading batch dimension. The garbage rows are
bit-identical across repeated calls. Probe script:

```python
import torch
from transformers.integrations.metal_quantization import (
    _affine_quantize_tensor, _affine_dequantize_tensor, _get_metal_kernel)

kernel = _get_metal_kernel()
torch.manual_seed(0)
K, N, GROUP, BITS = 2048, 1024, 64, 4
wp, sc, bi = _affine_quantize_tensor(torch.randn(N, K), GROUP, BITS)
w_deq = _affine_dequantize_tensor(wp, sc, bi, GROUP, BITS).to("mps")
wp, sc, bi = wp.to("mps"), sc.to("mps").bfloat16(), bi.to("mps").bfloat16()

for shape in [(2, K), (2, 1, K), (2, 4, K)]:
    x = torch.randn(*shape, device="mps", dtype=torch.bfloat16)
    y = kernel.affine_qmm_t(x, wp, sc, bi, GROUP, BITS)
    ref = (x.reshape(-1, K).float() @ w_deq.t()).reshape(*shape[:-1], N)
    err = ((y.float() - ref).abs().amax(-1) / ref.abs().amax(-1)).flatten()
    print(shape, [f"{e:.3f}" for e in err.tolist()])
# (2, 2048)       ['0.005', '0.006']            <- 2D fine
# (2, 1, 2048)    ['0.006', '9.874']            <- 3D: only batch elem 0 computed
# (2, 4, 2048)    ['0.005', ..., '9.9e+30', ...] <- rows 0-3 fine, 4-7 garbage
```

**Fix (verified).** Flattening 3D inputs to 2D around the kernel call makes batched
generate fully correct — with this patch all batch-4 rows match the batch-1 output
exactly:

```python
def forward(self, input):                      # MetalLinear.forward
    if input.dim() == 3:
        b, s, k = input.shape
        return _orig_forward(self, input.reshape(b * s, k)).reshape(b, s, -1)
    return _orig_forward(self, input)
```

Either `MetalLinear.forward` should flatten to 2D before calling `affine_qmm_t`
(one line, no kernel change needed — happy to open a PR), or the kernel should handle
the batch dimension of 3D inputs.

**Possibly relevant:** the kernel repo ships builds for torch 2.8 / 2.9 / 2.10 only
(`build/torch{28,29,210}-metal-aarch64-darwin`); under torch 2.13 the `kernels` loader
selects the `torch210` build. Reproduced on that combination.

**Impact:** any batched `generate()` with a Metal-quantized model on MPS silently
corrupts all sequences after the first — no error, no warning. Found while benchmarking
nvidia/canary-qwen-2.5b (speech LLM): int4 batch-32 transcription WER went from 2.1%
(batch 1) to 98% (batch 32). Batch 1 is unaffected.

### Expected behavior

Every row of a batched `generate()` should produce the same output as batch size 1 for
the same (duplicated, greedy-decoded) prompt — as batch element 0 already does, as the
2D-flatten patch produces, and as the unquantized bf16 model does for all rows.


Comments:
cc @sunmarc for quants

Hi! I'd like to work on this issue if it's still available.
I'll investigate the Quark API changes, verify compatibility across supported versions, and prepare a fix with tests if appropriate.
If no one is already working on it, I'll open a draft PR once I have a working implementation.