From f6712354c5bd910409477e76e0ba935ca0330f73 Mon Sep 17 00:00:00 2001 From: "renovate[bot]" <29139614+renovate[bot]@users.noreply.github.com> Date: Wed, 29 Apr 2026 03:05:33 +0000 Subject: [PATCH 01/29] fix(deps): update rust crate toml to v1 --- Cargo.lock | 32 ++++++++++++++++++++++++++++---- crates/higgs-bench/Cargo.toml | 2 +- 2 files changed, 29 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d89e9853..6c5e44f1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1196,7 +1196,7 @@ dependencies = [ "atomic", "pear", "serde", - "toml", + "toml 0.8.23", "uncased", "version_check", ] @@ -1617,7 +1617,7 @@ dependencies = [ "serde_json", "sysinfo", "tokio", - "toml", + "toml 1.1.2+spec-1.1.0", ] [[package]] @@ -3587,6 +3587,15 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_spanned" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6662b5879511e06e8999a8a235d848113e942c9124f211511b16466ee2995f26" +dependencies = [ + "serde_core", +] + [[package]] name = "serde_urlencoded" version = "0.7.1" @@ -4162,11 +4171,26 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc1beb996b9d83529a9e75c17a1686767d148d70663143c7854d8b4a09ced362" dependencies = [ "serde", - "serde_spanned", + "serde_spanned 0.6.9", "toml_datetime 0.6.11", "toml_edit 0.22.27", ] +[[package]] +name = "toml" +version = "1.1.2+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81f3d15e84cbcd896376e6730314d59fb5a87f31e4b038454184435cd57defee" +dependencies = [ + "indexmap", + "serde_core", + "serde_spanned 1.1.1", + "toml_datetime 1.1.1+spec-1.1.0", + "toml_parser", + "toml_writer", + "winnow 1.0.2", +] + [[package]] name = "toml_datetime" version = "0.6.11" @@ -4193,7 +4217,7 @@ checksum = "41fe8c660ae4257887cf66394862d21dbca4a6ddd26f04a3560410406a2f819a" dependencies = [ "indexmap", "serde", - "serde_spanned", + "serde_spanned 0.6.9", "toml_datetime 0.6.11", "toml_write", "winnow 0.7.15", diff --git a/crates/higgs-bench/Cargo.toml b/crates/higgs-bench/Cargo.toml index 907b2359..48352e0c 100644 --- a/crates/higgs-bench/Cargo.toml +++ b/crates/higgs-bench/Cargo.toml @@ -21,7 +21,7 @@ serde = { workspace = true } serde_json = { workspace = true } sysinfo = "0.32" tokio = { workspace = true } -toml = "0.8" +toml = "1.0" [build-dependencies] built = { version = "0.7", features = ["git2"] } From 50daaef27a8c42f9fed26c0798aa5824369abf52 Mon Sep 17 00:00:00 2001 From: "renovate[bot]" <29139614+renovate[bot]@users.noreply.github.com> Date: Mon, 4 May 2026 14:50:27 +0000 Subject: [PATCH 02/29] chore(deps): update rust crate tokio to v1.52.2 --- Cargo.lock | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 81fb968e..52217d01 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -88,7 +88,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -99,7 +99,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -1060,7 +1060,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -1141,7 +1141,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2008,7 +2008,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2536,7 +2536,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -3378,7 +3378,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -3437,7 +3437,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -3687,7 +3687,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -3854,7 +3854,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -4095,9 +4095,9 @@ dependencies = [ [[package]] name = "tokio" -version = "1.52.1" +version = "1.52.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b67dee974fe86fd92cc45b7a95fdd2f99a36a6d7b0d431a231178d3d670bbcc6" +checksum = "110a78583f19d5cdb2c5ccf321d1290344e71313c6c37d43520d386027d18386" dependencies = [ "bytes", "libc", @@ -4828,7 +4828,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] From 8659d3874037595dbf592ee421f7eb64ae8296fe Mon Sep 17 00:00:00 2001 From: "renovate[bot]" <29139614+renovate[bot]@users.noreply.github.com> Date: Tue, 5 May 2026 15:38:01 +0000 Subject: [PATCH 03/29] chore(deps): update taiki-e/install-action digest to cca35ed --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2cf818ef..cae64399 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -120,7 +120,7 @@ jobs: toolchain: stable - name: Install cargo-llvm-cov - uses: taiki-e/install-action@51cd0b8c0499559d9a4d75c0f5c67bec3a894ec8 # v2 + uses: taiki-e/install-action@cca35edeb1d01366c2843b68fc3ca441446d73d3 # v2 with: tool: cargo-llvm-cov From 442063ea6af56936e9ae06f9b0bd089a18288429 Mon Sep 17 00:00:00 2001 From: "renovate[bot]" <29139614+renovate[bot]@users.noreply.github.com> Date: Tue, 5 May 2026 18:00:35 +0000 Subject: [PATCH 04/29] chore(deps): update rust crate tower-http to v0.6.9 --- Cargo.lock | 40 +++++++++++++++------------------------- 1 file changed, 15 insertions(+), 25 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 81fb968e..3c6b643e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -88,7 +88,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -99,7 +99,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -1060,7 +1060,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -1141,7 +1141,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -1990,16 +1990,6 @@ version = "2.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" -[[package]] -name = "iri-string" -version = "0.7.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "25e659a4bb38e810ebc252e53b5814ff908a8c58c2a9ce2fae1bbec24cbf4e20" -dependencies = [ - "memchr", - "serde", -] - [[package]] name = "is-terminal" version = "0.4.17" @@ -2008,7 +1998,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2536,7 +2526,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -3378,7 +3368,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -3437,7 +3427,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -3687,7 +3677,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -3851,10 +3841,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.4.2", + "getrandom 0.3.4", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -4251,9 +4241,9 @@ dependencies = [ [[package]] name = "tower-http" -version = "0.6.8" +version = "0.6.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4e6559d53cc268e5031cd8429d05415bc4cb4aefc4aa5d6cc35fbf5b924a1f8" +checksum = "a28f0d049ccfaa566e14e9663d304d8577427b368cb4710a20528690287a738b" dependencies = [ "base64 0.22.1", "bitflags 2.11.1", @@ -4261,7 +4251,6 @@ dependencies = [ "futures-util", "http", "http-body", - "iri-string", "mime", "pin-project-lite", "tokio", @@ -4269,6 +4258,7 @@ dependencies = [ "tower-layer", "tower-service", "tracing", + "url", "uuid", ] @@ -4828,7 +4818,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] From 229c1110f2da4c3d93ac79c8f28c2435e0c97aad Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Wed, 6 May 2026 14:12:21 +0200 Subject: [PATCH 05/29] feat(cache): AnyCache::trim_by dispatcher for spec-decode rollback (#143) Adds AnyCache::trim_by to roll back KV layers for speculative decode while leaving hybrid Arrays state untouched.\n\nCI: https://github.com/panbanda/higgs/actions/runs/25312580791 --- crates/higgs-models/src/lib.rs | 78 ++++++++++++++++++++++++++++++++++ 1 file changed, 78 insertions(+) diff --git a/crates/higgs-models/src/lib.rs b/crates/higgs-models/src/lib.rs index 6aec6383..407f9fbf 100644 --- a/crates/higgs-models/src/lib.rs +++ b/crates/higgs-models/src/lib.rs @@ -91,6 +91,29 @@ pub enum AnyCache { Hybrid(Vec>), } +impl AnyCache { + /// Trim every layer cache by `count` tokens, discarding the most recent + /// entries. Used after speculative-decode verify to roll back rejected + /// draft tokens. Hybrid SSM (recurrent) layers are intentionally left + /// untouched — their state cannot be trimmed by offset alone. + pub fn trim_by(&mut self, count: usize) { + match self { + Self::KV(layers) => { + for layer in layers.iter_mut().flatten() { + layer.trim_by(count); + } + } + Self::Hybrid(layers) => { + for layer in layers.iter_mut().flatten() { + if let LayerCache::KV(kv) = layer { + kv.trim_by(count); + } + } + } + } + } +} + /// Unified model wrapper dispatching to the correct architecture. pub enum AnyModel { /// Standard transformer architectures: Llama, Mistral, Qwen2/2.5, Qwen3. @@ -1189,6 +1212,7 @@ fn remap_quantized_key(key: &str) -> Option { #[allow(clippy::panic, clippy::unwrap_used, clippy::indexing_slicing)] mod tests { use super::*; + use crate::cache::KeyValueCache; fn params(temp: f32, top_p: f32) -> SamplingParams { SamplingParams { @@ -1695,4 +1719,58 @@ mod tests { assert!((vals[1] - 2.0).abs() < 1e-5); assert!((vals[2] - 4.5).abs() < 1e-5); } + + // --- AnyCache::trim_by tests --- + + #[test] + fn any_cache_trim_by_kv_dispatches_to_each_layer() { + // Two KV layers, both at offset 0; trim_by saturates to 0. + // Verifies the dispatcher iterates None and Some(_) layers without panic. + let mut cache = AnyCache::KV(vec![ + Some(cache::SteppingKeyValueCache::new()), + None, + Some(cache::SteppingKeyValueCache::new()), + ]); + cache.trim_by(5); + if let AnyCache::KV(layers) = &cache { + assert_eq!(layers.len(), 3); + for layer in layers.iter().flatten() { + assert_eq!(layer.offset(), 0); + } + } else { + panic!("expected KV variant"); + } + } + + #[test] + fn any_cache_trim_by_hybrid_skips_arrays_layers() { + // Hybrid mixes LayerCache::KV (trimmable) and LayerCache::Arrays (recurrent, + // intentionally untouched). Verifies the dispatcher reaches into KV layers + // and leaves Arrays alone. + let mut arrays = qwen3_next::ArraysCache::new(); + arrays.offset = 7; + let mut cache = AnyCache::Hybrid(vec![ + Some(LayerCache::KV(cache::SteppingKeyValueCache::new())), + Some(LayerCache::Arrays(arrays)), + None, + ]); + cache.trim_by(3); + if let AnyCache::Hybrid(layers) = &cache { + assert_eq!(layers.len(), 3); + // KV layer trimmed (saturated at 0 since starting offset was 0) + if let Some(LayerCache::KV(kv)) = layers.first().and_then(|l| l.as_ref()) { + assert_eq!(kv.offset(), 0); + } else { + panic!("expected first layer to be KV variant"); + } + // Arrays layer offset unchanged (recurrent state, can't trim by offset) + if let Some(LayerCache::Arrays(a)) = layers.get(1).and_then(|l| l.as_ref()) { + assert_eq!(a.offset, 7, "Arrays layer offset must NOT be trimmed"); + } else { + panic!("expected second layer to be Arrays variant"); + } + } else { + panic!("expected Hybrid variant"); + } + } } From cc186160756c9d5741fc18baecbacd3666eeacf3 Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Wed, 6 May 2026 14:52:13 +0200 Subject: [PATCH 06/29] feat(qwen3_next): mixed-bit Qwen3.5 GDN BA loading fallback (#148) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(qwen3_next): mixed-bit Qwen3.5 GDN BA loading fallback Adds a fallback path for loading Qwen3.5 models with mixed-bit GDN projection weights (some layers q4, some q8 — common in unsloth's dynamic-quant variants). The default fused-projection loader fuses `in_proj_a` + `in_proj_b` into a single matmul; mixed-bit weights have incompatible shapes and the fusion fails. Behaviour: 1. Detect via `is_mixed_bit_gdn_ba_fusion_error` — matches a `ModelError::ShapeMismatch` whose message contains both `in_proj_ba` and `requires separate GDN projections`. 2. On detection, retry the load with `args.use_separate_gdn_projections = true`, taking the `load_qwen3_5_moe_weights_direct` path. Forward dispatches go from 2 to 4 GDN ops per layer — slightly slower but correct. 3. Forced separate (via `args.use_separate_gdn_projections` config or `HIGGS_SEPARATE_GDN_PROJ` env var) skips the fused attempt entirely. Also adds: * `qwen3_5_quantization_config` — parses `{group_size, bits}` from the per-layer `quantization` map in `config.json`. * `qwen3_5_mixed_ba_quantization_layers` — scans for the layers where `in_proj_a` and `in_proj_b` differ in bits or group_size. * `can_concatenate_axis0` — guard used inside `load_qwen3_5_moe_weights_fused` to emit the diagnostic `ShapeMismatch` error rather than panicking on the concat. * `load_qwen3_5_model_with_gdn_fallback` — private helper called by both `load_qwen3_5_model` (dense) and `load_qwen3_5_moe_model` (MoE), unifying the fallback path. Adaptations from feat/magic-canvas → origin/main: * The dense `load_qwen3_5_model` previously only honoured the env var; now it honours `args.use_separate_gdn_projections` too, matching the MoE path. Strict improvement: the config flag is set only by the env var or by mixed-bit detection. * No `unwrap()`, no `as` casts (use `i32::try_from`); match arms enumerate variants. No file-level allows added. Verification on origin/main (rustc 1.95.0): * `cargo check -p higgs-models` — clean * `cargo clippy --all-targets --all-features -- -D warnings` — clean * `cargo fmt --check` — clean * `cargo test -p higgs-models --lib` — 333/333 pass (3 new) Source: feat/magic-canvas commit `061e500c`. Direct cherry-pick had 5 conflict regions because origin/main has evolved the load functions independently; this is a manual surgical port that preserves origin/main's structure while adding the fallback behaviour. Co-Authored-By: Claude Opus 4.7 (1M context) * fix(qwen3_next): preserve explicit GDN projection config --------- Co-authored-by: Claude Opus 4.7 (1M context) Co-authored-by: Jonathan Reyes --- crates/higgs-models/src/qwen3_next.rs | 306 ++++++++++++++++++++++++-- 1 file changed, 283 insertions(+), 23 deletions(-) diff --git a/crates/higgs-models/src/qwen3_next.rs b/crates/higgs-models/src/qwen3_next.rs index a23d4543..736a7e70 100644 --- a/crates/higgs-models/src/qwen3_next.rs +++ b/crates/higgs-models/src/qwen3_next.rs @@ -3722,14 +3722,29 @@ fn load_qwen3_5_moe_text_config_args>( .or_insert(serde_json::Value::from(0)); } - // When HIGGS_SEPARATE_GDN_PROJ is set, construct the model with separate - // GDN projection fields so the direct weight loader can match them. - // Otherwise, construct with fused fields (weights are rearranged at load time). - let use_separate = std::env::var("HIGGS_SEPARATE_GDN_PROJ").is_ok(); + // When HIGGS_SEPARATE_GDN_PROJ is set, or when per-layer GDN BA quantization + // disagrees on bit-width / group_size between in_proj_a and in_proj_b (common + // in Unsloth dynamic quants), construct the model with separate GDN + // projection fields so the direct weight loader can match them. Otherwise, + // construct with fused fields (weights are rearranged at load time). + let mixed_ba_layers = qwen3_5_mixed_ba_quantization_layers(&config, text_config); + let config_requests_separate = map + .get("use_separate_gdn_projections") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false); + let use_separate = config_requests_separate + || std::env::var("HIGGS_SEPARATE_GDN_PROJ").is_ok() + || !mixed_ba_layers.is_empty(); map.insert( "use_separate_gdn_projections".to_owned(), serde_json::Value::from(use_separate), ); + if !mixed_ba_layers.is_empty() { + tracing::info!( + layers = ?mixed_ba_layers, + "Detected mixed-bit GDN BA projections; using separate GDN projections" + ); + } // Detect per-layer gate quantization override from top-level quantization config if let Some(gate_q) = gate_quantization_override(&config) { @@ -3739,6 +3754,59 @@ fn load_qwen3_5_moe_text_config_args>( Ok(serde_json::from_value(obj)?) } +/// Parse a `{group_size, bits}` quantization spec from a JSON node. +fn qwen3_5_quantization_config(value: &serde_json::Value) -> Option { + Some(QuantizationConfig { + group_size: i32::try_from(value.get("group_size")?.as_i64()?).ok()?, + bits: i32::try_from(value.get("bits")?.as_i64()?).ok()?, + }) +} + +/// Scan the per-layer `quantization` map and return layer indices where the GDN +/// `in_proj_a` and `in_proj_b` projections disagree on bit-width or group size. +/// Such layers cannot be fused into a single `in_proj_ba` matrix without +/// dequantizing, so the loader must fall back to separate GDN projections. +fn qwen3_5_mixed_ba_quantization_layers( + config: &serde_json::Value, + text_config: &serde_json::Value, +) -> Vec { + let Some(quant) = config.get("quantization") else { + return Vec::new(); + }; + let Some(default_quant) = qwen3_5_quantization_config(quant) else { + return Vec::new(); + }; + let Some(num_hidden_layers) = text_config + .get("num_hidden_layers") + .and_then(serde_json::Value::as_i64) + .and_then(|n| i32::try_from(n).ok()) + else { + return Vec::new(); + }; + + (0..num_hidden_layers) + .filter(|layer_idx| { + let prefixes = [ + format!("language_model.model.layers.{layer_idx}.linear_attn"), + format!("model.layers.{layer_idx}.linear_attn"), + ]; + let projection_quantization = |projection: &str| { + prefixes + .iter() + .find_map(|prefix| { + quant + .get(format!("{prefix}.{projection}")) + .and_then(qwen3_5_quantization_config) + }) + .unwrap_or_else(|| default_quant.clone()) + }; + let a_quant = projection_quantization("in_proj_a"); + let b_quant = projection_quantization("in_proj_b"); + a_quant.bits != b_quant.bits || a_quant.group_size != b_quant.group_size + }) + .collect() +} + /// Load a Qwen3.5 dense model (VLM wrapper around `Qwen3Next` architecture). /// /// Reads `text_config` for model args, strips `language_model.` prefix from @@ -3766,15 +3834,7 @@ pub fn load_qwen3_5_model>(model_dir: P) -> Result>( head_v_dim: args.linear_value_head_dim, }; gdn_dims.validate()?; - let mut model = Qwen3NextCausalLM::new(args.clone())?; - // Load weights with GDN projection rearrangement: flat (qkv,z,b,a) // → per-head-grouped (qkvz,ba) for fused 2-dispatch forward path. - // Respect use_separate_gdn_projections config flag or HIGGS_SEPARATE_GDN_PROJ env var. - let use_separate = + // Respects use_separate_gdn_projections (set by HIGGS_SEPARATE_GDN_PROJ env + // 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)?; + + tracing::info!("Qwen3.5-MoE model loaded successfully"); + Ok(model) +} + +/// Build a `Qwen3NextCausalLM` and load weights, choosing fused or separate GDN +/// projections. When the config (or env var) requests separate projections, use +/// the direct loader. Otherwise try the fused loader; if it reports a mixed-bit +/// `in_proj_ba` shape mismatch, rebuild the model with separate projections and +/// retry via the direct loader. +fn load_qwen3_5_model_with_gdn_fallback( + model_path: &Path, + mut args: Qwen3NextModelArgs, + gdn_dims: &GdnDims, +) -> Result { + let force_separate = args.use_separate_gdn_projections || std::env::var("HIGGS_SEPARATE_GDN_PROJ").is_ok(); - if use_separate { + 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)?; tracing::info!("Using SEPARATE GDN projections (4 dispatches per layer)"); - } else { - load_qwen3_5_moe_weights_fused(&mut model, model_path, &gdn_dims)?; - tracing::info!("Using FUSED GDN projections (2 dispatches per layer)"); + return Ok(model); } - tracing::info!("Qwen3.5-MoE model loaded successfully"); - Ok(model) + let mut fused_model = Qwen3NextCausalLM::new(args.clone())?; + match load_qwen3_5_moe_weights_fused(&mut fused_model, model_path, gdn_dims) { + Ok(()) => { + tracing::info!("Using FUSED GDN projections (2 dispatches per layer)"); + Ok(fused_model) + } + Err(err) if is_mixed_bit_gdn_ba_fusion_error(&err) => { + tracing::warn!( + error = %err, + "Detected mixed-bit GDN BA projection shapes; retrying with separate GDN projections" + ); + args.use_separate_gdn_projections = true; + let mut separate_model = Qwen3NextCausalLM::new(args)?; + load_qwen3_5_moe_weights_direct(&mut separate_model, model_path)?; + tracing::info!( + "Using SEPARATE GDN projections (4 dispatches per layer, mixed-bit fallback)" + ); + Ok(separate_model) + } + Err(err) => Err(err), + } +} + +/// Returns true when the supplied error is the mixed-bit BA fusion error raised +/// by [`load_qwen3_5_moe_weights_fused`] when `in_proj_a` and `in_proj_b` have +/// incompatible packed inner shapes. +fn is_mixed_bit_gdn_ba_fusion_error(err: &ModelError) -> bool { + matches!( + err, + ModelError::ShapeMismatch(message) + if message.contains("in_proj_ba") + && message.contains("requires separate GDN projections") + ) } /// GDN dimension info extracted from model args before move. @@ -3918,6 +4026,25 @@ fn concat_and_permute(a: &Array, b: &Array, perm: &[i32]) -> Result bool { + a_shape.len() == b_shape.len() + && a_shape + .iter() + .zip(b_shape.iter()) + .enumerate() + .all(|(axis, (lhs, rhs))| axis == 0 || lhs == rhs) +} + +fn can_concatenate_axis0(a: &Array, b: &Array) -> bool { + let a_shape = a.shape(); + let b_shape = b.shape(); + can_concatenate_axis0_shapes(a_shape, b_shape) +} + /// Load Qwen3.5-MoE weights with GDN projection fusion. /// /// Direct weight loader: strip `language_model.` prefix, no rearrangement. @@ -4052,6 +4179,13 @@ fn load_qwen3_5_moe_weights_fused( format!("Incomplete GDN projection pair for key: {combined_key}"), ))); }; + if combined_key.contains("in_proj_ba") && !can_concatenate_axis0(a, b) { + return Err(crate::error::ModelError::ShapeMismatch(format!( + "Mixed-bit BA fusion requires separate GDN projections for key {combined_key}: {:?} vs {:?}", + a.shape(), + b.shape() + ))); + } let Some(param) = params.get_mut(combined_key.as_str()) else { return Err(crate::error::ModelError::Io(std::io::Error::other( format!("Fused target key not found in model params: {combined_key}"), @@ -12825,6 +12959,132 @@ mod tests { assert_eq!(args.mtp_num_hidden_layers, 1); } + #[test] + fn test_load_qwen35_mixed_ba_quantization_forces_separate_gdn() { + let dir = tempfile::tempdir().unwrap(); + let config = format!( + r#"{{ + "text_config": {}, + "tie_word_embeddings": false, + "quantization": {{ + "group_size": 64, + "bits": 2, + "mode": "affine", + "language_model.model.layers.1.linear_attn.in_proj_a": {{ + "group_size": 64, + "bits": 5, + "mode": "affine" + }} + }} + }}"#, + qwen35_dense_text_config() + ); + std::fs::write(dir.path().join("config.json"), config).unwrap(); + + let args = load_qwen3_5_moe_text_config_args(dir.path()).unwrap(); + + assert!( + args.use_separate_gdn_projections, + "mixed-bit in_proj_a/in_proj_b must force separate GDN projections" + ); + } + + #[test] + fn test_load_qwen35_mixed_ba_quantization_supports_unprefixed_layer_keys() { + let dir = tempfile::tempdir().unwrap(); + let config = format!( + r#"{{ + "text_config": {}, + "tie_word_embeddings": false, + "quantization": {{ + "group_size": 64, + "bits": 2, + "mode": "affine", + "model.layers.1.linear_attn.in_proj_a": {{ + "group_size": 64, + "bits": 5, + "mode": "affine" + }} + }} + }}"#, + qwen35_dense_text_config() + ); + std::fs::write(dir.path().join("config.json"), config).unwrap(); + + let args = load_qwen3_5_moe_text_config_args(dir.path()).unwrap(); + + assert!( + args.use_separate_gdn_projections, + "unprefixed mixed-bit in_proj_a/in_proj_b must force separate GDN projections" + ); + } + + #[test] + fn test_load_qwen35_matching_ba_quantization_keeps_fused_gdn() { + let dir = tempfile::tempdir().unwrap(); + let config = format!( + r#"{{ + "text_config": {}, + "tie_word_embeddings": false, + "quantization": {{ + "group_size": 64, + "bits": 2, + "mode": "affine", + "language_model.model.layers.1.linear_attn.in_proj_a": {{ + "group_size": 64, + "bits": 5, + "mode": "affine" + }}, + "language_model.model.layers.1.linear_attn.in_proj_b": {{ + "group_size": 64, + "bits": 5, + "mode": "affine" + }} + }} + }}"#, + qwen35_dense_text_config() + ); + std::fs::write(dir.path().join("config.json"), config).unwrap(); + + let args = load_qwen3_5_moe_text_config_args(dir.path()).unwrap(); + + assert!( + !args.use_separate_gdn_projections, + "matching BA overrides should keep the fused GDN loader path" + ); + } + + #[test] + fn test_load_qwen35_explicit_separate_gdn_config_is_preserved() { + let dir = tempfile::tempdir().unwrap(); + let mut text_config = qwen35_dense_text_config().trim_end_matches('}').to_owned(); + text_config.push_str( + r#", + "use_separate_gdn_projections": true + }"#, + ); + write_qwen35_config(dir.path(), &text_config); + + let args = load_qwen3_5_moe_text_config_args(dir.path()).unwrap(); + + assert!( + args.use_separate_gdn_projections, + "explicit use_separate_gdn_projections=true must not be overwritten" + ); + } + + #[test] + fn test_can_concatenate_axis0_detects_quantized_inner_shape_mismatch() { + assert!( + !can_concatenate_axis0_shapes(&[48, 320], &[48, 800]), + "different packed inner dims must block BA fusion" + ); + assert!( + can_concatenate_axis0_shapes(&[48, 320], &[96, 320]), + "axis-0 size may differ because fusion concatenates rows" + ); + } + /// GQA ratio: `num_v_heads` must be divisible by `num_k_heads`. /// This validates the assumption used in test/bench GDN recurrence loops. #[test] From 60d7cb48f28ddff5254f8aece898a24a1a475d66 Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Wed, 6 May 2026 15:08:17 +0200 Subject: [PATCH 07/29] =?UTF-8?q?perf(models):=20opt-in=20fused=20MoE=20ga?= =?UTF-8?q?te+up=20=E2=80=94=203=E2=86=922=20expert=20matmuls=20per=20laye?= =?UTF-8?q?r=20(#141)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * perf(models): fused MoE gate+up — 3→2 expert matmuls per layer SwitchMlpWeights::forward_gather_fused() lazy-concatenates gate+up weights on first call, then dispatches a single gather_qmm instead of two separate calls. FfnBlock::forward() now routes through the fused path instead of forward_gather_global_sort(). Measured on 35B-A3B-3bit M4 base: - S=1 decode: 27ms → 17ms (−37%) - S=16 verify: 253ms → 112ms (−56%) - MoE/layer at K=1: 0.47ms (down from ~0.68ms) Co-Authored-By: Claude Sonnet 4.6 (1M context) * style: cargo fmt qwen3_next.rs Reflow `let fw/fs/fb = ops::concatenate_axis(..)` from broken-line indentation back onto single lines so `cargo fmt --all -- --check` passes in CI. * fix(clippy): backtick MoE/gather_qmm doc + safe top_k u32 cast Two errors flagged by `-D clippy::doc-markdown` and `-D clippy::cast-sign-loss`/`-D clippy::as-conversions`: - Backtick `MoE` and `gather_qmm` in the `fused_gate_up` doc comment. - Replace `top_k as u32` with the same `u32::try_from(top_k).map_err(...)` pattern already used by `forward_gather_global_sort`. * fix(qwen3_next): gate MoE gate-up fusion behind opt-in --------- Co-authored-by: Claude Sonnet 4.6 (1M context) Co-authored-by: Jonathan Reyes --- crates/higgs-models/src/qwen3_next.rs | 207 ++++++++++++++++++++++++-- 1 file changed, 198 insertions(+), 9 deletions(-) diff --git a/crates/higgs-models/src/qwen3_next.rs b/crates/higgs-models/src/qwen3_next.rs index 736a7e70..3b0b1723 100644 --- a/crates/higgs-models/src/qwen3_next.rs +++ b/crates/higgs-models/src/qwen3_next.rs @@ -1015,6 +1015,7 @@ static DECODE_GEMV_ENABLED: OnceLock = OnceLock::new(); static QGEMV_NSG_OVERRIDE: OnceLock> = OnceLock::new(); static DENSE_FFN_GEMV_MODE: OnceLock = OnceLock::new(); static DENSE_FFN_FUSE_GATE_UP: OnceLock = OnceLock::new(); +static MOE_FFN_FUSE_GATE_UP: OnceLock = OnceLock::new(); #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum DenseFfnGemvMode { @@ -1087,6 +1088,10 @@ fn dense_ffn_fuse_gate_up() -> bool { }) } +fn moe_ffn_fuse_gate_up() -> bool { + *MOE_FFN_FUSE_GATE_UP.get_or_init(|| truthy_env_var("HIGGS_MOE_FFN_GATE_UP")) +} + fn qgemv_config_cache_enabled() -> bool { *QGEMV_CONFIG_CACHE_ENABLED.get_or_init(|| truthy_env_var("HIGGS_CACHE_QGEMV_CONFIGS")) } @@ -1696,6 +1701,8 @@ pub(crate) struct SwitchMlpWeights { up_proj: QLinear, #[param] down_proj: QLinear, + /// Lazily fused gate+up weights for `MoE` `gather_qmm` (3→2 calls per layer). + fused_gate_up: Option<(Array, Array, Array, i32)>, } impl SwitchMlpWeights { @@ -1705,6 +1712,7 @@ impl SwitchMlpWeights { gate_proj, up_proj, down_proj, + fused_gate_up: None, }) } @@ -1867,6 +1875,106 @@ impl SwitchMlpWeights { // Reshape back to [B, L, top_k, D] out_unsorted.reshape(&[b, l, top_k, d]) } + + /// Like `forward_gather_global_sort` but fuses gate+up into a single + /// `gather_qmm` call (3→2 per layer). Lazy-inits fused weights on first call. + /// Production routing gates this behind `HIGGS_MOE_FFN_GATE_UP` because the + /// fused cache duplicates the resident gate/up tensors. + pub(crate) fn forward_gather_fused( + &mut self, + x: &Array, + indices: &Array, + ) -> Result { + // Lazy-init: concatenate gate+up weights along axis 1 (intermediate dim). + // MoE weights are [num_experts, intermediate_packed, hidden]. + if self.fused_gate_up.is_none() { + let intermediate = *self + .gate_proj + .weight + .shape() + .get(1) + .ok_or_else(|| Exception::custom("gate_proj weight missing dim 1"))?; + let fw = ops::concatenate_axis(&[&*self.gate_proj.weight, &*self.up_proj.weight], 1)?; + let fs = ops::concatenate_axis(&[&*self.gate_proj.scales, &*self.up_proj.scales], 1)?; + let fb = ops::concatenate_axis(&[&*self.gate_proj.biases, &*self.up_proj.biases], 1)?; + fw.eval()?; + fs.eval()?; + fb.eval()?; + self.fused_gate_up = Some((fw, fs, fb, intermediate)); + } + let (fw, fs, fb, intermediate) = self + .fused_gate_up + .as_ref() + .ok_or_else(|| Exception::custom("fused_gate_up missing after init"))?; + + // --- Global sort (same as forward_gather_global_sort) --- + let x_shape = x.shape(); + let err = || Exception::custom("forward_gather_fused input must be [B, L, D]"); + let b = *x_shape.first().ok_or_else(err)?; + let l = *x_shape.get(1).ok_or_else(err)?; + let d = *x_shape.get(2).ok_or_else(err)?; + let top_k = *indices + .shape() + .last() + .ok_or_else(|| Exception::custom("indices must have last dim"))?; + + let idx_flat = indices.flatten(None, None)?; + let order = ops::argsort_axis(&idx_flat, 0)?; + let inv_order = ops::argsort_axis(&order, 0)?; + + let top_k_u32 = + u32::try_from(top_k).map_err(|_| Exception::custom("top_k must fit in u32"))?; + let top_k_arr = Array::from_slice(&[top_k_u32], &[1]); + let token_idx = order.floor_divide(&top_k_arr)?; + + let x_flat = x.reshape(&[b * l, 1, d])?; + let x_sorted = x_flat.take_axis(&token_idx, 0)?; + let idx_sorted = idx_flat.take_axis(&order, 0)?; + + // --- Fused gate+up: ONE gather_qmm instead of TWO --- + let fused_out = gather_qmm( + &x_sorted, + fw, + fs, + fb, + &idx_sorted, + true, + self.gate_proj.group_size, + self.gate_proj.bits, + true, + )?; + // Split at intermediate boundary → gate_out, up_out + let parts = fused_out.split_axis(&[*intermediate], Some(-1))?; + let gate_out = parts + .first() + .ok_or_else(|| Exception::custom("fused split failed"))?; + let up_out = parts + .get(1) + .ok_or_else(|| Exception::custom("fused split failed"))?; + let activated = swiglu(gate_out, up_out)?; + + // --- down_proj: unchanged --- + let down_out = gather_qmm( + &activated, + &self.down_proj.weight, + &self.down_proj.scales, + &self.down_proj.biases, + &idx_sorted, + true, + self.down_proj.group_size, + self.down_proj.bits, + true, + )?; + + // down_out: [N, 1, D] -> squeeze M -> [N, D] + let out_flat = down_out.squeeze_axes(&[-2])?; + + // --- Unsort: restore original token order --- + let out_unsorted = out_flat.take_axis(&inv_order, 0)?; + + // Reshape back to [B, L, top_k, D] + out_unsorted.reshape(&[b, l, top_k, d]) + } } // --------------------------------------------------------------------------- @@ -2667,14 +2775,6 @@ impl FfnBlock { .gate .as_ref() .ok_or_else(|| Exception::custom("MoE gate missing"))?; - let switch_ref = self - .switch_mlp - .as_ref() - .ok_or_else(|| Exception::custom("MoE switch_mlp missing"))?; - let se_ref = self - .shared_expert - .as_ref() - .ok_or_else(|| Exception::custom("MoE shared_expert missing"))?; let seg_ref = self .shared_expert_gate .as_ref() @@ -2698,13 +2798,26 @@ impl FfnBlock { raw_scores }; - let y = switch_ref.forward_gather_global_sort(x, &inds)?; + let switch_ref = self + .switch_mlp + .as_mut() + .ok_or_else(|| Exception::custom("MoE switch_mlp missing"))?; + let y = if moe_ffn_fuse_gate_up() { + switch_ref.forward_gather_fused(x, &inds)? + } else { + switch_ref.forward_gather_global_sort(x, &inds)? + }; let expert_sum = y .multiply(&scores.expand_dims(-1)?)? .sum_axes(&[-2], false)?; + let se_ref = self + .shared_expert + .as_ref() + .ok_or_else(|| Exception::custom("MoE shared_expert missing"))?; let shared_y = se_ref.forward(x)?; + let shared_gate_val = nn::sigmoid(&seg_ref.forward(x)?)?; let shared_out = shared_y.multiply(&shared_gate_val)?; @@ -5081,6 +5194,81 @@ mod tests { ); } + #[test] + fn test_moe_gate_up_fusion_parity() { + // Fused gate+up (2 gather_qmm) must match unfused (3 gather_qmm). + // Uses random weights + distinct per-token inputs to stress sort/unsort. + let num_experts = 8; + let hidden = 128; + let intermediate = 64; + let top_k = 3; + let b = 1; + let l = 16; + + let mut block = SwitchMlpWeights::new(64, 4).unwrap(); + + let gate_w = mlx_rs::random::uniform::( + -1.0, + 1.0, + &[num_experts, intermediate, hidden], + None, + ) + .unwrap(); + let (gw, gs, gb) = quantize_weights(&gate_w, 64, 4); + *block.gate_proj.weight = gw; + *block.gate_proj.scales = gs; + *block.gate_proj.biases = gb; + + let up_w = mlx_rs::random::uniform::( + -1.0, + 1.0, + &[num_experts, intermediate, hidden], + None, + ) + .unwrap(); + let (uw, us, ub) = quantize_weights(&up_w, 64, 4); + *block.up_proj.weight = uw; + *block.up_proj.scales = us; + *block.up_proj.biases = ub; + + let down_w = mlx_rs::random::uniform::( + -1.0, + 1.0, + &[num_experts, hidden, intermediate], + None, + ) + .unwrap(); + let (dw, ds, db) = quantize_weights(&down_w, 64, 4); + *block.down_proj.weight = dw; + *block.down_proj.scales = ds; + *block.down_proj.biases = db; + + let x = mlx_rs::random::uniform::(-1.0, 1.0, &[b, l, hidden], None).unwrap(); + let idx_data: Vec = (0..(b * l * top_k) as u32) + .map(|i| i % num_experts as u32) + .collect(); + let indices = Array::from_slice(&idx_data, &[b, l, top_k]); + x.eval().unwrap(); + indices.eval().unwrap(); + + // Reference: unfused 3-call path + let reference = block.forward_gather_global_sort(&x, &indices).unwrap(); + // Fused: 2-call path + let fused = block.forward_gather_fused(&x, &indices).unwrap(); + reference.eval().unwrap(); + fused.eval().unwrap(); + + assert_eq!(reference.shape(), fused.shape()); + assert_eq!(fused.shape(), &[b, l, top_k, hidden]); + + let diff = reference.subtract(&fused).unwrap().abs().unwrap(); + let max_diff: f32 = diff.max(None).unwrap().item(); + assert!( + max_diff < 1e-5, + "fused gate+up differs from unfused by {max_diff}" + ); + } + #[test] fn test_switch_mlp_forward_gather_shapes() { // Verify forward_gather produces the correct output shape with the @@ -6479,6 +6667,7 @@ mod tests { gate_proj: make_switch_ql(d, d_inter), up_proj: make_switch_ql(d, d_inter), down_proj: make_switch_ql(d_inter, d), + fused_gate_up: None, }, shared_expert: Qwen3NextMLP { gate_proj: make_ql(d, shared_inter * 2, gs, bits), From fe43aabe44104ab285e19e8ff73c724c0875cbe0 Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Wed, 6 May 2026 16:47:21 +0200 Subject: [PATCH 08/29] feat(bonsai-q1): packed engine scaffold with upstream MLX guard (#142) * feat(bonsai_q1): add upstream-guarded packed engine * fix(bonsai-q1): address review feedback --------- Co-authored-by: Jonathan Reyes --- Cargo.lock | 12 + crates/higgs-engine/src/model_loader.rs | 100 +- crates/higgs-models/Cargo.toml | 2 + crates/higgs-models/src/bonsai_q1.rs | 1204 +++++++++++++++++++++++ crates/higgs-models/src/cache.rs | 8 + crates/higgs-models/src/deepseek_v2.rs | 21 +- crates/higgs-models/src/lib.rs | 51 +- crates/higgs-models/src/yarn.rs | 223 +++++ docs/BONSAI_Q1.md | 18 + 9 files changed, 1625 insertions(+), 14 deletions(-) create mode 100644 crates/higgs-models/src/bonsai_q1.rs create mode 100644 crates/higgs-models/src/yarn.rs create mode 100644 docs/BONSAI_Q1.md diff --git a/Cargo.lock b/Cargo.lock index 7299bd78..942dd4b1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1645,10 +1645,12 @@ name = "higgs-models" version = "1.1.1" dependencies = [ "criterion", + "half", "image", "mlx-rs", "mlx-sys", "rand 0.10.1", + "safetensors", "serde", "serde_json", "tempfile", @@ -3460,6 +3462,16 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "safetensors" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44560c11236a6130a46ce36c836a62936dc81ebf8c36a37947423571be0e55b6" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "same-file" version = "1.0.6" diff --git a/crates/higgs-engine/src/model_loader.rs b/crates/higgs-engine/src/model_loader.rs index aa8119fa..bd2f91d1 100644 --- a/crates/higgs-engine/src/model_loader.rs +++ b/crates/higgs-engine/src/model_loader.rs @@ -1,6 +1,8 @@ use std::path::{Path, PathBuf}; -use higgs_models::{AnyModel, load_tokenizer as shared_load_tokenizer, registry, transformer}; +use higgs_models::{ + AnyModel, error::ModelError, load_tokenizer as shared_load_tokenizer, registry, transformer, +}; use crate::error::EngineError; @@ -36,6 +38,17 @@ pub fn load_model>(model_dir: P) -> Result match config.model_type.as_str() { "qwen2" | "qwen3" | "llama" | "mistral" => { + // Packed 1.25-bpw Bonsai-Q1 checkpoints declare model_type="qwen3" + // but the weights are quantized to bits=1. Keep detection ahead of + // the fp16/Q4 transformer loader so users get an explicit error + // while the workspace remains on upstream oxideai/mlx-rs. + if is_bonsai_q1(&config.model_dir)? { + return Err(EngineError::Model(ModelError::UnsupportedModel( + "Bonsai-Q1 requires MLX bits=1 affine quantization support; \ + the workspace stays on upstream oxideai/mlx-rs until that support lands" + .to_owned(), + ))); + } let model = transformer::load_model(&config.model_dir).map_err(EngineError::Model)?; Ok(AnyModel::Transformer(model)) } @@ -90,6 +103,39 @@ pub fn load_model>(model_dir: P) -> Result } } +/// Peek into `config.json` to detect packed 1-bit Bonsai-Q1 checkpoints. +/// +/// Returns `true` for Qwen3-shaped `quantization.bits == 1` checkpoints using +/// the expected group size. Returns `false` for any other model type or +/// quantization config. A missing / malformed `config.json` propagates as an +/// IO / JSON error — we never mask it. +fn is_bonsai_q1(dir: &Path) -> Result { + let cfg_path = dir.join("config.json"); + let txt = std::fs::read_to_string(&cfg_path).map_err(|e| { + EngineError::Model(higgs_models::error::ModelError::Io(std::io::Error::new( + e.kind(), + format!("{}: {e}", cfg_path.display()), + ))) + })?; + let cfg: serde_json::Value = serde_json::from_str(&txt) + .map_err(|e| EngineError::Model(higgs_models::error::ModelError::Json(e)))?; + let bonsai_group_size = u64::try_from(higgs_models::bonsai_q1::GROUP_SIZE) + .map_err(|e| EngineError::Model(ModelError::ShapeMismatch(e.to_string())))?; + Ok( + cfg.get("model_type").and_then(serde_json::Value::as_str) == Some("qwen3") + && cfg + .get("quantization") + .and_then(|q| q.get("bits")) + .and_then(serde_json::Value::as_u64) + == Some(1) + && cfg + .get("quantization") + .and_then(|q| q.get("group_size")) + .and_then(serde_json::Value::as_u64) + == Some(bonsai_group_size), + ) +} + /// Load a tokenizer from a model directory. pub fn load_tokenizer>(model_dir: P) -> Result { shared_load_tokenizer(model_dir).map_err(|e| EngineError::Tokenization(e.to_string())) @@ -230,6 +276,58 @@ mod tests { )); } + #[test] + fn is_bonsai_q1_requires_qwen3_model_type_and_group_size() { + let (qwen3_dir, _qwen3_result) = config_from_raw( + r#"{ + "model_type": "qwen3", + "quantization": {"bits": 1, "group_size": 128} + }"#, + ); + assert!(is_bonsai_q1(qwen3_dir.path()).unwrap()); + + let (llama_dir, _llama_result) = config_from_raw( + r#"{ + "model_type": "llama", + "quantization": {"bits": 1, "group_size": 128} + }"#, + ); + assert!(!is_bonsai_q1(llama_dir.path()).unwrap()); + + let (wrong_group_dir, _wrong_group_result) = config_from_raw( + r#"{ + "model_type": "qwen3", + "quantization": {"bits": 1, "group_size": 64} + }"#, + ); + assert!(!is_bonsai_q1(wrong_group_dir.path()).unwrap()); + + let (q4_dir, _q4_result) = config_from_raw( + r#"{ + "model_type": "qwen3", + "quantization": {"bits": 4, "group_size": 128} + }"#, + ); + assert!( + !is_bonsai_q1(q4_dir.path()).unwrap(), + "regular Q4 Qwen3 must not be misclassified as Bonsai-Q1" + ); + } + + #[test] + fn load_model_rejects_bonsai_q1_without_runtime_support() { + let (dir, _result) = config_from_raw( + r#"{ + "model_type": "qwen3", + "quantization": {"bits": 1, "group_size": 128} + }"#, + ); + match load_model(dir.path()) { + Err(err) => assert!(err.to_string().contains("Bonsai-Q1 requires MLX bits=1")), + Ok(_) => panic!("Expected unsupported Bonsai-Q1 runtime error"), + } + } + #[test] fn load_tokenizer_missing_tokenizer_json() { let dir = tempfile::tempdir().unwrap(); diff --git a/crates/higgs-models/Cargo.toml b/crates/higgs-models/Cargo.toml index b2007b0a..47a370aa 100644 --- a/crates/higgs-models/Cargo.toml +++ b/crates/higgs-models/Cargo.toml @@ -11,8 +11,10 @@ homepage.workspace = true workspace = true [dependencies] +half = "2.4" mlx-rs.workspace = true mlx-sys.workspace = true +safetensors = "0.4" serde.workspace = true serde_json.workspace = true tokenizers.workspace = true diff --git a/crates/higgs-models/src/bonsai_q1.rs b/crates/higgs-models/src/bonsai_q1.rs new file mode 100644 index 00000000..93fe10e2 --- /dev/null +++ b/crates/higgs-models/src/bonsai_q1.rs @@ -0,0 +1,1204 @@ +//! Bonsai-Q1 target-capable engine: packed 1.25-bpw weight storage. +//! +//! Unlike `DiffusionEngine::load_q1` which dequantizes to fp32 at load (32 GB +//! residency on 8B), this engine holds MLX's `Q1_0_g128` affine encoding +//! verbatim: `w[row, col] = scales[row, col/128] * bit(col) + biases[row, +//! col/128]`. Dequant happens inline inside the MLX quantized matmul kernel +//! once upstream MLX provides bits=1 affine support. +//! +//! Residency: ~1.25 GB for Bonsai-8B-mlx-1bit, ~260 MB for Bonsai-1.7B-mlx-1bit. +//! +//! Scope: Rust-side loader and engine implementation. Runtime routing is held +//! back in `higgs-engine` until the upstream MLX dependency supports bits=1 +//! affine quantization. + +#![allow( + clippy::too_many_arguments, + clippy::too_many_lines, + // Quantization math uses small bounded dims (head_dim, GROUP_SIZE=128, vocab) and + // bit-packed u32→f32 conversions where precision/sign loss is intentional. + clippy::cast_possible_truncation, + clippy::cast_possible_wrap, + clippy::cast_precision_loss, + clippy::cast_sign_loss, + clippy::as_conversions, + // Dequant kernel + safetensors loader index into manually-bounds-checked slices. + clippy::indexing_slicing, + // Decode loop reuses names (q, k, v, t0) across rope/sdpa/o_proj stages by design. + clippy::shadow_unrelated, + clippy::shadow_reuse, + clippy::shadow_same, + // Loader unwraps on safetensors slices after explicit shape validation; load failure + // paths return ShapeMismatch via map_err elsewhere. + clippy::unwrap_used, + clippy::map_unwrap_or, + // YarnRoPE / Q1 / KV abbreviations are domain terms, not items to backtick. + clippy::doc_markdown, + clippy::doc_lazy_continuation, + clippy::missing_const_for_fn, + clippy::manual_flatten, + clippy::if_then_some_else_none, + clippy::suboptimal_flops, +)] + +use half::f16; +use std::path::Path; + +use mlx_rs::{Array, Dtype, error::Exception, fast, ops, ops::indexing::IndexOp}; +use safetensors::SafeTensors; + +use crate::{ + cache::{KeyValueCache, SteppingKeyValueCache}, + error::ModelError, + utils::{cached_scaled_dot_product_attention, create_attention_mask}, + yarn::{apply_yarn_rope, compute_yarn_freqs, yarn_get_mscale}, +}; + +/// Load and materialize a Bonsai-Q1 model from `model_dir` onto the GPU. +/// +/// Adapts [`BonsaiQ1Engine::load`]'s `Result<_, String>` into [`ModelError`] so +/// the engine surface in `higgs-engine::model_loader` can route it through the +/// same `EngineError::Model` path used by all other architectures. +pub fn load_bonsai_q1>(model_dir: P) -> Result { + let engine = BonsaiQ1Engine::load(model_dir).map_err(ModelError::ShapeMismatch)?; + engine.to_gpu().map_err(ModelError::Mlx) +} + +pub const GROUP_SIZE: usize = 128; +const BITS: i32 = 1; +const GROUP_SIZE_I32: i32 = GROUP_SIZE as i32; + +/// Packed 1-bit linear layer with affine per-group dequant. +/// +/// Layout (matches MLX 1-bit `QuantizedLinear`, `PrismML` fork): +/// - `w_packed`: `[out_features, in_features/32]` u32, bit `col%32` of word +/// `col/32` is the raw 1-bit weight for column `col`. +/// - `scales`, `biases`: `[out_features, in_features/128]` f16, one per group +/// of 128 input columns. +/// +/// Effective: 1 bit/weight + 32 bits/group / 128 weights = **1.25 bpw**. +pub struct PackedQ1Linear { + pub w_packed: Vec, + pub scales: Vec, + pub biases: Vec, + pub out_features: usize, + pub in_features: usize, +} + +impl PackedQ1Linear { + pub const fn resident_bytes(&self) -> usize { + self.w_packed.len() * 4 + self.scales.len() * 2 + self.biases.len() * 2 + } + + /// Dequantize a single row to fp32 (reference path for correctness tests). + /// + /// Not used on the hot path — P2 replaces this with a Metal kernel that + /// fuses dequant into the matmul. + pub fn dequant_row_to_fp32(&self, row: usize, out: &mut [f32]) { + debug_assert_eq!(out.len(), self.in_features); + let n_groups = self.in_features / GROUP_SIZE; + let packed_cols = self.in_features / 32; + let w_row = &self.w_packed[row * packed_cols..(row + 1) * packed_cols]; + let s_row = &self.scales[row * n_groups..(row + 1) * n_groups]; + let b_row = &self.biases[row * n_groups..(row + 1) * n_groups]; + for col in 0..self.in_features { + let word = w_row[col / 32]; + let bit = ((word >> (col % 32)) & 1) as f32; + let group = col / GROUP_SIZE; + out[col] = s_row[group].to_f32().mul_add(bit, b_row[group].to_f32()); + } + } +} + +pub struct BonsaiQ1LayerWeights { + pub q_proj: PackedQ1Linear, + pub k_proj: PackedQ1Linear, + pub v_proj: PackedQ1Linear, + pub o_proj: PackedQ1Linear, + pub gate_proj: PackedQ1Linear, + pub up_proj: PackedQ1Linear, + pub down_proj: PackedQ1Linear, + pub q_norm: Vec, + pub k_norm: Vec, + pub input_norm: Vec, + pub post_attn_norm: Vec, +} + +impl BonsaiQ1LayerWeights { + pub fn resident_bytes(&self) -> usize { + self.q_proj.resident_bytes() + + self.k_proj.resident_bytes() + + self.v_proj.resident_bytes() + + self.o_proj.resident_bytes() + + self.gate_proj.resident_bytes() + + self.up_proj.resident_bytes() + + self.down_proj.resident_bytes() + + (self.q_norm.len() + + self.k_norm.len() + + self.input_norm.len() + + self.post_attn_norm.len()) + * 2 + } +} + +#[derive(Debug, Clone)] +pub struct BonsaiQ1Config { + pub hidden: usize, + pub layers: usize, + pub heads: usize, + pub kv_heads: usize, + pub head_dim: usize, + pub inter: usize, + pub vocab: usize, + pub rms_norm_eps: f32, + pub rope_theta: f64, + /// YARN scaling factor if present (Bonsai-8B uses `factor=4.0, original=16384`). + pub rope_yarn_factor: Option, + pub rope_original_max_seq: Option, + pub tie_word_embeddings: bool, +} + +pub struct BonsaiQ1Engine { + pub config: BonsaiQ1Config, + pub layers: Vec, + /// Token embedding stored packed (dequants inline at embed lookup time). + pub embed: PackedQ1Linear, + /// Untied LM head for 8B (`tie_word_embeddings: false`). None for 1.7B. + pub lm_head: Option, + pub final_norm: Vec, +} + +impl BonsaiQ1Engine { + pub const fn num_layers(&self) -> usize { + self.layers.len() + } + + pub fn resident_bytes(&self) -> usize { + let layer_bytes: usize = self + .layers + .iter() + .map(BonsaiQ1LayerWeights::resident_bytes) + .sum(); + let lm_head_bytes = self + .lm_head + .as_ref() + .map_or(0, PackedQ1Linear::resident_bytes); + layer_bytes + self.embed.resident_bytes() + lm_head_bytes + self.final_norm.len() * 2 + } + + /// Load from a `HuggingFace` directory containing `config.json` + + /// `model.safetensors` in MLX 1-bit affine-quant format. + #[allow(clippy::too_many_lines)] + pub fn load>(model_dir: P) -> Result { + let dir = model_dir.as_ref(); + + let cfg_txt = std::fs::read_to_string(dir.join("config.json")) + .map_err(|e| format!("config.json: {e}"))?; + let cfg: serde_json::Value = + serde_json::from_str(&cfg_txt).map_err(|e| format!("config.json parse: {e}"))?; + + let u64_of = |k: &str| -> Result { + cfg[k] + .as_u64() + .ok_or_else(|| format!("config.json missing u64 '{k}'")) + }; + let hidden = u64_of("hidden_size")? as usize; + let heads = u64_of("num_attention_heads")? as usize; + let kv_heads = u64_of("num_key_value_heads")? as usize; + let head_dim = cfg["head_dim"].as_u64().map_or(128, |v| v as usize); + let inter = u64_of("intermediate_size")? as usize; + let layers_n = u64_of("num_hidden_layers")? as usize; + let vocab = u64_of("vocab_size")? as usize; + + let rms_norm_eps = cfg["rms_norm_eps"].as_f64().unwrap_or(1e-6) as f32; + let rope_theta = cfg["rope_theta"].as_f64().unwrap_or(1_000_000.0); + let tie_word_embeddings = cfg["tie_word_embeddings"].as_bool().unwrap_or(false); + + let (rope_yarn_factor, rope_original_max_seq) = cfg + .get("rope_scaling") + .and_then(|rs| { + (rs.get("rope_type").and_then(|v| v.as_str()) == Some("yarn")).then(|| { + let f = rs.get("factor").and_then(serde_json::Value::as_f64); + let o = rs + .get("original_max_position_embeddings") + .and_then(serde_json::Value::as_u64) + .map(|v| v as usize); + (f, o) + }) + }) + .unwrap_or((None, None)); + + let quant = cfg + .get("quantization") + .ok_or("missing quantization block")?; + let q_bits = quant.get("bits").and_then(serde_json::Value::as_u64); + let q_group = quant.get("group_size").and_then(serde_json::Value::as_u64); + if q_bits != Some(1) || q_group != Some(GROUP_SIZE as u64) { + return Err(format!( + "expected quantization {{bits:1, group_size:{GROUP_SIZE}}}, got bits={q_bits:?} \ + group_size={q_group:?}" + )); + } + + let st_path = dir.join("model.safetensors"); + let st_data = std::fs::read(&st_path).map_err(|e| format!("read safetensors: {e}"))?; + let tensors = SafeTensors::deserialize(&st_data) + .map_err(|e| format!("deserialize safetensors: {e}"))?; + + let config = BonsaiQ1Config { + hidden, + layers: layers_n, + heads, + kv_heads, + head_dim, + inter, + vocab, + rms_norm_eps, + rope_theta, + rope_yarn_factor, + rope_original_max_seq, + tie_word_embeddings, + }; + + let q_dim = heads * head_dim; + let kv_dim = kv_heads * head_dim; + + let embed = load_packed( + &tensors, + "model.embed_tokens", + vocab, + hidden, + "embed_tokens", + )?; + let lm_head = if tie_word_embeddings { + None + } else { + Some(load_packed(&tensors, "lm_head", vocab, hidden, "lm_head")?) + }; + let final_norm = load_f16(&tensors, "model.norm.weight")?; + if final_norm.len() != hidden { + return Err(format!( + "final_norm len {} != hidden {hidden}", + final_norm.len() + )); + } + + let mut layers = Vec::with_capacity(layers_n); + for i in 0..layers_n { + let p = format!("model.layers.{i}"); + let attn = format!("{p}.self_attn"); + let mlp = format!("{p}.mlp"); + + let layer = BonsaiQ1LayerWeights { + q_proj: load_packed(&tensors, &format!("{attn}.q_proj"), q_dim, hidden, "q_proj")?, + k_proj: load_packed( + &tensors, + &format!("{attn}.k_proj"), + kv_dim, + hidden, + "k_proj", + )?, + v_proj: load_packed( + &tensors, + &format!("{attn}.v_proj"), + kv_dim, + hidden, + "v_proj", + )?, + o_proj: load_packed(&tensors, &format!("{attn}.o_proj"), hidden, q_dim, "o_proj")?, + gate_proj: load_packed( + &tensors, + &format!("{mlp}.gate_proj"), + inter, + hidden, + "gate_proj", + )?, + up_proj: load_packed( + &tensors, + &format!("{mlp}.up_proj"), + inter, + hidden, + "up_proj", + )?, + down_proj: load_packed( + &tensors, + &format!("{mlp}.down_proj"), + hidden, + inter, + "down_proj", + )?, + q_norm: load_f16(&tensors, &format!("{attn}.q_norm.weight"))?, + k_norm: load_f16(&tensors, &format!("{attn}.k_norm.weight"))?, + input_norm: load_f16(&tensors, &format!("{p}.input_layernorm.weight"))?, + post_attn_norm: load_f16( + &tensors, + &format!("{p}.post_attention_layernorm.weight"), + )?, + }; + layers.push(layer); + } + + let engine = Self { + config, + layers, + embed, + lm_head, + final_norm, + }; + let resident_mb = engine.resident_bytes() as f64 / (1024.0 * 1024.0); + tracing::info!( + layers = engine.config.layers, + hidden = engine.config.hidden, + heads = engine.config.heads, + kv_heads = engine.config.kv_heads, + head_dim = engine.config.head_dim, + inter = engine.config.inter, + vocab = engine.config.vocab, + tied_embed = engine.config.tie_word_embeddings, + packed_resident_mb = format!("{resident_mb:.1}"), + "BonsaiQ1Engine::load", + ); + Ok(engine) + } +} + +// --------------------------------------------------------------------------- +// GPU-ready mirror — built once from the packed engine. +// --------------------------------------------------------------------------- + +/// MLX-resident 1-bit linear: weight as uint32 packed, scales/biases as f16, +/// same shape as `PackedQ1Linear` but ready for `ops::quantized_matmul`. +pub struct BonsaiQ1GpuLinear { + pub w: Array, + pub scales: Array, + pub biases: Array, + pub out_features: i32, + pub in_features: i32, +} + +impl BonsaiQ1GpuLinear { + fn from_packed(p: &PackedQ1Linear) -> Result { + let out = i32::try_from(p.out_features) + .map_err(|_| Exception::custom("out_features overflows i32"))?; + let inf = i32::try_from(p.in_features) + .map_err(|_| Exception::custom("in_features overflows i32"))?; + let packed_cols = inf / 32; + let n_groups = inf / GROUP_SIZE_I32; + + let w = Array::from_slice(&p.w_packed, &[out, packed_cols]); + let scales_f32: Vec = p.scales.iter().map(|h| h.to_f32()).collect(); + let biases_f32: Vec = p.biases.iter().map(|h| h.to_f32()).collect(); + let scales = Array::from_slice(&scales_f32, &[out, n_groups]).as_dtype(Dtype::Float16)?; + let biases = Array::from_slice(&biases_f32, &[out, n_groups]).as_dtype(Dtype::Float16)?; + + Ok(Self { + w, + scales, + biases, + out_features: out, + in_features: inf, + }) + } + + /// `y = x @ dequant(w, scales, biases).T` via fused bits=1 qmm. + pub fn forward(&self, x: &Array) -> Result { + ops::quantized_matmul( + x, + &self.w, + &self.scales, + &self.biases, + true, + GROUP_SIZE_I32, + BITS, + ) + } +} + +pub struct BonsaiQ1GpuLayer { + pub q_proj: BonsaiQ1GpuLinear, + pub k_proj: BonsaiQ1GpuLinear, + pub v_proj: BonsaiQ1GpuLinear, + pub o_proj: BonsaiQ1GpuLinear, + pub gate_proj: BonsaiQ1GpuLinear, + pub up_proj: BonsaiQ1GpuLinear, + pub down_proj: BonsaiQ1GpuLinear, + pub q_norm: Array, + pub k_norm: Array, + pub input_norm: Array, + pub post_attn_norm: Array, +} + +pub struct BonsaiQ1Gpu { + pub config: BonsaiQ1Config, + pub layers: Vec, + pub embed: BonsaiQ1GpuLinear, + pub lm_head: Option, + pub final_norm: Array, + /// YARN-scaled `RoPE` frequencies (per `head_dim/2`). None if no YARN. + pub yarn_freqs: Option, + pub yarn_mscale: f32, + pub attention_scale: f32, +} + +fn f16_vec_to_array(weights: &[f16]) -> Result { + let f32s: Vec = weights.iter().map(|h| h.to_f32()).collect(); + let len = + i32::try_from(weights.len()).map_err(|_| Exception::custom("norm len overflows i32"))?; + Array::from_slice(&f32s, &[len]).as_dtype(Dtype::Float16) +} + +impl BonsaiQ1Engine { + /// Consume the packed engine and materialize MLX arrays. + /// + /// Frees the `Vec` / `Vec` residency once copied to MLX. + pub fn to_gpu(self) -> Result { + let mut gpu_layers = Vec::with_capacity(self.layers.len()); + for layer in &self.layers { + gpu_layers.push(BonsaiQ1GpuLayer { + q_proj: BonsaiQ1GpuLinear::from_packed(&layer.q_proj)?, + k_proj: BonsaiQ1GpuLinear::from_packed(&layer.k_proj)?, + v_proj: BonsaiQ1GpuLinear::from_packed(&layer.v_proj)?, + o_proj: BonsaiQ1GpuLinear::from_packed(&layer.o_proj)?, + gate_proj: BonsaiQ1GpuLinear::from_packed(&layer.gate_proj)?, + up_proj: BonsaiQ1GpuLinear::from_packed(&layer.up_proj)?, + down_proj: BonsaiQ1GpuLinear::from_packed(&layer.down_proj)?, + q_norm: f16_vec_to_array(&layer.q_norm)?, + k_norm: f16_vec_to_array(&layer.k_norm)?, + input_norm: f16_vec_to_array(&layer.input_norm)?, + post_attn_norm: f16_vec_to_array(&layer.post_attn_norm)?, + }); + } + + let embed = BonsaiQ1GpuLinear::from_packed(&self.embed)?; + let lm_head = self + .lm_head + .as_ref() + .map(BonsaiQ1GpuLinear::from_packed) + .transpose()?; + let final_norm = f16_vec_to_array(&self.final_norm)?; + + // YARN precompute. + let head_dim_i = i32::try_from(self.config.head_dim) + .map_err(|_| Exception::custom("head_dim overflows i32"))?; + let base = self.config.rope_theta as f32; + let (yarn_freqs, yarn_mscale) = match self.config.rope_yarn_factor { + Some(factor) if factor > 1.0 => { + let orig_seq = self.config.rope_original_max_seq.ok_or_else(|| { + Exception::custom( + "rope_yarn_factor > 1.0 requires \ + rope_scaling.original_max_position_embeddings", + ) + })?; + let orig = i32::try_from(orig_seq) + .map_err(|_| Exception::custom("orig_max_seq overflows i32"))?; + let factor_f = factor as f32; + let freqs = compute_yarn_freqs(head_dim_i, base, factor_f, orig, 32.0, 1.0); + (Some(freqs), yarn_get_mscale(factor_f, 1.0)) + } + _ => (None, 1.0), + }; + + let head_dim_f = head_dim_i as f32; + let attention_scale = head_dim_f.sqrt().recip(); + + Ok(BonsaiQ1Gpu { + config: self.config, + layers: gpu_layers, + embed, + lm_head, + final_norm, + yarn_freqs, + yarn_mscale, + attention_scale, + }) + } +} + +impl BonsaiQ1Gpu { + pub fn num_layers(&self) -> usize { + self.layers.len() + } + + /// Gather embedding rows for a token-ID tensor. + /// + /// Uses MLX dequantize after gathering the selected packed rows. This path + /// requires bits=1 affine support in the active MLX runtime. + fn embed_rows(&self, ids: &Array) -> Result { + let shape = ids.shape().to_vec(); + let flat = ids.flatten(None, None)?; + let w = self.embed.w.take_axis(&flat, 0)?; + let s = self.embed.scales.take_axis(&flat, 0)?; + let b = self.embed.biases.take_axis(&flat, 0)?; + let out = ops::dequantize(&w, &s, &b, GROUP_SIZE_I32, BITS)?; + let mut ret_shape: Vec = shape; + ret_shape.push(-1); + out.reshape(&ret_shape) + } + + fn apply_rope(&self, x: &Array, offset: i32) -> Result { + let head_dim = i32::try_from(self.config.head_dim) + .map_err(|_| Exception::custom("head_dim overflows i32"))?; + let offset_array = Array::from_int(offset); + apply_yarn_rope( + x, + head_dim, + self.config.rope_theta as f32, + self.yarn_freqs.as_ref(), + self.yarn_mscale, + &offset_array, + false, // Qwen3 layout + ) + } + + /// Run the decoder trunk and return final-normed hidden `[B, T, hidden]`. + /// Shared body for `forward` (last-position logits) and + /// `forward_all_logits` (all-position logits, used by spec-decode verify). + /// + /// Body lives in [`forward_trunk_free`] so `compile_with_state` can wrap + /// it via a free-fn pointer (the `Copy + 'static` closure constraint + /// forbids capturing `&self`). + fn forward_trunk( + &self, + inputs: &Array, + cache: &mut Vec>, + ) -> Result { + forward_trunk_free(self, cache, inputs) + } + + /// Apply LM head (or tied embed) to `[B, T, hidden]` → `[B, T, vocab]`. + fn project_logits(&self, h: &Array) -> Result { + let logits = match &self.lm_head { + Some(head) => head.forward(h)?, + None => self.embed.forward(h)?, + }; + // Logits are returned as f32 by API contract (callers do as_slice:: + // for argmax / softmax). The trunk now stays in fp16 throughout (after + // the apply_yarn_rope dtype fix), so we cast here at the boundary. + logits.as_dtype(Dtype::Float32) + } + + /// Causal forward. Returns logits `[B, 1, vocab]` for the last position + /// (mlx_lm convention). + pub fn forward( + &self, + inputs: &Array, + cache: &mut Vec>, + ) -> Result { + let h = self.forward_trunk(inputs, cache)?; + let t = *h + .shape() + .get(1) + .ok_or_else(|| Exception::custom("trunk hidden missing T dim"))?; + let last = if t > 1 { h.index((.., -1.., ..)) } else { h }; + self.project_logits(&last) + } + + /// Causal forward returning logits at **every** position `[B, T, vocab]`. + /// Used by speculative-decode target verify: given the draft prefix, + /// obtain one logits row per proposed token in a single forward pass. + pub fn forward_all_logits( + &self, + inputs: &Array, + cache: &mut Vec>, + ) -> Result { + let h = self.forward_trunk(inputs, cache)?; + self.project_logits(&h) + } + + /// Profiled variant of `forward`: same result, but attributes per-section + /// wall time into `times`. Forces `.eval()` after every section (kills + /// lazy batching — that's the point: ratios matter, absolutes don't). + /// + /// Used by `bench_bonsai_q1_decode_breakdown` to answer the + /// dispatch-bound-vs-matmul-bound question for Bonsai-8B AR parity. + pub fn forward_profiled( + &self, + inputs: &Array, + cache: &mut Vec>, + times: &mut SectionTimes, + ) -> Result { + let h = self.forward_trunk_profiled(inputs, cache, times)?; + let t0 = std::time::Instant::now(); + let t = *h + .shape() + .get(1) + .ok_or_else(|| Exception::custom("trunk hidden missing T dim"))?; + let last = if t > 1 { h.index((.., -1.., ..)) } else { h }; + let logits = self.project_logits(&last)?; + logits.eval()?; + times.add("lm_head", t0.elapsed().as_nanos()); + Ok(logits) + } + + /// Profiled mirror of `forward_trunk`. Inserts `eval + record` at each + /// semantic section boundary. Sections are grouped by operation type + /// (qkv projections together, mlp up+gate together, etc.) — per-layer + /// noise is collapsed into section totals across all layers. + #[allow(non_snake_case)] + fn forward_trunk_profiled( + &self, + inputs: &Array, + cache: &mut Vec>, + times: &mut SectionTimes, + ) -> Result { + use std::time::Instant; + + let shape = inputs.shape(); + let B = *shape + .first() + .ok_or_else(|| Exception::custom("inputs must have >= 2 dims"))?; + let T = *shape + .get(1) + .ok_or_else(|| Exception::custom("inputs must have >= 2 dims"))?; + + if cache.is_empty() { + *cache = (0..self.layers.len()) + .map(|_| Some(SteppingKeyValueCache::new())) + .collect(); + } else if cache.len() != self.layers.len() { + return Err(Exception::custom(format!( + "cache len {} != num_layers {}", + cache.len(), + self.layers.len() + ))); + } + + // Sync point: make sure prior work isn't folded into embed_rows time. + inputs.eval()?; + + let t0 = Instant::now(); + let mut h = self.embed_rows(inputs)?; + h.eval()?; + times.add("embed_rows", t0.elapsed().as_nanos()); + + let mask = create_attention_mask(&h, cache, None)?; + + let heads = i32::try_from(self.config.heads) + .map_err(|_| Exception::custom("heads overflows i32"))?; + let kv_heads = i32::try_from(self.config.kv_heads) + .map_err(|_| Exception::custom("kv_heads overflows i32"))?; + let rms_eps = self.config.rms_norm_eps; + + for (layer, layer_cache) in self.layers.iter().zip(cache.iter_mut()) { + let t0 = Instant::now(); + let normed = fast::rms_norm(&h, &layer.input_norm, rms_eps)?; + normed.eval()?; + times.add("input_norm", t0.elapsed().as_nanos()); + + // qkv projections — 3× quantized_matmul on the same input. + let t0 = Instant::now(); + let q = layer.q_proj.forward(&normed)?; + let k = layer.k_proj.forward(&normed)?; + let v = layer.v_proj.forward(&normed)?; + q.eval()?; + k.eval()?; + v.eval()?; + times.add("qkv_proj", t0.elapsed().as_nanos()); + + // Reshape to [B, L, n_heads, head_dim] then transpose to + // [B, n_heads, L, head_dim]. Metadata-only; lumped with qk_norm. + let q = q + .reshape(&[B, T, heads, -1])? + .transpose_axes(&[0, 2, 1, 3])?; + let k = k + .reshape(&[B, T, kv_heads, -1])? + .transpose_axes(&[0, 2, 1, 3])?; + let v = v + .reshape(&[B, T, kv_heads, -1])? + .transpose_axes(&[0, 2, 1, 3])?; + + let t0 = Instant::now(); + let q = fast::rms_norm(&q, &layer.q_norm, rms_eps)?; + let k = fast::rms_norm(&k, &layer.k_norm, rms_eps)?; + q.eval()?; + k.eval()?; + times.add("qk_norm", t0.elapsed().as_nanos()); + + let offset = layer_cache.as_ref().map_or(0, KeyValueCache::offset); + let t0 = Instant::now(); + let q = self.apply_rope(&q, offset)?; + let k = self.apply_rope(&k, offset)?; + q.eval()?; + k.eval()?; + times.add("rope", t0.elapsed().as_nanos()); + + let mask_arr = match &mask { + Some(crate::utils::AttentionMask::Array(a)) => Some(a), + _ => None, + }; + let mask_arr_opt: Option<&Array> = mask_arr; + + let t0 = Instant::now(); + let attn_out = match layer_cache.as_mut() { + Some(c) => cached_scaled_dot_product_attention( + q, + c, + k, + v, + self.attention_scale, + mask_arr_opt, + )?, + None => fast::scaled_dot_product_attention( + q, + k, + v, + self.attention_scale, + mask_arr_opt.map(mlx_rs::fast::ScaledDotProductAttentionMask::Array), + None::<&Array>, + )?, + }; + attn_out.eval()?; + times.add("sdpa_kv", t0.elapsed().as_nanos()); + + let attn_out = attn_out + .transpose_axes(&[0, 2, 1, 3])? + .reshape(&[B, T, -1])?; + + let t0 = Instant::now(); + let attn_out = layer.o_proj.forward(&attn_out)?; + attn_out.eval()?; + times.add("o_proj", t0.elapsed().as_nanos()); + + let t0 = Instant::now(); + let h_post_attn = h.add(&attn_out)?; + h_post_attn.eval()?; + times.add("residual", t0.elapsed().as_nanos()); + + let t0 = Instant::now(); + let normed_post = fast::rms_norm(&h_post_attn, &layer.post_attn_norm, rms_eps)?; + normed_post.eval()?; + times.add("post_attn_norm", t0.elapsed().as_nanos()); + + let t0 = Instant::now(); + let gate = layer.gate_proj.forward(&normed_post)?; + let up = layer.up_proj.forward(&normed_post)?; + gate.eval()?; + up.eval()?; + times.add("mlp_up_gate", t0.elapsed().as_nanos()); + + let t0 = Instant::now(); + let mlp_hidden = mlx_rs::nn::silu(&gate)?.multiply(&up)?; + mlp_hidden.eval()?; + times.add("silu_mul", t0.elapsed().as_nanos()); + + let t0 = Instant::now(); + let mlp_out = layer.down_proj.forward(&mlp_hidden)?; + mlp_out.eval()?; + times.add("mlp_down", t0.elapsed().as_nanos()); + + let t0 = Instant::now(); + h = h_post_attn.add(&mlp_out)?; + h.eval()?; + times.add("residual", t0.elapsed().as_nanos()); + } + + let t0 = Instant::now(); + let out = fast::rms_norm(&h, &self.final_norm, rms_eps)?; + out.eval()?; + times.add("final_norm", t0.elapsed().as_nanos()); + Ok(out) + } +} + +/// Free-function body of the decoder trunk. +/// +/// Lives at module scope (not as a method) so a **function pointer** to +/// [`decode_step_free`] satisfies `compile_with_state`'s +/// `F: Copy + 'static` bound — a closure capturing `&self` would not. +/// All `self.xxx` access is replaced with `gpu.xxx`; `embed_rows`, +/// `apply_rope`, and `project_logits` are called as methods on `gpu` +/// (they are already `&self`-only, so no further plumbing is needed). +#[allow(non_snake_case)] +pub fn forward_trunk_free( + gpu: &BonsaiQ1Gpu, + cache: &mut Vec>, + inputs: &Array, +) -> Result { + let shape = inputs.shape(); + let B = *shape + .first() + .ok_or_else(|| Exception::custom("inputs must have >= 2 dims"))?; + let T = *shape + .get(1) + .ok_or_else(|| Exception::custom("inputs must have >= 2 dims"))?; + + if cache.is_empty() { + *cache = (0..gpu.layers.len()) + .map(|_| Some(SteppingKeyValueCache::new())) + .collect(); + } else if cache.len() != gpu.layers.len() { + return Err(Exception::custom(format!( + "cache len {} != num_layers {}", + cache.len(), + gpu.layers.len() + ))); + } + + let mut h = gpu.embed_rows(inputs)?; // [B, L, hidden] + + let mask = create_attention_mask(&h, cache, None)?; + + let heads = + i32::try_from(gpu.config.heads).map_err(|_| Exception::custom("heads overflows i32"))?; + let kv_heads = i32::try_from(gpu.config.kv_heads) + .map_err(|_| Exception::custom("kv_heads overflows i32"))?; + let rms_eps = gpu.config.rms_norm_eps; + + for (layer, layer_cache) in gpu.layers.iter().zip(cache.iter_mut()) { + let normed = fast::rms_norm(&h, &layer.input_norm, rms_eps)?; + + let q = layer.q_proj.forward(&normed)?; + let k = layer.k_proj.forward(&normed)?; + let v = layer.v_proj.forward(&normed)?; + + let q = q + .reshape(&[B, T, heads, -1])? + .transpose_axes(&[0, 2, 1, 3])?; + let k = k + .reshape(&[B, T, kv_heads, -1])? + .transpose_axes(&[0, 2, 1, 3])?; + let v = v + .reshape(&[B, T, kv_heads, -1])? + .transpose_axes(&[0, 2, 1, 3])?; + + let q = fast::rms_norm(&q, &layer.q_norm, rms_eps)?; + let k = fast::rms_norm(&k, &layer.k_norm, rms_eps)?; + + let offset = layer_cache.as_ref().map_or(0, KeyValueCache::offset); + let q = gpu.apply_rope(&q, offset)?; + let k = gpu.apply_rope(&k, offset)?; + + let mask_arr = match &mask { + Some(crate::utils::AttentionMask::Array(a)) => Some(a), + _ => None, + }; + let mask_arr_opt: Option<&Array> = mask_arr; + + let attn_out = match layer_cache.as_mut() { + Some(c) => { + cached_scaled_dot_product_attention(q, c, k, v, gpu.attention_scale, mask_arr_opt)? + } + None => fast::scaled_dot_product_attention( + q, + k, + v, + gpu.attention_scale, + mask_arr_opt.map(mlx_rs::fast::ScaledDotProductAttentionMask::Array), + None::<&Array>, + )?, + }; + + let attn_out = attn_out + .transpose_axes(&[0, 2, 1, 3])? + .reshape(&[B, T, -1])?; + let attn_out = layer.o_proj.forward(&attn_out)?; + let h_post_attn = h.add(&attn_out)?; + + let normed_post = fast::rms_norm(&h_post_attn, &layer.post_attn_norm, rms_eps)?; + let gate = layer.gate_proj.forward(&normed_post)?; + let up = layer.up_proj.forward(&normed_post)?; + let mlp_hidden = mlx_rs::nn::silu(&gate)?.multiply(&up)?; + let mlp_out = layer.down_proj.forward(&mlp_hidden)?; + + h = h_post_attn.add(&mlp_out)?; + } + + fast::rms_norm(&h, &gpu.final_norm, rms_eps) +} + +/// Owned state wrapper for [`compile_with_state`]-driven decoding. +/// +/// `compile_with_state` takes the state by `&mut U` where `U: Updatable`. +/// Wrapping the model **and** the per-layer KV cache in one owned struct +/// lets us hand-roll a single `Updatable` impl whose positional iteration +/// order covers both — safer than fighting lifetimes on `(&mut gpu, cache)` +/// tuples. See session-25 recap for the design rationale. +/// +/// Expected construction: after prefill, move `gpu` and the filled cache +/// vector into this struct, run the decode loop with a compiled step, +/// then destructure back out when done. +pub struct BonsaiQ1DecodeState { + pub gpu: BonsaiQ1Gpu, + pub cache: Vec>, +} + +/// Number of updatable `Array`s per decoder layer: +/// - `input_norm` + 3×(w,s,b) qkv + `q_norm` + `k_norm` + 3×(w,s,b) o_proj +/// ... wait: 1 + 3×3 + 2 + 3 + 1 + 3×3 = 1+9+2+3+1+9 = **25**. +/// Corresponds to the array push order in [`BonsaiQ1DecodeState::updatable_states`]. +const PER_LAYER_UPDATABLE: usize = 25; + +impl mlx_rs::utils::Updatable for BonsaiQ1DecodeState { + fn updatable_states_len(&self) -> usize { + let mut n = 3 // embed (w, scales, biases) + + self.gpu.layers.len() * PER_LAYER_UPDATABLE + + 1; // final_norm + if self.gpu.lm_head.is_some() { + n += 3; + } + if self.gpu.yarn_freqs.is_some() { + n += 1; + } + for slot in &self.cache { + if let Some(c) = slot { + if c.keys().is_some() { + n += 1; + } + if c.values().is_some() { + n += 1; + } + } + } + n + } + + fn updatable_states(&self) -> impl IntoIterator { + let mut v: Vec<&Array> = Vec::with_capacity(self.updatable_states_len()); + v.push(&self.gpu.embed.w); + v.push(&self.gpu.embed.scales); + v.push(&self.gpu.embed.biases); + for layer in &self.gpu.layers { + v.push(&layer.input_norm); + v.push(&layer.q_proj.w); + v.push(&layer.q_proj.scales); + v.push(&layer.q_proj.biases); + v.push(&layer.k_proj.w); + v.push(&layer.k_proj.scales); + v.push(&layer.k_proj.biases); + v.push(&layer.v_proj.w); + v.push(&layer.v_proj.scales); + v.push(&layer.v_proj.biases); + v.push(&layer.q_norm); + v.push(&layer.k_norm); + v.push(&layer.o_proj.w); + v.push(&layer.o_proj.scales); + v.push(&layer.o_proj.biases); + v.push(&layer.post_attn_norm); + v.push(&layer.gate_proj.w); + v.push(&layer.gate_proj.scales); + v.push(&layer.gate_proj.biases); + v.push(&layer.up_proj.w); + v.push(&layer.up_proj.scales); + v.push(&layer.up_proj.biases); + v.push(&layer.down_proj.w); + v.push(&layer.down_proj.scales); + v.push(&layer.down_proj.biases); + } + v.push(&self.gpu.final_norm); + if let Some(lm) = self.gpu.lm_head.as_ref() { + v.push(&lm.w); + v.push(&lm.scales); + v.push(&lm.biases); + } + if let Some(y) = self.gpu.yarn_freqs.as_ref() { + v.push(y); + } + for slot in &self.cache { + if let Some(c) = slot { + if let Some(k) = c.keys() { + v.push(k); + } + if let Some(val) = c.values() { + v.push(val); + } + } + } + v + } + + fn updatable_states_mut(&mut self) -> impl IntoIterator { + let mut v: Vec<&mut Array> = Vec::with_capacity(self.updatable_states_len()); + v.push(&mut self.gpu.embed.w); + v.push(&mut self.gpu.embed.scales); + v.push(&mut self.gpu.embed.biases); + for layer in &mut self.gpu.layers { + v.push(&mut layer.input_norm); + v.push(&mut layer.q_proj.w); + v.push(&mut layer.q_proj.scales); + v.push(&mut layer.q_proj.biases); + v.push(&mut layer.k_proj.w); + v.push(&mut layer.k_proj.scales); + v.push(&mut layer.k_proj.biases); + v.push(&mut layer.v_proj.w); + v.push(&mut layer.v_proj.scales); + v.push(&mut layer.v_proj.biases); + v.push(&mut layer.q_norm); + v.push(&mut layer.k_norm); + v.push(&mut layer.o_proj.w); + v.push(&mut layer.o_proj.scales); + v.push(&mut layer.o_proj.biases); + v.push(&mut layer.post_attn_norm); + v.push(&mut layer.gate_proj.w); + v.push(&mut layer.gate_proj.scales); + v.push(&mut layer.gate_proj.biases); + v.push(&mut layer.up_proj.w); + v.push(&mut layer.up_proj.scales); + v.push(&mut layer.up_proj.biases); + v.push(&mut layer.down_proj.w); + v.push(&mut layer.down_proj.scales); + v.push(&mut layer.down_proj.biases); + } + v.push(&mut self.gpu.final_norm); + if let Some(lm) = self.gpu.lm_head.as_mut() { + v.push(&mut lm.w); + v.push(&mut lm.scales); + v.push(&mut lm.biases); + } + if let Some(y) = self.gpu.yarn_freqs.as_mut() { + v.push(y); + } + for slot in &mut self.cache { + if let Some(c) = slot { + let (k_opt, v_opt) = c.key_value_arrays_mut(); + if let Some(k) = k_opt { + v.push(k); + } + if let Some(val) = v_opt { + v.push(val); + } + } + } + v + } +} + +/// Free-fn decode step compatible with `compile_with_state`. +/// +/// `state.cache` **must** be populated by a prefill call before this runs: +/// compile-wrap is applied only in decode, and shape consistency across +/// steps (for the MLX per-shape trace cache) requires +/// [`SteppingKeyValueCache::reserve_max_tokens`] ahead of the first +/// `update_dense`. +pub fn decode_step_free( + state: &mut BonsaiQ1DecodeState, + inputs: &Array, +) -> Result { + let h = forward_trunk_free(&state.gpu, &mut state.cache, inputs)?; + let t = *h + .shape() + .get(1) + .ok_or_else(|| Exception::custom("trunk hidden missing T dim"))?; + let last = if t > 1 { h.index((.., -1.., ..)) } else { h }; + state.gpu.project_logits(&last) +} + +/// Per-section wall-time accumulator for the Bonsai-Q1 forward pass. +/// +/// Exists only to attribute the 45 ms/tok Bonsai-8B AR decode cost to +/// individual sections (embed / norms / qkv / rope / sdpa / o_proj / mlp / lm_head). +/// Each section's compute is force-`.eval()`'d to prevent MLX lazy batching +/// from pooling multiple sections into one materialization — ratios between +/// sections are meaningful even though absolutes will be slower than the +/// unprofiled path. +#[derive(Debug, Default, Clone)] +pub struct SectionTimes { + totals: std::collections::BTreeMap<&'static str, (u128, u64)>, +} + +impl SectionTimes { + pub fn new() -> Self { + Self::default() + } + + pub fn add(&mut self, name: &'static str, ns: u128) { + let e = self.totals.entry(name).or_insert((0, 0)); + e.0 += ns; + e.1 += 1; + } + + /// Total across all sections (ns). + pub fn total_ns(&self) -> u128 { + self.totals.values().map(|(t, _)| *t).sum() + } + + /// Section totals: `(name, total_ns, call_count)`, sorted by ns descending. + pub fn entries(&self) -> Vec<(&'static str, u128, u64)> { + let mut v: Vec<_> = self.totals.iter().map(|(k, (t, n))| (*k, *t, *n)).collect(); + v.sort_by_key(|b| std::cmp::Reverse(b.1)); + v + } +} + +fn load_packed( + tensors: &SafeTensors<'_>, + prefix: &str, + out_features: usize, + in_features: usize, + who: &str, +) -> Result { + if in_features % GROUP_SIZE != 0 { + return Err(format!( + "{who}: in_features {in_features} not divisible by group_size {GROUP_SIZE}" + )); + } + let packed_cols = in_features / 32; + let n_groups = in_features / GROUP_SIZE; + + let w_view = tensors + .tensor(&format!("{prefix}.weight")) + .map_err(|e| format!("{who}: {prefix}.weight: {e}"))?; + let s_view = tensors + .tensor(&format!("{prefix}.scales")) + .map_err(|e| format!("{who}: {prefix}.scales: {e}"))?; + let b_view = tensors + .tensor(&format!("{prefix}.biases")) + .map_err(|e| format!("{who}: {prefix}.biases: {e}"))?; + + let w_bytes = w_view.data(); + let s_bytes = s_view.data(); + let b_bytes = b_view.data(); + + let expected_w_bytes = out_features * packed_cols * 4; + if w_bytes.len() != expected_w_bytes { + return Err(format!( + "{who}: weight byte-size mismatch: got {} expected {}", + w_bytes.len(), + expected_w_bytes, + )); + } + let expected_sb_bytes = out_features * n_groups * 2; + if s_bytes.len() != expected_sb_bytes { + return Err(format!( + "{who}: scales byte-size mismatch: got {} expected {}", + s_bytes.len(), + expected_sb_bytes, + )); + } + if b_bytes.len() != expected_sb_bytes { + return Err(format!( + "{who}: biases byte-size mismatch: got {} expected {}", + b_bytes.len(), + expected_sb_bytes, + )); + } + + let w_packed: Vec = w_bytes + .chunks_exact(4) + .map(|c| u32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(); + let scales = bytes_to_f16_vec(s_bytes); + let biases = bytes_to_f16_vec(b_bytes); + + Ok(PackedQ1Linear { + w_packed, + scales, + biases, + out_features, + in_features, + }) +} + +fn load_f16(tensors: &SafeTensors<'_>, name: &str) -> Result, String> { + let view = tensors.tensor(name).map_err(|e| format!("{name}: {e}"))?; + Ok(bytes_to_f16_vec(view.data())) +} + +fn bytes_to_f16_vec(b: &[u8]) -> Vec { + b.chunks_exact(2) + .map(|c| f16::from_bits(u16::from_le_bytes([c[0], c[1]]))) + .collect() +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- diff --git a/crates/higgs-models/src/cache.rs b/crates/higgs-models/src/cache.rs index 79eabd2a..14f5bed8 100644 --- a/crates/higgs-models/src/cache.rs +++ b/crates/higgs-models/src/cache.rs @@ -453,6 +453,14 @@ impl SteppingKeyValueCache { self.values.as_ref() } + /// Simultaneous mutable access to the key and value arrays. + /// + /// Re-borrows both optional fields from a single `&mut` split to satisfy + /// the borrow checker when both must yield from one iterator. + pub const fn key_value_arrays_mut(&mut self) -> (Option<&mut Array>, Option<&mut Array>) { + (self.keys.as_mut(), self.values.as_mut()) + } + /// Create a pre-filled cache from existing K/V arrays. /// /// Sets `offset = keys.shape()[2]` so the next `update_dense` triggers a diff --git a/crates/higgs-models/src/deepseek_v2.rs b/crates/higgs-models/src/deepseek_v2.rs index e3b09cad..625264fc 100644 --- a/crates/higgs-models/src/deepseek_v2.rs +++ b/crates/higgs-models/src/deepseek_v2.rs @@ -120,9 +120,14 @@ impl DeepSeekV2ModelArgs { // YaRN RoPE helpers // --------------------------------------------------------------------------- +#[allow( + clippy::as_conversions, + clippy::cast_precision_loss, + clippy::cast_possible_truncation +)] fn yarn_find_correction_dim(num_rotations: f32, dim: i32, base: f32, max_pos: i32) -> f32 { - let dim_f = f32::from(i16::try_from(dim).unwrap_or(i16::MAX)); - let max_pos_f = f32::from(i16::try_from(max_pos).unwrap_or(i16::MAX)); + let dim_f = dim as f32; + let max_pos_f = max_pos as f32; (dim_f * (max_pos_f / (num_rotations * 2.0 * PI)).ln()) / (2.0 * base.ln()) } @@ -154,6 +159,7 @@ fn yarn_get_mscale(scale: f32, mscale: f32) -> f32 { /// Precompute `YaRN`-interpolated `RoPE` frequencies. #[allow( clippy::as_conversions, + clippy::cast_precision_loss, clippy::cast_possible_truncation, clippy::cast_sign_loss, clippy::indexing_slicing @@ -167,14 +173,14 @@ fn compute_yarn_freqs( beta_slow: f32, ) -> Array { let half_dim = dim / 2; - let dim_f = f32::from(i16::try_from(dim).unwrap_or(i16::MAX)); + let dim_f = dim as f32; // freq_extra = base^(arange(0, dim, 2) / dim) -- standard theta // freq_inter = scaling_factor * freq_extra -- extended theta let mut freq_extra = Vec::with_capacity(half_dim as usize); let mut freq_inter = Vec::with_capacity(half_dim as usize); for i in 0..half_dim { - let exp = f32::from(i16::try_from(2 * i).unwrap_or(0)) / dim_f; + let exp = (2 * i) as f32 / dim_f; let theta = base.powf(exp); freq_extra.push(theta); freq_inter.push(scaling_factor * theta); @@ -183,8 +189,8 @@ fn compute_yarn_freqs( let (low, high) = yarn_find_correction_range(beta_fast, beta_slow, dim, base, orig_max_pos); // Linear ramp mask: 0 at low, 1 at high - let low_f = f32::from(i16::try_from(low).unwrap_or(0)); - let high_f = f32::from(i16::try_from(high).unwrap_or(0)); + let low_f = low as f32; + let high_f = high as f32; let range = if (high_f - low_f).abs() < 0.001 { high_f - low_f + 0.001 } else { @@ -195,7 +201,7 @@ fn compute_yarn_freqs( // freq_mask = 1 - ramp (high mask = use freq_extra, low mask = use freq_inter) let mut freqs = Vec::with_capacity(half_dim as usize); for i in 0..half_dim as usize { - let idx_f = f32::from(i16::try_from(i).unwrap_or(0)); + let idx_f = i as f32; let ramp = ((idx_f - low_f) / range).clamp(0.0, 1.0); let mask = 1.0 - ramp; let inter = freq_inter[i]; @@ -300,6 +306,7 @@ impl DeepSeekV2Attention { // YaRN RoPE #[allow( clippy::as_conversions, + clippy::cast_precision_loss, clippy::cast_possible_truncation, clippy::option_if_let_else )] diff --git a/crates/higgs-models/src/lib.rs b/crates/higgs-models/src/lib.rs index 407f9fbf..fe52d245 100644 --- a/crates/higgs-models/src/lib.rs +++ b/crates/higgs-models/src/lib.rs @@ -1,3 +1,4 @@ +pub mod bonsai_q1; pub mod cache; pub mod deepseek_v2; pub mod error; @@ -13,9 +14,11 @@ pub mod starcoder2; pub mod transformer; pub mod turboquant; pub mod utils; +pub mod yarn; use std::collections::{HashMap, HashSet}; use std::path::Path; +use std::sync::atomic::{AtomicBool, Ordering}; use mlx_rs::module::ModuleParametersExt; use mlx_rs::ops::indexing::IndexOp; @@ -27,6 +30,8 @@ use serde_json::Value; use crate::error::ModelError; use crate::turboquant::KvCacheConfig; +static BONSAI_IGNORED_MASK_WARNED: AtomicBool = AtomicBool::new(false); + // --------------------------------------------------------------------------- // SamplingParams -- configurable sampling parameters // --------------------------------------------------------------------------- @@ -132,6 +137,8 @@ pub enum AnyModel { LlavaQwen2(llava_qwen2::LlavaQwen2Model), /// DeepSeek-V2 with Multi-head Latent Attention and sparse `MoE`. DeepSeekV2(deepseek_v2::DeepSeekV2CausalLM), + /// Bonsai-Q1: packed 1.25-bpw Qwen3-shaped target (1.7B / 8B). + BonsaiQ1(bonsai_q1::BonsaiQ1Gpu), } fn checked_head_dim(hidden_size: i32, num_attention_heads: i32) -> Result { @@ -213,6 +220,16 @@ impl AnyModel { (Self::LlavaQwen2(m), AnyCache::KV(c)) => m.forward_text(inputs, mask, c), (Self::DeepSeekV2(m), AnyCache::KV(c)) => m.forward(inputs, mask, c), (Self::Qwen3Next(m), AnyCache::Hybrid(c)) => m.forward(inputs, mask, c), + // BonsaiQ1 builds its causal mask internally; any externally-provided + // mask is ignored (causal-only semantics). + (Self::BonsaiQ1(m), AnyCache::KV(c)) => { + if mask.is_some() && !BONSAI_IGNORED_MASK_WARNED.swap(true, Ordering::Relaxed) { + tracing::warn!( + "BonsaiQ1 ignores externally provided masks and builds its own causal mask" + ); + } + m.forward(inputs, c) + } _ => Err(Exception::custom("Model/cache type mismatch")), } } @@ -233,6 +250,7 @@ impl AnyModel { (Self::LlavaQwen2(m), AnyCache::KV(c)) => m.forward_text_hidden(inputs, mask, c), (Self::DeepSeekV2(m), AnyCache::KV(c)) => m.forward_hidden(inputs, mask, c), (Self::Qwen3Next(m), AnyCache::Hybrid(c)) => m.forward_hidden(inputs, mask, c), + (Self::BonsaiQ1(m), AnyCache::KV(c)) => bonsai_q1::forward_trunk_free(m, c, inputs), _ => Err(Exception::custom("Model/cache type mismatch")), } } @@ -347,7 +365,8 @@ impl AnyModel { | Self::Phi3(_) | Self::Starcoder2(_) | Self::LlavaQwen2(_) - | Self::DeepSeekV2(_) => Err(Exception::custom( + | Self::DeepSeekV2(_) + | Self::BonsaiQ1(_) => Err(Exception::custom( "Batched forward only supported for Transformer models", )), } @@ -373,7 +392,8 @@ impl AnyModel { | Self::Phi3(_) | Self::Starcoder2(_) | Self::LlavaQwen2(_) - | Self::DeepSeekV2(_) => None, + | Self::DeepSeekV2(_) + | Self::BonsaiQ1(_) => None, } } @@ -394,7 +414,8 @@ impl AnyModel { | Self::Phi3(_) | Self::Starcoder2(_) | Self::LlavaQwen2(_) - | Self::DeepSeekV2(_) => Err(Exception::custom("MTP not supported for this model")), + | Self::DeepSeekV2(_) + | Self::BonsaiQ1(_) => Err(Exception::custom("MTP not supported for this model")), } } @@ -413,7 +434,8 @@ impl AnyModel { | Self::Phi3(_) | Self::Starcoder2(_) | Self::LlavaQwen2(_) - | Self::DeepSeekV2(_) => Err(Exception::custom("MTP not supported for this model")), + | Self::DeepSeekV2(_) + | Self::BonsaiQ1(_) => Err(Exception::custom("MTP not supported for this model")), } } @@ -435,7 +457,7 @@ impl AnyModel { } /// The model's hidden dimension. - pub const fn hidden_size(&self) -> i32 { + pub fn hidden_size(&self) -> i32 { match self { Self::Transformer(m) => m.args.hidden_size, Self::Qwen3Moe(m) => m.args.hidden_size, @@ -445,6 +467,7 @@ impl AnyModel { Self::Starcoder2(m) => m.args.hidden_size, Self::LlavaQwen2(m) => m.hidden_size(), Self::DeepSeekV2(m) => m.args.hidden_size, + Self::BonsaiQ1(m) => i32::try_from(m.config.hidden).unwrap_or(i32::MAX), } } @@ -481,6 +504,10 @@ impl AnyModel { m.args.num_key_value_heads, m.args.qk_nope_head_dim + m.args.qk_rope_head_dim, )), + Self::BonsaiQ1(m) => Ok(( + i32::try_from(m.config.kv_heads).map_err(|e| Exception::custom(e.to_string()))?, + i32::try_from(m.config.head_dim).map_err(|e| Exception::custom(e.to_string()))?, + )), } } @@ -489,6 +516,7 @@ impl AnyModel { self.make_cache_with_config(KvCacheConfig::default()) } + #[allow(clippy::too_many_lines)] pub fn make_cache_with_config( &self, kv_cache_config: KvCacheConfig, @@ -583,6 +611,16 @@ impl AnyModel { Ok(AnyCache::Hybrid(m.make_cache())) } } + Self::BonsaiQ1(m) => { + if kv_cache_config.is_turboquant() { + return Err(Exception::custom( + "TurboQuant is not supported for BonsaiQ1 (1-bit packed engine)", + )); + } + let layers = + i32::try_from(m.config.layers).map_err(|e| Exception::custom(e.to_string()))?; + Ok(make_kv_cache(layers)) + } } } @@ -601,7 +639,8 @@ impl AnyModel { | Self::Gemma2(_) | Self::Phi3(_) | Self::Starcoder2(_) - | Self::DeepSeekV2(_) => None, + | Self::DeepSeekV2(_) + | Self::BonsaiQ1(_) => None, } } diff --git a/crates/higgs-models/src/yarn.rs b/crates/higgs-models/src/yarn.rs new file mode 100644 index 00000000..5081b8a0 --- /dev/null +++ b/crates/higgs-models/src/yarn.rs @@ -0,0 +1,223 @@ +//! YaRN RoPE helpers shared across models. +//! +//! Extracted verbatim from the original site in `deepseek_v2.rs`. The +//! `apply_yarn_rope` wrapper adds a `traditional` flag so Qwen3-family models +//! (Bonsai) can reuse the same freq precomputation with `traditional=false`, +//! while DeepSeek stays on `traditional=true`. + +#![allow(clippy::doc_markdown)] // YaRN, RoPE, etc. are domain terms, not items. + +use std::f32::consts::PI; + +use mlx_rs::{Array, error::Exception, fast}; + +#[allow( + clippy::as_conversions, + clippy::cast_precision_loss, + clippy::cast_possible_truncation +)] +fn yarn_find_correction_dim(num_rotations: f32, dim: i32, base: f32, max_pos: i32) -> f32 { + let dim_f = dim as f32; + let max_pos_f = max_pos as f32; + (dim_f * (max_pos_f / (num_rotations * 2.0 * PI)).ln()) / (2.0 * base.ln()) +} + +#[allow( + clippy::as_conversions, + clippy::cast_possible_truncation, + clippy::cast_sign_loss +)] +fn yarn_find_correction_range( + low_rot: f32, + high_rot: f32, + dim: i32, + base: f32, + max_pos: i32, +) -> (i32, i32) { + let low = yarn_find_correction_dim(low_rot, dim, base, max_pos).floor() as i32; + let high = yarn_find_correction_dim(high_rot, dim, base, max_pos).ceil() as i32; + (low.max(0), high.min(dim - 1)) +} + +pub(crate) fn yarn_get_mscale(scale: f32, mscale: f32) -> f32 { + if scale <= 1.0 { + 1.0 + } else { + (0.1 * mscale).mul_add(scale.ln(), 1.0) + } +} + +#[allow( + clippy::as_conversions, + clippy::cast_precision_loss, + clippy::cast_possible_truncation, + clippy::cast_sign_loss, + clippy::indexing_slicing +)] +pub(crate) fn compute_yarn_freqs( + dim: i32, + base: f32, + scaling_factor: f32, + orig_max_pos: i32, + beta_fast: f32, + beta_slow: f32, +) -> Array { + let half_dim = dim / 2; + let dim_f = dim as f32; + + let mut freq_extra = Vec::with_capacity(half_dim as usize); + let mut freq_inter = Vec::with_capacity(half_dim as usize); + for i in 0..half_dim { + let exp = (2 * i) as f32 / dim_f; + let theta = base.powf(exp); + freq_extra.push(theta); + freq_inter.push(scaling_factor * theta); + } + + let (low, high) = yarn_find_correction_range(beta_fast, beta_slow, dim, base, orig_max_pos); + + let low_f = low as f32; + let high_f = high as f32; + let range = if (high_f - low_f).abs() < 0.001 { + high_f - low_f + 0.001 + } else { + high_f - low_f + }; + + let mut freqs = Vec::with_capacity(half_dim as usize); + for i in 0..half_dim as usize { + let idx_f = i as f32; + let ramp = ((idx_f - low_f) / range).clamp(0.0, 1.0); + let mask = 1.0 - ramp; + let inter = freq_inter[i]; + let extra = freq_extra[i]; + let denom = inter * mask + extra * (1.0 - mask); + freqs.push((inter * extra) / denom); + } + + Array::from_slice(&freqs, &[half_dim]) +} + +/// Apply YaRN-scaled RoPE. +/// +/// When `mscale != 1.0`, inputs are pre-scaled before rotation (matches the +/// DeepSeek reference). `traditional=false` matches the Qwen3 / LLaMA rope +/// layout; `traditional=true` matches DeepSeek's packed complex layout. +/// +/// `offset` is a scalar `Array` (not an `i32`) so the value is not baked into +/// compiled traces — required for `compile_with_state` wrapping of decode. +pub(crate) fn apply_yarn_rope( + x: &Array, + dim: i32, + base: f32, + yarn_freqs: Option<&Array>, + mscale: f32, + offset: &Array, + traditional: bool, +) -> Result { + let x_scaled = if (mscale - 1.0).abs() > f32::EPSILON { + // Match x's dtype to avoid silent upcast (fp16 → f32) that bleeds into + // the entire attention path (rope, sdpa, o_proj inputs). For Bonsai + // with rope_yarn_factor>1, mscale ≈ 1.14, so this branch fires every + // rope call; without the cast the whole decode runs in f32 and pays + // ~28 ms/step on 8B. See bisect_decode v6 vs v7. + let scalar = Array::from_f32(mscale).as_dtype(x.dtype())?; + x.multiply(&scalar)? + } else { + x.clone() + }; + yarn_freqs.map_or_else( + || { + fast::rope_dynamic( + &x_scaled, + dim, + traditional, + base, + 1.0, + offset, + None::<&Array>, + ) + }, + |freqs| { + fast::rope_dynamic( + &x_scaled, + dim, + traditional, + None::, + 1.0, + offset, + Some(freqs), + ) + }, + ) +} + +#[cfg(test)] +#[allow( + clippy::panic, + clippy::unwrap_used, + clippy::expect_used, + clippy::indexing_slicing, + clippy::as_conversions, + clippy::cast_precision_loss, + clippy::cast_possible_truncation, + clippy::cast_sign_loss, + clippy::cast_lossless +)] +mod tests { + use super::*; + use mlx_rs::random; + + /// Parity: dynamic-offset rope (via `apply_yarn_rope`) must match + /// static-offset `fast::rope` for every offset in 0..64, both with and + /// without precomputed YaRN freqs. Guards the B1 step 1 migration to + /// `fast::rope_dynamic` (prerequisite for `compile_with_state`-wrapped + /// decode — the static `offset: i32` was being baked into the compile + /// trace, forcing a recompile every step). + #[test] + #[ignore = "passes targeted (`cargo test yarn::`) but fails when run after other MLX tests in the same process — global Metal/RNG state contamination, pre-existing harness limitation"] + fn rope_dynamic_matches_static_offset_0_to_64() { + random::seed(71).unwrap(); + // [B=2, H=4, T=1, head_dim=16] — matches decode shape (T=1). + let head_dim: i32 = 16; + let base: f32 = 10_000.0; + let x = random::uniform::<_, f32>(0.0, 1.0, &[2, 4, 1, head_dim], None).unwrap(); + + // Case A: no yarn_freqs (base path). + for offset in 0_i32..64 { + let off_arr = Array::from_int(offset); + let got = apply_yarn_rope(&x, head_dim, base, None, 1.0, &off_arr, false).unwrap(); + let want = fast::rope(&x, head_dim, false, base, 1.0, offset, None::<&Array>).unwrap(); + let diff = (&got - &want) + .abs() + .unwrap() + .max(None) + .unwrap() + .item::(); + assert!( + diff < 1e-5, + "offset={offset} no-freqs: max_diff={diff} >= 1e-5" + ); + } + + // Case B: with precomputed yarn_freqs (Bonsai path). + let freqs = compute_yarn_freqs(head_dim, base, 1.0, 2048, 32.0, 1.0); + for offset in 0_i32..64 { + let off_arr = Array::from_int(offset); + let got = + apply_yarn_rope(&x, head_dim, base, Some(&freqs), 1.0, &off_arr, false).unwrap(); + let want = + fast::rope(&x, head_dim, false, None::, 1.0, offset, Some(&freqs)).unwrap(); + let diff = (&got - &want) + .abs() + .unwrap() + .max(None) + .unwrap() + .item::(); + assert!( + diff < 1e-5, + "offset={offset} with-freqs: max_diff={diff} >= 1e-5" + ); + } + } +} diff --git a/docs/BONSAI_Q1.md b/docs/BONSAI_Q1.md new file mode 100644 index 00000000..7211d7d8 --- /dev/null +++ b/docs/BONSAI_Q1.md @@ -0,0 +1,18 @@ +# Bonsai-Q1 + +Bonsai-Q1 checkpoints are Qwen3-shaped models with MLX 1-bit affine +quantization metadata: + +- `model_type = "qwen3"` +- `quantization.bits = 1` +- `quantization.group_size = 128` + +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. + +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. From 4a099d4e3bc0dc5a788c1e3894a3665cde4ce65b Mon Sep 17 00:00:00 2001 From: "renovate[bot]" <29139614+renovate[bot]@users.noreply.github.com> Date: Wed, 6 May 2026 08:48:02 -0600 Subject: [PATCH 09/29] fix(deps): update rust crate sysinfo to 0.38 (#132) Co-authored-by: renovate[bot] <29139614+renovate[bot]@users.noreply.github.com> --- Cargo.lock | 131 +++++++++++++++++++--------------- crates/higgs-bench/Cargo.toml | 2 +- 2 files changed, 75 insertions(+), 58 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 942dd4b1..fd319b78 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -88,7 +88,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -99,7 +99,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -1060,7 +1060,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -1141,7 +1141,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -1778,7 +1778,7 @@ dependencies = [ "js-sys", "log", "wasm-bindgen", - "windows-core 0.62.2", + "windows-core", ] [[package]] @@ -2000,7 +2000,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -2528,7 +2528,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -2622,12 +2622,31 @@ dependencies = [ "objc2-encode", ] +[[package]] +name = "objc2-core-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" +dependencies = [ + "bitflags 2.11.1", +] + [[package]] name = "objc2-encode" version = "4.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ef25abbcd74fb2609453eb695bd2f860d389e457f67dc17cafc8b8cbc89d0c33" +[[package]] +name = "objc2-io-kit" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33fafba39597d6dc1fb709123dfa8289d39406734be322956a69f0931c73bb15" +dependencies = [ + "libc", + "objc2-core-foundation", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -3370,7 +3389,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3429,7 +3448,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3698,7 +3717,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -3822,15 +3841,15 @@ dependencies = [ [[package]] name = "sysinfo" -version = "0.32.1" +version = "0.38.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c33cd241af0f2e9e3b5c32163b873b29956890b5342e6745b917ce9d490f4af" +checksum = "92ab6a2f8bfe508deb3c6406578252e491d299cbbf3bc0529ecc3313aee4a52f" dependencies = [ - "core-foundation-sys", "libc", "memchr", "ntapi", - "rayon", + "objc2-core-foundation", + "objc2-io-kit", "windows", ] @@ -3862,10 +3881,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.3.4", + "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -4854,7 +4873,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -4865,24 +4884,23 @@ checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" [[package]] name = "windows" -version = "0.57.0" +version = "0.62.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "12342cb4d8e3b046f3d80effd474a7a02447231330ef77d71daa6fbc40681143" +checksum = "527fadee13e0c05939a6a05d5bd6eec6cd2e3dbd648b9f8e447c6518133d8580" dependencies = [ - "windows-core 0.57.0", - "windows-targets 0.52.6", + "windows-collections", + "windows-core", + "windows-future", + "windows-numerics", ] [[package]] -name = "windows-core" -version = "0.57.0" +name = "windows-collections" +version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d2ed2439a290666cd67ecce2b0ffaad89c2a56b976b736e6ece670297897832d" +checksum = "23b2d95af1a8a14a3c7367e1ed4fc9c20e0a26e79551b1454d72583c97cc6610" dependencies = [ - "windows-implement 0.57.0", - "windows-interface 0.57.0", - "windows-result 0.1.2", - "windows-targets 0.52.6", + "windows-core", ] [[package]] @@ -4891,22 +4909,22 @@ version = "0.62.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" dependencies = [ - "windows-implement 0.60.2", - "windows-interface 0.59.3", + "windows-implement", + "windows-interface", "windows-link", - "windows-result 0.4.1", + "windows-result", "windows-strings", ] [[package]] -name = "windows-implement" -version = "0.57.0" +name = "windows-future" +version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9107ddc059d5b6fbfbffdfa7a7fe3e22a226def0b2608f72e9d552763d3e1ad7" +checksum = "e1d6f90251fe18a279739e78025bd6ddc52a7e22f921070ccdc67dde84c605cb" dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.117", + "windows-core", + "windows-link", + "windows-threading", ] [[package]] @@ -4920,17 +4938,6 @@ dependencies = [ "syn 2.0.117", ] -[[package]] -name = "windows-interface" -version = "0.57.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "29bee4b38ea3cde66011baa44dba677c432a78593e202392d1e9070cf2a7fca7" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.117", -] - [[package]] name = "windows-interface" version = "0.59.3" @@ -4949,23 +4956,24 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" [[package]] -name = "windows-registry" -version = "0.6.1" +name = "windows-numerics" +version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "02752bf7fbdcce7f2a27a742f798510f3e5ad88dbe84871e5168e2120c3d5720" +checksum = "6e2e40844ac143cdb44aead537bbf727de9b044e107a0f1220392177d15b0f26" dependencies = [ + "windows-core", "windows-link", - "windows-result 0.4.1", - "windows-strings", ] [[package]] -name = "windows-result" -version = "0.1.2" +name = "windows-registry" +version = "0.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e383302e8ec8515204254685643de10811af0ed97ea37210dc26fb0032647f8" +checksum = "02752bf7fbdcce7f2a27a742f798510f3e5ad88dbe84871e5168e2120c3d5720" dependencies = [ - "windows-targets 0.52.6", + "windows-link", + "windows-result", + "windows-strings", ] [[package]] @@ -5079,6 +5087,15 @@ dependencies = [ "windows_x86_64_msvc 0.53.1", ] +[[package]] +name = "windows-threading" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3949bd5b99cafdf1c7ca86b43ca564028dfe27d66958f2470940f73d86d75b37" +dependencies = [ + "windows-link", +] + [[package]] name = "windows_aarch64_gnullvm" version = "0.42.2" diff --git a/crates/higgs-bench/Cargo.toml b/crates/higgs-bench/Cargo.toml index 7c6e8d18..17f022c9 100644 --- a/crates/higgs-bench/Cargo.toml +++ b/crates/higgs-bench/Cargo.toml @@ -19,7 +19,7 @@ futures = { workspace = true } reqwest = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } -sysinfo = "0.32" +sysinfo = "0.38" tokio = { workspace = true } toml = "1.0" From 372020ec5198b47bfb29babbaf60b4250507a76e Mon Sep 17 00:00:00 2001 From: Jonathan Reyes <150944+panbanda@users.noreply.github.com> Date: Wed, 6 May 2026 09:01:03 -0600 Subject: [PATCH 10/29] fix: default Qwen3.6 to non-thinking mode (#150) * fix: default qwen3.6 to non-thinking mode * fix: address qwen thinking review feedback --- crates/higgs/src/lib.rs | 1 + crates/higgs/src/reasoning.rs | 123 +++++++++++++++++++++++++++ crates/higgs/src/routes/anthropic.rs | 34 ++++++-- crates/higgs/src/routes/chat.rs | 85 +++--------------- crates/higgs/src/types/openai.rs | 44 +++++++++- 5 files changed, 201 insertions(+), 86 deletions(-) create mode 100644 crates/higgs/src/reasoning.rs diff --git a/crates/higgs/src/lib.rs b/crates/higgs/src/lib.rs index 20d0a089..7534c586 100644 --- a/crates/higgs/src/lib.rs +++ b/crates/higgs/src/lib.rs @@ -11,6 +11,7 @@ pub mod metrics_log; pub mod model_download; pub mod model_resolver; pub mod proxy; +pub(crate) mod reasoning; pub mod router; pub mod routes; #[doc(hidden)] diff --git a/crates/higgs/src/reasoning.rs b/crates/higgs/src/reasoning.rs new file mode 100644 index 00000000..795ccea7 --- /dev/null +++ b/crates/higgs/src/reasoning.rs @@ -0,0 +1,123 @@ +use crate::types::openai::ReasoningConfig; + +fn model_defaults_to_non_thinking(model_names: &[&str]) -> bool { + model_names.iter().any(|model_name| { + let normalized = model_name.to_ascii_lowercase(); + normalized.match_indices("qwen3.6").any(|(idx, _)| { + let after = idx + "qwen3.6".len(); + let before_is_boundary = idx == 0 + || normalized + .as_bytes() + .get(idx - 1) + .is_some_and(|b| !b.is_ascii_alphanumeric()); + let after_is_boundary = after == normalized.len() + || normalized + .as_bytes() + .get(after) + .is_some_and(|b| !b.is_ascii_digit()); + before_is_boundary && after_is_boundary + }) + }) +} + +pub fn effective_thinking_enabled( + engine_default: bool, + model_names: &[&str], + reasoning: Option<&ReasoningConfig>, +) -> bool { + if !engine_default { + return false; + } + + match reasoning.and_then(|r| r.effort.as_deref()) { + Some(effort) if effort.is_empty() || effort.eq_ignore_ascii_case("none") => false, + Some(_) => true, + None => !model_defaults_to_non_thinking(model_names), + } +} + +#[cfg(test)] +#[allow(clippy::panic, clippy::unwrap_used)] +mod tests { + use super::*; + + #[test] + fn defaults_qwen35_on() { + assert!(effective_thinking_enabled( + true, + &["mlx-community/Qwen3.5-foo"], + None, + )); + } + + #[test] + fn defaults_qwen36_off_from_route_name() { + assert!(!effective_thinking_enabled( + true, + &["mlx-community/Qwen3.6-35B-A3B-4bit"], + None, + )); + } + + #[test] + fn defaults_qwen36_off_from_engine_name_even_when_aliased() { + assert!(!effective_thinking_enabled( + true, + &["qwen", "mlx-community/Qwen3.6-35B-A3B-4bit"], + None, + )); + } + + #[test] + fn qwen365_does_not_use_qwen36_default() { + assert!(effective_thinking_enabled( + true, + &["mlx-community/Qwen3.65-35B-A3B-4bit"], + None, + )); + } + + #[test] + fn honors_reasoning_none() { + assert!(!effective_thinking_enabled( + true, + &["mlx-community/Qwen3.5-foo"], + Some(&ReasoningConfig { + effort: Some("none".to_owned()), + }), + )); + } + + #[test] + fn honors_empty_reasoning_as_not_explicit() { + assert!(!effective_thinking_enabled( + true, + &["mlx-community/Qwen3.5-foo"], + Some(&ReasoningConfig { + effort: Some(String::new()), + }), + )); + } + + #[test] + fn honors_explicit_reasoning_request() { + assert!(effective_thinking_enabled( + true, + &["mlx-community/Qwen3.6-35B-A3B-4bit"], + Some(&ReasoningConfig { + effort: Some("low".to_owned()), + }), + )); + } + + #[test] + fn engine_default_off_overrides_explicit_request() { + assert!(!effective_thinking_enabled( + false, + &["mlx-community/Qwen3.5-foo"], + Some(&ReasoningConfig { + effort: Some("low".to_owned()), + }), + )); + } +} diff --git a/crates/higgs/src/routes/anthropic.rs b/crates/higgs/src/routes/anthropic.rs index e949b93a..0fecbe43 100644 --- a/crates/higgs/src/routes/anthropic.rs +++ b/crates/higgs/src/routes/anthropic.rs @@ -234,20 +234,25 @@ async fn create_message_non_streaming( let engine_messages = anthropic_messages_to_engine(&req.messages, req.system.as_ref()); let tools = req.tools.as_deref(); + let thinking_enabled = crate::reasoning::effective_thinking_enabled( + engine.enable_thinking(), + &[engine.model_name(), req.model.as_str()], + None, + ); let prompt_tokens = engine - .prepare_chat_prompt(&engine_messages, tools) + .prepare_chat_prompt_with_thinking(&engine_messages, tools, thinking_enabled) .map_err(ServerError::Engine)?; - let thinking_enabled = engine.enable_thinking(); let output = tokio::task::spawn_blocking(move || { - engine.generate( + engine.generate_with_thinking( &prompt_tokens, max_tokens, &sampling, &stop_sequences, false, None, + thinking_enabled, None, None, ) @@ -315,9 +320,14 @@ fn create_message_stream( let engine_messages = anthropic_messages_to_engine(&req.messages, req.system.as_ref()); let tools = req.tools.as_deref(); + let thinking_enabled = crate::reasoning::effective_thinking_enabled( + engine.enable_thinking(), + &[engine.model_name(), req.model.as_str()], + None, + ); let prompt_tokens = engine - .prepare_chat_prompt(&engine_messages, tools) + .prepare_chat_prompt_with_thinking(&engine_messages, tools, thinking_enabled) .map_err(ServerError::Engine)?; let msg_id = format!("msg_{}", uuid::Uuid::new_v4().simple()); @@ -325,13 +335,11 @@ fn create_message_stream( let prompt_token_count = u32::try_from(prompt_tokens.len()) .map_err(|_| ServerError::BadRequest("Token count overflow".to_owned()))?; - let thinking_enabled = engine.enable_thinking(); - // Spawn generation before creating the stream so prefill starts immediately let (tx, mut rx) = tokio::sync::mpsc::channel(32); tokio::task::spawn_blocking(move || { - let result = engine.generate_streaming( + let result = engine.generate_streaming_with_thinking( &prompt_tokens, max_tokens, &sampling, @@ -339,6 +347,7 @@ fn create_message_stream( false, None, &tx, + thinking_enabled, None, None, ); @@ -509,12 +518,19 @@ pub async fn count_tokens( .map_err(ServerError::ModelNotFound)?; match resolved { - ResolvedRoute::Higgs { engine, .. } => { + ResolvedRoute::Higgs { + engine, model_name, .. + } => { let engine_messages = anthropic_messages_to_engine(&req.messages, req.system.as_ref()); let tools = req.tools.as_deref(); + let thinking_enabled = crate::reasoning::effective_thinking_enabled( + engine.enable_thinking(), + &[engine.model_name(), model_name.as_str()], + None, + ); let tokens = engine - .prepare_chat_prompt(&engine_messages, tools) + .prepare_chat_prompt_with_thinking(&engine_messages, tools, thinking_enabled) .map_err(ServerError::Engine)?; let count = u32::try_from(tokens.len()) diff --git a/crates/higgs/src/routes/chat.rs b/crates/higgs/src/routes/chat.rs index d1840939..dd4936fa 100644 --- a/crates/higgs/src/routes/chat.rs +++ b/crates/higgs/src/routes/chat.rs @@ -22,32 +22,12 @@ use crate::{ state::{Engine, SharedState}, types::openai::{ ChatCompletionChoice, ChatCompletionDelta, ChatCompletionMessage, ChatCompletionRequest, - ChatCompletionResponse, ChoiceLogprobs, CompletionUsage, MessageContent, ReasoningConfig, - StopSequence, TokenLogprob, ToolCall, ToolCallFunction, TopLogprob, + ChatCompletionResponse, ChoiceLogprobs, CompletionUsage, MessageContent, StopSequence, + TokenLogprob, ToolCall, ToolCallFunction, TopLogprob, }, }; use higgs_models::SamplingParams; -fn model_defaults_to_non_thinking(model_name: &str) -> bool { - model_name.to_ascii_lowercase().contains("qwen3.6") -} - -fn effective_thinking_enabled( - engine_default: bool, - model_name: &str, - reasoning: Option<&ReasoningConfig>, -) -> bool { - if !engine_default { - return false; - } - - match reasoning.and_then(|r| r.effort.as_deref()) { - Some(effort) if effort.eq_ignore_ascii_case("none") => false, - Some(_) => true, - None => !model_defaults_to_non_thinking(model_name), - } -} - #[allow(clippy::too_many_lines)] pub async fn chat_completions( State(state): State, @@ -288,9 +268,9 @@ async fn chat_completions_non_streaming( let messages = convert_messages(&effective_messages); let tools = req.tools.as_deref(); - let thinking_enabled = effective_thinking_enabled( + let thinking_enabled = crate::reasoning::effective_thinking_enabled( engine.enable_thinking(), - engine.model_name(), + &[engine.model_name(), req.model.as_str()], req.reasoning.as_ref(), ); @@ -348,18 +328,14 @@ async fn chat_completions_non_streaming( // When thinking mode is enabled, the template already opened `` in the prompt, // so the generated text starts inside the think block. Prepend `` so the parser // can find the matching `` and split reasoning from visible content. - let parse_input = if thinking_enabled { - if output_text.contains("") { + let (raw_text, reasoning_content) = if thinking_enabled { + let parse_input = if output_text.contains("") { format!("{output_text}") } else { // Model was length-stopped mid-thinking — close the tag so the // parser can extract reasoning instead of leaking raw ``. format!("{output_text}") - } - } else { - output_text.clone() - }; - let (raw_text, reasoning_content) = if thinking_enabled { + }; let reasoning_result = higgs_engine::reasoning_parser::parse_reasoning(&parse_input); let raw_text = if reasoning_result.reasoning.is_some() { reasoning_result.text @@ -368,7 +344,7 @@ async fn chat_completions_non_streaming( }; (raw_text, reasoning_result.reasoning) } else { - (parse_input, None) + (output_text, None) }; let (content, tool_calls, finish_reason) = if has_tools { @@ -469,9 +445,9 @@ fn chat_completions_stream( }; let messages = convert_messages(&effective_messages); - let thinking_enabled_stream = effective_thinking_enabled( + let thinking_enabled_stream = crate::reasoning::effective_thinking_enabled( engine.enable_thinking(), - engine.model_name(), + &[engine.model_name(), req.model.as_str()], req.reasoning.as_ref(), ); @@ -852,7 +828,6 @@ fn current_unix_timestamp() -> i64 { #[allow(clippy::panic, clippy::unwrap_used)] mod tests { use super::*; - use crate::types::openai::ReasoningConfig; fn simple_message(role: &str, content: Option<&str>) -> ChatCompletionMessage { ChatCompletionMessage { @@ -971,44 +946,4 @@ mod tests { let ts = current_unix_timestamp(); assert!(ts > 1_700_000_000, "timestamp too old: {ts}"); } - - #[test] - fn test_effective_thinking_enabled_defaults_qwen35_on() { - assert!(effective_thinking_enabled( - true, - "mlx-community/Qwen3.5-foo", - None - )); - } - - #[test] - fn test_effective_thinking_enabled_defaults_qwen36_off() { - assert!(!effective_thinking_enabled( - true, - "mlx-community/Qwen3.6-35B-A3B-4bit", - None, - )); - } - - #[test] - fn test_effective_thinking_enabled_honors_reasoning_none() { - assert!(!effective_thinking_enabled( - true, - "mlx-community/Qwen3.5-foo", - Some(&ReasoningConfig { - effort: Some("none".to_owned()), - }), - )); - } - - #[test] - fn test_effective_thinking_enabled_honors_explicit_reasoning_request() { - assert!(effective_thinking_enabled( - true, - "mlx-community/Qwen3.6-35B-A3B-4bit", - Some(&ReasoningConfig { - effort: Some("low".to_owned()), - }), - )); - } } diff --git a/crates/higgs/src/types/openai.rs b/crates/higgs/src/types/openai.rs index 83080fd3..f2ca3a42 100644 --- a/crates/higgs/src/types/openai.rs +++ b/crates/higgs/src/types/openai.rs @@ -5,7 +5,10 @@ use serde::{Deserialize, Serialize}; pub struct ChatCompletionRequest { pub model: String, pub messages: Vec, - #[serde(default)] + /// Maximum number of tokens to generate. + /// + /// Accepts `max_completion_tokens` and `max_output_tokens` aliases. + #[serde(default, alias = "max_completion_tokens", alias = "max_output_tokens")] pub max_tokens: Option, #[serde(default)] pub temperature: Option, @@ -284,7 +287,10 @@ pub struct ToolCallFunctionDelta { pub struct CompletionRequest { pub model: String, pub prompt: String, - #[serde(default)] + /// Maximum number of tokens to generate. + /// + /// Accepts `max_completion_tokens` and `max_output_tokens` aliases. + #[serde(default, alias = "max_completion_tokens", alias = "max_output_tokens")] pub max_tokens: Option, #[serde(default)] pub temperature: Option, @@ -728,6 +734,18 @@ mod tests { assert_eq!(req.max_tokens, Some(0)); } + #[test] + fn test_chat_request_accepts_max_completion_tokens_alias() { + let req = chat_request_with(r#""max_completion_tokens": 100"#); + assert_eq!(req.max_tokens, Some(100)); + } + + #[test] + fn test_chat_request_accepts_max_output_tokens_alias() { + let req = chat_request_with(r#""max_output_tokens": 100"#); + assert_eq!(req.max_tokens, Some(100)); + } + #[test] fn test_chat_request_with_temperature_zero() { let req = chat_request_with(r#""temperature": 0.0"#); @@ -806,6 +824,28 @@ mod tests { assert!(req.stop.is_some()); } + #[test] + fn test_completion_request_accepts_max_completion_tokens_alias() { + let json = r#"{ + "model": "m", + "prompt": "test", + "max_completion_tokens": 100 + }"#; + let req: CompletionRequest = serde_json::from_str(json).unwrap(); + assert_eq!(req.max_tokens, Some(100)); + } + + #[test] + fn test_completion_request_accepts_max_output_tokens_alias() { + let json = r#"{ + "model": "m", + "prompt": "test", + "max_output_tokens": 100 + }"#; + let req: CompletionRequest = serde_json::from_str(json).unwrap(); + assert_eq!(req.max_tokens, Some(100)); + } + #[test] fn test_extra_unknown_fields_silently_ignored() { let json = r#"{ From 65a2dc9e7615acb531730195a3130706f14aed5a Mon Sep 17 00:00:00 2001 From: Jonathan Reyes <150944+panbanda@users.noreply.github.com> Date: Wed, 6 May 2026 10:37:36 -0600 Subject: [PATCH 11/29] fix(release): use explicit higgs-bench package version (#152) --- crates/higgs-bench/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/higgs-bench/Cargo.toml b/crates/higgs-bench/Cargo.toml index 17f022c9..073256fd 100644 --- a/crates/higgs-bench/Cargo.toml +++ b/crates/higgs-bench/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "higgs-bench" -version = { workspace = true } +version = "1.0.0" edition = { workspace = true } license = { workspace = true } rust-version = { workspace = true } From a177ef6ce08508a08fa20ab9be4175d3a3a77476 Mon Sep 17 00:00:00 2001 From: "renovate[bot]" <29139614+renovate[bot]@users.noreply.github.com> Date: Wed, 6 May 2026 10:38:00 -0600 Subject: [PATCH 12/29] fix(deps): update rust crate safetensors to 0.7 (#151) Co-authored-by: renovate[bot] <29139614+renovate[bot]@users.noreply.github.com> --- Cargo.lock | 7 +++++-- crates/higgs-models/Cargo.toml | 2 +- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index fd319b78..a52b65db 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1515,6 +1515,8 @@ dependencies = [ "allocator-api2", "equivalent", "foldhash 0.2.0", + "serde", + "serde_core", ] [[package]] @@ -3483,10 +3485,11 @@ checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" [[package]] name = "safetensors" -version = "0.4.5" +version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44560c11236a6130a46ce36c836a62936dc81ebf8c36a37947423571be0e55b6" +checksum = "675656c1eabb620b921efea4f9199f97fc86e36dd6ffd1fbbe48d0f59a4987f5" dependencies = [ + "hashbrown 0.16.1", "serde", "serde_json", ] diff --git a/crates/higgs-models/Cargo.toml b/crates/higgs-models/Cargo.toml index 47a370aa..3aa2224b 100644 --- a/crates/higgs-models/Cargo.toml +++ b/crates/higgs-models/Cargo.toml @@ -14,7 +14,7 @@ workspace = true half = "2.4" mlx-rs.workspace = true mlx-sys.workspace = true -safetensors = "0.4" +safetensors = "0.7" serde.workspace = true serde_json.workspace = true tokenizers.workspace = true From de774d661ed37b5e00ac3f583b2847b083f5a2c8 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Wed, 6 May 2026 12:47:34 -0600 Subject: [PATCH 13/29] chore: release main (#153) Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- .release-please-manifest.json | 6 +++--- Cargo.lock | 6 +++--- crates/higgs-engine/CHANGELOG.md | 7 +++++++ crates/higgs-engine/Cargo.toml | 2 +- crates/higgs-models/CHANGELOG.md | 19 +++++++++++++++++++ crates/higgs-models/Cargo.toml | 2 +- crates/higgs/CHANGELOG.md | 7 +++++++ crates/higgs/Cargo.toml | 2 +- 8 files changed, 42 insertions(+), 9 deletions(-) diff --git a/.release-please-manifest.json b/.release-please-manifest.json index d71415d6..fbcbdcaa 100644 --- a/.release-please-manifest.json +++ b/.release-please-manifest.json @@ -1,5 +1,5 @@ { - "crates/higgs-models": "1.1.1", - "crates/higgs-engine": "1.1.1", - "crates/higgs": "1.1.1" + "crates/higgs-models": "1.2.0", + "crates/higgs-engine": "1.2.0", + "crates/higgs": "1.2.0" } diff --git a/Cargo.lock b/Cargo.lock index a52b65db..c4be9d9b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1564,7 +1564,7 @@ dependencies = [ [[package]] name = "higgs" -version = "1.1.1" +version = "1.2.0" dependencies = [ "async-stream", "axum", @@ -1624,7 +1624,7 @@ dependencies = [ [[package]] name = "higgs-engine" -version = "1.1.1" +version = "1.2.0" dependencies = [ "half", "higgs-models", @@ -1644,7 +1644,7 @@ dependencies = [ [[package]] name = "higgs-models" -version = "1.1.1" +version = "1.2.0" dependencies = [ "criterion", "half", diff --git a/crates/higgs-engine/CHANGELOG.md b/crates/higgs-engine/CHANGELOG.md index ddb928b2..b98e8404 100644 --- a/crates/higgs-engine/CHANGELOG.md +++ b/crates/higgs-engine/CHANGELOG.md @@ -1,5 +1,12 @@ # Changelog +## [1.2.0](https://github.com/panbanda/higgs/compare/higgs-engine-v1.1.1...higgs-engine-v1.2.0) (2026-05-06) + + +### Features + +* **bonsai-q1:** packed engine scaffold with upstream MLX guard ([#142](https://github.com/panbanda/higgs/issues/142)) ([fe43aab](https://github.com/panbanda/higgs/commit/fe43aabe44104ab285e19e8ff73c724c0875cbe0)) + ## [1.1.1](https://github.com/panbanda/higgs/compare/higgs-engine-v1.1.0...higgs-engine-v1.1.1) (2026-04-28) diff --git a/crates/higgs-engine/Cargo.toml b/crates/higgs-engine/Cargo.toml index c579c63d..9bc38182 100644 --- a/crates/higgs-engine/Cargo.toml +++ b/crates/higgs-engine/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "higgs-engine" -version = "1.1.1" +version = "1.2.0" description = "Inference engine (tokenization, generation loop, prompt caching) for higgs" edition.workspace = true license.workspace = true diff --git a/crates/higgs-models/CHANGELOG.md b/crates/higgs-models/CHANGELOG.md index 04f37de2..6375c01a 100644 --- a/crates/higgs-models/CHANGELOG.md +++ b/crates/higgs-models/CHANGELOG.md @@ -1,5 +1,24 @@ # Changelog +## [1.2.0](https://github.com/panbanda/higgs/compare/higgs-models-v1.1.1...higgs-models-v1.2.0) (2026-05-06) + + +### Features + +* **bonsai-q1:** packed engine scaffold with upstream MLX guard ([#142](https://github.com/panbanda/higgs/issues/142)) ([fe43aab](https://github.com/panbanda/higgs/commit/fe43aabe44104ab285e19e8ff73c724c0875cbe0)) +* **cache:** AnyCache::trim_by dispatcher for spec-decode rollback ([#143](https://github.com/panbanda/higgs/issues/143)) ([229c111](https://github.com/panbanda/higgs/commit/229c1110f2da4c3d93ac79c8f28c2435e0c97aad)) +* **qwen3_next:** mixed-bit Qwen3.5 GDN BA loading fallback ([#148](https://github.com/panbanda/higgs/issues/148)) ([cc18616](https://github.com/panbanda/higgs/commit/cc186160756c9d5741fc18baecbacd3666eeacf3)) + + +### Bug Fixes + +* **deps:** update rust crate safetensors to 0.7 ([#151](https://github.com/panbanda/higgs/issues/151)) ([a177ef6](https://github.com/panbanda/higgs/commit/a177ef6ce08508a08fa20ab9be4175d3a3a77476)) + + +### Performance Improvements + +* **models:** opt-in fused MoE gate+up — 3→2 expert matmuls per layer ([#141](https://github.com/panbanda/higgs/issues/141)) ([60d7cb4](https://github.com/panbanda/higgs/commit/60d7cb48f28ddff5254f8aece898a24a1a475d66)) + ## [1.1.1](https://github.com/panbanda/higgs/compare/higgs-models-v1.1.0...higgs-models-v1.1.1) (2026-04-28) diff --git a/crates/higgs-models/Cargo.toml b/crates/higgs-models/Cargo.toml index 3aa2224b..843ea42f 100644 --- a/crates/higgs-models/Cargo.toml +++ b/crates/higgs-models/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "higgs-models" -version = "1.1.1" +version = "1.2.0" description = "Model architectures and weight loading for higgs (LLaMA, Mistral, Qwen2/3)" edition.workspace = true license.workspace = true diff --git a/crates/higgs/CHANGELOG.md b/crates/higgs/CHANGELOG.md index 09500a37..4a73a74d 100644 --- a/crates/higgs/CHANGELOG.md +++ b/crates/higgs/CHANGELOG.md @@ -1,5 +1,12 @@ # Changelog +## [1.2.0](https://github.com/panbanda/higgs/compare/higgs-v1.1.1...higgs-v1.2.0) (2026-05-06) + + +### Bug Fixes + +* default Qwen3.6 to non-thinking mode ([#150](https://github.com/panbanda/higgs/issues/150)) ([372020e](https://github.com/panbanda/higgs/commit/372020ec5198b47bfb29babbaf60b4250507a76e)) + ## [1.1.1](https://github.com/panbanda/higgs/compare/higgs-v1.1.0...higgs-v1.1.1) (2026-04-28) diff --git a/crates/higgs/Cargo.toml b/crates/higgs/Cargo.toml index 728716fc..2f8c8ea3 100644 --- a/crates/higgs/Cargo.toml +++ b/crates/higgs/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "higgs" -version = "1.1.1" +version = "1.2.0" description = "OpenAI and Anthropic-compatible inference server for Apple Silicon, built on mlx-rs" edition.workspace = true license.workspace = true From ae6f8b02e72868e77059ea5bbe94ed2ae306d130 Mon Sep 17 00:00:00 2001 From: "renovate[bot]" <29139614+renovate[bot]@users.noreply.github.com> Date: Thu, 7 May 2026 07:30:46 -0600 Subject: [PATCH 14/29] chore(deps): update rust crate tower-http to v0.6.10 (#154) Co-authored-by: renovate[bot] <29139614+renovate[bot]@users.noreply.github.com> --- Cargo.lock | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c4be9d9b..395c7f00 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -88,7 +88,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -99,7 +99,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -1060,7 +1060,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -1141,7 +1141,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2002,7 +2002,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2530,7 +2530,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -3391,7 +3391,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -3450,7 +3450,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -3720,7 +3720,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -3884,10 +3884,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.4.2", + "getrandom 0.3.4", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -4299,9 +4299,9 @@ dependencies = [ [[package]] name = "tower-http" -version = "0.6.9" +version = "0.6.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a28f0d049ccfaa566e14e9663d304d8577427b368cb4710a20528690287a738b" +checksum = "68d6fdd9f81c2819c9a8b0e0cd91660e7746a8e6ea2ba7c6b2b057985f6bcb51" dependencies = [ "base64 0.22.1", "bitflags 2.11.1", @@ -4876,7 +4876,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] From f8a5b2cfdd1b3e920841dc0e1dd4b8b72ccc2048 Mon Sep 17 00:00:00 2001 From: "renovate[bot]" <29139614+renovate[bot]@users.noreply.github.com> Date: Fri, 8 May 2026 14:39:13 +0000 Subject: [PATCH 15/29] chore(deps): update rust crate tokio to v1.52.3 --- Cargo.lock | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 395c7f00..70e3f4bc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -88,7 +88,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -99,7 +99,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -1141,7 +1141,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -2002,7 +2002,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -3050,7 +3050,7 @@ dependencies = [ "once_cell", "socket2", "tracing", - "windows-sys 0.60.2", + "windows-sys 0.59.0", ] [[package]] @@ -3391,7 +3391,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -3450,7 +3450,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -3720,7 +3720,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -3887,7 +3887,7 @@ dependencies = [ "getrandom 0.3.4", "once_cell", "rustix", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -4128,9 +4128,9 @@ dependencies = [ [[package]] name = "tokio" -version = "1.52.2" +version = "1.52.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "110a78583f19d5cdb2c5ccf321d1290344e71313c6c37d43520d386027d18386" +checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" dependencies = [ "bytes", "libc", @@ -4876,7 +4876,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] From e78c4ce3cd68cdedca0551f6ab23589c6ac0e9a4 Mon Sep 17 00:00:00 2001 From: "renovate[bot]" <29139614+renovate[bot]@users.noreply.github.com> Date: Mon, 11 May 2026 10:56:44 +0000 Subject: [PATCH 16/29] chore(deps): update rust crate nix to v0.31.3 --- Cargo.lock | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 395c7f00..36ca4812 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -800,7 +800,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e0b1fab2ae45819af2d0731d60f2afe17227ebb1a1538a236da84c93e9a60162" dependencies = [ "dispatch2", - "nix 0.31.2", + "nix 0.31.3", "windows-sys 0.61.2", ] @@ -1586,7 +1586,7 @@ dependencies = [ "hyper", "json5", "mlx-rs", - "nix 0.31.2", + "nix 0.31.3", "ratatui", "regex", "reqwest", @@ -2489,9 +2489,9 @@ dependencies = [ [[package]] name = "nix" -version = "0.31.2" +version = "0.31.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5d6d0705320c1e6ba1d912b5e37cf18071b6c2e9b7fa8215a1e8a7651966f5d3" +checksum = "cf20d2fde8ff38632c426f1165ed7436270b44f199fc55284c38276f9db47c3d" dependencies = [ "bitflags 2.11.1", "cfg-if", From 0780e2f1a2646e5673cdcde1cbb9c3682faa9645 Mon Sep 17 00:00:00 2001 From: "renovate[bot]" <29139614+renovate[bot]@users.noreply.github.com> Date: Mon, 18 May 2026 17:55:31 +0000 Subject: [PATCH 17/29] chore(deps): update rust crate tower-http to v0.6.11 --- Cargo.lock | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 395c7f00..d62c5b2f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -88,7 +88,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -99,7 +99,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -1141,7 +1141,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -2002,7 +2002,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -3050,7 +3050,7 @@ dependencies = [ "once_cell", "socket2", "tracing", - "windows-sys 0.60.2", + "windows-sys 0.59.0", ] [[package]] @@ -3391,7 +3391,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -3450,7 +3450,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -3720,7 +3720,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -3887,7 +3887,7 @@ dependencies = [ "getrandom 0.3.4", "once_cell", "rustix", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -4299,9 +4299,9 @@ dependencies = [ [[package]] name = "tower-http" -version = "0.6.10" +version = "0.6.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68d6fdd9f81c2819c9a8b0e0cd91660e7746a8e6ea2ba7c6b2b057985f6bcb51" +checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" dependencies = [ "base64 0.22.1", "bitflags 2.11.1", @@ -4876,7 +4876,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] From e5d3a5b80efba15a10bf106308a58bed12ba7843 Mon Sep 17 00:00:00 2001 From: "renovate[bot]" <29139614+renovate[bot]@users.noreply.github.com> Date: Wed, 20 May 2026 01:07:23 +0000 Subject: [PATCH 18/29] chore(deps): update rust crate minijinja to v2.20.0 --- Cargo.lock | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 395c7f00..99cc7e74 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2340,9 +2340,9 @@ checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" [[package]] name = "minijinja" -version = "2.19.0" +version = "2.20.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "805bfd7352166bae857ee569628b52bcd85a1cecf7810861ebceb1686b72b75d" +checksum = "2929e494b2280e1e18959bb2e121da03347ae896896fdfaceaab43c88a02803f" dependencies = [ "memo-map", "serde", From 06411a45eb8babb24015372d46d7ec2ca043e5a7 Mon Sep 17 00:00:00 2001 From: "renovate[bot]" <29139614+renovate[bot]@users.noreply.github.com> Date: Wed, 20 May 2026 01:07:31 +0000 Subject: [PATCH 19/29] chore(deps): update rust crate minijinja-contrib to v2.20.0 --- Cargo.lock | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 395c7f00..0b98af8f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2340,9 +2340,9 @@ checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" [[package]] name = "minijinja" -version = "2.19.0" +version = "2.20.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "805bfd7352166bae857ee569628b52bcd85a1cecf7810861ebceb1686b72b75d" +checksum = "2929e494b2280e1e18959bb2e121da03347ae896896fdfaceaab43c88a02803f" dependencies = [ "memo-map", "serde", @@ -2350,9 +2350,9 @@ dependencies = [ [[package]] name = "minijinja-contrib" -version = "2.19.0" +version = "2.20.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "45092d80391870622fcf3bd82f5d2af18f99533ea60debb4bc9db0c76f0e809a" +checksum = "99df5123c54391e2a228014c1dbbd85a3dab08a25e776c810526f2f47542b3de" dependencies = [ "minijinja", "serde", From f9634e70ef061e531a0e36cc8b49c046956e5425 Mon Sep 17 00:00:00 2001 From: "renovate[bot]" <29139614+renovate[bot]@users.noreply.github.com> Date: Thu, 21 May 2026 22:50:41 +0000 Subject: [PATCH 20/29] chore(deps): update rust crate built to v0.8.1 --- Cargo.lock | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 395c7f00..ce567820 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -346,9 +346,9 @@ dependencies = [ [[package]] name = "built" -version = "0.8.0" +version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f4ad8f11f288f48ca24471bbd51ac257aaeaaa07adae295591266b792902ae64" +checksum = "5c0e531d93d39c34eef561e929e8a7f86d77a5af08aac4f6d6e39976c51858e9" dependencies = [ "git2", ] @@ -1420,15 +1420,14 @@ dependencies = [ [[package]] name = "git2" -version = "0.20.4" +version = "0.21.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7b88256088d75a56f8ecfa070513a775dd9107f6530ef14919dac831af9cfe2b" +checksum = "ddddbf932745a6be37109b6112d3ee09696106f848449069d3a57bba937ab82e" dependencies = [ "bitflags 2.11.1", "libc", "libgit2-sys", "log", - "url", ] [[package]] @@ -2157,9 +2156,9 @@ checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" [[package]] name = "libgit2-sys" -version = "0.18.3+1.9.2" +version = "0.18.4+1.9.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c9b3acc4b91781bb0b3386669d325163746af5f6e4f73e6d2d630e09a35f3487" +checksum = "9b26f66f35e1871b22efcf7191564123d2a446ca0538cde63c23adfefa9b15b7" dependencies = [ "cc", "libc", @@ -3884,7 +3883,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.3.4", + "getrandom 0.4.2", "once_cell", "rustix", "windows-sys 0.52.0", From b187fa6b860d9ec3b3c6020837f28df9c56d1070 Mon Sep 17 00:00:00 2001 From: "renovate[bot]" <29139614+renovate[bot]@users.noreply.github.com> Date: Fri, 22 May 2026 13:55:06 +0000 Subject: [PATCH 21/29] chore(deps): update github/codeql-action action to v4.36.0 --- .github/workflows/scorecard.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/scorecard.yml b/.github/workflows/scorecard.yml index b5d809c5..50e9e8d9 100644 --- a/.github/workflows/scorecard.yml +++ b/.github/workflows/scorecard.yml @@ -42,6 +42,6 @@ jobs: retention-days: 5 - name: Upload results to GitHub Security tab - uses: github/codeql-action/upload-sarif@e46ed2cbd01164d986452f91f178727624ae40d7 # v4.35.3 + uses: github/codeql-action/upload-sarif@7211b7c8077ea37d8641b6271f6a365a22a5fbfa # v4.36.0 with: sarif_file: results.sarif From 7397fc85e531f8eee03201e359d93988ace42c4b Mon Sep 17 00:00:00 2001 From: Jonathan Reyes <150944+panbanda@users.noreply.github.com> Date: Thu, 28 May 2026 18:22:10 -0600 Subject: [PATCH 22/29] Add configurable MTP draft depth (#161) * Add configurable MTP draft depth * Address MTP PR feedback * Support dense MTP sidecar checkpoints * Fix dense MTP formatting * Add architecture-neutral speculative drafts * Fix speculative decode lint * Sanitize local benchmark paths * Optimize MTP speculative decoding * Harden audit and benchmark metadata * Address speculative review blockers * Resolve remaining model review threads * Enable CodeRabbit approval workflow --- .coderabbit.yaml | 1 + .gitignore | 4 + benchmarks/bench_all.py | 232 ---- benchmarks/bench_h2h.py | 646 ---------- benchmarks/bench_mlx_tuning.py | 732 ----------- benchmarks/bench_moe_sort.py | 205 --- benchmarks/bench_ppl_tq.py | 462 ------- benchmarks/bench_prefix_cache.py | 300 ----- benchmarks/bench_prefix_cache_turns.py | 243 ---- benchmarks/bench_tq_configs.py | 302 ----- benchmarks/bench_ttft_quick.py | 64 - benchmarks/models.toml | 9 + benchmarks/test_bench_mlx_tuning.py | 121 -- crates/higgs-bench/Cargo.toml | 4 + crates/higgs-bench/src/bin/bench_decode.rs | 9 +- .../higgs-bench/src/bin/bench_speculative.rs | 534 ++++++++ crates/higgs-bench/src/lib.rs | 141 ++- crates/higgs-bench/src/speculative.rs | 234 ++++ crates/higgs-engine/src/mlx_tuning.rs | 113 +- crates/higgs-engine/src/mtp.rs | 683 +++++++++- crates/higgs-engine/src/simple.rs | 467 +++++-- crates/higgs-models/src/deepseek_v2.rs | 15 + crates/higgs-models/src/gemma2.rs | 39 + crates/higgs-models/src/lib.rs | 148 ++- crates/higgs-models/src/llava_qwen2.rs | 10 + crates/higgs-models/src/phi3.rs | 17 + crates/higgs-models/src/qwen3_moe.rs | 15 + crates/higgs-models/src/qwen3_next.rs | 1103 +++++++++++++++-- crates/higgs-models/src/starcoder2.rs | 17 + crates/higgs-models/src/transformer.rs | 21 + crates/higgs/src/error.rs | 23 +- docs/benchmarking.md | 97 +- docs/configuration.md | 1 + 33 files changed, 3378 insertions(+), 3634 deletions(-) delete mode 100644 benchmarks/bench_all.py delete mode 100755 benchmarks/bench_h2h.py delete mode 100644 benchmarks/bench_mlx_tuning.py delete mode 100644 benchmarks/bench_moe_sort.py delete mode 100644 benchmarks/bench_ppl_tq.py delete mode 100644 benchmarks/bench_prefix_cache.py delete mode 100644 benchmarks/bench_prefix_cache_turns.py delete mode 100644 benchmarks/bench_tq_configs.py delete mode 100644 benchmarks/bench_ttft_quick.py delete mode 100644 benchmarks/test_bench_mlx_tuning.py create mode 100644 crates/higgs-bench/src/bin/bench_speculative.rs create mode 100644 crates/higgs-bench/src/speculative.rs diff --git a/.coderabbit.yaml b/.coderabbit.yaml index 6220afc5..86d450ec 100644 --- a/.coderabbit.yaml +++ b/.coderabbit.yaml @@ -1,3 +1,4 @@ reviews: + request_changes_workflow: true commit_status: false fail_commit_status: false diff --git a/.gitignore b/.gitignore index d9752bf5..452f4a85 100644 --- a/.gitignore +++ b/.gitignore @@ -15,5 +15,9 @@ bench/ **/__pycache__/ .claude/ .planning/ +graphify-out/ +GRAPH_REPORT.md +graph.json benchmarks/*.txt benchmarks/*.safetensors +benchmarks/*.py diff --git a/benchmarks/bench_all.py b/benchmarks/bench_all.py deleted file mode 100644 index f6a7e030..00000000 --- a/benchmarks/bench_all.py +++ /dev/null @@ -1,232 +0,0 @@ -#!/usr/bin/env python3 -"""Benchmark all local MLX models with higgs. Measures TTFT, prefill, decode tok/s.""" - -import json, os, signal, subprocess, sys, time, urllib.request - -HIGGS = "./target/release/higgs" -PORT = 8899 -BASE = f"http://localhost:{PORT}" -MAX_TOKENS = 200 - -LMS = os.path.expanduser("~/.cache/lm-studio/models") -MODELS = [ - (f"{LMS}/mlx-community/Qwen3-0.6B-4bit",), - (f"{LMS}/mlx-community/Qwen3-4B-Instruct-2507-4bit",), - (f"{LMS}/mlx-community/Qwen3-8B-4bit",), - (f"{LMS}/mlx-community/Qwen3-14B-4bit",), - (f"{LMS}/mlx-community/Qwen3.5-0.8B-8bit",), - (f"{LMS}/mlx-community/Qwen3.5-2B-MLX-8bit",), - (f"{LMS}/mlx-community/Qwen3.5-4B-MLX-4bit",), - (f"{LMS}/mlx-community/Qwen3.5-9B-MLX-4bit",), - (f"{LMS}/NexVeridian/Qwen3.5-35B-A3B-3bit",), - (f"{LMS}/mlx-community/DeepSeek-Coder-V2-Lite-Instruct-4bit-mlx",), -] - -PROMPTS = { - "short": "Explain what a neural network is in one paragraph.", - "medium": ( - "Write a detailed technical explanation of how transformer architectures work, " - "covering attention mechanisms, positional encoding, layer normalization, " - "feed-forward networks, and the differences between encoder and decoder architectures. " - "Include discussion of multi-head attention, scaled dot-product attention, and how " - "these components work together. Also explain the training process including " - "backpropagation through the attention mechanism. Discuss the key innovations that " - "made transformers superior to RNNs and LSTMs for sequence modeling tasks. Cover the " - "evolution from the original Attention Is All You Need paper through modern variants " - "like GPT, BERT, and their derivatives. Explain how context windows work and the " - "computational complexity of self-attention." - ), - "long": ( - "Write an extremely comprehensive and detailed technical guide covering the following " - "topics in depth. For each topic, provide multiple paragraphs with specific technical " - "details, examples, and explanations:\n\n" - "1. COMPILER DESIGN: Explain lexical analysis, parsing (LL, LR, LALR), abstract syntax " - "trees, semantic analysis, intermediate representations (SSA form, three-address code), " - "optimization passes (constant folding, dead code elimination, loop unrolling, register " - "allocation via graph coloring), and code generation for modern CPU architectures.\n\n" - "2. OPERATING SYSTEMS: Cover process scheduling algorithms (CFS, MLFQ, lottery scheduling), " - "virtual memory management (page tables, TLB, huge pages, NUMA), file systems (ext4, btrfs, " - "ZFS internals), I/O scheduling, interrupt handling, system calls, and the differences " - "between monolithic and microkernel designs.\n\n" - "3. DISTRIBUTED SYSTEMS: Explain consensus protocols (Paxos, Raft, PBFT), distributed hash " - "tables, vector clocks, CRDTs, the CAP theorem and its practical implications, leader " - "election algorithms, distributed transactions (2PC, 3PC, saga pattern), and how systems " - "like Spanner, CockroachDB, and TiKV achieve global consistency.\n\n" - "4. CRYPTOGRAPHY: Cover symmetric encryption (AES internals, modes of operation), asymmetric " - "encryption (RSA, elliptic curves, key exchange), hash functions (SHA-256 internals), " - "digital signatures, zero-knowledge proofs, homomorphic encryption, and post-quantum " - "cryptography approaches.\n\n" - "5. DATABASE INTERNALS: Explain B-tree and LSM-tree storage engines, write-ahead logging, " - "MVCC, query optimization (cost-based vs rule-based), join algorithms (nested loop, hash, " - "sort-merge), buffer pool management, and how modern databases handle concurrent " - "transactions with different isolation levels.\n\n" - "6. NETWORKING: Cover TCP congestion control (Reno, CUBIC, BBR), QUIC protocol design, " - "BGP routing, DNS resolution chain, HTTP/3 internals, TLS certificate chains, NAT " - "traversal techniques, and software-defined networking concepts.\n\n" - "7. MACHINE LEARNING SYSTEMS: Explain backpropagation mathematics, gradient descent " - "variants (SGD, Adam, LAMB), mixed-precision training, data parallelism vs model " - "parallelism vs pipeline parallelism, attention mechanism computation, KV-cache " - "optimization, quantization methods (GPTQ, AWQ, GGML), and inference optimization.\n\n" - "Be thorough and technical throughout." - ), -} - -server_proc = None - -def log(msg): - print(msg, flush=True) - -def start_server(model_path): - global server_proc - env = os.environ.copy() - env["HIGGS_ENABLE_THINKING"] = "0" - server_proc = subprocess.Popen( - [HIGGS, "serve", "--model", model_path, "--port", str(PORT)], - stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, - env=env, preexec_fn=os.setsid, - ) - # Wait for server - for i in range(120): - try: - urllib.request.urlopen(f"{BASE}/v1/models", timeout=2) - return True - except: - time.sleep(1) - if server_proc.poll() is not None: - return False - return False - -def kill_server(): - global server_proc - if server_proc: - try: - os.killpg(os.getpgid(server_proc.pid), signal.SIGTERM) - except: - pass - try: - server_proc.wait(timeout=5) - except: - try: - os.killpg(os.getpgid(server_proc.pid), signal.SIGKILL) - except: - pass - server_proc = None - time.sleep(2) - -def bench(model_name, prompt, label): - body = json.dumps({ - "model": model_name, - "messages": [{"role": "user", "content": prompt}], - "max_tokens": MAX_TOKENS, - "temperature": 0, - "stream": True, - }).encode() - - req = urllib.request.Request( - f"{BASE}/v1/chat/completions", - data=body, - headers={"Content-Type": "application/json"}, - ) - - start = time.perf_counter() - first_token_time = None - prompt_tokens = 0 - completion_tokens = 0 - - try: - with urllib.request.urlopen(req, timeout=300) as resp: - buf = b"" - while True: - chunk = resp.read(1) - if not chunk: - break - buf += chunk - while b"\n" in buf: - line, buf = buf.split(b"\n", 1) - line = line.decode("utf-8", errors="replace").strip() - if not line.startswith("data: "): - continue - data = line[6:] - if data == "[DONE]": - break - try: - obj = json.loads(data) - except: - continue - choices = obj.get("choices", []) - if choices: - delta = choices[0].get("delta", {}) - content = delta.get("content", "") - if content and first_token_time is None: - first_token_time = time.perf_counter() - usage = obj.get("usage") - if usage: - prompt_tokens = usage.get("prompt_tokens", 0) - completion_tokens = usage.get("completion_tokens", 0) - except Exception as e: - log(f" {label:8s}: ERROR - {e}") - return - - end = time.perf_counter() - ttft_ms = (first_token_time - start) * 1000 if first_token_time else -1 - decode_s = end - first_token_time if first_token_time else 0.001 - decode_tps = max(completion_tokens - 1, 0) / decode_s if decode_s > 0 else 0 - prefill_tps = prompt_tokens / (ttft_ms / 1000) if ttft_ms > 0 and prompt_tokens > 0 else 0 - - log( - f" {label:8s}: prompt={prompt_tokens:>5d}tok " - f"TTFT={ttft_ms:>7.0f}ms prefill={prefill_tps:>8.1f}tok/s " - f"decode={decode_tps:>6.1f}tok/s gen={completion_tokens:>3d}tok " - f"total={end - start:>5.1f}s" - ) - -if __name__ == "__main__": - results_file = f"bench_results_{time.strftime('%Y%m%d_%H%M%S')}.txt" - - # Tee output to file - class Tee: - def __init__(self, f): self.file = f; self.stdout = sys.stdout - def write(self, s): self.stdout.write(s); self.file.write(s) - def flush(self): self.stdout.flush(); self.file.flush() - - rf = open(results_file, "w") - sys.stdout = Tee(rf) - - log("=" * 80) - log(f"HIGGS BENCHMARK - {time.strftime('%Y-%m-%d %H:%M:%S')}") - log(f"Binary: {HIGGS} Max tokens: {MAX_TOKENS} Thinking: disabled") - log("=" * 80) - log("") - - try: - for (model_path,) in MODELS: - model_name = os.path.basename(model_path) - if not os.path.isdir(model_path): - log(f"SKIP {model_name} (not found)") - log("") - continue - - kill_server() - log(f"--- {model_name} ---") - - if not start_server(model_path): - log(f" FAILED to start server") - log("") - kill_server() - continue - - log(f" Server ready (PID {server_proc.pid})") - - for label, prompt in PROMPTS.items(): - bench(model_name, prompt, label) - - kill_server() - log("") - except KeyboardInterrupt: - log("\nInterrupted by user") - finally: - kill_server() - log("=" * 80) - log(f"DONE - results in {results_file}") - log("=" * 80) - rf.close() diff --git a/benchmarks/bench_h2h.py b/benchmarks/bench_h2h.py deleted file mode 100755 index f795e114..00000000 --- a/benchmarks/bench_h2h.py +++ /dev/null @@ -1,646 +0,0 @@ -#!/usr/bin/env python3 -"""Head-to-head benchmark: Higgs vs oMLX on Apple Silicon. - -Compares TTFT, decode tok/s, and multi-turn behavior on the same models -using the same OpenAI-compatible streaming API. - -Usage: - python3 bench_h2h.py # all models - python3 bench_h2h.py --models 35B # just 35B - python3 bench_h2h.py --skip-multiturn # single-turn only - python3 bench_h2h.py --turns 10 # more turns -""" - -import argparse -import json -import os -import signal -import socket -import subprocess -import sys -import time -import urllib.error -import urllib.request - -# --------------------------------------------------------------------------- -# Config -# --------------------------------------------------------------------------- - -HIGGS_BIN = "./target/release/higgs" -OMLX_CLI = "/Applications/oMLX.app/Contents/MacOS/omlx-cli" -LMS = os.path.expanduser("~/.cache/lm-studio/models") - -HIGGS_PORT = 8899 -OMLX_PORT = 8000 # avoid 8080 in case oMLX GUI is running - -MAX_TOKENS = 100 -COOLDOWN = 5 # seconds between server swaps -WARMUP_TOKENS = 10 - -MODELS = { - "35B": { - "path": f"{LMS}/NexVeridian/Qwen3.5-35B-A3B-3bit", - "label": "Qwen3.5-35B-A3B-3bit (MoE)", - "higgs_name": "Qwen3.5-35B-A3B-3bit", - }, - "27B": { - "path": f"{LMS}/mlx-community/Qwen3.5-27B-Claude-4.6-Opus-Distilled-MLX-4bit", - "label": "Qwen3.5-27B-4bit (Dense)", - "higgs_name": "Qwen3.5-27B-Claude-4.6-Opus-Distilled-MLX-4bit", - }, - "DSV2": { - "path": f"{LMS}/mlx-community/DeepSeek-Coder-V2-Lite-Instruct-4bit-mlx", - "label": "DeepSeek-V2-Lite-4bit (MoE)", - "higgs_name": "DeepSeek-Coder-V2-Lite-Instruct-4bit-mlx", - }, -} - -PROMPTS = { - "short": "What is 2+2? Answer in one word.", - "medium": ( - "Write a detailed technical explanation of how transformer architectures work, " - "covering attention mechanisms, positional encoding, layer normalization, " - "feed-forward networks, and the differences between encoder and decoder architectures. " - "Include discussion of multi-head attention, scaled dot-product attention, and how " - "these components work together." - ), - "long": ( - "Write an extremely comprehensive and detailed technical guide covering the following " - "topics in depth:\n\n" - "1. COMPILER DESIGN: Explain lexical analysis, parsing (LL, LR, LALR), abstract syntax " - "trees, semantic analysis, intermediate representations (SSA form, three-address code), " - "optimization passes (constant folding, dead code elimination, loop unrolling, register " - "allocation via graph coloring), and code generation for modern CPU architectures.\n\n" - "2. OPERATING SYSTEMS: Cover process scheduling algorithms (CFS, MLFQ, lottery scheduling), " - "virtual memory management (page tables, TLB, huge pages, NUMA), file systems (ext4, btrfs, " - "ZFS internals), I/O scheduling, interrupt handling, system calls.\n\n" - "3. DISTRIBUTED SYSTEMS: Explain consensus protocols (Paxos, Raft, PBFT), distributed hash " - "tables, vector clocks, CRDTs, the CAP theorem, leader election algorithms, distributed " - "transactions (2PC, 3PC, saga pattern).\n\n" - "4. CRYPTOGRAPHY: Cover symmetric encryption (AES internals, modes of operation), asymmetric " - "encryption (RSA, elliptic curves, key exchange), hash functions (SHA-256 internals), " - "digital signatures, zero-knowledge proofs.\n\n" - "5. DATABASE INTERNALS: Explain B-tree and LSM-tree storage engines, write-ahead logging, " - "MVCC, query optimization, join algorithms, buffer pool management.\n\n" - "Be thorough and technical throughout." - ), -} - -SYSTEM_PROMPT = ( - "You are a highly skilled software architect with deep expertise in distributed systems, " - "database design, and cloud-native applications. You provide thorough, well-reasoned " - "technical advice with step-by-step reasoning and concrete examples." -) - -TURN_QUESTIONS = [ - "Explain the CAP theorem and its practical implications for system design.", - "How would you design a rate limiter for a distributed API gateway?", - "Compare event sourcing with traditional CRUD. When would you pick each?", - "What are the key differences between Raft and Paxos consensus protocols?", - "Design a notification system that handles 1M users with real-time delivery.", - "How does MVCC work in PostgreSQL? Walk me through a concurrent update scenario.", - "What strategies would you use to migrate a monolith to microservices safely?", -] - -# --------------------------------------------------------------------------- -# Server management -# --------------------------------------------------------------------------- - -server_proc = None - - -def kill_by_port(port): - """Kill all processes listening on the given port.""" - try: - out = subprocess.check_output( - ["lsof", "-ti", f":{port}"], text=True - ).strip() - if out: - pids = out.splitlines() - for pid in pids: - try: - os.kill(int(pid), signal.SIGKILL) - except (ProcessLookupError, PermissionError, ValueError): - pass - except subprocess.CalledProcessError: - pass # no process on port - - -def port_is_open(port): - """Return True when a local TCP port accepts connections.""" - try: - with socket.create_connection(("127.0.0.1", port), timeout=1): - return True - except OSError: - return False - - -def wait_for_port_closed(port, timeout=10): - """Poll until a local TCP port stops accepting connections.""" - deadline = time.time() + timeout - while time.time() < deadline: - if not port_is_open(port): - return True - time.sleep(0.2) - return not port_is_open(port) - - -def kill_server(): - global server_proc - if server_proc: - try: - os.killpg(os.getpgid(server_proc.pid), signal.SIGTERM) - except (ProcessLookupError, PermissionError): - pass - try: - server_proc.wait(timeout=10) - except subprocess.TimeoutExpired: - try: - os.killpg(os.getpgid(server_proc.pid), signal.SIGKILL) - except (ProcessLookupError, PermissionError): - pass - server_proc = None - # Kill any stray servers on our ports (catches leaked processes) - for port in (HIGGS_PORT, OMLX_PORT): - kill_by_port(port) - wait_for_port_closed(port, timeout=COOLDOWN) - - -def wait_for_server(port, api_key=None, timeout=180): - """Wait until /v1/models responds. Returns first model name.""" - url = f"http://127.0.0.1:{port}/v1/models" - deadline = time.time() + timeout - while time.time() < deadline: - try: - req = urllib.request.Request(url) - if api_key: - req.add_header("Authorization", f"Bearer {api_key}") - with urllib.request.urlopen(req, timeout=5) as resp: - data = json.loads(resp.read()) - models = data.get("data", []) - if models: - return models[0]["id"] - except (urllib.error.URLError, ConnectionRefusedError, OSError, KeyError, IndexError): - pass - time.sleep(1) - return None - - -def start_higgs(model_path, expected_name=None): - """Start Higgs server, return model name or None.""" - global server_proc - # Kill anything still on the port before starting - kill_by_port(HIGGS_PORT) - wait_for_port_closed(HIGGS_PORT) - env = {**os.environ, "HIGGS_ENABLE_THINKING": "0"} - server_proc = subprocess.Popen( - [HIGGS_BIN, "serve", "--model", model_path, "--port", str(HIGGS_PORT)], - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - env=env, - preexec_fn=os.setsid, - ) - model_name = wait_for_server(HIGGS_PORT) - if model_name and expected_name and expected_name not in model_name: - log(f" WARNING: Expected model '{expected_name}' but server reports '{model_name}'") - kill_server() - return None - return model_name - - -def start_omlx(model_parent_dir, expected_name=None): - """Start oMLX CLI server, return model name or None. - - oMLX discovers models one level deep from --model-dir, so pass the - immediate parent of the model directory (e.g. NexVeridian/ or mlx-community/). - """ - global server_proc - # Kill anything still on the port before starting - kill_by_port(OMLX_PORT) - wait_for_port_closed(OMLX_PORT) - server_proc = subprocess.Popen( - [ - OMLX_CLI, "serve", - "--model-dir", model_parent_dir, - "--port", str(OMLX_PORT), - "--no-cache", - "--max-num-seqs", "1", - "--log-level", "warning", - ], - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - preexec_fn=os.setsid, - ) - model_name = wait_for_server(OMLX_PORT, api_key="omlx") - if model_name and expected_name and expected_name not in model_name: - log(f" WARNING: Expected model '{expected_name}' but oMLX reports '{model_name}'") - kill_server() - return None - return model_name - - -def get_rss_mb(): - """Get RSS of server process tree in MB (includes child processes).""" - if not server_proc: - return 0 - try: - # Get RSS of entire process group (parent + children) - pgid = os.getpgid(server_proc.pid) - out = subprocess.check_output( - ["ps", "-o", "rss=", "-g", str(pgid)], text=True - ).strip() - total_kb = sum(int(line.strip()) for line in out.splitlines() if line.strip()) - return total_kb / 1024 # KB -> MB - except (subprocess.CalledProcessError, ValueError, ProcessLookupError): - return 0 - - -# --------------------------------------------------------------------------- -# Benchmark helpers -# --------------------------------------------------------------------------- - - -def stream_chat(port, messages, max_tokens=MAX_TOKENS, model_name="test", - api_key=None, temperature=0.0): - """Send streaming chat completion, measure TTFT + decode.""" - payload = json.dumps({ - "model": model_name, - "messages": messages, - "max_tokens": max_tokens, - "temperature": temperature, - "stream": True, - }).encode() - - req = urllib.request.Request( - f"http://127.0.0.1:{port}/v1/chat/completions", - data=payload, - headers={"Content-Type": "application/json"}, - ) - if api_key: - req.add_header("Authorization", f"Bearer {api_key}") - - t0 = time.perf_counter() - first_token_time = None - tokens = [] - prompt_tokens = 0 - completion_tokens = 0 - - try: - with urllib.request.urlopen(req, timeout=600) as resp: - buf = b"" - while True: - chunk = resp.read(1) - if not chunk: - break - buf += chunk - while b"\n" in buf: - line, buf = buf.split(b"\n", 1) - line = line.decode("utf-8", errors="replace").strip() - if not line.startswith("data: "): - continue - data_str = line[6:] - if data_str == "[DONE]": - break - try: - obj = json.loads(data_str) - except json.JSONDecodeError: - continue - choices = obj.get("choices", []) - if choices: - delta = choices[0].get("delta", {}) - content = delta.get("content", "") - if content and first_token_time is None: - first_token_time = time.perf_counter() - if content: - tokens.append(content) - usage = obj.get("usage") - if usage: - prompt_tokens = usage.get("prompt_tokens", prompt_tokens) - completion_tokens = usage.get("completion_tokens", completion_tokens) - except Exception as e: - return {"error": str(e)} - - end = time.perf_counter() - if first_token_time is None: - first_token_time = end - - ttft_s = first_token_time - t0 - decode_s = end - first_token_time - completion_tokens_estimated = False - if completion_tokens == 0: - output_chars = sum(len(chunk) for chunk in tokens) - completion_tokens = max(1, int(output_chars / 3.5)) if output_chars else 0 - completion_tokens_estimated = output_chars > 0 - n_completion = completion_tokens - decode_tps = max(n_completion - 1, 0) / decode_s if decode_s > 0.001 else 0 - - # Estimate prompt tokens if server didn't report (oMLX SSE doesn't) - if not prompt_tokens: - prompt_chars = sum(len(m.get("content", "")) for m in messages) - prompt_tokens = max(1, int(prompt_chars / 3.5)) # rough estimate - - return { - "ttft_ms": ttft_s * 1000, - "decode_tps": decode_tps, - "prompt_tokens": prompt_tokens, - "completion_tokens": n_completion, - "completion_tokens_estimated": completion_tokens_estimated, - "total_ms": (end - t0) * 1000, - "output": "".join(tokens), - "rss_mb": get_rss_mb(), - } - - -def fmt_result(r): - if "error" in r: - return f"ERROR: {r['error']}" - return ( - f"TTFT={r['ttft_ms']:>7.0f}ms " - f"decode={r['decode_tps']:>5.1f}tok/s " - f"prompt={r['prompt_tokens']:>4d}tok " - f"gen={r['completion_tokens']:>3d}tok" - + ("~" if r.get("completion_tokens_estimated") else "") - + " " - f"total={r['total_ms']/1000:>5.1f}s " - f"RSS={r['rss_mb']:>5.0f}MB" - ) - - -# --------------------------------------------------------------------------- -# Test phases -# --------------------------------------------------------------------------- - - -def run_single_turn(port, model_name, api_key=None): - """Run short/medium/long single-turn prompts.""" - results = {} - for label, prompt in PROMPTS.items(): - messages = [{"role": "user", "content": prompt}] - r = stream_chat(port, messages, model_name=model_name, api_key=api_key) - results[label] = r - log(f" {label:8s}: {fmt_result(r)}") - return results - - -def run_multi_turn(port, model_name, num_turns, api_key=None): - """Run multi-turn conversation, measuring TTFT degradation.""" - results = [] - messages = [{"role": "system", "content": SYSTEM_PROMPT}] - - for i in range(min(num_turns, len(TURN_QUESTIONS))): - messages.append({"role": "user", "content": TURN_QUESTIONS[i]}) - r = stream_chat(port, messages, max_tokens=80, model_name=model_name, - api_key=api_key) - results.append(r) - - if "error" in r: - log(f" turn {i+1}: {fmt_result(r)}") - break - - # Add assistant response to conversation history - messages.append({"role": "assistant", "content": r["output"]}) - - ctx_tokens = r["prompt_tokens"] - log(f" turn {i+1}: {fmt_result(r)} ctx~{ctx_tokens}") - - return results - - -# --------------------------------------------------------------------------- -# Main -# --------------------------------------------------------------------------- - -output_lines = [] - - -def log(msg): - print(msg, flush=True) - output_lines.append(msg) - - -def run_for_backend(backend, model_key, model_info, num_turns, skip_multiturn): - """Run all benchmarks for one backend + model combo.""" - model_path = model_info["path"] - - log(f"\n --- {backend} ---") - - expected = model_info.get("higgs_name", os.path.basename(model_path)) - - try: - if backend == "higgs": - port = HIGGS_PORT - api_key = None - log(f" Starting Higgs on :{port} ...") - model_name = start_higgs(model_path, expected_name=expected) - else: - port = OMLX_PORT - api_key = "omlx" - log(f" Starting oMLX on :{port} (--no-cache) ...") - model_name = start_omlx(os.path.dirname(model_path), expected_name=expected) - - if not model_name: - log(" FAILED to start server") - return None - - log(f" Server ready: model={model_name} RSS={get_rss_mb():.0f}MB") - - # Use the right model name for requests - if backend == "omlx": - # oMLX discovers models by directory basename - model_name = os.path.basename(model_path) - - # Warmup - log(" Warmup...") - warmup = stream_chat( - port, - [{"role": "user", "content": "Say hi."}], - max_tokens=WARMUP_TOKENS, - model_name=model_name, - api_key=api_key, - ) - if "error" in warmup: - log(f" Warmup failed: {warmup['error']}") - return None - log(f" Warmup done. RSS={get_rss_mb():.0f}MB") - - all_results = {"backend": backend, "model": model_key} - - # Phase 1: Single-turn - log(f"\n [Single-turn]") - all_results["single_turn"] = run_single_turn(port, model_name, api_key) - - # Phase 2: Multi-turn - if not skip_multiturn: - log(f"\n [Multi-turn, {num_turns} turns]") - all_results["multi_turn"] = run_multi_turn( - port, model_name, num_turns, api_key - ) - - return all_results - finally: - kill_server() - - -def print_comparison(higgs_results, omlx_results, model_label): - """Print side-by-side comparison table.""" - log(f"\n{'='*80}") - log(f"COMPARISON: {model_label}") - log(f"{'='*80}") - - # Single-turn - log(f"\n {'':12s} {'TTFT (ms)':>20s} {'Decode (tok/s)':>20s}") - log(f" {'Prompt':12s} {'Higgs':>9s} {'oMLX':>9s} {'Higgs':>9s} {'oMLX':>9s} {'TTFT':>7s} {'Decode':>7s}") - log(f" {'-'*74}") - - h_st = higgs_results.get("single_turn", {}) - o_st = omlx_results.get("single_turn", {}) - - for label in ("short", "medium", "long"): - h = h_st.get(label, {}) - o = o_st.get(label, {}) - - if "error" in h or "error" in o: - log(f" {label:12s} (error in one or both backends)") - continue - - h_ttft = h.get("ttft_ms", 0) - o_ttft = o.get("ttft_ms", 0) - h_dec = h.get("decode_tps", 0) - o_dec = o.get("decode_tps", 0) - - ttft_ratio = f"{o_ttft/h_ttft:.2f}x" if h_ttft > 0 else "—" - dec_ratio = f"{h_dec/o_dec:.2f}x" if o_dec > 0 else "—" - - log( - f" {label:12s} {h_ttft:>8.0f} {o_ttft:>8.0f} " - f"{h_dec:>8.1f} {o_dec:>8.1f} " - f"{ttft_ratio:>7s} {dec_ratio:>7s}" - ) - - # Multi-turn - h_mt = higgs_results.get("multi_turn", []) - o_mt = omlx_results.get("multi_turn", []) - - if h_mt and o_mt: - log(f"\n Multi-turn TTFT progression:") - log(f" {'Turn':>5s} {'Higgs TTFT':>11s} {'oMLX TTFT':>11s} {'Higgs dec':>10s} {'oMLX dec':>10s}") - log(f" {'-'*55}") - - for i in range(min(len(h_mt), len(o_mt))): - h = h_mt[i] - o = o_mt[i] - if "error" in h or "error" in o: - log(f" {i+1:>5d} (error)") - break - log( - f" {i+1:>5d} {h['ttft_ms']:>9.0f}ms {o['ttft_ms']:>9.0f}ms " - f"{h['decode_tps']:>8.1f}/s {o['decode_tps']:>8.1f}/s" - ) - - -def main(): - parser = argparse.ArgumentParser(description="Higgs vs oMLX head-to-head") - parser.add_argument("--models", nargs="+", default=list(MODELS.keys()), - help="Model keys to test (35B, 27B, DSV2)") - parser.add_argument("--turns", type=int, default=5, - help="Number of multi-turn conversation turns") - parser.add_argument("--skip-multiturn", action="store_true", - help="Skip multi-turn tests") - parser.add_argument("--higgs-only", action="store_true") - parser.add_argument("--omlx-only", action="store_true") - args = parser.parse_args() - - # Validate - if not os.path.isfile(HIGGS_BIN) and not args.omlx_only: - log(f"Higgs binary not found: {HIGGS_BIN}") - sys.exit(1) - if not os.path.isfile(OMLX_CLI) and not args.higgs_only: - log(f"oMLX CLI not found: {OMLX_CLI}") - sys.exit(1) - - selected = {k: v for k, v in MODELS.items() if k in args.models} - available = {k: v for k, v in selected.items() if os.path.isdir(v["path"])} - - if not available: - log("No models found on disk!") - log(f"Checked: {[v['path'] for v in selected.values()]}") - sys.exit(1) - - log("=" * 80) - log(f"HEAD-TO-HEAD: Higgs vs oMLX — {time.strftime('%Y-%m-%d %H:%M:%S')}") - log(f"Max tokens: {MAX_TOKENS} Turns: {args.turns} Cooldown: {COOLDOWN}s") - log(f"Models: {', '.join(available.keys())}") - log(f"oMLX: --no-cache (paged SSD cache disabled for fair comparison)") - log("=" * 80) - - all_comparisons = [] - - for model_key, model_info in available.items(): - log(f"\n{'#'*80}") - log(f"# MODEL: {model_info['label']}") - log(f"# Path: {model_info['path']}") - log(f"{'#'*80}") - - higgs_r = None - omlx_r = None - - # Run Higgs first (it's our project, gives oMLX more cooldown) - if not args.omlx_only: - try: - higgs_r = run_for_backend( - "higgs", model_key, model_info, args.turns, args.skip_multiturn - ) - except Exception as e: - log(f" Higgs error: {e}") - kill_server() - - # Run oMLX - if not args.higgs_only: - try: - omlx_r = run_for_backend( - "omlx", model_key, model_info, args.turns, args.skip_multiturn - ) - except Exception as e: - log(f" oMLX error: {e}") - kill_server() - - if higgs_r and omlx_r: - all_comparisons.append((model_key, model_info["label"], higgs_r, omlx_r)) - print_comparison(higgs_r, omlx_r, model_info["label"]) - - # Final summary - if all_comparisons: - log(f"\n{'='*80}") - log("FINAL SUMMARY") - log(f"{'='*80}") - log(f"\n {'Model':<40s} {'TTFT Higgs':>10s} {'TTFT oMLX':>10s} {'Dec Higgs':>10s} {'Dec oMLX':>10s}") - log(f" {'-'*84}") - for model_key, label, h, o in all_comparisons: - h_med = h.get("single_turn", {}).get("medium", {}) - o_med = o.get("single_turn", {}).get("medium", {}) - if "error" not in h_med and "error" not in o_med: - log( - f" {label:<40s} " - f"{h_med.get('ttft_ms',0):>8.0f}ms " - f"{o_med.get('ttft_ms',0):>8.0f}ms " - f"{h_med.get('decode_tps',0):>8.1f}/s " - f"{o_med.get('decode_tps',0):>8.1f}/s" - ) - - kill_server() - - # Save results - ts = time.strftime("%Y%m%d_%H%M%S") - outfile = f"bench_h2h_{ts}.txt" - with open(outfile, "w") as f: - f.write("\n".join(output_lines) + "\n") - log(f"\nResults saved to {outfile}") - - -if __name__ == "__main__": - try: - main() - except KeyboardInterrupt: - log("\nInterrupted") - kill_server() - sys.exit(1) diff --git a/benchmarks/bench_mlx_tuning.py b/benchmarks/bench_mlx_tuning.py deleted file mode 100644 index 9ff2687c..00000000 --- a/benchmarks/bench_mlx_tuning.py +++ /dev/null @@ -1,732 +0,0 @@ -#!/usr/bin/env python3 -"""Benchmark MLX tuning iterations for Higgs. - -This harness is designed to capture the tradeoffs that matter on Apple Silicon: -1. TTFT across short/medium/long prompts -2. Decode throughput -3. Long-context retrieval accuracy -4. Structured-output correctness -5. Prefix-cache speedup on multi-turn conversations - -It runs five optimization iterations and produces both raw metrics and a -composite score so the best profile is obvious. - -Usage: - python3 benchmarks/bench_mlx_tuning.py -""" - -from __future__ import annotations - -import argparse -import json -import os -import signal -import statistics -import subprocess -import sys -import time -import urllib.request -from dataclasses import dataclass -from typing import Any - -HIGGS_BIN = os.environ.get("HIGGS_BIN", "./target/release/higgs") -PORT = 8099 -BASE = f"http://127.0.0.1:{PORT}" -MAX_TOKENS = 96 -CACHE_SPEEDUP_CAP = 32.0 - -SHORT_PROMPT = "What is 17 + 25? Reply with digits only." -MEDIUM_PROMPT = ( - "Explain how KV cache reuse affects time-to-first-token and decode throughput " - "for autoregressive transformer inference on Apple Silicon. Keep the answer " - "technical and concise." -) -LONG_PROMPT = ( - "Write a technical note about optimizing LLM inference on unified-memory Apple " - "Silicon systems. Cover TTFT, decode throughput, prompt length sensitivity, " - "prefix cache reuse, chunked prefill, speculative decode, and quantized KV " - "caches. Include concrete engineering tradeoffs and failure modes.\n\n" - + "\n".join( - f"Section {idx}: Repeated background detail about scheduler fairness, " - f"prompt staging, and kernel launch overhead on MLX devices." - for idx in range(1, 91) - ) -) - -LONG_CONTEXT_FILLER = "\n".join( - f"Paragraph {idx}: Cape Town cluster notes about unified memory pressure, " - f"prefill staging, and context reuse across serving workloads." - for idx in range(1, 181) -) -LONG_CONTEXT_NEEDLE = "CAPE-TOWN-7419" - -PREFIX_CACHE_DOC = "\n".join( - f"Policy {idx}: Route latency data through the prefill pipeline, keep the " - f"region failover target as cape town, and retain prefix blocks for reuse." - for idx in range(1, 161) -) - -QA_CASES = [ - { - "prompt": "Reply with digits only. What is 37 * 19?", - "expected": "703", - }, - { - "prompt": "Reply with one lowercase word only. Which word appears twice in 'alpha beta gamma beta delta'?", - "expected": "beta", - }, - { - "prompt": "Reply with digits only. How many vowels are in the word instrumentation?", - "expected": "6", - }, - { - "prompt": "Reply with lowercase letters only. Reverse the string stressed.", - "expected": "desserts", - }, - { - "prompt": "Reply with comma-separated digits only. Sort 7,1,9,1 ascending.", - "expected": "1,1,7,9", - }, -] - -STRUCTURED_OUTPUT_SCHEMA = { - "type": "json_schema", - "json_schema": { - "name": "mlx_report", - "strict": True, - "schema": { - "type": "object", - "properties": { - "model_family": {"type": "string"}, - "iteration": {"type": "integer"}, - "prefill_focus": {"type": "boolean"}, - "kv_bits": {"type": "integer"}, - "primary_goal": {"type": "string"}, - }, - "required": [ - "model_family", - "iteration", - "prefill_focus", - "kv_bits", - "primary_goal", - ], - "additionalProperties": False, - }, - }, -} - - -@dataclass -class Iteration: - slug: str - label: str - env: dict[str, str] - args: list[str] - notes: str - - -ITERATIONS = [ - Iteration( - slug="baseline", - label="1. Baseline", - env={"HIGGS_MLX_PROFILE": "baseline"}, - args=[], - notes="Current conservative defaults", - ), - Iteration( - slug="latency", - label="2. Latency Profile", - env={"HIGGS_MLX_PROFILE": "latency"}, - args=[], - notes="Favor single-pass prefill and speculative decode", - ), - Iteration( - slug="balanced", - label="3. Balanced Profile", - env={"HIGGS_MLX_PROFILE": "balanced"}, - args=[], - notes="Model-aware chunking plus larger paged KV budget", - ), - Iteration( - slug="throughput", - label="4. Throughput Profile", - env={"HIGGS_MLX_PROFILE": "throughput"}, - args=[], - notes="Bigger decode-oriented chunks and paged KV budget", - ), - Iteration( - slug="throughput_turboquant", - label="5. Throughput + Safe TurboQuant", - env={"HIGGS_MLX_PROFILE": "throughput"}, - args=[ - "--kv-cache", - "turboquant", - "--kv-bits", - "3", - "--kv-key-bits", - "2", - "--kv-value-bits", - "3", - "--kv-adaptive-dense-layers", - "8", - ], - notes="Adds quality-preserving KV quantization after MLX runtime tuning", - ), -] - - -server_proc: subprocess.Popen[bytes] | None = None - - -def log(msg: str) -> None: - print(msg, flush=True) - - -def api_request( - endpoint: str, - body: dict[str, Any], - timeout: int = 300, -) -> tuple[dict[str, Any], float]: - data = json.dumps(body).encode() - req = urllib.request.Request( - f"{BASE}/v1/{endpoint}", - data=data, - headers={"Content-Type": "application/json"}, - ) - started = time.perf_counter() - with urllib.request.urlopen(req, timeout=timeout) as resp: - payload = json.loads(resp.read()) - return payload, time.perf_counter() - started - - -def stream_chat( - model: str, - messages: list[dict[str, str]], - max_tokens: int = MAX_TOKENS, - response_format: dict[str, Any] | None = None, - temperature: float = 0.0, - timeout: int = 600, -) -> dict[str, Any]: - body: dict[str, Any] = { - "model": model, - "messages": messages, - "max_tokens": max_tokens, - "temperature": temperature, - "stream": True, - } - if response_format is not None: - body["response_format"] = response_format - - req = urllib.request.Request( - f"{BASE}/v1/chat/completions", - data=json.dumps(body).encode(), - headers={"Content-Type": "application/json"}, - ) - - started = time.perf_counter() - first_token_time = None - output_chunks: list[str] = [] - prompt_tokens = 0 - completion_tokens = 0 - - with urllib.request.urlopen(req, timeout=timeout) as resp: - while True: - line = resp.readline() - if not line: - break - line = line.decode("utf-8", errors="replace").strip() - if not line.startswith("data: "): - continue - payload = line[6:] - if payload == "[DONE]": - break - try: - obj = json.loads(payload) - except json.JSONDecodeError: - continue - choice = obj.get("choices", [{}])[0] - delta = choice.get("delta", {}) - content = delta.get("content", "") - if content and first_token_time is None: - first_token_time = time.perf_counter() - if content: - output_chunks.append(content) - usage = obj.get("usage") or {} - prompt_tokens = usage.get("prompt_tokens", prompt_tokens) - completion_tokens = usage.get("completion_tokens", completion_tokens) - - ended = time.perf_counter() - if first_token_time is None: - first_token_time = ended - - if completion_tokens == 0: - output_chars = sum(len(chunk) for chunk in output_chunks) - completion_tokens = max(1, int(output_chars / 3.5)) if output_chars else 0 - - ttft_s = first_token_time - started - decode_s = max(ended - first_token_time, 0.001) - decode_tps = max(completion_tokens - 1, 0) / decode_s - - return { - "ttft_ms": ttft_s * 1000, - "decode_tps": decode_tps, - "total_ms": (ended - started) * 1000, - "prompt_tokens": prompt_tokens, - "completion_tokens": completion_tokens, - "output": "".join(output_chunks), - "rss_mb": get_rss_mb(), - } - - -def chat( - model: str, - messages: list[dict[str, str]], - max_tokens: int = MAX_TOKENS, - response_format: dict[str, Any] | None = None, - temperature: float = 0.0, - timeout: int = 300, -) -> dict[str, Any]: - body: dict[str, Any] = { - "model": model, - "messages": messages, - "max_tokens": max_tokens, - "temperature": temperature, - } - if response_format is not None: - body["response_format"] = response_format - payload, elapsed = api_request("chat/completions", body, timeout=timeout) - choice = payload.get("choices", [{}])[0] - usage = payload.get("usage", {}) - return { - "output": choice.get("message", {}).get("content", ""), - "prompt_tokens": usage.get("prompt_tokens", 0), - "completion_tokens": usage.get("completion_tokens", 0), - "total_ms": elapsed * 1000, - "rss_mb": get_rss_mb(), - } - - -def start_server(model_path: str, iteration: Iteration) -> subprocess.Popen[bytes]: - env = { - **os.environ, - **iteration.env, - "HIGGS_ENABLE_THINKING": "0", - "HIGGS_NO_CONFIG": "1", - } - cmd = [ - HIGGS_BIN, - "serve", - "--model", - model_path, - "--port", - str(PORT), - *iteration.args, - ] - proc = subprocess.Popen( - cmd, - env=env, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - preexec_fn=os.setsid, - ) - for _ in range(120): - try: - urllib.request.urlopen(f"{BASE}/v1/models", timeout=2) - return proc - except Exception: - time.sleep(1) - if proc.poll() is not None: - raise RuntimeError("server exited before becoming ready") - proc.kill() - proc.wait(timeout=5) - raise RuntimeError("server failed to start within 120 seconds") - - -def stop_server(proc: subprocess.Popen[bytes] | None) -> None: - if proc is None: - return - if proc.poll() is None: - try: - os.killpg(os.getpgid(proc.pid), signal.SIGTERM) - except ProcessLookupError: - pass - try: - proc.wait(timeout=10) - except subprocess.TimeoutExpired: - try: - os.killpg(os.getpgid(proc.pid), signal.SIGKILL) - except ProcessLookupError: - pass - time.sleep(2) - - -def get_model_id() -> str: - req = urllib.request.Request(f"{BASE}/v1/models") - with urllib.request.urlopen(req, timeout=5) as resp: - payload = json.loads(resp.read()) - return payload["data"][0]["id"] - - -def get_rss_mb() -> float: - global server_proc - if server_proc is None: - return 0.0 - try: - pgid = os.getpgid(server_proc.pid) - out = subprocess.check_output(["ps", "-o", "rss=", "-g", str(pgid)], text=True).strip() - total_kb = sum(int(line.strip()) for line in out.splitlines() if line.strip()) - return total_kb / 1024 - except (subprocess.CalledProcessError, ProcessLookupError, ValueError): - return 0.0 - - -def median(values: list[float]) -> float: - return statistics.median(values) if values else 0.0 - - -def normalize_text(text: str) -> str: - return " ".join(text.strip().split()).lower() - - -def clamp_cache_speedup(speedup: float, cap: float = CACHE_SPEEDUP_CAP) -> float: - return min(speedup, cap) - - -def compute_accuracy_score(result: dict[str, Any]) -> float: - qa_acc = result["qa"]["accuracy"] - long_acc = result["long_context"]["accuracy"] - structured_acc = result["structured_output"]["accuracy"] - cache_acc = result["prefix_cache"]["accuracy"] - return (qa_acc * 0.45) + (long_acc * 0.25) + (structured_acc * 0.15) + (cache_acc * 0.15) - - -def compute_speed_score(result: dict[str, Any], best_ttft: float, best_decode: float) -> float: - ttft_score = best_ttft / result["prompt_sweep"]["weighted_ttft_ms"] if best_ttft else 0.0 - decode_score = ( - result["prompt_sweep"]["weighted_decode_tps"] / best_decode if best_decode else 0.0 - ) - return (ttft_score * 0.55) + (decode_score * 0.45) - - -def compute_iteration_score(result: dict[str, Any], best_ttft: float, best_decode: float, best_cache: float) -> dict[str, float]: - accuracy = compute_accuracy_score(result) - speed = compute_speed_score(result, best_ttft, best_decode) - cache_speedup = ( - clamp_cache_speedup(result["prefix_cache"]["speedup"]) - if result["prefix_cache"]["passed"] - else 0.0 - ) - cache_score = cache_speedup / best_cache if best_cache else 0.0 - return { - "accuracy": accuracy, - "speed": speed, - "cache": cache_score, - "composite": 100.0 * ((accuracy * 0.45) + (speed * 0.45) + (cache_score * 0.10)), - } - - -def rank_results_by_score(results: list[dict[str, Any]]) -> list[dict[str, Any]]: - return sorted(results, key=lambda r: r["score"]["composite"], reverse=True) - - -def prompt_sweep(model: str, repeats: int) -> dict[str, Any]: - prompts = { - "short": SHORT_PROMPT, - "medium": MEDIUM_PROMPT, - "long": LONG_PROMPT, - } - weights = {"short": 0.2, "medium": 0.3, "long": 0.5} - metrics: dict[str, Any] = {} - - for label, prompt in prompts.items(): - runs = [] - for attempt in range(repeats): - if attempt == 0: - _ = stream_chat( - model, - [{"role": "user", "content": f"[warmup {label}] {prompt}"}], - max_tokens=8, - ) - result = stream_chat(model, [{"role": "user", "content": prompt}], max_tokens=64) - runs.append(result) - metrics[label] = { - "ttft_ms": median([run["ttft_ms"] for run in runs]), - "decode_tps": median([run["decode_tps"] for run in runs]), - "prompt_tokens": median([run["prompt_tokens"] for run in runs]), - "completion_tokens": median([run["completion_tokens"] for run in runs]), - } - - weighted_ttft = sum(metrics[label]["ttft_ms"] * weights[label] for label in prompts) - weighted_decode = sum(metrics[label]["decode_tps"] * weights[label] for label in prompts) - - return { - "by_prompt": metrics, - "weighted_ttft_ms": weighted_ttft, - "weighted_decode_tps": weighted_decode, - } - - -def run_qa_suite(model: str) -> dict[str, Any]: - results = [] - passed = 0 - for case in QA_CASES: - response = chat( - model, - [{"role": "user", "content": case["prompt"]}], - max_tokens=16, - ) - output = normalize_text(response["output"]) - success = output == case["expected"] - passed += int(success) - results.append( - { - "prompt": case["prompt"], - "expected": case["expected"], - "output": output, - "passed": success, - } - ) - return { - "passed": passed, - "total": len(QA_CASES), - "accuracy": passed / len(QA_CASES), - "results": results, - } - - -def run_long_context_needle(model: str) -> dict[str, Any]: - prompt = ( - "Read the following deployment notes carefully.\n\n" - f"{LONG_CONTEXT_FILLER}\n\n" - f"Important hidden code: {LONG_CONTEXT_NEEDLE}\n\n" - f"{LONG_CONTEXT_FILLER}\n\n" - "Question: reply with the deployment code only." - ) - response = chat(model, [{"role": "user", "content": prompt}], max_tokens=8) - output = normalize_text(response["output"]).replace(" ", "") - expected = LONG_CONTEXT_NEEDLE.lower() - return { - "expected": expected, - "output": output, - "passed": expected in output, - "accuracy": 1.0 if expected in output else 0.0, - "prompt_tokens": response["prompt_tokens"], - } - - -def run_structured_output(model: str, iteration_index: int) -> dict[str, Any]: - prompt = ( - "Return structured JSON only. Facts: model_family=qwen, iteration=" - f"{iteration_index}, prefill_focus=true, kv_bits=3, primary_goal=latency." - ) - response = chat( - model, - [{"role": "user", "content": prompt}], - max_tokens=64, - response_format=STRUCTURED_OUTPUT_SCHEMA, - ) - try: - parsed = json.loads(response["output"]) - except json.JSONDecodeError: - return {"passed": False, "accuracy": 0.0, "raw": response["output"]} - - passed = parsed == { - "model_family": "qwen", - "iteration": iteration_index, - "prefill_focus": True, - "kv_bits": 3, - "primary_goal": "latency", - } - return {"passed": passed, "accuracy": 1.0 if passed else 0.0, "json": parsed} - - -def run_prefix_cache_suite(model: str) -> dict[str, Any]: - prefix_messages = [ - { - "role": "system", - "content": ( - "You are reviewing an operations handbook. Study the material and " - "reply READY only.\n\n" - + PREFIX_CACHE_DOC - ), - }, - {"role": "user", "content": "Read the handbook and reply READY only."}, - ] - followup_messages = [ - *prefix_messages, - {"role": "assistant", "content": "READY"}, - { - "role": "user", - "content": "What is the failover region? Reply with two words only.", - }, - ] - - cold = stream_chat(model, followup_messages, max_tokens=8) - _ = stream_chat(model, prefix_messages, max_tokens=4) - warm = stream_chat(model, followup_messages, max_tokens=8) - - warm_output = normalize_text(warm["output"]) - passed = "cape town" in warm_output - speedup = cold["ttft_ms"] / warm["ttft_ms"] if warm["ttft_ms"] > 0 else 0.0 - return { - "cold_ttft_ms": cold["ttft_ms"], - "warm_ttft_ms": warm["ttft_ms"], - "speedup": speedup, - "answer": warm["output"], - "passed": passed, - "accuracy": 1.0 if passed else 0.0, - } - - -def benchmark_iteration( - model_path: str, - iteration_index: int, - iteration: Iteration, - repeats: int, -) -> dict[str, Any]: - global server_proc - - log(f"\n{'=' * 80}") - log(f"{iteration.label}") - log(f"Notes: {iteration.notes}") - if iteration.args: - log(f"Args: {' '.join(iteration.args)}") - log(f"{'=' * 80}") - - server_proc = start_server(model_path, iteration) - try: - model = get_model_id() - log(f"Model: {model}") - log("Warmup...") - _ = stream_chat(model, [{"role": "user", "content": "Say ready."}], max_tokens=4) - log(f"RSS after warmup: {get_rss_mb():.0f} MB") - - sweep = prompt_sweep(model, repeats) - qa = run_qa_suite(model) - long_ctx = run_long_context_needle(model) - structured = run_structured_output(model, iteration_index) - prefix_cache = run_prefix_cache_suite(model) - - log( - "Weighted prompt metrics: " - f"TTFT={sweep['weighted_ttft_ms']:.0f} ms " - f"decode={sweep['weighted_decode_tps']:.1f} tok/s" - ) - log( - "Accuracy checks: " - f"qa={qa['passed']}/{qa['total']} " - f"needle={'pass' if long_ctx['passed'] else 'fail'} " - f"json={'pass' if structured['passed'] else 'fail'} " - f"cache={'pass' if prefix_cache['passed'] else 'fail'}" - ) - log( - "Prefix cache: " - f"cold={prefix_cache['cold_ttft_ms']:.0f} ms " - f"warm={prefix_cache['warm_ttft_ms']:.0f} ms " - f"speedup={prefix_cache['speedup']:.2f}x" - ) - - return { - "iteration": iteration.slug, - "label": iteration.label, - "notes": iteration.notes, - "args": iteration.args, - "prompt_sweep": sweep, - "qa": qa, - "long_context": long_ctx, - "structured_output": structured, - "prefix_cache": prefix_cache, - "rss_mb": get_rss_mb(), - } - finally: - stop_server(server_proc) - server_proc = None - - -def score_results(results: list[dict[str, Any]]) -> None: - best_ttft = min(r["prompt_sweep"]["weighted_ttft_ms"] for r in results) - best_decode = max(r["prompt_sweep"]["weighted_decode_tps"] for r in results) - best_cache = max( - ( - min(r["prefix_cache"]["speedup"], CACHE_SPEEDUP_CAP) - for r in results - if r["prefix_cache"]["passed"] - ), - default=1.0, - ) - - for result in results: - result["score"] = compute_iteration_score(result, best_ttft, best_decode, best_cache) - - -def print_summary(results: list[dict[str, Any]]) -> None: - ordered = rank_results_by_score(results) - log(f"\n{'#' * 80}") - log("FINAL SUMMARY") - log(f"{'#' * 80}") - log( - f"{'Iteration':32s} {'Score':>8s} {'TTFT':>10s} {'Decode':>10s} " - f"{'QA':>6s} {'Needle':>8s} {'JSON':>6s} {'Cache':>8s}" - ) - log("-" * 96) - for result in ordered: - log( - f"{result['label'][:32]:32s} " - f"{result['score']['composite']:>7.1f} " - f"{result['prompt_sweep']['weighted_ttft_ms']:>9.0f} " - f"{result['prompt_sweep']['weighted_decode_tps']:>9.1f} " - f"{result['qa']['passed']:>2d}/{result['qa']['total']:<3d} " - f"{'pass' if result['long_context']['passed'] else 'fail':>8s} " - f"{'pass' if result['structured_output']['passed'] else 'fail':>6s} " - f"{result['prefix_cache']['speedup']:>7.2f}x" - ) - - winner = ordered[0] - log("") - log( - f"Winner: {winner['label']} " - f"(score={winner['score']['composite']:.1f}, " - f"TTFT={winner['prompt_sweep']['weighted_ttft_ms']:.0f} ms, " - f"decode={winner['prompt_sweep']['weighted_decode_tps']:.1f} tok/s)" - ) - - -def main() -> None: - parser = argparse.ArgumentParser(description="Benchmark Higgs MLX tuning profiles") - parser.add_argument("model_path", help="Local model path or resolved HF cache directory") - parser.add_argument("--repeats", type=int, default=2, help="Prompt-sweep repeats per iteration") - parser.add_argument( - "--output-json", - default=f"bench_mlx_tuning_{time.strftime('%Y%m%d_%H%M%S')}.json", - help="Path to write raw benchmark results", - ) - args = parser.parse_args() - - if not os.path.isfile(HIGGS_BIN): - raise SystemExit(f"Higgs binary not found: {HIGGS_BIN}") - - log("=" * 80) - log(f"MLX TUNING BENCHMARK — {time.strftime('%Y-%m-%d %H:%M:%S')}") - log(f"Binary: {HIGGS_BIN}") - log(f"Model: {args.model_path}") - log(f"Repeats: {args.repeats}") - log("=" * 80) - - results = [] - for idx, iteration in enumerate(ITERATIONS, start=1): - results.append(benchmark_iteration(args.model_path, idx, iteration, args.repeats)) - - score_results(results) - print_summary(results) - - with open(args.output_json, "w", encoding="utf-8") as f: - json.dump({"model_path": args.model_path, "results": results}, f, indent=2) - log(f"\nRaw results written to {args.output_json}") - - -if __name__ == "__main__": - try: - main() - except KeyboardInterrupt: - stop_server(server_proc) - sys.exit(130) diff --git a/benchmarks/bench_moe_sort.py b/benchmarks/bench_moe_sort.py deleted file mode 100644 index 98d31def..00000000 --- a/benchmarks/bench_moe_sort.py +++ /dev/null @@ -1,205 +0,0 @@ -""" -Profile MoE dispatch: per-token sort (Higgs) vs global batch sort (mlx-lm). - -Tests gather_qmm at DeepSeek-V2-Lite scale: - hidden=2048, intermediate=1408, 64 experts, top_k=6, 4-bit quantized -""" - -import statistics -import time -import mlx.core as mx -import mlx.nn -import numpy as np - -# DeepSeek-V2-Lite MoE dimensions -HIDDEN = 2048 -INTERMEDIATE = 1408 -NUM_EXPERTS = 64 -TOP_K = 6 -GROUP_SIZE = 64 -BITS = 4 - -WARMUP = 3 -ITERS = 10 - - -def make_weights(out_dim, in_dim): - """Create quantized expert weights [num_experts, out_dim, in_dim].""" - w_full = mx.random.normal((NUM_EXPERTS, out_dim, in_dim)) - w, scales, biases = mx.quantize(w_full, group_size=GROUP_SIZE, bits=BITS) - mx.eval(w, scales, biases) - return w, scales, biases - - -def make_inputs(B, L): - """Create input tensor and random top-k expert indices.""" - x = mx.random.normal((B, L, HIDDEN)) - indices = mx.random.randint(0, NUM_EXPERTS, (B, L, TOP_K)).astype(mx.uint32) - mx.eval(x, indices) - return x, indices - - -# -- Approach 1: Higgs (per-token sort or no sort) -- - -def forward_gather_higgs(x, indices, gate_w, gate_s, gate_b, - up_w, up_s, up_b, down_w, down_s, down_b, - do_sort=False): - """Higgs approach: expand x, call gather_qmm with per-token indices.""" - B, L, D = x.shape - - if do_sort: - # Per-token sort (Qwen3Next path) - indices = mx.sort(indices, axis=-1) - - x_exp = x.reshape(B, L, 1, 1, D) - - gate_out = mx.gather_qmm( - x_exp, gate_w, gate_s, gate_b, - rhs_indices=indices, transpose=True, - group_size=GROUP_SIZE, bits=BITS, sorted_indices=do_sort, - ) - up_out = mx.gather_qmm( - x_exp, up_w, up_s, up_b, - rhs_indices=indices, transpose=True, - group_size=GROUP_SIZE, bits=BITS, sorted_indices=do_sort, - ) - activated = mlx.nn.silu(gate_out) * up_out - - down_out = mx.gather_qmm( - activated, down_w, down_s, down_b, - rhs_indices=indices, transpose=True, - group_size=GROUP_SIZE, bits=BITS, sorted_indices=do_sort, - ) - return down_out.squeeze(-2) - - -# -- Approach 2: mlx-lm (global batch sort) -- - -def _gather_sort(x, indices): - """mlx-lm's global sort: flatten all tokens, sort by expert index.""" - *_, M = indices.shape # top_k - indices_flat = indices.flatten() - order = mx.argsort(indices_flat) - inv_order = mx.argsort(order) - x_sorted = x.reshape(-1, 1, x.shape[-1])[order // M] - return x_sorted, indices_flat[order], inv_order - - -def _scatter_unsort(x, inv_order, shape): - x = x[inv_order] - x = mx.unflatten(x, 0, shape) - return x - - -def forward_gather_mlxlm(x, indices, gate_w, gate_s, gate_b, - up_w, up_s, up_b, down_w, down_s, down_b): - """mlx-lm approach: global sort, then gather_qmm with sorted_indices=True. - - Matches SwitchGLU.__call__ from mlx_lm/models/switch_layers.py exactly. - """ - orig_shape = indices.shape # [B, L, top_k] - - # Step 1: expand x (same as SwitchGLU) - x_exp = x.reshape(x.shape[0], x.shape[1], 1, 1, x.shape[2]) - - # Step 2: global sort — flatten all tokens, reorder by expert index - # x_sorted: [B*L*top_k, 1, D], idx_sorted: [B*L*top_k], inv_order: [B*L*top_k] - x_sorted, idx_sorted, inv_order = _gather_sort(x_exp, indices) - - gate_out = mx.gather_qmm( - x_sorted, gate_w, gate_s, gate_b, - rhs_indices=idx_sorted, transpose=True, - group_size=GROUP_SIZE, bits=BITS, sorted_indices=True, - ) - up_out = mx.gather_qmm( - x_sorted, up_w, up_s, up_b, - rhs_indices=idx_sorted, transpose=True, - group_size=GROUP_SIZE, bits=BITS, sorted_indices=True, - ) - activated = mlx.nn.silu(gate_out) * up_out - - down_out = mx.gather_qmm( - activated, down_w, down_s, down_b, - rhs_indices=idx_sorted, transpose=True, - group_size=GROUP_SIZE, bits=BITS, sorted_indices=True, - ) - - # Step 3: unsort back to [B, L, top_k, 1, D], then squeeze - out = _scatter_unsort(down_out, inv_order, orig_shape) # [B, L, top_k, 1, D] - return out.squeeze(-2) # [B, L, top_k, D] - - -def bench(fn, label, *args, **kwargs): - # Warmup - for _ in range(WARMUP): - r = fn(*args, **kwargs) - mx.eval(r) - - # Timed - times = [] - for _ in range(ITERS): - t0 = time.perf_counter() - r = fn(*args, **kwargs) - mx.eval(r) - times.append(time.perf_counter() - t0) - - med = statistics.median(times) - return med - - -def main(): - print("DeepSeek-V2-Lite MoE dispatch profiling") - print(f" {NUM_EXPERTS} experts, top_k={TOP_K}, {BITS}-bit, " - f"hidden={HIDDEN}, intermediate={INTERMEDIATE}") - print(f" {WARMUP} warmup, {ITERS} iters, reporting median\n") - - # Create weights once - gate_w, gate_s, gate_b = make_weights(INTERMEDIATE, HIDDEN) - up_w, up_s, up_b = make_weights(INTERMEDIATE, HIDDEN) - down_w, down_s, down_b = make_weights(HIDDEN, INTERMEDIATE) - w_args = (gate_w, gate_s, gate_b, up_w, up_s, up_b, down_w, down_s, down_b) - - seq_lens = [1, 32, 128, 512, 1024, 2048] - - print(f"{'SeqLen':>8} | {'Higgs(nosort)':>14} | {'Higgs(ptsort)':>14} | {'mlx-lm(global)':>14} | {'Speedup':>10}") - print("-" * 75) - - for L in seq_lens: - B = 1 - x, indices = make_inputs(B, L) - - t_nosort = bench(forward_gather_higgs, f"nosort-{L}", - x, indices, *w_args, do_sort=False) - t_ptsort = bench(forward_gather_higgs, f"ptsort-{L}", - x, indices, *w_args, do_sort=True) - t_global = bench(forward_gather_mlxlm, f"global-{L}", - x, indices, *w_args) - - fastest_higgs = min(t_nosort, t_ptsort) - speedup = fastest_higgs / t_global if t_global > 0 else float('inf') - - print(f"{L:>8} | {t_nosort*1000:>11.2f} ms | {t_ptsort*1000:>11.2f} ms | " - f"{t_global*1000:>11.2f} ms | {speedup:>8.2f}x") - - # Also time the global sort plus the gather it currently requires. - print("\n--- Sort + gather overhead ---") - for L in [128, 512, 2048]: - x, indices = make_inputs(1, L) - x_exp = x.reshape(1, L, 1, 1, HIDDEN) - - # Time global sort - for _ in range(WARMUP): - r = _gather_sort(x_exp, indices) - mx.eval(*r) - times = [] - for _ in range(ITERS): - t0 = time.perf_counter() - r = _gather_sort(x_exp, indices) - mx.eval(*r) - times.append(time.perf_counter() - t0) - med = statistics.median(times) - print(f" L={L:>5}: global sort = {med*1000:.2f} ms") - - -if __name__ == "__main__": - main() diff --git a/benchmarks/bench_ppl_tq.py b/benchmarks/bench_ppl_tq.py deleted file mode 100644 index f0cec78a..00000000 --- a/benchmarks/bench_ppl_tq.py +++ /dev/null @@ -1,462 +0,0 @@ -#!/usr/bin/env python3 -"""Perplexity & quality benchmark: Higgs TurboQuant vs baseline. - -Measures: - 1. Ground-truth PPL via mlx-lm (no server, direct model) - 2. KLD between baseline and TQ via Higgs server logprobs - 3. Decode speed at various context lengths (up to 32K) - 4. Sparse V skip rate estimation - -Usage: - python3 benchmarks/bench_ppl_tq.py [--ctx 2048] [--stride 512] [--chunks 50] - -Requires: mlx-lm, datasets -""" - -import argparse -import json -import math -import os -import signal -import subprocess -import sys -import time -import urllib.request - -import mlx.core as mx -import numpy as np - -HIGGS_BIN = os.environ.get("HIGGS_BIN", "./target/release/higgs") - - -# --------------------------------------------------------------------------- -# Part 1: Ground-truth perplexity via mlx-lm -# --------------------------------------------------------------------------- - -def compute_ppl_mlx(model_path, ctx_len=2048, stride=512, max_chunks=50): - """Compute perplexity using mlx-lm direct inference (no server).""" - from datasets import load_dataset - from mlx_lm import load, stream_generate - - print(f"\n{'='*60}") - print("PART 1: Ground-truth PPL via mlx-lm") - print(f" Model: {model_path}") - print(f" Context: {ctx_len}, Stride: {stride}, Max chunks: {max_chunks}") - print(f"{'='*60}") - - # Load dataset - ds = load_dataset("wikitext", "wikitext-2-raw-v1", split="test") - text = "\n\n".join([x["text"] for x in ds if x["text"].strip()]) - - # Load model + tokenizer - print("Loading model...") - model, tokenizer = load(model_path) - print(f" Model loaded. Peak memory: {mx.get_peak_memory() / 1e9:.1f} GB") - - # Tokenize full text - tokens = tokenizer.encode(text) - print(f" Total tokens: {len(tokens)}") - - # Sliding window PPL - nlls = [] - n_tokens = 0 - chunk_count = 0 - - for start in range(0, len(tokens) - ctx_len, stride): - if chunk_count >= max_chunks: - break - - chunk = tokens[start : start + ctx_len] - input_ids = mx.array(chunk[:-1])[None, :] # (1, T-1) - targets = mx.array(chunk[1:]) # (T-1,) - - # Forward pass to get logits - logits = model(input_ids) - logits = logits.squeeze(0) # (T-1, vocab) - - # Only score tokens in the non-overlapping region (after stride) - if start > 0: - score_start = ctx_len - stride - else: - score_start = 0 - - score_logits = logits[score_start:] - score_targets = targets[score_start:] - - # Log-softmax (numerically stable) → gather target logprobs - log_probs = score_logits - mx.logsumexp(score_logits, axis=-1, keepdims=True) - - target_logprobs = mx.take_along_axis( - log_probs, score_targets[:, None], axis=1 - ).squeeze(-1) - - nll = -target_logprobs.sum().item() - count = score_targets.shape[0] - nlls.append(nll) - n_tokens += count - chunk_count += 1 - - ppl_so_far = math.exp(sum(nlls) / n_tokens) - mem = mx.get_peak_memory() / 1e9 - - if chunk_count % 10 == 0 or chunk_count <= 3: - print(f" Chunk {chunk_count}: PPL={ppl_so_far:.3f} " - f"(tokens={n_tokens}, mem={mem:.1f}GB)") - - mx.eval(target_logprobs) # force eval to free graph - - if n_tokens == 0: - print("\n FINAL PPL: N/A (no chunks scored — text shorter than context?)") - final_ppl = float("inf") - else: - final_ppl = math.exp(sum(nlls) / n_tokens) - print(f"\n FINAL PPL: {final_ppl:.4f} ({n_tokens} tokens, {chunk_count} chunks)") - print(f" Peak memory: {mx.get_peak_memory() / 1e9:.1f} GB") - - # Free model - del model - mx.clear_cache() - - return final_ppl, n_tokens - - -# --------------------------------------------------------------------------- -# Part 2: Server-based KLD and decode benchmarks -# --------------------------------------------------------------------------- - -def api(port, endpoint, body, timeout=300): - data = json.dumps(body).encode() - req = urllib.request.Request( - f"http://127.0.0.1:{port}/v1/{endpoint}", - data=data, - headers={"Content-Type": "application/json"}, - ) - t0 = time.perf_counter() - with urllib.request.urlopen(req, timeout=timeout) as resp: - result = json.loads(resp.read()) - elapsed = time.perf_counter() - t0 - return result, elapsed - - -def start_server(model_path, port, kv_mode=None, kv_bits=3): - cmd = [HIGGS_BIN, "serve", "--model", model_path, "--port", str(port)] - if kv_mode: - cmd += ["--kv-cache", kv_mode, "--kv-bits", str(kv_bits), "--kv-seed", "0"] - env = {**os.environ, "HIGGS_ENABLE_THINKING": "0", "HIGGS_NO_CONFIG": "1"} - proc = subprocess.Popen( - cmd, - env=env, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - for _ in range(90): - try: - urllib.request.urlopen(f"http://127.0.0.1:{port}/v1/models", timeout=2) - return proc - except Exception: - time.sleep(1) - proc.kill() - proc.wait(timeout=5) - raise RuntimeError(f"Server on port {port} failed to start") - - -def stop_server(proc): - if proc and proc.poll() is None: - proc.send_signal(signal.SIGTERM) - try: - proc.wait(timeout=5) - except subprocess.TimeoutExpired: - proc.kill() - - -def get_model_id(port): - req = urllib.request.Request(f"http://127.0.0.1:{port}/v1/models") - with urllib.request.urlopen(req, timeout=5) as resp: - data = json.loads(resp.read()) - return data["data"][0]["id"] - - -WORD = "the quick brown fox jumps over the lazy dog " # ~10 tokens - - -def chat(port, model, messages, max_tokens=256, temperature=0, logprobs=False, - top_logprobs=None): - body = { - "model": model, - "messages": messages, - "max_tokens": max_tokens, - "temperature": temperature, - } - if logprobs: - body["logprobs"] = True - if top_logprobs is not None: - body["top_logprobs"] = top_logprobs - result, elapsed = api(port, "chat/completions", body) - usage = result.get("usage", {}) - choice = result.get("choices", [{}])[0] - content = choice.get("message", {}).get("content", "") - lp = choice.get("logprobs", None) - return content, usage, elapsed, lp - - - -def test_quality_sequential(model_path, port, bits): - """Quality comparison using one server at a time (safe for 32GB RAM).""" - print(f"\n{'='*60}") - print("PART 2: Output quality — baseline vs TQ (sequential)") - print(f"{'='*60}") - - prompts = [ - "Explain how a hash table works step by step.", - "What causes the seasons on Earth?", - "Write a Python function to find the longest common subsequence.", - "Describe the process of photosynthesis in detail.", - "What is the difference between TCP and UDP?", - "Explain the theory of general relativity in simple terms.", - "How does a neural network learn through backpropagation?", - "What are the main causes of climate change?", - "Describe the water cycle and its importance.", - "How does public key cryptography work?", - ] - - # Step 1: Collect baseline outputs - print("\n Collecting baseline outputs...") - proc = start_server(model_path, port) - base_outputs = [] - try: - model_id = get_model_id(port) - for i, prompt in enumerate(prompts): - content, _usage, _elapsed, _lp = chat( - port, model_id, [{"role": "user", "content": prompt}], - max_tokens=64, logprobs=True, top_logprobs=5, - ) - base_outputs.append(content) - print(f" [{i+1}/{len(prompts)}] {len(content)} chars") - finally: - stop_server(proc) - print(" Baseline server stopped.") - time.sleep(3) - - # Step 2: Collect TQ outputs - print(f"\n Collecting TQ-{bits}bit outputs...") - proc = start_server(model_path, port, kv_mode="turboquant", kv_bits=bits) - tq_outputs = [] - try: - model_id = get_model_id(port) - for i, prompt in enumerate(prompts): - content, _usage, _elapsed, _lp = chat( - port, model_id, [{"role": "user", "content": prompt}], - max_tokens=64, logprobs=True, top_logprobs=5, - ) - tq_outputs.append(content) - print(f" [{i+1}/{len(prompts)}] {len(content)} chars") - finally: - stop_server(proc) - print(" TQ server stopped.") - time.sleep(3) - - # Step 3: Compare - results = [] - print(f"\n {'Prompt':<52} | {'Jaccard':>8} | {'Verdict':>10}") - print(" " + "-" * 75) - for prompt, c_base, c_tq in zip(prompts, base_outputs, tq_outputs): - w_base = set(c_base.lower().split()) - w_tq = set(c_tq.lower().split()) - jaccard = ( - len(w_base & w_tq) / len(w_base | w_tq) if (w_base or w_tq) else 0 - ) - verdict = "MATCH" if jaccard > 0.5 else "DIVERGED" - print(f" {prompt[:50]:<52} | {jaccard:>7.2f} | {verdict:>10}") - results.append({"prompt": prompt[:50], "jaccard": jaccard}) - - avg = sum(r["jaccard"] for r in results) / len(results) if results else 0 - print(f"\n Average Jaccard: {avg:.3f}") - return results, avg - - -def test_decode_context_sweep(port, model, label, context_sizes=None): - """Decode speed at various context lengths. - - Uses two non-streaming requests per context to separate TTFT from decode: - 1. max_tokens=1 → elapsed ≈ TTFT (prefill + 1 token) - 2. max_tokens=128 → elapsed = prefill + N tokens - decode_tps ≈ (ctoks - 1) / (elapsed_128 - elapsed_1) - """ - if context_sizes is None: - context_sizes = [100, 1000, 4000, 8000, 16000, 24000, 32000] - - print(f"\n{'='*60}") - print(f"PART 3: Decode speed vs context length — {label}") - print(f"{'='*60}") - print(f" {'Context':>10} | {'Decode tok/s':>14} | {'TTFT ms':>10} | {'Gen toks':>10}") - print(" " + "-" * 55) - - results = [] - for ctx_tokens in context_sizes: - repeat = max(1, ctx_tokens // 10) - prompt = WORD * repeat - - try: - # Request 1: TTFT only (max_tokens=1, unique prefix) - pfx1 = f"[a{time.time():.6f}] Summarize: " - _, u1, t1, _ = chat(port, model, - [{"role": "user", "content": pfx1 + prompt}], - max_tokens=1) - ptoks = u1.get("prompt_tokens", 0) - ttft_s = t1 - - # Request 2: Full generation (max_tokens=128, different unique prefix) - pfx2 = f"[b{time.time():.6f}] Summarize: " - _, u2, t2, _ = chat(port, model, - [{"role": "user", "content": pfx2 + prompt}], - max_tokens=128) - ctoks = u2.get("completion_tokens", 0) - - if ctoks > 1: - # decode_time ≈ full_time - ttft - decode_time = max(t2 - ttft_s, 0.01) - tps = (ctoks - 1) / decode_time - print(f" {ptoks:>10} | {tps:>11.1f} t/s | " - f"{ttft_s*1000:>7.0f} ms | {ctoks:>10}") - results.append({"ctx": ptoks, "tps": tps, - "ttft": ttft_s * 1000}) - else: - print(f" {ptoks:>10} | {'NO OUTPUT':>14} | " - f"{ttft_s*1000:>7.0f} ms | {ctoks:>10}") - results.append({"ctx": ptoks, "tps": 0, "ttft": ttft_s * 1000}) - - except Exception as e: - print(f" {ctx_tokens:>10} | {'FAILED':>14} | {str(e)[:40]:>40}") - results.append({"ctx": ctx_tokens, "tps": 0, "error": str(e)[:50]}) - if "memory" in str(e).lower() or "timeout" in str(e).lower(): - print(" *** Stopping sweep — memory or timeout error ***") - break - - return results - - -# --------------------------------------------------------------------------- -# Main -# --------------------------------------------------------------------------- - -def main(): - parser = argparse.ArgumentParser(description="TurboQuant PPL & quality benchmark") - parser.add_argument("model_path") - parser.add_argument("--ctx", type=int, default=2048, help="Context length for PPL") - parser.add_argument("--stride", type=int, default=512, help="Stride for PPL sliding window") - parser.add_argument("--chunks", type=int, default=50, help="Max chunks for PPL eval") - parser.add_argument("--bits", type=int, default=3, help="TQ bit width") - parser.add_argument("--port", type=int, default=8097, help="Base port") - parser.add_argument("--skip-ppl", action="store_true", help="Skip mlx-lm PPL (memory intensive)") - parser.add_argument("--skip-server", action="store_true", help="Skip server-based tests") - args = parser.parse_args() - - print("=" * 60) - print("HIGGS TURBOQUANT — PPL & QUALITY BENCHMARK") - print(f"Model: {args.model_path}") - print(f"TQ bits: {args.bits}") - print(f"PPL context: {args.ctx}, stride: {args.stride}, chunks: {args.chunks}") - print("=" * 60) - - # Part 1: Ground-truth PPL - ppl = None - if not args.skip_ppl: - try: - ppl, n_tokens = compute_ppl_mlx( - args.model_path, - ctx_len=args.ctx, - stride=args.stride, - max_chunks=args.chunks, - ) - except Exception as e: - print(f"\n PPL computation failed: {e}") - import traceback - traceback.print_exc() - else: - print("\n[Skipping mlx-lm PPL — --skip-ppl flag]") - - if args.skip_server: - print("\n[Skipping server tests — --skip-server flag]") - return - - # Part 2 & 3: Server-based tests (sequential — one server at a time) - port = args.port - - # --- Baseline server --- - print(f"\n{'='*60}") - print("Starting BASELINE server (no TQ)...") - proc_base = None - try: - proc_base = start_server(args.model_path, port) - model_base = get_model_id(port) - print(f" Ready: {model_base}") - - # Decode sweep — baseline - sweep_base = test_decode_context_sweep( - port, model_base, "baseline", - context_sizes=[100, 1000, 4000, 8000, 16000], - ) - finally: - stop_server(proc_base) - print("Baseline server stopped.") - time.sleep(2) # let port free - - # --- TQ server --- - print(f"\n{'='*60}") - print(f"Starting TURBOQUANT server ({args.bits}-bit)...") - proc_tq = None - try: - proc_tq = start_server(args.model_path, port, kv_mode="turboquant", kv_bits=args.bits) - model_tq = get_model_id(port) - print(f" Ready: {model_tq}") - - # Decode sweep — TQ (push further since KV is compressed) - sweep_tq = test_decode_context_sweep( - port, model_tq, f"turboquant-{args.bits}bit", - context_sizes=[100, 1000, 4000, 8000, 16000, 24000, 32000], - ) - finally: - stop_server(proc_tq) - print("TQ server stopped.") - time.sleep(2) - - # --- Quality comparison: single-server sequential (safe for 32GB) --- - kld_results, avg_jaccard = test_quality_sequential( - args.model_path, port, args.bits - ) - - # --- Final summary --- - print(f"\n{'='*60}") - print("FINAL SUMMARY") - print(f"{'='*60}") - - if ppl is not None: - print(f"\n Ground-truth PPL (mlx-lm): {ppl:.4f}") - - if sweep_base and sweep_tq: - print(f"\n Decode speed comparison:") - print(f" {'Context':>10} | {'Baseline':>12} | {'TQ {}-bit':>12} | {'Speedup':>8}".format(args.bits)) - print(" " + "-" * 50) - base_map = {r["ctx"]: r["tps"] for r in sweep_base} - for r in sweep_tq: - ctx = r["ctx"] - tps_tq = r["tps"] - tps_base = base_map.get(ctx, 0) - if tps_base > 0 and tps_tq > 0: - speedup = tps_tq / tps_base - print(f" {ctx:>10} | {tps_base:>9.1f} t/s | {tps_tq:>9.1f} t/s | {speedup:>7.2f}x") - elif tps_tq > 0: - print(f" {ctx:>10} | {'N/A':>12} | {tps_tq:>9.1f} t/s | {'':>8}") - else: - err = r.get("error", "failed") - print(f" {ctx:>10} | {'':>12} | {err:>12} | {'':>8}") - - if avg_jaccard > 0: - print(f"\n Output quality (Jaccard similarity): {avg_jaccard:.3f}") - print(f" > 0.5 = good, > 0.7 = very similar, < 0.3 = diverged") - - print() - - -if __name__ == "__main__": - main() diff --git a/benchmarks/bench_prefix_cache.py b/benchmarks/bench_prefix_cache.py deleted file mode 100644 index 09232a9b..00000000 --- a/benchmarks/bench_prefix_cache.py +++ /dev/null @@ -1,300 +0,0 @@ -#!/usr/bin/env python3 -"""Benchmark paged prefix cache: TTFT on cache miss vs cache hit.""" - -import json -import subprocess -import sys -import time -import signal -import os -import urllib.request -import urllib.error - -SERVER = "http://127.0.0.1:8080" -HIGGS = "./target/release/higgs" - -# Large system prompt (~1500 words, ~2000 tokens) to exercise prefix caching -SYSTEM_PROMPT = """You are a highly skilled software architect with deep expertise in distributed systems, \ -database design, and cloud-native applications. You provide thorough, well-reasoned technical advice. \ -When answering questions, you consider trade-offs, scalability implications, and maintainability. \ -You draw on experience with microservices, event-driven architectures, and modern DevOps practices. \ -You always explain your reasoning step by step and provide concrete examples when possible. \ -Your responses are structured with clear headings and bullet points for readability. - -You are familiar with AWS, GCP, Azure, Kubernetes, Docker, Terraform, and various CI/CD pipelines. \ -You understand the CAP theorem, ACID properties, eventual consistency, and distributed consensus algorithms. \ -You can discuss the pros and cons of SQL vs NoSQL databases, message queues vs event streams, \ -synchronous vs asynchronous communication patterns, and monolithic vs microservice architectures. \ -You are also knowledgeable about security best practices, observability (logging, metrics, tracing), \ -and performance optimization techniques. When discussing code, you prefer clean, idiomatic solutions \ -that follow established patterns in the relevant ecosystem. - -Here is your detailed knowledge base that you must reference when answering: - -SECTION 1 - DATABASE DESIGN PRINCIPLES: -When designing databases, always consider the access patterns first. For OLTP workloads, normalize to 3NF \ -and use appropriate indexes. For OLAP workloads, consider star or snowflake schemas. Partitioning strategies \ -should be based on query patterns: range partitioning for time-series data, hash partitioning for uniform \ -distribution, and list partitioning for categorical data. Always plan for data growth and implement proper \ -archival strategies. Connection pooling is essential - use PgBouncer for PostgreSQL, ProxySQL for MySQL. \ -Read replicas should be used for read-heavy workloads, but be aware of replication lag implications. \ -For multi-region deployments, consider CockroachDB, YugabyteDB, or Spanner for global consistency. - -SECTION 2 - MICROSERVICES PATTERNS: -Service boundaries should align with business domains (Domain-Driven Design). Use the Strangler Fig pattern \ -for migrating from monoliths. Implement Circuit Breakers (Hystrix/Resilience4j) for fault tolerance. \ -Use the Saga pattern for distributed transactions - prefer choreography over orchestration for loose coupling. \ -API Gateway pattern (Kong, Envoy) handles cross-cutting concerns. Service mesh (Istio, Linkerd) provides \ -observability and traffic management. Event sourcing with CQRS is appropriate when audit trails are required \ -or when read and write models differ significantly. Use the Outbox pattern for reliable event publishing. \ -Idempotency keys are essential for safe retries. Health checks should include both liveness and readiness probes. - -SECTION 3 - CACHING STRATEGIES: -Implement caching at multiple levels: CDN for static assets, application-level for computed results, \ -database query cache for frequent queries. Use Cache-Aside (Lazy Loading) as the default pattern. \ -Write-Through caching ensures consistency but adds latency. Write-Behind (Write-Back) improves write \ -performance but risks data loss. Cache invalidation strategies: TTL-based for simplicity, event-driven \ -for consistency. Redis Cluster for horizontal scaling, Redis Sentinel for high availability. \ -Implement cache warming for predictable traffic patterns. Use bloom filters to prevent cache penetration. \ -Consider local caches (Caffeine, Guava) before distributed caches to reduce network overhead. - -SECTION 4 - SECURITY BEST PRACTICES: -Implement defense in depth: network segmentation, WAF, rate limiting, input validation, output encoding. \ -Use OAuth 2.0 with PKCE for public clients, client credentials for service-to-service. JWT tokens should \ -be short-lived (15 minutes) with refresh token rotation. Store secrets in Vault or AWS Secrets Manager, \ -never in code or environment variables. Implement mutual TLS for service-to-service communication. \ -Use prepared statements to prevent SQL injection. Apply the principle of least privilege for IAM roles. \ -Enable audit logging for all sensitive operations. Implement CORS properly - never use wildcard in production. \ -Regular dependency scanning with Snyk or Dependabot. Container image scanning with Trivy or Aqua. - -SECTION 5 - OBSERVABILITY: -The three pillars: metrics (Prometheus/Datadog), logs (ELK/Loki), traces (Jaeger/Zipkin). \ -Use structured logging (JSON) with correlation IDs across services. Implement RED metrics for services \ -(Rate, Errors, Duration) and USE metrics for resources (Utilization, Saturation, Errors). \ -Set up alerts based on SLOs, not arbitrary thresholds. Use percentiles (p50, p95, p99) instead of averages \ -for latency metrics. Implement distributed tracing with OpenTelemetry for cross-service visibility. \ -Create dashboards that answer specific questions rather than showing all available metrics. \ -Log aggregation should include proper retention policies and cost management. Error tracking with \ -Sentry or Bugsnag for application-level issues. Synthetic monitoring for critical user journeys. - -SECTION 6 - PERFORMANCE OPTIMIZATION: -Profile before optimizing - use flame graphs for CPU, heap dumps for memory. Database: EXPLAIN ANALYZE \ -all slow queries, add covering indexes, optimize JOIN order. Application: minimize allocations, use \ -connection pooling, batch operations where possible. Network: enable HTTP/2, use gRPC for internal \ -services, implement request coalescing. Frontend: lazy loading, code splitting, image optimization. \ -Load testing with k6 or Locust before production releases. Capacity planning based on projected growth \ -with 2x headroom. Horizontal scaling is preferred over vertical - design for statelessness. \ -Use async processing for non-critical paths. Implement backpressure mechanisms to prevent cascading failures. - -You are pragmatic and focus on delivering value rather than over-engineering solutions. Keep answers concise \ -but thorough. Always mention relevant trade-offs.""" - -USER_PROMPTS = [ - "Explain the CAP theorem in one paragraph.", - "What is eventual consistency? Keep it brief.", - "Compare Redis and Memcached in three sentences.", -] - - -def wait_for_server(timeout=120): - """Wait until the server responds to /v1/models. Returns model name.""" - deadline = time.time() + timeout - while time.time() < deadline: - try: - req = urllib.request.Request(f"{SERVER}/v1/models") - with urllib.request.urlopen(req, timeout=5) as resp: - if resp.status == 200: - data = json.loads(resp.read()) - return data["data"][0]["id"] - except (urllib.error.URLError, ConnectionRefusedError, OSError, KeyError, IndexError): - pass - time.sleep(1) - return None - - -def stream_chat(messages, max_tokens=50, model_name="test"): - """Send a streaming chat completion and measure TTFT + collect output.""" - payload = json.dumps({ - "model": model_name, - "messages": messages, - "max_tokens": max_tokens, - "temperature": 0.0, - "stream": True, - }).encode() - - req = urllib.request.Request( - f"{SERVER}/v1/chat/completions", - data=payload, - headers={"Content-Type": "application/json"}, - ) - - t0 = time.perf_counter() - first_token_time = None - tokens = [] - - with urllib.request.urlopen(req, timeout=300) as resp: - buffer = b"" - while True: - chunk = resp.read(1) - if not chunk: - break - buffer += chunk - while b"\n" in buffer: - line, buffer = buffer.split(b"\n", 1) - line = line.strip() - if not line or line == b"data: [DONE]": - continue - if line.startswith(b"data: "): - try: - data = json.loads(line[6:]) - choices = data.get("choices", []) - if not choices: - continue - delta = choices[0].get("delta", {}) - content = delta.get("content", "") - if content and first_token_time is None: - first_token_time = time.perf_counter() - if content: - tokens.append(content) - except (json.JSONDecodeError, KeyError, IndexError): - pass - - total_time = time.perf_counter() - t0 - ttft = (first_token_time - t0) if first_token_time else total_time - - return { - "ttft_ms": ttft * 1000, - "total_ms": total_time * 1000, - "output": "".join(tokens), - "num_tokens": len(tokens), - } - - -def bench_model(model_dir, label): - """Benchmark a single model: cache miss, cache hit, different prompt.""" - print(f"\n{'='*70}") - print(f"Model: {label}") - print(f"Path: {model_dir}") - print(f"{'='*70}") - - # Start server - env = {**os.environ, "HIGGS_ENABLE_THINKING": "0"} - proc = subprocess.Popen( - [HIGGS, "serve", "--model", model_dir, "--port", "8080"], - env=env, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - preexec_fn=os.setsid, - ) - - try: - print("Waiting for server...", end="", flush=True) - model_name = wait_for_server() - if not model_name: - print(" TIMEOUT") - return None - print(f" ready (model={model_name})") - - results = [] - - # Warmup: one throwaway request to stabilize GPU - print("Warmup...", end="", flush=True) - warmup_msgs = [ - {"role": "system", "content": "Be brief."}, - {"role": "user", "content": "Say hello."}, - ] - stream_chat(warmup_msgs, max_tokens=10, model_name=model_name) - print(" done") - - for i, user_msg in enumerate(USER_PROMPTS): - messages = [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "user", "content": user_msg}, - ] - tag = "MISS" if i == 0 else "HIT" - print(f"\n [{tag}] Request {i+1}: \"{user_msg}\"") - - r = stream_chat(messages, max_tokens=80, model_name=model_name) - results.append(r) - - print(f" TTFT: {r['ttft_ms']:8.1f} ms") - print(f" Total: {r['total_ms']:8.1f} ms ({r['num_tokens']} tokens)") - print(f" Output: {r['output'][:120]}...") - - # Summary - miss_ttft = results[0]["ttft_ms"] - hit_ttfts = [r["ttft_ms"] for r in results[1:]] - avg_hit = sum(hit_ttfts) / len(hit_ttfts) - speedup = miss_ttft / avg_hit if avg_hit > 0 else 0 - - print(f"\n --- Summary ---") - print(f" Cache MISS TTFT: {miss_ttft:8.1f} ms") - print(f" Cache HIT TTFT: {avg_hit:8.1f} ms (avg of {len(hit_ttfts)} hits)") - print(f" Speedup: {speedup:8.2f}x") - - return { - "model": label, - "miss_ttft_ms": miss_ttft, - "avg_hit_ttft_ms": avg_hit, - "speedup": speedup, - "results": results, - } - - finally: - try: - os.killpg(os.getpgid(proc.pid), signal.SIGTERM) - proc.wait(timeout=10) - except (ProcessLookupError, subprocess.TimeoutExpired): - try: - os.killpg(os.getpgid(proc.pid), signal.SIGKILL) - except ProcessLookupError: - pass - time.sleep(2) - - -def main(): - models = [ - ( - os.path.expanduser("~/.cache/lm-studio/models/NexVeridian/Qwen3.5-35B-A3B-3bit"), - "Qwen3.5-35B-A3B-3bit (MoE)", - ), - ( - os.path.expanduser("~/.cache/lm-studio/models/mlx-community/Qwen3.5-27B-Claude-4.6-Opus-Distilled-MLX-4bit"), - "Qwen3.5-27B-4bit (Dense)", - ), - ( - os.path.expanduser("~/.cache/lm-studio/models/mlx-community/DeepSeek-Coder-V2-Lite-Instruct-4bit-mlx"), - "DeepSeek-V2-Lite-4bit (MoE)", - ), - ] - - # Filter to models that exist - models = [(d, l) for d, l in models if os.path.isdir(d)] - if not models: - print("No models found!") - sys.exit(1) - - print(f"Benchmarking {len(models)} models") - print(f"System prompt: ~{len(SYSTEM_PROMPT.split())} words") - - all_results = [] - for model_dir, label in models: - result = bench_model(model_dir, label) - if result: - all_results.append(result) - - # Final table - if all_results: - print(f"\n{'='*70}") - print("FINAL RESULTS") - print(f"{'='*70}") - print(f"{'Model':<35} {'Miss TTFT':>10} {'Hit TTFT':>10} {'Speedup':>8}") - print("-" * 70) - for r in all_results: - print( - f"{r['model']:<35} {r['miss_ttft_ms']:>8.0f}ms {r['avg_hit_ttft_ms']:>8.0f}ms {r['speedup']:>7.2f}x" - ) - - -if __name__ == "__main__": - main() diff --git a/benchmarks/bench_prefix_cache_turns.py b/benchmarks/bench_prefix_cache_turns.py deleted file mode 100644 index 6b912828..00000000 --- a/benchmarks/bench_prefix_cache_turns.py +++ /dev/null @@ -1,243 +0,0 @@ -#!/usr/bin/env python3 -"""Benchmark: does prefix cache speedup degrade with conversation turns?""" - -import json -import os -import signal -import subprocess -import sys -import time -import urllib.request -import urllib.error - -SERVER = "http://127.0.0.1:8080" -HIGGS = "./target/release/higgs" - -# Same large system prompt from main benchmark -SYSTEM_PROMPT = """You are a highly skilled software architect with deep expertise in distributed systems, \ -database design, and cloud-native applications. You provide thorough, well-reasoned technical advice. \ -When answering questions, you consider trade-offs, scalability implications, and maintainability. \ -You draw on experience with microservices, event-driven architectures, and modern DevOps practices. \ -You always explain your reasoning step by step and provide concrete examples when possible. - -Here is your detailed knowledge base: - -SECTION 1 - DATABASE DESIGN PRINCIPLES: -When designing databases, always consider the access patterns first. For OLTP workloads, normalize to 3NF \ -and use appropriate indexes. For OLAP workloads, consider star or snowflake schemas. Partitioning strategies \ -should be based on query patterns: range partitioning for time-series data, hash partitioning for uniform \ -distribution. Connection pooling is essential. Read replicas for read-heavy workloads. - -SECTION 2 - MICROSERVICES PATTERNS: -Service boundaries should align with business domains (Domain-Driven Design). Use the Strangler Fig pattern \ -for migrating from monoliths. Implement Circuit Breakers for fault tolerance. Use the Saga pattern for \ -distributed transactions. API Gateway handles cross-cutting concerns. Service mesh provides observability. \ -Event sourcing with CQRS when audit trails are required. Outbox pattern for reliable event publishing. - -SECTION 3 - CACHING STRATEGIES: -Implement caching at multiple levels: CDN for static assets, application-level for computed results. \ -Use Cache-Aside as the default pattern. Write-Through ensures consistency but adds latency. \ -Redis Cluster for horizontal scaling, Redis Sentinel for high availability. Bloom filters prevent \ -cache penetration. Consider local caches before distributed caches to reduce network overhead. - -SECTION 4 - SECURITY BEST PRACTICES: -Implement defense in depth: network segmentation, WAF, rate limiting, input validation. \ -Use OAuth 2.0 with PKCE for public clients. JWT tokens should be short-lived with refresh rotation. \ -Store secrets in Vault, never in code. Mutual TLS for service-to-service. Prepared statements for SQL. \ -Principle of least privilege for IAM. Audit logging for sensitive operations. - -SECTION 5 - OBSERVABILITY: -Three pillars: metrics, logs, traces. Structured logging with correlation IDs. RED metrics for services. \ -Alerts based on SLOs not thresholds. Percentiles not averages for latency. OpenTelemetry for tracing. \ -Error tracking with Sentry. Synthetic monitoring for critical journeys. - -Keep answers concise. Always mention trade-offs.""" - -# Simulated multi-turn conversation -TURNS = [ - ("What database would you recommend for a time-series IoT platform?", - "For a time-series IoT platform, I'd recommend TimescaleDB as the primary choice. It extends PostgreSQL with automatic partitioning by time, native compression achieving 90%+ reduction, and continuous aggregates for real-time rollups. Key advantages: you keep full SQL compatibility, can JOIN with relational data, and benefit from PostgreSQL's ecosystem. For extreme scale (millions of writes/sec), consider InfluxDB or QuestDB which sacrifice SQL flexibility for raw ingestion speed. The trade-off is clear: TimescaleDB gives you versatility, while purpose-built TSDBs give you peak throughput."), - ("How should I handle the ingestion pipeline?", - "For IoT ingestion, implement a buffered pipeline: devices publish to MQTT (Mosquitto/EMQX), a bridge forwards to Kafka for durability and backpressure, then Kafka consumers batch-insert into TimescaleDB. Key design decisions: use Kafka partitioning by device_id for ordering guarantees, set batch sizes of 1000-5000 rows for optimal insert throughput, and implement dead-letter queues for malformed data. Add schema validation at the MQTT bridge level to reject bad payloads early. For backpressure, Kafka naturally handles this - consumers process at their own pace while producers buffer in topics."), - ("What about real-time alerting on the data?", - "Layer your alerting: use TimescaleDB continuous aggregates for threshold-based alerts (e.g., avg temperature > X over 5min windows), and Kafka Streams or Flink for complex event processing (e.g., detecting anomaly patterns across multiple sensors). For the alert pipeline: Kafka topic for raw events, stream processor evaluates rules, alert events go to a separate topic, then a notification service dispatches via PagerDuty/Slack/email. Implement alert deduplication and suppression to avoid alert fatigue. Store alert history in PostgreSQL for audit trails. Use Grafana with TimescaleDB datasource for visualization dashboards."), - ("How do I scale this to millions of devices?", - "Scaling to millions of devices requires horizontal scaling at every layer. MQTT: use EMQX cluster with shared subscriptions, each node handles ~500K concurrent connections. Kafka: partition by device_id hash, scale consumers with consumer groups. TimescaleDB: use distributed hypertables across multiple nodes, partition by both time and device_id. Add a device registry service for metadata. Implement connection pooling with PgBouncer. For cost optimization, tier your storage: hot data (recent 7 days) on SSD-backed TimescaleDB, warm data (30 days) on cheaper storage, cold data archived to S3/Parquet for analytics. Use read replicas for dashboard queries to isolate from write path."), - ("What monitoring should I set up for this infrastructure?", - "Implement observability across all layers using the RED/USE framework. Infrastructure: Prometheus with node_exporter for CPU/memory/disk, kube-state-metrics for Kubernetes. MQTT: monitor connected clients, message rate, subscription count. Kafka: consumer lag (critical - use Burrow), broker throughput, partition skew. TimescaleDB: query latency p95/p99, connection pool utilization, chunk compression ratio, replication lag. Application: request rate, error rate, latency histograms per endpoint. Create four dashboards: system health overview, ingestion pipeline throughput, database performance, and alert system health. Set SLOs: 99.9% ingestion success rate, p99 query latency < 500ms, alert delivery within 60 seconds."), -] - -NEW_QUESTIONS = [ - "Now, what if I need to add machine learning predictions on the incoming data?", - "Should I use GraphQL or REST for the device management API?", -] - - -def wait_for_server(timeout=120): - deadline = time.time() + timeout - while time.time() < deadline: - try: - req = urllib.request.Request(f"{SERVER}/v1/models") - with urllib.request.urlopen(req, timeout=5) as resp: - if resp.status == 200: - data = json.loads(resp.read()) - return data["data"][0]["id"] - except (urllib.error.URLError, ConnectionRefusedError, OSError, KeyError, IndexError): - pass - time.sleep(1) - return None - - -def stream_chat(messages, max_tokens=50, model_name="test"): - payload = json.dumps({ - "model": model_name, - "messages": messages, - "max_tokens": max_tokens, - "temperature": 0.0, - "stream": True, - }).encode() - - req = urllib.request.Request( - f"{SERVER}/v1/chat/completions", - data=payload, - headers={"Content-Type": "application/json"}, - ) - - t0 = time.perf_counter() - first_token_time = None - tokens = [] - - with urllib.request.urlopen(req, timeout=300) as resp: - buffer = b"" - while True: - chunk = resp.read(1) - if not chunk: - break - buffer += chunk - while b"\n" in buffer: - line, buffer = buffer.split(b"\n", 1) - line = line.strip() - if not line or line == b"data: [DONE]": - continue - if line.startswith(b"data: "): - try: - data = json.loads(line[6:]) - choices = data.get("choices", []) - if not choices: - continue - delta = choices[0].get("delta", {}) - content = delta.get("content", "") - if content and first_token_time is None: - first_token_time = time.perf_counter() - if content: - tokens.append(content) - except (json.JSONDecodeError, KeyError, IndexError): - pass - - total_time = time.perf_counter() - t0 - ttft = (first_token_time - t0) if first_token_time else total_time - return ttft * 1000, "".join(tokens), len(tokens) - - -def count_message_tokens(messages): - """Rough token estimate: ~1.3 tokens per word.""" - total_words = sum(len(m["content"].split()) for m in messages) - return int(total_words * 1.3) - - -def bench_model(model_dir, label): - print(f"\n{'='*70}") - print(f"Model: {label}") - print(f"{'='*70}") - - env = {**os.environ, "HIGGS_ENABLE_THINKING": "0"} - proc = subprocess.Popen( - [HIGGS, "serve", "--model", model_dir, "--port", "8080"], - env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, - preexec_fn=os.setsid, - ) - - try: - print("Waiting for server...", end="", flush=True) - model_name = wait_for_server() - if not model_name: - print(" TIMEOUT") - return - print(f" ready") - - # Warmup - stream_chat([{"role": "system", "content": "Be brief."}, - {"role": "user", "content": "Hi."}], - max_tokens=5, model_name=model_name) - - # Build conversation history incrementally - # First: cache miss with system prompt + first question - messages = [{"role": "system", "content": SYSTEM_PROMPT}] - - print(f"\n {'Turn':<6} {'Ctx tokens':>10} {'TTFT (ms)':>10} {'vs Miss':>8} Question") - print(f" {'-'*64}") - - miss_ttft = None - - for turn_idx, (question, fake_answer) in enumerate(TURNS): - messages.append({"role": "user", "content": question}) - est_tokens = count_message_tokens(messages) - - ttft, output, ntok = stream_chat(messages, max_tokens=60, model_name=model_name) - - if turn_idx == 0: - miss_ttft = ttft - tag = "MISS" - ratio = "1.00x" - else: - tag = "HIT" - ratio = f"{miss_ttft / ttft:.1f}x" if ttft > 0 else "inf" - - print(f" {turn_idx+1:<6} {est_tokens:>10} {ttft:>10.0f} {ratio:>8} {question[:50]}...") - - # Add the fake assistant response to history for next turn - messages.append({"role": "assistant", "content": fake_answer}) - - # Now send 2 new questions with the FULL history — cache should hit on prefix - print(f"\n --- New questions with full history ({len(TURNS)} prior turns) ---") - for q in NEW_QUESTIONS: - messages_copy = list(messages) # copy to not pollute - messages_copy.append({"role": "user", "content": q}) - est_tokens = count_message_tokens(messages_copy) - - ttft, output, ntok = stream_chat(messages_copy, max_tokens=60, model_name=model_name) - ratio = f"{miss_ttft / ttft:.1f}x" if ttft > 0 else "inf" - print(f" {'new':<6} {est_tokens:>10} {ttft:>10.0f} {ratio:>8} {q[:50]}...") - - finally: - try: - os.killpg(os.getpgid(proc.pid), signal.SIGTERM) - proc.wait(timeout=10) - except (ProcessLookupError, subprocess.TimeoutExpired): - try: - os.killpg(os.getpgid(proc.pid), signal.SIGKILL) - proc.wait(timeout=5) - except (ProcessLookupError, subprocess.TimeoutExpired): - pass - time.sleep(2) - - -def main(): - models = [ - (os.path.expanduser("~/.cache/lm-studio/models/NexVeridian/Qwen3.5-35B-A3B-3bit"), - "Qwen3.5-35B-A3B-3bit (MoE)"), - (os.path.expanduser("~/.cache/lm-studio/models/mlx-community/DeepSeek-Coder-V2-Lite-Instruct-4bit-mlx"), - "DeepSeek-V2-Lite-4bit (MoE)"), - ] - models = [(d, l) for d, l in models if os.path.isdir(d)] - - print("Prefix cache degradation test: TTFT across conversation turns") - print(f"System prompt: ~{len(SYSTEM_PROMPT.split())} words") - print(f"Simulated turns: {len(TURNS)} + {len(NEW_QUESTIONS)} new questions") - - for model_dir, label in models: - bench_model(model_dir, label) - - -if __name__ == "__main__": - main() diff --git a/benchmarks/bench_tq_configs.py b/benchmarks/bench_tq_configs.py deleted file mode 100644 index d4c871c8..00000000 --- a/benchmarks/bench_tq_configs.py +++ /dev/null @@ -1,302 +0,0 @@ -#!/usr/bin/env python3 -"""Benchmark prefill, decode, and quality across TurboQuant configurations. - -Runs each config sequentially (one server at a time) to stay safe on 32GB. - -Usage: - python3 benchmarks/bench_tq_configs.py -""" - -import argparse -import json -import os -import signal -import subprocess -import sys -import time -import urllib.request - -HIGGS_BIN = os.environ.get("HIGGS_BIN", "./target/release/higgs") -WORD = "the quick brown fox jumps over the lazy dog " # ~10 tokens -PORT = 8097 - - -def api(port, endpoint, body, timeout=300): - data = json.dumps(body).encode() - req = urllib.request.Request( - f"http://127.0.0.1:{port}/v1/{endpoint}", - data=data, - headers={"Content-Type": "application/json"}, - ) - t0 = time.perf_counter() - with urllib.request.urlopen(req, timeout=timeout) as resp: - result = json.loads(resp.read()) - elapsed = time.perf_counter() - t0 - return result, elapsed - - -def start_server(model_path, port, extra_args=None): - cmd = [HIGGS_BIN, "serve", "--model", model_path, "--port", str(port)] - if extra_args: - cmd += extra_args - env = {**os.environ, "HIGGS_ENABLE_THINKING": "0", "HIGGS_NO_CONFIG": "1"} - proc = subprocess.Popen( - cmd, - env=env, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - for _ in range(90): - try: - urllib.request.urlopen(f"http://127.0.0.1:{port}/v1/models", timeout=2) - return proc - except Exception: - time.sleep(1) - proc.kill() - proc.wait(timeout=5) - raise RuntimeError(f"Server on port {port} failed to start") - - -def stop_server(proc): - if proc and proc.poll() is None: - proc.send_signal(signal.SIGTERM) - try: - proc.wait(timeout=5) - except subprocess.TimeoutExpired: - proc.kill() - - -def get_model_id(port): - req = urllib.request.Request(f"http://127.0.0.1:{port}/v1/models") - with urllib.request.urlopen(req, timeout=5) as resp: - data = json.loads(resp.read()) - return data["data"][0]["id"] - - -def chat(port, model, messages, max_tokens=256, temperature=0): - body = { - "model": model, - "messages": messages, - "max_tokens": max_tokens, - "temperature": temperature, - } - result, elapsed = api(port, "chat/completions", body) - usage = result.get("usage", {}) - choice = result.get("choices", [{}])[0] - content = choice.get("message", {}).get("content", "") - return content, usage, elapsed - - -# ---- Benchmarks per config ---- - -def measure_ttft(port, model, prompt): - """TTFT ≈ time for max_tokens=1 request.""" - pfx = f"[ttft{time.time():.6f}] " - _, usage, elapsed = chat(port, model, - [{"role": "user", "content": pfx + prompt}], max_tokens=1) - return elapsed, usage.get("prompt_tokens", 0) - - -def measure_decode(port, model, prompt, gen_tokens=128): - """Decode speed = (gen_tokens - 1) / (full_time - ttft).""" - # TTFT - ttft_s, ptoks = measure_ttft(port, model, prompt) - - # Full generation - pfx2 = f"[dec{time.time():.6f}] Summarize: " - _, usage2, elapsed2 = chat(port, model, - [{"role": "user", "content": pfx2 + prompt}], max_tokens=gen_tokens) - ctoks = usage2.get("completion_tokens", 0) - - if ctoks > 1: - decode_time = max(elapsed2 - ttft_s, 0.01) - tps = (ctoks - 1) / decode_time - else: - tps = 0 - - return {"ptoks": ptoks, "ctoks": ctoks, "ttft_ms": ttft_s * 1000, - "tps": tps, "total_s": elapsed2} - - -def run_quality_prompts(port, model): - """Generate 10 prompts, return list of outputs.""" - prompts = [ - "Explain how a hash table works step by step.", - "What causes the seasons on Earth?", - "Write a Python function to find the longest common subsequence.", - "Describe the process of photosynthesis in detail.", - "What is the difference between TCP and UDP?", - "Explain the theory of general relativity in simple terms.", - "How does a neural network learn through backpropagation?", - "What are the main causes of climate change?", - "Describe the water cycle and its importance.", - "How does public key cryptography work?", - ] - outputs = [] - for prompt in prompts: - content, _, _ = chat(port, model, - [{"role": "user", "content": prompt}], max_tokens=64) - outputs.append(content) - return prompts, outputs - - -def bench_config(model_path, label, extra_args=None): - """Run full benchmark suite for one configuration.""" - print(f"\n{'='*60}") - print(f"CONFIG: {label}") - if extra_args: - print(f" Args: {' '.join(extra_args)}") - print(f"{'='*60}") - - results = {"label": label, "args": extra_args or []} - proc = start_server(model_path, PORT, extra_args) - try: - model_id = get_model_id(PORT) - print(f" Server ready: {model_id}") - - # Warmup - chat(PORT, model_id, [{"role": "user", "content": "hi"}], max_tokens=2) - - # Prefill + Decode at various context sizes - context_sizes = [100, 1000, 4000] - print(f"\n {'Context':>8} | {'TTFT ms':>10} | {'Decode t/s':>12} | {'Gen toks':>10}") - print(f" {'-'*50}") - - sweep = [] - for ctx in context_sizes: - prompt = WORD * max(1, ctx // 10) - try: - r = measure_decode(PORT, model_id, prompt, gen_tokens=64) - print(f" {r['ptoks']:>8} | {r['ttft_ms']:>7.0f} ms | {r['tps']:>9.1f} t/s | {r['ctoks']:>10}") - sweep.append(r) - except Exception as e: - print(f" {ctx:>8} | {'FAILED':>40} | {str(e)[:30]}") - sweep.append({"ptoks": ctx, "tps": 0, "error": str(e)[:50]}) - - results["sweep"] = sweep - - # Quality: 10 prompts - print("\n Generating quality prompts...") - _prompts, outputs = run_quality_prompts(PORT, model_id) - results["outputs"] = outputs - print(f" Got {len(outputs)} outputs, avg {sum(len(o) for o in outputs)//len(outputs)} chars") - finally: - stop_server(proc) - print(" Server stopped.") - time.sleep(3) - - return results - - -def jaccard(a, b): - wa = set(a.lower().split()) - wb = set(b.lower().split()) - return len(wa & wb) / len(wa | wb) if (wa or wb) else 0 - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("model_path") - args = parser.parse_args() - - print("=" * 60) - print("TURBOQUANT MULTI-CONFIG BENCHMARK") - print(f"Model: {args.model_path}") - print("=" * 60) - - # Define configurations - configs = [ - ("baseline (no TQ)", None), - ("TQ default (bits=3, norm ON)", [ - "--kv-cache", "turboquant", "--kv-bits", "3"]), - ("TQ no-norm-correction", [ - "--kv-cache", "turboquant", "--kv-bits", "3", - "--kv-no-norm-correction"]), - ("TQ asymmetric (K=4, V=3)", [ - "--kv-cache", "turboquant", "--kv-bits", "3", - "--kv-key-bits", "4", "--kv-value-bits", "3"]), - ("TQ layer-adaptive (8 dense)", [ - "--kv-cache", "turboquant", "--kv-bits", "3", - "--kv-adaptive-dense-layers", "8"]), - ] - - all_results = [] - for label, extra_args in configs: - try: - r = bench_config(args.model_path, label, extra_args) - all_results.append(r) - except Exception as e: - print(f"\n CONFIG FAILED: {e}") - all_results.append({"label": label, "error": str(e)}) - - # ---- Summary ---- - print(f"\n\n{'='*60}") - print("SUMMARY") - print(f"{'='*60}") - - # Decode speed comparison - print(f"\n Decode speed (tok/s):") - header = f" {'Context':>8}" - for r in all_results: - name = r["label"][:18] - header += f" | {name:>18}" - print(header) - print(f" {'-'*(10 + 21 * len(all_results))}") - - if all_results and "sweep" in all_results[0]: - for i in range(len(all_results[0].get("sweep", []))): - row = f" {all_results[0]['sweep'][i].get('ptoks', '?'):>8}" - for r in all_results: - sweep = r.get("sweep", []) - if i < len(sweep): - tps = sweep[i].get("tps", 0) - if tps > 0: - row += f" | {tps:>15.1f} t/s" - else: - row += f" | {'FAIL':>18}" - else: - row += f" | {'N/A':>18}" - print(row) - - # TTFT comparison - print(f"\n TTFT (ms):") - header = f" {'Context':>8}" - for r in all_results: - name = r["label"][:18] - header += f" | {name:>18}" - print(header) - print(f" {'-'*(10 + 21 * len(all_results))}") - - if all_results and "sweep" in all_results[0]: - for i in range(len(all_results[0].get("sweep", []))): - row = f" {all_results[0]['sweep'][i].get('ptoks', '?'):>8}" - for r in all_results: - sweep = r.get("sweep", []) - if i < len(sweep): - ttft = sweep[i].get("ttft_ms", 0) - if ttft > 0: - row += f" | {ttft:>15.0f} ms" - else: - row += f" | {'FAIL':>18}" - else: - row += f" | {'N/A':>18}" - print(row) - - # Quality comparison (all configs vs baseline) - baseline_outputs = all_results[0].get("outputs", []) if all_results else [] - if baseline_outputs: - print(f"\n Quality (Jaccard vs baseline):") - for r in all_results[1:]: - outputs = r.get("outputs", []) - if outputs and len(outputs) == len(baseline_outputs): - scores = [jaccard(a, b) for a, b in zip(baseline_outputs, outputs)] - avg = sum(scores) / len(scores) - print(f" {r['label']:<35} avg={avg:.3f} min={min(scores):.3f}") - else: - print(f" {r['label']:<35} N/A") - - print() - - -if __name__ == "__main__": - main() diff --git a/benchmarks/bench_ttft_quick.py b/benchmarks/bench_ttft_quick.py deleted file mode 100644 index cf464bce..00000000 --- a/benchmarks/bench_ttft_quick.py +++ /dev/null @@ -1,64 +0,0 @@ -"""Quick TTFT benchmark at different prompt lengths.""" -import statistics -import time, json, urllib.request - -BASE = "http://localhost:9999/v1" -MODEL = "DeepSeek-Coder-V2-Lite-Instruct-4bit-mlx" - -# Generate prompts of increasing length -WORD = "the quick brown fox jumps over the lazy dog " # ~10 tokens -WARMUP = 1 -ITERS = 3 - -def ttft(prompt, max_tokens=1): - body = json.dumps({ - "model": MODEL, - "messages": [{"role": "user", "content": prompt}], - "max_tokens": max_tokens, - "temperature": 0, - }).encode() - req = urllib.request.Request( - f"{BASE}/chat/completions", - data=body, - headers={"Content-Type": "application/json"}, - ) - t0 = time.perf_counter() - with urllib.request.urlopen(req, timeout=120) as resp: - d = json.loads(resp.read()) - elapsed = time.perf_counter() - t0 - usage = d.get("usage", {}) - prompt_tokens = usage.get("prompt_tokens", "?") - return elapsed, prompt_tokens - -def bench_ttft(label, prompt): - # Warmup with a unique prompt to avoid prefix cache - for i in range(WARMUP): - ttft(f"[warmup {i}] {prompt}") - - times = [] - ptoks = None - for i in range(ITERS): - # Unique prefix each iter to defeat prefix cache - t, pt = ttft(f"[run {i} {time.time()}] {prompt}") - times.append(t) - ptoks = pt - - med = statistics.median(times) - return med, ptoks - -print(f"TTFT benchmark — {MODEL}") -print(f" {WARMUP} warmup, {ITERS} iters, reporting median\n") - -prompt_sizes = [ - ("short", WORD * 3), - ("medium", WORD * 50), - ("long", WORD * 200), - ("very_long", WORD * 500), -] - -print(f"{'Label':>12} | {'Tokens':>8} | {'TTFT':>10}") -print("-" * 40) - -for label, prompt in prompt_sizes: - med, ptoks = bench_ttft(label, prompt) - print(f"{label:>12} | {ptoks:>8} | {med*1000:>7.0f} ms") diff --git a/benchmarks/models.toml b/benchmarks/models.toml index aaa9baf6..497b2ef7 100644 --- a/benchmarks/models.toml +++ b/benchmarks/models.toml @@ -44,3 +44,12 @@ quantization = "3bit" approx_size_gb = 14.0 context = 32768 tags = ["large", "moe"] + +[[models]] +key = "qwen3.6-27B-mtp-8bit" +label = "Qwen3.6-27B MTP 8-bit" +path = "trevon/Qwen3.6-27B-mtp" +quantization = "8bit" +approx_size_gb = 27.0 +context = 32768 +tags = ["large", "dense", "mtp", "qwen3next"] diff --git a/benchmarks/test_bench_mlx_tuning.py b/benchmarks/test_bench_mlx_tuning.py deleted file mode 100644 index a8353943..00000000 --- a/benchmarks/test_bench_mlx_tuning.py +++ /dev/null @@ -1,121 +0,0 @@ -from __future__ import annotations - -import importlib.util -import sys -from pathlib import Path -import unittest - - -SCRIPT_PATH = Path(__file__).with_name("bench_mlx_tuning.py") -SPEC = importlib.util.spec_from_file_location("bench_mlx_tuning", SCRIPT_PATH) -assert SPEC is not None and SPEC.loader is not None -bench_mlx_tuning = importlib.util.module_from_spec(SPEC) -sys.modules["bench_mlx_tuning"] = bench_mlx_tuning -SPEC.loader.exec_module(bench_mlx_tuning) - - -def result_fixture( - weighted_ttft_ms: float, - weighted_decode_tps: float, - qa_accuracy: float, - long_accuracy: float, - structured_accuracy: float, - cache_passed: bool, - cache_speedup: float, -) -> dict: - return { - "prompt_sweep": { - "weighted_ttft_ms": weighted_ttft_ms, - "weighted_decode_tps": weighted_decode_tps, - }, - "qa": {"accuracy": qa_accuracy}, - "long_context": {"accuracy": long_accuracy}, - "structured_output": {"accuracy": structured_accuracy}, - "prefix_cache": { - "passed": cache_passed, - "speedup": cache_speedup, - "accuracy": 1.0 if cache_passed else 0.0, - }, - } - - -class BenchMlxTuningTests(unittest.TestCase): - def test_normalize_text(self) -> None: - text = " Leading\tand\ntrailing MIXED Case " - normalized = bench_mlx_tuning.normalize_text(text) - self.assertEqual(normalized, "leading and trailing mixed case") - - def test_cache_speedup_is_capped(self) -> None: - self.assertEqual(bench_mlx_tuning.clamp_cache_speedup(96.0), 32.0) - self.assertEqual(bench_mlx_tuning.clamp_cache_speedup(12.0), 12.0) - - def test_compute_iteration_score_applies_formula(self) -> None: - result = result_fixture( - weighted_ttft_ms=100.0, - weighted_decode_tps=100.0, - qa_accuracy=1.0, - long_accuracy=0.5, - structured_accuracy=0.0, - cache_passed=True, - cache_speedup=96.0, - ) - - score = bench_mlx_tuning.compute_iteration_score( - result, - best_ttft=100.0, - best_decode=50.0, - best_cache=16.0, - ) - expected_accuracy = (1.0 * 0.45) + (0.5 * 0.25) + (0.0 * 0.15) + (1.0 * 0.15) - expected_speed = (100.0 / 100.0) * 0.55 + (100.0 / 50.0) * 0.45 - expected_cache = (32.0 / 16.0) - expected_composite = 100.0 * ( - (expected_accuracy * 0.45) + (expected_speed * 0.45) + (expected_cache * 0.10) - ) - - self.assertAlmostEqual(score["accuracy"], expected_accuracy, places=9) - self.assertAlmostEqual(score["speed"], expected_speed, places=9) - self.assertAlmostEqual(score["cache"], expected_cache, places=9) - self.assertAlmostEqual(score["composite"], expected_composite, places=9) - - def test_score_results_marks_all_results(self) -> None: - results = [ - { - "prompt_sweep": { - "weighted_ttft_ms": 200.0, - "weighted_decode_tps": 80.0, - }, - "qa": {"accuracy": 0.4}, - "long_context": {"accuracy": 0.0}, - "structured_output": {"accuracy": 1.0}, - "prefix_cache": {"passed": False, "accuracy": 0.0, "speedup": 12.0}, - }, - { - "prompt_sweep": { - "weighted_ttft_ms": 100.0, - "weighted_decode_tps": 160.0, - }, - "qa": {"accuracy": 0.8}, - "long_context": {"accuracy": 1.0}, - "structured_output": {"accuracy": 1.0}, - "prefix_cache": {"passed": True, "accuracy": 1.0, "speedup": 96.0}, - }, - ] - - bench_mlx_tuning.score_results(results) - self.assertIn("score", results[0]) - self.assertIn("score", results[1]) - self.assertNotEqual(results[0]["score"]["composite"], results[1]["score"]["composite"]) - - def test_rank_results_by_score(self) -> None: - results = [ - {"score": {"composite": 33.0}}, - {"score": {"composite": 72.0}}, - {"score": {"composite": 51.0}}, - ] - ranked = bench_mlx_tuning.rank_results_by_score(results) - self.assertEqual(ranked[0]["score"]["composite"], 72.0) - - -if __name__ == "__main__": - unittest.main() diff --git a/crates/higgs-bench/Cargo.toml b/crates/higgs-bench/Cargo.toml index 073256fd..7caff309 100644 --- a/crates/higgs-bench/Cargo.toml +++ b/crates/higgs-bench/Cargo.toml @@ -30,6 +30,10 @@ built = { version = "0.8", features = ["git2"] } name = "bench_decode" path = "src/bin/bench_decode.rs" +[[bin]] +name = "bench_speculative" +path = "src/bin/bench_speculative.rs" + [[bin]] name = "bench_summarize" path = "src/bin/bench_summarize.rs" diff --git a/crates/higgs-bench/src/bin/bench_decode.rs b/crates/higgs-bench/src/bin/bench_decode.rs index 67b87f92..b9d97606 100644 --- a/crates/higgs-bench/src/bin/bench_decode.rs +++ b/crates/higgs-bench/src/bin/bench_decode.rs @@ -23,7 +23,7 @@ use clap::Parser; use futures::StreamExt; use higgs_bench::{ BenchOutput, ModelInfo, OutputFormat, RunMetadata, default_manifest_path, format_json, - format_markdown, models, persist_result, server, stats, + format_markdown, models, path_for_output, persist_result, public_model_ref, server, stats, }; use serde::Serialize; @@ -159,9 +159,10 @@ async fn run(args: Args) -> Result<()> { let manifest_path = args.manifest.clone().unwrap_or_else(default_manifest_path); let model = models::find_by_key(&manifest_path, &args.model)?; + let public_model_path = public_model_ref(&model.path, ""); metadata.model = Some(ModelInfo { key: model.key.clone(), - path: model.path.clone(), + path: public_model_path.clone(), quantization: model.quantization.clone(), approx_size_gb: model.approx_size_gb, }); @@ -218,7 +219,7 @@ async fn run(args: Args) -> Result<()> { host: args.host.clone(), port: args.port, model_key: model.key.clone(), - model_path: model.path.clone(), + model_path: public_model_path, max_tokens: args.max_tokens, warmup: args.warmup, trials: args.trials, @@ -237,7 +238,7 @@ async fn run(args: Args) -> Result<()> { }; let path = persist_result(&output)?; - eprintln!("[persisted] {}", path.display()); + eprintln!("[persisted] {}", path_for_output(&path)); let rendered = match args.format { OutputFormat::Json => format_json(&output)?, diff --git a/crates/higgs-bench/src/bin/bench_speculative.rs b/crates/higgs-bench/src/bin/bench_speculative.rs new file mode 100644 index 00000000..6e9ba477 --- /dev/null +++ b/crates/higgs-bench/src/bin/bench_speculative.rs @@ -0,0 +1,534 @@ +#![allow( + clippy::expect_used, + clippy::unwrap_used, + clippy::print_stdout, + clippy::print_stderr, + clippy::as_conversions, + clippy::cast_precision_loss, + clippy::cast_possible_truncation, + clippy::cast_sign_loss +)] +//! `bench_speculative` starts a fresh higgs server per trial and compares +//! greedy decode against MTP and prompt-lookup speculative modes. + +use std::fs; +use std::path::{Path, PathBuf}; +use std::process::{Child, Command, ExitCode, Stdio}; +use std::time::{Duration, Instant}; + +use anyhow::{Context, Result}; +use clap::Parser; +use higgs_bench::{ + BenchOutput, ModelInfo, OutputFormat, RunMetadata, default_manifest_path, format_json, + format_markdown, models, path_for_output, persist_result, public_model_ref, results_dir, + server, speculative, stats, +}; +use serde::Serialize; + +const DEFAULT_PROMPT: &str = "Write a concise technical explanation of speculative decoding for local LLM inference. Include acceptance rate, verification cost, and why greedy decode is the easiest correctness target."; + +#[derive(Debug, Parser)] +#[command( + name = "bench_speculative", + about = "Sweep Higgs speculative decode modes with a fresh server per trial", + version +)] +struct Args { + /// Model key from `benchmarks/models.toml`. + #[arg(long, conflicts_with = "model_path")] + model: Option, + + /// Direct model path or Hugging Face repo ID passed to `higgs serve --model`. + #[arg(long, conflicts_with = "model")] + model_path: Option, + + /// Request model name. Defaults to the HF repo ID or local path basename. + #[arg(long)] + model_name: Option, + + /// Override the manifest path. + #[arg(long)] + manifest: Option, + + /// Higgs binary to launch for each trial. + #[arg(long, default_value = "./target/release/higgs")] + higgs_bin: PathBuf, + + #[arg(long, default_value = "127.0.0.1")] + host: String, + + #[arg(long, default_value_t = 8098)] + port: u16, + + /// Maximum completion tokens per request. + #[arg(long, default_value_t = 192)] + max_tokens: u32, + + /// Number of measured repeats per trial mode. + #[arg(long, default_value_t = 1)] + repeats: u32, + + /// Comma-separated trial modes: `baseline`, `mtp_default`, + /// `prompt_lookup`, `prompt_lookup_unchecked`, or numeric MTP draft + /// depths such as `1,2,3`. + #[arg(long, default_value = "baseline,1,2,3")] + trials: String, + + #[arg(long)] + prompt: Option, + + #[arg(long, default_value_t = 300)] + startup_timeout_secs: u64, + + #[arg(long, default_value_t = 600)] + request_timeout_secs: u64, + + #[arg(long, value_enum, default_value_t = OutputFormat::Json)] + format: OutputFormat, +} + +#[derive(Debug, Serialize)] +struct Params { + host: String, + port: u16, + model_key: Option, + model_ref: String, + request_model: String, + max_tokens: u32, + repeats: u32, + trials: Vec, + prompt: String, + startup_timeout_secs: u64, + request_timeout_secs: u64, +} + +#[derive(Debug, Serialize, Clone)] +struct TrialRun { + elapsed_s: f64, + completion_tokens: u32, + tok_s: f64, + content_prefix: String, + telemetry: String, +} + +#[derive(Debug, Serialize)] +struct TrialSummary { + label: String, + runs: Vec, + elapsed_s_mean: f64, + tok_s_mean: f64, + tok_s_median: f64, + tok_s_p95: f64, + tok_s_stdev: f64, + speedup_vs_baseline: Option, +} + +#[derive(Debug, Serialize)] +struct Results { + trials: Vec, +} + +struct ResolvedModel { + key: Option, + serve_path: String, + public_ref: String, + request_model: String, + metadata: ModelInfo, +} + +fn main() -> ExitCode { + let args = Args::parse(); + let runtime = match tokio::runtime::Runtime::new() { + Ok(rt) => rt, + Err(e) => { + eprintln!("error: failed to start tokio runtime: {e}"); + return ExitCode::from(2); + } + }; + + match runtime.block_on(run(args)) { + Ok(()) => ExitCode::SUCCESS, + Err(e) => { + eprintln!("error: {e:#}"); + ExitCode::from(1) + } + } +} + +async fn run(args: Args) -> Result<()> { + if args.repeats == 0 { + anyhow::bail!("--repeats must be >= 1"); + } + + let mut metadata = RunMetadata::capture("bench_speculative"); + let started = Instant::now(); + + let manifest_path = args.manifest.clone().unwrap_or_else(default_manifest_path); + let model = resolve_model(&args, &manifest_path)?; + metadata.model = Some(model.metadata.clone()); + + let trial_specs = speculative::parse_trial_specs(&args.trials)?; + let prompt = args + .prompt + .clone() + .unwrap_or_else(|| DEFAULT_PROMPT.to_owned()); + let base_url = format!("http://{}:{}", args.host, args.port); + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(args.request_timeout_secs)) + .build() + .context("build request client")?; + + let mut summaries = Vec::with_capacity(trial_specs.len()); + let mut baseline_tok_s: Option = None; + + for spec in &trial_specs { + let mut runs = Vec::with_capacity(args.repeats as usize); + for repeat_idx in 0..args.repeats { + eprintln!( + "[{} repeat {}/{}]", + spec.label, + repeat_idx + 1, + args.repeats + ); + let run = + run_trial(&args, spec, repeat_idx, &client, &base_url, &model, &prompt).await?; + runs.push(run); + } + + let tok_s: Vec = runs.iter().map(|run| run.tok_s).collect(); + let elapsed: Vec = runs.iter().map(|run| run.elapsed_s).collect(); + let tok_s_mean = stats::mean(&tok_s); + if spec.label == "baseline_mtp_off" { + baseline_tok_s = Some(tok_s_mean); + } + let speedup_vs_baseline = baseline_tok_s + .filter(|baseline| *baseline > 0.0) + .map(|baseline| tok_s_mean / baseline); + + summaries.push(TrialSummary { + label: spec.label.clone(), + runs, + elapsed_s_mean: stats::mean(&elapsed), + tok_s_mean, + tok_s_median: stats::median(&tok_s), + tok_s_p95: stats::p95(&tok_s), + tok_s_stdev: stats::stdev(&tok_s), + speedup_vs_baseline, + }); + } + + metadata.duration_ms = started.elapsed().as_millis() as u64; + + let params = Params { + host: args.host, + port: args.port, + model_key: model.key, + model_ref: model.public_ref, + request_model: model.request_model, + max_tokens: args.max_tokens, + repeats: args.repeats, + trials: trial_specs, + prompt, + startup_timeout_secs: args.startup_timeout_secs, + request_timeout_secs: args.request_timeout_secs, + }; + + let output = BenchOutput { + metadata, + params, + results: Results { trials: summaries }, + }; + + let path = persist_result(&output)?; + eprintln!("[persisted] {}", path_for_output(&path)); + + let rendered = match args.format { + OutputFormat::Json => format_json(&output)?, + OutputFormat::Markdown => format_markdown(&output)?, + }; + println!("{rendered}"); + + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +async fn run_trial( + args: &Args, + spec: &speculative::TrialSpec, + repeat_idx: u32, + client: &reqwest::Client, + base_url: &str, + model: &ResolvedModel, + prompt: &str, +) -> Result { + let log_path = log_path(&spec.label, repeat_idx)?; + let mut child = start_higgs_server(args, spec, &model.serve_path, &log_path) + .with_context(|| format!("start higgs server for trial {}", spec.label))?; + + let result = async { + server::wait_until_ready(base_url, Duration::from_secs(args.startup_timeout_secs)) + .await + .with_context(|| format!("higgs server not ready at {base_url}"))?; + request_completion( + client, + base_url, + &model.request_model, + prompt, + args.max_tokens, + ) + .await + } + .await; + + stop_child(&mut child); + let telemetry = read_filtered_telemetry(&log_path); + let payload = result?; + let elapsed_s = payload.elapsed.as_secs_f64(); + let tok_s = if elapsed_s > 0.0 { + f64::from(payload.completion_tokens) / elapsed_s + } else { + 0.0 + }; + + Ok(TrialRun { + elapsed_s, + completion_tokens: payload.completion_tokens, + tok_s, + content_prefix: payload.content.chars().take(120).collect(), + telemetry, + }) +} + +struct CompletionPayload { + elapsed: Duration, + completion_tokens: u32, + content: String, +} + +async fn request_completion( + client: &reqwest::Client, + base_url: &str, + model_name: &str, + prompt: &str, + max_tokens: u32, +) -> Result { + let body = serde_json::json!({ + "model": model_name, + "messages": [{"role": "user", "content": prompt}], + "temperature": 0, + "reasoning": { "effort": "none" }, + "max_tokens": max_tokens, + "stream": false, + }); + + let url = format!("{base_url}/v1/chat/completions"); + let started = Instant::now(); + let resp = client + .post(&url) + .json(&body) + .send() + .await + .with_context(|| format!("POST {url}"))?; + let elapsed = started.elapsed(); + + if !resp.status().is_success() { + let status = resp.status(); + let text = resp.text().await.unwrap_or_default(); + anyhow::bail!("{url} returned HTTP {status}: {text}"); + } + + let payload: serde_json::Value = resp + .json() + .await + .context("parse completion response JSON")?; + let completion_tokens = payload + .get("usage") + .and_then(|usage| usage.get("completion_tokens")) + .and_then(serde_json::Value::as_u64) + .and_then(|tokens| u32::try_from(tokens).ok()) + .unwrap_or(0); + let content = payload + .get("choices") + .and_then(|choices| choices.get(0)) + .and_then(|choice| choice.get("message")) + .and_then(|message| message.get("content")) + .and_then(serde_json::Value::as_str) + .unwrap_or_default() + .to_owned(); + + Ok(CompletionPayload { + elapsed, + completion_tokens, + content, + }) +} + +fn resolve_model(args: &Args, manifest_path: &Path) -> Result { + match (&args.model, &args.model_path) { + (Some(key), None) => { + let model = models::find_by_key(manifest_path, key)?; + let request_model = args + .model_name + .clone() + .unwrap_or_else(|| speculative::derive_model_name(&model.path)); + let public_ref = public_model_ref(&model.path, &request_model); + Ok(ResolvedModel { + key: Some(model.key.clone()), + serve_path: model.path.clone(), + public_ref: public_ref.clone(), + request_model, + metadata: ModelInfo { + key: model.key, + path: public_ref, + quantization: model.quantization, + approx_size_gb: model.approx_size_gb, + }, + }) + } + (None, Some(path)) => { + let request_model = args + .model_name + .clone() + .unwrap_or_else(|| speculative::derive_model_name(path)); + let public_ref = public_model_ref(path, &request_model); + Ok(ResolvedModel { + key: None, + serve_path: path.clone(), + public_ref: public_ref.clone(), + request_model, + metadata: ModelInfo { + key: "direct".to_owned(), + path: public_ref, + quantization: "unknown".to_owned(), + approx_size_gb: 0.0, + }, + }) + } + (None, None) => anyhow::bail!("pass either --model or --model-path "), + (Some(_), Some(_)) => anyhow::bail!("pass only one of --model or --model-path"), + } +} + +fn start_higgs_server( + args: &Args, + spec: &speculative::TrialSpec, + model_path: &str, + log_path: &Path, +) -> Result { + if let Some(parent) = log_path.parent() { + fs::create_dir_all(parent) + .with_context(|| format!("create log dir {}", parent.display()))?; + } + let log = fs::File::create(log_path) + .with_context(|| format!("create server log {}", log_path.display()))?; + let log_stderr = log + .try_clone() + .with_context(|| format!("clone server log {}", log_path.display()))?; + + let mut cmd = Command::new(&args.higgs_bin); + cmd.arg("serve") + .arg("--model") + .arg(model_path) + .arg("--host") + .arg(&args.host) + .arg("--port") + .arg(args.port.to_string()) + .arg("--mlx-profile") + .arg("throughput") + .stdout(Stdio::from(log)) + .stderr(Stdio::from(log_stderr)); + + clear_speculative_env(&mut cmd); + for (key, value) in &spec.env { + cmd.env(key, value); + } + + cmd.spawn() + .with_context(|| format!("spawn {}", args.higgs_bin.display())) +} + +fn stop_child(child: &mut Child) { + if matches!(child.try_wait(), Ok(Some(_))) { + return; + } + if let Err(e) = child.kill() { + eprintln!("warning: failed to stop higgs server: {e}"); + } + if let Err(e) = child.wait() { + eprintln!("warning: failed to wait for higgs server exit: {e}"); + } +} + +fn log_path(label: &str, repeat_idx: u32) -> Result { + let logs_dir = results_dir().join("bench_speculative").join("logs"); + fs::create_dir_all(&logs_dir) + .with_context(|| format!("create log dir {}", logs_dir.display()))?; + let ts = chrono::Utc::now().format("%Y%m%dT%H%M%S%.3fZ"); + Ok(logs_dir.join(format!( + "{}__{}__{}.log", + sanitize_filename(label), + repeat_idx, + ts + ))) +} + +fn sanitize_filename(value: &str) -> String { + value + .chars() + .map(|c| { + if c.is_ascii_alphanumeric() || matches!(c, '.' | '-' | '_') { + c + } else { + '_' + } + }) + .collect() +} + +fn clear_speculative_env(cmd: &mut Command) { + for key in SPECULATIVE_ENV_KEYS { + cmd.env_remove(key); + } +} + +const SPECULATIVE_ENV_KEYS: &[&str] = &[ + "HIGGS_MTP", + "HIGGS_MTP_DRAFT_N_MAX", + "HIGGS_PROMPT_LOOKUP", + "HIGGS_PROMPT_LOOKUP_UNCHECKED", + "HIGGS_MTP_PRIME_PREFILL", + "HIGGS_MTP_MIRROR_VERIFY", +]; + +fn read_filtered_telemetry(path: &Path) -> String { + let Ok(body) = fs::read_to_string(path) else { + return String::new(); + }; + body.lines() + .filter(|line| { + line.contains("MTP decode complete") || line.contains("Prompt-lookup decode complete") + }) + .collect::>() + .join("\n") +} + +#[cfg(test)] +mod tests { + use super::{SPECULATIVE_ENV_KEYS, clear_speculative_env}; + use std::process::Command; + + #[test] + fn clear_speculative_env_marks_all_flags_for_removal() { + let mut cmd = Command::new("higgs"); + + clear_speculative_env(&mut cmd); + + for key in SPECULATIVE_ENV_KEYS { + assert!( + cmd.get_envs() + .any(|(name, value)| name == *key && value.is_none()), + "expected {key} to be explicitly removed" + ); + } + } +} diff --git a/crates/higgs-bench/src/lib.rs b/crates/higgs-bench/src/lib.rs index d1e8a288..9d8605ce 100644 --- a/crates/higgs-bench/src/lib.rs +++ b/crates/higgs-bench/src/lib.rs @@ -13,6 +13,7 @@ pub mod models; pub mod server; +pub mod speculative; pub mod stats; use std::fs; @@ -71,7 +72,10 @@ impl HostInfo { sys.refresh_memory(); sys.refresh_cpu_all(); - let hostname = System::host_name().unwrap_or_else(|| "unknown".into()); + let hostname = hostname_for_output( + System::host_name().unwrap_or_else(|| "unknown".into()), + include_hostname(), + ); let os_name = System::name().unwrap_or_else(|| "unknown".into()); let os_version = System::os_version().unwrap_or_else(|| "?".into()); let kernel = System::kernel_version().unwrap_or_else(|| "?".into()); @@ -107,6 +111,137 @@ fn detect_gpu() -> Option { .then(|| "Apple Silicon (MLX)".into()) } +fn include_hostname() -> bool { + matches!( + std::env::var("HIGGS_BENCH_INCLUDE_HOSTNAME") + .ok() + .as_deref(), + Some("1" | "true" | "TRUE" | "yes" | "YES") + ) +} + +fn hostname_for_output(hostname: String, include_hostname: bool) -> String { + if include_hostname { + hostname + } else { + "redacted".to_owned() + } +} + +fn redact_arg_for_output(arg: &str) -> String { + if let Some((prefix, value)) = arg.split_once('=') { + if is_local_path_like(value) { + return format!("{prefix}={}", redacted_local_path_value(value)); + } + } + + if is_local_path_like(arg) { + redacted_local_path_value(arg) + } else { + arg.to_owned() + } +} + +fn is_local_path_like(value: &str) -> bool { + let path = Path::new(value); + path.is_absolute() + || value.starts_with("./") + || value.starts_with("../") + || value.starts_with("~/") + || value.get(1..3) == Some(":\\") +} + +fn redacted_local_path_value(value: &str) -> String { + if looks_like_hf_cache_path(value) { + let model_name = speculative::derive_model_name(value); + if model_name != value && !is_local_path_like(&model_name) { + return model_name; + } + } + + let trimmed = value.trim_end_matches(['/', '\\']); + let basename = Path::new(trimmed) + .file_name() + .and_then(|name| name.to_str()) + .filter(|name| !name.is_empty()) + .unwrap_or("path"); + format!("") +} + +fn looks_like_hf_cache_path(value: &str) -> bool { + (value.contains("/models--") || value.contains("\\models--")) + && (value.contains("/snapshots/") || value.contains("\\snapshots\\")) +} + +/// Returns a model reference suitable for benchmark output. +/// +/// Hugging Face repo IDs are preserved. Local paths are reduced to either the +/// explicit request model name or a derived basename/cache repo so benchmark +/// JSON and markdown do not expose a developer's filesystem layout. +#[must_use] +pub fn public_model_ref(path: &str, request_model: &str) -> String { + if is_local_path_like(path) { + if request_model.is_empty() { + speculative::derive_model_name(path) + } else { + request_model.to_owned() + } + } else { + path.to_owned() + } +} + +/// Formats generated artifact paths without exposing the absolute workspace. +#[must_use] +pub fn path_for_output(path: &Path) -> String { + if let Some(root) = workspace_root() { + if let Ok(relative) = path.strip_prefix(root) { + return relative.display().to_string(); + } + } + redact_arg_for_output(&path.display().to_string()) +} + +#[cfg(test)] +mod tests { + use super::{hostname_for_output, public_model_ref, redact_arg_for_output}; + + #[test] + fn benchmark_metadata_redacts_hostname_by_default() { + assert_eq!( + hostname_for_output("developer-laptop".to_owned(), false), + "redacted" + ); + } + + #[test] + fn benchmark_metadata_redacts_absolute_path_args() { + assert_eq!( + redact_arg_for_output("/Users/alice/models/Qwen3.6-27B-mtp"), + "" + ); + assert_eq!( + redact_arg_for_output("--manifest=/Users/alice/dev/higgs/benchmarks/models.toml"), + "--manifest=" + ); + } + + #[test] + fn model_refs_hide_local_absolute_paths() { + assert_eq!( + public_model_ref( + "/Users/alice/.cache/huggingface/hub/models--org--Qwen3.6-27B-mtp/snapshots/abcdef", + "" + ), + "org/Qwen3.6-27B-mtp" + ); + assert_eq!( + public_model_ref("/Users/alice/models/private-qwen", "local-qwen"), + "local-qwen" + ); + } +} + /// The model under test, captured into bench output for reproducibility. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ModelInfo { @@ -167,7 +302,9 @@ impl RunMetadata { host: HostInfo::capture(), mlx_version: None, model: None, - args: std::env::args().collect(), + args: std::env::args() + .map(|arg| redact_arg_for_output(&arg)) + .collect(), } } } diff --git a/crates/higgs-bench/src/speculative.rs b/crates/higgs-bench/src/speculative.rs new file mode 100644 index 00000000..1af712e1 --- /dev/null +++ b/crates/higgs-bench/src/speculative.rs @@ -0,0 +1,234 @@ +//! Shared helpers for speculative decoding benchmarks. + +use std::collections::BTreeMap; +use std::path::Path; + +use anyhow::{Context, Result}; +use serde::{Deserialize, Serialize}; + +/// One speculative benchmark mode and its environment overrides. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct TrialSpec { + pub label: String, + pub env: BTreeMap, +} + +/// Derives the request model name from either a cache snapshot path or an +/// ordinary model path/repo ID. +#[must_use] +pub fn derive_model_name(model_path: &str) -> String { + if looks_like_huggingface_repo_id(model_path) { + return model_path.to_owned(); + } + + let path = Path::new(model_path.trim_end_matches('/')); + let parts: Vec<_> = path + .components() + .map(|component| component.as_os_str().to_string_lossy().into_owned()) + .collect(); + + if let Some(snapshot_idx) = parts.iter().position(|part| part == "snapshots") { + if snapshot_idx > 0 { + if let Some(cache_name) = parts.get(snapshot_idx - 1) { + if let Some(cache_repo) = cache_name.strip_prefix("models--") { + let model_name = cache_repo.replace("--", "/"); + if !model_name.is_empty() { + return model_name; + } + } + } + } + } + + path.file_name() + .and_then(|name| name.to_str()) + .filter(|name| !name.is_empty()) + .map_or_else(|| model_path.to_owned(), ToOwned::to_owned) +} + +fn looks_like_huggingface_repo_id(model_path: &str) -> bool { + if model_path.starts_with('/') + || model_path.starts_with("./") + || model_path.starts_with("../") + || model_path.starts_with("~/") + || model_path.contains('\\') + { + return false; + } + let mut parts = model_path.split('/'); + let Some(owner) = parts.next() else { + return false; + }; + let Some(name) = parts.next() else { + return false; + }; + if parts.next().is_some() || owner.is_empty() || name.is_empty() { + return false; + } + + !matches!( + owner, + "model" | "models" | "checkpoint" | "checkpoints" | "cache" | "data" | "target" + ) +} + +/// Parses a comma-separated trial list. +/// +/// Supported items are `baseline`, `mtp_default`, numeric MTP draft depths, +/// `prompt_lookup`, and `prompt_lookup_unchecked`. +pub fn parse_trial_specs(input: &str) -> Result> { + let mut trials = Vec::new(); + for raw in input.split(',') { + let trial = raw.trim(); + if trial.is_empty() { + continue; + } + trials.push(parse_trial_spec(trial).with_context(|| format!("parse trial '{trial}'"))?); + } + + if trials.is_empty() { + anyhow::bail!("at least one trial must be specified"); + } + + Ok(trials) +} + +fn parse_trial_spec(trial: &str) -> Result { + match trial { + "baseline" => Ok(trial_spec("baseline_mtp_off", [("HIGGS_MTP", "0")])), + "mtp_default" | "default" => Ok(trial_spec("mtp_default", [("HIGGS_MTP", "1")])), + "prompt_lookup" | "plookup" => Ok(trial_spec( + "prompt_lookup", + [("HIGGS_MTP", "0"), ("HIGGS_PROMPT_LOOKUP", "1")], + )), + "prompt_lookup_unchecked" | "plookup_unchecked" => Ok(trial_spec( + "prompt_lookup_unchecked", + [ + ("HIGGS_MTP", "0"), + ("HIGGS_PROMPT_LOOKUP", "1"), + ("HIGGS_PROMPT_LOOKUP_UNCHECKED", "1"), + ], + )), + raw_depth => { + let depth = raw_depth + .parse::() + .with_context(|| "expected a known trial name or numeric draft depth")?; + if depth == 0 { + anyhow::bail!("MTP draft depth must be >= 1"); + } + Ok(trial_spec( + format!("mtp_draft_{depth}"), + [ + ("HIGGS_MTP", "1".to_owned()), + ("HIGGS_MTP_DRAFT_N_MAX", depth.to_string()), + ], + )) + } + } +} + +fn trial_spec(label: impl Into, env: I) -> TrialSpec +where + K: Into, + V: Into, + I: IntoIterator, +{ + TrialSpec { + label: label.into(), + env: env + .into_iter() + .map(|(key, value)| (key.into(), value.into())) + .collect(), + } +} + +#[cfg(test)] +mod tests { + use super::{TrialSpec, derive_model_name, parse_trial_specs}; + + fn trial_at(trials: &[TrialSpec], index: usize) -> anyhow::Result<&TrialSpec> { + trials + .get(index) + .ok_or_else(|| anyhow::anyhow!("missing trial at index {index}")) + } + + #[test] + fn derive_model_name_from_huggingface_snapshot_path() { + let model = + derive_model_name("/cache/hub/models--org--Qwen3.6-27B-mtp/snapshots/abcdef123456"); + + assert_eq!(model, "org/Qwen3.6-27B-mtp"); + } + + #[test] + fn derive_model_name_keeps_huggingface_repo_id() { + let model = derive_model_name("org/Qwen3.6-27B-mtp"); + + assert_eq!(model, "org/Qwen3.6-27B-mtp"); + } + + #[test] + fn derive_model_name_uses_basename_for_common_relative_model_dirs() { + let model = derive_model_name("models/local-qwen"); + + assert_eq!(model, "local-qwen"); + } + + #[test] + fn parse_trial_specs_sets_expected_env_overrides() -> anyhow::Result<()> { + let trials = + parse_trial_specs("baseline,mtp_default,2,prompt_lookup,prompt_lookup_unchecked")?; + + let baseline = trial_at(&trials, 0)?; + let mtp_default = trial_at(&trials, 1)?; + let mtp_draft_2 = trial_at(&trials, 2)?; + let prompt_lookup = trial_at(&trials, 3)?; + let prompt_lookup_unchecked = trial_at(&trials, 4)?; + + assert_eq!(baseline.label, "baseline_mtp_off"); + assert_eq!(baseline.env.get("HIGGS_MTP").map(String::as_str), Some("0")); + + assert_eq!(mtp_default.label, "mtp_default"); + assert_eq!( + mtp_default.env.get("HIGGS_MTP").map(String::as_str), + Some("1") + ); + + assert_eq!(mtp_draft_2.label, "mtp_draft_2"); + assert_eq!( + mtp_draft_2.env.get("HIGGS_MTP").map(String::as_str), + Some("1") + ); + assert_eq!( + mtp_draft_2 + .env + .get("HIGGS_MTP_DRAFT_N_MAX") + .map(String::as_str), + Some("2") + ); + + assert_eq!(prompt_lookup.label, "prompt_lookup"); + assert_eq!( + prompt_lookup.env.get("HIGGS_MTP").map(String::as_str), + Some("0") + ); + assert_eq!( + prompt_lookup + .env + .get("HIGGS_PROMPT_LOOKUP") + .map(String::as_str), + Some("1") + ); + + assert_eq!(prompt_lookup_unchecked.label, "prompt_lookup_unchecked"); + assert_eq!( + prompt_lookup_unchecked + .env + .get("HIGGS_PROMPT_LOOKUP_UNCHECKED") + .map(String::as_str), + Some("1") + ); + + Ok(()) + } +} diff --git a/crates/higgs-engine/src/mlx_tuning.rs b/crates/higgs-engine/src/mlx_tuning.rs index 88e3a9f5..0717b120 100644 --- a/crates/higgs-engine/src/mlx_tuning.rs +++ b/crates/higgs-engine/src/mlx_tuning.rs @@ -9,6 +9,7 @@ const DEFAULT_CHUNKED_PREFILL_CHUNK_SIZE: i32 = 512; const DEFAULT_PAGED_KV_TARGET_BYTES: usize = 512 * 1024 * 1024; const MIN_PAGED_KV_TARGET_BYTES: usize = 256 * 1024 * 1024; const MAX_PAGED_KV_TARGET_BYTES: usize = 2 * 1024 * 1024 * 1024; +const MAX_MTP_DRAFT_N_MAX: usize = 8; fn parse_positive_chunked_prefill_value(raw: Option<&str>, default: i32) -> i32 { raw.and_then(|s| s.parse::().ok()) @@ -16,6 +17,12 @@ fn parse_positive_chunked_prefill_value(raw: Option<&str>, default: i32) -> i32 .unwrap_or(default) } +fn parse_mtp_draft_n_max(raw: Option<&str>, default: usize) -> usize { + raw.and_then(|s| s.parse::().ok()) + .filter(|v| *v > 0) + .map_or(default, |v| v.min(MAX_MTP_DRAFT_N_MAX)) +} + fn parse_enabled_flag(raw: Option<&str>) -> Option { match raw.map(str::trim).map(str::to_ascii_lowercase).as_deref() { Some("1" | "true" | "on" | "yes") => Some(true), @@ -24,6 +31,13 @@ fn parse_enabled_flag(raw: Option<&str>) -> Option { } } +const fn default_mtp_draft_n_max(size_class: ModelSizeClass) -> usize { + match size_class { + ModelSizeClass::Huge => 2, + ModelSizeClass::Small | ModelSizeClass::Medium | ModelSizeClass::Large => 1, + } +} + /// User-requested MLX profile before auto-resolution. /// /// `RequestedMlxProfile` is used by CLI/config/env precedence and is resolved to @@ -193,6 +207,11 @@ pub struct MlxRuntimeTuning { chunked_prefill_chunk_size: i32, clear_cache_after_prefill: bool, enable_mtp: bool, + /// Maximum MTP draft tokens per speculative cycle. + /// + /// Controlled by `HIGGS_MTP_DRAFT_N_MAX`; defaults to 2 for huge + /// checkpoints and 1 otherwise, clamped to 1..=8. + mtp_draft_n_max: usize, paged_kv_target_bytes: usize, } @@ -224,6 +243,10 @@ impl MlxRuntimeTuning { .unwrap_or(tuning.clear_cache_after_prefill); tuning.enable_mtp = parse_enabled_flag(std::env::var("HIGGS_MTP").ok().as_deref()) .unwrap_or(tuning.enable_mtp); + tuning.mtp_draft_n_max = parse_mtp_draft_n_max( + std::env::var("HIGGS_MTP_DRAFT_N_MAX").ok().as_deref(), + tuning.mtp_draft_n_max, + ); tuning } @@ -239,6 +262,7 @@ impl MlxRuntimeTuning { let (balanced_threshold, balanced_chunk) = 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); match resolved_profile { ResolvedMlxProfile::Baseline => Self { @@ -248,6 +272,7 @@ impl MlxRuntimeTuning { chunked_prefill_chunk_size: DEFAULT_CHUNKED_PREFILL_CHUNK_SIZE, clear_cache_after_prefill: false, enable_mtp: false, + mtp_draft_n_max: 1, paged_kv_target_bytes: DEFAULT_PAGED_KV_TARGET_BYTES, }, ResolvedMlxProfile::Latency => Self { @@ -257,6 +282,7 @@ impl MlxRuntimeTuning { chunked_prefill_chunk_size: balanced_chunk.max(768), clear_cache_after_prefill: false, enable_mtp: true, + mtp_draft_n_max: default_mtp_draft_n_max, paged_kv_target_bytes: clamp_paged_kv_target_bytes( balanced_paged_kv.saturating_mul(9) / 8, ), @@ -268,6 +294,7 @@ impl MlxRuntimeTuning { chunked_prefill_chunk_size: balanced_chunk, clear_cache_after_prefill: false, enable_mtp: true, + mtp_draft_n_max: default_mtp_draft_n_max, paged_kv_target_bytes: balanced_paged_kv, }, ResolvedMlxProfile::Throughput => Self { @@ -277,6 +304,7 @@ impl MlxRuntimeTuning { chunked_prefill_chunk_size: balanced_chunk.max(1024), clear_cache_after_prefill: false, enable_mtp: true, + mtp_draft_n_max: default_mtp_draft_n_max, paged_kv_target_bytes: clamp_paged_kv_target_bytes( balanced_paged_kv.saturating_mul(5) / 4, ), @@ -311,6 +339,13 @@ impl MlxRuntimeTuning { self.enable_mtp } + /// Maximum MTP draft tokens per speculative cycle. + /// + /// Configurable via `HIGGS_MTP_DRAFT_N_MAX`. + pub const fn mtp_draft_n_max(&self) -> usize { + self.mtp_draft_n_max + } + pub const fn paged_kv_target_bytes(&self) -> usize { self.paged_kv_target_bytes } @@ -379,7 +414,7 @@ fn heuristic_paged_kv_target_bytes( size_class: ModelSizeClass, is_moe: bool, ) -> usize { - let Some(max_recommended) = mlx_max_recommended_working_set_size() else { + let Some(max_recommended) = configured_max_working_set_bytes() else { return DEFAULT_PAGED_KV_TARGET_BYTES; }; @@ -408,6 +443,13 @@ fn heuristic_paged_kv_target_bytes( clamp_paged_kv_target_bytes(available / divisor) } +fn configured_max_working_set_bytes() -> Option { + std::env::var("HIGGS_MLX_MAX_WORKING_SET_BYTES") + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|value| *value > 0) +} + fn clamp_paged_kv_target_bytes(bytes: usize) -> usize { bytes.clamp(MIN_PAGED_KV_TARGET_BYTES, MAX_PAGED_KV_TARGET_BYTES) } @@ -467,15 +509,21 @@ fn model_weight_bytes(model_dir: &Path) -> Option { stack.push(path); continue; } - if file_type.is_file() + if (file_type.is_file() || file_type.is_symlink()) && path .extension() .is_some_and(|ext| ext.eq_ignore_ascii_case("safetensors")) { - match dir_entry.metadata() { - Ok(meta) => { + let metadata = if file_type.is_symlink() { + std::fs::metadata(&path) + } else { + dir_entry.metadata() + }; + match metadata { + Ok(meta) if meta.is_file() => { total = total.saturating_add(meta.len()); } + Ok(_) => {} Err(err) => { tracing::warn!( path = %path.display(), @@ -491,33 +539,12 @@ fn model_weight_bytes(model_dir: &Path) -> Option { Some(total).filter(|sum| *sum > 0) } -#[allow(unsafe_code)] -fn mlx_max_recommended_working_set_size() -> Option { - unsafe { - let mut info = mlx_sys::mlx_device_info_new(); - let mut dev = mlx_sys::mlx_device_new(); - mlx_sys::mlx_get_default_device(&raw mut dev); - let mut max_rec = None; - if mlx_sys::mlx_device_info_get(&raw mut info, dev) == 0 { - let mut value: usize = 0; - let key = c"max_recommended_working_set_size"; - if mlx_sys::mlx_device_info_get_size(&raw mut value, info, key.as_ptr()) == 0 - && value > 0 - { - max_rec = Some(value); - } - } - mlx_sys::mlx_device_info_free(info); - mlx_sys::mlx_device_free(dev); - max_rec - } -} - #[cfg(test)] mod tests { use super::{ - MlxRuntimeTuning, ModelMetadata, RequestedMlxProfile, ResolvedMlxProfile, - parse_enabled_flag, parse_positive_chunked_prefill_value, resolve_effective_mlx_profile, + MlxRuntimeTuning, ModelMetadata, ModelSizeClass, RequestedMlxProfile, ResolvedMlxProfile, + default_mtp_draft_n_max, model_weight_bytes, parse_enabled_flag, parse_mtp_draft_n_max, + parse_positive_chunked_prefill_value, resolve_effective_mlx_profile, resolve_profile_from_metadata, resolve_runtime_tuning, }; use std::fs; @@ -661,6 +688,36 @@ mod tests { assert_eq!(parse_positive_chunked_prefill_value(None, 32), 32); } + #[test] + fn test_parse_mtp_draft_n_max_clamps_invalid_and_large_values() { + assert_eq!(parse_mtp_draft_n_max(Some("1"), 3), 1); + assert_eq!(parse_mtp_draft_n_max(Some("3"), 1), 3); + assert_eq!(parse_mtp_draft_n_max(Some("0"), 3), 3); + assert_eq!(parse_mtp_draft_n_max(Some("bad"), 3), 3); + assert_eq!(parse_mtp_draft_n_max(Some("99"), 3), 8); + assert_eq!(parse_mtp_draft_n_max(None, 3), 3); + } + + #[test] + fn test_default_mtp_draft_depth_uses_two_only_for_huge_models() { + assert_eq!(default_mtp_draft_n_max(ModelSizeClass::Small), 1); + assert_eq!(default_mtp_draft_n_max(ModelSizeClass::Medium), 1); + assert_eq!(default_mtp_draft_n_max(ModelSizeClass::Large), 1); + assert_eq!(default_mtp_draft_n_max(ModelSizeClass::Huge), 2); + } + + #[cfg(unix)] + #[test] + fn test_model_weight_bytes_follows_hf_snapshot_symlinks() -> std::io::Result<()> { + let temp = TempDir::new().map_err(std::io::Error::other)?; + let blob = temp.path().join("blob"); + fs::write(&blob, [0u8; 13])?; + std::os::unix::fs::symlink(&blob, temp.path().join("model.safetensors"))?; + + assert_eq!(model_weight_bytes(temp.path()), Some(13)); + Ok(()) + } + #[test] fn test_parse_enabled_flag_ignores_unknown_values() { assert_eq!(parse_enabled_flag(Some("TRUE")), Some(true)); diff --git a/crates/higgs-engine/src/mtp.rs b/crates/higgs-engine/src/mtp.rs index f05049ae..ec8dc1e6 100644 --- a/crates/higgs-engine/src/mtp.rs +++ b/crates/higgs-engine/src/mtp.rs @@ -1,12 +1,17 @@ //! MTP (Multi-Token Prediction) speculative decode. //! -//! Uses the model's built-in MTP head to draft one extra token per cycle, -//! then verifies by processing the confirmed token through the backbone. +//! Uses the model's built-in MTP head to draft tokens, then verifies them by +//! processing the verifier window through the backbone in one batch and rolling +//! back to the committed prefix on rejection. //! //! Expected speedup: ~1.5x on dense models at ~80% acceptance rate. use higgs_models::{AnyCache, AnyModel, MtpCache}; -use mlx_rs::{Array, argmax_axis, ops::indexing::IndexOp, transforms::eval}; +use mlx_rs::{ + Array, argmax_axis, + ops::{self, concatenate_axis, indexing::IndexOp}, + transforms::eval, +}; use crate::error::EngineError; @@ -14,92 +19,588 @@ const fn draft_matches_target(draft_token_id: u32, target_id: u32) -> bool { draft_token_id == target_id } +/// Aggregate MTP decode counters. +/// +/// Tracks per-cycle telemetry for MTP speculative decoding. +#[derive(Debug, Default, Clone)] +pub struct MtpStats { + /// Number of speculative decode cycles executed. + cycles: u32, + /// Total speculative tokens drafted by the MTP head. + drafted: u32, + /// Drafted tokens that matched the backbone verifier. + accepted_drafts: u32, + /// Tokens emitted by MTP cycles, including confirmed tokens and accepted drafts. + emitted: u32, +} + +impl MtpStats { + pub fn record_cycle(&mut self, drafted_count: usize, emitted_count: usize) { + let drafted = u32::try_from(drafted_count).unwrap_or(u32::MAX); + let emitted = u32::try_from(emitted_count).unwrap_or(u32::MAX); + self.cycles = self.cycles.saturating_add(1); + self.drafted = self.drafted.saturating_add(drafted); + self.emitted = self.emitted.saturating_add(emitted); + self.accepted_drafts = self + .accepted_drafts + .saturating_add(emitted.saturating_sub(1).min(drafted)); + } + + pub const fn cycles(&self) -> u32 { + self.cycles + } + + pub const fn drafted(&self) -> u32 { + self.drafted + } + + pub const fn accepted_drafts(&self) -> u32 { + self.accepted_drafts + } + + pub const fn emitted(&self) -> u32 { + self.emitted + } + + #[allow(clippy::cast_precision_loss)] + pub fn acceptance_rate_percent(&self) -> f64 { + if self.drafted == 0 { + 0.0 + } else { + f64::from(self.accepted_drafts) * 100.0 / f64::from(self.drafted) + } + } +} + /// Result of a single MTP speculative decode cycle. pub struct MtpCycleResult { - /// Token IDs accepted this cycle (1 or 2). + /// Token IDs accepted this cycle (the confirmed token plus accepted drafts). pub tokens: Vec, /// Hidden state at the last accepted position (for next MTP draft). pub hidden: Array, /// The next confirmed token to process in the following cycle. pub next_token_id: u32, + /// Number of speculative draft tokens produced this cycle. + pub drafted: usize, + /// Number of speculative draft tokens accepted this cycle. + pub accepted_drafts: usize, +} + +/// Prompt-lookup speculative decode settings. +#[derive(Debug, Clone, Copy)] +pub struct PromptLookupConfig { + pub max_drafts: usize, + pub max_ngram: usize, + pub max_window: usize, +} + +impl Default for PromptLookupConfig { + fn default() -> Self { + Self { + max_drafts: 6, + max_ngram: 8, + max_window: 2048, + } + } +} + +/// Result of one architecture-neutral prompt-lookup speculative cycle. +pub struct PromptLookupCycleResult { + /// Token IDs accepted this cycle (the confirmed token plus accepted drafts). + pub tokens: Vec, + /// The next confirmed token to process in the following cycle. + pub next_token_id: u32, + /// Number of prompt-lookup draft tokens proposed this cycle. + pub drafted: usize, + /// Number of prompt-lookup draft tokens accepted this cycle. + pub accepted_drafts: usize, +} + +fn greedy_token_id(logits: &Array) -> Result { + let token_arr = argmax_axis!(&logits.index((.., -1, ..)), -1).map_err(EngineError::Mlx)?; + eval([&token_arr]).map_err(EngineError::Mlx)?; + Ok(token_arr.item()) +} + +fn greedy_token_ids(logits: &Array) -> Result, EngineError> { + let token_arr = argmax_axis!(logits, -1).map_err(EngineError::Mlx)?; + eval([&token_arr]).map_err(EngineError::Mlx)?; + Ok(token_arr.as_slice::().to_vec()) +} + +fn parse_enabled_flag(raw: Option<&str>) -> Option { + match raw.map(str::trim).map(str::to_ascii_lowercase).as_deref() { + Some("1" | "true" | "on" | "yes") => Some(true), + Some("0" | "false" | "off" | "no") => Some(false), + _ => None, + } +} + +fn mtp_mirror_verify_enabled() -> bool { + parse_enabled_flag(std::env::var("HIGGS_MTP_MIRROR_VERIFY").ok().as_deref()).unwrap_or(false) +} + +fn accepted_draft_prefix_len(drafts: &[u32], verifier_targets: &[u32]) -> usize { + drafts + .iter() + .zip(verifier_targets.iter()) + .take_while(|(draft, target)| draft_matches_target(**draft, **target)) + .count() +} + +fn emitted_tokens(confirmed_token_id: u32, drafts: &[u32], accepted_drafts: usize) -> Vec { + let mut tokens = Vec::with_capacity(accepted_drafts.saturating_add(1)); + tokens.push(confirmed_token_id); + tokens.extend(drafts.iter().take(accepted_drafts).copied()); + tokens +} + +pub fn prompt_lookup_draft( + context: &[u32], + max_drafts: usize, + max_ngram: usize, + max_window: usize, +) -> Vec { + if context.is_empty() || max_drafts == 0 || max_ngram == 0 { + return Vec::new(); + } + + let end = context.len(); + let capped_ngram = max_ngram.min(end); + let search_start = end.saturating_sub(max_window.max(1)); + + for ngram in (1..=capped_ngram).rev() { + let Some(suffix) = context.get(end - ngram..end) else { + continue; + }; + let search_end = end.saturating_sub(ngram); + + for pos in (search_start..search_end).rev() { + let match_end = pos + ngram; + if context.get(pos..match_end) != Some(suffix) { + continue; + } + + let draft_start = match_end; + if draft_start >= end { + continue; + } + let draft_end = draft_start.saturating_add(max_drafts).min(end); + if let Some(draft) = context.get(draft_start..draft_end) { + return draft.to_vec(); + } + } + } + + Vec::new() +} + +/// Run one prompt-lookup speculative decode cycle. +/// +/// This is architecture-neutral: the draft provider only copies tokens from +/// prior prompt/history, and the model verifies `[confirmed + drafts]` in one +/// forward pass using all-position logits. +pub fn prompt_lookup_cycle( + model: &mut AnyModel, + cache: &mut AnyCache, + history_before_confirmed: &[u32], + confirmed_token_id: u32, + config: PromptLookupConfig, +) -> Result { + let mut lookup_context = Vec::with_capacity(history_before_confirmed.len().saturating_add(1)); + lookup_context.extend_from_slice(history_before_confirmed); + lookup_context.push(confirmed_token_id); + let drafts = prompt_lookup_draft( + &lookup_context, + config.max_drafts, + config.max_ngram, + config.max_window, + ); + + let base_cache = cache.clone(); + let mut verify_tokens = Vec::with_capacity(drafts.len().saturating_add(1)); + verify_tokens.push(confirmed_token_id); + verify_tokens.extend(drafts.iter().copied()); + + let logits = model + .forward_all_logits(&token_input(&verify_tokens)?, None, cache) + .map_err(EngineError::Mlx)?; + let verifier_targets = greedy_token_ids(&logits)?; + if verifier_targets.len() < verify_tokens.len() { + return Err(EngineError::Generation(format!( + "prompt-lookup verifier returned {} target ids for {} input tokens", + verifier_targets.len(), + verify_tokens.len() + ))); + } + + let accepted_drafts = accepted_draft_prefix_len(&drafts, &verifier_targets); + let tokens = emitted_tokens(confirmed_token_id, &drafts, accepted_drafts); + + let next_token_id = if accepted_drafts == drafts.len() { + *verifier_targets.get(accepted_drafts).ok_or_else(|| { + EngineError::Generation(format!( + "prompt-lookup verifier missing target at accepted index {accepted_drafts}" + )) + })? + } else { + *cache = base_cache; + let replay_logits = model + .forward_all_logits(&token_input(&tokens)?, None, cache) + .map_err(EngineError::Mlx)?; + let replay_targets = greedy_token_ids(&replay_logits)?; + *replay_targets.get(accepted_drafts).ok_or_else(|| { + EngineError::Generation(format!( + "prompt-lookup replay returned {} target ids for accepted index {}", + replay_targets.len(), + accepted_drafts + )) + })? + }; + + Ok(PromptLookupCycleResult { + tokens, + next_token_id, + drafted: drafts.len(), + accepted_drafts, + }) +} + +/// Run one unchecked prompt-lookup cycle. +/// +/// This path copies draft tokens from prompt/history without per-token verifier +/// logits. It still advances the target model cache over the emitted span and +/// samples the next token from the final position, but it is not guaranteed to +/// reproduce greedy decode if the copied tokens would have been rejected. +pub fn unchecked_prompt_lookup_cycle( + model: &mut AnyModel, + cache: &mut AnyCache, + history_before_confirmed: &[u32], + confirmed_token_id: u32, + config: PromptLookupConfig, +) -> Result { + let mut lookup_context = Vec::with_capacity(history_before_confirmed.len().saturating_add(1)); + lookup_context.extend_from_slice(history_before_confirmed); + lookup_context.push(confirmed_token_id); + let drafts = prompt_lookup_draft( + &lookup_context, + config.max_drafts, + config.max_ngram, + config.max_window, + ); + + let mut tokens = Vec::with_capacity(drafts.len().saturating_add(1)); + tokens.push(confirmed_token_id); + tokens.extend(drafts.iter().copied()); + + let logits = model + .forward_last_token(&token_input(&tokens)?, None, cache) + .map_err(EngineError::Mlx)?; + let next_token_id = greedy_token_id(&logits)?; + + Ok(PromptLookupCycleResult { + tokens, + next_token_id, + drafted: drafts.len(), + accepted_drafts: drafts.len(), + }) +} + +fn token_input(tokens: &[u32]) -> Result { + let mut input = Vec::with_capacity(tokens.len()); + for &token in tokens { + input.push( + i32::try_from(token) + .map_err(|_| EngineError::Generation("token id exceeds i32 range".to_owned()))?, + ); + } + let len = i32::try_from(input.len()) + .map_err(|_| EngineError::Generation("token batch too large".to_owned()))?; + Ok(Array::from_slice(&input, &[1, len])) +} + +fn hidden_row(hidden: &Array, row: usize) -> Result { + let row_i32 = i32::try_from(row) + .map_err(|_| EngineError::Generation("hidden row index too large".to_owned()))?; + Ok(hidden.index((.., row_i32..row_i32 + 1, ..))) +} + +fn hidden_rows(hidden: &Array, start: usize, end: usize) -> Result { + let start_i32 = i32::try_from(start) + .map_err(|_| EngineError::Generation("hidden row start index too large".to_owned()))?; + let end_i32 = i32::try_from(end) + .map_err(|_| EngineError::Generation("hidden row end index too large".to_owned()))?; + Ok(hidden.index((.., start_i32..end_i32, ..))) +} + +fn zero_hidden_row_like(hidden: &Array) -> Result { + let shape = hidden.shape(); + let batch = *shape + .first() + .ok_or_else(|| EngineError::Generation("hidden tensor missing batch dim".to_owned()))?; + let hidden_dim = *shape + .get(2) + .ok_or_else(|| EngineError::Generation("hidden tensor missing hidden dim".to_owned()))?; + ops::zeros_dtype(&[batch, 1, hidden_dim], hidden.dtype()).map_err(EngineError::Mlx) +} + +fn shifted_hidden_rows( + initial_hidden: &Array, + hidden: &Array, + count: usize, +) -> Result { + if count == 0 { + return Err(EngineError::Generation( + "cannot build shifted hidden rows for empty token batch".to_owned(), + )); + } + if count == 1 { + return Ok(initial_hidden.clone()); + } + + let tail = hidden_rows(hidden, 0, count - 1)?; + concatenate_axis(&[initial_hidden, &tail], 1).map_err(EngineError::Mlx) +} + +/// Prime an MTP cache from a backbone hidden sequence. +/// +/// `hidden` must contain the raw backbone hidden states for `tokens`. +/// The first MTP row uses a zero previous-hidden row, matching llama.cpp's +/// draft-mtp prompt mirroring behavior. +pub fn prime_mtp_cache( + model: &mut AnyModel, + mtp_cache: &mut MtpCache, + tokens: &[u32], + hidden: &Array, +) -> Result<(), EngineError> { + if tokens.is_empty() { + return Ok(()); + } + + let zero = zero_hidden_row_like(hidden)?; + let shifted = shifted_hidden_rows(&zero, hidden, tokens.len())?; + model + .mtp_advance_many(&shifted, tokens, mtp_cache) + .map_err(EngineError::Mlx) +} + +/// Mirror one accepted backbone token into an already-primed MTP cache. +pub fn mirror_mtp_token( + model: &mut AnyModel, + mtp_cache: &mut MtpCache, + previous_hidden: &Array, + token: u32, +) -> Result<(), EngineError> { + model + .mtp_advance_many(previous_hidden, &[token], mtp_cache) + .map_err(EngineError::Mlx) +} + +fn backbone_verify_batch( + model: &mut AnyModel, + cache: &mut AnyCache, + tokens: &[u32], +) -> Result<(Array, Vec), EngineError> { + let input = token_input(tokens)?; + let (hidden, logits) = model + .forward_with_hidden(&input, None, cache) + .map_err(EngineError::Mlx)?; + let target_ids = greedy_token_ids(&logits)?; + Ok((hidden, target_ids)) +} + +fn commit_mtp_cache( + model: &mut AnyModel, + mtp_cache: &mut MtpCache, + confirmed_mtp_cache: MtpCache, + accepted_hidden_rows: &Array, + drafts: &[u32], + accepted_drafts: usize, +) -> Result<(), EngineError> { + *mtp_cache = confirmed_mtp_cache; + + if accepted_drafts > 0 { + let accepted = drafts.get(..accepted_drafts).ok_or_else(|| { + EngineError::Generation(format!( + "MTP cache commit missing accepted draft prefix len {accepted_drafts}" + )) + })?; + let hidden_before = hidden_rows(accepted_hidden_rows, 0, accepted_drafts)?; + model + .mtp_advance_many(&hidden_before, accepted, mtp_cache) + .map_err(EngineError::Mlx)?; + } + + Ok(()) +} + +fn trim_mtp_cache_by(mtp_cache: &mut MtpCache, rejected: usize) { + if rejected == 0 { + return; + } + + for layer in mtp_cache { + layer.trim_by(rejected); + } +} + +fn mirror_verified_mtp_cache( + model: &mut AnyModel, + mtp_cache: &mut MtpCache, + base_mtp_cache: MtpCache, + previous_hidden: &Array, + verify_hidden: &Array, + verify_tokens: &[u32], + accepted_token_count: usize, +) -> Result<(), EngineError> { + let mut mirrored = base_mtp_cache; + let shifted = shifted_hidden_rows(previous_hidden, verify_hidden, verify_tokens.len())?; + model + .mtp_advance_many(&shifted, verify_tokens, &mut mirrored) + .map_err(EngineError::Mlx)?; + + let rejected = verify_tokens.len().saturating_sub(accepted_token_count); + trim_mtp_cache_by(&mut mirrored, rejected); + *mtp_cache = mirrored; + + Ok(()) } /// Run one MTP speculative decode cycle. /// /// Given the backbone's hidden state at position t and the confirmed token t+1: -/// 1. MTP drafts token t+2 -/// 2. Process confirmed token through backbone → get prediction + hidden -/// 3. Check if backbone's prediction matches the draft -/// 4. If accepted: process draft through backbone too, return both tokens -/// 5. If rejected: return just confirmed (draft was never processed, no rollback needed) +/// 1. MTP drafts up to `draft_n_max` future tokens. +/// 2. The backbone verifies the confirmed token plus all drafts in one batch. +/// 3. The caches are kept on full acceptance or rebuilt from the accepted prefix +/// after a rejection. +#[allow(clippy::too_many_lines)] pub fn mtp_cycle( model: &mut AnyModel, cache: &mut AnyCache, mtp_cache: &mut MtpCache, hidden: &Array, confirmed_token_id: u32, + draft_n_max: usize, ) -> Result { - // Step 1: MTP draft — predict token at position t+2 - let draft_logits = model - .mtp_draft(hidden, confirmed_token_id, mtp_cache) - .map_err(EngineError::Mlx)?; - let draft_token_arr = - argmax_axis!(&draft_logits.index((.., -1, ..)), -1).map_err(EngineError::Mlx)?; - eval([&draft_token_arr]).map_err(EngineError::Mlx)?; - let draft_token_id: u32 = draft_token_arr.item(); - - // Step 2: Process confirmed token through backbone - let confirmed_input = Array::from_slice(&[confirmed_token_id.cast_signed()], &[1, 1]); - let (confirmed_hidden, confirmed_logits) = model - .forward_with_hidden(&confirmed_input, None, cache) - .map_err(EngineError::Mlx)?; + let draft_limit = draft_n_max.max(1); + let base_cache = cache.clone(); + let base_mtp_cache = mtp_cache.clone(); + let mut speculative_mtp_cache = mtp_cache.clone(); + let mut confirmed_mtp_cache: Option = None; + let mut speculative_hidden = hidden.clone(); + let mut speculative_token = confirmed_token_id; + let mut drafts = Vec::with_capacity(draft_limit); - let target_arr = - argmax_axis!(&confirmed_logits.index((.., -1, ..)), -1).map_err(EngineError::Mlx)?; - let h_confirmed = confirmed_hidden.index((.., -1.., ..)); - eval([&target_arr, &h_confirmed]).map_err(EngineError::Mlx)?; - let target_id: u32 = target_arr.item(); - - if draft_matches_target(draft_token_id, target_id) { - // ACCEPT: draft matches backbone prediction. - // Keep the MTP cache aligned with the newly accepted draft token too. - model - .mtp_advance(&h_confirmed, draft_token_id, mtp_cache) + for draft_idx in 0..draft_limit { + let (next_hidden, draft_logits) = model + .mtp_draft_with_hidden( + &speculative_hidden, + speculative_token, + &mut speculative_mtp_cache, + ) .map_err(EngineError::Mlx)?; + let draft_token_id = greedy_token_id(&draft_logits)?; + drafts.push(draft_token_id); + speculative_hidden = next_hidden; + speculative_token = draft_token_id; + if draft_idx == 0 { + confirmed_mtp_cache = Some(speculative_mtp_cache.clone()); + } + } - // Process draft token to advance cache and get bonus prediction. - let draft_input = Array::from_slice(&[draft_token_id.cast_signed()], &[1, 1]); - let (draft_hidden, bonus_logits) = model - .forward_with_hidden(&draft_input, None, cache) - .map_err(EngineError::Mlx)?; + let first_draft = *drafts + .first() + .ok_or_else(|| EngineError::Generation("MTP produced no draft tokens".to_owned()))?; + + let mut verify_tokens = Vec::with_capacity(drafts.len().saturating_add(1)); + verify_tokens.push(confirmed_token_id); + verify_tokens.extend(drafts.iter().copied()); - let bonus_token = - argmax_axis!(&bonus_logits.index((.., -1, ..)), -1).map_err(EngineError::Mlx)?; - let h_last = draft_hidden.index((.., -1.., ..)); - eval([&bonus_token, &h_last]).map_err(EngineError::Mlx)?; - let bonus_id: u32 = bonus_token.item(); - - Ok(MtpCycleResult { - tokens: vec![confirmed_token_id, draft_token_id], - hidden: h_last, - next_token_id: bonus_id, - }) + let (verify_hidden, verifier_targets) = backbone_verify_batch(model, cache, &verify_tokens)?; + let verify_hidden_for_mtp = verify_hidden.clone(); + if verifier_targets.len() < verify_tokens.len() { + return Err(EngineError::Generation(format!( + "batched MTP verifier returned {} target ids for {} input tokens", + verifier_targets.len(), + verify_tokens.len() + ))); + } + + let first_target = *verifier_targets + .first() + .ok_or_else(|| EngineError::Generation("MTP verifier returned no targets".to_owned()))?; + let accepted_drafts = if draft_matches_target(first_draft, first_target) { + accepted_draft_prefix_len(&drafts, &verifier_targets) } else { - // REJECT: keep the MTP cache entry for the confirmed token. - // `mtp_draft()` advanced the speculative head using the confirmed token, - // not the rejected draft token, so rolling it back would drop real history. + 0 + }; + let tokens = emitted_tokens(confirmed_token_id, &drafts, accepted_drafts); - Ok(MtpCycleResult { - tokens: vec![confirmed_token_id], - hidden: h_confirmed, - next_token_id: target_id, - }) + let (accepted_hidden_rows, next_token_id) = if accepted_drafts == drafts.len() { + let next = *verifier_targets.get(accepted_drafts).ok_or_else(|| { + EngineError::Generation(format!( + "MTP verifier missing target at accepted index {accepted_drafts}" + )) + })?; + (verify_hidden, next) + } else { + *cache = base_cache; + let (replay_hidden, replay_targets) = backbone_verify_batch(model, cache, &tokens)?; + let next = *replay_targets.get(accepted_drafts).ok_or_else(|| { + EngineError::Generation(format!( + "MTP replay returned {} target ids for accepted index {}", + replay_targets.len(), + accepted_drafts + )) + })?; + (replay_hidden, next) + }; + + let h_last = hidden_row(&accepted_hidden_rows, accepted_drafts)?; + if mtp_mirror_verify_enabled() { + mirror_verified_mtp_cache( + model, + mtp_cache, + base_mtp_cache, + hidden, + &verify_hidden_for_mtp, + &verify_tokens, + tokens.len(), + )?; + } else { + commit_mtp_cache( + model, + mtp_cache, + confirmed_mtp_cache.ok_or_else(|| { + EngineError::Generation("MTP produced no cache checkpoint".to_owned()) + })?, + &accepted_hidden_rows, + &drafts, + accepted_drafts, + )?; } + + if accepted_drafts < drafts.len() && tokens.is_empty() { + return Err(EngineError::Generation( + "MTP accepted no committed tokens".to_owned(), + )); + } + + Ok(MtpCycleResult { + tokens, + hidden: h_last, + next_token_id, + drafted: drafts.len(), + accepted_drafts, + }) } #[cfg(test)] mod tests { - use super::draft_matches_target; + use super::{ + MtpStats, accepted_draft_prefix_len, draft_matches_target, emitted_tokens, + prompt_lookup_draft, + }; #[test] fn draft_match_helper_accepts_identical_tokens() { @@ -110,4 +611,68 @@ mod tests { fn draft_match_helper_rejects_different_tokens() { assert!(!draft_matches_target(17, 18)); } + + #[test] + fn mtp_stats_tracks_drafted_and_bonus_acceptance_rate() { + let mut stats = MtpStats::default(); + stats.record_cycle(3, 4); + stats.record_cycle(2, 1); + + assert_eq!(stats.cycles(), 2); + assert_eq!(stats.drafted(), 5); + assert_eq!(stats.emitted(), 5); + assert_eq!(stats.accepted_drafts(), 3); + assert!((stats.acceptance_rate_percent() - 60.0).abs() < f64::EPSILON); + } + + #[test] + fn accepted_draft_prefix_len_stops_at_first_mismatch() { + let drafts = [10, 20, 30]; + let verifier_targets = [10, 21, 30, 40]; + + assert_eq!(accepted_draft_prefix_len(&drafts, &verifier_targets), 1); + } + + #[test] + fn accepted_draft_prefix_len_accepts_full_prefix() { + let drafts = [10, 20, 30]; + let verifier_targets = [10, 20, 30, 40]; + + assert_eq!(accepted_draft_prefix_len(&drafts, &verifier_targets), 3); + } + + #[test] + fn emitted_tokens_includes_confirmed_and_accepted_drafts() { + let drafts = [10, 20, 30]; + + assert_eq!(emitted_tokens(7, &drafts, 2), vec![7, 10, 20]); + } + + #[test] + fn prompt_lookup_drafts_from_longest_prior_suffix_match() { + let context = [1, 2, 3, 4, 5, 1, 2]; + + assert_eq!(prompt_lookup_draft(&context, 3, 4, 64), vec![3, 4, 5]); + } + + #[test] + fn prompt_lookup_caps_drafts() { + let context = [9, 8, 7, 6, 9, 8]; + + assert_eq!(prompt_lookup_draft(&context, 1, 3, 64), vec![7]); + } + + #[test] + fn prompt_lookup_ignores_current_tail_self_match() { + let context = [1, 2, 3, 4]; + + assert!(prompt_lookup_draft(&context, 3, 4, 64).is_empty()); + } + + #[test] + fn prompt_lookup_respects_search_window() { + let context = [1, 2, 3, 4, 5, 1, 2]; + + assert!(prompt_lookup_draft(&context, 3, 4, 3).is_empty()); + } } diff --git a/crates/higgs-engine/src/simple.rs b/crates/higgs-engine/src/simple.rs index 88c0425f..c4bf0ed2 100644 --- a/crates/higgs-engine/src/simple.rs +++ b/crates/higgs-engine/src/simple.rs @@ -56,6 +56,39 @@ fn experimental_paged_kv_enabled() -> bool { .unwrap_or(false) } +fn prompt_lookup_enabled() -> bool { + parse_enabled_flag(std::env::var("HIGGS_PROMPT_LOOKUP").ok().as_deref()).unwrap_or(false) +} + +fn unchecked_prompt_lookup_enabled() -> bool { + parse_enabled_flag( + std::env::var("HIGGS_PROMPT_LOOKUP_UNCHECKED") + .ok() + .as_deref(), + ) + .unwrap_or(false) +} + +fn mtp_prefill_priming_enabled() -> bool { + parse_enabled_flag(std::env::var("HIGGS_MTP_PRIME_PREFILL").ok().as_deref()).unwrap_or(true) +} + +fn parse_env_usize(name: &str, default: usize) -> usize { + std::env::var(name) + .ok() + .and_then(|raw| raw.trim().parse::().ok()) + .unwrap_or(default) +} + +fn prompt_lookup_config() -> crate::mtp::PromptLookupConfig { + let defaults = crate::mtp::PromptLookupConfig::default(); + crate::mtp::PromptLookupConfig { + max_drafts: parse_env_usize("HIGGS_PROMPT_LOOKUP_DRAFT_N_MAX", defaults.max_drafts), + max_ngram: parse_env_usize("HIGGS_PROMPT_LOOKUP_NGRAM_MAX", defaults.max_ngram), + max_window: parse_env_usize("HIGGS_PROMPT_LOOKUP_WINDOW", defaults.max_window), + } +} + fn estimate_paged_kv_blocks( target_bytes: usize, num_kv_heads: usize, @@ -217,6 +250,7 @@ pub struct SimpleEngine { struct PreparedGeneration<'a> { model: MutexGuard<'a, AnyModel>, cache: AnyCache, + actual_prompt_tokens: Vec, prompt_array: Array, prompt_len: u32, pixel_values: Option, @@ -328,6 +362,7 @@ impl SimpleEngine { chunked_prefill_chunk_size = tuning.chunked_prefill_chunk_size(), clear_cache_after_prefill = tuning.clear_cache_after_prefill(), mtp_enabled = tuning.enable_mtp(), + mtp_draft_n_max = tuning.mtp_draft_n_max(), paged_kv_target_mb = tuning.paged_kv_target_bytes() / (1024 * 1024), "Engine ready" ); @@ -341,6 +376,7 @@ impl SimpleEngine { chunked_prefill_chunk_size = tuning.chunked_prefill_chunk_size(), clear_cache_after_prefill = tuning.clear_cache_after_prefill(), mtp_enabled = tuning.enable_mtp(), + mtp_draft_n_max = tuning.mtp_draft_n_max(), "Engine ready" ); tracing::debug!("Experimental paged KV disabled; session cache allocation skipped"); @@ -547,6 +583,7 @@ impl SimpleEngine { Ok(PreparedGeneration { model, cache, + actual_prompt_tokens, prompt_array, prompt_len, pixel_values, @@ -563,7 +600,9 @@ impl SimpleEngine { params: &SamplingParams, logprob_top_n: Option, constraint: Option<&crate::constrained::ConstrainedGenerator>, - ) -> Result<(Array, Option), EngineError> { + capture_hidden: bool, + ) -> Result<(Array, Option, Option), EngineError> { + let mut prefill_hidden = None; let logits = if let Some(ref pixel_values) = prepared.pixel_values { // Multimodal path: full forward (VLMs need all tokens for vision) prepared @@ -576,7 +615,14 @@ impl SimpleEngine { let seq_len = prepared.prompt_array.shape().get(1).copied().unwrap_or(0); let chunked_threshold = self.tuning.chunked_prefill_threshold(); let chunked_size = self.tuning.chunked_prefill_chunk_size(); - if seq_len > chunked_threshold { + if capture_hidden && seq_len <= chunked_threshold { + let (hidden, logits) = prepared + .model + .forward_with_hidden(&prepared.prompt_array, None, &mut prepared.cache) + .map_err(EngineError::Mlx)?; + prefill_hidden = Some(hidden); + logits + } else if seq_len > chunked_threshold { prepared .model .forward_chunked(&prepared.prompt_array, &mut prepared.cache, chunked_size) @@ -619,6 +665,9 @@ impl SimpleEngine { if let Some(ref lp) = logprob_data { eval_targets.extend(lp.eval_targets()); } + if let Some(ref hidden) = prefill_hidden { + eval_targets.push(hidden); + } eval(eval_targets).map_err(EngineError::Mlx)?; } @@ -645,7 +694,7 @@ impl SimpleEngine { "simple_post_prefill", ); - Ok((current_token, logprob_data)) + Ok((current_token, logprob_data, prefill_hidden)) } /// Decode a single step: forward pass on the current token, apply penalties @@ -766,6 +815,12 @@ impl SimpleEngine { .map_err(|_| EngineError::Generation("Too many tokens generated".to_owned())) } + fn hidden_row_from_sequence(hidden: &Array, row_index: usize) -> Result { + let row_i32 = i32::try_from(row_index) + .map_err(|_| EngineError::Generation("hidden row index too large".to_owned()))?; + Ok(hidden.index((.., row_i32..row_i32 + 1, ..))) + } + // ========================================================================= // Session Management (Batched Generation) // ========================================================================= @@ -937,12 +992,22 @@ impl SimpleEngine { let mut prepared = self.prepare_generation(prompt_tokens, pixel_values)?; let prompt_len = prepared.prompt_len; - let (current_token, first_logprob_data) = self.run_prefill( + #[allow(clippy::float_cmp)] + let capture_mtp_prefill = mtp_prefill_priming_enabled() + && self.tuning.enable_mtp() + && prepared.model.has_mtp() + && prepared.pixel_values.is_none() + && constraint.is_none() + && !logprobs + && params.temperature == 0.0; + + let (current_token, first_logprob_data, prefill_hidden) = self.run_prefill( prompt_tokens, &mut prepared, params, logprob_top_n, constraint.as_ref(), + capture_mtp_prefill, )?; // Capture T1 (already eval'd inside run_prefill). @@ -990,6 +1055,24 @@ impl SimpleEngine { }); } + // Architecture-neutral speculative decode: prompt-lookup drafting plus + // batched verifier logits. Explicitly opt-in while benchmark data is + // collected because the normal path has a pipelined single-token loop. + #[allow(clippy::float_cmp)] + if prompt_lookup_enabled() && constraint.is_none() && !logprobs && params.temperature == 0.0 + { + return self.prompt_lookup_generate( + &mut prepared.model, + &mut prepared.cache, + first_token_id, + max_tokens, + prompt_len, + &mut tokens, + stop_sequences, + enable_thinking, + ); + } + // MTP speculative decode: enabled by the resolved MLX runtime tuning. // Only for greedy (temperature == 0), no constraints, no logprobs. #[allow(clippy::float_cmp)] @@ -999,9 +1082,12 @@ impl SimpleEngine { && !logprobs && params.temperature == 0.0 { + let actual_prompt_tokens = prepared.actual_prompt_tokens.clone(); return self.mtp_generate( &mut prepared.model, &mut prepared.cache, + &actual_prompt_tokens, + prefill_hidden.as_ref(), first_token_id, max_tokens, prompt_len, @@ -1249,10 +1335,237 @@ impl SimpleEngine { } } + #[allow(clippy::cast_precision_loss)] + fn log_mtp_decode_stats( + stats: &crate::mtp::MtpStats, + elapsed: std::time::Duration, + reason: &str, + ) { + tracing::info!( + reason, + cycles = stats.cycles(), + drafted = stats.drafted(), + accepted_drafts = stats.accepted_drafts(), + emitted = stats.emitted(), + accept_rate = format!("{:.1}%", stats.acceptance_rate_percent()), + tok_per_s = format!("{:.1}", f64::from(stats.emitted()) / elapsed.as_secs_f64()), + "MTP decode complete" + ); + } + + #[allow(clippy::cast_precision_loss)] + fn log_prompt_lookup_decode_stats( + stats: &crate::mtp::MtpStats, + elapsed: std::time::Duration, + reason: &str, + ) { + tracing::info!( + reason, + cycles = stats.cycles(), + drafted = stats.drafted(), + accepted_drafts = stats.accepted_drafts(), + emitted = stats.emitted(), + accept_rate = format!("{:.1}%", stats.acceptance_rate_percent()), + tok_per_s = format!("{:.1}", f64::from(stats.emitted()) / elapsed.as_secs_f64()), + "Prompt-lookup decode complete" + ); + } + + /// Architecture-neutral prompt-lookup speculative decode loop. + /// + /// The draft provider copies likely next tokens from repeated prompt/history + /// spans, then verifies the whole candidate window in one model pass. + #[allow( + clippy::too_many_arguments, + clippy::as_conversions, + clippy::cast_precision_loss + )] + fn prompt_lookup_generate( + &self, + model: &mut higgs_models::AnyModel, + cache: &mut higgs_models::AnyCache, + first_token_id: u32, + max_tokens: u32, + prompt_len: u32, + tokens: &mut Vec, + stop_sequences: &[String], + enable_thinking: bool, + ) -> Result { + let has_stop_sequences = !stop_sequences.is_empty(); + let unchecked_lookup = unchecked_prompt_lookup_enabled(); + let first_token_i32 = i32::try_from(first_token_id) + .map_err(|_| EngineError::Generation("token id exceeds i32 range".to_owned()))?; + let first_input = Array::from_slice(&[first_token_i32], &[1, 1]); + let logits = model + .forward_all_logits(&first_input, None, cache) + .map_err(EngineError::Mlx)?; + let next_arr = + mlx_rs::argmax_axis!(&logits.index((.., -1, ..)), -1).map_err(EngineError::Mlx)?; + eval([&next_arr]).map_err(EngineError::Mlx)?; + + let mut confirmed_token_id: u32 = next_arr.item(); + let base_config = prompt_lookup_config(); + let mut stats = crate::mtp::MtpStats::default(); + let t_start = std::time::Instant::now(); + + const THINKING_BUDGET: u32 = 256; + let think_close_token = if enable_thinking { + self.think_close_token + } else { + None + }; + let mut thinking_tokens: u32 = u32::from(think_close_token.is_some()); + let mut seen_think_close = + think_close_token.is_some_and(|close_id| first_token_id == close_id); + + loop { + let completion_len = Self::completion_len(tokens)?; + if completion_len >= max_tokens { + let elapsed = t_start.elapsed(); + Self::log_prompt_lookup_decode_stats(&stats, elapsed, "length"); + return Ok(GenerationOutput { + text: self.decode_tokens(tokens)?, + finish_reason: "length".to_owned(), + prompt_tokens: prompt_len, + completion_tokens: completion_len, + token_logprobs: None, + }); + } + + let remaining = usize::try_from(max_tokens.saturating_sub(completion_len)) + .map_err(|_| EngineError::Generation("max_tokens overflow".to_owned()))?; + let config = crate::mtp::PromptLookupConfig { + max_drafts: base_config.max_drafts.min(remaining.saturating_sub(1)), + ..base_config + }; + + let result = if unchecked_lookup { + crate::mtp::unchecked_prompt_lookup_cycle( + model, + cache, + tokens, + confirmed_token_id, + config, + )? + } else { + crate::mtp::prompt_lookup_cycle(model, cache, tokens, confirmed_token_id, config)? + }; + stats.record_cycle(result.drafted, result.tokens.len()); + + for &tok in &result.tokens { + if let Some(close_id) = think_close_token { + if !seen_think_close { + if tok == close_id { + seen_think_close = true; + } else { + thinking_tokens += 1; + if thinking_tokens >= THINKING_BUDGET { + tokens.push(close_id); + seen_think_close = true; + tracing::info!( + budget = THINKING_BUDGET, + "Prompt lookup: thinking budget reached, forcing " + ); + if self.eos_token_ids.contains(&close_id) { + let elapsed = t_start.elapsed(); + Self::log_prompt_lookup_decode_stats(&stats, elapsed, "stop"); + return Ok(GenerationOutput { + text: self.decode_tokens(tokens)?, + finish_reason: "stop".to_owned(), + prompt_tokens: prompt_len, + completion_tokens: Self::completion_len(tokens)?, + token_logprobs: None, + }); + } + + if has_stop_sequences { + let text = self.decode_tokens(tokens)?; + if let Some(truncated) = + check_stop_sequences(&text, stop_sequences) + { + let elapsed = t_start.elapsed(); + Self::log_prompt_lookup_decode_stats( + &stats, elapsed, "stop", + ); + return Ok(GenerationOutput { + text: truncated, + finish_reason: "stop".to_owned(), + prompt_tokens: prompt_len, + completion_tokens: Self::completion_len(tokens)?, + token_logprobs: None, + }); + } + } + + let forced_completion_len = Self::completion_len(tokens)?; + if forced_completion_len >= max_tokens { + let elapsed = t_start.elapsed(); + Self::log_prompt_lookup_decode_stats(&stats, elapsed, "length"); + return Ok(GenerationOutput { + text: self.decode_tokens(tokens)?, + finish_reason: "length".to_owned(), + prompt_tokens: prompt_len, + completion_tokens: forced_completion_len, + token_logprobs: None, + }); + } + break; + } + } + } + } + + tokens.push(tok); + + if self.eos_token_ids.contains(&tok) { + let elapsed = t_start.elapsed(); + Self::log_prompt_lookup_decode_stats(&stats, elapsed, "stop"); + return Ok(GenerationOutput { + text: self.decode_tokens(tokens)?, + finish_reason: "stop".to_owned(), + prompt_tokens: prompt_len, + completion_tokens: Self::completion_len(tokens)?, + token_logprobs: None, + }); + } + } + + if has_stop_sequences { + let text = self.decode_tokens(tokens)?; + if let Some(truncated) = check_stop_sequences(&text, stop_sequences) { + let elapsed = t_start.elapsed(); + Self::log_prompt_lookup_decode_stats(&stats, elapsed, "stop"); + return Ok(GenerationOutput { + text: truncated, + finish_reason: "stop".to_owned(), + prompt_tokens: prompt_len, + completion_tokens: Self::completion_len(tokens)?, + token_logprobs: None, + }); + } + } + + let final_completion_len = Self::completion_len(tokens)?; + if final_completion_len >= max_tokens { + let elapsed = t_start.elapsed(); + Self::log_prompt_lookup_decode_stats(&stats, elapsed, "length"); + return Ok(GenerationOutput { + text: self.decode_tokens(tokens)?, + finish_reason: "length".to_owned(), + prompt_tokens: prompt_len, + completion_tokens: final_completion_len, + token_logprobs: None, + }); + } + + confirmed_token_id = result.next_token_id; + } + } + /// MTP speculative decode loop. /// /// Runs the backbone to get the initial hidden state, then loops calling - /// `mtp_cycle()` which drafts one extra token per cycle for ~1.5x speedup. + /// `mtp_cycle()` which drafts multiple tokens per cycle for speculative speedup. #[allow( clippy::too_many_arguments, clippy::as_conversions, @@ -1262,6 +1575,8 @@ impl SimpleEngine { &self, model: &mut higgs_models::AnyModel, cache: &mut higgs_models::AnyCache, + actual_prompt_tokens: &[u32], + prefill_hidden: Option<&Array>, first_token_id: u32, max_tokens: u32, prompt_len: u32, @@ -1275,10 +1590,10 @@ impl SimpleEngine { let mut mtp_cache = model .make_mtp_cache() .ok_or_else(|| EngineError::Generation("MTP cache creation failed".into()))?; + if let Some(hidden) = prefill_hidden { + crate::mtp::prime_mtp_cache(model, &mut mtp_cache, actual_prompt_tokens, hidden)?; + } - // Get initial hidden state: re-run backbone on first token to obtain h_t. - // This single-token forward is cheap and gives us the hidden state needed - // for the first MTP draft. let first_input = Array::from_slice(&[first_token_id as i32], &[1, 1]); let (hidden, logits) = model .forward_with_hidden(&first_input, None, cache) @@ -1286,12 +1601,18 @@ impl SimpleEngine { let next_arr = mlx_rs::argmax_axis!(&logits.index((.., -1, ..)), -1).map_err(EngineError::Mlx)?; let h = hidden.index((.., -1.., ..)); + if let Some(previous_hidden) = prefill_hidden + .filter(|_| !actual_prompt_tokens.is_empty()) + .map(|prefill| Self::hidden_row_from_sequence(prefill, actual_prompt_tokens.len() - 1)) + .transpose()? + { + crate::mtp::mirror_mtp_token(model, &mut mtp_cache, &previous_hidden, first_token_id)?; + } eval([&next_arr, &h]).map_err(EngineError::Mlx)?; let mut current_hidden = h; let mut confirmed_token_id: u32 = next_arr.item(); - let mut accepted: u32 = 0; - let mut total_cycles: u32 = 0; + let mut mtp_stats = crate::mtp::MtpStats::default(); let t_start = std::time::Instant::now(); // Thinking budget: force after N tokens if model hasn't closed it. @@ -1313,9 +1634,10 @@ impl SimpleEngine { &mut mtp_cache, ¤t_hidden, confirmed_token_id, + self.tuning.mtp_draft_n_max(), )?; - total_cycles += 1; + mtp_stats.record_cycle(result.drafted, result.tokens.len()); for &tok in &result.tokens { // Thinking budget enforcement @@ -1327,7 +1649,6 @@ impl SimpleEngine { thinking_tokens += 1; if thinking_tokens >= THINKING_BUDGET { tokens.push(close_id); - accepted += 1; seen_think_close = true; tracing::info!( budget = THINKING_BUDGET, @@ -1335,20 +1656,7 @@ impl SimpleEngine { ); if self.eos_token_ids.contains(&close_id) { let elapsed = t_start.elapsed(); - tracing::info!( - tokens = accepted, - cycles = total_cycles, - accept_rate = format!( - "{:.1}%", - (f64::from(accepted) / f64::from(total_cycles) - 1.0) - * 100.0 - ), - tok_per_s = format!( - "{:.1}", - f64::from(accepted) / elapsed.as_secs_f64() - ), - "MTP decode complete" - ); + Self::log_mtp_decode_stats(&mtp_stats, elapsed, "stop"); return Ok(GenerationOutput { text: self.decode_tokens(tokens)?, finish_reason: "stop".to_owned(), @@ -1363,6 +1671,8 @@ impl SimpleEngine { if let Some(truncated) = check_stop_sequences(&text, stop_sequences) { + let elapsed = t_start.elapsed(); + Self::log_mtp_decode_stats(&mtp_stats, elapsed, "stop"); return Ok(GenerationOutput { text: truncated, finish_reason: "stop".to_owned(), @@ -1376,20 +1686,7 @@ impl SimpleEngine { let completion_len = Self::completion_len(tokens)?; if completion_len >= max_tokens { let elapsed = t_start.elapsed(); - tracing::info!( - tokens = accepted, - cycles = total_cycles, - accept_rate = format!( - "{:.1}%", - (f64::from(accepted) / f64::from(total_cycles) - 1.0) - * 100.0 - ), - tok_per_s = format!( - "{:.1}", - f64::from(accepted) / elapsed.as_secs_f64() - ), - "MTP decode complete (length limit)" - ); + Self::log_mtp_decode_stats(&mtp_stats, elapsed, "length"); return Ok(GenerationOutput { text: self.decode_tokens(tokens)?, finish_reason: "length".to_owned(), @@ -1407,20 +1704,10 @@ impl SimpleEngine { } tokens.push(tok); - accepted += 1; if self.eos_token_ids.contains(&tok) { let elapsed = t_start.elapsed(); - tracing::info!( - tokens = accepted, - cycles = total_cycles, - accept_rate = format!( - "{:.1}%", - (f64::from(accepted) / f64::from(total_cycles) - 1.0) * 100.0 - ), - tok_per_s = format!("{:.1}", f64::from(accepted) / elapsed.as_secs_f64()), - "MTP decode complete" - ); + Self::log_mtp_decode_stats(&mtp_stats, elapsed, "stop"); return Ok(GenerationOutput { text: self.decode_tokens(tokens)?, finish_reason: "stop".to_owned(), @@ -1434,6 +1721,8 @@ impl SimpleEngine { if has_stop_sequences { let text = self.decode_tokens(tokens)?; if let Some(truncated) = check_stop_sequences(&text, stop_sequences) { + let elapsed = t_start.elapsed(); + Self::log_mtp_decode_stats(&mtp_stats, elapsed, "stop"); return Ok(GenerationOutput { text: truncated, finish_reason: "stop".to_owned(), @@ -1447,16 +1736,7 @@ impl SimpleEngine { let completion_len = Self::completion_len(tokens)?; if completion_len >= max_tokens { let elapsed = t_start.elapsed(); - tracing::info!( - tokens = accepted, - cycles = total_cycles, - accept_rate = format!( - "{:.1}%", - (f64::from(accepted) / f64::from(total_cycles) - 1.0) * 100.0 - ), - tok_per_s = format!("{:.1}", f64::from(accepted) / elapsed.as_secs_f64()), - "MTP decode complete (length limit)" - ); + Self::log_mtp_decode_stats(&mtp_stats, elapsed, "length"); return Ok(GenerationOutput { text: self.decode_tokens(tokens)?, finish_reason: "length".to_owned(), @@ -1485,6 +1765,8 @@ impl SimpleEngine { &self, model: &mut higgs_models::AnyModel, cache: &mut higgs_models::AnyCache, + actual_prompt_tokens: &[u32], + prefill_hidden: Option<&Array>, first_token_id: u32, max_tokens: u32, prompt_len: u32, @@ -1499,6 +1781,9 @@ impl SimpleEngine { let mut mtp_cache = model .make_mtp_cache() .ok_or_else(|| EngineError::Generation("MTP cache creation failed".into()))?; + if let Some(hidden) = prefill_hidden { + crate::mtp::prime_mtp_cache(model, &mut mtp_cache, actual_prompt_tokens, hidden)?; + } let first_input = Array::from_slice(&[first_token_id as i32], &[1, 1]); let (hidden, logits) = model @@ -1507,12 +1792,18 @@ impl SimpleEngine { let next_arr = mlx_rs::argmax_axis!(&logits.index((.., -1, ..)), -1).map_err(EngineError::Mlx)?; let h = hidden.index((.., -1.., ..)); + if let Some(previous_hidden) = prefill_hidden + .filter(|_| !actual_prompt_tokens.is_empty()) + .map(|prefill| Self::hidden_row_from_sequence(prefill, actual_prompt_tokens.len() - 1)) + .transpose()? + { + crate::mtp::mirror_mtp_token(model, &mut mtp_cache, &previous_hidden, first_token_id)?; + } eval([&next_arr, &h]).map_err(EngineError::Mlx)?; let mut current_hidden = h; let mut confirmed_token_id: u32 = next_arr.item(); - let mut accepted: u32 = 0; - let mut total_cycles: u32 = 0; + let mut mtp_stats = crate::mtp::MtpStats::default(); let t_start = std::time::Instant::now(); const THINKING_BUDGET: u32 = 256; @@ -1532,9 +1823,10 @@ impl SimpleEngine { &mut mtp_cache, ¤t_hidden, confirmed_token_id, + self.tuning.mtp_draft_n_max(), )?; - total_cycles += 1; + mtp_stats.record_cycle(result.drafted, result.tokens.len()); for &tok in &result.tokens { // Thinking budget enforcement @@ -1546,7 +1838,6 @@ impl SimpleEngine { thinking_tokens += 1; if thinking_tokens >= THINKING_BUDGET { tokens.push(close_id); - accepted += 1; seen_think_close = true; tracing::info!( budget = THINKING_BUDGET, @@ -1598,19 +1889,10 @@ impl SimpleEngine { if step_finished { let elapsed = t_start.elapsed(); - tracing::info!( - tokens = accepted, - cycles = total_cycles, - accept_rate = format!( - "{:.1}%", - (f64::from(accepted) / f64::from(total_cycles) - 1.0) - * 100.0 - ), - tok_per_s = format!( - "{:.1}", - f64::from(accepted) / elapsed.as_secs_f64() - ), - "MTP streaming decode complete" + Self::log_mtp_decode_stats( + &mtp_stats, + elapsed, + finish_reason.as_deref().unwrap_or("client"), ); } @@ -1639,7 +1921,6 @@ impl SimpleEngine { } tokens.push(tok); - accepted += 1; let is_eos = self.eos_token_ids.contains(&tok); let completion_len = Self::completion_len(tokens)?; @@ -1679,15 +1960,10 @@ impl SimpleEngine { if step_finished { let elapsed = t_start.elapsed(); - tracing::info!( - tokens = accepted, - cycles = total_cycles, - accept_rate = format!( - "{:.1}%", - (f64::from(accepted) / f64::from(total_cycles) - 1.0) * 100.0 - ), - tok_per_s = format!("{:.1}", f64::from(accepted) / elapsed.as_secs_f64()), - "MTP streaming decode complete" + Self::log_mtp_decode_stats( + &mtp_stats, + elapsed, + finish_reason.as_deref().unwrap_or("client"), ); } @@ -1817,12 +2093,22 @@ impl SimpleEngine { let mut prepared = self.prepare_generation(prompt_tokens, pixel_values)?; let prompt_len = prepared.prompt_len; - let (current_token, first_logprob_data) = self.run_prefill( + #[allow(clippy::float_cmp)] + let capture_mtp_prefill = mtp_prefill_priming_enabled() + && self.tuning.enable_mtp() + && prepared.model.has_mtp() + && prepared.pixel_values.is_none() + && constraint.is_none() + && !logprobs + && params.temperature == 0.0; + + let (current_token, first_logprob_data, prefill_hidden) = self.run_prefill( prompt_tokens, &mut prepared, params, logprob_top_n, constraint.as_ref(), + capture_mtp_prefill, )?; let mut all_tokens: Vec = Vec::new(); @@ -1883,9 +2169,12 @@ impl SimpleEngine { && !logprobs && params.temperature == 0.0 { + let actual_prompt_tokens = prepared.actual_prompt_tokens.clone(); return self.mtp_generate_streaming( &mut prepared.model, &mut prepared.cache, + &actual_prompt_tokens, + prefill_hidden.as_ref(), first_token_id, max_tokens, prompt_len, diff --git a/crates/higgs-models/src/deepseek_v2.rs b/crates/higgs-models/src/deepseek_v2.rs index 625264fc..884902d7 100644 --- a/crates/higgs-models/src/deepseek_v2.rs +++ b/crates/higgs-models/src/deepseek_v2.rs @@ -858,6 +858,21 @@ impl DeepSeekV2CausalLM { None => self.model.embed_tokens.as_linear(&h_last), } } + + /// Run the model and return logits for every input position. + #[allow(non_snake_case)] + pub fn forward_all_logits( + &mut self, + inputs: &Array, + mask: Option<&Array>, + kv_cache: &mut Vec>, + ) -> Result { + let h = self.forward_hidden(inputs, mask, kv_cache)?; + match self.lm_head.as_ref() { + Some(head) => head.forward(&h), + None => self.model.embed_tokens.as_linear(&h), + } + } } // --------------------------------------------------------------------------- diff --git a/crates/higgs-models/src/gemma2.rs b/crates/higgs-models/src/gemma2.rs index bed2add2..6c2efb9b 100644 --- a/crates/higgs-models/src/gemma2.rs +++ b/crates/higgs-models/src/gemma2.rs @@ -845,6 +845,45 @@ impl Gemma2CausalLM { Ok(logits) } + /// Forward pass producing logits for every input position. + pub fn forward_all_logits( + &mut self, + inputs: &Array, + mask: Option<&Array>, + kv_cache: &mut Vec>, + ) -> Result { + let hidden_all = self.forward_hidden(inputs, mask, kv_cache)?; + let mut logits = match self.lm_head.as_mut() { + Some(head) => head.forward(&hidden_all)?, + None => match &mut self.model.embed_tokens { + MaybeQuantized::Original(embed) => embed.as_linear(&hidden_all)?, + MaybeQuantized::Quantized(q_embed) => q_embed.as_linear(&hidden_all)?, + }, + }; + + if let Some(cap) = self.args.final_logit_softcapping { + let needs_refresh = self + .cached_final_inv_cap + .as_ref() + .is_none_or(|cached| cached.dtype() != logits.dtype()); + if needs_refresh { + self.cached_final_inv_cap = Some(array!(1.0 / cap).as_dtype(logits.dtype())?); + self.cached_final_cap = Some(array!(cap).as_dtype(logits.dtype())?); + } + let final_inv_cap = self + .cached_final_inv_cap + .as_ref() + .ok_or_else(|| Exception::custom("cached_final_inv_cap not initialized"))?; + let final_cap = self + .cached_final_cap + .as_ref() + .ok_or_else(|| Exception::custom("cached_final_cap not initialized"))?; + logits = ops::tanh(&logits.multiply(final_inv_cap)?)?.multiply(final_cap)?; + } + + Ok(logits) + } + pub fn forward_hidden( &mut self, inputs: &Array, diff --git a/crates/higgs-models/src/lib.rs b/crates/higgs-models/src/lib.rs index fe52d245..8ae2716a 100644 --- a/crates/higgs-models/src/lib.rs +++ b/crates/higgs-models/src/lib.rs @@ -281,6 +281,35 @@ impl AnyModel { } } + /// Forward pass producing logits for every input position. + /// + /// Speculative verifiers use this for a candidate window where each + /// position's logits predict the following token. + pub fn forward_all_logits( + &mut self, + inputs: &Array, + mask: Option<&Array>, + cache: &mut AnyCache, + ) -> Result { + match (self, cache) { + (Self::Transformer(m), AnyCache::KV(c)) => m.forward_all_logits(inputs, mask, c), + (Self::Qwen3Moe(m), AnyCache::KV(c)) => m.forward_all_logits(inputs, mask, c), + (Self::Gemma2(m), AnyCache::KV(c)) => m.forward_all_logits(inputs, mask, c), + (Self::Phi3(m), AnyCache::KV(c)) => m.forward_all_logits(inputs, mask, c), + (Self::Starcoder2(m), AnyCache::KV(c)) => m.forward_all_logits(inputs, mask, c), + (Self::LlavaQwen2(m), AnyCache::KV(c)) => m.forward_text_all_logits(inputs, mask, c), + (Self::DeepSeekV2(m), AnyCache::KV(c)) => m.forward_all_logits(inputs, mask, c), + (Self::Qwen3Next(m), AnyCache::Hybrid(c)) => { + let (_, logits) = m.forward_with_hidden(inputs, mask, c)?; + Ok(logits) + } + (Self::BonsaiQ1(_), AnyCache::KV(_)) => Err(Exception::custom( + "forward_all_logits not supported for BonsaiQ1", + )), + _ => Err(Exception::custom("Model/cache type mismatch")), + } + } + /// Chunked prefill: process the prompt in `chunk_size`-token segments. /// /// Produces identical logits to `forward()` but evaluates the compute graph @@ -419,6 +448,26 @@ impl AnyModel { } } + /// Run the MTP head and return both speculative hidden state and logits. + pub fn mtp_draft_with_hidden( + &mut self, + hidden: &Array, + next_token_id: u32, + mtp_cache: &mut MtpCache, + ) -> Result<(Array, Array), Exception> { + match self { + Self::Qwen3Next(m) => m.mtp_draft_with_hidden(hidden, next_token_id, mtp_cache), + Self::Transformer(_) + | Self::Qwen3Moe(_) + | Self::Gemma2(_) + | Self::Phi3(_) + | Self::Starcoder2(_) + | Self::LlavaQwen2(_) + | Self::DeepSeekV2(_) + | Self::BonsaiQ1(_) => Err(Exception::custom("MTP not supported for this model")), + } + } + /// Advance the MTP head/cache for an accepted token without projecting logits. pub fn mtp_advance( &mut self, @@ -439,6 +488,26 @@ impl AnyModel { } } + /// Advance the MTP head/cache for accepted tokens without projecting logits. + pub fn mtp_advance_many( + &mut self, + hidden: &Array, + next_token_ids: &[u32], + mtp_cache: &mut MtpCache, + ) -> Result<(), Exception> { + match self { + Self::Qwen3Next(m) => m.mtp_advance_many(hidden, next_token_ids, mtp_cache), + Self::Transformer(_) + | Self::Qwen3Moe(_) + | Self::Gemma2(_) + | Self::Phi3(_) + | Self::Starcoder2(_) + | Self::LlavaQwen2(_) + | Self::DeepSeekV2(_) + | Self::BonsaiQ1(_) => Err(Exception::custom("MTP not supported for this model")), + } + } + /// Forward pass returning both hidden states and logits for all positions. /// /// Used by MTP speculative decode for the verify pass. @@ -1066,6 +1135,8 @@ pub struct WeightMapIndex { pub weight_map: HashMap, } +const AUXILIARY_SAFETENSORS_FILES: &[&str] = &["mtp.safetensors", "model-mtp.safetensors"]; + /// Load a tokenizer from a model directory. pub fn load_tokenizer>(model_dir: P) -> Result { let file = model_dir.as_ref().join("tokenizer.json"); @@ -1203,6 +1274,33 @@ pub fn load_quantized_safetensors_weights_with_prefix( /// Collect safetensors file paths from a model directory. fn collect_safetensors_files(model_path: &Path) -> Result, ModelError> { + fn existing_auxiliary_files(model_path: &Path) -> Vec { + AUXILIARY_SAFETENSORS_FILES + .iter() + .map(|file_name| model_path.join(file_name)) + .filter(|file_path| file_path.exists()) + .collect() + } + + fn append_existing_auxiliary_files( + model_path: &Path, + files: &mut Vec, + ) -> Result<(), ModelError> { + let auxiliary_files = existing_auxiliary_files(model_path); + if auxiliary_files.len() > 1 { + return Err(ModelError::UnsupportedModel( + "ambiguous MTP sidecars: both mtp.safetensors and model-mtp.safetensors are present; remove one".to_owned(), + )); + } + + if let Some(file_path) = auxiliary_files.into_iter().next() { + if !files.iter().any(|path| path == &file_path) { + files.push(file_path); + } + } + Ok(()) + } + let index_path = model_path.join("model.safetensors.index.json"); if index_path.exists() { let json = std::fs::read_to_string(&index_path)?; @@ -1213,11 +1311,16 @@ fn collect_safetensors_files(model_path: &Path) -> Result( + &mut self, + inputs: &Array, + mask: Option<&Array>, + cache: &mut Vec>, + ) -> Result { + self.language_model.forward_all_logits(inputs, mask, cache) + } + /// Encode an image through the vision tower and projector. /// /// Input: `pixel_values` with shape `[1, H, W, 3]` (NHWC). diff --git a/crates/higgs-models/src/phi3.rs b/crates/higgs-models/src/phi3.rs index c2a49ebe..e3352010 100644 --- a/crates/higgs-models/src/phi3.rs +++ b/crates/higgs-models/src/phi3.rs @@ -520,6 +520,23 @@ impl Phi3CausalLM { } } + #[allow(non_snake_case)] + pub fn forward_all_logits( + &mut self, + inputs: &Array, + mask: Option<&Array>, + kv_cache: &mut Vec>, + ) -> Result { + let out = self.forward_hidden(inputs, mask, kv_cache)?; + match self.lm_head.as_mut() { + Some(head) => head.forward(&out), + None => match &mut self.model.embed_tokens { + MaybeQuantized::Original(embed) => embed.as_linear(&out), + MaybeQuantized::Quantized(q_embed) => q_embed.as_linear(&out), + }, + } + } + pub fn forward_hidden( &mut self, inputs: &Array, diff --git a/crates/higgs-models/src/qwen3_moe.rs b/crates/higgs-models/src/qwen3_moe.rs index 0480275a..91a5075d 100644 --- a/crates/higgs-models/src/qwen3_moe.rs +++ b/crates/higgs-models/src/qwen3_moe.rs @@ -557,6 +557,21 @@ impl Qwen3MoeCausalLM { None => self.model.embed_tokens.as_linear(&lm_input), } } + + /// Forward pass producing logits for every input position. + #[allow(non_snake_case)] + pub fn forward_all_logits( + &mut self, + inputs: &Array, + mask: Option<&Array>, + kv_cache: &mut Vec>, + ) -> Result { + let h = self.forward_hidden(inputs, mask, kv_cache)?; + match self.lm_head.as_ref() { + Some(head) => head.forward(&h), + None => self.model.embed_tokens.as_linear(&h), + } + } } // --------------------------------------------------------------------------- diff --git a/crates/higgs-models/src/qwen3_next.rs b/crates/higgs-models/src/qwen3_next.rs index 3b0b1723..ba3e5da4 100644 --- a/crates/higgs-models/src/qwen3_next.rs +++ b/crates/higgs-models/src/qwen3_next.rs @@ -207,6 +207,14 @@ pub struct Qwen3NextModelArgs { /// 0 = no MTP head, 1 = one transformer layer for next-next-token prediction. #[serde(default)] pub mtp_num_hidden_layers: i32, + + /// Use dense projection tensors for the MTP head. + /// + /// Some Qwen3.5 MTP sidecars ship full-precision `mtp.*.weight` tensors + /// rather than quantized `weight/scales/biases` triples. This is set by the + /// loader after inspecting checkpoint keys; it is not expected in configs. + #[serde(default)] + pub use_dense_mtp: bool, } // --------------------------------------------------------------------------- @@ -215,12 +223,12 @@ pub struct Qwen3NextModelArgs { type QuantizedParams = (Param, Param, Param); -pub(crate) fn init_quantized_params() -> Result { - Ok(( - Param::new(Array::zeros::(&[1])?), - Param::new(Array::zeros::(&[1])?), - Param::new(Array::zeros::(&[1])?), - )) +pub(crate) fn init_quantized_params() -> QuantizedParams { + fn placeholder() -> Param { + Param::new(Array::from_slice(&[0.0_f32], &[1])) + } + + (placeholder(), placeholder(), placeholder()) } pub(crate) fn quantized_forward( @@ -249,8 +257,9 @@ pub(crate) struct QLinear { } impl QLinear { + #[allow(clippy::unnecessary_wraps)] pub(crate) fn new(group_size: i32, bits: i32) -> Result { - let (weight, scales, biases) = init_quantized_params()?; + let (weight, scales, biases) = init_quantized_params(); Ok(Self { weight, scales, @@ -288,6 +297,42 @@ impl QLinear { } } +/// Dense linear layer with a single weight tensor and no bias. +#[derive(Debug, Clone, ModuleParameters)] +struct DenseLinearNoBias { + #[param] + weight: Param, +} + +fn dense_linear_no_bias_forward(weight: &Array, x: &Array) -> Result { + let shape = x.shape().to_vec(); + let in_features = *shape + .last() + .ok_or_else(|| Exception::custom("empty input"))?; + let batch: i32 = shape.iter().take(shape.len() - 1).product(); + let x2d = x.reshape(&[batch, in_features])?; + let w = weight.as_dtype(x.dtype())?; + let out2d = x2d.matmul(&w.transpose()?)?; + let out_features = *out2d.shape().last().unwrap_or(&0); + let mut out_shape = shape; + if let Some(last) = out_shape.last_mut() { + *last = out_features; + } + out2d.reshape(&out_shape) +} + +impl DenseLinearNoBias { + fn new() -> Self { + Self { + weight: Param::new(Array::from_slice(&[0.0_f32], &[1])), + } + } + + fn forward(&self, x: &Array) -> Result { + dense_linear_no_bias_forward(&self.weight, x) + } +} + /// Quantized embedding stored as raw weight/scales/biases arrays. #[derive(Debug, Clone, ModuleParameters)] pub(crate) struct QEmbedding { @@ -302,8 +347,9 @@ pub(crate) struct QEmbedding { } impl QEmbedding { + #[allow(clippy::unnecessary_wraps)] pub(crate) fn new(group_size: i32, bits: i32) -> Result { - let (weight, scales, biases) = init_quantized_params()?; + let (weight, scales, biases) = init_quantized_params(); Ok(Self { weight, scales, @@ -1454,10 +1500,8 @@ impl Qwen3NextAttention { let offset = cache.offset(); queries = apply_rope(&queries, &self.rope, offset)?; keys = apply_rope(&keys, &self.rope, offset)?; - let view = cache.update_and_view(keys, values)?; let try_tq_decode = mask.is_none() && L == 1; - let output = match view { crate::cache::KvCacheView::TurboQuant(tq_view) if try_tq_decode => { let scores = tq_view.decode_scores(&queries, self.num_attention_heads)?; @@ -1484,7 +1528,6 @@ impl Qwen3NextAttention { .reshape(&[B, L, -1])? } }; - if L == 1 && async_layer_state_eval_enabled() { mlx_rs::transforms::async_eval(cache.eval_targets())?; } @@ -1534,6 +1577,149 @@ impl Qwen3NextAttention { } } +#[derive(Debug, Clone, ModuleParameters)] +struct DenseQwen3NextAttention { + #[param] + q_proj: DenseLinearNoBias, + #[param] + k_proj: DenseLinearNoBias, + #[param] + v_proj: DenseLinearNoBias, + #[param] + o_proj: DenseLinearNoBias, + #[param] + q_norm: nn::RmsNorm, + #[param] + k_norm: nn::RmsNorm, + #[param] + rope: nn::Rope, + num_attention_heads: i32, + num_key_value_heads: i32, + scale: f32, +} + +impl DenseQwen3NextAttention { + fn new(args: &Qwen3NextModelArgs) -> Result { + let head_dim = args.head_dim; + let head_dim_f32 = f32::from( + i16::try_from(head_dim).map_err(|_| Exception::custom("head_dim out of i16 range"))?, + ); + let scale = head_dim_f32.sqrt().recip(); + #[allow(clippy::as_conversions, clippy::cast_possible_truncation)] + let partial_dim = (head_dim_f32 * args.partial_rotary_factor).round() as i32; + + Ok(Self { + q_proj: DenseLinearNoBias::new(), + k_proj: DenseLinearNoBias::new(), + v_proj: DenseLinearNoBias::new(), + o_proj: DenseLinearNoBias::new(), + q_norm: nn::RmsNormBuilder::new(head_dim) + .eps(args.rms_norm_eps) + .build()?, + k_norm: nn::RmsNormBuilder::new(head_dim) + .eps(args.rms_norm_eps) + .build()?, + rope: nn::RopeBuilder::new(partial_dim) + .traditional(false) + .base(args.rope_theta) + .scale(1.0) + .build() + .map_err(|e| Exception::custom(format!("Failed to build RoPE: {e}")))?, + num_attention_heads: args.num_attention_heads, + num_key_value_heads: args.num_key_value_heads, + scale, + }) + } + + #[allow(non_snake_case)] + fn forward( + &mut self, + x: &Array, + mask: Option<&AttentionMask>, + cache: &mut SteppingKeyValueCache, + ) -> Result { + let shape = x.shape(); + let B = *shape + .first() + .ok_or_else(|| Exception::custom("Input must have >= 2 dims"))?; + let L = *shape + .get(1) + .ok_or_else(|| Exception::custom("Input must have >= 2 dims"))?; + + let q_proj_output = self.q_proj.forward(x)?; + let q_reshaped = q_proj_output.reshape(&[B, L, self.num_attention_heads, -1])?; + let q_halves = q_reshaped.split(2, Some(-1))?; + let queries_pre = q_halves + .first() + .ok_or_else(|| Exception::custom("split produced empty result"))?; + let gate = q_halves + .get(1) + .ok_or_else(|| Exception::custom("split produced empty result"))? + .reshape(&[B, L, -1])?; + + let keys_raw = self.k_proj.forward(x)?; + let values_raw = self.v_proj.forward(x)?; + + let mut queries = self + .q_norm + .forward(queries_pre)? + .transpose_axes(&[0, 2, 1, 3])?; + let mut keys = self + .k_norm + .forward(&keys_raw.reshape(&[B, L, self.num_key_value_heads, -1])?)? + .transpose_axes(&[0, 2, 1, 3])?; + let values = values_raw + .reshape(&[B, L, self.num_key_value_heads, -1])? + .transpose_axes(&[0, 2, 1, 3])?; + + let offset = cache.offset(); + queries = apply_rope(&queries, &self.rope, offset)?; + keys = apply_rope(&keys, &self.rope, offset)?; + + let view = cache.update_and_view(keys, values)?; + let try_tq_decode = mask.is_none() && L == 1; + + let output = match view { + crate::cache::KvCacheView::TurboQuant(tq_view) if try_tq_decode => { + let scores = tq_view.decode_scores(&queries, self.num_attention_heads)?; + let scale_arr = Array::from_f32(self.scale).as_dtype(scores.dtype())?; + let weights = ops::softmax_axis(&scores.multiply(&scale_arr)?, -1, true)?; + tq_view + .decode_values(&weights, self.num_attention_heads)? + .transpose_axes(&[0, 2, 1, 3])? + .reshape(&[B, L, -1])? + } + other @ (crate::cache::KvCacheView::Dense { .. } + | crate::cache::KvCacheView::TurboQuant(_)) => { + let (cached_keys, cached_values) = other.into_dense()?; + let sdpa_mask = mask.map(fast::ScaledDotProductAttentionMask::from); + fast::scaled_dot_product_attention( + queries, + cached_keys, + cached_values, + self.scale, + sdpa_mask, + None::<&Array>, + )? + .transpose_axes(&[0, 2, 1, 3])? + .reshape(&[B, L, -1])? + } + }; + + if L == 1 && async_layer_state_eval_enabled() { + mlx_rs::transforms::async_eval(cache.eval_targets())?; + } + + let gated = sigmoid_mul(&gate, &output)?; + let out = self.o_proj.forward(&gated)?; + if L == 1 { + mlx_rs::stop_gradient(&out) + } else { + Ok(out) + } + } +} + // --------------------------------------------------------------------------- // Qwen3NextMLP (standard SwiGLU) // --------------------------------------------------------------------------- @@ -1577,6 +1763,33 @@ impl Qwen3NextMLP { } } +#[derive(Debug, Clone, ModuleParameters)] +struct DenseQwen3NextMLP { + #[param] + gate_proj: DenseLinearNoBias, + #[param] + down_proj: DenseLinearNoBias, + #[param] + up_proj: DenseLinearNoBias, +} + +impl DenseQwen3NextMLP { + fn new() -> Self { + Self { + gate_proj: DenseLinearNoBias::new(), + down_proj: DenseLinearNoBias::new(), + up_proj: DenseLinearNoBias::new(), + } + } + + fn forward(&self, x: &Array) -> Result { + let gate_out = self.gate_proj.forward(x)?; + let up_out = self.up_proj.forward(x)?; + let activated = swiglu(&gate_out, &up_out)?; + self.down_proj.forward(&activated) + } +} + // --------------------------------------------------------------------------- // MTP (Multi-Token Prediction) head // --------------------------------------------------------------------------- @@ -1594,6 +1807,20 @@ struct MtpTransformerLayer { mlp: Qwen3NextMLP, } +/// Single dense MTP transformer layer for sidecar checkpoints that store +/// full-precision projection weights without quantization metadata. +#[derive(Debug, Clone, ModuleParameters)] +struct DenseMtpTransformerLayer { + #[param] + self_attn: DenseQwen3NextAttention, + #[param] + input_layernorm: nn::RmsNorm, + #[param] + post_attention_layernorm: nn::RmsNorm, + #[param] + mlp: DenseQwen3NextMLP, +} + /// Multi-Token Prediction head. /// /// Predicts the token at position t+2 given: @@ -1616,6 +1843,20 @@ pub struct MtpHead { norm: nn::RmsNorm, } +#[derive(Debug, Clone, ModuleParameters)] +struct DenseMtpHead { + #[param] + pre_fc_norm_hidden: nn::RmsNorm, + #[param] + pre_fc_norm_embedding: nn::RmsNorm, + #[param] + fc: MtpFc, + #[param] + layers: Vec, + #[param] + norm: nn::RmsNorm, +} + /// MTP fusion projection — kept in full precision (fp16) for accuracy. /// /// mlx-lm's `quant_predicate` excludes `mtp.fc` from quantization because @@ -1634,22 +1875,7 @@ impl MtpFc { } fn forward(&self, x: &Array) -> Result { - // Dense matmul: x @ W^T (weight shape [out_features, in_features]) - // Reshape to 2D for matmul, then restore batch dims. - let shape = x.shape().to_vec(); - let in_features = *shape - .last() - .ok_or_else(|| Exception::custom("empty input"))?; - let batch: i32 = shape.iter().take(shape.len() - 1).product(); - let x2d = x.reshape(&[batch, in_features])?; - let w = (*self.weight).as_dtype(x.dtype())?; - let out2d = x2d.matmul(&w.transpose()?)?; - let out_features = *out2d.shape().last().unwrap_or(&0); - let mut out_shape = shape; - if let Some(last) = out_shape.last_mut() { - *last = out_features; - } - out2d.reshape(&out_shape) + dense_linear_no_bias_forward(&self.weight, x) } } @@ -1689,6 +1915,42 @@ impl MtpHead { } } +impl DenseMtpHead { + fn new(args: &Qwen3NextModelArgs) -> Result { + let n = usize::try_from(args.mtp_num_hidden_layers) + .map_err(|_| Exception::custom("mtp_num_hidden_layers must be non-negative"))?; + + let layers = (0..n) + .map(|_| { + Ok(DenseMtpTransformerLayer { + self_attn: DenseQwen3NextAttention::new(args)?, + input_layernorm: nn::RmsNormBuilder::new(args.hidden_size) + .eps(args.rms_norm_eps) + .build()?, + post_attention_layernorm: nn::RmsNormBuilder::new(args.hidden_size) + .eps(args.rms_norm_eps) + .build()?, + mlp: DenseQwen3NextMLP::new(), + }) + }) + .collect::, Exception>>()?; + + Ok(Self { + pre_fc_norm_hidden: nn::RmsNormBuilder::new(args.hidden_size) + .eps(args.rms_norm_eps) + .build()?, + pre_fc_norm_embedding: nn::RmsNormBuilder::new(args.hidden_size) + .eps(args.rms_norm_eps) + .build()?, + fc: MtpFc::new()?, + layers, + norm: nn::RmsNormBuilder::new(args.hidden_size) + .eps(args.rms_norm_eps) + .build()?, + }) + } +} + // --------------------------------------------------------------------------- // SwitchMLP weights (stacked expert weights for MoE) // --------------------------------------------------------------------------- @@ -2237,40 +2499,51 @@ impl GatedDeltaNet { let conv_kernel_size = args.linear_conv_kernel_dim; let use_sep = args.use_separate_gdn_projections; + let in_proj_qkvz = QLinear::new(ql, qb)?; + let in_proj_ba = QLinear::new(ql, qb)?; + let in_proj_qkv = if use_sep { + Some(QLinear::new(ql, qb)?) + } else { + None + }; + let in_proj_z = if use_sep { + Some(QLinear::new(ql, qb)?) + } else { + None + }; + let in_proj_a = if use_sep { + Some(QLinear::new(ql, qb)?) + } else { + None + }; + let in_proj_b = if use_sep { + Some(QLinear::new(ql, qb)?) + } else { + None + }; + let conv1d = nn::Conv1dBuilder::new(conv_dim, conv_dim, conv_kernel_size) + .bias(false) + .groups(conv_dim) + .padding(0) + .build()?; + let norm = nn::RmsNormBuilder::new(head_v_dim) + .eps(args.rms_norm_eps) + .build()?; + let out_proj = QLinear::new(ql, qb)?; + let a_log = Param::new(Array::zeros::(&[num_v_heads])?); + let dt_bias = Param::new(Array::zeros::(&[num_v_heads])?); Ok(Self { - in_proj_qkvz: QLinear::new(ql, qb)?, - in_proj_ba: QLinear::new(ql, qb)?, - in_proj_qkv: if use_sep { - Some(QLinear::new(ql, qb)?) - } else { - None - }, - in_proj_z: if use_sep { - Some(QLinear::new(ql, qb)?) - } else { - None - }, - in_proj_a: if use_sep { - Some(QLinear::new(ql, qb)?) - } else { - None - }, - in_proj_b: if use_sep { - Some(QLinear::new(ql, qb)?) - } else { - None - }, - conv1d: nn::Conv1dBuilder::new(conv_dim, conv_dim, conv_kernel_size) - .bias(false) - .groups(conv_dim) - .padding(0) - .build()?, - norm: nn::RmsNormBuilder::new(head_v_dim) - .eps(args.rms_norm_eps) - .build()?, - out_proj: QLinear::new(ql, qb)?, - A_log: Param::new(Array::zeros::(&[num_v_heads])?), - dt_bias: Param::new(Array::zeros::(&[num_v_heads])?), + in_proj_qkvz, + in_proj_ba, + in_proj_qkv, + in_proj_z, + in_proj_a, + in_proj_b, + conv1d, + norm, + out_proj, + A_log: a_log, + dt_bias, num_k_heads, num_v_heads, head_k_dim, @@ -2284,20 +2557,28 @@ impl GatedDeltaNet { i16::try_from(head_k_dim) .map_err(|_| Exception::custom("head_k_dim out of i16 range"))?, ); - let s = dim_f32.sqrt().recip(); - let w = Array::ones::(&[head_k_dim])?.multiply(Array::from_f32(s * s))?; - w.eval()?; - w + let value = dim_f32.sqrt().recip().powi(2); + let values = vec![ + value; + usize::try_from(head_k_dim).map_err(|_| Exception::custom( + "head_k_dim out of usize range" + ))? + ]; + Array::from_slice(&values, &[head_k_dim]) }, qk_norm_weight_k: { let dim_f32 = f32::from( i16::try_from(head_k_dim) .map_err(|_| Exception::custom("head_k_dim out of i16 range"))?, ); - let s = dim_f32.sqrt().recip(); - let w = Array::ones::(&[head_k_dim])?.multiply(Array::from_f32(s))?; - w.eval()?; - w + let value = dim_f32.sqrt().recip(); + let values = vec![ + value; + usize::try_from(head_k_dim).map_err(|_| Exception::custom( + "head_k_dim out of usize range" + ))? + ]; + Array::from_slice(&values, &[head_k_dim]) }, conv_weight_t: None, }) @@ -2363,6 +2644,46 @@ impl GatedDeltaNet { silu_direct(&conv_flat.reshape(&[batch, 1, self.conv_dim])?) } + fn chronological_conv_state( + &self, + cache: &mut ArraysCache, + batch: i32, + dtype: Dtype, + ) -> Result { + let n_keep = self.conv_kernel_size - 1; + let Some(state) = cache.conv_state.take() else { + return ops::zeros_dtype(&[batch, n_keep, self.conv_dim], dtype); + }; + + if n_keep <= 0 { + return Ok(state); + } + + let available = cache.offset.clamp(0, n_keep); + if available == 0 { + return Ok(state); + } + if available == n_keep && cache.conv_pos == n_keep - 1 { + return Ok(state); + } + + let start = (cache.conv_pos - available + 1).rem_euclid(n_keep); + let ordered_tail = if start + available <= n_keep { + state.index((.., start..start + available, ..)) + } else { + let first = state.index((.., start.., ..)); + let second = state.index((.., ..(start + available - n_keep), ..)); + ops::concatenate_axis(&[&first, &second], 1)? + }; + + if available == n_keep { + return Ok(ordered_tail); + } + + let pad = ops::zeros_dtype(&[batch, n_keep - available, self.conv_dim], state.dtype())?; + ops::concatenate_axis(&[&pad, &ordered_tail], 1) + } + #[allow(non_snake_case, clippy::too_many_lines)] fn forward( &mut self, @@ -2441,10 +2762,7 @@ impl GatedDeltaNet { let conv_out = if S == 1 { self.decode_conv1d_step(&mixed_qkv, cache, B)? } else { - let conv_state = match cache.conv_state.take() { - Some(state) => state, - None => ops::zeros_dtype(&[B, n_keep, self.conv_dim], inputs.dtype())?, - }; + let conv_state = self.chronological_conv_state(cache, B, inputs.dtype())?; let conv_input = ops::concatenate_axis(&[&conv_state, &mixed_qkv], 1)?; let conv_input_len = *conv_input .shape() @@ -3015,6 +3333,8 @@ pub struct Qwen3NextCausalLM { lm_head: Option, #[param] mtp: Option, + #[param] + dense_mtp: Option, } // Manual RoPE implementation for arbitrary positions @@ -3128,17 +3448,23 @@ impl Qwen3NextCausalLM { } else { Some(QLinear::new(ql, qb)?) }; - let mtp = if args.mtp_num_hidden_layers > 0 { + let mtp = if args.mtp_num_hidden_layers > 0 && !args.use_dense_mtp { Some(MtpHead::new(&args, ql, qb)?) } else { None }; + let dense_mtp = if args.mtp_num_hidden_layers > 0 && args.use_dense_mtp { + Some(DenseMtpHead::new(&args)?) + } else { + None + }; Ok(Self { args, model, lm_head, mtp, + dense_mtp, }) } @@ -3323,11 +3649,10 @@ impl Qwen3NextCausalLM { h = h2.add(mlp_out)?; } - // Eval every 8 layers during prefill to bound lazy graph size. - // Without this, 40 layers × ~15 ops × T tokens accumulates a huge - // graph that MLX must analyze at once, increasing scheduler overhead - // and peak memory. Decode (T=1) is unaffected. - if T > 1 && (layer_idx + 1) % 8 == 0 { + // Eval every 8 layers during long prefill chunks to bound lazy + // graph size. Short speculative verifier windows are intentionally + // left fused; otherwise MTP pays several eval barriers per cycle. + if should_eval_between_prefill_layers(T, layer_idx) { mlx_rs::transforms::eval([&h])?; } } @@ -3489,18 +3814,22 @@ impl Qwen3NextCausalLM { /// Whether this model has an MTP head loaded. pub const fn has_mtp(&self) -> bool { - self.mtp.is_some() + self.mtp.is_some() || self.dense_mtp.is_some() } /// Create a fresh KV cache for the MTP head (one entry per MTP layer). /// Returns `None` if the model has no MTP head. pub fn make_mtp_cache(&self) -> Option> { - self.mtp.as_ref().map(|mtp| { - mtp.layers - .iter() + let layer_count = self + .mtp + .as_ref() + .map(|mtp| mtp.layers.len()) + .or_else(|| self.dense_mtp.as_ref().map(|mtp| mtp.layers.len()))?; + Some( + (0..layer_count) .map(|_| SteppingKeyValueCache::new()) - .collect() - }) + .collect(), + ) } /// Look up the embedding for a token id. Shape: `[1, 1, hidden_size]`. @@ -3511,6 +3840,37 @@ impl Qwen3NextCausalLM { self.model.embed_tokens.forward(&ids) } + fn embed_tokens_from_ids(&self, token_ids: &[u32]) -> Result { + let token_i32s: Vec = token_ids + .iter() + .map(|&token_id| { + i32::try_from(token_id).map_err(|_| Exception::custom("token_id exceeds i32 range")) + }) + .collect::>()?; + let len = i32::try_from(token_i32s.len()) + .map_err(|_| Exception::custom("token id batch exceeds i32 range"))?; + let ids_array = Array::from_slice(&token_i32s, &[1, len]); + self.model.embed_tokens.forward(&ids_array) + } + + fn mtp_attention_mask( + seq_len: i32, + mtp_cache: &[SteppingKeyValueCache], + ) -> Result, Exception> { + if seq_len <= 1 { + return Ok(None); + } + let offset = mtp_cache.first().map_or(0, SteppingKeyValueCache::offset); + if offset > 0 { + Ok(Some(AttentionMask::Array(create_causal_mask( + seq_len, + Some(offset), + )?))) + } else { + Ok(Some(AttentionMask::Causal)) + } + } + /// Run the MTP head to produce draft logits for position t+2. /// /// - `hidden` — backbone hidden state at position t, shape `[B, 1, D]`. @@ -3524,7 +3884,7 @@ impl Qwen3NextCausalLM { next_token_id: u32, mtp_cache: &mut [SteppingKeyValueCache], ) -> Result { - if self.mtp.is_none() { + if !self.has_mtp() { return Err(Exception::custom("MTP head not loaded")); } @@ -3532,12 +3892,7 @@ impl Qwen3NextCausalLM { let next_embed = self.embed_token(next_token_id)?; // Scope the mutable borrow: run MTP forward, defer lm_head projection. - Ok({ - let mtp = self - .mtp - .as_mut() - .ok_or_else(|| Exception::custom("MTP head not loaded"))?; - + if let Some(mtp) = self.mtp.as_mut() { let h_norm = mtp.pre_fc_norm_hidden.forward(hidden)?; let e_norm = mtp.pre_fc_norm_embedding.forward(&next_embed)?; let concat = ops::concatenate_axis(&[&e_norm, &h_norm], -1)?; @@ -3552,8 +3907,29 @@ impl Qwen3NextCausalLM { x = h2.add(mlp_out)?; } - mtp.norm.forward(&x)? - }) + return mtp.norm.forward(&x); + } + + let mtp = self + .dense_mtp + .as_mut() + .ok_or_else(|| Exception::custom("MTP head not loaded"))?; + + let h_norm = mtp.pre_fc_norm_hidden.forward(hidden)?; + let e_norm = mtp.pre_fc_norm_embedding.forward(&next_embed)?; + let concat = ops::concatenate_axis(&[&e_norm, &h_norm], -1)?; + let mut x = mtp.fc.forward(&concat)?; + + for (layer, kv) in mtp.layers.iter_mut().zip(mtp_cache.iter_mut()) { + let normed = layer.input_layernorm.forward(&x)?; + let attn_out = layer.self_attn.forward(&normed, None, kv)?; + let h2 = x.add(attn_out)?; + let normed_post = layer.post_attention_layernorm.forward(&h2)?; + let mlp_out = layer.mlp.forward(&normed_post)?; + x = h2.add(mlp_out)?; + } + + mtp.norm.forward(&x) } /// Run the MTP head to produce draft logits for position t+2. @@ -3569,13 +3945,29 @@ impl Qwen3NextCausalLM { next_token_id: u32, mtp_cache: &mut [SteppingKeyValueCache], ) -> Result { + let (_, logits) = self.mtp_draft_with_hidden(hidden, next_token_id, mtp_cache)?; + Ok(logits) + } + + /// Run the MTP head and return both its hidden state and draft logits. + /// + /// The hidden state is useful for chained speculative drafting. Final + /// committed MTP cache state is still replayed with backbone hidden states + /// after verification. + pub fn mtp_draft_with_hidden( + &mut self, + hidden: &Array, + next_token_id: u32, + mtp_cache: &mut [SteppingKeyValueCache], + ) -> Result<(Array, Array), Exception> { let normed = self.mtp_step_hidden(hidden, next_token_id, mtp_cache)?; // Now lm_head/embed_tokens can be borrowed immutably. - match self.lm_head.as_ref() { + let logits = match self.lm_head.as_ref() { Some(head) => head.forward(&normed), None => self.model.embed_tokens.as_linear(&normed), - } + }?; + Ok((normed, logits)) } /// Advance the MTP cache for a newly accepted token without computing logits. @@ -3589,6 +3981,114 @@ impl Qwen3NextCausalLM { Ok(()) } + /// Advance the MTP cache for multiple accepted tokens in one sequence pass. + pub fn mtp_advance_many( + &mut self, + hidden: &Array, + next_token_ids: &[u32], + mtp_cache: &mut [SteppingKeyValueCache], + ) -> Result<(), Exception> { + if next_token_ids.is_empty() { + return Ok(()); + } + if !self.has_mtp() { + return Err(Exception::custom("MTP head not loaded")); + } + + let seq_len = i32::try_from(next_token_ids.len()) + .map_err(|_| Exception::custom("MTP advance token batch exceeds i32 range"))?; + let expected_layers = self + .mtp + .as_ref() + .map(|mtp| mtp.layers.len()) + .or_else(|| self.dense_mtp.as_ref().map(|mtp| mtp.layers.len())) + .ok_or_else(|| Exception::custom("MTP head not loaded"))?; + Self::validate_mtp_advance_many_inputs(hidden, mtp_cache, expected_layers, seq_len)?; + + let next_embed = self.embed_tokens_from_ids(next_token_ids)?; + let mask = Self::mtp_attention_mask(seq_len, mtp_cache)?; + let mask_ref = mask.as_ref(); + + if let Some(mtp) = self.mtp.as_mut() { + let h_norm = mtp.pre_fc_norm_hidden.forward(hidden)?; + let e_norm = mtp.pre_fc_norm_embedding.forward(&next_embed)?; + let concat = ops::concatenate_axis(&[&e_norm, &h_norm], -1)?; + let mut x = mtp.fc.forward(&concat)?; + + for (layer, kv) in mtp.layers.iter_mut().zip(mtp_cache.iter_mut()) { + let normed = layer.input_layernorm.forward(&x)?; + let attn_out = layer.self_attn.forward(&normed, mask_ref, kv)?; + let h2 = x.add(attn_out)?; + let normed_post = layer.post_attention_layernorm.forward(&h2)?; + let mlp_out = layer.mlp.forward(&normed_post)?; + x = h2.add(mlp_out)?; + } + + let _ = mtp.norm.forward(&x)?; + return Ok(()); + } + + let mtp = self + .dense_mtp + .as_mut() + .ok_or_else(|| Exception::custom("MTP head not loaded"))?; + + let h_norm = mtp.pre_fc_norm_hidden.forward(hidden)?; + let e_norm = mtp.pre_fc_norm_embedding.forward(&next_embed)?; + let concat = ops::concatenate_axis(&[&e_norm, &h_norm], -1)?; + let mut x = mtp.fc.forward(&concat)?; + + for (layer, kv) in mtp.layers.iter_mut().zip(mtp_cache.iter_mut()) { + let normed = layer.input_layernorm.forward(&x)?; + let attn_out = layer.self_attn.forward(&normed, mask_ref, kv)?; + let h2 = x.add(attn_out)?; + let normed_post = layer.post_attention_layernorm.forward(&h2)?; + let mlp_out = layer.mlp.forward(&normed_post)?; + x = h2.add(mlp_out)?; + } + + let _ = mtp.norm.forward(&x)?; + Ok(()) + } + + fn validate_mtp_advance_many_inputs( + hidden: &Array, + mtp_cache: &[SteppingKeyValueCache], + expected_layers: usize, + seq_len: i32, + ) -> Result<(), Exception> { + Self::validate_mtp_advance_many_shape( + hidden.shape(), + mtp_cache.len(), + expected_layers, + seq_len, + ) + } + + fn validate_mtp_advance_many_shape( + hidden_shape: &[i32], + cache_layers: usize, + expected_layers: usize, + seq_len: i32, + ) -> Result<(), Exception> { + if cache_layers != expected_layers { + return Err(Exception::custom(format!( + "mtp_cache length ({cache_layers}) must match MTP layer count ({expected_layers})" + ))); + } + + let hidden_seq_len = *hidden_shape + .get(1) + .ok_or_else(|| Exception::custom("hidden must be [B, T, D]"))?; + if hidden_seq_len != seq_len { + return Err(Exception::custom(format!( + "hidden sequence length ({hidden_seq_len}) must match next_token_ids length ({seq_len})" + ))); + } + + Ok(()) + } + /// Forward pass returning BOTH raw hidden states and logits for all positions. /// /// Used by MTP speculative decode: the verify pass needs **raw** (pre-norm) @@ -3613,6 +4113,31 @@ impl Qwen3NextCausalLM { } } +const PREFILL_LAYER_EVAL_INTERVAL: usize = 8; +const PREFILL_LAYER_EVAL_MIN_SEQ_LEN: i32 = 17; + +const fn should_eval_between_prefill_layers(seq_len: i32, layer_idx: usize) -> bool { + seq_len >= PREFILL_LAYER_EVAL_MIN_SEQ_LEN + && (layer_idx + 1).is_multiple_of(PREFILL_LAYER_EVAL_INTERVAL) +} + +#[cfg(test)] +mod prefill_eval_tests { + use super::should_eval_between_prefill_layers; + + #[test] + fn skips_layer_eval_barriers_for_short_speculative_windows() { + assert!(!should_eval_between_prefill_layers(3, 7)); + assert!(!should_eval_between_prefill_layers(8, 7)); + } + + #[test] + fn keeps_layer_eval_barriers_for_long_prefill_chunks() { + assert!(should_eval_between_prefill_layers(128, 7)); + assert!(!should_eval_between_prefill_layers(128, 6)); + } +} + // --------------------------------------------------------------------------- // Loading // --------------------------------------------------------------------------- @@ -3687,7 +4212,50 @@ where ))) } -fn checkpoint_has_mtp_weights(model_path: &Path) -> Result { +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum MtpWeightLayout { + None, + Quantized, + Dense, +} + +fn is_mtp_key(key: &str) -> bool { + key.starts_with("mtp.") || key.contains(".mtp.") +} + +fn mtp_weight_layout_from_keys<'a>(keys: impl IntoIterator) -> MtpWeightLayout { + let mut has_mtp = false; + let mut has_unprefixed_mtp = false; + let mut has_quantized_aux = false; + + for key in keys { + if !is_mtp_key(key) { + continue; + } + has_mtp = true; + has_unprefixed_mtp |= key.starts_with("mtp."); + has_quantized_aux |= key.ends_with(".scales") || key.ends_with(".biases"); + } + + if has_quantized_aux { + MtpWeightLayout::Quantized + } else if has_unprefixed_mtp { + MtpWeightLayout::Dense + } else if has_mtp { + MtpWeightLayout::Quantized + } else { + MtpWeightLayout::None + } +} + +fn checkpoint_mtp_weight_layout(model_path: &Path) -> Result { + fn safetensors_file_mtp_weight_layout(file_path: &Path) -> Result { + let bytes = std::fs::read(file_path)?; + let metadata = safetensors::SafeTensors::deserialize(&bytes) + .map_err(|e| ModelError::Io(std::io::Error::other(e.to_string())))?; + Ok(mtp_weight_layout_from_keys(metadata.names())) + } + let index_path = model_path.join("model.safetensors.index.json"); if index_path.exists() { let file = std::fs::File::open(index_path)?; @@ -3700,23 +4268,45 @@ fn checkpoint_has_mtp_weights(model_path: &Path) -> Result { "model.safetensors.index.json missing weight_map".into(), )); }; - return Ok(weight_map - .keys() - .any(|key| key.starts_with("mtp.") || key.contains(".mtp."))); + let index_layout = mtp_weight_layout_from_keys(weight_map.keys().map(String::as_str)); + if index_layout != MtpWeightLayout::None { + return Ok(index_layout); + } + + let auxiliary_files: Vec<_> = crate::AUXILIARY_SAFETENSORS_FILES + .iter() + .map(|file_name| model_path.join(file_name)) + .filter(|file_path| file_path.exists()) + .collect(); + if auxiliary_files.len() > 1 { + return Err(ModelError::UnsupportedModel( + "ambiguous MTP sidecars: both mtp.safetensors and model-mtp.safetensors are present; remove one".to_owned(), + )); + } + + for file_path in auxiliary_files { + let layout = safetensors_file_mtp_weight_layout(&file_path)?; + if layout != MtpWeightLayout::None { + return Ok(layout); + } + } + + return Ok(MtpWeightLayout::None); } for file_path in crate::collect_safetensors_files(model_path)? { - let loaded = Array::load_safetensors(&file_path) - .map_err(|e| ModelError::Io(std::io::Error::other(e.to_string())))?; - if loaded - .keys() - .any(|key| key.starts_with("mtp.") || key.contains(".mtp.")) - { - return Ok(true); + let layout = safetensors_file_mtp_weight_layout(&file_path)?; + if layout != MtpWeightLayout::None { + return Ok(layout); } } - Ok(false) + Ok(MtpWeightLayout::None) +} + +#[cfg(test)] +fn checkpoint_has_mtp_weights(model_path: &Path) -> Result { + Ok(checkpoint_mtp_weight_layout(model_path)? != MtpWeightLayout::None) } fn maybe_disable_mtp_without_checkpoint_weights( @@ -3727,8 +4317,16 @@ fn maybe_disable_mtp_without_checkpoint_weights( return Ok(()); } - if checkpoint_has_mtp_weights(model_path)? { - return Ok(()); + match checkpoint_mtp_weight_layout(model_path)? { + MtpWeightLayout::Quantized => { + args.use_dense_mtp = false; + return Ok(()); + } + MtpWeightLayout::Dense => { + args.use_dense_mtp = true; + return Ok(()); + } + MtpWeightLayout::None => {} } tracing::warn!( @@ -4158,6 +4756,60 @@ fn can_concatenate_axis0(a: &Array, b: &Array) -> bool { can_concatenate_axis0_shapes(a_shape, b_shape) } +fn qwen35_checkpoint_param_key(key: &str) -> Option<&str> { + if key.starts_with("mtp.") { + Some(key) + } else { + key.strip_prefix("language_model.") + } +} + +fn dense_mtp_param_key(stripped: &str) -> Option { + stripped + .strip_prefix("mtp.") + .map(|rest| format!("dense_mtp.{rest}")) +} + +fn qwen35_target_param_key( + params: &HashMap, &mut Array>, + stripped: &str, +) -> Option<(String, bool)> { + if params.contains_key(stripped) { + Some((stripped.to_owned(), false)) + } else { + dense_mtp_param_key(stripped) + .filter(|dense_key| params.contains_key(dense_key.as_str())) + .map(|dense_key| (dense_key, true)) + } +} + +fn dense_mtp_rmsnorm_weight_key(stripped: &str) -> bool { + stripped.starts_with("mtp.") + && stripped.ends_with(".weight") + && (stripped.contains(".input_layernorm.") + || stripped.contains(".post_attention_layernorm.") + || stripped.contains(".q_norm.") + || stripped.contains(".k_norm.") + || stripped == "mtp.norm.weight" + || stripped == "mtp.pre_fc_norm_hidden.weight" + || stripped == "mtp.pre_fc_norm_embedding.weight") +} + +fn qwen35_loaded_value( + stripped: &str, + value: Array, + dense_mtp_target: bool, +) -> Result { + if dense_mtp_target && dense_mtp_rmsnorm_weight_key(stripped) { + let one = Array::from_f32(1.0) + .as_dtype(value.dtype()) + .map_err(crate::error::ModelError::Mlx)?; + value.add(&one).map_err(crate::error::ModelError::Mlx) + } else { + Ok(value) + } +} + /// Load Qwen3.5-MoE weights with GDN projection fusion. /// /// Direct weight loader: strip `language_model.` prefix, no rearrangement. @@ -4168,7 +4820,6 @@ fn load_qwen3_5_moe_weights_direct( ) -> Result<(), crate::error::ModelError> { let safetensors_files = crate::collect_safetensors_files(model_path)?; let mut params = model.parameters_mut().flatten(); - let prefix = "language_model."; let mut matched = 0usize; let mut unmatched = Vec::new(); @@ -4177,12 +4828,18 @@ fn load_qwen3_5_moe_weights_direct( .map_err(|e| crate::error::ModelError::Io(std::io::Error::other(e.to_string())))?; for (key, value) in loaded { - let Some(stripped) = key.strip_prefix(prefix) else { + let Some(stripped) = qwen35_checkpoint_param_key(&key) else { unmatched.push(key); continue; }; - if let Some(param) = params.get_mut(stripped) { - **param = value; + if let Some((target_key, dense_mtp_target)) = qwen35_target_param_key(¶ms, stripped) + { + if let Some(param) = params.get_mut(target_key.as_str()) { + **param = qwen35_loaded_value(stripped, value, dense_mtp_target)?; + } else { + unmatched.push(key); + continue; + } matched += 1; } else { unmatched.push(key); @@ -4220,6 +4877,7 @@ fn load_qwen3_5_moe_weights_direct( /// Rearranges flat (qkv,z,b,a) projections to per-head-grouped (qkvz,ba) /// so the model uses the fused 2-dispatch forward path instead of 4 separate. +#[allow(clippy::too_many_lines)] fn load_qwen3_5_moe_weights_fused( model: &mut M, model_path: &Path, @@ -4238,7 +4896,6 @@ fn load_qwen3_5_moe_weights_fused( // Key format: "model.layers.N.linear_attn.in_proj_qkvz.{weight|scales|biases}" let mut gdn_parts: HashMap, Option)> = HashMap::new(); - let prefix = "language_model."; let gdn_remap: &[(&str, &str, &str)] = &[ ("in_proj_qkv", "in_proj_z", "in_proj_qkvz"), ("in_proj_b", "in_proj_a", "in_proj_ba"), @@ -4249,7 +4906,7 @@ fn load_qwen3_5_moe_weights_fused( .map_err(|e| crate::error::ModelError::Io(std::io::Error::other(e.to_string())))?; for (key, value) in loaded { - let Some(stripped) = key.strip_prefix(prefix) else { + let Some(stripped) = qwen35_checkpoint_param_key(&key) else { continue; }; @@ -4277,8 +4934,12 @@ fn load_qwen3_5_moe_weights_fused( } if !handled { - if let Some(param) = params.get_mut(stripped) { - **param = value; + if let Some((target_key, dense_mtp_target)) = + qwen35_target_param_key(¶ms, stripped) + { + if let Some(param) = params.get_mut(target_key.as_str()) { + **param = qwen35_loaded_value(stripped, value, dense_mtp_target)?; + } } } } @@ -4895,6 +5556,47 @@ mod tests { }); } + #[test] + #[ignore = "requires real model weights; placeholder test tensors cannot run MTP forward"] + fn test_mtp_advance_many_appends_accepted_token_states() { + let stream = Stream::new(); + mlx_rs::with_new_default_stream(stream, || { + let mut args = valid_causal_lm_args(); + args.mtp_num_hidden_layers = 1; + let mut model = Qwen3NextCausalLM::new(args).unwrap(); + let mut mtp_cache = model.make_mtp_cache().unwrap(); + let hidden = Array::zeros::(&[1, 2, model.args.hidden_size]).unwrap(); + + model + .mtp_advance_many(&hidden, &[1, 2], &mut mtp_cache) + .unwrap(); + + assert_eq!(mtp_cache[0].offset(), 2); + }); + } + + #[test] + fn test_mtp_advance_many_rejects_cache_layer_mismatch() { + let err = + Qwen3NextCausalLM::validate_mtp_advance_many_shape(&[1, 1, 256], 0, 1, 1).unwrap_err(); + + assert!( + err.to_string().contains("mtp_cache length"), + "unexpected error: {err}" + ); + } + + #[test] + fn test_mtp_advance_many_rejects_hidden_sequence_mismatch() { + let err = + Qwen3NextCausalLM::validate_mtp_advance_many_shape(&[1, 2, 256], 1, 1, 1).unwrap_err(); + + assert!( + err.to_string().contains("hidden sequence length"), + "unexpected error: {err}" + ); + } + #[test] fn test_load_model_args_happy_path() { let dir = tempfile::tempdir().unwrap(); @@ -5599,6 +6301,17 @@ mod tests { ref_state = conv_in.index((.., 1.., ..)); cache.offset += 1; } + + let ordered = gdn + .chronological_conv_state(&mut cache, 1, Dtype::Float16) + .unwrap(); + mlx_rs::transforms::eval([&ordered, &ref_state]).unwrap(); + let diff = ordered.subtract(&ref_state).unwrap().abs().unwrap(); + let max_diff: f32 = diff.max(None).unwrap().item(); + assert!( + max_diff < 2e-3, + "linearized ring-buffer conv state differs from chronological state by {max_diff}" + ); } #[test] @@ -6496,13 +7209,16 @@ mod tests { #[ignore = "benchmark, requires GPU"] fn bench_gather_qmm_loaded_vs_random() { use mlx_rs::Dtype; - let model_dir = "/Users/panbanda/.cache/huggingface/hub/models--mlx-community--Qwen3-Coder-Next-4bit/snapshots/7b9321eabb85ce79625cac3f61ea691e4ea984b5"; - let shard = format!("{}/model-00001-of-00009.safetensors", model_dir); - let path = std::path::Path::new(&shard); - if !path.exists() { - eprintln!("Skipping: model not found"); + let Some(model_dir) = std::env::var_os("HIGGS_QWEN3_NEXT_BENCH_MODEL") else { + eprintln!("Skipping: set HIGGS_QWEN3_NEXT_BENCH_MODEL to a local model directory"); return; - } + }; + let model_dir = std::path::PathBuf::from(model_dir); + let Some(shard) = find_gather_qmm_bench_shard(&model_dir) else { + eprintln!("Skipping: no safetensors shard with switch_mlp gate weights found"); + return; + }; + let path = shard.as_path(); // Load one safetensors shard let loaded = Array::load_safetensors(path).unwrap(); @@ -6609,6 +7325,32 @@ mod tests { ); } + fn find_gather_qmm_bench_shard(model_dir: &std::path::Path) -> Option { + let mut candidates: Vec<_> = std::fs::read_dir(model_dir) + .ok()? + .filter_map(Result::ok) + .map(|entry| entry.path()) + .filter(|path| path.extension().is_some_and(|ext| ext == "safetensors")) + .collect(); + candidates.sort_by(|a, b| { + let a_name = a.file_name().and_then(|name| name.to_str()).unwrap_or(""); + let b_name = b.file_name().and_then(|name| name.to_str()).unwrap_or(""); + let a_shard = a_name.starts_with("model-") && a_name.contains("-of-"); + let b_shard = b_name.starts_with("model-") && b_name.contains("-of-"); + b_shard.cmp(&a_shard).then_with(|| a_name.cmp(b_name)) + }); + + candidates.into_iter().find(|path| { + Array::load_safetensors(path).ok().is_some_and(|loaded| { + loaded.keys().any(|key| { + key.contains("switch_mlp") + && key.contains("gate_proj") + && key.contains(".weight") + }) + }) + }) + } + /// Isolate what causes the module vs inline performance gap. /// Tests three variants at 48 layers: /// A) Module forward with multiply-by-zero attention (baseline slow path) @@ -7471,8 +8213,8 @@ mod tests { }); } - /// Benchmark 36 GDN layers using bare Arrays (matching Python bench_gdn_real_python.py). - /// Isolates GDN ops from the model framework to compare GPU time vs Python. + /// Benchmark 36 GDN layers using bare Arrays. + /// Isolates GDN ops from the model framework to compare direct GPU time. #[test] #[ignore = "requires GPU"] fn bench_gdn_layers() { @@ -12009,13 +12751,17 @@ mod tests { #[test] #[ignore = "requires model files on disk"] fn bench_actual_model_forward() { - let model_path = "/Users/panbanda/.cache/huggingface/hub/models--mlx-community--Qwen3-Coder-Next-4bit/snapshots/7b9321eabb85ce79625cac3f61ea691e4ea984b5"; - if !std::path::Path::new(model_path).exists() { - println!("Model not found at {model_path}, skipping"); + let Some(model_path) = std::env::var_os("HIGGS_QWEN3_NEXT_BENCH_MODEL") else { + println!("Skipping: set HIGGS_QWEN3_NEXT_BENCH_MODEL to a local model directory"); + return; + }; + let model_path = std::path::PathBuf::from(model_path); + if !model_path.exists() { + println!("Model not found at {}, skipping", model_path.display()); return; } - let mut model = load_qwen3_next_model(model_path).unwrap(); + let mut model = load_qwen3_next_model(&model_path).unwrap(); let mut cache: Vec> = Vec::new(); // Prefill with a short prompt @@ -13084,6 +13830,14 @@ mod tests { std::fs::write(dir.join("model.safetensors.index.json"), index).unwrap(); } + fn write_safetensors_file(dir: &std::path::Path, file_name: &str, key: &str) { + let data = [0_u8; 4]; + let tensor = + safetensors::tensor::TensorView::new(safetensors::tensor::Dtype::F32, vec![1], &data) + .unwrap(); + safetensors::serialize_to_file([(key, tensor)], None, &dir.join(file_name)).unwrap(); + } + #[test] fn test_load_qwen35_moe_text_config_moe_sets_decoder_sparse_step() { let dir = tempfile::tempdir().unwrap(); @@ -13120,6 +13874,97 @@ mod tests { assert!(checkpoint_has_mtp_weights(dir.path()).unwrap()); } + #[test] + fn test_checkpoint_has_mtp_weights_detects_auxiliary_mtp_file() { + let dir = tempfile::tempdir().unwrap(); + write_weight_index( + dir.path(), + &["language_model.model.layers.0.input_layernorm.weight"], + ); + write_safetensors_file( + dir.path(), + "model-mtp.safetensors", + "language_model.mtp.layers.0.self_attn.q_proj.weight", + ); + + assert!(checkpoint_has_mtp_weights(dir.path()).unwrap()); + } + + #[test] + fn test_checkpoint_mtp_weight_layout_detects_quantized_indexed_keys() { + let dir = tempfile::tempdir().unwrap(); + write_weight_index( + dir.path(), + &[ + "language_model.mtp.layers.0.self_attn.q_proj.weight", + "language_model.mtp.layers.0.self_attn.q_proj.scales", + "language_model.mtp.layers.0.self_attn.q_proj.biases", + ], + ); + + assert_eq!( + checkpoint_mtp_weight_layout(dir.path()).unwrap(), + MtpWeightLayout::Quantized + ); + } + + #[test] + fn test_checkpoint_mtp_weight_layout_detects_dense_auxiliary_mtp_file() { + let dir = tempfile::tempdir().unwrap(); + write_weight_index( + dir.path(), + &["language_model.model.layers.0.input_layernorm.weight"], + ); + write_safetensors_file( + dir.path(), + "model-mtp.safetensors", + "mtp.layers.0.self_attn.q_proj.weight", + ); + + assert_eq!( + checkpoint_mtp_weight_layout(dir.path()).unwrap(), + MtpWeightLayout::Dense + ); + } + + #[test] + fn test_qwen35_checkpoint_key_accepts_unprefixed_mtp_sidecar() { + assert_eq!( + qwen35_checkpoint_param_key("mtp.layers.0.self_attn.q_proj.weight"), + Some("mtp.layers.0.self_attn.q_proj.weight") + ); + assert_eq!( + qwen35_checkpoint_param_key("language_model.model.layers.0.input_layernorm.weight"), + Some("model.layers.0.input_layernorm.weight") + ); + assert_eq!(qwen35_checkpoint_param_key("vision_tower.foo"), None); + } + + #[test] + fn test_dense_mtp_param_key_remaps_mtp_namespace() { + assert_eq!( + dense_mtp_param_key("mtp.layers.0.self_attn.q_proj.weight").as_deref(), + Some("dense_mtp.layers.0.self_attn.q_proj.weight") + ); + assert_eq!(dense_mtp_param_key("model.layers.0.foo"), None); + } + + #[test] + fn test_dense_mtp_rmsnorm_weight_keys_require_plus_one() { + assert!(dense_mtp_rmsnorm_weight_key( + "mtp.layers.0.input_layernorm.weight" + )); + assert!(dense_mtp_rmsnorm_weight_key( + "mtp.layers.0.self_attn.q_norm.weight" + )); + assert!(dense_mtp_rmsnorm_weight_key( + "mtp.pre_fc_norm_hidden.weight" + )); + assert!(!dense_mtp_rmsnorm_weight_key( + "mtp.layers.0.self_attn.q_proj.weight" + )); + } + #[test] fn test_maybe_disable_mtp_without_checkpoint_weights_turns_off_missing_mtp() { let dir = tempfile::tempdir().unwrap(); diff --git a/crates/higgs-models/src/starcoder2.rs b/crates/higgs-models/src/starcoder2.rs index 2c19d748..2cce4dae 100644 --- a/crates/higgs-models/src/starcoder2.rs +++ b/crates/higgs-models/src/starcoder2.rs @@ -586,6 +586,23 @@ impl Starcoder2CausalLM { } } + #[allow(non_snake_case)] + pub fn forward_all_logits( + &mut self, + inputs: &Array, + mask: Option<&Array>, + kv_cache: &mut Vec>, + ) -> Result { + let out = self.forward_hidden(inputs, mask, kv_cache)?; + match self.lm_head.as_mut() { + Some(head) => head.forward(&out), + None => match &mut self.model.embed_tokens { + MaybeQuantized::Original(embed) => embed.as_linear(&out), + MaybeQuantized::Quantized(q_embed) => q_embed.as_linear(&out), + }, + } + } + pub fn forward_hidden( &mut self, inputs: &Array, diff --git a/crates/higgs-models/src/transformer.rs b/crates/higgs-models/src/transformer.rs index 61ee991d..ba393e9b 100644 --- a/crates/higgs-models/src/transformer.rs +++ b/crates/higgs-models/src/transformer.rs @@ -601,6 +601,17 @@ impl Model { self.apply_lm_head(&last) } + /// Run a forward pass producing logits for every input position. + pub fn forward_all_logits( + &mut self, + inputs: &Array, + mask: Option<&Array>, + kv_cache: &mut Vec>, + ) -> Result { + let hidden = self.forward_hidden(inputs, mask, kv_cache)?; + self.apply_lm_head_all(&hidden) + } + /// Get the hidden size. pub const fn hidden_size(&self) -> i32 { self.args.hidden_size @@ -847,6 +858,16 @@ impl Model { }, } } + + fn apply_lm_head_all(&mut self, hidden: &Array) -> Result { + match self.lm_head.as_mut() { + Some(head) => head.forward(hidden), + None => match &mut self.model.embed_tokens { + MaybeQuantized::Original(embed) => embed.as_linear(hidden), + MaybeQuantized::Quantized(q_embed) => q_embed.as_linear(hidden), + }, + } + } } // --- Loading --- diff --git a/crates/higgs/src/error.rs b/crates/higgs/src/error.rs index 5c5f78ec..58f1c698 100644 --- a/crates/higgs/src/error.rs +++ b/crates/higgs/src/error.rs @@ -72,7 +72,7 @@ impl IntoResponse for ServerError { ( StatusCode::BAD_GATEWAY, "proxy_error", - format!("Upstream provider error: {msg}"), + "Upstream provider error".to_owned(), ) } }; @@ -225,13 +225,26 @@ mod tests { assert_eq!(status, StatusCode::BAD_GATEWAY); assert_eq!(body["error"]["type"].as_str().unwrap(), "proxy_error"); - let message = body["error"]["message"].as_str().unwrap(); - assert!( - message.contains("connection refused"), - "expected proxy detail in message: {message}" + assert_eq!( + body["error"]["message"].as_str().unwrap(), + "Upstream provider error" ); } + #[tokio::test] + async fn test_proxy_error_masks_upstream_details() { + let error = ServerError::ProxyError( + "upstream returned HTTP 401: invalid api key provider-secret-token".to_owned(), + ); + let resp = error.into_response(); + let (status, body) = response_status_and_body(resp).await; + + assert_eq!(status, StatusCode::BAD_GATEWAY); + let message = body["error"]["message"].as_str().unwrap(); + assert_eq!(message, "Upstream provider error"); + assert!(!message.contains("provider-secret-token")); + } + #[tokio::test] async fn test_error_response_content_type_is_json() { let error = ServerError::BadRequest("test".to_owned()); diff --git a/docs/benchmarking.md b/docs/benchmarking.md index 66c08985..b7c09d96 100644 --- a/docs/benchmarking.md +++ b/docs/benchmarking.md @@ -10,20 +10,77 @@ This document collects the benchmark methodology and the benchmark-driven defaul ## MLX Tuning Harness -Use the benchmark harness below to compare five serving iterations on the same local model: +Use the Rust benchmark crate for checked-in benchmark workflows. Python +benchmark scripts are treated as local scratch files and are ignored by git. ```bash -python3 benchmarks/bench_mlx_tuning.py ~/.cache/lm-studio/models/mlx-community/Qwen3.6-35B-A3B-4bit +cargo run --release -p higgs-bench --bin bench_decode -- \ + --model qwen3-1.7B-4bit --port 8899 \ + --max-tokens 200 --warmup 1 --trials 5 \ + --temperature 0 ``` -The harness evaluates: +`bench_decode` evaluates: + +- TTFT through the streaming API +- decode throughput from server-reported token usage +- reproducible JSON/Markdown output with redacted host metadata and local + filesystem paths by default + +## MTP Draft-Depth Sweep + +Use the focused MTP sweep to compare baseline greedy decode with MTP disabled +against draft depths 1, 2, and 3: + +```bash +cargo run --release -p higgs-bench --bin bench_speculative -- \ + --model-path trevon/Qwen3.6-27B-mtp \ + --trials baseline,mtp_default,1,2,3,prompt_lookup \ + --max-tokens 96 --repeats 1 +``` -- TTFT across short, medium, and long prompts -- decode throughput -- short QA accuracy -- long-context retrieval accuracy -- structured-output correctness -- prefix-cache speedup on multi-turn conversations +The sweep sets `temperature=0`, starts a fresh Higgs server per trial, and +reports completion tokens per second plus filtered MTP/prompt-lookup telemetry +for each setting. Use `--model ` to target `benchmarks/models.toml`, or +`--model-path ` for ad-hoc local runs. Local paths are not +printed in the benchmark metadata; pass `--model-name ` when using a +snapshot path whose request model name cannot be derived automatically. + +Set `RUST_LOG=info` when you want the persisted JSON to include Higgs' internal +MTP decode telemetry (`cycles`, drafted tokens, accepted drafts, acceptance +rate, and decode-only tok/s). + +### Qwen3.6 MTP Notes + +The Qwen3Next MTP path mirrors llama.cpp's merged `draft-mtp` design in the +places that matter for speed and correctness: + +- the verifier processes `[confirmed + drafts]` in one backbone batch +- the MTP cache is primed from prompt/first-token backbone hidden states +- accepted draft tokens are advanced into the MTP cache in one sequence pass +- Qwen3Next GDN convolution state is linearized before multi-token verifier + windows, so batched verifier logits match sequential greedy decode + +`HIGGS_MTP_PRIME_PREFILL=0` disables prompt/first-token MTP cache priming for +experiments. `HIGGS_MTP_MIRROR_VERIFY=1` enables full verifier-window MTP cache +mirroring; on the Qwen3.6 27B MTP 8-bit benchmark below it was slightly slower +than the default accepted-prefix replay path, so it remains opt-in. + +Measured on M4 Max 128GB, `temperature=0`, 96 completion tokens, prompt: +`Write a concise technical explanation of speculative decoding for local LLM inference...` + +| Runtime | Mode | Request tok/s | Decode-only tok/s | Speedup vs runtime baseline | +| --- | ---: | ---: | ---: | ---: | +| Higgs | baseline MTP off | 14.32 | n/a | 1.00x | +| Higgs | MTP draft depth 2 | 22.89 | 28.0 | 1.60x request / 1.96x vs request baseline | +| llama.cpp `b1-d374e71` | baseline | n/a | 15.9 | 1.00x | +| llama.cpp `b1-d374e71` | MTP draft depth 1 | n/a | 25.0 | 1.57x | +| llama.cpp `b1-d374e71` | MTP draft depth 2 | n/a | 24.3 | 1.53x | + +The Higgs request-level number includes HTTP and prompt processing; llama.cpp's +CLI line reports generation only. The closest decode-only comparison from this +run is Higgs MTP depth 2 at `28.0 tok/s` versus llama.cpp's best measured MTP +setting at `25.0 tok/s`. ## Iterations @@ -80,7 +137,10 @@ top-level keys: `metadata`, `params`, `results`. ``` `metadata.git_commit` and `git_dirty` are captured at compile time via -the `built` crate; you must rebuild the bench binary to refresh them. +the `built` crate; you must rebuild the bench binary to refresh them. Benchmark +metadata redacts hostnames by default and reduces local absolute paths in +`args`, model refs, and persisted artifact messages to public model IDs, +basenames, or relative `target/bench-results/...` paths. Every binary supports two output formats: @@ -138,6 +198,23 @@ SSE chunk count for backends that don't emit usage. The bench also sends `reasoning: { effort: "none" }` so decode timing reflects time-to-generate, not time-to-visible-answer for thinking-mode models. +### `bench_speculative` + +Starts a fresh Higgs server per speculative mode and compares baseline greedy +decode with MTP draft depths and architecture-neutral prompt lookup. + +```bash +cargo run --release -p higgs-bench --bin bench_speculative -- \ + --model qwen3.6-27B-mtp-8bit \ + --trials baseline,mtp_default,1,2,3,prompt_lookup,prompt_lookup_unchecked \ + --max-tokens 96 --repeats 3 --format markdown +``` + +`results.trials[*].speedup_vs_baseline` is computed from mean completion +tokens/sec when a `baseline` trial appears before the speculative trial. Server +logs are captured under `target/bench-results/bench_speculative/logs/`, and the +JSON result stores only filtered speculative telemetry lines. + ### `bench_summarize` Walks `target/bench-results/`, picks the latest result per diff --git a/docs/configuration.md b/docs/configuration.md index a63b0387..2cab445f 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -38,6 +38,7 @@ This document collects the full CLI, environment, and config-file reference for - `HIGGS_CHUNKED_PREFILL_THRESHOLD` enables chunked prefill above a token threshold. - `HIGGS_CHUNKED_PREFILL_CHUNK_SIZE` controls chunk size during chunked prefill. - `HIGGS_MTP=0|1` overrides the tuning profile's speculative decode choice when conditions allow. +- `HIGGS_MTP_DRAFT_N_MAX` controls the maximum MTP draft tokens per speculative cycle. The default is `3` when MTP is enabled, clamped to `1..=8`. - `HIGGS_CLEAR_CACHE_AFTER_PREFILL` overrides the selected MLX profile behavior for cache clearing. - `HIGGS_TURBOQUANT_MIN_TOKENS` overrides the TurboQuant activation threshold. The default is `2048`. - `HIGGS_EXPERIMENTAL_PAGED_KV=1` enables the experimental paged-KV path. From eb075bfd174c874f530e63abd212efdacc21c1bc Mon Sep 17 00:00:00 2001 From: "renovate[bot]" <29139614+renovate[bot]@users.noreply.github.com> Date: Fri, 29 May 2026 13:50:23 +0000 Subject: [PATCH 23/29] chore(deps): update taiki-e/install-action digest to 873c745 --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index cae64399..e8494485 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -120,7 +120,7 @@ jobs: toolchain: stable - name: Install cargo-llvm-cov - uses: taiki-e/install-action@cca35edeb1d01366c2843b68fc3ca441446d73d3 # v2 + uses: taiki-e/install-action@873c7452cadb7c034694a1282227095d93fbdf92 # v2 with: tool: cargo-llvm-cov From 465a01e6ea36435461b59c5ed84da792fd828fef Mon Sep 17 00:00:00 2001 From: Jonathan Reyes Date: Fri, 29 May 2026 15:25:46 -0600 Subject: [PATCH 24/29] Optimize MTP follow-up modes --- crates/higgs-bench/src/bin/bench_decode.rs | 6 +- .../higgs-bench/src/bin/bench_speculative.rs | 13 +- crates/higgs-bench/src/lib.rs | 30 ++- crates/higgs-bench/src/speculative.rs | 64 ++++++- crates/higgs-engine/src/mlx_tuning.rs | 43 ++++- crates/higgs-engine/src/mtp.rs | 181 ++++++++++++++++-- crates/higgs-engine/src/simple.rs | 158 +++++++++++++-- docs/benchmarking.md | 42 ++-- docs/configuration.md | 4 +- 9 files changed, 469 insertions(+), 72 deletions(-) diff --git a/crates/higgs-bench/src/bin/bench_decode.rs b/crates/higgs-bench/src/bin/bench_decode.rs index b9d97606..7f6f48fe 100644 --- a/crates/higgs-bench/src/bin/bench_decode.rs +++ b/crates/higgs-bench/src/bin/bench_decode.rs @@ -22,8 +22,9 @@ use anyhow::{Context, Result}; use clap::Parser; use futures::StreamExt; use higgs_bench::{ - BenchOutput, ModelInfo, OutputFormat, RunMetadata, default_manifest_path, format_json, - format_markdown, models, path_for_output, persist_result, public_model_ref, server, stats, + BENCH_SCHEMA_VERSION, BenchOutput, ModelInfo, OutputFormat, RunMetadata, default_manifest_path, + format_json, format_markdown, models, path_for_output, persist_result, public_model_ref, + server, stats, }; use serde::Serialize; @@ -232,6 +233,7 @@ async fn run(args: Args) -> Result<()> { }; let output = BenchOutput { + schema_version: BENCH_SCHEMA_VERSION, metadata, params, results, diff --git a/crates/higgs-bench/src/bin/bench_speculative.rs b/crates/higgs-bench/src/bin/bench_speculative.rs index 6e9ba477..9b6a0436 100644 --- a/crates/higgs-bench/src/bin/bench_speculative.rs +++ b/crates/higgs-bench/src/bin/bench_speculative.rs @@ -19,9 +19,9 @@ use std::time::{Duration, Instant}; use anyhow::{Context, Result}; use clap::Parser; use higgs_bench::{ - BenchOutput, ModelInfo, OutputFormat, RunMetadata, default_manifest_path, format_json, - format_markdown, models, path_for_output, persist_result, public_model_ref, results_dir, - server, speculative, stats, + BENCH_SCHEMA_VERSION, BenchOutput, ModelInfo, OutputFormat, RunMetadata, default_manifest_path, + format_json, format_markdown, models, path_for_output, persist_result, public_model_ref, + results_dir, server, speculative, stats, }; use serde::Serialize; @@ -69,8 +69,8 @@ struct Args { repeats: u32, /// Comma-separated trial modes: `baseline`, `mtp_default`, - /// `prompt_lookup`, `prompt_lookup_unchecked`, or numeric MTP draft - /// depths such as `1,2,3`. + /// `mtp_adaptive`, `mtp_hybrid`, `prompt_lookup`, + /// `prompt_lookup_unchecked`, or numeric MTP draft depths such as `1,2,3`. #[arg(long, default_value = "baseline,1,2,3")] trials: String, @@ -234,6 +234,7 @@ async fn run(args: Args) -> Result<()> { }; let output = BenchOutput { + schema_version: BENCH_SCHEMA_VERSION, metadata, params, results: Results { trials: summaries }, @@ -494,6 +495,8 @@ fn clear_speculative_env(cmd: &mut Command) { const SPECULATIVE_ENV_KEYS: &[&str] = &[ "HIGGS_MTP", "HIGGS_MTP_DRAFT_N_MAX", + "HIGGS_MTP_ADAPTIVE_DRAFT", + "HIGGS_MTP_PROMPT_LOOKUP", "HIGGS_PROMPT_LOOKUP", "HIGGS_PROMPT_LOOKUP_UNCHECKED", "HIGGS_MTP_PRIME_PREFILL", diff --git a/crates/higgs-bench/src/lib.rs b/crates/higgs-bench/src/lib.rs index 9d8605ce..f3470906 100644 --- a/crates/higgs-bench/src/lib.rs +++ b/crates/higgs-bench/src/lib.rs @@ -32,6 +32,9 @@ mod built_info { /// Bench-crate version (from `CARGO_PKG_VERSION` at compile time). pub const BENCH_VERSION: &str = env!("CARGO_PKG_VERSION"); +/// Version of the JSON/Markdown benchmark output envelope. +pub const BENCH_SCHEMA_VERSION: u32 = 1; + /// Returns the short git commit hash captured at compile time. #[must_use] pub fn git_commit_short() -> String { @@ -204,7 +207,9 @@ pub fn path_for_output(path: &Path) -> String { #[cfg(test)] mod tests { - use super::{hostname_for_output, public_model_ref, redact_arg_for_output}; + use super::{ + BENCH_SCHEMA_VERSION, hostname_for_output, public_model_ref, redact_arg_for_output, + }; #[test] fn benchmark_metadata_redacts_hostname_by_default() { @@ -240,6 +245,26 @@ mod tests { "local-qwen" ); } + + #[test] + fn bench_output_includes_schema_version() -> anyhow::Result<()> { + let output = super::BenchOutput { + schema_version: BENCH_SCHEMA_VERSION, + metadata: super::RunMetadata::capture("unit"), + params: serde_json::json!({}), + results: serde_json::json!({}), + }; + let json = super::format_json(&output)?; + let value: serde_json::Value = serde_json::from_str(&json)?; + + assert_eq!( + value + .get("schema_version") + .and_then(serde_json::Value::as_u64), + Some(1) + ); + Ok(()) + } } /// The model under test, captured into bench output for reproducibility. @@ -316,6 +341,8 @@ where P: Serialize, R: Serialize, { + /// JSON/Markdown envelope schema version. + pub schema_version: u32, /// Reproducibility metadata (host, model, git, argv, timing). pub metadata: RunMetadata, /// Bench-specific parameters (CLI flags, prompt, etc.). @@ -368,6 +395,7 @@ where meta.git_commit_short, if meta.git_dirty { " (dirty)" } else { "" } )?; + writeln!(s, "| schema_version | {} |", output.schema_version)?; writeln!(s, "| started_at | {} |", meta.started_at.to_rfc3339())?; writeln!(s, "| duration_ms | {} |", meta.duration_ms)?; if let Some(model) = &meta.model { diff --git a/crates/higgs-bench/src/speculative.rs b/crates/higgs-bench/src/speculative.rs index 1af712e1..630ef6b4 100644 --- a/crates/higgs-bench/src/speculative.rs +++ b/crates/higgs-bench/src/speculative.rs @@ -74,8 +74,10 @@ fn looks_like_huggingface_repo_id(model_path: &str) -> bool { /// Parses a comma-separated trial list. /// -/// Supported items are `baseline`, `mtp_default`, numeric MTP draft depths, -/// `prompt_lookup`, and `prompt_lookup_unchecked`. +/// Supported items are `baseline`, `mtp_default`, `mtp_adaptive`, +/// `mtp_hybrid` (prompt lookup plus adaptive MTP), numeric MTP draft depths, +/// `prompt_lookup`, and +/// `prompt_lookup_unchecked`. pub fn parse_trial_specs(input: &str) -> Result> { let mut trials = Vec::new(); for raw in input.split(',') { @@ -97,6 +99,18 @@ fn parse_trial_spec(trial: &str) -> Result { match trial { "baseline" => Ok(trial_spec("baseline_mtp_off", [("HIGGS_MTP", "0")])), "mtp_default" | "default" => Ok(trial_spec("mtp_default", [("HIGGS_MTP", "1")])), + "mtp_adaptive" | "adaptive" => Ok(trial_spec( + "mtp_adaptive", + [("HIGGS_MTP", "1"), ("HIGGS_MTP_ADAPTIVE_DRAFT", "1")], + )), + "mtp_hybrid" | "hybrid" => Ok(trial_spec( + "mtp_hybrid", + [ + ("HIGGS_MTP", "1"), + ("HIGGS_MTP_ADAPTIVE_DRAFT", "1"), + ("HIGGS_MTP_PROMPT_LOOKUP", "1"), + ], + )), "prompt_lookup" | "plookup" => Ok(trial_spec( "prompt_lookup", [("HIGGS_MTP", "0"), ("HIGGS_PROMPT_LOOKUP", "1")], @@ -176,14 +190,17 @@ mod tests { #[test] fn parse_trial_specs_sets_expected_env_overrides() -> anyhow::Result<()> { - let trials = - parse_trial_specs("baseline,mtp_default,2,prompt_lookup,prompt_lookup_unchecked")?; + let trials = parse_trial_specs( + "baseline,mtp_default,mtp_adaptive,mtp_hybrid,2,prompt_lookup,prompt_lookup_unchecked", + )?; let baseline = trial_at(&trials, 0)?; let mtp_default = trial_at(&trials, 1)?; - let mtp_draft_2 = trial_at(&trials, 2)?; - let prompt_lookup = trial_at(&trials, 3)?; - let prompt_lookup_unchecked = trial_at(&trials, 4)?; + let mtp_adaptive = trial_at(&trials, 2)?; + let mtp_hybrid = trial_at(&trials, 3)?; + let mtp_draft_2 = trial_at(&trials, 4)?; + let prompt_lookup = trial_at(&trials, 5)?; + let prompt_lookup_unchecked = trial_at(&trials, 6)?; assert_eq!(baseline.label, "baseline_mtp_off"); assert_eq!(baseline.env.get("HIGGS_MTP").map(String::as_str), Some("0")); @@ -194,6 +211,39 @@ mod tests { Some("1") ); + assert_eq!(mtp_adaptive.label, "mtp_adaptive"); + assert_eq!( + mtp_adaptive.env.get("HIGGS_MTP").map(String::as_str), + Some("1") + ); + assert_eq!( + mtp_adaptive + .env + .get("HIGGS_MTP_ADAPTIVE_DRAFT") + .map(String::as_str), + Some("1") + ); + + assert_eq!(mtp_hybrid.label, "mtp_hybrid"); + assert_eq!( + mtp_hybrid.env.get("HIGGS_MTP").map(String::as_str), + Some("1") + ); + assert_eq!( + mtp_hybrid + .env + .get("HIGGS_MTP_ADAPTIVE_DRAFT") + .map(String::as_str), + Some("1") + ); + assert_eq!( + mtp_hybrid + .env + .get("HIGGS_MTP_PROMPT_LOOKUP") + .map(String::as_str), + Some("1") + ); + assert_eq!(mtp_draft_2.label, "mtp_draft_2"); assert_eq!( mtp_draft_2.env.get("HIGGS_MTP").map(String::as_str), diff --git a/crates/higgs-engine/src/mlx_tuning.rs b/crates/higgs-engine/src/mlx_tuning.rs index 0717b120..953dcb39 100644 --- a/crates/higgs-engine/src/mlx_tuning.rs +++ b/crates/higgs-engine/src/mlx_tuning.rs @@ -414,7 +414,21 @@ fn heuristic_paged_kv_target_bytes( size_class: ModelSizeClass, is_moe: bool, ) -> usize { - let Some(max_recommended) = configured_max_working_set_bytes() else { + heuristic_paged_kv_target_bytes_with_max( + metadata, + size_class, + is_moe, + configured_max_working_set_bytes(), + ) +} + +fn heuristic_paged_kv_target_bytes_with_max( + metadata: &ModelMetadata, + size_class: ModelSizeClass, + is_moe: bool, + max_working_set_bytes: Option, +) -> usize { + let Some(max_recommended) = max_working_set_bytes else { return DEFAULT_PAGED_KV_TARGET_BYTES; }; @@ -426,7 +440,7 @@ fn heuristic_paged_kv_target_bytes( }); if available == 0 { - return DEFAULT_PAGED_KV_TARGET_BYTES; + return DEFAULT_PAGED_KV_TARGET_BYTES.min(max_recommended); } let divisor = if is_moe { @@ -440,7 +454,7 @@ fn heuristic_paged_kv_target_bytes( } }; - clamp_paged_kv_target_bytes(available / divisor) + clamp_paged_kv_target_bytes(available / divisor).min(max_recommended) } fn configured_max_working_set_bytes() -> Option { @@ -543,9 +557,9 @@ fn model_weight_bytes(model_dir: &Path) -> Option { mod tests { use super::{ MlxRuntimeTuning, ModelMetadata, ModelSizeClass, RequestedMlxProfile, ResolvedMlxProfile, - default_mtp_draft_n_max, model_weight_bytes, parse_enabled_flag, parse_mtp_draft_n_max, - parse_positive_chunked_prefill_value, resolve_effective_mlx_profile, - resolve_profile_from_metadata, resolve_runtime_tuning, + default_mtp_draft_n_max, heuristic_paged_kv_target_bytes_with_max, model_weight_bytes, + parse_enabled_flag, parse_mtp_draft_n_max, parse_positive_chunked_prefill_value, + resolve_effective_mlx_profile, resolve_profile_from_metadata, resolve_runtime_tuning, }; use std::fs; use tempfile::TempDir; @@ -727,6 +741,23 @@ mod tests { assert_eq!(parse_enabled_flag(None), None); } + #[test] + fn test_paged_kv_target_respects_configured_working_set_cap() { + let metadata = ModelMetadata { + weight_bytes: Some(134_217_728), + ..ModelMetadata::default() + }; + + let target = heuristic_paged_kv_target_bytes_with_max( + &metadata, + ModelSizeClass::Small, + false, + Some(134_217_728), + ); + + assert_eq!(target, 134_217_728); + } + #[test] fn test_from_model_dir_defaults_when_config_and_weights_missing() -> std::io::Result<()> { let temp = TempDir::new().map_err(std::io::Error::other)?; diff --git a/crates/higgs-engine/src/mtp.rs b/crates/higgs-engine/src/mtp.rs index ec8dc1e6..d29ef9c4 100644 --- a/crates/higgs-engine/src/mtp.rs +++ b/crates/higgs-engine/src/mtp.rs @@ -35,15 +35,21 @@ pub struct MtpStats { } impl MtpStats { - pub fn record_cycle(&mut self, drafted_count: usize, emitted_count: usize) { + pub fn record_cycle( + &mut self, + drafted_count: usize, + emitted_count: usize, + accepted_drafts_count: usize, + ) { let drafted = u32::try_from(drafted_count).unwrap_or(u32::MAX); let emitted = u32::try_from(emitted_count).unwrap_or(u32::MAX); + let accepted_drafts = u32::try_from(accepted_drafts_count) + .unwrap_or(u32::MAX) + .min(drafted); self.cycles = self.cycles.saturating_add(1); self.drafted = self.drafted.saturating_add(drafted); self.emitted = self.emitted.saturating_add(emitted); - self.accepted_drafts = self - .accepted_drafts - .saturating_add(emitted.saturating_sub(1).min(drafted)); + self.accepted_drafts = self.accepted_drafts.saturating_add(accepted_drafts); } pub const fn cycles(&self) -> u32 { @@ -72,6 +78,48 @@ impl MtpStats { } } +/// Small adaptive controller for choosing the next MTP draft depth. +#[derive(Debug, Clone)] +pub struct AdaptiveDraftDepth { + current: usize, + min: usize, + max: usize, +} + +impl AdaptiveDraftDepth { + #[must_use] + pub fn new(initial: usize, max_depth: usize) -> Self { + let capped_max = max_depth.max(1); + Self { + current: initial.clamp(1, capped_max), + min: 1, + max: capped_max, + } + } + + #[must_use] + pub const fn current(&self) -> usize { + self.current + } + + pub const fn observe(&mut self, accepted_drafts: usize, drafted: usize) { + if drafted == 0 { + self.current = self.min; + return; + } + + if accepted_drafts == drafted && self.current < self.max { + self.current += 1; + } else if accepted_drafts.saturating_mul(4) <= drafted && self.current > self.min { + self.current -= 1; + } else if accepted_drafts.saturating_mul(4) >= drafted.saturating_mul(3) + && self.current < self.max + { + self.current += 1; + } + } +} + /// Result of a single MTP speculative decode cycle. pub struct MtpCycleResult { /// Token IDs accepted this cycle (the confirmed token plus accepted drafts). @@ -116,6 +164,92 @@ pub struct PromptLookupCycleResult { pub accepted_drafts: usize, } +/// Run one prompt-lookup draft inside an MTP decode loop. +/// +/// This verifies copied prompt/history tokens with the backbone and mirrors the +/// accepted verifier span into the MTP cache, so the next cycle can continue +/// with either prompt lookup or the model's MTP head. +pub fn mtp_prompt_lookup_cycle( + model: &mut AnyModel, + cache: &mut AnyCache, + mtp_cache: &mut MtpCache, + previous_hidden: &Array, + history_before_confirmed: &[u32], + confirmed_token_id: u32, + config: PromptLookupConfig, +) -> Result, EngineError> { + let mut lookup_context = Vec::with_capacity(history_before_confirmed.len().saturating_add(1)); + lookup_context.extend_from_slice(history_before_confirmed); + lookup_context.push(confirmed_token_id); + let drafts = prompt_lookup_draft( + &lookup_context, + config.max_drafts, + config.max_ngram, + config.max_window, + ); + if drafts.is_empty() { + return Ok(None); + } + + let base_cache = cache.clone(); + let base_mtp_cache = mtp_cache.clone(); + let mut verify_tokens = Vec::with_capacity(drafts.len().saturating_add(1)); + verify_tokens.push(confirmed_token_id); + verify_tokens.extend(drafts.iter().copied()); + + let (verify_hidden, verifier_targets) = backbone_verify_batch(model, cache, &verify_tokens)?; + let verify_hidden_for_mtp = verify_hidden.clone(); + if verifier_targets.len() < verify_tokens.len() { + return Err(EngineError::Generation(format!( + "hybrid prompt-lookup verifier returned {} target ids for {} input tokens", + verifier_targets.len(), + verify_tokens.len() + ))); + } + + let accepted_drafts = accepted_draft_prefix_len(&drafts, &verifier_targets); + let tokens = emitted_tokens(confirmed_token_id, &drafts, accepted_drafts); + + let (accepted_hidden_rows, next_token_id) = if accepted_drafts == drafts.len() { + let next = *verifier_targets.get(accepted_drafts).ok_or_else(|| { + EngineError::Generation(format!( + "hybrid prompt-lookup verifier missing target at accepted index {accepted_drafts}" + )) + })?; + (verify_hidden, next) + } else { + *cache = base_cache; + let (replay_hidden, replay_targets) = backbone_verify_batch(model, cache, &tokens)?; + let next = *replay_targets.get(accepted_drafts).ok_or_else(|| { + EngineError::Generation(format!( + "hybrid prompt-lookup replay returned {} target ids for accepted index {}", + replay_targets.len(), + accepted_drafts + )) + })?; + (replay_hidden, next) + }; + + let h_last = hidden_row(&accepted_hidden_rows, accepted_drafts)?; + mirror_verified_mtp_cache( + model, + mtp_cache, + base_mtp_cache, + previous_hidden, + &verify_hidden_for_mtp, + &verify_tokens, + tokens.len(), + )?; + + Ok(Some(MtpCycleResult { + tokens, + hidden: h_last, + next_token_id, + drafted: drafts.len(), + accepted_drafts, + })) +} + fn greedy_token_id(logits: &Array) -> Result { let token_arr = argmax_axis!(&logits.index((.., -1, ..)), -1).map_err(EngineError::Mlx)?; eval([&token_arr]).map_err(EngineError::Mlx)?; @@ -580,11 +714,10 @@ pub fn mtp_cycle( )?; } - if accepted_drafts < drafts.len() && tokens.is_empty() { - return Err(EngineError::Generation( - "MTP accepted no committed tokens".to_owned(), - )); - } + debug_assert!( + !tokens.is_empty(), + "MTP must always emit the confirmed token" + ); Ok(MtpCycleResult { tokens, @@ -598,8 +731,8 @@ pub fn mtp_cycle( #[cfg(test)] mod tests { use super::{ - MtpStats, accepted_draft_prefix_len, draft_matches_target, emitted_tokens, - prompt_lookup_draft, + AdaptiveDraftDepth, MtpStats, accepted_draft_prefix_len, draft_matches_target, + emitted_tokens, prompt_lookup_draft, }; #[test] @@ -613,16 +746,30 @@ mod tests { } #[test] - fn mtp_stats_tracks_drafted_and_bonus_acceptance_rate() { + fn mtp_stats_tracks_explicit_accepted_draft_count() { let mut stats = MtpStats::default(); - stats.record_cycle(3, 4); - stats.record_cycle(2, 1); + stats.record_cycle(3, 2, 2); + stats.record_cycle(2, 1, 0); assert_eq!(stats.cycles(), 2); assert_eq!(stats.drafted(), 5); - assert_eq!(stats.emitted(), 5); - assert_eq!(stats.accepted_drafts(), 3); - assert!((stats.acceptance_rate_percent() - 60.0).abs() < f64::EPSILON); + assert_eq!(stats.emitted(), 3); + assert_eq!(stats.accepted_drafts(), 2); + assert!((stats.acceptance_rate_percent() - 40.0).abs() < f64::EPSILON); + } + + #[test] + fn adaptive_draft_depth_grows_on_full_acceptance_and_backs_off_on_rejection() { + let mut depth = AdaptiveDraftDepth::new(2, 4); + + depth.observe(2, 2); + assert_eq!(depth.current(), 3); + + depth.observe(3, 3); + assert_eq!(depth.current(), 4); + + depth.observe(0, 4); + assert_eq!(depth.current(), 3); } #[test] diff --git a/crates/higgs-engine/src/simple.rs b/crates/higgs-engine/src/simple.rs index c4bf0ed2..9b37e73d 100644 --- a/crates/higgs-engine/src/simple.rs +++ b/crates/higgs-engine/src/simple.rs @@ -69,6 +69,18 @@ fn unchecked_prompt_lookup_enabled() -> bool { .unwrap_or(false) } +fn mtp_adaptive_draft_enabled() -> bool { + parse_enabled_flag(std::env::var("HIGGS_MTP_ADAPTIVE_DRAFT").ok().as_deref()).unwrap_or(false) +} + +fn mtp_prompt_lookup_enabled() -> bool { + parse_enabled_flag(std::env::var("HIGGS_MTP_PROMPT_LOOKUP").ok().as_deref()).unwrap_or(false) +} + +fn adaptive_draft_depth_for_cap(configured_max: usize) -> crate::mtp::AdaptiveDraftDepth { + crate::mtp::AdaptiveDraftDepth::new(configured_max, configured_max) +} + fn mtp_prefill_priming_enabled() -> bool { parse_enabled_flag(std::env::var("HIGGS_MTP_PRIME_PREFILL").ok().as_deref()).unwrap_or(true) } @@ -1450,7 +1462,7 @@ impl SimpleEngine { } else { crate::mtp::prompt_lookup_cycle(model, cache, tokens, confirmed_token_id, config)? }; - stats.record_cycle(result.drafted, result.tokens.len()); + stats.record_cycle(result.drafted, result.tokens.len(), result.accepted_drafts); for &tok in &result.tokens { if let Some(close_id) = think_close_token { @@ -1613,6 +1625,10 @@ impl SimpleEngine { let mut current_hidden = h; let mut confirmed_token_id: u32 = next_arr.item(); let mut mtp_stats = crate::mtp::MtpStats::default(); + let mut adaptive_depth = mtp_adaptive_draft_enabled() + .then(|| adaptive_draft_depth_for_cap(self.tuning.mtp_draft_n_max())); + let hybrid_prompt_lookup = mtp_prompt_lookup_enabled(); + let hybrid_prompt_lookup_config = prompt_lookup_config(); let t_start = std::time::Instant::now(); // Thinking budget: force after N tokens if model hasn't closed it. @@ -1628,16 +1644,60 @@ impl SimpleEngine { think_close_token.is_some_and(|close_id| first_token_id == close_id); loop { - let result = crate::mtp::mtp_cycle( - model, - cache, - &mut mtp_cache, - ¤t_hidden, - confirmed_token_id, - self.tuning.mtp_draft_n_max(), - )?; + let cycle_completion_len = Self::completion_len(tokens)?; + let remaining = usize::try_from(max_tokens.saturating_sub(cycle_completion_len)) + .map_err(|_| EngineError::Generation("max_tokens overflow".to_owned()))?; + let draft_depth = adaptive_depth + .as_ref() + .map_or_else( + || self.tuning.mtp_draft_n_max(), + crate::mtp::AdaptiveDraftDepth::current, + ) + .min(remaining.saturating_sub(1).max(1)); + let prompt_config = crate::mtp::PromptLookupConfig { + max_drafts: hybrid_prompt_lookup_config + .max_drafts + .min(remaining.saturating_sub(1)), + ..hybrid_prompt_lookup_config + }; + let result = if hybrid_prompt_lookup && prompt_config.max_drafts > 0 { + crate::mtp::mtp_prompt_lookup_cycle( + model, + cache, + &mut mtp_cache, + ¤t_hidden, + tokens, + confirmed_token_id, + prompt_config, + )? + .map_or_else( + || { + crate::mtp::mtp_cycle( + model, + cache, + &mut mtp_cache, + ¤t_hidden, + confirmed_token_id, + draft_depth, + ) + }, + Ok, + )? + } else { + crate::mtp::mtp_cycle( + model, + cache, + &mut mtp_cache, + ¤t_hidden, + confirmed_token_id, + draft_depth, + )? + }; - mtp_stats.record_cycle(result.drafted, result.tokens.len()); + mtp_stats.record_cycle(result.drafted, result.tokens.len(), result.accepted_drafts); + if let Some(depth) = &mut adaptive_depth { + depth.observe(result.accepted_drafts, result.drafted); + } for &tok in &result.tokens { // Thinking budget enforcement @@ -1804,6 +1864,10 @@ impl SimpleEngine { let mut current_hidden = h; let mut confirmed_token_id: u32 = next_arr.item(); let mut mtp_stats = crate::mtp::MtpStats::default(); + let mut adaptive_depth = mtp_adaptive_draft_enabled() + .then(|| adaptive_draft_depth_for_cap(self.tuning.mtp_draft_n_max())); + let hybrid_prompt_lookup = mtp_prompt_lookup_enabled(); + let hybrid_prompt_lookup_config = prompt_lookup_config(); let t_start = std::time::Instant::now(); const THINKING_BUDGET: u32 = 256; @@ -1817,16 +1881,60 @@ impl SimpleEngine { think_close_token.is_some_and(|close_id| first_token_id == close_id); loop { - let result = crate::mtp::mtp_cycle( - model, - cache, - &mut mtp_cache, - ¤t_hidden, - confirmed_token_id, - self.tuning.mtp_draft_n_max(), - )?; + let cycle_completion_len = Self::completion_len(tokens)?; + let remaining = usize::try_from(max_tokens.saturating_sub(cycle_completion_len)) + .map_err(|_| EngineError::Generation("max_tokens overflow".to_owned()))?; + let draft_depth = adaptive_depth + .as_ref() + .map_or_else( + || self.tuning.mtp_draft_n_max(), + crate::mtp::AdaptiveDraftDepth::current, + ) + .min(remaining.saturating_sub(1).max(1)); + let prompt_config = crate::mtp::PromptLookupConfig { + max_drafts: hybrid_prompt_lookup_config + .max_drafts + .min(remaining.saturating_sub(1)), + ..hybrid_prompt_lookup_config + }; + let result = if hybrid_prompt_lookup && prompt_config.max_drafts > 0 { + crate::mtp::mtp_prompt_lookup_cycle( + model, + cache, + &mut mtp_cache, + ¤t_hidden, + tokens, + confirmed_token_id, + prompt_config, + )? + .map_or_else( + || { + crate::mtp::mtp_cycle( + model, + cache, + &mut mtp_cache, + ¤t_hidden, + confirmed_token_id, + draft_depth, + ) + }, + Ok, + )? + } else { + crate::mtp::mtp_cycle( + model, + cache, + &mut mtp_cache, + ¤t_hidden, + confirmed_token_id, + draft_depth, + )? + }; - mtp_stats.record_cycle(result.drafted, result.tokens.len()); + mtp_stats.record_cycle(result.drafted, result.tokens.len(), result.accepted_drafts); + if let Some(depth) = &mut adaptive_depth { + depth.observe(result.accepted_drafts, result.drafted); + } for &tok in &result.tokens { // Thinking budget enforcement @@ -2479,7 +2587,8 @@ fn detect_thinking_support(model_dir: &Path) -> bool { #[allow(clippy::panic, clippy::unwrap_used)] mod tests { use super::{ - check_stop_sequences, derive_model_name, estimate_paged_kv_blocks, parse_enabled_flag, + adaptive_draft_depth_for_cap, check_stop_sequences, derive_model_name, + estimate_paged_kv_blocks, parse_enabled_flag, }; use std::path::Path; @@ -2516,6 +2625,15 @@ mod tests { assert_eq!(name, "my-model"); } + #[test] + fn adaptive_draft_depth_respects_configured_cap() { + let mut depth = adaptive_draft_depth_for_cap(1); + + depth.observe(1, 1); + + assert_eq!(depth.current(), 1); + } + /// Create a temp dir, write config.json with the given content, and return /// the result of `extract_eos_tokens`. fn eos_from_config(json: &str) -> Vec { diff --git a/docs/benchmarking.md b/docs/benchmarking.md index b7c09d96..76549f6a 100644 --- a/docs/benchmarking.md +++ b/docs/benchmarking.md @@ -66,21 +66,32 @@ experiments. `HIGGS_MTP_MIRROR_VERIFY=1` enables full verifier-window MTP cache mirroring; on the Qwen3.6 27B MTP 8-bit benchmark below it was slightly slower than the default accepted-prefix replay path, so it remains opt-in. +`HIGGS_MTP_ADAPTIVE_DRAFT=1` enables a lightweight controller that grows the +draft window after high verifier acceptance and backs off after rejections. +`HIGGS_MTP_PROMPT_LOOKUP=1` enables a hybrid path that verifies repeated +prompt/history spans inside the MTP loop and mirrors accepted verifier rows into +the MTP cache, so it can fall back to normal MTP-head cycles on models without +useful repeated spans. + Measured on M4 Max 128GB, `temperature=0`, 96 completion tokens, prompt: `Write a concise technical explanation of speculative decoding for local LLM inference...` | Runtime | Mode | Request tok/s | Decode-only tok/s | Speedup vs runtime baseline | | --- | ---: | ---: | ---: | ---: | -| Higgs | baseline MTP off | 14.32 | n/a | 1.00x | -| Higgs | MTP draft depth 2 | 22.89 | 28.0 | 1.60x request / 1.96x vs request baseline | -| llama.cpp `b1-d374e71` | baseline | n/a | 15.9 | 1.00x | -| llama.cpp `b1-d374e71` | MTP draft depth 1 | n/a | 25.0 | 1.57x | -| llama.cpp `b1-d374e71` | MTP draft depth 2 | n/a | 24.3 | 1.53x | - -The Higgs request-level number includes HTTP and prompt processing; llama.cpp's -CLI line reports generation only. The closest decode-only comparison from this -run is Higgs MTP depth 2 at `28.0 tok/s` versus llama.cpp's best measured MTP -setting at `25.0 tok/s`. +| Higgs | baseline MTP off | 14.14 | n/a | 1.00x | +| Higgs | MTP default | 22.75 | n/a | 1.61x request | +| Higgs | MTP adaptive | 22.11 | n/a | 1.56x request | +| Higgs | MTP hybrid prompt lookup + adaptive | 18.45 | n/a | 1.30x request | +| Higgs | MTP draft depth 2 | 22.79 | n/a | 1.61x request | +| llama.cpp `b9410-031ddb2e0` | baseline | 14.62 | 15.64 | 1.00x | +| llama.cpp `b9410-031ddb2e0` | `draft-mtp`, depth 2 | 21.63 | 24.16 | 1.48x request / 1.55x decode | + +The Higgs numbers are from `bench_speculative`, which starts a fresh server per +mode and reports end-to-end request tok/s. The llama.cpp rows use the +OpenAI-compatible server with Qwen thinking disabled via +`chat_template_kwargs.enable_thinking=false`, prompt cache disabled, and +`draft-mtp` depth 2. On this run, Higgs MTP draft depth 2 was `1.05x` faster +than llama.cpp `draft-mtp` depth 2 at the request level. ## Iterations @@ -201,12 +212,13 @@ not time-to-visible-answer for thinking-mode models. ### `bench_speculative` Starts a fresh Higgs server per speculative mode and compares baseline greedy -decode with MTP draft depths and architecture-neutral prompt lookup. +decode with MTP draft depths, adaptive MTP, hybrid prompt-lookup+MTP, and +architecture-neutral prompt lookup. ```bash cargo run --release -p higgs-bench --bin bench_speculative -- \ --model qwen3.6-27B-mtp-8bit \ - --trials baseline,mtp_default,1,2,3,prompt_lookup,prompt_lookup_unchecked \ + --trials baseline,mtp_default,mtp_adaptive,mtp_hybrid,1,2,3,prompt_lookup,prompt_lookup_unchecked \ --max-tokens 96 --repeats 3 --format markdown ``` @@ -215,6 +227,10 @@ tokens/sec when a `baseline` trial appears before the speculative trial. Server logs are captured under `target/bench-results/bench_speculative/logs/`, and the JSON result stores only filtered speculative telemetry lines. +The `mtp_hybrid` trial intentionally combines MTP heads, adaptive draft depth, +and MTP-local prompt lookup. Use `prompt_lookup` for architecture-neutral prompt +lookup without MTP heads. + ### `bench_summarize` Walks `target/bench-results/`, picks the latest result per @@ -233,7 +249,7 @@ cargo run --release -p higgs-bench --bin bench_summarize 3. Look up the model with `higgs_bench::models::find_by_key(...)` and set `metadata.model`. 4. Define `Params` and `Results` structs (must implement `Serialize`). -5. Build `BenchOutput { metadata, params, results }` and call +5. Build `BenchOutput { schema_version: higgs_bench::BENCH_SCHEMA_VERSION, metadata, params, results }` and call `higgs_bench::persist_result(&output)` plus `higgs_bench::format_json` / `format_markdown` based on `--format`. diff --git a/docs/configuration.md b/docs/configuration.md index 2cab445f..bf50abbf 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -38,7 +38,9 @@ This document collects the full CLI, environment, and config-file reference for - `HIGGS_CHUNKED_PREFILL_THRESHOLD` enables chunked prefill above a token threshold. - `HIGGS_CHUNKED_PREFILL_CHUNK_SIZE` controls chunk size during chunked prefill. - `HIGGS_MTP=0|1` overrides the tuning profile's speculative decode choice when conditions allow. -- `HIGGS_MTP_DRAFT_N_MAX` controls the maximum MTP draft tokens per speculative cycle. The default is `3` when MTP is enabled, clamped to `1..=8`. +- `HIGGS_MTP_DRAFT_N_MAX` controls the maximum MTP draft tokens per speculative cycle. The default is `2` for huge checkpoints and `1` otherwise, clamped to `1..=8`. +- `HIGGS_MTP_ADAPTIVE_DRAFT=1` lets the decode loop increase or decrease the MTP draft window based on recent verifier acceptance. +- `HIGGS_MTP_PROMPT_LOOKUP=1` enables a hybrid MTP loop that tries verified prompt-lookup drafts when the prompt/history has a repeated suffix, then keeps the MTP cache synchronized for later MTP-head cycles. - `HIGGS_CLEAR_CACHE_AFTER_PREFILL` overrides the selected MLX profile behavior for cache clearing. - `HIGGS_TURBOQUANT_MIN_TOKENS` overrides the TurboQuant activation threshold. The default is `2048`. - `HIGGS_EXPERIMENTAL_PAGED_KV=1` enables the experimental paged-KV path. From 1d15599ea86fccb022d786a56b2904131078b150 Mon Sep 17 00:00:00 2001 From: Jonathan Reyes Date: Sun, 31 May 2026 12:44:29 -0600 Subject: [PATCH 25/29] feat(mtp): release speculative decoding optimizations --- crates/higgs-engine/src/lib.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/crates/higgs-engine/src/lib.rs b/crates/higgs-engine/src/lib.rs index 705a82b4..45622ec6 100644 --- a/crates/higgs-engine/src/lib.rs +++ b/crates/higgs-engine/src/lib.rs @@ -1,3 +1,5 @@ +//! Engine crate for Higgs local inference, including MLX-backed MTP decoding. + pub mod batch_engine; pub mod cache; pub mod chat_template; From e398b31f9dda4b582e25ea785e7c95fc9e698abb Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Sun, 31 May 2026 18:49:33 +0000 Subject: [PATCH 26/29] chore: release main --- .release-please-manifest.json | 6 +++--- Cargo.lock | 6 +++--- crates/higgs-engine/CHANGELOG.md | 8 ++++++++ crates/higgs-engine/Cargo.toml | 2 +- crates/higgs-models/CHANGELOG.md | 7 +++++++ crates/higgs-models/Cargo.toml | 2 +- crates/higgs/CHANGELOG.md | 7 +++++++ crates/higgs/Cargo.toml | 2 +- 8 files changed, 31 insertions(+), 9 deletions(-) diff --git a/.release-please-manifest.json b/.release-please-manifest.json index fbcbdcaa..4d1cb705 100644 --- a/.release-please-manifest.json +++ b/.release-please-manifest.json @@ -1,5 +1,5 @@ { - "crates/higgs-models": "1.2.0", - "crates/higgs-engine": "1.2.0", - "crates/higgs": "1.2.0" + "crates/higgs-models": "1.3.0", + "crates/higgs-engine": "1.3.0", + "crates/higgs": "1.3.0" } diff --git a/Cargo.lock b/Cargo.lock index a454d989..44e043da 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1563,7 +1563,7 @@ dependencies = [ [[package]] name = "higgs" -version = "1.2.0" +version = "1.3.0" dependencies = [ "async-stream", "axum", @@ -1623,7 +1623,7 @@ dependencies = [ [[package]] name = "higgs-engine" -version = "1.2.0" +version = "1.3.0" dependencies = [ "half", "higgs-models", @@ -1643,7 +1643,7 @@ dependencies = [ [[package]] name = "higgs-models" -version = "1.2.0" +version = "1.3.0" dependencies = [ "criterion", "half", diff --git a/crates/higgs-engine/CHANGELOG.md b/crates/higgs-engine/CHANGELOG.md index b98e8404..e241c67e 100644 --- a/crates/higgs-engine/CHANGELOG.md +++ b/crates/higgs-engine/CHANGELOG.md @@ -1,5 +1,13 @@ # Changelog +## [1.3.0](https://github.com/panbanda/higgs/compare/higgs-engine-v1.2.0...higgs-engine-v1.3.0) (2026-05-31) + + +### Features + +* **mtp:** release speculative decoding optimizations ([0e5e458](https://github.com/panbanda/higgs/commit/0e5e458ce2b59ab5f1a6aaaa1fe91af6f298fe39)) +* **mtp:** release speculative decoding optimizations ([1d15599](https://github.com/panbanda/higgs/commit/1d15599ea86fccb022d786a56b2904131078b150)) + ## [1.2.0](https://github.com/panbanda/higgs/compare/higgs-engine-v1.1.1...higgs-engine-v1.2.0) (2026-05-06) diff --git a/crates/higgs-engine/Cargo.toml b/crates/higgs-engine/Cargo.toml index 9bc38182..9d241ac3 100644 --- a/crates/higgs-engine/Cargo.toml +++ b/crates/higgs-engine/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "higgs-engine" -version = "1.2.0" +version = "1.3.0" description = "Inference engine (tokenization, generation loop, prompt caching) for higgs" edition.workspace = true license.workspace = true diff --git a/crates/higgs-models/CHANGELOG.md b/crates/higgs-models/CHANGELOG.md index 6375c01a..8b596165 100644 --- a/crates/higgs-models/CHANGELOG.md +++ b/crates/higgs-models/CHANGELOG.md @@ -1,5 +1,12 @@ # Changelog +## [1.3.0](https://github.com/panbanda/higgs/compare/higgs-models-v1.2.0...higgs-models-v1.3.0) (2026-05-31) + + +### Miscellaneous Chores + +* **higgs-models:** Synchronize workspace versions + ## [1.2.0](https://github.com/panbanda/higgs/compare/higgs-models-v1.1.1...higgs-models-v1.2.0) (2026-05-06) diff --git a/crates/higgs-models/Cargo.toml b/crates/higgs-models/Cargo.toml index 843ea42f..796ce8b2 100644 --- a/crates/higgs-models/Cargo.toml +++ b/crates/higgs-models/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "higgs-models" -version = "1.2.0" +version = "1.3.0" description = "Model architectures and weight loading for higgs (LLaMA, Mistral, Qwen2/3)" edition.workspace = true license.workspace = true diff --git a/crates/higgs/CHANGELOG.md b/crates/higgs/CHANGELOG.md index 4a73a74d..b24020fb 100644 --- a/crates/higgs/CHANGELOG.md +++ b/crates/higgs/CHANGELOG.md @@ -1,5 +1,12 @@ # Changelog +## [1.3.0](https://github.com/panbanda/higgs/compare/higgs-v1.2.0...higgs-v1.3.0) (2026-05-31) + + +### Miscellaneous Chores + +* **higgs:** Synchronize workspace versions + ## [1.2.0](https://github.com/panbanda/higgs/compare/higgs-v1.1.1...higgs-v1.2.0) (2026-05-06) diff --git a/crates/higgs/Cargo.toml b/crates/higgs/Cargo.toml index 2f8c8ea3..2763df66 100644 --- a/crates/higgs/Cargo.toml +++ b/crates/higgs/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "higgs" -version = "1.2.0" +version = "1.3.0" description = "OpenAI and Anthropic-compatible inference server for Apple Silicon, built on mlx-rs" edition.workspace = true license.workspace = true From f02cc4dbd709fe952f482fcb22b00ecf435dc385 Mon Sep 17 00:00:00 2001 From: Jonathan Reyes Date: Sun, 31 May 2026 12:51:22 -0600 Subject: [PATCH 27/29] chore: trigger release pr checks From 386347136f73325d3a3ea7a6c5c00ec467ebb5c2 Mon Sep 17 00:00:00 2001 From: Peppi Littera Date: Tue, 2 Jun 2026 17:09:15 +0200 Subject: [PATCH 28/29] feat(bonsai): run Bonsai-Q1 bits=1 on vanilla MLX via JIT Metal kernels MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Stock oxideai/mlx-rs ships no bits=1 affine kernel (MLX gates affine quant to bits>=2), so ops::quantized_matmul / ops::dequantize with bits=1 failed at runtime with 'Unable to load kernel affine_dequantize_*_b_1'. Rather than fork mlx-rs (which forces a full from-source mlx-c rebuild), add the kernels from this crate using mlx-c's runtime JIT facility (mlx_fast_metal_kernel), already compiled into mlx-sys and reached over the -sys FFI. We stay on the official oxideai pin with no extra native recompile; kernels JIT-compile at first use. New crates/higgs-models/src/metal_kernel.rs: - bonsai_q1_qmv: fused 1-bit quantized matvec (decode hot path) — one simdgroup per output row over the packed weights, simd_sum reduce. - bonsai_q1_dequant: packed -> dense f16 (embedding gather + prefill matmul). FFI plumbing mirrors the proven qgemv_4bit pattern in qwen3_next; the unpack/ affine math mirrors PackedQ1Linear::dequant_row_to_fp32. Wired into BonsaiQ1GpuLinear::forward (decode uses qmv, prefill dequant+matmul) and BonsaiQ1Gpu::embed_rows. model_loader comment corrected and the obsolete bits=1 guard test repurposed to assert routing into the packed engine. Verified on vanilla MLX (no fork, no native rebuild): - Oracle unit tests vs CPU reference pass; clippy (nursery) + fmt clean. - Bonsai-1.7B-mlx-1bit: ~113 tok/s decode (8.82 ms/step), coherent output. - Bonsai-8B-mlx-1bit: ~28 tok/s decode (35.5 ms/step), coherent output; exercises K-chunking + nsg=16 (inter=12288) and the untied lm_head. Co-Authored-By: Claude Opus 4.8 (1M context) --- crates/higgs-engine/src/model_loader.rs | 32 +- crates/higgs-models/src/bonsai_q1.rs | 173 ++++++-- crates/higgs-models/src/lib.rs | 2 + crates/higgs-models/src/metal_kernel.rs | 498 ++++++++++++++++++++++++ 4 files changed, 672 insertions(+), 33 deletions(-) create mode 100644 crates/higgs-models/src/metal_kernel.rs diff --git a/crates/higgs-engine/src/model_loader.rs b/crates/higgs-engine/src/model_loader.rs index bd2f91d1..ef3c9882 100644 --- a/crates/higgs-engine/src/model_loader.rs +++ b/crates/higgs-engine/src/model_loader.rs @@ -39,15 +39,14 @@ pub fn load_model>(model_dir: P) -> Result match config.model_type.as_str() { "qwen2" | "qwen3" | "llama" | "mistral" => { // Packed 1.25-bpw Bonsai-Q1 checkpoints declare model_type="qwen3" - // but the weights are quantized to bits=1. Keep detection ahead of - // the fp16/Q4 transformer loader so users get an explicit error - // while the workspace remains on upstream oxideai/mlx-rs. + // but the weights are quantized to bits=1. Route them to the + // dedicated packed engine, whose bits=1 matvec/dequant run through + // runtime JIT Metal kernels (higgs-models::metal_kernel) — so it + // runs on stock oxideai/mlx-rs with no forked bits=1 MLX kernel. if is_bonsai_q1(&config.model_dir)? { - return Err(EngineError::Model(ModelError::UnsupportedModel( - "Bonsai-Q1 requires MLX bits=1 affine quantization support; \ - the workspace stays on upstream oxideai/mlx-rs until that support lands" - .to_owned(), - ))); + let gpu = higgs_models::bonsai_q1::load_bonsai_q1(&config.model_dir) + .map_err(EngineError::Model)?; + return Ok(AnyModel::BonsaiQ1(gpu)); } let model = transformer::load_model(&config.model_dir).map_err(EngineError::Model)?; Ok(AnyModel::Transformer(model)) @@ -315,7 +314,12 @@ mod tests { } #[test] - fn load_model_rejects_bonsai_q1_without_runtime_support() { + fn load_model_routes_bonsai_q1_to_packed_engine() { + // A bits=1 / group=128 qwen3 config now routes to the packed Bonsai-Q1 + // engine (its bits=1 kernels live in higgs-models::metal_kernel) instead + // of being rejected up front. With no weights in the dir the load still + // fails inside the engine — but it must no longer be gated out, and the + // old "requires MLX bits=1" guard error must be gone. let (dir, _result) = config_from_raw( r#"{ "model_type": "qwen3", @@ -323,8 +327,14 @@ mod tests { }"#, ); match load_model(dir.path()) { - Err(err) => assert!(err.to_string().contains("Bonsai-Q1 requires MLX bits=1")), - Ok(_) => panic!("Expected unsupported Bonsai-Q1 runtime error"), + Ok(_) => panic!("expected load failure: config-only dir has no weights"), + Err(EngineError::Model(ModelError::UnsupportedModel(_))) => { + panic!("Bonsai-Q1 must route to the packed engine, not be rejected as unsupported") + } + Err(err) => assert!( + !err.to_string().contains("requires MLX bits=1"), + "stale bits=1 guard error should be gone, got: {err}" + ), } } diff --git a/crates/higgs-models/src/bonsai_q1.rs b/crates/higgs-models/src/bonsai_q1.rs index 93fe10e2..131e263d 100644 --- a/crates/higgs-models/src/bonsai_q1.rs +++ b/crates/higgs-models/src/bonsai_q1.rs @@ -3,14 +3,15 @@ //! Unlike `DiffusionEngine::load_q1` which dequantizes to fp32 at load (32 GB //! residency on 8B), this engine holds MLX's `Q1_0_g128` affine encoding //! verbatim: `w[row, col] = scales[row, col/128] * bit(col) + biases[row, -//! col/128]`. Dequant happens inline inside the MLX quantized matmul kernel -//! once upstream MLX provides bits=1 affine support. +//! col/128]`. Because stock `oxideai/mlx-rs` ships no bits=1 affine kernel, the +//! matvec/dequant run through runtime JIT Metal kernels in +//! [`crate::metal_kernel`] (decode uses a fused matvec over the packed weights; +//! prefill/embedding dequantize to dense f16). //! //! Residency: ~1.25 GB for Bonsai-8B-mlx-1bit, ~260 MB for Bonsai-1.7B-mlx-1bit. //! -//! Scope: Rust-side loader and engine implementation. Runtime routing is held -//! back in `higgs-engine` until the upstream MLX dependency supports bits=1 -//! affine quantization. +//! Scope: Rust-side loader and engine implementation. Routing is enabled in +//! `higgs-engine::model_loader` since the kernels run on the stock MLX pin. #![allow( clippy::too_many_arguments, @@ -65,7 +66,6 @@ pub fn load_bonsai_q1>(model_dir: P) -> Result 1) dequantizes to + /// dense f16 and uses a regular matmul (the transient dense weight amortizes + /// over the M rows). Both run on stock `oxideai/mlx-rs` via + /// [`crate::metal_kernel`] — no `bits=1` MLX kernel required. pub fn forward(&self, x: &Array) -> Result { - ops::quantized_matmul( - x, - &self.w, - &self.scales, - &self.biases, - true, - GROUP_SIZE_I32, - BITS, - ) + let total: i32 = x.shape().iter().product(); + let m = if self.in_features > 0 { + total / self.in_features + } else { + 0 + }; + if m == 1 { + crate::metal_kernel::bonsai_q1_qmv( + x, + &self.w, + &self.scales, + &self.biases, + GROUP_SIZE_I32, + ) + } else { + let wd = crate::metal_kernel::bonsai_q1_dequant( + &self.w, + &self.scales, + &self.biases, + GROUP_SIZE_I32, + )?; + ops::matmul(x, &wd.transpose_axes(&[1, 0])?) + } } } @@ -521,15 +543,15 @@ impl BonsaiQ1Gpu { /// Gather embedding rows for a token-ID tensor. /// - /// Uses MLX dequantize after gathering the selected packed rows. This path - /// requires bits=1 affine support in the active MLX runtime. + /// Gathers the selected packed rows, then dequantizes them to dense f16 via + /// the [`crate::metal_kernel`] bits=1 kernel (runs on stock `oxideai/mlx-rs`). fn embed_rows(&self, ids: &Array) -> Result { let shape = ids.shape().to_vec(); let flat = ids.flatten(None, None)?; let w = self.embed.w.take_axis(&flat, 0)?; let s = self.embed.scales.take_axis(&flat, 0)?; let b = self.embed.biases.take_axis(&flat, 0)?; - let out = ops::dequantize(&w, &s, &b, GROUP_SIZE_I32, BITS)?; + let out = crate::metal_kernel::bonsai_q1_dequant(&w, &s, &b, GROUP_SIZE_I32)?; let mut ret_shape: Vec = shape; ret_shape.push(-1); out.reshape(&ret_shape) @@ -1202,3 +1224,110 @@ fn bytes_to_f16_vec(b: &[u8]) -> Vec { // --------------------------------------------------------------------------- // Tests // --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + use crate::metal_kernel::{bonsai_q1_dequant, bonsai_q1_qmv}; + + /// Deterministic PRNG (SplitMix-ish LCG). Tests prove the kernels match the + /// CPU reference over pseudo-random data, not against hand-picked constants. + fn lcg(state: &mut u64) -> u32 { + *state = state + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + (*state >> 32) as u32 + } + + fn make_packed(out_features: usize, in_features: usize, seed: u64) -> PackedQ1Linear { + let packed_cols = in_features / 32; + let n_groups = in_features / GROUP_SIZE; + let mut st = seed; + let w_packed: Vec = (0..out_features * packed_cols) + .map(|_| lcg(&mut st)) + .collect(); + // Per-(row,group) scales/biases so a wrong group index produces a clear + // mismatch. Magnitudes are small and signed, like real affine params. + let scales: Vec = (0..out_features * n_groups) + .map(|i| f16::from_f32(0.05 + 0.013 * ((i % 7) as f32))) + .collect(); + let biases: Vec = (0..out_features * n_groups) + .map(|i| f16::from_f32(-0.03 + 0.011 * ((i % 5) as f32))) + .collect(); + PackedQ1Linear { + w_packed, + scales, + biases, + out_features, + in_features, + } + } + + /// CPU reference: full dense dequant via the documented per-row path. + fn dense_reference(p: &PackedQ1Linear) -> Vec { + let mut wd = vec![0.0f32; p.out_features * p.in_features]; + for r in 0..p.out_features { + let (lo, hi) = (r * p.in_features, (r + 1) * p.in_features); + p.dequant_row_to_fp32(r, &mut wd[lo..hi]); + } + wd + } + + #[test] + fn dequant_kernel_matches_cpu_reference() { + let (out_f, in_f) = (96usize, 256usize); // 256 cols => 2 groups of 128 + let p = make_packed(out_f, in_f, 0xDEAD_BEEF); + let gpu = BonsaiQ1GpuLinear::from_packed(&p).unwrap(); + + let wd = bonsai_q1_dequant(&gpu.w, &gpu.scales, &gpu.biases, GROUP_SIZE_I32).unwrap(); + wd.eval().unwrap(); + let got = wd.as_slice::(); + let want = dense_reference(&p); + + assert_eq!(got.len(), want.len()); + for (i, (g, w)) in got.iter().zip(want.iter()).enumerate() { + let gv = g.to_f32(); + assert!( + (gv - w).abs() <= 2e-3, + "dequant mismatch at {i}: got {gv} want {w}" + ); + } + } + + #[test] + fn qmv_kernel_matches_cpu_reference() { + let (out_f, in_f) = (96usize, 256usize); + let p = make_packed(out_f, in_f, 0x1234_5678); + let gpu = BonsaiQ1GpuLinear::from_packed(&p).unwrap(); + + // x in [-1, 1], deterministic; the kernel reads it as f16, so the + // reference uses the f16-rounded values for an apples-to-apples compare. + let mut st = 0xABCD_EF01_u64; + let x_f32: Vec = (0..in_f) + .map(|_| (lcg(&mut st) as f32 / u32::MAX as f32).mul_add(2.0, -1.0)) + .collect(); + let x = Array::from_slice(&x_f32, &[1, in_f as i32]) + .as_dtype(Dtype::Float16) + .unwrap(); + let x_ref: Vec = x_f32.iter().map(|&v| f16::from_f32(v).to_f32()).collect(); + + let y = bonsai_q1_qmv(&x, &gpu.w, &gpu.scales, &gpu.biases, GROUP_SIZE_I32).unwrap(); + y.eval().unwrap(); + let got = y.as_slice::(); + assert_eq!(got.len(), out_f); + + let wd = dense_reference(&p); + for r in 0..out_f { + let mut acc = 0.0f32; + for c in 0..in_f { + acc += x_ref[c] * wd[r * in_f + c]; + } + let gv = got[r].to_f32(); + let tol = 1e-2 * acc.abs().max(1.0); + assert!( + (gv - acc).abs() <= tol, + "qmv mismatch at row {r}: got {gv} want {acc}" + ); + } + } +} diff --git a/crates/higgs-models/src/lib.rs b/crates/higgs-models/src/lib.rs index 8ae2716a..ed81b5c1 100644 --- a/crates/higgs-models/src/lib.rs +++ b/crates/higgs-models/src/lib.rs @@ -4,6 +4,8 @@ pub mod deepseek_v2; pub mod error; pub mod gemma2; pub mod llava_qwen2; +/// Internal: runtime JIT Metal kernels (Bonsai-Q1 bits=1 matvec/dequant). +mod metal_kernel; pub mod phi3; pub mod qwen3_moe; pub mod qwen3_next; diff --git a/crates/higgs-models/src/metal_kernel.rs b/crates/higgs-models/src/metal_kernel.rs new file mode 100644 index 00000000..61f86bf1 --- /dev/null +++ b/crates/higgs-models/src/metal_kernel.rs @@ -0,0 +1,498 @@ +//! Runtime JIT Metal kernels for Bonsai-Q1 (1-bit affine quantization). +//! +//! Upstream `oxideai/mlx-rs` ships no `bits=1` affine kernels (MLX gates affine +//! quant to `bits >= 2`), so `ops::quantized_matmul`/`ops::dequantize` with +//! `bits=1` fail at runtime with `Unable to load kernel affine_dequantize_*_b_1`. +//! +//! Rather than fork mlx-rs (which forces a full from-source mlx-c rebuild), we +//! add the missing kernels *from this crate* using the runtime JIT facility that +//! mlx-c already exposes (`mlx_fast_metal_kernel_*`) and that `mlx-sys` compiles +//! in. The kernels below are JIT-compiled by Metal at first use and cached by +//! MLX internally per template instantiation. This keeps us on the stock +//! `oxideai/mlx-rs` pin with no extra native recompile. +//! +//! The FFI plumbing (kernel handle wrapper, `Array` <-> `mlx_array`, vector +//! construction, error capture) mirrors the proven `qgemv_4bit` path in +//! [`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`. + +use std::ffi::{CStr, CString, c_char, c_void}; +use std::sync::OnceLock; + +use mlx_rs::{Array, Stream, error::Exception}; + +// --------------------------------------------------------------------------- +// FFI error capture (per-thread, mirrors qwen3_next). +// --------------------------------------------------------------------------- + +thread_local! { + static FFI_LAST_ERROR: std::cell::RefCell> = + const { std::cell::RefCell::new(None) }; +} + +/// Error handler registered once with MLX to capture error messages on the +/// calling thread. +#[allow(unsafe_code)] +unsafe extern "C" fn ffi_error_handler(msg: *const c_char, _data: *mut c_void) { + let s = unsafe { CStr::from_ptr(msg) } + .to_string_lossy() + .into_owned(); + FFI_LAST_ERROR.with(|cell| *cell.borrow_mut() = Some(s)); +} + +fn ensure_ffi_error_handler() { + static REGISTERED: OnceLock<()> = OnceLock::new(); + REGISTERED.get_or_init(|| { + #[allow(unsafe_code)] + unsafe { + mlx_sys::mlx_set_error_handler(Some(ffi_error_handler), std::ptr::null_mut(), None); + } + }); +} + +fn take_last_error() -> String { + FFI_LAST_ERROR + .with(|cell| cell.borrow_mut().take()) + .unwrap_or_else(|| "(no MLX error message captured)".to_owned()) +} + +// --------------------------------------------------------------------------- +// Cached kernel handle. +// --------------------------------------------------------------------------- + +/// Wraps a compiled `mlx_fast_metal_kernel`, freed on drop. +struct CachedMetalKernel(mlx_sys::mlx_fast_metal_kernel); + +// SAFETY: the handle is created once and only ever read (passed by value to +// `mlx_fast_metal_kernel_apply`); no interior mutability is shared across threads. +#[allow(unsafe_code)] +unsafe impl Send for CachedMetalKernel {} +#[allow(unsafe_code)] +unsafe impl Sync for CachedMetalKernel {} + +impl Drop for CachedMetalKernel { + fn drop(&mut self) { + #[allow(unsafe_code)] + unsafe { + mlx_sys::mlx_fast_metal_kernel_free(self.0); + } + } +} + +/// Number of simdgroups per threadgroup for the fused matvec. More simdgroups +/// help large-K layers (fewer chunk barriers). Overridable for tuning. +fn qmv_nsg(k_dim: i32) -> i32 { + static OVERRIDE: OnceLock> = OnceLock::new(); + let ovr = *OVERRIDE.get_or_init(|| { + std::env::var("HIGGS_BONSAI_QMV_NSG") + .ok() + .and_then(|s| s.parse::().ok()) + .filter(|n| matches!(n, 4 | 8 | 16 | 32)) + }); + ovr.unwrap_or(if k_dim > 8192 { 16 } else { 8 }) +} + +/// Build the vector-of-strings that names kernel inputs/outputs. +#[allow(unsafe_code)] +fn cstr_vec(names: &[&CStr]) -> mlx_sys::mlx_vector_string { + let ptrs: Vec<*const c_char> = names.iter().map(|s| s.as_ptr()).collect(); + unsafe { mlx_sys::mlx_vector_string_new_data(ptrs.as_ptr().cast_mut(), ptrs.len()) } +} + +// --------------------------------------------------------------------------- +// Fused 1-bit quantized matvec (decode hot path). +// +// y = x @ dequant(W).T for a single token (M = 1). +// Mirrors qgemv_4bit but unpacks 32 1-bit weights per uint32 word. +// One simdgroup per output row; x staged in threadgroup memory; simd_sum reduce. +// --------------------------------------------------------------------------- + +const QMV_KERNEL_SOURCE: &str = r" +constexpr int CHUNK = (K <= 8192) ? K : 8192; + +threadgroup OutT x_sh[CHUNK]; + +auto tg = threadgroup_position_in_grid.x; +auto sg = simdgroup_index_in_threadgroup; +auto lane = thread_index_in_simdgroup; +auto tid = thread_index_in_threadgroup; +auto n_sg = simdgroups_per_threadgroup; +uint tg_sz = n_sg * 32u; + +int row = tg * int(n_sg) + int(sg); +bool valid = (row < n_param); + +float acc = 0.0f; + +for (int k_off = 0; k_off < K; k_off += CHUNK) { + int k_end = min(k_off + CHUNK, K); + int k_len = k_end - k_off; + + for (uint i = tid; i < uint(k_len); i += tg_sz) { + x_sh[i] = x[k_off + i]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (valid) { + int wp_off = k_off / 32; + int wp_end = k_end / 32; + auto w_row = w + row * KPacked; + + for (int idx = wp_off + int(lane); idx < wp_end; idx += 32) { + uint packed = w_row[idx]; + int kl = (idx - wp_off) * 32; + + float dot_val = 0.0f; + float sum_x = 0.0f; + for (uint j = 0u; j < 32u; ++j) { + float xv = float(x_sh[kl + int(j)]); + float bit = float((packed >> j) & 1u); + dot_val += bit * xv; + sum_x += xv; + } + + int g = idx * 32 / GroupSize; + float s_val = float(sc[row * NumGroups + g]); + float b_val = float(bi[row * NumGroups + g]); + acc += s_val * dot_val + b_val * sum_x; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); +} + +if (valid) { + acc = simd_sum(acc); + if (lane == 0) { + y[row] = OutT(acc); + } +} +"; + +#[allow(unsafe_code)] +fn create_qmv_kernel() -> mlx_sys::mlx_fast_metal_kernel { + let in_vec = cstr_vec(&[c"w", c"sc", c"bi", c"x", c"n_param"]); + let out_vec = cstr_vec(&[c"y"]); + let source = CString::new(QMV_KERNEL_SOURCE).unwrap_or_default(); + unsafe { + let kernel = mlx_sys::mlx_fast_metal_kernel_new( + c"higgs_bonsai_q1_qmv".as_ptr(), + in_vec, + out_vec, + source.as_ptr(), + c"".as_ptr(), + false, // ensure_row_contiguous + false, // atomic_outputs + ); + mlx_sys::mlx_vector_string_free(in_vec); + mlx_sys::mlx_vector_string_free(out_vec); + kernel + } +} + +#[allow(unsafe_code)] +fn configure_qmv_kernel( + out_dtype: mlx_sys::mlx_dtype, + n_rows: i32, + k_dim: i32, + group_size: i32, +) -> mlx_sys::mlx_fast_metal_kernel_config { + unsafe { + let config = mlx_sys::mlx_fast_metal_kernel_config_new(); + mlx_sys::mlx_fast_metal_kernel_config_add_template_arg_dtype( + config, + c"OutT".as_ptr(), + out_dtype, + ); + mlx_sys::mlx_fast_metal_kernel_config_add_template_arg_int(config, c"K".as_ptr(), k_dim); + mlx_sys::mlx_fast_metal_kernel_config_add_template_arg_int( + config, + c"GroupSize".as_ptr(), + group_size, + ); + mlx_sys::mlx_fast_metal_kernel_config_add_template_arg_int( + config, + c"KPacked".as_ptr(), + k_dim / 32, + ); + mlx_sys::mlx_fast_metal_kernel_config_add_template_arg_int( + config, + c"NumGroups".as_ptr(), + k_dim / group_size, + ); + + let nsg = qmv_nsg(k_dim); + let n_tgs = (n_rows + nsg - 1) / nsg; + mlx_sys::mlx_fast_metal_kernel_config_set_grid(config, n_tgs * 32, nsg, 1); + mlx_sys::mlx_fast_metal_kernel_config_set_thread_group(config, 32, nsg, 1); + + let y_shape = [1, n_rows]; + mlx_sys::mlx_fast_metal_kernel_config_add_output_arg( + config, + y_shape.as_ptr(), + y_shape.len(), + out_dtype, + ); + config + } +} + +/// Fused 1-bit quantized matvec: `y = x @ dequant(weight).T` for a single token. +/// +/// `x` must hold exactly `in_features` elements (M = 1). `weight` is the packed +/// `[out_features, in_features/32]` uint32 matrix; `scales`/`biases` are +/// `[out_features, in_features/group_size]`. Output dtype matches `x`. +#[allow(unsafe_code)] +pub fn bonsai_q1_qmv( + x: &Array, + weight: &Array, + scales: &Array, + biases: &Array, + group_size: i32, +) -> Result { + ensure_ffi_error_handler(); + + let x_shape = x.shape(); + let weight_shape = weight.shape(); + let n_rows = weight_shape + .first() + .copied() + .ok_or_else(|| Exception::custom("bonsai_q1_qmv: weight has no rows"))?; + let k_packed = weight_shape + .get(1) + .copied() + .ok_or_else(|| Exception::custom("bonsai_q1_qmv: weight has no columns"))?; + let k_dim = k_packed * 32; // 32 one-bit weights per uint32 word + + 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 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 n_scalar = unsafe { mlx_sys::mlx_array_new_int(n_rows) }; + let input_ptrs = [ + w_flat.as_ptr(), + s_flat.as_ptr(), + b_flat.as_ptr(), + x_flat.as_ptr(), + n_scalar, + ]; + let inputs_vec = + unsafe { mlx_sys::mlx_vector_array_new_data(input_ptrs.as_ptr(), input_ptrs.len()) }; + + let mut outputs_vec = unsafe { mlx_sys::mlx_vector_array_new() }; + let status = unsafe { + mlx_sys::mlx_fast_metal_kernel_apply( + &raw mut outputs_vec, + cached.0, + inputs_vec, + config, + stream.as_ptr(), + ) + }; + + let result = if status != 0 { + Err(Exception::custom(format!( + "bonsai_q1_qmv failed: {}", + take_last_error() + ))) + } else { + let mut y_ptr = unsafe { mlx_sys::mlx_array_new() }; + unsafe { mlx_sys::mlx_vector_array_get(&raw mut y_ptr, outputs_vec, 0) }; + let y = unsafe { Array::from_ptr(y_ptr) }; + let trim_to = x_shape.len().saturating_sub(1); + let mut out_shape = x_shape + .get(..trim_to) + .ok_or_else(|| Exception::custom("bonsai_q1_qmv: x_shape too small"))? + .to_vec(); + out_shape.push(n_rows); + y.reshape(&out_shape) + }; + + unsafe { + mlx_sys::mlx_fast_metal_kernel_config_free(config); + mlx_sys::mlx_vector_array_free(inputs_vec); + mlx_sys::mlx_vector_array_free(outputs_vec); + mlx_sys::mlx_array_free(n_scalar); + } + result +} + +static QMV_KERNEL: OnceLock = OnceLock::new(); + +// --------------------------------------------------------------------------- +// 1-bit dequantize to dense (embedding gather + prefill matmul path). +// +// wd[n, c] = scales[n, c/G] * bit(w[n, c/32], c%32) + biases[n, c/G]. +// One thread per packed uint32 word (writes 32 dense outputs). +// --------------------------------------------------------------------------- + +const DEQUANT_KERNEL_SOURCE: &str = r" +uint gid = thread_position_in_grid.x; +if (gid >= uint(NWords)) { return; } + +uint n = gid / uint(KPacked); +uint idx = gid % uint(KPacked); +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)]); + +uint base = n * uint(K) + idx * 32u; +for (uint j = 0u; j < 32u; ++j) { + float bit = float((packed >> j) & 1u); + wd[base + j] = OutT(s_val * bit + b_val); +} +"; + +#[allow(unsafe_code)] +fn create_dequant_kernel() -> mlx_sys::mlx_fast_metal_kernel { + let in_vec = cstr_vec(&[c"w", c"sc", c"bi"]); + let out_vec = cstr_vec(&[c"wd"]); + let source = CString::new(DEQUANT_KERNEL_SOURCE).unwrap_or_default(); + unsafe { + let kernel = mlx_sys::mlx_fast_metal_kernel_new( + c"higgs_bonsai_q1_dequant".as_ptr(), + in_vec, + out_vec, + source.as_ptr(), + c"".as_ptr(), + false, + false, + ); + mlx_sys::mlx_vector_string_free(in_vec); + mlx_sys::mlx_vector_string_free(out_vec); + kernel + } +} + +#[allow(unsafe_code)] +fn configure_dequant_kernel( + out_dtype: mlx_sys::mlx_dtype, + n_rows: i32, + k_dim: i32, + group_size: i32, +) -> mlx_sys::mlx_fast_metal_kernel_config { + let k_packed = k_dim / 32; + let n_words = n_rows * k_packed; + unsafe { + let config = mlx_sys::mlx_fast_metal_kernel_config_new(); + mlx_sys::mlx_fast_metal_kernel_config_add_template_arg_dtype( + config, + c"OutT".as_ptr(), + out_dtype, + ); + mlx_sys::mlx_fast_metal_kernel_config_add_template_arg_int(config, c"K".as_ptr(), k_dim); + mlx_sys::mlx_fast_metal_kernel_config_add_template_arg_int( + config, + c"KPacked".as_ptr(), + k_packed, + ); + mlx_sys::mlx_fast_metal_kernel_config_add_template_arg_int( + config, + c"GroupSize".as_ptr(), + group_size, + ); + mlx_sys::mlx_fast_metal_kernel_config_add_template_arg_int( + config, + c"NumGroups".as_ptr(), + k_dim / group_size, + ); + mlx_sys::mlx_fast_metal_kernel_config_add_template_arg_int( + config, + c"NWords".as_ptr(), + n_words, + ); + + let tg: i32 = 256; + let grid = ((n_words + tg - 1) / tg) * tg; + mlx_sys::mlx_fast_metal_kernel_config_set_grid(config, grid, 1, 1); + mlx_sys::mlx_fast_metal_kernel_config_set_thread_group(config, tg, 1, 1); + + let wd_shape = [n_rows, k_dim]; + mlx_sys::mlx_fast_metal_kernel_config_add_output_arg( + config, + wd_shape.as_ptr(), + wd_shape.len(), + out_dtype, + ); + config + } +} + +/// Dequantize a packed 1-bit matrix to a dense `[out_features, in_features]` +/// array (dtype matches `scales`). Used for embedding gather and the prefill +/// (M > 1) matmul path. +#[allow(unsafe_code)] +pub fn bonsai_q1_dequant( + weight: &Array, + scales: &Array, + biases: &Array, + group_size: i32, +) -> Result { + ensure_ffi_error_handler(); + + let weight_shape = weight.shape(); + let n_rows = weight_shape + .first() + .copied() + .ok_or_else(|| Exception::custom("bonsai_q1_dequant: weight has no rows"))?; + let k_packed = weight_shape + .get(1) + .copied() + .ok_or_else(|| Exception::custom("bonsai_q1_dequant: weight has no columns"))?; + let k_dim = k_packed * 32; + + let w_flat = weight.reshape(&[-1])?; + let s_flat = scales.flatten(None, None)?; + let b_flat = 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 input_ptrs = [w_flat.as_ptr(), s_flat.as_ptr(), b_flat.as_ptr()]; + let inputs_vec = + unsafe { mlx_sys::mlx_vector_array_new_data(input_ptrs.as_ptr(), input_ptrs.len()) }; + + let mut outputs_vec = unsafe { mlx_sys::mlx_vector_array_new() }; + let status = unsafe { + mlx_sys::mlx_fast_metal_kernel_apply( + &raw mut outputs_vec, + cached.0, + inputs_vec, + config, + stream.as_ptr(), + ) + }; + + let result = if status != 0 { + Err(Exception::custom(format!( + "bonsai_q1_dequant failed: {}", + take_last_error() + ))) + } else { + let mut wd_ptr = unsafe { mlx_sys::mlx_array_new() }; + unsafe { mlx_sys::mlx_vector_array_get(&raw mut wd_ptr, outputs_vec, 0) }; + Ok(unsafe { Array::from_ptr(wd_ptr) }) + }; + + unsafe { + mlx_sys::mlx_fast_metal_kernel_config_free(config); + mlx_sys::mlx_vector_array_free(inputs_vec); + mlx_sys::mlx_vector_array_free(outputs_vec); + } + result +} + +static DEQUANT_KERNEL: OnceLock = OnceLock::new(); From 3694971e08f2c659451a0809fbf4973cf772d16f Mon Sep 17 00:00:00 2001 From: Peppi Littera Date: Tue, 2 Jun 2026 17:47:16 +0200 Subject: [PATCH 29/29] perf(bonsai): fuse q/k/v and gate/up projections into single matvecs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Concatenate q/k/v (and gate/up) packed weights row-wise at load so each layer issues one fused matvec dispatch instead of three (and two), then splits the output by slicing. Cuts per-token graph-build cost (forward_ms ~19% lower). GPU eval is unchanged — it is memory/glue-bound, not projection-dispatch-bound — so the end-to-end win is small (~1%); the main value is fewer ops for the upcoming compiled-decode path. Output stays coherent on 1.7B and 8B; oracle kernel tests still pass. Also added an #[ignore] decode-section profiling test (forward_profiled) used to localize the bottleneck. Co-Authored-By: Claude Opus 4.8 (1M context) --- crates/higgs-models/src/bonsai_q1.rs | 221 ++++++++++++++++++--------- 1 file changed, 145 insertions(+), 76 deletions(-) diff --git a/crates/higgs-models/src/bonsai_q1.rs b/crates/higgs-models/src/bonsai_q1.rs index 131e263d..b0e7706a 100644 --- a/crates/higgs-models/src/bonsai_q1.rs +++ b/crates/higgs-models/src/bonsai_q1.rs @@ -434,15 +434,33 @@ impl BonsaiQ1GpuLinear { ops::matmul(x, &wd.transpose_axes(&[1, 0])?) } } + + /// Fuse several linears along the output (row) axis into one, so a single + /// matvec dispatch replaces N. All parts must share `in_features` and group + /// layout (they do — q/k/v and gate/up all project from the same hidden). + fn concat_rows(parts: &[&Self]) -> Result { + let ws: Vec<&Array> = parts.iter().map(|p| &p.w).collect(); + let ss: Vec<&Array> = parts.iter().map(|p| &p.scales).collect(); + let bs: Vec<&Array> = parts.iter().map(|p| &p.biases).collect(); + // Axis 0 (row concat). NOTE: `ops::concatenate` (no axis) flattens to 1-D + // like numpy's axis=None — must use the explicit axis variant. + Ok(Self { + w: ops::concatenate_axis(&ws, 0)?, + scales: ops::concatenate_axis(&ss, 0)?, + biases: ops::concatenate_axis(&bs, 0)?, + out_features: parts.iter().map(|p| p.out_features).sum(), + in_features: parts[0].in_features, + }) + } } pub struct BonsaiQ1GpuLayer { - pub q_proj: BonsaiQ1GpuLinear, - pub k_proj: BonsaiQ1GpuLinear, - pub v_proj: BonsaiQ1GpuLinear, + /// Fused q+k+v projection (rows concatenated): one matvec instead of three. + /// Output split is `[heads*head_dim | kv_heads*head_dim | kv_heads*head_dim]`. + pub qkv_proj: BonsaiQ1GpuLinear, pub o_proj: BonsaiQ1GpuLinear, - pub gate_proj: BonsaiQ1GpuLinear, - pub up_proj: BonsaiQ1GpuLinear, + /// Fused gate+up projection (rows concatenated): output split `[inter | inter]`. + pub gate_up_proj: BonsaiQ1GpuLinear, pub down_proj: BonsaiQ1GpuLinear, pub q_norm: Array, pub k_norm: Array, @@ -476,13 +494,15 @@ impl BonsaiQ1Engine { pub fn to_gpu(self) -> Result { let mut gpu_layers = Vec::with_capacity(self.layers.len()); for layer in &self.layers { + let q = BonsaiQ1GpuLinear::from_packed(&layer.q_proj)?; + let k = BonsaiQ1GpuLinear::from_packed(&layer.k_proj)?; + let v = BonsaiQ1GpuLinear::from_packed(&layer.v_proj)?; + let gate = BonsaiQ1GpuLinear::from_packed(&layer.gate_proj)?; + let up = BonsaiQ1GpuLinear::from_packed(&layer.up_proj)?; gpu_layers.push(BonsaiQ1GpuLayer { - q_proj: BonsaiQ1GpuLinear::from_packed(&layer.q_proj)?, - k_proj: BonsaiQ1GpuLinear::from_packed(&layer.k_proj)?, - v_proj: BonsaiQ1GpuLinear::from_packed(&layer.v_proj)?, + qkv_proj: BonsaiQ1GpuLinear::concat_rows(&[&q, &k, &v])?, o_proj: BonsaiQ1GpuLinear::from_packed(&layer.o_proj)?, - gate_proj: BonsaiQ1GpuLinear::from_packed(&layer.gate_proj)?, - up_proj: BonsaiQ1GpuLinear::from_packed(&layer.up_proj)?, + gate_up_proj: BonsaiQ1GpuLinear::concat_rows(&[&gate, &up])?, down_proj: BonsaiQ1GpuLinear::from_packed(&layer.down_proj)?, q_norm: f16_vec_to_array(&layer.q_norm)?, k_norm: f16_vec_to_array(&layer.k_norm)?, @@ -699,6 +719,10 @@ impl BonsaiQ1Gpu { .map_err(|_| Exception::custom("heads overflows i32"))?; let kv_heads = i32::try_from(self.config.kv_heads) .map_err(|_| Exception::custom("kv_heads overflows i32"))?; + let head_dim_i = i32::try_from(self.config.head_dim) + .map_err(|_| Exception::custom("head_dim overflows i32"))?; + let inter_i = i32::try_from(self.config.inter) + .map_err(|_| Exception::custom("inter overflows i32"))?; let rms_eps = self.config.rms_norm_eps; for (layer, layer_cache) in self.layers.iter().zip(cache.iter_mut()) { @@ -707,26 +731,26 @@ impl BonsaiQ1Gpu { normed.eval()?; times.add("input_norm", t0.elapsed().as_nanos()); - // qkv projections — 3× quantized_matmul on the same input. + // Fused qkv projection — one matvec instead of three. let t0 = Instant::now(); - let q = layer.q_proj.forward(&normed)?; - let k = layer.k_proj.forward(&normed)?; - let v = layer.v_proj.forward(&normed)?; - q.eval()?; - k.eval()?; - v.eval()?; + let qkv = layer.qkv_proj.forward(&normed)?.reshape(&[ + B, + T, + heads + 2 * kv_heads, + head_dim_i, + ])?; + qkv.eval()?; times.add("qkv_proj", t0.elapsed().as_nanos()); - // Reshape to [B, L, n_heads, head_dim] then transpose to - // [B, n_heads, L, head_dim]. Metadata-only; lumped with qk_norm. - let q = q - .reshape(&[B, T, heads, -1])? + // Split heads along axis 2 then transpose to [B, n_heads, L, head_dim]. + let q = qkv + .index((.., .., ..heads, ..)) .transpose_axes(&[0, 2, 1, 3])?; - let k = k - .reshape(&[B, T, kv_heads, -1])? + let k = qkv + .index((.., .., heads..heads + kv_heads, ..)) .transpose_axes(&[0, 2, 1, 3])?; - let v = v - .reshape(&[B, T, kv_heads, -1])? + let v = qkv + .index((.., .., heads + kv_heads..heads + 2 * kv_heads, ..)) .transpose_axes(&[0, 2, 1, 3])?; let t0 = Instant::now(); @@ -792,11 +816,11 @@ impl BonsaiQ1Gpu { times.add("post_attn_norm", t0.elapsed().as_nanos()); let t0 = Instant::now(); - let gate = layer.gate_proj.forward(&normed_post)?; - let up = layer.up_proj.forward(&normed_post)?; - gate.eval()?; - up.eval()?; + let gate_up = layer.gate_up_proj.forward(&normed_post)?; + gate_up.eval()?; times.add("mlp_up_gate", t0.elapsed().as_nanos()); + let gate = gate_up.index((.., .., ..inter_i)); + let up = gate_up.index((.., .., inter_i..2 * inter_i)); let t0 = Instant::now(); let mlp_hidden = mlx_rs::nn::silu(&gate)?.multiply(&up)?; @@ -864,23 +888,30 @@ pub fn forward_trunk_free( i32::try_from(gpu.config.heads).map_err(|_| Exception::custom("heads overflows i32"))?; let kv_heads = i32::try_from(gpu.config.kv_heads) .map_err(|_| Exception::custom("kv_heads overflows i32"))?; + let head_dim_i = i32::try_from(gpu.config.head_dim) + .map_err(|_| Exception::custom("head_dim overflows i32"))?; + let inter_i = + i32::try_from(gpu.config.inter).map_err(|_| Exception::custom("inter overflows i32"))?; let rms_eps = gpu.config.rms_norm_eps; for (layer, layer_cache) in gpu.layers.iter().zip(cache.iter_mut()) { let normed = fast::rms_norm(&h, &layer.input_norm, rms_eps)?; - let q = layer.q_proj.forward(&normed)?; - let k = layer.k_proj.forward(&normed)?; - let v = layer.v_proj.forward(&normed)?; - - let q = q - .reshape(&[B, T, heads, -1])? + // One fused matvec for q/k/v, then split heads along axis 2 (views — no + // copy on the decode T=1 path) and transpose to [B, n_heads, T, head_dim]. + let qkv = + layer + .qkv_proj + .forward(&normed)? + .reshape(&[B, T, heads + 2 * kv_heads, head_dim_i])?; + let q = qkv + .index((.., .., ..heads, ..)) .transpose_axes(&[0, 2, 1, 3])?; - let k = k - .reshape(&[B, T, kv_heads, -1])? + let k = qkv + .index((.., .., heads..heads + kv_heads, ..)) .transpose_axes(&[0, 2, 1, 3])?; - let v = v - .reshape(&[B, T, kv_heads, -1])? + let v = qkv + .index((.., .., heads + kv_heads..heads + 2 * kv_heads, ..)) .transpose_axes(&[0, 2, 1, 3])?; let q = fast::rms_norm(&q, &layer.q_norm, rms_eps)?; @@ -917,8 +948,10 @@ pub fn forward_trunk_free( let h_post_attn = h.add(&attn_out)?; let normed_post = fast::rms_norm(&h_post_attn, &layer.post_attn_norm, rms_eps)?; - let gate = layer.gate_proj.forward(&normed_post)?; - let up = layer.up_proj.forward(&normed_post)?; + // One fused matvec for gate+up, then split along the last axis. + let gate_up = layer.gate_up_proj.forward(&normed_post)?; + let gate = gate_up.index((.., .., ..inter_i)); + let up = gate_up.index((.., .., inter_i..2 * inter_i)); let mlp_hidden = mlx_rs::nn::silu(&gate)?.multiply(&up)?; let mlp_out = layer.down_proj.forward(&mlp_hidden)?; @@ -944,11 +977,11 @@ pub struct BonsaiQ1DecodeState { pub cache: Vec>, } -/// Number of updatable `Array`s per decoder layer: -/// - `input_norm` + 3×(w,s,b) qkv + `q_norm` + `k_norm` + 3×(w,s,b) o_proj -/// ... wait: 1 + 3×3 + 2 + 3 + 1 + 3×3 = 1+9+2+3+1+9 = **25**. -/// Corresponds to the array push order in [`BonsaiQ1DecodeState::updatable_states`]. -const PER_LAYER_UPDATABLE: usize = 25; +/// Number of updatable `Array`s per decoder layer, matching the push order in +/// [`BonsaiQ1DecodeState::updatable_states`]: `input_norm` + (w,s,b) qkv_proj + +/// `q_norm` + `k_norm` + (w,s,b) o_proj + `post_attn_norm` + (w,s,b) gate_up_proj +/// + (w,s,b) down_proj = 1 + 3 + 2 + 3 + 1 + 3 + 3 = **16**. +const PER_LAYER_UPDATABLE: usize = 16; impl mlx_rs::utils::Updatable for BonsaiQ1DecodeState { fn updatable_states_len(&self) -> usize { @@ -981,27 +1014,18 @@ impl mlx_rs::utils::Updatable for BonsaiQ1DecodeState { v.push(&self.gpu.embed.biases); for layer in &self.gpu.layers { v.push(&layer.input_norm); - v.push(&layer.q_proj.w); - v.push(&layer.q_proj.scales); - v.push(&layer.q_proj.biases); - v.push(&layer.k_proj.w); - v.push(&layer.k_proj.scales); - v.push(&layer.k_proj.biases); - v.push(&layer.v_proj.w); - v.push(&layer.v_proj.scales); - v.push(&layer.v_proj.biases); + v.push(&layer.qkv_proj.w); + v.push(&layer.qkv_proj.scales); + v.push(&layer.qkv_proj.biases); v.push(&layer.q_norm); v.push(&layer.k_norm); v.push(&layer.o_proj.w); v.push(&layer.o_proj.scales); v.push(&layer.o_proj.biases); v.push(&layer.post_attn_norm); - v.push(&layer.gate_proj.w); - v.push(&layer.gate_proj.scales); - v.push(&layer.gate_proj.biases); - v.push(&layer.up_proj.w); - v.push(&layer.up_proj.scales); - v.push(&layer.up_proj.biases); + v.push(&layer.gate_up_proj.w); + v.push(&layer.gate_up_proj.scales); + v.push(&layer.gate_up_proj.biases); v.push(&layer.down_proj.w); v.push(&layer.down_proj.scales); v.push(&layer.down_proj.biases); @@ -1035,27 +1059,18 @@ impl mlx_rs::utils::Updatable for BonsaiQ1DecodeState { v.push(&mut self.gpu.embed.biases); for layer in &mut self.gpu.layers { v.push(&mut layer.input_norm); - v.push(&mut layer.q_proj.w); - v.push(&mut layer.q_proj.scales); - v.push(&mut layer.q_proj.biases); - v.push(&mut layer.k_proj.w); - v.push(&mut layer.k_proj.scales); - v.push(&mut layer.k_proj.biases); - v.push(&mut layer.v_proj.w); - v.push(&mut layer.v_proj.scales); - v.push(&mut layer.v_proj.biases); + v.push(&mut layer.qkv_proj.w); + v.push(&mut layer.qkv_proj.scales); + v.push(&mut layer.qkv_proj.biases); v.push(&mut layer.q_norm); v.push(&mut layer.k_norm); v.push(&mut layer.o_proj.w); v.push(&mut layer.o_proj.scales); v.push(&mut layer.o_proj.biases); v.push(&mut layer.post_attn_norm); - v.push(&mut layer.gate_proj.w); - v.push(&mut layer.gate_proj.scales); - v.push(&mut layer.gate_proj.biases); - v.push(&mut layer.up_proj.w); - v.push(&mut layer.up_proj.scales); - v.push(&mut layer.up_proj.biases); + v.push(&mut layer.gate_up_proj.w); + v.push(&mut layer.gate_up_proj.scales); + v.push(&mut layer.gate_up_proj.biases); v.push(&mut layer.down_proj.w); v.push(&mut layer.down_proj.scales); v.push(&mut layer.down_proj.biases); @@ -1330,4 +1345,58 @@ mod tests { ); } } + + /// Diagnostic: per-section decode timing on a real model. Forced per-section + /// evals inflate absolute numbers (lazy batching is killed) — the *ratios* + /// localize where decode time goes (projections vs attention vs lm_head). + #[test] + #[ignore = "profiling: set HIGGS_BONSAI_PROFILE_DIR to a Bonsai-Q1 model dir"] + #[allow(clippy::print_stderr, clippy::expect_used, clippy::cast_lossless)] + fn profile_decode_sections() { + let Ok(dir) = std::env::var("HIGGS_BONSAI_PROFILE_DIR") else { + eprintln!("skip: set HIGGS_BONSAI_PROFILE_DIR"); + return; + }; + let gpu = load_bonsai_q1(&dir).expect("load Bonsai-Q1"); + let mut cache: Vec> = Vec::new(); + + // Prefill a short prompt to populate the KV cache. + let prompt: Vec = (1..=16).collect(); + let pin = Array::from_slice(&prompt, &[1, 16]); + let mut warm = SectionTimes::new(); + gpu.forward_profiled(&pin, &mut cache, &mut warm) + .expect("prefill"); + + // Decode steps; skip the first few (JIT warmup), accumulate the rest. + let measured = 30; + let mut times = SectionTimes::new(); + let mut tok = 100i32; + for step in 0..(measured + 5) { + let din = Array::from_slice(&[tok], &[1, 1]); + if step < 5 { + let mut w = SectionTimes::new(); + gpu.forward_profiled(&din, &mut cache, &mut w) + .expect("decode warm"); + } else { + gpu.forward_profiled(&din, &mut cache, &mut times) + .expect("decode"); + } + tok = (tok * 7 + 13) % 60000 + 1; + } + + let total = times.total_ns() as f64; + eprintln!("=== decode section breakdown ({measured} steps) ==="); + for (name, ns, _cnt) in times.entries() { + eprintln!( + "{name:>16}: {:>7.3} ms {:>5.1}%", + ns as f64 / 1e6, + 100.0 * ns as f64 / total + ); + } + eprintln!( + "TOTAL {:.2} ms / {measured} steps = {:.3} ms/step (absolutes inflated by forced evals; use ratios)", + total / 1e6, + total / 1e6 / measured as f64 + ); + } }