Add 1-bit affine quantization support (Metal)

[khosravipasha]

Add 1-bit affine quantization support (Metal)

Proposed changes

This PR adds 1-bit support to MLX's affine quantization mode, extending the supported bit-widths from {2, 3, 4, 5, 6, 8} to {1, 2, 3, 4, 5, 6, 8}.

MLX already supports affine quantization at 2, 3, 4, 5, 6, and 8 bits via w_hat = scale * w_q + bias. This PR extends that same framework to 1-bit, adding full kernel support for 1-bit affine dequantization and quantized matmul across CPU and Metal backends.

This assumes the model has already been quantized externally (e.g. during training) — the contribution here is efficient packing and inference on Apple Silicon. It supports packing for both affine and symmetric 1-bit weights:

Affine 1-bit — weights have arbitrary per-group min/max:

scale = w_max - w_min,  bias = w_min
bit 0 → w_min,  bit 1 → w_max

Symmetric 1-bit — weights are {-d, +d} per group, automatically handled by the affine formula above since w_min = -d, w_max = +d:

scale = w_max - w_min = 2d,  bias = w_min = -d
bit 0 → 0·(2d) + (-d) = -d
bit 1 → 1·(2d) + (-d) = +d

A dedicated symmetric 1-bit mode (scale only, no bias) could save memory and skip the bias addition in the matmul kernels, but for now both cases run through the same affine path.

What's included

• CPU backend: 1-bit quantize, dequantize, and quantized matmul (qmm dispatch)
• Metal backend: Full 1-bit support in all quantized kernels — both non-NAX (quantized.h) and NAX (quantized_nax.h, quantized_nax.metal) paths, plus quantize/dequantize kernels
• Python bindings: Updated mx.quantize(w, bits=1), mx.dequantize(...), and mx.quantized_matmul(...) documentation
• Unit tests: Added 1-bit to test_quantize_dequantize, test_qmm, and a dedicated test_1bit_quantize_dequantize covering round-trip accuracy, zero handling, and quantized matmul correctness. Full test suite passes (672 tests, 0 failures).
• No CUDA support: 1-bit is not yet supported on the CUDA backend. The CUDA dispatch_bits does not include a case 1: path. The new 1-bit test is added to cuda_skip.py.

Expected model-level performance (hypothetical, 8B parameter model, Apple M4 Pro 48 GB)

Based on the kernel-level benchmarks below, a hypothetical 8B parameter model at 1-bit would see roughly (varying by group size due to scale/bias metadata overhead):

Configuration	Memory	Expected Throughput
FP16 (baseline)	~15.3 GB	~15 tok/s
1-bit (group size 128)	~1.3 GB	~100–130 tok/s
1-bit (group size 64)	~1.6 GB	~90–115 tok/s
1-bit (group size 32)	~2.0 GB	~80–100 tok/s

We verified one scenario (group size 128, all weights quantized) and observed throughput in the ballpark of the estimates above. The primary purpose of this table is to give a sense of the runtime speed that 1-bit quantization enables. These are back-of-the-envelope numbers — actual end-to-end performance will vary depending on which layers are quantized, group size, attention overhead, and other non-quantized computation.

Kernel Corretness validation (KL divergence, 8B parameter model, WikiText-2)

To validate matmul kernel correctness, we compared two runs of the same 1-bit quantized model: one using the quantized matmul kernels (weights stay packed in 1-bit), and the other with the 1-bit weights dequantized to FP16 first and run through standard FP16 matmul. This is not a comparison between an FP16 model and its quantized version — both sides use identical weight values, so any divergence would indicate a kernel bug. Both the prompt processing (qmm) and token generation (qmv) paths were tested.

Prompt processing (qmm path) — 20 WikiText-2 chunks:

Metric	Value
Forward KL(P||Q) Mean	0.000024
Reverse KL(Q||P) Mean	0.000017
Mean Top-1 Agreement	99.85%
Min Top-1 Agreement	99.29%

Token generation (qmv path) — 113 autoregressive steps (single-token qmv) across 5 prompts:

Metric	Value
Forward KL(P||Q) Mean	0.000067
Reverse KL(Q||P) Mean	-0.000038
Mean Top-1 Agreement	100.0%
Min Top-1 Agreement	100%

Both forward and reverse KL are near-zero, confirming the quantized kernels produce results consistent with the dequantized FP16 reference in both qmm and qmv code paths.

Changes

• mlx/backend/cpu/quantized.cpp - 1-bit quantization logic and qmm dispatch
• mlx/backend/metal/kernels/quantized.h - Metal 1-bit load_vector, qdot, qdot_safe, qouter, dequantize
• mlx/backend/metal/kernels/quantized_nax.h - Same for NAX kernels
• mlx/ops.cpp - Validation to accept bits=1
• python/src/ops.cpp - Updated docstring table
• python/tests/test_quantized.py - Added 1-bit to existing tests + dedicated 1-bit test
• python/tests/cuda_skip.py - Skip 1-bit test on CUDA
• benchmarks/python/comparative/bench_mlx.py - Added 1-bit entries to quant_matmul dict; auto-quantizes weight from --size args
• benchmarks/python/comparative/compare.py - Added quant_matmul benchmark entries comparing 1/2/4/8-bit across qmv and qmm paths

Notes

1. The Metal qmv_quad_impl kernel has a minor edge case with 1-bit when the inner dimension is < 128. In practice this should never come up — virtually all models have dimensions well above 128.

2. If all weights in a group are exactly 0, the affine 1-bit quantization computes scale = eps (floored) and bias = 0, which dequantizes all values to near-zero (correct behavior).

3. Kernel-level quantized_matmul benchmarks (Apple M4 Pro 48 GB, GPU, NAX path, group_size=128, 1000 calls, weight shape in parentheses):

   qmv path (M=1, single-token generation, memory-bandwidth bound):

   Layer	FP16	1-bit	2-bit	4-bit	1-bit speedup vs FP16
   attn_proj (4096×4096)	167 µs	27 µs	39 µs	38 µs	6.2×
   ffn_gate (11008×4096)	391 µs	46 µs	72 µs	110 µs	8.6×
   ffn_down (4096×11008)	409 µs	59 µs	75 µs	98 µs	6.9×

   qmm path (M=32, prompt processing, more compute-bound):

   Layer	FP16	1-bit	2-bit	4-bit	1-bit speedup vs FP16
   attn_proj (4096×4096)	308 µs	178 µs	179 µs	182 µs	1.7×
   ffn_gate (11008×4096)	841 µs	430 µs	438 µs	433 µs	2.0×
   ffn_down (4096×11008)	649 µs	441 µs	435 µs	440 µs	1.5×

   1-bit entries have been added to benchmarks/python/comparative/bench_mlx.py and compare.py. To reproduce (from repo root):

   # run all quant_matmul benchmarks (1/2/4/8-bit, qmv M=1, qmm M=32 & M=512)
   python benchmarks/python/comparative/compare.py --filter quant_matmul
   
   # or run individual benchmarks
   python benchmarks/python/comparative/bench_mlx.py quant_matmul_t_128_1 --size 1x4096 --size 4096x4096
   python benchmarks/python/comparative/bench_mlx.py quant_matmul_t_128_4 --size 1x4096 --size 4096x4096
   
   # unit tests
   python -m pytest python/tests/test_quantized.py::TestQuantized::test_1bit_quantize_dequantize -v
   python -m pytest python/tests/test_quantized.py::TestQuantized::test_quantize_dequantize -v
   python -m pytest python/tests/test_quantized.py::TestQuantized::test_qmm -v

Questions for reviewers

1. NAX vs non-NAX testing: All benchmarks and the full test suite were run on macOS 26.2 (M4 Pro 48 GB), where NAX is active. The non-NAX path was partially validated by rebuilding with -DCMAKE_CXX_FLAGS=-DMLX_METAL_NO_NAX — unit tests pass, but full benchmarking was only done on the NAX path. We only have access to an M4. Is the -DMLX_METAL_NO_NAX build flag sufficient to validate the non-NAX path, or would you recommend testing on actual older hardware (M1/M2/M3)?

2. Test coverage: The full test suite passes (672 tests, 0 failures), including dedicated 1-bit tests for both symmetric and asymmetric weight round-trip accuracy, zero handling, and quantized matmul correctness across both qmm and qmv paths. Is there any additional testing you'd like to see before merging?

Future work

• Dedicated symmetric 1-bit mode
• CUDA support

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[Copilot]

Pull request overview

This PR adds 1-bit affine quantization support to MLX, extending the existing quantization framework from {2, 3, 4, 5, 6, 8} bits to include 1-bit. The implementation provides efficient packing and inference for 1-bit quantized weights on Apple Silicon (Metal) and CPU backends.

Changes:

• Added 1-bit support to affine quantization with formula: scale = w_max - w_min, bias = w_min where bit 0 → w_min, bit 1 → w_max
• Implemented full Metal kernel support for 1-bit in both NAX and non-NAX paths across all quantized operations (quantize, dequantize, qmm, qmv)
• Extended CPU backend with 1-bit quantization, dequantization, and quantized matmul dispatch
• Added comprehensive test coverage for 1-bit symmetric/asymmetric weights, zero handling, and quantized matmul correctness
• Updated Python bindings documentation and validation to accept bits=1
• Added 1-bit benchmark entries for performance comparison across different group sizes
• Excluded CUDA backend from 1-bit support (added to cuda_skip.py)

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
python/tests/test_quantized.py	Added 1-bit to existing parameterized tests and new dedicated test_1bit_quantize_dequantize with symmetric/asymmetric weights, zero handling, and qmm/qmv correctness tests
python/tests/cuda_skip.py	Added 1-bit test to CUDA skip list since CUDA backend doesn't support 1-bit
python/src/ops.cpp	Updated quantization mode documentation table to include 1-bit in supported bits
mlx/ops.cpp	Modified validation to accept bits >= 1 and added 1-bit quantization formula (scale = w_max - w_min, bias = w_min)
mlx/backend/metal/kernels/quantized_nax.metal	Added 1-bit kernel instantiation macro for NAX path
mlx/backend/metal/kernels/quantized_nax.h	Implemented 1-bit versions of load_vector, load_vector_safe, qdot, qdot_safe, qouter, and dequantize for NAX optimized kernels
mlx/backend/metal/kernels/quantized.metal	Added 1-bit kernel instantiation macro and quantize/dequantize logic for non-NAX path
mlx/backend/metal/kernels/quantized.h	Implemented 1-bit versions of all quantization primitives for non-NAX kernels
mlx/backend/cpu/quantized.cpp	Added 1-bit case to qmm dispatch and quantization logic matching Metal implementation
benchmarks/python/comparative/compare.py	Added compare_mlx_quant function and 1-bit benchmark entries for qmv and qmm paths
benchmarks/python/comparative/bench_mlx.py	Added 1-bit entries to quant_matmul dictionary for all group sizes and transpose modes, plus auto-quantization logic

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[angeloskath]

Hi @khosravipasha that is pretty cool. I am not sure we want to support 1-bit quants natively in MLX, even 2 bits are not really used out there.

The options I see are:

1. You could leave it as an open PR for people to chime in if they want this or think it would be useful in any way
2. You could make it an extension that we 'd be happy to link to at MLX Community Projects #654 and anybody could simply pip install it and use it.

[khosravipasha]

@angeloskath Thanks for the comment, we actually do have a native 1-bit modelx just released today, so would be amazing to have it in the official mlx repo.

Sorry for delayed response, we were waiting to come out of stealth before resending the PR.

Checkout our launch: https://huggingface.co/prism-ml
For now we are hosting it in out public fork: https://github.com/PrismML-Eng/mlx

I made a few changes I will send a the fresh code in a bit, we also have changes on mlx-swift for iPhones.

[angeloskath]

Did a first pass and left some comments.

The main one is on the change in qmv_fast and why it is needed it isn't clear from the PR.

There is also linter errors and so on. But let's fix these and get it merged.

[khosravipasha]

Thanks for the feedback, addressed the easy ones, the main one is to figure out a good way to handle qmv_fast changes (4B model has shapes that are not divisible by 2028 nor 1024, that causes block_size for loops to have some left overs; see detail in the comment)

Had a question about NAX vs non-NAX, is there something special need to do, not too familiar what their difference is, I mainly tested in Mac M4 Pro. Want to make sure works well on older Mac (M1-M3).

[angeloskath]

Had a question about NAX vs non-NAX

NAX refers to the neural accelerators for M5. I can run some benchmarks. There isn't anything overly special per se, just the matmuls are faster so we need to make sure we can keep feeding the NAX quickly enough so dequantizing should be efficient.

There could be a lot of room for tuning this but let's get something working that isn't too far from bf16 for a largeish matrix and we should be good to go.

[khosravipasha]

NAX refers to the neural accelerators for M5. I can run some benchmarks. There isn't anything overly special per se, just the matmuls are faster so we need to make sure we can keep feeding the NAX quickly enough so dequantizing should be efficient.
There could be a lot of room for tuning this but let's get something working that isn't too far from bf16 for a largeish matrix and we should be good to go.

I see thanks, yeah more neural chips on M5 chips seems exciting from what I have read, have not had a chance to try them yet myself. For now mostly care about correctness, and not being very slow. Definetly can be tuned,

I saw Ivan ran speed benchmarks with M5 Max already, looks very fast
https://x.com/ivanfioravanti/status/2039077744114319461

[khosravipasha]

Few tests were failing, some are CUDA which are not supported, I tried skipping cuda tests since there is no support need to see what else to add

The windows CPU build failing was surprising since same test passes on other archs, does mlx support windows too?

  ======================================================================
  FAIL: test_1bit_quantize_dequantize (test_quantized.TestQuantized.test_1bit_quantize_dequantize) (gs=32, case='quantized_matmul_asymmetric')
  Test 1-bit affine quantization.
  ----------------------------------------------------------------------
  Traceback (most recent call last):
    File "D:\a\mlx\mlx\python\tests\test_quantized.py", line 264, in test_1bit_quantize_dequantize
      self.assertLess((y_q - y_hat).abs().max(), 1e-5)
      ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  AssertionError: array(1.14441e-05, dtype=float32) not less than 1e-05

[khosravipasha]

@angeloskath did you got a chance to do another pass. Was not sure which path for the extra epilogue is good to do.

[Copilot]

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated 2 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[khosravipasha]

@angeloskath

tried something similar got 15% speed on 4B on token generation, see comment below:

nice what made it faster? above comment is the only thing I see (other comments might still be in draft?)

fixed the merge conflict in the last push and merged with a more recent main (skip_cuda was removed and is now inline so made the change for 1-bit to skip cuda unit tests)

[seongyun1104]

Tested on M1 (8GB, macOS 26.5), source build from this branch.

CPU backend looks good — 29/29 tests pass and I cross-checked quantize/dequantize/quantized_matmul against an independent numpy reference implementation, all match exactly (max|Δ|=0) across gs={32,64,128} with various weight distributions.

On Metal I found an issue with 1-bit + gs=32. dequantize is fine but quantized_matmul returns garbage (inf/1e35). gs=64 and gs=128 work correctly (~1e-6 error). Existing bits (2,4,8) at gs=32 are also fine.

The PR tests don't catch this because they use binary-valued weights that happen to work at gs=32. It shows up with continuous (e.g. gaussian) weights.

Root cause is in qmv_fast_impl:

constexpr int packs_per_thread = bits == 2 ? 1 : 2;

For bits=1: pack_factor = 32, so values_per_thread = 32 * 2 = 64. With gs=32 that gives scale_step_per_thread = 32 / 64 = 0, which is a division by zero in the scale indexing a few lines below.

Changing it to bits <= 2 ? 1 : 2 fixes it — gives values_per_thread=32, scale_step=1. After that, all 29 tests pass on Metal and the cross-check with gaussian weights at gs=32 comes back clean (~3e-6).

This is related to the block size discussion above — with packs_per_thread=1 the block size drops to 1024, so it still won't work for dims not divisible by 1024.

Reproducer (before fix):

import mlx.core as mx, numpy as np
w = mx.array(np.random.default_rng(42).standard_normal((128, 512)).astype(np.float32) * 0.1)
x = mx.random.normal(shape=(1, 512))
wq, s, b = mx.quantize(w, group_size=32, bits=1)
y_q = mx.quantized_matmul(x, wq, s, b, True, 32, 1)
y_ref = x @ mx.dequantize(wq, s, b, 32, 1).T
mx.eval(y_q, y_ref)
print(float(mx.abs(y_q - y_ref).max().item()))  # ~14.8, should be ~1e-6

[khosravipasha]

@seongyun1104
Thanks for notice, good catch, is the fix only that 1 line of code?
I can try to test it with our models to see if it affects speed. Or pack our models with group 32 and check end-to-end correctness.

[seongyun1104]

Yes, just the one line:

- constexpr int packs_per_thread = bits == 2 ? 1 : 2;
+ constexpr int packs_per_thread = bits <= 2 ? 1 : 2;

For end-to-end check, a group_size=32 packed model is what reveals it — the existing tests use binary-valued weights that round-trip cleanly even with the bug, so gs=32 with continuous distributions is needed to surface it.

Happy to verify on M1 if useful.

[khosravipasha]

@seongyun1104 thanks for the pointer, pushed the fix into this PR branch, also seems to get some speed up for some cases.

@angeloskath
Any other blockers for merging this?

[seongyun1104]

@khosravipasha Thanks for picking it up — nice surprise that it nets a speedup in some cases too.

[zcbenz]

Thanks for keeping the PR updated!

btw in future please consider creating PRs from personal accounts, it is not possible for us to push to or rebase branch in org account even when you enable the option on GitHub.

[zcbenz]

This fails on Windows CPU build with array(0.00, dtype=float32) not less than 1e-05, seems to be bug of MLX, maybe try:

Suggested change

	self.assertLess((y_q - y_hat).abs().max(), 1e-5)
	self.assertLess((y_q - y_hat).abs().max().item(), 1e-5)

[khosravipasha]

@zcbenz Thanks for the review, we will take a look.

btw in future please consider creating PRs from personal accounts, it is not possible for us to push to or rebase branch in org account even when you enable the option on GitHub.

Oh did not know that, this was my personal account when I initially submitted the PR, but I guess the fork needs to be in a personal account too

[khosravipasha]

@zcbenz Thanks for the review! We merged the latest main plus applying suggested (fixes skipped unsupported 1-bit CUDA test cases, and added .item() to the scalar assertions). The CUDA failure comes from a missing 1-bit dispatch case; it doesn’t seem to use the generic fallback.

Metal quantization tests and the affected Mac CPU tests pass. Windows/CUDA still need CI rerun to confirm its all good.

Quick M5 Max benchmarks, measured before the upstream merge: 512 prompt tokens, 128 generated tokens, three runs after warmup.

Model	Prompt tok/s	Generation tok/s
1.7B	11,502	489
4B	5,227	217
8B	3,278	244

[zcbenz]

Can you try setting error to 2e-5 for Windows?

  FAIL: test_1bit_quantize_dequantize (test_quantized.TestQuantized.test_1bit_quantize_dequantize) (gs=32, case='quantized_matmul_asymmetric')
  Test 1-bit affine quantization.
  ----------------------------------------------------------------------
  Traceback (most recent call last):
    File "D:\a\mlx\mlx\python\tests\test_quantized.py", line 320, in test_1bit_quantize_dequantize
      self.assertLess((y_q - y_hat).abs().max().item(), 1e-5)
  AssertionError: 1.1444091796875e-05 not less than 1e-05

[khosravipasha]

hmm curious why this happens only on Window, added these two lines

tol = 2e-5 if platform.system() == "Windows" else 1e-5
self.assertLess((y_q - y_hat).abs().max().item(), tol)

[zcbenz]

Some CPU ops are probably having precision issues on Windows.