lablup / lablup/mlxcel

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

Open
#1,370 0 comments 0 reactions 0 assignees View on GitHub
arch:dense area:core area:docs area:inference area:models modelsize:small modeltype:text priority:medium status:ready type:enhancement
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.