From 73e577fc181080b1ec58853dae24c5dcc1ccb328 Mon Sep 17 00:00:00 2001 From: Prabod Rathnayaka Date: Fri, 31 Jul 2026 18:53:53 +1000 Subject: [PATCH] =?UTF-8?q?sync:=20engine=200.2.0=20=E2=80=94=20CUDA/GB10?= =?UTF-8?q?=20backend,=20backend-aware=20hub,=200.2.0=20bump?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Open-core source sync for the 0.2.0 release: the base-convert CLI/hub (backend-aware catalog resolution, catalog derive, format changes), the public C header + Swift/Node/Python bindings, the GB10 benchmark scripts, and the generic default-* profiles — all bumped to 0.2.0. Excludes the private moat: all tuned quant profiles (gemma4-moe-*, *-q4mix, *-q6mix, dense-q4mix-cuda, qwen3-moe-*, qwen35-moe-*) and tools/internal. Verified: cargo build --locked + clippy -D warnings + tests all pass; no include_str! of an excluded profile; leak check clean. --- base-convert/Cargo.lock | 16 +- base-convert/Cargo.toml | 2 +- base-convert/crates/base-arch/src/lib.rs | 33 +- base-convert/crates/base-convert/src/hub.rs | 73 +++- base-convert/crates/base-convert/src/main.rs | 4 + base-convert/crates/base-format/src/header.rs | 5 + base-convert/crates/base-format/src/writer.rs | 19 +- .../crates/base-format/tests/roundtrip.rs | 33 ++ base-convert/crates/base-hub/src/catalog.rs | 30 +- base-convert/crates/base-hub/src/gen.rs | 143 +++++++ base-convert/crates/base-hub/src/lib.rs | 1 + base-convert/crates/base-hub/src/registry.rs | 374 ++++++++++++++++-- base-convert/profiles/default-q4-embq6.json | 128 ++++++ base-convert/profiles/default-q8.json | 87 +++- benchmarks/scripts/fetch_3way_models.sh | 123 ++++++ .../scripts/gb10_basert_prefill_decode.sh | 63 +++ benchmarks/scripts/gb10_serving_vs_vllm.sh | 230 +++++++++++ benchmarks/scripts/gb10_two_way.sh | 22 ++ .../scripts/gb10_vllm_prefill_decode.sh | 101 +++++ benchmarks/scripts/reclaim_page_cache.sh | 27 ++ benchmarks/scripts/three_way_benchmark.sh | 82 +++- bindings/node/src/index.ts | 24 ++ bindings/python/baseRT/__init__.py | 2 +- bindings/python/setup.py | 2 +- .../swift/Sources/CBaseRT/include/baseRT.h | 235 ++++++++++- include/baseRT/baseRT.h | 235 ++++++++++- 26 files changed, 2016 insertions(+), 78 deletions(-) create mode 100644 base-convert/crates/base-hub/src/gen.rs create mode 100644 base-convert/profiles/default-q4-embq6.json create mode 100644 benchmarks/scripts/fetch_3way_models.sh create mode 100644 benchmarks/scripts/gb10_basert_prefill_decode.sh create mode 100644 benchmarks/scripts/gb10_serving_vs_vllm.sh create mode 100644 benchmarks/scripts/gb10_two_way.sh create mode 100644 benchmarks/scripts/gb10_vllm_prefill_decode.sh create mode 100644 benchmarks/scripts/reclaim_page_cache.sh diff --git a/base-convert/Cargo.lock b/base-convert/Cargo.lock index 1cc3fc9..03ffa9d 100644 --- a/base-convert/Cargo.lock +++ b/base-convert/Cargo.lock @@ -66,7 +66,7 @@ checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" [[package]] name = "base-arch" -version = "0.1.7" +version = "0.2.0" dependencies = [ "anyhow", "base-format", @@ -77,7 +77,7 @@ dependencies = [ [[package]] name = "base-awq" -version = "0.1.7" +version = "0.2.0" dependencies = [ "anyhow", "base-format", @@ -91,7 +91,7 @@ dependencies = [ [[package]] name = "base-convert" -version = "0.1.7" +version = "0.2.0" dependencies = [ "anyhow", "base-arch", @@ -114,7 +114,7 @@ dependencies = [ [[package]] name = "base-format" -version = "0.1.7" +version = "0.2.0" dependencies = [ "anyhow", "bitflags", @@ -130,7 +130,7 @@ dependencies = [ [[package]] name = "base-hub" -version = "0.1.7" +version = "0.2.0" dependencies = [ "anyhow", "base-format", @@ -145,7 +145,7 @@ dependencies = [ [[package]] name = "base-quant" -version = "0.1.7" +version = "0.2.0" dependencies = [ "anyhow", "base-format", @@ -157,7 +157,7 @@ dependencies = [ [[package]] name = "base-readers" -version = "0.1.7" +version = "0.2.0" dependencies = [ "anyhow", "base-format", @@ -171,7 +171,7 @@ dependencies = [ [[package]] name = "base-sign" -version = "0.1.7" +version = "0.2.0" dependencies = [ "anyhow", "base-format", diff --git a/base-convert/Cargo.toml b/base-convert/Cargo.toml index 9bee6df..f548a12 100644 --- a/base-convert/Cargo.toml +++ b/base-convert/Cargo.toml @@ -12,7 +12,7 @@ members = [ ] [workspace.package] -version = "0.1.7" +version = "0.2.0" edition = "2021" license = "Apache-2.0" repository = "https://github.com/basecompute/baseRT" diff --git a/base-convert/crates/base-arch/src/lib.rs b/base-convert/crates/base-arch/src/lib.rs index 9e6cf90..7b9ac68 100644 --- a/base-convert/crates/base-arch/src/lib.rs +++ b/base-convert/crates/base-arch/src/lib.rs @@ -81,15 +81,20 @@ pub fn hf_mapper_for_model_type(model_type: &str) -> Option<&'static dyn HfMappe // (canonical_arch="llama" → llama model class). Validated end-to-end on // Mistral-7B-Instruct-v0.3 and Ministral-8B-Instruct-2410. // - // Phi-3 is NOT enabled. The converter side is ready (SplittingProvider - // splits its fused qkv_proj/gate_up_proj; weights convert; HD=96 - // attention is fine — regression-tested) and the generation_config eos - // merge below makes it stop cleanly. BUT Phi-3 chat output degenerates - // (raw completions are coherent; chat floods/repeats and is incoherent at - // both Q4 and Q8, temp 0 and 0.7) — a chat-path issue (same class as - // SmolLM2) that's not yet root-caused. Re-enable "phi3" once that's - // fixed. (Phi-3.5 additionally needs LongRoPE — engine is linear-only.) - "llama" | "mistral" => Some(&llama::LlamaHfMapper), + // Phi-3-mini (4k, standard RoPE) is Llama-shaped and reuses the Llama + // mapper: the SplittingProvider (base-convert) slices its fused + // self_attn.qkv_proj / mlp.gate_up_proj into the canonical split names + // the mapper consumes, HD=96 is fine (regression-tested), and the + // generation_config eos merge makes it stop cleanly on <|end|> (32007). + // The chat-path flood that used to gate it (same class as SmolLM2) is + // resolved by two landed tokenizer fixes — the added-token lstrip/rstrip + // handling in split_special_tokens (Phi-3's `<|end|>` rstrip absorbs the + // trailing newline; tokenizer.cpp) and the GPT2* add_bos default + // (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), "qwen2" | "qwen3" => Some(&qwen::QwenHfMapper), "qwen2_moe" | "qwen3_moe" => Some(&qwen::QwenMoeHfMapper), // Qwen3.5 / 3.6: hybrid Gated-DeltaNet + full-attention decoder @@ -107,7 +112,13 @@ pub fn hf_mapper_for_model_type(model_type: &str) -> Option<&'static dyn HfMappe // actually google/gemma-3n-E2B-it. The canonical Gemma 4 lives // under google/gemma-4-{E2B,E4B}-it and uses model_type=gemma4 // with text_config.model_type=gemma4_text. - "gemma4" | "gemma4_text" => Some(&gemma::Gemma4HfMapper), + // gemma4_unified (gemma-4-12B-it): the encoder-free multimodal + // variant — a standard gemma4 text stack under + // `model.language_model.*` plus ~10 small modality-projection + // tensors (embed_vision/embed_audio/vision_embedder) that the + // text conversion skips like any other non-text tower. Config is + // gemma4-shaped (uniform head_dim, rope_parameters, layer_types). + "gemma4" | "gemma4_text" | "gemma4_unified" => Some(&gemma::Gemma4HfMapper), _ => None, } } @@ -118,6 +129,7 @@ pub fn hf_mapper_for_model_type(model_type: &str) -> Option<&'static dyn HfMappe pub const SUPPORTED_HF_MODEL_TYPES: &[&str] = &[ "llama", "mistral", + "phi3", "qwen2", "qwen3", "qwen2_moe", @@ -135,6 +147,7 @@ pub const SUPPORTED_HF_MODEL_TYPES: &[&str] = &[ "gemma3_text", "gemma4", "gemma4_text", + "gemma4_unified", ]; pub trait GgufMapper: Sync { diff --git a/base-convert/crates/base-convert/src/hub.rs b/base-convert/crates/base-convert/src/hub.rs index 7a22859..b9f078d 100644 --- a/base-convert/crates/base-convert/src/hub.rs +++ b/base-convert/crates/base-convert/src/hub.rs @@ -175,12 +175,83 @@ fn installed_single_path(reg: &MergedRegistry, id: &str) -> Result { + // Multiple variants installed — e.g. a universal `default-q4` cached + // before the backend-qualified catalog entries existed, plus a native + // `cuda-q4mix` pulled after. Rather than erroring, apply the same + // backend preference the catalog resolver uses (backend-native > + // universal) so bare `serve ` keeps Just Working. The installed + // variant dir name carries the backend slot (`cuda-*`/`rocm-*`/`cpu-*` + // vs the universal `default-*`); on a backend that can't load a + // foreign-native bundle it simply won't be present. + let backend = base_hub::registry::CatalogRegistry::client_backend(); + let native_prefix = format!("{backend}-"); + let native: Vec<&ModelEntry> = + many.iter().copied().filter(|r| r.variant.starts_with(&native_prefix)).collect(); + let pick = if native.len() == 1 { + Some(native[0]) + } else if native.is_empty() { + let uni: Vec<&ModelEntry> = + many.iter().copied().filter(|r| r.variant.starts_with("default-")).collect(); + (uni.len() == 1).then(|| uni[0]) + } else { + None + }; + if let Some(one) = pick { + return Ok(Some( + one.path + .clone() + .with_context(|| format!("installed model `{id}` has no artifact path"))?, + )); + } let variants = many.iter().map(|r| r.variant.as_str()).collect::>().join(", "); bail!("model `{id}` has multiple installed variants ({variants}) — specify one as `{id}:`") } } } +/// Installed artifact for `id` matching the requested variant `want`. Tries an +/// EXACT variant-dir match first (honors `:default-q4` / `:cuda-q4mix`); failing +/// that, matches by BIT WIDTH so the documented short `:q4` finds an installed +/// `cuda-q4mix` or `default-q4` (the exact dir is named for the quant slot, not +/// the bare bits), preferring the backend-native slot over universal — the same +/// rule the catalog resolver's quant_matches + backend preference apply. Returns +/// `None` when nothing installed matches (caller then auto-pulls). +fn installed_best_variant(reg: &MergedRegistry, id: &str, want: &str) -> Option { + if let Some(p) = reg.local.installed_path(id, want) { + return Some(p); // exact variant name present + } + let want_bits = base_hub::registry::quant_bits(want).unwrap_or(want); + let backend = base_hub::registry::CatalogRegistry::client_backend(); + let native_prefix = format!("{backend}-"); + // A foreign-native slot (e.g. `cuda-q4mix` on macOS) must be EXCLUDED, not + // just deprioritized: on a copied/shared cache it can't be loaded, and a bit- + // width tie would otherwise pick the lexicographically-earlier CUDA artifact + // over the runnable universal one. Runnable = universal (`default-*`), this + // client's native slot (`-*`), or a bare bit-width dir (no slot). + const KNOWN_BACKENDS: [&str; 4] = ["cuda", "rocm", "cpu", "metal"]; + let runnable = |v: &str| -> bool { + match v.split_once('-') { + Some((slot, _)) if KNOWN_BACKENDS.contains(&slot) => slot == backend, // native only + _ => true, // default-*, bare bits, etc. + } + }; + let installed = reg.local.list().ok()?; + let mut matches: Vec<&ModelEntry> = installed + .iter() + .filter(|r| { + r.id == id + && runnable(&r.variant) + && base_hub::registry::quant_bits(&r.variant).unwrap_or(r.variant.as_str()) == want_bits + }) + .collect(); + if matches.is_empty() { + return None; + } + // Backend-native slot first (stable so a deterministic pick when both exist). + matches.sort_by_key(|r| u8::from(!r.variant.starts_with(&native_prefix))); + matches[0].path.clone() +} + /// Fetch a not-yet-installed model on demand, then return its artifact path. /// Prefers the pre-converted basecompute mirror; otherwise converts the source /// repo on pull. Progress (download + quantization) is shown by `cmd_pull`. @@ -245,7 +316,7 @@ fn resolve_hub_model(token: &str, default_variant: Option<&str>) -> Result Option { || mm_name.starts_with("visual.") || mm_name.starts_with("embed_audio") || mm_name.starts_with("embed_vision") + // gemma4_unified (gemma-4-12B-it) names its encoder-free patch + // embedder `vision_embedder.*` (pos_embedding, patch_dense, + // patch_ln1/2, pos_norm) rather than `vision_tower.*`. + || mm_name.starts_with("vision_embedder") || mm_name.starts_with("multi_modal_projector") { let canonical = base_arch::gemma::map_gemma4_mmproj_name(mm_name).unwrap_or_else(|| mm_name.to_string()); diff --git a/base-convert/crates/base-format/src/header.rs b/base-convert/crates/base-format/src/header.rs index 4f51495..ab05f9f 100644 --- a/base-convert/crates/base-format/src/header.rs +++ b/base-convert/crates/base-format/src/header.rs @@ -183,6 +183,11 @@ pub enum TargetBackend { /// layouts. CudaSm89, CudaSm90, + /// NVIDIA GB10 (DGX Spark, consumer Blackwell). Serializes to + /// "cuda_sm121" — one of the two tags the C++ reader's CUDA build + /// accepts and Metal builds reject; used for bundles that carry + /// CUDA-only kernel requirements (base_q6 MoE experts today). + CudaSm121, /// AMD CDNA3 (MI300). MFMA tile layout. RocmCdna3, /// CPU AVX2 / NEON fallback paths. Row-major contiguous packing. diff --git a/base-convert/crates/base-format/src/writer.rs b/base-convert/crates/base-format/src/writer.rs index 59c5d41..6767383 100644 --- a/base-convert/crates/base-format/src/writer.rs +++ b/base-convert/crates/base-format/src/writer.rs @@ -1,4 +1,4 @@ -use crate::header::{Header, MmprojBundle, TensorEntry}; +use crate::header::{Header, MmprojBundle, TensorDtype, TensorEntry}; use crate::slots::{write_slots, Slot}; use crate::{Error, Result, BLOB_ALIGNMENT, FORMAT_VERSION, MAGIC, PREFIX_LEN}; use std::fs::File; @@ -127,6 +127,23 @@ impl BaseWriter { } self.header.tensors = entries; + // Stamp `target_backend` from CONTENT instead of trusting the + // caller's default (every construction site used to hardcode + // Metal, so CUDA-only bundles carried a `metal` tag and failed a + // kernel lookup only after a full download + load). The one + // CUDA-only content class today: base_q6 MoE expert slabs + // (`*_exps.*` tensors) — Metal ships no q6 MoE kernels. bf16-scale + // q8/q4 bundles stay `metal` (universal): both backends carry the + // `_sbf16` kernel families. + if self + .header + .tensors + .iter() + .any(|t| t.name.contains("_exps.") && t.dtype == TensorDtype::BaseQ6) + { + self.header.target_backend = crate::header::TargetBackend::CudaSm121; + } + // Multimodal sub-bundle entries land in the same weights blob, // continuing past the LM tensors. Their entries go into // `header.mmproj.tensors` so the runtime can decide whether to diff --git a/base-convert/crates/base-format/tests/roundtrip.rs b/base-convert/crates/base-format/tests/roundtrip.rs index 903811c..cbe07b2 100644 --- a/base-convert/crates/base-format/tests/roundtrip.rs +++ b/base-convert/crates/base-format/tests/roundtrip.rs @@ -456,3 +456,36 @@ fn rejects_unknown_version() { let err = expect_err(BaseReader::open(tmp.path())); assert!(matches!(err, base_format::Error::UnsupportedVersion(999, _))); } + +#[test] +fn target_backend_stamped_from_content() { + // q6 MoE expert slabs are CUDA-only (no Metal q6 MoE kernels): the + // writer must override the caller's Metal default so Metal builds + // reject the bundle at open, not at kernel lookup. + let tmp = tempfile::NamedTempFile::new().unwrap(); + let header = make_header(); + let mut writer = BaseWriter::create(tmp.path(), header).unwrap(); + let mut e = entry("layers.0.ffn_gate_exps.weight", vec![4, 8]); + e.dtype = TensorDtype::BaseQ6; + writer.add_tensor(TensorPayload { + entry: e, + data: vec![0u8; 24], + }); + writer.finish().unwrap(); + let reader = BaseReader::open(tmp.path()).unwrap(); + assert_eq!(reader.header().target_backend, TargetBackend::CudaSm121); + + // ...but a DENSE q6 tensor (or q6 outside the expert slabs) stays + // universal — Metal has dense q6 kernels. + let tmp2 = tempfile::NamedTempFile::new().unwrap(); + let mut writer = BaseWriter::create(tmp2.path(), make_header()).unwrap(); + let mut e = entry("layers.0.self_attn.q_proj.weight", vec![4, 8]); + e.dtype = TensorDtype::BaseQ6; + writer.add_tensor(TensorPayload { + entry: e, + data: vec![0u8; 24], + }); + writer.finish().unwrap(); + let reader = BaseReader::open(tmp2.path()).unwrap(); + assert_eq!(reader.header().target_backend, TargetBackend::Metal); +} diff --git a/base-convert/crates/base-hub/src/catalog.rs b/base-convert/crates/base-hub/src/catalog.rs index f670203..469c5b3 100644 --- a/base-convert/crates/base-hub/src/catalog.rs +++ b/base-convert/crates/base-hub/src/catalog.rs @@ -68,6 +68,14 @@ pub struct CatalogEntry { /// Optional integrity check for the downloaded `.base`. #[serde(default, skip_serializing_if = "Option::is_none")] pub sha256: Option, + /// Backend requirement. `None` = universal (runs on every backend — + /// f16/bf16 weights and bf16-scale q4/q8 bundles). `Some("cuda")` / + /// `Some("metal")` restricts resolution to clients of that backend + /// (e.g. base_q6-expert MoE bundles are CUDA-only). Checked at + /// resolve time so an incompatible pull refuses BEFORE the download, + /// not at load after 30 GB. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub backend: Option, } fn default_file() -> String { @@ -179,6 +187,15 @@ impl Catalog { .or_else(|| self.models.iter().find(|e| e.id.eq_ignore_ascii_case(id))) } + /// True if ANY row (any quant/backend) carries this id. Lets the resolver + /// tell "this is a catalog model we can't run on this backend" (refuse + /// pre-download) apart from "unknown id" (fall through to a raw HF repo). + pub fn has_id(&self, id: &str) -> bool { + self.models + .iter() + .any(|e| e.id == id || e.id.eq_ignore_ascii_case(id)) + } + /// Check the catalog is internally consistent: every entry well-formed, and /// no duplicate `(id, quant)` pair (which would shadow in `find`/listing). /// Run by tests/CI to catch a malformed catalog edit before it ships and @@ -198,8 +215,17 @@ impl Catalog { if e.quant.is_empty() { anyhow::bail!("{}: empty quant", e.id); } - if !seen.insert((e.id.as_str(), e.quant.as_str())) { - anyhow::bail!("duplicate catalog entry: {} [{}]", e.id, e.quant); + // A model may carry a universal row AND a per-backend variant of the + // same quant (the resolver prefers the backend-native one); those are + // distinguished by `backend`, so the uniqueness key includes it. Two + // rows identical in (id, quant, backend) are still a real duplicate. + if !seen.insert((e.id.as_str(), e.quant.as_str(), e.backend.as_deref())) { + anyhow::bail!( + "duplicate catalog entry: {} [{}] backend={:?}", + e.id, + e.quant, + e.backend + ); } } Ok(()) diff --git a/base-convert/crates/base-hub/src/gen.rs b/base-convert/crates/base-hub/src/gen.rs new file mode 100644 index 0000000..12445e1 --- /dev/null +++ b/base-convert/crates/base-hub/src/gen.rs @@ -0,0 +1,143 @@ +//! Generate catalog entries from `.base` files. +//! +//! The hub catalog has grown one hand-edited JSON row per (model, quant, +//! backend). That does not scale: the `size`/`sha256`/`backend` fields must +//! stay byte-exact with what clients download, and hand-maintaining them is +//! error-prone. Every `.base` header already records the `arch` and the +//! `target_backend` the bundle was packed for; combined with the file's size +//! and sha256 that is all a [`CatalogEntry`] needs except the publish +//! coordinates (`id`, `hf_repo`, quant identity). This module derives an entry +//! straight from the file so a publisher generates rows instead of editing +//! JSON — the catalog becomes a build artifact. + +use crate::catalog::CatalogEntry; +use anyhow::{Context, Result}; +use base_format::TargetBackend; +use std::io::Read; +use std::path::Path; + +/// Map a bundle's packed `target_backend` to the catalog `backend` field the +/// resolver filters on. `None` = universal (resolves on every backend). The +/// Apple/MLX-affine packing is the portable base, so `Metal` maps to universal; +/// the CUDA / ROCm / CPU tile layouts are backend-locked and only resolve for a +/// client of that backend. +pub fn backend_tag(target: TargetBackend) -> Result> { + Ok(match target { + TargetBackend::Metal => None, + // GB10 (sm121) serializes to `cuda_sm121`, which the CUDA runtime accepts + // alongside a generic `cuda`; we advertise the generic tag. sm89/sm90 are + // NOT accepted by the current reader, so refuse to generate an entry that + // would download and then fail to load. + TargetBackend::CudaSm121 => Some("cuda".to_string()), + TargetBackend::CudaSm89 | TargetBackend::CudaSm90 => anyhow::bail!( + "{target:?} is not supported for catalog publishing yet — the runtime accepts \ + cuda / cuda_sm121; rebuild the bundle with --target=cuda_sm121" + ), + TargetBackend::RocmCdna3 => Some("rocm".to_string()), + TargetBackend::CpuAvx2 | TargetBackend::CpuNeon => Some("cpu".to_string()), + }) +} + +/// sha256 of a file, streamed in 1 MiB chunks (bounded memory, no mmap dep). +fn sha256_file(path: &Path) -> Result { + use sha2::{Digest, Sha256}; + let mut f = std::fs::File::open(path).with_context(|| format!("open {}", path.display()))?; + let mut hasher = Sha256::new(); + let mut buf = vec![0u8; 1024 * 1024]; + loop { + let n = f.read(&mut buf)?; + if n == 0 { + break; + } + hasher.update(&buf[..n]); + } + Ok(hasher + .finalize() + .iter() + .map(|b| format!("{b:02x}")) + .collect()) +} + +/// Derive a catalog entry from a local `.base` file. The publisher supplies only +/// the coordinates the header cannot know — `id` and `hf_repo`. Everything the +/// resolver, cache layout, and integrity check depend on — `arch`, `backend`, +/// the quant IDENTITY, `size`, `sha256` — is read/derived from the bytes, so a +/// generated entry can never drift from what clients download. +/// +/// The quant identity is `-` where `` is the header's quant +/// bit-token (`q4`/`q8`/…) and `` is `default` for a universal bundle or +/// the backend tag for a backend-locked one (`cuda-q4`). Keeping the bits as the +/// LAST segment lets `quant_tag`/`pull id:q4` match it; using a distinct slot +/// per backend gives universal and native variants distinct cache dirs and a +/// distinct catalog uniqueness key, so both can coexist for one model. +pub fn entry_from_base(path: &Path, id: &str, hf_repo: &str) -> Result { + let header = base_format::BaseReader::read_header(path) + .with_context(|| format!("read .base header {}", path.display()))?; + 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 quant = match &backend { + Some(b) => format!("{b}-{bits}"), + None => format!("default-{bits}"), + }; + let size = std::fs::metadata(path) + .with_context(|| format!("stat {}", path.display()))? + .len(); + let sha256 = sha256_file(path)?; + let file = path + .file_name() + .and_then(|s| s.to_str()) + .unwrap_or("model.base") + .to_string(); + Ok(CatalogEntry { + id: id.to_string(), + hf_repo: hf_repo.to_string(), + file, + revision: "main".to_string(), + source_repo: None, + arch: Some(header.arch.clone()), + quant, + size: Some(size), + sha256: Some(sha256), + backend, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn backend_tag_maps_target_to_catalog_field() { + assert_eq!(backend_tag(TargetBackend::Metal).unwrap(), None); // universal base + assert_eq!( + backend_tag(TargetBackend::CudaSm121).unwrap().as_deref(), + Some("cuda") + ); + assert_eq!( + backend_tag(TargetBackend::RocmCdna3).unwrap().as_deref(), + Some("rocm") + ); + assert_eq!( + backend_tag(TargetBackend::CpuNeon).unwrap().as_deref(), + Some("cpu") + ); + // sm89 / sm90 aren't accepted by the runtime → refuse to generate an + // entry that would download and then fail to load. + assert!(backend_tag(TargetBackend::CudaSm89).is_err()); + assert!(backend_tag(TargetBackend::CudaSm90).is_err()); + } + + #[test] + fn sha256_file_matches_known_vector() { + // sha256("") = e3b0c442... + let tmp = tempfile::tempdir().unwrap(); + let p = tmp.path().join("empty.base"); + std::fs::write(&p, b"").unwrap(); + assert_eq!( + sha256_file(&p).unwrap(), + "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855" + ); + } +} diff --git a/base-convert/crates/base-hub/src/lib.rs b/base-convert/crates/base-hub/src/lib.rs index 9ed3398..1cb59e0 100644 --- a/base-convert/crates/base-hub/src/lib.rs +++ b/base-convert/crates/base-hub/src/lib.rs @@ -9,6 +9,7 @@ pub mod cache; pub mod catalog; pub mod fetch; +pub mod gen; pub mod registry; pub use cache::{models_dir, HubSidecar}; diff --git a/base-convert/crates/base-hub/src/registry.rs b/base-convert/crates/base-hub/src/registry.rs index 6d2f145..0fb81a5 100644 --- a/base-convert/crates/base-hub/src/registry.rs +++ b/base-convert/crates/base-hub/src/registry.rs @@ -11,6 +11,28 @@ use crate::catalog::Catalog; use anyhow::{Context, Result}; use std::path::{Path, PathBuf}; +/// Extract the quant bit-width token (`q4`, `q8`, `q6`, …) embedded in a quant +/// identity: `default-q4` → `q4`, `cuda-q4mix` → `q4`, `q8` → `q8`. A `q` +/// followed by digits, where the digits don't run into more digits, so `q4` +/// never spuriously matches inside `q40`. Returns `None` when there is no such +/// token (an opaque identity that only matches by exact string). +pub fn quant_bits(quant: &str) -> Option<&str> { + let b = quant.as_bytes(); + let mut i = 0; + while i < b.len() { + if b[i] == b'q' && i + 1 < b.len() && b[i + 1].is_ascii_digit() { + let start = i; + let mut j = i + 1; + while j < b.len() && b[j].is_ascii_digit() { + j += 1; + } + return Some(&quant[start..j]); + } + i += 1; + } + None +} + /// Where a resolved id ultimately comes from. #[derive(Debug, Clone, PartialEq, Eq)] pub enum ModelRef { @@ -206,9 +228,27 @@ impl CatalogRegistry { Self { catalog } } - pub fn resolve(&self, id: &str) -> Option { - let e = self.catalog.find(id)?; - Some(ModelRef::Catalog { + /// The backend this client build serves. The `basert` CLI ships per + /// platform: Apple builds pair with the Metal runtime, everything + /// else with CUDA. Entries whose `backend` field names a different + /// backend are hidden from resolution (and marked in `list`). + pub fn client_backend() -> &'static str { + if cfg!(target_os = "macos") { + "metal" + } else { + "cuda" + } + } + + fn backend_ok(e: &crate::catalog::CatalogEntry) -> bool { + e.backend + .as_deref() + .map(|b| b == Self::client_backend()) + .unwrap_or(true) + } + + fn entry_to_ref(e: &crate::catalog::CatalogEntry) -> ModelRef { + ModelRef::Catalog { id: e.id.clone(), hf_repo: e.hf_repo.clone(), file: e.file.clone(), @@ -217,7 +257,107 @@ impl CatalogRegistry { arch: e.arch.clone(), size: e.size, sha256: e.sha256.clone(), - }) + } + } + + /// True if `entry_quant` satisfies the caller's `want` quant token, + /// comparing on the embedded bit-width so a bare `q4`, a universal + /// `default-q4`, and a backend-native `cuda-q4mix` all match `want=q4` + /// (see [`quant_bits`]). Falls back to an exact string match for identities + /// with no `qN` token. + fn quant_matches(entry_quant: &str, want: &str) -> bool { + let w = quant_bits(want).unwrap_or(want); + match quant_bits(entry_quant) { + Some(b) => b == w, + None => entry_quant == want, + } + } + + /// Quant-agnostic, backend-aware resolve. + pub fn resolve(&self, id: &str) -> Option { + self.resolve_variant(id, None) + } + + /// Backend- and quant-aware catalog resolution. Among the entries sharing + /// `id` (one per quant, and now optionally a per-backend variant of each), + /// keep only those THIS client can run — `backend == client_backend()` or + /// the universal `backend == None` — and, when the caller names a quant, + /// only that quant. A backend-native bundle (e.g. the CUDA `q4mix` packing) + /// is PREFERRED over the universal one, so a CUDA client fetches the + /// CUDA-native `.base` instead of the portable packing the runtime may not + /// even be able to load; it falls back to the universal entry when no + /// backend-native variant is published. Selection happens BEFORE download. + pub fn resolve_variant(&self, id: &str, want_quant: Option<&str>) -> Option { + self.resolve_with_status(id, want_quant).0 + } + + /// Backend- and quant-aware resolution, returning `(resolved, backend_locked)`. + /// + /// Among the rows sharing `id` (one per quant, plus optional per-backend + /// variants), keep those THIS client can run — `backend == client_backend()` + /// or the universal `backend == None` — matching the requested quant by its + /// bit-width (so `q4` selects `default-q4` OR the native `cuda-q4mix`). A + /// backend-native bundle is PREFERRED over the universal fallback so a CUDA + /// client fetches the CUDA-native `.base` instead of the portable packing the + /// runtime may not even load; an exact-id row beats a case-insensitive alias + /// so the requested identity stays authoritative. For a quant-agnostic + /// resolve the requested quant defaults to the first id row's bits, keeping + /// `resolve(id)` quant-stable. + /// + /// `backend_locked` is true when a matching `(id, quant)` row exists but ONLY + /// for a foreign backend: the caller must refuse pre-download rather than + /// fall through to a convert-on-pull of a bundle this client can't run. When + /// the id+quant simply isn't published, both fields are `(None, false)` and + /// the caller may convert-on-pull from the source repo. + pub fn resolve_with_status(&self, id: &str, want_quant: Option<&str>) -> (Option, bool) { + let is_exact = |e: &crate::catalog::CatalogEntry| e.id == id; + let is_id = |e: &crate::catalog::CatalogEntry| e.id == id || e.id.eq_ignore_ascii_case(id); + + // Default quant for a bare resolve = the first id row's bits (exact id + // preferred), so resolve(id) keeps returning that quant, now with the + // backend-native variant of it when one is published. + let first = self + .catalog + .models + .iter() + .find(|e| is_exact(e)) + .or_else(|| self.catalog.models.iter().find(|e| is_id(e))); + let first = match first { + Some(f) => f, + None => return (None, false), // unknown id — caller may raw-HF it + }; + let want_bits: Option = match want_quant { + Some(w) => Some(quant_bits(w).unwrap_or(w).to_string()), + None => quant_bits(&first.quant).map(|s| s.to_string()), + }; + let quant_ok = |e: &crate::catalog::CatalogEntry| match want_bits.as_deref() { + None => true, + Some(w) => Self::quant_matches(&e.quant, w), + }; + + let matched: Vec<&crate::catalog::CatalogEntry> = + self.catalog.models.iter().filter(|e| is_id(e) && quant_ok(e)).collect(); + if matched.is_empty() { + return (None, false); // this id+quant isn't published — convert-on-pull + } + let mut runnable: Vec<&crate::catalog::CatalogEntry> = + matched.iter().copied().filter(|e| Self::backend_ok(e)).collect(); + if runnable.is_empty() { + // Published, but only for a backend this client can't run. + if let Some(e) = matched.first() { + eprintln!( + "hub: '{}' requires the {} backend (this client is {}) — not resolving", + id, + e.backend.as_deref().unwrap_or("?"), + Self::client_backend() + ); + } + return (None, true); + } + // Rank: exact-id before case-insensitive alias; then backend-native + // before universal; stable so catalog order breaks any remaining tie. + runnable.sort_by_key(|e| (u8::from(!is_exact(e)), u8::from(e.backend.is_none()))); + (Some(Self::entry_to_ref(runnable[0])), false) } } @@ -227,6 +367,12 @@ impl Registry for CatalogRegistry { .catalog .models .iter() + // Only advertise rows THIS client can actually pull — a backend- + // qualified row (backend=cuda) must not appear in `list --remote` on + // Metal, where resolution would filter it out anyway (it'd promise an + // artifact the client can't select). Universal rows (backend=None) + // pass on every backend. Mirrors resolve_variant's backend_ok gate. + .filter(|e| Self::backend_ok(e)) .map(|e| ModelEntry { id: e.id.clone(), variant: e.quant.clone(), @@ -293,10 +439,37 @@ impl MergedRegistry { want_quant: Option<&str>, force: bool, ) -> Result { + // Resolve the catalog entry FIRST (backend- + quant-aware) so the + // installed-shortcut checks the RIGHT variant: a CUDA client must not be + // handed a cached universal `default-q4` when the catalog directs it to + // the native `cuda-q4mix`, and a backend-locked model must be refused + // BEFORE any multi-GB download or convert-on-pull. + let (cat, backend_locked) = self.catalog.resolve_with_status(id, want_quant); + if let Some(cref) = cat { + if !force { + if let ModelRef::Catalog { variant, .. } = &cref { + if let Some(path) = self.local.installed_path(id, variant) { + return Ok(ModelRef::Local { + id: id.to_string(), + variant: variant.clone(), + path, + }); + } + } + } + return Ok(cref); + } + if backend_locked { + // Published for this id+quant, but only for another backend. Refuse + // rather than converting-on-pull a bundle this client can't run. + anyhow::bail!( + "hub: {id:?} has no bundle for the {} backend this client runs", + CatalogRegistry::client_backend() + ); + } + // Not published for this id+quant. Honour an already-installed variant + // (legacy `default-` layout) before converting-on-pull. if !force { - // Variants installed by every pull path are named `default-` - // (or the catalog entry's own quant). When the caller names a quant, - // honour exactly that one. if let Some(want) = want_quant { let variant = format!("default-{want}"); if let Some(path) = self.local.installed_path(id, &variant) { @@ -306,21 +479,10 @@ impl MergedRegistry { path, }); } - } else if let Some(ModelRef::Catalog { variant, .. }) = self.catalog.resolve(id) { - if let Some(path) = self.local.installed_path(id, &variant) { - return Ok(ModelRef::Local { - id: id.to_string(), - variant, - path, - }); - } } } - if let Some(r) = self.catalog.resolve(id) { - return Ok(r); - } - // Fall through to a raw HF repo. Require an `org/model` shape so a - // typo'd catalog id doesn't silently become a (failing) HF fetch. + // Fall through to a raw HF repo (convert-on-pull). Require an `org/model` + // shape so a typo'd catalog id doesn't silently become a failing fetch. if id.split('/').filter(|s| !s.is_empty()).count() < 2 { anyhow::bail!( "unknown model id {id:?}: not in the catalog and not an `org/model` HF repo" @@ -395,7 +557,16 @@ mod tests { #[test] fn resolve_installed_shortcut_is_quant_aware() { let tmp = tempfile::tempdir().unwrap(); - let reg = MergedRegistry::new(tmp.path(), catalog_with_one()); + // Catalog publishes BOTH quants under one id. + let cat = CatalogRegistry::from_catalog( + Catalog::from_json( + r#"{"schema":1,"updated":"x","models":[ + {"id":"basecompute/demo","hf_repo":"basecompute/demo","file":"demo-Q4.base","arch":"llama","quant":"default-q4"}, + {"id":"basecompute/demo","hf_repo":"basecompute/demo","file":"demo-Q8.base","arch":"llama","quant":"default-q8"}]}"#, + ) + .unwrap(), + ); + let reg = MergedRegistry::new(tmp.path(), cat); // Only the q4 variant is on disk. let vdir = cache::variant_dir(tmp.path(), "basecompute/demo", "default-q4").unwrap(); std::fs::create_dir_all(&vdir).unwrap(); @@ -406,14 +577,167 @@ mod tests { reg.resolve("basecompute/demo", "main", Some("q4"), false).unwrap(), ModelRef::Local { .. } )); - // Asking for q8 must NOT return the installed q4 — it falls through to - // the catalog so the right quant gets fetched. + // Asking for q8 must NOT return the installed q4 — quant-aware catalog + // resolution fetches the RIGHT quant (`demo-Q8.base`), not the wrong one. + match reg + .resolve("basecompute/demo", "main", Some("q8"), false) + .unwrap() + { + ModelRef::Catalog { file, variant, .. } => { + assert_eq!(file, "demo-Q8.base"); + assert_eq!(variant, "default-q8"); + } + other => panic!("expected the q8 catalog entry, got {other:?}"), + } + } + + // Build a CatalogRegistry from inline JSON (test helper). + fn catalog_json(models: &str) -> CatalogRegistry { + CatalogRegistry::from_catalog( + Catalog::from_json(&format!(r#"{{"schema":1,"updated":"x","models":[{models}]}}"#)) + .unwrap(), + ) + } + + #[test] + fn quant_bits_extracts_bit_token() { + assert_eq!(quant_bits("default-q4"), Some("q4")); + assert_eq!(quant_bits("cuda-q4mix"), Some("q4")); + assert_eq!(quant_bits("q8"), Some("q8")); + assert_eq!(quant_bits("cuda-q4"), Some("q4")); + assert_eq!(quant_bits("bf16"), None); // no qN token + assert_eq!(quant_bits("q40"), Some("q40")); // whole digit run, not "q4" + } + + #[test] + fn resolve_prefers_backend_native_over_universal() { + // Universal `default-q4` and a backend-native `-q4` (distinct cache + // identity, matched by bits). The client fetches the native bundle. + let be = CatalogRegistry::client_backend(); // "cuda" on Linux CI, "metal" on Apple + let uni = r#"{"id":"basecompute/hybrid","hf_repo":"basecompute/hybrid","file":"hybrid-Q4.base","arch":"qwen35","quant":"default-q4"}"#.to_string(); + let nat = format!( + r#"{{"id":"basecompute/hybrid","hf_repo":"basecompute/hybrid","file":"hybrid-Q4-{be}.base","arch":"qwen35","quant":"{be}-q4","backend":"{be}"}}"# + ); + // Native listed after the universal, and before it — the preference must + // hold regardless of catalog order. + for rows in [format!("{uni},{nat}"), format!("{nat},{uni}")] { + let reg = catalog_json(&rows); + reg.catalog.validate().unwrap(); // both rows coexist (backend in key) + match reg.resolve_variant("basecompute/hybrid", Some("q4")).unwrap() { + ModelRef::Catalog { file, variant, .. } => { + assert_eq!(file, format!("hybrid-Q4-{be}.base")); + assert_eq!(variant, format!("{be}-q4")); // distinct cache dir + } + other => panic!("expected the backend-native entry, got {other:?}"), + } + } + } + + #[test] + fn resolve_falls_back_to_universal_when_no_native() { + let reg = catalog_json( + r#"{"id":"basecompute/demo","hf_repo":"basecompute/demo","file":"demo-Q4.base","arch":"llama","quant":"default-q4"}"#, + ); + match reg.resolve_variant("basecompute/demo", Some("q4")).unwrap() { + ModelRef::Catalog { file, .. } => assert_eq!(file, "demo-Q4.base"), + other => panic!("expected the universal entry, got {other:?}"), + } + // A quant that isn't published: absent (not backend-locked) → the merged + // resolver may convert-on-pull. + assert_eq!(reg.resolve_with_status("basecompute/demo", Some("q8")), (None, false)); + } + + #[test] + fn resolve_backend_locked_reports_status_not_absent() { + let foreign = if CatalogRegistry::client_backend() == "cuda" { "metal" } else { "cuda" }; + let reg = catalog_json(&format!( + r#"{{"id":"basecompute/locked","hf_repo":"basecompute/locked","file":"locked-Q4.base","arch":"llama","quant":"{foreign}-q4","backend":"{foreign}"}}"# + )); + // Published for this id+quant, but only for a foreign backend → refuse, + // and flag it distinctly from "absent" so the caller errors (below). + assert_eq!(reg.resolve_with_status("basecompute/locked", Some("q4")), (None, true)); + } + + #[test] + fn merged_resolve_refuses_backend_locked_no_hf_fallthrough() { + let foreign = if CatalogRegistry::client_backend() == "cuda" { "metal" } else { "cuda" }; + let tmp = tempfile::tempdir().unwrap(); + let reg = MergedRegistry::new( + tmp.path(), + catalog_json(&format!( + r#"{{"id":"basecompute/locked","hf_repo":"basecompute/locked","file":"l-Q4.base","arch":"llama","quant":"{foreign}-q4","backend":"{foreign}"}}"# + )), + ); + // Must ERROR (backend-locked), not fall through to a raw HF download of a + // bundle this client can't run. + assert!(reg.resolve("basecompute/locked", "main", Some("q4"), false).is_err()); + } + + #[test] + fn merged_resolve_shortcut_uses_backend_resolved_variant() { + // A CUDA-native variant installed under its own cache dir is found; a + // cached universal default-q4 does NOT shadow the native preference. + let be = CatalogRegistry::client_backend(); + let tmp = tempfile::tempdir().unwrap(); + let reg = MergedRegistry::new( + tmp.path(), + catalog_json(&format!( + r#"{{"id":"basecompute/h","hf_repo":"basecompute/h","file":"h-Q4.base","arch":"qwen35","quant":"default-q4"}}, + {{"id":"basecompute/h","hf_repo":"basecompute/h","file":"h-Q4-{be}.base","arch":"qwen35","quant":"{be}-q4","backend":"{be}"}}"# + )), + ); + // Only the universal is on disk → resolve directs to the native CATALOG + // ref (download), NOT the cached universal. + let uni = cache::variant_dir(tmp.path(), "basecompute/h", "default-q4").unwrap(); + std::fs::create_dir_all(&uni).unwrap(); + std::fs::write(cache::base_artifact_path(&uni), b"universal").unwrap(); + match reg.resolve("basecompute/h", "main", Some("q4"), false).unwrap() { + ModelRef::Catalog { variant, .. } => assert_eq!(variant, format!("{be}-q4")), + other => panic!("cached universal must not shadow the native pick, got {other:?}"), + } + // Now install the native variant → the shortcut returns it as Local. + let nat = cache::variant_dir(tmp.path(), "basecompute/h", &format!("{be}-q4")).unwrap(); + std::fs::create_dir_all(&nat).unwrap(); + std::fs::write(cache::base_artifact_path(&nat), b"native").unwrap(); assert!(matches!( - reg.resolve("basecompute/demo", "main", Some("q8"), false).unwrap(), - ModelRef::Catalog { .. } + reg.resolve("basecompute/h", "main", Some("q4"), false).unwrap(), + ModelRef::Local { variant, .. } if variant == format!("{be}-q4") )); } + #[test] + fn resolve_quant_agnostic_keeps_default_quant() { + // default-q4 listed first, an optional q8 later. resolve(id) (no quant) + // must stay on the q4 family, not jump to q8. + let reg = catalog_json( + r#"{"id":"basecompute/m","hf_repo":"basecompute/m","file":"m-Q4.base","arch":"llama","quant":"default-q4"}, + {"id":"basecompute/m","hf_repo":"basecompute/m","file":"m-Q8.base","arch":"llama","quant":"default-q8"}"#, + ); + match reg.resolve("basecompute/m").unwrap() { + ModelRef::Catalog { file, .. } => assert_eq!(file, "m-Q4.base"), + other => panic!("resolve(id) must keep the default (q4), got {other:?}"), + } + // Explicit q8 still selects q8. + match reg.resolve_variant("basecompute/m", Some("q8")).unwrap() { + ModelRef::Catalog { file, .. } => assert_eq!(file, "m-Q8.base"), + other => panic!("expected q8, got {other:?}"), + } + } + + #[test] + fn resolve_prefers_exact_id_over_case_alias() { + // Two case-distinct ids: an exact request must resolve to the exact id's + // row, not a differently-cased alias listed first. + let reg = catalog_json( + r#"{"id":"basecompute/CamelModel","hf_repo":"basecompute/alias","file":"alias-Q4.base","arch":"llama","quant":"default-q4"}, + {"id":"basecompute/camelmodel","hf_repo":"basecompute/exact","file":"exact-Q4.base","arch":"llama","quant":"default-q4"}"#, + ); + match reg.resolve_variant("basecompute/camelmodel", Some("q4")).unwrap() { + ModelRef::Catalog { hf_repo, .. } => assert_eq!(hf_repo, "basecompute/exact"), + other => panic!("exact id must win, got {other:?}"), + } + } + #[test] fn local_scan_walks_tree_and_skips_staging() { let tmp = tempfile::tempdir().unwrap(); diff --git a/base-convert/profiles/default-q4-embq6.json b/base-convert/profiles/default-q4-embq6.json new file mode 100644 index 0000000..52e3f8d --- /dev/null +++ b/base-convert/profiles/default-q4-embq6.json @@ -0,0 +1,128 @@ +{ + "name": "default-q4-embq6", + "arch": "*", + "rules": [ + { + "pattern": "model.embed_tokens.weight", + "dtype": "base_q6", + "scale_dtype": "bf16", + "group_size": 64 + }, + { + "pattern": "embed_tokens.weight", + "dtype": "base_q6", + "scale_dtype": "bf16", + "group_size": 64 + }, + { + "pattern": "**.input_layernorm.weight", + "dtype": "f16" + }, + { + "pattern": "**.input_norm.weight", + "dtype": "f16" + }, + { + "pattern": "**.post_attention_layernorm.weight", + "dtype": "f16" + }, + { + "pattern": "**.post_attn_norm.weight", + "dtype": "f16" + }, + { + "pattern": "**.post_attention_norm.weight", + "dtype": "f16" + }, + { + "pattern": "**.post_ffw_norm.weight", + "dtype": "f16" + }, + { + "pattern": "**.pre_ffw_norm.weight", + "dtype": "f16" + }, + { + "pattern": "**.per_layer_post_norm.weight", + "dtype": "f16" + }, + { + "pattern": "per_layer_proj_norm.weight", + "dtype": "f16" + }, + { + "pattern": "**.q_norm.weight", + "dtype": "f16" + }, + { + "pattern": "**.k_norm.weight", + "dtype": "f16" + }, + { + "pattern": "**.attn_q_norm.weight", + "dtype": "f16" + }, + { + "pattern": "**.attn_k_norm.weight", + "dtype": "f16" + }, + { + "pattern": "**.layer_out_scale.weight", + "dtype": "f16" + }, + { + "pattern": "**.ffn_norm.weight", + "dtype": "f16" + }, + { + "pattern": "**.attn_norm.weight", + "dtype": "f16" + }, + { + "pattern": "model.norm.weight", + "dtype": "f16" + }, + { + "pattern": "final_norm.weight", + "dtype": "f16" + }, + { + "pattern": "output_norm.weight", + "dtype": "f16" + }, + { + "pattern": "**.{q,k,v,o}_proj.weight", + "dtype": "base_q4", + "scale_dtype": "bf16", + "group_size": 64 + }, + { + "pattern": "**.{gate,up,down}_proj.weight", + "dtype": "base_q4", + "scale_dtype": "bf16", + "group_size": 64 + }, + { + "pattern": "lm_head.weight", + "dtype": "base_q4", + "scale_dtype": "bf16", + "group_size": 64 + }, + { + "pattern": "output.weight", + "dtype": "base_q4", + "scale_dtype": "bf16", + "group_size": 64 + }, + { + "pattern": "**.weight", + "dtype": "base_q4", + "scale_dtype": "bf16", + "group_size": 64 + }, + { + "pattern": "**.bias", + "dtype": "f16" + } + ] +} diff --git a/base-convert/profiles/default-q8.json b/base-convert/profiles/default-q8.json index fc2efda..6793b2e 100644 --- a/base-convert/profiles/default-q8.json +++ b/base-convert/profiles/default-q8.json @@ -2,16 +2,81 @@ "name": "default-q8", "arch": "*", "rules": [ - {"pattern": "model.embed_tokens.weight", "dtype": "f16"}, - {"pattern": "**.input_layernorm.weight", "dtype": "f16"}, - {"pattern": "**.post_attention_layernorm.weight", "dtype": "f16"}, - {"pattern": "**.q_norm.weight", "dtype": "f16"}, - {"pattern": "**.k_norm.weight", "dtype": "f16"}, - {"pattern": "model.norm.weight", "dtype": "f16"}, - {"pattern": "**.{q,k,v,o}_proj.weight", "dtype": "base_q8", "scale_dtype": "bf16", "group_size": 128}, - {"pattern": "**.{gate,up,down}_proj.weight", "dtype": "base_q8", "scale_dtype": "bf16", "group_size": 128}, - {"pattern": "lm_head.weight", "dtype": "base_q8", "scale_dtype": "bf16", "group_size": 128}, - {"pattern": "**.weight", "dtype": "base_q8", "scale_dtype": "bf16", "group_size": 128}, - {"pattern": "**.bias", "dtype": "f16"} + { + "pattern": "model.embed_tokens.weight", + "dtype": "f16" + }, + { + "pattern": "**.input_layernorm.weight", + "dtype": "f16" + }, + { + "pattern": "**.post_attention_layernorm.weight", + "dtype": "f16" + }, + { + "pattern": "**.q_norm.weight", + "dtype": "f16" + }, + { + "pattern": "**.k_norm.weight", + "dtype": "f16" + }, + { + "pattern": "**.pre_feedforward_layernorm.weight", + "dtype": "f16" + }, + { + "pattern": "**.post_feedforward_layernorm.weight", + "dtype": "f16" + }, + { + "pattern": "**.pre_ffw_norm.weight", + "dtype": "f16" + }, + { + "pattern": "**.post_ffw_norm.weight", + "dtype": "f16" + }, + { + "pattern": "**.per_layer_post_norm.weight", + "dtype": "f16" + }, + { + "pattern": "model.norm.weight", + "dtype": "f16" + }, + { + "pattern": "**.{q,k,v,o}_proj.weight", + "dtype": "base_q8", + "scale_dtype": "bf16", + "group_size": 128 + }, + { + "pattern": "**.{gate,up,down}_proj.weight", + "dtype": "base_q8", + "scale_dtype": "bf16", + "group_size": 128 + }, + { + "pattern": "lm_head.weight", + "dtype": "base_q8", + "scale_dtype": "bf16", + "group_size": 128 + }, + { + "pattern": "**.post_attention_norm.weight", + "dtype": "f16" + }, + { + "pattern": "**.weight", + "dtype": "base_q8", + "scale_dtype": "bf16", + "group_size": 128 + }, + { + "pattern": "**.bias", + "dtype": "f16" + } ] } diff --git a/benchmarks/scripts/fetch_3way_models.sh b/benchmarks/scripts/fetch_3way_models.sh new file mode 100644 index 0000000..9428432 --- /dev/null +++ b/benchmarks/scripts/fetch_3way_models.sh @@ -0,0 +1,123 @@ +#!/bin/bash +# Fetch the model artifacts the three-way benchmark needs, into MODELS_DIR, +# from pinned public sources. Idempotent: files already present are skipped. +# +# After this, run: MODELS="" benchmarks/scripts/three_way_benchmark.sh +# +# Per engine: +# baseRT : .base pulled from the PUBLISHED basecompute HF catalog +# (basecompute/) — no local conversions. Every baseRT model +# in three_way_benchmark.sh MODEL_TABLE has a published repo, +# including the MoE (Qwen3-30B, Gemma-4-26B) and hybrid-GDN +# (Qwen3.5 / Qwen3.6) rows. +# llama.cpp: Q4/Q8 GGUF from bartowski / ggml-org / unsloth +# mlx-lm : nothing to fetch — mlx_lm.benchmark pulls the HF repo on first use +# +# Env: MODELS_DIR (default ../../models), SET=q4|q8|all (default all). +set -uo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +MODELS_DIR="${MODELS_DIR:-$(cd "$SCRIPT_DIR/../.." && pwd)/models}" +SET="${SET:-all}" +mkdir -p "$MODELS_DIR" + +# rows: local_filename | hf_repo | remote_filename +# (local_filename matches the name used in three_way_benchmark.sh MODEL_TABLE; +# a local_filename may contain a subpath, e.g. gguf/Qwen3.5-2B/…) +# +# baseRT throughput is independent of the exact quant scheme / instruct-vs-base +# variant (same architecture and size), so the .base saved under the table's +# label may pair with a Q4_K_M / Q8_0 GGUF and a 4/8-bit MLX repo. + +# ── baseRT .base — all from the published basecompute catalog ────────────── +Q4_BASE=( + # dense + "Llama-3.2-1B-Q4.base|basecompute/Llama-3.2-1B-Instruct|Llama-3.2-1B-Instruct-Q4.base" + "Llama-3.2-3B-Q4.base|basecompute/Llama-3.2-3B-Instruct|Llama-3.2-3B-Instruct-Q4.base" + "Qwen3-0.6B-Q4_0.base|basecompute/Qwen3-0.6B|Qwen3-0.6B-Q4.base" + "gemma-4-E2B-it-Q4_0.base|basecompute/gemma-4-E2B-it|gemma-4-E2B-it-Q4.base" + # MoE + "Qwen3-30B-A3B-Q4.base|basecompute/Qwen3-30B-A3B-Instruct-2507|Qwen3-30B-A3B-Instruct-2507-Q4.base" + "Gemma-4-26B-A4B-Q4.base|basecompute/gemma-4-26B-A4B-it|gemma-4-26B-A4B-it-Q4.base" + # hybrid-GDN (Qwen3.5 / Qwen3.6) + "Qwen3.5-2B-Base-Q4.base|basecompute/Qwen3.5-2B-Base|Qwen3.5-2B-Base-Q4.base" + "Qwen3.5-35B-A3B-Q4.base|basecompute/Qwen3.5-35B-A3B|Qwen3.5-35B-A3B-Q4.base" + "Qwen3.6-27B-Q4.base|basecompute/Qwen3.6-27B|Qwen3.6-27B-Q4.base" + "Qwen3.6-35B-A3B-Q4.base|basecompute/Qwen3.6-35B-A3B|Qwen3.6-35B-A3B-Q4.base" +) +Q8_BASE=( + # dense + "Llama-3.2-1B-Q8.base|basecompute/Llama-3.2-1B-Instruct|Llama-3.2-1B-Instruct-Q8.base" + "Llama-3.2-3B-Q8.base|basecompute/Llama-3.2-3B-Instruct|Llama-3.2-3B-Instruct-Q8.base" + "Qwen3-0.6B-Q8.base|basecompute/Qwen3-0.6B|Qwen3-0.6B-Q8.base" + "gemma-4-E2B-it-Q8.base|basecompute/gemma-4-E2B-it|gemma-4-E2B-it-Q8.base" + # MoE. The Instruct-2507 30B has no published Q8; the Thinking-2507 variant + # does and is the same architecture/size, so throughput is identical. + "Qwen3-30B-A3B-Q8.base|basecompute/Qwen3-30B-A3B-Thinking-2507|Qwen3-30B-A3B-Thinking-2507-Q8.base" + "Gemma-4-26B-A4B-Q8.base|basecompute/gemma-4-26B-A4B-it|gemma-4-26B-A4B-it-Q8.base" + # hybrid-GDN + "Qwen3.5-2B-Base-Q8.base|basecompute/Qwen3.5-2B-Base|Qwen3.5-2B-Base-Q8.base" + "Qwen3.6-27B-Q8.base|basecompute/Qwen3.6-27B|Qwen3.6-27B-Q8.base" +) + +# ── llama.cpp GGUF ───────────────────────────────────────────────────────── +Q4_GGUF=( + # dense + "Llama-3.2-1B-Instruct-Q4_0.gguf|bartowski/Llama-3.2-1B-Instruct-GGUF|Llama-3.2-1B-Instruct-Q4_0.gguf" + "Llama-3.2-3B-Instruct-Q4_0.gguf|bartowski/Llama-3.2-3B-Instruct-GGUF|Llama-3.2-3B-Instruct-Q4_0.gguf" + "Qwen3-0.6B-Q4_0.gguf|ggml-org/Qwen3-0.6B-GGUF|Qwen3-0.6B-Q4_0.gguf" + "gemma-4-E2B-it-Q4_0.gguf|unsloth/gemma-4-E2B-it-GGUF|gemma-4-E2B-it-Q4_0.gguf" + # MoE + "Qwen_Qwen3-30B-A3B-Instruct-2507-Q4_0.gguf|bartowski/Qwen_Qwen3-30B-A3B-Instruct-2507-GGUF|Qwen_Qwen3-30B-A3B-Instruct-2507-Q4_0.gguf" + "google_gemma-4-26B-A4B-it-Q4_0.gguf|bartowski/google_gemma-4-26B-A4B-it-GGUF|google_gemma-4-26B-A4B-it-Q4_0.gguf" + # hybrid-GDN + "gguf/Qwen3.5-2B/Qwen_Qwen3.5-2B-Q4_K_M.gguf|bartowski/Qwen_Qwen3.5-2B-GGUF|Qwen_Qwen3.5-2B-Q4_K_M.gguf" + "gguf/Qwen3.5-35B/Qwen_Qwen3.5-35B-A3B-Q4_K_M.gguf|bartowski/Qwen_Qwen3.5-35B-A3B-GGUF|Qwen_Qwen3.5-35B-A3B-Q4_K_M.gguf" + "gguf/Qwen3.6-27B/Qwen3.6-27B-Q4_K_M.gguf|unsloth/Qwen3.6-27B-GGUF|Qwen3.6-27B-Q4_K_M.gguf" + "gguf/Qwen3.6-35B/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf|unsloth/Qwen3.6-35B-A3B-GGUF|Qwen3.6-35B-A3B-UD-Q4_K_M.gguf" +) +Q8_GGUF=( + # dense + "Llama-3.2-1B-Instruct-Q8_0.gguf|bartowski/Llama-3.2-1B-Instruct-GGUF|Llama-3.2-1B-Instruct-Q8_0.gguf" + "Llama-3.2-3B-Instruct-Q8_0.gguf|bartowski/Llama-3.2-3B-Instruct-GGUF|Llama-3.2-3B-Instruct-Q8_0.gguf" + "Qwen3-0.6B-Q8_0.gguf|ggml-org/Qwen3-0.6B-GGUF|Qwen3-0.6B-Q8_0.gguf" + "unsloth/gemma-4-E2B-it-GGUF/gemma-4-E2B-it-Q8_0.gguf|unsloth/gemma-4-E2B-it-GGUF|gemma-4-E2B-it-Q8_0.gguf" + # MoE (see baseRT note: 30B Q8 uses the Thinking-2507 variant) + "Qwen_Qwen3-30B-A3B-Thinking-2507-Q8_0.gguf|bartowski/Qwen_Qwen3-30B-A3B-Thinking-2507-GGUF|Qwen_Qwen3-30B-A3B-Thinking-2507-Q8_0.gguf" + "google_gemma-4-26B-A4B-it-Q8_0.gguf|bartowski/google_gemma-4-26B-A4B-it-GGUF|google_gemma-4-26B-A4B-it-Q8_0.gguf" + # hybrid-GDN + "gguf/Qwen3.5-2B/Qwen_Qwen3.5-2B-Q8_0.gguf|bartowski/Qwen_Qwen3.5-2B-GGUF|Qwen_Qwen3.5-2B-Q8_0.gguf" + "gguf/Qwen3.6-27B/Qwen3.6-27B-Q8_0.gguf|unsloth/Qwen3.6-27B-GGUF|Qwen3.6-27B-Q8_0.gguf" +) + +fetch() { # local_name|repo|remote_name + local row="$1"; IFS='|' read -r local repo remote <<< "$row" + local dest="$MODELS_DIR/$local" + if [ -f "$dest" ]; then echo " have $local"; return; fi + echo " fetch $local <- $repo/$remote" + mkdir -p "$(dirname "$dest")" + local tmp; tmp="$(mktemp -d)" + if hf download "$repo" "$remote" --local-dir "$tmp" >/dev/null 2>&1; then + cp "$tmp/$remote" "$dest" + else + echo " WARN: failed to fetch $repo/$remote (skipping)" + fi + rm -rf "$tmp" +} + +echo "=== fetch 3-way models (set=$SET) into $MODELS_DIR ===" +# Inline expansion (no array indirection — works on macOS bash 3.2). +case "$SET" in + q4) ROWS=( "${Q4_BASE[@]}" "${Q4_GGUF[@]}" ) ;; + q8) ROWS=( "${Q8_BASE[@]}" "${Q8_GGUF[@]}" ) ;; + all) ROWS=( "${Q4_BASE[@]}" "${Q4_GGUF[@]}" "${Q8_BASE[@]}" "${Q8_GGUF[@]}" ) ;; + *) echo "unknown SET=$SET (use q4|q8|all)"; exit 1 ;; +esac +for r in "${ROWS[@]}"; do fetch "$r"; done +echo "mlx-lm repos are pulled on first run by mlx_lm.benchmark (no local fetch)." +echo "Done. Labels (see MODEL_TABLE in three_way_benchmark.sh for the full set):" +echo " Q4: Llama-3.2-1B-Q4 Llama-3.2-3B-Q4 Qwen3-0.6B-Q4 Gemma-4-E2B-Q4 \\" +echo " Qwen3-30B-A3B-Q4 Gemma-4-26B-A4B-Q4 Qwen3.5-2B-Q4 Qwen3.5-35B-A3B-Q4 \\" +echo " Qwen3.6-27B-Q4 Qwen3.6-35B-Q4" +echo " Q8: Llama-3.2-1B-Q8 Llama-3.2-3B-Q8 Qwen3-0.6B-Q8 Gemma-4-E2B-Q8 \\" +echo " Qwen3-30B-A3B-Q8 Gemma-4-26B-A4B-Q8 Qwen3.5-2B-Q8 Qwen3.6-27B-Q8" +echo "Run: MODELS=\"\" $SCRIPT_DIR/three_way_benchmark.sh" diff --git a/benchmarks/scripts/gb10_basert_prefill_decode.sh b/benchmarks/scripts/gb10_basert_prefill_decode.sh new file mode 100644 index 0000000..44110cf --- /dev/null +++ b/benchmarks/scripts/gb10_basert_prefill_decode.sh @@ -0,0 +1,63 @@ +#!/usr/bin/env bash +# baseRT single-stream prefill/decode on GB10 via baseRT_cuda_bench, for the +# canonical model set at BOTH quants (cuda-q4mix = 4-bit, cuda-q8 = 8-bit). The +# 8-bit rows are the fair counterpart to vLLM fp8 (gb10_vllm_prefill_decode.sh); +# the 4-bit rows pair with llama.cpp Q4 (bench_2way). pp512 + tg128, r=5, +# page-cache reclaimed first. Writes benchmarks/gb10/basert-prefill-decode.csv. +set -uo pipefail +cd "$(git -C "$(dirname "$0")" rev-parse --show-toplevel)" +BIN="${BASERT_BIN:-build-cuda/baseRT_cuda_bench}" +OUT="${OUT:-benchmarks/gb10}" +M="$HOME/models" +RECLAIM="$(dirname "$0")/reclaim_page_cache.sh" + +# key -> "q4mix_bundle q8_bundle" +declare -A BUN=( + [gemma-3-1b]="gemma-3-1b-it-cuda-q4mix gemma-3-1b-it-cuda-q8" + [gemma-4-e2b]="gemma-4-E2B-it-cuda-q4mix gemma-4-E2B-it-cuda-q8" + [gemma-4-26b]="gemma-4-26B-A4B-it-cuda-q4mix gemma-4-26B-A4B-it-cuda-q8" + [llama-3.2-1b]="Llama-3.2-1B-Instruct-cuda-q4mix Llama-3.2-1B-Instruct-cuda-q8" + [llama-3.2-3b]="Llama-3.2-3B-Instruct-cuda-q4mix Llama-3.2-3B-Instruct-cuda-q8" + [qwen3-0.6b]="Qwen3-0.6B-cuda-q4mix Qwen3-0.6B-cuda-q8" + [qwen3-30b-a3b]="Qwen3-30B-A3B-Instruct-2507-cuda-q4mix Qwen3-30B-A3B-Instruct-2507-cuda-q8" + [qwen3.5-2b]="Qwen3.5-2B-cuda-q4mix Qwen3.5-2B-cuda-q8" + [qwen3.5-35b-a3b]="Qwen3.5-35B-A3B-cuda-q4mix Qwen3.5-35B-A3B-cuda-q8" + [qwen3.6-27b]="Qwen3.6-27B-cuda-q4mix Qwen3.6-27B-cuda-q8" + [qwen3.6-35b]="Qwen3.6-35B-A3B-cuda-q4mix Qwen3.6-35B-A3B-cuda-q8" +) +KEYS=("$@"); [ ${#KEYS[@]} -eq 0 ] && KEYS=(gemma-3-1b gemma-4-e2b gemma-4-26b \ + llama-3.2-1b llama-3.2-3b qwen3-0.6b qwen3-30b-a3b \ + qwen3.5-2b qwen3.5-35b-a3b qwen3.6-27b qwen3.6-35b) + +mkdir -p "$OUT" # the header redirect + later mv/tee fail silently (no set -e) if OUT is a fresh dir +CSV="$OUT/basert-prefill-decode.csv" +[ -f "$CSV" ] || echo "model,quant,pp512_tok_s,tg128_tok_s" > "$CSV" +run1() { # bundle -> "pp tg" + local out pp tg + # Reclaim page cache before EACH bundle: GB10 unified memory decode drops + # 20-25% under high cache, and earlier bundles in this 22-run loop leave the + # cache hot — so a once-at-start reclaim would bias later rows. Keeps rows + # comparable (matches the "reclaim before every engine invocation" claim). + # Redirect BOTH streams: the helper prints "(reclaiming ...)" to STDOUT, which + # this captured function must not leak into its `pp tg` result. + bash "$RECLAIM" 40 >/dev/null 2>&1 || true + out=$("$BIN" "$M/$1.base" -p 512 -n 128 -r 5 2>&1) + pp=$(echo "$out" | awk -F'|' '$4 ~ /pp512/ {gsub(/[^0-9.]/,"",$5);print substr($5,1,index($5,".")+2)}' | tail -1) + tg=$(echo "$out" | awk -F'|' '$4 ~ /tg128/ {gsub(/[^0-9.]/,"",$5);print substr($5,1,index($5,".")+2)}' | tail -1) + echo "$pp $tg" +} +for k in "${KEYS[@]}"; do + read -r q4b q8b <<< "${BUN[$k]}" + for pair in "q4mix:$q4b" "q8:$q8b"; do + quant="${pair%%:*}"; bun="${pair#*:}" + [ -f "$M/$bun.base" ] || { echo "::skip $k/$quant ($bun missing)"; continue; } + read -r pp tg <<< "$(run1 "$bun")" + # Only replace the existing (model,quant) row once BOTH rates parsed — an OOM + # or a bench that emitted no pp512/tg128 row leaves them empty, and blowing + # away a prior valid measurement with a blank is worse than keeping the old. + if [ -z "$pp" ] || [ -z "$tg" ]; then echo "::skip $k/$quant (no measurement: pp='$pp' tg='$tg')"; continue; fi + awk -F, -v k="$k" -v q="$quant" 'NR==1 || !($1==k && $2==q)' "$CSV" > "$CSV.tmp" && mv "$CSV.tmp" "$CSV" + echo "$k,$quant,$pp,$tg" | tee -a "$CSV" + done +done +echo "=== wrote $CSV ===" diff --git a/benchmarks/scripts/gb10_serving_vs_vllm.sh b/benchmarks/scripts/gb10_serving_vs_vllm.sh new file mode 100644 index 0000000..148f8c4 --- /dev/null +++ b/benchmarks/scripts/gb10_serving_vs_vllm.sh @@ -0,0 +1,230 @@ +#!/usr/bin/env bash +# Reproducible baseRT-vs-vLLM serving sweep on GB10 (DGX Spark). +# +# For each model it launches vLLM (docker) and baseRT serve in turn (never +# together — they'd contend for the GPU), drives BOTH through the same client +# (bench/bench.sh -> bench_serving.py) at c1/8/16/32, and emits a matched +# out-tok/s comparison table. This is the committed replacement for the old +# session-scratchpad run32b.sh / vllm_sweep.sh (which produced the July numbers +# but were never checked in). +# +# Usage: benchmarks/scripts/gb10_serving_vs_vllm.sh [model_key ...] +# model_key: one of the keys in the MODELS table below (default: all). +# Env: +# CONC="1 8 16 32" concurrency points +# OUT=benchmarks/gb10 output dir (writes serving-vs-vllm.csv; the .md summary +# is regenerated from the CSV separately, not by this script) +# VLLM_IMAGE=hellohal2064/vllm-dgx-spark-gb10:latest +# BASERT_SERVE=build-rel/basert-serve +# HF_HUB=/mnt/nas-models/hf-hub +set -uo pipefail +cd "$(git -C "$(dirname "$0")" rev-parse --show-toplevel)" + +CONC="${CONC:-1 8 16 32}" +OUT="${OUT:-benchmarks/gb10}" +VLLM_IMAGE="${VLLM_IMAGE:-hellohal2064/vllm-dgx-spark-gb10:latest}" +BASERT_SERVE="${BASERT_SERVE:-build-rel/basert-serve}" +HF_HUB="${HF_HUB:-/mnt/nas-models/hf-hub}" +PORT=8000 +RESDIR="$OUT/serving_raw" +mkdir -p "$RESDIR" + +# key | baseRT .base | vLLM HF repo dir (the models--ORG--NAME cache dir; the +# whole repo is mounted so snapshot symlinks into blobs/ resolve in-container). +# FAIR 8-bit lane: baseRT serves its int8 (cuda-q8) bundle vs vLLM fp8 — both +# ~8-bit, comparable footprint. (q4mix is baseRT's lighter shipping format; the +# 4-bit lane is a separate baseRT-vs-llama.cpp table, no fair vLLM 4-bit here.) +# Canonical 11-model set, matching the Metal 3-way sweep (bench_3way). +declare -A BASE=( + [gemma-3-1b]="$HOME/models/gemma-3-1b-it-cuda-q8.base" + [gemma-4-e2b]="$HOME/models/gemma-4-E2B-it-cuda-q8.base" + [gemma-4-26b]="$HOME/models/gemma-4-26B-A4B-it-cuda-q8.base" + [llama-3.2-1b]="$HOME/models/Llama-3.2-1B-Instruct-cuda-q8.base" + [llama-3.2-3b]="$HOME/models/Llama-3.2-3B-Instruct-cuda-q8.base" + [qwen3-0.6b]="$HOME/models/Qwen3-0.6B-cuda-q8.base" + [qwen3-30b-a3b]="$HOME/models/Qwen3-30B-A3B-Instruct-2507-cuda-q8.base" + [qwen3.5-2b]="$HOME/models/Qwen3.5-2B-cuda-q8.base" + [qwen3.5-35b-a3b]="$HOME/models/Qwen3.5-35B-A3B-cuda-q8.base" + [qwen3.6-27b]="$HOME/models/Qwen3.6-27B-cuda-q8.base" + [qwen3.6-35b]="$HOME/models/Qwen3.6-35B-A3B-cuda-q8.base" +) +declare -A HFREPO=( + [gemma-3-1b]="$HF_HUB/models--unsloth--gemma-3-1b-it" + [gemma-4-e2b]="$HF_HUB/models--google--gemma-4-E2B-it" + [gemma-4-26b]="$HF_HUB/models--google--gemma-4-26B-A4B-it" + [llama-3.2-1b]="$HF_HUB/models--meta-llama--Llama-3.2-1B-Instruct" + [llama-3.2-3b]="$HF_HUB/models--meta-llama--Llama-3.2-3B-Instruct" + [qwen3-0.6b]="$HF_HUB/models--Qwen--Qwen3-0.6B" + [qwen3-30b-a3b]="$HF_HUB/models--Qwen--Qwen3-30B-A3B-Instruct-2507" + [qwen3.5-2b]="$HF_HUB/models--Qwen--Qwen3.5-2B" + [qwen3.5-35b-a3b]="$HF_HUB/models--Qwen--Qwen3.5-35B-A3B" + [qwen3.6-27b]="$HF_HUB/models--Qwen--Qwen3.6-27B" + [qwen3.6-35b]="$HF_HUB/models--Qwen--Qwen3.6-35B-A3B" +) +# vLLM image per model. The GB10-optimized image (default) is Qwen-SPECIALISED: +# its entrypoint bakes in `--reasoning-parser qwen3 --tool-call-parser qwen3_coder +# --enable-auto-tool-choice`, which CRASH on non-Qwen archs at startup (observed: +# llama/mistral containers exit 1 before serving). It also has older Transformers +# that don't recognise the qwen3_5/qwen3_6 hybrid arch. So: +# - Dense Qwen3 + Qwen3-MoE (30B): GB10 image (the fairest vLLM-on-GB10 baseline; +# qwen parsers apply cleanly). +# - Qwen3.5/3.6 hybrids: stock v0.19.1 (resolves Qwen3_5ForConditionalGeneration). +# - Llama / Mistral / Gemma-4: stock v0.20.0 (no qwen parsers; gemma-4 arch is +# 0.20-only). Gemma-4 is multimodal, so the stock branch raises max-num-batched- +# tokens past its per-image MM token count (2496) or vLLM refuses to start. +VLLM_IMAGE_HYBRID="${VLLM_IMAGE_HYBRID:-vllm/vllm-openai:v0.19.1-cu130}" +VLLM_IMAGE_STOCK="${VLLM_IMAGE_STOCK:-vllm/vllm-openai:v0.20.0-cu130}" +declare -A VIMG=( + [qwen3.5-2b]="$VLLM_IMAGE_HYBRID" + [qwen3.5-35b-a3b]="$VLLM_IMAGE_HYBRID" + [qwen3.6-27b]="$VLLM_IMAGE_HYBRID" + [qwen3.6-35b]="$VLLM_IMAGE_HYBRID" + [gemma-3-1b]="$VLLM_IMAGE_STOCK" + [gemma-4-e2b]="$VLLM_IMAGE_STOCK" + [gemma-4-26b]="$VLLM_IMAGE_STOCK" + [llama-3.2-1b]="$VLLM_IMAGE_STOCK" + [llama-3.2-3b]="$VLLM_IMAGE_STOCK" +) +# qwen3-0.6b + qwen3-30b-a3b (dense/MoE Qwen) use the default GB10 image. + +KEYS=("$@"); [ ${#KEYS[@]} -eq 0 ] && KEYS=(gemma-3-1b gemma-4-e2b gemma-4-26b \ + llama-3.2-1b llama-3.2-3b qwen3-0.6b qwen3-30b-a3b \ + qwen3.5-2b qwen3.5-35b-a3b qwen3.6-27b qwen3.6-35b) + +VLLM_CID=""; BASERT_PID="" +cleanup() { + [ -n "$VLLM_CID" ] && docker rm -f "$VLLM_CID" >/dev/null 2>&1 + if [ -n "$BASERT_PID" ]; then + kill "$BASERT_PID" >/dev/null 2>&1 + # WAIT for the process to actually exit — a bare kill returns immediately, + # but baseRT's CUDA context (model weights + KV, up to ~100GB on GB10's + # unified memory) is only released when the process fully tears down. If the + # NEXT model's vLLM launches before that, it sees a near-full GPU and + # OOM-crashes on startup (observed: models 2..N of a sweep silently skipped + # with 'vLLM never became ready', the container gone by the time we log it). + for _ in $(seq 1 30); do kill -0 "$BASERT_PID" 2>/dev/null || break; sleep 1; done + kill -9 "$BASERT_PID" >/dev/null 2>&1; wait "$BASERT_PID" 2>/dev/null + fi + VLLM_CID=""; BASERT_PID="" + # A few seconds for the CUDA driver to reclaim the freed allocations before + # the next engine probes memory. Cheap insurance vs a whole wasted model run. + sleep 8 +} +trap cleanup EXIT + +wait_ready() { # timeout_s + local t=0 + while [ "$t" -lt "$1" ]; do + curl -s -m 3 "http://127.0.0.1:$PORT/v1/models" 2>/dev/null | grep -q '"id"' && return 0 + sleep 3; t=$((t+3)) + done + return 1 +} + +# extract out_tok_per_s for a given tput/cN label from a bench JSONL +ots() { # jsonl label + python3 - "$1" "$2" <<'PY' +import json,sys +path,label=sys.argv[1],sys.argv[2] +try: + for line in open(path): + line=line.strip() + if not line or line[0] != '{': continue + d=json.loads(line) + if d.get("label")==label: + # Drop a data point with ANY failed request — bench_serving.py still + # reports out_tok_per_s over the wall window even when requests + # errored (common at high concurrency), which would understate/ + # corrupt the throughput row. Only a clean burst counts. + print("" if d.get("failed",0) else d.get("out_tok_per_s","")); break + else: print("") +except FileNotFoundError: print("") +PY +} + +CSV="$OUT/serving-vs-vllm.csv" +# Append-safe: only write the header if the CSV doesn't already exist, so a +# subsequent run over additional model keys extends the table instead of +# wiping the earlier rows. +[ -f "$CSV" ] || echo "model,concurrency,baseRT_out_tok_s,vllm_out_tok_s,baseRT_over_vllm_pct" > "$CSV" + +served_model_id() { curl -s -m 5 "http://127.0.0.1:$PORT/v1/models" 2>/dev/null \ + | python3 -c "import sys,json;print(json.load(sys.stdin)['data'][0]['id'])" 2>/dev/null; } + +for k in "${KEYS[@]}"; do + base="${BASE[$k]:-}"; repo="${HFREPO[$k]:-}" + [ -f "$base" ] || { echo "::skip $k — baseRT bundle missing: $base"; continue; } + [ -d "$repo/snapshots" ] || { echo "::skip $k — HF repo missing: $repo"; continue; } + # Deterministic snapshot: the revision refs/main points at, not the + # lexicographically-first hash `ls | head -1` would pick. + snaprel="snapshots/$( { cat "$repo/refs/main" 2>/dev/null || ls -t "$repo/snapshots" | head -1; } )" + echo "############ $k (baseRT=$base vLLM=$repo/$snaprel) ############" + + # ---- vLLM. Both images mount the whole repo (snapshot symlinks into blobs/ + # resolve) and serve under the mounted path (query /v1/models for the id). + # GB10 image = env-driven entrypoint; stock v0.19.1 = standard `vllm serve` + # arg style. ---- + img="${VIMG[$k]:-$VLLM_IMAGE}" + echo ">>> launching vLLM ($img)" + cleanup + if [ "$img" = "$VLLM_IMAGE" ]; then + VLLM_CID=$(docker run -d --rm --gpus all --network host \ + -e MODEL_PATH="/models/repo/$snaprel" -e HOST=0.0.0.0 -e PORT="$PORT" \ + -e MAX_MODEL_LEN=8192 -e GPU_MEMORY_UTIL=0.85 -e ATTENTION_BACKEND=FLASH_ATTN \ + -v "$repo":/models/repo:ro \ + "$img" --quantization fp8 --max-num-seqs 32 --enable-prefix-caching 2>/dev/null) + else + # --max-num-batched-tokens 8192: required for multimodal Gemma-4 (its per-image + # MM token count 2496 > the default 2048 budget → vLLM refuses to start); + # harmless (a larger prefill batch budget) for the text-only Llama/Mistral. + VLLM_CID=$(docker run -d --rm --gpus all --network host \ + -v "$repo":/models/repo:ro \ + "$img" --model "/models/repo/$snaprel" --port "$PORT" --quantization fp8 \ + --max-model-len 8192 --max-num-seqs 32 --enable-prefix-caching \ + --max-num-batched-tokens 8192 --gpu-memory-utilization 0.85 2>/dev/null) + fi + if wait_ready 900; then + mid=$(served_model_id) + bench/bench.sh --url "http://127.0.0.1:$PORT" --model "$mid" --tag "vllm-$k" \ + --concurrency "$CONC" --out "$RESDIR" >"$RESDIR/vllm-$k.log" 2>&1 + else echo " vLLM never became ready — skipping $k"; docker logs "$VLLM_CID" 2>&1 | tail -20; cleanup; continue; fi + cleanup + + # ---- baseRT ---- + # Port must be free — a stale server on :PORT would make wait_ready pass + # against the WRONG process and mislabel its results. + if curl -s -m 2 "http://127.0.0.1:$PORT/v1/models" 2>/dev/null | grep -q '"id"'; then + echo " :$PORT already serving before baseRT launch — skipping $k"; cleanup; continue; fi + echo ">>> launching baseRT serve" + "$BASERT_SERVE" "$base" --port "$PORT" --continuous-batching 32 --prefix-cache \ + --max-context 8192 >"$RESDIR/baseRT-$k.serve.log" 2>&1 & + BASERT_PID=$! + if wait_ready 300; then + # Confirm CB actually engaged — for models where it isn't supported + # (e.g. hybrid-MoE routes to the serial path) the row is baseRT's real + # serving throughput but NOT continuous batching; label it so the CB-vs- + # serial distinction is auditable rather than silently conflated. + if grep -qiE 'continuous batching .*(disabled|unsupported|serial)|sequence_create.*UNSUPPORTED' \ + "$RESDIR/baseRT-$k.serve.log" 2>/dev/null; then + echo " note: $k baseRT is on the SERIAL path (CB not engaged) — throughput won't scale with concurrency" + fi + bench/bench.sh --url "http://127.0.0.1:$PORT" --model default --tag "baseRT-$k" \ + --concurrency "$CONC" --out "$RESDIR" >"$RESDIR/baseRT-$k.log" 2>&1 + else echo " baseRT never became ready — skipping $k"; cleanup; continue; fi + cleanup + + # ---- compare ---- + # Idempotent per model: drop any prior rows for this key so a rerun REPLACES + # them instead of appending duplicates. Field-based (awk) so the '.' in keys + # like qwen3.5-2b isn't treated as a regex wildcard; the header ($1=model) and + # other models are preserved. + awk -F, -v k="$k" 'NR==1 || $1!=k' "$CSV" > "$CSV.tmp" && mv "$CSV.tmp" "$CSV" + for c in $CONC; do + b=$(ots "$RESDIR/baseRT-$k.jsonl" "tput/c$c") + v=$(ots "$RESDIR/vllm-$k.jsonl" "tput/c$c") + pct=""; [ -n "$b" ] && [ -n "$v" ] && [ "$v" != "0.0" ] && pct=$(python3 -c "print(round(100*$b/$v))") + echo "$k,$c,$b,$v,$pct" | tee -a "$CSV" + done +done + +echo "=== wrote $CSV (regenerate serving-vs-vllm.md from it; this script writes the CSV only) ===" diff --git a/benchmarks/scripts/gb10_two_way.sh b/benchmarks/scripts/gb10_two_way.sh new file mode 100644 index 0000000..660e22f --- /dev/null +++ b/benchmarks/scripts/gb10_two_way.sh @@ -0,0 +1,22 @@ +#!/bin/bash +# GB10 (CUDA) two-way wrapper around three_way_benchmark.sh: +# baseRT (build-cuda/baseRT_cuda_bench) vs llama.cpp CUDA; mlx-lm is not +# available on Linux/CUDA so its leg is disabled via MLX_BENCH=false +# (fails instantly per row and is skipped in results). +# +# Same MODEL_TABLE / protocol as the M4 Pro run: pp 128..2048, tg128, 5 reps. +# Fetch artifacts first: SET=all benchmarks/scripts/fetch_3way_models.sh +set -u +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" + +export BASERT_BIN="${BASERT_BIN:-$REPO_ROOT/build-cuda/baseRT_cuda_bench}" +export LLAMA_BENCH="${LLAMA_BENCH:-$HOME/llama.cpp/build/bin/llama-bench}" +export MLX_BENCH=false +export RESULTS="${RESULTS:-$REPO_ROOT/benchmarks/gb10/bench_2way.csv}" +export SUMMARY="${SUMMARY:-$REPO_ROOT/benchmarks/gb10/bench_2way.md}" +# GB10 unified memory: heavy page cache depresses large-model decode 20-25%. +# Reclaim before every engine invocation when cache exceeds 40GB. +export PRE_RUN_CMD="${PRE_RUN_CMD:-bash $SCRIPT_DIR/reclaim_page_cache.sh 40}" + +exec bash "$SCRIPT_DIR/three_way_benchmark.sh" diff --git a/benchmarks/scripts/gb10_vllm_prefill_decode.sh b/benchmarks/scripts/gb10_vllm_prefill_decode.sh new file mode 100644 index 0000000..aa66442 --- /dev/null +++ b/benchmarks/scripts/gb10_vllm_prefill_decode.sh @@ -0,0 +1,101 @@ +#!/usr/bin/env bash +# vLLM single-stream prefill/decode via `vllm bench latency` (engine-layer, no +# server) — the apples-to-apples counterpart to bench_2way's baseRT/llama.cpp +# pp512 + tg128. batch-size 1. Three latency points isolate each phase by +# CANCELLING the fixed per-call overhead (framework/sampling/detokenize): +# la = latency(in=512, out=1) ~= prefill(512) + 1 decode + overhead +# lb = latency(in=1024, out=1) ~= prefill(1024) + 1 decode + overhead +# lc = latency(in=512, out=128) ~= prefill(512) + 128 decode + overhead +# pp512 = 512 / (lb - la) # marginal cost of the 513..1024 prompt tokens (overhead + 1 decode +# # cancel). This EQUALS pp512-from-scratch only if prefill is ~linear +# # in sequence length; treat it as the near-512 marginal prefill rate. +# tg128 = 127 / (lc - la) # marginal cost of 127 decode tokens (prefill + overhead cancel) +# e2e = 128 / lc # single-request end-to-end gen rate (prefill amortized) +# Dense uses the GB10-optimized image; qwen3_5/qwen3_6 hybrids use stock v0.19.1 +# (GB10 image's Transformers can't load them). vLLM = fp8 (engine-vs-native- +# format vs baseRT/llama.cpp q4/q8, same caveat as the throughput table). +# +# Usage: gb10_vllm_prefill_decode.sh [model_key ...] (GPU must be free) +set -uo pipefail +cd "$(git -C "$(dirname "$0")" rev-parse --show-toplevel)" + +OUT="${OUT:-benchmarks/gb10}" +HF_HUB="${HF_HUB:-/mnt/nas-models/hf-hub}" +VLLM_IMAGE="${VLLM_IMAGE:-hellohal2064/vllm-dgx-spark-gb10:latest}" +VLLM_IMAGE_HYBRID="${VLLM_IMAGE_HYBRID:-vllm/vllm-openai:v0.19.1-cu130}" +VLLM_IMAGE_STOCK="${VLLM_IMAGE_STOCK:-vllm/vllm-openai:v0.20.0-cu130}" +INLEN="${INLEN:-512}"; OUTLEN="${OUTLEN:-128}" +# Absolute path — docker -v rejects a relative host path (reads it as a named volume). +mkdir -p "$OUT/serving_raw/vllm_lat" +JDIR="$(cd "$OUT/serving_raw/vllm_lat" && pwd)" + +declare -A HFREPO=( + [gemma-3-1b]="$HF_HUB/models--unsloth--gemma-3-1b-it" + [gemma-4-e2b]="$HF_HUB/models--google--gemma-4-E2B-it" + [gemma-4-26b]="$HF_HUB/models--google--gemma-4-26B-A4B-it" + [llama-3.2-1b]="$HF_HUB/models--meta-llama--Llama-3.2-1B-Instruct" + [llama-3.2-3b]="$HF_HUB/models--meta-llama--Llama-3.2-3B-Instruct" + [qwen3-0.6b]="$HF_HUB/models--Qwen--Qwen3-0.6B" + [qwen3-30b-a3b]="$HF_HUB/models--Qwen--Qwen3-30B-A3B-Instruct-2507" + [qwen3.5-2b]="$HF_HUB/models--Qwen--Qwen3.5-2B" [qwen3.5-35b-a3b]="$HF_HUB/models--Qwen--Qwen3.5-35B-A3B" + [qwen3.6-27b]="$HF_HUB/models--Qwen--Qwen3.6-27B" [qwen3.6-35b]="$HF_HUB/models--Qwen--Qwen3.6-35B-A3B" +) +# bench latency OVERRIDES the entrypoint, so the GB10 image's baked qwen parsers +# don't apply here (unlike the serving script) — but gemma-4 still needs the 0.20 +# image for the arch. Route Gemma/Llama to stock 0.20 for consistency with the +# serving run's vLLM side; dense/MoE Qwen use the default GB10 image. +declare -A VIMG=( + [qwen3.5-2b]="$VLLM_IMAGE_HYBRID" [qwen3.5-35b-a3b]="$VLLM_IMAGE_HYBRID" + [qwen3.6-27b]="$VLLM_IMAGE_HYBRID" [qwen3.6-35b]="$VLLM_IMAGE_HYBRID" + [gemma-3-1b]="$VLLM_IMAGE_STOCK" [gemma-4-e2b]="$VLLM_IMAGE_STOCK" + [gemma-4-26b]="$VLLM_IMAGE_STOCK" + [llama-3.2-1b]="$VLLM_IMAGE_STOCK" [llama-3.2-3b]="$VLLM_IMAGE_STOCK" +) +KEYS=("$@"); [ ${#KEYS[@]} -eq 0 ] && KEYS=(gemma-3-1b gemma-4-e2b gemma-4-26b \ + llama-3.2-1b llama-3.2-3b qwen3-0.6b qwen3-30b-a3b \ + qwen3.5-2b qwen3.5-35b-a3b qwen3.6-27b qwen3.6-35b) + +# run vllm bench latency once; echo avg_latency seconds (empty on failure) +lat() { # key repo snaprel inlen outlen -> avg_latency seconds + local img="${VIMG[$1]:-$VLLM_IMAGE}" + local tag="in$4.out$5" + local jf="$JDIR/$1.$tag.json" + rm -f "$jf" # discard any stale JSON so a failed docker run (OOM/bad image) reads empty, not last run's value + docker run --rm --gpus all --network host -v "$2":/models/repo:ro -v "$JDIR":/out \ + --entrypoint vllm "$img" bench latency \ + --model "/models/repo/$3" --quantization fp8 --max-model-len 8192 \ + --max-num-batched-tokens 8192 \ + --input-len "$4" --output-len "$5" --batch-size 1 \ + --num-iters-warmup 3 --num-iters 5 --output-json "/out/$1.$tag.json" >"$JDIR/$1.$tag.log" 2>&1 + python3 -c "import json;print(json.load(open('$jf')).get('avg_latency',''))" 2>/dev/null +} + +CSV="$OUT/vllm-prefill-decode.csv" +HDR="model,vllm_prefill_pp${INLEN}_tok_s,vllm_decode_tg${OUTLEN}_tok_s,vllm_e2e_tg${OUTLEN}_tok_s" +# Recreate the file when absent OR when its header doesn't match these INLEN/ +# OUTLEN — an override must not silently mix, e.g., tg64 values under a tg128 +# header (the column names encode the lengths). Same lengths → append-safe reuse. +[ -f "$CSV" ] && [ "$(head -1 "$CSV")" = "$HDR" ] || echo "$HDR" > "$CSV" +for k in "${KEYS[@]}"; do + repo="${HFREPO[$k]:-}"; [ -d "$repo/snapshots" ] || { echo "::skip $k (no HF repo)"; continue; } + snap="snapshots/$( { cat "$repo/refs/main" 2>/dev/null || ls -t "$repo/snapshots" | head -1; } )" + echo ">>> $k: vllm bench latency (in$INLEN/out1, in$((INLEN*2))/out1, in$INLEN/out$OUTLEN)" + la=$(lat "$k" "$repo" "$snap" "$INLEN" 1) # prefill INLEN + 1 decode + overhead + lb=$(lat "$k" "$repo" "$snap" $((INLEN*2)) 1) # prefill 2*INLEN + 1 decode + overhead + lc=$(lat "$k" "$repo" "$snap" "$INLEN" "$OUTLEN") # prefill INLEN + OUTLEN decode + overhead + pp=""; tg=""; e2e="" + # CLEAN prefill: (lb-la) is the marginal cost of INLEN extra prompt tokens, so + # the fixed per-call overhead AND the shared 1 decode cancel -> pure prefill. + [ -n "$la" ] && [ -n "$lb" ] && awk "BEGIN{exit !($lb>$la)}" && \ + pp=$(python3 -c "print(round($INLEN/($lb-$la),1))" 2>/dev/null) + [ -n "$la" ] && [ -n "$lc" ] && awk "BEGIN{exit !($lc>$la)}" && \ + tg=$(python3 -c "print(round(($OUTLEN-1)/($lc-$la),1))" 2>/dev/null) + [ -n "$lc" ] && e2e=$(python3 -c "print(round($OUTLEN/$lc,1))" 2>/dev/null) + # Replace this model's row ONLY after a real measurement — if all three latency + # runs failed (docker OOM, missing image, absent JSON) la/lb/lc are empty and we + # keep the previously-collected row rather than overwrite it with an all-blank one. + if [ -z "$la" ] && [ -z "$lb" ] && [ -z "$lc" ]; then echo "::skip $k (all latency runs failed — keeping prior row)"; continue; fi + awk -F, -v k="$k" 'NR==1 || $1!=k' "$CSV" > "$CSV.tmp" && mv "$CSV.tmp" "$CSV" + echo "$k,$pp,$tg,$e2e" | tee -a "$CSV" +done +echo "=== wrote $CSV ===" diff --git a/benchmarks/scripts/reclaim_page_cache.sh b/benchmarks/scripts/reclaim_page_cache.sh new file mode 100644 index 0000000..3bb9029 --- /dev/null +++ b/benchmarks/scripts/reclaim_page_cache.sh @@ -0,0 +1,27 @@ +#!/usr/bin/env bash +# Force page-cache reclaim when buff/cache is high. On unified-memory boxes +# (GB10) heavy page cache depresses large-model decode 20-25% — measured +# 2026-07-16: Qwen3-30B q6 tg128 57.4 with 114GB cache vs 74.0 clean, same +# binary + file. Used as PRE_RUN_CMD by gb10_two_way.sh. No sudo needed: +# allocating + touching anonymous memory forces the kernel to drop clean +# cache pages. Skips when cache is already low (threshold GB, default 40). +THRESH_GB="${1:-40}" +cache_kb=$(awk '/^Cached:/{print $2}' /proc/meminfo) +if [ "$cache_kb" -lt $((THRESH_GB*1024*1024)) ]; then exit 0; fi +avail_kb=$(awk '/MemAvailable/{print $2}' /proc/meminfo) +target_gb=$(( avail_kb/1024/1024 - 10 )) +[ "$target_gb" -gt 100 ] && target_gb=100 +[ "$target_gb" -lt 8 ] && exit 0 +echo " (reclaiming page cache: $((cache_kb/1024/1024))GB cached, touching ${target_gb}GB)" +python3 - "$target_gb" <<'PY' +import ctypes, sys +n = int(sys.argv[1]) // 2 +chunks = [] +try: + for _ in range(n): + b = ctypes.create_string_buffer(2*1024**3) + ctypes.memset(b, 1, 2*1024**3) + chunks.append(b) +except MemoryError: + pass +PY diff --git a/benchmarks/scripts/three_way_benchmark.sh b/benchmarks/scripts/three_way_benchmark.sh index 4d92954..ab3e1e2 100755 --- a/benchmarks/scripts/three_way_benchmark.sh +++ b/benchmarks/scripts/three_way_benchmark.sh @@ -5,6 +5,14 @@ # - Pairs each model across the three formats (baseRT .base, GGUF, MLX HF repo). # - Writes a flat CSV and a markdown summary. # +# Reproduce (e.g. on an M3): fetch the artifacts from pinned public sources, +# then run a subset. baseRT .base come from the basecompute HF catalog, GGUFs +# from bartowski/ggml-org, MLX repos auto-pull on first use: +# SET=q8 benchmarks/scripts/fetch_3way_models.sh # q8 only; SET=all (default) is huge +# MODELS="Llama-3.2-1B-Q8 Llama-3.2-3B-Q8 Qwen3-0.6B-Q8 Gemma-4-E2B-Q8" \ +# RESULTS=/tmp/q8.csv SUMMARY=/tmp/q8.md benchmarks/scripts/three_way_benchmark.sh +# MoE Q8 (Qwen3-30B/Gemma-4-26B) needs >24GB and OOMs on M3/M4 Pro — dense only. +# # Env knobs: # BASERT_BIN : path to basert-bench (default: ../../build/basert-bench) # MODELS_DIR : local model directory (default: ../../models) @@ -28,7 +36,7 @@ MODELS_DIR="${MODELS_DIR:-$REPO_ROOT/models}" LLAMA_BENCH="${LLAMA_BENCH:-llama-bench}" # Default MLX wrapper: assumes uv is available and pulls mlx-lm in a throwaway env. # Set MLX_BENCH=mlx_lm.benchmark if you have it on PATH. -MLX_BENCH="${MLX_BENCH:-uv run --with mlx-lm --with mlx --no-project mlx_lm.benchmark}" +MLX_BENCH="${MLX_BENCH:-uv run --python 3.12 --with mlx-lm --with mlx --no-project mlx_lm.benchmark}" PP_VALS="${PP_VALS:-128 256 512 1024 2048}" TG_VAL="${TG_VAL:-128}" @@ -37,18 +45,63 @@ RESULTS="${RESULTS:-/tmp/bench_3way.csv}" SUMMARY="${SUMMARY:-/tmp/bench_3way.md}" COOLDOWN="${COOLDOWN:-20}" +# KV-cache precision, pinned identically on baseRT and llama.cpp so the decode +# comparison is like-for-like. baseRT's Auto resolves to Q8_0 and llama-bench +# defaults to f16, which is not a matched comparison; standardize on f16 (16). +# Verified on M1 Max that baseRT decode is KV-precision-insensitive (Q8 vs f16 +# tg128 within noise even at pp2048), so this changes fairness, not baseRT's +# numbers. (16=f16, 8=Q8_0, 4=Q4_0.) mlx-lm uses its own default (f16). +KV_BITS="${KV_BITS:-16}" +case "$KV_BITS" in + 16) LLAMA_CACHE_TYPE="f16" ;; + 8) LLAMA_CACHE_TYPE="q8_0" ;; + 4) LLAMA_CACHE_TYPE="q4_0" ;; + *) echo "unknown KV_BITS=$KV_BITS (use 16|8|4)"; exit 1 ;; +esac + # MODEL_TABLE entries: label|baseRT_filename|gguf_filename|mlx_hf_repo # Each component may be "-" to skip that engine for the model. MODEL_TABLE=( "Llama-3.2-1B-Q4|Llama-3.2-1B-Q4.base|Llama-3.2-1B-Instruct-Q4_0.gguf|mlx-community/Llama-3.2-1B-Instruct-4bit" "Llama-3.2-3B-Q4|Llama-3.2-3B-Q4.base|Llama-3.2-3B-Instruct-Q4_0.gguf|mlx-community/Llama-3.2-3B-Instruct-4bit" "Qwen3-0.6B-Q4|Qwen3-0.6B-Q4_0.base|Qwen3-0.6B-Q4_0.gguf|mlx-community/Qwen3-0.6B-4bit" - "Gemma-3-1B-Q4|gemma-3-1b-it-Q4_K_M.base|gemma-3-1b-it-Q4_K_M.gguf|mlx-community/gemma-3-1b-it-4bit" - "Gemma-4-E2B-Q4|gemma-4-E2B-it-Q4_0.base|gemma-4-E2B-it-Q4_0.gguf|mlx-community/gemma-3n-E2B-it-4bit" - # MoE rows. llama.cpp side: Qwen has Q4_0 from bartowski; Gemma's only - # widely-available 24GB (m5 / m1-max-64G class); 2B runs everywhere. + "Qwen3.5-2B-Q4|Qwen3.5-2B-Base-Q4.base|gguf/Qwen3.5-2B/Qwen_Qwen3.5-2B-Q4_K_M.gguf|mlx-community/Qwen3.5-2B-4bit" + "Qwen3.5-2B-Q8|Qwen3.5-2B-Base-Q8.base|gguf/Qwen3.5-2B/Qwen_Qwen3.5-2B-Q8_0.gguf|mlx-community/Qwen3.5-2B-8bit" + "Qwen3.5-35B-A3B-Q4|Qwen3.5-35B-A3B-Q4.base|gguf/Qwen3.5-35B/Qwen_Qwen3.5-35B-A3B-Q4_K_M.gguf|mlx-community/Qwen3.5-35B-A3B-4bit" + "Qwen3.6-27B-Q4|Qwen3.6-27B-Q4.base|gguf/Qwen3.6-27B/Qwen3.6-27B-Q4_K_M.gguf|mlx-community/Qwen3.6-27B-4bit" + "Qwen3.6-35B-Q4|Qwen3.6-35B-A3B-Q4.base|gguf/Qwen3.6-35B/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf|mlx-community/Qwen3.6-35B-A3B-4bit" + "Qwen3.6-27B-Q8|Qwen3.6-27B-Q8.base|gguf/Qwen3.6-27B/Qwen3.6-27B-Q8_0.gguf|mlx-community/Qwen3.6-27B-8bit" + "Gemma-4-26B-A4B-Q4|Gemma-4-26B-A4B-Q4.base|google_gemma-4-26B-A4B-it-Q4_0.gguf|mlx-community/gemma-4-26b-a4b-it-4bit" + + # ── Q8 (8-bit) rows ──────────────────────────────────────────────────── + # The MoE Q8 weights (Qwen3-30B-A3B ~32GB, Gemma-4-26B-A4B ~28GB) exceed the + # 24GB working set on an M4 Pro and OOM at inference — those two rows only + # run on big-memory boxes (M1 Max 64GB, GB10 121GB); elsewhere the missing + # files skip them. + # baseRT .base from the basecompute catalog (Q8); llama.cpp Q8_0 GGUF from + # bartowski / ggml-org; mlx-community 8-bit repos. + "Llama-3.2-1B-Q8|Llama-3.2-1B-Q8.base|Llama-3.2-1B-Instruct-Q8_0.gguf|mlx-community/Llama-3.2-1B-Instruct-8bit" + "Llama-3.2-3B-Q8|Llama-3.2-3B-Q8.base|Llama-3.2-3B-Instruct-Q8_0.gguf|mlx-community/Llama-3.2-3B-Instruct-8bit" + "Qwen3-0.6B-Q8|Qwen3-0.6B-Q8.base|Qwen3-0.6B-Q8_0.gguf|mlx-community/Qwen3-0.6B-8bit" +# "Gemma-3-1B-Q8|gemma-3-1b-it-Q8.base|gemma-3-1b-it-Q8_0.gguf|mlx-community/gemma-3-1b-it-8bit" + "Gemma-4-E2B-Q8|gemma-4-E2B-it-Q8.base|unsloth/gemma-4-E2B-it-GGUF/gemma-4-E2B-it-Q8_0.gguf|mlx-community/gemma-4-e2b-it-8bit" + # MoE Q8 (big-memory boxes only — see note above). The Qwen 30B Q8 row uses + # the Thinking-2507 variant across baseRT + GGUF (only published Q8 for this + # architecture); mlx pairs the same-arch Instruct-2507 8-bit repo. + "Qwen3-30B-A3B-Q8|Qwen3-30B-A3B-Q8.base|Qwen_Qwen3-30B-A3B-Thinking-2507-Q8_0.gguf|mlx-community/Qwen3-30B-A3B-Instruct-2507-8bit" + "Gemma-4-26B-A4B-Q8|Gemma-4-26B-A4B-Q8.base|google_gemma-4-26B-A4B-it-Q8_0.gguf|mlx-community/gemma-4-26b-a4b-it-8bit" ) # Optional subset: e.g. MODELS="Qwen3-0.6B-Q4 Llama-3.2-1B-Q4" @@ -60,7 +113,11 @@ echo "model,engine,test,tok_per_sec,stddev" > "$RESULTS" # ── parsers (Python — robust to BSD/GNU awk + pipe-table edge cases) ─────── # Engine outputs are written to a temp file; Python reads via env path so the # heredoc (which IS python3's stdin source) doesn't fight with the data. -PARSER_PY="$(mktemp /tmp/three_way_parser.XXXXXX.py)" +# NB: the XXXXXX placeholder must be the trailing component — macOS/BSD mktemp +# does not substitute X's when a suffix (.py) follows, so it would create a +# literal, non-unique file that a killed run leaves behind and breaks the next +# run. python3 runs a file regardless of extension, so no suffix is needed. +PARSER_PY="$(mktemp "${TMPDIR:-/tmp}/three_way_parser.XXXXXX")" cat > "$PARSER_PY" <<'PY' import os, re, sys, statistics engine, label, tg, results_path, output_path = sys.argv[1:6] @@ -123,8 +180,9 @@ parse_mlx() { MLX_PP_TAG="$3" emit_rows "mlx-lm" "$2" "$TG_VAL" <<< "$1"; } run_baseRT() { local model_path="$1" label="$2" pp="$3" tg="$4" echo " baseRT pp${pp}/tg${tg}..." + ${PRE_RUN_CMD:-true} local out - out=$("$BASERT_BIN" "$model_path" -p "$pp" -n "$tg" -r "$REPS" 2>&1) || { + out=$(BASERT_KV_BITS="$KV_BITS" "$BASERT_BIN" "$model_path" -p "$pp" -n "$tg" -r "$REPS" 2>&1) || { echo " ERROR: basert-bench failed"; return } parse_baseRT "$out" "$label" @@ -133,8 +191,10 @@ run_baseRT() { run_llama() { local model_path="$1" label="$2" pp="$3" tg="$4" echo " llama.cpp pp${pp}/tg${tg}..." + ${PRE_RUN_CMD:-true} local out - out=$("$LLAMA_BENCH" -m "$model_path" -p "$pp" -n "$tg" -r "$REPS" 2>&1) || { + out=$("$LLAMA_BENCH" -m "$model_path" -p "$pp" -n "$tg" -r "$REPS" \ + -ctk "$LLAMA_CACHE_TYPE" -ctv "$LLAMA_CACHE_TYPE" 2>&1) || { echo " ERROR: llama-bench failed"; return } parse_llama "$out" "$label" @@ -144,7 +204,7 @@ run_mlx() { local repo="$1" label="$2" pp="$3" tg="$4" echo " mlx-lm pp${pp}/tg${tg}..." local out - out=$($MLX_BENCH --model "$repo" -p "$pp" -g "$tg" -n "$REPS" 2>&1) || { + out=$($MLX_BENCH --model "$repo" --prompt-tokens "$pp" --generation-tokens "$tg" --num-trials "$REPS" 2>&1) || { echo " ERROR: mlx_lm.benchmark failed"; return } parse_mlx "$out" "$label" "$pp" @@ -191,7 +251,7 @@ for row in "${MODEL_TABLE[@]}"; do done # ── markdown summary ─────────────────────────────────────────────────────── -SUMMARY_PY="$(mktemp /tmp/three_way_summary.XXXXXX.py)" +SUMMARY_PY="$(mktemp "${TMPDIR:-/tmp}/three_way_summary.XXXXXX")" cat > "$SUMMARY_PY" <<'PY' import csv, sys, collections, statistics, platform, subprocess results_path, summary_path = sys.argv[1], sys.argv[2] diff --git a/bindings/node/src/index.ts b/bindings/node/src/index.ts index 6c30ad5..9742a48 100644 --- a/bindings/node/src/index.ts +++ b/bindings/node/src/index.ts @@ -313,6 +313,10 @@ interface BaseRTLib { baseRT_set_kv_bits: (bits: number) => void; baseRT_set_paged_kv: (enable: number) => void; baseRT_set_max_batch_size: (n: number) => void; + baseRT_set_prefill_chunk: (n: number) => void; + baseRT_set_paged_weights: (mode: number) => void; + baseRT_set_baked_decode: (enable: number) => void; + baseRT_set_gpu_wait_timeout_ms: (ms: number) => void; baseRT_encode: ( model: unknown, text: string, @@ -444,6 +448,10 @@ function lib(): BaseRTLib { baseRT_set_kv_bits: k.func("void baseRT_set_kv_bits(int)"), baseRT_set_paged_kv: k.func("void baseRT_set_paged_kv(int)"), baseRT_set_max_batch_size: k.func("void baseRT_set_max_batch_size(int)"), + baseRT_set_prefill_chunk: k.func("void baseRT_set_prefill_chunk(int)"), + baseRT_set_paged_weights: k.func("void baseRT_set_paged_weights(int)"), + baseRT_set_baked_decode: k.func("void baseRT_set_baked_decode(int)"), + baseRT_set_gpu_wait_timeout_ms: k.func("void baseRT_set_gpu_wait_timeout_ms(double)"), baseRT_encode: k.func( "int baseRT_encode(void *, const char *, _Out_ uint32_t *, int)" ), @@ -555,6 +563,22 @@ export const Engine = { setMaxBatchSize(n: number): void { lib().baseRT_set_max_batch_size(n); }, + /** Override prefill chunk size (tokens/GEMM batch); 0 = per-chip default. */ + setPrefillChunk(n: number): void { + lib().baseRT_set_prefill_chunk(n); + }, + /** Paged-weights load policy: 0=auto, 1=force, 2=no-retry (issue #113). */ + setPagedWeights(mode: number): void { + lib().baseRT_set_paged_weights(mode); + }, + /** Toggle the baked-decode fast path: true=on (default), false=force live. */ + setBakedDecode(enable: boolean): void { + lib().baseRT_set_baked_decode(enable ? 1 : 0); + }, + /** GPU-wait timeout (ms); 0 = block indefinitely (default). Metal only. */ + setGpuWaitTimeoutMs(ms: number): void { + lib().baseRT_set_gpu_wait_timeout_ms(ms); + }, /** Last thread-local error string, or null. */ lastError(): string | null { return lib().baseRT_get_error() || null; diff --git a/bindings/python/baseRT/__init__.py b/bindings/python/baseRT/__init__.py index 5a43eb4..3d5ae8f 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.1.0" +__version__ = "0.2.0" # --------------------------------------------------------------------------- # Library loading diff --git a/bindings/python/setup.py b/bindings/python/setup.py index e48e6d9..83aadfc 100644 --- a/bindings/python/setup.py +++ b/bindings/python/setup.py @@ -6,7 +6,7 @@ setup( name="baseRT", - version="0.1.0", + version="0.2.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/swift/Sources/CBaseRT/include/baseRT.h b/bindings/swift/Sources/CBaseRT/include/baseRT.h index af8d477..de6344d 100644 --- a/bindings/swift/Sources/CBaseRT/include/baseRT.h +++ b/bindings/swift/Sources/CBaseRT/include/baseRT.h @@ -63,14 +63,14 @@ extern "C" { // === Versioning === #define BASERT_VERSION_MAJOR 0 -#define BASERT_VERSION_MINOR 1 -#define BASERT_VERSION_PATCH 7 +#define BASERT_VERSION_MINOR 2 +#define BASERT_VERSION_PATCH 0 /// Compile-time version, packed as `(MAJOR<<16) | (MINOR<<8) | PATCH`. /// Useful for `#if BASERT_VERSION >= 0x000200` feature checks. #define BASERT_VERSION ((BASERT_VERSION_MAJOR << 16) | (BASERT_VERSION_MINOR << 8) | BASERT_VERSION_PATCH) -/// Runtime-resolved version string ("0.1.0"). Matches the linked +/// Runtime-resolved version string ("0.2.0"). Matches the linked /// library; useful for diagnostics when a binding loads a different /// `.dylib` than it was compiled against. const char *baseRT_version_string(void); @@ -131,6 +131,39 @@ void baseRT_set_max_batch_size(int n); /// called before baseRT_load_model. Drive it via the baseRT_prefix_* API. void baseRT_set_prefix_cache(int enable); +/// Override the prefill chunk size (tokens per prefill GEMM batch). +/// n = 0 → per-chip default (recommended) +/// n >= 16 → clamp to this; shrinks the batch_* scratch footprint +/// (~linearly) so an oversized model fits a tighter GPU working-set +/// budget, at a prefill-throughput cost. Values outside [16, chip max] +/// are ignored. Process-wide; read at load. Set before baseRT_load_model. +void baseRT_set_prefill_chunk(int n); + +/// Paged-weights load policy (issue #113): per-tensor pinned weight buffers +/// for models past the OS wired-page budget. +/// mode = 0 → auto (default): normal load, retry paged on GPU-OOM +/// mode = 1 → force paged-weights on the first attempt (oversized/testing) +/// mode = 2 → disable the retry (fail hard on OOM, pre-#113 behavior) +/// Process-wide; read at load. Set before baseRT_load_model. +void baseRT_set_paged_weights(int mode); + +/// Toggle the baked-decode fast path (DispatchTable replay). +/// enable = 1 → on (default): replay the baked table when eligible +/// enable = 0 → force the live fused decode path (baked-vs-live A/B) +/// Process-wide; read per decode step. +void baseRT_set_baked_decode(int enable); + +/// GPU-wait timeout in milliseconds (Metal backend). +/// ms > 0 → (default: 300000, i.e. 5 min) return a loggable error if a +/// committed command buffer doesn't reach a terminal status in time. The +/// CPU unblocks; the wedged GPU work is NOT torn down (only a GPU +/// reset/reboot reclaims it). Chosen far above any legitimate single +/// command buffer so only a real wedge trips it. +/// ms = 0 → block indefinitely in the GPU wait (opt out of the bound). +/// No-op on non-Metal backends. Process-wide; read once per wait. The initial +/// value also honors the BASERT_GPU_WAIT_TIMEOUT_MS environment variable. +void baseRT_set_gpu_wait_timeout_ms(double ms); + /// Free all resources associated with a model. void baseRT_free_model(baseRT_model_t model); @@ -202,6 +235,26 @@ typedef bool (*baseRT_token_callback)(uint32_t token_id, const char *text, void BaseRTGenerationStats baseRT_generate(baseRT_model_t model, const uint32_t *prompt_tokens, int n_prompt, int max_tokens, BaseRTSamplingConfig sampling, baseRT_token_callback callback, void *user_data); +/// Generate tokens from a prompt WITH serial RadixCache prefix reuse. +/// +/// Semantically identical to `baseRT_generate` (single default sequence, same +/// greedy/sampled output), but when the model was loaded with BOTH `--paged-kv` +/// and `--prefix-cache` it reuses the longest cached whole-block prompt prefix: +/// it resets the default sequence, matches the prompt against the RadixCache, +/// seeds the shared prefix blocks into the default sequence, prefills ONLY the +/// divergent suffix `[matched_tokens, n_prompt)`, then inserts the full prompt +/// back into the cache for later reuse on generation end. +/// +/// Greedy output is BIT-IDENTICAL to a cold `baseRT_generate` (the seeded blocks +/// hold the same KV a cold prefill would have produced). When the prefix cache +/// or paged-KV is disabled this is an exact passthrough to `baseRT_generate`. +/// +/// Attention-KV (non-hybrid) models only — hybrid-GDN models fall back to the +/// plain path (their recurrent state is not block-shareable). +BaseRTGenerationStats baseRT_generate_cached(baseRT_model_t model, const uint32_t *prompt_tokens, int n_prompt, + int max_tokens, BaseRTSamplingConfig sampling, + baseRT_token_callback callback, void *user_data); + // === Multi-sequence generation (paged-KV only) === /// Opaque per-sequence handle. @@ -238,6 +291,27 @@ BaseRTGenerationStats baseRT_sequence_generate_continue(baseRT_sequence_t seq, c /// Release the sequence's blocks back to the pool and free the handle. void baseRT_sequence_free(baseRT_sequence_t seq); +/// Fused batch step with per-sequence trailing PAD counts (serving grid +/// alignment): counts[i] includes pads[i] throwaway tokens whose rows are +/// computed but whose argmax row is skipped (the output token comes from the +/// last REAL row). The caller must roll each padded sequence's KV back by +/// pads[i] after the call (baseRT_sequence_rollback). Greedy only. +int baseRT_batch_step_fused_pads(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, const uint32_t *in_tokens, + const int *in_token_counts, const int *pads, uint32_t *out_tokens); + +/// Warm the batched-decode fast paths for batch sizes up to `max_batch`: +/// each B runs two throwaway pure-decode ticks so shape-keyed caches (baked +/// dispatch tables, stream-captured CUDA graphs, cuBLAS plans) are built at +/// startup instead of on the first real requests — the same boot-time graph +/// warmup vLLM performs. Requires --paged-kv; call after load, before serving. +int baseRT_batch_warmup(baseRT_model_t model, int max_batch); + +/// 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). +int baseRT_sequence_rollback(baseRT_sequence_t seq, int length); + /// Batched decode: drives ONE batched decode step across N sequences. Each /// sequence writes its `new_tokens[i]` to its own KV cache slot, and attention /// reads each sequence's KV via its own block table. Throughput comes from @@ -330,6 +404,37 @@ int baseRT_batch_step(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, int baseRT_batch_step_fused(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, const uint32_t *in_tokens, const int *in_token_counts, uint32_t *out_tokens); +/// Fused batch step with per-sequence trailing PAD counts (serving grid +/// alignment): counts[i] includes pads[i] throwaway tokens whose rows are +/// computed but whose argmax row is skipped (the output token comes from the +/// last REAL row). The caller must roll each padded sequence's KV back by +/// pads[i] after the call (baseRT_sequence_rollback). Greedy only. +/// +/// Backend note: the shape-padding fast path exists for CUDA-graph capture; +/// on backends without stream capture (Metal) the pads are STRIPPED before +/// dispatch — semantics identical (no pad KV is written, so the caller's +/// rollback is a no-op), no wasted compute. +int baseRT_batch_step_fused_pads(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, const uint32_t *in_tokens, + const int *in_token_counts, const int *pads, uint32_t *out_tokens); + +/// Warm the batched-decode fast paths for batch sizes up to `max_batch`: +/// each B runs throwaway pure-decode ticks so shape-keyed caches (baked +/// dispatch tables, captured graphs, PSO/plan builds) are built at startup +/// instead of on the first real requests — the same boot-time warmup vLLM +/// performs. Requires --paged-kv; call after load, before serving. +int baseRT_batch_warmup(baseRT_model_t model, int max_batch); + +/// Max prompt tokens the fused (varlen) prefill can process in one packed batch. +/// The continuous-batching engine caps per-tick admitted prompt tokens by this +/// so a burst of long prompts doesn't overflow the packed prefill. 0 on null. +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). +int baseRT_sequence_rollback(baseRT_sequence_t seq, int length); + /// Multi-step autoregressive driver for baseRT_batch_step_fused. Step 0 /// ingests the mixed-length input from `first_in_tokens` / `first_in_token_counts` /// (one row per seq, total length `sum(first_in_token_counts)`). Subsequent @@ -406,6 +511,13 @@ BaseRTPrefixMatch baseRT_prefix_match(baseRT_model_t model, const uint32_t *toke /// paged / the sequence isn't empty. int baseRT_sequence_seed_prefix(baseRT_sequence_t seq, const int *blocks, int n_blocks, int n_tokens); +/// Paged-KV block (page) size in tokens for this model, or 0 when the model +/// was not loaded with --paged-kv. This is the block-alignment granularity for +/// baseRT_prefix_match/_seed_prefix (matched_tokens == matched_blocks * +/// page_size); the continuous-batching hybrid-GDN prefix-reuse path uses it to +/// pick the block-aligned GDN snapshot boundary at admit. +int baseRT_page_size(baseRT_model_t model); + /// Publish `seq`'s KV blocks for the block-aligned prefix of `tokens` into the /// prefix cache so later requests can reuse them. Idempotent for an already- /// cached prefix (no double refcount). No-op when the cache is disabled. @@ -415,8 +527,20 @@ int baseRT_prefix_insert(baseRT_model_t model, const uint32_t *tokens, int n_tok /// Release the lock a baseRT_prefix_match took on a prefix and free the match's /// bookkeeping. Call exactly once per non-zero handle, after the sequence that /// reused the prefix has been inserted/retired. No-op for handle==0. +/// +/// Use this ONLY when the match's blocks WERE seeded into a sequence +/// (baseRT_sequence_seed_prefix): the sequence owns those blocks and drops the +/// match's ownership incref when it resets/frees. If the match was NOT seeded +/// (you decided not to reuse it), call baseRT_prefix_release instead — unlock +/// alone would leak the increfed blocks. void baseRT_prefix_unlock(baseRT_model_t model, uint64_t handle); +/// Abandon a baseRT_prefix_match WITHOUT seeding it: drops the ownership incref +/// on each matched block (which no sequence adopted) AND releases the trie lock. +/// Call exactly once per non-zero handle when you matched a prefix but chose not +/// to seed it (e.g. a boundary mismatch). No-op for handle==0. +void baseRT_prefix_release(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 @@ -473,6 +597,25 @@ void baseRT_grammar_free(baseRT_grammar_t grammar); /// would feed the second sample through a terminated grammar (garbage). void baseRT_grammar_reset(baseRT_grammar_t grammar); +/// Grammar stepping for the continuous-batching server (xgrammar backend). +/// The server applies the bitmask to a lane's logits row on the host, then +/// accepts the sampled token to advance the grammar. A legacy-NPDA grammar +/// reports `bitmask_size == 0` — the caller must keep it on the serial path. +/// 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. +/// 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. +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); +int baseRT_grammar_is_terminated(baseRT_grammar_t grammar); +int baseRT_grammar_is_completed(baseRT_grammar_t grammar); + /// Generate tokens with grammar constraint. /// Grammar masks invalid tokens at each step, guaranteeing output conforms to the grammar. BaseRTGenerationStats baseRT_generate_grammar(baseRT_model_t model, const uint32_t *prompt_tokens, int n_prompt, @@ -632,11 +775,16 @@ void baseRT_reset(baseRT_model_t model); /// `current_length` prefix (not the unused tail), so the file size grows /// linearly with how much was prefilled+decoded. Returns 0 on success and /// a negative error code on failure; check `baseRT_get_error` for details. +/// Hybrid linear-attention models (Qwen 3.5/3.6) are REJECTED: the format +/// holds attention KV only, not the Gated-DeltaNet recurrent state. int baseRT_save_state(baseRT_model_t model, const char *path); /// Inverse of `baseRT_save_state`. The cache must have been allocated for /// a model with matching shape; mismatched files are rejected. After load, /// `baseRT_get_position` reflects the restored token count. +/// Hybrid linear-attention models (Qwen 3.5/3.6) are REJECTED: the file +/// holds attention KV only, and restoring it without the matching +/// Gated-DeltaNet recurrent state would yield a corrupt hybrid state. int baseRT_load_state(baseRT_model_t model, const char *path); /// Install a LoRA adapter on this model. The adapter file is a `.base` @@ -664,8 +812,80 @@ const char *baseRT_lora_id(baseRT_model_t model); /// shared chat-template prefix from a prior request — keeps the cached /// prefill of the common prefix while discarding the prior turn's /// user-message tail and assistant reply. +/// Hybrid linear-attention models (Qwen 3.5/3.6): the recurrent state +/// cannot be rewound to an arbitrary position. This call keeps its "KV +/// length == to_position" promise only when `to_position` exactly matches +/// the recurrent-state snapshot (see `baseRT_set_prefill_snapshot`); any +/// other position degrades to a FULL reset (equivalent to `baseRT_reset`) +/// — the caller must then prefill the entire prompt again. Use +/// `baseRT_try_rollback` to detect what happened, or to resume from a +/// snapshot that sits before the requested position. void baseRT_rollback(baseRT_model_t model, int to_position); +/// Rollback that reports the position actually achieved. Non-hybrid +/// models land on `min(to_position, current KV length)` — a target past +/// the cache end cannot be "achieved" by a rollback and is clamped so +/// callers prefilling from the returned position never skip tokens. Hybrid linear-attention models can only resume +/// from their recurrent-state snapshot (see +/// `baseRT_set_prefill_snapshot`): when the snapshot sits at or before +/// `to_position` the state is restored there and the SNAPSHOT position is +/// returned — the caller must prefill the prompt from that position +/// onward. When the snapshot lies past `to_position` (divergent history) +/// the call returns -1 and leaves the model state UNTOUCHED — fall back +/// to `baseRT_reset` + a full prefill. `to_position == 0` always succeeds +/// as a full reset. +int baseRT_try_rollback(baseRT_model_t model, int to_position); + +/// Hybrid linear-attention models only (no-op otherwise): ask prompt +/// prefills to capture the reuse snapshot once absolute KV position `pos` +/// has been processed, instead of at the prompt end. Chat servers pass +/// the rendered-history boundary (the prompt minus the generation +/// scaffold): the scaffold tokens never reappear in the next request's +/// render, so a prompt-end snapshot would never match, while the history +/// boundary is exactly where the next request's shared prefix ends. +/// PERSISTENT: stays armed until replaced by the next call (so n>1 +/// multi-choice requests re-snapshot the same boundary on every +/// full-prefill choice); pass -1 to clear. Out-of-range values fall back +/// to the prompt-end snapshot. Standalone `baseRT_prefill[_image/_audio]` +/// calls always snapshot at their prompt end (hints apply to +/// generate/generate_continue prefills only). +void baseRT_set_prefill_snapshot(baseRT_model_t model, int pos); + +/// Portable GDN reuse-snapshot blob (hybrid linear-attention models only). +/// The engine keeps a single most-recent boundary snapshot; a server-side +/// keyed store keeps several (one per distinct prior prompt) and loads the +/// best prefix match back before baseRT_try_rollback restores it. All three +/// are no-ops / return 0 / -1 on non-hybrid models. +/// baseRT_gdn_snapshot_size : fixed blob byte length for this model +/// (0 if not a hybrid model). Allocate this much for _capture. +/// baseRT_gdn_snapshot_capture : serialize the CURRENT snapshot (the one a +/// just-completed request's prompt prefill recorded) into `out` (capacity +/// `cap`). Returns bytes written, or -1 if there is no snapshot / cap is +/// too small / not hybrid. +/// baseRT_gdn_snapshot_load : deserialize `blob` back into the engine's +/// snapshot slot (NOT live state — a following baseRT_try_rollback applies +/// it). Returns the snapshot's KV position, or -1 on a length/model +/// mismatch. +int baseRT_gdn_snapshot_size(baseRT_model_t model); +int baseRT_gdn_snapshot_capture(baseRT_model_t model, uint8_t *out, int cap); +int baseRT_gdn_snapshot_load(baseRT_model_t model, const uint8_t *blob, int len); + +/// Per-sequence GDN snapshot (F6 M4: batched continuous-batching prefix reuse). +/// Capture/restore a CB sequence's OWN Gated-DeltaNet lane (its per-lane pool +/// slot) directly to/from a blob — distinct from the model-level snapshot APIs +/// above, which serve the single-sequence path via the lane-0 shadow. The blob +/// uses the same wire format and `baseRT_gdn_snapshot_size` byte length. +/// +/// baseRT_sequence_gdn_capture : serialize the sequence's lane state, +/// stamping its current KV length as the resume position. Call it when the +/// lane's state is at the intended (block-aligned) boundary. Returns bytes +/// written, or -1 (not hybrid / bad slot / cap too small). +/// baseRT_sequence_gdn_restore : deserialize `blob` into the sequence's lane +/// LIVE state. Returns the encoded position (the caller then sets the +/// sequence's KV length and prefills the suffix), or -1 on a mismatch. +int baseRT_sequence_gdn_capture(baseRT_sequence_t seq, uint8_t *out, int cap); +int baseRT_sequence_gdn_restore(baseRT_sequence_t seq, const uint8_t *blob, int len); + /// Generate tokens continuing from current KV cache state (no reset). /// Use for multi-turn chat: prefill new tokens only, then decode. BaseRTGenerationStats baseRT_generate_continue(baseRT_model_t model, const uint32_t *new_tokens, int n_new, @@ -706,12 +926,21 @@ const char *baseRT_chat_template_jinja(baseRT_model_t model); /// 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); + +/// BOS token id, for callers that need to prepend BOS to raw token +/// sequences (e.g. perplexity windows on BOS-sensitive models). +uint32_t baseRT_bos_id(baseRT_model_t model); const char *baseRT_eos_token(baseRT_model_t model); /// Primary end-of-sequence token id (the one the continuous-batching engine and /// other token-id consumers stop on). Returns 0 on a null handle. uint32_t baseRT_eos_token_id(baseRT_model_t model); +/// Max prompt tokens the fused (varlen) prefill can process in one packed batch. +/// The continuous-batching engine caps per-tick admitted prompt tokens by this +/// so a burst of long prompts doesn't overflow the packed prefill. 0 on null. +int baseRT_max_prefill_chunk(baseRT_model_t model); + // === Token counting === /// Count tokens in text without allocating an output buffer. diff --git a/include/baseRT/baseRT.h b/include/baseRT/baseRT.h index af8d477..de6344d 100644 --- a/include/baseRT/baseRT.h +++ b/include/baseRT/baseRT.h @@ -63,14 +63,14 @@ extern "C" { // === Versioning === #define BASERT_VERSION_MAJOR 0 -#define BASERT_VERSION_MINOR 1 -#define BASERT_VERSION_PATCH 7 +#define BASERT_VERSION_MINOR 2 +#define BASERT_VERSION_PATCH 0 /// Compile-time version, packed as `(MAJOR<<16) | (MINOR<<8) | PATCH`. /// Useful for `#if BASERT_VERSION >= 0x000200` feature checks. #define BASERT_VERSION ((BASERT_VERSION_MAJOR << 16) | (BASERT_VERSION_MINOR << 8) | BASERT_VERSION_PATCH) -/// Runtime-resolved version string ("0.1.0"). Matches the linked +/// Runtime-resolved version string ("0.2.0"). Matches the linked /// library; useful for diagnostics when a binding loads a different /// `.dylib` than it was compiled against. const char *baseRT_version_string(void); @@ -131,6 +131,39 @@ void baseRT_set_max_batch_size(int n); /// called before baseRT_load_model. Drive it via the baseRT_prefix_* API. void baseRT_set_prefix_cache(int enable); +/// Override the prefill chunk size (tokens per prefill GEMM batch). +/// n = 0 → per-chip default (recommended) +/// n >= 16 → clamp to this; shrinks the batch_* scratch footprint +/// (~linearly) so an oversized model fits a tighter GPU working-set +/// budget, at a prefill-throughput cost. Values outside [16, chip max] +/// are ignored. Process-wide; read at load. Set before baseRT_load_model. +void baseRT_set_prefill_chunk(int n); + +/// Paged-weights load policy (issue #113): per-tensor pinned weight buffers +/// for models past the OS wired-page budget. +/// mode = 0 → auto (default): normal load, retry paged on GPU-OOM +/// mode = 1 → force paged-weights on the first attempt (oversized/testing) +/// mode = 2 → disable the retry (fail hard on OOM, pre-#113 behavior) +/// Process-wide; read at load. Set before baseRT_load_model. +void baseRT_set_paged_weights(int mode); + +/// Toggle the baked-decode fast path (DispatchTable replay). +/// enable = 1 → on (default): replay the baked table when eligible +/// enable = 0 → force the live fused decode path (baked-vs-live A/B) +/// Process-wide; read per decode step. +void baseRT_set_baked_decode(int enable); + +/// GPU-wait timeout in milliseconds (Metal backend). +/// ms > 0 → (default: 300000, i.e. 5 min) return a loggable error if a +/// committed command buffer doesn't reach a terminal status in time. The +/// CPU unblocks; the wedged GPU work is NOT torn down (only a GPU +/// reset/reboot reclaims it). Chosen far above any legitimate single +/// command buffer so only a real wedge trips it. +/// ms = 0 → block indefinitely in the GPU wait (opt out of the bound). +/// No-op on non-Metal backends. Process-wide; read once per wait. The initial +/// value also honors the BASERT_GPU_WAIT_TIMEOUT_MS environment variable. +void baseRT_set_gpu_wait_timeout_ms(double ms); + /// Free all resources associated with a model. void baseRT_free_model(baseRT_model_t model); @@ -202,6 +235,26 @@ typedef bool (*baseRT_token_callback)(uint32_t token_id, const char *text, void BaseRTGenerationStats baseRT_generate(baseRT_model_t model, const uint32_t *prompt_tokens, int n_prompt, int max_tokens, BaseRTSamplingConfig sampling, baseRT_token_callback callback, void *user_data); +/// Generate tokens from a prompt WITH serial RadixCache prefix reuse. +/// +/// Semantically identical to `baseRT_generate` (single default sequence, same +/// greedy/sampled output), but when the model was loaded with BOTH `--paged-kv` +/// and `--prefix-cache` it reuses the longest cached whole-block prompt prefix: +/// it resets the default sequence, matches the prompt against the RadixCache, +/// seeds the shared prefix blocks into the default sequence, prefills ONLY the +/// divergent suffix `[matched_tokens, n_prompt)`, then inserts the full prompt +/// back into the cache for later reuse on generation end. +/// +/// Greedy output is BIT-IDENTICAL to a cold `baseRT_generate` (the seeded blocks +/// hold the same KV a cold prefill would have produced). When the prefix cache +/// or paged-KV is disabled this is an exact passthrough to `baseRT_generate`. +/// +/// Attention-KV (non-hybrid) models only — hybrid-GDN models fall back to the +/// plain path (their recurrent state is not block-shareable). +BaseRTGenerationStats baseRT_generate_cached(baseRT_model_t model, const uint32_t *prompt_tokens, int n_prompt, + int max_tokens, BaseRTSamplingConfig sampling, + baseRT_token_callback callback, void *user_data); + // === Multi-sequence generation (paged-KV only) === /// Opaque per-sequence handle. @@ -238,6 +291,27 @@ BaseRTGenerationStats baseRT_sequence_generate_continue(baseRT_sequence_t seq, c /// Release the sequence's blocks back to the pool and free the handle. void baseRT_sequence_free(baseRT_sequence_t seq); +/// Fused batch step with per-sequence trailing PAD counts (serving grid +/// alignment): counts[i] includes pads[i] throwaway tokens whose rows are +/// computed but whose argmax row is skipped (the output token comes from the +/// last REAL row). The caller must roll each padded sequence's KV back by +/// pads[i] after the call (baseRT_sequence_rollback). Greedy only. +int baseRT_batch_step_fused_pads(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, const uint32_t *in_tokens, + const int *in_token_counts, const int *pads, uint32_t *out_tokens); + +/// Warm the batched-decode fast paths for batch sizes up to `max_batch`: +/// each B runs two throwaway pure-decode ticks so shape-keyed caches (baked +/// dispatch tables, stream-captured CUDA graphs, cuBLAS plans) are built at +/// startup instead of on the first real requests — the same boot-time graph +/// warmup vLLM performs. Requires --paged-kv; call after load, before serving. +int baseRT_batch_warmup(baseRT_model_t model, int max_batch); + +/// 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). +int baseRT_sequence_rollback(baseRT_sequence_t seq, int length); + /// Batched decode: drives ONE batched decode step across N sequences. Each /// sequence writes its `new_tokens[i]` to its own KV cache slot, and attention /// reads each sequence's KV via its own block table. Throughput comes from @@ -330,6 +404,37 @@ int baseRT_batch_step(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, int baseRT_batch_step_fused(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, const uint32_t *in_tokens, const int *in_token_counts, uint32_t *out_tokens); +/// Fused batch step with per-sequence trailing PAD counts (serving grid +/// alignment): counts[i] includes pads[i] throwaway tokens whose rows are +/// computed but whose argmax row is skipped (the output token comes from the +/// last REAL row). The caller must roll each padded sequence's KV back by +/// pads[i] after the call (baseRT_sequence_rollback). Greedy only. +/// +/// Backend note: the shape-padding fast path exists for CUDA-graph capture; +/// on backends without stream capture (Metal) the pads are STRIPPED before +/// dispatch — semantics identical (no pad KV is written, so the caller's +/// rollback is a no-op), no wasted compute. +int baseRT_batch_step_fused_pads(baseRT_model_t model, baseRT_sequence_t *seqs, int n_seqs, const uint32_t *in_tokens, + const int *in_token_counts, const int *pads, uint32_t *out_tokens); + +/// Warm the batched-decode fast paths for batch sizes up to `max_batch`: +/// each B runs throwaway pure-decode ticks so shape-keyed caches (baked +/// dispatch tables, captured graphs, PSO/plan builds) are built at startup +/// instead of on the first real requests — the same boot-time warmup vLLM +/// performs. Requires --paged-kv; call after load, before serving. +int baseRT_batch_warmup(baseRT_model_t model, int max_batch); + +/// Max prompt tokens the fused (varlen) prefill can process in one packed batch. +/// The continuous-batching engine caps per-tick admitted prompt tokens by this +/// so a burst of long prompts doesn't overflow the packed prefill. 0 on null. +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). +int baseRT_sequence_rollback(baseRT_sequence_t seq, int length); + /// Multi-step autoregressive driver for baseRT_batch_step_fused. Step 0 /// ingests the mixed-length input from `first_in_tokens` / `first_in_token_counts` /// (one row per seq, total length `sum(first_in_token_counts)`). Subsequent @@ -406,6 +511,13 @@ BaseRTPrefixMatch baseRT_prefix_match(baseRT_model_t model, const uint32_t *toke /// paged / the sequence isn't empty. int baseRT_sequence_seed_prefix(baseRT_sequence_t seq, const int *blocks, int n_blocks, int n_tokens); +/// Paged-KV block (page) size in tokens for this model, or 0 when the model +/// was not loaded with --paged-kv. This is the block-alignment granularity for +/// baseRT_prefix_match/_seed_prefix (matched_tokens == matched_blocks * +/// page_size); the continuous-batching hybrid-GDN prefix-reuse path uses it to +/// pick the block-aligned GDN snapshot boundary at admit. +int baseRT_page_size(baseRT_model_t model); + /// Publish `seq`'s KV blocks for the block-aligned prefix of `tokens` into the /// prefix cache so later requests can reuse them. Idempotent for an already- /// cached prefix (no double refcount). No-op when the cache is disabled. @@ -415,8 +527,20 @@ int baseRT_prefix_insert(baseRT_model_t model, const uint32_t *tokens, int n_tok /// Release the lock a baseRT_prefix_match took on a prefix and free the match's /// bookkeeping. Call exactly once per non-zero handle, after the sequence that /// reused the prefix has been inserted/retired. No-op for handle==0. +/// +/// Use this ONLY when the match's blocks WERE seeded into a sequence +/// (baseRT_sequence_seed_prefix): the sequence owns those blocks and drops the +/// match's ownership incref when it resets/frees. If the match was NOT seeded +/// (you decided not to reuse it), call baseRT_prefix_release instead — unlock +/// alone would leak the increfed blocks. void baseRT_prefix_unlock(baseRT_model_t model, uint64_t handle); +/// Abandon a baseRT_prefix_match WITHOUT seeding it: drops the ownership incref +/// on each matched block (which no sequence adopted) AND releases the trie lock. +/// Call exactly once per non-zero handle when you matched a prefix but chose not +/// to seed it (e.g. a boundary mismatch). No-op for handle==0. +void baseRT_prefix_release(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 @@ -473,6 +597,25 @@ void baseRT_grammar_free(baseRT_grammar_t grammar); /// would feed the second sample through a terminated grammar (garbage). void baseRT_grammar_reset(baseRT_grammar_t grammar); +/// Grammar stepping for the continuous-batching server (xgrammar backend). +/// The server applies the bitmask to a lane's logits row on the host, then +/// accepts the sampled token to advance the grammar. A legacy-NPDA grammar +/// reports `bitmask_size == 0` — the caller must keep it on the serial path. +/// 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. +/// 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. +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); +int baseRT_grammar_is_terminated(baseRT_grammar_t grammar); +int baseRT_grammar_is_completed(baseRT_grammar_t grammar); + /// Generate tokens with grammar constraint. /// Grammar masks invalid tokens at each step, guaranteeing output conforms to the grammar. BaseRTGenerationStats baseRT_generate_grammar(baseRT_model_t model, const uint32_t *prompt_tokens, int n_prompt, @@ -632,11 +775,16 @@ void baseRT_reset(baseRT_model_t model); /// `current_length` prefix (not the unused tail), so the file size grows /// linearly with how much was prefilled+decoded. Returns 0 on success and /// a negative error code on failure; check `baseRT_get_error` for details. +/// Hybrid linear-attention models (Qwen 3.5/3.6) are REJECTED: the format +/// holds attention KV only, not the Gated-DeltaNet recurrent state. int baseRT_save_state(baseRT_model_t model, const char *path); /// Inverse of `baseRT_save_state`. The cache must have been allocated for /// a model with matching shape; mismatched files are rejected. After load, /// `baseRT_get_position` reflects the restored token count. +/// Hybrid linear-attention models (Qwen 3.5/3.6) are REJECTED: the file +/// holds attention KV only, and restoring it without the matching +/// Gated-DeltaNet recurrent state would yield a corrupt hybrid state. int baseRT_load_state(baseRT_model_t model, const char *path); /// Install a LoRA adapter on this model. The adapter file is a `.base` @@ -664,8 +812,80 @@ const char *baseRT_lora_id(baseRT_model_t model); /// shared chat-template prefix from a prior request — keeps the cached /// prefill of the common prefix while discarding the prior turn's /// user-message tail and assistant reply. +/// Hybrid linear-attention models (Qwen 3.5/3.6): the recurrent state +/// cannot be rewound to an arbitrary position. This call keeps its "KV +/// length == to_position" promise only when `to_position` exactly matches +/// the recurrent-state snapshot (see `baseRT_set_prefill_snapshot`); any +/// other position degrades to a FULL reset (equivalent to `baseRT_reset`) +/// — the caller must then prefill the entire prompt again. Use +/// `baseRT_try_rollback` to detect what happened, or to resume from a +/// snapshot that sits before the requested position. void baseRT_rollback(baseRT_model_t model, int to_position); +/// Rollback that reports the position actually achieved. Non-hybrid +/// models land on `min(to_position, current KV length)` — a target past +/// the cache end cannot be "achieved" by a rollback and is clamped so +/// callers prefilling from the returned position never skip tokens. Hybrid linear-attention models can only resume +/// from their recurrent-state snapshot (see +/// `baseRT_set_prefill_snapshot`): when the snapshot sits at or before +/// `to_position` the state is restored there and the SNAPSHOT position is +/// returned — the caller must prefill the prompt from that position +/// onward. When the snapshot lies past `to_position` (divergent history) +/// the call returns -1 and leaves the model state UNTOUCHED — fall back +/// to `baseRT_reset` + a full prefill. `to_position == 0` always succeeds +/// as a full reset. +int baseRT_try_rollback(baseRT_model_t model, int to_position); + +/// Hybrid linear-attention models only (no-op otherwise): ask prompt +/// prefills to capture the reuse snapshot once absolute KV position `pos` +/// has been processed, instead of at the prompt end. Chat servers pass +/// the rendered-history boundary (the prompt minus the generation +/// scaffold): the scaffold tokens never reappear in the next request's +/// render, so a prompt-end snapshot would never match, while the history +/// boundary is exactly where the next request's shared prefix ends. +/// PERSISTENT: stays armed until replaced by the next call (so n>1 +/// multi-choice requests re-snapshot the same boundary on every +/// full-prefill choice); pass -1 to clear. Out-of-range values fall back +/// to the prompt-end snapshot. Standalone `baseRT_prefill[_image/_audio]` +/// calls always snapshot at their prompt end (hints apply to +/// generate/generate_continue prefills only). +void baseRT_set_prefill_snapshot(baseRT_model_t model, int pos); + +/// Portable GDN reuse-snapshot blob (hybrid linear-attention models only). +/// The engine keeps a single most-recent boundary snapshot; a server-side +/// keyed store keeps several (one per distinct prior prompt) and loads the +/// best prefix match back before baseRT_try_rollback restores it. All three +/// are no-ops / return 0 / -1 on non-hybrid models. +/// baseRT_gdn_snapshot_size : fixed blob byte length for this model +/// (0 if not a hybrid model). Allocate this much for _capture. +/// baseRT_gdn_snapshot_capture : serialize the CURRENT snapshot (the one a +/// just-completed request's prompt prefill recorded) into `out` (capacity +/// `cap`). Returns bytes written, or -1 if there is no snapshot / cap is +/// too small / not hybrid. +/// baseRT_gdn_snapshot_load : deserialize `blob` back into the engine's +/// snapshot slot (NOT live state — a following baseRT_try_rollback applies +/// it). Returns the snapshot's KV position, or -1 on a length/model +/// mismatch. +int baseRT_gdn_snapshot_size(baseRT_model_t model); +int baseRT_gdn_snapshot_capture(baseRT_model_t model, uint8_t *out, int cap); +int baseRT_gdn_snapshot_load(baseRT_model_t model, const uint8_t *blob, int len); + +/// Per-sequence GDN snapshot (F6 M4: batched continuous-batching prefix reuse). +/// Capture/restore a CB sequence's OWN Gated-DeltaNet lane (its per-lane pool +/// slot) directly to/from a blob — distinct from the model-level snapshot APIs +/// above, which serve the single-sequence path via the lane-0 shadow. The blob +/// uses the same wire format and `baseRT_gdn_snapshot_size` byte length. +/// +/// baseRT_sequence_gdn_capture : serialize the sequence's lane state, +/// stamping its current KV length as the resume position. Call it when the +/// lane's state is at the intended (block-aligned) boundary. Returns bytes +/// written, or -1 (not hybrid / bad slot / cap too small). +/// baseRT_sequence_gdn_restore : deserialize `blob` into the sequence's lane +/// LIVE state. Returns the encoded position (the caller then sets the +/// sequence's KV length and prefills the suffix), or -1 on a mismatch. +int baseRT_sequence_gdn_capture(baseRT_sequence_t seq, uint8_t *out, int cap); +int baseRT_sequence_gdn_restore(baseRT_sequence_t seq, const uint8_t *blob, int len); + /// Generate tokens continuing from current KV cache state (no reset). /// Use for multi-turn chat: prefill new tokens only, then decode. BaseRTGenerationStats baseRT_generate_continue(baseRT_model_t model, const uint32_t *new_tokens, int n_new, @@ -706,12 +926,21 @@ const char *baseRT_chat_template_jinja(baseRT_model_t model); /// 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); + +/// BOS token id, for callers that need to prepend BOS to raw token +/// sequences (e.g. perplexity windows on BOS-sensitive models). +uint32_t baseRT_bos_id(baseRT_model_t model); const char *baseRT_eos_token(baseRT_model_t model); /// Primary end-of-sequence token id (the one the continuous-batching engine and /// other token-id consumers stop on). Returns 0 on a null handle. uint32_t baseRT_eos_token_id(baseRT_model_t model); +/// Max prompt tokens the fused (varlen) prefill can process in one packed batch. +/// The continuous-batching engine caps per-tick admitted prompt tokens by this +/// so a burst of long prompts doesn't overflow the packed prefill. 0 on null. +int baseRT_max_prefill_chunk(baseRT_model_t model); + // === Token counting === /// Count tokens in text without allocating an output buffer.