atlas: measured-cost quantization for MLX — give it a RAM budget, it finds the optimal (bit-width, group-size) plan

[Matth21]

I built a quantizer for MLX that replaces heuristics with measurement.

How it works: it fake-quantizes one transformer block at a time, measures KL divergence on the output logits over a fixed calibration set, and builds a cost table for every (bits, group_size) config — bits ∈ {3,4,5,6} × gs ∈ {32,64,128}. A validated additivity assumption turns budget-constrained allocation into an exact per-block Lagrangian rule, so one overnight profiling run gives you the entire 3–5 bit/w Pareto frontier. After that, any budget is milliseconds.

The product gesture:

atlas Qwen/Qwen2.5-7B-Instruct --budget-gb 4

→ best quality that fits in 4 GB. Pre-computed cost tables for TinyLlama-1.1B and Qwen2.5-7B ship in the repo, so first run takes minutes.

Results (disk size / total params, full Wikitext-2 sliding window, both TinyLlama-1.1B and Qwen2.5-7B):

• Strictly dominates uniform MLX quantization: e.g. +18.7% PPL at 3.62 bit/w vs +36% at 3.93 for the best uniform setting on Qwen
• Matches or beats llama.cpp K-quants at ≥4.4 bit/w; loses to Q3_K_M below 4 bits — that gap is the affine format vs K-quant superblocks, not the allocation. mxfp4 support is my next step to close it.

Honest footnote: v1 of this work used an entropy heuristic. Controlled baselines (random + inverted rankings) showed the entropy signal was exactly zero — the paper's Section 6 documents the correction openly.

Repo (code, tests, all experimental data): https://github.com/Matth21/atlas
Paper: https://doi.org/10.5281/zenodo.21190586

Feedback very welcome — especially from anyone who knows the quantized-matmul kernel internals and has opinions on richer low-bit formats in MLX.