Support Bonsai-27B affine Q1

[dusterbloom]

What

• route Qwen3.5 affine 1-bit linear, tied-head, embedding, and fused dense-FFN operations through Higgs' Bonsai Metal kernels
• reduce the fast Q1 decode tile from 64 to 32 values per lane to improve occupancy
• validate symmetric Q1 metadata and derive bias = -scale / 2 in Metal instead of retaining redundant bias tensors
• keep narrow Q1 matrix multiplies (up to 8 flattened rows) packed by mapping rows onto the Metal grid; wider prefill still uses the dense fallback
• clear unused MLX allocator cache after Q1 prefill and document the text-only limitation and tuning overrides

Why

prism-ml/Bonsai-27B-mlx-1bit declares model_type = "qwen3_5", so it loads through qwen3_next, not the Qwen3-shaped Bonsai engine. Its generic affine paths previously called upstream MLX with bits = 1, but the pinned MLX revision has no affine Q1 kernels.

The initial working path exposed three additional costs:

1. VPT=64 made the decode shader register-pressure limited on a base M4.
2. Bonsai's symmetric metadata stored a full bias tensor even though every value can be derived from its scale; genuinely affine tensors still need the general path.
3. Any Q1 forward with more than one row dequantized the full matrix before matmul. Small speculative-verifier windows could therefore allocate multi-GiB temporaries. The narrow-M path now stays packed. This follows the Q1 direction in MLX PR #3161 without requiring an MLX fork.

Measured impact

Measured with the exact checkpoint on a 10-GPU-core base M4:

• QMV tiling A/B, 32-step release decode: 13.19 -> 17.43 tok/s (VPT=64 -> 32, +32%)
• symmetric-bias A/B, 32-step release decode: 17.44 -> 18.58 tok/s (57.33 -> 53.81 ms/token, +6.5%)
• loader accounting after GDN fusion: 396 bias tensors / 412,262,400 bytes (393.2 MiB) removed; 6 genuinely affine tensors / 7.58 MiB retained
• final single-token regression run: 18.31 tok/s, within the measured 18.3-18.6 tok/s band
• short-request physical footprint before bias elision: 6.668 GB -> 4.594 GB with post-prefill cache clearing; peak remained about 6.65 GB because wide prefill still uses dense temporaries

A temporary, evaluation-only Q4 MTP sidecar adapter (not included here) confirmed the narrow-M bottleneck with nonzero acceptance telemetry: depth-1 MTP improved from 0.99 to 10.11 tok/s after packed Q1 verification at the same 53.7% acceptance. A longer 256-token run was still slower than greedy (11.03 vs 18.82 tok/s) because rejected drafts require recurrent-cache replay. This PR therefore does not enable or claim a speculative-decoding speedup.

Validation

• exact Bonsai-27B checkpoint loaded all 64 layers and fused all 288 GDN tensor pairs
• exact 32-step release decode: 54.63 ms/token, 18.31 tok/s
• packed Q1 M=1...9 oracle against dense dequantization for affine and symmetric metadata, including BF16 M=2 / K=1024 and [2,4,K] leading dimensions
• cargo test -p higgs-models --lib -- --test-threads=1 (430 passed, 26 ignored)
• cargo clippy -p higgs-models --lib -- -D warnings
• cargo check -p higgs-models
• cargo fmt --all

Limitations / non-goals

• The checkpoint includes a vision tower, but Higgs currently exposes only its Qwen3.5 text backbone.
• Q1 forwards wider than 8 rows still dequantize to the input dtype before MLX matmul; HIGGS_BONSAI_QMM_MAX_ROWS controls the threshold.
• PrismML dSpark is not integrated here. The published Bonsai-27B dSpark drafter is GGUF, and its block-4 verification would still pay Higgs' hybrid recurrent-cache rollback/replay cost.

Summary by CodeRabbit

• New Features

  • Added support for Bonsai Q1 (1-bit) quantized Qwen3 and Qwen3.5 models.
  • Improved execution of packed matrix operations for faster decoding and prefill.
  • Added optimized handling for symmetric quantization, reducing memory usage.
  • Added support for batched Q1 matrix operations.

• Bug Fixes

  • Improved cache handling after prefill for Qwen3.5 Q1 models.
  • Preserved correct behavior for non-symmetric quantized weights.

• Documentation

  • Documented Bonsai Q1 support, configuration options, performance behavior, and current multimodal limitations.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: f9837ace-cb24-422e-a427-1c93655790bf

📥 Commits

Reviewing files that changed from the base of the PR and between 9a7437e and 363dc40.

📒 Files selected for processing (4)

• crates/higgs-engine/src/mlx_tuning.rs
• crates/higgs-models/src/metal_kernel.rs
• crates/higgs-models/src/qwen3_next.rs
• docs/BONSAI_Q1.md

---

📝 Walkthrough

Walkthrough

Bonsai Q1 support now handles symmetric affine checkpoints through sentinel bias representation, specialized Metal kernels, Qwen3.5 loader and execution paths, and related tests and documentation. Qwen3.5 Q1 metadata also enables prefill cache clearing.

Changes

Bonsai Q1 support

Layer / File(s)	Summary
Q1 metadata-driven runtime tuning crates/higgs-engine/src/mlx_tuning.rs	Model metadata records quantization bits and enables prefill cache clearing for Qwen3.5 Q1 models.
Symmetric Q1 Metal kernels crates/higgs-models/src/metal_kernel.rs	Q1 matvec, batched fast matvec, matrix multiplication, and dequantization support symmetric bias sentinels and batched outputs.
Q1 dispatch and execution paths crates/higgs-models/src/qwen3_next.rs	Q1 execution validates tensor layouts, selects packed or dense paths, handles symmetric embeddings, and avoids incompatible fused FFN execution.
Symmetric bias compaction and validation crates/higgs-models/src/qwen3_next.rs, docs/BONSAI_Q1.md	Qwen3.5 loaders optionally compact exact symmetric Q1 biases, tests cover the resulting paths, and documentation describes the supported execution modes and controls.

Estimated code review effort: 4 (Complex) | ~60 minutes

Possibly related PRs

• panbanda/higgs#96: Introduces the MLX tuning/profile system extended here for quantization-aware cache clearing.
• panbanda/higgs#182: Adds the Bonsai-Q1 JIT kernels extended here with symmetric affine-bias handling.

Suggested labels: risk: medium

Suggested reviewers: panbanda

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the PR’s main goal of adding Bonsai-27B affine Q1 support.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}