From b21263851b0dbf9f2839c9338fc785b062eb3d2e Mon Sep 17 00:00:00 2001 From: baseRT sync Date: Thu, 8 Oct 2026 23:51:50 +1100 Subject: [PATCH] sync: open-core source for 0.3.0 Generated by tools/sync_public.py from the internal tree at 6f1215286. Only allow-listed paths are included; tuned quant profiles and internal tooling are excluded by rule. --- base-convert/Cargo.lock | 16 +- base-convert/Cargo.toml | 2 +- base-convert/crates/base-arch/src/dflash.rs | 1178 ++++++++++++ base-convert/crates/base-arch/src/gemma.rs | 8 +- base-convert/crates/base-arch/src/lib.rs | 240 ++- base-convert/crates/base-arch/src/llama.rs | 174 +- base-convert/crates/base-arch/src/qwen.rs | 72 +- base-convert/crates/base-awq/src/lib.rs | 36 +- base-convert/crates/base-awq/src/sidecar.rs | 33 +- .../crates/base-convert/src/gpt_oss.rs | 14 + base-convert/crates/base-convert/src/hub.rs | 52 +- base-convert/crates/base-convert/src/lora.rs | 474 +++++ base-convert/crates/base-convert/src/main.rs | 1674 ++++++++++++++++- .../base-convert/tests/mlx_passthrough.rs | 295 +++ base-convert/crates/base-format/src/header.rs | 37 +- base-convert/crates/base-format/src/lib.rs | 4 +- base-convert/crates/base-format/src/reader.rs | 67 +- base-convert/crates/base-format/src/writer.rs | 178 +- .../crates/base-format/tests/roundtrip.rs | 48 + base-convert/crates/base-hub/catalog.json | 930 ++++++++- base-convert/crates/base-hub/src/catalog.rs | 19 + base-convert/crates/base-hub/src/download.rs | 27 +- base-convert/crates/base-hub/src/gen.rs | 44 +- base-convert/crates/base-hub/src/parts.rs | 8 +- base-convert/crates/base-hub/src/scan.rs | 115 +- .../base-quant/tests/canonical_pipeline.rs | 1 + base-convert/crates/base-readers/src/hf.rs | 363 +++- base-convert/crates/base-sign/src/lib.rs | 31 +- .../crates/base-sign/tests/sign_roundtrip.rs | 1 + bindings/node/package.json | 2 +- bindings/python/baseRT/__init__.py | 2 +- bindings/python/setup.py | 2 +- bindings/rust/baseRT-sys/Cargo.toml | 2 +- bindings/rust/baseRT-sys/src/lib.rs | 324 +++- bindings/rust/baseRT/Cargo.toml | 2 +- .../swift/Sources/CBaseRT/include/baseRT.h | 604 +++++- include/baseRT/baseRT.h | 604 +++++- include/baseRT/baseRT_tick.h | 973 ++++++++++ 38 files changed, 8310 insertions(+), 346 deletions(-) create mode 100644 base-convert/crates/base-arch/src/dflash.rs create mode 100644 base-convert/crates/base-convert/src/lora.rs create mode 100644 base-convert/crates/base-convert/tests/mlx_passthrough.rs create mode 100644 include/baseRT/baseRT_tick.h diff --git a/base-convert/Cargo.lock b/base-convert/Cargo.lock index 163d3f5..0bd9ad9 100644 --- a/base-convert/Cargo.lock +++ b/base-convert/Cargo.lock @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "base-arch" -version = "0.2.6" +version = "0.3.0" dependencies = [ "anyhow", "base-format", @@ -162,7 +162,7 @@ dependencies = [ [[package]] name = "base-awq" -version = "0.2.6" +version = "0.3.0" dependencies = [ "anyhow", "base-format", @@ -176,7 +176,7 @@ dependencies = [ [[package]] name = "base-convert" -version = "0.2.6" +version = "0.3.0" dependencies = [ "anyhow", "base-arch", @@ -199,7 +199,7 @@ dependencies = [ [[package]] name = "base-format" -version = "0.2.6" +version = "0.3.0" dependencies = [ "anyhow", "bitflags", @@ -215,7 +215,7 @@ dependencies = [ [[package]] name = "base-hub" -version = "0.2.6" +version = "0.3.0" dependencies = [ "anyhow", "base-format", @@ -234,7 +234,7 @@ dependencies = [ [[package]] name = "base-quant" -version = "0.2.6" +version = "0.3.0" dependencies = [ "anyhow", "base-format", @@ -246,7 +246,7 @@ dependencies = [ [[package]] name = "base-readers" -version = "0.2.6" +version = "0.3.0" dependencies = [ "anyhow", "base-format", @@ -260,7 +260,7 @@ dependencies = [ [[package]] name = "base-sign" -version = "0.2.6" +version = "0.3.0" dependencies = [ "anyhow", "base-format", diff --git a/base-convert/Cargo.toml b/base-convert/Cargo.toml index 2db62b0..c3b3d93 100644 --- a/base-convert/Cargo.toml +++ b/base-convert/Cargo.toml @@ -12,7 +12,7 @@ members = [ ] [workspace.package] -version = "0.2.6" +version = "0.3.0" edition = "2021" license = "Apache-2.0" repository = "https://github.com/basecompute/baseRT" diff --git a/base-convert/crates/base-arch/src/dflash.rs b/base-convert/crates/base-arch/src/dflash.rs new file mode 100644 index 0000000..0358b7f --- /dev/null +++ b/base-convert/crates/base-arch/src/dflash.rs @@ -0,0 +1,1178 @@ +//! Block drafters (DFlash / DSpark) — speculative-decoding drafters that +//! ship as their own HF checkpoints and convert to a SIDECAR `.base` bundle +//! (arch `dflash` / `dspark`) the runtime loads against its target +//! (`baseRT_load_drafter`). +//! +//! A block drafter is a short qwen3-shaped decoder (5-6 layers of the +//! target's width) plus a context projection: `fc` ([dim, n_taps * dim]) +//! over the concatenation of the target's hidden states at +//! `target_layer_ids` (hidden_states[id + 1]), then `hidden_norm`. The +//! checkpoint's `model_type` is the backbone's (`qwen3`), so the drafter +//! is recognised by its `architectures` entry (`DFlashDraftModel`, +//! `DFlash2DraftModel`, `Qwen3DSparkModel`) — converting it as a Qwen3 would +//! drop `fc` / `hidden_norm` and write a broken LM. +//! +//! Tensor policy: every tensor stays f16 (the drafter's acceptance rate is +//! the whole point of the scheme, and the runtime's block step is bound by +//! the target's verify anyway). `fc.weight` is split by input column into +//! one `fc.{i}.weight` per tap (base-convert's `FcSplitProvider`), so the +//! runtime runs one GEMM per tap block of the tap-major input. +//! +//! Header config: the drafter's own geometry in the HF-style keys the +//! runtime parses (hidden_size, num_hidden_layers, heads, head_dim, +//! intermediate_size, rms_norm_eps, rope_theta, vocab_size, sliding +//! window as a period) plus the drafting fields under +//! [`ArchConfig::extra_config`]: `speculator_kind`, `block_size`, +//! `mask_token_id`, `logits_start`, `target_layer_ids`, +//! `num_target_layers`, `markov_rank`, `confidence_head`, and DFlash 2's +//! `conv_kernel_size` / `conv_group_size` / `selector_rank` / +//! `selector_top_k` when present. + +use crate::{ArchConfig, HfMapper, ValueTransform}; +use anyhow::{anyhow, bail, Result}; + +/// Scale folded into every `fc.{i}.weight` (see [`ValueTransform::Scale`]): +/// the fc output feeds only `hidden_norm`, so the drafter's math is +/// unchanged while the f16 output stays far from overflow. Recorded in the +/// header as `fc_scale`. +pub const FC_SCALE_DIV: u32 = 256; + +fn fc_transform(canonical: &str) -> Option { + let rest = canonical.strip_prefix("fc.")?; + let (idx, tail) = rest.split_once('.')?; + (tail == "weight" && idx.parse::().is_ok()).then_some(ValueTransform::DivBy(FC_SCALE_DIV)) +} + +pub struct DflashHfMapper; +pub struct DsparkHfMapper; + +/// Canonical arch strings of the block-drafter family. +pub fn is_block_drafter_arch(arch: &str) -> bool { + arch == "dflash" || arch == "dspark" +} + +/// Pick a drafter mapper from `config.json`'s `architectures`; None for an +/// ordinary LM checkpoint (which routes by `model_type` as before). +pub fn hf_mapper_for_architectures(config: &serde_json::Value) -> Option<&'static dyn HfMapper> { + let archs = config.get("architectures")?.as_array()?; + for a in archs.iter().filter_map(|v| v.as_str()) { + if a.starts_with("DFlash") { + return Some(&DflashHfMapper); + } + if a.contains("DSpark") { + return Some(&DsparkHfMapper); + } + } + None +} + +fn u64_at(c: &serde_json::Value, dc: Option<&serde_json::Value>, key: &str) -> Option { + dc.and_then(|d| d.get(key)) + .or_else(|| c.get(key)) + .and_then(|v| v.as_u64()) +} + +/// Shared config extraction. DFlash nests its drafting fields under +/// `dflash_config`; DSpark keeps them at the top level. Either way the +/// backbone fields are the standard HF ones. +/// RedHat's `speculators` packaging of a block drafter +/// (`speculators_model_type` "dflash" / "dflash2" / "dspark"): the layer config nested +/// under `transformer_layer_config`, the drafting fields at the top level, +/// taps as `aux_hidden_state_layer_ids` in hidden_states indices (z-lab's +/// `target_layer_ids` + 1), causality of the sliding layers stated as +/// `sliding_window_non_causal`, and `sample_from_anchor` choosing whether the +/// anchor row drafts. Flattened into the z-lab form read below. +fn flatten_speculators_block(c: &serde_json::Value) -> Result { + let mut flat = c + .get("transformer_layer_config") + .cloned() + .ok_or_else(|| anyhow!("speculators config without transformer_layer_config"))?; + let o = flat + .as_object_mut() + .ok_or_else(|| anyhow!("transformer_layer_config is not an object"))?; + for k in [ + "architectures", + "block_size", + "mask_token_id", + "conv_kernel_size", + "conv_group_size", + "selector_rank", + "selector_top_k", + "markov_rank", + "markov_head_type", + "enable_confidence_head", + "confidence_head_with_markov", + ] { + if let Some(v) = c.get(k).filter(|v| !v.is_null()) { + o.insert(k.to_string(), v.clone()); + } + } + // A reduced draft vocabulary (RedHat's Qwen3-30B-A3B / Nemotron DFlash: + // 32000 rows of lm_head plus a `d2t` offset table back to target ids) + // is carried through; drafter_config_from_hf decides which kinds take it. + let vocab = o.get("vocab_size").and_then(|v| v.as_u64()).unwrap_or(0); + if let Some(dv) = c.get("draft_vocab_size").and_then(|v| v.as_u64()) { + if dv == 0 || dv > vocab { + bail!("draft_vocab_size {dv} is outside 1..={vocab}"); + } + if dv < vocab { + o.insert("draft_vocab_size".into(), serde_json::json!(dv)); + } + } + let ids: Vec = c + .get("aux_hidden_state_layer_ids") + .and_then(|v| v.as_array()) + .map(|a| a.iter().filter_map(|x| x.as_i64()).collect()) + .unwrap_or_default(); + if ids.iter().any(|&i| i < 1) { + bail!("speculators aux_hidden_state_layer_ids {ids:?} name the embedding output; not supported"); + } + o.insert( + "target_layer_ids".into(), + serde_json::json!(ids.iter().map(|i| i - 1).collect::>()), + ); + // Sliding layers attend causally unless the packaging says otherwise; + // full layers are bidirectional either way (the z-lab rule, stated + // explicitly here). + if let Some(nc) = c.get("sliding_window_non_causal").and_then(|v| v.as_bool()) { + if nc { + o.insert("is_causal".into(), serde_json::json!(false)); + } + } + if let Some(a) = c.get("sample_from_anchor").and_then(|v| v.as_bool()) { + o.insert( + "logits_start".into(), + serde_json::json!(if a { 0 } else { 1 }), + ); + } + Ok(flat) +} + +fn drafter_config_from_hf(c: &serde_json::Value, kind: &str) -> Result { + let flat; + let speculators = c + .get("speculators_model_type") + .and_then(|v| v.as_str()) + .is_some_and(|t| t.starts_with("dflash") || t == "dspark"); + let c = if speculators { + flat = flatten_speculators_block(c)?; + &flat + } else { + c + }; + let backbone = c.get("model_type").and_then(|v| v.as_str()).unwrap_or(""); + // Gemma 4 backbone (deepseek-ai/dspark_gemma4_*: `Gemma4DSparkModel`, + // model_type gemma4_text): the Gemma 4 layer config, recorded as + // `drafter_backbone` so the runtime runs the block rows through its + // Gemma 4 encoder and the context pass in Gemma 4 form. + let gemma4 = backbone.starts_with("gemma4"); + // RedHat's speculators DFlash heads for Qwen3-30B-A3B and Nemotron say + // `llama` but carry Qwen3's per-head q/k norms. The label alone cannot + // tell those from a genuinely norm-free Llama backbone, which the runtime + // cannot run (it loads every dflash / dspark sidecar as Qwen3 and demands + // q/k norms on every layer), so the tensors decide: + // [`check_block_drafter_qk_norms`] refuses a sidecar without them. + if backbone != "qwen3" && backbone != "llama" && !gemma4 { + bail!( + "block drafter backbone model_type {backbone:?} is not qwen3, llama or gemma4 — the engine runs \ + llama/qwen3- and gemma4-shaped drafters only" + ); + } + if gemma4 && kind != "dspark" { + bail!("a gemma4-backbone block drafter is wired for DSpark only"); + } + let mut cfg = if gemma4 { + use crate::HfMapper; + crate::gemma::Gemma4HfMapper.config_from_hf(c)? + } else { + crate::llama::hf_generic_config(c)? + }; + // The logit head is the target's (DFlash) or the drafter's own + // `lm_head.weight` (DSpark): never synthesize a tied copy. + cfg.tie_word_embeddings = false; + + let dc = c.get("dflash_config"); + let target_layer_ids: Vec = dc + .and_then(|d| d.get("target_layer_ids")) + .or_else(|| c.get("target_layer_ids")) + .and_then(|v| v.as_array()) + .map(|a| a.iter().filter_map(|x| x.as_u64()).collect()) + .unwrap_or_default(); + if target_layer_ids.is_empty() { + bail!("config.json: target_layer_ids missing (dflash_config or top level) — the drafter's taps"); + } + if target_layer_ids.len() > 8 { + bail!( + "config.json: {} target_layer_ids; the engine taps at most 8 layers per forward", + target_layer_ids.len() + ); + } + // Two DSpark packagings. DeepSeek's DeepSpec (`Qwen3DSparkModel`) ships its + // own embed / lm_head and predicts from block row 0 (anchor as position + // 0). SpecForge's (`DSparkDraftModel`, `projector_type: "dspark"`; e.g. + // RadixArk/Qwen3.8-27B-DSpark) is a DFlash backbone: drafts from rows 1.. + // (`draft_hidden[:, -block_size + 1:]`) through the TARGET's lm_head. + let archs: Vec<&str> = c + .get("architectures") + .and_then(|v| v.as_array()) + .map(|a| a.iter().filter_map(|v| v.as_str()).collect()) + .unwrap_or_default(); + // RedHat's speculators DSpark reuses the `DSparkDraftModel` name but ships + // its own embed / lm_head (DeepSpec-style), so it is not SpecForge's. + let specforge_dspark = kind == "dspark" && archs.contains(&"DSparkDraftModel") && !speculators; + let ds_cfg = c.get("dspark_config"); + let markov_type = ds_cfg + .and_then(|d| d.get("markov_head_type")) + .or_else(|| dc.and_then(|d| d.get("markov_head_type"))) + .or_else(|| c.get("markov_head_type")) + .and_then(|v| v.as_str()) + .unwrap_or("vanilla"); + if kind == "dspark" && markov_type != "vanilla" { + bail!("DSpark markov_head_type {markov_type:?} is not implemented (vanilla only)"); + } + let default_block = if kind == "dspark" { 7 } else { 16 }; + let block_size = u64_at(c, dc, "block_size").unwrap_or(default_block); + if !(2..=16).contains(&block_size) { + bail!("config.json: block_size {block_size} outside 2..16 (the verify feed width)"); + } + let mask_token_id = u64_at(c, dc, "mask_token_id").ok_or_else(|| { + anyhow!("config.json: mask_token_id missing — the token the draft block is padded with") + })?; + // DFlash reserves block row 0 as the anchor and drafts from the mask + // rows (logits_start 1); DeepSpec-native DSpark heads predict from + // row 0 (anchor-as-position-0). A checkpoint may say so explicitly. + let logits_start = + u64_at(c, dc, "logits_start").unwrap_or(if kind == "dspark" && !specforge_dspark { + 0 + } else { + 1 + }); + let num_target_layers = u64_at(c, dc, "num_target_layers").unwrap_or(0); + let markov_rank = c + .get("markov_rank") + .or_else(|| ds_cfg.and_then(|d| d.get("markov_rank"))) + .and_then(|v| v.as_u64()) + .unwrap_or(0); + let confidence_head = c + .get("enable_confidence_head") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + + // Sliding-window schedule → the runtime's period form (layer i is + // global iff (i + 1) % pattern == 0). [S,S,S,S,S,F] → 6; all sliding → + // n_layers + 1 (never global). + let layer_types: Vec = c + .get("layer_types") + .and_then(|v| v.as_array()) + .map(|a| { + a.iter() + .filter_map(|x| x.as_str().map(str::to_string)) + .collect() + }) + .unwrap_or_default(); + let window = c + .get("sliding_window") + .and_then(|v| v.as_u64()) + .unwrap_or(0) as u32; + let use_sliding = c + .get("use_sliding_window") + .and_then(|v| v.as_bool()) + .unwrap_or(true); + let sliding: Vec = layer_types + .iter() + .map(|t| t == "sliding_attention") + .collect(); + if use_sliding && window > 0 && sliding.iter().any(|&s| s) { + let n = sliding.len(); + let pattern = (1..=n + 1).find(|&p| (0..n).all(|i| sliding[i] == ((i + 1) % p != 0))); + match pattern { + Some(p) => { + cfg.sliding_window = window; + cfg.sliding_window_pattern = p as u32; + } + None => bail!( + "config.json: irregular sliding-window schedule {layer_types:?} — the runtime expresses windows \ + as a period" + ), + } + } + + use serde_json::json; + let x = &mut cfg.extra_config; + x.insert("speculator_kind".into(), json!(kind)); + x.insert("block_size".into(), json!(block_size)); + x.insert("mask_token_id".into(), json!(mask_token_id)); + x.insert("logits_start".into(), json!(logits_start)); + x.insert("target_layer_ids".into(), json!(target_layer_ids)); + x.insert("fc_scale".into(), json!(1.0 / FC_SCALE_DIV as f64)); + if num_target_layers > 0 { + x.insert("num_target_layers".into(), json!(num_target_layers)); + } + if markov_rank > 0 { + x.insert("markov_rank".into(), json!(markov_rank)); + } + if confidence_head { + x.insert("confidence_head".into(), json!(true)); + } + if specforge_dspark { + // The runtime then lets the logit head fall through to the target's. + x.insert("drafter_head".into(), json!("target")); + } + if gemma4 { + x.insert("drafter_backbone".into(), json!("gemma4")); + } + if !layer_types.is_empty() { + x.insert("drafter_layer_types".into(), json!(layer_types)); + } + // Which block layers attend causally (z-lab/dflash Qwen3DFlashAttention): + // an explicit `is_causal` applies to every layer; absent, a + // `sliding_attention` layer is causal and a full one bidirectional. The + // sliding-window DFlash drafters (Qwen3.6-27B, Qwen3.5-35B-A3B) were + // trained that way, and running their sliding layers bidirectionally lets + // every block row read the rows after it. DSpark's reference attends the + // block bidirectionally throughout. + if kind == "dflash" { + let n = cfg.num_hidden_layers as usize; + let causal: Vec = match c.get("is_causal").and_then(|v| v.as_bool()) { + Some(all) => vec![all; n], + None => (0..n) + .map(|i| layer_types.get(i).is_some_and(|t| t == "sliding_attention")) + .collect(), + }; + // Always written for DFlash, all-false included: a bundle WITHOUT the + // key predates it, and the loader then applies the sliding rule + // itself — so the key's absence must never mean "all bidirectional". + x.insert("drafter_causal_layers".into(), json!(causal)); + } + // Logit / embedding scalings the reference applies when configured + // (`_draft_value`: dflash_config first, then top level). None of the + // published drafters sets them; refused rather than silently ignored. + for k in ["input_embedding_scale", "output_multiplier"] { + let v = dc + .and_then(|d| d.get(k)) + .or_else(|| c.get(k)) + .and_then(|v| v.as_f64()); + if v.is_some_and(|v| v != 1.0) { + bail!( + "config.json: {k} {} is not implemented for block drafters", + v.unwrap() + ); + } + } + // The reference caps the block logits (compute_logits: cap * tanh(x / cap); + // the z-lab Gemma 4 drafters carry the target's 30). Recorded for the + // runtime, which caps the rows its selector / Markov chain / host argmax + // read (a plain argmax is unchanged by a monotonic cap). + let softcap = dc + .and_then(|d| d.get("final_logit_softcapping")) + .or_else(|| c.get("final_logit_softcapping")) + .and_then(|v| v.as_f64()) + .filter(|v| *v > 0.0); + if let Some(cap) = softcap { + x.insert("drafter_logit_softcap".into(), json!(cap)); + } + // DFlash 2 (grouped dynamic conv + candidate selector): recorded for + // the runtime that implements them; DFlash 1 behaviour when absent. + for k in [ + "conv_kernel_size", + "conv_group_size", + "selector_rank", + "selector_top_k", + ] { + if let Some(v) = u64_at(c, dc, k) { + x.insert(k.into(), json!(v)); + } + } + // A reduced draft vocabulary: the drafter's own lm_head has + // `draft_vocab_size` rows and `d2t` maps each back to a target id; the + // embedding stays the full `vocab_size`. Only a plain DFlash head takes + // it — DSpark's Markov head and DFlash 2's selector are indexed by the + // full vocabulary. + if let Some(dv) = c.get("draft_vocab_size").and_then(|v| v.as_u64()) { + if dv < cfg.vocab_size as u64 { + if kind != "dflash" || x.contains_key("selector_rank") { + bail!("a reduced draft vocabulary ({dv} of {}) is supported for DFlash (1) drafters only", cfg.vocab_size); + } + x.insert("draft_vocab_size".into(), json!(dv)); + } + } + Ok(cfg) +} + +/// Refuse a llama/qwen3-backbone block drafter whose decoder layers lack the +/// per-head q/k norms. The runtime's `arch_from_config` maps every dflash / +/// dspark sidecar to Qwen3, and its drafter manifest requires +/// `attention.q_norm` / `attention.k_norm` on every layer, so a norm-free +/// (true Llama) backbone would convert cleanly and then fail to load with a +/// missing-tensor error. `canonical` is the converter's mapped tensor names +/// (`layers.N.self_attn.q_norm.weight`). Gemma 4 backbones map through their +/// own tensor names and always carry the norms, so they are not checked here. +pub fn check_block_drafter_qk_norms<'a>( + cfg: &ArchConfig, + canonical: impl IntoIterator, +) -> Result<()> { + if cfg + .extra_config + .get("drafter_backbone") + .and_then(|v| v.as_str()) + == Some("gemma4") + { + return Ok(()); + } + let have: std::collections::HashSet<&str> = canonical.into_iter().collect(); + let missing: Vec = (0..cfg.num_hidden_layers) + .flat_map(|l| ["q_norm", "k_norm"].map(|n| format!("layers.{l}.self_attn.{n}.weight"))) + .filter(|n| !have.contains(n.as_str())) + .collect(); + if missing.is_empty() { + return Ok(()); + } + let shown = missing + .iter() + .take(4) + .cloned() + .collect::>() + .join(", "); + let more = if missing.len() > 4 { + format!(" and {} more", missing.len() - 4) + } else { + String::new() + }; + bail!( + "block drafter layers lack per-head q/k norms (missing {shown}{more}): the runtime runs \ + dflash / dspark drafters through its Qwen3 encoder, which needs them on every layer — \ + a norm-free Llama drafter backbone is not supported" + ) +} + +impl HfMapper for DflashHfMapper { + fn canonical_arch(&self) -> &'static str { + "dflash" + } + fn config_from_hf(&self, c: &serde_json::Value) -> Result { + let mut cfg = drafter_config_from_hf(c, "dflash")?; + // DFlash 2 (`DFlash2DraftModel`): two grouped dynamic convolutions per + // layer and a candidate selector over the block's top-k. A checkpoint + // missing any of their four parameters would load as DFlash 1 and + // draft with arithmetic it was not trained with, so all four are + // required, and `dflash_version` 2 tells the loader to demand the + // tensors they size. + let v2 = c + .get("architectures") + .and_then(|v| v.as_array()) + .is_some_and(|a| { + a.iter() + .filter_map(|v| v.as_str()) + .any(|a| a.starts_with("DFlash2")) + }); + if v2 { + let x = &cfg.extra_config; + let get = |k: &str| x.get(k).and_then(|v| v.as_u64()).unwrap_or(0); + let missing: Vec<&str> = [ + "conv_kernel_size", + "conv_group_size", + "selector_rank", + "selector_top_k", + ] + .into_iter() + .filter(|k| get(k) == 0) + .collect(); + if !missing.is_empty() { + bail!("DFlash2 config.json: {missing:?} missing or zero"); + } + let hidden = cfg.hidden_size as u64; + let gs = get("conv_group_size"); + if hidden % gs != 0 { + bail!("DFlash2 config.json: conv_group_size {gs} does not divide hidden_size {hidden}"); + } + if get("selector_top_k") > 64 { + bail!( + "DFlash2 config.json: selector_top_k {} > 64 (the selector keeps at most 64 candidates)", + get("selector_top_k") + ); + } + cfg.extra_config + .insert("dflash_version".into(), serde_json::json!(2)); + } + Ok(cfg) + } + fn value_transform(&self, canonical: &str) -> Option { + fc_transform(canonical) + } +} + +impl HfMapper for DsparkHfMapper { + fn canonical_arch(&self) -> &'static str { + "dspark" + } + fn config_from_hf(&self, c: &serde_json::Value) -> Result { + drafter_config_from_hf(c, "dspark") + } + fn value_transform(&self, canonical: &str) -> Option { + fc_transform(canonical) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn base() -> serde_json::Value { + json!({ + "model_type": "qwen3", "hidden_size": 2560, "num_hidden_layers": 5, + "num_attention_heads": 32, "num_key_value_heads": 8, "head_dim": 128, + "intermediate_size": 9728, "vocab_size": 151936, "rms_norm_eps": 1e-6, + "rope_theta": 1000000.0, "tie_word_embeddings": true + }) + } + + #[test] + fn dflash_config_reads_nested_fields_and_windows() { + let mut c = base(); + c["architectures"] = json!(["DFlashDraftModel"]); + c["dflash_config"] = json!({"block_size": 16, "mask_token_id": 151669, "target_layer_ids": [1, 9, 17, 25, 33]}); + c["layer_types"] = json!([ + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention" + ]); + c["sliding_window"] = json!(4096); + assert!(matches!( + hf_mapper_for_architectures(&c).map(|m| m.canonical_arch()), + Some("dflash") + )); + let cfg = DflashHfMapper.config_from_hf(&c).unwrap(); + assert!( + !cfg.tie_word_embeddings, + "never synthesize a head for a drafter" + ); + assert_eq!(cfg.sliding_window, 4096); + assert_eq!(cfg.sliding_window_pattern, 5); + assert_eq!(cfg.extra_config["block_size"], json!(16)); + assert_eq!(cfg.extra_config["logits_start"], json!(1)); + assert_eq!( + cfg.extra_config["target_layer_ids"], + json!([1, 9, 17, 25, 33]) + ); + assert_eq!(cfg.extra_config["speculator_kind"], json!("dflash")); + // No `is_causal`: the sliding layers attend causally, the full one + // bidirectionally (the reference's rule). + assert_eq!( + cfg.extra_config["drafter_causal_layers"], + json!([true, true, true, true, false]) + ); + // All-sliding (DFlash 2 heads): a period no layer reaches. + c["layer_types"] = json!(vec!["sliding_attention"; 5]); + let cfg = DflashHfMapper.config_from_hf(&c).unwrap(); + assert_eq!(cfg.sliding_window_pattern, 6); + // An explicit is_causal=false overrides the sliding rule, and is + // written out (all false) rather than left absent. + c["is_causal"] = json!(false); + let cfg = DflashHfMapper.config_from_hf(&c).unwrap(); + assert_eq!( + cfg.extra_config["drafter_causal_layers"], + json!(vec![false; 5]) + ); + } + + #[test] + fn dflash2_requires_its_conv_and_selector_geometry() { + let mut c = base(); + c["architectures"] = json!(["DFlash2DraftModel"]); + c["is_causal"] = json!(false); + c["dflash_config"] = json!({ + "block_size": 8, "mask_token_id": 151669, "target_layer_ids": [1, 9, 17, 25, 33], + "conv_kernel_size": 2, "conv_group_size": 16, "selector_rank": 256, "selector_top_k": 16 + }); + let cfg = DflashHfMapper.config_from_hf(&c).unwrap(); + assert_eq!(cfg.extra_config["dflash_version"], json!(2)); + assert_eq!(cfg.extra_config["conv_group_size"], json!(16)); + assert_eq!(cfg.extra_config["selector_top_k"], json!(16)); + c["dflash_config"]["selector_rank"] = serde_json::Value::Null; + assert!( + DflashHfMapper.config_from_hf(&c).is_err(), + "selector_rank missing" + ); + c["dflash_config"]["selector_rank"] = json!(256); + c["dflash_config"]["conv_group_size"] = json!(7); + assert!( + DflashHfMapper.config_from_hf(&c).is_err(), + "group does not divide hidden" + ); + // A DFlash 1 checkpoint never gets the version tag. + c["architectures"] = json!(["DFlashDraftModel"]); + c["dflash_config"]["conv_group_size"] = json!(16); + let cfg = DflashHfMapper.config_from_hf(&c).unwrap(); + assert!(!cfg.extra_config.contains_key("dflash_version")); + // Configured logit scalings are refused, not ignored. + c["dflash_config"]["output_multiplier"] = json!(0.5); + assert!( + DflashHfMapper.config_from_hf(&c).is_err(), + "output_multiplier" + ); + } + + #[test] + fn dspark_config_is_top_level_and_anchor_at_zero() { + let mut c = base(); + c["architectures"] = json!(["Qwen3DSparkModel"]); + c["block_size"] = json!(7); + c["mask_token_id"] = json!(151669); + c["target_layer_ids"] = json!([1, 9, 17, 25, 33]); + c["markov_rank"] = json!(256); + c["enable_confidence_head"] = json!(true); + c["sliding_window"] = serde_json::Value::Null; + assert!(matches!( + hf_mapper_for_architectures(&c).map(|m| m.canonical_arch()), + Some("dspark") + )); + let cfg = DsparkHfMapper.config_from_hf(&c).unwrap(); + assert_eq!(cfg.extra_config["logits_start"], json!(0)); + assert_eq!(cfg.extra_config["markov_rank"], json!(256)); + assert_eq!(cfg.extra_config["confidence_head"], json!(true)); + assert_eq!(cfg.sliding_window_pattern, 0); + } + + #[test] + fn specforge_dspark_drafts_from_row_one_through_the_target_head() { + let mut c = base(); + c["architectures"] = json!(["DSparkDraftModel"]); + c["block_size"] = json!(7); + c["dspark_config"] = json!({"markov_rank": 256, "markov_head_type": "vanilla", "mask_token_id": 1, + "target_layer_ids": [5, 19, 33, 47, 61]}); + c["dflash_config"] = c["dspark_config"].clone(); + c["enable_confidence_head"] = json!(true); + c["rope_parameters"] = json!({"rope_theta": 1.0e7, "rope_type": "yarn", "factor": 32.0, + "beta_fast": 32.0, "beta_slow": 1.0, "original_max_position_embeddings": 8192}); + let cfg = DsparkHfMapper.config_from_hf(&c).unwrap(); + assert_eq!(cfg.extra_config["logits_start"], json!(1)); + assert_eq!(cfg.extra_config["drafter_head"], json!("target")); + assert_eq!(cfg.extra_config["markov_rank"], json!(256)); + assert_eq!(cfg.rope_scaling_type, "yarn"); + assert_eq!(cfg.rope_yarn_beta_fast, 32.0); + assert!(cfg.rope_yarn_truncate); + c["dspark_config"]["markov_head_type"] = json!("gated"); + assert!( + DsparkHfMapper.config_from_hf(&c).is_err(), + "gated Markov head is not implemented" + ); + } + + #[test] + fn gemma4_dspark_records_its_backbone() { + // deepseek-ai/dspark_gemma4_12b_block7: `Gemma4DSparkModel`, a + // gemma4_text decoder (global K=V layers, proportional rope) with + // DeepSpec's mixed rope_parameters (top-level keys beside the + // per-layer-type dicts). + let c: serde_json::Value = serde_json::from_str(r#"{"architectures": ["Gemma4DSparkModel"], "attention_bias": false, "attention_dropout": 0.0, "attention_k_eq_v": true, "block_size": 7, "bos_token_id": 2, "confidence_head_with_markov": true, "dtype": "bfloat16", "enable_confidence_head": true, "enable_moe_block": false, "eos_token_id": 1, "final_logit_softcapping": 30.0, "global_head_dim": 512, "head_dim": 256, "hidden_activation": "gelu_pytorch_tanh", "hidden_size": 3840, "hidden_size_per_layer_input": 0, "initializer_range": 0.02, "intermediate_size": 15360, "layer_types": ["full_attention", "full_attention", "full_attention", "full_attention", "full_attention"], "markov_head_type": "vanilla", "markov_rank": 256, "mask_token_id": 4, "max_position_embeddings": 262144, "model_type": "gemma4_text", "moe_intermediate_size": null, "num_anchors": 512, "num_attention_heads": 16, "num_experts": null, "num_global_key_value_heads": 1, "num_hidden_layers": 5, "num_key_value_heads": 8, "num_kv_shared_layers": 0, "num_target_layers": 48, "pad_token_id": 0, "rms_norm_eps": 1e-06, "rope_parameters": {"full_attention": {"partial_rotary_factor": 0.25, "rope_theta": 1000000.0, "rope_type": "proportional"}, "rope_theta": null, "rope_type": "default", "sliding_attention": {"rope_theta": 10000.0, "rope_type": "default"}}, "sliding_window": 1024, "target_layer_ids": [5, 17, 29, 41, 46], "target_model_type": "gemma4_unified", "target_text_model_type": "gemma4_unified_text", "tie_word_embeddings": false, "top_k_experts": null, "transformers_version": "5.10.2", "use_bidirectional_attention": "vision", "use_cache": true, "use_double_wide_mlp": false, "vocab_size": 262144, "vocab_size_per_layer_input": 262144}"#).unwrap(); + let m = hf_mapper_for_architectures(&c).expect("Gemma4DSparkModel maps to a drafter"); + let cfg = m.config_from_hf(&c).unwrap(); + assert_eq!(cfg.extra_config["drafter_backbone"], json!("gemma4")); + assert_eq!(cfg.extra_config["speculator_kind"], json!("dspark")); + assert_eq!(cfg.extra_config["logits_start"], json!(0)); + assert_eq!(cfg.extra_config["markov_rank"], json!(256)); + assert_eq!( + cfg.extra_config["target_layer_ids"], + json!([5, 17, 29, 41, 46]) + ); + assert!( + !cfg.extra_config.contains_key("drafter_head"), + "carries its own head" + ); + let mut q = c.clone(); + q["model_type"] = json!("gemma3"); + assert!( + m.config_from_hf(&q).is_err(), + "only llama / qwen3 / gemma4 backbones" + ); + } + + #[test] + fn speculators_dspark_keeps_its_own_head_and_drafts_from_the_anchor() { + // RedHatAI/Qwen3-4B-speculator.dspark: `DSparkDraftModel` like SpecForge, + // but its own embed / lm_head and sample_from_anchor = true (row 0). + let c = json!({ + "architectures": ["DSparkDraftModel"], "speculators_model_type": "dspark", + "aux_hidden_state_layer_ids": [1, 9, 17, 25, 33], "block_size": 7, "mask_token_id": 151669, + "draft_vocab_size": 151936, "sample_from_anchor": true, "sliding_window_non_causal": false, + "markov_rank": 256, "markov_head_type": "vanilla", "enable_confidence_head": true, + "confidence_head_with_markov": true, + "transformer_layer_config": { + "model_type": "qwen3", "hidden_size": 2560, "num_hidden_layers": 5, "num_attention_heads": 32, + "num_key_value_heads": 8, "head_dim": 128, "intermediate_size": 9728, "vocab_size": 151936, + "rms_norm_eps": 1e-6, "rope_parameters": {"rope_theta": 1000000, "rope_type": "default"}, + "layer_types": ["sliding_attention", "sliding_attention", "sliding_attention", + "sliding_attention", "sliding_attention"], + "sliding_window": 2048 + } + }); + let cfg = DsparkHfMapper.config_from_hf(&c).unwrap(); + assert_eq!(cfg.extra_config["logits_start"], json!(0)); + assert!( + !cfg.extra_config.contains_key("drafter_head"), + "its own lm_head, not the target's" + ); + assert_eq!(cfg.extra_config["markov_rank"], json!(256)); + assert_eq!(cfg.extra_config["confidence_head"], json!(true)); + assert_eq!( + cfg.extra_config["target_layer_ids"], + json!([0, 8, 16, 24, 32]) + ); + } + + #[test] + fn speculators_block_drafter_is_flattened() { + let mut c = json!({ + "architectures": ["DFlash2DraftModel"], "speculators_model_type": "dflash2", + "aux_hidden_state_layer_ids": [2, 10, 18, 26, 34], "block_size": 8, "mask_token_id": 151669, + "draft_vocab_size": 151936, "sample_from_anchor": false, "sliding_window_non_causal": true, + "conv_kernel_size": 2, "conv_group_size": 16, "selector_rank": 256, "selector_top_k": 16, + "transformer_layer_config": { + "model_type": "qwen3", "hidden_size": 4096, "num_hidden_layers": 5, "num_attention_heads": 32, + "num_key_value_heads": 8, "head_dim": 128, "intermediate_size": 12288, "vocab_size": 151936, + "rms_norm_eps": 1e-6, "rope_parameters": {"rope_theta": 1000000, "rope_type": "default"}, + "layer_types": ["sliding_attention", "sliding_attention", "sliding_attention", + "sliding_attention", "sliding_attention"], + "sliding_window": 2048, "use_sliding_window": true + } + }); + let cfg = DflashHfMapper.config_from_hf(&c).unwrap(); + assert_eq!( + cfg.extra_config["target_layer_ids"], + json!([1, 9, 17, 25, 33]) + ); + assert_eq!(cfg.extra_config["dflash_version"], json!(2)); + assert_eq!(cfg.extra_config["logits_start"], json!(1)); + assert_eq!( + cfg.extra_config["drafter_causal_layers"], + json!(vec![false; 5]) + ); + assert_eq!(cfg.rope_theta, 1.0e6); + // sliding_window_non_causal false: the sliding layers are causal. + c["sliding_window_non_causal"] = json!(false); + c["architectures"] = json!(["DFlashDraftModel"]); + let cfg = DflashHfMapper.config_from_hf(&c).unwrap(); + assert_eq!( + cfg.extra_config["drafter_causal_layers"], + json!(vec![true; 5]) + ); + // A reduced draft vocabulary rides through for a DFlash (1) head... + for k in [ + "conv_kernel_size", + "conv_group_size", + "selector_rank", + "selector_top_k", + ] { + c.as_object_mut().unwrap().remove(k); + } + c["speculators_model_type"] = json!("dflash"); + c["draft_vocab_size"] = json!(32000); + let cfg = DflashHfMapper.config_from_hf(&c).unwrap(); + assert_eq!(cfg.extra_config["draft_vocab_size"], json!(32000)); + assert_eq!( + cfg.vocab_size, 151936, + "the embedding keeps the full vocabulary" + ); + // ...and is refused where the full vocabulary indexes a head. + c["architectures"] = json!(["DFlash2DraftModel"]); + for (k, v) in [ + ("conv_kernel_size", 2), + ("conv_group_size", 16), + ("selector_rank", 256), + ("selector_top_k", 16), + ] { + c[k] = json!(v); + } + assert!( + DflashHfMapper.config_from_hf(&c).is_err(), + "DFlash 2 with a reduced vocabulary" + ); + c["draft_vocab_size"] = json!(200000); + assert!( + DflashHfMapper.config_from_hf(&c).is_err(), + "larger than the vocabulary" + ); + } + + fn layer_names(layers: u32, with_qk_norms: bool) -> Vec { + let mut v = vec!["fc.0.weight".to_string(), "hidden_norm.weight".to_string()]; + for l in 0..layers { + for t in ["q_proj", "k_proj", "v_proj", "o_proj"] { + v.push(format!("layers.{l}.self_attn.{t}.weight")); + } + if with_qk_norms { + v.push(format!("layers.{l}.self_attn.q_norm.weight")); + v.push(format!("layers.{l}.self_attn.k_norm.weight")); + } + } + v + } + + #[test] + fn llama_labelled_backbone_needs_qk_norms() { + // RedHatAI/Qwen3-30B-A3B-speculator.dflash: transformer_layer_config + // says `llama`, yet every layer ships self_attn.q_norm / k_norm. + let mut c = base(); + c["model_type"] = json!("llama"); + c["dflash_config"] = json!({"mask_token_id": 1, "target_layer_ids": [1, 12]}); + let cfg = DflashHfMapper.config_from_hf(&c).unwrap(); + let names = layer_names(cfg.num_hidden_layers, true); + check_block_drafter_qk_norms(&cfg, names.iter().map(String::as_str)) + .expect("Qwen-shaped llama-labelled layers are accepted"); + // A genuine Llama backbone (no q/k norms) would load as Qwen3 and fail + // on its first missing norm: refused at conversion instead. + let names = layer_names(cfg.num_hidden_layers, false); + let err = check_block_drafter_qk_norms(&cfg, names.iter().map(String::as_str)) + .unwrap_err() + .to_string(); + assert!(err.contains("layers.0.self_attn.q_norm.weight"), "{err}"); + assert!(err.contains("layers.0.self_attn.k_norm.weight"), "{err}"); + assert!(err.contains("and 6 more"), "{err}"); + assert!( + err.contains("norm-free Llama drafter backbone is not supported"), + "{err}" + ); + // One layer short is enough to refuse. + let names: Vec = layer_names(cfg.num_hidden_layers, true) + .into_iter() + .filter(|n| n != "layers.4.self_attn.k_norm.weight") + .collect(); + let err = check_block_drafter_qk_norms(&cfg, names.iter().map(String::as_str)) + .unwrap_err() + .to_string(); + assert!( + err.contains("missing layers.4.self_attn.k_norm.weight)"), + "{err}" + ); + } + + #[test] + fn fc_blocks_carry_the_norm_invariant_scale() { + assert_eq!( + fc_transform("fc.0.weight"), + Some(ValueTransform::DivBy(FC_SCALE_DIV)) + ); + assert_eq!( + fc_transform("fc.7.weight"), + Some(ValueTransform::DivBy(256)) + ); + assert!(fc_transform("fc.weight").is_none()); + assert!(fc_transform("hidden_norm.weight").is_none()); + assert!(fc_transform("layers.0.mlp.gate_proj.weight").is_none()); + } + + #[test] + fn non_qwen3_backbone_and_missing_taps_are_refused() { + let mut c = base(); + c["dflash_config"] = json!({"mask_token_id": 1}); + assert!( + DflashHfMapper.config_from_hf(&c).is_err(), + "no target_layer_ids" + ); + c["dflash_config"] = json!({"mask_token_id": 1, "target_layer_ids": [1]}); + c["model_type"] = json!("gemma3"); + assert!(DflashHfMapper.config_from_hf(&c).is_err(), "gemma backbone"); + c["model_type"] = json!("llama"); + assert!(DflashHfMapper.config_from_hf(&c).is_ok(), "llama backbone"); + assert!( + hf_mapper_for_architectures(&json!({"architectures": ["Qwen3ForCausalLM"]})).is_none() + ); + } +} + +// ── EAGLE-3 heads ──────────────────────────────────────────────────────────── +// +// An EAGLE-3 drafter (SpecForge / EAGLE `LlamaForCausalLMEagle3`) is ONE +// llama-shaped decoder layer whose attention input is the concatenation of +// the normed token embedding and the normed feature `fc(concat of three +// target hidden states)` (so q/k/v project from 2·dim), a residual over the +// feature, its own output norm and a REDUCED-vocabulary lm_head +// (`draft_vocab_size` rows) with a `d2t` offset table back to target ids. +// The embedding is the target's (a copy when the checkpoint ships one). +// Sidecar arch `eagle3`: `midlayer.*` → `layers.0.*`, the q/k/v projections +// split by input column into `*_proj_emb` / `*_proj_hid`, `fc` split per tap +// (no scale fold — its output is also the residual), `d2t` carried as f32, +// `t2d` dropped (derivable). + +pub struct Eagle3HfMapper; + +pub fn is_eagle3_arch(arch: &str) -> bool { + arch == "eagle3" +} + +/// RedHat's `speculators` packaging (`speculators_model_type: "eagle3"`, +/// `Eagle3Speculator`): the layer geometry nests under +/// `transformer_layer_config`, the EAGLE fields sit at the top level, and +/// `eagle_aux_hidden_state_layer_ids` are hidden_states indices (vLLM captures +/// "before layer i"), where SGLang / SpecForge configs name the layer whose +/// OUTPUT is tapped (SGLang adds one). Flattened here into the SpecForge +/// form the rest of the converter and the runtime read. +fn flatten_speculators_eagle3(c: &serde_json::Value) -> Result { + let mut flat = c + .get("transformer_layer_config") + .cloned() + .ok_or_else(|| anyhow!("speculators config without transformer_layer_config"))?; + let o = flat + .as_object_mut() + .ok_or_else(|| anyhow!("transformer_layer_config is not an object"))?; + for k in [ + "draft_vocab_size", + "norm_before_residual", + "norm_before_fc", + "fc_norm", + "target_hidden_size", + ] { + if let Some(v) = c.get(k).filter(|v| !v.is_null()) { + o.insert(k.to_string(), v.clone()); + } + } + // speculators' `norm_output` is Eagle 3.1's "feed the POST-norm hidden back + // across steps" — not the SpecForge flag of the same name (norm before the + // lm_head, which speculators always does). + if c.get("norm_output") + .and_then(|v| v.as_bool()) + .unwrap_or(false) + { + bail!("speculators EAGLE-3 with norm_output (Eagle 3.1 post-norm chaining) is not supported yet"); + } + if let Some(ids) = c + .get("eagle_aux_hidden_state_layer_ids") + .and_then(|v| v.as_array()) + { + let ids: Vec = ids.iter().filter_map(|x| x.as_i64()).collect(); + if ids.iter().any(|&i| i < 1) { + bail!("speculators aux layer ids {ids:?} name the embedding output; not supported"); + } + let out: Vec = ids.iter().map(|i| i - 1).collect(); + o.insert( + "eagle_aux_hidden_state_layer_ids".into(), + serde_json::json!(out), + ); + } + Ok(flat) +} + +fn eagle3_config_from_hf(c: &serde_json::Value) -> Result { + let flat; + let c = if c.get("speculators_model_type").and_then(|v| v.as_str()) == Some("eagle3") { + flat = flatten_speculators_eagle3(c)?; + &flat + } else { + c + }; + let backbone = c.get("model_type").and_then(|v| v.as_str()).unwrap_or(""); + if backbone != "llama" { + bail!("EAGLE-3 head model_type {backbone:?} is not llama — the engine runs the llama-shaped midlayer only"); + } + let mut cfg = crate::llama::hf_generic_config(c)?; + cfg.tie_word_embeddings = false; + if cfg.num_hidden_layers != 1 { + bail!( + "EAGLE-3 head has {} layers; one midlayer expected", + cfg.num_hidden_layers + ); + } + let draft_vocab = c + .get("draft_vocab_size") + .or_else(|| c.get("truncated_vocab_size")) + .and_then(|v| v.as_u64()) + .ok_or_else(|| anyhow!("config.json: draft_vocab_size missing (the reduced lm_head)"))?; + if c.get("fc_norm").and_then(|v| v.as_bool()).unwrap_or(false) { + bail!("EAGLE-3 head with fc_norm (per-tap norms before fc) is not supported yet"); + } + // norm_before_fc: one RMSNorm over the concatenated taps before `fc` + // (RedHat's gpt-oss heads, weight `input_norm` [n_taps * hidden]); + // norm_before_residual: the attention residual is the NORMED feature. + // Both recorded for the runtime. + let norm_before_fc = c + .get("norm_before_fc") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + let norm_before_residual = c + .get("norm_before_residual") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + // norm_output=false (no final norm before the lm_head): the runtime + // always applies output_norm and the loader requires it, so such a head + // would either fail to load or draft from mismatched logits. + if !c + .get("norm_output") + .and_then(|v| v.as_bool()) + .unwrap_or(true) + { + bail!("EAGLE-3 head with norm_output=false is not supported yet (the runtime always normalizes before the lm_head)"); + } + let target_hidden = c + .get("target_hidden_size") + .and_then(|v| v.as_u64()) + .unwrap_or(cfg.hidden_size as u64); + if target_hidden != cfg.hidden_size as u64 { + bail!( + "EAGLE-3 head hidden {} != target hidden {}: cross-width heads are not supported", + cfg.hidden_size, + target_hidden + ); + } + use serde_json::json; + let x = &mut cfg.extra_config; + x.insert("speculator_kind".into(), json!("eagle3")); + x.insert("draft_vocab_size".into(), json!(draft_vocab)); + if norm_before_residual { + x.insert("norm_before_residual".into(), json!(true)); + } + if norm_before_fc { + x.insert("norm_before_fc".into(), json!(true)); + } + x.insert( + "norm_output".into(), + json!(c + .get("norm_output") + .and_then(|v| v.as_bool()) + .unwrap_or(true)), + ); + // Explicit aux layers (layer ids whose OUTPUT is tapped); absent → the + // runtime's default [1, L/2 - 1, L - 4] on the target's layer count. + if let Some(ids) = c + .get("eagle_config") + .and_then(|e| e.get("eagle_aux_hidden_state_layer_ids")) + .or_else(|| c.get("eagle_aux_hidden_state_layer_ids")) + .and_then(|v| v.as_array()) + { + let ids: Vec = ids.iter().filter_map(|x| x.as_u64()).collect(); + if ids.len() != 3 { + bail!("EAGLE-3 aux layer ids must name 3 layers, got {ids:?}"); + } + x.insert("aux_layer_ids".into(), json!(ids)); + } + Ok(cfg) +} + +impl HfMapper for Eagle3HfMapper { + fn canonical_arch(&self) -> &'static str { + "eagle3" + } + fn config_from_hf(&self, c: &serde_json::Value) -> Result { + eagle3_config_from_hf(c) + } +} + +/// Pick the EAGLE-3 mapper from `architectures` (`LlamaForCausalLMEagle3`, +/// `Eagle3LlamaForCausalLM`). +pub fn hf_mapper_for_eagle3(config: &serde_json::Value) -> Option<&'static dyn HfMapper> { + let archs = config.get("architectures")?.as_array()?; + let named = archs + .iter() + .filter_map(|v| v.as_str()) + .any(|a| a.contains("Eagle3") || a.contains("EAGLE3")); + // nebius' EAGLE-3 heads (nebius/EAGLE3-gpt-oss-20b) say only + // `LlamaForCausalLM`: a ONE-layer llama with a reduced draft vocabulary + // is an EAGLE-3 head, never a language model. + let unnamed_head = config.get("num_hidden_layers").and_then(|v| v.as_u64()) == Some(1) + && (config.get("draft_vocab_size").is_some() + || config.get("truncated_vocab_size").is_some()); + (named || unnamed_head).then_some(&Eagle3HfMapper) +} + +#[cfg(test)] +mod eagle3_tests { + use super::*; + use serde_json::json; + + #[test] + fn eagle3_config_reads_draft_vocab_and_aux_layers() { + let mut c = json!({ + "architectures": ["Eagle3LlamaForCausalLM"], "model_type": "llama", "hidden_size": 2560, + "num_hidden_layers": 1, "num_attention_heads": 32, "num_key_value_heads": 8, "head_dim": 128, + "intermediate_size": 9728, "vocab_size": 151936, "rms_norm_eps": 1e-6, "rope_theta": 1000000.0, + "draft_vocab_size": 32000, "tie_word_embeddings": false + }); + assert!(matches!( + hf_mapper_for_eagle3(&c).map(|m| m.canonical_arch()), + Some("eagle3") + )); + let cfg = Eagle3HfMapper.config_from_hf(&c).unwrap(); + assert_eq!(cfg.extra_config["draft_vocab_size"], json!(32000)); + assert!(!cfg.extra_config.contains_key("aux_layer_ids")); + c["eagle_config"] = json!({"eagle_aux_hidden_state_layer_ids": [1, 15, 28]}); + let cfg = Eagle3HfMapper.config_from_hf(&c).unwrap(); + assert_eq!(cfg.extra_config["aux_layer_ids"], json!([1, 15, 28])); + c["num_hidden_layers"] = json!(2); + assert!( + Eagle3HfMapper.config_from_hf(&c).is_err(), + "one midlayer only" + ); + assert!(hf_mapper_for_eagle3(&json!({"architectures": ["LlamaForCausalLM"]})).is_none()); + } + + #[test] + fn eagle3_norm_before_fc_is_recorded() { + // RedHatAI/gpt-oss-20b-speculator.eagle3: one RMSNorm over the + // concatenated taps before fc (weight `input_norm` [3H]). + let c = json!({ + "architectures": ["Eagle3DraftModel"], "speculators_model_type": "eagle3", + "draft_vocab_size": 64000, "norm_before_fc": true, "norm_before_residual": true, + "transformer_layer_config": { + "model_type": "llama", "hidden_size": 2880, "num_hidden_layers": 1, "num_attention_heads": 64, + "num_key_value_heads": 8, "head_dim": 64, "intermediate_size": 2880, "vocab_size": 201088, + "rms_norm_eps": 1e-5, "rope_theta": 10000.0 + } + }); + let cfg = Eagle3HfMapper.config_from_hf(&c).unwrap(); + assert_eq!(cfg.extra_config["norm_before_fc"], json!(true)); + assert_eq!(cfg.extra_config["norm_before_residual"], json!(true)); + } + + #[test] + fn unnamed_one_layer_head_is_eagle3() { + // nebius/EAGLE3-gpt-oss-20b: architectures says only LlamaForCausalLM, + // the reduced vocabulary is `truncated_vocab_size`, the rope is YaRN. + let c = json!({ + "architectures": ["LlamaForCausalLM"], "model_type": "llama", "num_hidden_layers": 1, + "truncated_vocab_size": 64000, "draft_vocab_size": 64000, "hidden_size": 2880, + "num_attention_heads": 64, "num_key_value_heads": 8, "head_dim": 64, "intermediate_size": 11520, + "vocab_size": 201088, "rms_norm_eps": 1e-5, "rope_theta": 150000, + "rope_scaling": {"beta_fast": 32.0, "beta_slow": 1.0, "factor": 32.0, + "original_max_position_embeddings": 4096, "rope_type": "yarn", "truncate": false} + }); + assert!(hf_mapper_for_eagle3(&c).is_some()); + let cfg = Eagle3HfMapper.config_from_hf(&c).unwrap(); + assert_eq!(cfg.extra_config["draft_vocab_size"], json!(64000)); + assert_eq!(cfg.rope_scaling_type, "yarn"); + // A real one-layer LM with no draft vocabulary is not a head. + let mut lm = c.clone(); + lm.as_object_mut().unwrap().remove("truncated_vocab_size"); + lm.as_object_mut().unwrap().remove("draft_vocab_size"); + assert!(hf_mapper_for_eagle3(&lm).is_none()); + } + + #[test] + fn speculators_eagle3_is_flattened_with_hidden_state_taps() { + let c = json!({ + "architectures": ["Eagle3Speculator"], "speculators_model_type": "eagle3", + "draft_vocab_size": 32000, "norm_before_residual": true, + "eagle_aux_hidden_state_layer_ids": [2, 18, 33], + "transformer_layer_config": { + "model_type": "llama", "hidden_size": 4096, "num_hidden_layers": 1, + "num_attention_heads": 32, "num_key_value_heads": 8, "head_dim": 128, + "intermediate_size": 12288, "vocab_size": 151936, "rms_norm_eps": 1e-6, + "rope_theta": 1000000 + } + }); + assert!(hf_mapper_for_eagle3(&c).is_some()); + let cfg = Eagle3HfMapper.config_from_hf(&c).unwrap(); + assert_eq!(cfg.hidden_size, 4096); + assert_eq!(cfg.extra_config["draft_vocab_size"], json!(32000)); + assert_eq!(cfg.extra_config["norm_before_residual"], json!(true)); + // hidden_states indices [2, 18, 33] -> the loader's output-layer form. + assert_eq!(cfg.extra_config["aux_layer_ids"], json!([1, 17, 32])); + let mut bad = c.clone(); + bad["norm_output"] = json!(true); + assert!( + Eagle3HfMapper.config_from_hf(&bad).is_err(), + "Eagle 3.1 post-norm chaining" + ); + } +} diff --git a/base-convert/crates/base-arch/src/gemma.rs b/base-convert/crates/base-arch/src/gemma.rs index fae169e..4dfe7f6 100644 --- a/base-convert/crates/base-arch/src/gemma.rs +++ b/base-convert/crates/base-arch/src/gemma.rs @@ -41,9 +41,11 @@ impl crate::HfMapper for Gemma3HfMapper { fn config_from_hf(&self, c: &serde_json::Value) -> Result { // Gemma HF often nests the text config under "text_config". Prefer - // that sub-object when present. + // that sub-object when present — but token ids stay at the wrapper + // level, so backfill them from the full config. let source = c.get("text_config").unwrap_or(c); let mut config = crate::llama::hf_generic_config(source)?; + crate::backfill_wrapper_token_ids(&mut config, c); config.tie_word_embeddings = source .get("tie_word_embeddings") .and_then(|v| v.as_bool()) @@ -92,9 +94,11 @@ impl crate::HfMapper for Gemma4HfMapper { fn config_from_hf(&self, c: &serde_json::Value) -> Result { // Gemma-4 multimodal configs nest text under "text_config". For // the primary model weights we use that sub-config; audio/vision - // tower configs are siblings handled by the mmproj path. + // tower configs are siblings handled by the mmproj path. Token ids + // stay at the wrapper level — backfill from the full config. let source = c.get("text_config").unwrap_or(c); let mut config = crate::llama::hf_generic_config(source)?; + crate::backfill_wrapper_token_ids(&mut config, c); config.tie_word_embeddings = source .get("tie_word_embeddings") .and_then(|v| v.as_bool()) diff --git a/base-convert/crates/base-arch/src/lib.rs b/base-convert/crates/base-arch/src/lib.rs index d5fafa1..2807b2c 100644 --- a/base-convert/crates/base-arch/src/lib.rs +++ b/base-convert/crates/base-arch/src/lib.rs @@ -7,6 +7,7 @@ //! consistently regardless of source format. pub mod bert; +pub mod dflash; pub mod gemma; pub mod glm; pub mod gpt_oss; @@ -34,6 +35,43 @@ pub struct LayerNames { } /// Dispatch table: which arch module to use for a given GGUF arch string. +/// Token ids (bos/eos) live at the WRAPPER level of multimodal HF configs — +/// `text_config` carries none — so a mapper that reads only the text +/// sub-config silently drops them. gemma-3-4b lost its BOS this way (a +/// BOS-mandatory family: scored without BOS it degrades to ~1b quality, +/// which is how the bug was found). Call after `hf_generic_config(text)` +/// with the FULL config to backfill from the wrapper level. +pub fn backfill_wrapper_token_ids(config: &mut ArchConfig, full: &serde_json::Value) { + let ids = |k: &str| -> Vec { + match full.get(k) { + Some(v) if v.is_u64() => vec![v.as_u64().unwrap() as u32], + Some(v) => v + .as_array() + .map(|a| { + a.iter() + .filter_map(|x| x.as_u64().map(|n| n as u32)) + .collect() + }) + .unwrap_or_default(), + None => Vec::new(), + } + }; + if config.bos_token_id == 0 { + if let Some(&b) = ids("bos_token_id").first() { + config.bos_token_id = b; + } + } + if config.eos_token_id == 0 { + let eos = ids("eos_token_id"); + if let Some(&e) = eos.first() { + config.eos_token_id = e; + if eos.len() > 1 && config.eos_token_ids.is_empty() { + config.eos_token_ids = eos[1..].to_vec(); + } + } + } +} + pub fn source_mapper_for_gguf(arch: &str) -> Option<&'static dyn GgufMapper> { match arch { "llama" => Some(&llama::LlamaMapper), @@ -68,6 +106,12 @@ pub enum ValueTransform { /// code; the value the scan kernel wants is `A`, which must be /// negative for the recurrence to decay. NegExp, + /// `x -> x / n`. Block drafters fold 1/256 into the context projection + /// `fc`: its output only ever feeds `hidden_norm`, which is invariant to + /// a uniform scale (eps aside), and the un-scaled output overflows f16 — + /// the target's residual stream carries massive activations (|x| up to + /// ~5e3 on Qwen3-4B) and the fc sum over five taps reaches ~8e4. + DivBy(u32), } impl ValueTransform { @@ -78,6 +122,12 @@ impl ValueTransform { *v = -v.exp(); } } + ValueTransform::DivBy(n) => { + let s = 1.0f32 / n as f32; + for v in values.iter_mut() { + *v *= s; + } + } } } } @@ -96,6 +146,18 @@ pub trait HfMapper: Sync { 0.0 } + /// [`norm_shift`](Self::norm_shift) for an MLX (mlx-lm) source, whose + /// stored norm values need not follow the HF checkpoint's convention. + /// mlx-lm's own converters may bake a family's `(1 + weight)` offset + /// into the stored tensor so its runtime uses the plain `nn.RMSNorm`; + /// an HF-shaped +1 on top of that doubles it, and the bundle decodes + /// garbage while `--validate` (which mirrors the converter) passes. + /// Default: the HF rule, for families whose MLX export keeps the HF + /// values. + fn mlx_norm_shift(&self, canonical: &str) -> f32 { + self.norm_shift(canonical) + } + /// RoPE row-permutation head count for a canonical tensor at HF→.base /// conversion time, or None when the tensor needs no permutation. /// @@ -171,8 +233,12 @@ pub fn hf_mapper_for_model_type(model_type: &str) -> Option<&'static dyn HfMappe // (tokenizer_defaults.h). REVALIDATED on microsoft/Phi-3-mini-4k-instruct // (converted Q8): chat output is coherent across probes, no flood. // Phi-3.5 stays OUT — it needs LongRoPE, which the engine (linear scaling - // only) does not implement. - "llama" | "mistral" | "phi3" => Some(&llama::LlamaHfMapper), + // only) does not implement. HF ships Phi-3.5 under the same broad + // `phi3` model_type, so phi3 routes through a guarded wrapper that + // fails the conversion when `rope_scaling` (LongRoPE) is present + // instead of silently converting with incorrect rotary scaling. + "llama" | "mistral" => Some(&llama::LlamaHfMapper), + "phi3" => Some(&llama::Phi3HfMapper), "qwen2" | "qwen3" => Some(&qwen::QwenHfMapper), "qwen2_moe" | "qwen3_moe" => Some(&qwen::QwenMoeHfMapper), // Qwen3.5 / 3.6: hybrid Gated-DeltaNet + full-attention decoder @@ -326,8 +392,71 @@ pub trait GgufMapper: Sync { } } +/// Generation defaults the MODEL publishes, from the checkpoint's +/// `generation_config.json` — the file HF model authors use to say how their +/// model should be sampled. Qwen3.8-27B ships `top_k: 20, top_p: 0.95`; +/// Qwen2.5-0.5B-Instruct ships `temperature: 0.7, top_p: 0.8, top_k: 20`. They +/// differ per model, which is exactly why a serving runtime cannot express +/// this as one global constant. +/// +/// `None` means the checkpoint made no claim, and the runtime keeps its own +/// default. Every field is validated on the way in (see `from_generation_config`): +/// a checkpoint is free to publish nonsense and we would rather ignore a field +/// than bake a bad value into a bundle that outlives the conversion. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct SamplingDefaults { + pub temperature: Option, + pub top_p: Option, + pub top_k: Option, + pub min_p: Option, + /// Recorded for visibility but DELIBERATELY NOT applied as a default: a + /// repetition penalty damages dialect structure (baseRT PR #417 — the + /// model swerves to unpenalized near-synonyms and tool calls leak as + /// visible text), and several checkpoints publish 1.1. A caller that wants + /// one asks for it explicitly. + pub repetition_penalty: Option, + /// `do_sample: false` means the author wants greedy decoding. + pub do_sample: Option, +} + +impl SamplingDefaults { + /// Parse an HF `generation_config.json` body. Out-of-range values are + /// dropped field-by-field rather than failing the conversion — a bad + /// `temperature` should not cost us a good `top_k`. + pub fn from_generation_config(gc: &serde_json::Value) -> Self { + let f32_in = |key: &str, lo: f32, hi: f32| -> Option { + let v = gc.get(key)?.as_f64()? as f32; + (v.is_finite() && v >= lo && v <= hi).then_some(v) + }; + Self { + // 0 is legal and means greedy; above ~2 is a sampling experiment, + // not a recommendation. + temperature: f32_in("temperature", 0.0, 4.0), + // 0 would admit nothing; 1.0 means "no truncation" and is a claim + // worth honouring as much as 0.95 is. + top_p: f32_in("top_p", 0.001, 1.0), + top_k: gc + .get("top_k") + .and_then(|v| v.as_u64()) + .filter(|k| *k <= 1_000_000) + .map(|k| k as u32), + min_p: f32_in("min_p", 0.0, 1.0), + repetition_penalty: f32_in("repetition_penalty", 0.5, 2.0), + do_sample: gc.get("do_sample").and_then(|v| v.as_bool()), + } + } + + pub fn is_empty(&self) -> bool { + *self == Self::default() + } +} + #[derive(Debug, Clone, Default, PartialEq)] pub struct ArchConfig { + /// Open-namespace header keys a mapper adds beyond the typed fields + /// (block drafters: `speculator_kind`, `block_size`, `mask_token_id`, + /// `target_layer_ids`, …). Merged into the header config verbatim. + pub extra_config: std::collections::BTreeMap, pub hidden_size: u32, pub num_hidden_layers: u32, pub num_attention_heads: u32, @@ -386,6 +515,10 @@ pub struct ArchConfig { /// `Tokenizer::add_eos_id` so generation honors any of them. pub eos_token_ids: Vec, + /// Sampling defaults published by the checkpoint (see `SamplingDefaults`). + /// Empty when it published none. + pub sampling_defaults: SamplingDefaults, + // ── Gemma-4-specific fields (zero/empty for other archs) ───────── /// Explicit attention scale used at Q·K^T (Gemma 4 uses 1.0 instead /// of the standard 1/sqrt(head_dim)). 0.0 = derive 1/sqrt(head_dim) @@ -617,6 +750,29 @@ impl ArchConfig { pub fn to_config_map(&self) -> std::collections::BTreeMap { use serde_json::json; let mut m = std::collections::BTreeMap::new(); + // Model-published sampling defaults. Emitted only when the checkpoint + // actually claimed one, so "absent" stays distinguishable from "the + // author asked for this value" — the runtime needs that difference to + // know when its own default applies. + let sd = &self.sampling_defaults; + if let Some(v) = sd.temperature { + m.insert("default_temperature".into(), json!(v)); + } + if let Some(v) = sd.top_p { + m.insert("default_top_p".into(), json!(v)); + } + if let Some(v) = sd.top_k { + m.insert("default_top_k".into(), json!(v)); + } + if let Some(v) = sd.min_p { + m.insert("default_min_p".into(), json!(v)); + } + if let Some(v) = sd.repetition_penalty { + m.insert("default_repetition_penalty".into(), json!(v)); + } + if let Some(v) = sd.do_sample { + m.insert("default_do_sample".into(), json!(v)); + } m.insert("hidden_size".into(), json!(self.hidden_size)); m.insert("num_hidden_layers".into(), json!(self.num_hidden_layers)); m.insert( @@ -938,14 +1094,94 @@ impl ArchConfig { json!(self.index_share_for_mtp_iteration), ); } + for (k, v) in &self.extra_config { + m.insert(k.clone(), v.clone()); + } m } } #[cfg(test)] mod tests { + + /// Real published values, verbatim from the checkpoints. Qwen3.8-27B is + /// the model whose full-tail sampling produced runaway repeats. + #[test] + fn parses_real_generation_configs() { + let qwen38 = serde_json::json!({ + "bos_token_id": 248044, "do_sample": true, + "eos_token_id": [248046, 248044], "pad_token_id": 248044, + "temperature": 1.0, "top_k": 20, "top_p": 0.95, + "min_p": 0, "repetition_penalty": 1.0 + }); + let sd = SamplingDefaults::from_generation_config(&qwen38); + assert_eq!(sd.top_k, Some(20), "Qwen3.8 asks for top_k 20"); + assert_eq!(sd.top_p, Some(0.95)); + assert_eq!(sd.temperature, Some(1.0)); + assert_eq!(sd.min_p, Some(0.0)); + assert_eq!(sd.do_sample, Some(true)); + + let qwen25 = serde_json::json!({ + "do_sample": true, "repetition_penalty": 1.1, + "temperature": 0.7, "top_p": 0.8, "top_k": 20 + }); + let sd = SamplingDefaults::from_generation_config(&qwen25); + assert_eq!(sd.temperature, Some(0.7), "defaults differ per model"); + assert_eq!(sd.top_p, Some(0.8)); + assert_eq!(sd.repetition_penalty, Some(1.1)); + } + + /// A checkpoint free-texting nonsense must not poison a bundle that + /// outlives the conversion — drop the bad field, keep the good ones. + #[test] + fn out_of_range_fields_are_dropped_individually() { + let junk = serde_json::json!({ + "temperature": 99.0, "top_p": 5.0, "top_k": 20, "min_p": -1.0 + }); + let sd = SamplingDefaults::from_generation_config(&junk); + assert_eq!(sd.temperature, None, "99.0 is not a recommendation"); + assert_eq!(sd.top_p, None, "top_p > 1 is meaningless"); + assert_eq!(sd.min_p, None); + assert_eq!(sd.top_k, Some(20), "the sane field survives"); + } + + /// No generation_config keys at all = no claim, and the runtime keeps its + /// own default. "Absent" must stay distinguishable from "author chose". + #[test] + fn absent_keys_make_no_claim() { + let sd = SamplingDefaults::from_generation_config(&serde_json::json!({"eos_token_id": 7})); + assert!(sd.is_empty()); + assert!(!ArchConfig::default() + .to_config_map() + .contains_key("default_top_k")); + } + + /// Only claimed fields reach the bundle header. + #[test] + fn only_claimed_defaults_are_emitted() { + let mut c = ArchConfig::default(); + c.sampling_defaults.top_k = Some(20); + let m = c.to_config_map(); + assert_eq!(m.get("default_top_k"), Some(&serde_json::json!(20))); + assert!(!m.contains_key("default_top_p"), "unclaimed stays absent"); + } use super::*; + /// Explicitly unsupported architectures stay fail-closed at the dispatch + /// layer (capability-audit): Gemma 3n's AltUp/Laurel/per-layer-FFN stack + /// has no engine execution path, so neither its model_type nor any + /// variant spelling may resolve to a mapper. + #[test] + fn gemma3n_has_no_mapper() { + for t in ["gemma3n", "gemma3n_text", "gemma-3n"] { + assert!( + hf_mapper_for_model_type(t).is_none(), + "{t} must not convert" + ); + assert!(!SUPPORTED_HF_MODEL_TYPES.contains(&t)); + } + } + /// `SUPPORTED_HF_MODEL_TYPES` must list exactly what the dispatch match /// accepts — every advertised type resolves, and nothing else creeps in. #[test] diff --git a/base-convert/crates/base-arch/src/llama.rs b/base-convert/crates/base-arch/src/llama.rs index eb7c229..b331334 100644 --- a/base-convert/crates/base-arch/src/llama.rs +++ b/base-convert/crates/base-arch/src/llama.rs @@ -24,6 +24,7 @@ use std::collections::BTreeMap; pub struct LlamaMapper; pub struct LlamaHfMapper; +pub struct Phi3HfMapper; impl crate::HfMapper for LlamaHfMapper { fn canonical_arch(&self) -> &'static str { @@ -44,6 +45,40 @@ impl crate::HfMapper for LlamaHfMapper { } } +/// Phi-3 through the Llama mapper, minus the variants the engine cannot +/// run. HF ships every Phi-3 generation under the same `model_type: +/// "phi3"`, but only the standard-RoPE checkpoints (Phi-3-mini-4k) are +/// Llama-shaped at runtime: the 128k-context and all Phi-3.5 checkpoints +/// carry `rope_scaling.type = "longrope"` (per-frequency short/long factor +/// arrays), which the engine does not implement. Converting one anyway +/// would produce a bundle that loads and generates with the wrong rotary +/// scaling — so any `rope_scaling` on a phi3 config fails the conversion +/// instead of falling through to the generic parse. +impl crate::HfMapper for Phi3HfMapper { + fn canonical_arch(&self) -> &'static str { + LlamaHfMapper.canonical_arch() + } + fn config_from_hf(&self, c: &serde_json::Value) -> Result { + if let Some(rs) = c.get("rope_scaling").filter(|v| !v.is_null()) { + let ty = rs + .get("rope_type") + .or_else(|| rs.get("type")) + .and_then(|v| v.as_str()) + .unwrap_or(""); + anyhow::bail!( + "phi3 checkpoint declares rope_scaling type {ty:?}: this is a Phi-3.5 / \ + long-context variant (LongRoPE), which the engine does not implement — \ + converting it would run with incorrect rotary scaling. Only standard-RoPE \ + Phi-3 checkpoints (e.g. microsoft/Phi-3-mini-4k-instruct) are supported." + ); + } + LlamaHfMapper.config_from_hf(c) + } + fn rope_permute_heads(&self, canonical: &str, cfg: &crate::ArchConfig) -> Option { + LlamaHfMapper.rope_permute_heads(canonical, cfg) + } +} + pub(crate) fn hf_generic_config(c: &serde_json::Value) -> Result { let u32_key = |k: &str| { c.get(k) @@ -78,8 +113,31 @@ pub(crate) fn hf_generic_config(c: &serde_json::Value) -> Result Result 0.0) + .unwrap_or(32.0) + } else { + 0.0 + }; + let rope_yarn_beta_slow = if is_yarn { + Some(rs_f32("beta_slow")) + .filter(|v| *v > 0.0) + .unwrap_or(1.0) + } else { + 0.0 + }; + let rope_yarn_truncate = is_yarn + && rope_scaling + .and_then(|v| v.get("truncate")) + .and_then(|v| v.as_bool()) + .unwrap_or(true); let rms_norm_eps = f32_key("rms_norm_eps").unwrap_or(1e-6); let tie_word_embeddings = bool_key("tie_word_embeddings").unwrap_or(false); @@ -153,6 +234,9 @@ pub(crate) fn hf_generic_config(c: &serde_json::Value) -> Result` lands in `eos_token_id`; `<|eom_id|>` and /// `<|eot_id|>` go into `eos_token_ids` for runtime multi-EOS. @@ -382,6 +493,65 @@ mod tests { assert_eq!(c.eos_token_ids, vec![128008u32, 128009u32]); } + /// Phi-3.5 (and Phi-3-*-128k) share `model_type: "phi3"` with the + /// supported 4k checkpoints but declare LongRoPE via `rope_scaling`. + /// The guarded mapper must reject them at conversion time — the + /// generic parse would otherwise emit a bundle that generates with + /// incorrect rotary scaling. + #[test] + fn phi35_longrope_fails_closed() { + let cfg = serde_json::json!({ + "hidden_size": 3072, + "num_hidden_layers": 32, + "num_attention_heads": 32, + "num_key_value_heads": 32, + "intermediate_size": 8192, + "vocab_size": 32064, + "rope_theta": 10_000.0, + "rms_norm_eps": 1e-5, + "bos_token_id": 1, + "eos_token_id": 32000, + "rope_scaling": { + "type": "longrope", + "short_factor": [1.0, 1.03, 1.05], + "long_factor": [1.0, 1.5, 2.0], + }, + }); + let err = crate::HfMapper::config_from_hf(&Phi3HfMapper, &cfg).unwrap_err(); + let msg = format!("{err}"); + assert!(msg.contains("longrope"), "error names the variant: {msg}"); + // Older configs spell the key `type`; current HF uses `rope_type`. + let mut cfg2 = cfg.clone(); + cfg2["rope_scaling"] = serde_json::json!({"rope_type": "longrope", "factor": 4.0}); + assert!(crate::HfMapper::config_from_hf(&Phi3HfMapper, &cfg2).is_err()); + // Even an unnamed scaling block fails closed rather than parsing. + cfg2["rope_scaling"] = serde_json::json!({"factor": 4.0}); + assert!(crate::HfMapper::config_from_hf(&Phi3HfMapper, &cfg2).is_err()); + } + + /// Phi-3-mini-4k (`rope_scaling: null`) still converts through the + /// guarded mapper exactly as it did through the plain Llama mapper. + #[test] + fn phi3_mini_4k_still_converts() { + let cfg = serde_json::json!({ + "hidden_size": 3072, + "num_hidden_layers": 32, + "num_attention_heads": 32, + "num_key_value_heads": 32, + "intermediate_size": 8192, + "vocab_size": 32064, + "rope_theta": 10_000.0, + "rms_norm_eps": 1e-5, + "bos_token_id": 1, + "eos_token_id": 32000, + "rope_scaling": serde_json::Value::Null, + }); + let c = crate::HfMapper::config_from_hf(&Phi3HfMapper, &cfg).unwrap(); + assert_eq!(c.hidden_size, 3072); + assert_eq!(c.eos_token_id, 32000); + assert_eq!(c.rope_scaling_type, "", "no scaling recorded"); + } + #[test] fn eos_token_id_scalar_form_unchanged() { let cfg = serde_json::json!({ diff --git a/base-convert/crates/base-arch/src/qwen.rs b/base-convert/crates/base-arch/src/qwen.rs index 0679c1d..4fb4dea 100644 --- a/base-convert/crates/base-arch/src/qwen.rs +++ b/base-convert/crates/base-arch/src/qwen.rs @@ -34,6 +34,10 @@ fn qwen35_config_from_hf(c: &serde_json::Value) -> Result { // nested object when present. let tc = c.get("text_config").unwrap_or(c); let mut config = crate::llama::hf_generic_config(tc)?; + // Multimodal wrappers keep bos/eos at the wrapper level (see + // backfill_wrapper_token_ids — this is the bug that silently dropped + // gemma-3-4b's BOS). + crate::backfill_wrapper_token_ids(&mut config, c); let u32_v = |v: &serde_json::Value, k: &str| v.get(k).and_then(|x| x.as_u64()).map(|n| n as u32); @@ -88,13 +92,31 @@ fn qwen35_config_from_hf(c: &serde_json::Value) -> Result { /// (`linear_attn.norm.weight`), which is a Qwen3NextRMSNormGated and uses plain /// `weight` with no offset. fn qwen35_norm_shift(canonical: &str) -> f32 { - if canonical.ends_with("norm.weight") && !canonical.ends_with(".linear_attn.norm.weight") { + // The MTP head's two pre-fc norms are Qwen3NextRMSNorm too (zero-centered + // gamma, HF values near 0) but their names end in `_embedding` / + // `_hidden`, not `norm.weight`. + let mtp_pre_fc = canonical.ends_with("pre_fc_norm_embedding.weight") + || canonical.ends_with("pre_fc_norm_hidden.weight"); + if mtp_pre_fc + || (canonical.ends_with("norm.weight") && !canonical.ends_with(".linear_attn.norm.weight")) + { 1.0 } else { 0.0 } } +/// MLX sources of Qwen3.5 / 3.6 / 3.8 carry NO norm shift: mlx-lm's +/// `qwen3_next` converter already stores every Qwen3NextRMSNorm as +/// `1 + weight` so its runtime can use the plain `nn.RMSNorm` (measured on +/// `mlx-community/Qwen3.8-27B-4bit`: `input_layernorm` mean 0.97, +/// `q_norm` 1.36, `model.norm` 1.94 — and `linear_attn.norm`, plain in HF +/// too, 0.87). Those are exactly the values the `.base` runtime wants, so +/// the HF +1 above must not be applied again (BAS-959: it was, every norm +/// sat near 2, the bundle decoded garbage and the MTP head accepted +/// nothing). +const QWEN35_MLX_NORM_SHIFT: f32 = 0.0; + impl crate::HfMapper for Qwen35HfMapper { fn canonical_arch(&self) -> &'static str { "qwen35" @@ -105,6 +127,9 @@ impl crate::HfMapper for Qwen35HfMapper { fn norm_shift(&self, canonical: &str) -> f32 { qwen35_norm_shift(canonical) } + fn mlx_norm_shift(&self, _canonical: &str) -> f32 { + QWEN35_MLX_NORM_SHIFT + } } impl crate::HfMapper for Qwen35MoeHfMapper { @@ -114,6 +139,9 @@ impl crate::HfMapper for Qwen35MoeHfMapper { fn norm_shift(&self, canonical: &str) -> f32 { qwen35_norm_shift(canonical) } + fn mlx_norm_shift(&self, _canonical: &str) -> f32 { + QWEN35_MLX_NORM_SHIFT + } fn config_from_hf(&self, c: &serde_json::Value) -> Result { // The real Qwen3.5/3.6-MoE text_config carries NO `intermediate_size` // (the HF config class deletes the attribute) — the dense-FFN slot is @@ -464,6 +492,11 @@ mod tests { "layers.3.self_attn.q_norm.weight", "layers.3.self_attn.k_norm.weight", "final_norm.weight", + // The MTP head's norms, under the speculator prefix. + "mtp.layers.0.input_norm.weight", + "mtp.norm.weight", + "mtp.pre_fc_norm_embedding.weight", + "mtp.pre_fc_norm_hidden.weight", ] { assert_eq!(qwen35_norm_shift(s), 1.0, "should shift: {s}"); } @@ -473,11 +506,48 @@ mod tests { "layers.0.linear_attn.in_proj_qkv.weight", "layers.0.mlp.gate_proj.weight", "embed_tokens.weight", + "mtp.fc_embed.weight", + "mtp.fc_hidden.weight", ] { assert_eq!(qwen35_norm_shift(s), 0.0, "should not shift: {s}"); } } + /// An MLX source already stores `1 + weight` (mlx-lm bakes the offset), + /// so the MLX rule is "never shift" for every tensor the HF rule shifts + /// — dense and MoE mappers alike (BAS-959). + #[test] + fn qwen35_mlx_sources_are_never_shifted() { + use crate::HfMapper; + let shifted_by_hf = [ + "layers.0.input_norm.weight", + "layers.3.post_attn_norm.weight", + "layers.3.self_attn.q_norm.weight", + "layers.3.self_attn.k_norm.weight", + "final_norm.weight", + "mtp.norm.weight", + "mtp.pre_fc_norm_embedding.weight", + "mtp.pre_fc_norm_hidden.weight", + "layers.0.linear_attn.norm.weight", + "embed_tokens.weight", + ]; + for m in [ + &Qwen35HfMapper as &dyn HfMapper, + &Qwen35MoeHfMapper as &dyn HfMapper, + ] { + for s in shifted_by_hf { + assert_eq!( + m.mlx_norm_shift(s), + 0.0, + "{}: MLX source must not shift {s}", + m.canonical_arch() + ); + } + // The HF rule is untouched. + assert_eq!(m.norm_shift("layers.0.input_norm.weight"), 1.0); + } + } + /// Qwen3.5-35B-A3B / Qwen3.6-35B-A3B (identical config shape): the REAL /// checkpoint's text_config has NO `intermediate_size` (the HF config /// class deletes it — the dense-FFN slot is the shared expert), GQA GDN diff --git a/base-convert/crates/base-awq/src/lib.rs b/base-convert/crates/base-awq/src/lib.rs index 389b0ed..de850b1 100644 --- a/base-convert/crates/base-awq/src/lib.rs +++ b/base-convert/crates/base-awq/src/lib.rs @@ -108,8 +108,15 @@ pub struct AwqPlan { impl AwqConfig { /// AWQ alpha search per the original paper. Iterates alpha_grid, /// for each α scales weights by `absmax^α`, RTN-quantizes, - /// dequantizes, undoes the scale, and measures reconstruction - /// MSE. Returns the best α + its scale vector. + /// dequantizes, undoes the scale, and measures the + /// ACTIVATION-WEIGHTED reconstruction error — each channel's error + /// counts in proportion to its calibration absmax, i.e. + /// ‖(Q(W·s)·s⁻¹ − W)·diag(a)‖². Returns the best α + its scales. + /// + /// The weighting is the point of AWQ: with plain (unweighted) MSE + /// the α=0 identity always wins, because RTN already minimizes + /// per-group weight-space MSE — which is exactly the bug this + /// implementation shipped with (α=0.00 on every real tensor). /// /// `weights` is row-major `[out, in]`. `absmax_per_input_channel` /// has length = in_features and contains the per-channel max @@ -152,6 +159,7 @@ impl AwqConfig { in_features, out_features, &inverse_scales, + absmax_per_input_channel, bits, group_size, symmetric, @@ -218,7 +226,11 @@ fn scale_columns(weights: &[f32], in_features: usize, scales: &[f32]) -> Vec = vec![1.0; n_in]; - let plain_mse = - rtn_reconstruction_mse(&weights, &weights, n_in, n_out, &identity, 4, 64, false); + let plain_mse = rtn_reconstruction_mse( + &weights, &weights, n_in, n_out, &identity, &absmax, 4, 64, false, + ); // Tiny tolerance for floating-point order-of-ops. assert!( plan.mse <= plain_mse * 1.0001, @@ -313,8 +330,9 @@ mod tests { let plan = AwqConfig::default().search(&weights, n_in, &absmax, 2, 32, false); let identity: Vec = vec![1.0; n_in]; - let plain_mse = - rtn_reconstruction_mse(&weights, &weights, n_in, n_out, &identity, 2, 32, false); + let plain_mse = rtn_reconstruction_mse( + &weights, &weights, n_in, n_out, &identity, &absmax, 2, 32, false, + ); assert!( plan.mse < plain_mse, "AWQ should help: awq={} plain={}", diff --git a/base-convert/crates/base-awq/src/sidecar.rs b/base-convert/crates/base-awq/src/sidecar.rs index 3f3885b..f265013 100644 --- a/base-convert/crates/base-awq/src/sidecar.rs +++ b/base-convert/crates/base-awq/src/sidecar.rs @@ -35,10 +35,37 @@ impl AwqProfile { /// Lookup the per-input-channel absmax vector for a tensor name. /// Returns None if the profile lacks an entry — callers fall back /// to plain RTN for that tensor. + /// + /// The calibrate tool records RUNTIME tensor names + /// (`layers.N.attention.q.weight`, `layers.N.ffn.down.weight`) while + /// the converter looks up CANONICAL names + /// (`layers.N.self_attn.q_proj.weight`, `layers.N.mlp.down_proj.weight`), + /// so an exact miss falls back to the canonical→runtime translation. + /// `ffn.gate` serves BOTH `gate_proj` and `up_proj`: the runtime + /// captures one absmax per matmul INPUT, and gate/up consume the same + /// post-norm hidden state. pub fn absmax(&self, tensor_name: &str) -> Option<&[f32]> { - self.per_tensor_absmax - .get(tensor_name) - .map(|v| v.as_slice()) + if let Some(v) = self.per_tensor_absmax.get(tensor_name) { + return Some(v.as_slice()); + } + const MAP: [(&str, &str); 7] = [ + (".self_attn.q_proj.weight", ".attention.q.weight"), + (".self_attn.k_proj.weight", ".attention.k.weight"), + (".self_attn.v_proj.weight", ".attention.v.weight"), + (".self_attn.o_proj.weight", ".attention.output.weight"), + (".mlp.down_proj.weight", ".ffn.down.weight"), + (".mlp.gate_proj.weight", ".ffn.gate.weight"), + (".mlp.up_proj.weight", ".ffn.gate.weight"), + ]; + for (canon, runtime) in MAP { + if let Some(prefix) = tensor_name.strip_suffix(canon) { + let alt = format!("{prefix}{runtime}"); + if let Some(v) = self.per_tensor_absmax.get(&alt) { + return Some(v.as_slice()); + } + } + } + None } /// Validate the profile against an expected source fingerprint diff --git a/base-convert/crates/base-convert/src/gpt_oss.rs b/base-convert/crates/base-convert/src/gpt_oss.rs index d42fadf..2e10c69 100644 --- a/base-convert/crates/base-convert/src/gpt_oss.rs +++ b/base-convert/crates/base-convert/src/gpt_oss.rs @@ -235,6 +235,19 @@ pub(crate) fn convert_gpt_oss( // fine and runs past `<|call|>`. let gc: serde_json::Value = serde_json::from_slice(&bytes) .with_context(|| format!("parsing {}", gc_path.display()))?; + // Same file also carries the author's recommended SAMPLING + // settings. gpt-oss is routed here before the generic format + // detection, so it never reaches the shared + // apply_generation_config() and would otherwise be the one family + // served with runtime fallbacks instead of its published defaults. + config.sampling_defaults = base_arch::SamplingDefaults::from_generation_config(&gc); + if !config.sampling_defaults.is_empty() { + let sd = &config.sampling_defaults; + eprintln!( + " sampling: model defaults temp={:?} top_p={:?} top_k={:?} min_p={:?}", + sd.temperature, sd.top_p, sd.top_k, sd.min_p + ); + } let ids: Vec = match gc.get("eos_token_id") { Some(serde_json::Value::Number(n)) => { n.as_u64().map(|x| vec![x as u32]).unwrap_or_default() @@ -348,6 +361,7 @@ pub(crate) fn convert_gpt_oss( .collect(), tensors: vec![], mmproj: None, + speculator: None, calibration: None, sig: None, provenance: None, // filled below via the writer's header diff --git a/base-convert/crates/base-convert/src/hub.rs b/base-convert/crates/base-convert/src/hub.rs index 28cce67..69b5420 100644 --- a/base-convert/crates/base-convert/src/hub.rs +++ b/base-convert/crates/base-convert/src/hub.rs @@ -386,6 +386,13 @@ fn resolve_model_args(rest: &[String], default_variant: Option<&str>) -> Result< Ok(resolve_hub_model(arg, default_variant)? .to_string_lossy() .into_owned()) + } else if let Some((kind, id)) = speculate_drafter_id(arg) { + // `--speculate dflash:` takes the SAME road as the model + // it drafts for: catalogue lookup, pull, convert-on-pull, + // variant selection, hub cache. A drafter was the one model a + // user had to fetch and convert by hand and then name by path. + let path = resolve_hub_model(id, default_variant)?; + Ok(format!("{kind}:{}", path.to_string_lossy())) } else { Ok(arg.clone()) } @@ -393,6 +400,19 @@ fn resolve_model_args(rest: &[String], default_variant: Option<&str>) -> Result< .collect() } +/// `:` for the strategies that take a drafter artifact — +/// `dflash:basecompute/Qwen3-4B-DFlash`, `eagle3:org/head:default-q4`. Returns +/// None for a path (`dflash:./x.base`, which keeps working unchanged), for a +/// strategy that takes no artifact (`prompt-lookup`, `mtp-head`), and for +/// anything that is not id-shaped. +fn speculate_drafter_id(token: &str) -> Option<(&str, &str)> { + let (kind, rest) = token.split_once(':')?; + if !matches!(kind, "dflash" | "dspark" | "eagle3" | "draft-model") { + return None; + } + looks_like_hub_id(rest).then_some((kind, rest)) +} + /// Pull a `--variant ` / `--variant=` selector out of the forwarded args, /// returning the variant (if any) and the args with it removed — the runtime /// binary never sees a flag it doesn't understand. The selector applies to @@ -469,7 +489,8 @@ pub fn dispatch_external(argv: Vec) -> Result<()> { let is_computearena = cmd == "computearena"; let (candidates, rest) = if is_computearena { // ComputeArena is independently distributed and owns runtime - // selection. Preserve its arguments and select the BaseRT adapter. + // selection. Preserve its arguments and select the BaseRT adapter — + // no variant resolution. let mut forwarded = Vec::with_capacity(rest.len() + 1); forwarded.push("basert".to_string()); forwarded.extend_from_slice(rest); @@ -826,9 +847,17 @@ fn pull_and_convert( awq_profile: None, allow_quant_from_quant: false, no_mlx_passthrough: false, + quantize_drafter: false, + speculator_lm_head: false, + mtp_draft_ranking: None, + mtp_draft_vocab: None, // Convert-on-pull always produces a canonical-quant bundle; // k-quant passthrough stays an explicit `convert` opt-in. kquant_passthrough: false, + // Convert-on-pull follows the converter default: a tied model gets + // its head materialized so the profile head rule sets logit + // precision (`plan_tied_lm_head`); the opt-out is a `convert` flag. + no_synth_tied_head: false, // Convert-on-pull sources HF checkpoints, where the perception // tower already lives in the same snapshot — the mmproj flags only // apply to GGUF sources, which ship the tower separately. @@ -1836,6 +1865,27 @@ mod tests { assert!(!looks_like_hub_id("")); } + #[test] + fn speculate_drafter_id_splits_ids_but_not_paths() { + assert_eq!( + speculate_drafter_id("dflash:basecompute/Qwen3-4B-DFlash"), + Some(("dflash", "basecompute/Qwen3-4B-DFlash")) + ); + // An inline variant rides along: resolve_hub_model parses it. + assert_eq!( + speculate_drafter_id("eagle3:org/head:default-q4"), + Some(("eagle3", "org/head:default-q4")) + ); + // Paths keep working: the artifact may be a local build. + assert_eq!(speculate_drafter_id("dflash:./drafters/x.base"), None); + assert_eq!(speculate_drafter_id("dspark:/abs/x.base"), None); + assert_eq!(speculate_drafter_id("dflash:x.base"), None); + // Strategies without an artifact, and unknown prefixes. + assert_eq!(speculate_drafter_id("prompt-lookup"), None); + assert_eq!(speculate_drafter_id("mtp-head"), None); + assert_eq!(speculate_drafter_id("nope:org/model"), None); + } + #[test] fn resolve_model_args_passes_through_non_ids() { // Flags and plain paths are untouched (and never hit the registry). diff --git a/base-convert/crates/base-convert/src/lora.rs b/base-convert/crates/base-convert/src/lora.rs new file mode 100644 index 0000000..e76ceaf --- /dev/null +++ b/base-convert/crates/base-convert/src/lora.rs @@ -0,0 +1,474 @@ +//! `basert convert-lora`: PEFT adapter directory → adapter `.base` bundle +//! (capability-audit: the LoRA endpoints existed with no way to author the +//! adapter files they load). +//! +//! Reads `adapter_config.json` + `adapter_model.safetensors` and emits the +//! `lora..{A,B}` tensor pairs the runtime loader consumes +//! (`src/lora/lora_loader.cpp`), plus `lora.rank` / `lora.alpha` / +//! `lora.max_seq` metadata. B is stored RAW — the loader pre-scales it by +//! alpha/rank at load time. +//! +//! Targets are keyed by the RUNTIME DISPATCH names (`layers.N.attention.q +//! .weight`, `layers.N.ffn.gate.weight`, …) — what `dispatch_gemm` receives +//! and `apply_delta` looks up — not the converter-canonical bundle names. +//! Two families of GEMM exist per projection group: +//! +//! - the FUSED path (qkv in one GEMM carrying the q name with +//! N = q+k+v; gate+up in one GEMM carrying the gate name with +//! N = 2*ffn), used by quantized decode/prefill; and +//! - the SPLIT path (per-projection GEMMs), used by mixed-dtype configs. +//! +//! The two are mutually exclusive at dispatch time, so we emit BOTH forms: +//! stacked block-diagonal pairs for the fused names AND per-projection +//! pairs for the split names. Whichever path runs finds its entry; the +//! other idles. Stacking needs the true per-projection output widths, which +//! come from `--base ` (also validating that the adapter and +//! model agree on the hidden size). +//! +//! Scope: the Llama-shaped families (llama / mistral / phi3 — the runtime +//! names above). Other architectures refuse loudly rather than emitting +//! entries that would silently never match a dispatch name. + +use anyhow::{bail, Context, Result}; +use base_format::{Header, TensorDtype, TensorEntry, TensorPayload}; +use base_readers::safetensors::{SafetensorsFile, StDtype}; +use std::collections::BTreeMap; +use std::path::PathBuf; + +#[derive(clap::Parser, Debug)] +pub struct ConvertLoraArgs { + /// PEFT adapter directory (adapter_config.json + adapter_model.safetensors). + pub input: PathBuf, + /// The `.base` model bundle this adapter patches. Supplies the fused-GEMM + /// output widths for stacking and validates the hidden size. + #[arg(long)] + pub base: PathBuf, + /// Output adapter `.base` file. Defaults to `/adapter.base`. + #[arg(short, long)] + pub output: Option, + /// `lora.max_seq`: the largest per-dispatch row count the runtime's + /// delta scratch is sized for (prefill chunks must fit). + #[arg(long, default_value_t = 4096)] + pub max_seq: u32, +} + +/// One projection's A/B pair in f32 (row-major). +struct Pair { + a: Vec, // [r, in] + b: Vec, // [out, r] + r: usize, + d_in: usize, + d_out: usize, +} + +fn to_f32(st: &SafetensorsFile, name: &str) -> Result<(Vec, Vec)> { + let info = st + .tensors + .iter() + .find(|t| t.name == name) + .with_context(|| format!("adapter tensor {name} missing"))? + .clone(); + let bytes = st.tensor_bytes(&info); + let n: usize = info.shape.iter().product::() as usize; + let v = match info.dtype { + StDtype::F32 => bytes + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(), + StDtype::F16 => bytes + .chunks_exact(2) + .map(|c| half::f16::from_le_bytes([c[0], c[1]]).to_f32()) + .collect(), + StDtype::Bf16 => bytes + .chunks_exact(2) + .map(|c| f32::from_bits(u32::from(u16::from_le_bytes([c[0], c[1]])) << 16)) + .collect(), + other => bail!("adapter tensor {name}: unsupported dtype {other:?}"), + }; + let v: Vec = v; + if v.len() != n { + bail!("adapter tensor {name}: byte length disagrees with shape"); + } + Ok((info.shape.clone(), v)) +} + +fn f32s_to_f16_bytes(v: &[f32]) -> Vec { + let mut out = Vec::with_capacity(v.len() * 2); + for &x in v { + out.extend_from_slice(&half::f16::from_f32(x).to_le_bytes()); + } + out +} + +/// Shapes of the base model's canonical projection tensors, per layer: +/// (name suffix → [out, in]). +type BaseInfo = ( + BTreeMap>, + BTreeMap, +); + +fn base_shapes(base: &PathBuf) -> Result { + let f = std::fs::read(base).with_context(|| format!("read {}", base.display()))?; + if f.len() < 16 || &f[0..4] != b"BASE" { + bail!("{} is not a .base bundle", base.display()); + } + let hlen = u64::from_le_bytes(f[8..16].try_into().unwrap()) as usize; + let h: serde_json::Value = serde_json::from_slice(&f[16..16 + hlen])?; + let arch = h["arch"].as_str().unwrap_or(""); + if arch != "llama" { + bail!( + "convert-lora currently maps the Llama-shaped runtime dispatch names; the base model's \ + arch is {arch:?}. Refusing to emit entries that would silently never match." + ); + } + let mut m = BTreeMap::new(); + for t in h["tensors"] + .as_array() + .context(".base header has no tensors")? + { + let name = t["name"].as_str().unwrap_or(""); + if name.contains(".self_attn.") || name.contains(".mlp.") { + let shape: Vec = t["shape"] + .as_array() + .map(|a| a.iter().filter_map(|x| x.as_u64()).collect()) + .unwrap_or_default(); + m.insert(name.to_string(), shape); + } + } + // The engine's weight store refuses a degenerate config (dim=0 …), so + // the adapter carries the BASE model's core dims — which doubles as + // provenance for which model this adapter patches. + let mut cfg = BTreeMap::new(); + if let Some(c) = h["config"].as_object() { + for k in [ + "architecture", + "hidden_size", + "num_hidden_layers", + "num_attention_heads", + "num_key_value_heads", + "head_dim", + "vocab_size", + "intermediate_size", + ] { + if let Some(v) = c.get(k) { + cfg.insert(k.to_string(), v.clone()); + } + } + } + Ok((m, cfg)) +} + +/// Stack member pairs block-diagonally into one fused-target pair. +/// `members`: (pair, out_offset) in fused-row order; `total_out` covers the +/// full fused GEMM N so absent members contribute zero rows. +fn stack(members: &[(&Pair, usize)], total_out: usize, d_in: usize) -> Pair { + let eff_r: usize = members.iter().map(|(p, _)| p.r).sum(); + let mut a = vec![0f32; eff_r * d_in]; + let mut b = vec![0f32; total_out * eff_r]; + let mut r_off = 0usize; + for (p, out_off) in members { + for r in 0..p.r { + a[(r_off + r) * d_in..(r_off + r + 1) * d_in] + .copy_from_slice(&p.a[r * d_in..(r + 1) * d_in]); + } + for o in 0..p.d_out { + for r in 0..p.r { + b[(out_off + o) * eff_r + (r_off + r)] = p.b[o * p.r + r]; + } + } + r_off += p.r; + } + Pair { + a, + b, + r: eff_r, + d_in, + d_out: total_out, + } +} + +pub fn cmd_convert_lora(args: ConvertLoraArgs) -> Result<()> { + let cfg_path = args.input.join("adapter_config.json"); + let cfg: serde_json::Value = serde_json::from_str( + &std::fs::read_to_string(&cfg_path) + .with_context(|| format!("read {}", cfg_path.display()))?, + )?; + let rank = cfg["r"].as_u64().context("adapter_config.json missing r")? as usize; + let alpha = cfg["lora_alpha"] + .as_f64() + .context("adapter_config.json missing lora_alpha")? as f32; + let st = SafetensorsFile::open(args.input.join("adapter_model.safetensors"))?; + let (shapes, base_cfg) = base_shapes(&args.base)?; + let out = args + .output + .clone() + .unwrap_or_else(|| args.input.join("adapter.base")); + + // Collect per-(layer, projection) pairs from the PEFT names. + // `base_model.model.model.layers.N.self_attn.q_proj.lora_A.weight`. + let mut pairs: BTreeMap<(u32, String), Pair> = BTreeMap::new(); + let names: Vec = st.tensors.iter().map(|t| t.name.clone()).collect(); + for name in &names { + let Some(rest) = name.split(".layers.").nth(1) else { + continue; + }; + if !name.ends_with(".lora_A.weight") { + continue; + } + let mut it = rest.splitn(2, '.'); + let layer: u32 = it.next().unwrap_or("").parse().context("layer index")?; + let proj_path = it.next().unwrap_or(""); // "self_attn.q_proj.lora_A.weight" + let proj = proj_path + .trim_end_matches(".lora_A.weight") + .rsplit('.') + .next() + .unwrap_or("") + .to_string(); // "q_proj" + let b_name = name.replace(".lora_A.weight", ".lora_B.weight"); + let (a_shape, a) = to_f32(&st, name)?; + let (b_shape, b) = to_f32(&st, &b_name)?; + if a_shape.len() != 2 + || b_shape.len() != 2 + || a_shape[0] as usize != rank + || b_shape[1] as usize != rank + { + bail!( + "{name}: expected A [r,in] / B [out,r] with r={rank}, got {a_shape:?}/{b_shape:?}" + ); + } + pairs.insert( + (layer, proj), + Pair { + a, + b, + r: rank, + d_in: a_shape[1] as usize, + d_out: b_shape[0] as usize, + }, + ); + } + if pairs.is_empty() { + bail!("no lora_A/lora_B pairs found in adapter_model.safetensors"); + } + + // Validate against the base model and emit runtime-target pairs. + let layers: std::collections::BTreeSet = pairs.keys().map(|(l, _)| *l).collect(); + let mut targets: Vec<(String, Pair)> = Vec::new(); + let shape_of = |layer: u32, canon: &str| -> Option<(usize, usize)> { + shapes + .get(&format!("layers.{layer}.{canon}.weight")) + .and_then(|s| (s.len() == 2).then(|| (s[0] as usize, s[1] as usize))) + }; + for layer in &layers { + let get = |p: &str| pairs.get(&(*layer, p.to_string())); + // Per-projection sanity + split-path entries. + let singles = [ + ("q_proj", "attention.q"), + ("k_proj", "attention.k"), + ("v_proj", "attention.v"), + ("o_proj", "attention.output"), + ("gate_proj", "ffn.gate"), + ("up_proj", "ffn.up"), + ("down_proj", "ffn.down"), + ]; + for (peft, runtime) in singles { + let Some(p) = get(peft) else { continue }; + let canon = match peft { + "o_proj" => "self_attn.o_proj".to_string(), + "q_proj" | "k_proj" | "v_proj" => format!("self_attn.{peft}"), + _ => format!("mlp.{peft}"), + }; + if let Some((out_n, in_k)) = shape_of(*layer, &canon) { + if p.d_in != in_k || p.d_out != out_n { + bail!( + "layer {layer} {peft}: adapter is [{}x{}] but the base tensor is [{out_n}x{in_k}] — \ + wrong-model adapter?", + p.d_out, p.d_in + ); + } + } + targets.push(( + format!("layers.{layer}.{runtime}.weight"), + Pair { + a: p.a.clone(), + b: p.b.clone(), + r: p.r, + d_in: p.d_in, + d_out: p.d_out, + }, + )); + } + // Fused qkv (dispatch name = the q name, N = q+k+v) — emitted only + // when the fused dispatch name differs from the split one, which it + // does not for q ("attention.q" carries both) … it DOES differ in N: + // the fused GEMM's tensor_name is STILL "layers.N.attention.q.weight" + // but with N = q+k+v, so the SAME key must carry the STACKED pair for + // the fused path to apply all three deltas — and the split path's + // q-GEMM then finds a [q+k+v, ..]-shaped B whose N mismatches and + // (by the loader's shape guard) skips. One key cannot serve both. + // The fused path is what every quantized bundle runs, so the stacked + // pair WINS the key; the split-path singles above keep k/v/o and the + // FFN entries live. + let (q, k, v) = (get("q_proj"), get("k_proj"), get("v_proj")); + if q.is_some() || k.is_some() || v.is_some() { + let (qn, d) = shape_of(*layer, "self_attn.q_proj") + .with_context(|| format!("layer {layer}: base has no q_proj shape"))?; + let (kn, _) = shape_of(*layer, "self_attn.k_proj") + .with_context(|| format!("layer {layer}: base has no k_proj shape"))?; + let (vn, _) = shape_of(*layer, "self_attn.v_proj") + .with_context(|| format!("layer {layer}: base has no v_proj shape"))?; + let mut members: Vec<(&Pair, usize)> = Vec::new(); + if let Some(p) = q { + members.push((p, 0)); + } + if let Some(p) = k { + members.push((p, qn)); + } + if let Some(p) = v { + members.push((p, qn + kn)); + } + let stacked = stack(&members, qn + kn + vn, d); + // Replace the single-q entry for this key (see comment above). + targets.retain(|(n, _)| n != &format!("layers.{layer}.attention.q.weight")); + targets.push((format!("layers.{layer}.attention.q.weight"), stacked)); + } + // Fused gate+up (dispatch name = the gate name, N = 2*ffn). + let (g, u) = (get("gate_proj"), get("up_proj")); + if g.is_some() || u.is_some() { + let (gn, d) = shape_of(*layer, "mlp.gate_proj") + .with_context(|| format!("layer {layer}: base has no gate_proj shape"))?; + let (un, _) = shape_of(*layer, "mlp.up_proj") + .with_context(|| format!("layer {layer}: base has no up_proj shape"))?; + let mut members: Vec<(&Pair, usize)> = Vec::new(); + if let Some(p) = g { + members.push((p, 0)); + } + if let Some(p) = u { + members.push((p, gn)); + } + let stacked = stack(&members, gn + un, d); + targets.retain(|(n, _)| n != &format!("layers.{layer}.ffn.gate.weight")); + targets.push((format!("layers.{layer}.ffn.gate.weight"), stacked)); + } + } + + // Write the adapter bundle: f16 A/B pairs + flat lora.* metadata. + let mut metadata = BTreeMap::new(); + metadata.insert("lora.rank".to_string(), serde_json::json!(rank)); + metadata.insert("lora.alpha".to_string(), serde_json::json!(alpha)); + metadata.insert("lora.max_seq".to_string(), serde_json::json!(args.max_seq)); + let header = Header { + schema: 1, + arch: "lora-adapter".to_string(), + quant_scheme: base_format::QuantScheme::F16, + min_hw: "apple_m1".to_string(), + created: crate::chrono_now(), + base_rt_version: env!("CARGO_PKG_VERSION").to_string(), + source: base_format::SourceInfo { + format: "peft_safetensors".to_string(), + sha256: String::new(), + filename: args + .input + .file_name() + .and_then(|n| n.to_str()) + .unwrap_or("adapter") + .to_string(), + }, + tokenizer: base_format::TokenizerBlob { + fields: BTreeMap::new(), + }, + config: base_format::ModelConfig { fields: base_cfg }, + metadata, + target_backend: base_format::TargetBackend::Metal, + quant_profile: String::new(), + alignment: Default::default(), + flags: base_format::HeaderFlags::HAS_LORA, + layers: Vec::new(), + tensors: Vec::new(), + calibration: None, + mmproj: None, + speculator: None, + // A LoRA adapter carries its own f16 A/B pairs verbatim from the PEFT + // checkpoint; there is no source-tensor transplant to record. + provenance: None, + sig: None, + }; + let mut w = base_format::BaseWriter::create(&out, header)?; + let mut n_pairs = 0usize; + for (target, p) in &targets { + for (suffix, data, shape) in [ + ("A", &p.a, vec![p.r as u64, p.d_in as u64]), + ("B", &p.b, vec![p.d_out as u64, p.r as u64]), + ] { + w.add_tensor(TensorPayload { + entry: TensorEntry { + name: format!("lora.{target}.{suffix}"), + dtype: TensorDtype::F16, + shape, + offset: 0, + length: 0, + scale_offset: None, + scale_length: None, + bias_offset: None, + bias_length: None, + awq_scale_offset: None, + awq_scale_length: None, + group_size: None, + scale_dtype: None, + symmetric: false, + layout: None, + residency: None, + compute_region: Default::default(), + flags: base_format::TensorFlags::empty(), + checksum_xxh64: None, + source_ggml_type: None, + }, + data: f32s_to_f16_bytes(data), + }); + } + n_pairs += 1; + } + w.finish()?; + println!( + "wrote {} ({} targets, rank {}, alpha {})", + out.display(), + n_pairs, + rank, + alpha + ); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Block-diagonal stacking: member deltas must land in their own output + /// rows and rank columns; absent members contribute zero rows. + #[test] + fn stack_places_members_block_diagonally() { + let q = Pair { + a: vec![1.0, 2.0], + b: vec![10.0, 20.0], + r: 1, + d_in: 2, + d_out: 2, + }; + let v = Pair { + a: vec![3.0, 4.0], + b: vec![30.0], + r: 1, + d_in: 2, + d_out: 1, + }; + // Fused layout [q(2) | k(1, absent) | v(1)] → total_out 4. + let s = stack(&[(&q, 0), (&v, 3)], 4, 2); + assert_eq!(s.r, 2); + assert_eq!(s.d_out, 4); + // A rows: q's then v's. + assert_eq!(s.a, vec![1.0, 2.0, 3.0, 4.0]); + // B [4 rows x 2 ranks]: q rows use col 0, absent k row all-zero, + // v row uses col 1. + assert_eq!(s.b, vec![10.0, 0.0, 20.0, 0.0, 0.0, 0.0, 0.0, 30.0]); + } +} diff --git a/base-convert/crates/base-convert/src/main.rs b/base-convert/crates/base-convert/src/main.rs index 41a128c..91fdcd6 100644 --- a/base-convert/crates/base-convert/src/main.rs +++ b/base-convert/crates/base-convert/src/main.rs @@ -4,6 +4,7 @@ use std::path::PathBuf; mod gpt_oss; mod hub; +mod lora; /// Hand-written top-level help: the runtime commands are dispatched via /// clap's external-subcommand catch-all (invisible to clap's auto listing), @@ -54,6 +55,9 @@ struct Args { enum Cmd { /// Convert a source model to `.base`. Convert(ConvertArgs), + /// Convert a PEFT LoRA adapter directory to an adapter `.base` + /// (loadable via `/v1/lora/load` / `baseRT_lora_load`). + ConvertLora(lora::ConvertLoraArgs), /// Sign an existing unsigned `.base` file. Sign(SignArgs), /// Verify a signed `.base` file. @@ -148,6 +152,15 @@ struct ConvertArgs { /// transplant path existed. #[arg(long)] no_mlx_passthrough: bool, + /// Block-drafter / EAGLE-3 sidecars only: quantize the drafter's own + /// transformer layers (`layers.*` q/k/v/o and FFN projections) with + /// `--target` instead of keeping the whole sidecar f16. The heads that + /// define the scheme (fc taps, hidden_norm, Markov / confidence heads, + /// the EAGLE-3 lm_head and d2t) and every norm stay f16. A DFlash-b16 + /// head for Qwen3-4B is 1 GB f16 — streamed on every block step; q8 + /// halves it, q4 quarters it, at a small cost in acceptance. + #[arg(long)] + quantize_drafter: bool, /// GGUF sources only: copy Q4_K / Q5_K / Q6_K super-blocks into the /// bundle VERBATIM instead of dequantizing and re-packing them. /// @@ -167,6 +180,52 @@ struct ConvertArgs { #[arg(long)] kquant_passthrough: bool, + /// Do NOT materialize a separate `lm_head.weight` for a tied model. + /// + /// By default a `tie_word_embeddings` source gets a second copy of its + /// embedding appended under `lm_head.weight`, so the profile's head rule + /// (q6/q8, or the `**.weight` catch-all) sets the LOGIT-projection + /// precision while `embed_tokens.weight` keeps its own (usually f16) + /// rule — see `plan_tied_lm_head`. Relative to a bundle converted before + /// that existed, the logits of Qwen3-0.6B/1.7B/4B or Llama-3.2 move from + /// f16 to the head rule's dtype on reconvert. This flag reproduces the + /// older layout: the runtime projects logits through the embedding + /// tensor at the embedding's precision. Every conversion prints which + /// of the two it did and which profile rule the head hit. + /// + /// Archs whose runtime never reads `output.weight` (Gemma 3/4) never + /// synthesize a head regardless of this flag. + #[arg(long)] + no_synth_tied_head: bool, + + /// Give a bundle's MTP speculator head its OWN copy of the lm_head + /// (`mtp.lm_head.weight`, quantized by the profile's rule for that + /// name — the `**.weight` catch-all otherwise) instead of sharing the + /// target's. On q4mix profiles the target's head is q8 for accuracy, + /// and every chained draft of the speculator then streams the full q8 + /// head (the 27B's 1.27 GB is 8 of a draft's 10 ms on a GB10); the + /// drafts are verified against the target anyway, so the head can use + /// a q4 copy. Costs the copy's bytes (a q4 head is half the q8 one). + #[arg(long)] + speculator_lm_head: bool, + + /// Ship a drafting-vocabulary ranking with the bundle's MTP speculator + /// head: a ranking JSON (`{"ranked_ids": [...]}`, token ids most frequent + /// first — tools/mtp_draft_ranking.py writes one) stored as + /// `mtp.draft_rank` (f32 ids) in the speculator section. The runtime's + /// chain then scores only the top `--mtp-draft-vocab` rows of the lm_head + /// per drafted token (plus every special token) instead of the whole + /// vocabulary — FR-Spec; the target's verify is unchanged, so output is + /// too. Refused on a bundle without an MTP speculator. + #[arg(long, value_name = "RANKING_JSON")] + mtp_draft_ranking: Option, + + /// With `--mtp-draft-ranking`: the drafting vocabulary's default size + /// (`draft_vocab_size` in the speculator config; BASERT_MTP_DRAFT_VOCAB + /// overrides it at run time). Unset = the runtime's default. + #[arg(long, value_name = "K", requires = "mtp_draft_ranking")] + mtp_draft_vocab: Option, + /// GGUF sources only: fold a companion `mmproj-*.gguf` perception tower /// into the bundle, so one `.base` carries both the text model and the /// vision encoder. @@ -417,6 +476,7 @@ fn main() -> Result<()> { let args = Args::parse(); match args.cmd { Cmd::Convert(a) => cmd_convert(a), + Cmd::ConvertLora(a) => lora::cmd_convert_lora(a), Cmd::Sign(a) => cmd_sign(a), Cmd::Verify(a) => cmd_verify(a), Cmd::Inspect(a) => cmd_inspect(a), @@ -489,7 +549,7 @@ fn cmd_convert(args: ConvertArgs) -> Result<()> { let fmt = base_readers::detect_format(&args.input) .with_context(|| format!("detecting format for {:?}", args.input))?; - match fmt { + let result = match fmt { SourceFormat::Gguf => convert_gguf( &args.input, &output, @@ -499,7 +559,35 @@ fn cmd_convert(args: ConvertArgs) -> Result<()> { ), SourceFormat::HfSafetensors => convert_hf(&args.input, &output, &ctx), SourceFormat::MlxSafetensors => convert_mlx(&args.input, &output, &ctx), + }; + if result.is_ok() + && ctx.awq_profile.is_some() + && ctx.awq_applied.load(std::sync::atomic::Ordering::Relaxed) == 0 + { + // Two different failures end in the same "0 rotated": no name in the + // sidecar matched any tensor, or names matched but every absmax had + // the wrong length. The fix differs (re-name vs re-capture), so say + // which one happened. + let shape_skipped = ctx + .awq_shape_skipped + .load(std::sync::atomic::Ordering::Relaxed); + if shape_skipped == 0 { + eprintln!( + " WARNING: --awq-profile matched ZERO tensors — every sidecar \ + lookup missed and plain RTN was used throughout. The sidecar's \ + tensor names do not correspond to this model." + ); + } else { + eprintln!( + " WARNING: --awq-profile rotated ZERO tensors — {shape_skipped} sidecar \ + lookups hit but every absmax length differed from the tensor's \ + in_features, so plain RTN was used throughout. The sidecar's names \ + correspond to this model but its stats were captured at a different \ + geometry (another size or variant); re-run calibration on this model." + ); + } } + result } /// Build the typed per-layer descriptors for the header. Homogeneous @@ -624,7 +712,14 @@ fn convert_gguf( let mapper = source_mapper_for_gguf(arch) .ok_or_else(|| anyhow::anyhow!("arch {:?} not supported yet", arch))?; - let config = mapper.config_from_gguf(&gguf.metadata)?; + let mut config = mapper.config_from_gguf(&gguf.metadata)?; + + config.tie_word_embeddings = infer_gguf_tied_embeddings( + mapper, + config.tie_word_embeddings, + gguf.tensors.iter().map(|t| t.name.as_str()), + ); + let config = config; eprintln!( " config: hidden={}, layers={}, heads={}/{}, ffn={}, vocab={}", @@ -696,6 +791,7 @@ fn convert_gguf( layers: layer_descriptors_from_config(&config), tensors: vec![], mmproj: None, + speculator: None, calibration: None, provenance: None, sig: None, @@ -733,11 +829,51 @@ fn convert_gguf( let mut kept = 0usize; let mut passthrough_tensors = 0usize; let mut unpermuted = 0usize; - for info in gguf.tensors.iter() { + // Resolve names first so a tie_word_embeddings model can append a synthetic + // `lm_head.weight` sourced from the embedding (see plan_tied_lm_head). Index + // into gguf.tensors rather than holding references so the extra entry can + // alias an earlier tensor. + let mut work: Vec<(usize, String)> = Vec::new(); + for (i, info) in gguf.tensors.iter().enumerate() { let Some(canonical) = mapper.map_tensor_name(&info.name) else { dropped += 1; continue; }; + work.push((i, canonical)); + } + { + let pairs: Vec<(String, String)> = work + .iter() + .map(|(i, c)| (gguf.tensors[*i].name.clone(), c.clone())) + .collect(); + let plan = plan_tied_lm_head( + &pairs, + config.tie_word_embeddings, + mapper.canonical_arch(), + ctx.synth_tied_head, + ); + let src_idx = match &plan { + TiedHeadPlan::Synthesize { src } => gguf.tensors.iter().position(|t| t.name == *src), + _ => None, + }; + // The passthrough branch below copies k-quant super-blocks verbatim, + // so a k-quant embedding yields a head at the SAME precision as the + // embedding — the profile's head rule never runs. Say so. + let copy = match src_idx { + Some(idx) if kquant_passthrough && is_kquant(gguf.tensors[idx].ggml_type) => { + TiedHeadCopy::Verbatim { + ggml_type: ggml_type_name(gguf.tensors[idx].ggml_type), + } + } + _ => TiedHeadCopy::Quantized, + }; + report_tied_head_plan(&plan, ctx, mapper.canonical_arch(), copy); + if let Some(idx) = src_idx { + work.push((idx, "lm_head.weight".to_string())); + } + } + for (idx, canonical) in work { + let info = &gguf.tensors[idx]; kept += 1; let raw = gguf @@ -1428,45 +1564,31 @@ fn kquant_passthrough_entry( } /// Convert from an HF safetensors directory. -fn convert_hf(input: &std::path::Path, output: &std::path::Path, ctx: &QuantContext) -> Result<()> { - use base_arch::hf_mapper_for_model_type; - use base_readers::hf::HfDir; - let hf = HfDir::open(input)?; - let model_type = hf - .model_type() - .ok_or_else(|| anyhow::anyhow!("config.json missing model_type"))?; - eprintln!(" arch: {}", model_type); - let mapper = hf_mapper_for_model_type(model_type) - .ok_or_else(|| anyhow::anyhow!("HF model_type {:?} not supported yet", model_type))?; - let config = mapper.config_from_hf(&hf.config)?; - // Whisper is an encoder-decoder speech model with its own - // engine-native tensor naming, flat `whisper.*` token metadata, and - // an all-f16 tensor policy — none of which fit the decoder-only - // quantize path below. Dedicated path. - if mapper.canonical_arch() == "whisper" { - return convert_whisper(input, output, ctx, &hf, config); - } - eprintln!( - " config: hidden={}, layers={}, heads={}/{}, ffn={}, vocab={}", - config.hidden_size, - config.num_hidden_layers, - config.num_attention_heads, - config.num_kv_heads, - config.intermediate_size, - config.vocab_size - ); - // Merge the generation stop set from generation_config.json. HF declares - // the authoritative end-of-generation ids there; config.json often carries - // only a single scalar eos. Phi-3 is the canonical case: config.json has - // eos_token_id=32000 (<|endoftext|>) but generation_config.json has - // [32000, 32001, 32007] where <|end|>(32007) ends every chat turn. Without - // the full set the runtime never stops on <|end|>, so the model floods / - // repeats past its turn (finish_reason=length forever). Keep config.json's - // primary eos first; append any extra ids from generation_config.json. - let mut config = config; +/// Fold a checkpoint's `generation_config.json` into `config`: the generation +/// STOP SET and the author's recommended SAMPLING settings. +/// +/// Shared by every directory-shaped conversion path. It used to live inline in +/// `convert_hf`, which meant an MLX checkpoint carrying the same file — and the +/// source downloader keeps it — silently lost both, and was served with runtime +/// fallbacks instead of what its author published. +fn apply_generation_config(input: &std::path::Path, config: &mut base_arch::ArchConfig) { let gc_path = input.join("generation_config.json"); if let Ok(text) = std::fs::read_to_string(&gc_path) { if let Ok(gc) = serde_json::from_str::(&text) { + // The same file also carries the author's recommended SAMPLING + // settings, which a serving runtime otherwise has to guess at (and + // guessed wrong: Qwen3.8 asks for top_k 20 / top_p 0.95, and + // sampling its whole 150K-id tail instead is what tipped agent + // turns into verbatim repeat loops). + config.sampling_defaults = base_arch::SamplingDefaults::from_generation_config(&gc); + if !config.sampling_defaults.is_empty() { + let sd = &config.sampling_defaults; + eprintln!( + " sampling: model defaults temp={:?} top_p={:?} top_k={:?} min_p={:?} \ + (repetition_penalty={:?} recorded, not applied)", + sd.temperature, sd.top_p, sd.top_k, sd.min_p, sd.repetition_penalty + ); + } let mut gen_ids: Vec = Vec::new(); match &gc["eos_token_id"] { serde_json::Value::Number(n) => { @@ -1509,7 +1631,71 @@ fn convert_hf(input: &std::path::Path, output: &std::path::Path, ctx: &QuantCont } } } +} +fn convert_hf(input: &std::path::Path, output: &std::path::Path, ctx: &QuantContext) -> Result<()> { + use base_arch::hf_mapper_for_model_type; + use base_readers::hf::HfDir; + let hf = HfDir::open(input)?; + // RedHat's speculators packaging keeps model_type inside + // `transformer_layer_config` (the EAGLE-3 layer's own config). + let nested_model_type = hf + .config + .get("transformer_layer_config") + .and_then(|t| t.get("model_type")) + .and_then(|v| v.as_str()) + .map(str::to_string); + let model_type: &str = hf + .model_type() + .or(nested_model_type.as_deref()) + .ok_or_else(|| anyhow::anyhow!("config.json missing model_type"))?; + // Block drafters (DFlash / DSpark) carry the BACKBONE's model_type + // (`qwen3`); their `architectures` entry is what tells them apart from + // the LM they draft for. + let mapper: &dyn base_arch::HfMapper = if let Some(m) = + base_arch::dflash::hf_mapper_for_architectures(&hf.config) + { + eprintln!( + " arch: {} (block drafter sidecar: {})", + model_type, + m.canonical_arch() + ); + m + } else if let Some(m) = base_arch::dflash::hf_mapper_for_eagle3(&hf.config) { + eprintln!(" arch: {} (EAGLE-3 head sidecar)", model_type); + m + } else { + eprintln!(" arch: {}", model_type); + hf_mapper_for_model_type(model_type) + .ok_or_else(|| anyhow::anyhow!("HF model_type {:?} not supported yet", model_type))? + }; + let config = mapper.config_from_hf(&hf.config)?; + // Whisper is an encoder-decoder speech model with its own + // engine-native tensor naming, flat `whisper.*` token metadata, and + // an all-f16 tensor policy — none of which fit the decoder-only + // quantize path below. Dedicated path. + if mapper.canonical_arch() == "whisper" { + return convert_whisper(input, output, ctx, &hf, config); + } + eprintln!( + " config: hidden={}, layers={}, heads={}/{}, ffn={}, vocab={}", + config.hidden_size, + config.num_hidden_layers, + config.num_attention_heads, + config.num_kv_heads, + config.intermediate_size, + config.vocab_size + ); + // Merge the generation stop set from generation_config.json. HF declares + // the authoritative end-of-generation ids there; config.json often carries + // only a single scalar eos. Phi-3 is the canonical case: config.json has + // eos_token_id=32000 (<|endoftext|>) but generation_config.json has + // [32000, 32001, 32007] where <|end|>(32007) ends every chat turn. Without + // the full set the runtime never stops on <|end|>, so the model floods / + // repeats past its turn (finish_reason=length forever). Keep config.json's + // primary eos first; append any extra ids from generation_config.json. + let mut config = config; + apply_generation_config(input, &mut config); let provider = HfTensorProvider { hf: &hf }; let mmproj_cfg = mmproj_config_from_hf(&hf); let config_for_permute = config.clone(); @@ -1525,6 +1711,13 @@ fn convert_hf(input: &std::path::Path, output: &std::path::Path, ctx: &QuantCont if nvfp4_source { eprintln!(" source: NVFP4-quantized checkpoint (mirror policy: transplant + carry)"); } + if let Some([bo, bi]) = hf.fp8_block_size() { + eprintln!( + " source: block-FP8 checkpoint ({} e4m3 weights × [{bo}, {bi}] scales, \ + dequantized to f32 before the target scheme)", + hf.fp8_weight_count() + ); + } convert_generic( input, output, @@ -1776,6 +1969,7 @@ fn convert_whisper( .collect(), tensors: vec![], mmproj: None, + speculator: None, calibration: None, sig: None, provenance: None, @@ -1968,7 +2162,11 @@ fn convert_mlx( } let mapper = hf_mapper_for_model_type(model_type) .ok_or_else(|| anyhow::anyhow!("model_type {:?} not supported yet", model_type))?; - let config = mapper.config_from_hf(&mlx.hf.config)?; + let mut config = mapper.config_from_hf(&mlx.hf.config)?; + // An MLX checkpoint ships the same generation_config.json as its HF source + // (the source downloader keeps it), so it gets the same stop set and the + // same published sampling defaults. + apply_generation_config(input, &mut config); eprintln!( " config: hidden={}, layers={}, heads={}/{}, ffn={}, vocab={}", config.hidden_size, @@ -2009,7 +2207,10 @@ fn convert_mlx( names.clone(), &tokenizer_from_hf(&mlx.hf), mmproj_cfg, - &|n| mapper.norm_shift(n), + // The MLX rule, not the HF one: mlx-lm may have baked a family's + // norm offset into the stored values already (Qwen3.5+ does), and + // shifting those again is BAS-959. + &|n| mapper.mlx_norm_shift(n), &|n| mapper.rope_permute_heads(n, &config_for_permute), &|n| mapper.value_transform(n), mapper.shape_fastest_first(), @@ -2021,7 +2222,7 @@ fn convert_mlx( )?; if ctx.validate { validate_mlx_bundle(output, &mlx, &names, mapper.canonical_arch(), &|n| { - mapper.norm_shift(n) + mapper.mlx_norm_shift(n) })?; } Ok(()) @@ -2054,6 +2255,12 @@ fn validate_mlx_bundle( use base_format::{BaseReader, TensorDtype}; let reader = BaseReader::open(output).context("re-opening written .base for --validate")?; let (mut n_pass, mut n_f16, mut n_repack, mut n_skip) = (0usize, 0usize, 0usize, 0usize); + // 1-D tensors the converter offset from the source. The comparison + // below mirrors the converter's own rule, so it cannot judge the rule; + // the count makes the rule's effect visible (BAS-959 passed "byte- + // identical" with every norm shifted — it should read 0 shifted for + // a family whose MLX export already carries its offset). + let mut n_shifted = 0usize; for n in source_names { let Some(Canonical::Main(canonical)) = to_canonical_name(n, canonical_arch) else { continue; @@ -2124,6 +2331,7 @@ fn validate_mlx_bundle( if entry.shape.len() == 1 { let s = norm_shift(&canonical); if s != 0.0 { + n_shifted += 1; for v in f32s.iter_mut() { *v += s; } @@ -2142,8 +2350,8 @@ fn validate_mlx_bundle( } } eprintln!( - " validate: OK — byte-identical to source ({n_pass} passthrough, {n_f16} f16, \ - {n_repack} repacked, {n_skip} skipped)" + " validate: OK — byte-identical to source ({n_pass} passthrough, {n_f16} f16 of which \ + {n_shifted} norm-shifted, {n_repack} repacked, {n_skip} skipped)" ); Ok(()) } @@ -2546,6 +2754,9 @@ fn bf16_lossless(vals: &[f32]) -> bool { fn vt_label(vt: base_arch::ValueTransform) -> &'static str { match vt { base_arch::ValueTransform::NegExp => "neg_exp", + // The divisor is the drafter mappers' fixed FC_SCALE_DIV (256); the + // header's `fc_scale` records the value. + base_arch::ValueTransform::DivBy(_) => "div_by_fc_scale", } } @@ -3206,6 +3417,195 @@ impl TensorProvider for GlmKvbSplitProvider<'_> { } } +/// Column-block split of a fused projection into per-input GEMM operands, +/// so the runtime feeds each input through a plain GEMM instead of +/// assembling a concatenated activation: +/// +/// * Qwen3.5/3.6/3.8 MTP head: `mtp.fc.weight` ([dim, 2*dim], the +/// projection of `concat(norm(embed(t+1)), norm(h_t))`) → +/// `mtp.fc_embed.weight` + `mtp.fc_hidden.weight` ([dim, dim] each). +/// * Block drafters (DFlash / DSpark sidecars): `fc.weight` +/// ([dim, n_taps * dim], the projection of the concatenated target taps) +/// → `fc.{i}.weight` per tap, [dim, dim] each, in target_layer_ids order. +/// +/// Values are copied column-wise from the bf16 source; every part stays +/// f16 like the rest of the speculator / drafter tensors. +struct FcSplitProvider<'a> { + inner: &'a dyn TensorProvider, + /// virtual name -> (fused source name, column block index, block count) + splits: std::collections::BTreeMap, + rewritten: Vec, +} +impl<'a> FcSplitProvider<'a> { + fn build( + inner: &'a dyn TensorProvider, + source_names: &[String], + canonical_arch: &str, + expected_taps: Option, + ) -> Result { + let mut splits = std::collections::BTreeMap::new(); + let mut rewritten = Vec::with_capacity(source_names.len()); + let mtp = canonical_arch.starts_with("qwen35"); + let eagle3 = base_arch::dflash::is_eagle3_arch(canonical_arch); + let drafter = base_arch::dflash::is_block_drafter_arch(canonical_arch) || eagle3; + for n in source_names { + // EAGLE-3: q/k/v project from [normed embedding ; normed feature] + // (2·dim columns) — split into the embedding half and the + // feature half so the runtime runs two GEMMs over its two inputs. + // (`layers.0.` is the speculators packaging's name for the midlayer.) + if eagle3 + && (n.starts_with("midlayer.self_attn.") || n.starts_with("layers.0.self_attn.")) + && n.ends_with("_proj.weight") + && !n.contains("o_proj") + { + let shape = inner.source_shape(n)?; + if shape.len() != 2 || shape[1] % 2 != 0 { + bail!("{n}: expected a [out, 2*dim] EAGLE-3 projection, got {shape:?}"); + } + let e = n.replace("_proj.weight", "_proj_emb.weight"); + let h = n.replace("_proj.weight", "_proj_hid.weight"); + splits.insert(e.clone(), (n.clone(), 0, 2)); + splits.insert(h.clone(), (n.clone(), 1, 2)); + rewritten.push(e); + rewritten.push(h); + continue; + } + if mtp && (n == "mtp.fc.weight" || n.ends_with(".mtp.fc.weight")) { + let e = n.replace("mtp.fc.weight", "mtp.fc_embed.weight"); + let h = n.replace("mtp.fc.weight", "mtp.fc_hidden.weight"); + splits.insert(e.clone(), (n.clone(), 0, 2)); + splits.insert(h.clone(), (n.clone(), 1, 2)); + rewritten.push(e); + rewritten.push(h); + } else if drafter && n == "fc.weight" { + let shape = inner.source_shape(n)?; + if shape.len() != 2 + || shape[0] == 0 + || shape[1] % shape[0] != 0 + || shape[1] / shape[0] > 8 + { + bail!( + "{n}: expected a [dim, n_taps * dim] drafter fc (1..8 taps), got {shape:?}" + ); + } + let parts = (shape[1] / shape[0]) as usize; + // The runtime derives n_taps from the CONFIG and binds only + // fc.0..fc.{n_taps-1}. A checkpoint whose fc carries more + // blocks than the config declares taps would convert, load and + // draft — with the extra trained projections silently dropped, + // i.e. a different proposal model than the one that was + // trained. Fewer blocks than taps fails later at bind; catch + // both here, where the source is still in hand. + if let Some(taps) = expected_taps { + if parts != taps { + bail!( + "{n}: the checkpoint splits into {parts} tap blocks but the config declares {taps} \ + taps — convert the drafter against the config it was trained with" + ); + } + } + for i in 0..parts { + let part = format!("fc.{i}.weight"); + splits.insert(part.clone(), (n.clone(), i, parts)); + rewritten.push(part); + } + } else { + rewritten.push(n.clone()); + } + } + Ok(FcSplitProvider { + inner, + splits, + rewritten, + }) + } + fn rewritten_names(&self) -> Vec { + self.rewritten.clone() + } + fn part_shape(&self, src: &str, parts: usize) -> Result<(usize, usize)> { + let shape = self.inner.source_shape(src)?; + if shape.len() != 2 || shape[1] as usize % parts != 0 { + bail!("{src}: expected a [dim, {parts}*dim] fc, got {shape:?}"); + } + Ok((shape[0] as usize, shape[1] as usize / parts)) + } +} +impl TensorProvider for FcSplitProvider<'_> { + fn source_shape(&self, name: &str) -> Result> { + match self.splits.get(name) { + Some((src, _, parts)) => { + let (rows, cols) = self.part_shape(src, *parts)?; + Ok(vec![rows as u64, cols as u64]) + } + None => self.inner.source_shape(name), + } + } + fn to_f32(&self, name: &str) -> Result> { + let Some((src, which, parts)) = self.splits.get(name) else { + return self.inner.to_f32(name); + }; + let (rows, cols) = self.part_shape(src, *parts)?; + let data = self.inner.to_f32(src)?; + if data.len() != rows * parts * cols { + bail!( + "{src}: {} elements, expected {}", + data.len(), + rows * parts * cols + ); + } + let mut out = Vec::with_capacity(rows * cols); + for r in 0..rows { + let start = r * parts * cols + which * cols; + out.extend_from_slice(&data[start..start + cols]); + } + Ok(out) + } + fn packed_base_q4(&self, name: &str, group_size: u32) -> Result> { + match self.splits.get(name) { + Some(_) => Ok(None), + None => self.inner.packed_base_q4(name, group_size), + } + } + fn packed_nvfp4(&self, name: &str) -> Result> { + match self.splits.get(name) { + Some(_) => Ok(None), + None => self.inner.packed_nvfp4(name), + } + } + fn source_is_quantized(&self, name: &str) -> bool { + match self.splits.get(name) { + Some(_) => false, + None => self.inner.source_is_quantized(name), + } + } +} + +/// Tensors the runtime reads at f16 rather than through a quant GEMM +/// route, so a quantized source must dequantize them rather than hand +/// them to the quantized transplant: the MoE router (`mlp.router` / +/// HF `mlp.gate`, canonical names the runtime aliases to +/// `ffn_gate_inp`), the Qwen3.5-MoE per-token shared-expert gate +/// (`ffn_shexp_gate`, a [1, dim] row) and the GLM MLA absorb +/// projections. `mlp.gate_proj` (a dense FFN) does not match. +fn f16_consumer(canonical: &str) -> bool { + canonical.ends_with(".mlp.router.weight") + || canonical.ends_with(".mlp.gate.weight") + || canonical.ends_with(".mlp.shared_expert_gate.weight") + || canonical.ends_with(".k_b_proj.weight") + || canonical.ends_with(".v_b_proj.weight") +} + +/// Whether a tensor the source stores quantized is dequantized and +/// stored as f16 rather than handed to the quantized transplant: an +/// `f16_consumer` from an MLX source, whose affine tensors read back to +/// f32 exactly. A modelopt NVFP4 source is quantized too, but its packed +/// e2m1 codes have no f32 reader (`HfDir::tensor_to_f32` rejects U8), so +/// forcing one of its tensors to f16 failed the whole conversion; its +/// quantized tensors stay on the transplant. +fn dequant_to_f16(source_format: &str, src_quantized: bool, canonical: &str) -> bool { + src_quantized && source_format == "mlx_safetensors" && f16_consumer(canonical) +} + #[allow(clippy::too_many_arguments)] fn convert_generic( input: &std::path::Path, @@ -3299,6 +3699,38 @@ fn convert_generic( let source_names = kvb.rewritten_names(); let provider: &dyn TensorProvider = &kvb; + // Fused input projections split by input column so the runtime feeds + // each input through a plain GEMM: the Qwen3.5 MTP head's fc (two + // halves), a block drafter's fc (one block per target tap). + // How many tap blocks the config says the drafter has: EAGLE-3 is always + // three, a block drafter names them in `target_layer_ids`. + let expected_taps: Option = if base_arch::dflash::is_eagle3_arch(canonical_arch) { + Some(3) + } else if base_arch::dflash::is_block_drafter_arch(canonical_arch) { + config + .extra_config + .get("target_layer_ids") + .and_then(|v| v.as_array()) + .map(|a| a.len()) + } else { + None + }; + let fcs = FcSplitProvider::build(provider, &source_names, canonical_arch, expected_taps)?; + if !fcs.splits.is_empty() { + if base_arch::dflash::is_block_drafter_arch(canonical_arch) + || base_arch::dflash::is_eagle3_arch(canonical_arch) + { + eprintln!( + " drafter: split fused projections into {} column blocks (f16)", + fcs.splits.len() + ); + } else { + eprintln!(" mtp: split mtp.fc.weight into fc_embed / fc_hidden (f16)"); + } + } + let source_names = fcs.rewritten_names(); + let provider: &dyn TensorProvider = &fcs; + // Map source names → canonical. Use the same llama-style map that // GGUF uses for blk.N.* tensors, plus a HF-style pass-through for // `model.layers.N.*` already-canonical names. Multimodal towers @@ -3306,20 +3738,73 @@ fn convert_generic( // whether to materialize them. let mut mapped: Vec<(String, String)> = Vec::new(); let mut mmproj_mapped: Vec<(String, String)> = Vec::new(); + let mut speculator_mapped: Vec<(String, String)> = Vec::new(); let mut dropped_names: Vec = Vec::new(); + // A Gemma 4-backbone block drafter's decoder layers take the Gemma 4 + // tensor names (four norms, layer_scalar, no v_proj on K=V layers); + // its fc / hidden_norm / heads keep the block-drafter names. + let gemma4_drafter = base_arch::dflash::is_block_drafter_arch(canonical_arch) + && config + .extra_config + .get("drafter_backbone") + .and_then(|v| v.as_str()) + == Some("gemma4"); for n in &source_names { - match to_canonical_name(n, canonical_arch) { + let canon = if gemma4_drafter && n.starts_with("layers.") { + to_canonical_name(&format!("model.{n}"), "gemma4") + } else { + to_canonical_name(n, canonical_arch) + }; + match canon { Some(Canonical::Main(c)) => mapped.push((n.clone(), c)), Some(Canonical::Mmproj(c)) => mmproj_mapped.push((n.clone(), c)), + Some(Canonical::Speculator(c)) => speculator_mapped.push((n.clone(), c)), None => dropped_names.push(n.clone()), } } + // A llama-labelled block drafter is accepted by its config only because + // RedHat's Qwen-shaped heads carry that label; the runtime needs q/k + // norms on every layer, so check the tensors before writing a sidecar + // that could never load. + if base_arch::dflash::is_block_drafter_arch(canonical_arch) { + base_arch::dflash::check_block_drafter_qk_norms( + &config, + mapped.iter().map(|(_, c)| c.as_str()), + )?; + } + let plan = plan_tied_lm_head( + &mapped, + config.tie_word_embeddings, + canonical_arch, + ctx.synth_tied_head, + ); + report_tied_head_plan(&plan, ctx, canonical_arch, TiedHeadCopy::Quantized); + if let TiedHeadPlan::Synthesize { src } = plan { + mapped.push((src, "lm_head.weight".to_string())); + } + // --speculator-lm-head: the MTP head's own lm_head, from the same + // source rows as the target's (the runtime's speculator view resolves + // `mtp.lm_head.weight` before falling through to the shared head). + if ctx.speculator_lm_head && !speculator_mapped.is_empty() { + if let Some((src, _)) = mapped.iter().find(|(_, c)| c == "lm_head.weight") { + eprintln!( + " mtp: own lm_head for the speculator head (mtp.lm_head.weight from {src})" + ); + speculator_mapped.push((src.clone(), "mtp.lm_head.weight".to_string())); + } else { + eprintln!(" mtp: --speculator-lm-head: no lm_head source found, head keeps the shared one"); + } + } eprintln!( - " mapped: {} tensors kept, {} mmproj, {} dropped", + " mapped: {} tensors kept, {} mmproj, {} speculator, {} dropped", mapped.len(), mmproj_mapped.len(), + speculator_mapped.len(), dropped_names.len() ); + if ctx.mtp_draft_ranking.is_some() && speculator_mapped.is_empty() { + bail!("--mtp-draft-ranking: this checkpoint has no MTP speculator head to draft with"); + } let mut tok_fields = tokenizer_fields.clone(); if !tok_fields.contains_key("tokenizer_type") { @@ -3334,6 +3819,37 @@ fn convert_generic( // for the late layers — model output is garbage. Auto-derive from the // actual MLX/HF tensor shapes the source ships. let mut config = config; + // DeepSpec's DSpark trains with the target's embedding and head COPIED + // in and frozen: a sidecar whose embed_tokens and lm_head are the same + // bytes carries two copies of a tied target's embedding (Gemma 4 12B: + // 2 x 2 GB f16, the head streamed on every block step). Drop both and + // draft through the target's own (quantized, already resident) table — + // the tensors resolve through the sidecar's fallback to the target. + if base_arch::dflash::is_block_drafter_arch(canonical_arch) + && config + .extra_config + .get("speculator_kind") + .and_then(|v| v.as_str()) + == Some("dspark") + && !config.extra_config.contains_key("drafter_head") + { + let src_of = |c: &str| mapped.iter().find(|(_, m)| m == c).map(|(s, _)| s.clone()); + if let (Some(e), Some(h)) = (src_of("embed_tokens.weight"), src_of("lm_head.weight")) { + if provider.source_shape(&e)? == provider.source_shape(&h)? + && provider.to_f32(&e)? == provider.to_f32(&h)? + { + mapped.retain(|(_, c)| c != "embed_tokens.weight" && c != "lm_head.weight"); + config + .extra_config + .insert("drafter_head".into(), serde_json::json!("target")); + config.tie_word_embeddings = true; + eprintln!( + " drafter: embed_tokens == lm_head (the target's frozen tied embedding) — \ + both dropped, drafting through the target's table" + ); + } + } + } if config.per_layer_ffn.is_empty() && config.num_hidden_layers > 0 { let mut per_layer_ffn = vec![0u32; config.num_hidden_layers as usize]; let mut found_any = false; @@ -3422,6 +3938,7 @@ fn convert_generic( layers: layer_descriptors_from_config(&config), tensors: vec![], mmproj: None, + speculator: None, calibration: None, provenance: None, sig: None, @@ -3430,6 +3947,8 @@ fn convert_generic( if config.tie_word_embeddings { header.flags |= HeaderFlags::TIED_EMBEDDINGS; } + // (HAS_SPECULATOR is raised by the writer when the speculator section + // is installed.) let has_moe = source_names .iter() .any(|n| n.contains("experts") || n.contains("_exps") || n.contains("_shexp")); @@ -3469,7 +3988,80 @@ fn convert_generic( let mut n_carried = 0usize; let mut prov_tensors: std::collections::BTreeMap = Default::default(); + // Speculator tensors ride the same per-tensor pipeline (norm shift, + // stacked experts, provenance) after every LM tensor, but stay f16 and + // land in the speculator sub-bundle: the head's acceptance rate is what + // the whole scheme buys, and a quantized MTP head measured near-zero + // acceptance on MoE targets. + let spec_names: std::collections::HashSet = + speculator_mapped.iter().map(|(_, c)| c.clone()).collect(); + if !speculator_mapped.is_empty() { + writer.set_speculator_arch(format!("{canonical_arch}-mtp")); + let mut n_spec_layers = 0u64; + for (_, c) in &speculator_mapped { + if let Some(r) = c.strip_prefix("mtp.layers.") { + if let Some((i, _)) = r.split_once('.') { + if let Ok(i) = i.parse::() { + n_spec_layers = n_spec_layers.max(i + 1); + } + } + } + } + let spec_moe = speculator_mapped + .iter() + .any(|(_, c)| c.contains("_exps") || c.contains(".experts.")); + let mut spec_cfg = std::collections::BTreeMap::new(); + spec_cfg.insert("kind".to_string(), serde_json::json!("mtp")); + spec_cfg.insert("num_layers".to_string(), serde_json::json!(n_spec_layers)); + spec_cfg.insert("moe".to_string(), serde_json::json!(spec_moe)); + spec_cfg.insert("shared_embeddings".to_string(), serde_json::json!(true)); + // `--speculator-lm-head` gave the head its own `mtp.lm_head.weight`: + // say so, or a reader of this field falls back to the target's head. + let private_lm_head = speculator_mapped + .iter() + .any(|(_, c)| c == "mtp.lm_head.weight"); + spec_cfg.insert( + "shared_lm_head".to_string(), + serde_json::json!(!private_lm_head), + ); + spec_cfg.insert( + "fc_inputs".to_string(), + serde_json::json!(["embedding", "hidden"]), + ); + if let (Some(_), Some(k)) = (&ctx.mtp_draft_ranking, ctx.mtp_draft_vocab) { + spec_cfg.insert("draft_vocab_size".to_string(), serde_json::json!(k)); + } + writer.set_speculator_config(spec_cfg); + pb.set_length((mapped.len() + speculator_mapped.len()) as u64); + } + let mapped: Vec<(String, String)> = mapped + .iter() + .cloned() + .chain(speculator_mapped.iter().cloned()) + .collect(); + // Block-drafter sidecars stay f16 end to end (acceptance rate is what + // the scheme buys; the block step is small next to the target's verify). + let is_drafter_arch = base_arch::dflash::is_block_drafter_arch(canonical_arch) + || base_arch::dflash::is_eagle3_arch(canonical_arch); for (src_name, canonical) in &mapped { + // --quantize-drafter: the drafter's transformer layers follow + // --target; everything else in the sidecar (fc taps, norms, heads, + // d2t) stays f16. + // DFlash 2's conv `base_kernel` ([2, K, D]) is read as raw halves, + // not a projection — it stays f16 (its kernel_projection quantizes). + let drafter_f16 = is_drafter_arch + && !(ctx.quantize_drafter + && canonical.starts_with("layers.") + && !canonical.ends_with(".base_kernel")); + let is_spec = spec_names.contains(canonical); + // --quantize-drafter also quantizes an in-bundle speculator's transformer + // layer (`mtp.layers.*`): a Qwen3.8-27B MTP head is ~700 MB f16 and + // streams on every chained draft; its fc / norms stay f16. The head's + // own lm_head (--speculator-lm-head) is quantized too — the whole + // point of the copy is a cheaper head stream than the target's. + let spec_quantized = ctx.quantize_drafter + && is_spec + && (canonical.starts_with("mtp.layers.") || canonical == "mtp.lm_head.weight"); pb.set_message(canonical.clone()); let src_shape = provider.source_shape(src_name)?; // The reduction dim is the source's last (HF C-order [out, in]); @@ -3514,18 +4106,32 @@ fn convert_generic( // NVFP4 sidecars: the per-tensor global scale and the calibrated // activation scale. Always f32 — they parameterize the dequant // itself, and the verify gate compares them value-exactly. - let is_sidecar = - canonical.ends_with(".weight_scale_2") || canonical.ends_with(".input_scale"); + // (EAGLE-3's `d2t` offset table rides the same f32 path: token ids + // are exact in f32, not in f16.) + let is_sidecar = canonical.ends_with(".weight_scale_2") + || canonical.ends_with(".input_scale") + || canonical == "d2t"; // See GGUF path above for rationale on the f16/GPU route. Qwen3.5 GDN // conv1d + per-head gate projections also route here (f16 on GPU). let is_norm_like = !is_sidecar && (shape.len() == 1 || is_gdn_f16(canonical)); let is_embedding = canonical == "embed_tokens.weight" || canonical == "lm_head.weight"; // GLM MLA k_b/v_b projections and router weights are raw-f16 // consumers and must never enter the quantized transplant path. - let force_f16 = canonical_arch == "glm_dsa" - && (canonical.ends_with(".k_b_proj.weight") - || canonical.ends_with(".v_b_proj.weight") - || canonical.ends_with(".mlp.router.weight")); + // Under the mirror policy the same holds for every tensor the + // runtime reads at f16 that a quantized source keeps at a scheme + // the transplant cannot carry: mlx-community stores the Qwen3.6-MoE + // router and shared-expert gate at 8-bit inside a 4-bit checkpoint, + // and requantizing them to base-q4 is neither bit-identical nor + // what `--mlx-passthrough` promises — they dequantize exactly and + // store as f16 (BAS-960). See `dequant_to_f16` for why MLX only. + let src_quantized = mirror && provider.source_is_quantized(src_name); + let force_f16 = (is_spec && !spec_quantized) + || drafter_f16 + || (canonical_arch == "glm_dsa" + && (canonical.ends_with(".k_b_proj.weight") + || canonical.ends_with(".v_b_proj.weight") + || canonical.ends_with(".mlp.router.weight"))) + || dequant_to_f16(source_format, src_quantized, canonical); // Mirror policy: this tensor is unquantized in the (quantized) // source checkpoint, so it must stay unquantized in the bundle. let src_unquantized = mirror && !provider.source_is_quantized(src_name); @@ -3907,7 +4513,11 @@ fn convert_generic( (entry, data) }; - writer.add_tensor(TensorPayload { entry, data }); + if is_spec { + writer.add_speculator_tensor(TensorPayload { entry, data }); + } else { + writer.add_tensor(TensorPayload { entry, data }); + } // Record where this tensor came from and what was done to it, so // verification tooling can check the bundle against the source @@ -3928,6 +4538,13 @@ fn convert_generic( if let Some(spec) = splitting.splits.get(src_name) { rec.src = vec![spec.src.clone()]; rec.rows = Some([spec.row_off, spec.row_cnt]); + } else if let Some((fc_src, block, parts)) = fcs.splits.get(src_name) { + // A column block of a fused projection: the checkpoint holds + // only the fused source (the split name is virtual), so record + // the real source and which input columns this block took. + let (_, width) = fcs.part_shape(fc_src, *parts)?; + rec.src = vec![fc_src.clone()]; + rec.cols = Some([(*block * width) as u64, width as u64]); } else if let Some(parts) = stacking.stacks.get(src_name) { rec.stack = Some(base_format::StackRef { pattern: src_name.replace(".experts.", ".experts.{e}."), @@ -3940,6 +4557,49 @@ fn convert_generic( pb.inc(1); } pb.finish_and_clear(); + // --mtp-draft-ranking: the ranked ids as `mtp.draft_rank` (f32 — ids are + // exact in f32, the same encoding as an EAGLE-3 d2t table), after every + // other speculator tensor. + if let Some(ranked) = &ctx.mtp_draft_ranking { + let vocab = config.vocab_size as u64; + if let Some(bad) = ranked.iter().find(|&&id| id as u64 >= vocab) { + bail!("--mtp-draft-ranking: token {bad} is outside the {vocab}-token vocabulary"); + } + let data: Vec = ranked + .iter() + .flat_map(|&id| (id as f32).to_le_bytes()) + .collect(); + let entry = base_format::TensorEntry { + name: "mtp.draft_rank".to_string(), + dtype: TensorDtype::F32, + shape: vec![ranked.len() as u64], + offset: 0, + length: data.len() as u64, + scale_offset: None, + scale_length: None, + bias_offset: None, + bias_length: None, + awq_scale_offset: None, + awq_scale_length: None, + group_size: None, + layout: None, + residency: Some(base_format::ResidencyHint::Cold), + compute_region: ComputeRegion::Cpu, + scale_dtype: None, + symmetric: false, + flags: TensorFlags::empty(), + checksum_xxh64: None, + source_ggml_type: None, + }; + writer.add_speculator_tensor(TensorPayload { entry, data }); + eprintln!( + " mtp: drafting ranking of {} ids (mtp.draft_rank){}", + ranked.len(), + ctx.mtp_draft_vocab + .map(|k| format!(", default K {k}")) + .unwrap_or_default() + ); + } if n_passthrough > 0 || n_carried > 0 { eprintln!( " quantized {} tensors ({} transplanted from the source's identical \ @@ -3952,11 +4612,29 @@ fn convert_generic( } else { eprintln!(" quantized {} tensors", mapped.len()); } + // The source-tensor lists name checkpoint tensors: a virtual fc-split + // name (mtp.fc_embed.weight, fc.0.weight, ...) maps back to the fused + // source it came from, once. + let real_sources = |names: &[(String, String)]| -> Vec { + let mut out: Vec = Vec::new(); + for (s, _) in names { + let real = fcs + .splits + .get(s) + .map(|(src, _, _)| src.clone()) + .unwrap_or_else(|| s.clone()); + if !out.contains(&real) { + out.push(real); + } + } + out + }; writer.set_provenance(base_format::Provenance { schema: 1, mirror, dropped: dropped_names, - mmproj: mmproj_mapped.iter().map(|(s, _)| s.clone()).collect(), + mmproj: real_sources(&mmproj_mapped), + speculator: real_sources(&speculator_mapped), tensors: prov_tensors, }); @@ -4150,6 +4828,216 @@ fn convert_generic( Ok(()) } +/// Decide whether to materialize a standalone `lm_head.weight` for a +/// tie_word_embeddings model, returning the (source, canonical) pair to append +/// to the tensor work-list. +/// +/// A tied checkpoint ships one matrix, used both to look up token embeddings +/// and — transposed — as the output projection. The quant profiles pin +/// `embed_tokens.weight` to f16 because a quantized lookup table costs quality +/// for no speed (a lookup touches one row per token). But under tying that same +/// f16 rule silently sets the precision of the LOGIT PROJECTION, which streams +/// the WHOLE table on every decode step: on Llama-3.2-1B that is 128256x2048 f16 +/// = 525 MB/token, measured at 34% of decode GPU time in a rocprofv3 trace even +/// though the kernel itself runs at 189 GB/s (~73% of a Radeon 8060S's peak). +/// +/// Emitting a second, quantized copy under `lm_head.weight` lets the profile's +/// existing `lm_head.weight` rule apply to it. The runtimes that consume it +/// (see `runtime_reads_output_weight`) prefer `output.weight` — the runtime +/// name of `lm_head.weight` — whenever it resolves and fall back to the +/// embedding otherwise, so the decode path picks up the quantized copy while +/// the embedding lookup keeps its f16 fidelity. Costs disk: +148 MB at base_q4 +/// for a 128k x 2048 vocab. +/// +/// Nothing is planned when the model is untied, or when the source already +/// carries an explicit head — some repackaged repos untie and ship +/// `lm_head.weight`, and duplicating it would write the tensor twice. +/// `RuntimeIgnoresHead` covers the archs whose runtime never looks the head up +/// (a synthesized copy there is dead weight), and `OptedOut` is +/// `--no-synth-tied-head`. Pure: the convert-time line is printed by +/// `report_tied_head_plan` at the call site, which knows how the copy will be +/// packed. +fn plan_tied_lm_head( + mapped: &[(String, String)], + tied: bool, + canonical_arch: &str, + synth_enabled: bool, +) -> TiedHeadPlan { + if !tied { + return TiedHeadPlan::Nothing; + } + if mapped + .iter() + .any(|(_, c)| c == "lm_head.weight" || c == "output.weight") + { + return TiedHeadPlan::Nothing; + } + if !runtime_reads_output_weight(canonical_arch) { + return TiedHeadPlan::RuntimeIgnoresHead; + } + if !synth_enabled { + return TiedHeadPlan::OptedOut; + } + match mapped.iter().find(|(_, c)| c == "embed_tokens.weight") { + Some((src, _)) => TiedHeadPlan::Synthesize { src: src.clone() }, + None => TiedHeadPlan::Nothing, + } +} + +/// Outcome of `plan_tied_lm_head`. +#[derive(Debug, Clone, PartialEq, Eq)] +enum TiedHeadPlan { + /// Append `lm_head.weight` as a second copy of this source tensor. + Synthesize { src: String }, + /// Tied, but this arch's runtime never resolves `output.weight`. + RuntimeIgnoresHead, + /// Tied, but `--no-synth-tied-head` was passed. + OptedOut, + /// Untied, the source already ships a head, or it has no embedding. + Nothing, +} + +/// Whether the runtime that will execute `canonical_arch` resolves its logit +/// projection through `output.weight` — the runtime name of `lm_head.weight` +/// (`kBaseTopRules` in src/weight/base_weight_store.cpp). Only those archs +/// can consume a synthesized head; for the others it is bytes on disk the +/// runtime never touches. +/// +/// Derived from the per-model prefill encoders, which are the only readers of +/// the logit tensor (prefill.cpp dispatches on descriptor flags and falls +/// through to `encode_prefill_llama`): +/// +/// - `llama.cpp` (llama, qwen2/3, qwen MoE, mistral, phi — the fallthrough) +/// and `qwen3_5.cpp` (qwen35, qwen35moe) call +/// `weights.get_tensor("output.weight")` FIRST and only fall back to +/// `token_embedding.weight` when it is absent. Consumers. +/// - `muse_glimmer.cpp` keys on `arch.tie_embeddings`, which +/// `arch_muse_glimmer()` hardcodes false: it reads `output.weight` +/// unconditionally. Consumer (a tied Muse checkpoint would NEED the copy). +/// - `gemma.cpp` and `gemma4.cpp` key on `arch.tie_embeddings`, which +/// `arch_gemma()` / `arch_gemma4()` hardcode TRUE +/// (src/core/arch_descriptor.h): the logit name is always +/// `token_embedding.weight` and `output.weight` is never looked up. A +/// synthesized head is dead weight — ~335 MB at base_q4 for gemma-3-4b's +/// 262k-row vocab. Not consumers. +/// - `bert.cpp` has no logit projection; whisper converts on its own path +/// (`convert_whisper`) and never reaches `plan_tied_lm_head`. Not consumers. +fn runtime_reads_output_weight(canonical_arch: &str) -> bool { + match canonical_arch { + "llama" | "qwen" | "qwen_moe" | "qwen3_moe" | "qwen35" | "qwen35moe" | "muse_glimmer" => { + true + } + "gemma3" | "gemma4" | "nomic-bert" | "whisper" => false, + // The runtime's arch_from_config matches these families by + // substring, so a new "gemma*"/"*bert*" mapper lands on the same + // encoders as above; anything else falls through to the llama + // encoder, which reads output.weight. + other => !other.starts_with("gemma") && !other.contains("bert"), + } +} + +/// GGUF carries no `tie_word_embeddings` key, so `config_from_gguf` leaves it +/// false and the header would claim an untied model (and drop the +/// TIED_EMBEDDINGS flag) for e.g. Llama-3.2, which IS tied. Infer it the way +/// llama.cpp does: a causal LM must have an output projection, so a source +/// with no head tensor is necessarily reusing its embedding as one. Skipped +/// for embedding-only archs (BERT), which legitimately have no head and must +/// not get a synthesized one. A config that already says tied stays tied. +fn infer_gguf_tied_embeddings<'a>( + mapper: &dyn base_arch::GgufMapper, + config_tied: bool, + tensor_names: impl Iterator, +) -> bool { + if config_tied || mapper.canonical_arch().contains("bert") { + return config_tied; + } + let has_head = tensor_names.into_iter().any(|n| { + matches!( + mapper.map_tensor_name(n).as_deref(), + Some("lm_head.weight") | Some("output.weight") + ) + }); + !has_head +} + +/// How a synthesized head will physically be written, for the convert-time +/// line: dequantized and re-packed by the profile / `--target` like any other +/// tensor, or (GGUF `--kquant-passthrough` on a k-quant embedding) copied +/// verbatim, in which case no profile rule touches it. +enum TiedHeadCopy { + Quantized, + Verbatim { ggml_type: &'static str }, +} + +/// The first profile rule matching `name` — what `QuantProfile::resolve` +/// applies, with the pattern kept so the decision can be shown. +fn profile_rule_for<'a>( + profile: &'a base_quant::QuantProfile, + name: &str, +) -> Option<&'a base_quant::RuleEntry> { + profile + .rules + .iter() + .find(|r| base_quant::profile::pattern_matches(&r.pattern, name)) +} + +/// One line saying what happened to the tied head and why, so a precision +/// change on reconvert is visible in the log rather than only in the bundle. +fn report_tied_head_plan( + plan: &TiedHeadPlan, + ctx: &QuantContext, + canonical_arch: &str, + copy: TiedHeadCopy, +) { + let describe = |name: &str| -> String { + match &ctx.profile { + Some(p) => match profile_rule_for(p, name) { + Some(r) => { + let resolved = p.resolve(name).expect("rule matched"); + // gs is meaningless for the float dtypes (resolves to 1). + let gs = if resolved.group_size > 1 { + format!(" gs={}", resolved.group_size) + } else { + String::new() + }; + format!( + "profile {:?} rule `{}` -> {:?}{gs}", + p.name, r.pattern, resolved.dtype + ) + } + None => format!("profile {:?} has NO rule for {name}", p.name), + }, + None => format!("--target {:?}", ctx.target), + } + }; + match plan { + TiedHeadPlan::Synthesize { src } => match copy { + TiedHeadCopy::Quantized => eprintln!( + " tied: materializing lm_head.weight from {src}: {}; \ + embed_tokens.weight keeps {}. Logits now project at the head's \ + precision (--no-synth-tied-head keeps them at the embedding's).", + describe("lm_head.weight"), + describe("embed_tokens.weight"), + ), + TiedHeadCopy::Verbatim { ggml_type } => eprintln!( + " tied: materializing lm_head.weight from {src} as a VERBATIM \ + {ggml_type} copy (--kquant-passthrough): no profile rule applies, \ + so logits project at the embedding's own precision." + ), + }, + TiedHeadPlan::RuntimeIgnoresHead => eprintln!( + " tied: no lm_head.weight synthesized — the {canonical_arch} runtime \ + projects logits through embed_tokens.weight and never reads output.weight." + ), + TiedHeadPlan::OptedOut => eprintln!( + " tied: --no-synth-tied-head: logits project through embed_tokens.weight \ + at its own precision ({}).", + describe("embed_tokens.weight"), + ), + TiedHeadPlan::Nothing => {} + } +} + /// Routing decision for an HF/MLX tensor name. #[derive(Debug, Clone)] enum Canonical { @@ -4160,6 +5048,91 @@ enum Canonical { /// for now we keep the HF prefix verbatim so the future runtime can /// dispatch on the canonical multimodal layout. Mmproj(String), + /// Goes into the speculator sub-bundle (header.speculator.tensors): + /// the checkpoint's own MTP head, named `mtp.`. + Speculator(String), +} + +/// EAGLE-3 head sidecar tensor → canonical name: `midlayer.*` is the one +/// decoder layer (`layers.0.*` through the llama main-layer rules; the +/// split q/k/v parts and `hidden_norm` pass through under it), `norm` is +/// the head's output norm, `d2t` the draft→target offset table (f32), `t2d` +/// is derivable and dropped; `fc.weight` is split into `fc.{i}.weight` +/// before mapping; an `embed_tokens` copy and the reduced `lm_head` keep +/// their names. +fn eagle3_canonical(name: &str) -> Option { + if name == "fc.weight" || name == "t2d" { + return None; + } + // SpecForge / AngelSlim call the one layer `midlayer`; RedHat's + // speculators packaging calls it `layers.0`. + if let Some(rest) = name + .strip_prefix("midlayer.") + .or_else(|| name.strip_prefix("layers.0.")) + { + return match to_canonical_name(&format!("model.layers.0.{rest}"), "llama")? { + Canonical::Main(c) => Some(c), + _ => None, + }; + } + Some(match name { + "norm.weight" => "final_norm.weight".to_string(), + // norm_before_fc's weight over the concatenated taps [n_taps * H] + // (RedHat's gpt-oss heads); a distinct name so no per-layer + // `input_norm` alias can claim it. + "input_norm.weight" => "fc_input_norm.weight".to_string(), + other => other.to_string(), + }) +} + +/// Block-drafter (DFlash / DSpark) sidecar tensor → canonical name. The +/// decoder layer goes through the qwen3 main-layer rules (so the runtime's +/// per-layer aliases apply as for the target's own layers); `norm.weight` +/// is the drafter's output norm (`final_norm.weight`, like an HF LM's); +/// DSpark's Markov / confidence heads get short canonical names; the fc +/// per-tap halves, hidden_norm, an own embed_tokens / lm_head, and DFlash +/// 2's conv / selector tensors keep their names. The fused `fc.weight` is +/// never written: FcSplitProvider splits it into `fc.{i}.weight` first. +fn block_drafter_canonical(name: &str) -> Option { + // `t2d` (a reduced-vocabulary head's target→draft mask) is derivable + // from `d2t`, which keeps its name. + if name == "fc.weight" || name == "t2d" { + return None; + } + if let Some(rest) = name.strip_prefix("layers.") { + return match to_canonical_name(&format!("model.layers.{rest}"), "qwen")? { + Canonical::Main(c) => Some(c), + _ => None, + }; + } + Some(match name { + "norm.weight" => "final_norm.weight".to_string(), + "markov_head.markov_w1.weight" => "markov.w1.weight".to_string(), + "markov_head.markov_w2.weight" => "markov.w2.weight".to_string(), + "confidence_head.proj.weight" => "confidence.proj.weight".to_string(), + "confidence_head.proj.bias" => "confidence.proj.bias".to_string(), + other => other.to_string(), + }) +} + +/// Qwen3.5/3.6/3.8 MTP head tensor → speculator name. The head is one +/// decoder layer of the target's geometry (`mtp.layers.0.*`) plus `fc`, +/// the two pre-fc norms and its own output norm. The decoder layer goes +/// through the main-layer rules (HF renames, expert stacking canonical +/// forms) so the runtime's per-layer aliases apply under the `mtp.` prefix +/// unchanged; the head's own tensors keep their HF names under the prefix +/// (`mtp.norm.weight` stays — it is the head's output norm, not the model's +/// final norm). +fn mtp_speculator_name(rest: &str, arch: &str) -> Option { + let inner = if rest.starts_with("layers.") { + match to_canonical_name(&format!("model.{rest}"), arch)? { + Canonical::Main(c) => c, + _ => return None, + } + } else { + rest.to_string() + }; + Some(Canonical::Speculator(format!("mtp.{inner}"))) } /// nomic-bert HF safetensors → canonical `.base` names. Targets the @@ -4233,13 +5206,27 @@ fn is_gdn_f16(canonical: &str) -> bool { /// mostly strip the `model.` prefix. For anything not in HF convention, /// fall back to the llama-style GGUF mapper. fn to_canonical_name(name: &str, arch: &str) -> Option { - // Qwen3.5/3.6 checkpoints ship a Multi-Token-Prediction head (`mtp.*` — - // one extra decoder layer + fc/norms, its own experts on the MoE - // variants). It's a speculative-decoding auxiliary the runtime doesn't - // execute; keeping it would bloat the .base (256 extra expert stacks on - // 35B-A3B) and land unknown-named tensors in the main bundle. Drop it. - if arch.starts_with("qwen35") && (name == "mtp" || name.starts_with("mtp.")) { - return None; + // Block-drafter sidecars (DFlash / DSpark): bare HF names (no `model.` + // wrapper), a qwen3 decoder under `layers.N.*`, and the drafter's own + // fc / hidden_norm / heads. Everything is a main tensor of the sidecar. + if base_arch::dflash::is_block_drafter_arch(arch) { + return block_drafter_canonical(name).map(Canonical::Main); + } + if base_arch::dflash::is_eagle3_arch(arch) { + return eagle3_canonical(name).map(Canonical::Main); + } + // Qwen3.5/3.6/3.8 checkpoints ship a Multi-Token-Prediction head (`mtp.*` + // — one extra decoder layer + fc / pre-fc norms / output norm, its own + // experts on the MoE variants). It is the bundle's speculator: routed + // into the speculator sub-bundle (kept f16, skipped by a runtime that + // does not speculate) under `mtp.`. + if arch.starts_with("qwen35") { + if name == "mtp" { + return None; + } + if let Some(rest) = name.strip_prefix("mtp.") { + return mtp_speculator_name(rest, arch); + } } // Multimodal wrappers (Gemma-4) put the LLM under `language_model.model.*` // and the audio/vision towers under siblings. The towers ship in @@ -4320,11 +5307,15 @@ fn to_canonical_name(name: &str, arch: &str) -> Option { .unwrap_or(name); // MTP can also arrive wrapped (`model.mtp.*`, or `model.language_model. - // mtp.*` on VL-style checkpoints); the bare-prefix drop at the top of - // this function runs before wrapper stripping and misses those, which - // would canonicalize the dead MTP decoder/expert stacks into the bundle. - if arch.starts_with("qwen35") && (stripped == "mtp" || stripped.starts_with("mtp.")) { - return None; + // mtp.*` on VL-style checkpoints); the bare-prefix route at the top of + // this function runs before wrapper stripping and misses those. + if arch.starts_with("qwen35") { + if stripped == "mtp" { + return None; + } + if let Some(rest) = stripped.strip_prefix("mtp.") { + return mtp_speculator_name(rest, arch); + } } // GLM 5.2 (glm_dsa) MLX/HF checkpoints — dedicated rename table so the @@ -4686,6 +5677,7 @@ fn convert_synthetic_with_ctx(output: &std::path::Path, ctx: &QuantContext) -> R .collect(), tensors: vec![], mmproj: None, + speculator: None, calibration: None, provenance: None, sig: None, @@ -4958,6 +5950,7 @@ fn convert_synthetic(output: &std::path::Path, target: TargetScheme) -> Result<( .collect(), tensors: vec![], mmproj: None, + speculator: None, calibration: None, provenance: None, sig: None, @@ -5185,9 +6178,64 @@ fn pack_bf16(weights: &[f32]) -> base_quant::Packed { /// /// Holds the loaded profile + optional AWQ sidecar; produces /// per-tensor packed bytes for canonical-quant bundles. +/// A drafting-vocabulary ranking JSON (tools/mtp_draft_ranking.py): the +/// `ranked_ids` array, most frequent first. Duplicates are refused — the +/// runtime takes the first K distinct ids, and a repeated id would quietly +/// shrink every K. +fn read_draft_ranking(path: &std::path::Path) -> Result> { + let text = std::fs::read_to_string(path) + .with_context(|| format!("reading --mtp-draft-ranking {}", path.display()))?; + let v: serde_json::Value = serde_json::from_str(&text) + .with_context(|| format!("--mtp-draft-ranking {} is not JSON", path.display()))?; + let arr = v + .get("ranked_ids") + .and_then(|a| a.as_array()) + .with_context(|| { + format!( + "--mtp-draft-ranking {}: no `ranked_ids` array", + path.display() + ) + })?; + let mut ids = Vec::with_capacity(arr.len()); + let mut seen = std::collections::HashSet::with_capacity(arr.len()); + for x in arr { + let id = x + .as_u64() + .filter(|&i| i <= u32::MAX as u64) + .with_context(|| { + format!( + "--mtp-draft-ranking {}: `{x}` is not a token id", + path.display() + ) + })? as u32; + if !seen.insert(id) { + bail!( + "--mtp-draft-ranking {}: token {id} is ranked twice", + path.display() + ); + } + ids.push(id); + } + if ids.is_empty() { + bail!( + "--mtp-draft-ranking {}: `ranked_ids` is empty", + path.display() + ); + } + Ok(ids) +} + struct QuantContext { profile: Option, awq_profile: Option, + /// Tensors AWQ actually rotated; 0 with a sidecar present means the run + /// silently degraded to plain RTN — `awq_shape_skipped` says why. + awq_applied: std::sync::atomic::AtomicUsize, + /// Tensors whose sidecar lookup HIT but was skipped because the absmax + /// length did not match the tensor's in_features. Distinguishes "the + /// names do not correspond to this model" (0 hits) from "the names + /// correspond but the stats were captured at another geometry". + awq_shape_skipped: std::sync::atomic::AtomicUsize, awq_config: base_awq::AwqConfig, /// Fallback target when no profile is set. target: TargetScheme, @@ -5196,8 +6244,19 @@ struct QuantContext { /// Transplant MLX affine-q4 tensors into `base_q4` verbatim instead /// of requantizing them through f32. q4_passthrough: bool, + /// --quantize-drafter (sidecars): drafter `layers.*` follow the target scheme. + quantize_drafter: bool, + /// `--speculator-lm-head`: the MTP head gets its own `mtp.lm_head.weight`. + speculator_lm_head: bool, + /// `--mtp-draft-ranking`: ranked token ids for the MTP head's drafting + /// vocabulary (`mtp.draft_rank`), and `--mtp-draft-vocab`, its default K. + mtp_draft_ranking: Option>, + mtp_draft_vocab: Option, /// Copy GGUF Q4_K/Q5_K/Q6_K super-blocks through verbatim. kquant_passthrough: bool, + /// Materialize `lm_head.weight` for tied models (`plan_tied_lm_head`); + /// false under `--no-synth-tied-head`. + synth_tied_head: bool, /// MLX sources only: reuse the source's packed q4 payloads /// verbatim instead of dequant→requant (see `--mlx-passthrough`). mlx_passthrough: bool, @@ -5249,14 +6308,25 @@ impl QuantContext { if args.imatrix && args.awq_profile.is_none() { bail!("--imatrix requires --awq-profile (the calibration absmax supplies the channel weights)"); } + let mtp_draft_ranking = match &args.mtp_draft_ranking { + Some(p) => Some(read_draft_ranking(p)?), + None => None, + }; Ok(Self { profile, awq_profile, + awq_applied: std::sync::atomic::AtomicUsize::new(0), + awq_shape_skipped: std::sync::atomic::AtomicUsize::new(0), awq_config: base_awq::AwqConfig::default(), target: args.target, allow_quant_from_quant: args.allow_quant_from_quant, q4_passthrough: !args.no_mlx_passthrough, + quantize_drafter: args.quantize_drafter, + speculator_lm_head: args.speculator_lm_head, + mtp_draft_ranking, + mtp_draft_vocab: args.mtp_draft_vocab, kquant_passthrough: args.kquant_passthrough, + synth_tied_head: !args.no_synth_tied_head, mlx_passthrough: args.mlx_passthrough, validate: args.validate, direct_write: args.direct_write, @@ -5443,6 +6513,8 @@ impl QuantContext { // tensor in the `.base` header so the runtime can // recover the original output. eprintln!(" awq: {name} α={:.2} mse={:.4e}", plan.alpha, plan.mse); + self.awq_applied + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); base_awq::awq_apply(weights, in_feat, &plan.scales) } else { eprintln!( @@ -5450,6 +6522,8 @@ impl QuantContext { absmax.len(), in_feat ); + self.awq_shape_skipped + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); weights.to_vec() } } else { @@ -5532,6 +6606,18 @@ fn cmd_inspect(args: InspectArgs) -> Result<()> { println!(" {dtype:<8} {count:4} tensors {bytes:>12} bytes"); } + if let Some(spec) = &h.speculator { + println!( + "speculator: {} ({} tensors)", + spec.arch, + spec.tensors.len() + ); + for (k, v) in &spec.config { + println!(" {k} = {v}"); + } + let spec_bytes: u64 = spec.tensors.iter().map(|t| t.length).sum(); + println!(" bytes = {spec_bytes}"); + } let slots = reader.slots()?; println!("n_slots: {}", slots.len()); for s in &slots { @@ -5545,18 +6631,22 @@ fn cmd_inspect(args: InspectArgs) -> Result<()> { if args.verify_checksums { eprintln!("verifying all tensor checksums..."); - for t in h.tensors.iter() { + // The sub-bundles (mmproj towers, the speculator head) carry + // checksums from the same writer: verify them too. + let mut n = 0usize; + for t in reader.all_entries() { reader .verify_tensor(&t.name) .with_context(|| format!("verifying {:?}", t.name))?; + n += 1; } - eprintln!("all checksums OK"); + eprintln!("all checksums OK ({n} tensors incl. sub-bundles)"); } Ok(()) } /// Minimal "current timestamp" without pulling in chrono. ISO 8601 UTC. -fn chrono_now() -> String { +pub(crate) fn chrono_now() -> String { use std::time::{SystemTime, UNIX_EPOCH}; let secs = SystemTime::now() .duration_since(UNIX_EPOCH) @@ -5572,14 +6662,33 @@ mod canonical_name_tests { fn main_canon(name: &str, arch: &str) -> Option { match to_canonical_name(name, arch)? { Canonical::Main(s) => Some(s), - Canonical::Mmproj(_) => None, + Canonical::Mmproj(_) | Canonical::Speculator(_) => None, + } + } + + #[test] + fn block_drafter_qk_norm_check_sees_the_mapped_names() { + // check_block_drafter_qk_norms looks for these exact canonical names; + // if the drafter mapping renamed them, every sidecar (RedHat's + // llama-labelled Qwen heads included) would be refused. + for n in ["q_norm", "k_norm"] { + let src = format!("layers.3.self_attn.{n}.weight"); + assert_eq!(main_canon(&src, "dflash"), Some(src.clone())); + assert_eq!(main_canon(&src, "dspark"), Some(src)); } } fn mmproj_canon(name: &str, arch: &str) -> Option { match to_canonical_name(name, arch)? { Canonical::Mmproj(s) => Some(s), - Canonical::Main(_) => None, + Canonical::Main(_) | Canonical::Speculator(_) => None, + } + } + + fn spec_canon(name: &str, arch: &str) -> Option { + match to_canonical_name(name, arch)? { + Canonical::Speculator(s) => Some(s), + Canonical::Main(_) | Canonical::Mmproj(_) => None, } } @@ -5718,6 +6827,48 @@ mod canonical_name_tests { // MLX Qwen3-MoE: `model.layers.N.mlp.switch_mlp.X_proj.weight` // (per-expert stack) must canonicalize to `layers.N.ffn_X_exps.weight` // without doubling the `.weight` suffix. + /// A quantized source hands these to the dequant → f16 path rather + /// than the quantized transplant (BAS-960: mlx-community keeps the + /// Qwen3.6-MoE router and shared-expert gate at 8-bit in a 4-bit + /// checkpoint). The dense FFN gate is not a router. + #[test] + fn f16_consumers_are_the_router_shared_gate_and_mla_absorb() { + for n in [ + "layers.3.mlp.router.weight", + "layers.3.mlp.gate.weight", + "layers.3.mlp.shared_expert_gate.weight", + "layers.3.self_attn.k_b_proj.weight", + "layers.3.self_attn.v_b_proj.weight", + ] { + assert!(super::f16_consumer(n), "{n}"); + } + for n in [ + "layers.3.mlp.gate_proj.weight", + "layers.3.ffn_gate_exps.weight", + "layers.3.mlp.shared_expert.gate_proj.weight", + "layers.3.mlp.gate.e_score_correction_bias", + "layers.3.mlp.router.bias", + ] { + assert!(!super::f16_consumer(n), "{n}"); + } + } + + /// Only an MLX source dequantizes them: an NVFP4 source's packed codes + /// have no f32 reader, so its router stays on the transplant. + #[test] + fn only_mlx_sources_dequantize_f16_consumers() { + let router = "layers.3.mlp.gate.weight"; + assert!(super::dequant_to_f16("mlx_safetensors", true, router)); + assert!(!super::dequant_to_f16("nvfp4_safetensors", true, router)); + assert!(!super::dequant_to_f16("hf_safetensors", true, router)); + assert!(!super::dequant_to_f16("mlx_safetensors", false, router)); + assert!(!super::dequant_to_f16( + "mlx_safetensors", + true, + "layers.3.mlp.gate_proj.weight" + )); + } + #[test] fn mlx_qwen3_moe_split_experts_no_double_weight() { assert_eq!( @@ -5867,41 +7018,168 @@ mod canonical_name_tests { } // The shared-expert path (Qwen2-MoE / some Qwen3 variants) uses - // Qwen3.5/3.6 MTP (multi-token-prediction) head tensors must be DROPPED - // (`mtp.*` — an auxiliary decoder layer the runtime never executes; on - // 35B-A3B it carries 256 extra expert stacks that would bloat the .base). + // Block-drafter sidecars: bare HF names, the qwen3 layer through the + // main-layer rules, the drafter's own tensors under short canonical names. #[test] - fn qwen35_mtp_tensors_dropped() { - for n in [ - "mtp.fc.weight", - "mtp.norm.weight", - "mtp.pre_fc_norm_embedding.weight", - "mtp.layers.0.mlp.experts.gate_up_proj", - "mtp.layers.0.self_attn.q_proj.weight", + fn block_drafter_names_are_main_tensors() { + for arch in ["dflash", "dspark"] { + for (src, want) in [ + ("fc.0.weight", "fc.0.weight"), + ("fc.4.weight", "fc.4.weight"), + ("hidden_norm.weight", "hidden_norm.weight"), + ("norm.weight", "final_norm.weight"), + ("embed_tokens.weight", "embed_tokens.weight"), + ("lm_head.weight", "lm_head.weight"), + ("markov_head.markov_w1.weight", "markov.w1.weight"), + ("markov_head.markov_w2.weight", "markov.w2.weight"), + ("confidence_head.proj.bias", "confidence.proj.bias"), + ( + "layers.0.self_attn.q_proj.weight", + "layers.0.self_attn.q_proj.weight", + ), + ( + "layers.0.self_attn.k_norm.weight", + "layers.0.self_attn.k_norm.weight", + ), + ( + "layers.2.input_layernorm.weight", + "layers.2.input_norm.weight", + ), + ( + "layers.2.mlp.gate_proj.weight", + "layers.2.mlp.gate_proj.weight", + ), + // DFlash 2 extras ride along untouched for the runtime that reads them. + ( + "layers.1.attention_conv.kernel_projection.weight", + "layers.1.attention_conv.kernel_projection.weight", + ), + ] { + assert_eq!( + main_canon(src, arch).as_deref(), + Some(want), + "{arch}: {src}" + ); + } + // The fused fc is split before mapping; a leftover is dropped, not written. + assert!(main_canon("fc.weight", arch).is_none()); + } + } + + #[test] + fn eagle3_names_are_main_tensors() { + for (src, want) in [ + ("fc.0.weight", "fc.0.weight"), + ("fc.2.weight", "fc.2.weight"), + ("d2t", "d2t"), + ("norm.weight", "final_norm.weight"), + ("lm_head.weight", "lm_head.weight"), + ("embed_tokens.weight", "embed_tokens.weight"), + ("midlayer.hidden_norm.weight", "layers.0.hidden_norm.weight"), + ( + "midlayer.input_layernorm.weight", + "layers.0.input_norm.weight", + ), + ( + "midlayer.self_attn.q_proj_emb.weight", + "layers.0.self_attn.q_proj_emb.weight", + ), + ( + "midlayer.self_attn.k_proj_hid.weight", + "layers.0.self_attn.k_proj_hid.weight", + ), + ( + "midlayer.self_attn.o_proj.weight", + "layers.0.self_attn.o_proj.weight", + ), + ( + "midlayer.mlp.gate_proj.weight", + "layers.0.mlp.gate_proj.weight", + ), + ] { + assert_eq!( + main_canon(src, "eagle3").as_deref(), + Some(want), + "eagle3: {src}" + ); + } + assert!(main_canon("t2d", "eagle3").is_none(), "t2d is derivable"); + assert!( + main_canon("fc.weight", "eagle3").is_none(), + "the fused fc is split first" + ); + } + + // Qwen3.5/3.6/3.8 MTP (multi-token-prediction) head tensors route into + // the speculator sub-bundle as `mtp.`: the decoder layer takes + // the main-layer canonical forms (so the runtime's per-layer aliases + // apply under the prefix), the head's own fc / norms keep their names. + #[test] + fn qwen35_mtp_tensors_route_to_speculator() { + for (src, want) in [ + ("mtp.fc.weight", "mtp.fc.weight"), + ("mtp.fc_embed.weight", "mtp.fc_embed.weight"), + ("mtp.fc_hidden.weight", "mtp.fc_hidden.weight"), + ("mtp.norm.weight", "mtp.norm.weight"), + ( + "mtp.pre_fc_norm_embedding.weight", + "mtp.pre_fc_norm_embedding.weight", + ), + ( + "mtp.pre_fc_norm_hidden.weight", + "mtp.pre_fc_norm_hidden.weight", + ), + ( + "mtp.layers.0.self_attn.q_proj.weight", + "mtp.layers.0.self_attn.q_proj.weight", + ), + // `input_layernorm` takes the main layers' `input_norm` rename. + ( + "mtp.layers.0.input_layernorm.weight", + "mtp.layers.0.input_norm.weight", + ), + ( + "mtp.layers.0.mlp.experts.gate_up_proj", + "mtp.layers.0.ffn_gate_up_exps.weight", + ), ] { + assert_eq!( + spec_canon(src, "qwen35moe").as_deref(), + Some(want), + "speculator route: {src}" + ); + assert_eq!( + spec_canon(src, "qwen35").as_deref(), + Some(want), + "speculator route: {src}" + ); assert!( - to_canonical_name(n, "qwen35moe").is_none(), - "should drop: {n}" + main_canon(src, "qwen35").is_none(), + "never the main bundle: {src}" ); - assert!(to_canonical_name(n, "qwen35").is_none(), "should drop: {n}"); } // Wrapped variants (`model.mtp.*`, VL-style `model.language_model. - // mtp.*`) must be dropped too — the drop re-applies after wrapper + // mtp.*`) route the same way — the route re-applies after wrapper // prefix stripping. - for n in [ - "model.mtp.fc.weight", - "model.mtp.layers.0.mlp.experts.gate_up_proj", - "model.language_model.mtp.fc.weight", - "language_model.mtp.norm.weight", + for (src, want) in [ + ("model.mtp.fc.weight", "mtp.fc.weight"), + ( + "model.mtp.layers.0.mlp.experts.gate_up_proj", + "mtp.layers.0.ffn_gate_up_exps.weight", + ), + ("model.language_model.mtp.fc.weight", "mtp.fc.weight"), + ("language_model.mtp.norm.weight", "mtp.norm.weight"), ] { - assert!( - to_canonical_name(n, "qwen35moe").is_none(), - "should drop wrapped: {n}" + assert_eq!( + spec_canon(src, "qwen35moe").as_deref(), + Some(want), + "wrapped: {src}" ); } // Non-qwen35 arches keep their behavior for names that merely start - // with "mtp" — the drop is arch-gated. + // with "mtp" — the route is arch-gated. assert!(to_canonical_name("model.layers.0.mlp.gate_proj.weight", "qwen35moe").is_some()); + assert!(spec_canon("mtp.fc.weight", "llama").is_none()); } // `mlp.shared_expert.X_proj.weight`. The experts-rename rules must NOT @@ -6234,6 +7512,38 @@ fn rope_permute_rows(f32s: &[f32], rows: usize, cols: usize, n_heads: u32) -> Ve out } +#[cfg(test)] +mod draft_ranking_tests { + use super::read_draft_ranking; + + fn ranking(text: &str) -> anyhow::Result> { + let dir = std::env::temp_dir().join(format!("draft-rank-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let p = dir.join(format!("r{}.json", text.len())); + std::fs::write(&p, text).unwrap(); + let r = read_draft_ranking(&p); + let _ = std::fs::remove_file(&p); + r + } + + #[test] + fn reads_ranked_ids_in_order() { + let ids = + ranking(r#"{"kind":"mtp_draft_ranking","ranked_ids":[13,0,271,248046]}"#).unwrap(); + assert_eq!(ids, vec![13, 0, 271, 248046]); + } + + /// A repeated id would silently shrink every K the runtime takes. + #[test] + fn refuses_duplicates_and_non_ids() { + assert!(ranking(r#"{"ranked_ids":[5,6,5]}"#).is_err()); + assert!(ranking(r#"{"ranked_ids":[5,-1]}"#).is_err()); + assert!(ranking(r#"{"ranked_ids":[5,1.5]}"#).is_err()); + assert!(ranking(r#"{"ranked_ids":[]}"#).is_err()); + assert!(ranking(r#"{"ids":[1,2]}"#).is_err()); + } +} + #[cfg(test)] mod rope_permute_tests { use super::rope_permute_rows; @@ -6433,3 +7743,179 @@ mod gguf_passthrough_tests { } } } + +#[cfg(test)] +mod tied_head_tests { + use super::{ + infer_gguf_tied_embeddings, plan_tied_lm_head, runtime_reads_output_weight, TiedHeadPlan, + }; + use base_arch::source_mapper_for_gguf; + + fn pairs(v: &[(&str, &str)]) -> Vec<(String, String)> { + v.iter() + .map(|(s, c)| (s.to_string(), c.to_string())) + .collect() + } + + /// HF-style tied checkpoint: one matrix, no head. + fn tied_hf() -> Vec<(String, String)> { + pairs(&[ + ("model.embed_tokens.weight", "embed_tokens.weight"), + ("model.norm.weight", "final_norm.weight"), + ]) + } + + // ── GGUF `has_head` inference ────────────────────────────────────── + + #[test] + fn gguf_with_output_weight_is_untied() { + let m = source_mapper_for_gguf("llama").unwrap(); + let names = ["token_embd.weight", "output.weight", "output_norm.weight"]; + assert!(!infer_gguf_tied_embeddings(m, false, names.into_iter())); + // And the mapped list carries the head, so nothing is synthesized + // even if a caller claimed tied. + let mapped = pairs(&[ + ("token_embd.weight", "embed_tokens.weight"), + ("output.weight", "lm_head.weight"), + ]); + assert_eq!( + plan_tied_lm_head(&mapped, true, "llama", true), + TiedHeadPlan::Nothing + ); + } + + #[test] + fn gguf_without_head_is_tied() { + let m = source_mapper_for_gguf("llama").unwrap(); + let names = [ + "token_embd.weight", + "output_norm.weight", + "blk.0.attn_q.weight", + ]; + assert!(infer_gguf_tied_embeddings(m, false, names.into_iter())); + let mapped = pairs(&[ + ("token_embd.weight", "embed_tokens.weight"), + ("output_norm.weight", "final_norm.weight"), + ]); + assert_eq!( + plan_tied_lm_head(&mapped, true, "llama", true), + TiedHeadPlan::Synthesize { + src: "token_embd.weight".into() + } + ); + } + + #[test] + fn gguf_config_that_says_tied_stays_tied() { + let m = source_mapper_for_gguf("qwen3").unwrap(); + let names = ["token_embd.weight", "output.weight"]; + assert!(infer_gguf_tied_embeddings(m, true, names.into_iter())); + } + + #[test] + fn nomic_bert_gguf_without_head_is_not_tied() { + let m = source_mapper_for_gguf("nomic-bert").unwrap(); + let names = [ + "token_embd.weight", + "token_embd_norm.weight", + "blk.0.attn_qkv.weight", + ]; + assert!(!infer_gguf_tied_embeddings(m, false, names.into_iter())); + assert!(!runtime_reads_output_weight("nomic-bert")); + } + + // ── plan_tied_lm_head ────────────────────────────────────────────── + + #[test] + fn hf_checkpoint_that_ties_and_ships_lm_head_is_not_duplicated() { + let mut mapped = tied_hf(); + mapped.push(("lm_head.weight".into(), "lm_head.weight".into())); + assert_eq!( + plan_tied_lm_head(&mapped, true, "llama", true), + TiedHeadPlan::Nothing + ); + assert_eq!( + plan_tied_lm_head(&mapped, true, "qwen", true), + TiedHeadPlan::Nothing + ); + } + + #[test] + fn untied_or_embeddingless_plans_nothing() { + assert_eq!( + plan_tied_lm_head(&tied_hf(), false, "llama", true), + TiedHeadPlan::Nothing + ); + let no_embed = pairs(&[("model.norm.weight", "final_norm.weight")]); + assert_eq!( + plan_tied_lm_head(&no_embed, true, "llama", true), + TiedHeadPlan::Nothing + ); + } + + /// The arch gate mirrors the runtime: gemma.cpp / gemma4.cpp select the + /// logit tensor by arch.tie_embeddings (hardcoded true) and never look + /// output.weight up, so a synthesized head is dead weight there. + #[test] + fn gemma_runtime_never_reads_output_weight() { + for arch in ["gemma3", "gemma4"] { + assert!(!runtime_reads_output_weight(arch), "{arch}"); + assert_eq!( + plan_tied_lm_head(&tied_hf(), true, arch, true), + TiedHeadPlan::RuntimeIgnoresHead, + "{arch}" + ); + // The opt-out flag is moot there; the arch gate reports first. + assert_eq!( + plan_tied_lm_head(&tied_hf(), true, arch, false), + TiedHeadPlan::RuntimeIgnoresHead, + "{arch}" + ); + } + // A hypothetical new gemma variant lands on the same encoders. + assert!(!runtime_reads_output_weight("gemma5")); + } + + /// llama.cpp / qwen3_5.cpp prefer output.weight when it resolves; + /// muse_glimmer.cpp reads it unconditionally. + #[test] + fn head_consumers_get_a_synthesized_head() { + for arch in [ + "llama", + "qwen", + "qwen_moe", + "qwen3_moe", + "qwen35", + "qwen35moe", + "muse_glimmer", + ] { + assert!(runtime_reads_output_weight(arch), "{arch}"); + assert_eq!( + plan_tied_lm_head(&tied_hf(), true, arch, true), + TiedHeadPlan::Synthesize { + src: "model.embed_tokens.weight".into() + }, + "{arch}" + ); + } + // Unknown archs fall through to encode_prefill_llama in the runtime. + assert!(runtime_reads_output_weight("mistral")); + } + + #[test] + fn no_synth_tied_head_opts_out() { + assert_eq!( + plan_tied_lm_head(&tied_hf(), true, "llama", false), + TiedHeadPlan::OptedOut + ); + assert_eq!( + plan_tied_lm_head(&tied_hf(), true, "qwen35", false), + TiedHeadPlan::OptedOut + ); + // Opt-out never invents a head for an untied model either. + assert_eq!( + plan_tied_lm_head(&tied_hf(), false, "llama", false), + TiedHeadPlan::Nothing + ); + } +} diff --git a/base-convert/crates/base-convert/tests/mlx_passthrough.rs b/base-convert/crates/base-convert/tests/mlx_passthrough.rs new file mode 100644 index 0000000..9f95fd5 --- /dev/null +++ b/base-convert/crates/base-convert/tests/mlx_passthrough.rs @@ -0,0 +1,295 @@ +//! `--mlx-passthrough` on a synthetic MLX Qwen3-MoE checkpoint whose +//! router and shared-expert gate are 8-bit inside a 4-bit source — the +//! layout mlx-community ships for Qwen3.6-35B-A3B (BAS-960). The 4-bit +//! experts must transplant byte-verbatim; the 8-bit tensors the runtime +//! reads at f16 must dequantize exactly and store as f16, and the +//! `--validate` gate must accept the result. + +use base_format::{BaseReader, TensorDtype}; +use std::path::Path; +use std::process::Command; + +const HIDDEN: u64 = 64; +const HEADS: u64 = 2; +const KV_HEADS: u64 = 1; +const HEAD_DIM: u64 = 32; +const EXPERTS: u64 = 4; +const MOE_FFN: u64 = 64; +const VOCAB: u64 = 128; +const GS: u64 = 64; + +fn bin() -> &'static str { + env!("CARGO_BIN_EXE_basert") +} + +/// One MLX affine-quantized tensor: packed codes, f16 scales and biases, +/// and the values they decode to (`q * scale + bias`, f32 math), in +/// row-major order over `shape`. +struct Quant { + shape: Vec, + bits: u32, + packed: Vec, + scales: Vec, + biases: Vec, + values: Vec, +} + +/// Deterministic codes/scales/biases from a seed: the test needs exact +/// bytes to compare, not realistic weights. +fn quant(seed: u32, shape: &[u64], bits: u32) -> Quant { + let in_features = *shape.last().unwrap() as usize; + assert_eq!(in_features as u64 % GS, 0); + let rows: usize = shape[..shape.len() - 1].iter().product::() as usize; + let per_word = 32 / bits as usize; + let words_per_row = in_features / per_word; + let groups_per_row = in_features / GS as usize; + let mut state = seed.wrapping_mul(2654435761) | 1; + let mut next = || { + state ^= state << 13; + state ^= state >> 17; + state ^= state << 5; + state + }; + let mask = (1u32 << bits) - 1; + let mut packed = vec![0u32; rows * words_per_row]; + let mut scales = Vec::with_capacity(rows * groups_per_row); + let mut biases = Vec::with_capacity(rows * groups_per_row); + let mut values = Vec::with_capacity(rows * in_features); + for r in 0..rows { + let mut sc = Vec::new(); + let mut bi = Vec::new(); + for _ in 0..groups_per_row { + let s = half::f16::from_f32(0.001 + (next() % 1000) as f32 * 1e-4); + let b = half::f16::from_f32(-0.5 + (next() % 1000) as f32 * 1e-3); + scales.push(s.to_bits()); + biases.push(b.to_bits()); + sc.push(s.to_f32()); + bi.push(b.to_f32()); + } + for j in 0..in_features { + let q = next() & mask; + let word = r * words_per_row + j / per_word; + let shift = (j % per_word) * bits as usize; + packed[word] |= q << shift; + let g = j / GS as usize; + values.push(q as f32 * sc[g] + bi[g]); + } + } + Quant { + shape: shape.to_vec(), + bits, + packed, + scales, + biases, + values, + } +} + +struct St { + header: serde_json::Map, + blobs: Vec>, + offset: u64, +} + +impl St { + fn new() -> Self { + Self { + header: serde_json::Map::new(), + blobs: Vec::new(), + offset: 0, + } + } + fn push(&mut self, name: &str, dtype: &str, shape: &[u64], bytes: Vec) { + let end = self.offset + bytes.len() as u64; + self.header.insert( + name.to_string(), + serde_json::json!({ "dtype": dtype, "shape": shape, "data_offsets": [self.offset, end] }), + ); + self.offset = end; + self.blobs.push(bytes); + } + fn f16(&mut self, name: &str, shape: &[u64], v: f32) { + let n: u64 = shape.iter().product(); + let bytes: Vec = (0..n) + .flat_map(|_| half::f16::from_f32(v).to_le_bytes()) + .collect(); + self.push(name, "F16", shape, bytes); + } + fn quant(&mut self, name: &str, q: &Quant) { + let per_word = 32 / q.bits as u64; + let mut pshape = q.shape.clone(); + *pshape.last_mut().unwrap() /= per_word; + let mut sshape = q.shape.clone(); + *sshape.last_mut().unwrap() /= GS; + let u32s = |v: &[u32]| v.iter().flat_map(|x| x.to_le_bytes()).collect::>(); + let u16s = |v: &[u16]| v.iter().flat_map(|x| x.to_le_bytes()).collect::>(); + self.push(&format!("{name}.weight"), "U32", &pshape, u32s(&q.packed)); + self.push(&format!("{name}.scales"), "F16", &sshape, u16s(&q.scales)); + self.push(&format!("{name}.biases"), "F16", &sshape, u16s(&q.biases)); + } + fn write(self, path: &Path) { + let hdr = serde_json::to_vec(&serde_json::Value::Object(self.header)).unwrap(); + let mut file = Vec::new(); + file.extend_from_slice(&(hdr.len() as u64).to_le_bytes()); + file.extend_from_slice(&hdr); + for b in &self.blobs { + file.extend_from_slice(b); + } + std::fs::write(path, file).unwrap(); + } +} + +fn f16_bytes(values: &[f32]) -> Vec { + values + .iter() + .flat_map(|&v| half::f16::from_f32(v).to_le_bytes()) + .collect() +} + +#[test] +fn eight_bit_router_and_shared_gate_store_as_f16_beside_verbatim_q4_experts() { + let tmp = std::env::temp_dir().join(format!("basert-mlx-pt-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&tmp); + let dir = tmp.join("src"); + std::fs::create_dir_all(&dir).unwrap(); + std::fs::write( + dir.join("config.json"), + serde_json::json!({ + "model_type": "qwen3_moe", + "architectures": ["Qwen3MoeForCausalLM"], + "hidden_size": HIDDEN, + "num_hidden_layers": 1, + "num_attention_heads": HEADS, + "num_key_value_heads": KV_HEADS, + "head_dim": HEAD_DIM, + "intermediate_size": MOE_FFN, + "moe_intermediate_size": MOE_FFN, + "num_experts": EXPERTS, + "num_experts_per_tok": 2, + "norm_topk_prob": true, + "vocab_size": VOCAB, + "rms_norm_eps": 1e-6, + "rope_theta": 10000.0, + "max_position_embeddings": 512, + "tie_word_embeddings": false, + // The mlx-community layout: a 4-bit checkpoint whose router + // and shared-expert gate are overridden to 8-bit. + "quantization": { + "bits": 4, "group_size": GS, + "model.layers.0.mlp.gate": {"bits": 8, "group_size": GS}, + "model.layers.0.mlp.shared_expert_gate": {"bits": 8, "group_size": GS}, + }, + }) + .to_string(), + ) + .unwrap(); + + let q_dim = HEADS * HEAD_DIM; + let kv_dim = KV_HEADS * HEAD_DIM; + let router = quant(7, &[EXPERTS, HIDDEN], 8); + let shexp_gate = quant(8, &[1, HIDDEN], 8); + let exp_gate = quant(9, &[EXPERTS, MOE_FFN, HIDDEN], 4); + let exp_up = quant(10, &[EXPERTS, MOE_FFN, HIDDEN], 4); + let exp_down = quant(11, &[EXPERTS, HIDDEN, MOE_FFN], 4); + let mut st = St::new(); + st.quant("model.embed_tokens", &quant(1, &[VOCAB, HIDDEN], 4)); + st.quant("lm_head", &quant(2, &[VOCAB, HIDDEN], 4)); + st.f16("model.norm.weight", &[HIDDEN], 1.0); + let l = "model.layers.0"; + st.f16(&format!("{l}.input_layernorm.weight"), &[HIDDEN], 1.0); + st.f16( + &format!("{l}.post_attention_layernorm.weight"), + &[HIDDEN], + 1.0, + ); + st.f16(&format!("{l}.self_attn.q_norm.weight"), &[HEAD_DIM], 1.0); + st.f16(&format!("{l}.self_attn.k_norm.weight"), &[HEAD_DIM], 1.0); + st.quant( + &format!("{l}.self_attn.q_proj"), + &quant(3, &[q_dim, HIDDEN], 4), + ); + st.quant( + &format!("{l}.self_attn.k_proj"), + &quant(4, &[kv_dim, HIDDEN], 4), + ); + st.quant( + &format!("{l}.self_attn.v_proj"), + &quant(5, &[kv_dim, HIDDEN], 4), + ); + st.quant( + &format!("{l}.self_attn.o_proj"), + &quant(6, &[HIDDEN, q_dim], 4), + ); + st.quant(&format!("{l}.mlp.gate"), &router); + st.quant(&format!("{l}.mlp.shared_expert_gate"), &shexp_gate); + st.quant(&format!("{l}.mlp.switch_mlp.gate_proj"), &exp_gate); + st.quant(&format!("{l}.mlp.switch_mlp.up_proj"), &exp_up); + st.quant(&format!("{l}.mlp.switch_mlp.down_proj"), &exp_down); + st.write(&dir.join("model.safetensors")); + + let out = tmp.join("out.base"); + let output = Command::new(bin()) + .arg("convert") + .arg(&dir) + .arg("-o") + .arg(&out) + .arg("--mlx-passthrough") + .arg("--validate") + .output() + .expect("run basert convert"); + assert!( + output.status.success(), + "convert --mlx-passthrough --validate failed:\n{}", + String::from_utf8_lossy(&output.stderr) + ); + + let reader = BaseReader::open(&out).unwrap(); + let entry = |name: &str| { + reader + .header() + .tensors + .iter() + .find(|t| t.name == name) + .unwrap_or_else(|| { + panic!( + "{name} missing; have {:?}", + reader + .header() + .tensors + .iter() + .map(|t| &t.name) + .collect::>() + ) + }) + .clone() + }; + + // The 8-bit tensors the runtime reads at f16: exact dequant, f16. + for (name, q) in [ + ("layers.0.mlp.gate.weight", &router), + ("layers.0.mlp.shared_expert_gate.weight", &shexp_gate), + ] { + let e = entry(name); + assert_eq!(e.dtype, TensorDtype::F16, "{name} dtype"); + assert_eq!( + reader.tensor_bytes(name).unwrap(), + f16_bytes(&q.values), + "{name} payload" + ); + } + + // The 4-bit experts: transplanted packed codes, byte-verbatim. + for (name, q) in [ + ("layers.0.ffn_gate_exps.weight", &exp_gate), + ("layers.0.ffn_up_exps.weight", &exp_up), + ("layers.0.ffn_down_exps.weight", &exp_down), + ] { + let e = entry(name); + assert_eq!(e.dtype, TensorDtype::BaseQ4, "{name} dtype"); + let data = reader.tensor_bytes(name).unwrap(); + let scale_off = e.scale_offset.unwrap() as usize; + let want: Vec = q.packed.iter().flat_map(|w| w.to_le_bytes()).collect(); + assert_eq!(&data[..scale_off], &want[..], "{name} packed codes"); + } + let _ = std::fs::remove_dir_all(&tmp); +} diff --git a/base-convert/crates/base-format/src/header.rs b/base-convert/crates/base-format/src/header.rs index befca4b..230181b 100644 --- a/base-convert/crates/base-format/src/header.rs +++ b/base-convert/crates/base-format/src/header.rs @@ -18,7 +18,9 @@ bitflags! { const HAS_HYBRID = 1 << 3; /// Extension slots contain LoRA delta weights. const HAS_LORA = 1 << 4; - /// Extension slots contain a paired speculator model. + /// The bundle carries a speculator sub-bundle (`header.speculator`): + /// a drafter that conditions on this model's hidden states — an + /// in-checkpoint MTP head, or a converted EAGLE / DFlash drafter. const HAS_SPECULATOR = 1 << 5; /// Extension slots contain precompiled compute graphs. const HAS_COMPUTE_GRAPH = 1 << 6; @@ -525,6 +527,9 @@ pub struct Provenance { /// Source tensors routed into the mmproj sub-bundle. #[serde(default, skip_serializing_if = "Vec::is_empty")] pub mmproj: Vec, + /// Source tensors routed into the speculator sub-bundle. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub speculator: Vec, /// Bundle tensor name -> provenance record. pub tensors: BTreeMap, } @@ -557,6 +562,11 @@ pub struct TensorProvenance { /// qkv_proj / gate_up_proj splits). #[serde(default, skip_serializing_if = "Option::is_none")] pub rows: Option<[u64; 2]>, + /// `[col_offset, col_count]` input-column block of the source tensor + /// (the MTP head's `mtp.fc` split into fc_embed / fc_hidden, a block + /// drafter's per-tap `fc.{i}`, EAGLE-3's `*_proj_emb` / `*_proj_hid`). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cols: Option<[u64; 2]>, /// Stacked from `count` per-expert source tensors; `pattern` /// contains `{e}` where the expert index goes. #[serde(default, skip_serializing_if = "Option::is_none")] @@ -626,6 +636,8 @@ pub struct Header { #[serde(skip_serializing_if = "Option::is_none", default)] pub mmproj: Option, #[serde(skip_serializing_if = "Option::is_none", default)] + pub speculator: Option, + #[serde(skip_serializing_if = "Option::is_none", default)] pub calibration: Option, #[serde(skip_serializing_if = "Option::is_none", default)] pub sig: Option, @@ -663,6 +675,29 @@ pub struct MmprojBundle { pub tensors: Vec, } +/// Speculator sub-bundle: a drafter the runtime runs alongside the main +/// model for speculative decoding. Same layout convention as `mmproj` — +/// its tensors live in the same weights blob, listed here so a runtime +/// without speculation skips them. Kinds so far: +/// * `-mtp`: the checkpoint's own Multi-Token-Prediction head +/// (Qwen3.5/3.6/3.8 `mtp.*`, GLM `nextn`): one decoder layer of the +/// target's geometry plus `fc` / pre-fc norms / final norm, sharing the +/// target's embeddings and lm_head. Tensors are named `mtp.` +/// (the decoder layer as `mtp.layers.0.*`, canonicalized exactly like a +/// main layer, norm shifts baked). Kept unquantized (f16): the head's +/// acceptance rate is what the whole scheme buys, and quantizing it +/// measured near-zero acceptance on MoE targets. +/// +/// `config` is open-namespace: `kind`, `num_layers`, `moe`, +/// `shared_embeddings`, drafter-specific keys. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SpeculatorBundle { + pub arch: String, + #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] + pub config: BTreeMap, + pub tensors: Vec, +} + /// Per-layer kind. Drives runtime dispatch: which forward path to run /// (attention vs SSM), which KV / SSM buffers to allocate, and whether /// the layer feeds into an MoE FFN. diff --git a/base-convert/crates/base-format/src/lib.rs b/base-convert/crates/base-format/src/lib.rs index 00e1db0..b176f29 100644 --- a/base-convert/crates/base-format/src/lib.rs +++ b/base-convert/crates/base-format/src/lib.rs @@ -12,8 +12,8 @@ pub use error::{Error, Result}; pub use header::{ AlignmentConfig, CalibrationInfo, ComputeRegion, Header, HeaderFlags, LayerDescriptor, LayerKind, LayerPrecision, Layout, ModelConfig, Provenance, QuantScheme, ResidencyHint, - ScaleDtype, Signature, SourceInfo, StackRef, TargetBackend, TensorDtype, TensorEntry, - TensorFlags, TensorProvenance, TokenizerBlob, + ScaleDtype, Signature, SourceInfo, SpeculatorBundle, StackRef, TargetBackend, TensorDtype, + TensorEntry, TensorFlags, TensorProvenance, TokenizerBlob, }; pub use reader::BaseReader; pub use slots::{read_slots, write_slots, Slot, SlotFlags, SlotKind}; diff --git a/base-convert/crates/base-format/src/reader.rs b/base-convert/crates/base-format/src/reader.rs index d8bd623..cf5ecfe 100644 --- a/base-convert/crates/base-format/src/reader.rs +++ b/base-convert/crates/base-format/src/reader.rs @@ -116,15 +116,31 @@ impl BaseReader { } /// Get a zero-copy slice of a tensor's raw bytes. - pub fn tensor_bytes(&self, name: &str) -> Result<&[u8]> { - let entry = self - .header - .tensors - .iter() + /// Resolve a tensor entry by name across the main list AND the + /// sub-bundles (`header.mmproj.tensors`, `header.speculator.tensors`): + /// they share the weights blob, so every byte-level accessor and the + /// checksum verifier must see them too — scanning only `header.tensors` + /// leaves a corrupted MTP payload reporting "all checksums OK". + pub fn find_entry(&self, name: &str) -> Result<&crate::header::TensorEntry> { + self.all_entries() .find(|t| t.name == name) .ok_or_else(|| Error::TensorNotFound { name: name.to_string(), - })?; + }) + } + + /// Every tensor entry of the bundle: main list first, then the mmproj + /// tower(s), then the speculator head. + pub fn all_entries(&self) -> impl Iterator { + self.header + .tensors + .iter() + .chain(self.header.mmproj.iter().flat_map(|m| m.tensors.iter())) + .chain(self.header.speculator.iter().flat_map(|m| m.tensors.iter())) + } + + pub fn tensor_bytes(&self, name: &str) -> Result<&[u8]> { + let entry = self.find_entry(name)?; let abs_start = self.blob_offset + entry.offset; let abs_end = abs_start + entry.length; @@ -145,14 +161,7 @@ impl BaseReader { /// constructing zero-copy buffers (MTLBuffer.makeBufferWithBytesNoCopy, /// cudaHostRegister, etc.) without holding a slice. pub fn tensor_file_offset(&self, name: &str) -> Result { - let entry = self - .header - .tensors - .iter() - .find(|t| t.name == name) - .ok_or_else(|| Error::TensorNotFound { - name: name.to_string(), - })?; + let entry = self.find_entry(name)?; Ok(self.blob_offset + entry.offset) } @@ -162,14 +171,7 @@ impl BaseReader { /// invariant was violated and `makeBufferWithBytesNoCopy` would /// silently fall back to a copy. pub fn tensor_is_zero_copy_eligible(&self, name: &str) -> Result { - let entry = self - .header - .tensors - .iter() - .find(|t| t.name == name) - .ok_or_else(|| Error::TensorNotFound { - name: name.to_string(), - })?; + let entry = self.find_entry(name)?; let align = self.header.alignment.align_for(entry.compute_region); let abs = self.blob_offset + entry.offset; Ok(abs % align == 0) @@ -217,7 +219,17 @@ impl BaseReader { .iter() .flat_map(|m| m.tensors.iter()) .map(|t| self.blob_offset + t.offset + t.length); - let blob_end = main_end.chain(mmproj_end).max().unwrap_or(self.blob_offset); + let spec_end = self + .header + .speculator + .iter() + .flat_map(|m| m.tensors.iter()) + .map(|t| self.blob_offset + t.offset + t.length); + let blob_end = main_end + .chain(mmproj_end) + .chain(spec_end) + .max() + .unwrap_or(self.blob_offset); (blob_end + 7) & !7u64 } @@ -227,14 +239,7 @@ impl BaseReader { /// Err(ChecksumMismatch) when the recorded checksum disagrees with /// the current bytes. pub fn verify_tensor(&self, name: &str) -> Result<()> { - let entry = self - .header - .tensors - .iter() - .find(|t| t.name == name) - .ok_or_else(|| Error::TensorNotFound { - name: name.to_string(), - })?; + let entry = self.find_entry(name)?; let Some(expected) = entry.checksum_xxh64 else { return Ok(()); }; diff --git a/base-convert/crates/base-format/src/writer.rs b/base-convert/crates/base-format/src/writer.rs index c23220b..dcf9ac9 100644 --- a/base-convert/crates/base-format/src/writer.rs +++ b/base-convert/crates/base-format/src/writer.rs @@ -1,4 +1,6 @@ -use crate::header::{Header, MmprojBundle, TensorDtype, TensorEntry}; +use crate::header::{ + Header, HeaderFlags, MmprojBundle, SpeculatorBundle, TensorDtype, TensorEntry, +}; use crate::slots::{write_slots, Slot}; use crate::{Error, Result, BLOB_ALIGNMENT, FORMAT_VERSION, MAGIC, PREFIX_LEN}; use std::fs::File; @@ -21,15 +23,52 @@ pub struct TensorPayload { /// serialized (all offsets/lengths/checksums are already known) and the /// temp blob is streamed into the final file after it. The on-disk format /// is byte-identical to the buffered path. +/// Owns removal of the streamed-blob temp file. Held BY `BlobStream`, so the +/// temp is removed on every exit path from the moment it is created: an +/// `add_tensor` I/O error that makes `finish` return before it reaches +/// `finish_streaming`, a `BaseWriter` dropped without `finish` at all, a panic, +/// as well as the happy path. Previously the guard was constructed inside +/// `finish_streaming`, so those earlier exits stranded a model-sized +/// `.blobtmp` beside the failed output — worst of all for the expected +/// disk-full failure, where the orphan eats the space the retry needs. +struct TmpBlob(PathBuf); + +impl Drop for TmpBlob { + fn drop(&mut self) { + let _ = std::fs::remove_file(&self.0); + } +} + +/// Which header tensor list a streamed payload belongs to. +#[derive(Clone, Copy)] +enum Section { + Lm, + Mmproj, + Speculator, +} + struct BlobStream { file: BufWriter, - path: PathBuf, + /// The temp path, and the guard that removes it — `Some` ONLY on the + /// `create` path, which streams into a `.blobtmp` sibling. `create_direct` + /// streams into the FINAL file behind a reserved header region and has no + /// temp at all, so it leaves this `None`: arming the guard there would + /// delete the bundle the caller just wrote. + /// + /// `BlobStream` itself has no `Drop`, so `finish_streaming` can still move + /// `file` out of it by value; this field drops with whatever remains. + /// + /// Never read: it exists for its `Drop`. `expect` rather than `allow` so + /// the attribute itself warns if a later change does start reading it. + #[expect(dead_code)] + tmp: Option, /// Bytes written to the blob so far (blob-relative cursor). Tensor /// `entry.offset` values are relative to the blob start, matching the /// reader's `blob_offset + entry.offset` addressing. cursor: u64, lm_entries: Vec, mmproj_entries: Vec, + speculator_entries: Vec, } /// Writer for the `.base` single-file format. @@ -64,6 +103,12 @@ pub struct BaseWriter { /// so the runtime can populate vision/audio fields without a /// separate config file. mmproj_config: std::collections::BTreeMap, + /// Speculator sub-bundle payloads / arch / config — same convention as + /// mmproj: the same blob, listed under `header.speculator`. Streamed + /// after every LM and mmproj tensor. + speculator_payloads: Vec, + speculator_arch: Option, + speculator_config: std::collections::BTreeMap, slots: Vec, /// First error encountered while streaming a payload to the temp blob /// (add_tensor is infallible for API compatibility); reported at finish. @@ -101,15 +146,19 @@ impl BaseWriter> { header, stream: Some(BlobStream { file: BufWriter::new(tmp), - path: tmp_path, + tmp: Some(TmpBlob(tmp_path)), cursor: 0, lm_entries: Vec::new(), mmproj_entries: Vec::new(), + speculator_entries: Vec::new(), }), payloads: Vec::new(), mmproj_payloads: Vec::new(), mmproj_arch: None, mmproj_config: std::collections::BTreeMap::new(), + speculator_payloads: Vec::new(), + speculator_arch: None, + speculator_config: std::collections::BTreeMap::new(), slots: Vec::new(), stream_error: None, direct_reserve: None, @@ -147,15 +196,19 @@ impl BaseWriter> { header, stream: Some(BlobStream { file: BufWriter::new(blob), - path: path.to_path_buf(), + tmp: None, // direct-write: the "stream" IS the final file cursor: 0, lm_entries: Vec::new(), mmproj_entries: Vec::new(), + speculator_entries: Vec::new(), }), payloads: Vec::new(), mmproj_payloads: Vec::new(), mmproj_arch: None, mmproj_config: std::collections::BTreeMap::new(), + speculator_payloads: Vec::new(), + speculator_arch: None, + speculator_config: std::collections::BTreeMap::new(), slots: Vec::new(), stream_error: None, direct_reserve: Some(header_reserve), @@ -173,6 +226,9 @@ impl BaseWriter { mmproj_payloads: Vec::new(), mmproj_arch: None, mmproj_config: std::collections::BTreeMap::new(), + speculator_payloads: Vec::new(), + speculator_arch: None, + speculator_config: std::collections::BTreeMap::new(), slots: Vec::new(), stream_error: None, direct_reserve: None, @@ -186,7 +242,7 @@ impl BaseWriter { alignment: crate::header::AlignmentConfig, s: &mut BlobStream, payload: TensorPayload, - is_mmproj: bool, + section: Section, ) -> Result<()> { let align = alignment.align_for(payload.entry.compute_region); let aligned = align_up(s.cursor, align); @@ -201,10 +257,10 @@ impl BaseWriter { } s.file.write_all(&payload.data)?; s.cursor = aligned + payload.data.len() as u64; - if is_mmproj { - s.mmproj_entries.push(entry); - } else { - s.lm_entries.push(entry); + match section { + Section::Lm => s.lm_entries.push(entry), + Section::Mmproj => s.mmproj_entries.push(entry), + Section::Speculator => s.speculator_entries.push(entry), } Ok(()) } @@ -216,7 +272,7 @@ impl BaseWriter { // a write error here is surfaced by re-checking at finish via // the BufWriter's retained error state on flush.) let alignment = self.header.alignment; - if let Err(e) = Self::stream_payload(alignment, s, payload, false) { + if let Err(e) = Self::stream_payload(alignment, s, payload, Section::Lm) { // Stash the error to report at finish; keep the API simple. self.stream_error.get_or_insert(e); } @@ -235,7 +291,7 @@ impl BaseWriter { // add every LM tensor before the first mmproj tensor, so the // streamed order matches the buffered path. let alignment = self.header.alignment; - if let Err(e) = Self::stream_payload(alignment, s, payload, true) { + if let Err(e) = Self::stream_payload(alignment, s, payload, Section::Mmproj) { self.stream_error.get_or_insert(e); } } else { @@ -243,6 +299,37 @@ impl BaseWriter { } } + /// Add a tensor that belongs to the speculator sub-bundle (an MTP head + /// or a converted drafter). Same blob; its entry lands under + /// `header.speculator.tensors`. Callers add every LM and mmproj tensor + /// before the first speculator tensor so the streamed order matches + /// the buffered path. + pub fn add_speculator_tensor(&mut self, payload: TensorPayload) { + if let Some(s) = self.stream.as_mut() { + let alignment = self.header.alignment; + if let Err(e) = Self::stream_payload(alignment, s, payload, Section::Speculator) { + self.stream_error.get_or_insert(e); + } + } else { + self.speculator_payloads.push(payload); + } + } + + /// Set the speculator sub-bundle arch tag (e.g. "qwen35-mtp"). + /// Required when any `add_speculator_tensor` was called. + pub fn set_speculator_arch(&mut self, arch: impl Into) { + self.speculator_arch = Some(arch.into()); + } + + /// Set the speculator sub-bundle config block (`kind`, `num_layers`, + /// `moe`, `shared_embeddings`, drafter-specific keys). + pub fn set_speculator_config( + &mut self, + cfg: std::collections::BTreeMap, + ) { + self.speculator_config = cfg; + } + /// Set the mmproj sub-bundle arch tag (e.g. "gemma4_vision_audio"). /// Required when any `add_mmproj_tensor` was called. pub fn set_mmproj_arch(&mut self, arch: impl Into) { @@ -295,6 +382,21 @@ impl BaseWriter { } } + /// Publish the speculator sub-bundle in the header (arch, config, + /// entries) and raise HAS_SPECULATOR. + fn install_speculator(&mut self, entries: Vec) { + let arch = self + .speculator_arch + .clone() + .unwrap_or_else(|| "speculator".to_string()); + self.header.speculator = Some(SpeculatorBundle { + arch, + config: std::mem::take(&mut self.speculator_config), + tensors: entries, + }); + self.header.flags |= HeaderFlags::HAS_SPECULATOR; + } + pub fn finish(mut self) -> Result<()> { // Surface any error stashed while streaming payloads to the temp blob. if let Some(e) = self.stream_error.take() { @@ -360,6 +462,24 @@ impl BaseWriter { tensors: mmproj_entries, }); } + // Speculator sub-bundle entries continue past the mmproj tensors. + if !self.speculator_payloads.is_empty() { + let mut spec_entries: Vec = + Vec::with_capacity(self.speculator_payloads.len()); + for p in &self.speculator_payloads { + let align = alignment.align_for(p.entry.compute_region); + let aligned = align_up(blob_cursor, align); + let mut entry = p.entry.clone(); + entry.offset = aligned; + entry.length = p.data.len() as u64; + if entry.checksum_xxh64.is_none() { + entry.checksum_xxh64 = Some(xxhash_rust::xxh64::xxh64(&p.data, 0)); + } + spec_entries.push(entry); + blob_cursor = aligned + p.data.len() as u64; + } + self.install_speculator(spec_entries); + } // Serialize header (canonical JSON, sorted keys via BTreeMap). let header_json = self.header.to_canonical_json().map_err(Error::Json)?; @@ -412,6 +532,19 @@ impl BaseWriter { blob_written += p.data.len() as u64; } } + // Speculator data continues past the mmproj tensors. + if let Some(spec) = &self.header.speculator { + for (p, e) in self.speculator_payloads.iter().zip(spec.tensors.iter()) { + if e.offset > blob_written { + let pad = (e.offset - blob_written) as usize; + let zeros = vec![0u8; pad]; + self.inner.write_all(&zeros)?; + blob_written += pad as u64; + } + self.inner.write_all(&p.data)?; + blob_written += p.data.len() as u64; + } + } // Extension slots, if any. Pad to 8-byte boundary first so the // slots section starts aligned. @@ -448,19 +581,14 @@ impl BaseWriter { tensors: std::mem::take(&mut s.mmproj_entries), }); } - - // Remove the temp blob on EVERY exit from here on, not just the happy - // path. A failure below (the destination filling mid-copy is the - // realistic one) used to return through `?` and strand a model-sized - // `.blobtmp` next to the partial output — hundreds of GB that the user - // has to find by hand, and that makes every retry fail for space. - struct TmpGuard(std::path::PathBuf); - impl Drop for TmpGuard { - fn drop(&mut self) { - let _ = std::fs::remove_file(&self.0); - } + if !s.speculator_entries.is_empty() { + let entries = std::mem::take(&mut s.speculator_entries); + self.install_speculator(entries); } - let _tmp_guard = TmpGuard(s.path.clone()); + + // The temp blob is removed by `s.tmp` whichever way this function + // leaves — including the destination filling mid-copy, the realistic + // failure here. // Flush + rewind the temp blob for reading. s.file.flush()?; @@ -520,6 +648,10 @@ impl BaseWriter { tensors: std::mem::take(&mut s.mmproj_entries), }); } + if !s.speculator_entries.is_empty() { + let entries = std::mem::take(&mut s.speculator_entries); + self.install_speculator(entries); + } // Extension slots continue past the blob (8-byte aligned) on the // blob handle — it is already positioned at the blob end. diff --git a/base-convert/crates/base-format/tests/roundtrip.rs b/base-convert/crates/base-format/tests/roundtrip.rs index 475c25e..054c236 100644 --- a/base-convert/crates/base-format/tests/roundtrip.rs +++ b/base-convert/crates/base-format/tests/roundtrip.rs @@ -32,6 +32,7 @@ fn make_header() -> Header { layers: vec![], tensors: vec![], mmproj: None, + speculator: None, calibration: None, provenance: None, sig: None, @@ -496,3 +497,50 @@ fn target_backend_stamped_from_content() { let reader = BaseReader::open(tmp2.path()).unwrap(); assert_eq!(reader.header().target_backend, TargetBackend::Metal); } + +/// The speculator sub-bundle rides the same blob after the LM tensors and +/// is listed under `header.speculator`; the LM list, the flag, the config +/// and the byte payloads all round-trip, and the slot scan walks past it. +#[test] +fn roundtrip_speculator_section() { + let tmp = tempfile::NamedTempFile::new().unwrap(); + let path = tmp.path().to_path_buf(); + let mut writer = BaseWriter::create(&path, make_header()).unwrap(); + let lm: Vec = (0..64u8).collect(); + let head: Vec = (0..96u8).map(|b| 255 - b).collect(); + writer.add_tensor(TensorPayload { + entry: entry("layers.0.self_attn.q_proj.weight", vec![16]), + data: lm.clone(), + }); + writer.set_speculator_arch("test-mtp"); + let mut cfg = BTreeMap::new(); + cfg.insert("kind".to_string(), serde_json::json!("mtp")); + cfg.insert("num_layers".to_string(), serde_json::json!(1)); + writer.set_speculator_config(cfg); + writer.add_speculator_tensor(TensorPayload { + entry: entry("mtp.fc.weight", vec![24]), + data: head.clone(), + }); + writer.add_slot(Slot::new(SlotKind::Custom, vec![1, 2, 3])); + writer.finish().unwrap(); + + let reader = BaseReader::open(&path).unwrap(); + let h = reader.header(); + assert_eq!(h.tensors.len(), 1, "the LM list holds only LM tensors"); + assert!(h.flags.contains(HeaderFlags::HAS_SPECULATOR)); + let spec = h.speculator.as_ref().expect("speculator section"); + assert_eq!(spec.arch, "test-mtp"); + assert_eq!(spec.config["kind"], serde_json::json!("mtp")); + assert_eq!(spec.tensors.len(), 1); + assert_eq!(spec.tensors[0].name, "mtp.fc.weight"); + // Payload sits past the LM tensor in the shared blob. + let lm_e = &h.tensors[0]; + assert!(spec.tensors[0].offset >= lm_e.offset + lm_e.length); + let off = (reader.blob_offset() + spec.tensors[0].offset) as usize; + let bytes = std::fs::read(&path).unwrap(); + assert_eq!(&bytes[off..off + head.len()], head.as_slice()); + // The slot scan starts after the speculator payload, not inside it. + let slots = reader.slots().unwrap(); + assert_eq!(slots.len(), 1); + assert_eq!(slots[0].payload, vec![1, 2, 3]); +} diff --git a/base-convert/crates/base-hub/catalog.json b/base-convert/crates/base-hub/catalog.json index 34c7f07..a813239 100644 --- a/base-convert/crates/base-hub/catalog.json +++ b/base-convert/crates/base-hub/catalog.json @@ -1,6 +1,6 @@ { "schema": 1, - "updated": "2026-09-08", + "updated": "2026-09-24", "models": [ { "id": "basecompute/Llama-3.1-8B-Instruct", @@ -381,8 +381,9 @@ "source_repo": "Qwen/Qwen3.5-2B", "arch": "qwen35", "quant": "default-q4", - "size": 1266810880, - "sha256": "1c95f41730d5713fcb886a71d2481a29ff1a6e2fc9cea8a4d9ddbe29082fa262" + "size": 1599324160, + "sha256": "deb6319c18511baf709862c87f34669f17c6c6e904995050a92d3152d0d53188", + "speculator": "mtp" }, { "id": "basecompute/Qwen3.5-2B", @@ -392,8 +393,9 @@ "source_repo": "Qwen/Qwen3.5-2B", "arch": "qwen35", "quant": "default-q8", - "size": 2301460480, - "sha256": "d4cf325278e873b3e0aefc18c5eff0a863f12fee5a82f0ac2f2e5ebce4067f3c" + "size": 2896936960, + "sha256": "a144502c5280dc4f6191d36560b85e334383685cf2e7be6b8f6f845e6d8a4411", + "speculator": "mtp" }, { "id": "basecompute/Qwen3.5-2B", @@ -403,9 +405,10 @@ "source_repo": "Qwen/Qwen3.5-2B", "arch": "qwen35", "quant": "cuda-q4mix", - "size": 1592131584, - "sha256": "db457f5dd499be915e51e757d0e999e279def9f861f0ecc206708deaa5ec49ca", - "backend": "cuda" + "size": 2468659200, + "sha256": "899a740301babfd1fa65e9175ed713c1af9ff2d6799b0b42518084ee5cad6cb4", + "backend": "cuda", + "speculator": "mtp" }, { "id": "basecompute/Qwen3.5-2B", @@ -415,9 +418,10 @@ "source_repo": "Qwen/Qwen3.5-2B", "arch": "qwen35", "quant": "cuda-q8", - "size": 2301460480, - "sha256": "027907a8ce268e1a7ffc339cd2746db87c2e5f01867910284b19cbb52ec62d78", - "backend": "cuda" + "size": 2896936960, + "sha256": "3711edbe2dddcc4304d74cacd0c40478c2fc08590c870ff53197ea1f7fec4fc8", + "backend": "cuda", + "speculator": "mtp" }, { "id": "basecompute/Qwen3.5-2B-Base", @@ -427,8 +431,9 @@ "source_repo": "Qwen/Qwen3.5-2B-Base", "arch": "qwen35", "quant": "default-q4", - "size": 1266810880, - "sha256": "ef10005ccd4c6bd3e2d9fbcac31b7f4491eff2639584bdc8006e67226f63d986" + "size": 1599324160, + "sha256": "1d7b425349b17d3dcb36dc56573e92b188a01d5b613e01a92fa19d6fb24b1522", + "speculator": "mtp" }, { "id": "basecompute/Qwen3.5-2B-Base", @@ -438,8 +443,9 @@ "source_repo": "Qwen/Qwen3.5-2B-Base", "arch": "qwen35", "quant": "default-q8", - "size": 2301460480, - "sha256": "f546b5074d92f104453a782a44143f3a44ed538569befdc7e2ff22675c719a30" + "size": 2896936960, + "sha256": "10e026b1f0360fe431463e106f9090b3a197a437dbd1780ff7b63726cd2520fb", + "speculator": "mtp" }, { "id": "basecompute/Qwen3.5-35B-A3B", @@ -449,8 +455,9 @@ "source_repo": "Qwen/Qwen3.5-35B-A3B", "arch": "qwen35moe", "quant": "default-q4", - "size": 19776225280, - "sha256": "e62635e523e26c7a31df2d0d26cec485630b3b2172a9fd710b06e90c45174ff5" + "size": 20515667968, + "sha256": "a7fc993c9d2ea5706c4b36a977740c9ccc253d236a9df07398c49c77ebec627e", + "speculator": "mtp" }, { "id": "basecompute/Qwen3.5-35B-A3B", @@ -460,8 +467,9 @@ "source_repo": "Qwen/Qwen3.5-35B-A3B", "arch": "qwen35moe", "quant": "default-q8", - "size": 36228218880, - "sha256": "6f8878ff2332451efbf2258887bad0535a97aeb657cf0b7ce8492219881184ff" + "size": 37296832512, + "sha256": "9c3214ea3b6c862ce83ca25acc64253b4bc28cafeba80944659536791856f2e3", + "speculator": "mtp" }, { "id": "basecompute/Qwen3.5-35B-A3B", @@ -471,9 +479,10 @@ "source_repo": "Qwen/Qwen3.5-35B-A3B", "arch": "qwen35moe", "quant": "cuda-q4mix", - "size": 21123661824, - "sha256": "1f836346dfe62bddb1c9b2a4f514c7b2fcea47b63d6d5dd1d8f4cfee95b50363", - "backend": "cuda" + "size": 22086270976, + "sha256": "f010c9598a63c6c91019d1469611ce29c4583076689ae9a5dd7ac2540e0855f0", + "backend": "cuda", + "speculator": "mtp" }, { "id": "basecompute/Qwen3.5-35B-A3B", @@ -483,9 +492,10 @@ "source_repo": "Qwen/Qwen3.5-35B-A3B", "arch": "qwen35moe", "quant": "cuda-q8", - "size": 36357767168, - "sha256": "7dbbe60e58c96ab845eedf14ad9684abaa5175ba04d137a3da704b2a16afacbd", - "backend": "cuda" + "size": 37296832512, + "sha256": "a6a3372d9ce1b7c69b7844c47b02bdd7abe33080dccfb0c3a36b15f183123d6f", + "backend": "cuda", + "speculator": "mtp" }, { "id": "basecompute/Qwen3.6-27B", @@ -495,8 +505,9 @@ "source_repo": "Qwen/Qwen3.6-27B", "arch": "qwen35", "quant": "default-q4", - "size": 15447605248, - "sha256": "f24b31c552994f82e18f2a3b7a7352ff0c41cf070351e941816365bd02bca79c" + "size": 15955853312, + "sha256": "4472dd8e158ac0ce7056bafd180268ebd124509a21f8df06951b2986a7f7f06c", + "speculator": "mtp" }, { "id": "basecompute/Qwen3.6-27B", @@ -506,8 +517,9 @@ "source_repo": "Qwen/Qwen3.6-27B", "arch": "qwen35", "quant": "default-q8", - "size": 28256698368, - "sha256": "f65023ea2a5e1cc2a6091cc926a4156935e4ada58133d4bc84bc6d759f738039" + "size": 28876619776, + "sha256": "e778d52df4fa65091cbf520feceb1b2a1b5e18c46a432223657ef18f26277478", + "speculator": "mtp" }, { "id": "basecompute/Qwen3.6-27B", @@ -517,9 +529,10 @@ "source_repo": "Qwen/Qwen3.6-27B", "arch": "qwen35", "quant": "cuda-q4mix", - "size": 18159501312, - "sha256": "7aebf3944f2d0ba10b5b2d9ddbb9fd0bfe3f52f2eeea011da362416209709921", - "backend": "cuda" + "size": 19534544896, + "sha256": "2535e83f73c223f98379cf7a8a97fd385a440d3f6edc477e4831bda0b67aec51", + "backend": "cuda", + "speculator": "mtp" }, { "id": "basecompute/Qwen3.6-27B", @@ -529,9 +542,10 @@ "source_repo": "Qwen/Qwen3.6-27B", "arch": "qwen35", "quant": "cuda-q8", - "size": 28386312192, - "sha256": "d0be0cad29ccafde360cb92fa2298527abc35749e8ba390b2a5b7ce596ba938f", - "backend": "cuda" + "size": 28934225920, + "sha256": "7575c2273cb38948ce67fbf921568aa739c87e3a88058ab2c8930bff85e1fdd6", + "backend": "cuda", + "speculator": "mtp" }, { "id": "basecompute/Qwen3.6-35B-A3B", @@ -541,8 +555,9 @@ "source_repo": "Qwen/Qwen3.6-35B-A3B", "arch": "qwen35moe", "quant": "default-q4", - "size": 19776225280, - "sha256": "6819d2152fe025af5b1eb12d02e862c5277258a016e6b901c6c91791ba24c48b" + "size": 20515667968, + "sha256": "7171b37ffccdf16dfd920f55be766e6faffa31f64e04ca9cd0064943bf72881b", + "speculator": "mtp" }, { "id": "basecompute/Qwen3.6-35B-A3B", @@ -552,8 +567,9 @@ "source_repo": "Qwen/Qwen3.6-35B-A3B", "arch": "qwen35moe", "quant": "default-q8", - "size": 36228218880, - "sha256": "11d093a918e2a0b917430547dafcdcab2fe1c2d689ec86b5343382f72d2b2c3b" + "size": 37296832512, + "sha256": "a1ad0f6a89535f60d2e3862e2451297c67698858fc61e91ab13026372d4667b1", + "speculator": "mtp" }, { "id": "basecompute/Qwen3.6-35B-A3B", @@ -563,9 +579,10 @@ "source_repo": "Qwen/Qwen3.6-35B-A3B", "arch": "qwen35moe", "quant": "cuda-q4mix", - "size": 21123661824, - "sha256": "75951215950ba8feb94dfb86351505b6f6fc818c6dbedd64d946c6cd9214fe4a", - "backend": "cuda" + "size": 22086270976, + "sha256": "a275b616f539badb9d61dd0788d9fb33381c4fd576b5ec7e46dc4b12e3bce924", + "backend": "cuda", + "speculator": "mtp" }, { "id": "basecompute/Qwen3.6-35B-A3B", @@ -575,9 +592,10 @@ "source_repo": "Qwen/Qwen3.6-35B-A3B", "arch": "qwen35moe", "quant": "cuda-q8", - "size": 36357767168, - "sha256": "350f09147443f1ed5e7e8709e41d43fe7d886138ad93129c69c8881208951a04", - "backend": "cuda" + "size": 37296832512, + "sha256": "cd5e4334529e3e7bb4c43064e5691a68bb4d86f6ef12b1dd618427f8117058d5", + "backend": "cuda", + "speculator": "mtp" }, { "id": "basecompute/Qwen3.8-27B", @@ -587,8 +605,9 @@ "source_repo": "Qwen/Qwen3.8-27B", "arch": "qwen35", "quant": "default-q4", - "size": 17467801600, - "sha256": "069f4ceca264d9e3a69ee9af725249a924024a27d46680d37da6b5851c49a294" + "size": 17841094656, + "sha256": "03978674a487c567f0e5378b725b8ee1226bf554e8ecafa72943a4e7b99a0526", + "speculator": "mtp" }, { "id": "basecompute/Qwen3.8-27B", @@ -598,8 +617,9 @@ "source_repo": "Qwen/Qwen3.8-27B", "arch": "qwen35", "quant": "default-q8", - "size": 28386443264, - "sha256": "6d4d362bf410c55507e894c132e03ecc0b5fbdb038b09f999ed99f59569a79d0" + "size": 28934225920, + "sha256": "62cf0e6caa6270bcd3964963a259ac850c76a054e010eb19e50322fc30d8559e", + "speculator": "mtp" }, { "id": "basecompute/gemma-3-1b-it", @@ -1198,6 +1218,818 @@ "quant": "bf16", "size": 63165857792, "sha256": "f1bbee6ad05870c33d09a4ffeb2f6480b2fe13e34d4d8aaead8b302f115535e1" + }, + { + "id": "basecompute/NVIDIA-Nemotron-3-Super-120B-A12B-DFlash-RedHat", + "hf_repo": "basecompute/NVIDIA-Nemotron-3-Super-120B-A12B-DFlash-RedHat", + "file": "NVIDIA-Nemotron-3-Super-120B-A12B-DFlash-RedHat-F16.base", + "revision": "main", + "arch": "dflash", + "quant": "default-f16", + "size": 2191024128, + "sha256": "ce9d0f8f9c19de363402634f6242931482fecb030e7469d035d13640572b86ab", + "speculator": "dflash", + "source_repo": "RedHatAI/NVIDIA-Nemotron-3-Super-120B-A12B-speculator.dflash", + "speculator_for": "basecompute/NVIDIA-Nemotron-3-Super-120B-A12B", + "speculator_rank": 1 + }, + { + "id": "basecompute/NVIDIA-Nemotron-3-Super-120B-A12B-DFlash-RedHat", + "hf_repo": "basecompute/NVIDIA-Nemotron-3-Super-120B-A12B-DFlash-RedHat", + "file": "NVIDIA-Nemotron-3-Super-120B-A12B-DFlash-RedHat-Q4.base", + "revision": "main", + "arch": "dflash", + "quant": "default-q4", + "size": 1697374208, + "sha256": "70d90cca4b4f0950cfa7988553f578c8fc83b36ec53d7f2556a55fd69ca36470", + "speculator": "dflash", + "source_repo": "RedHatAI/NVIDIA-Nemotron-3-Super-120B-A12B-speculator.dflash", + "speculator_for": "basecompute/NVIDIA-Nemotron-3-Super-120B-A12B", + "speculator_rank": 1 + }, + { + "id": "basecompute/Qwen3-1.7B-EAGLE3", + "hf_repo": "basecompute/Qwen3-1.7B-EAGLE3", + "file": "Qwen3-1.7B-EAGLE3-F16.base", + "revision": "main", + "source_repo": "AngelSlim/Qwen3-1.7B_eagle3", + "arch": "eagle3", + "quant": "default-f16", + "size": 273928192, + "sha256": "9dc566c5c83c2dc60afc0e25daf7879790ce266f965ec7b51b23db7f8b0e093e", + "speculator": "eagle3", + "speculator_for": "basecompute/Qwen3-1.7B", + "speculator_rank": 1 + }, + { + "id": "basecompute/Qwen3-1.7B-EAGLE3", + "hf_repo": "basecompute/Qwen3-1.7B-EAGLE3", + "file": "Qwen3-1.7B-EAGLE3-Q4.base", + "revision": "main", + "source_repo": "AngelSlim/Qwen3-1.7B_eagle3", + "arch": "eagle3", + "quant": "default-q4", + "size": 189517824, + "sha256": "164f409e98b94e24335146c634d947322ef4b2a5198b45d9048bbc9e22236cb9", + "speculator": "eagle3", + "speculator_for": "basecompute/Qwen3-1.7B", + "speculator_rank": 1 + }, + { + "id": "basecompute/Qwen3-30B-A3B-DFlash-RedHat", + "hf_repo": "basecompute/Qwen3-30B-A3B-DFlash-RedHat", + "file": "Qwen3-30B-A3B-DFlash-RedHat-F16.base", + "revision": "main", + "arch": "dflash", + "quant": "default-f16", + "size": 1362120704, + "sha256": "f23059e95ec09f74454eed44597105ed5777311ceee718a331a3ee360e31d84f", + "speculator": "dflash", + "source_repo": "RedHatAI/Qwen3-30B-A3B-speculator.dflash", + "speculator_for": "basecompute/Qwen3-30B-A3B", + "speculator_rank": 1 + }, + { + "id": "basecompute/Qwen3-30B-A3B-DFlash-RedHat", + "hf_repo": "basecompute/Qwen3-30B-A3B-DFlash-RedHat", + "file": "Qwen3-30B-A3B-DFlash-RedHat-Q4.base", + "revision": "main", + "arch": "dflash", + "quant": "default-q4", + "size": 955142144, + "sha256": "f6c25d5b1c878c6263288bf7856b4d81734960c815d8b39e2f6064683c49ddbf", + "speculator": "dflash", + "source_repo": "RedHatAI/Qwen3-30B-A3B-speculator.dflash", + "speculator_for": "basecompute/Qwen3-30B-A3B", + "speculator_rank": 1 + }, + { + "id": "basecompute/Qwen3-30B-A3B-EAGLE3-RedHat", + "hf_repo": "basecompute/Qwen3-30B-A3B-EAGLE3-RedHat", + "file": "Qwen3-30B-A3B-EAGLE3-RedHat-F16.base", + "revision": "main", + "arch": "eagle3", + "quant": "default-f16", + "size": 1044238336, + "sha256": "2ef75d064dab80888210e6653289499c7d2bd45b34fc4913ff62163888b960bd", + "speculator": "eagle3", + "source_repo": "RedHatAI/Qwen3-30B-A3B-speculator.eagle3", + "speculator_for": "basecompute/Qwen3-30B-A3B", + "speculator_rank": 2 + }, + { + "id": "basecompute/Qwen3-30B-A3B-EAGLE3-RedHat", + "hf_repo": "basecompute/Qwen3-30B-A3B-EAGLE3-RedHat", + "file": "Qwen3-30B-A3B-EAGLE3-RedHat-Q4.base", + "revision": "main", + "arch": "eagle3", + "quant": "default-q4", + "size": 947769344, + "sha256": "0d59a69d4430ddb812d9bb8e210582bc1961e2bbe0023d04d4ee7cce40ecdf71", + "speculator": "eagle3", + "source_repo": "RedHatAI/Qwen3-30B-A3B-speculator.eagle3", + "speculator_for": "basecompute/Qwen3-30B-A3B", + "speculator_rank": 2 + }, + { + "id": "basecompute/Qwen3-30B-A3B-Instruct-2507-DFlash-RedHat", + "hf_repo": "basecompute/Qwen3-30B-A3B-Instruct-2507-DFlash-RedHat", + "file": "Qwen3-30B-A3B-Instruct-2507-DFlash-RedHat-F16.base", + "revision": "main", + "arch": "dflash", + "quant": "default-f16", + "size": 1362120704, + "sha256": "50a53bbd3ce60734693eeac3328404e74e2bf43acd2332953a58997c82960cf2", + "speculator": "dflash", + "source_repo": "RedHatAI/Qwen3-30B-A3B-Instruct-2507-speculator.dflash", + "speculator_for": "basecompute/Qwen3-30B-A3B-Instruct-2507", + "speculator_rank": 1 + }, + { + "id": "basecompute/Qwen3-30B-A3B-Instruct-2507-DFlash-RedHat", + "hf_repo": "basecompute/Qwen3-30B-A3B-Instruct-2507-DFlash-RedHat", + "file": "Qwen3-30B-A3B-Instruct-2507-DFlash-RedHat-Q4.base", + "revision": "main", + "arch": "dflash", + "quant": "default-q4", + "size": 955142144, + "sha256": "de2598e92f2f6cedae7a2afd7a6d9ff0552a6317add376e5a45279c4088154c5", + "speculator": "dflash", + "source_repo": "RedHatAI/Qwen3-30B-A3B-Instruct-2507-speculator.dflash", + "speculator_for": "basecompute/Qwen3-30B-A3B-Instruct-2507", + "speculator_rank": 1 + }, + { + "id": "basecompute/Qwen3-30B-A3B-Instruct-2507-EAGLE3-RedHat", + "hf_repo": "basecompute/Qwen3-30B-A3B-Instruct-2507-EAGLE3-RedHat", + "file": "Qwen3-30B-A3B-Instruct-2507-EAGLE3-RedHat-F16.base", + "revision": "main", + "arch": "eagle3", + "quant": "default-f16", + "size": 1044238336, + "sha256": "fb870cddc7e471c889f95f29d91ca97b642aaf862973587531cc294d557a5171", + "speculator": "eagle3", + "source_repo": "RedHatAI/Qwen3-30B-A3B-Instruct-2507-speculator.eagle3", + "speculator_for": "basecompute/Qwen3-30B-A3B-Instruct-2507", + "speculator_rank": 2 + }, + { + "id": "basecompute/Qwen3-30B-A3B-Instruct-2507-EAGLE3-RedHat", + "hf_repo": "basecompute/Qwen3-30B-A3B-Instruct-2507-EAGLE3-RedHat", + "file": "Qwen3-30B-A3B-Instruct-2507-EAGLE3-RedHat-Q4.base", + "revision": "main", + "arch": "eagle3", + "quant": "default-q4", + "size": 947769344, + "sha256": "da9ed5952ad03424665b0cab288b6ba576a18d9090c3ef03a839f16a3c541cfb", + "speculator": "eagle3", + "source_repo": "RedHatAI/Qwen3-30B-A3B-Instruct-2507-speculator.eagle3", + "speculator_for": "basecompute/Qwen3-30B-A3B-Instruct-2507", + "speculator_rank": 2 + }, + { + "id": "basecompute/Qwen3-30B-A3B-Thinking-2507-EAGLE3-RedHat", + "hf_repo": "basecompute/Qwen3-30B-A3B-Thinking-2507-EAGLE3-RedHat", + "file": "Qwen3-30B-A3B-Thinking-2507-EAGLE3-RedHat-F16.base", + "revision": "main", + "arch": "eagle3", + "quant": "default-f16", + "size": 1044238336, + "sha256": "db4398979639a8704119681918a340be865505c4f8f906c2c9de2a1678d95f8e", + "speculator": "eagle3", + "source_repo": "RedHatAI/Qwen3-30B-A3B-Thinking-2507-speculator.eagle3", + "speculator_for": "basecompute/Qwen3-30B-A3B-Thinking-2507", + "speculator_rank": 1 + }, + { + "id": "basecompute/Qwen3-30B-A3B-Thinking-2507-EAGLE3-RedHat", + "hf_repo": "basecompute/Qwen3-30B-A3B-Thinking-2507-EAGLE3-RedHat", + "file": "Qwen3-30B-A3B-Thinking-2507-EAGLE3-RedHat-Q4.base", + "revision": "main", + "arch": "eagle3", + "quant": "default-q4", + "size": 947769344, + "sha256": "124dbcf615eeda15c54755c019edd6cd819aa592fc79b0eef0492c194fe97cb9", + "speculator": "eagle3", + "source_repo": "RedHatAI/Qwen3-30B-A3B-Thinking-2507-speculator.eagle3", + "speculator_for": "basecompute/Qwen3-30B-A3B-Thinking-2507", + "speculator_rank": 1 + }, + { + "id": "basecompute/Qwen3-4B-DFlash", + "hf_repo": "basecompute/Qwen3-4B-DFlash", + "file": "Qwen3-4B-DFlash-F16.base", + "revision": "main", + "source_repo": "z-lab/Qwen3-4B-DFlash-b16", + "arch": "dflash", + "quant": "default-f16", + "size": 1075205120, + "sha256": "9d7154cb8c18c923d98dbf8d3ba7e5a7d674f2dbc22ddf785cc9a5c7282c0caf", + "speculator": "dflash", + "speculator_for": "basecompute/Qwen3-4B", + "speculator_rank": 3 + }, + { + "id": "basecompute/Qwen3-4B-DFlash", + "hf_repo": "basecompute/Qwen3-4B-DFlash", + "file": "Qwen3-4B-DFlash-Q4.base", + "revision": "main", + "source_repo": "z-lab/Qwen3-4B-DFlash-b16", + "arch": "dflash", + "quant": "default-q4", + "size": 349803520, + "sha256": "6fb99f09fe2c8c13773621da093a440b83ffc919986ab4060d99012fc9ad7912", + "speculator": "dflash", + "speculator_for": "basecompute/Qwen3-4B", + "speculator_rank": 3 + }, + { + "id": "basecompute/Qwen3-4B-DFlash-block7", + "hf_repo": "basecompute/Qwen3-4B-DFlash-block7", + "file": "Qwen3-4B-DFlash-block7-F16.base", + "revision": "main", + "source_repo": "deepseek-ai/dflash_qwen3_4b_block7", + "arch": "dspark", + "quant": "default-f16", + "size": 1075205120, + "sha256": "d773f1f65d9ee3d836ce551c3df346ddf4198b9ff49defab7309461ce1a0c3f8", + "speculator": "dspark", + "speculator_for": "basecompute/Qwen3-4B", + "speculator_rank": 5 + }, + { + "id": "basecompute/Qwen3-4B-DFlash-block7", + "hf_repo": "basecompute/Qwen3-4B-DFlash-block7", + "file": "Qwen3-4B-DFlash-block7-Q4.base", + "revision": "main", + "source_repo": "deepseek-ai/dflash_qwen3_4b_block7", + "arch": "dspark", + "quant": "default-q4", + "size": 349803520, + "sha256": "d125084052fd5555b83b80b97ce4f2a76187178b29a4ca89c0704b6da6923116", + "speculator": "dspark", + "speculator_for": "basecompute/Qwen3-4B", + "speculator_rank": 5 + }, + { + "id": "basecompute/Qwen3-4B-DSpark", + "hf_repo": "basecompute/Qwen3-4B-DSpark", + "file": "Qwen3-4B-DSpark-F16.base", + "revision": "main", + "source_repo": "deepseek-ai/dspark_qwen3_4b_block7", + "arch": "dspark", + "quant": "default-f16", + "size": 1230820352, + "sha256": "3b76715f6986cacbd4f62a9c92b5c59d1198ee311f57ac33365bb74673057e01", + "speculator": "dspark", + "speculator_for": "basecompute/Qwen3-4B", + "speculator_rank": 2 + }, + { + "id": "basecompute/Qwen3-4B-DSpark", + "hf_repo": "basecompute/Qwen3-4B-DSpark", + "file": "Qwen3-4B-DSpark-Q4.base", + "revision": "main", + "source_repo": "deepseek-ai/dspark_qwen3_4b_block7", + "arch": "dspark", + "quant": "default-q4", + "size": 505418752, + "sha256": "f7764fbb5dba8df0f31cf734ed89e8289a927357d7c2d6ce7fa9811a6bff848f", + "speculator": "dspark", + "speculator_for": "basecompute/Qwen3-4B", + "speculator_rank": 2 + }, + { + "id": "basecompute/Qwen3-4B-DSpark-RedHat", + "hf_repo": "basecompute/Qwen3-4B-DSpark-RedHat", + "file": "Qwen3-4B-DSpark-RedHat-F16.base", + "revision": "main", + "arch": "dspark", + "quant": "default-f16", + "size": 1230820352, + "sha256": "bcd29cc89ddb001f599c31a20c847383ad9e47cbc8590642badb826fb923b1ff", + "speculator": "dspark", + "source_repo": "RedHatAI/Qwen3-4B-speculator.dspark", + "speculator_for": "basecompute/Qwen3-4B", + "speculator_rank": 1 + }, + { + "id": "basecompute/Qwen3-4B-DSpark-RedHat", + "hf_repo": "basecompute/Qwen3-4B-DSpark-RedHat", + "file": "Qwen3-4B-DSpark-RedHat-Q4.base", + "revision": "main", + "arch": "dspark", + "quant": "default-q4", + "size": 505418752, + "sha256": "ebcea1fe37b68583ea1b1d761e45a4700fd8b3f0d6acff4d024066624172733d", + "speculator": "dspark", + "source_repo": "RedHatAI/Qwen3-4B-speculator.dspark", + "speculator_for": "basecompute/Qwen3-4B", + "speculator_rank": 1 + }, + { + "id": "basecompute/Qwen3-4B-EAGLE3", + "hf_repo": "basecompute/Qwen3-4B-EAGLE3", + "file": "Qwen3-4B-EAGLE3-F16.base", + "revision": "main", + "source_repo": "AngelSlim/Qwen3-4B_eagle3", + "arch": "eagle3", + "quant": "default-f16", + "size": 436720640, + "sha256": "1accf8f1a4028f4511f4b4bc824f28b2872b496abf180298bd9246d682302fbd", + "speculator": "eagle3", + "speculator_for": "basecompute/Qwen3-4B", + "speculator_rank": 4 + }, + { + "id": "basecompute/Qwen3-4B-EAGLE3", + "hf_repo": "basecompute/Qwen3-4B-EAGLE3", + "file": "Qwen3-4B-EAGLE3-Q4.base", + "revision": "main", + "source_repo": "AngelSlim/Qwen3-4B_eagle3", + "arch": "eagle3", + "quant": "default-q4", + "size": 269030400, + "sha256": "815dc165e3ca4578f1f21ac5ca91e768a2090bb60d1868be3c16e7b6ddd5718e", + "speculator": "eagle3", + "speculator_for": "basecompute/Qwen3-4B", + "speculator_rank": 4 + }, + { + "id": "basecompute/Qwen3-8B-DFlash", + "hf_repo": "basecompute/Qwen3-8B-DFlash", + "file": "Qwen3-8B-DFlash-F16.base", + "revision": "main", + "source_repo": "z-lab/Qwen3-8B-DFlash-b16", + "arch": "dflash", + "quant": "default-f16", + "size": 2097569792, + "sha256": "caf90ccbec6bb8fba5db4a82782e9e9c79170477b7f2245777a52b9d3d00495f", + "speculator": "dflash", + "speculator_for": "basecompute/Qwen3-8B", + "speculator_rank": 5 + }, + { + "id": "basecompute/Qwen3-8B-DFlash", + "hf_repo": "basecompute/Qwen3-8B-DFlash", + "file": "Qwen3-8B-DFlash-Q4.base", + "revision": "main", + "source_repo": "z-lab/Qwen3-8B-DFlash-b16", + "arch": "dflash", + "quant": "default-q4", + "size": 710828032, + "sha256": "c54556345758193bd4dc850d2a78d2e019d35e59359e05f691f329c8013b8aea", + "speculator": "dflash", + "speculator_for": "basecompute/Qwen3-8B", + "speculator_rank": 5 + }, + { + "id": "basecompute/Qwen3-8B-DFlash-RedHat", + "hf_repo": "basecompute/Qwen3-8B-DFlash-RedHat", + "file": "Qwen3-8B-DFlash-RedHat-F16.base", + "revision": "main", + "arch": "dflash", + "quant": "default-f16", + "size": 4586889216, + "sha256": "7868da707a312fbb476e232bf381e0e0e4468c5aa050f84fa2a21798ef3f4b7d", + "speculator": "dflash", + "source_repo": "RedHatAI/Qwen3-8B-speculator.dflash", + "speculator_for": "basecompute/Qwen3-8B", + "speculator_rank": 4 + }, + { + "id": "basecompute/Qwen3-8B-DFlash-RedHat", + "hf_repo": "basecompute/Qwen3-8B-DFlash-RedHat", + "file": "Qwen3-8B-DFlash-RedHat-Q4.base", + "revision": "main", + "arch": "dflash", + "quant": "default-q4", + "size": 3200147456, + "sha256": "25e57c4a86fd62acdf279fe551ef8166711a9c1d84f2e527a4b7e71f1dbfbfbb", + "speculator": "dflash", + "source_repo": "RedHatAI/Qwen3-8B-speculator.dflash", + "speculator_for": "basecompute/Qwen3-8B", + "speculator_rank": 4 + }, + { + "id": "basecompute/Qwen3-8B-DFlash-block7", + "hf_repo": "basecompute/Qwen3-8B-DFlash-block7", + "file": "Qwen3-8B-DFlash-block7-F16.base", + "revision": "main", + "source_repo": "deepseek-ai/dflash_qwen3_8b_block7", + "arch": "dspark", + "quant": "default-f16", + "size": 4586889216, + "sha256": "a4a137841f79dca6f6c3bd9b3285511a7a624f259a94216ba80a85c8401d2653", + "speculator": "dspark", + "speculator_for": "basecompute/Qwen3-8B", + "speculator_rank": 3 + }, + { + "id": "basecompute/Qwen3-8B-DFlash-block7", + "hf_repo": "basecompute/Qwen3-8B-DFlash-block7", + "file": "Qwen3-8B-DFlash-block7-Q4.base", + "revision": "main", + "source_repo": "deepseek-ai/dflash_qwen3_8b_block7", + "arch": "dspark", + "quant": "default-q4", + "size": 3200147456, + "sha256": "2e659d8818f25dc01c7bdf68ffa1ccbc96c7fcda4c3c30d220046f169bbe10be", + "speculator": "dspark", + "speculator_for": "basecompute/Qwen3-8B", + "speculator_rank": 3 + }, + { + "id": "basecompute/Qwen3-8B-DFlash2-RedHat", + "hf_repo": "basecompute/Qwen3-8B-DFlash2-RedHat", + "file": "Qwen3-8B-DFlash2-RedHat-F16.base", + "revision": "main", + "arch": "dflash", + "quant": "default-f16", + "size": 2339463168, + "sha256": "5c4c29a5d22fe82a2ad5fe9655786ec86b43ca16f3e67ac12783b33716267c6f", + "speculator": "dflash", + "source_repo": "RedHatAI/Qwen3-8B-speculator.dflash2", + "speculator_for": "basecompute/Qwen3-8B", + "speculator_rank": 2 + }, + { + "id": "basecompute/Qwen3-8B-DFlash2-RedHat", + "hf_repo": "basecompute/Qwen3-8B-DFlash2-RedHat", + "file": "Qwen3-8B-DFlash2-RedHat-Q4.base", + "revision": "main", + "arch": "dflash", + "quant": "default-q4", + "size": 892428288, + "sha256": "73e85feeb49b14e734f9fab1461e621c55d125d4da70e4d65a7a152472c48062", + "speculator": "dflash", + "source_repo": "RedHatAI/Qwen3-8B-speculator.dflash2", + "speculator_for": "basecompute/Qwen3-8B", + "speculator_rank": 2 + }, + { + "id": "basecompute/Qwen3-8B-DSpark", + "hf_repo": "basecompute/Qwen3-8B-DSpark", + "file": "Qwen3-8B-DSpark-F16.base", + "revision": "main", + "source_repo": "deepseek-ai/dspark_qwen3_8b_block7", + "arch": "dspark", + "quant": "default-f16", + "size": 4742504448, + "sha256": "c0f17a84f474fa25b053414baf47a46c78c4585fce1ca04b67dd7e0bfd91597d", + "speculator": "dspark", + "speculator_for": "basecompute/Qwen3-8B", + "speculator_rank": 1 + }, + { + "id": "basecompute/Qwen3-8B-DSpark", + "hf_repo": "basecompute/Qwen3-8B-DSpark", + "file": "Qwen3-8B-DSpark-Q4.base", + "revision": "main", + "source_repo": "deepseek-ai/dspark_qwen3_8b_block7", + "arch": "dspark", + "quant": "default-q4", + "size": 3355762688, + "sha256": "5701001b4c818ee8f8ac5b45e03ff626542ba4d78c2ee4632f5068b2caad21ff", + "speculator": "dspark", + "speculator_for": "basecompute/Qwen3-8B", + "speculator_rank": 1 + }, + { + "id": "basecompute/Qwen3-8B-EAGLE3", + "hf_repo": "basecompute/Qwen3-8B-EAGLE3", + "file": "Qwen3-8B-EAGLE3-F16.base", + "revision": "main", + "source_repo": "AngelSlim/Qwen3-8B_eagle3", + "arch": "eagle3", + "quant": "default-f16", + "size": 799268864, + "sha256": "2649a654f64f9100d1473b611a0dec20fdecbcdc10bca96d69ea590b13df7d94", + "speculator": "eagle3", + "speculator_for": "basecompute/Qwen3-8B", + "speculator_rank": 7 + }, + { + "id": "basecompute/Qwen3-8B-EAGLE3", + "hf_repo": "basecompute/Qwen3-8B-EAGLE3", + "file": "Qwen3-8B-EAGLE3-Q4.base", + "revision": "main", + "source_repo": "AngelSlim/Qwen3-8B_eagle3", + "arch": "eagle3", + "quant": "default-q4", + "size": 485744640, + "sha256": "09881137e7be7bd4d09f3a80ad4ce6a4ad0c2fe5d8213ad2c47966d98072bf4f", + "speculator": "eagle3", + "speculator_for": "basecompute/Qwen3-8B", + "speculator_rank": 7 + }, + { + "id": "basecompute/Qwen3-8B-EAGLE3-RedHat", + "hf_repo": "basecompute/Qwen3-8B-EAGLE3-RedHat", + "file": "Qwen3-8B-EAGLE3-RedHat-F16.base", + "revision": "main", + "arch": "eagle3", + "quant": "default-f16", + "size": 2043928576, + "sha256": "640e3c3a2db497d923604fdfb596fea94b4df116e10334071fc8545f88710ece", + "speculator": "eagle3", + "source_repo": "RedHatAI/Qwen3-8B-speculator.eagle3", + "speculator_for": "basecompute/Qwen3-8B", + "speculator_rank": 6 + }, + { + "id": "basecompute/Qwen3-8B-EAGLE3-RedHat", + "hf_repo": "basecompute/Qwen3-8B-EAGLE3-RedHat", + "file": "Qwen3-8B-EAGLE3-RedHat-Q4.base", + "revision": "main", + "arch": "eagle3", + "quant": "default-q4", + "size": 1730404352, + "sha256": "7975f3c2e332818467971cb66f6c57942e43eba732ac41f0b97e5d0349c403fd", + "speculator": "eagle3", + "source_repo": "RedHatAI/Qwen3-8B-speculator.eagle3", + "speculator_for": "basecompute/Qwen3-8B", + "speculator_rank": 6 + }, + { + "id": "basecompute/Qwen3.5-35B-A3B-DFlash", + "hf_repo": "basecompute/Qwen3.5-35B-A3B-DFlash", + "file": "Qwen3.5-35B-A3B-DFlash-F16.base", + "revision": "main", + "source_repo": "z-lab/Qwen3.5-35B-A3B-DFlash", + "arch": "dflash", + "quant": "default-f16", + "size": 772231168, + "sha256": "8d310eb1165e7396fc473341ca40433d133c5db5d202338292f76f907d2ed793", + "speculator": "dflash", + "speculator_for": "basecompute/Qwen3.5-35B-A3B", + "speculator_rank": 1 + }, + { + "id": "basecompute/Qwen3.5-35B-A3B-DFlash", + "hf_repo": "basecompute/Qwen3.5-35B-A3B-DFlash", + "file": "Qwen3.5-35B-A3B-DFlash-Q4.base", + "revision": "main", + "source_repo": "z-lab/Qwen3.5-35B-A3B-DFlash", + "arch": "dflash", + "quant": "default-q4", + "size": 265768960, + "sha256": "0a25730378642519cc3aa23c3968d87a659e3b0114461460dd16e9801b741352", + "speculator": "dflash", + "speculator_for": "basecompute/Qwen3.5-35B-A3B", + "speculator_rank": 1 + }, + { + "id": "basecompute/Qwen3.5-4B-DFlash", + "hf_repo": "basecompute/Qwen3.5-4B-DFlash", + "file": "Qwen3.5-4B-DFlash-F16.base", + "revision": "main", + "arch": "dflash", + "quant": "default-f16", + "size": 1269257216, + "sha256": "e43a83e9d6d10c7244e396a2428a219b37769bb504788b90553dcc3e32023a41", + "speculator": "dflash", + "source_repo": "z-lab/Qwen3.5-4B-DFlash", + "speculator_for": "basecompute/Qwen3.5-4B", + "speculator_rank": 1 + }, + { + "id": "basecompute/Qwen3.5-4B-DFlash", + "hf_repo": "basecompute/Qwen3.5-4B-DFlash", + "file": "Qwen3.5-4B-DFlash-Q4.base", + "revision": "main", + "arch": "dflash", + "quant": "default-q4", + "size": 432690176, + "sha256": "4452112ed55e8bb7e5e7852de6959ef266a5fcb5fc7ed60979a257370642fd9d", + "speculator": "dflash", + "source_repo": "z-lab/Qwen3.5-4B-DFlash", + "speculator_for": "basecompute/Qwen3.5-4B", + "speculator_rank": 1 + }, + { + "id": "basecompute/Qwen3.6-27B-DFlash", + "hf_repo": "basecompute/Qwen3.6-27B-DFlash", + "file": "Qwen3.6-27B-DFlash-F16.base", + "revision": "main", + "source_repo": "z-lab/Qwen3.6-27B-DFlash", + "arch": "dflash", + "quant": "default-f16", + "size": 3460720640, + "sha256": "2357b4cf4c0f9d3958df9fcc53416308d53b7e0319e64017ccd4a9d888a990e5", + "speculator": "dflash", + "speculator_for": "basecompute/Qwen3.6-27B", + "speculator_rank": 1 + }, + { + "id": "basecompute/Qwen3.6-27B-DFlash", + "hf_repo": "basecompute/Qwen3.6-27B-DFlash", + "file": "Qwen3.6-27B-DFlash-Q4.base", + "revision": "main", + "source_repo": "z-lab/Qwen3.6-27B-DFlash", + "arch": "dflash", + "quant": "default-q4", + "size": 1162045440, + "sha256": "cded6acf2bc0b5053ee3b9b82b10a83a16eadc3f840eedc30e33c7638bedb532", + "speculator": "dflash", + "speculator_for": "basecompute/Qwen3.6-27B", + "speculator_rank": 1 + }, + { + "id": "basecompute/Qwen3.6-35B-A3B-DFlash", + "hf_repo": "basecompute/Qwen3.6-35B-A3B-DFlash", + "file": "Qwen3.6-35B-A3B-DFlash-F16.base", + "revision": "main", + "source_repo": "z-lab/Qwen3.6-35B-A3B-DFlash", + "arch": "dflash", + "quant": "default-f16", + "size": 772231168, + "sha256": "44989d3f2129ec5fcd60cb3068b9126247de7577fc6fb7ac216184dbc2c149c3", + "speculator": "dflash", + "speculator_for": "basecompute/Qwen3.6-35B-A3B", + "speculator_rank": 1 + }, + { + "id": "basecompute/Qwen3.6-35B-A3B-DFlash", + "hf_repo": "basecompute/Qwen3.6-35B-A3B-DFlash", + "file": "Qwen3.6-35B-A3B-DFlash-Q4.base", + "revision": "main", + "source_repo": "z-lab/Qwen3.6-35B-A3B-DFlash", + "arch": "dflash", + "quant": "default-q4", + "size": 265768960, + "sha256": "56a5bc65f7019d08f231de8f51efb091c066fc9490c026d9188109c51921a485", + "speculator": "dflash", + "speculator_for": "basecompute/Qwen3.6-35B-A3B", + "speculator_rank": 1 + }, + { + "id": "basecompute/Qwen3.8-27B-DFlash2", + "hf_repo": "basecompute/Qwen3.8-27B-DFlash2", + "file": "Qwen3.8-27B-DFlash2-F16.base", + "revision": "main", + "arch": "dflash", + "quant": "default-f16", + "size": 3849185280, + "sha256": "3d56eae29f63180d25ebee784483bc04a68123068cbde8526f8c48402cf0390e", + "speculator": "dflash", + "source_repo": "z-lab/Qwen3.8-27B-DFlash2", + "speculator_for": "basecompute/Qwen3.8-27B", + "speculator_rank": 1 + }, + { + "id": "basecompute/Qwen3.8-27B-DFlash2", + "hf_repo": "basecompute/Qwen3.8-27B-DFlash2", + "file": "Qwen3.8-27B-DFlash2-Q4.base", + "revision": "main", + "arch": "dflash", + "quant": "default-q4", + "size": 1456302080, + "sha256": "e6ec0261177cba95a85216a6c62c0388472a40db601555c0fe9b0b58d26606d6", + "speculator": "dflash", + "source_repo": "z-lab/Qwen3.8-27B-DFlash2", + "speculator_for": "basecompute/Qwen3.8-27B", + "speculator_rank": 1 + }, + { + "id": "basecompute/Qwen3.8-27B-DSpark", + "hf_repo": "basecompute/Qwen3.8-27B-DSpark", + "file": "Qwen3.8-27B-DSpark-F16.base", + "revision": "main", + "arch": "dspark", + "quant": "default-f16", + "size": 3715033088, + "sha256": "e5e538ec4bec67766f35b72ccaedf84817b79ab5fc534b8ca00d54c10b077ce3", + "speculator": "dspark", + "source_repo": "RadixArk/Qwen3.8-27B-DSpark", + "speculator_for": "basecompute/Qwen3.8-27B", + "speculator_rank": 2 + }, + { + "id": "basecompute/Qwen3.8-27B-DSpark", + "hf_repo": "basecompute/Qwen3.8-27B-DSpark", + "file": "Qwen3.8-27B-DSpark-Q4.base", + "revision": "main", + "arch": "dspark", + "quant": "default-q4", + "size": 1416357888, + "sha256": "6c50c6614e68c91d59e5cab1ae988712b7b2ffde95a3cdc43eddb96f9b239360", + "speculator": "dspark", + "source_repo": "RadixArk/Qwen3.8-27B-DSpark", + "speculator_for": "basecompute/Qwen3.8-27B", + "speculator_rank": 2 + }, + { + "id": "basecompute/gemma-4-12b-it-DFlash", + "hf_repo": "basecompute/gemma-4-12b-it-DFlash", + "file": "gemma-4-12b-it-DFlash-F16.base", + "revision": "main", + "arch": "dflash", + "quant": "default-f16", + "size": 1455316480, + "sha256": "486ec750c1081bf9748bf28c026c68b9102ce7749292c11952f73b01182f606a", + "speculator": "dflash", + "source_repo": "z-lab/gemma4-12B-it-DFlash", + "speculator_for": "basecompute/gemma-4-12b-it", + "speculator_rank": 1 + }, + { + "id": "basecompute/gemma-4-12b-it-DFlash", + "hf_repo": "basecompute/gemma-4-12b-it-DFlash", + "file": "gemma-4-12b-it-DFlash-Q4.base", + "revision": "main", + "arch": "dflash", + "quant": "default-q4", + "size": 536911360, + "sha256": "e794c62b324c553eb532ffc66629126ddd29e2e00db71a6c1de6196663218349", + "speculator": "dflash", + "source_repo": "z-lab/gemma4-12B-it-DFlash", + "speculator_for": "basecompute/gemma-4-12b-it", + "speculator_rank": 1 + }, + { + "id": "basecompute/gemma-4-12b-it-DSpark", + "hf_repo": "basecompute/gemma-4-12b-it-DSpark", + "file": "gemma-4-12b-it-DSpark-F16.base", + "revision": "main", + "arch": "dspark", + "quant": "default-f16", + "size": 2834865664, + "sha256": "9fed56f9125286b1ac6b752e5b2bbb8735afb08d2512e0cf68a067e7e0b9863c", + "speculator": "dspark", + "source_repo": "deepseek-ai/dspark_gemma4_12b_block7", + "speculator_for": "basecompute/gemma-4-12b-it", + "speculator_rank": 2 + }, + { + "id": "basecompute/gemma-4-12b-it-DSpark", + "hf_repo": "basecompute/gemma-4-12b-it-DSpark", + "file": "gemma-4-12b-it-DSpark-Q4.base", + "revision": "main", + "arch": "dspark", + "quant": "default-q4", + "size": 1096769024, + "sha256": "fbf742e6e08ad05344817a254332a9b718f88f7a4d25a22127983f05b6ba01f2", + "speculator": "dspark", + "source_repo": "deepseek-ai/dspark_gemma4_12b_block7", + "speculator_for": "basecompute/gemma-4-12b-it", + "speculator_rank": 2 + }, + { + "id": "basecompute/gpt-oss-20b-DFlash", + "hf_repo": "basecompute/gpt-oss-20b-DFlash", + "file": "gpt-oss-20b-DFlash-F16.base", + "revision": "main", + "arch": "dflash", + "quant": "default-f16", + "size": 1570477696, + "sha256": "2ac654af8078cb44ec548f4709621a69b366072686104f43040dc70d7e1d3836", + "speculator": "dflash", + "source_repo": "z-lab/gpt-oss-20b-DFlash", + "speculator_for": "basecompute/gpt-oss-20b", + "speculator_rank": 2 + }, + { + "id": "basecompute/gpt-oss-20b-DFlash", + "hf_repo": "basecompute/gpt-oss-20b-DFlash", + "file": "gpt-oss-20b-DFlash-Q4.base", + "revision": "main", + "arch": "dflash", + "quant": "default-q4", + "size": 502503040, + "sha256": "1652d9d00ba08c52db82b8b950fc91a80650de4286312318521236e6a8d14c92", + "speculator": "dflash", + "source_repo": "z-lab/gpt-oss-20b-DFlash", + "speculator_for": "basecompute/gpt-oss-20b", + "speculator_rank": 2 + }, + { + "id": "basecompute/gpt-oss-20b-EAGLE3", + "hf_repo": "basecompute/gpt-oss-20b-EAGLE3", + "file": "gpt-oss-20b-EAGLE3-F16.base", + "revision": "main", + "arch": "eagle3", + "quant": "default-f16", + "size": 1709479552, + "sha256": "c3e0cd91d5e7bd46f48a40faf9527e96c2e2d5f8d6551b0386a59d3f2d2ba9ea", + "speculator": "eagle3", + "source_repo": "RedHatAI/gpt-oss-20b-speculator.eagle3", + "speculator_for": "basecompute/gpt-oss-20b", + "speculator_rank": 1 + }, + { + "id": "basecompute/gpt-oss-20b-EAGLE3", + "hf_repo": "basecompute/gpt-oss-20b-EAGLE3", + "file": "gpt-oss-20b-EAGLE3-Q4.base", + "revision": "main", + "arch": "eagle3", + "quant": "default-q4", + "size": 1614370432, + "sha256": "030b0d05cf67b6467d83eb874800a9b6034aa9e9a88383b2ff3fd57305dc77ec", + "speculator": "eagle3", + "source_repo": "RedHatAI/gpt-oss-20b-speculator.eagle3", + "speculator_for": "basecompute/gpt-oss-20b", + "speculator_rank": 1 } ] } diff --git a/base-convert/crates/base-hub/src/catalog.rs b/base-convert/crates/base-hub/src/catalog.rs index 3096998..11c86e4 100644 --- a/base-convert/crates/base-hub/src/catalog.rs +++ b/base-convert/crates/base-hub/src/catalog.rs @@ -83,6 +83,25 @@ pub struct CatalogEntry { /// not at load after 30 GB. #[serde(default, skip_serializing_if = "Option::is_none")] pub backend: Option, + /// Speculative-decoding role. On a target bundle: the in-bundle + /// speculator it carries (`"mtp"` — `basertd --speculate mtp-head` + /// works out of the box). On a drafter sidecar row: its kind + /// (`"dflash"` / `"dspark"` / `"eagle3"`), served with + /// `--speculate :`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub speculator: Option, + /// Sidecar rows only: the catalog id of the target the drafter was + /// trained for (it shares that target's embedding / lm_head and taps + /// its hidden states). `basertd --speculate auto` pairs an installed + /// drafter with its target through this. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub speculator_for: Option, + /// Sidecar rows only: this drafter's preference among the sidecars of + /// the same target, lower first. `--speculate auto` takes the + /// lowest-ranked installed one; unranked sidecars follow, in strategy + /// order (dspark, dflash, eagle3). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub speculator_rank: Option, } fn default_file() -> String { diff --git a/base-convert/crates/base-hub/src/download.rs b/base-convert/crates/base-hub/src/download.rs index 8994d49..4e2638b 100644 --- a/base-convert/crates/base-hub/src/download.rs +++ b/base-convert/crates/base-hub/src/download.rs @@ -25,11 +25,14 @@ //! than hanging the pull forever. //! //! Not locked against a second process pulling the same file at the same -//! time. Both would write identical bytes at identical offsets, so the result -//! is not corrupt, but the loser's final rename can fail once the winner has -//! moved the partial away. hf-hub 0.5 did not lock here either, so this is -//! not a regression; it is worth a lock file if concurrent pulls ever become -//! a normal thing to do. +//! time. Both write identical bytes at identical offsets, so the result is +//! never corrupt. The one observable symptom was the LOSER's final rename +//! failing once the winner had moved the shared partial away — which surfaced +//! as a bogus "no .safetensors shards" failure on a pull that had actually +//! succeeded. Finalize now treats "dst exists and the partial is gone" as +//! success, so concurrent pulls converge instead of one spuriously failing. +//! A real lock file would additionally stop the duplicated download work; that +//! is a bandwidth optimization, not a correctness fix. use anyhow::{bail, Context, Result}; use hf_hub::{HFClientSync, RepoTypeModel}; @@ -294,7 +297,19 @@ pub fn download_ranged( } drop(file); map.remove(); - std::fs::rename(&partial, dst).with_context(|| format!("finalizing {}", dst.display()))?; + // A second process pulling the same blob concurrently may have finished + // first and renamed ITS partial onto dst — taking ours with it, since both + // use the same partial path. Both wrote identical bytes at identical + // offsets, so the winner's file IS our result: succeed rather than failing + // a pull that actually produced the right file. (Observed for real: a + // second `pull` of the same repo made the first fail with "no .safetensors + // shards".) Only tolerated when dst is present and the partial is gone — + // any other rename error is still an error. + if let Err(e) = std::fs::rename(&partial, dst) { + if !(dst.exists() && !partial.exists()) { + return Err(e).with_context(|| format!("finalizing {}", dst.display())); + } + } bar.finish_and_clear(); Ok(()) } diff --git a/base-convert/crates/base-hub/src/gen.rs b/base-convert/crates/base-hub/src/gen.rs index f9f9dcb..8d93a7c 100644 --- a/base-convert/crates/base-hub/src/gen.rs +++ b/base-convert/crates/base-hub/src/gen.rs @@ -92,6 +92,23 @@ pub fn entry_from_base(path: &Path, id: &str, hf_repo: &str) -> Result Option { + if !matches!(header.arch.as_str(), "dflash" | "dspark" | "eagle3") { + return None; + } + let layer = header.tensors.iter().find(|t| { + t.name.starts_with("layers.") && t.name.ends_with(".weight") && t.shape.len() == 2 + })?; + let tag = format!("{:?}", layer.dtype).to_ascii_lowercase(); + Some(tag.strip_prefix("base").map(str::to_string).unwrap_or(tag)) +} + pub fn entry_from_header( header: &base_format::Header, id: &str, @@ -101,9 +118,11 @@ pub fn entry_from_header( sha256: String, ) -> Result { let backend = backend_tag(header.target_backend)?; - let bits = crate::registry::quant_bits(&header.quant_profile) - .map(str::to_string) - .unwrap_or_else(|| format!("{:?}", header.quant_scheme).to_ascii_lowercase()); + let bits = sidecar_bits(header).unwrap_or_else(|| { + crate::registry::quant_bits(&header.quant_profile) + .map(str::to_string) + .unwrap_or_else(|| format!("{:?}", header.quant_scheme).to_ascii_lowercase()) + }); let quant = match &backend { Some(b) => format!("{b}-{bits}"), None => format!("default-{bits}"), @@ -120,6 +139,25 @@ pub fn entry_from_header( sha256: Some(sha256), parts_sha256: None, backend, + // The selectable strategy kind: a target's in-bundle head says so in + // `speculator.config.kind` ("mtp"; its `arch` is model-specific, e.g. + // qwen35-mtp), a sidecar IS the drafter, so its own architecture is + // the kind (dflash / dspark / eagle3). + speculator: header + .speculator + .as_ref() + .and_then(|s| { + s.config + .get("kind") + .and_then(|v| v.as_str()) + .map(str::to_string) + }) + .or_else(|| match header.arch.as_str() { + "dflash" | "dspark" | "eagle3" => Some(header.arch.clone()), + _ => None, + }), + speculator_for: None, + speculator_rank: None, }) } diff --git a/base-convert/crates/base-hub/src/parts.rs b/base-convert/crates/base-hub/src/parts.rs index a29797d..473be40 100644 --- a/base-convert/crates/base-hub/src/parts.rs +++ b/base-convert/crates/base-hub/src/parts.rs @@ -769,8 +769,10 @@ fn discard_staged(fetcher: &dyn Fetcher, repo: &str, src: &Path) { const MAX_HEADER_LEN: u64 = 256 * 1024 * 1024; /// Flags that promise an extension-slot section after the weights blob. +/// HAS_SPECULATOR is not one of them: the speculator sub-bundle lives in +/// the weights blob (`header.speculator.tensors`, the mmproj convention) +/// and writes no slot. const SLOT_FLAGS: base_format::HeaderFlags = base_format::HeaderFlags::HAS_LORA - .union(base_format::HeaderFlags::HAS_SPECULATOR) .union(base_format::HeaderFlags::HAS_COMPUTE_GRAPH) .union(base_format::HeaderFlags::HAS_KV_WARMUP) .union(base_format::HeaderFlags::HAS_TRACE_REF) @@ -813,8 +815,10 @@ fn check_complete(path: &Path, len: u64) -> Result<()> { .checked_next_multiple_of(BLOB_ALIGNMENT) .ok_or_else(overflow)?; let mut blob_end = blob_offset; + // The sub-bundles (mmproj towers, the speculator head) share the blob. let mmproj = header.mmproj.iter().flat_map(|m| m.tensors.iter()); - for t in header.tensors.iter().chain(mmproj) { + let speculator = header.speculator.iter().flat_map(|m| m.tensors.iter()); + for t in header.tensors.iter().chain(mmproj).chain(speculator) { let regions = [ (Some(t.offset), Some(t.length)), (t.scale_offset, t.scale_length), diff --git a/base-convert/crates/base-hub/src/scan.rs b/base-convert/crates/base-hub/src/scan.rs index 401871b..021f468 100644 --- a/base-convert/crates/base-hub/src/scan.rs +++ b/base-convert/crates/base-hub/src/scan.rs @@ -118,13 +118,23 @@ impl ScanReport { } /// List every model repository under `org`, following pagination. +/// +/// Authenticated when a Hub token resolves (the same one the tree reads +/// send): an anonymous listing hides the org's private repositories, and a +/// scan that cannot see a repository drops its rows as "no longer +/// published" — every private drafter sidecar, on the last run. pub fn list_org_repos(org: &str) -> Result> { let mut out = Vec::new(); let mut url = format!("{}/api/models?author={org}&limit=100", endpoint()); + let token = crate::fetch::resolve_token(); // The Hub paginates with a `Link: <…>; rel="next"` header. Bounded so a // malformed or cyclic Link chain cannot spin forever. for _ in 0..50 { - let resp = ureq::get(&url) + let mut req = ureq::get(&url); + if let Some(t) = &token { + req = req.header("Authorization", &format!("Bearer {t}")); + } + let resp = req .call() .with_context(|| format!("listing models for {org}"))?; let next = resp @@ -198,6 +208,28 @@ fn inherited_source_repo(known: &Catalog, repo: &str) -> Option { .and_then(|m| m.source_repo.clone()) } +/// Which target a sidecar drafts for (`speculator_for`) is curated the same +/// way: a sidecar's header names its own architecture (dflash / dspark / +/// eagle3), never the model it was trained against. Carried across from any +/// known row of the same repo, as `source_repo` is. +fn inherited_speculator_for(known: &Catalog, repo: &str) -> Option { + known + .models + .iter() + .find(|m| m.hf_repo == repo && m.speculator_for.is_some()) + .and_then(|m| m.speculator_for.clone()) +} + +/// A sidecar's rank among its target's drafters is curated beside +/// `speculator_for` and carried across the same way. +fn inherited_speculator_rank(known: &Catalog, repo: &str) -> Option { + known + .models + .iter() + .find(|m| m.hf_repo == repo && m.speculator_rank.is_some()) + .and_then(|m| m.speculator_rank) +} + /// Read a published bundle's header without downloading the bundle. /// /// Two reads, because the header's length is in the header: the 16-byte prefix @@ -570,6 +602,8 @@ pub fn scan_org( // so publishing a new quant does not quietly drop it from // the ones already curated. entry.source_repo = inherited_source_repo(known, repo); + entry.speculator_for = inherited_speculator_for(known, repo); + entry.speculator_rank = inherited_speculator_rank(known, repo); entry.parts_sha256 = parts_sha256; report.entries.push(entry) } @@ -717,6 +751,9 @@ mod tests { sha256: Some("abc123".into()), parts_sha256: None, backend: None, + speculator: None, + speculator_for: None, + speculator_rank: None, }], }; let r = scan_org( @@ -782,6 +819,9 @@ mod tests { sha256: Some("q4sha".into()), parts_sha256: None, backend: None, + speculator: None, + speculator_for: None, + speculator_rank: None, }, CatalogEntry { id: "basecompute/other".into(), @@ -795,6 +835,9 @@ mod tests { sha256: Some("othersha".into()), parts_sha256: None, backend: None, + speculator: None, + speculator_for: None, + speculator_rank: None, }, ], }; @@ -808,6 +851,73 @@ mod tests { assert_eq!(inherited_source_repo(&known, "basecompute/unknown"), None); } + #[test] + fn a_rescanned_sidecar_keeps_its_curated_target() { + // A sidecar's header says what it is (its own arch), not what it + // drafts for: `speculator_for` lives only in the checked-in catalog, + // so a re-read of the same repo (a changed file, a new quant) must + // inherit it from a sibling row or the association is lost. + let known = Catalog { + schema: 1, + updated: String::new(), + models: vec![ + CatalogEntry { + id: "basecompute/Qwen3-4B-DFlash".into(), + hf_repo: "basecompute/Qwen3-4B-DFlash".into(), + file: "Qwen3-4B-DFlash-F16.base".into(), + revision: "main".into(), + source_repo: Some("z-lab/Qwen3-4B-DFlash".into()), + arch: Some("dflash".into()), + quant: "default-f16".into(), + size: Some(1), + sha256: Some("f16sha".into()), + parts_sha256: None, + backend: None, + speculator: Some("dflash".into()), + speculator_for: Some("basecompute/Qwen3-4B".into()), + speculator_rank: Some(2), + }, + CatalogEntry { + id: "basecompute/Qwen3-4B".into(), + hf_repo: "basecompute/Qwen3-4B".into(), + file: "Qwen3-4B-Q4.base".into(), + revision: "main".into(), + source_repo: Some("Qwen/Qwen3-4B".into()), + arch: Some("qwen".into()), + quant: "default-q4".into(), + size: Some(1), + sha256: Some("q4sha".into()), + parts_sha256: None, + backend: None, + speculator: None, + speculator_for: None, + speculator_rank: None, + }, + ], + }; + assert_eq!( + inherited_speculator_for(&known, "basecompute/Qwen3-4B-DFlash").as_deref(), + Some("basecompute/Qwen3-4B") + ); + assert_eq!( + inherited_speculator_rank(&known, "basecompute/Qwen3-4B-DFlash"), + Some(2) + ); + assert_eq!( + inherited_speculator_rank(&known, "basecompute/Qwen3-4B"), + None + ); + // A target repo never borrows a sidecar's association. + assert_eq!( + inherited_speculator_for(&known, "basecompute/Qwen3-4B"), + None + ); + assert_eq!( + inherited_speculator_for(&known, "basecompute/unknown"), + None + ); + } + #[test] fn a_split_bundle_is_catalogued_from_its_manifest() { use crate::parts::{synthetic_bundle, Manifest, ManifestPart}; @@ -1112,6 +1222,9 @@ mod tests { sha256: Some("abc123".into()), parts_sha256: None, backend: None, + speculator: None, + speculator_for: None, + speculator_rank: None, }], }; let r = scan_org( diff --git a/base-convert/crates/base-quant/tests/canonical_pipeline.rs b/base-convert/crates/base-quant/tests/canonical_pipeline.rs index 9137ef9..a5d6e5d 100644 --- a/base-convert/crates/base-quant/tests/canonical_pipeline.rs +++ b/base-convert/crates/base-quant/tests/canonical_pipeline.rs @@ -78,6 +78,7 @@ fn make_header() -> Header { layers: vec![], tensors: vec![], mmproj: None, + speculator: None, calibration: None, provenance: None, sig: None, diff --git a/base-convert/crates/base-readers/src/hf.rs b/base-convert/crates/base-readers/src/hf.rs index a5f8678..2c64ea0 100644 --- a/base-convert/crates/base-readers/src/hf.rs +++ b/base-convert/crates/base-readers/src/hf.rs @@ -14,14 +14,19 @@ //! ``` //! //! This reader opens all shards mmap-style and presents a unified -//! tensor listing. It does not dequantize — safetensors from HF is -//! almost always unquantized (F32/F16/BF16). MLX-quantized safetensors -//! go through `mlx` module instead. +//! tensor listing. Safetensors from HF is mostly unquantized +//! (F32/F16/BF16); MLX-quantized safetensors go through the `mlx` +//! module instead. The one quantized HF layout read here is block FP8 +//! (DeepSeek-V3 / GLM-5.3 style: `quantization_config.quant_method = +//! "fp8"`, e4m3 `.weight` + f32 `.weight_scale_inv`, one +//! scale per `weight_block_size` tile): `tensor_to_f32` dequantizes the +//! weight × its block scale, and the scale tensors are hidden from +//! `tensor_names` so the converter sees a plain unquantized checkpoint. use crate::safetensors::{SafetensorsFile, StDtype, StTensorInfo}; use anyhow::{bail, Context, Result}; use serde::Deserialize; -use std::collections::BTreeMap; +use std::collections::{BTreeMap, BTreeSet}; use std::path::{Path, PathBuf}; pub struct HfDir { @@ -40,6 +45,13 @@ pub struct HfDir { shards: Vec, /// name → (shard_idx, tensor_idx) lookup: BTreeMap, + /// Block-FP8 checkpoints: e4m3 weight name → its `weight_scale_inv`. + fp8_scales: BTreeMap, + /// The paired scale tensors, hidden from `tensor_names`. + fp8_scale_names: BTreeSet, + /// `quantization_config.weight_block_size` ([rows, cols]) of a + /// block-FP8 checkpoint. + fp8_block: Option<[u64; 2]>, } #[derive(Deserialize)] @@ -160,7 +172,7 @@ impl HfDir { } } - Ok(Self { + let mut hf = Self { model_dir: dir, config, tokenizer_json, @@ -168,11 +180,103 @@ impl HfDir { chat_template_jinja, shards, lookup, - }) + fp8_scales: BTreeMap::new(), + fp8_scale_names: BTreeSet::new(), + fp8_block: None, + }; + hf.pair_fp8_block_scales()?; + Ok(hf) } + /// Pair every e4m3 `.weight` with its `.weight_scale_inv` + /// and check each scale covers its weight at the declared block size, + /// so a malformed checkpoint fails at open rather than mid-convert. + fn pair_fp8_block_scales(&mut self) -> Result<()> { + let qc = self.config.get("quantization_config"); + let block = match qc { + Some(qc) if qc.get("quant_method").and_then(|v| v.as_str()) == Some("fp8") => { + if let Some(fmt) = qc.get("fmt").and_then(|v| v.as_str()) { + if fmt != "e4m3" { + bail!("fp8 checkpoint with fmt {fmt:?}: only e4m3 is supported"); + } + } + let bs: Vec = qc + .get("weight_block_size") + .and_then(|v| v.as_array()) + .map(|a| a.iter().filter_map(|x| x.as_u64()).collect()) + .unwrap_or_default(); + if bs.len() != 2 || bs.contains(&0) { + bail!( + "fp8 checkpoint: quantization_config.weight_block_size must be \ + [rows, cols], got {:?}", + qc.get("weight_block_size") + ); + } + Some([bs[0], bs[1]]) + } + _ => None, + }; + for (name, &(si, ti)) in &self.lookup { + let info = &self.shards[si].tensors[ti]; + if info.dtype != StDtype::F8E4m3 { + continue; + } + let Some(stem) = name.strip_suffix(".weight") else { + continue; + }; + let scale = format!("{stem}.weight_scale_inv"); + let Some(&(ssi, sti)) = self.lookup.get(&scale) else { + continue; + }; + let Some([bo, bi]) = block else { + bail!( + "{name} is F8_E4M3 with a weight_scale_inv sibling, but config.json \ + declares no fp8 quantization_config (weight_block_size unknown)" + ); + }; + if info.shape.len() != 2 { + bail!( + "fp8 block weight {name} has shape {:?}; only 2-D is supported", + info.shape + ); + } + let want = [info.shape[0].div_ceil(bo), info.shape[1].div_ceil(bi)]; + let sinfo = &self.shards[ssi].tensors[sti]; + if sinfo.shape != want { + bail!( + "{scale} has shape {:?}; {name} {:?} at block [{bo}, {bi}] needs {:?}", + sinfo.shape, + info.shape, + want + ); + } + self.fp8_scales.insert(name.clone(), scale.clone()); + self.fp8_scale_names.insert(scale); + } + if !self.fp8_scales.is_empty() { + self.fp8_block = block; + } + Ok(()) + } + + /// `[rows, cols]` block size when this is a block-FP8 checkpoint whose + /// e4m3 weights `tensor_to_f32` dequantizes. + pub fn fp8_block_size(&self) -> Option<[u64; 2]> { + self.fp8_block + } + + /// Number of e4m3 weights dequantized through a block scale. + pub fn fp8_weight_count(&self) -> usize { + self.fp8_scales.len() + } + + /// Every tensor the checkpoint holds, minus the block-FP8 scales that + /// `tensor_to_f32` folds into their weights. pub fn tensor_names(&self) -> impl Iterator { - self.lookup.keys().map(|s| s.as_str()) + self.lookup + .keys() + .filter(|n| !self.fp8_scale_names.contains(*n)) + .map(|s| s.as_str()) } pub fn tensor_info(&self, name: &str) -> Option<&StTensorInfo> { @@ -225,6 +329,33 @@ impl HfDir { .chunks_exact(2) .map(|c| bf16::from_le_bytes([c[0], c[1]]).to_f32()) .collect(), + // Index tables (EAGLE-3's d2t offsets) and masks: exact in f32 + // for the magnitudes they hold (token ids < 2^24). + StDtype::I64 => bytes + .chunks_exact(8) + .map(|c| i64::from_le_bytes(c.try_into().unwrap()) as f32) + .collect(), + StDtype::I32 => bytes + .chunks_exact(4) + .map(|c| i32::from_le_bytes(c.try_into().unwrap()) as f32) + .collect(), + StDtype::Bool => bytes + .iter() + .map(|&b| if b != 0 { 1.0 } else { 0.0 }) + .collect(), + StDtype::F8E4m3 => { + // Unscaled e4m3 is off by the block scale (often 1e-3..1e-1), + // so a weight without one is refused, never decoded raw. + let scale_name = self.fp8_scales.get(name).with_context(|| { + format!( + "tensor {name} is F8_E4M3 with no .weight_scale_inv block \ + scale to dequantize it" + ) + })?; + let [bo, bi] = self.fp8_block.expect("set whenever a scale is paired"); + let scale = self.tensor_to_f32(scale_name)?; + dequant_fp8_block(bytes, &info.shape, &scale, bo as usize, bi as usize) + } other => bail!( "tensor_to_f32: unsupported safetensors dtype {:?} (tensor {name}, {n} values)", other @@ -233,6 +364,46 @@ impl HfDir { } } +/// OCP e4m3fn (the `F8_E4M3` safetensors dtype): 1 sign, 4 exponent +/// (bias 7), 3 mantissa bits, no infinities; S.1111.111 is NaN, so the +/// largest finite magnitude is 448. +pub fn e4m3fn_to_f32(b: u8) -> f32 { + let sign = if b & 0x80 != 0 { -1.0 } else { 1.0 }; + let exp = ((b >> 3) & 0x0F) as i32; + let man = (b & 0x07) as f32; + if exp == 0x0F && b & 0x07 == 0x07 { + return f32::NAN; + } + if exp == 0 { + // Subnormal: 0.mmm × 2^-6. + return sign * (man / 8.0) * 2f32.powi(-6); + } + sign * (1.0 + man / 8.0) * 2f32.powi(exp - 7) +} + +/// Dequantize a row-major `[rows, cols]` e4m3 weight against its +/// `[ceil(rows/bo), ceil(cols/bi)]` block scale: `w[r][c] = e4m3(q[r][c]) +/// × scale[r/bo][c/bi]`. (The checkpoint calls it `weight_scale_inv`, +/// but it is the factor the stored codes are multiplied by — DeepSeek's +/// own `weight_dequant` is `x * s`.) +fn dequant_fp8_block(bytes: &[u8], shape: &[u64], scale: &[f32], bo: usize, bi: usize) -> Vec { + let (rows, cols) = (shape[0] as usize, shape[1] as usize); + let scale_cols = cols.div_ceil(bi); + debug_assert_eq!(bytes.len(), rows * cols); + debug_assert_eq!(scale.len(), rows.div_ceil(bo) * scale_cols); + let lut: [f32; 256] = std::array::from_fn(|i| e4m3fn_to_f32(i as u8)); + let mut out = Vec::with_capacity(rows * cols); + for (r, row) in bytes.chunks_exact(cols).enumerate() { + let srow = &scale[(r / bo) * scale_cols..][..scale_cols]; + out.extend( + row.iter() + .enumerate() + .map(|(c, &q)| lut[q as usize] * srow[c / bi]), + ); + } + out +} + fn read_optional_json(path: &Path) -> Result> { if !path.exists() { return Ok(None); @@ -353,3 +524,181 @@ mod sanitize_tests { assert_eq!(n, 0); } } + +#[cfg(test)] +mod fp8_tests { + use super::{e4m3fn_to_f32, HfDir}; + use std::io::Write; + use std::path::Path; + + #[test] + fn e4m3fn_edge_values() { + assert_eq!(e4m3fn_to_f32(0x00), 0.0); + assert_eq!(e4m3fn_to_f32(0x38), 1.0); + assert_eq!(e4m3fn_to_f32(0xB8), -1.0); + assert_eq!(e4m3fn_to_f32(0x7E), 448.0); // largest finite + assert_eq!(e4m3fn_to_f32(0xFE), -448.0); + assert_eq!(e4m3fn_to_f32(0x08), 2f32.powi(-6)); // smallest normal + assert_eq!(e4m3fn_to_f32(0x01), 2f32.powi(-9)); // smallest subnormal + assert_eq!(e4m3fn_to_f32(0x07), 7.0 * 2f32.powi(-9)); + assert_eq!(e4m3fn_to_f32(0x3C), 1.5); + assert!(e4m3fn_to_f32(0x7F).is_nan()); + assert!(e4m3fn_to_f32(0xFF).is_nan()); + // Every other code is finite. + let finite = (0u8..=255) + .filter(|&b| e4m3fn_to_f32(b).is_finite()) + .count(); + assert_eq!(finite, 254); + } + + fn write_safetensors(path: &Path, tensors: &[(&str, &str, &[u64], Vec)]) { + let mut header = serde_json::Map::new(); + let mut offset = 0u64; + for (name, dtype, shape, bytes) in tensors { + let end = offset + bytes.len() as u64; + header.insert( + name.to_string(), + serde_json::json!({ "dtype": dtype, "shape": shape, "data_offsets": [offset, end] }), + ); + offset = end; + } + let hdr = serde_json::to_vec(&serde_json::Value::Object(header)).unwrap(); + let mut f = std::fs::File::create(path).unwrap(); + f.write_all(&(hdr.len() as u64).to_le_bytes()).unwrap(); + f.write_all(&hdr).unwrap(); + for (_, _, _, bytes) in tensors { + f.write_all(bytes).unwrap(); + } + } + + fn refusal(r: anyhow::Result) -> String { + match r { + Ok(_) => panic!("must refuse"), + Err(e) => e.to_string(), + } + } + + fn f32_bytes(v: &[f32]) -> Vec { + v.iter().flat_map(|x| x.to_le_bytes()).collect() + } + + fn fp8_config(dir: &Path, block: serde_json::Value) { + let config = serde_json::json!({ + "model_type": "test", + "quantization_config": { + "quant_method": "fp8", "fmt": "e4m3", "activation_scheme": "dynamic", + "weight_block_size": block, + }, + }); + std::fs::write( + dir.join("config.json"), + serde_json::to_vec(&config).unwrap(), + ) + .unwrap(); + } + + /// A 3×5 weight at block [2, 2]: ragged edge tiles in both + /// dimensions, a distinct scale per tile, and a bf16 tensor alongside + /// that must pass through untouched. + #[test] + fn dequantizes_weight_times_block_scale() { + let tmp = tempfile::tempdir().unwrap(); + fp8_config(tmp.path(), serde_json::json!([2, 2])); + // Codes: 1.0 everywhere except a 1.5 and a -448. + let mut q = vec![0x38u8; 15]; + q[4] = 0x3C; // [0][4] = 1.5 + q[10] = 0xFE; // [2][0] = -448 + // scale [2, 3]: tile (r/2, c/2). + let scale = [1.0f32, 2.0, 4.0, 0.5, 0.25, 0.125]; + let norm: Vec = [1.0f32, 2.0] + .iter() + .flat_map(|x| half::bf16::from_f32(*x).to_le_bytes()) + .collect(); + write_safetensors( + &tmp.path().join("model.safetensors"), + &[ + ("l.weight", "F8_E4M3", &[3, 5], q), + ("l.weight_scale_inv", "F32", &[2, 3], f32_bytes(&scale)), + ("n.weight", "BF16", &[2], norm), + ], + ); + let hf = HfDir::open(tmp.path()).unwrap(); + assert_eq!(hf.fp8_block_size(), Some([2, 2])); + assert_eq!(hf.fp8_weight_count(), 1); + let names: Vec<&str> = hf.tensor_names().collect(); + assert_eq!( + names, + ["l.weight", "n.weight"], + "the scale is folded, not listed" + ); + let w = hf.tensor_to_f32("l.weight").unwrap(); + #[rustfmt::skip] + let want = [ + 1.0, 1.0, 2.0, 2.0, 4.0 * 1.5, + 1.0, 1.0, 2.0, 2.0, 4.0, + -448.0 * 0.5, 0.5, 0.25, 0.25, 0.125, + ]; + assert_eq!(w, want); + assert_eq!(hf.tensor_to_f32("n.weight").unwrap(), [1.0, 2.0]); + } + + #[test] + fn bf16_block_scales_are_accepted() { + let tmp = tempfile::tempdir().unwrap(); + fp8_config(tmp.path(), serde_json::json!([128, 128])); + let scale: Vec = half::bf16::from_f32(0.5).to_le_bytes().to_vec(); + write_safetensors( + &tmp.path().join("model.safetensors"), + &[ + ("l.weight", "F8_E4M3", &[2, 3], vec![0x38; 6]), + ("l.weight_scale_inv", "BF16", &[1, 1], scale), + ], + ); + let hf = HfDir::open(tmp.path()).unwrap(); + assert_eq!(hf.tensor_to_f32("l.weight").unwrap(), [0.5; 6]); + } + + #[test] + fn scale_shape_mismatch_fails_at_open() { + let tmp = tempfile::tempdir().unwrap(); + fp8_config(tmp.path(), serde_json::json!([2, 2])); + write_safetensors( + &tmp.path().join("model.safetensors"), + &[ + ("l.weight", "F8_E4M3", &[3, 5], vec![0x38; 15]), + ("l.weight_scale_inv", "F32", &[1, 3], f32_bytes(&[1.0; 3])), + ], + ); + let err = refusal(HfDir::open(tmp.path())); + assert!(err.contains("needs [2, 3]"), "{err}"); + } + + #[test] + fn unscaled_fp8_is_refused_not_decoded_raw() { + let tmp = tempfile::tempdir().unwrap(); + fp8_config(tmp.path(), serde_json::json!([2, 2])); + write_safetensors( + &tmp.path().join("model.safetensors"), + &[("l.weight", "F8_E4M3", &[2, 2], vec![0x38; 4])], + ); + let hf = HfDir::open(tmp.path()).unwrap(); + assert_eq!(hf.fp8_block_size(), None); + let err = refusal(hf.tensor_to_f32("l.weight")); + assert!(err.contains("no .weight_scale_inv"), "{err}"); + } + + #[test] + fn paired_scale_without_fp8_config_fails_at_open() { + let tmp = tempfile::tempdir().unwrap(); + std::fs::write(tmp.path().join("config.json"), br#"{"model_type": "test"}"#).unwrap(); + write_safetensors( + &tmp.path().join("model.safetensors"), + &[ + ("l.weight", "F8_E4M3", &[2, 2], vec![0x38; 4]), + ("l.weight_scale_inv", "F32", &[1, 1], f32_bytes(&[1.0])), + ], + ); + let err = refusal(HfDir::open(tmp.path())); + assert!(err.contains("weight_block_size unknown"), "{err}"); + } +} diff --git a/base-convert/crates/base-sign/src/lib.rs b/base-convert/crates/base-sign/src/lib.rs index 72d4b7a..cb8fed3 100644 --- a/base-convert/crates/base-sign/src/lib.rs +++ b/base-convert/crates/base-sign/src/lib.rs @@ -131,6 +131,21 @@ pub fn b64_decode(s: &str) -> Result> { /// This is intentionally a post-process step: keeps `BaseWriter` /// streaming-friendly (no memory-buffered blob) while centralizing /// signing logic here. +/// End of the weights blob relative to its start: max(offset + length) over +/// EVERY tensor list — the main tensors and the sub-bundles that share the +/// blob (mmproj towers, the speculator head). Scanning only `header.tensors` +/// would leave every speculator payload out of the signed digest, so a +/// modified MTP head still verified. +fn blob_end_rel(h: &base_format::Header) -> u64 { + h.tensors + .iter() + .chain(h.mmproj.iter().flat_map(|m| m.tensors.iter())) + .chain(h.speculator.iter().flat_map(|m| m.tensors.iter())) + .map(|t| t.offset + t.length) + .max() + .unwrap_or(0) +} + pub fn sign_base_file>( input: P, output: P, @@ -145,13 +160,7 @@ pub fn sign_base_file>( .with_context(|| format!("opening {:?}", input.as_ref()))?; // Compute end-of-blob as max(offset + length) over all tensors. - let blob_end_rel = reader - .header() - .tensors - .iter() - .map(|t| t.offset + t.length) - .max() - .unwrap_or(0); + let blob_end_rel = blob_end_rel(reader.header()); let blob_start = reader.blob_offset(); let blob_bytes: Vec = { let file_bytes = std::fs::read(input.as_ref())?; @@ -242,13 +251,7 @@ pub fn verify_base_file>(path: P, key: &VerifyingKey) sig_arr.copy_from_slice(&sig_bytes); let sig = Signature::from_bytes(&sig_arr); - let blob_end_rel = reader - .header() - .tensors - .iter() - .map(|t| t.offset + t.length) - .max() - .unwrap_or(0); + let blob_end_rel = blob_end_rel(reader.header()); let blob_start = reader.blob_offset(); let file_bytes = std::fs::read(path.as_ref())?; let blob_bytes = &file_bytes[blob_start as usize..(blob_start + blob_end_rel) as usize]; diff --git a/base-convert/crates/base-sign/tests/sign_roundtrip.rs b/base-convert/crates/base-sign/tests/sign_roundtrip.rs index 10d0461..58fd853 100644 --- a/base-convert/crates/base-sign/tests/sign_roundtrip.rs +++ b/base-convert/crates/base-sign/tests/sign_roundtrip.rs @@ -34,6 +34,7 @@ fn make_header() -> Header { layers: vec![], tensors: vec![], mmproj: None, + speculator: None, calibration: None, provenance: None, sig: None, diff --git a/bindings/node/package.json b/bindings/node/package.json index 9127c51..45f4f1c 100644 --- a/bindings/node/package.json +++ b/bindings/node/package.json @@ -1,6 +1,6 @@ { "name": "@baseRT/node", - "version": "0.2.6", + "version": "0.3.0", "private": true, "description": "Node.js bindings for BaseRT — LLM inference engine for Apple Silicon (Metal)", "main": "dist/index.js", diff --git a/bindings/python/baseRT/__init__.py b/bindings/python/baseRT/__init__.py index 2d6596f..0825091 100644 --- a/bindings/python/baseRT/__init__.py +++ b/bindings/python/baseRT/__init__.py @@ -21,7 +21,7 @@ from pathlib import Path from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union -__version__ = "0.2.6" +__version__ = "0.3.0" # --------------------------------------------------------------------------- # Library loading diff --git a/bindings/python/setup.py b/bindings/python/setup.py index 4cb8a9a..dceb288 100644 --- a/bindings/python/setup.py +++ b/bindings/python/setup.py @@ -6,7 +6,7 @@ setup( name="baseRT", - version="0.2.6", + version="0.3.0", description="Python bindings for the BaseRT LLM inference engine (Apple Silicon / Metal)", long_description=long_description, long_description_content_type="text/markdown", diff --git a/bindings/rust/baseRT-sys/Cargo.toml b/bindings/rust/baseRT-sys/Cargo.toml index 41f4bfe..72e0944 100644 --- a/bindings/rust/baseRT-sys/Cargo.toml +++ b/bindings/rust/baseRT-sys/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "baseRT-sys" -version = "0.2.6" +version = "0.3.0" edition = "2021" description = "Raw FFI bindings for the BaseRT LLM inference engine" license = "Apache-2.0" diff --git a/bindings/rust/baseRT-sys/src/lib.rs b/bindings/rust/baseRT-sys/src/lib.rs index bb6c674..7bd6f1b 100644 --- a/bindings/rust/baseRT-sys/src/lib.rs +++ b/bindings/rust/baseRT-sys/src/lib.rs @@ -265,9 +265,27 @@ pub struct BaseRTGenerationStats { pub decode_tokens_per_sec: c_float, } +/// Opaque per-sequence handle (independent KV state on the model's shared +/// paged-KV pool). Mirrors `baseRT_sequence_t` in baseRT.h. +pub type baseRT_sequence_t = *mut c_void; + +/// Result of a prefix-cache lookup (`baseRT_prefix_match`). `blocks` points +/// into engine-owned storage valid until `baseRT_prefix_unlock(handle)` / +/// `baseRT_prefix_release(handle)`. `handle == 0` means miss / cache +/// disabled — nothing to unlock. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct BaseRTPrefixMatch { + pub matched_tokens: c_int, + pub n_blocks: c_int, + pub blocks: *const c_int, + pub handle: u64, +} + /// Callback for streaming token output. Return `false` to stop generation. -pub type baseRT_token_callback = - Option bool>; +pub type baseRT_token_callback = Option< + unsafe extern "C" fn(token_id: u32, text: *const c_char, user_data: *mut c_void) -> bool, +>; /// Callback for streaming transcription segments. Return `false` to stop transcription. pub type baseRT_segment_callback = Option< @@ -289,6 +307,9 @@ extern "C" { ) -> baseRT_model_t; pub fn baseRT_free_model(model: baseRT_model_t); + /// Tokenizer-only load: metadata + vocab, no GPU tensors. Only the + /// tokenizer surface is valid on the returned handle. + pub fn baseRT_load_tokenizer_only(model_path: *const c_char) -> baseRT_model_t; // === Model info === @@ -316,6 +337,25 @@ extern "C" { kv_bits: c_int, paged_kv: c_int, ) -> c_int; + /// As above, plus the speculators loading beside those models: + /// `speculator_embedded[i]` non-zero is the target bundle's embedded head + /// (its KV pool only); zero (or a null array) a sidecar / draft model + /// (weights and KV pool), even when also served. `speculator_target[i]` + /// >= 0 is the index of the model a sidecar's baseRT_load_drafter reopens + /// for its embedding / lm_head (charged where the backend copies wrapped + /// weights); -1 (or a null array) for none. Speculators never cap the + /// window. + pub fn baseRT_suggest_max_context_spec( + model_paths: *const *const c_char, + n_models: c_int, + speculator_paths: *const *const c_char, + speculator_embedded: *const c_int, + speculator_target: *const c_int, + n_speculators: c_int, + max_batch: c_int, + kv_bits: c_int, + paged_kv: c_int, + ) -> c_int; /// 1 = the window fits the co-resident set, 0 = it does not, -1 = unknown. pub fn baseRT_context_window_fits( model_paths: *const *const c_char, @@ -335,8 +375,35 @@ extern "C" { out_tokens: *mut u32, max_tokens: c_int, ) -> c_int; + /// Special-token strings in `text` are ordinary text (client content), + /// never control ids. See the header. + pub fn baseRT_encode_plain( + model: baseRT_model_t, + text: *const c_char, + out_tokens: *mut u32, + max_tokens: c_int, + ) -> c_int; + /// `n_pieces` texts encoded as one: piece `i` parses marker strings + /// when `plain[i]` is zero and keeps them as text otherwise, with the + /// pre-tokenizer seeing the concatenation. See the header. + pub fn baseRT_encode_pieces( + model: baseRT_model_t, + texts: *const *const c_char, + plain: *const c_int, + n_pieces: c_int, + out_tokens: *mut u32, + max_tokens: c_int, + ) -> c_int; pub fn baseRT_decode_token(model: baseRT_model_t, token_id: u32) -> *const c_char; + /// Length-preserving stateless decode (byte-level tokens can contain + /// NUL); returns the true byte length, which may exceed `max_bytes`. + pub fn baseRT_decode_token_raw( + model: baseRT_model_t, + token_id: u32, + out: *mut c_char, + max_bytes: c_int, + ) -> c_int; // === Generation === @@ -458,6 +525,7 @@ extern "C" { ) -> c_int; pub fn baseRT_embedding_dim(model: baseRT_model_t) -> c_int; + pub fn baseRT_is_embedding_model(model: baseRT_model_t) -> bool; // === Chat templates === @@ -469,6 +537,60 @@ extern "C" { pub fn baseRT_chat_template(model: baseRT_model_t) -> *const c_char; + /// Raw Jinja `chat_template` source folded in from the HF tokenizer + /// (empty if the bundle carries none). + pub fn baseRT_chat_template_jinja(model: baseRT_model_t) -> *const c_char; + + /// Create a grammar from a JSON Schema (converted to GBNF internally); + /// NULL on error. Used to validate that a codec's tool-call grammar + /// compiles in the engine (the runtime constrained-decode path). + pub fn baseRT_grammar_create_from_schema( + model: baseRT_model_t, + json_schema: *const c_char, + ) -> *mut std::os::raw::c_void; + + /// Create a grammar from an xgrammar STRUCTURAL TAG document (JSON, not a + /// grammar string); NULL on error. The tag's content formats reach past + /// JSON Schema — `sequence`, `const_string`, `regex`, `or` — which is how + /// a dialect whose call body is not JSON (Gemma 4's `call:fn{k:v}`) still + /// gets an enforceable `tool_choice`. + pub fn baseRT_grammar_create_from_structural_tag( + model: baseRT_model_t, + tag_json: *const c_char, + ) -> *mut std::os::raw::c_void; + + /// Create a grammar from a GBNF string; NULL on parse error. + pub fn baseRT_grammar_create( + model: baseRT_model_t, + gbnf: *const c_char, + ) -> *mut std::os::raw::c_void; + + /// Step a grammar by one token; 0 = the grammar REJECTS it. Lets a test + /// assert what a compiled grammar actually admits, not merely that it + /// compiled — the difference between "the tag parses" and "the tag + /// permits the calls this dialect must be able to make". + pub fn baseRT_grammar_accept_token( + grammar: *mut std::os::raw::c_void, + token_id: u32, + ) -> std::os::raw::c_int; + + /// Reset a grammar's acceptance state to its initial stacks, so one + /// handle can check several candidate strings. + pub fn baseRT_grammar_reset(grammar: *mut std::os::raw::c_void); + + /// Free a grammar. + pub fn baseRT_grammar_free(grammar: *mut std::os::raw::c_void); + + /// Number of added/special tokens the model declares. + pub fn baseRT_special_token_count(model: baseRT_model_t) -> c_int; + + /// The `index`-th special token string; writes its id to `id_out`. + pub fn baseRT_special_token( + model: baseRT_model_t, + index: c_int, + id_out: *mut u32, + ) -> *const c_char; + // === Token counting === pub fn baseRT_token_count(model: baseRT_model_t, text: *const c_char) -> c_int; @@ -502,6 +624,136 @@ extern "C" { n_tokens: c_int, penalty: c_float, ); + + // === Pre-load configuration (call BEFORE baseRT_load_model) === + + pub fn baseRT_set_paged_kv(enable: c_int); + pub fn baseRT_set_baked_decode(enable: c_int); + pub fn baseRT_set_max_batch_size(n: c_int); + pub fn baseRT_set_prefix_cache(enable: c_int); + pub fn baseRT_set_prefill_chunk(n: c_int); + pub fn baseRT_set_kv_bits(bits: c_int); + pub fn baseRT_set_verbose(on: c_int); + + // === Error / capability introspection === + + pub fn baseRT_get_error_code() -> c_int; + pub fn baseRT_strerror(code: c_int) -> *const c_char; + pub fn baseRT_capabilities(model: baseRT_model_t) -> u32; + pub fn baseRT_eos_token_id(model: baseRT_model_t) -> u32; + pub fn baseRT_is_eos_token(model: baseRT_model_t, token: u32) -> i32; + pub fn baseRT_weights_identity(model: baseRT_model_t) -> u64; + pub fn baseRT_kv_bits_effective(model: baseRT_model_t) -> c_int; + pub fn baseRT_bos_id(model: baseRT_model_t) -> u32; + + // === Multi-sequence generation (paged-KV only) === + + pub fn baseRT_sequence_create(model: baseRT_model_t) -> baseRT_sequence_t; + pub fn baseRT_sequence_free(seq: baseRT_sequence_t); + pub fn baseRT_sequence_rollback(seq: baseRT_sequence_t, length: c_int) -> c_int; + + // === Batched prefill/decode === + + pub fn baseRT_batch_step( + model: baseRT_model_t, + seqs: *mut baseRT_sequence_t, + n_seqs: c_int, + in_tokens: *const u32, + in_token_counts: *const c_int, + out_tokens: *mut u32, + ) -> c_int; + pub fn baseRT_batch_step_fused( + model: baseRT_model_t, + seqs: *mut baseRT_sequence_t, + n_seqs: c_int, + in_tokens: *const u32, + in_token_counts: *const c_int, + out_tokens: *mut u32, + ) -> c_int; + pub fn baseRT_batch_step_fused_logits( + model: baseRT_model_t, + seqs: *mut baseRT_sequence_t, + n_seqs: c_int, + in_tokens: *const u32, + in_token_counts: *const c_int, + ) -> c_int; + pub fn baseRT_batch_warmup(model: baseRT_model_t, max_batch: c_int) -> c_int; + pub fn baseRT_max_prefill_chunk(model: baseRT_model_t) -> c_int; + + // === Host-side logits-row operations === + + pub fn baseRT_read_batch_logits( + model: baseRT_model_t, + n_seqs: c_int, + out_logits_f16: *mut c_void, + ) -> c_int; + pub fn baseRT_batch_logits_stride(model: baseRT_model_t) -> usize; + pub fn baseRT_mask_logits_row( + model: baseRT_model_t, + row: *mut c_void, + bitmask: *const i32, + ) -> c_int; + pub fn baseRT_sample_logits_row( + model: baseRT_model_t, + row: *const c_void, + cfg: *const BaseRTSamplingConfig, + prev_tokens: *const u32, + n_prev: c_int, + repeat_window: c_int, + seed_offset: u32, + ) -> u32; + pub fn baseRT_argmax_logits_row(model: baseRT_model_t, row: *const c_void) -> u32; + pub fn baseRT_logits_row_logprobs( + model: baseRT_model_t, + row: *const c_void, + token: u32, + top_k: c_int, + out_token_logprob: *mut c_float, + out_ids: *mut u32, + out_logprobs: *mut c_float, + ) -> c_int; + + // === Prefix cache — scheduler-driven primitives === + + pub fn baseRT_prefix_match( + model: baseRT_model_t, + tokens: *const u32, + n_tokens: c_int, + ) -> BaseRTPrefixMatch; + pub fn baseRT_sequence_seed_prefix( + seq: baseRT_sequence_t, + blocks: *const c_int, + n_blocks: c_int, + n_tokens: c_int, + ) -> c_int; + pub fn baseRT_page_size(model: baseRT_model_t) -> c_int; + pub fn baseRT_prefix_insert( + model: baseRT_model_t, + tokens: *const u32, + n_tokens: c_int, + seq: baseRT_sequence_t, + ) -> c_int; + pub fn baseRT_prefix_unlock(model: baseRT_model_t, handle: u64); + pub fn baseRT_prefix_release(model: baseRT_model_t, handle: u64); + pub fn baseRT_prefix_discard(model: baseRT_model_t, handle: u64); + pub fn baseRT_prefix_evict(model: baseRT_model_t, n_blocks: c_int) -> c_int; + pub fn baseRT_prefix_cache_stats( + model: baseRT_model_t, + out_hits: *mut u64, + out_misses: *mut u64, + out_reused_tokens: *mut u64, + out_blocks_cached: *mut c_int, + ); + + // === Recurrent (GDN) state snapshots — hybrid models only === + + pub fn baseRT_gdn_snapshot_size(model: baseRT_model_t) -> c_int; + pub fn baseRT_sequence_gdn_capture(seq: baseRT_sequence_t, out: *mut u8, cap: c_int) -> c_int; + pub fn baseRT_sequence_gdn_restore( + seq: baseRT_sequence_t, + blob: *const u8, + len: c_int, + ) -> c_int; } #[cfg(test)] @@ -595,7 +847,10 @@ mod tests { assert_eq!(&base.n_layers as *const _ as usize - base_ptr, 4); assert_eq!(&base.norm_eps as *const _ as usize - base_ptr, 40); assert_eq!(&base.rope_theta as *const _ as usize - base_ptr, 44); - assert_eq!(&base.sliding_window_pattern as *const _ as usize - base_ptr, 48); + assert_eq!( + &base.sliding_window_pattern as *const _ as usize - base_ptr, + 48 + ); // The field the old binding dropped — everything after shifts by 4. assert_eq!(&base.sliding_window as *const _ as usize - base_ptr, 52); assert_eq!(&base.rope_local_theta as *const _ as usize - base_ptr, 56); @@ -606,31 +861,55 @@ mod tests { assert_eq!(&base.swa_layers as *const _ as usize - base_ptr, 144); assert_eq!(&base.ffn_dims as *const _ as usize - base_ptr, 208); assert_eq!(&base.attn_output_gate as *const _ as usize - base_ptr, 1232); - assert_eq!(&base.linear_attn_layers as *const _ as usize - base_ptr, 1244); + assert_eq!( + &base.linear_attn_layers as *const _ as usize - base_ptr, + 1244 + ); assert_eq!(&base.gdn_num_k_heads as *const _ as usize - base_ptr, 1308); // Nemotron-H SSM block, inserted between gdn_* and the MoE block — // 20 bytes that shifted every field below it. assert_eq!(&base.ssm_state_size as *const _ as usize - base_ptr, 1328); assert_eq!(&base.ssm_num_heads as *const _ as usize - base_ptr, 1344); assert_eq!(&base.n_experts as *const _ as usize - base_ptr, 1348); - assert_eq!(&base.expert_weights_scale as *const _ as usize - base_ptr, 1368); + assert_eq!( + &base.expert_weights_scale as *const _ as usize - base_ptr, + 1368 + ); assert_eq!(&base.vision_n_layers as *const _ as usize - base_ptr, 1372); assert_eq!(&base.vision_arch as *const _ as usize - base_ptr, 1432); assert_eq!(&base.audio_n_layers as *const _ as usize - base_ptr, 1448); assert_eq!(&base.eoa_token_id as *const _ as usize - base_ptr, 1524); assert_eq!(&base.mrope_section as *const _ as usize - base_ptr, 1528); - assert_eq!(&base.mrope_interleaved as *const _ as usize - base_ptr, 1540); - assert_eq!(&base.rope_scaling_factor as *const _ as usize - base_ptr, 1544); - assert_eq!(&base.rope_orig_max_pos as *const _ as usize - base_ptr, 1556); - assert_eq!(&base.rope_scaling_type as *const _ as usize - base_ptr, 1560); + assert_eq!( + &base.mrope_interleaved as *const _ as usize - base_ptr, + 1540 + ); + assert_eq!( + &base.rope_scaling_factor as *const _ as usize - base_ptr, + 1544 + ); + assert_eq!( + &base.rope_orig_max_pos as *const _ as usize - base_ptr, + 1556 + ); + assert_eq!( + &base.rope_scaling_type as *const _ as usize - base_ptr, + 1560 + ); assert_eq!(&base.qk_scale_factor as *const _ as usize - base_ptr, 1564); assert_eq!(&base.nope_layers as *const _ as usize - base_ptr, 1576); assert_eq!(&base.embed_norm_eps as *const _ as usize - base_ptr, 1640); - assert_eq!(&base.vision_window_layers as *const _ as usize - base_ptr, 1644); + assert_eq!( + &base.vision_window_layers as *const _ as usize - base_ptr, + 1644 + ); assert_eq!(&base.video_token_id as *const _ as usize - base_ptr, 1724); assert_eq!(&base.q_lora_rank as *const _ as usize - base_ptr, 1728); assert_eq!(&base.indexer_top_k as *const _ as usize - base_ptr, 1768); - assert_eq!(&base.rope_yarn_beta_fast as *const _ as usize - base_ptr, 1772); + assert_eq!( + &base.rope_yarn_beta_fast as *const _ as usize - base_ptr, + 1772 + ); } #[test] @@ -661,8 +940,14 @@ mod tests { assert_eq!(&base.generated_tokens as *const _ as usize - base_ptr, 4); assert_eq!(&base.prefill_time_ms as *const _ as usize - base_ptr, 8); assert_eq!(&base.decode_time_ms as *const _ as usize - base_ptr, 12); - assert_eq!(&base.prefill_tokens_per_sec as *const _ as usize - base_ptr, 16); - assert_eq!(&base.decode_tokens_per_sec as *const _ as usize - base_ptr, 20); + assert_eq!( + &base.prefill_tokens_per_sec as *const _ as usize - base_ptr, + 16 + ); + assert_eq!( + &base.decode_tokens_per_sec as *const _ as usize - base_ptr, + 20 + ); } // ----------------------------------------------------------------------- @@ -685,6 +970,19 @@ mod tests { assert!(cfg.logit_bias_values.is_null()); } + #[test] + fn prefix_match_layout() { + // int matched_tokens (4) + int n_blocks (8) + const int *blocks + // (16) + uint64_t handle (24). Returned BY VALUE from + // baseRT_prefix_match, so the mirror must match the C layout + // exactly. + assert_eq!(mem::size_of::(), 24); + assert_eq!(mem::offset_of!(BaseRTPrefixMatch, matched_tokens), 0); + assert_eq!(mem::offset_of!(BaseRTPrefixMatch, n_blocks), 4); + assert_eq!(mem::offset_of!(BaseRTPrefixMatch, blocks), 8); + assert_eq!(mem::offset_of!(BaseRTPrefixMatch, handle), 16); + } + #[test] fn transcribe_stats_default_values() { let stats = BaseRTTranscribeStats::default(); diff --git a/bindings/rust/baseRT/Cargo.toml b/bindings/rust/baseRT/Cargo.toml index b42b65e..5121901 100644 --- a/bindings/rust/baseRT/Cargo.toml +++ b/bindings/rust/baseRT/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "baseRT" -version = "0.2.6" +version = "0.3.0" edition = "2021" description = "Safe Rust bindings for the BaseRT LLM inference engine (Apple Silicon)" license = "Apache-2.0" diff --git a/bindings/swift/Sources/CBaseRT/include/baseRT.h b/bindings/swift/Sources/CBaseRT/include/baseRT.h index 224dcd5..04823d8 100644 --- a/bindings/swift/Sources/CBaseRT/include/baseRT.h +++ b/bindings/swift/Sources/CBaseRT/include/baseRT.h @@ -63,8 +63,8 @@ extern "C" { // === Versioning === #define BASERT_VERSION_MAJOR 0 -#define BASERT_VERSION_MINOR 2 -#define BASERT_VERSION_PATCH 6 +#define BASERT_VERSION_MINOR 3 +#define BASERT_VERSION_PATCH 0 /// Compile-time version, packed as `(MAJOR<<16) | (MINOR<<8) | PATCH`. /// Useful for `#if BASERT_VERSION >= 0x000200` feature checks. @@ -84,8 +84,11 @@ typedef void *baseRT_model_t; /// Other source formats (GGUF, HF safetensors, MLX safetensors) must be /// converted offline first via `basert convert`. /// kernel_library_path: path to the compiled GPU kernel library (on the Metal -/// backend, baseRT.metallib), or NULL to auto-detect. Auto-detect order: a -/// kernel library next to the build, then a copy embedded in the loaded +/// backend, baseRT.metallib; on CUDA, baseRT_cuda_kernels.fatbin; on ROCm, +/// baseRT_rocm_kernels.hsaco), or NULL to auto-detect. Auto-detect order: +/// $BASERT_KERNEL_LIB (a file, or a directory holding the artifact), a +/// kernel library next to the build, next to the +/// executable, the current directory, then a copy embedded in the loaded /// binary itself (single-file distributions ship the shared library with the /// kernels linked in, so NULL just works). Named generically so non-Metal /// backends (CUDA/ROCm, future) can reuse the same parameter. @@ -97,6 +100,17 @@ typedef void *baseRT_model_t; /// Returns NULL on failure. baseRT_model_t baseRT_load_model(const char *model_path, const char *kernel_library_path, int max_context); +/// Load ONLY the bundle's tokenizer: metadata parse + vocab/merges, no +/// tensor upload, no KV/scratch allocation (a serving daemon tokenizes in +/// its own process while a worker owns the GPU state — basertd §3's codec +/// placement). The returned handle supports exactly the tokenizer surface: +/// baseRT_encode, baseRT_decode_token / _static / _raw, baseRT_token_count, +/// baseRT_eos_token_id, baseRT_bos_id / baseRT_bos_token / baseRT_eos_token, +/// baseRT_get_config, and baseRT_free_model. Calling generation or +/// sequence APIs on it is undefined. NULL on failure (see +/// baseRT_get_error). +baseRT_model_t baseRT_load_tokenizer_only(const char *model_path); + /// Load a model with per-call options instead of the process-wide /// baseRT_set_* pre-load setters. `opts` may be NULL (identical to /// baseRT_load_model); otherwise set opts->struct_size = @@ -138,8 +152,27 @@ uint32_t baseRT_capabilities(baseRT_model_t model); /// Returns BASERT_OK, or an error code (leaves the model on the OLD library on /// load failure). Not supported for whisper models. NOT thread-safe against /// in-flight inference — serialize the caller. +/// +/// INVALIDATES ALL DECODED STATE. Everything in a KV or recurrent slab was +/// produced by the old kernels, so resuming on it would mix kernel versions — +/// the one thing a tuning comparison must not do. On success the default +/// sequence's KV, recurrent state and prefix cache are reset for you (position +/// returns to 0). Caller-owned sequences (baseRT_sequence_create) are NOT +/// reachable from here and must be destroyed and recreated; using one across a +/// reload is undefined. int baseRT_reload_metallib(baseRT_model_t model, const char *metallib_path); +/// A JSON capability descriptor for this loaded model (schema field +/// `descriptor_version`): workload, modalities, serving, state, adaptation, +/// and tick-op support, each capability either `true` or a stable reason +/// string explaining why it is unavailable on this architecture + backend + +/// load-options triple. Computed once at load (verdicts frozen, like the +/// multi-row overcommit verdict) — immutable, race-free, never touches the +/// error state. The pointer is owned by the model handle and valid until +/// baseRT_free_model. Supersedes the five-bit mask above for every question +/// it cannot answer; the mask stays for cheap scheduler branching. +const char *baseRT_capability_descriptor(baseRT_model_t model); + /// Override the KV cache element width for the next baseRT_load_model call. /// bits = 0 → auto (per-model default; Q8_0 when head_dim%32==0) /// bits = 8 → force Q8_0 K and V slabs (1.88x smaller; tiny precision cost) @@ -148,6 +181,16 @@ int baseRT_reload_metallib(baseRT_model_t model, const char *metallib_path); /// baseRT_load_model. Other values are ignored. void baseRT_set_kv_bits(int bits); +/// The KV cache's element width as ACTUALLY allocated for `model`, in the +/// `baseRT_set_kv_bits` encoding: 16 (f16), 8 (Q8_0), 4 (Q4_0), 84 (K at +/// Q8_0, V at Q4_0), or 0 for a model with no KV cache. Differs from what +/// was set whenever the loader overrode it — hybrid GDN / SSM, gpt-oss and +/// MLA models force f16; a paged pool is uniform (84 → 8); a head_dim that +/// is not a multiple of 32 cannot take the quantized blocks — and the load +/// printed one `[baseRT] : --kv-bits N requested, ...` line saying +/// why. Serve this rather than the setting when reporting KV precision. +int baseRT_kv_bits_effective(baseRT_model_t model); + /// Enable engine diagnostics (RoPE/tokenizer/GPU/architecture dumps, the /// per-token dispatch-command count, "Warming up"). Off by default so end /// users see only model output. @@ -244,14 +287,19 @@ size_t baseRT_model_config_sizeof(void); /// Get total GPU memory used by model (bytes). size_t baseRT_model_memory(baseRT_model_t model); -/// The device memory budget the engine allocates within, in bytes, without -/// needing a loaded model. NOT installed RAM: on Metal this is the unified- -/// memory working set the OS recommends (~75% of physical RAM), on CUDA the -/// device's total memory. 0 when no supported device is present. +/// The device memory budget automatic context sizing plans within, in bytes, +/// without needing a loaded model. NOT installed RAM: on Metal this is 85% of +/// the unified-memory working set the OS recommends (itself a fraction of RAM +/// the OS picks per machine, roughly two thirds to three quarters) — the share +/// the engine keeps its pinned weights within, since past it Metal evicts from +/// the pinned set and reads slow down or go corrupt; on CUDA the device's total +/// memory. 0 when no supported device is present. /// -/// This is the number that matters for sizing decisions — it is what the load -/// path compares the model's working set against before warning that the OS -/// will start paging weights. +/// This is the number `baseRT_suggest_max_context*` and +/// `baseRT_context_window_fits` size against. It is a planning line, not a +/// hard limit: a load whose working set lands between it and the full working +/// set still succeeds, with the tail of its weights left to page-cache LRU (and +/// the load warns once the working set passes the full budget). size_t baseRT_device_memory_budget(void); /// Floor on a window derived by `baseRT_suggest_max_context`. Policy, not a @@ -298,6 +346,29 @@ int baseRT_suggest_max_context(const char *model_path, int max_batch, int kv_bit int baseRT_suggest_max_context_multi(const char *const *model_paths, int n_models, int max_batch, int kv_bits, int paged_kv); +/// As `baseRT_suggest_max_context_multi`, plus the speculators that will load +/// beside those models. A drafter runs at its target's window and lane count +/// with its own paged KV pool, so it spends the same budget: leaving it out +/// sizes the targets' pools into memory the drafter then cannot get. +/// `speculator_embedded[i]` non-zero marks speculator i as its target +/// bundle's own embedded head (header.speculator, e.g. MTP): its KV pool only, +/// its weights being part of the bundle. Zero — or a null array — is a +/// separately loaded bundle (a DFlash / DSpark / EAGLE-3 sidecar or a draft +/// model): its weights and its KV pool, even if the same file is also one of +/// `model_paths` (a draft model is loaded again beside the served copy). +/// `speculator_target[i]` >= 0 marks speculator i as a sidecar loaded with +/// `baseRT_load_drafter` against `model_paths[speculator_target[i]]`: that +/// loader reopens the target's bundle for its token embedding / lm_head, and a +/// backend that copies wrapped weights (CUDA's UMA path) holds those rows a +/// second time, so they are charged too. -1 — or a null array — is a +/// speculator with no such view (a draft model loads on its own); it is +/// ignored for an embedded head, whose view of its own bundle is always charged. +/// Speculators never cap the window, one that cannot be read is skipped, and +/// none is charged unless `paged_kv` (their loaders require a paged target). +int baseRT_suggest_max_context_spec(const char *const *model_paths, int n_models, const char *const *speculator_paths, + const int *speculator_embedded, const int *speculator_target, int n_speculators, + int max_batch, int kv_bits, int paged_kv); + /// Does `window` fit `model_paths` co-resident on this device, under the same /// budget the suggestion above derives from? 1 = yes, 0 = no, -1 = cannot tell /// (a bundle would not open, no trained window, unknown budget). @@ -325,6 +396,11 @@ typedef struct BaseRTMemoryStats { uint64_t kv_cache_used_bytes; uint64_t kv_cache_blocks_total; uint64_t kv_cache_blocks_used; + /// Tokens resident in the KV. Exact on a contiguous cache. On a PAGED cache + /// this is block-granular — `kv_cache_blocks_used * page_size`, so each + /// sequence's partially-filled last block counts whole — because the pool + /// tracks blocks, not per-sequence lengths, and the default lane's own + /// length reads 0 whenever batched sequences own the live state. uint64_t kv_cache_tokens_used; BaseRTKVCacheLayout kv_cache_layout; uint32_t reserved; @@ -352,8 +428,37 @@ const char *baseRT_strerror(BaseRTErrorCode code); // === Tokenization === /// Encode text to token IDs. Returns number of tokens written. +/// +/// A substring equal to one of the model's special tokens (`<|im_end|>`, +/// ``, ``, …) becomes that single control id — what a +/// rendered chat prompt needs, since its framing IS those markers. int baseRT_encode(baseRT_model_t model, const char *text, uint32_t *out_tokens, int max_tokens); +/// Encode text with special-token strings treated as ordinary text. +/// +/// The same substrings tokenize to their BPE pieces instead of control ids, +/// so a message body, tool result or document that QUOTES a marker cannot +/// forge a turn boundary or open a reasoning block. Use this for client +/// content and `baseRT_encode` for a dialect's own framing. No automatic +/// BOS-dedup either: a leading BOS string here is content. Returns number +/// of tokens written. +int baseRT_encode_plain(baseRT_model_t model, const char *text, uint32_t *out_tokens, int max_tokens); + +/// Encode a sequence of text pieces as ONE text. +/// +/// Piece `i` is framing when `plain[i]` is zero (marker strings parse to +/// control ids, as in `baseRT_encode`) and client content otherwise (they +/// stay ordinary text, as in `baseRT_encode_plain`). Unlike encoding the +/// pieces one by one, the pre-tokenizer and BPE see the concatenation, so a +/// boundary between two pieces tokenizes exactly as it does in a whole-text +/// encode of the same characters (a content piece starting with a newline +/// after a framing newline still merges into the `\n\n` token) and the +/// automatic BOS is handled once at the front. A chat codec renders a turn +/// with this: the role line and the wire around a message are framing, the +/// message is content. Returns number of tokens written. +int baseRT_encode_pieces(baseRT_model_t model, const char *const *texts, const int *plain, int n_pieces, + uint32_t *out_tokens, int max_tokens); + /// Decode a single token ID to text. Returns static string (do not free). /// /// Advances the tokenizer's incremental-decode state — call this from @@ -375,17 +480,22 @@ const char *baseRT_decode_token_static(baseRT_model_t model, uint32_t token_id); typedef struct baseRT_decode_stream *baseRT_decode_stream_t; baseRT_decode_stream_t baseRT_decode_stream_create(baseRT_model_t model); void baseRT_decode_stream_reset(baseRT_decode_stream_t stream); + +/// Reset a stream for a lane whose generation resumes after `prompt`, deriving +/// the channel resume state from its last framing token. Prefer this over +/// `baseRT_decode_stream_reset` on any lane that may carry a RAW (unframed) +/// prompt: the plain reset assumes a chat prompt, which ends inside a +/// `<|start|>assistant` header, and on a Muse vocabulary an unframed prompt +/// left in that state has every generated token buffered as protocol until a +/// 128-byte valve trips — a short completion returns nothing at all. Harmony +/// has its own fail-open path and is unaffected either way. `prompt` NULL or +/// `n_prompt` 0 behaves exactly like `baseRT_decode_stream_reset`. +void baseRT_decode_stream_reset_for_prompt(baseRT_decode_stream_t stream, const uint32_t *prompt, int n_prompt); /// Returns a pointer into the stream's own buffer, valid until the next /// call on the same stream. const char *baseRT_decode_stream_token(baseRT_decode_stream_t stream, uint32_t token_id); void baseRT_decode_stream_free(baseRT_decode_stream_t stream); -/// 1 when `token_id` ends the assistant turn: the bundle's eos_token_id OR -/// any of its extra end-of-turn ids (Llama 3's <|eot_id|>, gpt-oss's -/// <|call|> and <|endoftext|>, …). `baseRT_eos_token_id` reports only the -/// first; a lane that compares against that alone runs past the others. -int baseRT_is_eos_token(baseRT_model_t model, uint32_t token_id); - /// Length-preserving variant of `baseRT_decode_token_static` for callers /// that need the token's EXACT raw bytes. Byte-level BPE / byte-fallback /// tokens can decode to bytes containing 0x00, which the C-string variants @@ -588,7 +698,19 @@ int baseRT_max_prefill_chunk(baseRT_model_t model); /// Roll a sequence's KV state back to `length` tokens, returning any blocks /// past that point to the pool. `length` must be <= the current length; 0 /// resets the sequence to empty. Used by the serving engine's shape-padding -/// dummy lanes (their KV is discarded after every tick). +/// dummy lanes (their KV is discarded after every tick) and by speculative +/// verification (drop the rejected drafts). Recurrent-state models: a +/// hybrid-GDN lane rolls back only to a length inside its last +/// baseRT_batch_step_fused_logits_rows feed (the recurrence is rebuilt from +/// that feed's captured rows — deterministic, and bit-identical to a feed +/// that stopped there under bit-exact routing (BASERT_SPEC_BITEXACT=1); +/// the default speculative routing picks its small-M GEMM kernel by row +/// count, so the replayed prefix carries the longer feed's last-ulp +/// rounding instead); `length == current` is a no-op commit on EVERY +/// lane (it drops a hybrid-GDN lane's capture); any shorter target on a +/// Mamba-2 lane, or outside the captured feed on a hybrid-GDN lane, +/// returns BASERT_ERR_UNSUPPORTED — a no-op rollback is not a probe of +/// rollback support. int baseRT_sequence_rollback(baseRT_sequence_t seq, int length); /// Multi-step autoregressive driver for baseRT_batch_step_fused. Step 0 @@ -619,10 +741,289 @@ int baseRT_batch_step_fused_loop(baseRT_model_t model, baseRT_sequence_t *seqs, int baseRT_batch_step_fused_logits(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, const uint32_t *in_tokens, const int *in_token_counts); -/// Read back the [n_seqs, vocab] f16 logits left by the most recent -/// baseRT_batch_step_fused_logits into `out_logits_f16` (n_seqs * vocab halves, -/// row-major). Pure UMA copy, no dispatch. Returns vocab_size, or <0 on error. -/// `n_seqs` must match the batch of the preceding step and be <= max_batch_size. +/// Per-row sampling spec for baseRT_batch_step_fused_sample. temperature <= 0 +/// samples greedily (argmax) for that row; otherwise Gumbel-max over the +/// top-k / top-p / min-p survivors with `seed` driving the noise stream +/// (same seed + same logits = same token). +/// +/// Penalties ARE applied for a sampled row, over `history` — the caller's +/// recent-token window for that row — with the model's penalty-exemption mask +/// honoured, so a sampled row no longer has to leave the fused path to get +/// one. Logit bias and grammar still take the host logits path. +typedef struct BaseRTRowSampling { + float temperature; + float top_p; ///< 1.0 = disabled + float min_p; ///< 0.0 = disabled + int32_t top_k; ///< <= 0 = disabled + uint32_t seed; ///< Gumbel noise seed for this row + float repeat_penalty; ///< 1.0 (or 0.0) = disabled + float presence_penalty; ///< 0.0 = disabled + float frequency_penalty; ///< 0.0 = disabled + /// Recent tokens this row's penalties are computed over, most recent last. + /// NULL / 0 disables the penalties for the row regardless of the values + /// above: with no history there is nothing to penalize. + const uint32_t *history; + uint32_t history_count; +} BaseRTRowSampling; + +/// Fused batch step that samples every output row ON THE GPU per `rows[i]` +/// (greedy rows argmax, sampled rows Gumbel-max) and returns one token per +/// seq in `out_tokens` — no full-vocab logits readback and no per-row host +/// sampling. Same forward as baseRT_batch_step_fused_logits (so every +/// architecture that serves host logits serves this), followed by one +/// gumbel_topk_f16 dispatch per sampled row and one batched argmax. +/// Returns BASERT_ERR_UNSUPPORTED when the backend lacks the sampling +/// kernels (caller falls back to the logits path). +int baseRT_batch_step_fused_sample(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, const uint32_t *in_tokens, + const int *in_token_counts, const BaseRTRowSampling *rows, uint32_t *out_tokens); + +/// The sampling pass of baseRT_batch_step_fused_sample on its own: samples +/// logits rows [first_row, first_row + n_rows) that the last fused forward +/// left in the host logits buffer (baseRT_batch_step_fused_logits_rows / +/// _rows_taps — a speculative verify's rows), `rows[i]` driving row +/// first_row + i, one token per row into `out_tokens`. The pass noises the +/// rows in place: read them back first (baseRT_read_batch_logits) if the raw +/// values are still needed. BASERT_ERR_UNSUPPORTED when the backend lacks the +/// sampling kernels or the penalty windows do not fit — the rows are untouched +/// then and the caller samples them on the host. +int baseRT_sample_logits_rows(baseRT_model_t model, int first_row, int n_rows, const BaseRTRowSampling *rows, + uint32_t *out_tokens); + +/// Multi-row variant for speculative verification (basertd P2.4): seq i +/// contributes in_token_counts[i] input rows and leaves its LAST +/// out_rows[i] rows' logits (1 <= out_rows[i] <= in_token_counts[i]) in the +/// host logits buffer, ascending, seqs in order — sum(out_rows) rows total, +/// read back with baseRT_read_batch_logits(sum(out_rows)). The host checks +/// each drafted position against the target's logits from ONE forward and +/// rolls the sequence back to the accepted length. Served on the fused +/// VARLEN path and on the staged decomposition (MoE, hybrid GDN, residency +/// overcommit) alike; Mamba-2 hybrids return BASERT_ERR_UNSUPPORTED (no +/// per-lane SSM rollback). On the staged route a lane's rows must fit in +/// its last prefill chunk; a hybrid-GDN lane's whole feed must be one +/// chunk of at most 31 tokens, and the lane's next call must be either +/// baseRT_sequence_rollback (to any length inside the feed) or a plain step +/// (which commits the feed). +/// sum(out_rows) is bounded by the logits scratch (max_batch_size × 8). +int baseRT_batch_step_fused_logits_rows(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, + const uint32_t *in_tokens, const int *in_token_counts, const int *out_rows); + +/// Hidden-state taps for speculative drafters (MTP heads, EAGLE-3, DFlash +/// condition on the target's residual stream). Same step as +/// baseRT_batch_step_fused_logits_rows, plus: for each of the `n_tap_layers` +/// requested hidden_states indices in `tap_layers` (0 = embedding output, +/// k = residual stream after decoder layer k-1, n_layers = final pre-norm +/// hidden), the LAST tap_rows[i] fed rows of lane i are copied into the +/// engine's tap buffer, rows in lane order, one contiguous [rows, dim] f16 +/// block per tap (read back with baseRT_read_tap_rows, or consumed on the +/// device by an in-engine drafter). At most 8 distinct tap layers and +/// baseRT_tap_rows_capacity() tapped rows per call; a lane's tapped rows +/// must fit in its last prefill chunk. A 1-token lane with tap_rows 1 runs +/// its step serially (the batched decode tables cannot tap). tap_rows may +/// name 0 rows for lanes that need no taps. Taps are emitted by the llama +/// family, Gemma 4, Qwen3.5, gpt-oss and GLM-DSA encoders; a tapped step on any other +/// architecture returns BASERT_ERR_UNSUPPORTED (BASERT_ERR_OUT_OF_MEMORY +/// when the tap buffer cannot be allocated). +/// Greedy speculative verification without a logits readback: the same +/// multi-row forward as baseRT_batch_step_fused_logits_rows_taps (taps +/// optional: n_tap_layers = 0), every requested row argmaxed on the GPU; +/// `out_tokens` receives sum(out_rows) tokens in lane order. Row i of a +/// lane is the target's greedy choice after the lane's i-th fed token — +/// accept draft i while it equals row i. BASERT_ERR_UNSUPPORTED on bundles +/// that verify through host logits (MoE / hybrid / non-llama output +/// stages): fall back to the logits variant. +int baseRT_batch_step_fused_rows_argmax(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, + const uint32_t *in_tokens, const int *in_token_counts, const int *out_rows, + const int *tap_layers, int n_tap_layers, const int *tap_rows, + uint32_t *out_tokens); + +int baseRT_batch_step_fused_logits_rows_taps(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, + const uint32_t *in_tokens, const int *in_token_counts, const int *out_rows, + const int *tap_layers, int n_tap_layers, const int *tap_rows); + +/// Rows one tapped call may cover (min(max_prefill_chunk, 1024)); <0 on a +/// null model. +int baseRT_tap_rows_capacity(baseRT_model_t model); + +/// Tapped step with an INPUT OVERRIDE: `input_f16` ([sum(in_token_counts), +/// dim] f16, rows in feed order) replaces the embedding lookup — the +/// drafters' input is a projection of the target's hidden state, not a +/// token. Token ids still drive positions and KV. Rows are bounded by +/// baseRT_tap_rows_capacity(). Feeding a step's own tap-0 rows back through +/// this call reproduces its logits bit for bit. +int baseRT_batch_step_fused_logits_rows_taps_input(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, + const uint32_t *in_tokens, const int *in_token_counts, + const int *out_rows, const int *tap_layers, int n_tap_layers, + const int *tap_rows, const void *input_f16); + +// ── Speculator heads (bundle `header.speculator`) ────────────────────── + +/// 1 when the loaded bundle carries a speculator sub-bundle (an MTP head or +/// a converted drafter), 0 otherwise. +int baseRT_has_speculator(baseRT_model_t model); + +/// The speculator's `kind` ("mtp", ...) into `out`; returns its length, 0 +/// when the bundle has no speculator, <0 on error. +int baseRT_speculator_kind(baseRT_model_t model, char *out, int cap); + +/// Whether a loaded speculator head can actually run the MTP prologue: the +/// prologue is an encoder's embedding stage — the Qwen3.5 hybrid encoder's +/// (`qwen35_mtp_prologue`) or GLM-DSA's nextn one (`glm_mtp_prologue`) — so +/// `baseRT_mtp_chain` and the per-step path both refuse a head on any other +/// architecture. A scheduler should ask this +/// before granting an MTP strategy rather than registering one whose every +/// proposal then falls back. 1 = supported, 0 = not, <0 = not a head. +int baseRT_speculator_head_supported(baseRT_model_t head); + +/// Whether this bundle's encoder can emit hidden-state taps. Every block +/// drafter and EAGLE-3 sidecar consumes target taps, so a scheduler should +/// ask before granting one of those strategies: without taps the first +/// tapped forward fails with BASERT_ERR_UNSUPPORTED, which the tick reports +/// as a generation failure rather than falling back to plain decoding. +/// 1 = taps available, 0 = not, <0 = bad argument. +int baseRT_model_supports_taps(baseRT_model_t model); + +/// Load the bundle's speculator head as its own model handle: the head's +/// decoder layer(s) and norms from the speculator section, the target's +/// embeddings and lm_head shared. It is a paged model of the target's +/// context with its own sequences (baseRT_sequence_create), stepped with +/// the tapped/override steps above; free it with baseRT_free_model before +/// the target. NULL (with an error) when the bundle has no loadable head. +baseRT_model_t baseRT_load_speculator_head(baseRT_model_t target); + +/// MTP head input rows (host convenience for tests and drivers): out[r] = +/// fc_embed · rmsnorm(embed(next_tokens[r])) + fc_hidden · rmsnorm(hidden[r]) +/// (GLM-DSA nextn heads: eh_proj · [enorm(embed) | hnorm(output_norm(hidden))]) +/// — `hidden_f16` the target's final pre-norm hidden ([rows, dim], tap +/// n_layers) at each row's position, `out_f16` [rows, dim] to feed the head +/// as its input override. The head predicts the token AFTER next_tokens[r]. +int baseRT_mtp_prologue(baseRT_model_t head, const uint32_t *next_tokens, int rows, const void *hidden_f16, + void *out_f16); + +/// The whole MTP-head proposal of `n_seqs` head lanes in ONE encoding (a +/// GLM-DSA head runs it as staged per-step forwards, same contract): +/// each lane's committed rows (`counts[i]` of them, 1..15; `tokens` the +/// NEXT token of every row packed lane by lane, `hidden_f16` the target's +/// final pre-norm hidden of every row, [sum counts][dim] f16) go through +/// the prologue and one head forward, then `n_draft - 1` chained steps of +/// one row per lane draft from the head's own output hidden and previous +/// draft, without a host round trip between steps. `out_drafts` receives +/// n_seqs * n_draft ids, lane-major (`out_drafts[i * n_draft + s]` = lane +/// i's draft s; draft 0 is the head's prediction after the last committed +/// row). Each lane's head sequence ends `counts[i] + n_draft - 1` +/// positions longer; roll the chained positions back before the next +/// commit. Requires the head's --paged-kv context; n_draft 1..15. +int baseRT_mtp_chain(baseRT_model_t head, baseRT_sequence_t *seqs, int n_seqs, const uint32_t *tokens, + const int *counts, const void *hidden_f16, int n_draft, uint32_t *out_drafts); + +/// Copy tap `tap_index` (its position in the call's tap_layers) of the most +/// recent tapped step into `out_f16`: n_rows * dim halves, row-major, rows +/// in lane order as fed. Bounded by that step: `tap_index` below its tap +/// count and `n_rows` at most its tapped rows (BASERT_ERR_INVALID_ARGUMENT +/// otherwise). Returns dim, or <0 on error. +int baseRT_read_tap_rows(baseRT_model_t model, int tap_index, int n_rows, void *out_f16); + +/// Block drafters (DFlash / DSpark): load a drafter SIDECAR bundle (arch +/// "dflash" / "dspark", converted from the drafter checkpoint) against +/// `target`. The drafter shares the target's token embedding and lm_head +/// (resolved from the target's bundle when the sidecar lacks them), runs at +/// the target's paged/lane shape, and is stepped with baseRT_drafter_step +/// on its own sequences (baseRT_sequence_create). Free it with +/// baseRT_free_model before the target. NULL with an error when the sidecar +/// does not fit the target (dim / vocab / taps). +baseRT_model_t baseRT_load_drafter(baseRT_model_t target, const char *path); + +/// Drafter geometry: `kind` ("dflash" | "dspark"), the block width, the +/// mask token the block is padded with, the first block row whose logits +/// are a draft (DFlash 1: row 0 is the anchor; DSpark 0), and the TARGET +/// hidden_states indices to tap (target_layer_ids + 1) in the order the +/// drafter's context expects them. Returns the tap count, or <0. +int baseRT_drafter_info(baseRT_model_t drafter, char *kind, int kind_cap, int *block_size, uint32_t *mask_token_id, + int *logits_start, int *tap_layers, int tap_cap); + +/// One drafter step on `seq`: append `n_ctx` context rows — the target's +/// tapped hidden rows at the next `n_ctx` positions, `ctx_taps_f16` laid +/// out [n_taps][n_ctx][dim] (tap-major, the baseRT_drafter_info order; +/// each tap block is exactly what baseRT_read_tap_rows returns) — then, +/// when `block_len > 0`, run the block `block[0..block_len)` (the last +/// committed token followed by mask tokens) after the context with +/// bidirectional block attention and leave its `block_len` logits rows for +/// baseRT_read_batch_logits (row j = block position j; drafts are rows +/// logits_start.. — the argmax of row j predicts the token at block +/// position j for DSpark, or the token replacing mask j for DFlash). The +/// block's KV is dropped; the sequence ends at its context length. Either +/// half may be empty (context-only ingestion of a prompt slice, or a block +/// on an up-to-date context). n_ctx <= baseRT_tap_rows_capacity(drafter); +/// `block_len` is 0 or exactly the block width (the drafters are trained +/// on a fixed-width bidirectional block — a narrower block is refused). +int baseRT_drafter_step(baseRT_model_t drafter, baseRT_sequence_t seq, const void *ctx_taps_f16, int n_ctx, + const uint32_t *block, int block_len); + +/// DSpark heads carried by the sidecar: `markov_rank` (0 = no Markov head) +/// and whether a confidence head is present. +int baseRT_drafter_heads(baseRT_model_t drafter, int *markov_rank, int *has_confidence); + +/// DSpark Markov head, one block row: add `markov.w2 · markov.w1[prev]` to +/// row `row` of the last block step's logits (in place) and return its +/// argmax in `out_token` — the draft for that row given the token before +/// it (`prev` = the pending token for the first drafted row, then the +/// previous draft). Sequential: call per row in order after +/// baseRT_drafter_step. +int baseRT_drafter_markov_argmax(baseRT_model_t drafter, int row, uint32_t prev, uint32_t *out_token); + +/// DSpark confidence head over the last block step's rows 0..rows: the +/// probability that block row r's draft survives verification, given the +/// token before it (`prev_tokens[r]`). Adaptive block length keeps drafts +/// while the running product stays above a threshold. Returns `rows`, or +/// <0 (BASERT_ERR_UNSUPPORTED without a confidence head). +int baseRT_drafter_confidence(baseRT_model_t drafter, const uint32_t *prev_tokens, int rows, float *out_conf); + +/// EAGLE-3 heads (a sidecar of arch "eagle3", loaded with baseRT_load_drafter): +/// one llama layer drafting the next token from (the current token, a +/// feature of the target's hidden states at three layers). `tap_layers` +/// are the TARGET hidden_states indices to tap (3), `draft_vocab` the +/// reduced vocabulary the head predicts over. Returns the tap count. +int baseRT_eagle3_info(baseRT_model_t head, int *draft_vocab, int *tap_layers, int tap_cap); + +/// One head step over `rows` positions from the sequence's length: +/// `tokens[r]` is the token that FOLLOWED position r (the token the row +/// conditions on), and exactly one of `taps_f16` ([3][rows][dim], the +/// target's tapped rows at those positions, baseRT_read_tap_rows layout) or +/// `hidden_f16` ([rows][dim], the head's own output hidden — chaining) is +/// the feature. Afterwards baseRT_eagle3_draft gives row r's draft (the +/// argmax over the draft vocabulary mapped to a target id) and +/// baseRT_eagle3_read the raw rows. +int baseRT_eagle3_step(baseRT_model_t head, baseRT_sequence_t seq, const uint32_t *tokens, int rows, + const void *taps_f16, const void *hidden_f16); +int baseRT_eagle3_read(baseRT_model_t head, int rows, void *logits_f16, void *hidden_f16); +int baseRT_eagle3_draft(baseRT_model_t head, int row, uint32_t *out_target_token); + +/// The whole proposal of `n_seqs` head sequences in one call: each lane's +/// `counts[i]` committed rows (tokens packed lane by lane; `taps_f16` the +/// target's tapped rows for ALL packed rows, [3][sum rows][dim]) as one +/// forward, then `n_draft - 1` chained one-row steps per lane on the +/// device (the previous draft is the token, the head's own output hidden +/// the feature) with no host round trip between them. `out_drafts[i * +/// n_draft + s]` is lane i's draft s (target ids). The head sequences end +/// `counts[i] + n_draft - 1` positions longer; roll back to drop the +/// chained positions. n_draft in 1..15. +int baseRT_eagle3_chain(baseRT_model_t head, baseRT_sequence_t *seqs, int n_seqs, const uint32_t *tokens, + const int *counts, const void *taps_f16, int n_draft, uint32_t *out_drafts); + +/// Media variant (basertd P2.7): the packed prompt's `image_token_id` +/// placeholders (exactly `n_rows` of them, in order) take rows from the +/// engine-owned features buffer produced by the media tower (`feats_buffer` +/// is a `baseRT::Buffer*` obtained through the tick media surface). Fused +/// VARLEN live path only (hybrid/MoE/overcommit return UNSUPPORTED). Host +/// sampled: read the logits back with baseRT_read_batch_logits. +int baseRT_batch_step_fused_logits_media(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, + const uint32_t *in_tokens, const int *in_token_counts, void *feats_buffer, + int n_rows, uint32_t image_token_id, int grid_h, int grid_w); + +/// Read back the [n_rows, vocab] f16 logits left by the most recent +/// baseRT_batch_step_fused_logits(_rows) into `out_logits_f16` (n_rows * vocab +/// halves, row-major). Pure UMA copy, no dispatch. Returns vocab_size, or <0 +/// on error. `n_rows` is the preceding step's batch (one row per seq) or its +/// sum(out_rows), bounded by the logits scratch (max_batch_size × 8 rows). int baseRT_read_batch_logits(baseRT_model_t model, int n_seqs, void *out_logits_f16); // === Host-side logits-row operations === @@ -749,6 +1150,14 @@ void baseRT_prefix_unlock(baseRT_model_t model, uint64_t handle); /// to seed it (e.g. a boundary mismatch). No-op for handle==0. void baseRT_prefix_release(baseRT_model_t model, uint64_t handle); +/// Drop a baseRT_prefix_match taken only to LOOK at what the cache holds (a +/// metadata probe, e.g. basertd's space_match): releases the pin like +/// baseRT_prefix_release but settles nothing, so the probe counts as neither +/// a hit nor a miss in baseRT_prefix_cache_stats. A match that was a reuse +/// decision and was declined belongs in baseRT_prefix_release. No-op for +/// handle==0. +void baseRT_prefix_discard(baseRT_model_t model, uint64_t handle); + /// Evict least-recently-used UNLOCKED cached prefixes until at least `n_blocks` /// block-frees have been performed back to the pool. Returns the number freed /// (may be < n_blocks if the remaining prefixes are all locked by live @@ -772,13 +1181,50 @@ int baseRT_prefix_cache_save(baseRT_model_t model, const char *path); /// (no outstanding matches). Returns BASERT_OK or an error code. int baseRT_prefix_cache_load(baseRT_model_t model, const char *path); -/// Lifetime prefix-cache stats (any out-pointer may be NULL). `hits`/`misses` -/// count baseRT_prefix_match calls that did / didn't reuse >=1 block; -/// `reused_tokens` is the running total of prompt tokens served from cache; +/// Lifetime prefix-cache stats (any out-pointer may be NULL), all monotonic. +/// A baseRT_prefix_match that found nothing counts as a miss at once; one that +/// found blocks settles later, as a hit at baseRT_prefix_unlock (its blocks +/// were seeded) or as a miss at baseRT_prefix_release (abandoned); one +/// dropped by baseRT_prefix_discard (a probe) is in neither, as is an +/// outstanding match. `reused_tokens` is the running total of +/// prompt tokens served from cache by settled hits; /// `blocks_cached` is the current number of blocks held by the trie. void baseRT_prefix_cache_stats(baseRT_model_t model, uint64_t *out_hits, uint64_t *out_misses, uint64_t *out_reused_tokens, int *out_blocks_cached); +// === Boundary-snapshot blobs (hybrid GDN prefix reuse) === +// +// A hybrid (Gated-DeltaNet) model's cached KV blocks are only half of the +// state a warm seed needs: the recurrence at the SAME position is the other +// half. These calls attach an opaque per-boundary blob (a +// baseRT_sequence_gdn_capture snapshot) to a cached prefix at an exact +// block-aligned depth, discover the deepest attached boundary along a +// prompt's match path, and fetch the blob back for a +// baseRT_sequence_gdn_restore after seeding the KV. The cache stores bytes; +// validity-at-position is the caller's contract. Attachments ride node +// eviction (a dropped prefix drops its blobs), live under their own byte +// budget (BASERT_PREFIX_BLOB_BUDGET_MB, default 512, oldest evicted first), +// and are NOT persisted by baseRT_prefix_cache_save. + +/// Attach `blob` at exactly the block-aligned depth `n_tokens` of a prefix +/// already published with baseRT_prefix_insert (insert first, then attach). +/// Replaces an existing attachment at that depth. Best-effort: a disabled +/// cache, an uncovered/misaligned path, or a blob over budget is a no-op. +/// Returns BASERT_OK (attached or clean no-op) or an argument error. +int baseRT_prefix_attach_blob(baseRT_model_t model, const uint32_t *tokens, int n_tokens, const void *blob, + int blob_len); + +/// The deepest attached-blob depth (in tokens) along the exact-match path of +/// `tokens`, capped at `n_tokens`. 0 = none / cache disabled. Read-only (no +/// lock, no LRU update): use it to cap a hybrid match at a boundary a seed +/// can actually restore. +int baseRT_prefix_blob_match(baseRT_model_t model, const uint32_t *tokens, int n_tokens); + +/// Copy the blob attached at exactly depth `n_tokens` into `out` (capacity +/// `cap`). Returns the blob's byte length, or -1 when no blob is attached at +/// that exact depth / the buffer is too small / the cache is disabled. +int baseRT_prefix_blob_fetch(baseRT_model_t model, const uint32_t *tokens, int n_tokens, void *out, int cap); + // === Grammar-constrained decoding === /// Opaque grammar handle. @@ -793,6 +1239,19 @@ baseRT_grammar_t baseRT_grammar_create(baseRT_model_t model, const char *gbnf); /// Returns NULL on error. baseRT_grammar_t baseRT_grammar_create_from_schema(baseRT_model_t model, const char *json_schema); +/// Create a grammar from an xgrammar STRUCTURAL TAG (a JSON document, not a +/// grammar string). +/// +/// A structural tag constrains only the parts of the output that matter: free +/// text until a trigger appears, then the matching tag's schema until its end +/// tag, then free text again. This is what lets a model choose whether to call +/// a tool while guaranteeing any call it makes is well-formed and delimited — +/// which a whole-output schema grammar cannot express. +/// +/// Returns NULL on error (malformed tag, or a schema inside it that cannot be +/// compiled); the caller should decode unconstrained rather than fail. +baseRT_grammar_t baseRT_grammar_create_from_structural_tag(baseRT_model_t model, const char *tag_json); + /// Create a grammar for generic JSON output (any valid JSON object/array). baseRT_grammar_t baseRT_grammar_create_json(baseRT_model_t model); @@ -812,12 +1271,20 @@ void baseRT_grammar_reset(baseRT_grammar_t grammar); /// baseRT_grammar_bitmask_size : packed int32 words in the token bitmask /// (0 = not an xgrammar grammar; use the serial decode path instead). /// baseRT_grammar_fill_bitmask : fill `out_bitmask` (bitmask_size words) for -/// the CURRENT grammar state; a set bit = allowed token. 1 on success. +/// the CURRENT grammar state; a set bit = allowed token. Returns 1 when +/// the mask excludes some token and is to be applied, 0 when there is +/// nothing to apply: every token may follow (free text in a structural +/// tag, where the model's markers are text as they are unconstrained), +/// or the grammar has no matcher. Not an error indicator. /// baseRT_grammar_accept_token : advance the grammar by one token. 1 on ok. /// baseRT_grammar_is_terminated: 1 once the grammar reaches an end state. /// baseRT_grammar_is_completed : 1 once a full match is accepted (a /// structured value is complete). Decoding should stop on terminated OR -/// completed — matching the serial grammar loop. +/// completed — matching the serial grammar loop. A STRUCTURAL-TAG +/// grammar never reports completed: free text is admitted around (and, +/// unless at_least_one, instead of) its tagged regions, so "could end +/// here" holds at every free-text step and is no reason to stop — those +/// lanes end on EOS (or terminated). int baseRT_grammar_bitmask_size(baseRT_grammar_t grammar); int baseRT_grammar_fill_bitmask(baseRT_grammar_t grammar, int32_t *out_bitmask); int baseRT_grammar_accept_token(baseRT_grammar_t grammar, uint32_t token_id); @@ -836,6 +1303,13 @@ BaseRTGenerationStats baseRT_generate_grammar_continue(baseRT_model_t model, con baseRT_grammar_t grammar, baseRT_token_callback callback, void *user_data); +/// `baseRT_generate_resume` with a grammar constraint (see there): prefill +/// `[n_cached, n_prompt)` over the kept KV, seed history from the whole prompt. +BaseRTGenerationStats baseRT_generate_grammar_resume(baseRT_model_t model, const uint32_t *prompt_tokens, int n_prompt, + int n_cached, int max_tokens, BaseRTSamplingConfig sampling, + baseRT_grammar_t grammar, baseRT_token_callback callback, + void *user_data); + // === GPU sampling === /// Run a profiled decode step — returns per-layer GPU timing. @@ -844,6 +1318,13 @@ BaseRTGenerationStats baseRT_generate_grammar_continue(baseRT_model_t model, con /// timing_out: array of (n_layers + 3) floats [embedding, norm, layer0..N-1, logit, argmax] /// Returns number of timing entries written, or -1 on a failed step /// (details via baseRT_get_error). +/// Per-kernel profile of one baked BATCHED decode round (see the .cpp +/// comment). Requires a prior baked batched round at this B. Labels via +/// baseRT_profile_batched_label(model, n_seqs, index). +int baseRT_profile_batched_decode(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, const uint32_t *in_tokens, + float *timing_out, int max_entries); +const char *baseRT_profile_batched_label(baseRT_model_t model, int n_seqs, int index); + int baseRT_profile_decode_step(baseRT_model_t model, uint32_t token_id, int position, float *timing_out, int max_entries); @@ -872,6 +1353,11 @@ uint32_t baseRT_tensor_dtype(baseRT_model_t model, int index); /// Returns empty string out of range. const char *baseRT_tensor_raw_dtype(baseRT_model_t model, int index); +/// Byte size of the tensor's storage blob at index (weights only — MLX +/// companion `.scales`/`.biases` tensors are separate entries with their +/// own sizes). Returns 0 out of range. +size_t baseRT_tensor_nbytes(baseRT_model_t model, int index); + /// Whether the loaded model carries an mmproj sub-bundle (vision/audio /// tower weights). Returns 0/1. int baseRT_has_mmproj(baseRT_model_t model); @@ -1014,12 +1500,18 @@ int baseRT_load_state(baseRT_model_t model, const char *path); /// every forward pass that runs a GEMM with a tensor_name registered in /// the adapter has a post-GEMM low-rank delta applied (`y += B @ A @ x`). /// -/// While an adapter is active, single-sequence decode skips the baked -/// dispatch-table replay (the table carries no delta dispatches) and takes -/// the per-token immediate encode instead — correct output at reduced -/// decode throughput. Speculative decode is likewise disabled for the -/// duration. The batched multi-sequence API (`baseRT_batch_step*`) does -/// NOT apply adapters. +/// While an adapter is active, decode skips the baked dispatch-table replay +/// (the table carries no delta dispatches, and one recorded WITH them would go +/// stale on unload) and takes the immediate encode instead — correct output at +/// reduced throughput. Speculative decode is likewise disabled for the +/// duration. +/// +/// The batched multi-sequence entry points DO apply the adapter: +/// `baseRT_batch_decode_step`, `baseRT_batch_step` and the fused/logits +/// variants each install it for the pass, bypass baked replay and command +/// bucketing while it is active, and emit the delta GEMMs per layer. (This +/// paragraph previously said the opposite, which would lead a caller to +/// disable batching or write a serial fallback it does not need.) /// /// Calling `baseRT_lora_load` again replaces the active adapter (no /// stacking). Returns 0 on success, negative on failure (see @@ -1120,6 +1612,20 @@ BaseRTGenerationStats baseRT_generate_continue(baseRT_model_t model, const uint3 int max_tokens, BaseRTSamplingConfig sampling, baseRT_token_callback callback, void *user_data); +/// Resume generation over a KV the caller kept: the KV holds exactly the first +/// `n_cached` tokens of `prompt_tokens` (e.g. after baseRT_try_rollback to a +/// shared prefix), and only `[n_cached, n_prompt)` is prefilled. Unlike +/// `baseRT_generate_continue`, which sees only the tokens it is given, the +/// channel decoder and the sampling history (repetition / presence / +/// frequency penalties) are seeded from the WHOLE prompt, so the run samples +/// as a cold `baseRT_generate(prompt_tokens, n_prompt)` would. When the KV +/// length is not `n_cached`, or `n_cached` leaves no token to prefill, it runs +/// the prompt cold instead. `prompt_tokens` in the returned stats counts the +/// whole prompt. +BaseRTGenerationStats baseRT_generate_resume(baseRT_model_t model, const uint32_t *prompt_tokens, int n_prompt, + int n_cached, int max_tokens, BaseRTSamplingConfig sampling, + baseRT_token_callback callback, void *user_data); + // === Embeddings === /// Compute text embeddings from token IDs using the model's hidden states. @@ -1134,6 +1640,11 @@ int baseRT_embed_text(baseRT_model_t model, const char *text, float *out_embeddi /// Get the embedding dimension for a model. int baseRT_embedding_dim(baseRT_model_t model); +/// True for a BERT-style encoder-only embedding model (no autoregressive +/// decode, no paged KV) — the serving worker uses this to allow loading such a +/// model despite it having no paged decode KV. +bool baseRT_is_embedding_model(baseRT_model_t model); + // === Chat templates === /// Format a chat prompt using the model's native template. @@ -1151,6 +1662,18 @@ const char *baseRT_chat_template(baseRT_model_t model); /// the model is freed. const char *baseRT_chat_template_jinja(baseRT_model_t model); +/// Number of added/special tokens the model declares (chat scaffold, +/// reasoning/tool channel markers, etc.). Pairs with +/// `baseRT_special_token` so a codec can learn a model's markers from the +/// model itself rather than assuming a hardcoded dialect. +int baseRT_special_token_count(baseRT_model_t model); + +/// The `index`-th special token's string; writes its token id (or +/// UINT32_MAX if the string is not a single vocab entry) to `id_out`. +/// Returns "" for an out-of-range index. Returned pointer valid until the +/// next call on the same thread. +const char *baseRT_special_token(baseRT_model_t model, int index, uint32_t *id_out); + /// BOS / EOS token strings (what minja substitutes for `{{ bos_token }}` /// and `{{ eos_token }}` in HF chat templates). const char *baseRT_bos_token(baseRT_model_t model); @@ -1171,6 +1694,21 @@ const char *baseRT_eos_token(baseRT_model_t model); /// other token-id consumers stop on). Returns 0 on a null handle. uint32_t baseRT_eos_token_id(baseRT_model_t model); +/// Content-derived weights/bundle identity: FNV-1a-64 over the model +/// file's first MiB, stamped at load (basertd §4 materialization +/// identity input — two fine-tunes of one architecture never share a +/// version_tag). 0 on a null handle or unidentified weights. +uint64_t baseRT_weights_identity(baseRT_model_t model); + +/// Whether `token` is ANY of the model's stop tokens — the primary +/// `eos_token_id` plus the extra ids the loader registers (`<|eot|>`, +/// `<|eot_id|>`, ``, gpt-oss's `<|call|>` and `<|endoftext|>`, +/// …). Multi-EOS models (Muse Glimmer: +/// `<|eot|>` ends the turn while `eos_token_id` only ends the document) +/// need this; comparing against `baseRT_eos_token_id` alone never +/// terminates their chat turns. Returns 0 on a null handle. +int32_t baseRT_is_eos_token(baseRT_model_t model, uint32_t token); + // (baseRT_max_prefill_chunk is declared once, with the batched-decode API.) // === Token counting === diff --git a/include/baseRT/baseRT.h b/include/baseRT/baseRT.h index 224dcd5..04823d8 100644 --- a/include/baseRT/baseRT.h +++ b/include/baseRT/baseRT.h @@ -63,8 +63,8 @@ extern "C" { // === Versioning === #define BASERT_VERSION_MAJOR 0 -#define BASERT_VERSION_MINOR 2 -#define BASERT_VERSION_PATCH 6 +#define BASERT_VERSION_MINOR 3 +#define BASERT_VERSION_PATCH 0 /// Compile-time version, packed as `(MAJOR<<16) | (MINOR<<8) | PATCH`. /// Useful for `#if BASERT_VERSION >= 0x000200` feature checks. @@ -84,8 +84,11 @@ typedef void *baseRT_model_t; /// Other source formats (GGUF, HF safetensors, MLX safetensors) must be /// converted offline first via `basert convert`. /// kernel_library_path: path to the compiled GPU kernel library (on the Metal -/// backend, baseRT.metallib), or NULL to auto-detect. Auto-detect order: a -/// kernel library next to the build, then a copy embedded in the loaded +/// backend, baseRT.metallib; on CUDA, baseRT_cuda_kernels.fatbin; on ROCm, +/// baseRT_rocm_kernels.hsaco), or NULL to auto-detect. Auto-detect order: +/// $BASERT_KERNEL_LIB (a file, or a directory holding the artifact), a +/// kernel library next to the build, next to the +/// executable, the current directory, then a copy embedded in the loaded /// binary itself (single-file distributions ship the shared library with the /// kernels linked in, so NULL just works). Named generically so non-Metal /// backends (CUDA/ROCm, future) can reuse the same parameter. @@ -97,6 +100,17 @@ typedef void *baseRT_model_t; /// Returns NULL on failure. baseRT_model_t baseRT_load_model(const char *model_path, const char *kernel_library_path, int max_context); +/// Load ONLY the bundle's tokenizer: metadata parse + vocab/merges, no +/// tensor upload, no KV/scratch allocation (a serving daemon tokenizes in +/// its own process while a worker owns the GPU state — basertd §3's codec +/// placement). The returned handle supports exactly the tokenizer surface: +/// baseRT_encode, baseRT_decode_token / _static / _raw, baseRT_token_count, +/// baseRT_eos_token_id, baseRT_bos_id / baseRT_bos_token / baseRT_eos_token, +/// baseRT_get_config, and baseRT_free_model. Calling generation or +/// sequence APIs on it is undefined. NULL on failure (see +/// baseRT_get_error). +baseRT_model_t baseRT_load_tokenizer_only(const char *model_path); + /// Load a model with per-call options instead of the process-wide /// baseRT_set_* pre-load setters. `opts` may be NULL (identical to /// baseRT_load_model); otherwise set opts->struct_size = @@ -138,8 +152,27 @@ uint32_t baseRT_capabilities(baseRT_model_t model); /// Returns BASERT_OK, or an error code (leaves the model on the OLD library on /// load failure). Not supported for whisper models. NOT thread-safe against /// in-flight inference — serialize the caller. +/// +/// INVALIDATES ALL DECODED STATE. Everything in a KV or recurrent slab was +/// produced by the old kernels, so resuming on it would mix kernel versions — +/// the one thing a tuning comparison must not do. On success the default +/// sequence's KV, recurrent state and prefix cache are reset for you (position +/// returns to 0). Caller-owned sequences (baseRT_sequence_create) are NOT +/// reachable from here and must be destroyed and recreated; using one across a +/// reload is undefined. int baseRT_reload_metallib(baseRT_model_t model, const char *metallib_path); +/// A JSON capability descriptor for this loaded model (schema field +/// `descriptor_version`): workload, modalities, serving, state, adaptation, +/// and tick-op support, each capability either `true` or a stable reason +/// string explaining why it is unavailable on this architecture + backend + +/// load-options triple. Computed once at load (verdicts frozen, like the +/// multi-row overcommit verdict) — immutable, race-free, never touches the +/// error state. The pointer is owned by the model handle and valid until +/// baseRT_free_model. Supersedes the five-bit mask above for every question +/// it cannot answer; the mask stays for cheap scheduler branching. +const char *baseRT_capability_descriptor(baseRT_model_t model); + /// Override the KV cache element width for the next baseRT_load_model call. /// bits = 0 → auto (per-model default; Q8_0 when head_dim%32==0) /// bits = 8 → force Q8_0 K and V slabs (1.88x smaller; tiny precision cost) @@ -148,6 +181,16 @@ int baseRT_reload_metallib(baseRT_model_t model, const char *metallib_path); /// baseRT_load_model. Other values are ignored. void baseRT_set_kv_bits(int bits); +/// The KV cache's element width as ACTUALLY allocated for `model`, in the +/// `baseRT_set_kv_bits` encoding: 16 (f16), 8 (Q8_0), 4 (Q4_0), 84 (K at +/// Q8_0, V at Q4_0), or 0 for a model with no KV cache. Differs from what +/// was set whenever the loader overrode it — hybrid GDN / SSM, gpt-oss and +/// MLA models force f16; a paged pool is uniform (84 → 8); a head_dim that +/// is not a multiple of 32 cannot take the quantized blocks — and the load +/// printed one `[baseRT] : --kv-bits N requested, ...` line saying +/// why. Serve this rather than the setting when reporting KV precision. +int baseRT_kv_bits_effective(baseRT_model_t model); + /// Enable engine diagnostics (RoPE/tokenizer/GPU/architecture dumps, the /// per-token dispatch-command count, "Warming up"). Off by default so end /// users see only model output. @@ -244,14 +287,19 @@ size_t baseRT_model_config_sizeof(void); /// Get total GPU memory used by model (bytes). size_t baseRT_model_memory(baseRT_model_t model); -/// The device memory budget the engine allocates within, in bytes, without -/// needing a loaded model. NOT installed RAM: on Metal this is the unified- -/// memory working set the OS recommends (~75% of physical RAM), on CUDA the -/// device's total memory. 0 when no supported device is present. +/// The device memory budget automatic context sizing plans within, in bytes, +/// without needing a loaded model. NOT installed RAM: on Metal this is 85% of +/// the unified-memory working set the OS recommends (itself a fraction of RAM +/// the OS picks per machine, roughly two thirds to three quarters) — the share +/// the engine keeps its pinned weights within, since past it Metal evicts from +/// the pinned set and reads slow down or go corrupt; on CUDA the device's total +/// memory. 0 when no supported device is present. /// -/// This is the number that matters for sizing decisions — it is what the load -/// path compares the model's working set against before warning that the OS -/// will start paging weights. +/// This is the number `baseRT_suggest_max_context*` and +/// `baseRT_context_window_fits` size against. It is a planning line, not a +/// hard limit: a load whose working set lands between it and the full working +/// set still succeeds, with the tail of its weights left to page-cache LRU (and +/// the load warns once the working set passes the full budget). size_t baseRT_device_memory_budget(void); /// Floor on a window derived by `baseRT_suggest_max_context`. Policy, not a @@ -298,6 +346,29 @@ int baseRT_suggest_max_context(const char *model_path, int max_batch, int kv_bit int baseRT_suggest_max_context_multi(const char *const *model_paths, int n_models, int max_batch, int kv_bits, int paged_kv); +/// As `baseRT_suggest_max_context_multi`, plus the speculators that will load +/// beside those models. A drafter runs at its target's window and lane count +/// with its own paged KV pool, so it spends the same budget: leaving it out +/// sizes the targets' pools into memory the drafter then cannot get. +/// `speculator_embedded[i]` non-zero marks speculator i as its target +/// bundle's own embedded head (header.speculator, e.g. MTP): its KV pool only, +/// its weights being part of the bundle. Zero — or a null array — is a +/// separately loaded bundle (a DFlash / DSpark / EAGLE-3 sidecar or a draft +/// model): its weights and its KV pool, even if the same file is also one of +/// `model_paths` (a draft model is loaded again beside the served copy). +/// `speculator_target[i]` >= 0 marks speculator i as a sidecar loaded with +/// `baseRT_load_drafter` against `model_paths[speculator_target[i]]`: that +/// loader reopens the target's bundle for its token embedding / lm_head, and a +/// backend that copies wrapped weights (CUDA's UMA path) holds those rows a +/// second time, so they are charged too. -1 — or a null array — is a +/// speculator with no such view (a draft model loads on its own); it is +/// ignored for an embedded head, whose view of its own bundle is always charged. +/// Speculators never cap the window, one that cannot be read is skipped, and +/// none is charged unless `paged_kv` (their loaders require a paged target). +int baseRT_suggest_max_context_spec(const char *const *model_paths, int n_models, const char *const *speculator_paths, + const int *speculator_embedded, const int *speculator_target, int n_speculators, + int max_batch, int kv_bits, int paged_kv); + /// Does `window` fit `model_paths` co-resident on this device, under the same /// budget the suggestion above derives from? 1 = yes, 0 = no, -1 = cannot tell /// (a bundle would not open, no trained window, unknown budget). @@ -325,6 +396,11 @@ typedef struct BaseRTMemoryStats { uint64_t kv_cache_used_bytes; uint64_t kv_cache_blocks_total; uint64_t kv_cache_blocks_used; + /// Tokens resident in the KV. Exact on a contiguous cache. On a PAGED cache + /// this is block-granular — `kv_cache_blocks_used * page_size`, so each + /// sequence's partially-filled last block counts whole — because the pool + /// tracks blocks, not per-sequence lengths, and the default lane's own + /// length reads 0 whenever batched sequences own the live state. uint64_t kv_cache_tokens_used; BaseRTKVCacheLayout kv_cache_layout; uint32_t reserved; @@ -352,8 +428,37 @@ const char *baseRT_strerror(BaseRTErrorCode code); // === Tokenization === /// Encode text to token IDs. Returns number of tokens written. +/// +/// A substring equal to one of the model's special tokens (`<|im_end|>`, +/// ``, ``, …) becomes that single control id — what a +/// rendered chat prompt needs, since its framing IS those markers. int baseRT_encode(baseRT_model_t model, const char *text, uint32_t *out_tokens, int max_tokens); +/// Encode text with special-token strings treated as ordinary text. +/// +/// The same substrings tokenize to their BPE pieces instead of control ids, +/// so a message body, tool result or document that QUOTES a marker cannot +/// forge a turn boundary or open a reasoning block. Use this for client +/// content and `baseRT_encode` for a dialect's own framing. No automatic +/// BOS-dedup either: a leading BOS string here is content. Returns number +/// of tokens written. +int baseRT_encode_plain(baseRT_model_t model, const char *text, uint32_t *out_tokens, int max_tokens); + +/// Encode a sequence of text pieces as ONE text. +/// +/// Piece `i` is framing when `plain[i]` is zero (marker strings parse to +/// control ids, as in `baseRT_encode`) and client content otherwise (they +/// stay ordinary text, as in `baseRT_encode_plain`). Unlike encoding the +/// pieces one by one, the pre-tokenizer and BPE see the concatenation, so a +/// boundary between two pieces tokenizes exactly as it does in a whole-text +/// encode of the same characters (a content piece starting with a newline +/// after a framing newline still merges into the `\n\n` token) and the +/// automatic BOS is handled once at the front. A chat codec renders a turn +/// with this: the role line and the wire around a message are framing, the +/// message is content. Returns number of tokens written. +int baseRT_encode_pieces(baseRT_model_t model, const char *const *texts, const int *plain, int n_pieces, + uint32_t *out_tokens, int max_tokens); + /// Decode a single token ID to text. Returns static string (do not free). /// /// Advances the tokenizer's incremental-decode state — call this from @@ -375,17 +480,22 @@ const char *baseRT_decode_token_static(baseRT_model_t model, uint32_t token_id); typedef struct baseRT_decode_stream *baseRT_decode_stream_t; baseRT_decode_stream_t baseRT_decode_stream_create(baseRT_model_t model); void baseRT_decode_stream_reset(baseRT_decode_stream_t stream); + +/// Reset a stream for a lane whose generation resumes after `prompt`, deriving +/// the channel resume state from its last framing token. Prefer this over +/// `baseRT_decode_stream_reset` on any lane that may carry a RAW (unframed) +/// prompt: the plain reset assumes a chat prompt, which ends inside a +/// `<|start|>assistant` header, and on a Muse vocabulary an unframed prompt +/// left in that state has every generated token buffered as protocol until a +/// 128-byte valve trips — a short completion returns nothing at all. Harmony +/// has its own fail-open path and is unaffected either way. `prompt` NULL or +/// `n_prompt` 0 behaves exactly like `baseRT_decode_stream_reset`. +void baseRT_decode_stream_reset_for_prompt(baseRT_decode_stream_t stream, const uint32_t *prompt, int n_prompt); /// Returns a pointer into the stream's own buffer, valid until the next /// call on the same stream. const char *baseRT_decode_stream_token(baseRT_decode_stream_t stream, uint32_t token_id); void baseRT_decode_stream_free(baseRT_decode_stream_t stream); -/// 1 when `token_id` ends the assistant turn: the bundle's eos_token_id OR -/// any of its extra end-of-turn ids (Llama 3's <|eot_id|>, gpt-oss's -/// <|call|> and <|endoftext|>, …). `baseRT_eos_token_id` reports only the -/// first; a lane that compares against that alone runs past the others. -int baseRT_is_eos_token(baseRT_model_t model, uint32_t token_id); - /// Length-preserving variant of `baseRT_decode_token_static` for callers /// that need the token's EXACT raw bytes. Byte-level BPE / byte-fallback /// tokens can decode to bytes containing 0x00, which the C-string variants @@ -588,7 +698,19 @@ int baseRT_max_prefill_chunk(baseRT_model_t model); /// Roll a sequence's KV state back to `length` tokens, returning any blocks /// past that point to the pool. `length` must be <= the current length; 0 /// resets the sequence to empty. Used by the serving engine's shape-padding -/// dummy lanes (their KV is discarded after every tick). +/// dummy lanes (their KV is discarded after every tick) and by speculative +/// verification (drop the rejected drafts). Recurrent-state models: a +/// hybrid-GDN lane rolls back only to a length inside its last +/// baseRT_batch_step_fused_logits_rows feed (the recurrence is rebuilt from +/// that feed's captured rows — deterministic, and bit-identical to a feed +/// that stopped there under bit-exact routing (BASERT_SPEC_BITEXACT=1); +/// the default speculative routing picks its small-M GEMM kernel by row +/// count, so the replayed prefix carries the longer feed's last-ulp +/// rounding instead); `length == current` is a no-op commit on EVERY +/// lane (it drops a hybrid-GDN lane's capture); any shorter target on a +/// Mamba-2 lane, or outside the captured feed on a hybrid-GDN lane, +/// returns BASERT_ERR_UNSUPPORTED — a no-op rollback is not a probe of +/// rollback support. int baseRT_sequence_rollback(baseRT_sequence_t seq, int length); /// Multi-step autoregressive driver for baseRT_batch_step_fused. Step 0 @@ -619,10 +741,289 @@ int baseRT_batch_step_fused_loop(baseRT_model_t model, baseRT_sequence_t *seqs, int baseRT_batch_step_fused_logits(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, const uint32_t *in_tokens, const int *in_token_counts); -/// Read back the [n_seqs, vocab] f16 logits left by the most recent -/// baseRT_batch_step_fused_logits into `out_logits_f16` (n_seqs * vocab halves, -/// row-major). Pure UMA copy, no dispatch. Returns vocab_size, or <0 on error. -/// `n_seqs` must match the batch of the preceding step and be <= max_batch_size. +/// Per-row sampling spec for baseRT_batch_step_fused_sample. temperature <= 0 +/// samples greedily (argmax) for that row; otherwise Gumbel-max over the +/// top-k / top-p / min-p survivors with `seed` driving the noise stream +/// (same seed + same logits = same token). +/// +/// Penalties ARE applied for a sampled row, over `history` — the caller's +/// recent-token window for that row — with the model's penalty-exemption mask +/// honoured, so a sampled row no longer has to leave the fused path to get +/// one. Logit bias and grammar still take the host logits path. +typedef struct BaseRTRowSampling { + float temperature; + float top_p; ///< 1.0 = disabled + float min_p; ///< 0.0 = disabled + int32_t top_k; ///< <= 0 = disabled + uint32_t seed; ///< Gumbel noise seed for this row + float repeat_penalty; ///< 1.0 (or 0.0) = disabled + float presence_penalty; ///< 0.0 = disabled + float frequency_penalty; ///< 0.0 = disabled + /// Recent tokens this row's penalties are computed over, most recent last. + /// NULL / 0 disables the penalties for the row regardless of the values + /// above: with no history there is nothing to penalize. + const uint32_t *history; + uint32_t history_count; +} BaseRTRowSampling; + +/// Fused batch step that samples every output row ON THE GPU per `rows[i]` +/// (greedy rows argmax, sampled rows Gumbel-max) and returns one token per +/// seq in `out_tokens` — no full-vocab logits readback and no per-row host +/// sampling. Same forward as baseRT_batch_step_fused_logits (so every +/// architecture that serves host logits serves this), followed by one +/// gumbel_topk_f16 dispatch per sampled row and one batched argmax. +/// Returns BASERT_ERR_UNSUPPORTED when the backend lacks the sampling +/// kernels (caller falls back to the logits path). +int baseRT_batch_step_fused_sample(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, const uint32_t *in_tokens, + const int *in_token_counts, const BaseRTRowSampling *rows, uint32_t *out_tokens); + +/// The sampling pass of baseRT_batch_step_fused_sample on its own: samples +/// logits rows [first_row, first_row + n_rows) that the last fused forward +/// left in the host logits buffer (baseRT_batch_step_fused_logits_rows / +/// _rows_taps — a speculative verify's rows), `rows[i]` driving row +/// first_row + i, one token per row into `out_tokens`. The pass noises the +/// rows in place: read them back first (baseRT_read_batch_logits) if the raw +/// values are still needed. BASERT_ERR_UNSUPPORTED when the backend lacks the +/// sampling kernels or the penalty windows do not fit — the rows are untouched +/// then and the caller samples them on the host. +int baseRT_sample_logits_rows(baseRT_model_t model, int first_row, int n_rows, const BaseRTRowSampling *rows, + uint32_t *out_tokens); + +/// Multi-row variant for speculative verification (basertd P2.4): seq i +/// contributes in_token_counts[i] input rows and leaves its LAST +/// out_rows[i] rows' logits (1 <= out_rows[i] <= in_token_counts[i]) in the +/// host logits buffer, ascending, seqs in order — sum(out_rows) rows total, +/// read back with baseRT_read_batch_logits(sum(out_rows)). The host checks +/// each drafted position against the target's logits from ONE forward and +/// rolls the sequence back to the accepted length. Served on the fused +/// VARLEN path and on the staged decomposition (MoE, hybrid GDN, residency +/// overcommit) alike; Mamba-2 hybrids return BASERT_ERR_UNSUPPORTED (no +/// per-lane SSM rollback). On the staged route a lane's rows must fit in +/// its last prefill chunk; a hybrid-GDN lane's whole feed must be one +/// chunk of at most 31 tokens, and the lane's next call must be either +/// baseRT_sequence_rollback (to any length inside the feed) or a plain step +/// (which commits the feed). +/// sum(out_rows) is bounded by the logits scratch (max_batch_size × 8). +int baseRT_batch_step_fused_logits_rows(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, + const uint32_t *in_tokens, const int *in_token_counts, const int *out_rows); + +/// Hidden-state taps for speculative drafters (MTP heads, EAGLE-3, DFlash +/// condition on the target's residual stream). Same step as +/// baseRT_batch_step_fused_logits_rows, plus: for each of the `n_tap_layers` +/// requested hidden_states indices in `tap_layers` (0 = embedding output, +/// k = residual stream after decoder layer k-1, n_layers = final pre-norm +/// hidden), the LAST tap_rows[i] fed rows of lane i are copied into the +/// engine's tap buffer, rows in lane order, one contiguous [rows, dim] f16 +/// block per tap (read back with baseRT_read_tap_rows, or consumed on the +/// device by an in-engine drafter). At most 8 distinct tap layers and +/// baseRT_tap_rows_capacity() tapped rows per call; a lane's tapped rows +/// must fit in its last prefill chunk. A 1-token lane with tap_rows 1 runs +/// its step serially (the batched decode tables cannot tap). tap_rows may +/// name 0 rows for lanes that need no taps. Taps are emitted by the llama +/// family, Gemma 4, Qwen3.5, gpt-oss and GLM-DSA encoders; a tapped step on any other +/// architecture returns BASERT_ERR_UNSUPPORTED (BASERT_ERR_OUT_OF_MEMORY +/// when the tap buffer cannot be allocated). +/// Greedy speculative verification without a logits readback: the same +/// multi-row forward as baseRT_batch_step_fused_logits_rows_taps (taps +/// optional: n_tap_layers = 0), every requested row argmaxed on the GPU; +/// `out_tokens` receives sum(out_rows) tokens in lane order. Row i of a +/// lane is the target's greedy choice after the lane's i-th fed token — +/// accept draft i while it equals row i. BASERT_ERR_UNSUPPORTED on bundles +/// that verify through host logits (MoE / hybrid / non-llama output +/// stages): fall back to the logits variant. +int baseRT_batch_step_fused_rows_argmax(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, + const uint32_t *in_tokens, const int *in_token_counts, const int *out_rows, + const int *tap_layers, int n_tap_layers, const int *tap_rows, + uint32_t *out_tokens); + +int baseRT_batch_step_fused_logits_rows_taps(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, + const uint32_t *in_tokens, const int *in_token_counts, const int *out_rows, + const int *tap_layers, int n_tap_layers, const int *tap_rows); + +/// Rows one tapped call may cover (min(max_prefill_chunk, 1024)); <0 on a +/// null model. +int baseRT_tap_rows_capacity(baseRT_model_t model); + +/// Tapped step with an INPUT OVERRIDE: `input_f16` ([sum(in_token_counts), +/// dim] f16, rows in feed order) replaces the embedding lookup — the +/// drafters' input is a projection of the target's hidden state, not a +/// token. Token ids still drive positions and KV. Rows are bounded by +/// baseRT_tap_rows_capacity(). Feeding a step's own tap-0 rows back through +/// this call reproduces its logits bit for bit. +int baseRT_batch_step_fused_logits_rows_taps_input(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, + const uint32_t *in_tokens, const int *in_token_counts, + const int *out_rows, const int *tap_layers, int n_tap_layers, + const int *tap_rows, const void *input_f16); + +// ── Speculator heads (bundle `header.speculator`) ────────────────────── + +/// 1 when the loaded bundle carries a speculator sub-bundle (an MTP head or +/// a converted drafter), 0 otherwise. +int baseRT_has_speculator(baseRT_model_t model); + +/// The speculator's `kind` ("mtp", ...) into `out`; returns its length, 0 +/// when the bundle has no speculator, <0 on error. +int baseRT_speculator_kind(baseRT_model_t model, char *out, int cap); + +/// Whether a loaded speculator head can actually run the MTP prologue: the +/// prologue is an encoder's embedding stage — the Qwen3.5 hybrid encoder's +/// (`qwen35_mtp_prologue`) or GLM-DSA's nextn one (`glm_mtp_prologue`) — so +/// `baseRT_mtp_chain` and the per-step path both refuse a head on any other +/// architecture. A scheduler should ask this +/// before granting an MTP strategy rather than registering one whose every +/// proposal then falls back. 1 = supported, 0 = not, <0 = not a head. +int baseRT_speculator_head_supported(baseRT_model_t head); + +/// Whether this bundle's encoder can emit hidden-state taps. Every block +/// drafter and EAGLE-3 sidecar consumes target taps, so a scheduler should +/// ask before granting one of those strategies: without taps the first +/// tapped forward fails with BASERT_ERR_UNSUPPORTED, which the tick reports +/// as a generation failure rather than falling back to plain decoding. +/// 1 = taps available, 0 = not, <0 = bad argument. +int baseRT_model_supports_taps(baseRT_model_t model); + +/// Load the bundle's speculator head as its own model handle: the head's +/// decoder layer(s) and norms from the speculator section, the target's +/// embeddings and lm_head shared. It is a paged model of the target's +/// context with its own sequences (baseRT_sequence_create), stepped with +/// the tapped/override steps above; free it with baseRT_free_model before +/// the target. NULL (with an error) when the bundle has no loadable head. +baseRT_model_t baseRT_load_speculator_head(baseRT_model_t target); + +/// MTP head input rows (host convenience for tests and drivers): out[r] = +/// fc_embed · rmsnorm(embed(next_tokens[r])) + fc_hidden · rmsnorm(hidden[r]) +/// (GLM-DSA nextn heads: eh_proj · [enorm(embed) | hnorm(output_norm(hidden))]) +/// — `hidden_f16` the target's final pre-norm hidden ([rows, dim], tap +/// n_layers) at each row's position, `out_f16` [rows, dim] to feed the head +/// as its input override. The head predicts the token AFTER next_tokens[r]. +int baseRT_mtp_prologue(baseRT_model_t head, const uint32_t *next_tokens, int rows, const void *hidden_f16, + void *out_f16); + +/// The whole MTP-head proposal of `n_seqs` head lanes in ONE encoding (a +/// GLM-DSA head runs it as staged per-step forwards, same contract): +/// each lane's committed rows (`counts[i]` of them, 1..15; `tokens` the +/// NEXT token of every row packed lane by lane, `hidden_f16` the target's +/// final pre-norm hidden of every row, [sum counts][dim] f16) go through +/// the prologue and one head forward, then `n_draft - 1` chained steps of +/// one row per lane draft from the head's own output hidden and previous +/// draft, without a host round trip between steps. `out_drafts` receives +/// n_seqs * n_draft ids, lane-major (`out_drafts[i * n_draft + s]` = lane +/// i's draft s; draft 0 is the head's prediction after the last committed +/// row). Each lane's head sequence ends `counts[i] + n_draft - 1` +/// positions longer; roll the chained positions back before the next +/// commit. Requires the head's --paged-kv context; n_draft 1..15. +int baseRT_mtp_chain(baseRT_model_t head, baseRT_sequence_t *seqs, int n_seqs, const uint32_t *tokens, + const int *counts, const void *hidden_f16, int n_draft, uint32_t *out_drafts); + +/// Copy tap `tap_index` (its position in the call's tap_layers) of the most +/// recent tapped step into `out_f16`: n_rows * dim halves, row-major, rows +/// in lane order as fed. Bounded by that step: `tap_index` below its tap +/// count and `n_rows` at most its tapped rows (BASERT_ERR_INVALID_ARGUMENT +/// otherwise). Returns dim, or <0 on error. +int baseRT_read_tap_rows(baseRT_model_t model, int tap_index, int n_rows, void *out_f16); + +/// Block drafters (DFlash / DSpark): load a drafter SIDECAR bundle (arch +/// "dflash" / "dspark", converted from the drafter checkpoint) against +/// `target`. The drafter shares the target's token embedding and lm_head +/// (resolved from the target's bundle when the sidecar lacks them), runs at +/// the target's paged/lane shape, and is stepped with baseRT_drafter_step +/// on its own sequences (baseRT_sequence_create). Free it with +/// baseRT_free_model before the target. NULL with an error when the sidecar +/// does not fit the target (dim / vocab / taps). +baseRT_model_t baseRT_load_drafter(baseRT_model_t target, const char *path); + +/// Drafter geometry: `kind` ("dflash" | "dspark"), the block width, the +/// mask token the block is padded with, the first block row whose logits +/// are a draft (DFlash 1: row 0 is the anchor; DSpark 0), and the TARGET +/// hidden_states indices to tap (target_layer_ids + 1) in the order the +/// drafter's context expects them. Returns the tap count, or <0. +int baseRT_drafter_info(baseRT_model_t drafter, char *kind, int kind_cap, int *block_size, uint32_t *mask_token_id, + int *logits_start, int *tap_layers, int tap_cap); + +/// One drafter step on `seq`: append `n_ctx` context rows — the target's +/// tapped hidden rows at the next `n_ctx` positions, `ctx_taps_f16` laid +/// out [n_taps][n_ctx][dim] (tap-major, the baseRT_drafter_info order; +/// each tap block is exactly what baseRT_read_tap_rows returns) — then, +/// when `block_len > 0`, run the block `block[0..block_len)` (the last +/// committed token followed by mask tokens) after the context with +/// bidirectional block attention and leave its `block_len` logits rows for +/// baseRT_read_batch_logits (row j = block position j; drafts are rows +/// logits_start.. — the argmax of row j predicts the token at block +/// position j for DSpark, or the token replacing mask j for DFlash). The +/// block's KV is dropped; the sequence ends at its context length. Either +/// half may be empty (context-only ingestion of a prompt slice, or a block +/// on an up-to-date context). n_ctx <= baseRT_tap_rows_capacity(drafter); +/// `block_len` is 0 or exactly the block width (the drafters are trained +/// on a fixed-width bidirectional block — a narrower block is refused). +int baseRT_drafter_step(baseRT_model_t drafter, baseRT_sequence_t seq, const void *ctx_taps_f16, int n_ctx, + const uint32_t *block, int block_len); + +/// DSpark heads carried by the sidecar: `markov_rank` (0 = no Markov head) +/// and whether a confidence head is present. +int baseRT_drafter_heads(baseRT_model_t drafter, int *markov_rank, int *has_confidence); + +/// DSpark Markov head, one block row: add `markov.w2 · markov.w1[prev]` to +/// row `row` of the last block step's logits (in place) and return its +/// argmax in `out_token` — the draft for that row given the token before +/// it (`prev` = the pending token for the first drafted row, then the +/// previous draft). Sequential: call per row in order after +/// baseRT_drafter_step. +int baseRT_drafter_markov_argmax(baseRT_model_t drafter, int row, uint32_t prev, uint32_t *out_token); + +/// DSpark confidence head over the last block step's rows 0..rows: the +/// probability that block row r's draft survives verification, given the +/// token before it (`prev_tokens[r]`). Adaptive block length keeps drafts +/// while the running product stays above a threshold. Returns `rows`, or +/// <0 (BASERT_ERR_UNSUPPORTED without a confidence head). +int baseRT_drafter_confidence(baseRT_model_t drafter, const uint32_t *prev_tokens, int rows, float *out_conf); + +/// EAGLE-3 heads (a sidecar of arch "eagle3", loaded with baseRT_load_drafter): +/// one llama layer drafting the next token from (the current token, a +/// feature of the target's hidden states at three layers). `tap_layers` +/// are the TARGET hidden_states indices to tap (3), `draft_vocab` the +/// reduced vocabulary the head predicts over. Returns the tap count. +int baseRT_eagle3_info(baseRT_model_t head, int *draft_vocab, int *tap_layers, int tap_cap); + +/// One head step over `rows` positions from the sequence's length: +/// `tokens[r]` is the token that FOLLOWED position r (the token the row +/// conditions on), and exactly one of `taps_f16` ([3][rows][dim], the +/// target's tapped rows at those positions, baseRT_read_tap_rows layout) or +/// `hidden_f16` ([rows][dim], the head's own output hidden — chaining) is +/// the feature. Afterwards baseRT_eagle3_draft gives row r's draft (the +/// argmax over the draft vocabulary mapped to a target id) and +/// baseRT_eagle3_read the raw rows. +int baseRT_eagle3_step(baseRT_model_t head, baseRT_sequence_t seq, const uint32_t *tokens, int rows, + const void *taps_f16, const void *hidden_f16); +int baseRT_eagle3_read(baseRT_model_t head, int rows, void *logits_f16, void *hidden_f16); +int baseRT_eagle3_draft(baseRT_model_t head, int row, uint32_t *out_target_token); + +/// The whole proposal of `n_seqs` head sequences in one call: each lane's +/// `counts[i]` committed rows (tokens packed lane by lane; `taps_f16` the +/// target's tapped rows for ALL packed rows, [3][sum rows][dim]) as one +/// forward, then `n_draft - 1` chained one-row steps per lane on the +/// device (the previous draft is the token, the head's own output hidden +/// the feature) with no host round trip between them. `out_drafts[i * +/// n_draft + s]` is lane i's draft s (target ids). The head sequences end +/// `counts[i] + n_draft - 1` positions longer; roll back to drop the +/// chained positions. n_draft in 1..15. +int baseRT_eagle3_chain(baseRT_model_t head, baseRT_sequence_t *seqs, int n_seqs, const uint32_t *tokens, + const int *counts, const void *taps_f16, int n_draft, uint32_t *out_drafts); + +/// Media variant (basertd P2.7): the packed prompt's `image_token_id` +/// placeholders (exactly `n_rows` of them, in order) take rows from the +/// engine-owned features buffer produced by the media tower (`feats_buffer` +/// is a `baseRT::Buffer*` obtained through the tick media surface). Fused +/// VARLEN live path only (hybrid/MoE/overcommit return UNSUPPORTED). Host +/// sampled: read the logits back with baseRT_read_batch_logits. +int baseRT_batch_step_fused_logits_media(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, + const uint32_t *in_tokens, const int *in_token_counts, void *feats_buffer, + int n_rows, uint32_t image_token_id, int grid_h, int grid_w); + +/// Read back the [n_rows, vocab] f16 logits left by the most recent +/// baseRT_batch_step_fused_logits(_rows) into `out_logits_f16` (n_rows * vocab +/// halves, row-major). Pure UMA copy, no dispatch. Returns vocab_size, or <0 +/// on error. `n_rows` is the preceding step's batch (one row per seq) or its +/// sum(out_rows), bounded by the logits scratch (max_batch_size × 8 rows). int baseRT_read_batch_logits(baseRT_model_t model, int n_seqs, void *out_logits_f16); // === Host-side logits-row operations === @@ -749,6 +1150,14 @@ void baseRT_prefix_unlock(baseRT_model_t model, uint64_t handle); /// to seed it (e.g. a boundary mismatch). No-op for handle==0. void baseRT_prefix_release(baseRT_model_t model, uint64_t handle); +/// Drop a baseRT_prefix_match taken only to LOOK at what the cache holds (a +/// metadata probe, e.g. basertd's space_match): releases the pin like +/// baseRT_prefix_release but settles nothing, so the probe counts as neither +/// a hit nor a miss in baseRT_prefix_cache_stats. A match that was a reuse +/// decision and was declined belongs in baseRT_prefix_release. No-op for +/// handle==0. +void baseRT_prefix_discard(baseRT_model_t model, uint64_t handle); + /// Evict least-recently-used UNLOCKED cached prefixes until at least `n_blocks` /// block-frees have been performed back to the pool. Returns the number freed /// (may be < n_blocks if the remaining prefixes are all locked by live @@ -772,13 +1181,50 @@ int baseRT_prefix_cache_save(baseRT_model_t model, const char *path); /// (no outstanding matches). Returns BASERT_OK or an error code. int baseRT_prefix_cache_load(baseRT_model_t model, const char *path); -/// Lifetime prefix-cache stats (any out-pointer may be NULL). `hits`/`misses` -/// count baseRT_prefix_match calls that did / didn't reuse >=1 block; -/// `reused_tokens` is the running total of prompt tokens served from cache; +/// Lifetime prefix-cache stats (any out-pointer may be NULL), all monotonic. +/// A baseRT_prefix_match that found nothing counts as a miss at once; one that +/// found blocks settles later, as a hit at baseRT_prefix_unlock (its blocks +/// were seeded) or as a miss at baseRT_prefix_release (abandoned); one +/// dropped by baseRT_prefix_discard (a probe) is in neither, as is an +/// outstanding match. `reused_tokens` is the running total of +/// prompt tokens served from cache by settled hits; /// `blocks_cached` is the current number of blocks held by the trie. void baseRT_prefix_cache_stats(baseRT_model_t model, uint64_t *out_hits, uint64_t *out_misses, uint64_t *out_reused_tokens, int *out_blocks_cached); +// === Boundary-snapshot blobs (hybrid GDN prefix reuse) === +// +// A hybrid (Gated-DeltaNet) model's cached KV blocks are only half of the +// state a warm seed needs: the recurrence at the SAME position is the other +// half. These calls attach an opaque per-boundary blob (a +// baseRT_sequence_gdn_capture snapshot) to a cached prefix at an exact +// block-aligned depth, discover the deepest attached boundary along a +// prompt's match path, and fetch the blob back for a +// baseRT_sequence_gdn_restore after seeding the KV. The cache stores bytes; +// validity-at-position is the caller's contract. Attachments ride node +// eviction (a dropped prefix drops its blobs), live under their own byte +// budget (BASERT_PREFIX_BLOB_BUDGET_MB, default 512, oldest evicted first), +// and are NOT persisted by baseRT_prefix_cache_save. + +/// Attach `blob` at exactly the block-aligned depth `n_tokens` of a prefix +/// already published with baseRT_prefix_insert (insert first, then attach). +/// Replaces an existing attachment at that depth. Best-effort: a disabled +/// cache, an uncovered/misaligned path, or a blob over budget is a no-op. +/// Returns BASERT_OK (attached or clean no-op) or an argument error. +int baseRT_prefix_attach_blob(baseRT_model_t model, const uint32_t *tokens, int n_tokens, const void *blob, + int blob_len); + +/// The deepest attached-blob depth (in tokens) along the exact-match path of +/// `tokens`, capped at `n_tokens`. 0 = none / cache disabled. Read-only (no +/// lock, no LRU update): use it to cap a hybrid match at a boundary a seed +/// can actually restore. +int baseRT_prefix_blob_match(baseRT_model_t model, const uint32_t *tokens, int n_tokens); + +/// Copy the blob attached at exactly depth `n_tokens` into `out` (capacity +/// `cap`). Returns the blob's byte length, or -1 when no blob is attached at +/// that exact depth / the buffer is too small / the cache is disabled. +int baseRT_prefix_blob_fetch(baseRT_model_t model, const uint32_t *tokens, int n_tokens, void *out, int cap); + // === Grammar-constrained decoding === /// Opaque grammar handle. @@ -793,6 +1239,19 @@ baseRT_grammar_t baseRT_grammar_create(baseRT_model_t model, const char *gbnf); /// Returns NULL on error. baseRT_grammar_t baseRT_grammar_create_from_schema(baseRT_model_t model, const char *json_schema); +/// Create a grammar from an xgrammar STRUCTURAL TAG (a JSON document, not a +/// grammar string). +/// +/// A structural tag constrains only the parts of the output that matter: free +/// text until a trigger appears, then the matching tag's schema until its end +/// tag, then free text again. This is what lets a model choose whether to call +/// a tool while guaranteeing any call it makes is well-formed and delimited — +/// which a whole-output schema grammar cannot express. +/// +/// Returns NULL on error (malformed tag, or a schema inside it that cannot be +/// compiled); the caller should decode unconstrained rather than fail. +baseRT_grammar_t baseRT_grammar_create_from_structural_tag(baseRT_model_t model, const char *tag_json); + /// Create a grammar for generic JSON output (any valid JSON object/array). baseRT_grammar_t baseRT_grammar_create_json(baseRT_model_t model); @@ -812,12 +1271,20 @@ void baseRT_grammar_reset(baseRT_grammar_t grammar); /// baseRT_grammar_bitmask_size : packed int32 words in the token bitmask /// (0 = not an xgrammar grammar; use the serial decode path instead). /// baseRT_grammar_fill_bitmask : fill `out_bitmask` (bitmask_size words) for -/// the CURRENT grammar state; a set bit = allowed token. 1 on success. +/// the CURRENT grammar state; a set bit = allowed token. Returns 1 when +/// the mask excludes some token and is to be applied, 0 when there is +/// nothing to apply: every token may follow (free text in a structural +/// tag, where the model's markers are text as they are unconstrained), +/// or the grammar has no matcher. Not an error indicator. /// baseRT_grammar_accept_token : advance the grammar by one token. 1 on ok. /// baseRT_grammar_is_terminated: 1 once the grammar reaches an end state. /// baseRT_grammar_is_completed : 1 once a full match is accepted (a /// structured value is complete). Decoding should stop on terminated OR -/// completed — matching the serial grammar loop. +/// completed — matching the serial grammar loop. A STRUCTURAL-TAG +/// grammar never reports completed: free text is admitted around (and, +/// unless at_least_one, instead of) its tagged regions, so "could end +/// here" holds at every free-text step and is no reason to stop — those +/// lanes end on EOS (or terminated). int baseRT_grammar_bitmask_size(baseRT_grammar_t grammar); int baseRT_grammar_fill_bitmask(baseRT_grammar_t grammar, int32_t *out_bitmask); int baseRT_grammar_accept_token(baseRT_grammar_t grammar, uint32_t token_id); @@ -836,6 +1303,13 @@ BaseRTGenerationStats baseRT_generate_grammar_continue(baseRT_model_t model, con baseRT_grammar_t grammar, baseRT_token_callback callback, void *user_data); +/// `baseRT_generate_resume` with a grammar constraint (see there): prefill +/// `[n_cached, n_prompt)` over the kept KV, seed history from the whole prompt. +BaseRTGenerationStats baseRT_generate_grammar_resume(baseRT_model_t model, const uint32_t *prompt_tokens, int n_prompt, + int n_cached, int max_tokens, BaseRTSamplingConfig sampling, + baseRT_grammar_t grammar, baseRT_token_callback callback, + void *user_data); + // === GPU sampling === /// Run a profiled decode step — returns per-layer GPU timing. @@ -844,6 +1318,13 @@ BaseRTGenerationStats baseRT_generate_grammar_continue(baseRT_model_t model, con /// timing_out: array of (n_layers + 3) floats [embedding, norm, layer0..N-1, logit, argmax] /// Returns number of timing entries written, or -1 on a failed step /// (details via baseRT_get_error). +/// Per-kernel profile of one baked BATCHED decode round (see the .cpp +/// comment). Requires a prior baked batched round at this B. Labels via +/// baseRT_profile_batched_label(model, n_seqs, index). +int baseRT_profile_batched_decode(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, const uint32_t *in_tokens, + float *timing_out, int max_entries); +const char *baseRT_profile_batched_label(baseRT_model_t model, int n_seqs, int index); + int baseRT_profile_decode_step(baseRT_model_t model, uint32_t token_id, int position, float *timing_out, int max_entries); @@ -872,6 +1353,11 @@ uint32_t baseRT_tensor_dtype(baseRT_model_t model, int index); /// Returns empty string out of range. const char *baseRT_tensor_raw_dtype(baseRT_model_t model, int index); +/// Byte size of the tensor's storage blob at index (weights only — MLX +/// companion `.scales`/`.biases` tensors are separate entries with their +/// own sizes). Returns 0 out of range. +size_t baseRT_tensor_nbytes(baseRT_model_t model, int index); + /// Whether the loaded model carries an mmproj sub-bundle (vision/audio /// tower weights). Returns 0/1. int baseRT_has_mmproj(baseRT_model_t model); @@ -1014,12 +1500,18 @@ int baseRT_load_state(baseRT_model_t model, const char *path); /// every forward pass that runs a GEMM with a tensor_name registered in /// the adapter has a post-GEMM low-rank delta applied (`y += B @ A @ x`). /// -/// While an adapter is active, single-sequence decode skips the baked -/// dispatch-table replay (the table carries no delta dispatches) and takes -/// the per-token immediate encode instead — correct output at reduced -/// decode throughput. Speculative decode is likewise disabled for the -/// duration. The batched multi-sequence API (`baseRT_batch_step*`) does -/// NOT apply adapters. +/// While an adapter is active, decode skips the baked dispatch-table replay +/// (the table carries no delta dispatches, and one recorded WITH them would go +/// stale on unload) and takes the immediate encode instead — correct output at +/// reduced throughput. Speculative decode is likewise disabled for the +/// duration. +/// +/// The batched multi-sequence entry points DO apply the adapter: +/// `baseRT_batch_decode_step`, `baseRT_batch_step` and the fused/logits +/// variants each install it for the pass, bypass baked replay and command +/// bucketing while it is active, and emit the delta GEMMs per layer. (This +/// paragraph previously said the opposite, which would lead a caller to +/// disable batching or write a serial fallback it does not need.) /// /// Calling `baseRT_lora_load` again replaces the active adapter (no /// stacking). Returns 0 on success, negative on failure (see @@ -1120,6 +1612,20 @@ BaseRTGenerationStats baseRT_generate_continue(baseRT_model_t model, const uint3 int max_tokens, BaseRTSamplingConfig sampling, baseRT_token_callback callback, void *user_data); +/// Resume generation over a KV the caller kept: the KV holds exactly the first +/// `n_cached` tokens of `prompt_tokens` (e.g. after baseRT_try_rollback to a +/// shared prefix), and only `[n_cached, n_prompt)` is prefilled. Unlike +/// `baseRT_generate_continue`, which sees only the tokens it is given, the +/// channel decoder and the sampling history (repetition / presence / +/// frequency penalties) are seeded from the WHOLE prompt, so the run samples +/// as a cold `baseRT_generate(prompt_tokens, n_prompt)` would. When the KV +/// length is not `n_cached`, or `n_cached` leaves no token to prefill, it runs +/// the prompt cold instead. `prompt_tokens` in the returned stats counts the +/// whole prompt. +BaseRTGenerationStats baseRT_generate_resume(baseRT_model_t model, const uint32_t *prompt_tokens, int n_prompt, + int n_cached, int max_tokens, BaseRTSamplingConfig sampling, + baseRT_token_callback callback, void *user_data); + // === Embeddings === /// Compute text embeddings from token IDs using the model's hidden states. @@ -1134,6 +1640,11 @@ int baseRT_embed_text(baseRT_model_t model, const char *text, float *out_embeddi /// Get the embedding dimension for a model. int baseRT_embedding_dim(baseRT_model_t model); +/// True for a BERT-style encoder-only embedding model (no autoregressive +/// decode, no paged KV) — the serving worker uses this to allow loading such a +/// model despite it having no paged decode KV. +bool baseRT_is_embedding_model(baseRT_model_t model); + // === Chat templates === /// Format a chat prompt using the model's native template. @@ -1151,6 +1662,18 @@ const char *baseRT_chat_template(baseRT_model_t model); /// the model is freed. const char *baseRT_chat_template_jinja(baseRT_model_t model); +/// Number of added/special tokens the model declares (chat scaffold, +/// reasoning/tool channel markers, etc.). Pairs with +/// `baseRT_special_token` so a codec can learn a model's markers from the +/// model itself rather than assuming a hardcoded dialect. +int baseRT_special_token_count(baseRT_model_t model); + +/// The `index`-th special token's string; writes its token id (or +/// UINT32_MAX if the string is not a single vocab entry) to `id_out`. +/// Returns "" for an out-of-range index. Returned pointer valid until the +/// next call on the same thread. +const char *baseRT_special_token(baseRT_model_t model, int index, uint32_t *id_out); + /// BOS / EOS token strings (what minja substitutes for `{{ bos_token }}` /// and `{{ eos_token }}` in HF chat templates). const char *baseRT_bos_token(baseRT_model_t model); @@ -1171,6 +1694,21 @@ const char *baseRT_eos_token(baseRT_model_t model); /// other token-id consumers stop on). Returns 0 on a null handle. uint32_t baseRT_eos_token_id(baseRT_model_t model); +/// Content-derived weights/bundle identity: FNV-1a-64 over the model +/// file's first MiB, stamped at load (basertd §4 materialization +/// identity input — two fine-tunes of one architecture never share a +/// version_tag). 0 on a null handle or unidentified weights. +uint64_t baseRT_weights_identity(baseRT_model_t model); + +/// Whether `token` is ANY of the model's stop tokens — the primary +/// `eos_token_id` plus the extra ids the loader registers (`<|eot|>`, +/// `<|eot_id|>`, ``, gpt-oss's `<|call|>` and `<|endoftext|>`, +/// …). Multi-EOS models (Muse Glimmer: +/// `<|eot|>` ends the turn while `eos_token_id` only ends the document) +/// need this; comparing against `baseRT_eos_token_id` alone never +/// terminates their chat turns. Returns 0 on a null handle. +int32_t baseRT_is_eos_token(baseRT_model_t model, uint32_t token); + // (baseRT_max_prefill_chunk is declared once, with the batched-decode API.) // === Token counting === diff --git a/include/baseRT/baseRT_tick.h b/include/baseRT/baseRT_tick.h new file mode 100644 index 0000000..9a0e73b --- /dev/null +++ b/include/baseRT/baseRT_tick.h @@ -0,0 +1,973 @@ +// baseRT_tick.h — the basertd engine ABI additions (basertd §10 items 1–3, 8). +// +// Normative reference: docs/design/basertd-tick-schema.md (draft 7). +// This header is the in-process ABI between a Rust engine-worker and the +// linked libbaseRT: raw pointers and C strings are legal here; nothing in +// this header crosses a process boundary (the IPC frame format is a +// separate, serialized schema — schema doc §4). +// +// Layout rules (normative): +// - Every struct is 8-byte aligned with EXPLICIT padding; no field relies +// on implicit compiler padding. Layouts are identical across the +// compilers we target; basertd-abi's layout probe test enforces parity +// with the Rust #[repr(C)] mirrors. +// - Top-level structs carry `struct_size` and grow append-only. +// - Every record array is a BaseRTArray; decoding follows the four +// normative rules in schema §1.1 (zero-init + min(elem_size, local) +// copy; writer zero-pads; elem_size < min-prefix rejects the plan; +// 8-byte alignment and overflow/bounds checks before any element read). + +#ifndef BASERT_TICK_H +#define BASERT_TICK_H + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +// From baseRT.h; duplicated guard so this header is self-sufficient. +typedef void *baseRT_model_t; +typedef struct baseRT_sequence_s *baseRT_sequence_t; + +// --------------------------------------------------------------------------- +// Status codes +// --------------------------------------------------------------------------- + +typedef int32_t baseRT_status; + +enum { + BASERT_STATUS_OK = 0, + + // --- plan rejection class: tick refused, no state change (§1.5) --- + BASERT_REJECT_BAD_STRIDE = -100, // elem_size < record min-prefix + BASERT_REJECT_BAD_ALIGNMENT = -101, // array data not 8-byte aligned + BASERT_REJECT_BOUNDS = -102, // count*elem_size overflow / arena OOB + BASERT_REJECT_BAD_REF_GENERATION = -103, // stale ring/buffer generation + BASERT_REJECT_UNKNOWN_LANE = -104, + BASERT_REJECT_DUPLICATE_LANE = -105, // lane_tag twice in one array + BASERT_REJECT_ILLEGAL_COMBINATION = -106, // admit+retire, commit on non-HOST, ... + BASERT_REJECT_BUDGET = -107, // budget exceeds engine capacity + BASERT_REJECT_STALE_SEED = -108, // seed handle expired/consumed + BASERT_REJECT_HOST_RULES = -110, // HOST lane: max_new_tokens != 1, + // strategy without HOST_COMPATIBLE, + // decode with uncommitted logits + BASERT_REJECT_TRANSFER_LOCKED = -111, // sequence has in-flight state op + BASERT_REJECT_STALE_NONCE = -112, // commit's logits_nonce != emitted row + BASERT_NEEDS_REPLAN = -113, // shed envelope cannot serve the plan + BASERT_REJECT_BAD_STRUCT = -114, // struct_size below v1 minimum + + // --- synchronous op errors --- + BASERT_EBUSY = -120, // (sequence, space) already claimed + BASERT_ERR_BUFFER_TOO_SMALL = -121, + BASERT_ERR_STALE_SIZING = -122, // sizing_gen no longer current + BASERT_ERR_OUT_OF_BOUNDARY = -123, // trim off a snapshot boundary + BASERT_ERR_ENVELOPE_MISMATCH = -124, // restore/promote envelope vs target + BASERT_ERR_CHECKSUM = -125, + BASERT_ERR_IDENTITY_MISMATCH = -126, // compatibility identity mismatch + // Named OP_UNSUPPORTED (not ERR_UNSUPPORTED) because baseRT.h's + // BaseRTErrorCode already claims that identifier; the two headers + // compile together inside the engine. Wire value is unchanged. + BASERT_ERR_OP_UNSUPPORTED = -127, // per-kind rule (e.g. lossy blob demote) + BASERT_ERR_REGISTRATION_REFUSED = -128, // fail-closed strategy registration + BASERT_ERR_UNKNOWN_HANDLE = -129, // unknown op/seed/sequence handle + BASERT_ERR_SEED_UNSERVABLE = -130, // acquire: a space cannot serve prefix + BASERT_ERR_RING_LAYOUT = -131, // rings_attach: header/size mismatch + + // --- admission-result codes (schema §2.2): per-lane outcomes reported + // in admit_results while the tick executes — deliberately OUTSIDE + // the plan-rejection range, because a lane whose configuration + // matches no certificate is REJECTED in that tick's admit_results, + // never a whole-plan refusal --- + BASERT_REJECT_CERT_UNMATCHED = -140, // no certificate at/above floor + + // --- fatal --- + BASERT_ERR_FATAL = -1, // process presumed dead after this +}; + +// Per-lane fault codes (BaseRTLaneFault.code — schema §1.5: "lane stops +// with typed code; tick continues"). Distinct vocabulary from +// baseRT_status: these are recoverable, per-lane, and never imply plan +// rejection or state rollback. +enum { + BASERT_FAULT_UNSPECIFIED = 0, + BASERT_FAULT_GRAMMAR_OVERFLOW = 1, // grammar stack overflow + BASERT_FAULT_REF_SPAN_OOB = 2, // ref span out of bounds mid-lane + BASERT_FAULT_STRATEGY = 3, // speculation-strategy fault + BASERT_FAULT_OOM_TENTATIVE = 4, // tentative-namespace OOM (rollback taken) + BASERT_FAULT_INTERNAL = 5, // engine-internal, lane-recoverable + BASERT_FAULT_MEDIA = 6, // media bind/placeholder mismatch, split run +}; + +// --------------------------------------------------------------------------- +// Common value types +// --------------------------------------------------------------------------- + +// Stride-aware record array (schema §1.1). `elem_size` is the WRITER's +// sizeof(record); `data` must be 8-byte aligned. +typedef struct BaseRTArray { + const void *data; + uint32_t count; + uint32_t elem_size; +} BaseRTArray; // 16 bytes + +// Span of token ids in a shm ring: (ring_id, index, generation) triple plus +// length. Stale generations are detectable, never dereferenceable. +typedef struct BaseRTTokenRef { + uint32_t ring_id; + uint32_t index; // first slot in the ring + uint32_t count; // number of token ids + uint32_t _pad0; + uint64_t generation; +} BaseRTTokenRef; // 24 bytes + +// One row in a logits shm ring. The 64-bit generation doubles as the +// row's logits_nonce (schema §1.3: one identity, no separate field). +typedef struct BaseRTRingRef { + uint32_t ring_id; + uint32_t index; + uint64_t generation; +} BaseRTRingRef; // 16 bytes + +// Caller-owned host buffer. Caller allocates, caller frees, always; loaned +// to the engine from op submission to terminal state (schema §3.2). +typedef struct BaseRTBuf { + void *ptr; + uint64_t len; +} BaseRTBuf; // 16 bytes + +typedef struct BaseRTTokenRange { // [start, end) + uint64_t start; + uint64_t end; +} BaseRTTokenRange; // 16 bytes + +// --------------------------------------------------------------------------- +// Enums (carried as fixed-width integers in records) +// --------------------------------------------------------------------------- + +enum { // BaseRTLaneAdmit.sampling + BASERT_SAMPLING_GPU_GREEDY = 0, + BASERT_SAMPLING_GPU_GUMBEL = 1, + BASERT_SAMPLING_HOST = 2, +}; + +enum { // determinism_class + BASERT_DETERMINISM_BEST_EFFORT = 0, + BASERT_DETERMINISM_DETERMINISTIC = 1, +}; + +// Exactness classes, ordered so `minimum_exactness` floors compare +// numerically (basertd §4: approximate < distribution-exact < path-invariant). +enum { + BASERT_EXACTNESS_APPROXIMATE = 0, + BASERT_EXACTNESS_DISTRIBUTION_EXACT = 1, + BASERT_EXACTNESS_SEED_PATH_INVARIANT = 2, +}; + +enum { // BaseRTTickPlan.flags + BASERT_TICK_DRAIN = 1u << 0, + BASERT_TICK_CAPTURE_OK = 1u << 1, + // The plan carries a live `on_partial_emit` callback (appended fields — + // struct_size-gated): the tick invokes it as lanes produce tokens, so + // the host can stream them out MID-tick instead of waiting for the + // tick's terminal emit. The terminal emit stays authoritative and + // unchanged; partial emissions are display-path deltas. + BASERT_TICK_PARTIAL_EMITS = 1u << 2, +}; + +enum { // BaseRTTickEvents.tick_status + BASERT_TICK_OK = 0, + BASERT_TICK_SHED = 1, + BASERT_TICK_LANE_ERRORS = 2, + BASERT_TICK_FATAL = 3, +}; + +enum { // BaseRTAdmitResult.status + BASERT_ADMIT_UNSPECIFIED = 0, // absent (older writer) + BASERT_ADMIT_ADMITTED = 1, + BASERT_ADMIT_REJECTED = 2, +}; + +enum { // BaseRTLaneEmit.finish + BASERT_FINISH_NONE = 0, + BASERT_FINISH_EOS = 1, + BASERT_FINISH_LENGTH = 2, + BASERT_FINISH_GRAMMAR = 3, + BASERT_FINISH_CANCELLED = 4, + BASERT_FINISH_ERROR = 5, +}; + +enum { // BaseRTOpStatus.state + BASERT_OP_PENDING = 0, + BASERT_OP_RUNNING = 1, + BASERT_OP_DONE = 2, + BASERT_OP_FAILED = 3, + BASERT_OP_CANCELLED = 4, +}; + +enum { // demote/snapshot encoding + BASERT_ENCODING_LOSSLESS = 0, + BASERT_ENCODING_Q8 = 1, + BASERT_ENCODING_Q4 = 2, +}; + +enum { // state-space kinds (§7.2) + BASERT_SPACE_KIND_UNSPECIFIED = 0, + BASERT_SPACE_KIND_PAGED_TOKEN_KV = 1, + BASERT_SPACE_KIND_RING_KV = 2, + BASERT_SPACE_KIND_RECURRENT_BLOB = 3, + BASERT_SPACE_KIND_DEPTH_PAGED_KV = 4, + BASERT_SPACE_KIND_ENCODER_CACHE = 5, +}; + +enum { // BaseRTStateSpaceDesc.flags + // The space exists in the model (its state is real and evolves under + // ticks) but NO state op serves it yet: seeds cannot pin it, snapshot/ + // demote/restore/trim refuse BASERT_ERR_OP_UNSUPPORTED. A seed or park + // covering only the OTHER spaces would rebuild PARTIAL state, so the + // engine refuses seed acquisition on such bundles and the daemon must + // neither park nor resume them (basertd phase 2, P2.0). Cleared per + // space as its ops land (P2.5 for RecurrentBlob). + BASERT_SPACE_FLAG_OPS_UNAVAILABLE = 1u << 0, + // The space's state is NOT served by the reuse cache: a cache-based + // seed (baseRT_seed_acquire over published prefixes) can never cover + // it, so a bundle carrying such a space admits warm ONLY through + // adopted sequences (baseRT_seed_adopt over a restored sequence) — the + // hybrid GDN blob in v1 (no boundary-snapshot cache yet; basertd + // phase 2, P2.5). Cache seeds refuse SEED_UNSERVABLE on such bundles. + BASERT_SPACE_FLAG_NO_PREFIX_CACHE = 1u << 1, +}; + +enum { // BaseRTArtifact.role + BASERT_ARTIFACT_UNSPECIFIED = 0, + BASERT_ARTIFACT_DRAFT_WEIGHTS = 1, + BASERT_ARTIFACT_AUX_HEAD = 2, + BASERT_ARTIFACT_PROJECTION = 3, + BASERT_ARTIFACT_DRAFT_TOKENIZER = 4, +}; + +enum { // BaseRTTapSpec.tensor + BASERT_TAP_TENSOR_UNSPECIFIED = 0, + BASERT_TAP_TENSOR_HIDDEN = 1, + BASERT_TAP_TENSOR_ATTN_OUT = 2, + BASERT_TAP_TENSOR_EMBED = 3, +}; + +// Certificate applicability: sampling-mode bitset (schema §2.2). +enum { + BASERT_CERT_MODE_GREEDY = 1u << 0, + BASERT_CERT_MODE_GUMBEL = 1u << 1, + BASERT_CERT_MODE_HOST = 1u << 2, +}; + +// Certificate parameter-domain. 0 matches nothing (fail-closed default). +enum { + BASERT_CERT_DOMAIN_NONE = 0, + BASERT_CERT_DOMAIN_GREEDY_ONLY = 1, // temperature == 0 only + BASERT_CERT_DOMAIN_PENALTY_FREE = 2, // any temp, no penalties + BASERT_CERT_DOMAIN_ANY = 3, +}; + +enum { // strategy capability kinds: + BASERT_STRATEGY_CAP_UNSPECIFIED = 0, // an engine that does not + BASERT_STRATEGY_CAP_PROPOSAL_LINEAR = 1, // recognize a kind_id REFUSES + BASERT_STRATEGY_CAP_PROPOSAL_BLOCK = 2, // registration (fail-closed) + BASERT_STRATEGY_CAP_PROPOSAL_TREE = 3, + BASERT_STRATEGY_CAP_MASK_ANCESTOR = 4, // tree ancestor-only masks + BASERT_STRATEGY_CAP_HOST_COMPATIBLE = 5, // verification consumes no sampler +}; + +enum { // provenance taint reason bits + BASERT_TAINT_QUANTIZED_DEMOTION = 1u << 0, + BASERT_TAINT_DECOMPOSITION = 1u << 1, // decomposition-changing restore +}; + +enum { // resident tier (match results) + BASERT_TIER_GPU = 0, + BASERT_TIER_HOST = 1, + BASERT_TIER_GONE = 2, // metadata only; re-prefill +}; + +enum { // BaseRTShedEntry.kind + BASERT_SHED_UNSPECIFIED = 0, + BASERT_SHED_PREFILL_CHUNK = 1, + BASERT_SHED_CACHE_EVICTION = 2, +}; + +enum { // fork flags + BASERT_FORK_EAGER = 1u << 0, +}; + +enum { // BaseRTStateEnvelope.checksum_kind + // 0 is reserved-invalid: payload checksum validation on restore and + // promote is UNCONDITIONAL (schema §3.2 — corruption is detected, not + // trusted away); an envelope without a recognized checksum kind is a + // typed refusal, never a skipped check. + BASERT_CHECKSUM_XXH3_64 = 1, + // FNV-1a-64 over the payload (the canonical portable checksum the + // native engine writes; see the identity spec in basertd_tick.cpp). + BASERT_CHECKSUM_FNV1A64 = 2, +}; + +// --------------------------------------------------------------------------- +// §1.2 Tick plan +// --------------------------------------------------------------------------- + +typedef struct BaseRTSamplingParams { + float temperature; + float top_p; + float min_p; + uint32_t top_k; // 0 = disabled + float freq_penalty; + float presence_penalty; + float repeat_penalty; // 0 = disabled (means 1.0) + uint32_t flags; // BASERT_SAMPLING_FLAG_* (was _pad0; 0 = none) +} BaseRTSamplingParams; // 32 bytes + +enum { // BaseRTSamplingParams.flags + // The lane does not finish on an end-of-sequence token: it decodes to + // its max_new_tokens like any other token. A benchmark's fixed output + // length (vLLM's `ignore_eos`); unlike a logit bias on the stop set it + // leaves the lane on the GPU sampling and speculation paths, so what is + // measured is the path real traffic takes. + BASERT_SAMPLING_FLAG_IGNORE_EOS = 1u << 0, +}; + +// Shedding envelope: the engine may shed/evict ONLY within this envelope +// (schema §1.2). Shared verbatim by baseRT_cache_evict. +typedef struct BaseRTShedPolicy { + BaseRTArray victim_lanes; // of uint64_t lane_tag, eviction order + uint64_t evictable_cache_classes; // bitset over daemon-defined classes + uint64_t protected_quota_bytes; + uint64_t max_evict_bytes; +} BaseRTShedPolicy; // 40 bytes + +typedef struct BaseRTLaneAdmit { + uint64_t lane_tag; // opaque; never interpreted + BaseRTTokenRef prompt; + uint64_t seed_handle; // leased via baseRT_seed_acquire; 0 = cold + uint32_t sampling; // BASERT_SAMPLING_* + uint32_t logit_bias_handle; // 0 = none; created out-of-band (was _pad0) + BaseRTSamplingParams params; // used by GPU_GUMBEL; recorded regardless + uint64_t rng_counter_base; // counter-based RNG start (basertd §4) + uint32_t grammar_handle; // 0 = none; created out-of-band + uint32_t strategy_slot; // registered speculation slot; 0 = off + uint8_t determinism_class; // BASERT_DETERMINISM_* + uint8_t minimum_exactness; // floor for strategy certificates + uint8_t allow_approximate; // explicit consent; default 0 rejects + uint8_t want_logprobs; // 1 = emit each token's logprob (was _pad1) + uint32_t required_cert_id; // 0 = engine selects; else exact match + uint8_t host_sampler_identity[32]; // HOST lanes only; zero = cert-ineligible + // Appended (min-prefix stays 136, the v1 size): a preempted lane's + // continuation resumes its grammar where the generation stopped. The + // prompt's last `grammar_replay` tokens are what the grammar already + // accepted; the engine replays them into the fresh matcher before the + // lane's first draw (an admit whose replay the grammar refuses, or with + // no grammar, is rejected). 0 = a fresh generation. + uint32_t grammar_replay; + // The prompt's last `decode_replay` tokens were GENERATED by this very + // request before a preemption (the daemon's produced-so-far count). A + // resumed lane whose seed stops short of them does not prefill them as + // a slice: it feeds each as a one-token decode step whose output is + // discarded — the forward the uninterrupted lane ran at that position — + // so the continuation is schedule-exact (the §11 numeric floor is per + // execution shape; a prefill slice is not the shape that produced + // them). Prefill entries end where these begin; decode is legal from + // there. 0 = a fresh generation. + uint32_t decode_replay; +} BaseRTLaneAdmit; // 144 bytes + +// HOST-sampled feedback: EXACTLY ONE token — one logits row authorizes one +// autoregressive choice. No rng_counter field by design: the ENGINE advances +// the canonical position counter by exactly one per committed token. +typedef struct BaseRTLaneCommit { + uint64_t lane_tag; + uint32_t token_id; + uint32_t _pad0; + uint64_t logits_nonce; // = emitted row's ring-ref generation +} BaseRTLaneCommit; // 24 bytes + +typedef struct BaseRTLanePrefill { + uint64_t lane_tag; + uint32_t token_offset; // into the lane's admitted prompt + uint32_t token_count; +} BaseRTLanePrefill; // 16 bytes + +typedef struct BaseRTLaneDecode { + uint64_t lane_tag; + uint16_t max_new_tokens; // HOST lanes: MUST be 1 + // Tokens a speculative round MAY commit past max_new_tokens (0 = none): + // when max_new_tokens is the tick's decode grant rather than the end of + // the request, the scheduler sets this to the request's headroom, so the + // lane's last round drafts a full block instead of a truncated one (or a + // draft-less step). New rounds still start only while the grant lasts. + uint16_t overshoot; + uint32_t _pad1; +} BaseRTLaneDecode; // 16 bytes + +typedef struct BaseRTLaneRetire { + uint64_t lane_tag; + uint8_t publish_to_cache; + uint8_t _pad0[7]; +} BaseRTLaneRetire; // 16 bytes + +typedef struct BaseRTTickPlan { + uint64_t struct_size; // ABI guard, append-only growth + uint64_t plan_seq; // monotonic per bundle; echoed + uint32_t flags; // BASERT_TICK_* + uint32_t _pad0; + + uint32_t prefill_token_budget; // max prompt tokens this tick + uint32_t max_decode_lanes; // <= engine max_batch_size + + BaseRTArray admits; // of BaseRTLaneAdmit + BaseRTArray commits; // of BaseRTLaneCommit + BaseRTArray prefills; // of BaseRTLanePrefill + BaseRTArray decodes; // of BaseRTLaneDecode + BaseRTArray retires; // of BaseRTLaneRetire + + BaseRTShedPolicy shed_policy; + + // Appended (struct_size-gated, BASERT_TICK_PARTIAL_EMITS): invoked from + // WITHIN the tick, on the ticking thread, each time a lane commits + // freshly produced tokens — decode rounds and accepted speculative + // drafts alike. The pointer must stay valid for the duration of the + // call; the callback must not re-enter the engine. + void (*on_partial_emit)(void *user, uint64_t lane_tag, const uint32_t *tokens, uint32_t n_tokens); + void *partial_emit_user; +} BaseRTTickPlan; // 152 bytes legacy; 168 with the partial-emit tail + +// --------------------------------------------------------------------------- +// §1.3 Tick events +// --------------------------------------------------------------------------- + +typedef struct BaseRTAdmitResult { + uint64_t lane_tag; + uint32_t status; // BASERT_ADMIT_* + uint32_t reject_code; // baseRT_status value when REJECTED + uint32_t cert_id; // certificate in effect (0 = none) + uint8_t granted_class; // BASERT_EXACTNESS_* actually granted + uint8_t _pad0[3]; +} BaseRTAdmitResult; // 24 bytes + +typedef struct BaseRTSpecStats { + uint32_t proposed; + uint32_t accepted; + uint32_t cert_id_in_effect; // echo of the admission result + uint32_t _pad0; +} BaseRTSpecStats; // 16 bytes + +typedef struct BaseRTLaneEmit { + uint64_t lane_tag; + BaseRTTokenRef token_ref; // committed tokens (shm) + uint32_t n_tokens; + uint32_t finish; // BASERT_FINISH_* + BaseRTRingRef logits_row; // HOST lanes only; generation IS the + // logits_nonce (zeroed otherwise) + BaseRTSpecStats spec; +} BaseRTLaneEmit; // 72 bytes + +typedef struct BaseRTLaneFault { + uint64_t lane_tag; + uint32_t code; // typed, recoverable per-lane error + uint32_t _pad0; + uint64_t detail; +} BaseRTLaneFault; // 24 bytes + +typedef struct BaseRTShedEntry { + uint64_t lane_tag; // 0 for non-lane cache evictions + uint32_t kind; // BASERT_SHED_* + uint32_t reason; + uint64_t bytes; + uint64_t tokens; +} BaseRTShedEntry; // 32 bytes + +typedef struct BaseRTShedReport { // never silent (schema §1.4 rule 2) + uint32_t prefill_chunks_dropped; + uint32_t evictions_performed; + uint64_t bytes_evicted; + BaseRTArray entries; // of BaseRTShedEntry, with reasons +} BaseRTShedReport; // 32 bytes + +typedef struct BaseRTOpComplete { + uint64_t op; // baseRT_op_t + uint8_t state; // BASERT_OP_* (terminal) + uint8_t _pad0[3]; + uint32_t error; + uint64_t bytes_moved; +} BaseRTOpComplete; // 24 bytes + +typedef struct BaseRTTickTimings { // coarse, always-on (basertd §11.2) + uint64_t wall_ns; + uint64_t prefill_ns; + uint64_t decode_ns; + uint32_t graph_hits; + uint32_t graph_misses; +} BaseRTTickTimings; // 32 bytes + +typedef struct BaseRTMemCounters { // the budget manager's ledger + uint64_t allocated_bytes; + uint64_t host_retained_bytes; + uint64_t pool_blocks_total; + uint64_t pool_blocks_used; + uint64_t pool_bytes_evictable; + uint64_t tentative_bytes; // speculation namespaces +} BaseRTMemCounters; // 48 bytes + +// OpenAI caps `top_logprobs` at 20; the per-token record carries that many +// alternatives inline (bounded, fixed stride — no variable-length side arrays). +#define BASERT_MAX_TOP_LOGPROBS 20 + +// One per-token logprob for a lane that admitted with `want_logprobs` (its admit +// byte holds 1 + top_logprobs count). Records appear in token-emission order per +// lane, so the daemon zips them with the lane's emitted tokens +// (`BaseRTLaneEmit.token_ref`) in order. `logprob` is the natural-log +// probability of the emitted token under the row it was sampled from (after any +// grammar mask / logit_bias). `top_ids[0..n_top]` / `top_logprobs[0..n_top]` are +// the highest-probability alternatives at that position, descending. +typedef struct BaseRTLaneLogprob { + uint64_t lane_tag; + float logprob; + uint32_t n_top; // number of valid entries in top_ids/top_logprobs (0..20) + uint32_t top_ids[BASERT_MAX_TOP_LOGPROBS]; + float top_logprobs[BASERT_MAX_TOP_LOGPROBS]; +} BaseRTLaneLogprob; // 176 bytes + +typedef struct BaseRTTickEvents { + uint64_t struct_size; + uint64_t plan_seq; // echo + uint32_t tick_status; // BASERT_TICK_* + uint32_t _pad0; + + BaseRTArray admit_results; // of BaseRTAdmitResult + BaseRTArray emits; // of BaseRTLaneEmit + BaseRTArray faults; // of BaseRTLaneFault + + BaseRTShedReport shed; + BaseRTArray op_completions; // of BaseRTOpComplete + + BaseRTTickTimings timings; + BaseRTMemCounters mem; + BaseRTArray logprobs; // of BaseRTLaneLogprob (want_logprobs lanes only) +} BaseRTTickEvents; // 216 bytes + +// One call per tick per execution bundle. Blocking; returns when the tick's +// GPU work is complete and events are final. `*events_out` is an immutable +// view into an engine-owned double buffer, valid until the next baseRT_tick +// or baseRT_free_model call on the same bundle. +baseRT_status baseRT_tick(baseRT_model_t bundle, const BaseRTTickPlan *plan, const BaseRTTickEvents **events_out); + +// --------------------------------------------------------------------------- +// Ring attachment (basertd §10 item 5: caller-provided shm rings) +// --------------------------------------------------------------------------- +// +// The worker owns the shm mappings; the engine is handed base pointers. +// Ring memory layout (normative; basertd-shm's `SharedRing` is the +// reference implementation and the interop test): +// +// [64-byte ring header] magic u64 = BASERT_RING_MAGIC ("BRTRING1" LE) +// ring_id u32 | slot_bytes u32 | slots u32 +// (rest reserved, zero) +// [slots × slot_bytes] each slot: generation u64 | payload_len u32 | +// pad u32 | payload +// +// Exactly one writer per ring. The writer publishes payload bytes first +// and the generation last, with release fences between; readers validate +// the generation before AND after copying the payload (a recycled slot +// is detected, torn bytes are discarded). Generation 0 is the +// never-written sentinel and never matches. Writer cursors start at +// (index 0, generation 1) and generations are globally monotonic per +// ring — `BaseRTTokenRef`/`BaseRTRingRef` generations come from here. + +#define BASERT_RING_MAGIC 0x42525452494E4731ull /* "BRTRING1" LE */ + +enum { // BaseRTRingDesc.kind + BASERT_RING_KIND_UNSPECIFIED = 0, + BASERT_RING_KIND_TOKENS = 1, // payload = u32 token ids + BASERT_RING_KIND_LOGITS = 2, // payload = f32 logits row +}; + +enum { // BaseRTRingDesc.role — + BASERT_RING_ROLE_UNSPECIFIED = 0, // from the ENGINE's side + BASERT_RING_ROLE_ENGINE_READS = 1, // agent-staged prompts + BASERT_RING_ROLE_ENGINE_WRITES = 2, // emits / logits rows +}; + +typedef struct BaseRTRingDesc { + uint64_t struct_size; // ABI guard, append-only + uint32_t ring_id; + uint32_t kind; // BASERT_RING_KIND_* + uint32_t role; // BASERT_RING_ROLE_* + uint32_t slots; + uint32_t slot_bytes; // 16-byte slot header + payload + uint32_t _pad0; + void *base; // mapped ring memory (worker-owned + // mapping; ring header leads) + uint64_t len; // mapped bytes; >= 64 + slots*slot_bytes +} BaseRTRingDesc; // 48 bytes + +// Attach (or re-attach, replacing by ring_id) the bundle's rings. The +// engine validates each ring header against the descriptor, fail-closed +// with BASERT_ERR_RING_LAYOUT; on any failure NOTHING is attached (all +// or nothing, like plan validation). Mappings must outlive every +// subsequent call on the bundle until re-attach or baseRT_free_model — +// they are the worker's to unmap, never the engine's. Attaching resets +// the engine's writer cursors: attach exactly once per ring lifetime +// (fresh rings per connection, generation_base 1 — schema §4.2). +baseRT_status baseRT_rings_attach(baseRT_model_t bundle, const BaseRTRingDesc *descs, uint32_t count); + +// --------------------------------------------------------------------------- +// §10 item 2: state-space discovery +// --------------------------------------------------------------------------- + +typedef struct BaseRTStateSpaceDesc { + uint32_t space_id; + uint32_t kind; // BASERT_SPACE_KIND_* + uint64_t version_tag; // stale -> refuse restore -> re-prefill + uint64_t bytes_per_token; // paged kinds; 0 for blob kinds + uint64_t blob_bytes; // blob kinds; 0 for paged kinds + uint32_t page_size_tokens; // paged kinds; 0 otherwise + uint32_t fork_cost_class; // 0 unspecified; 1 refcount; 2 copy + uint64_t fork_cost_bytes; // what a lazy CoW copy will move + uint32_t snapshot_cadence; // 0 none; 1 tool-boundary; 2 interval + uint32_t snapshot_interval_tokens; // cadence == 2 only + uint32_t placement; // 0 local (Stage-1 PP field) + uint32_t flags; + char name[32]; // NUL-terminated debug name +} BaseRTStateSpaceDesc; // 96 bytes + +// Engine-owned descriptor list; valid until the next call on the same bundle. +baseRT_status baseRT_state_spaces(baseRT_model_t bundle, BaseRTArray *out_descs); + +// --------------------------------------------------------------------------- +// §2 Speculation strategy registration (at bundle load, before admission) +// --------------------------------------------------------------------------- + +typedef struct BaseRTArtifact { + uint32_t role; // BASERT_ARTIFACT_* + uint32_t _pad0; + uint8_t content_hash[32]; // claim; the open-handle hash is identity + uint64_t byte_size; + const char *load_path; +} BaseRTArtifact; // 56 bytes + +typedef struct BaseRTTapSpec { + uint32_t layer; + uint32_t tensor; // BASERT_TAP_TENSOR_* + uint32_t dtype; // engine dtype enum + uint32_t layout; +} BaseRTTapSpec; // 16 bytes + +// Extensible typed capability descriptor. Unknown kind_id -> registration +// refused, fail-closed (schema §2 scope note). +typedef struct BaseRTCapabilityReq { + uint32_t kind_id; // BASERT_STRATEGY_CAP_* or newer + uint32_t _pad0; + BaseRTArray params; // kind-defined records +} BaseRTCapabilityReq; // 24 bytes + +typedef struct BaseRTStrategyRegistration { + uint64_t struct_size; + const char *strategy_id; // "prompt-lookup", "draft-model", ... + const char *impl_version; // semver + uint8_t impl_hash[32]; // detects difference (not compat) + uint8_t config_hash[32]; // canonical config (basertd §4 identity) + + BaseRTArray artifacts; // of BaseRTArtifact + BaseRTArray taps; // of BaseRTTapSpec + BaseRTArray capabilities; // of BaseRTCapabilityReq + BaseRTArray target_archs; // of const char * + + uint32_t kernel_caps_required; // bitset over DeviceCaps-style bits + uint32_t _pad0; + uint64_t est_state_bytes; // drafter state, for the budget ledger + + uint8_t claimed_exactness; // verified, never trusted (§2.3) + uint8_t _pad1[3]; + uint32_t rng_contract_version; +} BaseRTStrategyRegistration; // 176 bytes + +typedef struct BaseRTExactnessCert { + uint32_t cert_id; // nonzero + uint8_t exactness; // BASERT_EXACTNESS_* + uint8_t sampling_modes; // bitset of BASERT_CERT_MODE_* + uint8_t grammar_allowed; + uint8_t logit_bias_allowed; + uint32_t param_domain; // BASERT_CERT_DOMAIN_* + uint32_t verification_shape_class; + BaseRTArray admissible_host_identities; // of uint8_t[32]; empty = never HOST +} BaseRTExactnessCert; // 32 bytes + +typedef struct BaseRTStrategyGrant { + uint64_t struct_size; + uint32_t strategy_slot; // referenced by BaseRTLaneAdmit + uint32_t _pad0; + BaseRTArray certificates; // of BaseRTExactnessCert; starts EMPTY + BaseRTArray state_spaces; // of BaseRTStateSpaceDesc (drafter state) + uint64_t reserved_bytes; // what the ledger was charged +} BaseRTStrategyGrant; // 56 bytes + +// Engine-owned grant; valid until the next registration or free on the +// bundle. Fails closed with a typed reason before the model admits work. +baseRT_status baseRT_strategy_register(baseRT_model_t bundle, const BaseRTStrategyRegistration *reg, + const BaseRTStrategyGrant **grant_out); + +// --------------------------------------------------------------------------- +// §3 Per-space state operations +// --------------------------------------------------------------------------- + +typedef uint64_t baseRT_op_t; // per-bundle, never reused + +typedef struct BaseRTOpStatus { + uint64_t struct_size; + uint8_t state; // BASERT_OP_* + uint8_t _pad0[3]; + uint32_t error; // typed, valid when FAILED + uint64_t bytes_moved; // progress, monotonic +} BaseRTOpStatus; // 24 bytes + +baseRT_status baseRT_op_poll(baseRT_model_t bundle, baseRT_op_t op, BaseRTOpStatus *out); +baseRT_status baseRT_op_cancel(baseRT_model_t bundle, baseRT_op_t op); // request only + +#define BASERT_ALL_SPACES 0xFFFFFFFFu + +typedef struct BaseRTMatchCandidate { + uint64_t prefix_len; // tokens + uint8_t provenance_digest[32]; // basertd §4 hash chain + uint32_t taint_bits; // BASERT_TAINT_* + uint8_t resident_tier; // BASERT_TIER_* + uint8_t _pad0[3]; +} BaseRTMatchCandidate; // 48 bytes + +typedef struct BaseRTSpaceMatch { + uint32_t space_id; + uint32_t _pad0; + BaseRTArray candidates; // of BaseRTMatchCandidate +} BaseRTSpaceMatch; // 24 bytes + +typedef struct BaseRTMatchResult { + uint64_t struct_size; + BaseRTArray spaces; // of BaseRTSpaceMatch +} BaseRTMatchResult; // 24 bytes + +// match: synchronous, metadata only, NO pinning. Advisory: a candidate can +// vanish between match and acquire. Result is engine-owned, next-call +// lifetime on the same bundle. +baseRT_status baseRT_space_match(baseRT_model_t bundle, + uint32_t space_id, // or BASERT_ALL_SPACES + BaseRTTokenRef span, + BaseRTArray media_deps, // of uint8_t[32] blob hashes + const BaseRTMatchResult **out); + +// seed acquire / release: the pinning step. Acquire re-validates the +// coordinator-chosen common prefix against every space and pins the +// underlying blocks/snapshots until the handle is consumed by exactly one +// BaseRTLaneAdmit or released. Unconsumed leases expire after the engine's +// declared TTL (invalidates the handle, not the data). +baseRT_status baseRT_seed_acquire(baseRT_model_t bundle, BaseRTTokenRef span, uint64_t prefix_len, + uint8_t determinism_class, // clean-variant filter + uint64_t *seed_handle); +baseRT_status baseRT_seed_release(baseRT_model_t bundle, uint64_t seed_handle); +// The declared seed-lease TTL, in ticks: a lease acquired at tick count T +// survives every tick that starts at or before T + ticks. A holder that +// keeps a lease across ticks (a pinned prefix) renews within it. +baseRT_status baseRT_seed_lease_ticks(baseRT_model_t bundle, uint64_t *ticks); + +// fork: synchronous, logical. Refcounts paged spaces; CoW-marks blob/ring +// spaces (copy at the child's first divergent step, or now under FORK_EAGER). +// EBUSY if any space of the parent has an in-flight op. +baseRT_status baseRT_seq_fork(baseRT_sequence_t parent, baseRT_sequence_t *child, uint32_t flags); + +// export sizing: MANDATORY before snapshot/demote. The result is valid for a +// specific content-generation; the export submission carries sizing_gen back +// and is rejected synchronously with BASERT_ERR_STALE_SIZING if the +// (sequence, space) advanced. Undersized dst rejects synchronously with +// BASERT_ERR_BUFFER_TOO_SMALL — no op handle, no partial writes, ever. +baseRT_status baseRT_space_export_size(baseRT_sequence_t seq, uint32_t space_id, BaseRTTokenRange range, + uint8_t encoding, uint64_t *required_bytes, uint64_t *sizing_gen); + +// Engine-authored, self-describing envelope wrapping EVERY exported payload +// (snapshot AND demote). restore/promote accept raw bytes ONLY through this +// envelope: the engine validates it against the TARGET (identity, kind, +// bounds) and the payload against the checksum before publishing. There is +// no caller-asserted "expected identity". Mismatch is a typed refusal. +#define BASERT_STATE_ENVELOPE_MAGIC 0x4554415453545242ull /* "BRTSTATE" LE */ + +typedef struct BaseRTStateEnvelope { + uint64_t magic; + uint32_t envelope_version; + uint32_t space_kind; // BASERT_SPACE_KIND_* + uint8_t compat_identity[32]; // basertd §4 identity 1, per space + uint64_t version_tag; + uint8_t provenance_digest[32]; // basertd §4 identity 2 + uint32_t taint_bits; // BASERT_TAINT_* — taint is sticky + // along the derivation chain + // (basertd §4): the digest travels + // WITH its reason bits, so a + // lossless export of tainted state + // can never restore as clean + uint32_t _pad0; + BaseRTTokenRange range; // blob kinds: {boundary, boundary} + uint8_t encoding; // BASERT_ENCODING_* + uint8_t _pad1[3]; + uint32_t checksum_kind; // BASERT_CHECKSUM_* + uint64_t payload_len; // bytes following this envelope + uint64_t content_checksum; +} BaseRTStateEnvelope; // 136 bytes + +// snapshot / restore: async, envelope-checked. dst/src are caller-owned but +// LOANED to the engine from submission to terminal state. +baseRT_status baseRT_space_snapshot(baseRT_sequence_t seq, uint32_t space_id, uint64_t boundary_pos, BaseRTBuf dst, + uint64_t sizing_gen, baseRT_op_t *op); +baseRT_status baseRT_space_restore(baseRT_sequence_t seq, uint32_t space_id, BaseRTBuf src, baseRT_op_t *op); + +// The newest DURABLE snapshot boundary <= cap for a space: the position at +// which an export equals state the prefix cache could have served, so an +// adopted (park) resume re-prefills exactly the gap a cache-warm resume +// re-prefills — byte-identical start state + identical feed schedule = +// bit-identical continuation. Blob spaces answer their cached capture +// position (0 = nothing durable yet — parking would not be resume-exact); +// the paged KV space answers the last whole-page length. Synchronous. +baseRT_status baseRT_space_snapshot_boundary(baseRT_sequence_t seq, uint32_t space_id, uint64_t cap, + uint64_t *boundary_out); + +// trim: synchronous metadata. Paged spaces: any length. Blob/ring spaces: +// snapshot boundaries only (BASERT_ERR_OUT_OF_BOUNDARY otherwise). +baseRT_status baseRT_space_trim(baseRT_sequence_t seq, uint32_t space_id, uint64_t new_len); + +// demote / promote: async tier movement. A lossy encoding sets the +// determinism-taint reason bit in the promoted-back state's provenance. +// RecurrentBlob demote is LOSSLESS-only (BASERT_ERR_UNSUPPORTED otherwise). +baseRT_status baseRT_space_demote(baseRT_sequence_t seq, uint32_t space_id, BaseRTTokenRange range, uint8_t encoding, + BaseRTBuf dst, uint64_t sizing_gen, baseRT_op_t *op); +baseRT_status baseRT_space_promote(baseRT_sequence_t seq, uint32_t space_id, BaseRTTokenRange range, BaseRTBuf src, + baseRT_op_t *op); + +// cache maintenance: policy-enveloped, like shedding. Pinned blocks (active +// ops, seed leases) are never eligible. Publication into the cache is the +// tick's retire.publish_to_cache, not here. +baseRT_status baseRT_cache_evict(baseRT_model_t bundle, const BaseRTShedPolicy *policy, uint64_t bytes_target, + uint64_t *bytes_freed); + +// The cache-entry key for enumerated eviction: the basertd §8 physical +// cache key — (compatibility key, provenance digest) — one record per +// physical variant to drop. +typedef struct BaseRTCacheKey { + uint8_t compat_key[32]; // per-space §8 chain key + uint8_t provenance_digest[32]; // variant selector +} BaseRTCacheKey; // 64 bytes + +baseRT_status baseRT_cache_evict_entries(baseRT_model_t bundle, + BaseRTArray entry_keys); // of BaseRTCacheKey + +// --------------------------------------------------------------------------- +// Tick-bundle sequence lifecycle (support surface for §3 state ops): +// restore targets need sequences that no lane owns yet, and callers +// driving state ops need the sequence a lane owns. In-process handles. +// --------------------------------------------------------------------------- + +// The sequence a live lane owns (NULL if the lane is unknown). +baseRT_sequence_t baseRT_tick_lane_sequence(baseRT_model_t bundle, uint64_t lane_tag); + +// Create a bare, empty sequence on the bundle's paged pool — a restore +// target. Freed with baseRT_tick_sequence_free (or bundle teardown). +baseRT_sequence_t baseRT_tick_sequence_create(baseRT_model_t bundle); +// Free a LIFECYCLE-OWNED bare sequence. Owned-only, fail-closed: +// unknown, cross-bundle, already-freed, or lane-owned handles are +// refused with BASERT_ERR_UNKNOWN_HANDLE (a lane's sequence is retired +// by its plan, never through this call). +baseRT_status baseRT_tick_sequence_free(baseRT_model_t bundle, baseRT_sequence_t seq); + +// Compile a JSON-Schema grammar (xgrammar backend) into the bundle's tick +// registry for grammar-constrained decoding (§6). Returns a nonzero handle +// a lane's admit references via `grammar_handle`, or 0 on failure (bad +// schema, or a legacy grammar with no host bitmask — the batched path +// requires xgrammar). Free with baseRT_tick_grammar_free at the end of the +// constrained generation. The matcher's acceptance state is per-generation, +// so a handle is used by one lane at a time. +uint32_t baseRT_tick_grammar_create(baseRT_model_t bundle, const char *json_schema); + +/// As `baseRT_tick_grammar_create`, but from an xgrammar STRUCTURAL TAG (see +/// `baseRT_grammar_create_from_structural_tag`): free text except inside the +/// tag's delimiters, where its schema holds. The handle is interchangeable +/// with a schema-created one — same registry, same host-bitmask path, same +/// free. 0 on failure, including a grammar with no host bitmask. +uint32_t baseRT_tick_grammar_create_structural(baseRT_model_t bundle, const char *tag_json); +// Free a grammar handle; BASERT_ERR_UNKNOWN_HANDLE if absent. +baseRT_status baseRT_tick_grammar_free(baseRT_model_t bundle, uint32_t handle); + +// Register a per-request logit_bias map (OpenAI `logit_bias`) into the bundle's +// tick registry: `tokens[i]` gets `+values[i]` added to its logit before the +// argmax/sample in `sample_row`. Returns a nonzero handle a lane's admit +// references via `logit_bias_handle`, or 0 on failure (n <= 0, or a null +// array). A lane carrying a bias handle is kept off the GPU-argmax fast path +// (the bias must reach the host logits row). Free at the end of the generation. +uint32_t baseRT_tick_logit_bias_create(baseRT_model_t bundle, const int32_t *tokens, const float *values, int n); +// Free a logit_bias handle; BASERT_ERR_UNKNOWN_HANDLE if absent. +baseRT_status baseRT_tick_logit_bias_free(baseRT_model_t bundle, uint32_t handle); + +// Adopt: lease a LIFECYCLE-OWNED bare sequence (a restore target whose +// state ops have run) as the seed of the next admission (basertd P2.5). +// `span` names the tokens the sequence's state covers (its length must +// equal the sequence's current length; the daemon validated the content +// against the artifact's stream digest — the engine records it as the +// lane's committed prefix). Unlike a cache seed the sequence is taken +// over WHOLE — every state space it holds (KV incl. a partial last block, +// recurrent blobs) — so it is the sound warm path for bundles whose +// spaces the cache cannot serve. The lease consumes the sequence at +// admission (the lane owns it from then on); an expired lease frees it. +baseRT_status baseRT_seed_adopt(baseRT_model_t bundle, baseRT_sequence_t seq, BaseRTTokenRef span, + uint64_t *seed_handle); + +// --------------------------------------------------------------------------- +// Media (basertd §4 media, §7.2 EncoderCache — phase 2 P2.7). Images enter +// a lane's prompt as runs of the model's image placeholder token; the +// tower's features for a blob are an engine-owned EncoderCache entry +// addressed by a media handle; a bind attaches handles to a lane BEFORE +// its placeholders prefill (by lane tag, admitted or not yet — the admit +// adopts pending binds; a retire drops the lane's remaining binds). +// --------------------------------------------------------------------------- + +typedef struct BaseRTMediaInfo { + uint32_t n_tokens; // placeholders this image expands to + uint32_t image_token_id; // the placeholder token + uint32_t boi_token_id; // begin-of-image wrapper (0 = none) + uint32_t eoi_token_id; // end-of-image wrapper (0 = none) + uint64_t preprocess_fp; // preprocessing-config fingerprint (basertd §4 + // media identity: patch/merge/pool geometry) +} BaseRTMediaInfo; // 24 bytes + +// Preprocess only (no tower): how many placeholders `image_path` expands +// to and the wrapper token ids. UNSUPPORTED on bundles without a tower. +baseRT_status baseRT_media_probe(baseRT_model_t bundle, const char *image_path, BaseRTMediaInfo *out); + +// Preprocess + tower: an EncoderCache entry (engine-owned features) for +// `image_path`. UNSUPPORTED on bundles whose tick path cannot splice +// media (hybrid-GDN: no per-lane M-RoPE in continuous batching yet; MoE). +baseRT_status baseRT_media_encode(baseRT_model_t bundle, const char *image_path, uint64_t *media_handle, + BaseRTMediaInfo *out); +baseRT_status baseRT_media_release(baseRT_model_t bundle, uint64_t media_handle); + +// Bind a media handle to a lane (by tag; before or after admit) at the +// prompt offset where its placeholder run starts. The lane's prefill +// chunk containing that run splices the features; a chunk that splits a +// run is a per-lane fault. Binds are consumed in offset order and +// released with the lane. +baseRT_status baseRT_media_bind(baseRT_model_t bundle, uint64_t lane_tag, uint64_t media_handle, uint32_t token_offset); + +// Publish a bare sequence's block-aligned KV prefix into the reuse +// cache under `tokens` — the resume bridge: restore into a fresh +// sequence, publish, then re-enter through the ordinary seeded +// admission path. Same whole-block contract as retire's +// publish_to_cache (sub-block tails publish nothing and still +// succeed). Consumes the sequence on BASERT_STATUS_OK; on error the +// sequence is untouched and still owned by the caller. +baseRT_status baseRT_tick_sequence_publish(baseRT_model_t bundle, baseRT_sequence_t seq, const uint32_t *tokens, + uint32_t n_tokens); + +#ifdef __cplusplus +} // extern "C" +#endif + +#endif // BASERT_TICK_H