Skip to content

feat(core): run affine 1-bit checkpoints with fused Metal kernels - #2035

Merged
inureyes merged 2 commits into
mainfrom
feature/issue-1370-one-bit-affine
Sep 29, 2026
Merged

inureyes merged 2 commits into
mainfrom
feature/issue-1370-one-bit-affine

Conversation

@inureyes

Copy link
Copy Markdown
Member

Summary

Affine 1-bit checkpoints (quantization.bits == 1, e.g. prism-ml/Bonsai-*-mlx-1bit) now load and generate. Before this change, main aborted on the first token of Bonsai-1.7B-mlx-1bit with libc++abi: terminating due to uncaught exception of type std::runtime_error: [metal::Device] Unable to load kernel affine_dequantize_float16_t_gs_128_b_1 (exit 134).

  • New cpp/mlx_cxx_one_bit.cpp: Metal qmv kernel (M < 16, factored scale * masked_sum + bias * total_sum), tiled simdgroup qmm kernel (M >= 16, f32 accumulation), and a dequantize graph used as the reference, the non-Metal fallback, and the MLXCEL_ONE_BIT_KERNEL=0 kill switch.
  • Routing lives in the bridge primitives (quantized_matmul in both layouts, quantized_linear_forward, _global_scale, dequantize, quantized_embedding), so every dense linear, embedding lookup and tied lm_head is covered, including callers outside layers.rs.
  • UnifiedLinear::quantized_weight() / as_quantized_weight() answer None for 1-bit weights (same mechanism as runtime LoRA), so every fused helper that calls MLX's quantized_matmul internally takes its graph fallback. The su-scaled-rope and compiled SwiGLUMLP paths decline explicitly.
  • Loader: 1 bit is accepted through a new LOADABLE_AFFINE_BITS; SUPPORTED_AFFINE_BITS stays the producer set (split-mtp, gemma4 overrides). 1-bit without .biases, 1-bit expert stacks (SwitchLinear) and 1-bit MLA (QuantizedMultiLinear) fail at load with the prefix named.
  • Docs: supported-models.md quantization row, MLXCEL_ONE_BIT_KERNEL in environment-variables.md.

Differences from the issue text

  • The Bonsai export ships f16 scales and biases, not f32. They are kept as shipped.
  • forward_fused_rope_append needs no decline: it consumes project_concat, which goes through UnifiedLinear::forward.
  • The explicit FFI entry points (one_bit_quantized_matmul, one_bit_dequantize) return Result and are used by tests; the load path reaches the same code through the bridge primitives, which keeps call sites unchanged.
  • CUDA qmv port: pending. No CUDA host is reachable (GB10 runner down), and an unexecuted JIT kernel running by default would risk silent wrong output. CUDA and ROCm run the dequantize graph, which the unit tests cover.
  • inspect memory estimate needed no change: it reads packed sizes from the safetensors header.

Test plan

  • cargo test --release -p mlxcel-core --lib one_bit_tests: 17 pass. Kernels and fallback are each checked against a host f64 oracle (M 1/3/7/16/17/18/20/33/40, K 256 to 4096, N 64 to 1003 incl. unaligned, G 32/64/128, f16/bf16/f32), kill switch is bit-identical to the graph, rejections name the prefix.
  • layers::tests (92), models::switch_layers::tests (26), cli_help_consistency, dead_doc_pointers, llama_compat_manifest pass. infer_bits_rejects_non_canonical_widths updated: it asserted that 1 bit is rejected, and now uses 7.
  • cargo fmt --all -- --check, cargo clippy --release -p mlxcel-core --lib --tests -- -D warnings, same for -p mlxcel.
  • Bonsai-1.7B-mlx-1bit, "The capital of France is", -n 64 -t 0: kernel and MLXCEL_ONE_BIT_KERNEL=0 outputs identical over all 64 tokens, answer "Paris".
  • Bonsai-8B-mlx-1bit, chess prompt, -n 128 -t 0: kernel and fallback identical over 128 tokens, fluent English.
  • mlxcel-server on the 1.7B, two concurrent greedy chat requests: "Paris is the capital of France." and a correct primary-colors answer.
  • Decode throughput vs mlx-community/Qwen3-8B-4bit: deferred. The machine was shared with other GPU jobs during this run, so no timing is reported here.

Closes #1370

@inureyes inureyes added status:review Under review type:enhancement New features, capabilities, or significant additions priority:medium Medium priority area:core mlxcel-core: MLX FFI, primitives, KV cache, layers area:inference Generation, sampling, decoding (incl. speculative, DRY) area:models Model architectures, weights, loading, metadata area:docs User and developer documentation labels Sep 29, 2026
inureyes added a commit that referenced this pull request Sep 29, 2026
@inureyes inureyes added status:done Completed and removed status:review Under review labels Sep 29, 2026
MLX implements affine quantized kernels for 2, 3, 4, 5, 6 and 8 bits only, so a checkpoint with `quantization.bits == 1` (prism-ml Bonsai, plain qwen3) loaded fine and then aborted on the first forward pass with "Unable to load kernel affine_dequantize_float16_t_gs_128_b_1".

The new cpp/mlx_cxx_one_bit.cpp adds a Metal qmv kernel for fewer than 16 rows, a tiled simdgroup qmm kernel for larger batches, and a dequantize graph used as the reference, the non-Metal fallback, and the MLXCEL_ONE_BIT_KERNEL=0 kill switch. The quantized_matmul, quantized_linear_forward(_global_scale), dequantize and quantized_embedding bridge primitives route an affine 1-bit triple there, so every dense linear, embedding lookup and tied lm_head is covered. The UnifiedLinear accessors that feed fused C++ helpers (which call MLX's quantized_matmul directly) now answer None for 1-bit weights, so those callers take their graph fallback.

The loader accepts 1 bit through a separate LOADABLE_AFFINE_BITS set, leaving SUPPORTED_AFFINE_BITS as the producer allowlist. 1-bit without .biases, 1-bit MoE expert stacks and 1-bit MLA per-head projections fail at load with the prefix named. CUDA and ROCm take the dequantize graph until a port can be validated on hardware.

Validated with 17 unit tests against a host f64 oracle and on Bonsai-1.7B and Bonsai-8B, where kernel and fallback greedy outputs are identical over 64 and 128 tokens.

Refs #1370
@inureyes
inureyes force-pushed the feature/issue-1370-one-bit-affine branch from ecab741 to b7c7437 Compare September 29, 2026 09:54
@inureyes
inureyes merged commit 55ceacf into main Sep 29, 2026
@inureyes
inureyes deleted the feature/issue-1370-one-bit-affine branch September 29, 2026 09:54
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:core mlxcel-core: MLX FFI, primitives, KV cache, layers area:docs User and developer documentation area:inference Generation, sampling, decoding (incl. speculative, DRY) area:models Model architectures, weights, loading, metadata priority:medium Medium priority status:done Completed type:enhancement New features, capabilities, or significant additions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat(core): load and run affine 1-bit checkpoints (quantization.bits == 1) with a fused Metal matvec

1 participant