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)

[fenilsonani]

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
• My own modified scripts

Tasks

• An officially supported task in the examples folder
• 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):

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:

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:

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.

[Rocketknight1]

cc @SunMarc for quants