Repository navigation
qmm: stage dequantized weights as int8/fp8 instead of fp16 to use the 8-bit matrix lanes #4627
Description
Summary. In MLX 0.32.3 every quantized matmul path (affine 2-8 bit, mxfp4, mxfp8, nvfp4) dequantizes packed weights into threadgroup memory and then runs an fp16/bf16 GEMM (QuantizedBlockLoader::load_unsafe, quantized.h:485-569, consumed by matmul2d in the NAX path quantized_nax.h:921-1054). The 8-bit compute lanes the hardware exposes are never used, so quantized prefill is capped at the fp16 lane.
Measured here on an Apple M6 (12 GPU cores), MPP matmul2d fed from device memory, 2048x4096x4096, best tile, MLX 0.32.3:
| lane | rate | vs fp16 |
|---|---|---|
| fp16 x fp16 -> f32 | 13.6 TF/s | 1.00x |
| int8 x int8 -> i32 | 26.9 TOPS | 1.98x |
| fp8 e4m3 (plain) | 28.7-30.6 TOPS | 2.11-2.25x |
| stock 4-bit QMM kernel | 18.9 TF/s | (dequant -> fp16) |
Proposal. Add a staged-8-bit variant of the qmm/NAX prefill path: keep the decode-into-threadgroup structure, but stage int8 (threadgroup uchar, i8 x i8 -> i32) or fp8 for 8-bit modes, and invoke matmul2d with the matching 8-bit operand type. Activations are dynamically per-row quantized to the same type.
Headroom measured with a standalone probe (staged-int8 kernel vs the stock 4-bit QMM kernel, same shapes): 1.53x at 2048x4096x4096 and 1.30x at 8192x16384x4096 on M6; ~2.1x on M5 Air and M5 Max. Honest caveat: stock q4 on M6 already runs at ~95% of its fp16 lane, so the M6 ceiling for a staged variant is ~1.42x (26.9 vs 18.9), not 2x. At model level the headroom is confirmed: a plain per-row-scaled FP8 arm (custom kernels, Qwen3-8B) measures 1.72x over bf16 / ~1.29x over stock 4-bit prefill at N=2048 with +0.71% worst-domain PPL against bf16.
Scale handling (affine). For bits <= 6 there are at most 64 distinct levels per group, so staging to int8 with one per-group scale (v8 = round((scale*q + bias)/s8)*s8) is near-lossless (<1% of group range, coherent within the group) and needs no epilogue. If a per-group epilogue is preferred, group sizes {32,64,128} already align with the qmm BK=64 tiling.
Note on the MX modes (mxfp4/mxfp8/nvfp4). These also decode to half today. MPP has a native UE8M0 scale-plane path, but measured per GPU core it runs at the prior generation's fp8 rate while the plain fp8 lane roughly doubled on this generation: plain fp8 2.55 vs MXFP8 1.30 TOPS/core (same harness; M5 Max 1.37 vs 1.19; M5 Air 1.41 vs 1.07). Until that driver-side gap is addressed (reported separately to Apple), a software-scaled or plain per-tensor fp8 path is the faster option on this generation.
Ask. Would a staged-8-bit qmm/NAX path be something the project would consider? We have a validated kernel pattern (threadgroup staging + matmul2d i8xi8->i32 with a fused per-row x per-N scale epilogue), correctness gates (bit-exact integer constants, wrong-packing negative control), and both microbenchmark and whole-model numbers ready. Happy to contribute a patch or share the probe scripts for reproduction first.
Activity
magnuslundstedt commentedon Oct 7, 2026
Following up on this issue with code and completed measurements (about +40% prefill speed on M6).
Env used:
64ea011) + the int8 patch; mlx-lm 0.32.0; Apple M6 (12 GPU cores) and M5 Air (10 GPU cores); macOS 27.0 (26A428) / macOS 27.0 beta; Xcode 27 beta.mainat264c14fe6+ the fp8 patch, deployment target 27.0 (Metal 4.1); mlx-lm 0.32.0; Apple M6 only; macOS 27.0 (26A428); Xcode 27 beta.MLX_ENABLE_TF32=0(see the TF32 note).staged-8bit-qmm, one commit against currentmain(53fbe5667): precisit/mlx@main...staged-8bit-qmmBottom line: staging packed weights into eight-bit compute operands raises quantized prefill by 1.39x (int8, M6) and 1.37-1.42x (fp8, MX weights, M6) at a worst-domain perplexity cost of +0.92% and +0.71/+0.78% against each artifact's own stock path. Decode is unchanged. The first integrated fp8 attempt failed on how it stored results, not on how it multiplied; the diagnosis and fix are below.
What we measured
int8 staging over 4-bit affine group-64 weights (Qwen3-8B, converted with stock mlx-lm): the quantized prefill matmuls run
i8xi8->i32with per-group scales applied in the epilogue, instead of the stock path's dequantize-to-bf16.fp8 staging over mxfp8 and mxfp4 weights: activations are quantized to E4M3 with one per-row power-of-two scale, weights are moved to their row's scale (exact for mxfp4; at most 0.012% of weights change for mxfp8, within 3.3e-7), and the matmul is
fp8 x fp8on memory operands with a float cooperative destination.MLX_ENABLE_TF32All staged runs used the pinned
TF32=0. We then ran MLX's default (TF32=1) on the same build and nothing moves: logits hashes are identical for all six arm/artifact combinations, per-chunk loss sums are equal as floats, and each arm against itself across the settings is within +/-0.5% (twelve cells). The flag only selects the NAX kernel for float32 operands, and no float32 GEMM executes in these graphs (all 253 projections are bf16 in/out; the perplexity head is logsumexp, not a GEMM). So the numbers hold for a default user.Two side notes from that check, offered rather than asked: (1) the staged paths sit ahead of the stock dispatch and do not consult the flag, so with float32 inputs and
TF32=0an opted-in staged path is still taken - see question 3; (2) on our M6,test_quantized.pyfails under the library default on unmodifiedmaintoo (upstream's own runner pinsTF32=0); happy to file that separately if useful.Scope
One model (Qwen3-8B); the fp8 arm was measured on M6 only (the fp8 lane is ~1.05x on M5, so there is no lane to stage into there); a per-row-tile dispatch for ragged chunks is described but not built. There is no 8-bit kernel on GPUs without the matrix unit; those and non-NAX shapes stay on the stock path, and with the switch off the build is stock to the bit. mxfp4 carries half the bytes of mxfp8, and its stock decode is about twice as fast in these runs (35.3 vs 18.9 tok/s); which artifact to use is a model and deployment choice.
Code and provenance: the branch contains only the change, its docs and its tests, and it is the tree these numbers were validated on (the staged outputs match the scored build on both GPUs). A precompiled Metal 4.1 library needs deployment target 27; a run-time-JIT build works on macOS 27 as well.
Questions
compute_8bit_, 1/2/4 for the int8 modes and 8 for fp8); what shape would the maintainers prefer?MLX_ENABLE_TF32: shouldTF32=0veto staging for float32 inputs? That is a semantics call for the maintainers.AI usage disclosure: the measurements come from an automated harness; this text was drafted with a coding agent and posted by the reporter.
magnuslundstedt commentedon Oct 9, 2026
The public write-up with the full methods and measurement receipts is now up: precisit.com/en/blog/apple-matrix-formats/
It includes the paired gates, quality bounds, run identities and a small dispatch check, and the same experimental code.
Also see the questions in my previous comment (interest, API shape for the compute selector, TF32 semantics). After more thought on the selector: int8 and fp8 are really one choice per operation, since affine weights only take int8 and mxfp4/mxfp8 only take fp8. So a single per-op selector (off, int8 variants, fp8) may fit better than the two switches I used. Happy to follow whatever is considered best practice.
One more design question we tested before proposing anything: could activations stay fp16 while only the weights are staged? They cannot reach the fast lane. On M6, fp16 x fp8 reached about 1.26x the fp16 lane without weight scales and 0.82x with them, while fp8 x fp8 is the 2.05x lane. So both operands are staged, and that is where the small quality cost lives: worst-case perplexity increase below 1% on the tested 8B model, for roughly 40% faster prefill. Similar in spirit to the TF32 trade-off, though this one is opt-in and off by default, and the bounds only cover the models we measured.