feat(core): load and run affine 1-bit checkpoints (`quantization.bits == 1`) with a fused Metal matvec
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Summary
Public 1-bit affine checkpoints exist on the Hub and mlxcel cannot run them. `prism-ml/Bonsai-1.7B-mlx-1bit`, `prism-ml/Bonsai-4B-mlx-1bit`, `prism-ml/Bonsai-8B-mlx-1bit` (1.28 GB) and `prism-ml/Bonsai-27B-mlx-1bit` are plain Qwen3 (`model_type: qwen3`, `architectures: ["Qwen3ForCausalLM"]`) with `"quantization": {"group_size": 128, "bits": 1}` and the standard MLX packed layout (`uint32` `.weight`, `.scales`, `.biases`) on every projection, the embedding table, and `lm_head`. MLX's quantized kernels exist only for 2, 3, 4, 5, 6 and 8 bits, so the triple mlxcel builds for these weights reaches `quantized_matmul` with `bits = 1` and the first forward pass throws inside C++. This issue adds a 1-bit matvec / matmul kernel and a dequantize graph, routes `UnifiedLinear` and `QuantizedEmbedding` to them when `bits == 1`, and makes the loader accept the width explicitly.
## Current behavior
The config validator admits `bits = 1` and the per-tensor reconciler returns it unchanged because the packed shapes are self-consistent (`packed_in * 32 == 1 * num_groups * 128`):
```rust
// src/lib/mlxcel-core/src/layers.rs:1290-1298
pub fn validate_quantization_params(group_size: i32, bits: i32) -> Result<(), String> {
if !(1..=32).contains(&bits) { ... }
// src/lib/mlxcel-core/src/layers.rs:1186-1197
let inferred_bits = numerator / denominator;
if inferred_bits == caller_bits { return Ok(caller_bits); } // 1 == 1: the {2,3,4,5,6,8} allowlist below is never consulted
if ![2, 3, 4, 5, 6, 8].contains(&inferred_bits) { return Err(...) }
```
`UnifiedLinear::from_weights_with_mode` (`layers.rs:1904-2000`) then stores `QuantizedWeight { bits: 1, mode: "affine" }` and `forward` (`layers.rs:2027`) calls `ffi::quantized_matmul(.., group_size, 1, "affine")`. The pinned MLX instantiates its affine Metal kernels for bits 2, 3, 4, 5, 6 and 8 only (`quantized.metal`, `instantiate_quantized_groups(2|3|4|5|6|8)`, and `quantized.h` static-asserts the same set), so the call throws. A C++ throw crosses the cxx bridge as an uncatchable abort (the #929 / #973 class), so the user sees a crash at the first request rather than a load error. Expected on `prism-ml/Bonsai-1.7B-mlx-1bit`; confirm as step 0 of the implementation and record the exact message in the PR.
Nearest existing code: the BitNet ternary kernel, `bitlinear_matmul` (`src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp:60-80` Metal, `:86+` CUDA port, FFI at `src/lib/mlxcel-core/src/lib.rs:1766`), which is the same shape of problem (a packed weight format MLX has no kernel for) solved with `mlx::core::fast::metal_kernel` plus a `cuda_kernel` port. BitNet's format is unrelated (2-bit ternary, 4 rows per byte, one per-tensor scale) and `src/models/bitnet.rs` stays as is.
## Expected behavior
Checkpoint layout (affine, `bits = 1`, `group_size G in {32, 64, 128}`), for a linear with `out` rows and `in` columns:
- `weight`: `uint32 [out, in / 32]`. Bit `j` (LSB first) of word `c` is column `i = 32 * c + j`.
- `scales`: `[out, in / G]`, `biases`: `[out, in / G]`, float (f32 in the Bonsai export; accept f16 / bf16 too, keep as shipped, do not promote, per the quantized-model rule).
- Dequant: `w[o, i] = bit(o, i) * scales[o, i / G] + biases[o, i / G]` where `bit(o, i) = (weight[o, i >> 5] >> (i & 31)) & 1`.
- Embedding tables use the same layout with `out = vocab`.
Matvec math (what the kernel computes; exact, no approximation):
```
y[b, o] = sum_g ( scales[o, g] * sum_{i in group g} bit(o, i) * x[b, i]
+ biases[o, g] * sum_{i in group g} x[b, i] )
```
i.e. per (row, group) one masked sum of activations and one plain sum of activations; the bias term factors out of the bit mask. With `G` a multiple of 32 the group index of a 32-bit word is `c * 32 / G` and never straddles a word.
Loader: `quantization.bits == 1` with `.biases` present is accepted as mode `"affine"`; `bits == 1` without `.biases` (block-float) is refused with a clear error. MoE expert stacks (`SwitchLinear` / `gather_qmm`) at 1 bit are out of scope and must fail at load with a message naming the prefix, not abort in `gather_qmm`.
Decode routing: all the fused quantized fast paths that hand a `QuantizedWeight` straight to an MLX `quantized_matmul` inside a C++ helper must decline `bits == 1` and fall back to the graph path: `forward_split_norm_rope_quantized` (`layers.rs:3027`, used by `src/models/qwen3.rs:123`), `forward_split_rope_quantized` (`layers.rs:2966`), `forward_fused_rope_append` (`layers.rs:2856`), and `quantized_linear_forward_global_scale` (`layers.rs:2043`).
Performance target: decode of `Bonsai-8B-mlx-1bit` on Apple Silicon at least as fast as `mlx-community/Qwen3-8B-4bit` decode on the same machine (the kernel reads 4x fewer weight bytes; matching 4-bit is the floor, not the goal).
## Implementation plan
1. `src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp`: add `one_bit_qmv` (Metal, `fast::metal_kernel`) and `one_bit_qmm` (Metal, simdgroup-matrix tiles) next to the BitLinear holders, with the same `KernelHolder` lazy-init pattern. Inputs `x [M, K]`, `weight [N, K/32] uint32`, `scales [N, K/G]`, `biases [N, K/G]`, output `[M, N]` in `x.dtype()`.
- qmv (M small): threadgroup of 64 threads = 2 simdgroups; each simdgroup owns `R` consecutive output rows (`R = 4` when `N <= 64 || K >= 2N`, else `8`); each lane owns 16 consecutive activations per step (`block_start = lane * 16`, stride `512`), loads them once into registers with their plain sum, then for each of its `R` rows reads the 16-bit half-word `ushort(weight[o, block_start >> 5] >> (block_start & 31))`, accumulates `selected_sum = sum_i (bit_i ? x_i : 0)` and does `acc[r] += selected_sum * scales[o, block_start / G] + total_sum * biases[o, block_start / G]`; `simd_sum` across lanes; lane 0 writes. Grid `(ceil(N / (2R)) * 64, M, 1)`, threadgroup `(64, 1, 1)`; a `ROW_VALID` template switch drops the bounds check when `N % (2R) == 0`. Template args must include `{"T", x.dtype()}` plus `G`, `R`, `aligned` (the dtype entry is the cache key, see CLAUDE.md "JIT kernel cache keys").
- qmm (M >= 16 and K % 512 == 0): 32x32x32 tiles, 128 threads; each thread loads 8 activations and expands 8 packed bits into `bias + (bit ? scale : 0)` into a threadgroup `weight_tile`, then `simdgroup_multiply_accumulate` over 8x8 fragments; grid `(ceil(N/32) * 128, ceil(M/32), 1)`.
- CUDA: port qmv with one warp per `R` rows and `__shfl_down_sync` reduction (mirror the BitLinear CUDA port); qmm may fall back to the dequantize graph on CUDA in the first PR.
- `one_bit_dequantize(weight, scales, biases, G)` as a plain MLX graph: `bits = (weight[..., None] >> arange(32, uint32)) & 1`, reshape to `[.., K]`, `astype(scales.dtype) * repeat(scales, G, -1) + repeat(biases, G, -1)`. This is the reference for tests and the non-Metal/non-CUDA fallback (`x @ dequant.T`).
- `one_bit_quantized_matmul(x, weight, scales, biases, G)`: reshape `x` to `[M, K]`, validate `weight.shape[1] * 32 == K`, `scales.shape[1] * G == K`, `G in {32, 64, 128}`, dispatch qmm / qmv / fallback, reshape back. Env kill switch `MLXCEL_ONE_BIT_KERNEL=0` forces the dequantize fallback (the A/B and the parity oracle).
2. `src/lib/mlxcel-core/src/lib.rs` (cxx bridge, `ext` or `kernels` section): declare `fn one_bit_quantized_matmul(x, weight, scales, biases, group_size) -> UniquePtr` and `fn one_bit_dequantize(weight, scales, biases, group_size)`; header in `mlx_cxx_bridge.h`.
3. `src/lib/mlxcel-core/src/layers.rs`:
- `infer_quantization_bits`: extend the allowlist to `[1, 2, 3, 4, 5, 6, 8]` (the early return already admits 1; make it explicit so the message stays truthful). `validate_quantization_biases`: a `bits == 1` triple without `.biases` is an error ("1-bit weights are affine-only").
- `UnifiedLinear::forward` and `dequantized_weight`: when `weight.bits == 1 && weight.mode == "affine"`, call `ffi::one_bit_quantized_matmul` (adding the dense `bias` afterwards) / `ffi::one_bit_dequantize`. `QuantizedWeight` gains a `fn is_one_bit(&self) -> bool` helper.
- `QuantizedEmbedding::forward`: gather `weight`, `scales`, `biases` rows with `ffi::take` by token id, then `one_bit_dequantize`; `as_linear`: `one_bit_quantized_matmul`. (`ffi::quantized_embedding` and `ffi::quantized_linear_forward` throw at 1 bit.) Bonsai-1.7B has `tie_word_embeddings: true`, so `as_linear` is on the decode path.
- The four fused entry points listed under Expected behavior return `None` when `fused_quantized_weight().is_one_bit()`.
- `QuantizedMultiLinear` (`layers.rs:3085`, MLA) and `SwitchLinear` (`src/models/switch_layers.rs:547`, range `(2..=8)`): reject `bits == 1` at load with the prefix in the message.
4. `src/models/qwen3.rs`: no change expected; it loads through `UnifiedLinear::from_weights` / `UnifiedEmbedding`. Smoke the fused-qkv concatenation (`QkvProjection::Fused`) with 1-bit planes: concatenating `weight` / `scales` / `biases` along the row axis is layout-preserving for this format too.
5. Memory estimation (`mlxcel inspect`, pre-load estimate): verify the weight-bytes model uses the packed tensor sizes from the safetensors header rather than `bits`; if it divides by `bits`, add 1 to the supported set.
6. `docs/supported-models.md` quantization table: add the row "affine 1-bit (`bits: 1`, `group_size` 32/64/128, scales + biases): supported on Metal (fused qmv / qmm) and CUDA (qmv), dequantize fallback elsewhere; dense linears and embeddings only". `mlxcel list` output unchanged (architecture is still `qwen3`).
## Validation
(a) Unit tests in the existing `#[cfg(test)] mod tests` of `src/lib/mlxcel-core/src/layers.rs` (line 5427, the file's convention) for the layer-level cases and in `src/lib/mlxcel-core/src/ffi_tests.rs` for the kernel cases:
- `one_bit_dequantize_matches_host_formula`: random `uint32 [8, 4]` words, `scales` / `biases` `[8, 1]` (`G = 128`); compare against a host loop implementing `bit * scale + bias` exactly.
- `one_bit_qmv_matches_dequant_matmul_{m1,m7}`: `K in {512, 2048, 4096}`, `N in {64, 96, 1000}` (covers both `R` choices and the unaligned row guard), f16 and bf16 activations; `|y_kernel - x @ dequant.T| <= 2e-2 * max|y|`.
- `one_bit_qmm_matches_dequant_matmul_m16_m33`: `K = 1024`, `N = 96`.
- `one_bit_qmv_group_sizes_32_64_128`: same check across `G`.
- `one_bit_kill_switch_routes_to_dequant`: `MLXCEL_ONE_BIT_KERNEL=0` output equals the dequant path bit-for-bit.
- `one_bit_embedding_gather_matches_dequant_rows`, `one_bit_embedding_as_linear_matches_qmv`.
- `unified_linear_routes_one_bit`: a `WeightMap` with a `bits = 1` triple loads as `Quantized { bits: 1 }` and `forward` matches the dequant path.
- `one_bit_without_biases_is_rejected`, `switch_linear_rejects_one_bit_with_prefix`, `fused_qkv_declines_one_bit`.
- `infer_quantization_bits_accepts_one` (`packed_in = 64`, `num_groups = 16`, `G = 128`, declared 1).
(b) Real checkpoints (`prism-ml/Bonsai-1.7B-mlx-1bit` first, 8B second):
```
./target/release/mlxcel download prism-ml/Bonsai-1.7B-mlx-1bit
./target/release/mlxcel generate -m models/Bonsai-1.7B-mlx-1bit -p "The capital of France is" -n 64
MLXCEL_ONE_BIT_KERNEL=0 ./target/release/mlxcel generate -m models/Bonsai-1.7B-mlx-1bit -p "The capital of France is" -n 64
./target/release/mlxcel generate -m models/Bonsai-8B-mlx-1bit -p "Explain the rules of chess in five sentences." -n 128
```
Acceptance: no abort; greedy output of the kernel path is token-identical to the dequantize-fallback path for the first 64 tokens (same dequant formula, f32 accumulation; if a near-tie flips a token past 32 tokens record it, it is the documented reduction-order class); output is fluent English ("Paris" completes the first prompt); `Bonsai-8B-mlx-1bit` decode tok/s on the same machine is at or above `mlx-community/Qwen3-8B-4bit`, recorded in the PR. Serve `Bonsai-1.7B-mlx-1bit` with `mlxcel-server` and run two concurrent chat requests (batched decode uses the same `UnifiedLinear::forward`; qmm path exercised by prefill).
## Acceptance criteria
- [ ] `prism-ml/Bonsai-1.7B-mlx-1bit` and `prism-ml/Bonsai-8B-mlx-1bit` load and generate fluent text on Metal via `mlxcel generate` and `mlxcel-server`
- [ ] Kernel path matches the dequantize path (token-exact greedy on the 1.7B prompt above; numeric tolerance in unit tests)
- [ ] `MLXCEL_ONE_BIT_KERNEL=0` fallback and the CUDA qmv port build and pass the same unit tests (CUDA on a GB10 node or recorded as pending with the reason)
- [ ] 1-bit without `.biases`, and 1-bit expert stacks, fail at load with a prefix-naming error instead of aborting
- [ ] All fused quantized fast paths decline `bits == 1` and fall back to the graph path
- [ ] Decode throughput on the 8B checkpoint at or above the 4-bit Qwen3-8B baseline, numbers in the PR
- [ ] docs/supported-models.md updated
- [ ] cargo test --workspace --profile test-fast --features metal,accelerate passes
- [ ] cargo clippy --workspace --all-targets -- -D warnings and cargo fmt --all -- --check pass
## Out of scope
- 1-bit MoE expert stacks (`gather_qmm`); no public checkpoint ships them.
- Producing 1-bit checkpoints (mlxcel has no conversion pipeline).
- The 2-bit ternary Bonsai variants (`prism-ml/Ternary-Bonsai-*-mlx-2bit`): `bits = 2` is a stock MLX affine width and should load today; verify separately and file a bug if not.
Contributor guide
Research direction
Start by running the Bonsai-1.7B-mlx-1bit generate command in the issue and record the current abort. Read the BitNet kernel in src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp, the FFI in src/lib/mlxcel-core/src/lib.rs, and routing in layers.rs; run the listed layers.rs and ffi_tests.rs cases. Done means affine 1-bit dense linears and embeddings load without aborting, match the dequantize fallback, and satisfy the checkpoint and performance checks.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- machine-learning, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100