Skip to content

Support Bonsai-27B affine Q1 #212

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dusterbloom wants to merge 5 commits into
base: main
Choose a base branch
from
Open

Support Bonsai-27B affine Q1 #212

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
293387e
support Bonsai-27B affine Q1
dusterbloom Jul 14, 2026
7871714
tune Bonsai Q1 performance and residency
dusterbloom Jul 14, 2026
3cef4d7
avoid exact float comparisons in Q1 test
dusterbloom Jul 14, 2026
2220c19
optimize symmetric Bonsai Q1 decode
dusterbloom Jul 14, 2026
363dc40
keep narrow Bonsai Q1 verifies packed
dusterbloom Jul 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 36 additions & 4 deletions crates/higgs-engine/src/mlx_tuning.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,7 @@ enum ModelSizeClass {
#[derive(Debug, Clone, Default)]
struct ModelMetadata {
model_type: Option<String>,
quantization_bits: Option<u8>,
num_hidden_layers: Option<usize>,
hidden_size: Option<usize>,
max_position_embeddings: Option<usize>,
Expand All @@ -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")
Expand Down Expand Up @@ -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)]
Expand Down Expand Up @@ -263,14 +273,15 @@ 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 {
requested_profile,
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,
Expand All @@ -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(
Expand All @@ -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,
Expand All @@ -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(
Expand Down Expand Up @@ -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}"))
Expand Down
118 changes: 91 additions & 27 deletions crates/higgs-models/src/metal_kernel.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
}
}
Expand Down Expand Up @@ -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();
Expand All @@ -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;
Expand Down Expand Up @@ -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 = [
Expand Down Expand Up @@ -374,42 +390,47 @@ 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];
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
Expand All @@ -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;
}
}
Expand All @@ -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;
}
Expand Down Expand Up @@ -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;
}
}
Expand All @@ -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);
}
}
";
Expand Down Expand Up @@ -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();
Expand All @@ -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(),
Expand Down Expand Up @@ -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 = [
Expand Down Expand Up @@ -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<Array, Exception> {
bonsai_q1_qmv_fast(x, weight, scales, biases, group_size)
}

// ---------------------------------------------------------------------------
// 1-bit dequantize to dense (embedding gather + prefill matmul path).
//
Expand All @@ -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) {
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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 =
Expand Down
Loading
Loading