Repository navigation
feat: int8 weight quantization with quantized matmul on the OpenXLA path (#449) #516
Description
Summary
The real NPU lever. Today the loader dequantizes MLX 4/8-bit checkpoints to f32, discarding the quantization. Instead, keep the weights quantized and emit a quantized matmul (dequant-in-graph or StableHLO quantized ops) so the bandwidth and the int8 compute survive into the graph, where an NPU's systolic int8 kernels (or Metal/CUDA int8) can use them. This is what makes the OpenXLA path viable on int8-native accelerators.
Scope
- Emitter: a quantized contraction path (per-tensor / per-group int8, matching the MLX-quantized checkpoint layout the loader already reads).
- Loader (
iree.rs): stop dequantizing to f32 for the quantized path; upload the quantized weights + scales/zeros and feed the quantized graph. - Keep f16/f32 paths intact; quantization is selected for quantized checkpoints.
Integration (required for done)
A real MLX-quantized checkpoint (e.g. an mlx-community 8-bit Llama) runs end-to-end through the quantized OpenXLA graph via the normal load path, on Metal and CPU, gated by the accuracy harness.
Validation
- Accuracy within the gate from S2 vs the f32 reference.
- tok/s and the memory-bandwidth reduction measured vs the f32 and f16 paths.
Depends on S1 and S2. Part of #513
Activity
docs: OpenXLA performance table and the transferable-precision decisi…
docs: OpenXLA performance table and the transferable-precision decisi…
docs: OpenXLA performance table and the transferable-precision decisi…
inureyes commentedon Jul 1, 2026
Deferred to a hardware-gated follow-up (2026-07-01).
Rationale (see epic #513 and ADR 0004): int8/int4 quantization is the NPU lever, but its payoff is memory bandwidth. A compute-bound Metal decode cannot demonstrate it (the f16 spike gave ~2x with f32 memory traffic, so Metal is not bandwidth-bound here), and measuring the payoff needs an actual NPU, which is not available. The transferable f16 win already landed in #514/#515.
Implementation, when picked up, requires three coupled changes:
bits-wide values from[out, in_packed]U32, applyq*scale + biaspergroup_size, then matmul).iree.rs): stop callingdequantize_affine(...) -> f32; upload the packed U32 + F16 scales + F16 biases as separate device buffers.csrc/xla_iree.c): the weight ABI is hard-wired to*const f32; make it per-weight-dtype.On Metal only token-exactness would be verifiable, not the speedup. Pick this up when an NPU plus its IREE HAL driver is available to validate the bandwidth payoff.
inureyes commentedon Jul 1, 2026
Picked up on a GB10 (Grace-Blackwell, CUDA via a source-built IREE runtime), the first int-native target available. Implemented in #568.
Done: the packed in-graph dequant path (
MLXCEL_XLA_QUANT=packed) keeps the MLX 4/8-bit weights resident packed (ui32+ f16 scales/biases) and dequantizes each weight in the StableHLO graph (bit-unpack, thenq*scale + bias), instead of dequantizing to f32 at load. The three coupled changes the deferral named all landed: the emitter dequant primitive, the per-dtype weight upload in the loader, and the per-weight-dtype C shim ABI.Finding (the hardware gate): on the GB10 the packed decode is token-exact with the f32/f16 path (the in-graph reconstruction is bit-identical to the host dequant) but about 4.3x slower (~1.6 vs ~6.7 tok/s), with GPU util rising 84% -> 96%. IREE's CUDA codegen does not fuse the unpack+dequant into the matmul: the decode step is ~678 dispatches and the reconstructed f32 weight is materialized to DRAM every step, so the graph pays more bandwidth + compute, not less. The aggressive-fusion / generalize-matmul / early-truncate-fusion flags leave the dispatch count unchanged (678 -> 677).
So the memory-bandwidth payoff is not realized by authoring the dequant in the portable graph alone; it needs the target to fuse dequant into the matmul (a quantized-matmul op, or an int8
dot_generallowering to the hardware int8 path). That is the same lever split the f16 profiling reached, now confirmed for int8. The packed path lands off by default as the correctness-verified foundation a fused quantized-matmul reuses.Follow-up: a fused quantized-matmul (int8 dot lowering, in IREE or on an NPU whose HAL driver lowers a quantized dot to its systolic int8 kernels) is what turns this into the bandwidth win. The bandwidth win is now gated on that fusion, not on hardware availability alone. Leaving this issue for the maintainer to close (implemented + gated) or keep open for the fused lever.
feat: int8/int4 packed weight quantization on the OpenXLA path (#516) (…
docs: OpenXLA performance table and the transferable-precision decisi…
feat: int8/int4 packed weight quantization on the OpenXLA path (lablu…