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..95d49309 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 = [ @@ -374,29 +390,34 @@ 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 (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 -// 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. +// 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" -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; 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]; @@ -404,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 @@ -435,7 +456,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; } } @@ -446,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; } @@ -475,7 +496,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; } } @@ -484,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); } } "; @@ -514,8 +535,10 @@ 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, ) -> mlx_sys::mlx_fast_metal_kernel_config { unsafe { let config = mlx_sys::mlx_fast_metal_kernel_config_new(); @@ -540,15 +563,20 @@ 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(); 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(), @@ -582,17 +610,26 @@ 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 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, m_rows, k_dim, group_size, symmetric); let n_scalar = unsafe { mlx_sys::mlx_array_new_int(n_rows) }; let input_ptrs = [ @@ -643,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). // @@ -660,7 +713,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) { @@ -696,6 +749,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; @@ -727,6 +781,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; @@ -769,13 +828,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 c6dfc2ed..3fb3d5d9 100644 --- a/crates/higgs-models/src/qwen3_next.rs +++ b/crates/higgs-models/src/qwen3_next.rs @@ -240,6 +240,30 @@ 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 +} + +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, @@ -248,7 +272,64 @@ 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. 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, + 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 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())?; + x.matmul(&dense.transpose()?) + } } /// Quantized linear layer stored as raw weight/scales/biases arrays. @@ -373,8 +454,17 @@ 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 = ops::dequantize(&w, &s, &b, self.group_size, self.bits)?; + let out = if self.bits == 1 { + 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; ret_shape.push(-1); out.reshape(&ret_shape) @@ -3098,6 +3188,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 @@ -4330,6 +4434,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(std::string::ToString::to_string) + .collect::>(); + let mut compacted = SymmetricQ1Compaction::default(); + + for bias_key in bias_keys { + let Some(scale_prefix) = bias_key.strip_suffix(".biases") else { + continue; + }; + let scale_key = format!("{scale_prefix}.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, @@ -4704,7 +4887,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) @@ -4761,19 +4950,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) @@ -4785,7 +4980,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)" ); @@ -5050,6 +5245,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(); @@ -5110,6 +5306,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())))?; @@ -5124,6 +5329,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; @@ -5237,6 +5443,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())))?; @@ -5290,6 +5505,281 @@ 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).abs() <= f32::EPSILON) + ); + assert!( + embedding_values[128..] + .iter() + .all(|value| (*value - 2.0).abs() <= f32::EPSILON) + ); + } + + #[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 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; + + 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 7211d7d8..28f5738a 100644 --- a/docs/BONSAI_Q1.md +++ b/docs/BONSAI_Q1.md @@ -1,18 +1,31 @@ # 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 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. 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 +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.