MTP head crashes with "Biases must be provided for affine quantization" when the target model's lm_head is mxfp8-quantized

[usenetj23]

Summary

headQuantParams() in src/mtp.zig calls transformer_mod.affineParamsFromGeometry() before checking the target model's actual quant_mode. That function infers bits/group_size purely from tensor shape and always returns mode = .affine when the shape math resolves to a plausible (bits, group_size) pair — it never inspects the scales tensor's dtype or the model config's declared quantization mode.

For a target model whose lm_head is genuinely mxfp8-quantized (scales dtype uint8, no .biases tensor), the shape math can coincidentally match a valid affine configuration. Concretely, for lm_head.weight shape [vocab, hidden*8/32] (U32-packed) and lm_head.scales shape [vocab, hidden/32] (U8), with hidden=5120:

• bits = (w_cols * 32) / in_dim = 8
• gs = in_dim / s_cols = 32

Both land in affineParamsFromGeometry's accepted bits/gs switch cases, so it returns {bits: 8, group_size: 32, mode: .affine} — a false positive. headQuantParams then uses this result directly (src/mtp.zig around line 1050-1061) without ever falling through to the config-driven branch that would have correctly returned mode: .mxfp8.

targetLmHead() then calls mlx_quantized_matmul(..., qp.mode.cstr(), ...) with mode="affine" against a tensor that has no .biases (mxfp8 has none by design — see the comment at the top of model.zig: "non-affine modes store NO .biases tensors"). This throws:

MLX error: [dequantize] Biases must be provided for affine quantization. at .../mlx-c/mlx/c/ops.cpp:1106

Reproduction

• Model: dealignai/Qwen3.6-27B-MXFP8-CRACK-MTP (base: Qwen/Qwen3.6-27B), MXFP8 8-bit/gs32 quantization, MTP weights embedded in the main safetensors shards (not shipped as a sidecar).
• Built a native sidecar with tests/build_mtp_sidecar.py Qwen/Qwen3.6-27B <out> (unmodified, against the original upstream bf16 repo) — this step works fine, sidecar loads cleanly ([mtp] loaded native MTP head (dense-mlp; ...) logs with no error).
• Launch: mlx-serve --model <dir-with-mtp-sidecar> --serve --mtp
• First chat completion request crashes with the error above, immediately after [mtp] loaded native MTP head — i.e. during the first MTP forward pass, not during sidecar loading.

Suggested fix

headQuantParams() (and possibly other affineParamsFromGeometry() call sites operating on target-model tensors, e.g. qLinearFwd's geometry probe at src/mtp.zig:908 — need to confirm whether that path is sidecar-only or can also touch target tensors) should check the target's actual config.quant_mode / the scales tensor's dtype (uint8 → non-affine) before trusting a geometry-only affine match, mirroring what computeQuantParams() already does correctly (dtype-uint8 check at transformer.zig:8815 gates the mxfp8 branch). A one-line guard — skip the affineParamsFromGeometry shortcut when mlx_array_dtype(scales) == .uint8 and config.quant_mode != .affine — should resolve it.

Environment

• mlx-serve 26.7.8 (MLX 0.31.2)
• macOS, Apple Silicon