From 293387e6e7d5c8213b59f1e40902fbc04e146afb Mon Sep 17 00:00:00 2001 From: Peppi Littera Date: Tue, 14 Jul 2026 22:19:05 +0200 Subject: [PATCH 1/5] support Bonsai-27B affine Q1 --- crates/higgs-models/src/qwen3_next.rs | 106 +++++++++++++++++++++++++- docs/BONSAI_Q1.md | 30 ++++---- 2 files changed, 121 insertions(+), 15 deletions(-) diff --git a/crates/higgs-models/src/qwen3_next.rs b/crates/higgs-models/src/qwen3_next.rs index c6dfc2ed..4361e0a7 100644 --- a/crates/higgs-models/src/qwen3_next.rs +++ b/crates/higgs-models/src/qwen3_next.rs @@ -248,7 +248,61 @@ pub(crate) fn quantized_forward( group_size: i32, bits: i32, ) -> Result { - ops::quantized_matmul(x, weight, scales, biases, true, group_size, bits) + if bits == 1 { + affine_q1_forward(x, weight, scales, biases, group_size) + } else { + ops::quantized_matmul(x, weight, scales, biases, true, group_size, bits) + } +} + +/// Affine 1-bit matrix multiplication using Higgs' runtime Metal kernels. +/// +/// Upstream MLX does not provide the affine `bits=1` kernels used by Bonsai +/// checkpoints. Decode uses the fused packed matvec; multi-token forwards +/// dequantize the current matrix to the input dtype and use regular MLX matmul. This +/// is shared by the Qwen3.5 hybrid path (Bonsai-27B) and its LM head. +fn affine_q1_forward( + x: &Array, + weight: &Array, + scales: &Array, + biases: &Array, + group_size: i32, +) -> Result { + let x_shape = x.shape(); + let input_dim = x_shape + .last() + .copied() + .ok_or_else(|| Exception::custom("1-bit affine input has no dimensions"))?; + let weight_shape = weight.shape(); + let packed_dim = weight_shape + .get(1) + .copied() + .ok_or_else(|| Exception::custom("1-bit affine weight must be a matrix"))?; + let expected_input_dim = packed_dim + .checked_mul(32) + .ok_or_else(|| Exception::custom("1-bit affine input dimension overflow"))?; + if input_dim != expected_input_dim { + return Err(Exception::custom(format!( + "1-bit affine input dim {input_dim} does not match packed weight dim {expected_input_dim}" + ))); + } + if group_size <= 0 || expected_input_dim % group_size != 0 { + return Err(Exception::custom(format!( + "invalid 1-bit affine group size {group_size} for input dim {expected_input_dim}" + ))); + } + + let row_count: i32 = x_shape + .iter() + .take(x_shape.len().saturating_sub(1)) + .product(); + if row_count == 1 { + crate::metal_kernel::bonsai_q1_qmv(x, weight, scales, biases, group_size) + } else { + let dense = crate::metal_kernel::bonsai_q1_dequant(weight, scales, biases, group_size)? + .as_dtype(x.dtype())?; + x.matmul(&dense.transpose()?) + } } /// Quantized linear layer stored as raw weight/scales/biases arrays. @@ -374,7 +428,11 @@ impl QEmbedding { let w = (*self.weight).take_axis(&flat, 0)?; let s = (*self.scales).take_axis(&flat, 0)?; let b = (*self.biases).take_axis(&flat, 0)?; - let out = ops::dequantize(&w, &s, &b, self.group_size, self.bits)?; + let out = if self.bits == 1 { + crate::metal_kernel::bonsai_q1_dequant(&w, &s, &b, self.group_size)? + } else { + ops::dequantize(&w, &s, &b, self.group_size, self.bits)? + }; let mut ret_shape: Vec = shape; ret_shape.push(-1); out.reshape(&ret_shape) @@ -5290,6 +5348,50 @@ mod tests { use super::*; use crate::cache::KeyValueCache; + #[test] + fn affine_q1_linear_and_embedding_paths_match_known_values() { + let group_size = 128; + let input_dim = 128; + let weight = Array::from_slice( + &[ + 0_u32, + 0, + 0, + 0, // row 0 dequantizes to bias = 1 + u32::MAX, + u32::MAX, + u32::MAX, + u32::MAX, // row 1 dequantizes to scale + bias = 2 + ], + &[2, input_dim / 32], + ); + let scales = Array::from_slice(&[2.0_f32, 3.0], &[2, 1]); + let biases = Array::from_slice(&[1.0_f32, -1.0], &[2, 1]); + + let decode = Array::from_slice(&vec![1.0_f32; input_dim as usize], &[1, 1, input_dim]); + let decode_out = affine_q1_forward(&decode, &weight, &scales, &biases, group_size).unwrap(); + + let mut prefill_values = vec![1.0_f32; input_dim as usize]; + prefill_values.extend(vec![2.0_f32; input_dim as usize]); + let prefill = Array::from_slice(&prefill_values, &[1, 2, input_dim]); + let prefill_out = + affine_q1_forward(&prefill, &weight, &scales, &biases, group_size).unwrap(); + + let mut embedding = QEmbedding::new(group_size, 1).unwrap(); + embedding.weight = Param::new(weight); + embedding.scales = Param::new(scales); + embedding.biases = Param::new(biases); + let ids = Array::from_slice(&[0_u32, 1], &[1, 2]); + let embedding_out = embedding.forward(&ids).unwrap(); + + mlx_rs::transforms::eval([&decode_out, &prefill_out, &embedding_out]).unwrap(); + assert_eq!(decode_out.as_slice::(), &[128.0, 256.0]); + assert_eq!(prefill_out.as_slice::(), &[128.0, 256.0, 256.0, 512.0]); + let embedding_values = embedding_out.as_slice::(); + assert!(embedding_values[..128].iter().all(|value| *value == 1.0)); + assert!(embedding_values[128..].iter().all(|value| *value == 2.0)); + } + #[test] fn test_config_deserialization() { let json = r#"{ diff --git a/docs/BONSAI_Q1.md b/docs/BONSAI_Q1.md index 7211d7d8..24dacb76 100644 --- a/docs/BONSAI_Q1.md +++ b/docs/BONSAI_Q1.md @@ -1,18 +1,22 @@ # Bonsai-Q1 -Bonsai-Q1 checkpoints are Qwen3-shaped models with MLX 1-bit affine -quantization metadata: +Higgs supports MLX affine 1-bit checkpoints with `quantization.bits = 1` and +`quantization.group_size = 128` on the pinned upstream `oxideai/mlx-rs` +revision. Upstream MLX does not ship the required 1-bit affine kernels, so Higgs +provides runtime JIT Metal kernels for packed matvec and dequantization. -- `model_type = "qwen3"` -- `quantization.bits = 1` -- `quantization.group_size = 128` +Two layouts are supported: -The Higgs workspace stays on the pinned upstream `oxideai/mlx-rs` dependency. -That upstream revision does not yet include the MLX bits=1 affine Metal kernels, -so `higgs-engine` detects Bonsai-Q1 configs and returns an explicit unsupported -model error instead of routing them into the regular Qwen3 transformer loader. +- Qwen3-shaped Bonsai checkpoints use the dedicated packed engine in + `crates/higgs-models/src/bonsai_q1.rs`. +- Qwen3.5 hybrid checkpoints, including Bonsai-27B, use the existing + `qwen3_next` architecture with its affine 1-bit operations dispatched to the + same Higgs Metal kernels. -The packed loader and engine live in `crates/higgs-models/src/bonsai_q1.rs` so -the Rust-side code can be reviewed independently. Runtime enablement should wait -until bits=1 affine quantization support lands upstream in the MLX dependency -chain. +Single-token decode stays packed. Embedding lookup and multi-token prefill +dequantize the selected matrix to the input dtype before using regular MLX +matmul. + +Qwen3.5 checkpoints packaged as multimodal models currently load the text +backbone only. Their vision tower is not exposed by Higgs, so image input remains +unsupported for those checkpoints. From 7871714416e73df4bf360035d470946e8a14ba6e Mon Sep 17 00:00:00 2001 From: Peppi Littera Date: Tue, 14 Jul 2026 23:18:08 +0200 Subject: [PATCH 2/5] tune Bonsai Q1 performance and residency --- crates/higgs-engine/src/mlx_tuning.rs | 40 ++++++++++++++++++++++--- crates/higgs-models/src/metal_kernel.rs | 13 ++++---- 2 files changed, 43 insertions(+), 10 deletions(-) diff --git a/crates/higgs-engine/src/mlx_tuning.rs b/crates/higgs-engine/src/mlx_tuning.rs index 953dcb39..499774f4 100644 --- a/crates/higgs-engine/src/mlx_tuning.rs +++ b/crates/higgs-engine/src/mlx_tuning.rs @@ -127,6 +127,7 @@ enum ModelSizeClass { #[derive(Debug, Clone, Default)] struct ModelMetadata { model_type: Option, + quantization_bits: Option, num_hidden_layers: Option, hidden_size: Option, max_position_embeddings: Option, @@ -143,6 +144,11 @@ impl ModelMetadata { Self { model_type: config_lookup_str(&config, "model_type").map(str::to_owned), + quantization_bits: config + .get("quantization") + .and_then(|quantization| quantization.get("bits")) + .and_then(serde_json::Value::as_u64) + .and_then(|bits| u8::try_from(bits).ok()), num_hidden_layers: config_lookup_u64(&config, "num_hidden_layers") .and_then(|v| usize::try_from(v).ok()), hidden_size: config_lookup_u64(&config, "hidden_size") @@ -192,6 +198,10 @@ impl ModelMetadata { fn is_long_context(&self) -> bool { self.max_position_embeddings.unwrap_or_default() >= 65_536 } + + fn should_clear_cache_after_prefill(&self) -> bool { + matches!(self.model_type.as_deref(), Some("qwen3_5")) && self.quantization_bits == Some(1) + } } #[derive(Debug, Clone)] @@ -263,6 +273,7 @@ impl MlxRuntimeTuning { balanced_chunked_prefill(size_class, is_long_context, is_moe); let balanced_paged_kv = heuristic_paged_kv_target_bytes(metadata, size_class, is_moe); let default_mtp_draft_n_max = default_mtp_draft_n_max(size_class); + let clear_cache_after_prefill = metadata.should_clear_cache_after_prefill(); match resolved_profile { ResolvedMlxProfile::Baseline => Self { @@ -270,7 +281,7 @@ impl MlxRuntimeTuning { resolved_profile, chunked_prefill_threshold: DEFAULT_CHUNKED_PREFILL_THRESHOLD, chunked_prefill_chunk_size: DEFAULT_CHUNKED_PREFILL_CHUNK_SIZE, - clear_cache_after_prefill: false, + clear_cache_after_prefill, enable_mtp: false, mtp_draft_n_max: 1, paged_kv_target_bytes: DEFAULT_PAGED_KV_TARGET_BYTES, @@ -280,7 +291,7 @@ impl MlxRuntimeTuning { resolved_profile, chunked_prefill_threshold: (balanced_threshold.saturating_mul(2)).min(4096), chunked_prefill_chunk_size: balanced_chunk.max(768), - clear_cache_after_prefill: false, + clear_cache_after_prefill, enable_mtp: true, mtp_draft_n_max: default_mtp_draft_n_max, paged_kv_target_bytes: clamp_paged_kv_target_bytes( @@ -292,7 +303,7 @@ impl MlxRuntimeTuning { resolved_profile, chunked_prefill_threshold: balanced_threshold, chunked_prefill_chunk_size: balanced_chunk, - clear_cache_after_prefill: false, + clear_cache_after_prefill, enable_mtp: true, mtp_draft_n_max: default_mtp_draft_n_max, paged_kv_target_bytes: balanced_paged_kv, @@ -302,7 +313,7 @@ impl MlxRuntimeTuning { resolved_profile, chunked_prefill_threshold: balanced_threshold.max(1024), chunked_prefill_chunk_size: balanced_chunk.max(1024), - clear_cache_after_prefill: false, + clear_cache_after_prefill, enable_mtp: true, mtp_draft_n_max: default_mtp_draft_n_max, paged_kv_target_bytes: clamp_paged_kv_target_bytes( @@ -611,6 +622,27 @@ mod tests { ); } + #[test] + fn test_qwen35_q1_clears_prefill_allocator_cache_by_default() { + let metadata = ModelMetadata { + model_type: Some("qwen3_5".to_owned()), + quantization_bits: Some(1), + ..ModelMetadata::default() + }; + let tuning = MlxRuntimeTuning::from_profile( + RequestedMlxProfile::Latency, + ResolvedMlxProfile::Latency, + &metadata, + ); + assert!(tuning.clear_cache_after_prefill()); + + let non_q1 = ModelMetadata { + quantization_bits: Some(2), + ..metadata + }; + assert!(!non_q1.should_clear_cache_after_prefill()); + } + fn write_json(path: &std::path::Path, value: &serde_json::Value) -> std::io::Result<()> { let bytes = serde_json::to_vec_pretty(value).map_err(|error| { std::io::Error::other(format!("failed to serialize JSON fixture: {error}")) diff --git a/crates/higgs-models/src/metal_kernel.rs b/crates/higgs-models/src/metal_kernel.rs index ce81ab06..da3b86ce 100644 --- a/crates/higgs-models/src/metal_kernel.rs +++ b/crates/higgs-models/src/metal_kernel.rs @@ -378,18 +378,19 @@ pub fn bonsai_q1_qmv( // // Ports MLX/PrismML `qmv_fast` tiling onto our uint32 packing: each simdgroup // computes RESULTS_PER_SIMDGROUP (4) output rows; each of its 32 lanes holds -// VPT (64) input values in registers (no threadgroup memory, no barriers) and -// reuses them across all 4 rows. block_size = 64 * 32 = 2048. The bits=1 affine +// VPT (32) input values in registers (no threadgroup memory, no barriers) and +// reuses them across all 4 rows. Keeping one packed word per lane reduces +// register pressure and raises occupancy for 1-bit weights. The bits=1 affine // math is identical to the legacy kernel — `scale * sum(bit*x) + bias * sum(x)` // — only the data movement differs. Group scales/biases are per-lane (a lane's -// 64 values lie in one 128-wide group); per-row partials are simd_sum-reduced. +// 32 values lie in one 128-wide group); per-row partials are simd_sum-reduced. // --------------------------------------------------------------------------- const FAST_QMV_KERNEL_SOURCE: &str = r" -constexpr int VPT = 64; // values_per_thread +constexpr int VPT = 32; // values_per_thread (one packed word per lane) constexpr int RPS = 4; // results_per_simdgroup -constexpr int WPT = VPT / 32; // packed uint32 words per thread (2) -constexpr int BLK = VPT * 32; // block_size = 2048 +constexpr int WPT = VPT / 32; // packed uint32 words per thread (1) +constexpr int BLK = VPT * 32; // block_size = 1024 uint tgx = threadgroup_position_in_grid.x; uint sg = simdgroup_index_in_threadgroup; From 3cef4d79addd0506e49f46d8b04bf808d3a8f54e Mon Sep 17 00:00:00 2001 From: Peppi Littera Date: Tue, 14 Jul 2026 23:25:19 +0200 Subject: [PATCH 3/5] avoid exact float comparisons in Q1 test --- crates/higgs-models/src/qwen3_next.rs | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/crates/higgs-models/src/qwen3_next.rs b/crates/higgs-models/src/qwen3_next.rs index 4361e0a7..a27ed7b8 100644 --- a/crates/higgs-models/src/qwen3_next.rs +++ b/crates/higgs-models/src/qwen3_next.rs @@ -5388,8 +5388,16 @@ mod tests { assert_eq!(decode_out.as_slice::(), &[128.0, 256.0]); assert_eq!(prefill_out.as_slice::(), &[128.0, 256.0, 256.0, 512.0]); let embedding_values = embedding_out.as_slice::(); - assert!(embedding_values[..128].iter().all(|value| *value == 1.0)); - assert!(embedding_values[128..].iter().all(|value| *value == 2.0)); + assert!( + embedding_values[..128] + .iter() + .all(|value| (*value - 1.0).abs() <= f32::EPSILON) + ); + assert!( + embedding_values[128..] + .iter() + .all(|value| (*value - 2.0).abs() <= f32::EPSILON) + ); } #[test] From 2220c19bc0edfc67d1d137fe987fd32abc6f004d Mon Sep 17 00:00:00 2001 From: Peppi Littera Date: Wed, 15 Jul 2026 00:06:30 +0200 Subject: [PATCH 4/5] optimize symmetric Bonsai Q1 decode --- crates/higgs-models/src/metal_kernel.rs | 58 +++++- crates/higgs-models/src/qwen3_next.rs | 236 +++++++++++++++++++++++- docs/BONSAI_Q1.md | 6 + 3 files changed, 283 insertions(+), 17 deletions(-) diff --git a/crates/higgs-models/src/metal_kernel.rs b/crates/higgs-models/src/metal_kernel.rs index da3b86ce..dcf7f315 100644 --- a/crates/higgs-models/src/metal_kernel.rs +++ b/crates/higgs-models/src/metal_kernel.rs @@ -16,6 +16,8 @@ //! [`crate::qwen3_next`]; the kernel math mirrors //! [`crate::bonsai_q1::PackedQ1Linear::dequant_row_to_fp32`]: //! `W[r,c] = scale[r, c/G] * bit + bias[r, c/G]`, `bit = (w[r, c/32] >> (c%32)) & 1`. +//! Checkpoints whose affine metadata is symmetric use an empty bias sentinel; +//! their kernels derive `bias = -scale / 2` and never read a bias buffer. use std::ffi::{CStr, CString, c_char, c_void}; use std::sync::OnceLock; @@ -154,7 +156,7 @@ for (int k_off = 0; k_off < K; k_off += CHUNK) { int g = idx * 32 / GroupSize; float s_val = float(sc[row * NumGroups + g]); - float b_val = float(bi[row * NumGroups + g]); + float b_val = Symmetric ? (-0.5f * s_val) : float(bi[row * NumGroups + g]); acc += s_val * dot_val + b_val * sum_x; } } @@ -197,6 +199,7 @@ fn configure_qmv_kernel( n_rows: i32, k_dim: i32, group_size: i32, + symmetric: bool, ) -> mlx_sys::mlx_fast_metal_kernel_config { unsafe { let config = mlx_sys::mlx_fast_metal_kernel_config_new(); @@ -221,6 +224,11 @@ fn configure_qmv_kernel( c"NumGroups".as_ptr(), k_dim / group_size, ); + mlx_sys::mlx_fast_metal_kernel_config_add_template_arg_int( + config, + c"Symmetric".as_ptr(), + i32::from(symmetric), + ); let nsg = qmv_nsg(k_dim); let n_tgs = (n_rows + nsg - 1) / nsg; @@ -266,13 +274,21 @@ pub fn bonsai_q1_qmv_legacy( let x_flat = x.reshape(&[k_dim])?; let w_flat = weight.reshape(&[-1])?; let s_flat = scales.flatten(None, None)?; - let b_flat = biases.flatten(None, None)?; + let symmetric = biases.size() == 0; + // FastMetal still binds the affine input signature. Reuse the scale array + // as a harmless dummy; the `Symmetric` template constant removes the bias + // load from the compiled kernel. + let b_flat = if symmetric { + s_flat.clone() + } else { + biases.flatten(None, None)? + }; let stream = Stream::task_local_or_default(); let out_dtype = unsafe { mlx_sys::mlx_array_dtype(x.as_ptr()) }; let cached = QMV_KERNEL.get_or_init(|| CachedMetalKernel(create_qmv_kernel())); - let config = configure_qmv_kernel(out_dtype, n_rows, k_dim, group_size); + let config = configure_qmv_kernel(out_dtype, n_rows, k_dim, group_size, symmetric); let n_scalar = unsafe { mlx_sys::mlx_array_new_int(n_rows) }; let input_ptrs = [ @@ -436,7 +452,7 @@ for (int k = 0; k < aligned_end; k += BLK) { } } float s_val = float(sc[row * NumGroups + g]); - float b_val = float(bi[row * NumGroups + g]); + float b_val = Symmetric ? (-0.5f * s_val) : float(bi[row * NumGroups + g]); result[r] += s_val * accum + b_val * sum; } } @@ -476,7 +492,7 @@ if (aligned_end < K) { } } float s_val = float(sc[row * NumGroups + g]); - float b_val = float(bi[row * NumGroups + g]); + float b_val = Symmetric ? (-0.5f * s_val) : float(bi[row * NumGroups + g]); result[r] += s_val * accum + b_val * sum; } } @@ -517,6 +533,7 @@ fn configure_fast_qmv_kernel( n_rows: i32, k_dim: i32, group_size: i32, + symmetric: bool, ) -> mlx_sys::mlx_fast_metal_kernel_config { unsafe { let config = mlx_sys::mlx_fast_metal_kernel_config_new(); @@ -541,6 +558,11 @@ fn configure_fast_qmv_kernel( c"NumGroups".as_ptr(), k_dim / group_size, ); + mlx_sys::mlx_fast_metal_kernel_config_add_template_arg_int( + config, + c"Symmetric".as_ptr(), + i32::from(symmetric), + ); // Each simdgroup computes 4 rows; nsg simdgroups per threadgroup. let nsg = fast_qmv_nsg(); @@ -587,13 +609,18 @@ pub fn bonsai_q1_qmv_fast( let x_flat = x.reshape(&[k_dim])?; let w_flat = weight.reshape(&[-1])?; let s_flat = scales.flatten(None, None)?; - let b_flat = biases.flatten(None, None)?; + let symmetric = biases.size() == 0; + let b_flat = if symmetric { + s_flat.clone() + } else { + biases.flatten(None, None)? + }; let stream = Stream::task_local_or_default(); let out_dtype = unsafe { mlx_sys::mlx_array_dtype(x.as_ptr()) }; let cached = FAST_QMV_KERNEL.get_or_init(|| CachedMetalKernel(create_fast_qmv_kernel())); - let config = configure_fast_qmv_kernel(out_dtype, n_rows, k_dim, group_size); + let config = configure_fast_qmv_kernel(out_dtype, n_rows, k_dim, group_size, symmetric); let n_scalar = unsafe { mlx_sys::mlx_array_new_int(n_rows) }; let input_ptrs = [ @@ -661,7 +688,7 @@ uint packed = w[gid]; int g = int(idx) * 32 / GroupSize; float s_val = float(sc[n * uint(NumGroups) + uint(g)]); -float b_val = float(bi[n * uint(NumGroups) + uint(g)]); +float b_val = Symmetric ? (-0.5f * s_val) : float(bi[n * uint(NumGroups) + uint(g)]); uint base = n * uint(K) + idx * 32u; for (uint j = 0u; j < 32u; ++j) { @@ -697,6 +724,7 @@ fn configure_dequant_kernel( n_rows: i32, k_dim: i32, group_size: i32, + symmetric: bool, ) -> mlx_sys::mlx_fast_metal_kernel_config { let k_packed = k_dim / 32; let n_words = n_rows * k_packed; @@ -728,6 +756,11 @@ fn configure_dequant_kernel( c"NWords".as_ptr(), n_words, ); + mlx_sys::mlx_fast_metal_kernel_config_add_template_arg_int( + config, + c"Symmetric".as_ptr(), + i32::from(symmetric), + ); let tg: i32 = 256; let grid = ((n_words + tg - 1) / tg) * tg; @@ -770,13 +803,18 @@ pub fn bonsai_q1_dequant( let w_flat = weight.reshape(&[-1])?; let s_flat = scales.flatten(None, None)?; - let b_flat = biases.flatten(None, None)?; + let symmetric = biases.size() == 0; + let b_flat = if symmetric { + s_flat.clone() + } else { + biases.flatten(None, None)? + }; let stream = Stream::task_local_or_default(); let out_dtype = unsafe { mlx_sys::mlx_array_dtype(scales.as_ptr()) }; let cached = DEQUANT_KERNEL.get_or_init(|| CachedMetalKernel(create_dequant_kernel())); - let config = configure_dequant_kernel(out_dtype, n_rows, k_dim, group_size); + let config = configure_dequant_kernel(out_dtype, n_rows, k_dim, group_size, symmetric); let input_ptrs = [w_flat.as_ptr(), s_flat.as_ptr(), b_flat.as_ptr()]; let inputs_vec = diff --git a/crates/higgs-models/src/qwen3_next.rs b/crates/higgs-models/src/qwen3_next.rs index a27ed7b8..b27fcb33 100644 --- a/crates/higgs-models/src/qwen3_next.rs +++ b/crates/higgs-models/src/qwen3_next.rs @@ -240,6 +240,19 @@ pub(crate) fn init_quantized_params() -> QuantizedParams { (placeholder(), placeholder(), placeholder()) } +/// Zero-sized marker stored in place of validated symmetric Q1 biases. +/// +/// Q1 affine weights with `bias = -scale / 2` need no resident bias buffer; +/// the runtime Metal kernels derive it from the scale. A zero-sized array is +/// distinct from the `[1]` unloaded-parameter placeholder used by the loader. +fn symmetric_q1_bias_sentinel() -> Array { + Array::from_slice::(&[], &[0]) +} + +fn has_symmetric_q1_biases(biases: &Array) -> bool { + biases.size() == 0 +} + pub(crate) fn quantized_forward( x: &Array, weight: &Array, @@ -427,10 +440,15 @@ impl QEmbedding { let flat = indices.flatten(None, None)?; let w = (*self.weight).take_axis(&flat, 0)?; let s = (*self.scales).take_axis(&flat, 0)?; - let b = (*self.biases).take_axis(&flat, 0)?; let out = if self.bits == 1 { - crate::metal_kernel::bonsai_q1_dequant(&w, &s, &b, self.group_size)? + if has_symmetric_q1_biases(&self.biases) { + crate::metal_kernel::bonsai_q1_dequant(&w, &s, &self.biases, self.group_size)? + } else { + let b = (*self.biases).take_axis(&flat, 0)?; + crate::metal_kernel::bonsai_q1_dequant(&w, &s, &b, self.group_size)? + } } else { + let b = (*self.biases).take_axis(&flat, 0)?; ops::dequantize(&w, &s, &b, self.group_size, self.bits)? }; let mut ret_shape: Vec = shape; @@ -3156,6 +3174,20 @@ impl FfnBlock { } fn dense_hidden_fused(&mut self, x: &Array, use_fused_gemv: bool) -> Result { + // The optional persistent fusion path expects materialized affine bias + // arrays. Symmetric Q1 deliberately drops them; keep the memory-saving + // representation and use the normal two-projection path instead. + if self + .gate_proj + .as_ref() + .is_some_and(|proj| has_symmetric_q1_biases(&proj.biases)) + || self + .up_proj + .as_ref() + .is_some_and(|proj| has_symmetric_q1_biases(&proj.biases)) + { + return self.dense_hidden_separate(x); + } if self.fused_gate_up.is_none() { let gp = self .gate_proj @@ -4388,6 +4420,85 @@ where ))) } +#[derive(Debug, Default, PartialEq, Eq)] +struct SymmetricQ1Compaction { + tensors: usize, + bytes: usize, +} + +fn symmetric_q1_compaction_enabled() -> bool { + !std::env::var("HIGGS_BONSAI_SYMMETRIC_Q1").is_ok_and(|raw| { + matches!( + raw.trim().to_ascii_lowercase().as_str(), + "0" | "false" | "off" | "no" + ) + }) +} + +/// Whether a Q1 affine bias tensor is exactly `-scale / 2` under the Float32 +/// arithmetic used by the Metal kernel. Any deviation keeps the original bias +/// tensor and therefore preserves the generic affine fallback exactly. +fn q1_biases_are_symmetric(scales: &Array, biases: &Array) -> Result { + if scales.shape() != biases.shape() + || scales.size() == 0 + || scales.shape() == [1] + || biases.shape() == [1] + { + return Ok(false); + } + + let scales_f32 = scales.as_dtype(Dtype::Float32).map_err(ModelError::Mlx)?; + let biases_f32 = biases.as_dtype(Dtype::Float32).map_err(ModelError::Mlx)?; + let expected = scales_f32 + .multiply(&Array::from_f32(-0.5)) + .map_err(ModelError::Mlx)?; + let equal = biases_f32 + .array_eq(&expected, None) + .map_err(ModelError::Mlx)?; + equal.try_item::().map_err(ModelError::Mlx) +} + +/// Validate every loaded Q1 scale/bias pair, then replace only symmetric bias +/// tensors with a zero-sized marker. Non-symmetric affine tensors remain fully +/// supported and continue through the existing bias-reading kernels. +fn compact_symmetric_q1_biases( + params: &mut HashMap, &mut Array>, +) -> Result { + let bias_keys = params + .keys() + .filter(|key| key.ends_with(".biases")) + .map(|key| key.to_string()) + .collect::>(); + let mut compacted = SymmetricQ1Compaction::default(); + + for bias_key in bias_keys { + let Some(scale_key) = bias_key.strip_suffix(".biases") else { + continue; + }; + let scale_key = format!("{scale_key}.scales"); + let Some(scales) = params + .get(scale_key.as_str()) + .map(|value| (**value).clone()) + else { + continue; + }; + let Some(biases) = params.get(bias_key.as_str()).map(|value| (**value).clone()) else { + continue; + }; + if !q1_biases_are_symmetric(&scales, &biases)? { + continue; + } + + compacted.tensors += 1; + compacted.bytes = compacted.bytes.saturating_add(biases.nbytes()); + if let Some(param) = params.get_mut(bias_key.as_str()) { + **param = symmetric_q1_bias_sentinel(); + } + } + + Ok(compacted) +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum MtpWeightLayout { None, @@ -4762,7 +4873,13 @@ pub fn load_qwen3_5_model>(model_dir: P) -> Result>( // var or mixed-bit BA detection in load_qwen3_5_moe_text_config_args), and // falls back to separate projections at runtime if fusion finds a // shape-incompatible BA pair. - let model = load_qwen3_5_model_with_gdn_fallback(model_path, args, &gdn_dims)?; + let model = load_qwen3_5_model_with_gdn_fallback(model_path, args, &gdn_dims, false)?; tracing::info!("Qwen3.5-MoE model loaded successfully"); Ok(model) @@ -4819,19 +4936,25 @@ fn load_qwen3_5_model_with_gdn_fallback( model_path: &Path, mut args: Qwen3NextModelArgs, gdn_dims: &GdnDims, + compact_symmetric_q1: bool, ) -> Result { let force_separate = args.use_separate_gdn_projections || std::env::var("HIGGS_SEPARATE_GDN_PROJ").is_ok(); if force_separate { args.use_separate_gdn_projections = true; let mut model = Qwen3NextCausalLM::new(args)?; - load_qwen3_5_moe_weights_direct(&mut model, model_path)?; + load_qwen3_5_moe_weights_direct(&mut model, model_path, compact_symmetric_q1)?; tracing::info!("Using SEPARATE GDN projections (4 dispatches per layer)"); return Ok(model); } let mut fused_model = Qwen3NextCausalLM::new(args.clone())?; - match load_qwen3_5_moe_weights_fused(&mut fused_model, model_path, gdn_dims) { + match load_qwen3_5_moe_weights_fused( + &mut fused_model, + model_path, + gdn_dims, + compact_symmetric_q1, + ) { Ok(()) => { tracing::info!("Using FUSED GDN projections (2 dispatches per layer)"); Ok(fused_model) @@ -4843,7 +4966,7 @@ fn load_qwen3_5_model_with_gdn_fallback( ); args.use_separate_gdn_projections = true; let mut separate_model = Qwen3NextCausalLM::new(args)?; - load_qwen3_5_moe_weights_direct(&mut separate_model, model_path)?; + load_qwen3_5_moe_weights_direct(&mut separate_model, model_path, compact_symmetric_q1)?; tracing::info!( "Using SEPARATE GDN projections (4 dispatches per layer, mixed-bit fallback)" ); @@ -5108,6 +5231,7 @@ fn load_qwen3_next_weights( fn load_qwen3_5_moe_weights_direct( model: &mut M, model_path: &Path, + compact_symmetric_q1: bool, ) -> Result<(), crate::error::ModelError> { let safetensors_files = crate::collect_safetensors_files(model_path)?; let mut params = model.parameters_mut().flatten(); @@ -5168,6 +5292,15 @@ fn load_qwen3_5_moe_weights_direct( )?; tracing::info!(param_count, matched, "Total model parameters loaded"); + if compact_symmetric_q1 { + let compacted = compact_symmetric_q1_biases(&mut params)?; + tracing::info!( + tensors = compacted.tensors, + bytes = compacted.bytes, + "Dropped validated symmetric Q1 bias tensors" + ); + } + model .eval() .map_err(|e| crate::error::ModelError::Io(std::io::Error::other(e.to_string())))?; @@ -5182,6 +5315,7 @@ fn load_qwen3_5_moe_weights_fused( model: &mut M, model_path: &Path, gdn_dims: &GdnDims, + compact_symmetric_q1: bool, ) -> Result<(), crate::error::ModelError> { use std::collections::HashMap; @@ -5295,6 +5429,15 @@ fn load_qwen3_5_moe_weights_fused( .map(|(name, value)| (std::rc::Rc::::clone(name), &**value)), )?; + if compact_symmetric_q1 { + let compacted = compact_symmetric_q1_biases(&mut params)?; + tracing::info!( + tensors = compacted.tensors, + bytes = compacted.bytes, + "Dropped validated symmetric Q1 bias tensors" + ); + } + model .eval() .map_err(|e| crate::error::ModelError::Io(std::io::Error::other(e.to_string())))?; @@ -5400,6 +5543,85 @@ mod tests { ); } + #[test] + fn symmetric_q1_linear_and_embedding_paths_derive_bias() { + let group_size = 128; + let input_dim = 128; + let weight = Array::from_slice( + &[ + 0_u32, + 0, + 0, + 0, // row 0: bit=0, scale=2 => -1 + u32::MAX, + u32::MAX, + u32::MAX, + u32::MAX, // row 1: bit=1, scale=4 => +2 + ], + &[2, input_dim / 32], + ); + let scales = Array::from_slice(&[2.0_f32, 4.0], &[2, 1]); + let no_biases = symmetric_q1_bias_sentinel(); + + let decode = Array::from_slice(&vec![1.0_f32; input_dim as usize], &[1, 1, input_dim]); + let decode_out = + affine_q1_forward(&decode, &weight, &scales, &no_biases, group_size).unwrap(); + + let mut prefill_values = vec![1.0_f32; input_dim as usize]; + prefill_values.extend(vec![2.0_f32; input_dim as usize]); + let prefill = Array::from_slice(&prefill_values, &[1, 2, input_dim]); + let prefill_out = + affine_q1_forward(&prefill, &weight, &scales, &no_biases, group_size).unwrap(); + + let mut embedding = QEmbedding::new(group_size, 1).unwrap(); + embedding.weight = Param::new(weight); + embedding.scales = Param::new(scales); + embedding.biases = Param::new(no_biases); + let ids = Array::from_slice(&[0_u32, 1], &[1, 2]); + let embedding_out = embedding.forward(&ids).unwrap(); + + mlx_rs::transforms::eval([&decode_out, &prefill_out, &embedding_out]).unwrap(); + assert_eq!(decode_out.as_slice::(), &[-128.0, 256.0]); + assert_eq!( + prefill_out.as_slice::(), + &[-128.0, 256.0, -256.0, 512.0] + ); + let embedding_values = embedding_out.as_slice::(); + assert!( + embedding_values[..128] + .iter() + .all(|value| (*value + 1.0).abs() <= f32::EPSILON) + ); + assert!( + embedding_values[128..] + .iter() + .all(|value| (*value - 2.0).abs() <= f32::EPSILON) + ); + } + + #[test] + fn symmetric_q1_bias_validation_and_compaction_preserve_affine_fallback() { + const FP16_MIN_SUBNORMAL: f32 = 5.960_464_5e-8; + + let mut scales = Array::from_slice(&[2.0_f32, 4.0, 6.0, 2.0 * FP16_MIN_SUBNORMAL], &[2, 2]); + let mut symmetric = + Array::from_slice(&[-1.0_f32, -2.0, -3.0, -FP16_MIN_SUBNORMAL], &[2, 2]); + assert!(q1_biases_are_symmetric(&scales, &symmetric).unwrap()); + + let asymmetric = Array::from_slice(&[-1.0_f32, -2.0, -3.0, 0.0], &[2, 2]); + assert!(!q1_biases_are_symmetric(&scales, &asymmetric).unwrap()); + + let mut params = HashMap::new(); + params.insert(std::rc::Rc::::from("layer.scales"), &mut scales); + params.insert(std::rc::Rc::::from("layer.biases"), &mut symmetric); + let compacted = compact_symmetric_q1_biases(&mut params).unwrap(); + drop(params); + + assert_eq!(compacted.tensors, 1); + assert_eq!(compacted.bytes, 4 * std::mem::size_of::()); + assert!(has_symmetric_q1_biases(&symmetric)); + } + #[test] fn test_config_deserialization() { let json = r#"{ diff --git a/docs/BONSAI_Q1.md b/docs/BONSAI_Q1.md index 24dacb76..4601295e 100644 --- a/docs/BONSAI_Q1.md +++ b/docs/BONSAI_Q1.md @@ -17,6 +17,12 @@ Single-token decode stays packed. Embedding lookup and multi-token prefill dequantize the selected matrix to the input dtype before using regular MLX matmul. +For Qwen3.5 Q1 checkpoints, the loader validates every affine scale/bias pair. +When a tensor is exactly symmetric (`bias = -scale / 2`), Higgs releases its +bias array and derives the bias in the Metal kernel. Any non-symmetric tensor +keeps the general affine path. Set `HIGGS_BONSAI_SYMMETRIC_Q1=0` to retain all +bias tensors for A/B debugging. + Qwen3.5 checkpoints packaged as multimodal models currently load the text backbone only. Their vision tower is not exposed by Higgs, so image input remains unsupported for those checkpoints. From 363dc4045b00a4e7f99093b742b3f5daefe59755 Mon Sep 17 00:00:00 2001 From: Peppi Littera Date: Wed, 15 Jul 2026 00:54:10 +0200 Subject: [PATCH 5/5] keep narrow Bonsai Q1 verifies packed --- crates/higgs-models/src/metal_kernel.rs | 57 +++++--- crates/higgs-models/src/qwen3_next.rs | 172 +++++++++++++++++++++++- docs/BONSAI_Q1.md | 7 +- 3 files changed, 211 insertions(+), 25 deletions(-) diff --git a/crates/higgs-models/src/metal_kernel.rs b/crates/higgs-models/src/metal_kernel.rs index dcf7f315..95d49309 100644 --- a/crates/higgs-models/src/metal_kernel.rs +++ b/crates/higgs-models/src/metal_kernel.rs @@ -390,16 +390,18 @@ pub fn bonsai_q1_qmv( } // --------------------------------------------------------------------------- -// `qmv_fast`-class 1-bit matvec (decode hot path). +// `qmv_fast`-class 1-bit narrow matrix multiply (decode / verify hot path). // // Ports MLX/PrismML `qmv_fast` tiling onto our uint32 packing: each simdgroup -// computes RESULTS_PER_SIMDGROUP (4) output rows; each of its 32 lanes holds -// VPT (32) input values in registers (no threadgroup memory, no barriers) and -// reuses them across all 4 rows. Keeping one packed word per lane reduces -// register pressure and raises occupancy for 1-bit weights. The bits=1 affine -// math is identical to the legacy kernel — `scale * sum(bit*x) + bias * sum(x)` -// — only the data movement differs. Group scales/biases are per-lane (a lane's -// 32 values lie in one 128-wide group); per-row partials are simd_sum-reduced. +// computes RESULTS_PER_SIMDGROUP (4) output rows for one input row; the grid's +// z dimension covers narrow M > 1 verifier batches without materializing the +// dense weight matrix. Each lane holds VPT (32) input values in registers and +// reuses them across all 4 output rows. Keeping one packed word per lane +// reduces register pressure and raises occupancy for 1-bit weights. The bits=1 +// affine math is identical to the legacy kernel — +// `scale * sum(bit*x) + bias * sum(x)` — only the data movement differs. +// Group scales/biases are per-lane (a lane's 32 values lie in one 128-wide +// group); per-row partials are simd_sum-reduced. // --------------------------------------------------------------------------- const FAST_QMV_KERNEL_SOURCE: &str = r" @@ -412,8 +414,10 @@ uint tgx = threadgroup_position_in_grid.x; uint sg = simdgroup_index_in_threadgroup; uint lid = thread_index_in_simdgroup; uint nsg = simdgroups_per_threadgroup; +uint batch = threadgroup_position_in_grid.z; int out_row = int(tgx) * (int(nsg) * RPS) + int(sg) * RPS; +auto x_row = x + int(batch) * K; float xt[VPT]; float result[RPS]; @@ -421,12 +425,12 @@ for (int r = 0; r < RPS; ++r) { result[r] = 0.0f; } int aligned_end = (K / BLK) * BLK; -// Main loop: full 2048-element blocks (covers every real Bonsai layer, since +// Main loop: full 1024-element blocks (covers every real Bonsai layer, since // all K are multiples of 2048). for (int k = 0; k < aligned_end; k += BLK) { int xbase = k + int(lid) * VPT; float sum = 0.0f; - for (int i = 0; i < VPT; ++i) { float v = float(x[xbase + i]); xt[i] = v; sum += v; } + for (int i = 0; i < VPT; ++i) { float v = float(x_row[xbase + i]); xt[i] = v; sum += v; } int wcol = (k / 32) + int(lid) * WPT; int g = xbase / GroupSize; // all VPT values fall in one group @@ -463,7 +467,7 @@ if (aligned_end < K) { bool in_bounds = xbase < K; float sum = 0.0f; for (int i = 0; i < VPT; ++i) { - float v = (in_bounds && (xbase + i) < K) ? float(x[xbase + i]) : 0.0f; + float v = (in_bounds && (xbase + i) < K) ? float(x_row[xbase + i]) : 0.0f; xt[i] = v; sum += v; } @@ -501,7 +505,7 @@ for (int r = 0; r < RPS; ++r) { int row = out_row + r; float v = simd_sum(result[r]); if (lid == 0u && row < n_param) { - y[row] = OutT(v); + y[int(batch) * n_param + row] = OutT(v); } } "; @@ -531,6 +535,7 @@ fn create_fast_qmv_kernel() -> mlx_sys::mlx_fast_metal_kernel { fn configure_fast_qmv_kernel( out_dtype: mlx_sys::mlx_dtype, n_rows: i32, + m_rows: i32, k_dim: i32, group_size: i32, symmetric: bool, @@ -568,10 +573,10 @@ fn configure_fast_qmv_kernel( let nsg = fast_qmv_nsg(); let rows_per_tg = nsg * 4; let n_tgs = (n_rows + rows_per_tg - 1) / rows_per_tg; - mlx_sys::mlx_fast_metal_kernel_config_set_grid(config, n_tgs * 32, nsg, 1); + mlx_sys::mlx_fast_metal_kernel_config_set_grid(config, n_tgs * 32, nsg, m_rows); mlx_sys::mlx_fast_metal_kernel_config_set_thread_group(config, 32, nsg, 1); - let y_shape = [1, n_rows]; + let y_shape = [m_rows, n_rows]; mlx_sys::mlx_fast_metal_kernel_config_add_output_arg( config, y_shape.as_ptr(), @@ -605,8 +610,12 @@ pub fn bonsai_q1_qmv_fast( .copied() .ok_or_else(|| Exception::custom("bonsai_q1_qmv_fast: weight has no columns"))?; let k_dim = k_packed * 32; + let m_rows: i32 = x_shape + .iter() + .take(x_shape.len().saturating_sub(1)) + .product(); - let x_flat = x.reshape(&[k_dim])?; + let x_flat = x.reshape(&[m_rows, k_dim])?; let w_flat = weight.reshape(&[-1])?; let s_flat = scales.flatten(None, None)?; let symmetric = biases.size() == 0; @@ -620,7 +629,7 @@ pub fn bonsai_q1_qmv_fast( let out_dtype = unsafe { mlx_sys::mlx_array_dtype(x.as_ptr()) }; let cached = FAST_QMV_KERNEL.get_or_init(|| CachedMetalKernel(create_fast_qmv_kernel())); - let config = configure_fast_qmv_kernel(out_dtype, n_rows, k_dim, group_size, symmetric); + let config = configure_fast_qmv_kernel(out_dtype, n_rows, m_rows, k_dim, group_size, symmetric); let n_scalar = unsafe { mlx_sys::mlx_array_new_int(n_rows) }; let input_ptrs = [ @@ -671,6 +680,22 @@ pub fn bonsai_q1_qmv_fast( result } +/// Packed affine Q1 matrix multiply for narrow verifier batches. +/// +/// This shares the decode-optimized kernel with [`bonsai_q1_qmv_fast`] but +/// dispatches one grid slice per flattened input row. It intentionally targets +/// small sequence lengths: weights stay packed and resident, avoiding the very +/// large temporary produced by full dequantization. +pub fn bonsai_q1_qmm( + x: &Array, + weight: &Array, + scales: &Array, + biases: &Array, + group_size: i32, +) -> Result { + bonsai_q1_qmv_fast(x, weight, scales, biases, group_size) +} + // --------------------------------------------------------------------------- // 1-bit dequantize to dense (embedding gather + prefill matmul path). // diff --git a/crates/higgs-models/src/qwen3_next.rs b/crates/higgs-models/src/qwen3_next.rs index b27fcb33..3fb3d5d9 100644 --- a/crates/higgs-models/src/qwen3_next.rs +++ b/crates/higgs-models/src/qwen3_next.rs @@ -253,6 +253,17 @@ fn has_symmetric_q1_biases(biases: &Array) -> bool { biases.size() == 0 } +fn bonsai_q1_qmm_max_rows() -> i32 { + static MAX_ROWS: OnceLock = OnceLock::new(); + *MAX_ROWS.get_or_init(|| { + std::env::var("HIGGS_BONSAI_QMM_MAX_ROWS") + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|rows| (0..=64).contains(rows)) + .unwrap_or(8) + }) +} + pub(crate) fn quantized_forward( x: &Array, weight: &Array, @@ -271,9 +282,10 @@ pub(crate) fn quantized_forward( /// Affine 1-bit matrix multiplication using Higgs' runtime Metal kernels. /// /// Upstream MLX does not provide the affine `bits=1` kernels used by Bonsai -/// checkpoints. Decode uses the fused packed matvec; multi-token forwards -/// dequantize the current matrix to the input dtype and use regular MLX matmul. This -/// is shared by the Qwen3.5 hybrid path (Bonsai-27B) and its LM head. +/// checkpoints. Decode uses the fused packed matvec. Narrow multi-token +/// verifier batches use the same packed kernel over a z-dimension batch; wider +/// prefill inputs retain the dense dequantize + MLX matmul fallback. This is +/// shared by the Qwen3.5 hybrid path (Bonsai-27B) and its LM head. fn affine_q1_forward( x: &Array, weight: &Array, @@ -311,6 +323,8 @@ fn affine_q1_forward( .product(); if row_count == 1 { crate::metal_kernel::bonsai_q1_qmv(x, weight, scales, biases, group_size) + } else if row_count > 0 && row_count <= bonsai_q1_qmm_max_rows() { + crate::metal_kernel::bonsai_q1_qmm(x, weight, scales, biases, group_size) } else { let dense = crate::metal_kernel::bonsai_q1_dequant(weight, scales, biases, group_size)? .as_dtype(x.dtype())?; @@ -4450,7 +4464,7 @@ fn q1_biases_are_symmetric(scales: &Array, biases: &Array) -> Result>(); let mut compacted = SymmetricQ1Compaction::default(); for bias_key in bias_keys { - let Some(scale_key) = bias_key.strip_suffix(".biases") else { + let Some(scale_prefix) = bias_key.strip_suffix(".biases") else { continue; }; - let scale_key = format!("{scale_key}.scales"); + let scale_key = format!("{scale_prefix}.scales"); let Some(scales) = params .get(scale_key.as_str()) .map(|value| (**value).clone()) @@ -5599,6 +5613,150 @@ mod tests { ); } + #[test] + fn packed_q1_qmm_matches_dense_reference_for_m1_through_m9() { + const GROUP_SIZE: i32 = 128; + const K: i32 = 128; + const N: i32 = 9; + + let mut packed = Vec::with_capacity((N * K / 32) as usize); + for row in 0..N { + let word = match row % 3 { + 0 => 0_u32, + 1 => u32::MAX, + _ => 0xAAAA_AAAA, + }; + packed.extend(std::iter::repeat_n(word, (K / 32) as usize)); + } + let weight = Array::from_slice(&packed, &[N, K / 32]); + let scales = Array::from_slice( + &(0..N) + .map(|row| 0.5_f32 + row as f32 * 0.125) + .collect::>(), + &[N, 1], + ); + let affine_biases = Array::from_slice( + &(0..N) + .map(|row| -0.25_f32 + row as f32 * 0.031_25) + .collect::>(), + &[N, 1], + ); + let symmetric_biases = symmetric_q1_bias_sentinel(); + + for biases in [&affine_biases, &symmetric_biases] { + let dense = + crate::metal_kernel::bonsai_q1_dequant(&weight, &scales, biases, GROUP_SIZE) + .unwrap(); + + for m in 1..=9 { + let values = (0..m) + .flat_map(|row| { + (0..K).map(move |col| { + 0.25_f32 + row as f32 * 0.5 + (col % 7) as f32 * 0.062_5 + }) + }) + .collect::>(); + let x = Array::from_slice(&values, &[m, K]); + let actual = if m <= 8 { + crate::metal_kernel::bonsai_q1_qmm(&x, &weight, &scales, biases, GROUP_SIZE) + .unwrap() + } else { + // M=9 exercises the dense fallback at the dispatch boundary. + affine_q1_forward(&x, &weight, &scales, biases, GROUP_SIZE).unwrap() + }; + let expected = x.matmul(&dense.transpose().unwrap()).unwrap(); + mlx_rs::transforms::eval([&actual, &expected]).unwrap(); + + assert_eq!(actual.shape(), &[m, N]); + for (index, (got, want)) in actual + .as_slice::() + .iter() + .zip(expected.as_slice::()) + .enumerate() + { + let tolerance = 1e-3_f32 * want.abs().max(1.0); + assert!( + (*got - *want).abs() <= tolerance, + "M={m} value {index}: packed={got}, dense={want}, tolerance={tolerance}" + ); + } + } + + let leading = Array::from_slice(&vec![0.5_f32; (8 * K) as usize], &[2, 4, K]); + let output = affine_q1_forward(&leading, &weight, &scales, biases, GROUP_SIZE).unwrap(); + mlx_rs::transforms::eval([&output]).unwrap(); + assert_eq!(output.shape(), &[2, 4, N]); + } + + // Exercise the fast kernel's full 1024-value block with the dtype used + // by the real Bonsai-27B backbone. + const MAIN_K: i32 = 1024; + const MAIN_M: i32 = 2; + let main_weight = Array::from_slice( + &(0..N * MAIN_K / 32) + .map(|index| { + if index % 2 == 0 { + 0x5555_5555_u32 + } else { + 0xAAAA_AAAA_u32 + } + }) + .collect::>(), + &[N, MAIN_K / 32], + ); + let main_scales = Array::from_slice( + &vec![0.75_f32; (N * MAIN_K / GROUP_SIZE) as usize], + &[N, MAIN_K / GROUP_SIZE], + ) + .as_dtype(mlx_rs::Dtype::Bfloat16) + .unwrap(); + let main_biases = Array::from_slice( + &vec![-0.375_f32; (N * MAIN_K / GROUP_SIZE) as usize], + &[N, MAIN_K / GROUP_SIZE], + ) + .as_dtype(mlx_rs::Dtype::Bfloat16) + .unwrap(); + let main_x = Array::from_slice( + &(0..MAIN_M * MAIN_K) + .map(|index| 0.125_f32 + (index % 11) as f32 * 0.031_25) + .collect::>(), + &[MAIN_M, MAIN_K], + ) + .as_dtype(mlx_rs::Dtype::Bfloat16) + .unwrap(); + let main_actual = crate::metal_kernel::bonsai_q1_qmm( + &main_x, + &main_weight, + &main_scales, + &main_biases, + GROUP_SIZE, + ) + .unwrap(); + let main_dense = crate::metal_kernel::bonsai_q1_dequant( + &main_weight, + &main_scales, + &main_biases, + GROUP_SIZE, + ) + .unwrap(); + let main_expected = main_x.matmul(&main_dense.transpose().unwrap()).unwrap(); + let main_actual_f32 = main_actual.as_dtype(mlx_rs::Dtype::Float32).unwrap(); + let main_expected_f32 = main_expected.as_dtype(mlx_rs::Dtype::Float32).unwrap(); + mlx_rs::transforms::eval([&main_actual_f32, &main_expected_f32]).unwrap(); + for (index, (got, want)) in main_actual_f32 + .as_slice::() + .iter() + .zip(main_expected_f32.as_slice::()) + .enumerate() + { + let tolerance = 0.02_f32 * want.abs().max(1.0); + assert!( + (*got - *want).abs() <= tolerance, + "BF16 main block value {index}: packed={got}, dense={want}, tolerance={tolerance}" + ); + } + } + #[test] fn symmetric_q1_bias_validation_and_compaction_preserve_affine_fallback() { const FP16_MIN_SUBNORMAL: f32 = 5.960_464_5e-8; diff --git a/docs/BONSAI_Q1.md b/docs/BONSAI_Q1.md index 4601295e..28f5738a 100644 --- a/docs/BONSAI_Q1.md +++ b/docs/BONSAI_Q1.md @@ -13,9 +13,12 @@ Two layouts are supported: `qwen3_next` architecture with its affine 1-bit operations dispatched to the same Higgs Metal kernels. -Single-token decode stays packed. Embedding lookup and multi-token prefill +Single-token decode and narrow multi-token forwards stay packed. For Qwen3.5, +the packed Metal path covers up to 8 flattened rows by default, including the +small verifier batches used by speculative decoding. Wider prefill inputs dequantize the selected matrix to the input dtype before using regular MLX -matmul. +matmul. Set `HIGGS_BONSAI_QMM_MAX_ROWS=0` to disable the narrow packed path, or +raise it up to 64 for A/B testing. For Qwen3.5 Q1 checkpoints, the loader validates every affine scale/bias pair. When a tensor is exactly symmetric (`bias = -scale / 2`), Higgs releases its