[magnuslundstedt]

• ☑️ I understand it is strictly prohibited to use AI to write PR description

• Implements the int8 half of mlx#4627 to stage the matmul's weights and activations as int8 instead of dequantizing to fp16; the fp8 half will follow in a separate PR. Enabled with the environment variable MLX_QMM_INT8, off by default; stock path unchanged.

• Validation lists: Speed improvements for quantized matmul. Measured on 36969c0, the PR is now rebased to 8329918.

Scenario	M5 10-core GPU	M6 12-core GPU
int8 prefill, N=8192	1.392x	1.349x
int8 prefill, N=4096	1.450x	1.413x
int8 prefill, N=2048	1.467x	1.420x
int8 prefill, N=512	1.531x	1.363x
int8 decode	0.999x, steps bit-identical	1.004x, steps bit-identical

(prefill M=2048, decode M=1)

• Exactness: with the switch off (default), every single-op cell is bit-identical to the unmodified build (12/12); decode steps are bit-identical on both hosts; the upstream test_quantized.py passes in every build, including the three new tests.

• Single-op table: Speed improvements (M = 2048, bf16 input, switch off over on, median of nine pairs; modes 1 / 2 / 4), as measured on base 36969c0.

shape	M6 mode 1 / 2 / 4	M5 mode 1 / 2 / 4
q/o 4096 x 4096	1.33x / 1.60x / 1.49x	1.33x / 1.68x / 1.56x
gate/up 4096 x 12288	1.31x / 1.63x / 1.52x	1.35x / 1.64x / 1.51x
down 12288 x 4096	1.49x / 1.85x / 1.73x	1.38x / 1.66x / 1.52x
thin N 5120 x 48	1.14x / 1.15x / 1.17x	1.28x / 1.27x / 1.28x

• Serving test using Qwen 27B (M5 Pro; lmstudio-community/Qwen3.8-27B-MLX-4bit, affine 4-bit g64, oMLX): prefill 1.44-1.51x at a measurable quality cost, depending on policy:

Policy (all 4-bit affine g64, head stock, >= 64 rows)	Cold prefill / stock speed	Quality, relative PPL, worst domain (literature)
Full staged int8 (attention mode 2; MLP/GDN mode 4)	1.436x geomean (1.408 / 1.490 / 1.411); 1.511x in a later session; 1.456x against the release wheel with native kernels	+2.08% [+1.74, +2.43]; science +1.53%, math +0.65%, code +0.10%
Attention stock (MLP/GDN mode 4)	1.466x	+0.85% / upper +1.13%
Attention + MLP stock (GDN mode 4)	1.106x	+0.31% / upper +0.48%

• Quality impact: the 27B worst-domain upper bound (full staged int8, +2.08%) fails the research's own +1% bar, literature and science degrade, math and code are near parity; task-level parity was not tested (no GSM8K or coding suite). For the smaller Qwen3-8B on the same path we measured similar prefill speedups, with a +0.56 / +0.76 / +0.59 / +0.18% relative PPL impact (worst upper 0.92%) on base commit. The mechanism is opt-in and off by default; per-layer policy lives in the serving harness, not in MLX.

• Note: the numbers are from before [Metal] Add staged affine_qmm_t_nax kernels with split-k #4682 which is now in main and makes M6 stock slower at large M, while flat/faster on M5. This is a separate issue, which I may report separately. This PR is not addressing that issue.

• Benchmark-script: qmm_8bit_bench.py (attached in comment to PR)

python qmm_8bit_bench.py --save-hashes h.json    # run on an unpatched build first
python qmm_8bit_bench.py --check-hashes h.json   # then on this PR

• Selector rules for new GEMM QMM path:

  • Environment variable MLX_QMM_INT8: default 0 (off); set to 1, 4 or 2 to enable, most to least precise; other values are ignored; read when the op is built.
  • quantized.cpp dispatch (affine, 2-D transposed weights, int8-capable hardware, K % 64 == 0, M >= 64; bits > 1, so 1-bit weights always use the stock path). Also K/64 <= 1024 (K ≤ 65,536), group sizes 32/64/128, non-batched. (the gate is a coverage choice, not a small-M gain claim)
  • In essence, these dispatch rules accelerate only affine-quantized linear layers during prefill on NAX hardware (M5/M6 today) while decode is unchanged.

• gather_qmm untouched

• SIMD fix (c8a1312) to remove undefined behavior for partial SIMD groups in affine_qmm_i8_quantize_x (padded threadgroup), sweep over 624 cells/host across 13 K values (9 not multiples of 32) are bit-identical; no arithmetic change; gates/tests unchanged

• test_quantized.py: 57 OK on stock build, 60 OK on head build. JIT build (MLX_METAL_JIT=ON) on the pushed head sweep bit-identical; JIT and non-JIT outputs agree.

• FP8: This PR adds only the int8 path. Staging QMM on fp8 instead of int8 (which additionally exploits the M6 fp8 lanes) will follow in a separate PR.

• AI usage disclosure:
  Written by me (human) and I understand it is strictly prohibited to use AI to write PR description. Agents did assist in providing the solution, performing the measurements and preparing the data tables.

[magnuslundstedt]

Benchmark script used for the single-op table. Save it as qmm_8bit_bench.py; run --save-hashes on an unpatched build, --check-hashes on this branch.

qmm_8bit_bench.py