Bonsai 1-bit prefill (M>5) falls back to fp16 dequant: ~16x memory blowup, ~150x slower

[duffercn]

Summary

For bits=1 models (e.g. prism-ml/Bonsai-27B-mlx-1bit), prefill (any forward pass with M > 5 tokens — i.e. essentially every real prompt) falls back to explicitly dequantizing the 1-bit weights to the input dtype (fp16/bf16) and running a dense matmul, instead of using a native quantized kernel. This causes a ~16x memory blowup for affected weight tensors and generation throughput roughly 150x slower than the reference implementation.

This is a known, intentional scope limit from #2248 (omlx/patches/bonsai_qmv.py):

# Prefill: M > _MAX_DECODE_M uses stock quantized_matmul which calls affine_dequantize.
# Stock MLX doesn't have affine_dequantize for bits=1 in its metallib.
# For bits=1 prefill, dequantize to float16 explicitly.
if M > _MAX_DECODE_M:
    if bits == 1:
        w_fp = _dequant_1bit(self.weight, self.scales, ..., x.dtype)
        out = x @ w_fp.T

with _MAX_DECODE_M = 5. #2248's decode-path kernels (the ones people are cheering about in #526) only cover M ≤ 5, so this fallback is hit on essentially every request's prompt processing — not an edge case.

Reproduction & measurements (2026-07-22, M4 16GB, oMLX 0.5.2)

Isolated test bypassing the server (oMLX's own bundled mlx 0.5.2 + mlx_lm 0.31.3 + omlx.patches.bonsai_qmv applied directly, no FastAPI/BatchedEngine in the loop), using mx.get_active_memory()/get_cache_memory()/get_peak_memory():

stage	active	cache	peak
after load()	4.21 GB	0	4.21 GB
after one generate() call (185 tokens)	4.21 GB	7.80 GB	17.45 GB

That single generate call took 2045 seconds (0.1 tok/s). Two subsequent generate calls in the same (now "warm") process stayed at the same 0.1 tok/s and did not grow memory further — consistent with the slow path being hit on every prefill, not a one-time JIT/compile cost.

For comparison, the PrismML-Eng/mlx prism branch (commit 88c9c20, the fork #2248 upstreamed some of but not the prefill path) running the same model via plain mlx_lm.generate on the same machine: load 4.21GB (identical), generation peak 4.49–4.54GB, 15.9–19.4 tok/s. Weight loading is identical between the two paths (same 4.21GB), confirming the entire gap is prefill-kernel-specific, not a loading/architecture difference.

Environment

• oMLX 0.5.2 (Homebrew, jundot/omlx)
• macOS, Apple Silicon (M4)
• Model: prism-ml/Bonsai-27B-mlx-1bit
• Bundled: mlx 0.5.2 (stock), mlx_lm 0.31.3

Suggested fix

Either:

1. Extend the native kernel from feat: Bonsai 1-bit / 2-bit decode kernels (upstream from Prism MLX fork) #2248 to cover prefill (M > 5) for bits=1, following the same approach the PrismML-Eng fork uses for its prefill path (commit 88c9c20, "1-bit qmv_fast use 1 pack/thread for occupancy" — the fork does not need to fall back to fp16 dequant for prefill), or
2. At minimum, chunk the fp16 dequant fallback (process in smaller M chunks) so it doesn't materialize the full dequantized weight tensor at once, to cap the memory blowup even if the speed gap remains.

Given #526 is still open despite #2248 being merged, I suspect this prefill gap is the reason — worth linking the two.