From 20e91712da8c188697b161b2d4b30c7f2d9b396f Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Mon, 7 Sep 2026 22:06:09 -0500 Subject: [PATCH] chore: delete the 25,984-line dead duplicate under examples/src `crates/mummu/examples/src/` was a stale copy of `crates/mummu/src/`, added by commit 769d218 ("stuff", 2026-08-31) outside the PR process. Cargo never built it, and it had already diverged from the real source -- so it was pure weight that silently polluted every grep over the crate. It cost the 2026-09-07 run a wrong read of where `TensorSnapshot` is used, and doubled the apparent site count when auditing `HybridKv` and `ParamSrc`. Proof the delete is inert: - no `main.rs` under the tree, so example auto-discovery never made it a target; - no `[[example]]` entry in any Cargo.toml (auto-discovery only) and no `include!` / `#[path]` reference from outside the tree; - `cargo metadata` reports the SAME 30 example targets before and after, none of them named `src`. 40 files, 25,984 deletions. `cargo check --workspace --all-targets` is green with 0 errors and 0 warnings. Co-Authored-By: Claude Opus 5 --- ROADMAP.md | 20 +- crates/mummu/examples/src/adapt.rs | 394 ---- crates/mummu/examples/src/attn_config.rs | 353 --- crates/mummu/examples/src/backend.rs | 938 -------- crates/mummu/examples/src/cache.rs | 277 --- crates/mummu/examples/src/chat.rs | 1488 ------------- crates/mummu/examples/src/decode.rs | 441 ---- crates/mummu/examples/src/gguf.rs | 2029 ----------------- crates/mummu/examples/src/gguf_iq_grids.rs | 325 --- crates/mummu/examples/src/hub.rs | 647 ------ crates/mummu/examples/src/import.rs | 574 ----- crates/mummu/examples/src/lib.rs | 45 - crates/mummu/examples/src/manage.rs | 372 ---- crates/mummu/examples/src/models/lfm2.rs | 830 ------- crates/mummu/examples/src/models/minilm.rs | 391 ---- crates/mummu/examples/src/models/mod.rs | 204 -- crates/mummu/examples/src/models/olmoe.rs | 1592 -------------- crates/mummu/examples/src/models/qwen2.rs | 831 ------- crates/mummu/examples/src/models/qwen3.rs | 738 ------- crates/mummu/examples/src/models/qwen35.rs | 1941 ---------------- crates/mummu/examples/src/nn/attention.rs | 458 ---- crates/mummu/examples/src/nn/conv.rs | 247 --- crates/mummu/examples/src/nn/mlp.rs | 112 - crates/mummu/examples/src/nn/mod.rs | 29 - crates/mummu/examples/src/nn/moe.rs | 2204 ------------------- crates/mummu/examples/src/nn/packed_gemv.rs | 588 ----- crates/mummu/examples/src/nn/rope.rs | 150 -- crates/mummu/examples/src/pack.rs | 1152 ---------- crates/mummu/examples/src/partition.rs | 480 ---- crates/mummu/examples/src/plan.rs | 333 --- crates/mummu/examples/src/prof.rs | 268 --- crates/mummu/examples/src/quant.rs | 158 -- crates/mummu/examples/src/registry.rs | 398 ---- crates/mummu/examples/src/safetensors.rs | 1192 ---------- crates/mummu/examples/src/template.rs | 413 ---- crates/mummu/examples/src/tier.rs | 528 ----- crates/mummu/examples/src/tok_config.rs | 934 -------- crates/mummu/examples/src/tokenizer.rs | 1033 --------- crates/mummu/examples/src/tune.rs | 253 --- crates/mummu/examples/src/vram.rs | 218 -- crates/mummu/examples/src/workingset.rs | 426 ---- 41 files changed, 11 insertions(+), 25993 deletions(-) delete mode 100644 crates/mummu/examples/src/adapt.rs delete mode 100644 crates/mummu/examples/src/attn_config.rs delete mode 100644 crates/mummu/examples/src/backend.rs delete mode 100644 crates/mummu/examples/src/cache.rs delete mode 100644 crates/mummu/examples/src/chat.rs delete mode 100644 crates/mummu/examples/src/decode.rs delete mode 100644 crates/mummu/examples/src/gguf.rs delete mode 100644 crates/mummu/examples/src/gguf_iq_grids.rs delete mode 100644 crates/mummu/examples/src/hub.rs delete mode 100644 crates/mummu/examples/src/import.rs delete mode 100644 crates/mummu/examples/src/lib.rs delete mode 100644 crates/mummu/examples/src/manage.rs delete mode 100644 crates/mummu/examples/src/models/lfm2.rs delete mode 100644 crates/mummu/examples/src/models/minilm.rs delete mode 100644 crates/mummu/examples/src/models/mod.rs delete mode 100644 crates/mummu/examples/src/models/olmoe.rs delete mode 100644 crates/mummu/examples/src/models/qwen2.rs delete mode 100644 crates/mummu/examples/src/models/qwen3.rs delete mode 100644 crates/mummu/examples/src/models/qwen35.rs delete mode 100644 crates/mummu/examples/src/nn/attention.rs delete mode 100644 crates/mummu/examples/src/nn/conv.rs delete mode 100644 crates/mummu/examples/src/nn/mlp.rs delete mode 100644 crates/mummu/examples/src/nn/mod.rs delete mode 100644 crates/mummu/examples/src/nn/moe.rs delete mode 100644 crates/mummu/examples/src/nn/packed_gemv.rs delete mode 100644 crates/mummu/examples/src/nn/rope.rs delete mode 100644 crates/mummu/examples/src/pack.rs delete mode 100644 crates/mummu/examples/src/partition.rs delete mode 100644 crates/mummu/examples/src/plan.rs delete mode 100644 crates/mummu/examples/src/prof.rs delete mode 100644 crates/mummu/examples/src/quant.rs delete mode 100644 crates/mummu/examples/src/registry.rs delete mode 100644 crates/mummu/examples/src/safetensors.rs delete mode 100644 crates/mummu/examples/src/template.rs delete mode 100644 crates/mummu/examples/src/tier.rs delete mode 100644 crates/mummu/examples/src/tok_config.rs delete mode 100644 crates/mummu/examples/src/tokenizer.rs delete mode 100644 crates/mummu/examples/src/tune.rs delete mode 100644 crates/mummu/examples/src/vram.rs delete mode 100644 crates/mummu/examples/src/workingset.rs diff --git a/ROADMAP.md b/ROADMAP.md index 927147c..689404d 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -1468,15 +1468,17 @@ a benchmark holds/improves its budget; README perf claims link an artifact. same user, and a broad `pkill` from any of them takes out another's linker (this run demonstrated the reverse: its own `pkill -9 rustc` killed the Nanna build's compiles). Until it is pinned down, treat a lone `signal: 9` at link as transient and retry before believing it. *(2026-09-07)* -- [ ] **`crates/mummu/examples/src/` is a 25,984-line dead duplicate of `crates/mummu/src/`** — added by - commit `769d218` ("stuff", 2026-08-31), which landed on `main` outside the PR process. Cargo never - builds it: a subdirectory of `examples/` is only a target when it contains `main.rs`, and - `find crates/mummu/examples/src -name main.rs` returns **0**. It has already diverged from the real - source (`diff crates/mummu/src/gguf.rs crates/mummu/examples/src/gguf.rs` reports they differ), so - it is stale weight that silently pollutes every `grep`/`rg` over the crate — it cost this run a - wrong read of where `TensorSnapshot` is used, and doubled the apparent site count when auditing - `HybridKv`/`ParamSrc`. Delete it, or promote whichever files were meant to be examples into real - example targets. *(2026-09-07)* +- [x] **`crates/mummu/examples/src/` was a 25,984-line dead duplicate of `crates/mummu/src/` — deleted.** + Added by commit `769d218` ("stuff", 2026-08-31), which landed on `main` outside the PR process. + Cargo never built it: a subdirectory of `examples/` is only a target when it contains `main.rs`, + and `find crates/mummu/examples/src -name main.rs` returned **0**. It had already diverged from the + real source (`diff crates/mummu/src/gguf.rs crates/mummu/examples/src/gguf.rs` reported a + difference), so it was stale weight that silently polluted every `grep`/`rg` over the crate — it + cost the 2026-09-07 run a wrong read of where `TensorSnapshot` is used and doubled the apparent + site count when auditing `HybridKv`/`ParamSrc`. Proof the delete is inert: no `[[example]]` entry + anywhere (auto-discovery only), no `include!`/`#[path]` reference from outside the tree, and + `cargo metadata` reports the **same 30 example targets before and after**, none named `src`. + 40 files, 25,984 deletions. *(2026-09-07)* ### P1 — Backends & device *(ex-laurelane)* - [x] Backend abstraction generic over `B: Backend`; one binary compiling BOTH `Wgpu` (Vulkan/DX12/Metal, diff --git a/crates/mummu/examples/src/adapt.rs b/crates/mummu/examples/src/adapt.rs deleted file mode 100644 index b32b7a1..0000000 --- a/crates/mummu/examples/src/adapt.rs +++ /dev/null @@ -1,394 +0,0 @@ -//! **Adaptive placement** — keep the device/host split earning its keep as -//! the machine changes underneath it. -//! -//! A placement decided once at load is a guess about a machine that will not -//! stay still: a game starts, another model loads, the desktop compositor -//! grows, a background job eats host RAM. The split that was optimal at load -//! becomes an out-of-memory crash or an idle GPU an hour later. -//! -//! # Why control on *effects*, not causes -//! -//! The tempting design is to introspect every cause — free VRAM, other -//! processes' usage, GPU utilization. That needs a different API per OS and -//! per vendor (DXGI `QueryVideoMemoryInfo`, NVML, sysfs), each of which can -//! be missing, stale, or lie about what a driver will actually hand out. -//! -//! This controller instead watches what actually matters and is always -//! observable: **did an allocation fail**, and **did throughput move**. Those -//! two signals capture every cause, including ones no query would reveal -//! (thermal throttling, a driver reserving memory, another process's spike). -//! Direct readings — host memory available, our own resident bytes — are used -//! where they are cheap, as *hints* that bound the search rather than as the -//! decision. -//! -//! # AIMD, and why that shape -//! -//! The device budget moves by **additive increase, multiplicative decrease**: -//! grow slowly while things are fine, cut hard the moment they are not. This -//! is the shape TCP congestion control settled on for the same problem — -//! a shared resource with an unknown, moving limit where overshoot is far -//! more expensive than undershoot. Here overshoot means an OOM that kills a -//! generation (or, as measured on this project, the whole process); undershoot -//! only means some layers run on the CPU. The asymmetry in cost justifies the -//! asymmetry in response. -//! -//! Two damping rules keep it from thrashing: -//! -//! - a **dwell time** — no change until the current placement has been -//! observed long enough to judge it; -//! - a **deadband** — throughput must move by more than measurement noise -//! before it counts as evidence. -//! -//! Moves are applied in bounded batches by the caller (see -//! `ExpertPool::apply_schedule`), so adapting never stalls a generation. - -use std::time::{Duration, Instant}; - -/// What the controller learned from one observation window. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum Adjust { - /// Budget unchanged — either nothing moved enough to act on, or the - /// placement has not been observed long enough yet. - Hold, - /// Grow the device budget to this many bytes (additive increase). - Grow(u64), - /// Shrink to this many bytes (multiplicative decrease). - Shrink(u64), -} - -/// One observation of how the current placement is doing. -#[derive(Debug, Clone, Copy)] -pub struct Sample { - /// Decode throughput since the last sample. The signal that matters. - pub tokens_per_sec: f64, - /// Did anything fail to allocate on the device since the last sample? - /// Overrides everything else: a placement that cannot allocate is not a - /// placement, however fast it looked. - pub device_alloc_failed: bool, - /// Host memory currently available, when the OS will say. `None` means - /// "unknown", never "plenty". - pub host_available_bytes: Option, - /// Bytes the model currently holds on the device. - pub device_bytes_in_use: u64, -} - -/// Tunables. The defaults are deliberately cautious: this controller runs -/// unattended against a machine that other software is also using. -#[derive(Debug, Clone, Copy)] -pub struct Policy { - /// Never plan below this device budget — under it, the device is not - /// worth the cross-device traffic. - pub floor_bytes: u64, - /// Never plan above this (the hardware bound, minus whatever headroom - /// the caller wants to leave the rest of the system). - pub ceiling_bytes: u64, - /// Additive increase step. - pub grow_step_bytes: u64, - /// Multiplicative decrease factor, in (0, 1). - pub shrink_factor: f64, - /// How long a placement must be observed before it is judged. - pub dwell: Duration, - /// Relative throughput change below which a sample is treated as noise. - pub deadband: f64, -} - -impl Policy { - /// A policy for a device of `total_bytes`, leaving `reserve_bytes` to the - /// rest of the system (desktop compositor, other processes). - #[must_use] - pub fn for_device(total_bytes: u64, reserve_bytes: u64) -> Self { - let ceiling = total_bytes.saturating_sub(reserve_bytes); - Self { - floor_bytes: (total_bytes / 8).min(ceiling), - ceiling_bytes: ceiling, - // ~6% of the ceiling per step: a dozen good windows to go from - // floor to ceiling, which is slow enough to notice a mistake. - grow_step_bytes: (ceiling / 16).max(1), - shrink_factor: 0.75, - dwell: Duration::from_secs(30), - deadband: 0.05, - } - } -} - -/// The controller. Feed it [`Sample`]s; it answers with an [`Adjust`]. -#[derive(Debug)] -pub struct Controller { - policy: Policy, - budget: u64, - /// Best throughput seen, and the budget that produced it — what "did this - /// change help?" is judged against. - best: Option<(f64, u64)>, - last_change: Instant, - /// Budgets that produced an allocation failure. Never grow back into one - /// blindly: the ceiling that actually matters is the one the machine - /// enforced, not the one the spec sheet advertises. - failed_at: Option, -} - -impl Controller { - #[must_use] - pub fn new(policy: Policy, initial_budget: u64) -> Self { - Self { - budget: initial_budget.clamp(policy.floor_bytes, policy.ceiling_bytes), - policy, - best: None, - last_change: Instant::now(), - failed_at: None, - } - } - - /// The device budget the placement planner should use right now. - #[must_use] - pub fn budget(&self) -> u64 { - self.budget - } - - /// Feed one observation. `now` is injected so the logic is testable - /// without sleeping. - pub fn observe(&mut self, sample: &Sample, now: Instant) -> Adjust { - // 1. An allocation failure is not a data point to weigh against - // throughput — it is a hard ceiling discovery. Act immediately, - // ignoring dwell: staying here risks the next OOM. - if sample.device_alloc_failed { - let ceiling = sample.device_bytes_in_use.min(self.budget); - self.failed_at = Some(match self.failed_at { - Some(prev) => prev.min(ceiling), - None => ceiling, - }); - let next = ((ceiling as f64 * self.policy.shrink_factor) as u64) - .max(self.policy.floor_bytes); - self.last_change = now; - self.best = None; // the old best was measured under a limit that no longer holds - if next < self.budget { - self.budget = next; - return Adjust::Shrink(next); - } - return Adjust::Hold; - } - - // 2. Judge a placement only once it has been observed long enough. - if now.duration_since(self.last_change) < self.policy.dwell { - return Adjust::Hold; - } - - // 3. Host pressure: if the OS says memory is short, the CPU side of - // the split is the problem, so pulling MORE onto the device helps. - // (Unknown is not permission — `None` does nothing.) - let host_pressure = sample - .host_available_bytes - .is_some_and(|avail| avail < self.policy.grow_step_bytes * 2); - - let improved = match self.best { - None => true, - Some((best_tps, _)) => { - let delta = (sample.tokens_per_sec - best_tps) / best_tps.max(1e-9); - delta > self.policy.deadband - } - }; - let regressed = match self.best { - None => false, - Some((best_tps, best_budget)) => { - let delta = (best_tps - sample.tokens_per_sec) / best_tps.max(1e-9); - delta > self.policy.deadband && self.budget != best_budget - } - }; - - if improved { - self.best = Some((sample.tokens_per_sec, self.budget)); - } - - // 4. A regression means the last move was wrong: go back to what was - // measurably better rather than continuing to explore. - if regressed && let Some((_, best_budget)) = self.best { - self.last_change = now; - let next = best_budget.clamp(self.policy.floor_bytes, self.policy.ceiling_bytes); - if next != self.budget { - self.budget = next; - return if next > self.budget { - Adjust::Grow(next) - } else { - Adjust::Shrink(next) - }; - } - return Adjust::Hold; - } - - // 5. Otherwise creep upward, but never back into a budget the machine - // already refused, and never past the ceiling. - let hard_ceiling = self - .failed_at - .map_or(self.policy.ceiling_bytes, |f| { - // Stay a decrease-step below the level that failed. - ((f as f64 * self.policy.shrink_factor) as u64).max(self.policy.floor_bytes) - }) - .min(self.policy.ceiling_bytes); - - let want = self.budget.saturating_add(self.policy.grow_step_bytes); - if (improved || host_pressure) && want <= hard_ceiling { - self.last_change = now; - self.budget = want; - return Adjust::Grow(want); - } - Adjust::Hold - } -} - -#[cfg(test)] -mod tests { - use super::*; - - fn policy() -> Policy { - Policy { - floor_bytes: 1_000, - ceiling_bytes: 10_000, - grow_step_bytes: 1_000, - shrink_factor: 0.5, - dwell: Duration::from_secs(10), - deadband: 0.05, - } - } - - fn sample(tps: f64) -> Sample { - Sample { - tokens_per_sec: tps, - device_alloc_failed: false, - host_available_bytes: None, - device_bytes_in_use: 5_000, - } - } - - #[test] - fn an_allocation_failure_shrinks_immediately_ignoring_dwell() { - let mut c = Controller::new(policy(), 8_000); - let now = Instant::now(); - // No dwell has passed at all — a hard limit must not wait for one. - let got = c.observe( - &Sample { - device_alloc_failed: true, - device_bytes_in_use: 8_000, - ..sample(10.0) - }, - now, - ); - assert_eq!(got, Adjust::Shrink(4_000)); - assert_eq!(c.budget(), 4_000); - } - - #[test] - fn it_never_grows_back_into_a_budget_the_machine_refused() { - let mut c = Controller::new(policy(), 8_000); - let t0 = Instant::now(); - c.observe( - &Sample { - device_alloc_failed: true, - device_bytes_in_use: 8_000, - ..sample(10.0) - }, - t0, - ); - // Now improve repeatedly for a long time; it may creep, but never to - // the level that failed. - let mut t = t0; - let mut tps = 10.0; - for _ in 0..20 { - t += Duration::from_secs(11); - tps *= 1.5; - c.observe(&sample(tps), t); - } - assert!( - c.budget() < 8_000, - "must stay below the refused budget, got {}", - c.budget() - ); - } - - #[test] - fn it_holds_until_the_placement_has_been_observed_long_enough() { - let mut c = Controller::new(policy(), 5_000); - let t = Instant::now(); - assert_eq!(c.observe(&sample(10.0), t), Adjust::Hold); - assert_eq!(c.observe(&sample(99.0), t + Duration::from_secs(5)), Adjust::Hold); - // Past the dwell, an improvement is actionable. - assert_eq!( - c.observe(&sample(99.0), t + Duration::from_secs(11)), - Adjust::Grow(6_000) - ); - } - - #[test] - fn noise_within_the_deadband_does_not_move_the_budget() { - let mut c = Controller::new(policy(), 5_000); - let t = Instant::now(); - c.observe(&sample(10.0), t + Duration::from_secs(11)); // sets a baseline, grows - let after_first = c.budget(); - // A 2% wobble is not evidence of anything. - let got = c.observe(&sample(10.2), t + Duration::from_secs(30)); - assert_eq!(got, Adjust::Hold, "2% is inside the deadband"); - assert_eq!(c.budget(), after_first); - } - - #[test] - fn a_regression_returns_to_the_budget_that_measured_best() { - let mut c = Controller::new(policy(), 5_000); - let t = Instant::now(); - // Establish a good result at 5_000, then grow to 6_000. - assert_eq!(c.observe(&sample(20.0), t + Duration::from_secs(11)), Adjust::Grow(6_000)); - // 6_000 turns out much worse -> go back to what was measurably better. - let got = c.observe(&sample(5.0), t + Duration::from_secs(30)); - assert!( - matches!(got, Adjust::Shrink(_) | Adjust::Hold), - "a regression must not keep growing: {got:?}" - ); - assert!(c.budget() <= 6_000); - } - - #[test] - fn host_memory_pressure_pulls_work_onto_the_device() { - let mut c = Controller::new(policy(), 5_000); - let t = Instant::now(); - c.observe(&sample(10.0), t + Duration::from_secs(11)); - let before = c.budget(); - // Throughput flat (no improvement), but the host is nearly out of - // memory — moving more onto the device is the relief valve. - let got = c.observe( - &Sample { - host_available_bytes: Some(100), - ..sample(10.0) - }, - t + Duration::from_secs(40), - ); - assert!(matches!(got, Adjust::Grow(_)), "host pressure should pull work to the device: {got:?}"); - assert!(c.budget() > before); - } - - #[test] - fn unknown_host_memory_is_not_treated_as_plenty() { - // `None` must behave like "no information", not like "lots free" — - // the difference decides whether an unsupported OS silently gets the - // aggressive path. - let mut c = Controller::new(policy(), 5_000); - let t = Instant::now(); - c.observe(&sample(10.0), t + Duration::from_secs(11)); - let before = c.budget(); - let got = c.observe(&sample(10.0), t + Duration::from_secs(40)); - assert_eq!(got, Adjust::Hold); - assert_eq!(c.budget(), before); - } - - #[test] - fn the_budget_stays_within_its_bounds() { - let p = policy(); - let mut c = Controller::new(p, 50_000); // above the ceiling - assert_eq!(c.budget(), p.ceiling_bytes, "clamped on construction"); - let mut t = Instant::now(); - let mut tps = 1.0; - for _ in 0..30 { - t += Duration::from_secs(11); - tps *= 1.5; - c.observe(&sample(tps), t); - assert!(c.budget() <= p.ceiling_bytes, "never above the ceiling"); - assert!(c.budget() >= p.floor_bytes, "never below the floor"); - } - } -} diff --git a/crates/mummu/examples/src/attn_config.rs b/crates/mummu/examples/src/attn_config.rs deleted file mode 100644 index fdc3e16..0000000 --- a/crates/mummu/examples/src/attn_config.rs +++ /dev/null @@ -1,353 +0,0 @@ -//! Attention-shaping configuration a checkpoint may declare that the shared -//! blocks do not implement: **RoPE frequency scaling** (`rope_scaling`) and -//! **sliding-window attention** (`sliding_window`). -//! -//! [`crate::nn::rope_tables`] computes plain rotary frequencies and -//! [`crate::nn::causal_mask`] a full causal mask. Every model currently in the -//! zoo is fine — their configs ship `rope_scaling: null` and no *enabled* -//! window — but this is the silent-wrong-answer class of gap: a checkpoint -//! that carries `rope_scaling` (Qwen2.5 past its 32 k native context via YaRN) -//! or an enabled `sliding_window` (Mistral-family; Gemma 2/3's alternating -//! layers) would load clean, pass the short-prompt parity probes, and degrade -//! numerically only far out in the context — exactly where nothing looks. -//! -//! So the cheap half first: **parse the fields and refuse the load, naming the -//! mode**. That converts silent degradation into an error a consumer can act -//! on. Implementing the modes is the second half (scaled `rope_tables` + a -//! windowed mask as config-driven variants of the same blocks), gated on a -//! long-context parity leg the short probes cannot see by construction. -//! -//! The trap this module exists to not fall into: **`sliding_window` being -//! present does not mean it is on**. Every Qwen2.5 checkpoint ships -//! `"sliding_window": 32768` together with `"use_sliding_window": false` — the -//! window is inert, and rejecting on the field's presence would refuse two -//! models Mummu has parity-verified for a month. Each family therefore decides -//! *enabled* by its own convention and passes the answer here; see -//! [`check_sliding_window`]. - -use crate::gguf::{GgufFile, GgufValue}; - -/// The `rope_scaling` object of an HF `config.json`, in both spellings that -/// appear in the wild: transformers ≥ 4.38 writes `rope_type`, older -/// checkpoints write `type`, and some carry both. -/// -/// Deliberately lenient about the *payload* (`factor` and friends stay -/// `Option`) and strict about the *mode*: an unknown mode must be refused, and -/// refusing it does not require understanding its parameters. -#[derive(Debug, Clone, Default, PartialEq, serde::Deserialize)] -pub struct RopeScaling { - /// transformers ≥ 4.38 spelling. - #[serde(default)] - pub rope_type: Option, - /// Pre-4.38 spelling. Kept separate rather than aliased so a checkpoint - /// carrying both disagreeing values is visible instead of arbitrated. - #[serde(default, rename = "type")] - pub legacy_type: Option, - /// Context-extension factor (YaRN / linear / dynamic-NTK). Unused until - /// the scaled tables land; parsed so the eventual implementation reads the - /// same struct the rejection does. - #[serde(default)] - pub factor: Option, - /// The context length the checkpoint was trained at, before scaling. - #[serde(default)] - pub original_max_position_embeddings: Option, - /// Everything else the object carries. Kept, not discarded, because the - /// *shape* of the leftovers is load-bearing: transformers lets a - /// Gemma-3-style config nest one RoPE object **per layer type** - /// (`{"full_attention": {…}, "sliding_attention": {…}}`), and a nested map - /// deserializes into this struct with every named field absent — which - /// would read as "plain rotary" and sail through [`Self::check`]. Seeing - /// the nested objects is what lets that be refused instead. - #[serde(flatten)] - pub extra: std::collections::BTreeMap, -} - -/// The spellings that mean "no scaling — plain rotary". `default` is what -/// transformers writes for an unscaled `rope_parameters`; `none`/`null` show -/// up in hand-written and converted configs. -const PLAIN_ROPE_TYPES: [&str; 3] = ["default", "none", "null"]; - -impl RopeScaling { - /// The declared mode, lowercased, or `"default"` when the object names - /// none. Both spellings are consulted; `rope_type` wins when they agree in - /// meaning, and a disagreement is reported by [`Self::check`] rather than - /// silently resolved. - #[must_use] - pub fn kind(&self) -> String { - self.rope_type - .as_deref() - .or(self.legacy_type.as_deref()) - .unwrap_or("default") - .trim() - .to_ascii_lowercase() - } - - /// Is this object just "plain rotary, no scaling"? - #[must_use] - pub fn is_plain(&self) -> bool { - let k = self.kind(); - PLAIN_ROPE_TYPES.contains(&k.as_str()) - } - - /// Refuse anything that is not plain rotary, naming the mode. - /// - /// `whose` labels the source in the message (`"qwen2 config.json"`, - /// `"GGUF qwen2.rope.scaling.type"`) so a consumer knows which file to - /// look at. Returns `Ok(())` for a plain/absent mode — the only case the - /// shared blocks actually compute. - pub fn check(&self, whose: &str) -> Result<(), String> { - debug_assert!(!whose.is_empty(), "check: `whose` must name a source"); - // A per-layer-type map (Gemma 3 and friends) names no mode of its own, - // so it would otherwise read as plain. Refuse on the nesting itself — - // whatever the sub-objects say, per-layer RoPE is not implemented. - let nested: Vec<&str> = self - .extra - .iter() - .filter(|(_, v)| v.is_object()) - .map(|(k, _)| k.as_str()) - .collect(); - if !nested.is_empty() { - return Err(format!( - "{whose}: rope_scaling nests a RoPE object per layer type ({}) — Mummu computes one rotary table for every layer, so a per-layer-type map cannot be honored", - nested.join(", ") - )); - } - // Both spellings present and disagreeing: neither is safe to believe. - if let (Some(new), Some(old)) = (self.rope_type.as_deref(), self.legacy_type.as_deref()) - && !new.trim().eq_ignore_ascii_case(old.trim()) - { - return Err(format!( - "{whose}: rope_scaling names two different modes — rope_type '{new}' vs type \ - '{old}'; refusing rather than guessing which one the weights were trained with" - )); - } - if self.is_plain() { - return Ok(()); - } - let kind = self.kind(); - let factor = self - .factor - .map_or_else(|| "unspecified".to_string(), |f| format!("{f}")); - Err(format!( - "{whose}: rope_scaling mode '{kind}' (factor {factor}) is not implemented — Mummu \ - computes plain rotary frequencies, so this checkpoint would load clean and degrade \ - numerically past its original context instead of failing. Run it at its unscaled \ - context, or wait for scaled RoPE (ROADMAP P2)" - )) - } - - /// Read `.rope.scaling.*` out of a GGUF header. `None` when the - /// header declares no scaling at all — llama.cpp omits the keys for an - /// unscaled model rather than writing `"none"`. - #[must_use] - pub fn from_gguf(f: &GgufFile, arch: &str) -> Option { - debug_assert!(!arch.is_empty(), "from_gguf: arch must be named"); - let ty = f - .get(&format!("{arch}.rope.scaling.type")) - .and_then(GgufValue::as_str) - .map(str::to_owned); - let factor = f - .get(&format!("{arch}.rope.scaling.factor")) - .and_then(GgufValue::as_f32) - .map(f64::from); - let original = f - .get(&format!("{arch}.rope.scaling.original_context_length")) - .and_then(GgufValue::as_u64) - .and_then(|v| usize::try_from(v).ok()); - if ty.is_none() && factor.is_none() && original.is_none() { - return None; - } - Some(Self { - rope_type: ty, - legacy_type: None, - factor, - original_max_position_embeddings: original, - extra: std::collections::BTreeMap::new(), - }) - } -} - -/// Refuse an **enabled** sliding window, naming the span. -/// -/// `enabled` is the family's own answer, not a guess from `window.is_some()`: -/// Qwen2/Qwen3 gate their `sliding_window` behind `use_sliding_window` and -/// ship it `false`, while a Gemma-style config has no such flag and a present -/// window is live. Passing `enabled: false` with a `Some(window)` is the -/// normal, correct call for every Qwen checkpoint in the zoo. -/// -/// A window at least as long as the trained context is also inert — it can -/// never clip a position the model is allowed to reach — so it is accepted -/// even when enabled, with `max_positions` supplying that ceiling. -pub fn check_sliding_window( - window: Option, - enabled: bool, - max_positions: Option, - whose: &str, -) -> Result<(), String> { - debug_assert!(!whose.is_empty(), "check_sliding_window: name a source"); - let Some(window) = window else { - return Ok(()); - }; - if !enabled { - return Ok(()); - } - if window == 0 { - return Err(format!( - "{whose}: sliding_window is enabled but zero — no query could attend to any key" - )); - } - // A window that spans the whole trained context masks nothing. - if max_positions.is_some_and(|max| window >= max) { - return Ok(()); - } - Err(format!( - "{whose}: sliding-window attention is enabled (window {window} tokens) but Mummu builds a \ - full causal mask — past {window} tokens this checkpoint would attend to keys the trained \ - model masks, and answer wrong without failing. Run it under {window} tokens of context, \ - or wait for windowed attention (ROADMAP P2)" - )) -} - -/// The sliding-window span a GGUF header declares (`.attention. -/// sliding_window`), if any. llama.cpp writes the key only for architectures -/// that use it, so `None` means "full attention" — the same convention the -/// `rope.scaling.*` keys follow. -#[must_use] -pub fn sliding_window_from_gguf(f: &GgufFile, arch: &str) -> Option { - debug_assert!( - !arch.is_empty(), - "sliding_window_from_gguf: arch must be named" - ); - f.get(&format!("{arch}.attention.sliding_window")) - .and_then(GgufValue::as_u64) - .and_then(|v| usize::try_from(v).ok()) - .filter(|&w| w > 0) -} - -#[cfg(test)] -mod tests { - use super::*; - - fn parse(json: &str) -> RopeScaling { - serde_json::from_str(json).expect("rope_scaling parses") - } - - #[test] - fn absent_mode_reads_as_plain_and_passes() { - let s = RopeScaling::default(); - assert_eq!(s.kind(), "default"); - assert!(s.is_plain()); - assert!(s.check("test").is_ok()); - } - - #[test] - fn plain_spellings_all_pass() { - for ty in ["default", "none", "NULL", " Default "] { - let s = parse(&format!(r#"{{"rope_type": "{ty}"}}"#)); - assert!(s.is_plain(), "{ty} should read as plain"); - assert!(s.check("test").is_ok()); - } - } - - #[test] - fn yarn_is_refused_and_the_message_names_the_mode_and_factor() { - let s = parse( - r#"{"rope_type": "yarn", "factor": 4.0, "original_max_position_embeddings": 32768}"#, - ); - assert_eq!(s.kind(), "yarn"); - assert_eq!(s.factor, Some(4.0)); - assert_eq!(s.original_max_position_embeddings, Some(32768)); - let err = s - .check("qwen2 config.json") - .expect_err("yarn must be refused"); - assert!(err.contains("qwen2 config.json"), "{err}"); - assert!(err.contains("yarn"), "{err}"); - assert!(err.contains('4'), "{err}"); - } - - #[test] - fn the_legacy_type_spelling_is_read_too() { - let s = parse(r#"{"type": "linear", "factor": 2.0}"#); - assert_eq!(s.kind(), "linear"); - assert!(s.check("test").is_err()); - } - - #[test] - fn llama3_and_dynamic_are_refused_by_name() { - for ty in ["llama3", "dynamic", "longrope", "someone_invented_this"] { - let s = parse(&format!(r#"{{"rope_type": "{ty}"}}"#)); - let err = s.check("test").expect_err("non-plain must be refused"); - assert!(err.contains(ty), "message should name '{ty}': {err}"); - } - } - - #[test] - fn two_disagreeing_spellings_are_refused_rather_than_arbitrated() { - let s = parse(r#"{"rope_type": "yarn", "type": "linear"}"#); - let err = s.check("test").expect_err("disagreement must be refused"); - assert!(err.contains("yarn") && err.contains("linear"), "{err}"); - // The same mode written twice is not a disagreement. - let same = parse(r#"{"rope_type": "Linear", "type": "linear"}"#); - let err = same - .check("test") - .expect_err("linear is still unimplemented"); - assert!(!err.contains("two different modes"), "{err}"); - } - - /// The hole this `extra` map exists to close: a Gemma-3-style per-layer - /// map names no mode, so without the nesting check it would deserialize to - /// an all-default struct and read as plain rotary. - #[test] - fn a_per_layer_type_rope_map_is_refused_rather_than_read_as_plain() { - let s = parse( - r#"{"full_attention": {"rope_type": "dynamic", "factor": 8.0}, - "sliding_attention": {"rope_type": "default"}}"#, - ); - assert_eq!(s.rope_type, None, "the nested map names no top-level mode"); - let err = s - .check("gemma3 config.json") - .expect_err("nesting must refuse"); - assert!(err.contains("full_attention"), "{err}"); - assert!(err.contains("sliding_attention"), "{err}"); - } - - /// Scalar extras (YaRN's `beta_fast`, an `attention_factor`) are NOT - /// nesting, and must not be mistaken for it — the mode still decides. - #[test] - fn scalar_extras_do_not_trip_the_nesting_check() { - let s = parse(r#"{"rope_type": "default", "beta_fast": 32, "attention_factor": 1.0}"#); - assert!(s.check("test").is_ok(), "scalar extras are harmless"); - let yarn = parse(r#"{"rope_type": "yarn", "factor": 4.0, "beta_slow": 1}"#); - let err = yarn.check("test").expect_err("yarn is still refused"); - assert!(err.contains("yarn"), "{err}"); - } - - #[test] - fn an_inert_window_loads_and_an_enabled_one_does_not() { - // The Qwen2.5 shape: present, but `use_sliding_window: false`. - assert!(check_sliding_window(Some(32768), false, Some(32768), "qwen2").is_ok()); - // Enabled and shorter than the trained context: refused, span named. - let err = check_sliding_window(Some(4096), true, Some(131_072), "mistral") - .expect_err("an enabled window must be refused"); - assert!(err.contains("4096"), "{err}"); - assert!(err.contains("mistral"), "{err}"); - // No window at all is the common case. - assert!(check_sliding_window(None, true, Some(4096), "x").is_ok()); - } - - #[test] - fn a_window_spanning_the_whole_context_masks_nothing() { - assert!(check_sliding_window(Some(4096), true, Some(4096), "x").is_ok()); - assert!(check_sliding_window(Some(8192), true, Some(4096), "x").is_ok()); - assert!(check_sliding_window(Some(4095), true, Some(4096), "x").is_err()); - // Unknown ceiling: cannot prove it is inert, so it is refused. - assert!(check_sliding_window(Some(4096), true, None, "x").is_err()); - } - - #[test] - fn a_zero_window_is_refused_as_degenerate() { - let err = check_sliding_window(Some(0), true, Some(4096), "x") - .expect_err("a zero window is nonsense"); - assert!(err.contains("zero"), "{err}"); - } -} diff --git a/crates/mummu/examples/src/backend.rs b/crates/mummu/examples/src/backend.rs deleted file mode 100644 index 4393dfc..0000000 --- a/crates/mummu/examples/src/backend.rs +++ /dev/null @@ -1,938 +0,0 @@ -//! Backend selection: one binary compiles BOTH backends and picks at runtime. -//! -//! The model code is generic over `B: Backend`; nothing here forces a device -//! choice on a consumer. What this module owns is the *default* policy proven -//! in laurelane: enumerate GPU adapters once with a cheap `wgpu` probe (no -//! device creation), run on **wgpu (GPU via Vulkan/DX12/Metal — no CUDA -//! toolchain)** when a hardware adapter is present, else **burn-flex (CPU)**. -//! No feature-split builds. -//! -//! The probe also records which adapters advertise `SHADER_F16` — the input -//! the hardware planner (P6) uses to decide whether the f16 element type is -//! viable on this machine. - -use once_cell::sync::OnceCell; - -/// The default GPU device (wgpu: Vulkan / DX12 / Metal). burn 0.22 selects -/// backends at runtime through [`burn::tensor::Device`]; with the workspace -/// `fusion` feature, fusion applies to supporting devices automatically. -#[must_use] -pub fn gpu_device() -> burn::tensor::Device { - burn::tensor::Device::wgpu(Default::default()) -} - -/// Raise every cubecl device-server thread ("DSD-*") above the compute -/// pools. Those threads encode command buffers, submit to the driver, and -/// signal readback-map completions — microseconds of CPU each — but at -/// normal priority they starve behind the trunk's spinning gemm workers: -/// measured, a remote FFN group whose kernels total well under 1 ms still -/// held its caller ~26 ms at the fence, and the wait tracked scheduler -/// quanta, not GPU time. Call after model load (the servers spawn on first -/// device use); repeat calls are cheap and idempotent. -pub fn boost_device_server_threads() { - #[cfg(windows)] - unsafe { - #[link(name = "kernel32.dll", kind = "raw-dylib", modifiers = "+verbatim")] - unsafe extern "system" { - fn CreateToolhelp32Snapshot(flags: u32, pid: u32) -> isize; - fn Thread32First(snap: isize, entry: *mut ThreadEntry32) -> i32; - fn Thread32Next(snap: isize, entry: *mut ThreadEntry32) -> i32; - fn OpenThread(access: u32, inherit: i32, tid: u32) -> isize; - fn GetThreadDescription(handle: isize, desc: *mut *mut u16) -> i32; - fn SetThreadPriority(handle: isize, priority: i32) -> i32; - fn CloseHandle(handle: isize) -> i32; - fn GetCurrentProcessId() -> u32; - fn LocalFree(mem: isize) -> isize; - } - #[repr(C)] - struct ThreadEntry32 { - size: u32, - usage: u32, - thread_id: u32, - owner_pid: u32, - base_pri: i32, - delta_pri: i32, - flags: u32, - } - const TH32CS_SNAPTHREAD: u32 = 0x4; - const THREAD_SET_INFORMATION: u32 = 0x20; - const THREAD_QUERY_LIMITED_INFORMATION: u32 = 0x800; - const ABOVE_NORMAL: i32 = 1; - - let pid = GetCurrentProcessId(); - let snap = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0); - if snap == -1 || snap == 0 { - return; - } - let mut entry = ThreadEntry32 { - size: size_of::() as u32, - usage: 0, - thread_id: 0, - owner_pid: 0, - base_pri: 0, - delta_pri: 0, - flags: 0, - }; - let mut boosted = 0u32; - let mut ok = Thread32First(snap, &raw mut entry); - while ok != 0 { - if entry.owner_pid == pid { - let h = OpenThread( - THREAD_SET_INFORMATION | THREAD_QUERY_LIMITED_INFORMATION, - 0, - entry.thread_id, - ); - if h != 0 { - let mut desc: *mut u16 = std::ptr::null_mut(); - if GetThreadDescription(h, &raw mut desc) >= 0 && !desc.is_null() { - let mut len = 0usize; - while *desc.add(len) != 0 { - len += 1; - } - let name = String::from_utf16_lossy(std::slice::from_raw_parts(desc, len)); - if name.starts_with("DSD") { - SetThreadPriority(h, ABOVE_NORMAL); - boosted += 1; - } - LocalFree(desc as isize); - } - CloseHandle(h); - } - } - ok = Thread32Next(snap, &raw mut entry); - } - CloseHandle(snap); - if boosted > 0 { - eprintln!("[mummu] raised {boosted} device-server thread(s) above the compute pools"); - } - } -} - -/// Move a tensor to `device`, staging through host memory when both ends are -/// GPUs. -/// -/// cubecl does not implement peer-to-peer transfer for wgpu: `comm_init` and -/// `send` on its server trait are `unimplemented!()`, so a direct -/// discrete-GPU -> integrated-GPU move panics with a bare "not implemented" -/// (cubecl-runtime `server/base.rs`). Staging through the host is the only -/// portable path, and it is what makes a placement spanning two GPUs work at -/// all. -/// -/// A same-device move is a no-op, and a move with the host at either end is a -/// single transfer, so this costs nothing on the common paths. -#[must_use] -pub fn move_to( - tensor: burn::tensor::Tensor, - device: &burn::tensor::Device, -) -> burn::tensor::Tensor { - let from = tensor.device(); - if from == *device { - return tensor; - } - if is_accelerator(&from) && is_accelerator(device) { - return tensor.to_device(&cpu_device()).to_device(device); - } - tensor.to_device(device) -} - -/// Is this an accelerator (as opposed to the host)? burn 0.22 selects -/// backends by runtime `Device` value and exposes no kind accessor, so the -/// debug form is the handle available. -fn is_accelerator(device: &burn::tensor::Device) -> bool { - let name = format!("{device:?}"); - name.contains("Wgpu") || name.contains("Cuda") -} - -/// The integrated GPU, when one exists. -/// -/// Addressed explicitly because [`gpu_device`] resolves to wgpu's default, -/// which is the *discrete* card on any machine that has one — so the -/// integrated adapter is invisible to a placement that only ever asks for -/// "the GPU", however much idle capacity it has. -#[must_use] -pub fn integrated_gpu_device() -> burn::tensor::Device { - burn::tensor::Device::wgpu(burn::tensor::DeviceKind::IntegratedGpu(0)) -} - -/// Does this machine expose an integrated GPU distinct from the discrete one? -#[must_use] -pub fn has_integrated_gpu() -> bool { - inventory() - .gpus - .iter() - .any(|g| g.device_type == wgpu::DeviceType::IntegratedGpu) -} - -/// The CPU device (burn-flex: pure-Rust SIMD + gemm). -#[must_use] -pub fn cpu_device() -> burn::tensor::Device { - burn::tensor::Device::flex() -} - -/// The CUDA device (feature `cuda`). NVRTC compiles kernels at runtime — the -/// WSL2-container GPU path where no correct Vulkan reaches the process. -#[cfg(feature = "cuda")] -#[must_use] -pub fn cuda_device() -> burn::tensor::Device { - burn::tensor::Device::cuda(0) -} - -/// The float dtype mummu creates tensors with **on `device`**. -/// -/// burn 0.22 moved the element type off the backend type and onto the device -/// as a runtime setting ([`burn::tensor::Device::configure`]), so the precision -/// a model runs in is a property of the device it was handed - not of a type -/// alias, and not of a process-wide constant. Reading it back here is what lets -/// one process hold an f16 GPU device beside an f32 host device, which the 0.21 -/// `Gpu`/`GpuF16` alias split could not express at all. -/// -/// Tensor-creation sites still name a dtype **explicitly** - the 0.21 rationale -/// is unchanged. What changed is where the answer comes from. -#[must_use] -pub fn float_dtype(device: &burn::tensor::Device) -> burn::tensor::DType { - let dtype: burn::tensor::DType = device.settings().float_dtype.into(); - debug_assert!( - dtype.is_float(), - "a device's float setting must be a float dtype, got {dtype:?}" - ); - dtype -} - -/// Int dtype counterpart of [`float_dtype`], read from the same device. -#[must_use] -pub fn int_dtype(device: &burn::tensor::Device) -> burn::tensor::DType { - let dtype: burn::tensor::DType = device.settings().int_dtype.into(); - debug_assert!( - dtype.is_int(), - "a device's int setting must be an int dtype, got {dtype:?}" - ); - dtype -} - -/// A GPU device configured to compute in **f16**. -/// -/// The 0.22 replacement for the `GpuF16` type alias. Device settings lock on -/// first use and cannot be changed afterwards, so this must run before any -/// tensor exists on the discrete GPU - which is why every f16 gate lives in its -/// own test binary, exactly as it did under the one-alias-per-process rule. -/// -/// # Errors -/// -/// `AlreadyInitialized` when the device has already computed something, and an -/// unsupported-dtype error when the adapter cannot do f16 at all (check -/// [`DeviceInventory::any_shader_f16`] first). -pub fn gpu_device_f16() -> Result { - let mut device = gpu_device(); - match device.configure((burn::tensor::FloatDType::F16, burn::tensor::IntDType::I32)) { - Ok(()) => {} - // Idempotent on purpose: several f16 gates share one process, and the - // second caller must get the same device rather than an error - but - // ONLY when the lock actually landed on f16. An f32-locked device is - // reported, never handed back wearing an f16 label (the 2026-07-11 - // mislabelling bug is exactly this branch going the other way). - Err(e) => { - if float_dtype(&device) != burn::tensor::DType::F16 { - return Err(e); - } - } - } - assert_eq!( - float_dtype(&device), - burn::tensor::DType::F16, - "gpu_device_f16 must return an f16 device or an error, never an f32 device" - ); - Ok(device) -} - -/// One enumerated GPU adapter, as reported by wgpu. -#[derive(Debug, Clone)] -pub struct GpuAdapter { - /// Driver-reported adapter name, e.g. `"NVIDIA GeForce RTX 4070 Ti SUPER"`. - pub name: String, - /// Graphics API carrying this adapter (`Vulkan`, `Dx12`, `Metal`, ...). - pub backend: wgpu::Backend, - /// Discrete / integrated / virtual — never [`wgpu::DeviceType::Cpu`] - /// (software rasterizers are filtered out of the inventory). - pub device_type: wgpu::DeviceType, - /// Does this adapter advertise `SHADER_F16` (native f16 shader arithmetic)? - pub shader_f16: bool, - /// Largest single buffer this adapter permits — a hard bound the placement - /// planner respects per tensor/shard. (True VRAM capacity is NOT exposed - /// portably by wgpu; querying it per-API via wgpu-hal is a P6 follow-up.) - pub max_buffer_bytes: u64, - /// Dedicated video memory (true VRAM capacity) — the P6 planner's fit - /// budget. wgpu exposes no portable query (gfx-rs/wgpu#2447), so this is - /// filled per-OS: DXGI `DedicatedVideoMemory` on Windows (matched to the - /// wgpu adapter by name, covering every API's view of the same card); - /// `None` where no query is implemented yet (Linux/macOS follow-ups). - /// Integrated GPUs legitimately report small values here — their real - /// budget is shared system RAM ([`CpuInfo::total_ram_bytes`]). - pub vram_bytes: Option, -} - -/// The host CPU as a compute device (the `burn-flex` target and the P6 -/// offload pool). -#[derive(Debug, Clone)] -pub struct CpuInfo { - /// Logical cores (SMT threads) available to this process; at least 1. - pub logical_cores: usize, - /// Total physical RAM; `None` where no query is implemented yet (macOS). - pub total_ram_bytes: Option, -} - -impl Default for CpuInfo { - fn default() -> Self { - Self { - logical_cores: 1, - total_ram_bytes: None, - } - } -} - -/// Every hardware GPU visible to wgpu plus the host CPU, enumerated once per -/// process — the device set the P6 hardware planner (and app settings UIs) -/// read. -#[derive(Debug, Clone, Default)] -pub struct DeviceInventory { - /// Hardware adapters across the primary graphics APIs. The same physical - /// card appears once per API that exposes it (e.g. Vulkan AND DX12) — - /// deliberate, because features like `SHADER_F16` differ per API. - pub gpus: Vec, - /// The host CPU (cores + RAM). - pub cpu: CpuInfo, -} - -impl DeviceInventory { - /// Is at least one hardware GPU present (on any API)? - #[must_use] - pub fn has_gpu(&self) -> bool { - !self.gpus.is_empty() - } - - /// Does any adapter advertise `SHADER_F16`? Gates [`gpu_device_f16`]. - #[must_use] - pub fn any_shader_f16(&self) -> bool { - self.gpus.iter().any(|g| g.shader_f16) - } -} - -/// Enumerate hardware adapters on `backends`. Cheap: adapter listing only, no -/// device creation. wgpu 29's enumeration is async, so block on it here. -fn enumerate(instance: &wgpu::Instance, backends: wgpu::Backends) -> Vec { - let vram = vram_by_adapter_name(); - pollster::block_on(instance.enumerate_adapters(backends)) - .into_iter() - .filter_map(|adapter| { - let info = adapter.get_info(); - if matches!(info.device_type, wgpu::DeviceType::Cpu) { - return None; // software rasterizer, not a hardware GPU - } - let shader_f16 = adapter.features().contains(wgpu::Features::SHADER_F16); - let vram_bytes = lookup_vram(&vram, &info.name); - Some(GpuAdapter { - name: info.name, - backend: info.backend, - device_type: info.device_type, - shader_f16, - max_buffer_bytes: adapter.limits().max_buffer_size, - vram_bytes, - }) - }) - .collect() -} - -/// Find `name`'s dedicated VRAM in the per-OS `(adapter name, bytes)` table. -/// Driver stacks decorate the same card's name slightly differently per API -/// (e.g. a `(TM)` suffix), so fall back to a case-insensitive prefix match in -/// either direction when the exact name misses. -fn lookup_vram(table: &[(String, u64)], name: &str) -> Option { - debug_assert!(!name.is_empty(), "wgpu adapters always carry a name"); - if let Some(&(_, bytes)) = table.iter().find(|(n, _)| n == name) { - return Some(bytes); - } - let lower = name.to_lowercase(); - table - .iter() - .find(|(n, _)| { - let n = n.to_lowercase(); - n.starts_with(&lower) || lower.starts_with(&n) - }) - .map(|&(_, bytes)| bytes) -} - -/// Minimal hand-bound slice of the DXGI 1.1 COM ABI — the one Windows API -/// that reports true VRAM capacity for every GPU regardless of which graphics -/// API wgpu reached it through. Bound by hand because `windows-sys` 0.60+ -/// dropped COM interface bindings and the full `windows` crate is a heavy -/// dependency for two vtable calls. The ABI is frozen (shipped with Windows 7, -/// 2009): vtables are declared slot-exact below and the dev-box unit test -/// (`windows_discrete_adapters_report_plausible_vram`) cross-checks the -/// numbers against reality. -#[cfg(windows)] -mod dxgi { - use core::ffi::c_void; - - #[repr(C)] - pub struct Guid { - data1: u32, - data2: u16, - data3: u16, - data4: [u8; 8], - } - - /// `IID_IDXGIFactory1` = `{770aae78-f26f-4dba-a829-253c83d1b387}`. - pub const IID_IDXGI_FACTORY1: Guid = Guid { - data1: 0x770a_ae78, - data2: 0xf26f, - data3: 0x4dba, - data4: [0xa8, 0x29, 0x25, 0x3c, 0x83, 0xd1, 0xb3, 0x87], - }; - - pub const DXGI_ERROR_NOT_FOUND: i32 = 0x887A_0002_u32 as i32; - - /// `IID_IDXGIAdapter3` = `{645967a4-1392-4310-a798-8053ce3e93fd}`. The - /// interface that reports the OS's *current* video-memory budget for - /// this process, which is what shrinks when another process takes VRAM. - pub const IID_IDXGI_ADAPTER3: Guid = Guid { - data1: 0x6459_67a4, - data2: 0x1392, - data3: 0x4310, - data4: [0xa7, 0x98, 0x80, 0x53, 0xce, 0x3e, 0x93, 0xfd], - }; - - /// `DXGI_MEMORY_SEGMENT_GROUP_LOCAL` — the adapter's own VRAM, as - /// opposed to system memory it may spill into. - pub const MEMORY_SEGMENT_LOCAL: u32 = 0; - - /// `DXGI_QUERY_VIDEO_MEMORY_INFO`, verbatim layout. - #[repr(C)] - #[derive(Default, Clone, Copy)] - pub struct QueryVideoMemoryInfo { - /// Bytes the OS is currently willing to let this process use. - pub budget: u64, - /// Bytes this process currently has resident. - pub current_usage: u64, - pub available_for_reservation: u64, - pub current_reservation: u64, - } - - /// `IDXGIAdapter3`'s vtable. Slot arithmetic, in order: IUnknown (3), - /// IDXGIObject (4), IDXGIAdapter (3), IDXGIAdapter1 (1: GetDesc1), - /// IDXGIAdapter2 (1: GetDesc2), then IDXGIAdapter3's own six — of which - /// `QueryVideoMemoryInfo` is the third. Getting this count wrong calls - /// the wrong function pointer, so it is spelled out rather than padded. - #[repr(C)] - pub struct Adapter3Vtbl { - _query_interface: usize, - _add_ref: usize, - pub release: unsafe extern "system" fn(*mut c_void) -> u32, - _idxgi_object: [usize; 4], - _idxgi_adapter: [usize; 3], - _get_desc1: usize, - _get_desc2: usize, - _register_teardown: usize, - _unregister_teardown: usize, - pub query_video_memory_info: - unsafe extern "system" fn(*mut c_void, u32, u32, *mut QueryVideoMemoryInfo) -> i32, - _set_video_memory_reservation: usize, - _register_budget_change: usize, - _unregister_budget_change: usize, - } - - /// `DXGI_ADAPTER_DESC1`, verbatim layout. - #[repr(C)] - pub struct AdapterDesc1 { - pub description: [u16; 128], - pub vendor_id: u32, - pub device_id: u32, - pub sub_sys_id: u32, - pub revision: u32, - pub dedicated_video_memory: usize, - pub dedicated_system_memory: usize, - pub shared_system_memory: usize, - pub adapter_luid: [u32; 2], - pub flags: u32, - } - - /// `IDXGIFactory1`'s vtable: IUnknown (3 slots) + IDXGIObject (4) + - /// IDXGIFactory (5) + IDXGIFactory1 (2). Uncalled slots are opaque. - #[repr(C)] - pub struct Factory1Vtbl { - _query_interface: usize, - _add_ref: usize, - pub release: unsafe extern "system" fn(*mut c_void) -> u32, - _idxgi_object: [usize; 4], - _idxgi_factory: [usize; 5], - pub enum_adapters1: unsafe extern "system" fn(*mut c_void, u32, *mut *mut c_void) -> i32, - _is_current: usize, - } - - /// `IDXGIAdapter1`'s vtable: IUnknown (3) + IDXGIObject (4) + - /// IDXGIAdapter (3) + IDXGIAdapter1 (1: GetDesc1). - #[repr(C)] - pub struct Adapter1Vtbl { - pub query_interface: - unsafe extern "system" fn(*mut c_void, *const Guid, *mut *mut c_void) -> i32, - _add_ref: usize, - pub release: unsafe extern "system" fn(*mut c_void) -> u32, - _idxgi_object: [usize; 4], - _idxgi_adapter: [usize; 3], - pub get_desc1: unsafe extern "system" fn(*mut c_void, *mut AdapterDesc1) -> i32, - } - - // raw-dylib: rustc generates the import stubs itself, so linking needs - // no Windows SDK .lib at all — the same mechanism windows-sys uses. - #[link(name = "dxgi.dll", kind = "raw-dylib", modifiers = "+verbatim")] - unsafe extern "system" { - pub fn CreateDXGIFactory1(riid: *const Guid, factory: *mut *mut c_void) -> i32; - } - - /// The vtable pointer every COM object starts with. - #[inline] - pub unsafe fn vtbl(object: *mut c_void) -> *const T { - debug_assert!(!object.is_null(), "COM object must be live"); - // SAFETY: caller guarantees `object` is a live COM interface pointer; - // its first field is the vtable pointer. - unsafe { *object.cast::<*const T>() } - } -} - -/// Every adapter's `(name, dedicated video memory)` as DXGI reports it: one -/// factory, a bounded adapter walk, everything released before returning. -/// Failures degrade to an empty table (VRAM stays `None`), never to a panic — -/// this runs once at inventory time. -#[cfg(windows)] -fn vram_by_adapter_name() -> Vec<(String, u64)> { - use dxgi::{Adapter1Vtbl, Factory1Vtbl}; - - /// More adapters than any real machine exposes; bounds the walk. - const MAX_ADAPTERS: u32 = 64; - - let mut factory: *mut core::ffi::c_void = core::ptr::null_mut(); - // SAFETY: CreateDXGIFactory1 writes a factory pointer on success; checked - // via the HRESULT before use. - let hr = unsafe { dxgi::CreateDXGIFactory1(&dxgi::IID_IDXGI_FACTORY1, &mut factory) }; - if hr < 0 || factory.is_null() { - return Vec::new(); - } - - let mut out = Vec::new(); - for index in 0..MAX_ADAPTERS { - let mut adapter: *mut core::ffi::c_void = core::ptr::null_mut(); - // SAFETY: factory is the live IDXGIFactory1 created above; - // EnumAdapters1 writes an adapter or returns DXGI_ERROR_NOT_FOUND. - let hr = unsafe { - ((*dxgi::vtbl::(factory)).enum_adapters1)(factory, index, &mut adapter) - }; - if hr == dxgi::DXGI_ERROR_NOT_FOUND { - break; - } - if hr < 0 || adapter.is_null() { - continue; - } - // SAFETY: adapter is the live IDXGIAdapter1 just handed out; GetDesc1 - // fills the struct; Release balances EnumAdapters1's reference. - let (hr, desc) = unsafe { - let mut desc = core::mem::zeroed::(); - let hr = ((*dxgi::vtbl::(adapter)).get_desc1)(adapter, &mut desc); - ((*dxgi::vtbl::(adapter)).release)(adapter); - (hr, desc) - }; - if hr < 0 { - continue; - } - let len = desc - .description - .iter() - .position(|&c| c == 0) - .unwrap_or(desc.description.len()); - let name = String::from_utf16_lossy(&desc.description[..len]); - let dedicated = desc.dedicated_video_memory as u64; - // Negative space: skip software adapters (they report ~0 dedicated - // VRAM; the "Microsoft Basic Render Driver") rather than risk mapping - // a real card's lookup onto them. - if !name.is_empty() && dedicated > 0 { - out.push((name, dedicated)); - } - } - // SAFETY: Release on the factory created above. - unsafe { ((*dxgi::vtbl::(factory)).release)(factory) }; - debug_assert!( - out.iter().all(|&(_, b)| b < 1u64 << 42), - "implausible VRAM in the DXGI table: {out:?}" - ); - out -} - -#[cfg(not(windows))] -fn vram_by_adapter_name() -> Vec<(String, u64)> { - Vec::new() // Linux (Vulkan memory heaps) and macOS are P6 follow-ups. -} - -/// How much VRAM the OS is *currently* willing to give this process, and how -/// much of it we already hold. -/// -/// This is the external-pressure signal: `budget` is not the card's size, it -/// is the driver's running allocation to this process, and it falls when -/// another process (a game, a second model, a browser compositing video) -/// takes memory. Placement uses it to decide how much of a model can stay -/// resident, and at what precision — see [`crate::mix`]. -/// -/// Distinct from the total in [`GpuAdapter::vram_bytes`], which never moves. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct VideoMemory { - /// Bytes the OS currently allows this process on the local segment. - pub budget: u64, - /// Bytes this process currently holds there. - pub current_usage: u64, -} - -impl VideoMemory { - /// Headroom before the driver starts demoting our allocations. Saturating: - /// usage legitimately exceeds budget when the OS has just cut it, and that - /// means *zero* headroom, not a huge negative one. - #[must_use] - pub fn headroom(self) -> u64 { - self.budget.saturating_sub(self.current_usage) - } -} - -/// Query the discrete GPU's current video-memory budget. -/// -/// Two things to know about the numbers. **`current_usage` is this process -/// only** — it reads 0 from a program that has allocated nothing, even while -/// the card is full. Other processes do not appear there; they appear as a -/// *smaller `budget`*, which is exactly the pressure signal we want. -/// **Adapter choice is by dedicated VRAM, not by budget**: an integrated -/// GPU's "local" segment is system RAM, so on this box the iGPU reports a -/// ~101 GiB budget and would win any largest-budget contest. -/// -/// `None` when the platform or driver will not say (non-Windows today, or a -/// pre-DXGI-1.4 adapter). Callers must treat `None` as "no information" and -/// hold their current placement rather than assuming either plenty or -/// pressure — guessing in either direction is worse than not adapting. -#[cfg(windows)] -#[must_use] -pub fn video_memory() -> Option { - use dxgi::{Adapter1Vtbl, Adapter3Vtbl, Factory1Vtbl}; - - const MAX_ADAPTERS: u32 = 64; - let mut factory: *mut core::ffi::c_void = core::ptr::null_mut(); - // SAFETY: writes a factory pointer on success, checked via the HRESULT. - let hr = unsafe { dxgi::CreateDXGIFactory1(&dxgi::IID_IDXGI_FACTORY1, &mut factory) }; - if hr < 0 || factory.is_null() { - return None; - } - - // (dedicated VRAM, budget) — the largest dedicated wins, see the note above. - let mut best: Option<(u64, VideoMemory)> = None; - for index in 0..MAX_ADAPTERS { - let mut adapter: *mut core::ffi::c_void = core::ptr::null_mut(); - // SAFETY: factory is the live IDXGIFactory1 created above. - let hr = unsafe { - ((*dxgi::vtbl::(factory)).enum_adapters1)(factory, index, &mut adapter) - }; - if hr == dxgi::DXGI_ERROR_NOT_FOUND { - break; - } - if hr < 0 || adapter.is_null() { - continue; - } - // SAFETY: `adapter` is a live IDXGIAdapter1. QueryInterface either - // hands back a live IDXGIAdapter3 or leaves the pointer null; both - // references are released before the loop continues. - let info = unsafe { - // The adapter's fixed VRAM, used only to tell discrete from - // integrated; the budget itself comes from IDXGIAdapter3. - let mut desc = core::mem::zeroed::(); - let dedicated = - if ((*dxgi::vtbl::(adapter)).get_desc1)(adapter, &mut desc) >= 0 { - desc.dedicated_video_memory as u64 - } else { - 0 - }; - let mut adapter3: *mut core::ffi::c_void = core::ptr::null_mut(); - let hr = ((*dxgi::vtbl::(adapter)).query_interface)( - adapter, - &dxgi::IID_IDXGI_ADAPTER3, - &mut adapter3, - ); - let info = if hr >= 0 && !adapter3.is_null() { - let mut info = dxgi::QueryVideoMemoryInfo::default(); - let vt = dxgi::vtbl::(adapter3); - let hr = ((*vt).query_video_memory_info)( - adapter3, - 0, // node 0: single-GPU adapters have exactly one - dxgi::MEMORY_SEGMENT_LOCAL, - &mut info, - ); - ((*vt).release)(adapter3); - (hr >= 0).then_some(info) - } else { - None - }; - ((*dxgi::vtbl::(adapter)).release)(adapter); - info.map(|i| (dedicated, i)) - }; - - if let Some((dedicated, info)) = info { - let seen = VideoMemory { - budget: info.budget, - current_usage: info.current_usage, - }; - if best.is_none_or(|(d, _)| dedicated > d) { - best = Some((dedicated, seen)); - } - } - } - // SAFETY: `factory` is live and owned here; this balances its creation. - unsafe { ((*dxgi::vtbl::(factory)).release)(factory) }; - best.map(|(_, v)| v) -} - -/// No portable budget query off Windows yet — Vulkan's -/// `VK_EXT_memory_budget` is the equivalent and a P6 follow-up. -#[cfg(not(windows))] -#[must_use] -pub fn video_memory() -> Option { - None -} - -/// Total physical RAM, per platform. Kept syscall-thin — this runs once, at -/// inventory time. -#[cfg(windows)] -fn total_ram_bytes() -> Option { - use windows_sys::Win32::System::SystemInformation::{GlobalMemoryStatusEx, MEMORYSTATUSEX}; - let mut status = MEMORYSTATUSEX { - dwLength: core::mem::size_of::() as u32, - dwMemoryLoad: 0, - ullTotalPhys: 0, - ullAvailPhys: 0, - ullTotalPageFile: 0, - ullAvailPageFile: 0, - ullTotalVirtual: 0, - ullAvailVirtual: 0, - ullAvailExtendedVirtual: 0, - }; - // SAFETY: `status` is a live, writable MEMORYSTATUSEX with dwLength set, - // exactly what the API contract requires. - let ok = unsafe { GlobalMemoryStatusEx(&mut status) }; - (ok != 0).then_some(status.ullTotalPhys) -} - -#[cfg(target_os = "linux")] -fn total_ram_bytes() -> Option { - let meminfo = std::fs::read_to_string("/proc/meminfo").ok()?; - let kib: u64 = meminfo - .lines() - .find(|l| l.starts_with("MemTotal:"))? - .split_whitespace() - .nth(1)? - .parse() - .ok()?; - Some(kib * 1024) -} - -#[cfg(not(any(windows, target_os = "linux")))] -fn total_ram_bytes() -> Option { - None // macOS et al.: a sysctl query is a P6 follow-up -} - -/// The host CPU: logical cores + total RAM. -fn cpu_info() -> CpuInfo { - let logical_cores = std::thread::available_parallelism() - .map(std::num::NonZero::get) - .unwrap_or(1); - let total_ram_bytes = total_ram_bytes(); - assert!( - logical_cores >= 1, - "a running process has at least one core" - ); - // Negative space: an answer below 64 MiB is a parse/API bug, not a machine. - debug_assert!( - total_ram_bytes.is_none_or(|b| b >= 64 << 20), - "implausible total RAM: {total_ram_bytes:?}" - ); - CpuInfo { - logical_cores, - total_ram_bytes, - } -} - -/// The process-lifetime device inventory. Enumerated once (first call pays -/// ~tens of milliseconds); every later call is a cache read. -pub fn inventory() -> &'static DeviceInventory { - static INVENTORY: OnceCell = OnceCell::new(); - INVENTORY.get_or_init(|| { - let instance = wgpu::Instance::default(); - let gpus = enumerate(&instance, wgpu::Backends::PRIMARY); - let inv = DeviceInventory { - gpus, - cpu: cpu_info(), - }; - // Positive space: every inventoried adapter is real hardware. - debug_assert!( - inv.gpus - .iter() - .all(|g| !matches!(g.device_type, wgpu::DeviceType::Cpu)), - "CPU adapters must be filtered out of the GPU inventory" - ); - inv - }) -} - -/// Default device policy: run on the GPU when a hardware adapter is present. -/// Stable for the process lifetime (backed by [`inventory`]). -#[must_use] -pub fn use_gpu() -> bool { - let gpu = inventory().has_gpu(); - // Negative space: the decision must agree with the inventory it came from. - debug_assert!(gpu == !inventory().gpus.is_empty()); - gpu -} - -/// Human label for where the default policy will run. -#[must_use] -pub fn device_label() -> &'static str { - if use_gpu() { - "GPU (wgpu)" - } else { - "CPU (flex)" - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn inventory_is_cached_and_consistent() { - let first = inventory(); - let second = inventory(); - // Same allocation: the OnceCell caches, never re-enumerates. - assert!(std::ptr::eq(first, second)); - assert_eq!(first.has_gpu(), !first.gpus.is_empty()); - } - - #[test] - fn no_cpu_adapters_in_inventory() { - assert!( - inventory() - .gpus - .iter() - .all(|g| !matches!(g.device_type, wgpu::DeviceType::Cpu)) - ); - } - - #[test] - fn device_label_matches_policy() { - let label = device_label(); - if use_gpu() { - assert_eq!(label, "GPU (wgpu)"); - } else { - assert_eq!(label, "CPU (flex)"); - } - } - - #[test] - fn f16_gate_requires_a_gpu() { - let inv = inventory(); - // Negative space: SHADER_F16 can't be advertised with no GPUs at all. - if inv.any_shader_f16() { - assert!(inv.has_gpu()); - } - } - - #[test] - fn cpu_inventory_reports_cores_and_plausible_ram() { - let cpu = &inventory().cpu; - assert!(cpu.logical_cores >= 1); - // Windows and Linux have a RAM query; its answer must be a real - // machine's (1 GiB ..= 64 TiB), not a unit slip. - if cfg!(any(windows, target_os = "linux")) { - let ram = cpu.total_ram_bytes.expect("RAM query exists here"); - assert!((1 << 30..=1u64 << 46).contains(&ram), "implausible: {ram}"); - } - } - - #[test] - fn windows_discrete_adapters_report_plausible_vram() { - // The DXGI walk exists on Windows: every *discrete* adapter must get - // a dedicated-VRAM figure inside a real card's range (256 MiB..4 TiB). - // Integrated/virtual adapters may report None (name mismatch) or a - // small carve-out — both fine, the planner budgets those via RAM. - for gpu in &inventory().gpus { - if cfg!(windows) && matches!(gpu.device_type, wgpu::DeviceType::DiscreteGpu) { - let vram = gpu - .vram_bytes - .unwrap_or_else(|| panic!("{}: no DXGI VRAM match", gpu.name)); - assert!( - (256 << 20..1u64 << 42).contains(&vram), - "{}: implausible VRAM {vram}", - gpu.name - ); - } - } - } - - #[test] - fn vram_lookup_matches_exact_then_prefix() { - let table = vec![ - ("NVIDIA GeForce RTX 4070 Ti SUPER".to_string(), 16 << 30), - ("AMD Radeon(TM) Graphics".to_string(), 512 << 20), - ]; - // Exact hit. - assert_eq!( - lookup_vram(&table, "NVIDIA GeForce RTX 4070 Ti SUPER"), - Some(16 << 30) - ); - // Per-API name decoration: prefix in either direction, any case. - assert_eq!( - lookup_vram(&table, "AMD Radeon(TM) Graphics (RADV)"), - Some(512 << 20) - ); - assert_eq!(lookup_vram(&table, "amd radeon(tm)"), Some(512 << 20)); - // Negative space: a different card must never borrow a table entry. - assert_eq!(lookup_vram(&table, "Intel(R) Arc(TM) A770"), None); - } - - #[test] - fn adapters_report_a_usable_buffer_bound() { - // Every real adapter permits at least the WebGPU floor (256 MiB); - // a smaller answer means the limits plumbing broke. - for gpu in &inventory().gpus { - assert!( - gpu.max_buffer_bytes >= 256 << 20, - "{}: max_buffer_bytes {} below the WebGPU floor", - gpu.name, - gpu.max_buffer_bytes - ); - } - } - - /// Print the inventory so the nightly log records what this machine has. - #[test] - fn report_inventory() { - for gpu in &inventory().gpus { - eprintln!( - "[mummu] {:?} / {} ({:?}): SHADER_F16 = {}, max buffer {:.1} GiB, VRAM {}", - gpu.backend, - gpu.name, - gpu.device_type, - gpu.shader_f16, - gpu.max_buffer_bytes as f64 / f64::from(1u32 << 30), - gpu.vram_bytes.map_or("unknown".into(), |b| format!( - "{:.1} GiB", - b as f64 / f64::from(1u32 << 30) - )), - ); - } - let cpu = &inventory().cpu; - eprintln!( - "[mummu] CPU: {} logical cores, RAM {:?} GiB", - cpu.logical_cores, - cpu.total_ram_bytes.map(|b| b >> 30), - ); - eprintln!("[mummu] policy: {}", device_label()); - } -} diff --git a/crates/mummu/examples/src/cache.rs b/crates/mummu/examples/src/cache.rs deleted file mode 100644 index 01afda4..0000000 --- a/crates/mummu/examples/src/cache.rs +++ /dev/null @@ -1,277 +0,0 @@ -//! Process-lifetime model caching. Loading a checkpoint costs seconds and -//! gigabytes, so consumers keep one [`ModelSlot`] static per (model, backend) -//! and pay the load once. Burn's `Param` is not `Sync`, so the loaded value -//! lives behind a `Mutex` and is only reachable inside [`ModelSlot::with`] / -//! [`ModelSlot::with_async`] — which also serializes inference, the right -//! default for a single GPU. -//! -//! The mutex is tokio's, because the async accessor holds its guard across -//! an await (a `std` guard is not `Send`, so it would not survive one). The -//! sync accessor takes the same lock with `blocking_lock`, which is exactly -//! what it says: callers outside a runtime — tests, examples, the parity -//! harness — block as they always did. One lock, not two: a second one would -//! be a second slot, and the model would load twice. -//! -//! Switching to a different checkpoint dir through the same slot drops the -//! old model (freeing its VRAM/RAM) and loads the new one — this is the -//! active-model-switch primitive the P8 management API builds on. - -use std::path::{Path, PathBuf}; - -use tokio::sync::Mutex; - -struct Entry { - key: PathBuf, - value: T, -} - -/// A one-model cache slot, keyed by checkpoint directory. -pub struct ModelSlot { - inner: Mutex>>, -} - -impl Default for ModelSlot { - fn default() -> Self { - Self::new() - } -} - -impl ModelSlot { - #[must_use] - pub const fn new() -> Self { - Self { - // `const_new`, not `new`: the slot is used as a `static`, so its - // constructor must be const (tokio's plain `new` is not). - inner: Mutex::const_new(None), - } - } - - /// Run `f` with the model for `key`, loading it first if the slot is - /// empty or holds a different checkpoint (the old model is dropped - /// before `load` runs, so peak memory stays one model per slot). - pub fn with( - &self, - key: &Path, - load: impl FnOnce(&Path) -> Result, - f: impl FnOnce(&T) -> R, - ) -> Result { - assert!(!key.as_os_str().is_empty(), "model cache: empty key"); - let mut guard = self.inner.blocking_lock(); - let hit = guard.as_ref().is_some_and(|e| e.key == key); - if !hit { - *guard = None; // free the old model before loading the new one - let value = load(key)?; - *guard = Some(Entry { - key: key.to_path_buf(), - value, - }); - } - let entry = guard.as_ref().expect("slot was just filled"); - debug_assert!(entry.key == key, "slot must hold the requested model"); - Ok(f(&entry.value)) - } - - /// Async access to the slot: loads if needed and returns a **guard** - /// holding the model, so the caller can `.await` across it. - /// - /// A closure-taking twin of [`Self::with`] cannot express this: the - /// future would borrow from the `&T` the closure receives, and - /// `FnOnce(&T) -> Fut` has no way to tie `Fut`'s lifetime to that - /// borrow. A guard says the same thing without the higher-ranked - /// gymnastics — hold it, await through it, drop it to release the slot - /// (which is what serializes generations and protects VRAM). - pub async fn acquire( - &self, - key: &Path, - load: impl FnOnce(&Path) -> Result, - ) -> Result, E> { - assert!(!key.as_os_str().is_empty(), "model cache: empty key"); - let mut guard = self.inner.lock().await; - let hit = guard.as_ref().is_some_and(|e| e.key == key); - if !hit { - *guard = None; // free the old model before loading the new one - // Loading is the one genuinely blocking thing on this path: - // minutes of CPU-bound work with no await in it, reading weights - // off disk and placing them across devices. Left on an async - // worker it starves the runtime — measured, a WebSocket heartbeat - // elsewhere in the process was not polled ONCE in 557 seconds and - // then fired 38 missed pings at the end, by which point the - // connection it existed to keep alive would already be reaped. - // - // `block_in_place` hands this worker's other tasks to a sibling - // thread for the duration, so only this call blocks. Deliberately - // narrow: the decode loop around it stays async and awaits. - let value = match tokio::runtime::Handle::try_current().map(|h| h.runtime_flavor()) { - Ok(tokio::runtime::RuntimeFlavor::MultiThread) => { - tokio::task::block_in_place(|| load(key))? - } - // A current-thread runtime (tests) has no sibling worker to - // hand the other tasks to, and `block_in_place` panics there. - _ => load(key)?, - }; - *guard = Some(Entry { - key: key.to_path_buf(), - value, - }); - } - debug_assert!( - guard.as_ref().is_some_and(|e| e.key == key), - "slot must hold the requested model" - ); - Ok(SlotGuard { guard }) - } - - /// Drop the cached model (freeing its VRAM/RAM). No-op when empty. - /// - /// Returns `false` when a generation currently holds the slot — the model - /// cannot be freed under it, and *saying so* beats the alternatives: - /// blocking here panics inside a tokio runtime ("cannot block the current - /// thread from within a runtime"), and waiting would park a worker behind - /// a decode that can run for minutes. Use [`Self::clear_async`] to wait. - pub fn clear(&self) -> bool { - match self.inner.try_lock() { - Ok(mut guard) => { - *guard = None; - true - } - Err(_) => false, - } - } - - /// [`Self::clear`], waiting for any in-flight generation to release the - /// slot first. - pub async fn clear_async(&self) { - *self.inner.lock().await = None; - } - - /// The checkpoint dir currently loaded, if any — for settings UIs. - /// - /// A **peek**: `None` when the slot is empty *or* busy serving a - /// generation. Callers use this to answer "is this already loaded?" and - /// "what is resident?", where waiting behind a multi-minute decode would - /// be worse than a conservative answer (and blocking inside a runtime - /// would panic outright). - #[must_use] - pub fn loaded_key(&self) -> Option { - self.inner - .try_lock() - .ok() - .and_then(|g| g.as_ref().map(|e| e.key.clone())) - } - - /// [`Self::loaded_key`], waiting for the slot instead of reporting busy. - pub async fn loaded_key_async(&self) -> Option { - self.inner.lock().await.as_ref().map(|e| e.key.clone()) - } -} - -/// A held model slot (see [`ModelSlot::acquire`]). Deref to the model; -/// dropping it releases the slot for the next generation. -pub struct SlotGuard<'a, T> { - guard: tokio::sync::MutexGuard<'a, Option>>, -} - -impl std::ops::Deref for SlotGuard<'_, T> { - type Target = T; - - fn deref(&self) -> &T { - &self - .guard - .as_ref() - .expect("a slot guard always holds a loaded model") - .value - } -} - -#[cfg(test)] -mod tests { - use super::*; - use std::convert::Infallible; - - #[test] - fn loads_once_and_reuses_for_the_same_key() { - let slot: ModelSlot = ModelSlot::new(); - let mut loads = 0; - for _ in 0..3 { - let got = slot - .with::<_, Infallible>( - Path::new("model-a"), - |k| { - loads += 1; - Ok(k.display().to_string()) - }, - |m| m.clone(), - ) - .unwrap(); - assert_eq!(got, "model-a"); - } - assert_eq!(loads, 1, "same key must load exactly once"); - assert_eq!(slot.loaded_key().as_deref(), Some(Path::new("model-a"))); - } - - #[test] - fn switching_key_reloads_and_replaces() { - let slot: ModelSlot = ModelSlot::new(); - let mut loads = 0; - let mut run = |key: &str| { - slot.with::<_, Infallible>( - Path::new(key), - |k| { - loads += 1; - Ok(k.display().to_string()) - }, - |m| m.clone(), - ) - .unwrap() - }; - assert_eq!(run("model-a"), "model-a"); - assert_eq!(run("model-b"), "model-b"); // switch: drop a, load b - assert_eq!(run("model-b"), "model-b"); // hit - assert_eq!(loads, 2); - assert_eq!(slot.loaded_key().as_deref(), Some(Path::new("model-b"))); - } - - #[test] - fn failed_load_leaves_the_slot_empty() { - let slot: ModelSlot = ModelSlot::new(); - let err = slot.with(Path::new("bad"), |_| Err("boom"), |m: &String| m.clone()); - assert_eq!(err, Err("boom")); - assert_eq!(slot.loaded_key(), None, "a failed load must not cache"); - } - - #[test] - fn clear_unloads() { - let slot: ModelSlot = ModelSlot::new(); - slot.with::<_, Infallible>(Path::new("m"), |_| Ok(7), |_| ()) - .unwrap(); - assert!(slot.loaded_key().is_some()); - slot.clear(); - assert_eq!(slot.loaded_key(), None); - } - - #[test] - #[should_panic(expected = "empty key")] - fn empty_key_is_rejected() { - let slot: ModelSlot = ModelSlot::new(); - let _ = slot.with::<_, Infallible>(Path::new(""), |_| Ok(1), |_| ()); - } - - /// The slot is usable as a `static` (the whole point). - static GLOBAL: ModelSlot = ModelSlot::new(); - - #[test] - fn works_as_a_static_across_threads() { - let handles: Vec<_> = (0..4) - .map(|_| { - std::thread::spawn(|| { - GLOBAL - .with::<_, Infallible>(Path::new("shared"), |_| Ok(41), |v| v + 1) - .unwrap() - }) - }) - .collect(); - for h in handles { - assert_eq!(h.join().unwrap(), 42); - } - } -} diff --git a/crates/mummu/examples/src/chat.rs b/crates/mummu/examples/src/chat.rs deleted file mode 100644 index 4fdff6c..0000000 --- a/crates/mummu/examples/src/chat.rs +++ /dev/null @@ -1,1488 +0,0 @@ -//! Explicit chat templates. Prompt wrapping is part of a model's contract — -//! an implicit or slightly-wrong template silently ruins output quality — so -//! templates are code here, never guessed: each per-model constructor is -//! byte-verified against the parity references (the Qwen2 template renders -//! the exact prompt committed in the Candle logits fixture). -//! -//! Both zoo LLMs speak ChatML; LFM2.5 additionally prefixes `<|startoftext|>`. -//! Tool use comes in two conventions, selected by the per-model constructor: -//! -//! - **Hermes** (Qwen2.5/Qwen3): tool signatures in a `` block of the -//! system turn, calls emitted as `{json}`, results -//! returned inside `` blocks of a user turn. Qwen3's -//! template additionally strips `` reasoning from assistant turns -//! at/before the last user query and injects no default system preamble — -//! [`ChatMl::qwen3`] carries those deltas; [`ChatMl::qwen2`] re-renders -//! history verbatim. -//! - **LFM** (LFM2.5, per its `chat_template.jinja` + model card): tool -//! signatures as bare JSON in a `List of tools: […]` line of the system -//! turn, calls emitted as a *Pythonic call list* between the -//! `<|tool_call_start|>`/`<|tool_call_end|>` special tokens — e.g. -//! `[get_weather(city="Paris")]` — results returned in a dedicated `tool` -//! role turn, and ``-prefixed reasoning stripped from every -//! assistant history turn but the last. - -/// Who is speaking in a [`Turn`]. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum Role { - System, - User, - Assistant, - /// A tool result going back to the model. Hermes-style templates render - /// these inside a *user* turn as `` blocks; LFM-style - /// templates give each one its own `tool` role turn. - Tool, -} - -/// Which tool-use convention a template speaks (see the module docs). -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum ToolCallStyle { - Hermes, - Lfm, -} - -/// How a template treats ``-prefixed reasoning when re-rendering -/// assistant history turns (each variant byte-verified against its family's -/// imported `chat_template` by `tests/template_gate.rs`). -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum ThinkStrip { - /// Re-render history verbatim — Qwen2.5's template has no think path. - Keep, - /// Strip from every assistant turn but the last: LFM2.5's - /// `keep_past_thinking=false` default. - PastAssistant, - /// Qwen3: strip from assistant turns at/before the last real user query; - /// later turns (mid tool loop) keep their reasoning, re-emitted in the - /// template's normalized `\n…\n\n\n` shape. - BeforeLastUserQuery, -} - -impl Role { - fn tag(self, style: ToolCallStyle) -> &'static str { - match (self, style) { - (Self::System, _) => "system", - (Self::User, _) => "user", - (Self::Assistant, _) => "assistant", - // Hermes: tool results ride in a user turn; LFM: a real tool turn. - (Self::Tool, ToolCallStyle::Hermes) => "user", - (Self::Tool, ToolCallStyle::Lfm) => "tool", - } - } -} - -/// One message in a conversation. -#[derive(Debug, Clone)] -pub struct Turn { - pub role: Role, - pub content: String, - /// The tool calls an assistant turn makes, kept **structurally** beside - /// the rendered `content`. The family renderers here never read it — - /// they re-emit `content`, which the `assistant_tool_calls*` - /// constructors already wrote in the family's own wire format. It exists - /// for renderers that need the calls unwrapped: a checkpoint's imported - /// Jinja template (see [`crate::template`]) receives `tool_calls` as data - /// and writes the markers *itself*, in whatever convention that template - /// speaks. Empty for every other role and for a plain assistant turn. - pub tool_calls: Vec, -} - -impl Turn { - #[must_use] - pub fn system(content: impl Into) -> Self { - Self { - role: Role::System, - content: content.into(), - tool_calls: Vec::new(), - } - } - - #[must_use] - pub fn user(content: impl Into) -> Self { - Self { - role: Role::User, - content: content.into(), - tool_calls: Vec::new(), - } - } - - #[must_use] - pub fn assistant(content: impl Into) -> Self { - Self { - role: Role::Assistant, - content: content.into(), - tool_calls: Vec::new(), - } - } - - /// An assistant turn that invokes tools: each call becomes a Hermes - /// `` block in the turn body (what the model itself would - /// have emitted), so histories containing calls re-render faithfully. - #[must_use] - pub fn assistant_tool_calls(calls: &[ToolCall]) -> Self { - assert!(!calls.is_empty(), "assistant_tool_calls: no calls"); - assert!( - calls.len() <= MAX_TOOL_CALLS, - "assistant_tool_calls: {} calls exceeds the {MAX_TOOL_CALLS} bound", - calls.len() - ); - let blocks: Vec = calls - .iter() - .map(|c| { - let json = python_json(c); - debug_assert!(!json.is_empty(), "a ToolCall always serializes"); - format!("\n{json}\n") - }) - .collect(); - Self { - role: Role::Assistant, - content: blocks.join("\n"), - tool_calls: calls.to_vec(), - } - } - - /// An assistant turn that invokes tools in LFM's Pythonic convention: - /// the calls render as one bracketed call list between the - /// `<|tool_call_start|>`/`<|tool_call_end|>` special tokens — exactly - /// what an LFM2.5 model emits — so histories re-render faithfully. - #[must_use] - pub fn assistant_tool_calls_lfm(calls: &[ToolCall]) -> Self { - assert!(!calls.is_empty(), "assistant_tool_calls_lfm: no calls"); - assert!( - calls.len() <= MAX_TOOL_CALLS, - "assistant_tool_calls_lfm: {} calls exceeds the {MAX_TOOL_CALLS} bound", - calls.len() - ); - Self { - role: Role::Assistant, - content: format!( - "<|tool_call_start|>{}<|tool_call_end|>", - pythonic_calls(calls) - ), - tool_calls: calls.to_vec(), - } - } - - /// A tool's result going back to the model. Hermes templates render it as - /// a `` block (consecutive ones merge into one user turn); - /// LFM templates give it its own `tool` role turn. - #[must_use] - pub fn tool_response(content: impl Into) -> Self { - Self { - role: Role::Tool, - content: content.into(), - tool_calls: Vec::new(), - } - } -} - -/// Deepest literal nesting the Pythonic renderer/parser will follow — far -/// past any real argument payload, and the recursion bound for both. -const MAX_VALUE_DEPTH: usize = 8; - -/// Serialize a value the way Python's `json.dumps` does by default — `", "` -/// between items, `": "` after keys. Prompt JSON renders this way -/// deliberately: it is byte-for-byte what `transformers.apply_chat_template` -/// produces through Jinja's `tojson` (and the spacing the models emit back in -/// their own `` JSON), pinned by the template byte gate -/// (`tests/template_gate.rs`). -fn python_json(value: &T) -> String { - struct PySeparators; - impl serde_json::ser::Formatter for PySeparators { - fn begin_array_value(&mut self, writer: &mut W, first: bool) -> std::io::Result<()> - where - W: ?Sized + std::io::Write, - { - if !first { - writer.write_all(b", ")?; - } - Ok(()) - } - - fn begin_object_key(&mut self, writer: &mut W, first: bool) -> std::io::Result<()> - where - W: ?Sized + std::io::Write, - { - if !first { - writer.write_all(b", ")?; - } - Ok(()) - } - - fn begin_object_value(&mut self, writer: &mut W) -> std::io::Result<()> - where - W: ?Sized + std::io::Write, - { - writer.write_all(b": ") - } - } - - let mut out = Vec::with_capacity(128); - let mut ser = serde_json::Serializer::with_formatter(&mut out, PySeparators); - let serialized = serde::Serialize::serialize(value, &mut ser).is_ok(); - debug_assert!(serialized, "prompt JSON values always serialize"); - debug_assert!( - !out.is_empty() || !serialized, - "a serialized value is non-empty" - ); - if !serialized { - return String::new(); - } - String::from_utf8(out).unwrap_or_default() -} - -/// Render tool calls as LFM's Pythonic call list: `[name(k=v, …), …]`. -/// JSON scalars map to Python spellings (`true`→`True`, `null`→`None`); -/// strings/lists/objects render as Python literals. -fn pythonic_calls(calls: &[ToolCall]) -> String { - assert!(!calls.is_empty(), "pythonic_calls: no calls"); - assert!(calls.len() <= MAX_TOOL_CALLS, "pythonic_calls: over bound"); - let mut out = String::from("["); - for (i, call) in calls.iter().enumerate() { - assert!(!call.name.is_empty(), "pythonic_calls: unnamed call"); - if i > 0 { - out.push_str(", "); - } - out.push_str(&call.name); - out.push('('); - match &call.arguments { - serde_json::Value::Object(args) => { - for (j, (key, value)) in args.iter().enumerate() { - if j > 0 { - out.push_str(", "); - } - out.push_str(key); - out.push('='); - python_literal(value, &mut out, 0); - } - } - serde_json::Value::Null => {} - other => panic!("pythonic_calls: arguments must be an object or null, got {other}"), - } - out.push(')'); - } - out.push(']'); - debug_assert!(out.starts_with('[') && out.ends_with(']'), "list shape"); - out -} - -/// Append one JSON value as a Python literal. Strings render double-quoted -/// with JSON-style escapes (valid Python), scalars as Python spellings. -fn python_literal(value: &serde_json::Value, out: &mut String, depth: usize) { - assert!(depth <= MAX_VALUE_DEPTH, "python_literal: value too deep"); - match value { - serde_json::Value::Null => out.push_str("None"), - serde_json::Value::Bool(true) => out.push_str("True"), - serde_json::Value::Bool(false) => out.push_str("False"), - serde_json::Value::Number(n) => out.push_str(&n.to_string()), - serde_json::Value::String(s) => python_string_literal(s, out), - serde_json::Value::Array(items) => { - out.push('['); - for (i, item) in items.iter().enumerate() { - if i > 0 { - out.push_str(", "); - } - python_literal(item, out, depth + 1); - } - out.push(']'); - } - serde_json::Value::Object(map) => { - out.push('{'); - for (i, (key, val)) in map.iter().enumerate() { - if i > 0 { - out.push_str(", "); - } - python_string_literal(key, out); - out.push_str(": "); - python_literal(val, out, depth + 1); - } - out.push('}'); - } - } -} - -/// Double-quoted Python string literal with JSON-compatible escapes. -fn python_string_literal(s: &str, out: &mut String) { - out.push('"'); - for c in s.chars() { - match c { - '"' => out.push_str("\\\""), - '\\' => out.push_str("\\\\"), - '\n' => out.push_str("\\n"), - '\t' => out.push_str("\\t"), - '\r' => out.push_str("\\r"), - c if (c as u32) < 0x20 => { - out.push_str(&format!("\\u{:04x}", c as u32)); - } - c => out.push(c), - } - } - out.push('"'); - debug_assert!(out.ends_with('"'), "string literal closes"); -} - -/// A callable tool signature, serialized into the system prompt exactly as -/// the Hermes-style templates expect: `{"type": "function", "function": -/// {"name": …, "description": …, "parameters": }}`. -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] -pub struct ToolSpec { - pub name: String, - pub description: String, - /// JSON schema of the arguments object. - pub parameters: serde_json::Value, -} - -/// The Hermes wire shape for one tool (field order matters for byte-stable -/// rendering, so this is a struct, not a `json!` map). -#[derive(serde::Serialize)] -struct ToolWire<'a> { - r#type: &'static str, - function: &'a ToolSpec, -} - -/// One tool invocation, as emitted by the model inside `` tags. -#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] -pub struct ToolCall { - pub name: String, - #[serde(default)] - pub arguments: serde_json::Value, -} - -/// Most tool calls a single response may contain (or a single assistant -/// history turn may carry) — far past anything a small model emits. -pub const MAX_TOOL_CALLS: usize = 64; - -/// Most tools one render will advertise. -pub(crate) const MAX_TOOLS: usize = 128; - -/// What went wrong extracting tool calls from a model response. -#[derive(Debug, thiserror::Error)] -pub enum ToolCallError { - #[error("tool call {index}: unclosed tool-call tag")] - Unclosed { index: usize }, - #[error("tool call {index}: {reason}")] - BadJson { index: usize, reason: String }, - #[error("tool call block {index}, byte {offset}: {reason}")] - Syntax { - /// Which `<|tool_call_start|>` block (0-based) failed to parse. - index: usize, - /// Byte offset *inside the block* where parsing stopped. - offset: usize, - reason: String, - }, - #[error("tool call arguments nested deeper than {MAX_VALUE_DEPTH} levels")] - TooDeep, - #[error("more than {MAX_TOOL_CALLS} tool calls in one response")] - TooMany, -} - -/// Extract Hermes-style tool calls from a model response: every -/// `…` block parses as a [`ToolCall`]; the text -/// outside the blocks (the model's prose, trimmed) comes back alongside. -/// Text with no blocks is simply `(vec![], text)` — not an error. -pub fn parse_tool_calls(text: &str) -> Result<(Vec, String), ToolCallError> { - const OPEN: &str = ""; - const CLOSE: &str = ""; - let mut calls = Vec::new(); - let mut prose = String::new(); - let mut rest = text; - while let Some(start) = rest.find(OPEN) { - if calls.len() == MAX_TOOL_CALLS { - return Err(ToolCallError::TooMany); - } - prose.push_str(&rest[..start]); - let after_open = &rest[start + OPEN.len()..]; - let Some(end) = after_open.find(CLOSE) else { - return Err(ToolCallError::Unclosed { index: calls.len() }); - }; - let body = after_open[..end].trim(); - let call: ToolCall = serde_json::from_str(body).map_err(|e| ToolCallError::BadJson { - index: calls.len(), - reason: e.to_string(), - })?; - calls.push(call); - rest = &after_open[end + CLOSE.len()..]; - } - prose.push_str(rest); - debug_assert!(calls.len() <= MAX_TOOL_CALLS, "bound enforced in the loop"); - Ok((calls, prose.trim().to_string())) -} - -/// Extract LFM-style tool calls from a model response: every -/// `<|tool_call_start|>…<|tool_call_end|>` block parses as a *Pythonic call -/// list* (`[name(k=v, …), …]`); the text outside the blocks (the model's -/// prose, trimmed) comes back alongside. Text with no blocks is simply -/// `(vec![], text)` — not an error. -pub fn parse_tool_calls_lfm(text: &str) -> Result<(Vec, String), ToolCallError> { - const OPEN: &str = "<|tool_call_start|>"; - const CLOSE: &str = "<|tool_call_end|>"; - let mut calls = Vec::new(); - let mut prose = String::new(); - let mut rest = text; - let mut block = 0; - while let Some(start) = rest.find(OPEN) { - prose.push_str(&rest[..start]); - let after_open = &rest[start + OPEN.len()..]; - let Some(end) = after_open.find(CLOSE) else { - return Err(ToolCallError::Unclosed { index: block }); - }; - let parsed = PythonicParser::new(after_open[..end].trim(), block).parse_call_list()?; - if calls.len() + parsed.len() > MAX_TOOL_CALLS { - return Err(ToolCallError::TooMany); - } - calls.extend(parsed); - rest = &after_open[end + CLOSE.len()..]; - block += 1; - } - prose.push_str(rest); - debug_assert!(calls.len() <= MAX_TOOL_CALLS, "bound enforced per block"); - Ok((calls, prose.trim().to_string())) -} - -/// A bounded recursive-descent parser for LFM's Pythonic call list. -/// -/// Grammar (whitespace-tolerant, trailing commas allowed): -/// ```text -/// calls := '[' [ call (',' call)* [','] ] ']' -/// call := ident '(' [ kwarg (',' kwarg)* [','] ] ')' -/// kwarg := ident '=' value -/// value := 'True' | 'False' | 'None' | 'true' | 'false' | 'null' -/// | number | string | '[' … ']' | '{' string ':' value, … '}' -/// ``` -/// Lowercase JSON spellings are accepted because LFM2.5 documents a JSON -/// fallback mode and small models mix the two. -struct PythonicParser<'a> { - src: &'a str, - pos: usize, - block: usize, -} - -impl<'a> PythonicParser<'a> { - fn new(src: &'a str, block: usize) -> Self { - Self { src, pos: 0, block } - } - - fn fail(&self, reason: impl Into) -> ToolCallError { - ToolCallError::Syntax { - index: self.block, - offset: self.pos, - reason: reason.into(), - } - } - - fn skip_ws(&mut self) { - while self.src[self.pos..].starts_with(|c: char| c.is_ascii_whitespace()) { - self.pos += 1; - } - } - - fn peek(&self) -> Option { - self.src[self.pos..].chars().next() - } - - fn eat(&mut self, c: char) -> bool { - if self.peek() == Some(c) { - self.pos += c.len_utf8(); - true - } else { - false - } - } - - fn expect(&mut self, c: char, what: &str) -> Result<(), ToolCallError> { - self.skip_ws(); - if self.eat(c) { - Ok(()) - } else { - Err(self.fail(format!("expected '{c}' {what}"))) - } - } - - /// The whole block: a bracketed list of calls, then end of input. - fn parse_call_list(mut self) -> Result, ToolCallError> { - assert!(self.pos == 0, "parse_call_list: parser already consumed"); - self.expect('[', "to open the call list")?; - let mut calls = Vec::new(); - loop { - self.skip_ws(); - if self.eat(']') { - break; - } - if calls.len() == MAX_TOOL_CALLS { - return Err(ToolCallError::TooMany); - } - calls.push(self.parse_call()?); - self.skip_ws(); - if !self.eat(',') && self.peek() != Some(']') { - return Err(self.fail("expected ',' or ']' after a call")); - } - } - self.skip_ws(); - if self.pos != self.src.len() { - return Err(self.fail("trailing text after the call list")); - } - debug_assert!(calls.len() <= MAX_TOOL_CALLS, "bound enforced in loop"); - Ok(calls) - } - - fn parse_call(&mut self) -> Result { - let name = self.parse_ident("a function name")?; - self.expect('(', "to open the arguments")?; - let mut args = serde_json::Map::new(); - loop { - self.skip_ws(); - if self.eat(')') { - break; - } - let key = self.parse_ident("an argument name")?; - self.expect('=', "between argument name and value")?; - self.skip_ws(); - let value = self.parse_value(0)?; - args.insert(key, value); - self.skip_ws(); - if !self.eat(',') && self.peek() != Some(')') { - return Err(self.fail("expected ',' or ')' after an argument")); - } - } - debug_assert!(!name.is_empty(), "parse_ident never returns empty"); - Ok(ToolCall { - name, - arguments: serde_json::Value::Object(args), - }) - } - - fn parse_ident(&mut self, what: &str) -> Result { - self.skip_ws(); - let start = self.pos; - if self - .peek() - .is_some_and(|c| c.is_ascii_alphabetic() || c == '_') - { - self.pos += 1; - while self - .peek() - .is_some_and(|c| c.is_ascii_alphanumeric() || c == '_') - { - self.pos += 1; - } - } - if self.pos == start { - return Err(self.fail(format!("expected {what}"))); - } - Ok(self.src[start..self.pos].to_string()) - } - - fn parse_value(&mut self, depth: usize) -> Result { - if depth > MAX_VALUE_DEPTH { - return Err(ToolCallError::TooDeep); - } - self.skip_ws(); - match self.peek() { - Some('"') | Some('\'') => Ok(serde_json::Value::String(self.parse_string()?)), - Some('[') => self.parse_list(depth), - Some('{') => self.parse_dict(depth), - Some(c) if c == '-' || c.is_ascii_digit() => self.parse_number(), - Some(c) if c.is_ascii_alphabetic() => { - let word = self.parse_ident("a literal")?; - match word.as_str() { - "True" | "true" => Ok(serde_json::Value::Bool(true)), - "False" | "false" => Ok(serde_json::Value::Bool(false)), - "None" | "null" => Ok(serde_json::Value::Null), - other => Err(self.fail(format!("unknown literal '{other}'"))), - } - } - _ => Err(self.fail("expected a value")), - } - } - - fn parse_list(&mut self, depth: usize) -> Result { - assert!(self.peek() == Some('['), "parse_list: caller checked '['"); - self.pos += 1; - let mut items = Vec::new(); - loop { - self.skip_ws(); - if self.eat(']') { - break; - } - items.push(self.parse_value(depth + 1)?); - self.skip_ws(); - if !self.eat(',') && self.peek() != Some(']') { - return Err(self.fail("expected ',' or ']' in a list")); - } - } - Ok(serde_json::Value::Array(items)) - } - - fn parse_dict(&mut self, depth: usize) -> Result { - assert!(self.peek() == Some('{'), "parse_dict: caller checked brace"); - self.pos += 1; - let mut map = serde_json::Map::new(); - loop { - self.skip_ws(); - if self.eat('}') { - break; - } - self.skip_ws(); - if !matches!(self.peek(), Some('"') | Some('\'')) { - return Err(self.fail("dict keys must be strings")); - } - let key = self.parse_string()?; - self.expect(':', "between dict key and value")?; - let value = self.parse_value(depth + 1)?; - map.insert(key, value); - self.skip_ws(); - if !self.eat(',') && self.peek() != Some('}') { - return Err(self.fail("expected ',' or '}' in a dict")); - } - } - Ok(serde_json::Value::Object(map)) - } - - /// A single- or double-quoted string with Python/JSON escapes. - fn parse_string(&mut self) -> Result { - let quote = self.peek().expect("parse_string: caller checked a quote"); - assert!(quote == '"' || quote == '\'', "caller checked the quote"); - self.pos += 1; - let mut out = String::new(); - loop { - let Some(c) = self.peek() else { - return Err(self.fail("unterminated string")); - }; - self.pos += c.len_utf8(); - match c { - c if c == quote => break, - '\\' => out.push(self.parse_escape()?), - c => out.push(c), - } - } - Ok(out) - } - - fn parse_escape(&mut self) -> Result { - let Some(c) = self.peek() else { - return Err(self.fail("dangling escape at end of string")); - }; - self.pos += c.len_utf8(); - match c { - '"' | '\'' | '\\' | '/' => Ok(c), - 'n' => Ok('\n'), - 't' => Ok('\t'), - 'r' => Ok('\r'), - 'u' => { - let hex = self - .src - .get(self.pos..self.pos + 4) - .ok_or_else(|| self.fail("truncated \\u escape"))?; - let code = u32::from_str_radix(hex, 16).map_err(|_| self.fail("bad \\u escape"))?; - self.pos += 4; - char::from_u32(code).ok_or_else(|| self.fail("\\u escape is not a scalar")) - } - other => Err(self.fail(format!("unsupported escape '\\{other}'"))), - } - } - - fn parse_number(&mut self) -> Result { - let start = self.pos; - self.eat('-'); - while self - .peek() - .is_some_and(|c| c.is_ascii_digit() || matches!(c, '.' | 'e' | 'E' | '+' | '-')) - { - self.pos += 1; - } - let text = &self.src[start..self.pos]; - debug_assert!(!text.is_empty(), "caller checked a digit or '-'"); - if let Ok(i) = text.parse::() { - return Ok(serde_json::Value::Number(i.into())); - } - let f = text - .parse::() - .map_err(|_| self.fail(format!("bad number '{text}'")))?; - serde_json::Number::from_f64(f) - .map(serde_json::Value::Number) - .ok_or_else(|| self.fail(format!("non-finite number '{text}'"))) - } -} - -/// Longest conversation a single render will wrap — a generous bound that -/// still catches an unbounded history being passed by mistake. -pub(crate) const MAX_TURNS: usize = 1024; - -/// The index Qwen3's template calls `last_query_index`: the last USER turn -/// whose content is not a pre-wrapped `` block (the legacy -/// convention for returning tool results inside a user turn). When no such -/// turn exists the template's fallback makes every turn "at/before" it. -fn last_user_query(turns: &[Turn]) -> usize { - debug_assert!(!turns.is_empty(), "render asserts a non-empty history"); - debug_assert!(turns.len() <= MAX_TURNS, "render asserts the bound"); - turns - .iter() - .rposition(|t| { - t.role == Role::User - && !(t.content.starts_with("") - && t.content.ends_with("")) - }) - .unwrap_or(turns.len().saturating_sub(1)) -} - -/// Qwen3's assistant-history reasoning rule, byte-for-byte from its -/// `chat_template` (pinned by `tests/template_gate.rs`): a turn at/before -/// the last user query renders only the text after its final `` -/// (leading newlines dropped — Python `lstrip('\n')`, NOT a full trim); a -/// later turn (mid tool loop) keeps its reasoning, re-emitted in the -/// template's normalized `\n…\n\n\n` shape. -fn qwen3_think_content( - content: &str, - index: usize, - last_user_query: usize, -) -> std::borrow::Cow<'_, str> { - let Some(last_close) = content.rfind("") else { - return std::borrow::Cow::Borrowed(content); - }; - let body = content[last_close + "".len()..].trim_start_matches('\n'); - debug_assert!(!body.starts_with('\n'), "leading newlines are stripped"); - if index <= last_user_query { - return std::borrow::Cow::Borrowed(body); - } - // The template's Python chain, literally: reasoning = - // split('')[0].rstrip('\n').split('')[-1], emitted - // through a final .strip('\n'). - let first_close = content.find("").unwrap_or(last_close); - debug_assert!(first_close <= last_close, "find precedes rfind"); - let before = content[..first_close].trim_end_matches('\n'); - let reasoning = match before.rfind("") { - Some(open) => &before[open + "".len()..], - None => before, - }; - let reasoning = reasoning.trim_matches('\n'); - std::borrow::Cow::Owned(format!("\n{reasoning}\n\n\n{body}")) -} - -/// The ChatML template family: `<|im_start|>role\ncontent<|im_end|>\n` per -/// turn, then an open assistant turn for the model to complete. `bos` is -/// prepended once when a model requires a start-of-text token; `tool_style` -/// picks the tool-use convention (see the module docs). -#[derive(Debug, Clone)] -pub struct ChatMl { - bos: Option<&'static str>, - tool_style: ToolCallStyle, - think_strip: ThinkStrip, - /// The system preamble `render_with_tools` synthesizes when the - /// conversation opens without a system turn — `None` for templates - /// (Qwen3, LFM) that inject nothing of their own. - tools_default_system: Option<&'static str>, -} - -impl ChatMl { - /// Qwen2 / Qwen2.5-Instruct: plain ChatML, no BOS, Hermes tool use. - #[must_use] - pub fn qwen2() -> Self { - Self { - bos: None, - tool_style: ToolCallStyle::Hermes, - think_strip: ThinkStrip::Keep, - tools_default_system: Some("You are a helpful assistant."), - } - } - - /// Qwen3 / Qwen3.5: Qwen2's ChatML + Hermes wire format, plus the two - /// behavioral deltas its template adds — `` reasoning stripped - /// from assistant turns at/before the last user query, and NO default - /// system preamble when tools are supplied without a system turn. - #[must_use] - pub fn qwen3() -> Self { - Self { - bos: None, - tool_style: ToolCallStyle::Hermes, - think_strip: ThinkStrip::BeforeLastUserQuery, - tools_default_system: None, - } - } - - /// LFM2 / LFM2.5-Instruct: ChatML behind `<|startoftext|>`, LFM - /// (Pythonic) tool use. - #[must_use] - pub fn lfm2() -> Self { - Self { - bos: Some("<|startoftext|>"), - tool_style: ToolCallStyle::Lfm, - think_strip: ThinkStrip::PastAssistant, - tools_default_system: None, - } - } - - /// Render a conversation into the raw prompt string, ending with the open - /// assistant turn the model completes. The caller tokenizes the result - /// with special tokens enabled by the tokenizer itself, not re-added. - #[must_use] - pub fn render(&self, turns: &[Turn]) -> String { - assert!(!turns.is_empty(), "chat render: no turns"); - assert!( - turns.len() <= MAX_TURNS, - "chat render: {} turns exceeds the {MAX_TURNS} bound", - turns.len() - ); - assert!( - turns.last().map(|t| t.role) != Some(Role::Assistant), - "chat render: the template opens the assistant turn itself; \ - a trailing assistant turn would double it" - ); - let last_assistant = turns.iter().rposition(|t| t.role == Role::Assistant); - let last_user_query = last_user_query(turns); - let mut out = String::from(self.bos.unwrap_or("")); - let mut i = 0; - while i < turns.len() { - if self.tool_style == ToolCallStyle::Hermes && turns[i].role == Role::Tool { - // Hermes: consecutive tool results merge into ONE user turn, - // each wrapped in its own block. - out.push_str("<|im_start|>user"); - while i < turns.len() && turns[i].role == Role::Tool { - out.push_str("\n\n"); - out.push_str(&turns[i].content); - out.push_str("\n"); - i += 1; - } - out.push_str("<|im_end|>\n"); - } else { - out.push_str("<|im_start|>"); - out.push_str(turns[i].role.tag(self.tool_style)); - out.push('\n'); - out.push_str(&self.turn_content(&turns[i], i, last_assistant, last_user_query)); - out.push_str("<|im_end|>\n"); - i += 1; - } - } - out.push_str("<|im_start|>assistant\n"); - debug_assert!(out.ends_with("assistant\n"), "render must open a turn"); - out - } - - /// What a turn's body renders as — the per-family [`ThinkStrip`] policy - /// applied to assistant history turns; everything else passes through. - fn turn_content<'a>( - &self, - turn: &'a Turn, - index: usize, - last_assistant: Option, - last_user_query: usize, - ) -> std::borrow::Cow<'a, str> { - debug_assert!(index <= MAX_TURNS, "index bounded by the render assert"); - debug_assert!( - turn.role != Role::Assistant || last_assistant.is_some(), - "an assistant turn implies a last-assistant index" - ); - if turn.role != Role::Assistant { - return std::borrow::Cow::Borrowed(&turn.content); - } - match self.think_strip { - ThinkStrip::Keep => std::borrow::Cow::Borrowed(&turn.content), - ThinkStrip::PastAssistant => { - let is_past = last_assistant.is_some_and(|l| index != l); - match turn.content.rfind("") { - Some(end) if is_past => { - std::borrow::Cow::Borrowed(turn.content[end + "".len()..].trim()) - } - _ => std::borrow::Cow::Borrowed(turn.content.as_str()), - } - } - ThinkStrip::BeforeLastUserQuery => { - qwen3_think_content(&turn.content, index, last_user_query) - } - } - } - - /// [`render`](Self::render) with function calling in this template's - /// convention. - /// - /// - **Hermes** (Qwen2.5/Qwen3): tool signatures advertised in a - /// `# Tools` section of the system turn — the exact wording and tag - /// structure those models ship in their chat template. A conversation - /// without a system turn gets a neutral "You are a helpful assistant." - /// preamble under [`ChatMl::qwen2`]; [`ChatMl::qwen3`] injects none, - /// matching its template. - /// - **LFM** (LFM2.5): bare tool JSON on a `List of tools: […]` line - /// appended to the system turn — the exact shape of LFM2.5's - /// `chat_template.jinja` + model card, which injects *no* default - /// preamble: without a system turn the tools line stands alone. - #[must_use] - pub fn render_with_tools(&self, tools: &[ToolSpec], turns: &[Turn]) -> String { - assert!( - !tools.is_empty(), - "render_with_tools: no tools — use render()" - ); - assert!( - tools.len() <= MAX_TOOLS, - "render_with_tools: {} tools exceeds the {MAX_TOOLS} bound", - tools.len() - ); - assert!( - tools.iter().all(|t| !t.name.is_empty()), - "render_with_tools: every tool needs a name" - ); - match self.tool_style { - ToolCallStyle::Hermes => self.render_with_tools_hermes(tools, turns), - ToolCallStyle::Lfm => self.render_with_tools_lfm(tools, turns), - } - } - - fn render_with_tools_hermes(&self, tools: &[ToolSpec], turns: &[Turn]) -> String { - debug_assert!(!tools.is_empty(), "checked by render_with_tools"); - debug_assert!(self.tool_style == ToolCallStyle::Hermes, "hermes only"); - let (preamble, rest) = match turns.first() { - Some(t) if t.role == Role::System => (Some(t.content.as_str()), &turns[1..]), - // Qwen2.5 synthesizes a preamble; Qwen3's template injects none. - _ => (self.tools_default_system, turns), - }; - let mut system = String::new(); - if let Some(preamble) = preamble { - system.push_str(preamble); - system.push_str("\n\n"); - } - system.push_str( - "# Tools\n\nYou may call one or more functions to assist with the user query.\n\n\ - You are provided with function signatures within XML tags:\n", - ); - for tool in tools { - let json = python_json(&ToolWire { - r#type: "function", - function: tool, - }); - debug_assert!(!json.is_empty(), "a ToolSpec always serializes"); - system.push('\n'); - system.push_str(&json); - } - system.push_str( - "\n\n\nFor each function call, return a json object with function name and \ - arguments within XML tags:\n\n{\"name\": \ - , \"arguments\": }\n", - ); - - let mut wrapped = Vec::with_capacity(rest.len() + 1); - wrapped.push(Turn::system(system)); - wrapped.extend_from_slice(rest); - self.render(&wrapped) - } - - fn render_with_tools_lfm(&self, tools: &[ToolSpec], turns: &[Turn]) -> String { - debug_assert!(!tools.is_empty(), "checked by render_with_tools"); - debug_assert!(self.tool_style == ToolCallStyle::Lfm, "lfm only"); - let (preamble, rest) = match turns.first() { - Some(t) if t.role == Role::System => (t.content.as_str(), &turns[1..]), - _ => ("", turns), - }; - let mut system = String::from(preamble); - if !system.is_empty() { - system.push('\n'); - } - system.push_str("List of tools: ["); - for (i, tool) in tools.iter().enumerate() { - if i > 0 { - system.push_str(", "); - } - let json = python_json(tool); - debug_assert!(!json.is_empty(), "a ToolSpec always serializes"); - system.push_str(&json); - } - system.push(']'); - - let mut wrapped = Vec::with_capacity(rest.len() + 1); - wrapped.push(Turn::system(system)); - wrapped.extend_from_slice(rest); - self.render(&wrapped) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn qwen2_render_matches_the_parity_verified_shape() { - // This exact string (with this system prompt + user text) is what the - // Candle fixture and the Ollama fp16 greedy leg were verified against. - let raw = ChatMl::qwen2().render(&[ - Turn::system("You are a helpful assistant."), - Turn::user("List the first five prime numbers."), - ]); - assert_eq!( - raw, - "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n\ - <|im_start|>user\nList the first five prime numbers.<|im_end|>\n\ - <|im_start|>assistant\n" - ); - } - - #[test] - fn lfm2_render_prefixes_bos_and_matches_the_parity_shape() { - let raw = ChatMl::lfm2().render(&[Turn::user("List the first five prime numbers.")]); - assert_eq!( - raw, - "<|startoftext|><|im_start|>user\nList the first five prime numbers.<|im_end|>\n\ - <|im_start|>assistant\n" - ); - } - - #[test] - fn multi_turn_history_renders_in_order() { - let raw = ChatMl::qwen2().render(&[ - Turn::user("Hi."), - Turn::assistant("Hello!"), - Turn::user("Bye."), - ]); - assert_eq!( - raw, - "<|im_start|>user\nHi.<|im_end|>\n\ - <|im_start|>assistant\nHello!<|im_end|>\n\ - <|im_start|>user\nBye.<|im_end|>\n\ - <|im_start|>assistant\n" - ); - } - - #[test] - #[should_panic(expected = "no turns")] - fn empty_conversation_is_rejected() { - let _ = ChatMl::qwen2().render(&[]); - } - - fn weather_tool() -> ToolSpec { - ToolSpec { - name: "get_weather".into(), - description: "Get the current weather for a city.".into(), - parameters: serde_json::json!({ - "type": "object", - "properties": { "city": { "type": "string" } }, - "required": ["city"] - }), - } - } - - /// The tools section must match the Qwen2.5/Qwen3 chat template's wording - /// and tag structure byte-for-byte (the model was trained on this text). - /// Inside a tool's `parameters` schema, keys serialize in INSERTION order - /// (serde_json `preserve_order`, a workspace feature) — matching how - /// Python/transformers renders the same schema from a dict. - #[test] - fn tools_render_matches_the_hermes_template_shape() { - let raw = ChatMl::qwen2().render_with_tools( - &[weather_tool()], - &[ - Turn::system("You are a helpful assistant."), - Turn::user("Weather in Paris?"), - ], - ); - let expected = "<|im_start|>system\nYou are a helpful assistant.\n\n# Tools\n\n\ - You may call one or more functions to assist with the user query.\n\n\ - You are provided with function signatures within XML tags:\n\ - \n\ - {\"type\": \"function\", \"function\": {\"name\": \"get_weather\", \"description\": \"Get the current weather for a city.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"city\": {\"type\": \"string\"}}, \"required\": [\"city\"]}}}\n\ - \n\n\ - For each function call, return a json object with function name and arguments within XML tags:\n\ - \n{\"name\": , \"arguments\": }\n<|im_end|>\n\ - <|im_start|>user\nWeather in Paris?<|im_end|>\n\ - <|im_start|>assistant\n"; - assert_eq!(raw, expected); - } - - #[test] - fn tools_render_without_a_system_turn_injects_a_neutral_preamble() { - let raw = ChatMl::qwen2().render_with_tools(&[weather_tool()], &[Turn::user("hi")]); - assert!(raw.starts_with("<|im_start|>system\nYou are a helpful assistant.\n\n# Tools")); - // The user turn survives un-consumed. - assert!(raw.contains("<|im_start|>user\nhi<|im_end|>")); - } - - #[test] - fn consecutive_tool_responses_merge_into_one_user_turn() { - let calls = [ToolCall { - name: "get_weather".into(), - arguments: serde_json::json!({"city": "Paris"}), - }]; - let raw = ChatMl::qwen2().render(&[ - Turn::user("Weather in Paris and Lyon?"), - Turn::assistant_tool_calls(&calls), - Turn::tool_response("{\"temp_c\": 21}"), - Turn::tool_response("{\"temp_c\": 24}"), - ]); - // The assistant history turn carries the block it emitted - // (Python json.dumps spacing — the shape the model itself emits). - assert!(raw.contains( - "<|im_start|>assistant\n\n\ - {\"name\": \"get_weather\", \"arguments\": {\"city\": \"Paris\"}}\n\ - <|im_end|>\n" - )); - // Both results ride in ONE user turn, each in its own block. - assert!(raw.contains( - "<|im_start|>user\n\ - \n{\"temp_c\": 21}\n\n\ - \n{\"temp_c\": 24}\n<|im_end|>\n" - )); - assert_eq!(raw.matches("<|im_start|>user").count(), 2); - } - - #[test] - fn parse_extracts_calls_and_prose() { - let text = "Let me check.\n\n{\"name\": \"get_weather\", \ - \"arguments\": {\"city\": \"Paris\"}}\n\n\ - \n{\"name\": \"get_weather\", \"arguments\": \ - {\"city\": \"Lyon\"}}\n"; - let (calls, prose) = parse_tool_calls(text).unwrap(); - assert_eq!(calls.len(), 2); - assert_eq!(calls[0].name, "get_weather"); - assert_eq!(calls[0].arguments["city"], "Paris"); - assert_eq!(calls[1].arguments["city"], "Lyon"); - assert_eq!(prose, "Let me check."); - } - - #[test] - fn parse_of_plain_text_is_empty_not_an_error() { - let (calls, prose) = parse_tool_calls("The answer is 4.").unwrap(); - assert!(calls.is_empty()); - assert_eq!(prose, "The answer is 4."); - } - - #[test] - fn parse_rejects_unclosed_and_bad_json() { - assert!(matches!( - parse_tool_calls("\n{\"name\": \"x\"}"), - Err(ToolCallError::Unclosed { index: 0 }) - )); - assert!(matches!( - parse_tool_calls("\nnot json\n"), - Err(ToolCallError::BadJson { index: 0, .. }) - )); - } - - /// The whole loop: a rendered history turn re-parses to the same calls. - #[test] - fn tool_calls_round_trip_through_render_and_parse() { - let calls = vec![ToolCall { - name: "lookup".into(), - arguments: serde_json::json!({"q": "primes", "k": 5}), - }]; - let turn = Turn::assistant_tool_calls(&calls); - let (parsed, prose) = parse_tool_calls(&turn.content).unwrap(); - assert_eq!(parsed, calls); - assert!(prose.is_empty()); - } - - #[test] - #[should_panic(expected = "double it")] - fn trailing_assistant_turn_is_rejected() { - let _ = ChatMl::qwen2().render(&[Turn::user("q"), Turn::assistant("half-done")]); - } - - // ---- LFM (Pythonic) tool use ---------------------------------------- - - /// The tools line must match LFM2.5's `chat_template.jinja` byte shape: - /// `List of tools: [{bare tool json}, …]` appended to the system turn - /// with a `\n`, tools comma-joined, bare (no Hermes `"type":"function"` - /// wrapper — the model card's examples show `{"name": …}` directly). - #[test] - fn lfm_tools_render_matches_the_lfm25_template_shape() { - let raw = ChatMl::lfm2().render_with_tools( - &[weather_tool()], - &[ - Turn::system("You are a helpful assistant."), - Turn::user("Weather in Paris?"), - ], - ); - let expected = "<|startoftext|><|im_start|>system\nYou are a helpful assistant.\n\ - List of tools: [{\"name\": \"get_weather\", \"description\": \"Get the current weather for a city.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"city\": {\"type\": \"string\"}}, \"required\": [\"city\"]}}]<|im_end|>\n\ - <|im_start|>user\nWeather in Paris?<|im_end|>\n\ - <|im_start|>assistant\n"; - assert_eq!(raw, expected); - } - - /// LFM2.5's template injects NO default preamble: without a system turn - /// the tools line stands alone as the whole system prompt. - #[test] - fn lfm_tools_render_without_a_system_turn_has_no_preamble() { - let raw = ChatMl::lfm2().render_with_tools(&[weather_tool()], &[Turn::user("hi")]); - assert!(raw.starts_with("<|startoftext|><|im_start|>system\nList of tools: [")); - assert!(raw.contains("<|im_start|>user\nhi<|im_end|>")); - } - - #[test] - fn lfm_tools_comma_join_in_one_list() { - let mut second = weather_tool(); - second.name = "get_time".into(); - let raw = ChatMl::lfm2().render_with_tools(&[weather_tool(), second], &[Turn::user("hi")]); - assert!(raw.contains("\"name\": \"get_weather\"")); - assert!(raw.contains("}, {\"name\": \"get_time\"")); - assert_eq!(raw.matches("List of tools: [").count(), 1); - } - - /// LFM tool results are real `tool` role turns — one each, no Hermes - /// merging, no `` wrapper. - #[test] - fn lfm_tool_responses_render_as_tool_turns() { - let calls = [ToolCall { - name: "get_weather".into(), - arguments: serde_json::json!({"city": "Paris"}), - }]; - let raw = ChatMl::lfm2().render(&[ - Turn::user("Weather in Paris and Lyon?"), - Turn::assistant_tool_calls_lfm(&calls), - Turn::tool_response("{\"temp_c\": 21}"), - Turn::tool_response("{\"temp_c\": 24}"), - ]); - assert!(raw.contains( - "<|im_start|>assistant\n<|tool_call_start|>[get_weather(city=\"Paris\")]\ - <|tool_call_end|><|im_end|>\n" - )); - assert!(raw.contains("<|im_start|>tool\n{\"temp_c\": 21}<|im_end|>\n")); - assert!(raw.contains("<|im_start|>tool\n{\"temp_c\": 24}<|im_end|>\n")); - assert!(!raw.contains("")); - assert_eq!(raw.matches("<|im_start|>tool\n").count(), 2); - } - - /// Hermes rendering is untouched by the LFM additions: tool turns still - /// merge into a user turn. - #[test] - fn hermes_tool_turns_still_merge_after_the_style_split() { - let raw = ChatMl::qwen2().render(&[ - Turn::user("q"), - Turn::tool_response("r1"), - Turn::tool_response("r2"), - ]); - assert_eq!(raw.matches("<|im_start|>user").count(), 2); - assert!(!raw.contains("<|im_start|>tool")); - } - - /// The LFM2.5 template strips `` reasoning from every assistant - /// history turn but the LAST (keep_past_thinking=false default). - #[test] - fn lfm_strips_past_thinking_but_keeps_the_last() { - let raw = ChatMl::lfm2().render(&[ - Turn::user("a?"), - Turn::assistant("hmm\n\nAlpha."), - Turn::user("b?"), - Turn::assistant("later thoughts\n\nBeta."), - Turn::user("c?"), - ]); - assert!(raw.contains("<|im_start|>assistant\nAlpha.<|im_end|>")); - assert!(raw.contains("later thoughts")); - assert!(!raw.contains("hmm")); - } - - /// Qwen2 (Hermes) does no thinking-stripping — not part of its template. - #[test] - fn hermes_keeps_past_thinking_verbatim() { - let raw = ChatMl::qwen2().render(&[ - Turn::user("a?"), - Turn::assistant("hmmAlpha."), - Turn::user("b?"), - ]); - assert!(raw.contains("hmmAlpha.")); - } - - /// Qwen3 strips `` reasoning from assistant turns at/before the - /// last user query — and its strip is Python `lstrip('\n')`, newlines - /// only, NOT LFM's full trim (a leading space survives). - #[test] - fn qwen3_strips_thinking_at_or_before_the_last_user_query() { - let raw = ChatMl::qwen3().render(&[ - Turn::user("a?"), - Turn::assistant("hmm\n\n Alpha."), - Turn::user("b?"), - ]); - assert!(raw.contains("<|im_start|>assistant\n Alpha.<|im_end|>")); - assert!(!raw.contains("")); - } - - /// An assistant turn AFTER the last user query (mid tool loop) keeps its - /// reasoning, re-emitted in the template's normalized - /// `\n…\n\n\n` shape. - #[test] - fn qwen3_normalizes_thinking_after_the_last_user_query() { - let raw = ChatMl::qwen3().render(&[ - Turn::user("Weather in Paris?"), - Turn::assistant("need the toolcall below"), - Turn::tool_response("{\"temp_c\": 21}"), - ]); - assert!(raw.contains( - "<|im_start|>assistant\n\nneed the tool\n\n\ncall below<|im_end|>" - )); - } - - /// A pre-wrapped `` user turn is NOT a user query: the - /// last real query stays earlier, so a later assistant turn keeps its - /// reasoning (the template's multi_step_tool rule). - #[test] - fn qwen3_ignores_tool_response_user_turns_when_finding_the_last_query() { - let raw = ChatMl::qwen3().render(&[ - Turn::user("a?"), - Turn::assistant("hmmreply"), - Turn::user("\nr\n"), - ]); - assert!(raw.contains("\nhmm\n\n\nreply")); - } - - /// Assistant turns without reasoning render verbatim under Qwen3 — - /// wherever they sit relative to the last user query. - #[test] - fn qwen3_renders_thoughtless_assistant_turns_verbatim() { - let raw = ChatMl::qwen3().render(&[ - Turn::user("a?"), - Turn::assistant("Alpha."), - Turn::user("b?"), - ]); - assert!(raw.contains("<|im_start|>assistant\nAlpha.<|im_end|>")); - } - - /// Qwen3 with tools and no system turn injects NO preamble — the system - /// turn opens directly with `# Tools` (Qwen2 keeps the neutral default). - #[test] - fn qwen3_tools_without_system_turn_has_no_preamble() { - let raw = ChatMl::qwen3().render_with_tools(&[weather_tool()], &[Turn::user("hi")]); - assert!(raw.starts_with("<|im_start|>system\n# Tools\n\n")); - assert!(!raw.contains("You are a helpful assistant.")); - - let qwen2 = ChatMl::qwen2().render_with_tools(&[weather_tool()], &[Turn::user("hi")]); - assert!( - qwen2.starts_with("<|im_start|>system\nYou are a helpful assistant.\n\n# Tools\n\n") - ); - } - - #[test] - fn pythonic_rendering_covers_the_scalar_spellings() { - let calls = [ToolCall { - name: "f".into(), - arguments: serde_json::json!({ - "s": "he said \"hi\"\n", - "i": -3, - "x": 1.5, - "yes": true, - "no": false, - "nothing": null, - "list": [1, "two"], - "map": {"k": true} - }), - }]; - let turn = Turn::assistant_tool_calls_lfm(&calls); - // serde_json `preserve_order` (workspace feature): object keys iterate - // in INSERTION order — the same order Python/transformers renders. - assert_eq!( - turn.content, - "<|tool_call_start|>[f(s=\"he said \\\"hi\\\"\\n\", i=-3, x=1.5, \ - yes=True, no=False, nothing=None, list=[1, \"two\"], \ - map={\"k\": True})]<|tool_call_end|>" - ); - } - - /// The model card's own example parses to the exact call. - #[test] - fn lfm_parse_handles_the_model_card_example() { - let text = "<|tool_call_start|>[get_candidate_status(candidate_id=\"12345\")]\ - <|tool_call_end|>"; - let (calls, prose) = parse_tool_calls_lfm(text).unwrap(); - assert_eq!(calls.len(), 1); - assert_eq!(calls[0].name, "get_candidate_status"); - assert_eq!(calls[0].arguments["candidate_id"], "12345"); - assert!(prose.is_empty()); - } - - #[test] - fn lfm_parse_extracts_multiple_calls_and_prose() { - let text = "Checking both.\n<|tool_call_start|>[get_weather(city=\"Paris\"), \ - get_weather(city='Lyon', units=None)]<|tool_call_end|> done"; - let (calls, prose) = parse_tool_calls_lfm(text).unwrap(); - assert_eq!(calls.len(), 2); - assert_eq!(calls[0].arguments["city"], "Paris"); - assert_eq!(calls[1].arguments["city"], "Lyon"); - assert_eq!(calls[1].arguments["units"], serde_json::Value::Null); - assert_eq!(prose, "Checking both.\n done"); - } - - #[test] - fn lfm_parse_of_plain_text_is_empty_not_an_error() { - let (calls, prose) = parse_tool_calls_lfm("The answer is 4.").unwrap(); - assert!(calls.is_empty()); - assert_eq!(prose, "The answer is 4."); - } - - #[test] - fn lfm_parse_accepts_python_and_json_literal_spellings() { - let text = "<|tool_call_start|>[f(a=True, b=false, c=null, d=None, \ - e=[1, 2.5, -3], g={\"k\": \"v\", 'k2': True,})]<|tool_call_end|>"; - let (calls, _) = parse_tool_calls_lfm(text).unwrap(); - let args = &calls[0].arguments; - assert_eq!(args["a"], true); - assert_eq!(args["b"], false); - assert_eq!(args["c"], serde_json::Value::Null); - assert_eq!(args["d"], serde_json::Value::Null); - assert_eq!(args["e"], serde_json::json!([1, 2.5, -3])); - assert_eq!(args["g"], serde_json::json!({"k": "v", "k2": true})); - } - - #[test] - fn lfm_parse_handles_string_escapes() { - let text = "<|tool_call_start|>[f(s='it\\'s \\\"q\\\" \\u00e9\\n')]<|tool_call_end|>"; - let (calls, _) = parse_tool_calls_lfm(text).unwrap(); - assert_eq!(calls[0].arguments["s"], "it's \"q\" \u{e9}\n"); - } - - #[test] - fn lfm_parse_rejects_malformed_blocks_loudly() { - // Unclosed special token. - assert!(matches!( - parse_tool_calls_lfm("<|tool_call_start|>[f()]"), - Err(ToolCallError::Unclosed { index: 0 }) - )); - // Not a call list. - assert!(matches!( - parse_tool_calls_lfm("<|tool_call_start|>f()<|tool_call_end|>"), - Err(ToolCallError::Syntax { index: 0, .. }) - )); - // Positional args are not in the grammar. - assert!(matches!( - parse_tool_calls_lfm("<|tool_call_start|>[f(\"paris\")]<|tool_call_end|>"), - Err(ToolCallError::Syntax { .. }) - )); - // Unterminated string. - assert!(matches!( - parse_tool_calls_lfm("<|tool_call_start|>[f(a=\"oops)]<|tool_call_end|>"), - Err(ToolCallError::Syntax { .. }) - )); - // Trailing junk after the list. - assert!(matches!( - parse_tool_calls_lfm("<|tool_call_start|>[f()] junk<|tool_call_end|>"), - Err(ToolCallError::Syntax { .. }) - )); - } - - #[test] - fn lfm_parse_bounds_depth_and_call_count() { - // 9 levels of list nesting exceeds MAX_VALUE_DEPTH = 8. - let deep = format!( - "<|tool_call_start|>[f(a={}1{})]<|tool_call_end|>", - "[".repeat(9), - "]".repeat(9) - ); - assert!(matches!( - parse_tool_calls_lfm(&deep), - Err(ToolCallError::TooDeep) - )); - let many = format!( - "<|tool_call_start|>[{}]<|tool_call_end|>", - vec!["f()"; MAX_TOOL_CALLS + 1].join(", ") - ); - assert!(matches!( - parse_tool_calls_lfm(&many), - Err(ToolCallError::TooMany) - )); - } - - /// The whole LFM loop: a rendered history turn re-parses to the same - /// calls, through the Pythonic spelling and back. - #[test] - fn lfm_tool_calls_round_trip_through_render_and_parse() { - let calls = vec![ToolCall { - name: "lookup".into(), - arguments: serde_json::json!({"q": "primes", "k": 5, "deep": {"a": [true, null]}}), - }]; - let turn = Turn::assistant_tool_calls_lfm(&calls); - let (parsed, prose) = parse_tool_calls_lfm(&turn.content).unwrap(); - assert_eq!(parsed, calls); - assert!(prose.is_empty()); - } -} diff --git a/crates/mummu/examples/src/decode.rs b/crates/mummu/examples/src/decode.rs deleted file mode 100644 index 5e4b731..0000000 --- a/crates/mummu/examples/src/decode.rs +++ /dev/null @@ -1,441 +0,0 @@ -//! Decode-loop primitives shared by every causal model: on-device argmax, a -//! top-k probe, temperature/top-p sampling with a deterministic in-house RNG, -//! and the streaming `generate_loop` driver with cooperative cancellation. - -use std::ops::ControlFlow; - -use burn::tensor::Tensor; - -/// Hard ceiling on the vocab a sampled step will read back to the CPU -/// (~4 MB of f32 at the bound); anything larger is a wiring bug, not a model. -const VOCAB_READBACK_BOUND: usize = 1 << 20; - -/// Candidate-set cap when sampling: top-p truncation happens *within* the -/// `top_k` highest-logit tokens, so the post-softmax walk is O(k log k), not -/// O(vocab log vocab). 1024 keeps >99.9% of realistic nucleus mass. -const DEFAULT_TOP_K: usize = 1024; - -/// Greedy next-token id from `[1, vocab]` logits. The argmax runs -/// **on-device** and only the single winning index is synced back — vs. -/// copying a whole ~150k-logit vector to the CPU every decode step. -pub async fn argmax_id(logits: Tensor<2>) -> Result { - debug_assert!(logits.dims()[0] == 1, "argmax_id expects [1, vocab] logits"); - let data = logits - .argmax(1) - .into_data_async() - .await - .map_err(|e| format!("argmax readback: {e:?}"))? - .convert::() - .to_vec::() - .map_err(|e| format!("argmax readback: {e:?}"))?; - debug_assert!(data.len() == 1, "argmax over [1, vocab] must yield one id"); - let id = data.first().copied().ok_or("argmax returned no data")?; - Ok(id as u32) -} - -/// Indices of the `k` largest values, descending (the parity probe's top-k). -#[must_use] -pub fn top_k_ids(v: &[f32], k: usize) -> Vec { - assert!(!v.is_empty(), "top_k_ids: empty logits"); - assert!(k >= 1, "top_k_ids: k must be >= 1"); - let mut idx: Vec = (0..v.len()).collect(); - idx.sort_unstable_by(|&a, &b| v[b].partial_cmp(&v[a]).unwrap_or(std::cmp::Ordering::Equal)); - idx.into_iter().take(k).map(|i| i as u32).collect() -} - -/// Sampling knobs for one generation. `temperature == 0` means exact greedy -/// (argmax stays on-device; nothing else is consulted). -#[derive(Debug, Clone)] -pub struct SamplerOptions { - /// 0 = greedy; higher flattens the distribution. Must be finite and >= 0. - pub temperature: f32, - /// Nucleus mass in (0, 1]: sample only from the smallest prefix of - /// probability-sorted candidates whose mass reaches this. - pub top_p: f32, - /// Candidate-set cap applied before top-p (>= 1). - pub top_k: usize, - /// RNG seed: the same (options, logits, seed) always picks the same token. - pub seed: u64, -} - -impl Default for SamplerOptions { - fn default() -> Self { - Self { - temperature: 0.0, - top_p: 1.0, - top_k: DEFAULT_TOP_K, - seed: 0, - } - } -} - -impl SamplerOptions { - /// Greedy decoding (temperature 0) — the parity-gate configuration. - #[must_use] - pub fn greedy() -> Self { - Self::default() - } - - fn validate(&self) { - assert!( - self.temperature.is_finite() && self.temperature >= 0.0, - "sampler: temperature must be finite and >= 0, got {}", - self.temperature - ); - assert!( - self.top_p > 0.0 && self.top_p <= 1.0, - "sampler: top_p must be in (0, 1], got {}", - self.top_p - ); - assert!(self.top_k >= 1, "sampler: top_k must be >= 1"); - } -} - -/// PCG-XSH-RR 32 (O'Neill): a tiny deterministic RNG so sampling is -/// reproducible from a seed without pulling in a rand dependency. -pub struct Pcg32 { - state: u64, - inc: u64, -} - -impl Pcg32 { - const MULT: u64 = 6_364_136_223_846_793_005; - - #[must_use] - pub fn new(seed: u64) -> Self { - // Fixed stream; the standard seeding dance (advance, add, advance). - let mut rng = Self { - state: 0, - inc: (54 << 1) | 1, - }; - rng.next_u32(); - rng.state = rng.state.wrapping_add(seed); - rng.next_u32(); - rng - } - - pub fn next_u32(&mut self) -> u32 { - let old = self.state; - self.state = old.wrapping_mul(Self::MULT).wrapping_add(self.inc); - let xorshifted = (((old >> 18) ^ old) >> 27) as u32; - let rot = (old >> 59) as u32; - xorshifted.rotate_right(rot) - } - - /// Uniform in [0, 1) with 24 bits of mantissa. - pub fn next_f32(&mut self) -> f32 { - (self.next_u32() >> 8) as f32 * (1.0 / (1 << 24) as f32) - } -} - -/// Sample a token id from raw logits with temperature + top-k + top-p. -/// Pure and deterministic given (logits, opts, rng state). `temperature == 0` -/// callers should use [`argmax_id`] instead (asserted here). -#[must_use] -pub fn sample_id(logits: &[f32], opts: &SamplerOptions, rng: &mut Pcg32) -> u32 { - opts.validate(); - assert!(!logits.is_empty(), "sample_id: empty logits"); - assert!( - logits.len() <= VOCAB_READBACK_BOUND, - "sample_id: vocab {} exceeds the readback bound", - logits.len() - ); - assert!( - opts.temperature > 0.0, - "sample_id: temperature 0 is the argmax path" - ); - - // Top-k prefilter: O(vocab) partial select, then sort just the candidates. - let k = opts.top_k.min(logits.len()); - let mut idx: Vec = (0..logits.len() as u32).collect(); - let by_logit_desc = |&a: &u32, &b: &u32| logits[b as usize].total_cmp(&logits[a as usize]); - if k < idx.len() { - idx.select_nth_unstable_by(k - 1, by_logit_desc); - idx.truncate(k); - } - idx.sort_unstable_by(by_logit_desc); - - // Temperature softmax over the candidates (max-subtracted: never overflows). - let max_logit = logits[idx[0] as usize]; - let mut probs: Vec = idx - .iter() - .map(|&i| ((logits[i as usize] - max_logit) / opts.temperature).exp()) - .collect(); - let total: f32 = probs.iter().sum(); - debug_assert!(total > 0.0, "softmax mass must be positive"); - for p in &mut probs { - *p /= total; - } - - // Nucleus: keep the smallest probability-sorted prefix with mass >= top_p - // (probs are already descending because idx is logit-sorted). - let mut cut = probs.len(); - let mut mass = 0.0_f32; - for (i, &p) in probs.iter().enumerate() { - mass += p; - if mass >= opts.top_p { - cut = i + 1; - break; - } - } - debug_assert!(cut >= 1, "nucleus must keep at least the top token"); - - // Draw within the (renormalized) nucleus by cumulative walk. - let nucleus_mass: f32 = probs[..cut].iter().sum(); - let mut u = rng.next_f32() * nucleus_mass; - let mut chosen = idx[cut - 1]; // fallback: rounding can leave u > 0 at the end - for (i, &p) in probs[..cut].iter().enumerate() { - if u < p { - chosen = idx[i]; - break; - } - u -= p; - } - assert!( - (chosen as usize) < logits.len(), - "sampled id out of the vocab" - ); - chosen -} - -/// The shared decode driver: prefill once via `step`, then one token per -/// iteration. Emits each accepted token through `on_token`; a `Break` return -/// cancels cooperatively *before* the next forward. EOS is never emitted. -/// `step(new_ids, past)` returns `[1, vocab]` logits for the last position. -pub async fn generate_loop( - mut step: impl FnMut(&[u32], usize) -> Tensor<2>, - prompt_ids: &[u32], - max_tokens: usize, - opts: &SamplerOptions, - is_eos: impl Fn(u32) -> bool, - mut on_token: impl FnMut(u32) -> ControlFlow<()>, -) -> Result, String> { - opts.validate(); - assert!(!prompt_ids.is_empty(), "generate_loop: empty prompt"); - assert!(max_tokens >= 1, "generate_loop: max_tokens must be >= 1"); - - let mut rng = Pcg32::new(opts.seed); - let greedy = opts.temperature == 0.0; - let mut logits = { - let _s = crate::prof::scope("prefill"); - step(prompt_ids, 0) - }; - let mut out: Vec = Vec::with_capacity(max_tokens); - for past in (prompt_ids.len()..).take(max_tokens) { - let vocab = logits.dims()[1] as u32; - // This span crosses an await, so it must not hold a scope guard - // (thread-local stack; the future may resume on another worker) — - // timed by hand and attributed with `record` instead. It is also - // where every enqueued-but-unfinished GPU op comes due: the readback - // is the sync point, so GPU-side FFN time surfaces HERE, not in the - // scopes that enqueued it. - let readback_started = std::time::Instant::now(); - let next = if greedy { - argmax_id(logits).await? - } else { - let v = logits - .into_data_async() - .await - .map_err(|e| format!("logits readback: {e:?}"))? - .convert::() - .to_vec::() - .map_err(|e| format!("logits readback: {e:?}"))?; - sample_id(&v, opts, &mut rng) - }; - crate::prof::record("logits_readback+sample", readback_started.elapsed()); - // A GPU argmax over NaN logits can return an out-of-range sentinel - // (observed: exactly `vocab` on f16 numeric collapse) — fail loudly - // instead of emitting garbage ids the tokenizer silently drops. - if next >= vocab { - return Err(format!( - "decode step {past}: id {next} is outside the {vocab}-token vocab — NaN logits / numeric collapse on this backend?" - )); - } - if is_eos(next) { - break; - } - out.push(next); - if on_token(next).is_break() { - break; - } - // Cooperative yield: a CPU-backend decode is a long stretch of - // blocking compute between awaits, and without this a single - // generation would monopolize its worker for the whole request. - tokio::task::yield_now().await; - logits = { - let _s = crate::prof::scope("step"); - step(&[next], past) - }; - } - debug_assert!(out.len() <= max_tokens); - Ok(out) -} - -#[cfg(test)] -mod tests { - use super::*; - use burn::tensor::Tensor; - - #[tokio::test] - async fn argmax_id_finds_the_peak() { - let device = crate::backend::cpu_device(); - let logits = Tensor::<1>::from_floats([0.1, -2.0, 7.5, 3.0], &device).reshape([1, 4]); - assert_eq!(argmax_id(logits).await.unwrap(), 2); - } - - #[test] - fn pcg32_is_deterministic_per_seed_and_in_unit_range() { - let (mut a, mut b) = (Pcg32::new(7), Pcg32::new(7)); - let seq_a: Vec = (0..8).map(|_| a.next_u32()).collect(); - let seq_b: Vec = (0..8).map(|_| b.next_u32()).collect(); - assert_eq!(seq_a, seq_b, "same seed must replay the same stream"); - - let mut c = Pcg32::new(8); - let seq_c: Vec = (0..8).map(|_| c.next_u32()).collect(); - assert_ne!(seq_a, seq_c, "different seeds must diverge"); - - let mut r = Pcg32::new(99); - for _ in 0..1000 { - let f = r.next_f32(); - assert!((0.0..1.0).contains(&f), "next_f32 out of [0,1): {f}"); - } - } - - #[test] - fn sample_id_peaked_logits_always_pick_the_peak() { - let logits = [0.0f32, 30.0, -5.0, 1.0]; - let opts = SamplerOptions { - temperature: 0.8, - top_p: 0.95, - ..SamplerOptions::default() - }; - for seed in 0..32 { - let mut rng = Pcg32::new(seed); - assert_eq!(sample_id(&logits, &opts, &mut rng), 1); - } - } - - #[test] - fn sample_id_top_k_one_is_argmax_at_any_temperature() { - let logits = [1.0f32, 3.0, 2.0, 2.9]; - let opts = SamplerOptions { - temperature: 10.0, - top_p: 1.0, - top_k: 1, - seed: 0, - }; - for seed in 0..16 { - let mut rng = Pcg32::new(seed); - assert_eq!( - sample_id( - &logits, - &SamplerOptions { - seed, - ..opts.clone() - }, - &mut rng - ), - 1 - ); - } - } - - #[test] - fn sample_id_tiny_top_p_degenerates_to_argmax() { - let logits = [1.0f32, 1.1, 0.9, 1.05]; - let opts = SamplerOptions { - temperature: 5.0, - top_p: 0.01, - ..SamplerOptions::default() - }; - for seed in 0..16 { - let mut rng = Pcg32::new(seed); - assert_eq!(sample_id(&logits, &opts, &mut rng), 1); - } - } - - #[test] - fn sample_id_high_temperature_spreads_over_candidates() { - let logits = [2.0f32, 2.0, 2.0, 2.0]; - let opts = SamplerOptions { - temperature: 1.0, - top_p: 1.0, - ..SamplerOptions::default() - }; - let picks: std::collections::HashSet = (0..64) - .map(|seed| sample_id(&logits, &opts, &mut Pcg32::new(seed))) - .collect(); - assert!( - picks.len() >= 3, - "uniform logits over 64 seeds should hit >= 3 of 4 ids, got {picks:?}" - ); - for &p in &picks { - assert!(p < 4); - } - } - - #[test] - #[should_panic(expected = "argmax path")] - fn sample_id_rejects_temperature_zero() { - let mut rng = Pcg32::new(0); - let _ = sample_id(&[1.0, 2.0], &SamplerOptions::greedy(), &mut rng); - } - - /// A fixed toy vocab where the "model" always prefers id 2, then id 3 - /// after seeing 2 — enough to drive the loop without weights. - fn toy_step(device: &burn::tensor::Device) -> impl FnMut(&[u32], usize) -> Tensor<2> { - let device = device.clone(); - move |new_ids: &[u32], _past: usize| { - let peak = if new_ids.last() == Some(&2) { 3 } else { 2 }; - let mut v = vec![0.0f32; 8]; - v[peak] = 9.0; - Tensor::<1>::from_floats(v.as_slice(), &device).reshape([1, 8]) - } - } - - #[tokio::test] - async fn generate_loop_greedy_follows_argmax_and_stops_at_eos() { - let device = crate::backend::cpu_device(); - let out = generate_loop( - toy_step(&device), - &[1], - 6, - &SamplerOptions::greedy(), - |id| id == 3, // treat the follow-up token as EOS - |_| std::ops::ControlFlow::Continue(()), - ) - .await.unwrap(); - assert_eq!(out, vec![2], "one token, then EOS never emitted"); - } - - #[tokio::test] - async fn generate_loop_cancels_cooperatively_between_tokens() { - let device = crate::backend::cpu_device(); - let mut streamed = Vec::new(); - let out = generate_loop( - toy_step(&device), - &[1], - 100, - &SamplerOptions::greedy(), - |_| false, // no EOS: only the callback can stop this - |id| { - streamed.push(id); - if streamed.len() == 2 { - std::ops::ControlFlow::Break(()) - } else { - std::ops::ControlFlow::Continue(()) - } - }, - ) - .await - .unwrap(); - assert_eq!(out.len(), 2, "break after the 2nd token stops the loop"); - assert_eq!(streamed, out, "every emitted token was streamed"); - } - - #[test] - fn top_k_ids_orders_descending() { - let v = [0.1f32, 5.0, -2.0, 3.0]; - assert_eq!(top_k_ids(&v, 3), vec![1, 3, 0]); - } -} diff --git a/crates/mummu/examples/src/gguf.rs b/crates/mummu/examples/src/gguf.rs deleted file mode 100644 index 8589aa8..0000000 --- a/crates/mummu/examples/src/gguf.rs +++ /dev/null @@ -1,2029 +0,0 @@ -//! GGUF container reader — the first slice of P3's "run what the ecosystem -//! ships" import path. GGUF (llama.cpp's format) is one file: a small header -//! of typed metadata key-values, a tensor table (name, shape, quantized -//! dtype, offset), then an aligned blob of tensor payloads. -//! -//! This module reads the *header*: every metadata value typed and bounded, -//! every tensor located and size-checked — and fails loudly on anything -//! malformed, oversized, or unknown. Tensor payloads are *located*, never -//! loaded here; dequantizing them into Burn tensors is the next slice. - -use std::fs::File; -use std::io::{BufReader, Read, Seek, Write}; -use std::path::Path; - -/// GGUF file magic, little-endian `"GGUF"`. -const MAGIC: [u8; 4] = *b"GGUF"; - -/// Versions this reader understands (v2 moved counts to u64; v3 is v2 plus a -/// big-endian variant this reader rejects by magic). -const SUPPORTED_VERSIONS: [u32; 2] = [2, 3]; - -/// Default payload alignment when `general.alignment` is absent. -const DEFAULT_ALIGNMENT: u64 = 32; - -/// Most metadata key-values a sane model file carries (real models: ~20-40). -const MAX_KVS: u64 = 4096; - -/// Most tensors a supported model carries (Qwen2.5-1.5B: 339). -const MAX_TENSORS: u64 = 65_536; - -/// Longest metadata string (chat templates run ~10 KiB; 1 MiB is generous). -const MAX_STRING_BYTES: u64 = 1 << 20; - -/// Longest metadata array (tokenizer vocab/merges run ~152k entries). -const MAX_ARRAY_LEN: u64 = 1 << 22; - -/// GGML allows at most 4 tensor dimensions. -const MAX_DIMS: u32 = 4; - -/// Largest dequantized-to-f32 payload either dequant will produce (a ~12B-param -/// model; the reference machine has 128 GB). Only -/// [`GgufFile::dequant_to_safetensors`] holds this much at once; -/// [`GgufFile::dequant_to_safetensors_file`] streams it and never buffers more -/// than one tensor. -const MAX_DEQUANT_BYTES: u64 = 48 << 30; - -/// Largest SINGLE dequantized tensor the streaming dequant will hold. The -/// widest real one is a vocabulary embedding — Qwen2.5-1.5B's is 933 MB at -/// f32, a 35B's would be ~3.1 GB — so 8 GiB is a corrupt header claiming an -/// absurd tensor. It is twice `safetensors::MAX_PART_BYTES` on purpose: that -/// path keeps the source dtype, this one widens everything to f32. -const MAX_TENSOR_F32_BYTES: u64 = 8 << 30; - -/// Bytes staged between the f32 values and the sink. Fixed and small: the -/// point of the streaming dequant is that peak allocation tracks the widest -/// TENSOR, never the payload, and this buffer must not reintroduce a second -/// copy of either. -const DEQUANT_STAGE_BYTES: usize = 1 << 20; - -/// What went wrong reading a GGUF header. -#[derive(Debug, thiserror::Error)] -pub enum GgufError { - #[error("gguf {path}: {source}")] - Io { - path: String, - source: std::io::Error, - }, - #[error("gguf {path}: bad magic {found:02x?} (big-endian GGUF is unsupported)")] - BadMagic { path: String, found: [u8; 4] }, - #[error("gguf {path}: unsupported version {version} (supported: {SUPPORTED_VERSIONS:?})")] - UnsupportedVersion { path: String, version: u32 }, - #[error("gguf {path}: {what} count {count} exceeds the {bound} bound")] - OverBound { - path: String, - what: &'static str, - count: u64, - bound: u64, - }, - #[error("gguf {path}: metadata '{key}': {reason}")] - BadValue { - path: String, - key: String, - reason: String, - }, - #[error("gguf {path}: tensor {index}: {reason}")] - BadTensor { - path: String, - index: usize, - reason: String, - }, -} - -/// `u64` -> `usize` as an error rather than a panic. -/// -/// Every caller has already bounded `value` against one of the `MAX_*` -/// ceilings or a tensor's own declared size, so on a 64-bit target this -/// cannot fail — but an import path is exactly where "cannot fail" should -/// still be an `Err` instead of an `.expect()`, so a 32-bit build degrades to -/// a clean error instead of aborting mid-load. (Same helper, same reasoning, -/// as `safetensors::to_usize`.) -fn to_usize(value: u64, path: &str, what: &'static str) -> Result { - debug_assert!( - value <= usize::MAX as u64, - "{what} fits usize on this target" - ); - debug_assert!(!path.is_empty(), "an error needs a file to name"); - usize::try_from(value).map_err(|_| GgufError::OverBound { - path: path.to_string(), - what, - count: value, - bound: usize::MAX as u64, - }) -} - -/// One typed metadata value. Arrays are homogeneous per the spec; nested -/// arrays are legal but bounded to one level of nesting in practice. -#[derive(Debug, Clone, PartialEq)] -pub enum GgufValue { - U8(u8), - I8(i8), - U16(u16), - I16(i16), - U32(u32), - I32(i32), - F32(f32), - Bool(bool), - Str(String), - Array(Vec), - U64(u64), - I64(i64), - F64(f64), -} - -impl GgufValue { - /// The value as a string, if it is one. - #[must_use] - pub fn as_str(&self) -> Option<&str> { - match self { - Self::Str(s) => Some(s), - _ => None, - } - } - - /// The value widened to u64, if it is any unsigned integer. - #[must_use] - pub fn as_u64(&self) -> Option { - match *self { - Self::U8(v) => Some(u64::from(v)), - Self::U16(v) => Some(u64::from(v)), - Self::U32(v) => Some(u64::from(v)), - Self::U64(v) => Some(v), - _ => None, - } - } - - /// The value widened to i64, if it is any integer (signed or unsigned) - /// that fits. - #[must_use] - pub fn as_i64(&self) -> Option { - match *self { - Self::I8(v) => Some(i64::from(v)), - Self::I16(v) => Some(i64::from(v)), - Self::I32(v) => Some(i64::from(v)), - Self::I64(v) => Some(v), - _ => self.as_u64().and_then(|v| i64::try_from(v).ok()), - } - } - - /// The value as f32, if it is one. - #[must_use] - pub fn as_f32(&self) -> Option { - match *self { - Self::F32(v) => Some(v), - _ => None, - } - } - - /// The value as a bool, if it is one. - #[must_use] - pub fn as_bool(&self) -> Option { - match *self { - Self::Bool(v) => Some(v), - _ => None, - } - } - - /// The value as an array slice, if it is one. - #[must_use] - pub fn as_array(&self) -> Option<&[GgufValue]> { - match self { - Self::Array(items) => Some(items), - _ => None, - } - } -} - -/// A GGML tensor dtype, as stored on disk. Quantized types pack fixed-size -/// blocks; `block_size`/`bytes_per_block` give the layout the dequant slice -/// (and size validation here) needs. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -#[allow(non_camel_case_types)] // the ecosystem's canonical spellings -pub enum GgmlType { - F32, - F16, - Q4_0, - Q4_1, - Q5_0, - Q5_1, - Q8_0, - Q8_1, - Q2_K, - Q3_K, - Q4_K, - Q5_K, - Q6_K, - Q8_K, - /// IQ4_XS: 4-bit non-linear (a 16-entry value table) with 6-bit - /// sub-scales — the workhorse of unsloth's UD dynamic quants. - IQ4_XS, - /// IQ4_NL: IQ4_XS's 16-entry table in a simple 32-element block. - IQ4_NL, - /// IQ2_XS: 2.3125 bpw — 8-value rows from a 512-entry codebook grid. - IQ2_XS, - /// IQ2_S: 2.5625 bpw — the 1024-entry grid with separate sign bytes. - IQ2_S, - /// IQ3_XXS: 3.0625 bpw — 4-value rows from a 256-entry grid. - IQ3_XXS, - /// IQ3_S: 3.4375 bpw — the 512-entry grid, high index bits in `qh`. - IQ3_S, - BF16, -} - -impl GgmlType { - /// Decode the on-disk type id; unknown ids are a loud error, never a guess. - fn from_id(id: u32) -> Option { - let ty = match id { - 0 => Self::F32, - 1 => Self::F16, - 2 => Self::Q4_0, - 3 => Self::Q4_1, - 6 => Self::Q5_0, - 7 => Self::Q5_1, - 8 => Self::Q8_0, - 9 => Self::Q8_1, - 10 => Self::Q2_K, - 11 => Self::Q3_K, - 12 => Self::Q4_K, - 13 => Self::Q5_K, - 14 => Self::Q6_K, - 15 => Self::Q8_K, - 17 => Self::IQ2_XS, - 18 => Self::IQ3_XXS, - 20 => Self::IQ4_NL, - 21 => Self::IQ3_S, - 22 => Self::IQ2_S, - 23 => Self::IQ4_XS, - 30 => Self::BF16, - _ => return None, - }; - Some(ty) - } - - /// Elements per quantization block (1 for plain float types). - #[must_use] - pub fn block_size(self) -> u64 { - match self { - Self::F32 | Self::F16 | Self::BF16 => 1, - Self::Q4_0 | Self::Q4_1 | Self::Q5_0 | Self::Q5_1 | Self::Q8_0 | Self::Q8_1 - | Self::IQ4_NL => 32, - Self::Q2_K | Self::Q3_K | Self::Q4_K | Self::Q5_K | Self::Q6_K | Self::Q8_K - | Self::IQ4_XS | Self::IQ2_XS | Self::IQ2_S | Self::IQ3_XXS | Self::IQ3_S => 256, - } - } - - /// Bytes one block occupies on disk (ggml's type sizes). - #[must_use] - pub fn bytes_per_block(self) -> u64 { - match self { - Self::F32 => 4, - Self::F16 | Self::BF16 => 2, - Self::Q4_0 => 18, // f16 d + 16 B qs - Self::Q4_1 => 20, // f16 d + f16 m + 16 B qs - Self::Q5_0 => 22, // f16 d + 4 B qh + 16 B qs - Self::Q5_1 => 24, // f16 d + f16 m + 4 B qh + 16 B qs - Self::Q8_0 => 34, // f16 d + 32 i8 - Self::Q8_1 => 36, // f16 d + f16 s + 32 i8 - Self::Q2_K => 84, // 16 B scales + 64 B qs + f16 d + f16 dmin - Self::Q3_K => 110, // 32 B hmask + 64 B qs + 12 B scales + f16 d - Self::Q4_K => 144, // f16 d + f16 dmin + 12 B scales + 128 B qs - Self::Q5_K => 176, // Q4_K + 32 B qh - Self::Q6_K => 210, // 128 B ql + 64 B qh + 16 i8 scales + f16 d - Self::Q8_K => 292, // f32 d + 256 i8 + 16 i16 bsums - Self::IQ4_XS => 136, // f16 d + u16 scales_h + 4 B scales_l + 128 B qs - Self::IQ4_NL => 18, // f16 d + 16 B qs - Self::IQ2_XS => 74, // f16 d + 32 u16 (grid|signs) + 8 B scales - Self::IQ2_S => 82, // f16 d + 64 B qs(idx lo + signs) + 8 B qh + 8 B scales - Self::IQ3_XXS => 98, // f16 d + 64 B idx + 32 B (scale|signs) words - Self::IQ3_S => 110, // f16 d + 64 B qs + 8 B qh + 32 B signs + 4 B scales - } - } -} - -/// One entry of the tensor table: where a tensor lives and what shape/dtype -/// it has. `offset` is relative to [`GgufFile::data_offset`] and is a -/// multiple of the file's alignment (validated on read). -#[derive(Debug, Clone)] -pub struct GgufTensorInfo { - pub name: String, - /// Dimensions in ggml order (fastest-varying first — the *reverse* of - /// the row-major order safetensors/PyTorch shapes use). - pub dims: Vec, - pub dtype: GgmlType, - pub offset: u64, -} - -impl GgufTensorInfo { - /// Total element count. - #[must_use] - pub fn element_count(&self) -> u64 { - self.dims.iter().product() - } - - /// Exact on-disk payload size. Element counts are validated to be whole - /// blocks at parse time, so this is always exact for parsed tensors. - #[must_use] - pub fn byte_len(&self) -> u64 { - let elements = self.element_count(); - debug_assert!( - elements.is_multiple_of(self.dtype.block_size()), - "parse validated whole blocks" - ); - elements / self.dtype.block_size() * self.dtype.bytes_per_block() - } -} - -/// How one GGUF tensor lands in the safetensors blob -/// ([`GgufFile::dequant_to_safetensors`]). -#[derive(Debug, Clone)] -pub enum GgufMap { - /// Target name; shape = the GGUF dims reversed (the row-major twin of - /// the ggml layout — right for everything but squeezed kernels). - Rename(String), - /// Target name + explicit row-major shape (same element count, same - /// bytes — e.g. un-squeezing a depthwise conv kernel back to - /// `[channels, 1, k]`). - Reshape(String, Vec), - /// Deliberately drop this tensor (e.g. qwen35's unused NextN/MTP block). - /// Distinct from an unmapped name, which stays a loud error. - Skip, -} - -/// Where one tensor lands in the output blob, decided before any payload -/// byte is read ([`GgufFile::plan_dequant`]). -#[derive(Debug)] -struct PlannedTensor { - /// Index into [`GgufFile::tensors`] — the payload to dequantize. - source: usize, - name: String, - shape: Vec, - /// Byte offset within the payload; contiguous and ascending. - start: u64, - /// Dequantized f32 byte length. - len: u64, -} - -/// A parsed GGUF header: typed metadata + the located tensor table. -#[derive(Debug)] -pub struct GgufFile { - /// Where this header was read from (payload reads re-open it). - pub path: std::path::PathBuf, - pub version: u32, - /// Metadata in file order (keys are unique per the spec). - pub metadata: Vec<(String, GgufValue)>, - pub tensors: Vec, - /// Payload alignment (`general.alignment`, default 32). - pub alignment: u64, - /// Absolute file offset where the aligned tensor payload blob begins. - pub data_offset: u64, -} - -impl GgufFile { - /// Read and validate a GGUF header (metadata + tensor table only — no - /// tensor payloads are loaded). - pub fn open(path: &Path) -> Result { - let file = File::open(path).map_err(|source| GgufError::Io { - path: path.display().to_string(), - source, - })?; - let mut r = Reader { - inner: BufReader::new(file), - path: path.display().to_string(), - }; - let parsed = r.read_file(path)?; - // Positive space: the header must end before the payload it locates. - assert!( - parsed.data_offset.is_multiple_of(parsed.alignment), - "data offset is aligned by construction" - ); - Ok(parsed) - } - - /// Read one tensor's payload and dequantize it to f32, in the on-disk - /// (ggml fastest-varying-first) element order. - pub fn read_tensor_f32(&self, name: &str) -> Result, GgufError> { - use std::io::SeekFrom; - let info = self.tensor(name).ok_or_else(|| GgufError::BadValue { - path: self.path.display().to_string(), - key: name.to_string(), - reason: "no such tensor".into(), - })?; - let mut file = File::open(&self.path).map_err(|source| GgufError::Io { - path: self.path.display().to_string(), - source, - })?; - file.seek(SeekFrom::Start(self.data_offset + info.offset)) - .map_err(|source| GgufError::Io { - path: self.path.display().to_string(), - source, - })?; - let byte_len = usize::try_from(info.byte_len()).map_err(|_| GgufError::OverBound { - path: self.path.display().to_string(), - what: "tensor payload bytes", - count: info.byte_len(), - bound: usize::MAX as u64, - })?; - let mut bytes = vec![0u8; byte_len]; - file.read_exact(&mut bytes) - .map_err(|source| GgufError::Io { - path: self.path.display().to_string(), - source, - })?; - let out = dequantize(info.dtype, &bytes).map_err(|reason| GgufError::BadTensor { - path: self.path.display().to_string(), - index: 0, - reason: format!("{name}: {reason}"), - })?; - assert_eq!( - out.len() as u64, - info.element_count(), - "dequant must yield exactly the tensor's elements" - ); - Ok(out) - } - - /// Dequantize every tensor to f32 and serialize the result as an - /// in-memory **safetensors** file — the bridge onto the exact store - /// pipeline (adapters, key remaps, checked load) the safetensors path - /// already trusts. `map` maps each GGUF tensor to the name (and - /// optionally an explicit shape) the blob should carry (HF-checkpoint - /// naming, so per-model remap tables apply unchanged); an unmapped - /// tensor is a loud error, never a skip — a name this crate doesn't - /// recognize means weights would silently vanish from the model. - /// - /// Default shapes are the GGUF dims **reversed**: ggml orders dims - /// fastest-varying-first, so the raw payload bytes are exactly the - /// row-major layout of the reversed shape — same bytes, HF convention. - /// [`GgufMap::Reshape`] overrides that for tensors whose checkpoint - /// shape differs by more than dim order (e.g. llama.cpp squeezes the - /// middle 1 out of depthwise-conv kernels). - /// - /// This form needs the whole f32 payload resident. For a model whose - /// dequantized size is a meaningful fraction of RAM, use - /// [`Self::dequant_to_safetensors_file`] — the two produce byte-identical - /// output. - pub fn dequant_to_safetensors( - &self, - map: &dyn Fn(&GgufTensorInfo) -> Option, - ) -> Result, GgufError> { - let mut blob = Vec::new(); - self.dequant_into(map, &mut blob)?; - assert!(blob.len() > 8, "the header length prefix is always written"); - Ok(blob) - } - - /// [`Self::dequant_to_safetensors`] straight to a file, never holding the - /// payload in RAM. Returns the payload bytes written (header excluded). - /// - /// This is the variant a real quantized checkpoint wants. The in-memory - /// form needs the whole f32 payload resident — ~28 GB for OLMoE-1B-7B — - /// *on top of* the model the load then builds from it, and that sum is - /// what a 128 GB box with other tenants actually fails to satisfy. - /// Writing to disk trades the spike for temp space and lets - /// `SafetensorsStore::from_file` page the weights in as it needs them. - /// (`safetensors::fuse_checkpoint_to_file` is the same trade on the - /// unquantized path.) - pub fn dequant_to_safetensors_file( - &self, - map: &dyn Fn(&GgufTensorInfo) -> Option, - out: &Path, - ) -> Result { - assert!(!out.as_os_str().is_empty(), "the sink file must be named"); - let io = |source: std::io::Error| GgufError::Io { - path: out.display().to_string(), - source, - }; - let file = File::create(out).map_err(io)?; - let mut sink = std::io::BufWriter::with_capacity(DEQUANT_STAGE_BYTES, file); - let written = self.dequant_into(map, &mut sink)?; - sink.flush().map_err(io)?; - sink.into_inner() - .map_err(|e| GgufError::Io { - path: out.display().to_string(), - source: e.into_error(), - })? - .sync_all() - .map_err(io)?; - Ok(written) - } - - /// The shared dequant: plan, then stream header + payload into `sink` in - /// output order, one tensor at a time. - /// - /// Planning first is not only what makes streaming possible — it moves - /// every *claim* check (unmapped name, reshape that changes the element - /// count, rename collision) ahead of the first payload byte, so a bad map - /// fails in milliseconds instead of after N tensors have been - /// dequantized. - fn dequant_into( - &self, - map: &dyn Fn(&GgufTensorInfo) -> Option, - sink: &mut W, - ) -> Result { - let path = self.path.display().to_string(); - let plan = self.plan_dequant(map)?; - let total: u64 = plan.iter().map(|p| p.len).sum(); - - // Header first: names, dtypes, shapes, and the contiguous offsets the - // copy pass will fill. - let mut header = String::from("{"); - for (i, p) in plan.iter().enumerate() { - if i > 0 { - header.push(','); - } - let json_name = serde_json::to_string(&p.name).map_err(|e| GgufError::BadTensor { - path: path.clone(), - index: p.source, - reason: format!("name is not encodable as JSON: {e}"), - })?; - header.push_str(&format!( - "{json_name}:{{\"dtype\":\"F32\",\"shape\":{:?},\"data_offsets\":[{},{}]}}", - p.shape, - p.start, - p.start + p.len, - )); - } - header.push('}'); - - let io = |source: std::io::Error| GgufError::Io { - path: path.clone(), - source, - }; - sink.write_all(&(header.len() as u64).to_le_bytes()) - .map_err(io)?; - sink.write_all(header.as_bytes()).map_err(io)?; - - // One small staging buffer, reused for every tensor: peak allocation - // is the widest tensor's f32 values (from `read_tensor_f32`) plus this - // 1 MiB, never a second copy of the payload. - let mut stage: Vec = Vec::with_capacity(DEQUANT_STAGE_BYTES); - let mut written = 0u64; - for p in &plan { - debug_assert_eq!(written, p.start, "tensors are written in output order"); - let values = self.read_tensor_f32(&self.tensors[p.source].name)?; - debug_assert_eq!( - values.len() as u64 * 4, - p.len, - "the plan sized this tensor from the same element count" - ); - for chunk in values.chunks(DEQUANT_STAGE_BYTES / 4) { - stage.clear(); - stage.extend(chunk.iter().flat_map(|v| v.to_le_bytes())); - sink.write_all(&stage).map_err(io)?; - written += stage.len() as u64; - } - } - assert_eq!( - written, total, - "every planned byte was written exactly once" - ); - assert!( - stage.capacity() <= DEQUANT_STAGE_BYTES, - "the staging buffer never grew past its bound" - ); - Ok(written) - } - - /// First pass: decide what every tensor is called, what shape it claims, - /// and where its bytes land — reading no payload at all. - /// - /// Offsets are contiguous and ascending in tensor-table order, which is - /// what lets the copy pass be a single forward stream. A - /// [`GgufMap::Skip`] tensor (e.g. qwen35's unused NextN block) is dropped - /// here: it earns no plan entry and counts toward neither the size bound - /// nor the blob. - fn plan_dequant( - &self, - map: &dyn Fn(&GgufTensorInfo) -> Option, - ) -> Result, GgufError> { - let bad = |index: usize, reason: String| GgufError::BadTensor { - path: self.path.display().to_string(), - index, - reason, - }; - - let mut plan: Vec = Vec::with_capacity(self.tensors.len()); - let mut names = std::collections::HashSet::with_capacity(self.tensors.len()); - let mut start = 0u64; - for (index, info) in self.tensors.iter().enumerate() { - let (name, shape) = match map(info) { - Some(GgufMap::Rename(name)) => { - (name, info.dims.iter().rev().copied().collect::>()) - } - Some(GgufMap::Reshape(name, shape)) => { - if shape.iter().product::() != info.element_count() { - return Err(bad( - index, - format!( - "reshape of '{}' to {shape:?} changes the element count", - info.name - ), - )); - } - (name, shape) - } - Some(GgufMap::Skip) => continue, - None => { - return Err(bad(index, format!("unmapped tensor name '{}'", info.name))); - } - }; - if !names.insert(name.clone()) { - return Err(bad(index, format!("rename collision on '{name}'"))); - } - let len = info.element_count() * 4; - if len > MAX_TENSOR_F32_BYTES { - return Err(GgufError::OverBound { - path: self.path.display().to_string(), - what: "single dequantized tensor bytes", - count: len, - bound: MAX_TENSOR_F32_BYTES, - }); - } - plan.push(PlannedTensor { - source: index, - name, - shape, - start, - len, - }); - start += len; - } - // The bound is on the payload actually planned (skips excluded); this - // reads no payload, so checking after the plan loop still fails before - // `dequant_into` copies a single byte. - if start > MAX_DEQUANT_BYTES { - return Err(GgufError::OverBound { - path: self.path.display().to_string(), - what: "dequantized f32 payload bytes", - count: start, - bound: MAX_DEQUANT_BYTES, - }); - } - Ok(plan) - } - - /// Look up a metadata value by exact key. - #[must_use] - pub fn get(&self, key: &str) -> Option<&GgufValue> { - self.metadata.iter().find(|(k, _)| k == key).map(|(_, v)| v) - } - - /// The model architecture (`general.architecture`), when present. - #[must_use] - pub fn architecture(&self) -> Option<&str> { - self.get("general.architecture").and_then(GgufValue::as_str) - } - - /// Look up a tensor by exact name. - #[must_use] - pub fn tensor(&self, name: &str) -> Option<&GgufTensorInfo> { - self.tensors.iter().find(|t| t.name == name) - } -} - -/// Sequential little-endian reader over the header bytes. -struct Reader { - inner: BufReader, - path: String, -} - -impl Reader { - fn io_err(&self, source: std::io::Error) -> GgufError { - GgufError::Io { - path: self.path.clone(), - source, - } - } - - fn bytes(&mut self) -> Result<[u8; N], GgufError> { - let mut buf = [0u8; N]; - self.inner - .read_exact(&mut buf) - .map_err(|e| self.io_err(e))?; - Ok(buf) - } - - fn u32(&mut self) -> Result { - Ok(u32::from_le_bytes(self.bytes()?)) - } - - fn u64(&mut self) -> Result { - Ok(u64::from_le_bytes(self.bytes()?)) - } - - /// A length-prefixed UTF-8 string, bounded by [`MAX_STRING_BYTES`]. - fn string(&mut self, what: &'static str) -> Result { - let len = self.u64()?; - if len > MAX_STRING_BYTES { - return Err(GgufError::OverBound { - path: self.path.clone(), - what, - count: len, - bound: MAX_STRING_BYTES, - }); - } - let mut buf = vec![0u8; to_usize(len, &self.path, what)?]; - self.inner - .read_exact(&mut buf) - .map_err(|e| self.io_err(e))?; - String::from_utf8(buf).map_err(|e| GgufError::BadValue { - path: self.path.clone(), - key: what.to_string(), - reason: format!("invalid UTF-8: {e}"), - }) - } - - /// One typed metadata value. `depth` bounds array nesting. - fn value(&mut self, key: &str, type_id: u32, depth: u32) -> Result { - let bad = |reason: String, path: &str| GgufError::BadValue { - path: path.to_string(), - key: key.to_string(), - reason, - }; - let v = match type_id { - 0 => GgufValue::U8(self.bytes::<1>()?[0]), - #[allow(clippy::cast_possible_wrap)] // bit-exact reinterpret is the format - 1 => GgufValue::I8(self.bytes::<1>()?[0] as i8), - 2 => GgufValue::U16(u16::from_le_bytes(self.bytes()?)), - 3 => GgufValue::I16(i16::from_le_bytes(self.bytes()?)), - 4 => GgufValue::U32(self.u32()?), - 5 => GgufValue::I32(i32::from_le_bytes(self.bytes()?)), - 6 => GgufValue::F32(f32::from_le_bytes(self.bytes()?)), - 7 => match self.bytes::<1>()?[0] { - 0 => GgufValue::Bool(false), - 1 => GgufValue::Bool(true), - other => return Err(bad(format!("bool byte {other}"), &self.path)), - }, - 8 => GgufValue::Str(self.string("metadata string")?), - 9 => { - if depth >= 2 { - return Err(bad("arrays nested deeper than 2".into(), &self.path)); - } - let elem_type = self.u32()?; - let len = self.u64()?; - if len > MAX_ARRAY_LEN { - return Err(GgufError::OverBound { - path: self.path.clone(), - what: "metadata array", - count: len, - bound: MAX_ARRAY_LEN, - }); - } - let mut items = Vec::with_capacity(to_usize(len, &self.path, "metadata array")?); - for _ in 0..len { - items.push(self.value(key, elem_type, depth + 1)?); - } - GgufValue::Array(items) - } - 10 => GgufValue::U64(self.u64()?), - 11 => GgufValue::I64(i64::from_le_bytes(self.bytes()?)), - 12 => GgufValue::F64(f64::from_le_bytes(self.bytes()?)), - other => return Err(bad(format!("unknown value type {other}"), &self.path)), - }; - Ok(v) - } - - /// The whole header: magic, version, metadata, tensor table, alignment. - fn read_file(&mut self, source_path: &Path) -> Result { - let magic = self.bytes::<4>()?; - if magic != MAGIC { - return Err(GgufError::BadMagic { - path: self.path.clone(), - found: magic, - }); - } - let version = self.u32()?; - if !SUPPORTED_VERSIONS.contains(&version) { - return Err(GgufError::UnsupportedVersion { - path: self.path.clone(), - version, - }); - } - let tensor_count = self.u64()?; - let kv_count = self.u64()?; - for (what, count, bound) in [ - ("tensor", tensor_count, MAX_TENSORS), - ("metadata kv", kv_count, MAX_KVS), - ] { - if count > bound { - return Err(GgufError::OverBound { - path: self.path.clone(), - what, - count, - bound, - }); - } - } - - let mut metadata = - Vec::with_capacity(to_usize(kv_count, &self.path, "metadata key-values")?); - for _ in 0..kv_count { - let key = self.string("metadata key")?; - let type_id = self.u32()?; - let value = self.value(&key, type_id, 0)?; - metadata.push((key, value)); - } - - let alignment = metadata - .iter() - .find(|(k, _)| k == "general.alignment") - .and_then(|(_, v)| v.as_u64()) - .unwrap_or(DEFAULT_ALIGNMENT); - if alignment == 0 || !alignment.is_power_of_two() { - return Err(GgufError::BadValue { - path: self.path.clone(), - key: "general.alignment".into(), - reason: format!("{alignment} is not a power of two"), - }); - } - - let tensors = self.tensor_table(tensor_count, alignment)?; - let header_end = self.inner.stream_position().map_err(|e| self.io_err(e))?; - let data_offset = header_end.div_ceil(alignment) * alignment; - debug_assert!(data_offset >= header_end, "padding never rewinds"); - Ok(GgufFile { - path: source_path.to_path_buf(), - version, - metadata, - tensors, - alignment, - data_offset, - }) - } - - /// The tensor table, with every entry's shape/dtype/offset validated. - fn tensor_table( - &mut self, - count: u64, - alignment: u64, - ) -> Result, GgufError> { - assert!(count <= MAX_TENSORS, "caller bounded the count"); - assert!(alignment.is_power_of_two(), "caller validated alignment"); - let count_usize = to_usize(count, &self.path, "tensor count")?; - let mut tensors: Vec = Vec::with_capacity(count_usize); - for index in 0..count_usize { - let bad = |reason: String, path: &str| GgufError::BadTensor { - path: path.to_string(), - index, - reason, - }; - let name = self.string("tensor name")?; - let n_dims = self.u32()?; - if n_dims == 0 || n_dims > MAX_DIMS { - return Err(bad(format!("{n_dims} dims (1..={MAX_DIMS})"), &self.path)); - } - let mut dims = Vec::with_capacity(n_dims as usize); - for _ in 0..n_dims { - dims.push(self.u64()?); - } - let type_id = self.u32()?; - let Some(dtype) = GgmlType::from_id(type_id) else { - return Err(bad(format!("unknown ggml type id {type_id}"), &self.path)); - }; - let offset = self.u64()?; - if !offset.is_multiple_of(alignment) { - return Err(bad( - format!("offset {offset} not {alignment}-aligned"), - &self.path, - )); - } - let elements: u64 = dims.iter().product(); - if elements == 0 || !elements.is_multiple_of(dtype.block_size()) { - return Err(bad( - format!("{elements} elements is not whole {dtype:?} blocks"), - &self.path, - )); - } - if tensors.iter().any(|t| t.name == name) { - return Err(bad(format!("duplicate tensor name '{name}'"), &self.path)); - } - tensors.push(GgufTensorInfo { - name, - dims, - dtype, - offset, - }); - } - Ok(tensors) - } -} - -// ---- Dequantization ------------------------------------------------------ -// -// Exact ports of ggml's reference dequantizers (ggml-quants.c) for every -// dtype llama.cpp stores model weights in: plain floats, the 32-element -// legacy blocks Q4_0/Q4_1/Q5_0/Q5_1/Q8_0, and the 256-element K-quant -// superblocks Q2_K–Q6_K. Layouts follow `GgmlType::bytes_per_block`. -// Q8_1/Q8_K are activation formats (dot-product scratch), never tensor -// storage — they stay loud errors. - -/// Dequantize a whole tensor payload to f32. `bytes` must be whole blocks of -/// `dtype` (guaranteed for payload slices sized by [`GgufTensorInfo::byte_len`]). -pub fn dequantize(dtype: GgmlType, bytes: &[u8]) -> Result, String> { - let bpb = usize::try_from(dtype.bytes_per_block()).map_err(|_| { - format!( - "{dtype:?} block is {} bytes, wider than usize", - dtype.bytes_per_block() - ) - })?; - if bytes.is_empty() || !bytes.len().is_multiple_of(bpb) { - return Err(format!( - "{} bytes is not whole {dtype:?} blocks of {bpb}", - bytes.len() - )); - } - let blocks = bytes.len() / bpb; - let block_elems = usize::try_from(dtype.block_size()).map_err(|_| { - format!( - "{dtype:?} block holds {} elements, more than usize", - dtype.block_size() - ) - })?; - let mut out = Vec::with_capacity(blocks * block_elems); - for block in bytes.chunks_exact(bpb) { - match dtype { - GgmlType::F32 => { - out.push(f32::from_le_bytes(block.try_into().map_err(|_| { - format!("F32 block is {} bytes, not 4", block.len()) - })?)) - } - GgmlType::F16 => out.push(f16_to_f32(u16::from_le_bytes([block[0], block[1]]))), - GgmlType::BF16 => { - out.push(f32::from_bits( - u32::from(u16::from_le_bytes([block[0], block[1]])) << 16, - )); - } - GgmlType::Q4_0 => dequant_q4_0(block, &mut out), - GgmlType::Q4_1 => dequant_q4_1(block, &mut out), - GgmlType::Q5_0 => dequant_q5_0(block, &mut out), - GgmlType::Q5_1 => dequant_q5_1(block, &mut out), - GgmlType::Q8_0 => dequant_q8_0(block, &mut out), - GgmlType::Q2_K => dequant_q2_k(block, &mut out), - GgmlType::Q3_K => dequant_q3_k(block, &mut out), - GgmlType::Q4_K => dequant_q4_k(block, &mut out), - GgmlType::Q5_K => dequant_q5_k(block, &mut out), - GgmlType::Q6_K => dequant_q6_k(block, &mut out), - GgmlType::IQ4_XS => dequant_iq4_xs(block, &mut out), - GgmlType::IQ4_NL => dequant_iq4_nl(block, &mut out), - GgmlType::IQ2_XS => dequant_iq2_xs(block, &mut out), - GgmlType::IQ2_S => dequant_iq2_s(block, &mut out), - GgmlType::IQ3_XXS => dequant_iq3_xxs(block, &mut out), - GgmlType::IQ3_S => dequant_iq3_s(block, &mut out), - other => return Err(format!("dequant for {other:?} is not implemented yet")), - } - } - assert_eq!(out.len(), blocks * block_elems, "whole blocks out"); - Ok(out) -} - -/// IEEE 754 half → f32 (no `half` dep on this path; exhaustive over u16 in -/// tests against the `half` crate the workspace already carries). -fn f16_to_f32(bits: u16) -> f32 { - f32::from(half::f16::from_bits(bits)) -} - -/// Q8_0: f16 scale + 32 signed bytes; `x = d * q`. -fn dequant_q8_0(block: &[u8], out: &mut Vec) { - assert_eq!(block.len(), 34, "Q8_0 block is 34 bytes"); - let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]])); - #[allow(clippy::cast_possible_wrap)] // bit-exact reinterpret is the format - out.extend(block[2..34].iter().map(|&q| d * f32::from(q as i8))); -} - -/// Q4_0: f16 scale + 16 bytes of 4-bit quants; `x = d·(q − 8)` — all 16 low -/// nibbles are elements 0..16, the high nibbles elements 16..32. -fn dequant_q4_0(block: &[u8], out: &mut Vec) { - assert_eq!(block.len(), 18, "Q4_0 block is 18 bytes"); - let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]])); - let qs = &block[2..18]; - out.extend(qs.iter().map(|&b| d * (f32::from(b & 0x0F) - 8.0))); - out.extend(qs.iter().map(|&b| d * (f32::from(b >> 4) - 8.0))); -} - -/// Q4_1: f16 scale + f16 min + 16 bytes of 4-bit quants; `x = d·q + m`. -fn dequant_q4_1(block: &[u8], out: &mut Vec) { - assert_eq!(block.len(), 20, "Q4_1 block is 20 bytes"); - let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]])); - let m = f16_to_f32(u16::from_le_bytes([block[2], block[3]])); - let qs = &block[4..20]; - out.extend(qs.iter().map(|&b| d * f32::from(b & 0x0F) + m)); - out.extend(qs.iter().map(|&b| d * f32::from(b >> 4) + m)); -} - -/// Q5_0: f16 scale + 4 B of packed 5th bits + 16 B of 4-bit quants; -/// `x = d·(q − 16)` with bit `j` of `qh` topping up element `j`. -fn dequant_q5_0(block: &[u8], out: &mut Vec) { - assert_eq!(block.len(), 22, "Q5_0 block is 22 bytes"); - let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]])); - let qh = u32::from_le_bytes([block[2], block[3], block[4], block[5]]); - let qs = &block[6..22]; - #[allow(clippy::cast_possible_truncation)] // masked to one nibble bit - out.extend(qs.iter().enumerate().map(|(j, &b)| { - let hi = ((qh >> j) << 4) as u8 & 0x10; - d * (f32::from((b & 0x0F) | hi) - 16.0) - })); - #[allow(clippy::cast_possible_truncation)] // masked to one nibble bit - out.extend(qs.iter().enumerate().map(|(j, &b)| { - let hi = (qh >> (j + 12)) as u8 & 0x10; - d * (f32::from((b >> 4) | hi) - 16.0) - })); -} - -/// Q5_1: f16 scale + f16 min + 4 B packed 5th bits + 16 B quants; `x = d·q + m`. -fn dequant_q5_1(block: &[u8], out: &mut Vec) { - assert_eq!(block.len(), 24, "Q5_1 block is 24 bytes"); - let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]])); - let m = f16_to_f32(u16::from_le_bytes([block[2], block[3]])); - let qh = u32::from_le_bytes([block[4], block[5], block[6], block[7]]); - let qs = &block[8..24]; - #[allow(clippy::cast_possible_truncation)] // masked to one nibble bit - out.extend(qs.iter().enumerate().map(|(j, &b)| { - let hi = ((qh >> j) << 4) as u8 & 0x10; - d * f32::from((b & 0x0F) | hi) + m - })); - #[allow(clippy::cast_possible_truncation)] // masked to one nibble bit - out.extend(qs.iter().enumerate().map(|(j, &b)| { - let hi = (qh >> (j + 12)) as u8 & 0x10; - d * f32::from((b >> 4) | hi) + m - })); -} - -/// Q2_K: 256-element superblock — 16 packed (scale, min) nibbles + 64 B of -/// 2-bit quants + f16 d + f16 dmin; `x = d·sc·q − dmin·m` over 16 sub-blocks -/// of 16. -fn dequant_q2_k(block: &[u8], out: &mut Vec) { - assert_eq!(block.len(), 84, "Q2_K superblock is 84 bytes"); - let scales = &block[0..16]; - let qs = &block[16..80]; - let d = f16_to_f32(u16::from_le_bytes([block[80], block[81]])); - let dmin = f16_to_f32(u16::from_le_bytes([block[82], block[83]])); - let mut is = 0; - // Two halves of 128 values; each half reads 32 quant bytes at 4 shifts. - for q in [&qs[0..32], &qs[32..64]] { - for shift in [0u8, 2, 4, 6] { - for part in [&q[0..16], &q[16..32]] { - let sc = scales[is]; - is += 1; - let dl = d * f32::from(sc & 0x0F); - let ml = dmin * f32::from(sc >> 4); - out.extend(part.iter().map(|&b| dl * f32::from((b >> shift) & 3) - ml)); - } - } - } - assert_eq!(is, 16, "16 sub-block scales consumed"); -} - -/// Unpack Q3_K's 12 packed scale bytes into 16 signed 6-bit sub-scales -/// (ggml's kmask bit dance, done bytewise). -fn q3_k_scales(packed: &[u8]) -> [i8; 16] { - assert_eq!(packed.len(), 12, "Q3_K scale block is 12 bytes"); - let mut sc = [0i8; 16]; - #[allow(clippy::cast_possible_wrap)] // 6-bit values reinterpret exactly - for j in 0..4 { - let hi = packed[8 + j]; // 2-bit tops for slots j, j+4, j+8, j+12 - sc[j] = ((packed[j] & 0x0F) | ((hi & 3) << 4)) as i8; - sc[j + 4] = ((packed[j + 4] & 0x0F) | (((hi >> 2) & 3) << 4)) as i8; - sc[j + 8] = ((packed[j] >> 4) | (((hi >> 4) & 3) << 4)) as i8; - sc[j + 12] = ((packed[j + 4] >> 4) | ((hi >> 6) << 4)) as i8; - } - sc -} - -/// Q3_K: 256-element superblock — 32 B high-bit mask + 64 B of 2-bit quants -/// + 12 B packed 6-bit sub-scales + f16 d; `x = d·(sc − 32)·(q − hm·4)`. -fn dequant_q3_k(block: &[u8], out: &mut Vec) { - assert_eq!(block.len(), 110, "Q3_K superblock is 110 bytes"); - let hmask = &block[0..32]; - let qs = &block[32..96]; - let scales = q3_k_scales(&block[96..108]); - let d = f16_to_f32(u16::from_le_bytes([block[108], block[109]])); - let mut is = 0; - let mut m: u8 = 1; - for q in [&qs[0..32], &qs[32..64]] { - for shift in [0u8, 2, 4, 6] { - for base in [0usize, 16] { - let dl = d * f32::from(i16::from(scales[is]) - 32); - is += 1; - for l in base..base + 16 { - let low = i16::from((q[l] >> shift) & 3); - let sub = if hmask[l] & m == 0 { 4 } else { 0 }; - out.push(dl * f32::from(low - sub)); - } - } - m <<= 1; // the hmask bit advances per (half, shift) pair - } - } - assert_eq!(is, 16, "16 sub-block scales consumed"); -} - -/// The Q4_K/Q5_K 6-bit (scale, min) pair for sub-block `j` — ggml's -/// `get_scale_min_k4`. -fn scale_min_k4(scales: &[u8], j: usize) -> (f32, f32) { - assert_eq!(scales.len(), 12, "K-quant scale block is 12 bytes"); - assert!(j < 8, "8 sub-blocks per superblock"); - let (sc, m) = if j < 4 { - (scales[j] & 63, scales[j + 4] & 63) - } else { - ( - (scales[j + 4] & 0x0F) | ((scales[j - 4] >> 6) << 4), - (scales[j + 4] >> 4) | ((scales[j] >> 6) << 4), - ) - }; - (f32::from(sc), f32::from(m)) -} - -/// Q4_K: 256-element superblock — f16 d + f16 dmin + 12 B packed 6-bit -/// (scale, min) pairs + 128 B of 4-bit quants; `x = d·sc·q − dmin·m`. -fn dequant_q4_k(block: &[u8], out: &mut Vec) { - assert_eq!(block.len(), 144, "Q4_K superblock is 144 bytes"); - let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]])); - let dmin = f16_to_f32(u16::from_le_bytes([block[2], block[3]])); - let scales = &block[4..16]; - let qs = &block[16..144]; - // 4 chunks of 64 values; each chunk reads 32 bytes — low nibbles first. - for chunk in 0..4 { - let (sc1, m1) = scale_min_k4(scales, chunk * 2); - let (sc2, m2) = scale_min_k4(scales, chunk * 2 + 1); - let q = &qs[chunk * 32..chunk * 32 + 32]; - out.extend(q.iter().map(|&b| d * sc1 * f32::from(b & 0x0F) - dmin * m1)); - out.extend(q.iter().map(|&b| d * sc2 * f32::from(b >> 4) - dmin * m2)); - } -} - -/// Q5_K: 256-element superblock — f16 d + f16 dmin + 12 B packed 6-bit -/// (scale, min) pairs + 32 B of 5th bits + 128 B of 4-bit quants; -/// `x = d·sc·q − dmin·m` with two `qh` bits per byte per chunk. -fn dequant_q5_k(block: &[u8], out: &mut Vec) { - assert_eq!(block.len(), 176, "Q5_K superblock is 176 bytes"); - let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]])); - let dmin = f16_to_f32(u16::from_le_bytes([block[2], block[3]])); - let scales = &block[4..16]; - let qh = &block[16..48]; - let qs = &block[48..176]; - let mut u1: u8 = 1; - let mut u2: u8 = 2; - // 4 chunks of 64 values; each chunk reads 32 quant bytes — low nibbles - // first — and one bit-plane pair of the shared 32-byte qh. - for chunk in 0..4 { - let (sc1, m1) = scale_min_k4(scales, chunk * 2); - let (sc2, m2) = scale_min_k4(scales, chunk * 2 + 1); - let q = &qs[chunk * 32..chunk * 32 + 32]; - out.extend(q.iter().zip(qh).map(|(&b, &h)| { - let top = if h & u1 == 0 { 0 } else { 16 }; - d * sc1 * f32::from((b & 0x0F) + top) - dmin * m1 - })); - out.extend(q.iter().zip(qh).map(|(&b, &h)| { - let top = if h & u2 == 0 { 0 } else { 16 }; - d * sc2 * f32::from((b >> 4) + top) - dmin * m2 - })); - u1 <<= 2; - u2 <<= 2; - } - assert_eq!(u1, 0, "four bit-plane pairs consumed"); // 1<<8 wraps to 0 -} - -/// Q6_K: 256-element superblock — 128 B low-4 + 64 B high-2 + 16 i8 -/// sub-scales + f16 d; `x = d·sc·(q − 32)`. -fn dequant_q6_k(block: &[u8], out: &mut Vec) { - assert_eq!(block.len(), 210, "Q6_K superblock is 210 bytes"); - let (ql_all, rest) = block.split_at(128); - let (qh_all, rest) = rest.split_at(64); - let (scales, d_bytes) = rest.split_at(16); - let d = f16_to_f32(u16::from_le_bytes([d_bytes[0], d_bytes[1]])); - let start = out.len(); - out.resize(start + 256, 0.0); - let y = &mut out[start..]; - // Two halves of 128 values, each consuming 64 ql / 32 qh / 8 scales. - for half_idx in 0..2 { - let ql = &ql_all[half_idx * 64..half_idx * 64 + 64]; - let qh = &qh_all[half_idx * 32..half_idx * 32 + 32]; - let sc = &scales[half_idx * 8..half_idx * 8 + 8]; - let base = half_idx * 128; - for l in 0..32 { - let is = l / 16; - let q1 = i16::from((ql[l] & 0x0F) | ((qh[l] & 3) << 4)) - 32; - let q2 = i16::from((ql[l + 32] & 0x0F) | (((qh[l] >> 2) & 3) << 4)) - 32; - let q3 = i16::from((ql[l] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32; - let q4 = i16::from((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32; - #[allow(clippy::cast_possible_wrap)] // i8 sub-scales are the format - let s = |i: usize| f32::from(sc[i] as i8); - y[base + l] = d * s(is) * f32::from(q1); - y[base + l + 32] = d * s(is + 2) * f32::from(q2); - y[base + l + 64] = d * s(is + 4) * f32::from(q3); - y[base + l + 96] = d * s(is + 6) * f32::from(q4); - } - } -} - -/// IQ4_NL/IQ4_XS's 16-entry non-linear value table (llama.cpp -/// `kvalues_iq4nl`, frozen with the format). -const KVALUES_IQ4NL: [f32; 16] = [ - -127.0, -104.0, -83.0, -65.0, -49.0, -35.0, -22.0, -10.0, 1.0, 13.0, 25.0, 38.0, 53.0, 69.0, - 89.0, 113.0, -]; - -/// IQ4_XS: 256-element superblock — f16 d + 8 six-bit sub-scales (low 4 bits -/// packed two-per-byte in `scales_l`, high 2 bits packed in `scales_h`) + -/// 128 B of 4-bit indices into [`KVALUES_IQ4NL`]; per 32-group, -/// `x = d·(sc − 32)·kvalues[q]` with the 16 low nibbles first -/// (llama.cpp `dequantize_row_iq4_xs`). -fn dequant_iq4_xs(block: &[u8], out: &mut Vec) { - assert_eq!(block.len(), 136, "IQ4_XS superblock is 136 bytes"); - let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]])); - let scales_h = u16::from_le_bytes([block[2], block[3]]); - let scales_l = &block[4..8]; - let qs = &block[8..136]; - for ib32 in 0..8usize { - let lo = (scales_l[ib32 / 2] >> (4 * (ib32 % 2))) & 0x0F; - let hi = ((scales_h >> (2 * ib32)) & 3) as u8; - let sc = i16::from(lo | (hi << 4)) - 32; - let dl = d * f32::from(sc); - let q = &qs[16 * ib32..16 * ib32 + 16]; - out.extend(q.iter().map(|&b| dl * KVALUES_IQ4NL[usize::from(b & 0x0F)])); - out.extend(q.iter().map(|&b| dl * KVALUES_IQ4NL[usize::from(b >> 4)])); - } -} - -/// IQ4_NL: 32-element block — f16 d + 16 B of 4-bit indices into -/// [`KVALUES_IQ4NL`]; low nibbles are elements 0..16, high 16..32 -/// (llama.cpp `dequantize_row_iq4_nl`). -fn dequant_iq4_nl(block: &[u8], out: &mut Vec) { - assert_eq!(block.len(), 18, "IQ4_NL block is 18 bytes"); - let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]])); - let qs = &block[2..18]; - out.extend(qs.iter().map(|&b| d * KVALUES_IQ4NL[usize::from(b & 0x0F)])); - out.extend(qs.iter().map(|&b| d * KVALUES_IQ4NL[usize::from(b >> 4)])); -} - -use crate::gguf_iq_grids::{IQ2S_GRID, IQ2XS_GRID, IQ3S_GRID, IQ3XXS_GRID, KSIGNS_IQ2XS}; - -/// Unpack a grid word's 8 byte-magnitudes with `signs`' 8 sign bits applied -/// (ggml's `kmask_iq2xs` is just bit `j`), scaled by `dl`. -fn push_signed_row8(out: &mut Vec, dl: f32, grid: u64, signs: u8) { - for j in 0..8 { - let mag = f32::from((grid >> (8 * j)) as u8); - let sign = if signs & (1 << j) != 0 { -1.0 } else { 1.0 }; - out.push(dl * mag * sign); - } -} - -/// IQ2_XS: 256-element superblock — f16 d + 32 u16 words (9-bit grid index + -/// 7-bit sign index) + 8 packed 4-bit sub-scales; -/// `x = d·(0.5 + sc)·0.25·grid·±1` (llama.cpp `dequantize_row_iq2_xs`). -fn dequant_iq2_xs(block: &[u8], out: &mut Vec) { - assert_eq!(block.len(), 74, "IQ2_XS superblock is 74 bytes"); - let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]])); - let qs: Vec = block[2..66] - .as_chunks::<2>() - .0 - .iter() - .map(|c| u16::from_le_bytes(*c)) - .collect(); - let scales = &block[66..74]; - for ib32 in 0..8usize { - let db = [ - d * (0.5 + f32::from(scales[ib32] & 0x0F)) * 0.25, - d * (0.5 + f32::from(scales[ib32] >> 4)) * 0.25, - ]; - for l in 0..4usize { - let word = qs[4 * ib32 + l]; - let grid = IQ2XS_GRID[usize::from(word & 511)]; - let signs = KSIGNS_IQ2XS[usize::from(word >> 9)]; - push_signed_row8(out, db[l / 2], grid, signs); - } - } -} - -/// IQ2_S: 256-element superblock — f16 d + 64 B `qs` (32 low index bytes, -/// then 32 sign bytes) + 8 B `qh` (index bits 8..10) + 8 packed sub-scales -/// (llama.cpp `dequantize_row_iq2_s`). -fn dequant_iq2_s(block: &[u8], out: &mut Vec) { - assert_eq!(block.len(), 82, "IQ2_S superblock is 82 bytes"); - let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]])); - let qs = &block[2..34]; // low 8 bits of grid indices, 4 per 32-group - let signs = &block[34..66]; // one sign byte per grid row - let qh = &block[66..74]; - let scales = &block[74..82]; - for ib32 in 0..8usize { - let db = [ - d * (0.5 + f32::from(scales[ib32] & 0x0F)) * 0.25, - d * (0.5 + f32::from(scales[ib32] >> 4)) * 0.25, - ]; - for l in 0..4usize { - let hi = (u16::from(qh[ib32]) << (8 - 2 * l)) & 0x300; - let grid = IQ2S_GRID[usize::from(u16::from(qs[4 * ib32 + l]) | hi)]; - push_signed_row8(out, db[l / 2], grid, signs[4 * ib32 + l]); - } - } -} - -/// Unpack a u32 grid word's 4 byte-magnitudes with 4 sign bits (offset -/// `bit0` into ggml's 8-bit sign mask), scaled by `dl`. -fn push_signed_row4(out: &mut Vec, dl: f32, grid: u32, signs: u8, bit0: u8) { - for j in 0..4u8 { - let mag = f32::from((grid >> (8 * j)) as u8); - let sign = if signs & (1 << (bit0 + j)) != 0 { -1.0 } else { 1.0 }; - out.push(dl * mag * sign); - } -} - -/// IQ3_XXS: 256-element superblock — f16 d + 64 index bytes (into the -/// 256-entry grid, 4 values each) + 8 u32 words carrying a 4-bit scale and -/// four 7-bit sign indices (llama.cpp `dequantize_row_iq3_xxs`). -fn dequant_iq3_xxs(block: &[u8], out: &mut Vec) { - assert_eq!(block.len(), 98, "IQ3_XXS superblock is 98 bytes"); - let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]])); - let qs = &block[2..66]; - let sas = &block[66..98]; // scales-and-signs, one u32 per 32-group - for ib32 in 0..8usize { - let aux32 = u32::from_le_bytes(sas[4 * ib32..4 * ib32 + 4].try_into().expect("4")); - let db = d * (0.5 + (aux32 >> 28) as f32) * 0.5; - for l in 0..4usize { - let signs = KSIGNS_IQ2XS[usize::try_from((aux32 >> (7 * l)) & 127).expect("7 bits")]; - let g1 = IQ3XXS_GRID[usize::from(qs[8 * ib32 + 2 * l])]; - let g2 = IQ3XXS_GRID[usize::from(qs[8 * ib32 + 2 * l + 1])]; - push_signed_row4(out, db, g1, signs, 0); - push_signed_row4(out, db, g2, signs, 4); - } - } -} - -/// IQ3_S: 256-element superblock — f16 d + 64 index bytes + 8 B `qh` (index -/// bit 8) + 32 sign bytes + 4 packed 4-bit sub-scales; -/// `x = d·(1 + 2·sc)·grid·±1` (llama.cpp `dequantize_row_iq3_s`). -fn dequant_iq3_s(block: &[u8], out: &mut Vec) { - assert_eq!(block.len(), 110, "IQ3_S superblock is 110 bytes"); - let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]])); - let qs = &block[2..66]; - let qh = &block[66..74]; - let signs = &block[74..106]; - let scales = &block[106..110]; - for ib32 in 0..8usize { - let db = d - * (1.0 - + 2.0 * f32::from((scales[ib32 / 2] >> (4 * (ib32 % 2))) & 0x0F)); - for l in 0..4usize { - let h = u16::from(qh[ib32]); - let i1 = usize::from(u16::from(qs[8 * ib32 + 2 * l]) | ((h << (8 - 2 * l)) & 256)); - let i2 = usize::from(u16::from(qs[8 * ib32 + 2 * l + 1]) | ((h << (7 - 2 * l)) & 256)); - let sign_byte = signs[4 * ib32 + l]; - push_signed_row4(out, db, IQ3S_GRID[i1], sign_byte, 0); - push_signed_row4(out, db, IQ3S_GRID[i2], sign_byte, 4); - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use std::io::Write; - - /// Minimal in-memory GGUF builder for tests. - struct TestGguf { - buf: Vec, - tensor_count: u64, - kv_count: u64, - kvs: Vec, - tensors: Vec, - } - - impl TestGguf { - fn new() -> Self { - Self { - buf: Vec::new(), - tensor_count: 0, - kv_count: 0, - kvs: Vec::new(), - tensors: Vec::new(), - } - } - - fn push_str(out: &mut Vec, s: &str) { - out.extend_from_slice(&(s.len() as u64).to_le_bytes()); - out.extend_from_slice(s.as_bytes()); - } - - fn kv_str(mut self, key: &str, value: &str) -> Self { - Self::push_str(&mut self.kvs, key); - self.kvs.extend_from_slice(&8u32.to_le_bytes()); - Self::push_str(&mut self.kvs, value); - self.kv_count += 1; - self - } - - fn kv_u32(mut self, key: &str, value: u32) -> Self { - Self::push_str(&mut self.kvs, key); - self.kvs.extend_from_slice(&4u32.to_le_bytes()); - self.kvs.extend_from_slice(&value.to_le_bytes()); - self.kv_count += 1; - self - } - - fn kv_str_array(mut self, key: &str, values: &[&str]) -> Self { - Self::push_str(&mut self.kvs, key); - self.kvs.extend_from_slice(&9u32.to_le_bytes()); - self.kvs.extend_from_slice(&8u32.to_le_bytes()); - self.kvs - .extend_from_slice(&(values.len() as u64).to_le_bytes()); - for v in values { - Self::push_str(&mut self.kvs, v); - } - self.kv_count += 1; - self - } - - fn tensor(mut self, name: &str, dims: &[u64], type_id: u32, offset: u64) -> Self { - Self::push_str(&mut self.tensors, name); - self.tensors - .extend_from_slice(&(dims.len() as u32).to_le_bytes()); - for d in dims { - self.tensors.extend_from_slice(&d.to_le_bytes()); - } - self.tensors.extend_from_slice(&type_id.to_le_bytes()); - self.tensors.extend_from_slice(&offset.to_le_bytes()); - self.tensor_count += 1; - self - } - - fn build(mut self) -> Vec { - self.buf.extend_from_slice(&MAGIC); - self.buf.extend_from_slice(&3u32.to_le_bytes()); - self.buf.extend_from_slice(&self.tensor_count.to_le_bytes()); - self.buf.extend_from_slice(&self.kv_count.to_le_bytes()); - self.buf.extend_from_slice(&self.kvs); - self.buf.extend_from_slice(&self.tensors); - self.buf - } - - /// Header + alignment padding + tensor payload bytes (offsets in the - /// tensor table are relative to the padded data start). - fn build_with_payload(self, payload: &[u8]) -> Vec { - let mut buf = self.build(); - let data_offset = (buf.len() as u64).div_ceil(DEFAULT_ALIGNMENT) * DEFAULT_ALIGNMENT; - buf.resize(usize::try_from(data_offset).expect("small test file"), 0); - buf.extend_from_slice(payload); - buf - } - } - - /// Write `bytes` to a fresh temp file and run `f` on the parse result - /// while the file still exists (payload reads re-open the path). - fn with_gguf_bytes(bytes: &[u8], f: impl FnOnce(Result) -> R) -> R { - use std::sync::atomic::{AtomicU64, Ordering}; - // Parallel tests in one process must never share a temp file. - static NEXT: AtomicU64 = AtomicU64::new(0); - let dir = std::env::temp_dir().join("mummu-gguf-tests"); - std::fs::create_dir_all(&dir).expect("temp dir"); - let path = dir.join(format!( - "t-{}-{}.gguf", - std::process::id(), - NEXT.fetch_add(1, Ordering::Relaxed) - )); - let mut file = File::create(&path).expect("temp file"); - file.write_all(bytes).expect("write"); - drop(file); - let result = f(GgufFile::open(&path)); - let _ = std::fs::remove_file(&path); - result - } - - /// A unique scratch path for tests that write their own output file. - fn scratch_path(tag: &str) -> std::path::PathBuf { - use std::sync::atomic::{AtomicU64, Ordering}; - static NEXT: AtomicU64 = AtomicU64::new(0); - let dir = std::env::temp_dir().join("mummu-gguf-tests"); - std::fs::create_dir_all(&dir).expect("temp dir"); - dir.join(format!( - "{tag}-{}-{}.safetensors", - std::process::id(), - NEXT.fetch_add(1, Ordering::Relaxed) - )) - } - - fn open_bytes(bytes: &[u8]) -> Result { - with_gguf_bytes(bytes, |r| r) - } - - #[test] - fn minimal_file_round_trips() { - let bytes = TestGguf::new() - .kv_str("general.architecture", "qwen2") - .kv_u32("qwen2.block_count", 28) - .kv_str_array("tokenizer.ggml.tokens", &["a", "b", "c"]) - .tensor("token_embd.weight", &[64, 2], 0, 0) - .tensor("blk.0.attn_q.weight", &[256], 12, 512) - .build(); - let f = open_bytes(&bytes).expect("parses"); - assert_eq!(f.version, 3); - assert_eq!(f.architecture(), Some("qwen2")); - assert_eq!( - f.get("qwen2.block_count").and_then(GgufValue::as_u64), - Some(28) - ); - assert_eq!( - f.get("tokenizer.ggml.tokens") - .and_then(GgufValue::as_array) - .map(<[GgufValue]>::len), - Some(3) - ); - let embd = f.tensor("token_embd.weight").expect("present"); - assert_eq!(embd.dims, vec![64, 2]); - assert_eq!(embd.dtype, GgmlType::F32); - assert_eq!(embd.byte_len(), 64 * 2 * 4); - let q = f.tensor("blk.0.attn_q.weight").expect("present"); - assert_eq!(q.dtype, GgmlType::Q4_K); - assert_eq!(q.byte_len(), 144); // one 256-element Q4_K superblock - assert_eq!(f.alignment, DEFAULT_ALIGNMENT); - assert!(f.data_offset.is_multiple_of(f.alignment)); - assert!(f.data_offset >= (bytes.len() as u64)); - } - - #[test] - fn bad_magic_is_rejected() { - let mut bytes = TestGguf::new().build(); - bytes[..4].copy_from_slice(b"FUGG"); - assert!(matches!( - open_bytes(&bytes), - Err(GgufError::BadMagic { .. }) - )); - } - - #[test] - fn unsupported_version_is_rejected() { - let mut bytes = TestGguf::new().build(); - bytes[4..8].copy_from_slice(&1u32.to_le_bytes()); - assert!(matches!( - open_bytes(&bytes), - Err(GgufError::UnsupportedVersion { version: 1, .. }) - )); - } - - #[test] - fn truncated_file_is_an_io_error_not_a_hang() { - let bytes = TestGguf::new() - .kv_str("general.architecture", "qwen2") - .build(); - assert!(matches!( - open_bytes(&bytes[..bytes.len() - 3]), - Err(GgufError::Io { .. }) - )); - } - - #[test] - fn oversized_counts_are_rejected() { - let mut bytes = TestGguf::new().build(); - // tensor_count lives at bytes 8..16. - bytes[8..16].copy_from_slice(&(MAX_TENSORS + 1).to_le_bytes()); - assert!(matches!( - open_bytes(&bytes), - Err(GgufError::OverBound { what: "tensor", .. }) - )); - } - - #[test] - fn unknown_value_type_and_ggml_type_are_rejected() { - let mut with_kv = TestGguf::new().kv_u32("some.key", 1).build(); - // The kv's type id (4 = u32) sits right after the 8-byte key string - // prefix + 8 bytes of key: magic(4)+ver(4)+counts(16)+len(8)+key(8). - with_kv[40..44].copy_from_slice(&99u32.to_le_bytes()); - assert!(matches!( - open_bytes(&with_kv), - Err(GgufError::BadValue { .. }) - )); - - let with_tensor = TestGguf::new().tensor("t", &[32], 63, 0).build(); - assert!(matches!( - open_bytes(&with_tensor), - Err(GgufError::BadTensor { .. }) - )); - } - - // ---- dequant --------------------------------------------------------- - - fn f16_bytes(v: f32) -> [u8; 2] { - half::f16::from_f32(v).to_bits().to_le_bytes() - } - - #[test] - fn float_widths_dequantize_exactly() { - let f32_bytes = 1.5f32.to_le_bytes(); - assert_eq!(dequantize(GgmlType::F32, &f32_bytes).unwrap(), vec![1.5]); - assert_eq!( - dequantize(GgmlType::F16, &f16_bytes(-0.25)).unwrap(), - vec![-0.25] - ); - // bf16 is the top half of the f32 bit pattern. - let bf16 = (2.0f32.to_bits() >> 16) as u16; - assert_eq!( - dequantize(GgmlType::BF16, &bf16.to_le_bytes()).unwrap(), - vec![2.0] - ); - } - - #[test] - fn q8_0_block_matches_hand_computation() { - let mut block = Vec::new(); - block.extend_from_slice(&f16_bytes(0.5)); - #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] - block.extend((0..32).map(|i| (i - 16) as i8 as u8)); - let out = dequantize(GgmlType::Q8_0, &block).unwrap(); - assert_eq!(out.len(), 32); - assert_eq!(out[0], 0.5 * -16.0); - assert_eq!(out[16], 0.0); - assert_eq!(out[31], 0.5 * 15.0); - } - - #[test] - fn q4_k_superblock_matches_hand_computation() { - let mut block = vec![0u8; 144]; - block[0..2].copy_from_slice(&f16_bytes(1.0)); // d - block[2..4].copy_from_slice(&f16_bytes(0.5)); // dmin - // Sub-block 0: sc=2, m=1 · sub-block 1: sc=3, m=0 (direct 6-bit slots). - block[4] = 2; - block[5] = 3; - block[8] = 1; - // Sub-block 4: packed slot — sc = scales[8] & 0xF = 1, m = scales[8] >> 4 = 2. - block[12] = 0x21; - // First quant byte of chunk 0: low nibble 1 (sub 0), high nibble 5 (sub 1). - block[16] = 0x51; - // First quant byte of chunk 2 (sub-blocks 4/5): low nibble 4. - block[16 + 64] = 0x04; - let out = dequantize(GgmlType::Q4_K, &block).unwrap(); - assert_eq!(out.len(), 256); - assert_eq!(out[0], 2.0 * 1.0 - 0.5 * 1.0); // d·sc0·q − dmin·m0 = 1.5 - assert_eq!(out[32], 3.0 * 5.0); // sub 1: m=0 - assert_eq!(out[128], 1.0 * 4.0 - 0.5 * 2.0); // sub 4 via packed scales - // A zero quant in sub-block 0 still subtracts the min. - assert_eq!(out[1], -0.5); - } - - #[test] - fn q6_k_superblock_matches_hand_computation() { - let mut block = vec![0u8; 210]; - block[0] = 0x0F; // ql[0]: low 4 bits = 15 - block[128] = 0b0000_0011; // qh[0]: high 2 bits = 3 for q1 - block[192] = 2; // scales[0] = 2 - block[194] = 1; // scales[2] = 1 - block[208..210].copy_from_slice(&f16_bytes(1.0)); // d - let out = dequantize(GgmlType::Q6_K, &block).unwrap(); - assert_eq!(out.len(), 256); - // q1 = (15 | 3<<4) − 32 = 31, scale 2 → 62. - assert_eq!(out[0], 62.0); - // q2 = (0 | 0) − 32 = −32, scale sc[2]=1 → −32. - assert_eq!(out[32], -32.0); - // Zero scale zeroes the value even though q3 = −32. - assert_eq!(out[64], 0.0); - } - - #[test] - fn iq4_xs_superblock_matches_hand_computation() { - let mut block = vec![0u8; 136]; - block[0..2].copy_from_slice(&f16_bytes(2.0)); // d - // Sub-scale 0 = 33 (low 4 bits = 1, high 2 bits = 2 → 1|2<<4 = 33): - // scales_l[0] low nibble = 1; scales_h bits 0..2 = 2. - block[4] = 0x01; - block[2..4].copy_from_slice(&2u16.to_le_bytes()); - // Sub-scale 1 = 0 → dl = 2·(0−32) = −64 for group 1. - // qs[0]: low nibble index 8 (→ +1), high nibble index 15 (→ +113). - block[8] = 0xF8; // low nibble = 8, high nibble = 15 - let out = dequantize(GgmlType::IQ4_XS, &block).unwrap(); - assert_eq!(out.len(), 256); - // Group 0: dl = 2·(33−32) = 2. Element 0 = 2·kvalues[8] = 2·1. - assert_eq!(out[0], 2.0); - // High nibbles are elements 16..32 of the group: 2·kvalues[15]. - assert_eq!(out[16], 2.0 * 113.0); - // Untouched qs bytes are index 0 → kvalues[0] = −127, scaled. - assert_eq!(out[1], 2.0 * -127.0); - // Group 1's zero sub-scale gives dl = −64: element 32 = −64·−127. - assert_eq!(out[32], -64.0 * -127.0); - } - - #[test] - fn iq4_nl_block_matches_hand_computation() { - let mut b = vec![0u8; 18]; - b[0..2].copy_from_slice(&f16_bytes(2.0)); - b[2] = 0xF8; // low nibble 8 → kvalues[8] = 1; high 15 → kvalues[15] = 113 - let out = dequantize(GgmlType::IQ4_NL, &b).unwrap(); - assert_eq!(out.len(), 32); - assert_eq!(out[0], 2.0); - assert_eq!(out[16], 2.0 * 113.0); - assert_eq!(out[1], 2.0 * -127.0); // zero nibble hits kvalues[0] - } - - #[test] - fn iq2_xs_superblock_matches_hand_computation() { - let mut b = vec![0u8; 74]; - b[0..2].copy_from_slice(&f16_bytes(1.0)); - // qs word 0: grid index 1 (0x…2b in byte 0), sign index 1 (KSIGNS[1] - // = 129: bits 0 and 7 negative). - b[2..4].copy_from_slice(&0x0201u16.to_le_bytes()); - b[66] = 0x21; // scales[0]: low nibble 1 → db0 = 0.375, high 2 → db1 = 0.625 - let out = dequantize(GgmlType::IQ2_XS, &b).unwrap(); - assert_eq!(out.len(), 256); - assert_eq!(out[0], -0.375 * 43.0); // grid byte 0 = 0x2b, sign bit 0 - assert_eq!(out[1], 0.375 * 8.0); - assert_eq!(out[7], -0.375 * 8.0); // sign bit 7 - // Row l = 2 (elements 16..24) rides db1; untouched words are - // grid 0 (all 8s), sign 0. - assert_eq!(out[16], 0.625 * 8.0); - } - - #[test] - fn iq2_s_superblock_matches_hand_computation() { - let mut b = vec![0u8; 82]; - b[0..2].copy_from_slice(&f16_bytes(2.0)); - b[34] = 129; // signs byte for row 0: bits 0 and 7 - b[74] = 0x01; // scales[0] low nibble 1 → db0 = 2·1.5·0.25 = 0.75 - let out = dequantize(GgmlType::IQ2_S, &b).unwrap(); - assert_eq!(out.len(), 256); - assert_eq!(out[0], -0.75 * 8.0); // grid 0 is all 8s - assert_eq!(out[1], 0.75 * 8.0); - assert_eq!(out[7], -0.75 * 8.0); - assert_eq!(out[8], 0.75 * 8.0); // row 1 still rides db0 - assert_eq!(out[16], 0.5 * 0.25 * 2.0 * 8.0); // rows 2..4: zero high nibble → 0.5·d·0.25 - } - - #[test] - fn iq3_xxs_superblock_matches_hand_computation() { - let mut b = vec![0u8; 98]; - b[0..2].copy_from_slice(&f16_bytes(2.0)); - b[2] = 1; // grid1 index 1 = 0x04040414: byte 0 = 20, rest 4 - // aux32 for group 0: scale bits 1 (db = 2·1.5·0.5 = 1.5), sign index 1. - b[66..70].copy_from_slice(&((1u32 << 28) | 1).to_le_bytes()); - let out = dequantize(GgmlType::IQ3_XXS, &b).unwrap(); - assert_eq!(out.len(), 256); - assert_eq!(out[0], -1.5 * 20.0); // KSIGNS[1] bit 0 - assert_eq!(out[1], 1.5 * 4.0); - assert_eq!(out[4], 1.5 * 4.0); // grid2 (index 0) is all 4s - assert_eq!(out[7], -1.5 * 4.0); // KSIGNS[1] bit 7 lands on grid2's row - } - - #[test] - fn iq3_s_superblock_matches_hand_computation() { - let mut b = vec![0u8; 110]; - b[0..2].copy_from_slice(&f16_bytes(1.0)); - b[2] = 1; // grid1 index 1 = 0x01010103: byte 0 = 3, rest 1 - b[74] = 1; // signs byte row 0: bit 0 - b[106] = 0x01; // scales[0] low nibble 1 → db = 1 + 2·1 = 3 - let out = dequantize(GgmlType::IQ3_S, &b).unwrap(); - assert_eq!(out.len(), 256); - assert_eq!(out[0], -3.0 * 3.0); - assert_eq!(out[1], 3.0); - assert_eq!(out[4], 3.0); // grid2 (index 0) all 1s, no sign - // Group 1 shares scales[0]'s high nibble (0) → db = 1. - assert_eq!(out[32], 1.0); - } - - #[test] - fn q4_0_and_q4_1_blocks_match_hand_computation() { - // Q4_0: symmetric around 8 — low nibbles are elements 0..16. - let mut b = vec![0u8; 18]; - b[0..2].copy_from_slice(&f16_bytes(2.0)); - b[2] = 0x31; // low nibble 1 → elem 0, high nibble 3 → elem 16 - let out = dequantize(GgmlType::Q4_0, &b).unwrap(); - assert_eq!(out.len(), 32); - assert_eq!(out[0], (1.0 - 8.0) * 2.0); - assert_eq!(out[16], (3.0 - 8.0) * 2.0); - assert_eq!(out[1], -16.0); // zero nibble is −8·d, not 0 - - // Q4_1: affine — the min shifts every element. - let mut b = vec![0u8; 20]; - b[0..2].copy_from_slice(&f16_bytes(2.0)); - b[2..4].copy_from_slice(&f16_bytes(1.0)); - b[4] = 0x31; - let out = dequantize(GgmlType::Q4_1, &b).unwrap(); - assert_eq!(out[0], 1.0 * 2.0 + 1.0); - assert_eq!(out[16], 3.0 * 2.0 + 1.0); - assert_eq!(out[1], 1.0); - } - - #[test] - fn q5_0_and_q5_1_high_bits_land_on_the_right_elements() { - // Q5_0: qh bit 0 tops element 0, bit 16 tops element 16. - let mut b = vec![0u8; 22]; - b[0..2].copy_from_slice(&f16_bytes(1.0)); - b[2..6].copy_from_slice(&(1u32 | (1 << 16)).to_le_bytes()); - b[6] = 0x21; // low nibble 1 → elem 0, high nibble 2 → elem 16 - let out = dequantize(GgmlType::Q5_0, &b).unwrap(); - assert_eq!(out.len(), 32); - assert_eq!(out[0], (1.0 + 16.0) - 16.0); - assert_eq!(out[16], (2.0 + 16.0) - 16.0); - assert_eq!(out[1], -16.0); // no high bit, zero nibble - - // Q5_1: same bit packing, affine. - let mut b = vec![0u8; 24]; - b[0..2].copy_from_slice(&f16_bytes(1.0)); - b[2..4].copy_from_slice(&f16_bytes(1.0)); - b[4..8].copy_from_slice(&1u32.to_le_bytes()); - b[8] = 0x01; - let out = dequantize(GgmlType::Q5_1, &b).unwrap(); - assert_eq!(out[0], (1.0 + 16.0) * 1.0 + 1.0); - assert_eq!(out[16], 1.0); // high nibble 0, qh bit 16 unset → just m - } - - #[test] - fn q2_k_superblock_matches_hand_computation() { - let mut b = vec![0u8; 84]; - b[80..82].copy_from_slice(&f16_bytes(1.0)); // d - b[82..84].copy_from_slice(&f16_bytes(0.5)); // dmin - b[0] = 0x12; // sub 0: sc=2, min=1 - b[1] = 0x01; // sub 1: sc=1, min=0 - b[8] = 0x0F; // sub 8 (second half, shift 0): sc=15, min=0 - b[16] = 3; // qs[0] bits 0–1 → elem 0 - b[32] = 1; // qs[16] → elem 16 (sub 1) - b[48] = 2; // qs[32] → elem 128 (second half) - let out = dequantize(GgmlType::Q2_K, &b).unwrap(); - assert_eq!(out.len(), 256); - assert_eq!(out[0], 2.0 * 3.0 - 0.5); // d·sc·q − dmin·m - assert_eq!(out[1], -0.5); // zero quant still subtracts the min - assert_eq!(out[16], 1.0); - assert_eq!(out[128], 15.0 * 2.0); // second half reads qs[32..] - } - - #[test] - fn q3_k_superblock_matches_hand_computation() { - let mut b = vec![0u8; 110]; - b[108..110].copy_from_slice(&f16_bytes(1.0)); // d - // scales: slot 0 = 2|32 = 34 → dl 2; slot 1 = 1|32 = 33 → dl 1; - // slot 8 = (0x32>>4)|0 = 3 → dl 3−32 = −29 (tests the >>4 packing). - b[96] = 0x32; - b[97] = 0x01; - b[104] = 0b10; // top bits of slot 0 - b[105] = 0b10; // top bits of slot 1 - b[0] = 1; // hmask[0] bit 0 → element 0 keeps its high bit (no −4) - b[16] = 1; // hmask[16] bit 0 → element 16 too - b[32] = 3; // qs[0] → elem 0 - b[48] = 2; // qs[16] → elem 16 - let out = dequantize(GgmlType::Q3_K, &b).unwrap(); - assert_eq!(out.len(), 256); - assert_eq!(out[0], 2.0 * 3.0); // high bit set → q unshifted - assert_eq!(out[1], 2.0 * -4.0); // high bit clear → q − 4 - assert_eq!(out[16], 1.0 * 2.0); - // Second half, shift 0 (is=8): hmask[0] bit 4 clear → (0−4)·(3−32). - assert_eq!(out[128], -4.0 * (3.0 - 32.0)); - } - - #[test] - fn q5_k_superblock_matches_hand_computation() { - let mut b = vec![0u8; 176]; - b[0..2].copy_from_slice(&f16_bytes(1.0)); // d - b[2..4].copy_from_slice(&f16_bytes(1.0)); // dmin - b[4] = 2; // sub 0 scale - b[5] = 3; // sub 1 scale - b[8] = 1; // sub 0 min - b[16] = 1; // qh[0] bit 0 → elem 0 gets +16 (u1 = 1) - b[48] = 0x21; // ql[0]: low 1 → elem 0, high 2 → elem 32 - let out = dequantize(GgmlType::Q5_K, &b).unwrap(); - assert_eq!(out.len(), 256); - assert_eq!(out[0], 2.0 * (1.0 + 16.0) - 1.0); - assert_eq!(out[1], -1.0); // zero quant still subtracts the min - assert_eq!(out[32], 3.0 * 2.0); // qh bit 1 unset → no +16; m1 = 0 - } - - #[test] - fn dequant_rejects_partial_blocks_and_unimplemented_types() { - assert!(dequantize(GgmlType::Q8_0, &[0u8; 33]).is_err()); - assert!(dequantize(GgmlType::Q8_0, &[]).is_err()); - // Q8_K is an activation format, never tensor storage. - assert!(dequantize(GgmlType::Q8_K, &[0u8; 292]).is_err()); - } - - #[test] - fn dequant_to_safetensors_reverses_dims_and_round_trips_bytes() { - // One F32 tensor with ggml dims [2, 3] and payload 1..=6. - let payload: Vec = (1..=6).flat_map(|v| (v as f32).to_le_bytes()).collect(); - let bytes = TestGguf::new() - .kv_str("general.architecture", "qwen2") - .tensor("token_embd.weight", &[2, 3], 0, 0) - .build_with_payload(&payload); - with_gguf_bytes(&bytes, |f| { - let f = f.expect("parses"); - let blob = f - .dequant_to_safetensors(&|i| Some(GgufMap::Rename(format!("model.{}", i.name)))) - .expect("serializes"); - let header_len = u64::from_le_bytes(blob[0..8].try_into().unwrap()); - let json: serde_json::Value = - serde_json::from_slice(&blob[8..8 + usize::try_from(header_len).unwrap()]) - .expect("header is valid JSON"); - let entry = &json["model.token_embd.weight"]; - assert_eq!(entry["dtype"], "F32"); - assert_eq!(entry["shape"], serde_json::json!([3, 2])); // reversed - assert_eq!(entry["data_offsets"], serde_json::json!([0, 24])); - // F32 → f32 is byte-identical. - assert_eq!( - &blob[8 + usize::try_from(header_len).unwrap()..], - &payload[..] - ); - - // An unmapped tensor is a loud error, never a silent skip. - assert!(matches!( - f.dequant_to_safetensors(&|_| None), - Err(GgufError::BadTensor { .. }) - )); - }); - } - - #[test] - fn dequant_to_safetensors_rejects_rename_collisions() { - let payload = [0u8; 64]; // two 8-element F32 tensors, offsets 0 and 32 - let bytes = TestGguf::new() - .tensor("a.weight", &[8], 0, 0) - .tensor("b.weight", &[8], 0, 32) - .build_with_payload(&payload); - with_gguf_bytes(&bytes, |f| { - let f = f.expect("parses"); - assert!(matches!( - f.dequant_to_safetensors(&|_| Some(GgufMap::Rename("same".into()))), - Err(GgufError::BadTensor { .. }) - )); - }); - } - - #[test] - fn dequanting_to_a_file_is_byte_identical_to_dequanting_in_memory() { - // Mixed dtypes and widths, so the pin covers the quantized path and - // the tensor-to-tensor offset arithmetic, not just one F32 copy. - // Q8_0 blocks are 34 B for 32 elements: two blocks = 68 B at offset 0, - // then a 6-element F32 tensor at the next 32-aligned offset. - let mut payload = vec![0u8; 96]; - for (i, b) in payload.iter_mut().enumerate().take(68) { - *b = u8::try_from(i % 251).expect("bounded by the modulus"); - } - payload.extend((1..=6u8).flat_map(|v| f32::from(v).to_le_bytes())); - let bytes = TestGguf::new() - .kv_str("general.architecture", "qwen2") - .tensor("blk.0.attn_q.weight", &[32, 2], 8, 0) // Q8_0 - .tensor("token_embd.weight", &[2, 3], 0, 96) // F32 - .build_with_payload(&payload); - - with_gguf_bytes(&bytes, |f| { - let f = f.expect("parses"); - let map = |i: &GgufTensorInfo| Some(GgufMap::Rename(format!("model.{}", i.name))); - let blob = f.dequant_to_safetensors(&map).expect("serializes"); - - let out = scratch_path("dequant-pin"); - let written = f - .dequant_to_safetensors_file(&map, &out) - .expect("streams to a file"); - let from_file = std::fs::read(&out).expect("reads back"); - let _ = std::fs::remove_file(&out); - - assert_eq!( - blob, from_file, - "the in-memory and to-file dequants must stay ONE importer" - ); - // The returned count is the payload, header excluded — so the - // caller can size the model without re-statting the file. - let header_len = u64::from_le_bytes(blob[0..8].try_into().expect("8 bytes")); - assert_eq!(written, blob.len() as u64 - 8 - header_len); - // 64 Q8_0 elements + 6 F32 elements, all at f32. - assert_eq!(written, (64 + 6) * 4); - }); - } - - #[test] - fn a_bad_map_is_rejected_before_any_payload_is_read() { - // The tensor table claims a payload far past the end of the file, so - // ANY read of it is an io error. A map error must still surface as the - // map error: planning happens first, by construction. - let bytes = TestGguf::new() - .tensor("a.weight", &[8], 0, 0) - .tensor("b.weight", &[1 << 20], 0, 32) - .build_with_payload(&[0u8; 64]); - with_gguf_bytes(&bytes, |f| { - let f = f.expect("parses"); - assert!(matches!( - f.dequant_to_safetensors( - &|i| (i.name != "b.weight").then(|| GgufMap::Rename(i.name.clone())) - ), - Err(GgufError::BadTensor { index: 1, .. }) - )); - // And the collision check too — it is the second claim planning - // makes, on a tensor whose payload is unreadable. - assert!(matches!( - f.dequant_to_safetensors(&|_| Some(GgufMap::Rename("same".into()))), - Err(GgufError::BadTensor { index: 1, .. }) - )); - }); - } - - #[test] - fn misaligned_offsets_partial_blocks_and_duplicates_are_rejected() { - let misaligned = TestGguf::new().tensor("t", &[32], 8, 7).build(); - assert!(matches!( - open_bytes(&misaligned), - Err(GgufError::BadTensor { .. }) - )); - // 100 elements is not whole 256-element Q4_K superblocks. - let partial = TestGguf::new().tensor("t", &[100], 12, 0).build(); - assert!(matches!( - open_bytes(&partial), - Err(GgufError::BadTensor { .. }) - )); - let dup = TestGguf::new() - .tensor("t", &[32], 8, 0) - .tensor("t", &[32], 8, 64) - .build(); - assert!(matches!(open_bytes(&dup), Err(GgufError::BadTensor { .. }))); - } -} diff --git a/crates/mummu/examples/src/gguf_iq_grids.rs b/crates/mummu/examples/src/gguf_iq_grids.rs deleted file mode 100644 index 0dc3108..0000000 --- a/crates/mummu/examples/src/gguf_iq_grids.rs +++ /dev/null @@ -1,325 +0,0 @@ -//! The IQ-quant codebook tables, extracted verbatim from ggml's -//! `ggml-common.h` (llama.cpp master, fetched 2026-08-21) by a regex -//! generator, never hand-transcribed. Each grid entry packs 8 (u64) or -//! 4 (u32) unsigned byte magnitudes; `KSIGNS_IQ2XS` maps a 7-bit index -//! to 8 sign bits (odd parity in bit 7). - -pub(crate) const KSIGNS_IQ2XS: [u8; 128] = [ - 0, 129, 130, 3, 132, 5, 6, 135, - 136, 9, 10, 139, 12, 141, 142, 15, - 144, 17, 18, 147, 20, 149, 150, 23, - 24, 153, 154, 27, 156, 29, 30, 159, - 160, 33, 34, 163, 36, 165, 166, 39, - 40, 169, 170, 43, 172, 45, 46, 175, - 48, 177, 178, 51, 180, 53, 54, 183, - 184, 57, 58, 187, 60, 189, 190, 63, - 192, 65, 66, 195, 68, 197, 198, 71, - 72, 201, 202, 75, 204, 77, 78, 207, - 80, 209, 210, 83, 212, 85, 86, 215, - 216, 89, 90, 219, 92, 221, 222, 95, - 96, 225, 226, 99, 228, 101, 102, 231, - 232, 105, 106, 235, 108, 237, 238, 111, - 240, 113, 114, 243, 116, 245, 246, 119, - 120, 249, 250, 123, 252, 125, 126, 255, -]; - -pub(crate) const IQ2XS_GRID: [u64; 512] = [ - 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x080808080819192b, - 0x0808080808192b19, 0x08080808082b0808, 0x08080808082b082b, 0x08080808082b1919, 0x08080808082b2b08, 0x0808080819080819, 0x0808080819081908, 0x080808081908192b, - 0x0808080819082b19, 0x0808080819190808, 0x080808081919082b, 0x0808080819191919, 0x0808080819192b08, 0x08080808192b0819, 0x08080808192b1908, 0x080808082b080808, - 0x080808082b08082b, 0x080808082b081919, 0x080808082b082b08, 0x080808082b190819, 0x080808082b191908, 0x080808082b192b19, 0x080808082b2b0808, 0x0808081908080819, - 0x0808081908081908, 0x080808190808192b, 0x0808081908082b19, 0x0808081908190808, 0x080808190819082b, 0x0808081908191919, 0x0808081908192b08, 0x0808081908192b2b, - 0x08080819082b0819, 0x08080819082b1908, 0x0808081919080808, 0x080808191908082b, 0x0808081919081919, 0x0808081919082b08, 0x0808081919190819, 0x0808081919191908, - 0x08080819192b0808, 0x08080819192b2b08, 0x080808192b080819, 0x080808192b081908, 0x080808192b190808, 0x0808082b08080808, 0x0808082b0808082b, 0x0808082b08081919, - 0x0808082b08082b08, 0x0808082b08190819, 0x0808082b08191908, 0x0808082b082b0808, 0x0808082b19080819, 0x0808082b19081908, 0x0808082b19190808, 0x0808082b19191919, - 0x0808082b2b080808, 0x0808082b2b082b2b, 0x0808190808080819, 0x0808190808081908, 0x080819080808192b, 0x0808190808082b19, 0x0808190808190808, 0x080819080819082b, - 0x0808190808191919, 0x0808190808192b08, 0x08081908082b0819, 0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819081919, 0x0808190819082b08, - 0x0808190819190819, 0x0808190819191908, 0x080819081919192b, 0x08081908192b0808, 0x080819082b080819, 0x080819082b081908, 0x080819082b190808, 0x0808191908080808, - 0x080819190808082b, 0x0808191908081919, 0x0808191908082b08, 0x0808191908190819, 0x0808191908191908, 0x08081919082b0808, 0x0808191919080819, 0x0808191919081908, - 0x0808191919190808, 0x08081919192b0819, 0x080819192b080808, 0x0808192b08080819, 0x0808192b08081908, 0x0808192b08190808, 0x0808192b082b192b, 0x0808192b19080808, - 0x0808192b1908082b, 0x0808192b2b081908, 0x08082b0808080808, 0x08082b080808082b, 0x08082b0808081919, 0x08082b0808082b08, 0x08082b0808082b2b, 0x08082b0808190819, - 0x08082b0808191908, 0x08082b08082b0808, 0x08082b08082b1919, 0x08082b0819080819, 0x08082b0819081908, 0x08082b0819190808, 0x08082b0819192b08, 0x08082b082b080808, - 0x08082b082b2b0808, 0x08082b082b2b2b2b, 0x08082b1908080819, 0x08082b1908081908, 0x08082b1908190808, 0x08082b1919080808, 0x08082b192b080819, 0x08082b192b082b19, - 0x08082b2b08080808, 0x08082b2b082b0808, 0x08082b2b082b2b08, 0x08082b2b2b19192b, 0x08082b2b2b2b0808, 0x0819080808080819, 0x0819080808081908, 0x081908080808192b, - 0x0819080808082b19, 0x0819080808190808, 0x081908080819082b, 0x0819080808191919, 0x0819080808192b08, 0x08190808082b0819, 0x08190808082b1908, 0x0819080819080808, - 0x081908081908082b, 0x0819080819081919, 0x0819080819082b08, 0x0819080819190819, 0x0819080819191908, 0x08190808192b0808, 0x08190808192b2b2b, 0x081908082b080819, - 0x081908082b081908, 0x081908082b190808, 0x0819081908080808, 0x081908190808082b, 0x0819081908081919, 0x0819081908082b08, 0x0819081908190819, 0x0819081908191908, - 0x08190819082b0808, 0x0819081919080819, 0x0819081919081908, 0x0819081919190808, 0x081908192b080808, 0x081908192b191908, 0x081908192b19192b, 0x0819082b08080819, - 0x0819082b08081908, 0x0819082b0808192b, 0x0819082b08190808, 0x0819082b19080808, 0x0819082b192b0808, 0x0819190808080808, 0x081919080808082b, 0x0819190808081919, - 0x0819190808082b08, 0x0819190808190819, 0x0819190808191908, 0x08191908082b0808, 0x0819190819080819, 0x0819190819081908, 0x0819190819082b19, 0x0819190819190808, - 0x08191908192b1908, 0x081919082b080808, 0x0819191908080819, 0x0819191908081908, 0x0819191908190808, 0x0819191919080808, 0x0819192b08080808, 0x0819192b08191908, - 0x0819192b19082b19, 0x08192b0808080819, 0x08192b0808081908, 0x08192b0808190808, 0x08192b080819082b, 0x08192b0819080808, 0x08192b0819191908, 0x08192b082b08192b, - 0x08192b1908080808, 0x08192b1908081919, 0x08192b19192b192b, 0x08192b2b19190819, 0x08192b2b2b2b2b19, 0x082b080808080808, 0x082b08080808082b, 0x082b080808081919, - 0x082b080808082b08, 0x082b080808082b2b, 0x082b080808190819, 0x082b080808191908, 0x082b0808082b0808, 0x082b080819080819, 0x082b080819081908, 0x082b080819190808, - 0x082b08082b080808, 0x082b08082b2b0808, 0x082b081908080819, 0x082b081908081908, 0x082b081908190808, 0x082b081919080808, 0x082b081919082b08, 0x082b0819192b1919, - 0x082b082b08080808, 0x082b082b082b082b, 0x082b082b2b080808, 0x082b082b2b2b2b08, 0x082b190808080819, 0x082b190808081908, 0x082b190808190808, 0x082b1908082b2b19, - 0x082b190819080808, 0x082b191908080808, 0x082b191919080819, 0x082b19191919082b, 0x082b19192b192b19, 0x082b192b08080819, 0x082b192b08192b2b, 0x082b192b2b2b192b, - 0x082b2b0808080808, 0x082b2b0808082b08, 0x082b2b0808082b2b, 0x082b2b08082b0808, 0x082b2b0819191919, 0x082b2b082b082b08, 0x082b2b082b2b082b, 0x082b2b19192b2b08, - 0x082b2b192b190808, 0x082b2b2b08082b08, 0x082b2b2b082b0808, 0x082b2b2b2b08082b, 0x082b2b2b2b082b08, 0x082b2b2b2b082b2b, 0x1908080808080819, 0x1908080808081908, - 0x190808080808192b, 0x1908080808082b19, 0x1908080808190808, 0x190808080819082b, 0x1908080808191919, 0x1908080808192b08, 0x19080808082b0819, 0x19080808082b1908, - 0x1908080819080808, 0x190808081908082b, 0x1908080819081919, 0x1908080819082b08, 0x1908080819082b2b, 0x1908080819190819, 0x1908080819191908, 0x19080808192b0808, - 0x19080808192b1919, 0x190808082b080819, 0x190808082b081908, 0x190808082b190808, 0x1908081908080808, 0x190808190808082b, 0x1908081908081919, 0x1908081908082b08, - 0x1908081908190819, 0x1908081908191908, 0x19080819082b0808, 0x1908081919080819, 0x1908081919081908, 0x1908081919190808, 0x190808192b080808, 0x190808192b081919, - 0x190808192b2b082b, 0x1908082b08080819, 0x1908082b08081908, 0x1908082b08190808, 0x1908082b0819082b, 0x1908082b082b2b19, 0x1908082b19080808, 0x1908190808080808, - 0x190819080808082b, 0x1908190808081919, 0x1908190808082b08, 0x1908190808190819, 0x1908190808191908, 0x1908190808192b19, 0x19081908082b0808, 0x1908190819080819, - 0x1908190819081908, 0x1908190819190808, 0x190819082b080808, 0x190819082b191908, 0x1908191908080819, 0x1908191908081908, 0x1908191908190808, 0x19081919082b1908, - 0x1908191919080808, 0x190819192b192b2b, 0x1908192b08080808, 0x1908192b08082b2b, 0x1908192b19081908, 0x1908192b19190808, 0x19082b0808080819, 0x19082b0808081908, - 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919, 0x19082b0819191908, 0x19082b08192b082b, 0x19082b1908080808, 0x19082b1908190819, 0x19082b1919081908, - 0x19082b1919190808, 0x19082b19192b2b19, 0x19082b2b08081908, 0x1919080808080808, 0x191908080808082b, 0x1919080808081919, 0x1919080808082b08, 0x1919080808190819, - 0x1919080808191908, 0x19190808082b0808, 0x19190808082b2b08, 0x1919080819080819, 0x1919080819081908, 0x1919080819190808, 0x191908082b080808, 0x1919081908080819, - 0x1919081908081908, 0x1919081908190808, 0x1919081908191919, 0x1919081919080808, 0x191908191908082b, 0x1919082b08080808, 0x1919082b19081908, 0x1919082b2b2b2b2b, - 0x1919190808080819, 0x1919190808081908, 0x1919190808190808, 0x19191908082b0819, 0x1919190819080808, 0x19191908192b0808, 0x191919082b080819, 0x191919082b2b0819, - 0x1919191908080808, 0x1919191908082b08, 0x191919192b080808, 0x191919192b082b08, 0x1919192b082b0819, 0x1919192b192b2b08, 0x1919192b2b2b0819, 0x19192b0808080808, - 0x19192b0808191908, 0x19192b0819080819, 0x19192b0819190808, 0x19192b082b192b19, 0x19192b1908192b2b, 0x19192b1919080808, 0x19192b191908082b, 0x19192b2b2b081919, - 0x192b080808080819, 0x192b080808081908, 0x192b080808190808, 0x192b080819080808, 0x192b080819191908, 0x192b0808192b082b, 0x192b08082b08192b, 0x192b08082b2b2b19, - 0x192b081908080808, 0x192b082b082b1908, 0x192b082b19082b2b, 0x192b082b2b19082b, 0x192b190808080808, 0x192b19080819192b, 0x192b191908190808, 0x192b191919080808, - 0x192b191919081919, 0x192b19192b2b1908, 0x192b2b0808080819, 0x192b2b08192b2b2b, 0x192b2b19082b1919, 0x192b2b2b0808192b, 0x192b2b2b19191908, 0x192b2b2b192b082b, - 0x2b08080808080808, 0x2b0808080808082b, 0x2b08080808081919, 0x2b08080808082b08, 0x2b08080808190819, 0x2b08080808191908, 0x2b080808082b0808, 0x2b080808082b2b2b, - 0x2b08080819080819, 0x2b08080819081908, 0x2b08080819190808, 0x2b0808082b080808, 0x2b0808082b08082b, 0x2b0808082b2b2b08, 0x2b0808082b2b2b2b, 0x2b08081908080819, - 0x2b08081908081908, 0x2b0808190808192b, 0x2b08081908190808, 0x2b08081919080808, 0x2b08081919190819, 0x2b08081919192b19, 0x2b08082b08080808, 0x2b08082b082b0808, - 0x2b08082b2b080808, 0x2b08082b2b08082b, 0x2b08082b2b2b0808, 0x2b08082b2b2b2b08, 0x2b08190808080819, 0x2b08190808081908, 0x2b08190808190808, 0x2b0819080819082b, - 0x2b08190808191919, 0x2b08190819080808, 0x2b081908192b0808, 0x2b0819082b082b19, 0x2b08191908080808, 0x2b08191919081908, 0x2b0819192b2b1919, 0x2b08192b08192b08, - 0x2b08192b192b2b2b, 0x2b082b0808080808, 0x2b082b0808082b08, 0x2b082b08082b1919, 0x2b082b0819192b2b, 0x2b082b082b080808, 0x2b082b082b08082b, 0x2b082b082b2b2b08, - 0x2b082b190808192b, 0x2b082b2b082b082b, 0x2b082b2b2b080808, 0x2b082b2b2b082b08, 0x2b082b2b2b19192b, 0x2b082b2b2b2b2b08, 0x2b19080808080819, 0x2b19080808081908, - 0x2b19080808190808, 0x2b19080819080808, 0x2b1908081919192b, 0x2b1908082b081908, 0x2b19081908080808, 0x2b190819082b082b, 0x2b190819192b1908, 0x2b19082b1919192b, - 0x2b19082b2b082b19, 0x2b19190808080808, 0x2b19190808081919, 0x2b19190819081908, 0x2b19190819190808, 0x2b19190819192b08, 0x2b191919082b2b19, 0x2b1919192b190808, - 0x2b1919192b19082b, 0x2b19192b19080819, 0x2b192b0819190819, 0x2b192b082b2b192b, 0x2b192b1919082b19, 0x2b192b2b08191919, 0x2b192b2b192b0808, 0x2b2b080808080808, - 0x2b2b08080808082b, 0x2b2b080808082b08, 0x2b2b080808082b2b, 0x2b2b0808082b0808, 0x2b2b0808082b2b2b, 0x2b2b08082b2b0808, 0x2b2b081919190819, 0x2b2b081919192b19, - 0x2b2b08192b2b192b, 0x2b2b082b08080808, 0x2b2b082b0808082b, 0x2b2b082b08082b08, 0x2b2b082b082b2b2b, 0x2b2b082b2b080808, 0x2b2b082b2b2b0808, 0x2b2b190819080808, - 0x2b2b19082b191919, 0x2b2b192b192b1919, 0x2b2b192b2b192b08, 0x2b2b2b0808082b2b, 0x2b2b2b08082b0808, 0x2b2b2b08082b082b, 0x2b2b2b08082b2b08, 0x2b2b2b082b2b0808, - 0x2b2b2b082b2b2b08, 0x2b2b2b1908081908, 0x2b2b2b192b081908, 0x2b2b2b192b08192b, 0x2b2b2b2b082b2b08, 0x2b2b2b2b082b2b2b, 0x2b2b2b2b2b190819, 0x2b2b2b2b2b2b2b2b, -]; - -pub(crate) const IQ2S_GRID: [u64; 1024] = [ - 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x080808080819192b, - 0x0808080808192b19, 0x08080808082b0808, 0x08080808082b082b, 0x08080808082b1919, 0x08080808082b2b08, 0x0808080819080819, 0x0808080819081908, 0x080808081908192b, - 0x0808080819082b19, 0x0808080819190808, 0x080808081919082b, 0x0808080819191919, 0x0808080819192b08, 0x08080808192b0819, 0x08080808192b1908, 0x08080808192b192b, - 0x08080808192b2b19, 0x080808082b080808, 0x080808082b08082b, 0x080808082b081919, 0x080808082b082b08, 0x080808082b190819, 0x080808082b191908, 0x080808082b2b0808, - 0x080808082b2b1919, 0x080808082b2b2b2b, 0x0808081908080819, 0x0808081908081908, 0x080808190808192b, 0x0808081908082b19, 0x0808081908190808, 0x080808190819082b, - 0x0808081908191919, 0x0808081908192b08, 0x08080819082b0819, 0x08080819082b1908, 0x0808081919080808, 0x080808191908082b, 0x0808081919081919, 0x0808081919082b08, - 0x0808081919190819, 0x0808081919191908, 0x080808191919192b, 0x0808081919192b19, 0x08080819192b0808, 0x08080819192b1919, 0x08080819192b2b08, 0x080808192b080819, - 0x080808192b081908, 0x080808192b190808, 0x080808192b19082b, 0x080808192b191919, 0x080808192b2b0819, 0x080808192b2b1908, 0x0808082b08080808, 0x0808082b0808082b, - 0x0808082b08081919, 0x0808082b08082b08, 0x0808082b08190819, 0x0808082b08191908, 0x0808082b082b0808, 0x0808082b082b2b2b, 0x0808082b19080819, 0x0808082b19081908, - 0x0808082b1908192b, 0x0808082b19082b19, 0x0808082b19190808, 0x0808082b19191919, 0x0808082b2b080808, 0x0808082b2b081919, 0x0808082b2b082b2b, 0x0808082b2b191908, - 0x0808082b2b2b082b, 0x0808190808080819, 0x0808190808081908, 0x080819080808192b, 0x0808190808082b19, 0x0808190808190808, 0x080819080819082b, 0x0808190808191919, - 0x0808190808192b08, 0x08081908082b0819, 0x08081908082b1908, 0x08081908082b192b, 0x08081908082b2b19, 0x0808190819080808, 0x080819081908082b, 0x0808190819081919, - 0x0808190819082b08, 0x0808190819082b2b, 0x0808190819190819, 0x0808190819191908, 0x080819081919192b, 0x0808190819192b19, 0x08081908192b0808, 0x08081908192b082b, - 0x08081908192b1919, 0x080819082b080819, 0x080819082b081908, 0x080819082b08192b, 0x080819082b082b19, 0x080819082b190808, 0x080819082b191919, 0x080819082b192b08, - 0x080819082b2b0819, 0x080819082b2b1908, 0x0808191908080808, 0x080819190808082b, 0x0808191908081919, 0x0808191908082b08, 0x0808191908082b2b, 0x0808191908190819, - 0x0808191908191908, 0x080819190819192b, 0x0808191908192b19, 0x08081919082b0808, 0x08081919082b1919, 0x08081919082b2b08, 0x0808191919080819, 0x0808191919081908, - 0x080819191908192b, 0x0808191919082b19, 0x0808191919190808, 0x080819191919082b, 0x0808191919191919, 0x0808191919192b08, 0x08081919192b0819, 0x08081919192b1908, - 0x080819192b080808, 0x080819192b08082b, 0x080819192b081919, 0x080819192b082b08, 0x080819192b190819, 0x080819192b191908, 0x080819192b2b0808, 0x0808192b08080819, - 0x0808192b08081908, 0x0808192b0808192b, 0x0808192b08082b19, 0x0808192b08190808, 0x0808192b08191919, 0x0808192b19080808, 0x0808192b19081919, 0x0808192b19082b08, - 0x0808192b19190819, 0x0808192b19191908, 0x0808192b192b0808, 0x0808192b2b080819, 0x0808192b2b081908, 0x0808192b2b190808, 0x08082b0808080808, 0x08082b080808082b, - 0x08082b0808081919, 0x08082b0808082b08, 0x08082b0808190819, 0x08082b0808191908, 0x08082b080819192b, 0x08082b0808192b19, 0x08082b08082b0808, 0x08082b08082b1919, - 0x08082b08082b2b2b, 0x08082b0819080819, 0x08082b0819081908, 0x08082b081908192b, 0x08082b0819082b19, 0x08082b0819190808, 0x08082b081919082b, 0x08082b0819191919, - 0x08082b0819192b08, 0x08082b08192b0819, 0x08082b08192b1908, 0x08082b082b080808, 0x08082b082b081919, 0x08082b082b191908, 0x08082b082b2b2b2b, 0x08082b1908080819, - 0x08082b1908081908, 0x08082b1908190808, 0x08082b190819082b, 0x08082b1908191919, 0x08082b1908192b08, 0x08082b19082b0819, 0x08082b1919080808, 0x08082b1919081919, - 0x08082b1919082b08, 0x08082b1919190819, 0x08082b1919191908, 0x08082b19192b0808, 0x08082b192b080819, 0x08082b192b190808, 0x08082b2b08080808, 0x08082b2b08190819, - 0x08082b2b08191908, 0x08082b2b082b082b, 0x08082b2b082b2b08, 0x08082b2b082b2b2b, 0x08082b2b19190808, 0x08082b2b2b192b19, 0x0819080808080819, 0x0819080808081908, - 0x081908080808192b, 0x0819080808082b19, 0x0819080808190808, 0x081908080819082b, 0x0819080808191919, 0x0819080808192b08, 0x08190808082b0819, 0x08190808082b1908, - 0x08190808082b192b, 0x0819080819080808, 0x081908081908082b, 0x0819080819081919, 0x0819080819082b08, 0x0819080819190819, 0x0819080819191908, 0x081908081919192b, - 0x0819080819192b19, 0x08190808192b0808, 0x08190808192b082b, 0x08190808192b1919, 0x08190808192b2b08, 0x081908082b080819, 0x081908082b081908, 0x081908082b08192b, - 0x081908082b190808, 0x081908082b191919, 0x081908082b192b08, 0x081908082b2b0819, 0x081908082b2b1908, 0x0819081908080808, 0x081908190808082b, 0x0819081908081919, - 0x0819081908082b08, 0x0819081908082b2b, 0x0819081908190819, 0x0819081908191908, 0x081908190819192b, 0x0819081908192b19, 0x08190819082b0808, 0x08190819082b082b, - 0x08190819082b1919, 0x08190819082b2b08, 0x0819081919080819, 0x0819081919081908, 0x081908191908192b, 0x0819081919082b19, 0x0819081919190808, 0x081908191919082b, - 0x0819081919191919, 0x0819081919192b08, 0x08190819192b0819, 0x08190819192b1908, 0x081908192b080808, 0x081908192b08082b, 0x081908192b081919, 0x081908192b082b08, - 0x081908192b190819, 0x081908192b191908, 0x0819082b08080819, 0x0819082b08081908, 0x0819082b08082b19, 0x0819082b08190808, 0x0819082b08191919, 0x0819082b082b0819, - 0x0819082b082b1908, 0x0819082b19080808, 0x0819082b19081919, 0x0819082b19190819, 0x0819082b19191908, 0x0819082b2b080819, 0x0819082b2b081908, 0x0819082b2b190808, - 0x0819190808080808, 0x081919080808082b, 0x0819190808081919, 0x0819190808082b08, 0x0819190808190819, 0x0819190808191908, 0x081919080819192b, 0x0819190808192b19, - 0x08191908082b0808, 0x08191908082b1919, 0x08191908082b2b08, 0x0819190819080819, 0x0819190819081908, 0x081919081908192b, 0x0819190819082b19, 0x0819190819190808, - 0x081919081919082b, 0x0819190819191919, 0x0819190819192b08, 0x08191908192b0819, 0x08191908192b1908, 0x081919082b080808, 0x081919082b08082b, 0x081919082b081919, - 0x081919082b082b08, 0x081919082b190819, 0x081919082b191908, 0x081919082b2b0808, 0x0819191908080819, 0x0819191908081908, 0x081919190808192b, 0x0819191908082b19, - 0x0819191908190808, 0x081919190819082b, 0x0819191908191919, 0x0819191908192b08, 0x08191919082b0819, 0x08191919082b1908, 0x0819191919080808, 0x081919191908082b, - 0x0819191919081919, 0x0819191919082b08, 0x0819191919190819, 0x0819191919191908, 0x08191919192b0808, 0x081919192b080819, 0x081919192b081908, 0x081919192b190808, - 0x0819192b08080808, 0x0819192b08081919, 0x0819192b08082b08, 0x0819192b08190819, 0x0819192b08191908, 0x0819192b082b0808, 0x0819192b19080819, 0x0819192b19081908, - 0x0819192b19190808, 0x0819192b2b080808, 0x0819192b2b2b2b2b, 0x08192b0808080819, 0x08192b0808081908, 0x08192b080808192b, 0x08192b0808082b19, 0x08192b0808190808, - 0x08192b0808191919, 0x08192b0808192b08, 0x08192b08082b0819, 0x08192b0819080808, 0x08192b081908082b, 0x08192b0819081919, 0x08192b0819082b08, 0x08192b0819190819, - 0x08192b0819191908, 0x08192b08192b0808, 0x08192b082b080819, 0x08192b082b081908, 0x08192b1908080808, 0x08192b190808082b, 0x08192b1908081919, 0x08192b1908082b08, - 0x08192b1908190819, 0x08192b1908191908, 0x08192b19082b0808, 0x08192b1919080819, 0x08192b1919081908, 0x08192b1919190808, 0x08192b19192b2b19, 0x08192b192b2b082b, - 0x08192b2b08081908, 0x08192b2b08190808, 0x08192b2b19080808, 0x08192b2b1919192b, 0x082b080808080808, 0x082b08080808082b, 0x082b080808081919, 0x082b080808082b08, - 0x082b080808190819, 0x082b080808191908, 0x082b08080819192b, 0x082b080808192b19, 0x082b0808082b0808, 0x082b0808082b1919, 0x082b0808082b2b2b, 0x082b080819080819, - 0x082b080819081908, 0x082b080819190808, 0x082b08081919082b, 0x082b080819191919, 0x082b0808192b1908, 0x082b08082b080808, 0x082b08082b082b2b, 0x082b08082b191908, - 0x082b08082b2b2b2b, 0x082b081908080819, 0x082b081908081908, 0x082b081908190808, 0x082b08190819082b, 0x082b081908191919, 0x082b0819082b0819, 0x082b081919080808, - 0x082b08191908082b, 0x082b081919081919, 0x082b081919190819, 0x082b081919191908, 0x082b0819192b0808, 0x082b08192b080819, 0x082b08192b081908, 0x082b08192b190808, - 0x082b082b08080808, 0x082b082b08082b2b, 0x082b082b082b082b, 0x082b082b082b2b08, 0x082b082b082b2b2b, 0x082b082b19081908, 0x082b082b19190808, 0x082b082b2b082b08, - 0x082b082b2b082b2b, 0x082b082b2b2b2b08, 0x082b190808080819, 0x082b190808081908, 0x082b19080808192b, 0x082b190808082b19, 0x082b190808190808, 0x082b190808191919, - 0x082b190808192b08, 0x082b1908082b0819, 0x082b1908082b1908, 0x082b190819080808, 0x082b19081908082b, 0x082b190819081919, 0x082b190819082b08, 0x082b190819190819, - 0x082b190819191908, 0x082b1908192b0808, 0x082b19082b080819, 0x082b19082b081908, 0x082b19082b190808, 0x082b191908080808, 0x082b191908081919, 0x082b191908082b08, - 0x082b191908190819, 0x082b191908191908, 0x082b1919082b0808, 0x082b191919080819, 0x082b191919081908, 0x082b191919190808, 0x082b1919192b192b, 0x082b19192b080808, - 0x082b192b08080819, 0x082b192b08081908, 0x082b192b08190808, 0x082b192b19080808, 0x082b192b19192b19, 0x082b2b0808080808, 0x082b2b0808081919, 0x082b2b0808190819, - 0x082b2b0808191908, 0x082b2b0819080819, 0x082b2b0819081908, 0x082b2b0819190808, 0x082b2b082b082b2b, 0x082b2b082b2b2b2b, 0x082b2b1908080819, 0x082b2b1908081908, - 0x082b2b1908190808, 0x082b2b192b191919, 0x082b2b2b08082b2b, 0x082b2b2b082b082b, 0x082b2b2b192b1908, 0x082b2b2b2b082b08, 0x082b2b2b2b082b2b, 0x1908080808080819, - 0x1908080808081908, 0x190808080808192b, 0x1908080808082b19, 0x1908080808190808, 0x190808080819082b, 0x1908080808191919, 0x1908080808192b08, 0x1908080808192b2b, - 0x19080808082b0819, 0x19080808082b1908, 0x19080808082b192b, 0x1908080819080808, 0x190808081908082b, 0x1908080819081919, 0x1908080819082b08, 0x1908080819082b2b, - 0x1908080819190819, 0x1908080819191908, 0x190808081919192b, 0x1908080819192b19, 0x19080808192b0808, 0x19080808192b082b, 0x19080808192b1919, 0x190808082b080819, - 0x190808082b081908, 0x190808082b190808, 0x190808082b191919, 0x190808082b192b08, 0x190808082b2b0819, 0x190808082b2b1908, 0x1908081908080808, 0x190808190808082b, - 0x1908081908081919, 0x1908081908082b08, 0x1908081908190819, 0x1908081908191908, 0x190808190819192b, 0x1908081908192b19, 0x19080819082b0808, 0x19080819082b082b, - 0x19080819082b1919, 0x1908081919080819, 0x1908081919081908, 0x190808191908192b, 0x1908081919082b19, 0x1908081919190808, 0x190808191919082b, 0x1908081919191919, - 0x1908081919192b08, 0x19080819192b0819, 0x19080819192b1908, 0x190808192b080808, 0x190808192b08082b, 0x190808192b081919, 0x190808192b082b08, 0x190808192b190819, - 0x190808192b191908, 0x190808192b2b0808, 0x1908082b08080819, 0x1908082b08081908, 0x1908082b08190808, 0x1908082b0819082b, 0x1908082b08191919, 0x1908082b08192b08, - 0x1908082b082b1908, 0x1908082b19080808, 0x1908082b19081919, 0x1908082b19082b08, 0x1908082b19190819, 0x1908082b19191908, 0x1908082b192b0808, 0x1908082b2b080819, - 0x1908082b2b081908, 0x1908190808080808, 0x190819080808082b, 0x1908190808081919, 0x1908190808082b08, 0x1908190808082b2b, 0x1908190808190819, 0x1908190808191908, - 0x190819080819192b, 0x1908190808192b19, 0x19081908082b0808, 0x19081908082b082b, 0x19081908082b1919, 0x19081908082b2b08, 0x1908190819080819, 0x1908190819081908, - 0x190819081908192b, 0x1908190819082b19, 0x1908190819190808, 0x190819081919082b, 0x1908190819191919, 0x1908190819192b08, 0x19081908192b0819, 0x19081908192b1908, - 0x190819082b080808, 0x190819082b08082b, 0x190819082b081919, 0x190819082b082b08, 0x190819082b190819, 0x190819082b191908, 0x190819082b2b0808, 0x1908191908080819, - 0x1908191908081908, 0x190819190808192b, 0x1908191908082b19, 0x1908191908190808, 0x190819190819082b, 0x1908191908191919, 0x1908191908192b08, 0x19081919082b0819, - 0x19081919082b1908, 0x1908191919080808, 0x190819191908082b, 0x1908191919081919, 0x1908191919082b08, 0x1908191919190819, 0x1908191919191908, 0x19081919192b0808, - 0x19081919192b2b2b, 0x190819192b080819, 0x190819192b081908, 0x190819192b190808, 0x1908192b08080808, 0x1908192b0808082b, 0x1908192b08081919, 0x1908192b08082b08, - 0x1908192b08190819, 0x1908192b08191908, 0x1908192b082b0808, 0x1908192b19080819, 0x1908192b19081908, 0x1908192b19190808, 0x1908192b2b080808, 0x1908192b2b2b1919, - 0x19082b0808080819, 0x19082b0808081908, 0x19082b0808082b19, 0x19082b0808190808, 0x19082b080819082b, 0x19082b0808191919, 0x19082b0808192b08, 0x19082b08082b0819, - 0x19082b08082b1908, 0x19082b0819080808, 0x19082b081908082b, 0x19082b0819081919, 0x19082b0819082b08, 0x19082b0819190819, 0x19082b0819191908, 0x19082b08192b0808, - 0x19082b082b081908, 0x19082b082b190808, 0x19082b1908080808, 0x19082b190808082b, 0x19082b1908081919, 0x19082b1908082b08, 0x19082b1908190819, 0x19082b1908191908, - 0x19082b19082b0808, 0x19082b1919080819, 0x19082b1919081908, 0x19082b1919190808, 0x19082b192b080808, 0x19082b192b19192b, 0x19082b2b08080819, 0x19082b2b08081908, - 0x19082b2b08190808, 0x19082b2b19080808, 0x1919080808080808, 0x191908080808082b, 0x1919080808081919, 0x1919080808082b08, 0x1919080808190819, 0x1919080808191908, - 0x191908080819192b, 0x1919080808192b19, 0x19190808082b0808, 0x19190808082b082b, 0x19190808082b1919, 0x19190808082b2b08, 0x1919080819080819, 0x1919080819081908, - 0x191908081908192b, 0x1919080819082b19, 0x1919080819190808, 0x191908081919082b, 0x1919080819191919, 0x1919080819192b08, 0x19190808192b0819, 0x19190808192b1908, - 0x191908082b080808, 0x191908082b08082b, 0x191908082b081919, 0x191908082b082b08, 0x191908082b190819, 0x191908082b191908, 0x1919081908080819, 0x1919081908081908, - 0x191908190808192b, 0x1919081908082b19, 0x1919081908190808, 0x191908190819082b, 0x1919081908191919, 0x1919081908192b08, 0x19190819082b0819, 0x19190819082b1908, - 0x1919081919080808, 0x191908191908082b, 0x1919081919081919, 0x1919081919082b08, 0x1919081919190819, 0x1919081919191908, 0x19190819192b0808, 0x191908192b080819, - 0x191908192b081908, 0x191908192b190808, 0x1919082b08080808, 0x1919082b08081919, 0x1919082b08082b08, 0x1919082b08190819, 0x1919082b08191908, 0x1919082b082b0808, - 0x1919082b19080819, 0x1919082b19081908, 0x1919082b19190808, 0x1919082b192b2b19, 0x1919082b2b080808, 0x1919190808080819, 0x1919190808081908, 0x191919080808192b, - 0x1919190808082b19, 0x1919190808190808, 0x191919080819082b, 0x1919190808191919, 0x1919190808192b08, 0x19191908082b0819, 0x19191908082b1908, 0x1919190819080808, - 0x191919081908082b, 0x1919190819081919, 0x1919190819082b08, 0x1919190819190819, 0x1919190819191908, 0x19191908192b0808, 0x191919082b080819, 0x191919082b081908, - 0x191919082b190808, 0x1919191908080808, 0x191919190808082b, 0x1919191908081919, 0x1919191908082b08, 0x1919191908190819, 0x1919191908191908, 0x19191919082b0808, - 0x1919191919080819, 0x1919191919081908, 0x1919191919190808, 0x191919192b080808, 0x1919192b08080819, 0x1919192b08081908, 0x1919192b08190808, 0x1919192b082b192b, - 0x1919192b19080808, 0x19192b0808080808, 0x19192b080808082b, 0x19192b0808081919, 0x19192b0808082b08, 0x19192b0808190819, 0x19192b0808191908, 0x19192b08082b0808, - 0x19192b0819080819, 0x19192b0819081908, 0x19192b0819190808, 0x19192b0819192b2b, 0x19192b082b080808, 0x19192b1908080819, 0x19192b1908081908, 0x19192b1908190808, - 0x19192b1919080808, 0x19192b2b08080808, 0x19192b2b08192b19, 0x19192b2b2b081919, 0x19192b2b2b2b2b08, 0x192b080808080819, 0x192b080808081908, 0x192b08080808192b, - 0x192b080808190808, 0x192b08080819082b, 0x192b080808191919, 0x192b080808192b08, 0x192b0808082b0819, 0x192b0808082b1908, 0x192b080819080808, 0x192b080819081919, - 0x192b080819082b08, 0x192b080819190819, 0x192b080819191908, 0x192b0808192b0808, 0x192b08082b081908, 0x192b08082b190808, 0x192b081908080808, 0x192b08190808082b, - 0x192b081908081919, 0x192b081908082b08, 0x192b081908190819, 0x192b081908191908, 0x192b0819082b0808, 0x192b081919080819, 0x192b081919081908, 0x192b081919190808, - 0x192b08192b080808, 0x192b08192b192b19, 0x192b082b08081908, 0x192b082b08190808, 0x192b082b19080808, 0x192b082b1919192b, 0x192b082b2b2b0819, 0x192b190808080808, - 0x192b190808081919, 0x192b190808082b08, 0x192b190808190819, 0x192b190808191908, 0x192b1908082b0808, 0x192b190819080819, 0x192b190819081908, 0x192b190819190808, - 0x192b19082b080808, 0x192b191908080819, 0x192b191908081908, 0x192b191908190808, 0x192b191919080808, 0x192b191919082b2b, 0x192b1919192b2b08, 0x192b19192b19082b, - 0x192b192b08080808, 0x192b192b2b191908, 0x192b2b0808080819, 0x192b2b0808081908, 0x192b2b0808190808, 0x192b2b08192b1919, 0x192b2b082b192b08, 0x192b2b1908080808, - 0x192b2b19082b2b2b, 0x192b2b2b1908082b, 0x192b2b2b2b2b0819, 0x2b08080808080808, 0x2b0808080808082b, 0x2b08080808081919, 0x2b08080808082b08, 0x2b08080808190819, - 0x2b08080808191908, 0x2b08080808192b19, 0x2b080808082b0808, 0x2b080808082b1919, 0x2b08080819080819, 0x2b08080819081908, 0x2b08080819190808, 0x2b0808081919082b, - 0x2b08080819191919, 0x2b08080819192b08, 0x2b080808192b0819, 0x2b0808082b080808, 0x2b0808082b081919, 0x2b0808082b190819, 0x2b0808082b191908, 0x2b08081908080819, - 0x2b08081908081908, 0x2b08081908082b19, 0x2b08081908190808, 0x2b0808190819082b, 0x2b08081908191919, 0x2b08081908192b08, 0x2b080819082b0819, 0x2b080819082b1908, - 0x2b08081919080808, 0x2b0808191908082b, 0x2b08081919081919, 0x2b08081919082b08, 0x2b08081919190819, 0x2b08081919191908, 0x2b0808192b080819, 0x2b0808192b081908, - 0x2b0808192b190808, 0x2b0808192b2b2b19, 0x2b08082b08080808, 0x2b08082b08081919, 0x2b08082b08082b2b, 0x2b08082b08190819, 0x2b08082b08191908, 0x2b08082b19080819, - 0x2b08082b19081908, 0x2b08082b19190808, 0x2b08190808080819, 0x2b08190808081908, 0x2b0819080808192b, 0x2b08190808082b19, 0x2b08190808190808, 0x2b0819080819082b, - 0x2b08190808191919, 0x2b08190808192b08, 0x2b081908082b0819, 0x2b08190819080808, 0x2b0819081908082b, 0x2b08190819081919, 0x2b08190819082b08, 0x2b08190819190819, - 0x2b08190819191908, 0x2b081908192b0808, 0x2b0819082b080819, 0x2b0819082b081908, 0x2b0819082b190808, 0x2b08191908080808, 0x2b0819190808082b, 0x2b08191908081919, - 0x2b08191908082b08, 0x2b08191908190819, 0x2b08191908191908, 0x2b081919082b0808, 0x2b08191919080819, 0x2b08191919081908, 0x2b08191919190808, 0x2b0819192b080808, - 0x2b0819192b082b2b, 0x2b08192b08080819, 0x2b08192b08081908, 0x2b08192b08190808, 0x2b08192b082b2b19, 0x2b08192b19080808, 0x2b082b0808080808, 0x2b082b0808081919, - 0x2b082b0808190819, 0x2b082b0808191908, 0x2b082b0819080819, 0x2b082b0819081908, 0x2b082b0819190808, 0x2b082b082b2b082b, 0x2b082b1908080819, 0x2b082b1908081908, - 0x2b082b1919080808, 0x2b082b19192b1919, 0x2b082b2b082b082b, 0x2b082b2b19192b08, 0x2b082b2b19192b2b, 0x2b082b2b2b08082b, 0x2b082b2b2b2b082b, 0x2b19080808080819, - 0x2b19080808081908, 0x2b19080808082b19, 0x2b19080808190808, 0x2b1908080819082b, 0x2b19080808191919, 0x2b19080808192b08, 0x2b190808082b1908, 0x2b19080819080808, - 0x2b1908081908082b, 0x2b19080819081919, 0x2b19080819082b08, 0x2b19080819190819, 0x2b19080819191908, 0x2b190808192b0808, 0x2b1908082b080819, 0x2b1908082b081908, - 0x2b1908082b190808, 0x2b19081908080808, 0x2b19081908081919, 0x2b19081908190819, 0x2b19081908191908, 0x2b19081919080819, 0x2b19081919081908, 0x2b19081919190808, - 0x2b19081919192b2b, 0x2b19082b08080819, 0x2b19082b08081908, 0x2b19082b08190808, 0x2b19082b19080808, 0x2b19082b2b2b192b, 0x2b19190808080808, 0x2b1919080808082b, - 0x2b19190808081919, 0x2b19190808082b08, 0x2b19190808190819, 0x2b19190808191908, 0x2b191908082b0808, 0x2b19190819080819, 0x2b19190819081908, 0x2b19190819190808, - 0x2b1919082b080808, 0x2b1919082b19192b, 0x2b19191908080819, 0x2b19191908081908, 0x2b19191908190808, 0x2b19191919080808, 0x2b1919192b192b08, 0x2b1919192b2b0819, - 0x2b19192b08080808, 0x2b19192b1908192b, 0x2b19192b192b1908, 0x2b192b0808080819, 0x2b192b0808081908, 0x2b192b0808190808, 0x2b192b08082b192b, 0x2b192b0819080808, - 0x2b192b082b2b2b19, 0x2b192b1908080808, 0x2b192b1919082b19, 0x2b192b191919082b, 0x2b192b2b2b190808, 0x2b2b080808080808, 0x2b2b080808081919, 0x2b2b080808082b2b, - 0x2b2b080808191908, 0x2b2b0808082b082b, 0x2b2b0808082b2b2b, 0x2b2b080819080819, 0x2b2b080819081908, 0x2b2b080819190808, 0x2b2b08082b2b082b, 0x2b2b08082b2b2b2b, - 0x2b2b081919080808, 0x2b2b0819192b1919, 0x2b2b082b0808082b, 0x2b2b082b08082b2b, 0x2b2b082b082b082b, 0x2b2b082b082b2b08, 0x2b2b082b082b2b2b, 0x2b2b082b2b08082b, - 0x2b2b082b2b082b08, 0x2b2b082b2b082b2b, 0x2b2b082b2b2b2b08, 0x2b2b190808080819, 0x2b2b190808081908, 0x2b2b190808190808, 0x2b2b190819080808, 0x2b2b19082b082b19, - 0x2b2b19082b2b1908, 0x2b2b191908080808, 0x2b2b191908192b19, 0x2b2b192b19190819, 0x2b2b2b0808082b2b, 0x2b2b2b08082b2b08, 0x2b2b2b082b2b082b, 0x2b2b2b1919191908, - 0x2b2b2b192b08192b, 0x2b2b2b2b08082b08, 0x2b2b2b2b08082b2b, 0x2b2b2b2b082b0808, 0x2b2b2b2b082b082b, 0x2b2b2b2b082b2b08, 0x2b2b2b2b2b082b08, 0x2b2b2b2b2b2b2b2b, -]; - -pub(crate) const IQ3XXS_GRID: [u32; 256] = [ - 0x04040404, 0x04040414, 0x04040424, 0x04040c0c, 0x04040c1c, 0x04040c3e, 0x04041404, 0x04041414, - 0x04041c0c, 0x04042414, 0x04043e1c, 0x04043e2c, 0x040c040c, 0x040c041c, 0x040c0c04, 0x040c0c14, - 0x040c140c, 0x040c142c, 0x040c1c04, 0x040c1c14, 0x040c240c, 0x040c2c24, 0x040c3e04, 0x04140404, - 0x04140414, 0x04140424, 0x04140c0c, 0x04141404, 0x04141414, 0x04141c0c, 0x04141c1c, 0x04141c3e, - 0x04142c0c, 0x04142c3e, 0x04143e2c, 0x041c040c, 0x041c043e, 0x041c0c04, 0x041c0c14, 0x041c142c, - 0x041c3e04, 0x04240c1c, 0x04241c3e, 0x04242424, 0x04242c3e, 0x04243e1c, 0x04243e2c, 0x042c040c, - 0x042c043e, 0x042c1c14, 0x042c2c14, 0x04341c2c, 0x04343424, 0x043e0c04, 0x043e0c24, 0x043e0c34, - 0x043e241c, 0x043e340c, 0x0c04040c, 0x0c04041c, 0x0c040c04, 0x0c040c14, 0x0c04140c, 0x0c04141c, - 0x0c041c04, 0x0c041c14, 0x0c041c24, 0x0c04243e, 0x0c042c04, 0x0c0c0404, 0x0c0c0414, 0x0c0c0c0c, - 0x0c0c1404, 0x0c0c1414, 0x0c14040c, 0x0c14041c, 0x0c140c04, 0x0c140c14, 0x0c14140c, 0x0c141c04, - 0x0c143e14, 0x0c1c0404, 0x0c1c0414, 0x0c1c1404, 0x0c1c1c0c, 0x0c1c2434, 0x0c1c3434, 0x0c24040c, - 0x0c24042c, 0x0c242c04, 0x0c2c1404, 0x0c2c1424, 0x0c2c2434, 0x0c2c3e0c, 0x0c34042c, 0x0c3e1414, - 0x0c3e2404, 0x14040404, 0x14040414, 0x14040c0c, 0x14040c1c, 0x14041404, 0x14041414, 0x14041434, - 0x14041c0c, 0x14042414, 0x140c040c, 0x140c041c, 0x140c042c, 0x140c0c04, 0x140c0c14, 0x140c140c, - 0x140c1c04, 0x140c341c, 0x140c343e, 0x140c3e04, 0x14140404, 0x14140414, 0x14140c0c, 0x14140c3e, - 0x14141404, 0x14141414, 0x14141c3e, 0x14142404, 0x14142c2c, 0x141c040c, 0x141c0c04, 0x141c0c24, - 0x141c3e04, 0x141c3e24, 0x14241c2c, 0x14242c1c, 0x142c041c, 0x142c143e, 0x142c240c, 0x142c3e24, - 0x143e040c, 0x143e041c, 0x143e0c34, 0x143e242c, 0x1c04040c, 0x1c040c04, 0x1c040c14, 0x1c04140c, - 0x1c04141c, 0x1c042c04, 0x1c04342c, 0x1c043e14, 0x1c0c0404, 0x1c0c0414, 0x1c0c1404, 0x1c0c1c0c, - 0x1c0c2424, 0x1c0c2434, 0x1c14040c, 0x1c14041c, 0x1c140c04, 0x1c14142c, 0x1c142c14, 0x1c143e14, - 0x1c1c0c0c, 0x1c1c1c1c, 0x1c241c04, 0x1c24243e, 0x1c243e14, 0x1c2c0404, 0x1c2c0434, 0x1c2c1414, - 0x1c2c2c2c, 0x1c340c24, 0x1c341c34, 0x1c34341c, 0x1c3e1c1c, 0x1c3e3404, 0x24040424, 0x24040c3e, - 0x24041c2c, 0x24041c3e, 0x24042c1c, 0x24042c3e, 0x240c3e24, 0x24141404, 0x24141c3e, 0x24142404, - 0x24143404, 0x24143434, 0x241c043e, 0x241c242c, 0x24240424, 0x24242c0c, 0x24243424, 0x242c142c, - 0x242c241c, 0x242c3e04, 0x243e042c, 0x243e0c04, 0x243e0c14, 0x243e1c04, 0x2c040c14, 0x2c04240c, - 0x2c043e04, 0x2c0c0404, 0x2c0c0434, 0x2c0c1434, 0x2c0c2c2c, 0x2c140c24, 0x2c141c14, 0x2c143e14, - 0x2c1c0414, 0x2c1c2c1c, 0x2c240c04, 0x2c24141c, 0x2c24143e, 0x2c243e14, 0x2c2c0414, 0x2c2c1c0c, - 0x2c342c04, 0x2c3e1424, 0x2c3e2414, 0x34041424, 0x34042424, 0x34042434, 0x34043424, 0x340c140c, - 0x340c340c, 0x34140c3e, 0x34143424, 0x341c1c04, 0x341c1c34, 0x34242424, 0x342c042c, 0x342c2c14, - 0x34341c1c, 0x343e041c, 0x343e140c, 0x3e04041c, 0x3e04042c, 0x3e04043e, 0x3e040c04, 0x3e041c14, - 0x3e042c14, 0x3e0c1434, 0x3e0c2404, 0x3e140c14, 0x3e14242c, 0x3e142c14, 0x3e1c0404, 0x3e1c0c2c, - 0x3e1c1c1c, 0x3e1c3404, 0x3e24140c, 0x3e24240c, 0x3e2c0404, 0x3e2c0414, 0x3e2c1424, 0x3e341c04, -]; - -pub(crate) const IQ3S_GRID: [u32; 512] = [ - 0x01010101, 0x01010103, 0x01010105, 0x0101010b, 0x0101010f, 0x01010301, 0x01010303, 0x01010305, - 0x01010309, 0x0101030d, 0x01010501, 0x01010503, 0x0101050b, 0x01010707, 0x01010901, 0x01010905, - 0x0101090b, 0x0101090f, 0x01010b03, 0x01010b07, 0x01010d01, 0x01010d05, 0x01010f03, 0x01010f09, - 0x01010f0f, 0x01030101, 0x01030103, 0x01030105, 0x01030109, 0x01030301, 0x01030303, 0x0103030b, - 0x01030501, 0x01030507, 0x0103050f, 0x01030703, 0x0103070b, 0x01030909, 0x01030d03, 0x01030d0b, - 0x01030f05, 0x01050101, 0x01050103, 0x0105010b, 0x0105010f, 0x01050301, 0x01050307, 0x0105030d, - 0x01050503, 0x0105050b, 0x01050701, 0x01050709, 0x01050905, 0x0105090b, 0x0105090f, 0x01050b03, - 0x01050b07, 0x01050f01, 0x01050f07, 0x01070107, 0x01070303, 0x0107030b, 0x01070501, 0x01070505, - 0x01070703, 0x01070707, 0x0107070d, 0x01070909, 0x01070b01, 0x01070b05, 0x01070d0f, 0x01070f03, - 0x01070f0b, 0x01090101, 0x01090307, 0x0109030f, 0x01090503, 0x01090509, 0x01090705, 0x01090901, - 0x01090907, 0x01090b03, 0x01090f01, 0x010b0105, 0x010b0109, 0x010b0501, 0x010b0505, 0x010b050d, - 0x010b0707, 0x010b0903, 0x010b090b, 0x010b090f, 0x010b0d0d, 0x010b0f07, 0x010d010d, 0x010d0303, - 0x010d0307, 0x010d0703, 0x010d0b05, 0x010d0f03, 0x010f0101, 0x010f0105, 0x010f0109, 0x010f0501, - 0x010f0505, 0x010f050d, 0x010f0707, 0x010f0b01, 0x010f0b09, 0x03010101, 0x03010103, 0x03010105, - 0x03010109, 0x03010301, 0x03010303, 0x03010307, 0x0301030b, 0x0301030f, 0x03010501, 0x03010505, - 0x03010703, 0x03010709, 0x0301070d, 0x03010b09, 0x03010b0d, 0x03010d03, 0x03010f05, 0x03030101, - 0x03030103, 0x03030107, 0x0303010d, 0x03030301, 0x03030309, 0x03030503, 0x03030701, 0x03030707, - 0x03030903, 0x03030b01, 0x03030b05, 0x03030f01, 0x03030f0d, 0x03050101, 0x03050305, 0x0305030b, - 0x0305030f, 0x03050501, 0x03050509, 0x03050705, 0x03050901, 0x03050907, 0x03050b0b, 0x03050d01, - 0x03050f05, 0x03070103, 0x03070109, 0x0307010f, 0x03070301, 0x03070307, 0x03070503, 0x0307050f, - 0x03070701, 0x03070709, 0x03070903, 0x03070d05, 0x03070f01, 0x03090107, 0x0309010b, 0x03090305, - 0x03090309, 0x03090703, 0x03090707, 0x03090905, 0x0309090d, 0x03090b01, 0x03090b09, 0x030b0103, - 0x030b0301, 0x030b0307, 0x030b0503, 0x030b0701, 0x030b0705, 0x030b0b03, 0x030d0501, 0x030d0509, - 0x030d050f, 0x030d0909, 0x030d090d, 0x030f0103, 0x030f0107, 0x030f0301, 0x030f0305, 0x030f0503, - 0x030f070b, 0x030f0903, 0x030f0d05, 0x030f0f01, 0x05010101, 0x05010103, 0x05010107, 0x0501010b, - 0x0501010f, 0x05010301, 0x05010305, 0x05010309, 0x0501030d, 0x05010503, 0x05010507, 0x0501050f, - 0x05010701, 0x05010705, 0x05010903, 0x05010907, 0x0501090b, 0x05010b01, 0x05010b05, 0x05010d0f, - 0x05010f01, 0x05010f07, 0x05010f0b, 0x05030101, 0x05030105, 0x05030301, 0x05030307, 0x0503030f, - 0x05030505, 0x0503050b, 0x05030703, 0x05030709, 0x05030905, 0x05030b03, 0x05050103, 0x05050109, - 0x0505010f, 0x05050503, 0x05050507, 0x05050701, 0x0505070f, 0x05050903, 0x05050b07, 0x05050b0f, - 0x05050f03, 0x05050f09, 0x05070101, 0x05070105, 0x0507010b, 0x05070303, 0x05070505, 0x05070509, - 0x05070703, 0x05070707, 0x05070905, 0x05070b01, 0x05070d0d, 0x05090103, 0x0509010f, 0x05090501, - 0x05090507, 0x05090705, 0x0509070b, 0x05090903, 0x05090f05, 0x05090f0b, 0x050b0109, 0x050b0303, - 0x050b0505, 0x050b070f, 0x050b0901, 0x050b0b07, 0x050b0f01, 0x050d0101, 0x050d0105, 0x050d010f, - 0x050d0503, 0x050d0b0b, 0x050d0d03, 0x050f010b, 0x050f0303, 0x050f050d, 0x050f0701, 0x050f0907, - 0x050f0b01, 0x07010105, 0x07010303, 0x07010307, 0x0701030b, 0x0701030f, 0x07010505, 0x07010703, - 0x07010707, 0x0701070b, 0x07010905, 0x07010909, 0x0701090f, 0x07010b03, 0x07010d07, 0x07010f03, - 0x07030103, 0x07030107, 0x0703010b, 0x07030309, 0x07030503, 0x07030507, 0x07030901, 0x07030d01, - 0x07030f05, 0x07030f0d, 0x07050101, 0x07050305, 0x07050501, 0x07050705, 0x07050709, 0x07050b01, - 0x07070103, 0x07070301, 0x07070309, 0x07070503, 0x07070507, 0x0707050f, 0x07070701, 0x07070903, - 0x07070907, 0x0707090f, 0x07070b0b, 0x07070f07, 0x07090107, 0x07090303, 0x0709030d, 0x07090505, - 0x07090703, 0x07090b05, 0x07090d01, 0x07090d09, 0x070b0103, 0x070b0301, 0x070b0305, 0x070b050b, - 0x070b0705, 0x070b0909, 0x070b0b0d, 0x070b0f07, 0x070d030d, 0x070d0903, 0x070f0103, 0x070f0107, - 0x070f0501, 0x070f0505, 0x070f070b, 0x09010101, 0x09010109, 0x09010305, 0x09010501, 0x09010509, - 0x0901050f, 0x09010705, 0x09010903, 0x09010b01, 0x09010f01, 0x09030105, 0x0903010f, 0x09030303, - 0x09030307, 0x09030505, 0x09030701, 0x0903070b, 0x09030907, 0x09030b03, 0x09030b0b, 0x09050103, - 0x09050107, 0x09050301, 0x0905030b, 0x09050503, 0x09050707, 0x09050901, 0x09050b0f, 0x09050d05, - 0x09050f01, 0x09070109, 0x09070303, 0x09070307, 0x09070501, 0x09070505, 0x09070703, 0x0907070b, - 0x09090101, 0x09090105, 0x09090509, 0x0909070f, 0x09090901, 0x09090f03, 0x090b010b, 0x090b010f, - 0x090b0503, 0x090b0d05, 0x090d0307, 0x090d0709, 0x090d0d01, 0x090f0301, 0x090f030b, 0x090f0701, - 0x090f0907, 0x090f0b03, 0x0b010105, 0x0b010301, 0x0b010309, 0x0b010505, 0x0b010901, 0x0b010909, - 0x0b01090f, 0x0b010b05, 0x0b010d0d, 0x0b010f09, 0x0b030103, 0x0b030107, 0x0b03010b, 0x0b030305, - 0x0b030503, 0x0b030705, 0x0b030f05, 0x0b050101, 0x0b050303, 0x0b050507, 0x0b050701, 0x0b05070d, - 0x0b050b07, 0x0b070105, 0x0b07010f, 0x0b070301, 0x0b07050f, 0x0b070909, 0x0b070b03, 0x0b070d0b, - 0x0b070f07, 0x0b090103, 0x0b090109, 0x0b090501, 0x0b090705, 0x0b09090d, 0x0b0b0305, 0x0b0b050d, - 0x0b0b0b03, 0x0b0b0b07, 0x0b0d0905, 0x0b0f0105, 0x0b0f0109, 0x0b0f0505, 0x0d010303, 0x0d010307, - 0x0d01030b, 0x0d010703, 0x0d010707, 0x0d010d01, 0x0d030101, 0x0d030501, 0x0d03050f, 0x0d030d09, - 0x0d050305, 0x0d050709, 0x0d050905, 0x0d050b0b, 0x0d050d05, 0x0d050f01, 0x0d070101, 0x0d070309, - 0x0d070503, 0x0d070901, 0x0d09050b, 0x0d090907, 0x0d090d05, 0x0d0b0101, 0x0d0b0107, 0x0d0b0709, - 0x0d0b0d01, 0x0d0d010b, 0x0d0d0901, 0x0d0f0303, 0x0d0f0307, 0x0f010101, 0x0f010109, 0x0f01010f, - 0x0f010501, 0x0f010505, 0x0f01070d, 0x0f010901, 0x0f010b09, 0x0f010d05, 0x0f030105, 0x0f030303, - 0x0f030509, 0x0f030907, 0x0f03090b, 0x0f050103, 0x0f050109, 0x0f050301, 0x0f05030d, 0x0f050503, - 0x0f050701, 0x0f050b03, 0x0f070105, 0x0f070705, 0x0f07070b, 0x0f070b07, 0x0f090103, 0x0f09010b, - 0x0f090307, 0x0f090501, 0x0f090b01, 0x0f0b0505, 0x0f0b0905, 0x0f0d0105, 0x0f0d0703, 0x0f0f0101, -]; - diff --git a/crates/mummu/examples/src/hub.rs b/crates/mummu/examples/src/hub.rs deleted file mode 100644 index 712bb6f..0000000 --- a/crates/mummu/examples/src/hub.rs +++ /dev/null @@ -1,647 +0,0 @@ -//! Model downloads: HuggingFace Hub (or any HTTP host) → the local model -//! cache. Streaming, **resumable** (a `.part` picks up where a killed -//! download stopped, via HTTP `Range`), **integrity-checked** (streamed -//! sha256 against the Hub's announced LFS `X-Linked-ETag`, length as the -//! fallback), and **sharded-checkpoint aware** (`model.safetensors.index.json` -//! → fetch every shard). Progress surfaces through a callback so app settings -//! UIs can show it (P8). Completed files are cache-first: an existing -//! destination is never re-fetched unless [`FetchOptions::verify_cached`] -//! asks for a re-hash. - -use std::io::{Read, Write}; -use std::path::{Path, PathBuf}; - -use sha2::{Digest, Sha256}; - -/// Streaming copy granularity: big enough to amortize syscalls, small enough -/// to keep progress callbacks responsive. -const CHUNK_BYTES: usize = 64 * 1024; - -/// Hard per-file ceiling — larger than any model shard we'd fetch (shards are -/// conventionally ≤ ~10 GB); anything past this is a wiring bug or a hostile -/// server, not a model. -const MAX_FILE_BYTES: u64 = 64 << 30; - -/// Ceiling on shards in an index — the largest public checkpoints ship tens. -const MAX_SHARDS: usize = 512; - -/// Everything that can go wrong fetching a model. -#[derive(Debug, thiserror::Error)] -pub enum HubError { - #[error("http {url}: {reason}")] - Http { url: String, reason: String }, - #[error("io {path}: {reason}")] - Io { path: PathBuf, reason: String }, - #[error("{url}: expected {expected} bytes, received {received}")] - Incomplete { - url: String, - expected: u64, - received: u64, - }, - #[error("shard index {path}: {reason}")] - BadIndex { path: PathBuf, reason: String }, - #[error("{url}: sha256 mismatch — announced {expected}, computed {computed}")] - Corrupt { - url: String, - expected: String, - computed: String, - }, -} - -/// Options for [`fetch_file_with`] / [`fetch_model_with`]. -#[derive(Debug, Clone, Copy, Default)] -pub struct FetchOptions { - /// Re-hash an already-complete destination against the server's announced - /// sha256 (one extra HEAD request per file); on a mismatch the corrupt - /// copy is deleted and re-fetched once. Off by default: a completed file - /// was already verified as it streamed in. - pub verify_cached: bool, -} - -/// Download progress for one file, reported after every chunk. -#[derive(Debug, Clone)] -pub struct Progress<'a> { - pub file: &'a str, - pub received_bytes: u64, - /// Total including any resumed prefix; `None` when the server omits it. - pub total_bytes: Option, -} - -/// `https://huggingface.co/{repo}/resolve/{revision}/{file}` — the Hub's -/// stable raw-file endpoint. -#[must_use] -pub fn hub_file_url(repo: &str, revision: &str, file: &str) -> String { - assert!( - !repo.is_empty() && repo.contains('/'), - "repo must be owner/name, got {repo:?}" - ); - assert!(!revision.is_empty(), "revision must be non-empty"); - format!("https://huggingface.co/{repo}/resolve/{revision}/{file}") -} - -/// The unique shard files referenced by a `*.index.json` (weight_map values, -/// deduped, sorted for a deterministic fetch order). -pub fn shards_from_index(index_json: &[u8], index_path: &Path) -> Result, HubError> { - let v: serde_json::Value = - serde_json::from_slice(index_json).map_err(|e| HubError::BadIndex { - path: index_path.to_path_buf(), - reason: e.to_string(), - })?; - let map = v["weight_map"] - .as_object() - .ok_or_else(|| HubError::BadIndex { - path: index_path.to_path_buf(), - reason: "no weight_map object".into(), - })?; - let mut shards: Vec = map - .values() - .filter_map(|s| s.as_str().map(str::to_string)) - .collect(); - shards.sort_unstable(); - shards.dedup(); - if shards.is_empty() || shards.len() > MAX_SHARDS { - return Err(HubError::BadIndex { - path: index_path.to_path_buf(), - reason: format!("{} shards (expected 1..={MAX_SHARDS})", shards.len()), - }); - } - Ok(shards) -} - -/// The in-flight twin of `dest` (`.part`). -fn part_path(dest: &Path) -> PathBuf { - let mut p = dest.as_os_str().to_owned(); - p.push(".part"); - PathBuf::from(p) -} - -/// Lowercase hex of a digest. -fn hex64(bytes: &[u8]) -> String { - use std::fmt::Write as _; - debug_assert_eq!(bytes.len(), 32, "sha256 digests are 32 bytes"); - let mut s = String::with_capacity(bytes.len() * 2); - for b in bytes { - let _ = write!(s, "{b:02x}"); - } - s -} - -/// A quoted 64-hex etag value is a content sha256 (the Hub's `X-Linked-ETag` -/// for LFS files). Git-style etags (40-hex sha1, or non-hex) parse to `None` — -/// they name a revision, not the bytes, and cannot verify a stream. -fn parse_sha256_etag(raw: &str) -> Option { - let v = raw.trim().trim_start_matches("W/").trim_matches('"'); - let is_sha256 = v.len() == 64 && v.bytes().all(|b| b.is_ascii_hexdigit()); - is_sha256.then(|| v.to_ascii_lowercase()) -} - -/// Ask the server for the file's content sha256: a redirect-stopped HEAD reads -/// the Hub's `X-Linked-ETag` from the `resolve/` endpoint itself, before the -/// CDN handoff would replace the headers. `Ok(None)` when nothing usable is -/// announced (non-LFS files, other hosts, HEAD rejected) — those downloads -/// stay length-verified only. Transport failures are loud: the GET would fail -/// the same way. -fn announced_sha256(url: &str) -> Result, HubError> { - assert!(url.starts_with("https://"), "refusing non-https url: {url}"); - let agent: ureq::Agent = ureq::Agent::config_builder() - .max_redirects(0) - .http_status_as_error(false) - .build() - .into(); - let resp = agent.head(url).call().map_err(|e| HubError::Http { - url: url.into(), - reason: e.to_string(), - })?; - Ok(["x-linked-etag", "etag"].iter().find_map(|h| { - resp.headers() - .get(*h) - .and_then(|v| v.to_str().ok()) - .and_then(parse_sha256_etag) - })) -} - -/// The sibling files that are fetched **only if the repo ships them**. -/// -/// `tokenizer_config.json` is what the import-validation gates read (EOS -/// agreement, added-token ids, tool-call convention); `chat_template.jinja` is -/// how checkpoints that keep their template out of `tokenizer_config.json` -/// ship it. Neither is universal on the Hub, so a 404 is a legitimate answer — -/// but not asking for them at all silently disarms those gates for every model -/// installed through this path. -const OPTIONAL_FILES: [&str; 2] = ["tokenizer_config.json", "chat_template.jinja"]; - -/// Does the repo actually ship `url`? -/// -/// One HEAD, redirects not followed: the Hub answers a `resolve` URL for a -/// present file with 200 (or a 302 to the CDN) and for an absent one with 404. -/// ONLY 404 counts as absent — anything else (403 on a gated repo, a 5xx) is -/// reported as present so the real fetch raises it properly rather than this -/// probe swallowing it as "the repo just doesn't have it". -fn repo_has_file(url: &str) -> Result { - assert!(url.starts_with("https://"), "refusing non-https url: {url}"); - let agent: ureq::Agent = ureq::Agent::config_builder() - .max_redirects(0) - .http_status_as_error(false) - .build() - .into(); - let resp = agent.head(url).call().map_err(|e| HubError::Http { - url: url.into(), - reason: e.to_string(), - })?; - let status = resp.status().as_u16(); - debug_assert!( - (100..600).contains(&status), - "http status in range: {status}" - ); - Ok(status != 404) -} - -/// Streaming sha256 of a file on disk, as lowercase hex. -fn sha256_hex_of_file(path: &Path) -> Result { - let io_err = |e: std::io::Error| HubError::Io { - path: path.to_path_buf(), - reason: e.to_string(), - }; - let mut f = std::fs::File::open(path).map_err(io_err)?; - let mut hasher = Sha256::new(); - let mut buf = vec![0u8; CHUNK_BYTES]; - let mut hashed = 0u64; - loop { - let n = f.read(&mut buf).map_err(io_err)?; - if n == 0 { - break; - } - hasher.update(&buf[..n]); - hashed += n as u64; - assert!( - hashed <= MAX_FILE_BYTES, - "{path:?}: exceeds the file bound while hashing" - ); - } - Ok(hex64(&hasher.finalize())) -} - -/// Feed the already-downloaded `.part` prefix into the stream hasher so a -/// resumed download still verifies as one whole file. -fn hash_part_prefix(part: &Path, resume_from: u64, hasher: &mut Sha256) -> Result<(), HubError> { - assert!(resume_from > 0, "no prefix to hash"); - let mut f = std::fs::File::open(part).map_err(|e| HubError::Io { - path: part.to_path_buf(), - reason: e.to_string(), - })?; - let mut remaining = resume_from; - let mut buf = vec![0u8; CHUNK_BYTES]; - while remaining > 0 { - let want = remaining.min(CHUNK_BYTES as u64) as usize; - let n = f.read(&mut buf[..want]).map_err(|e| HubError::Io { - path: part.to_path_buf(), - reason: e.to_string(), - })?; - // The prefix length came from this file's own metadata an instant ago. - assert!(n > 0, "{part:?}: prefix ended {remaining} bytes early"); - hasher.update(&buf[..n]); - remaining -= n as u64; - } - Ok(()) -} - -/// Fetch `url` into `dest`, streaming through `.part` and resuming any -/// earlier partial download. No-op when `dest` already exists (cache-first). -/// `on_progress` fires after every chunk with cumulative counts. The stream -/// is verified against the server's announced sha256 when there is one -/// (Hub LFS files), else by length. -pub fn fetch_file( - url: &str, - dest: &Path, - on_progress: impl FnMut(Progress<'_>), -) -> Result<(), HubError> { - fetch_file_with(url, dest, FetchOptions::default(), on_progress) -} - -/// [`fetch_file`] with explicit [`FetchOptions`]. With `verify_cached`, an -/// existing `dest` is re-hashed against the announced sha256; a mismatch -/// deletes the corrupt copy (and any stale `.part` that would poison a -/// resume) and re-fetches once — self-healing, never silent. -pub fn fetch_file_with( - url: &str, - dest: &Path, - opts: FetchOptions, - mut on_progress: impl FnMut(Progress<'_>), -) -> Result<(), HubError> { - assert!(url.starts_with("https://"), "refusing non-https url: {url}"); - if dest.exists() { - if !opts.verify_cached { - return Ok(()); // cache hit — never re-fetch a completed file - } - let Some(expected) = announced_sha256(url)? else { - return Ok(()); // nothing announced — nothing to re-verify against - }; - if sha256_hex_of_file(dest)? == expected { - return Ok(()); - } - for stale in [dest.to_path_buf(), part_path(dest)] { - if stale.exists() { - std::fs::remove_file(&stale).map_err(|e| HubError::Io { - path: stale.clone(), - reason: e.to_string(), - })?; - } - } - } - download(url, dest, &mut on_progress) -} - -/// The streaming GET behind [`fetch_file_with`]: resume, hash, length-check, -/// then atomically rename `.part` → `dest`. -fn download( - url: &str, - dest: &Path, - on_progress: &mut impl FnMut(Progress<'_>), -) -> Result<(), HubError> { - debug_assert!(!dest.exists(), "download() requires a vacant destination"); - let file_label = dest - .file_name() - .map(|n| n.to_string_lossy().into_owned()) - .unwrap_or_default(); - assert!(!file_label.is_empty(), "dest must name a file: {dest:?}"); - if let Some(parent) = dest.parent() { - std::fs::create_dir_all(parent).map_err(|e| HubError::Io { - path: parent.to_path_buf(), - reason: e.to_string(), - })?; - } - - let part = part_path(dest); - let resume_from = std::fs::metadata(&part).map(|m| m.len()).unwrap_or(0); - // One cheap HEAD up front: with an announced sha256 the whole stream - // (resumed prefix included) is verified; without one, length still is. - let expected_sha = announced_sha256(url)?; - - let mut req = ureq::get(url); - if resume_from > 0 { - req = req.header("Range", format!("bytes={resume_from}-")); - } - let mut resp = req.call().map_err(|e| HubError::Http { - url: url.into(), - reason: e.to_string(), - })?; - // A server that ignores Range (200 instead of 206) restarts the body from - // byte 0 — truncate our part file to match, never splice mismatched halves. - let resumed = resp.status() == 206 && resume_from > 0; - let body_len: Option = resp - .headers() - .get("content-length") - .and_then(|v| v.to_str().ok()) - .and_then(|s| s.parse().ok()); - let already = if resumed { resume_from } else { 0 }; - let total = body_len.map(|l| l + already); - if let Some(t) = total { - assert!( - t <= MAX_FILE_BYTES, - "{url}: {t} bytes exceeds the file bound" - ); - } - - let mut hasher = expected_sha.as_ref().map(|_| Sha256::new()); - if resumed && let Some(h) = hasher.as_mut() { - hash_part_prefix(&part, resume_from, h)?; - } - - let mut out = std::fs::OpenOptions::new() - .create(true) - .append(resumed) - .write(true) - .truncate(!resumed) - .open(&part) - .map_err(|e| HubError::Io { - path: part.clone(), - reason: e.to_string(), - })?; - - let mut reader = resp.body_mut().as_reader(); - let mut received = already; - let mut buf = vec![0u8; CHUNK_BYTES]; - loop { - let n = reader.read(&mut buf).map_err(|e| HubError::Http { - url: url.into(), - reason: e.to_string(), - })?; - if n == 0 { - break; - } - out.write_all(&buf[..n]).map_err(|e| HubError::Io { - path: part.clone(), - reason: e.to_string(), - })?; - if let Some(h) = hasher.as_mut() { - h.update(&buf[..n]); - } - received += n as u64; - assert!( - received <= MAX_FILE_BYTES, - "{url}: stream exceeded the file bound" - ); - on_progress(Progress { - file: &file_label, - received_bytes: received, - total_bytes: total, - }); - } - drop(out); - - if let Some(expected) = total - && received != expected - { - // Keep the .part for a future resume; report loudly. - return Err(HubError::Incomplete { - url: url.into(), - expected, - received, - }); - } - if let (Some(expected), Some(h)) = (expected_sha, hasher) { - let computed = hex64(&h.finalize()); - if computed != expected { - // A wrong-hash .part must never seed a resume — drop it. - std::fs::remove_file(&part).map_err(|e| HubError::Io { - path: part.clone(), - reason: e.to_string(), - })?; - return Err(HubError::Corrupt { - url: url.into(), - expected, - computed, - }); - } - } - std::fs::rename(&part, dest).map_err(|e| HubError::Io { - path: dest.to_path_buf(), - reason: e.to_string(), - }) -} - -/// Fetch a whole model from the Hub into `dest_dir`: `config.json`, -/// `tokenizer.json`, the weights — `model.safetensors` when the repo is -/// single-file, else every shard listed by `model.safetensors.index.json` — -/// and, when the repo ships them, the optional siblings in [`OPTIONAL_FILES`]. -/// Returns `dest_dir` ready for the per-model `load_from_dir`. -/// -/// `config.json` and `tokenizer.json` are required: a 404 on either is an -/// error. The optional siblings are best-effort, because a repo that does not -/// ship them is normal — but they are asked for, so a checkpoint installed -/// this way arrives with the files the import-validation gates need. -pub fn fetch_model( - repo: &str, - revision: &str, - dest_dir: &Path, - on_progress: impl FnMut(Progress<'_>), -) -> Result { - fetch_model_with( - repo, - revision, - dest_dir, - FetchOptions::default(), - on_progress, - ) -} - -/// [`fetch_model`] with explicit [`FetchOptions`] (e.g. re-verify cached -/// files' sha256 before trusting them). -pub fn fetch_model_with( - repo: &str, - revision: &str, - dest_dir: &Path, - opts: FetchOptions, - mut on_progress: impl FnMut(Progress<'_>), -) -> Result { - for file in ["config.json", "tokenizer.json"] { - fetch_file_with( - &hub_file_url(repo, revision, file), - &dest_dir.join(file), - opts, - &mut on_progress, - )?; - } - // Then the optional siblings. These MUST be fetched before the weights, - // because the single-file branch below returns early on success — putting - // them after it would fetch them for sharded checkpoints only. - for file in OPTIONAL_FILES { - let url = hub_file_url(repo, revision, file); - if repo_has_file(&url)? { - fetch_file_with(&url, &dest_dir.join(file), opts, &mut on_progress)?; - } - } - // Single-file first (the common case for the small-model tiers we target). - let single = fetch_file_with( - &hub_file_url(repo, revision, "model.safetensors"), - &dest_dir.join("model.safetensors"), - opts, - &mut on_progress, - ); - if single.is_ok() { - return Ok(dest_dir.to_path_buf()); - } - // Fall back to a sharded checkpoint; if there's no index either, report - // the original single-file error (the more useful signal). - let index_name = "model.safetensors.index.json"; - let index_dest = dest_dir.join(index_name); - if fetch_file_with( - &hub_file_url(repo, revision, index_name), - &index_dest, - opts, - &mut on_progress, - ) - .is_err() - { - return single.map(|()| dest_dir.to_path_buf()); - } - let index_bytes = std::fs::read(&index_dest).map_err(|e| HubError::Io { - path: index_dest.clone(), - reason: e.to_string(), - })?; - for shard in shards_from_index(&index_bytes, &index_dest)? { - fetch_file_with( - &hub_file_url(repo, revision, &shard), - &dest_dir.join(&shard), - opts, - &mut on_progress, - )?; - } - Ok(dest_dir.to_path_buf()) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn hub_url_has_the_resolve_shape() { - assert_eq!( - hub_file_url("Qwen/Qwen2.5-1.5B-Instruct", "main", "config.json"), - "https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct/resolve/main/config.json" - ); - } - - #[test] - #[should_panic(expected = "owner/name")] - fn bare_repo_names_are_rejected() { - let _ = hub_file_url("qwen", "main", "config.json"); - } - - /// The gates that read these files fail OPEN — `validate_checkpoint_dir` - /// returns `Ok(None)` when `tokenizer_config.json` is absent, so the - /// EOS-agreement and added-token-id checks simply do not run. That makes - /// "did we even ask the Hub for it?" the load-bearing question, and it is - /// what this list answers. - #[test] - fn optional_files_cover_the_siblings_the_import_gates_read() { - assert!(OPTIONAL_FILES.contains(&"tokenizer_config.json")); - assert!(OPTIONAL_FILES.contains(&"chat_template.jinja")); - for file in OPTIONAL_FILES { - assert!( - !["config.json", "tokenizer.json"].contains(&file), - "{file} is required, not optional — a 404 on it must stay an error" - ); - assert_eq!( - hub_file_url("allenai/OLMoE-1B-7B-0125-Instruct", "main", file), - format!( - "https://huggingface.co/allenai/OLMoE-1B-7B-0125-Instruct/resolve/main/{file}" - ) - ); - } - } - - #[test] - fn shard_index_dedupes_and_sorts() { - let idx = br#"{"metadata":{},"weight_map":{ - "a.weight":"model-00002-of-00002.safetensors", - "b.weight":"model-00001-of-00002.safetensors", - "c.weight":"model-00001-of-00002.safetensors"}}"#; - let shards = shards_from_index(idx, Path::new("x.index.json")).unwrap(); - assert_eq!( - shards, - vec![ - "model-00001-of-00002.safetensors", - "model-00002-of-00002.safetensors" - ] - ); - } - - #[test] - fn shard_index_without_weight_map_is_rejected() { - let err = shards_from_index(b"{}", Path::new("x.index.json")); - assert!(matches!(err, Err(HubError::BadIndex { .. }))); - } - - #[test] - fn sha256_etag_parsing_accepts_only_content_hashes() { - let sha = "a".repeat(64); - // The Hub quotes LFS etags; weak etags carry a W/ prefix. - assert_eq!(parse_sha256_etag(&format!("\"{sha}\"")), Some(sha.clone())); - assert_eq!( - parse_sha256_etag(&format!("W/\"{sha}\"")), - Some(sha.clone()) - ); - assert_eq!(parse_sha256_etag(&sha.to_uppercase()), Some(sha)); - // Git-style 40-hex sha1, non-hex, and empty values are not sha256s. - assert_eq!(parse_sha256_etag(&format!("\"{}\"", "b".repeat(40))), None); - assert_eq!(parse_sha256_etag(&format!("\"{}\"", "z".repeat(64))), None); - assert_eq!(parse_sha256_etag(""), None); - } - - #[test] - fn file_hash_matches_the_reference_vector() { - // FIPS 180-2 test vector: sha256("abc"). - let dir = std::env::temp_dir().join("mummu-hub-test-sha"); - std::fs::create_dir_all(&dir).unwrap(); - let p = dir.join("abc.txt"); - std::fs::write(&p, b"abc").unwrap(); - assert_eq!( - sha256_hex_of_file(&p).unwrap(), - "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad" - ); - std::fs::remove_dir_all(&dir).ok(); - } - - #[test] - fn prefix_plus_remainder_hash_equals_whole_file_hash() { - // The resume path hashes the .part prefix, then the streamed tail; - // together they must equal one pass over the whole file. - let dir = std::env::temp_dir().join("mummu-hub-test-prefix"); - std::fs::create_dir_all(&dir).unwrap(); - let whole: Vec = (0u32..200_000).map(|i| (i % 251) as u8).collect(); - let p = dir.join("whole.bin"); - std::fs::write(&p, &whole).unwrap(); - let reference = sha256_hex_of_file(&p).unwrap(); - - let split = whole.len() / 3; - let part = dir.join("whole.bin.part"); - std::fs::write(&part, &whole[..split]).unwrap(); - let mut h = Sha256::new(); - hash_part_prefix(&part, split as u64, &mut h).unwrap(); - h.update(&whole[split..]); - assert_eq!(hex64(&h.finalize()), reference); - std::fs::remove_dir_all(&dir).ok(); - } - - #[test] - fn part_path_appends_suffix() { - assert_eq!( - part_path(Path::new("m/model.safetensors")), - Path::new("m/model.safetensors.part") - ); - } - - #[test] - fn existing_dest_is_a_cache_hit_without_any_http() { - // A bogus URL proves no request is made when the file already exists. - let dir = std::env::temp_dir().join("mummu-hub-test-cache-hit"); - std::fs::create_dir_all(&dir).unwrap(); - let dest = dir.join("present.bin"); - std::fs::write(&dest, b"already here").unwrap(); - let mut calls = 0; - fetch_file("https://invalid.invalid/x", &dest, |_| calls += 1).unwrap(); - assert_eq!(calls, 0, "cache hit must not stream"); - std::fs::remove_dir_all(&dir).ok(); - } -} diff --git a/crates/mummu/examples/src/import.rs b/crates/mummu/examples/src/import.rs deleted file mode 100644 index a3e4821..0000000 --- a/crates/mummu/examples/src/import.rs +++ /dev/null @@ -1,574 +0,0 @@ -//! Weight import: checkpoint files → Burn modules, checked and loud. -//! -//! The pieces every model load shares (P3): a dtype-cast adapter (HF ships -//! bf16, which wgpu can't ingest directly), a weights-file picker -//! (**safetensors** preferred, the PyTorch state dict `pytorch_model.bin` as -//! the fallback for models never re-shipped as safetensors), and a -//! checked-load wrapper that **fails on missing or errored params** instead -//! of silently zero-initing — a partial load is a quietly broken model. - -use std::path::{Path, PathBuf}; -use std::rc::Rc; - -use burn::module::Module; -use burn::store::{ - ModuleAdapter, ModuleSnapshot, ModuleStore, PyTorchToBurnAdapter, SafetensorsStore, - TensorSnapshot, -}; -use burn::tensor::DType; - -use crate::gguf::{GgufFile, GgufMap, GgufTensorInfo}; - -/// Everything that can go wrong turning files on disk into a loaded model. -#[derive(Debug, thiserror::Error)] -pub enum ImportError { - #[error("required file missing: {0}")] - MissingFile(PathBuf), - #[error("parse {file}: {reason}")] - Parse { file: PathBuf, reason: String }, - #[error("load weights ({file}): {reason}")] - Load { file: PathBuf, reason: String }, - #[error( - "weight load incomplete ({file}): {applied} applied, {missing} missing, {errors} errors\n{report}" - )] - Incomplete { - file: PathBuf, - applied: usize, - missing: usize, - errors: usize, - report: String, - }, - /// A checkpoint's own metadata files disagree with each other (e.g. - /// `tokenizer_config.json` names an EOS `config.json` does not, or a - /// chat-template whose tool-call convention is not the one this family's - /// byte-verified renderer speaks). A repackaging bug that a checked *weight* - /// load cannot see; surfaced loudly at load rather than mis-stopping or - /// mis-templating at generate time. - #[error("inconsistent metadata ({file}): {reason}")] - Inconsistent { file: PathBuf, reason: String }, -} - -/// Why a freshly-imported model failed its post-load **sanity smoke** (one -/// forward, checked for liveness). These are the silent-broken-import failure -/// modes a checked *load* cannot see — the weights all applied, but the model -/// does not actually compute. -#[derive(Debug, thiserror::Error, PartialEq)] -pub enum SanityError { - /// Logits contain NaN or ±Inf — the classic signature of a wrong dtype - /// (e.g. an f16 overflow that should have been an f32 island) or corrupt - /// weight bytes that still deserialized to the right shape. - #[error( - "{count} non-finite logit(s) (first at index {first_index}) — bad dtype or corrupt weights" - )] - NonFinite { count: usize, first_index: usize }, - /// The logits width does not equal the expected vocabulary — a config / - /// tokenizer / checkpoint mismatch (the model and its tokenizer disagree - /// on how many tokens exist). - #[error("logits width {logits_len} != expected vocab {expected} — config/tokenizer mismatch")] - WrongVocab { logits_len: usize, expected: usize }, - /// Every logit is (near-)identical — a dead forward: zero-initialized - /// weights, or an all-masked/zeroed activation path that a checked load - /// still reports as fully applied. - #[error("degenerate logits (spread {spread:e} < {threshold:e}) — dead/zero-init forward")] - Degenerate { spread: f32, threshold: f32 }, -} - -/// The healthy result of a [`logit_sanity`] smoke: the winning token and how -/// sharply the distribution favors it (a coarse confidence signal for logs). -#[derive(Debug, Clone, Copy, PartialEq)] -pub struct SanitySmoke { - /// argmax token id. - pub top_id: u32, - /// The winning logit. - pub top_logit: f32, - /// max − min over the logits (the dynamic range; large for a live model). - pub spread: f32, -} - -/// Logits below this spread (max − min) are treated as a dead/uniform forward. -/// A live LM's first-token logits span tens; a zero-init forward spans ~0. The -/// gap is enormous, so a tiny threshold never false-positives a real model. -const DEGENERATE_SPREAD: f32 = 1e-4; - -/// Post-load **sanity smoke** over one forward's logits: finite, the expected -/// width, and not degenerate. This is *liveness*, not parity — an arbitrary -/// user import has no reference to compare against (catalog models get the P7 -/// parity gates); this only proves the model actually computes rather than -/// silently returning garbage a checked load reported as fully applied. -pub fn logit_sanity(logits: &[f32], expected_vocab: usize) -> Result { - assert!( - expected_vocab > 0, - "logit_sanity: expected_vocab must be positive" - ); - if logits.len() != expected_vocab { - return Err(SanityError::WrongVocab { - logits_len: logits.len(), - expected: expected_vocab, - }); - } - assert!( - !logits.is_empty(), - "logit_sanity: width matched a positive vocab" - ); - - if let Some((first_index, _)) = logits.iter().enumerate().find(|(_, l)| !l.is_finite()) { - let count = logits.iter().filter(|l| !l.is_finite()).count(); - return Err(SanityError::NonFinite { count, first_index }); - } - - // All finite: min/max define the spread and the argmax. - let (mut min, mut max, mut top_id) = (f32::INFINITY, f32::NEG_INFINITY, 0usize); - for (i, &l) in logits.iter().enumerate() { - if l < min { - min = l; - } - if l > max { - (max, top_id) = (l, i); - } - } - let spread = max - min; - if spread < DEGENERATE_SPREAD { - return Err(SanityError::Degenerate { - spread, - threshold: DEGENERATE_SPREAD, - }); - } - debug_assert!( - spread > 0.0 && max.is_finite(), - "healthy smoke has positive finite spread" - ); - Ok(SanitySmoke { - top_id: top_id as u32, - top_logit: max, - spread, - }) -} - -/// Cast bf16/f16/f64 float tensors to a target dtype on load. HF checkpoints -/// are commonly stored in **bf16**; `burn-store` keeps the source dtype, which -/// the wgpu backend can't ingest. Mirrors Burn's `HalfPrecisionAdapter` but -/// covers bf16 → f32/f16 too. Non-float tensors pass through untouched. -#[derive(Clone)] -pub struct CastFloatAdapter { - target: DType, -} - -impl CastFloatAdapter { - #[must_use] - pub fn new(target: DType) -> Self { - assert!( - matches!(target, DType::F16 | DType::F32 | DType::F64 | DType::BF16), - "CastFloatAdapter: target must be a float dtype, got {target:?}" - ); - Self { target } - } -} - -impl ModuleAdapter for CastFloatAdapter { - fn adapt(&self, snapshot: &TensorSnapshot) -> TensorSnapshot { - let is_float = matches!( - snapshot.dtype, - DType::BF16 | DType::F16 | DType::F32 | DType::F64 - ); - if !is_float || snapshot.dtype == self.target { - return snapshot.clone(); - } - let target = self.target; - let data_fn = snapshot.clone_data_fn(); - let cast = Rc::new(move || Ok(data_fn()?.convert_dtype(target))); - TensorSnapshot::from_closure( - cast, - target, - snapshot.shape.clone(), - snapshot.path_stack.clone().unwrap_or_default(), - snapshot.container_stack.clone().unwrap_or_default(), - snapshot.tensor_id.unwrap_or_default(), - ) - } - - fn clone_box(&self) -> Box { - Box::new(self.clone()) - } -} - -/// Load `store` (any format: safetensors, PyTorch state dict, …) into -/// `module`, refusing partial results: any missing param or per-tensor error -/// is an [`ImportError::Incomplete`] carrying the store's own readable -/// report. Unused checkpoint tensors are *allowed* (e.g. BERT's -/// intentionally-skipped `pooler.*`) — callers that care inspect the report. -pub fn load_checked( - module: &mut M, - store: &mut S, - weights_path: &Path, -) -> Result<(), ImportError> -where - M: Module + ModuleSnapshot, - S: ModuleStore, -{ - let report = module.load_from(store).map_err(|e| ImportError::Load { - file: weights_path.to_path_buf(), - reason: e.to_string(), - })?; - if !report.errors.is_empty() || !report.missing.is_empty() { - return Err(ImportError::Incomplete { - file: weights_path.to_path_buf(), - applied: report.applied.len(), - missing: report.missing.len(), - errors: report.errors.len(), - report: format!("{report}"), - }); - } - debug_assert!( - !report.applied.is_empty(), - "load_checked: a successful load must have applied at least one tensor" - ); - Ok(()) -} - -/// Largest f32 payload [`DequantSink::Auto`] will let a dequant hold in RAM. -/// -/// The in-memory sink's peak is ~2x the payload (the blob, plus the model -/// being built from it), so this ceiling is really a ~1 GiB peak — small -/// enough to be free on any machine that could run the model at all. Above -/// it, [`DequantSink::Scratch`] costs one f32 write plus an mmap read-back -/// and halves the peak; measured on Qwen2.5-1.5B (6.2 GB of f32), that trade -/// is inside the run-to-run noise of the load it is part of, so there is no -/// reason to keep paying the doubled peak for anything of real size. -pub const MAX_IN_MEMORY_DEQUANT_BYTES: u64 = 512 << 20; - -/// Where a GGUF dequant's f32 payload lands on its way into a -/// [`SafetensorsStore`]. -/// -/// The two sinks are byte-identical by construction — `gguf::dequant_into` is -/// ONE function behind both, pinned by a unit test — so this is a resource -/// trade, never a correctness one. [`Self::Memory`] peaks at ~2x the payload -/// and needs nothing; [`Self::Scratch`] peaks at ~1x (the payload becomes -/// file-backed mmap) and needs payload-sized free disk beside the weights. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum DequantSink { - /// Hold the whole f32 payload in RAM and hand it to the store as bytes. - Memory, - /// Stream it to a scratch file beside the weights; `burn-store` mmaps - /// that back and materializes tensors lazily, so the payload is never - /// charged to commit. - Scratch, - /// Pick by payload size: [`Self::Memory`] only up to - /// [`MAX_IN_MEMORY_DEQUANT_BYTES`], [`Self::Scratch`] above it. - Auto, -} - -impl DequantSink { - /// Resolve [`Self::Auto`] against a payload size; the other two answer - /// for themselves. Never returns `Auto`. - #[must_use] - pub fn resolve(self, total_f32_bytes: u64) -> Self { - let picked = match self { - Self::Auto if total_f32_bytes > MAX_IN_MEMORY_DEQUANT_BYTES => Self::Scratch, - Self::Auto => Self::Memory, - explicit => explicit, - }; - debug_assert!(picked != Self::Auto, "resolve must decide"); - debug_assert!( - picked != Self::Memory - || total_f32_bytes <= MAX_IN_MEMORY_DEQUANT_BYTES - || self == Self::Memory, - "Auto never picks Memory above the ceiling" - ); - picked - } -} - -/// Owns a scratch file for the life of one load. -/// -/// The file is as large as the weights it carries (~28 GB dequantizing -/// OLMoE-1B-7B's Q4_K_M), so leaving one behind on a failed load would -/// quietly fill the disk over a few retries. `Drop` removes it on every exit -/// path, success or `?` — and because the store reads it lazily, the guard -/// must outlive `load_checked`, which holding it as a local does. -#[derive(Debug)] -pub struct ScratchFile { - path: PathBuf, -} - -impl ScratchFile { - /// Placed beside the weights, so the scratch write lands on the same - /// volume as the model rather than on a small system temp drive. - /// - /// The name carries a process-unique counter as well as the pid: two - /// concurrent loads in one process must not choose the same file and - /// interleave their writes into it. - pub fn new(dir: &Path) -> Result { - use std::sync::atomic::{AtomicU64, Ordering}; - static NEXT: AtomicU64 = AtomicU64::new(0); - assert!(!dir.as_os_str().is_empty(), "scratch dir must be named"); - let path = dir.join(format!( - "mummu-scratch-{}-{}.safetensors.tmp", - std::process::id(), - NEXT.fetch_add(1, Ordering::Relaxed) - )); - // A leftover from a killed process must never be mistaken for ours. - if path.exists() { - std::fs::remove_file(&path).map_err(|e| ImportError::Parse { - file: path.clone(), - reason: format!("could not clear a stale scratch file: {e}"), - })?; - } - Ok(Self { path }) - } - - #[must_use] - pub fn path(&self) -> &Path { - debug_assert!(!self.path.as_os_str().is_empty(), "scratch path is named"); - &self.path - } -} - -impl Drop for ScratchFile { - fn drop(&mut self) { - // Best effort by construction: a Drop that can fail has nowhere to - // report to, and a stranded scratch file is not worth a panic. - let _ = std::fs::remove_file(&self.path); - } -} - -/// Dequantize `f` to f32 safetensors and return the store every GGUF loader -/// then adds its own key remaps to, plus the scratch guard (if any) that must -/// stay alive until the load finishes. -/// -/// This is the one place the four GGUF ports share: the sink choice, the -/// adapter chain, and — the part worth centralizing — taking the target float -/// dtype from `device`. Under burn 0.22 the float element type is a per-device -/// runtime setting rather than a backend type parameter, so the dtype a load -/// casts to is decided by the device the caller hands in. -pub fn gguf_store( - f: &GgufFile, - map: &dyn Fn(&GgufTensorInfo) -> Option, - sink: DequantSink, - device: &burn::tensor::Device, -) -> Result<(SafetensorsStore, Option), ImportError> { - let parse = |reason: String| ImportError::Parse { - file: f.path.clone(), - reason, - }; - let total_f32_bytes: u64 = f.tensors.iter().map(|t| t.element_count() * 4).sum(); - let target_float = crate::backend::float_dtype(device); - let adapter = PyTorchToBurnAdapter.chain(CastFloatAdapter::new(target_float)); - - match sink.resolve(total_f32_bytes) { - DequantSink::Memory => { - let blob = f - .dequant_to_safetensors(map) - .map_err(|e| parse(e.to_string()))?; - assert!(blob.len() > 8, "a parsed GGUF yields a non-empty blob"); - let store = SafetensorsStore::from_bytes(Some(blob)) - .with_from_adapter(adapter) - .allow_partial(true); - Ok((store, None)) - } - DequantSink::Scratch => { - // Beside the gguf, so the scratch write lands on the volume the - // weights already live on. - let scratch = ScratchFile::new(f.path.parent().unwrap_or(Path::new(".")))?; - let bytes = f - .dequant_to_safetensors_file(map, scratch.path()) - .map_err(|e| parse(e.to_string()))?; - assert!(bytes > 0, "a parsed GGUF yields a non-empty payload"); - let store = SafetensorsStore::from_file(scratch.path().to_path_buf()) - .with_from_adapter(adapter) - .allow_partial(true); - Ok((store, Some(scratch))) - } - DequantSink::Auto => unreachable!("resolve never returns Auto"), - } -} - -/// `dir/file`, or [`ImportError::MissingFile`] if absent. -pub fn required_file(dir: &Path, file: &str) -> Result { - assert!(!file.is_empty(), "required_file: empty file name"); - let path = dir.join(file); - if path.is_file() { - Ok(path) - } else { - Err(ImportError::MissingFile(path)) - } -} - -/// The weights checkpoint found in a model dir. -#[derive(Debug, Clone)] -pub enum WeightsFile { - /// `model.safetensors` — the primary format. - Safetensors(PathBuf), - /// `pytorch_model.bin` — the PyTorch state dict older checkpoints ship. - PytorchBin(PathBuf), -} - -/// Pick the weights file in `dir`: `model.safetensors` when present, else -/// `pytorch_model.bin`. Reports the *safetensors* name when neither exists -/// (it's the file a fresh download would produce). -pub fn weights_file(dir: &Path) -> Result { - let safetensors = dir.join("model.safetensors"); - if safetensors.is_file() { - return Ok(WeightsFile::Safetensors(safetensors)); - } - let bin = dir.join("pytorch_model.bin"); - if bin.is_file() { - return Ok(WeightsFile::PytorchBin(bin)); - } - Err(ImportError::MissingFile(safetensors)) -} - -#[cfg(test)] -mod tests { - use super::{DequantSink, MAX_IN_MEMORY_DEQUANT_BYTES, ScratchFile}; - - #[test] - fn auto_picks_the_scratch_sink_exactly_above_the_ceiling() { - // The boundary is the whole point: `Auto` must never leave a payload - // in RAM whose 2x peak is what this ceiling exists to cap. - assert_eq!(DequantSink::Auto.resolve(0), DequantSink::Memory); - assert_eq!( - DequantSink::Auto.resolve(MAX_IN_MEMORY_DEQUANT_BYTES), - DequantSink::Memory - ); - assert_eq!( - DequantSink::Auto.resolve(MAX_IN_MEMORY_DEQUANT_BYTES + 1), - DequantSink::Scratch - ); - // An explicit choice is honoured in BOTH directions, including the - // one `Auto` would never make — the A/B that set the ceiling needs it. - assert_eq!(DequantSink::Memory.resolve(u64::MAX), DequantSink::Memory); - assert_eq!(DequantSink::Scratch.resolve(0), DequantSink::Scratch); - } - - #[test] - fn scratch_files_are_unique_per_instance_and_deleted_on_drop() { - let dir = std::env::temp_dir().join("mummu-scratch-tests"); - std::fs::create_dir_all(&dir).expect("temp dir"); - - let a = ScratchFile::new(&dir).expect("names a scratch file"); - let b = ScratchFile::new(&dir).expect("names a second one"); - assert_ne!( - a.path(), - b.path(), - "two live loads must not share one scratch file" - ); - - let path = a.path().to_path_buf(); - std::fs::write(&path, b"payload").expect("writes"); - assert!(path.is_file()); - drop(a); - assert!(!path.exists(), "Drop removes the scratch file"); - - // And a stale file at the same name is cleared, not adopted: write - // one at b's path, then re-create a guard for it. - let stale = b.path().to_path_buf(); - std::fs::write(&stale, b"stale").expect("writes"); - drop(b); - assert!(!stale.exists()); - } - - use super::*; - - #[test] - fn required_file_finds_present_and_rejects_absent() { - let dir = std::env::temp_dir().join("mummu_import_test"); - std::fs::create_dir_all(&dir).unwrap(); - std::fs::write(dir.join("config.json"), b"{}").unwrap(); - assert!(required_file(&dir, "config.json").is_ok()); - let err = required_file(&dir, "nope.bin").unwrap_err(); - assert!(matches!(err, ImportError::MissingFile(_))); - std::fs::remove_dir_all(&dir).unwrap(); - } - - #[test] - #[should_panic(expected = "float dtype")] - fn cast_adapter_rejects_non_float_target() { - let _ = CastFloatAdapter::new(DType::I32); - } - - #[test] - fn logit_sanity_accepts_a_live_distribution() { - let logits = [0.1, 2.5, -1.0, 9.0, 0.3]; - let smoke = logit_sanity(&logits, 5).expect("live logits pass"); - assert_eq!(smoke.top_id, 3); - assert_eq!(smoke.top_logit, 9.0); - assert!( - (smoke.spread - 10.0).abs() < 1e-6, - "spread = max-min = 9-(-1)" - ); - } - - #[test] - fn logit_sanity_catches_non_finite() { - let nan = [1.0, f32::NAN, 3.0, 2.0]; - assert_eq!( - logit_sanity(&nan, 4), - Err(SanityError::NonFinite { - count: 1, - first_index: 1 - }) - ); - let inf = [f32::INFINITY, 0.0, f32::NEG_INFINITY]; - assert_eq!( - logit_sanity(&inf, 3), - Err(SanityError::NonFinite { - count: 2, - first_index: 0 - }) - ); - } - - #[test] - fn logit_sanity_catches_wrong_vocab() { - // The finiteness check must not run before the width check — a short - // vector is a mismatch, reported as such (not an index panic). - assert_eq!( - logit_sanity(&[1.0, 2.0, 3.0], 4), - Err(SanityError::WrongVocab { - logits_len: 3, - expected: 4 - }) - ); - } - - #[test] - fn logit_sanity_catches_degenerate_forward() { - // Zero-init / dead forward: all (near-)equal. - let flat = [0.5f32; 8]; - assert!(matches!( - logit_sanity(&flat, 8), - Err(SanityError::Degenerate { .. }) - )); - // A spread just above the threshold is live. - let mut live = [0.0f32; 8]; - live[7] = 1e-3; - assert!(logit_sanity(&live, 8).is_ok()); - } - - #[test] - #[should_panic(expected = "expected_vocab must be positive")] - fn logit_sanity_rejects_zero_vocab() { - let _ = logit_sanity(&[], 0); - } - - #[test] - fn weights_file_prefers_safetensors_falls_back_to_pytorch() { - let dir = std::env::temp_dir().join("mummu_weights_file_test"); - std::fs::create_dir_all(&dir).unwrap(); - // Neither present: missing, reported by the safetensors name. - assert!(matches!( - weights_file(&dir), - Err(ImportError::MissingFile(p)) if p.ends_with("model.safetensors") - )); - // Only the state dict: the PyTorch path. - std::fs::write(dir.join("pytorch_model.bin"), b"pt").unwrap(); - assert!(matches!(weights_file(&dir), Ok(WeightsFile::PytorchBin(_)))); - // Both: safetensors wins. - std::fs::write(dir.join("model.safetensors"), b"st").unwrap(); - assert!(matches!( - weights_file(&dir), - Ok(WeightsFile::Safetensors(_)) - )); - std::fs::remove_dir_all(&dir).unwrap(); - } -} diff --git a/crates/mummu/examples/src/lib.rs b/crates/mummu/examples/src/lib.rs deleted file mode 100644 index 223b4df..0000000 --- a/crates/mummu/examples/src/lib.rs +++ /dev/null @@ -1,45 +0,0 @@ -//! Mummu — a from-scratch [Burn](https://burn.dev) model runner. -//! -//! One library that imports open models (LLM, embeddings, later vision), -//! quantizes them to fit, and runs them natively in Rust on whatever hardware -//! is present — CPU, one GPU, or several. All model code is generic over -//! `B: Backend`; consumers (laurelane, Nanna) pick a device at runtime and -//! keep their domain glue out of this crate. - -// Burn's `fusion` feature wraps backends in deeply nested generic types. -#![recursion_limit = "512"] - -pub mod adapt; -pub mod attn_config; -pub mod backend; -pub mod cache; -pub mod chat; -pub mod decode; -pub mod gguf; -/// Generated IQ-quant codebook tables (see the module header). -mod gguf_iq_grids; -pub mod hub; -pub mod import; -pub mod manage; -/// Scheduler B — per-tensor precision placement (crate `mummu-mix`). -pub use mummu_mix as mix; -pub mod models; -pub mod nn; -pub mod pack; -pub mod partition; -pub mod plan; -pub mod prof; -pub mod quant; -pub mod registry; -pub mod safetensors; -/// Scheduler A — dividing work across devices (crate `mummu-schedule`). -pub use mummu_schedule as schedule; -/// Render a checkpoint's own imported chat template (feature `jinja-template`). -#[cfg(feature = "jinja-template")] -pub mod template; -pub mod tier; -pub mod tok_config; -pub mod vram; -pub mod workingset; -pub mod tokenizer; -pub mod tune; diff --git a/crates/mummu/examples/src/manage.rs b/crates/mummu/examples/src/manage.rs deleted file mode 100644 index 5968e81..0000000 --- a/crates/mummu/examples/src/manage.rs +++ /dev/null @@ -1,372 +0,0 @@ -//! Model-cache disk accounting (P8): report how much space each cached model -//! takes and validate a removal target so a consumer's settings UI can -//! reclaim disk without ever escaping the cache dir. App-agnostic and free of -//! async/UI types; ported from laurelane's unit-tested implementation. - -use std::path::{Component, Path, PathBuf}; - -/// One cached model directory and its size on disk. -#[derive(serde::Serialize, Debug, PartialEq, Eq)] -#[serde(rename_all = "camelCase")] -pub struct ModelDisk { - /// Cache subdir name, e.g. `"qwen2.5-1.5b"`. - pub name: String, - pub bytes: u64, -} - -/// A full disk report for a model cache dir. -#[derive(serde::Serialize, Debug, PartialEq, Eq)] -#[serde(rename_all = "camelCase")] -pub struct DiskReport { - /// Cached models, largest first. - pub models: Vec, - pub total_bytes: u64, - /// The cache dir being reported — the "your models live here" line. - pub location: String, -} - -/// Recursively sum the sizes of every regular file under `path`. Best-effort: -/// unreadable entries are skipped and symlinked dirs are not followed (only -/// real directories recurse), so a broken link can't send it off the rails. -#[must_use] -pub fn dir_size(path: &Path) -> u64 { - let mut total = 0; - let Ok(entries) = std::fs::read_dir(path) else { - return 0; - }; - for entry in entries.flatten() { - let Ok(ft) = entry.file_type() else { continue }; - if ft.is_dir() { - total += dir_size(&entry.path()); - } else if ft.is_file() - && let Ok(meta) = entry.metadata() - { - total += meta.len(); - } - } - total -} - -/// Immediate subdirs of `models_dir`, each with its recursive size, sorted -/// largest-first (ties broken by name). A missing dir is an empty list, not -/// an error — a fresh install has no downloaded models yet. -#[must_use] -pub fn cached_models(models_dir: &Path) -> Vec { - let mut out = Vec::new(); - let Ok(entries) = std::fs::read_dir(models_dir) else { - return out; - }; - for entry in entries.flatten() { - if entry.file_type().map(|t| t.is_dir()).unwrap_or(false) { - let name = entry.file_name().to_string_lossy().into_owned(); - out.push(ModelDisk { - name, - bytes: dir_size(&entry.path()), - }); - } - } - out.sort_by(|a, b| b.bytes.cmp(&a.bytes).then_with(|| a.name.cmp(&b.name))); - out -} - -/// A full report: cached models + their combined size + the cache location. -#[must_use] -pub fn report(models_dir: &Path) -> DiskReport { - let models = cached_models(models_dir); - let total_bytes = models.iter().map(|m| m.bytes).sum(); - debug_assert!( - models.windows(2).all(|w| w[0].bytes >= w[1].bytes), - "cached_models must sort largest-first" - ); - DiskReport { - models, - total_bytes, - location: models_dir.to_string_lossy().into_owned(), - } -} - -/// Is `name` a safe single cache-subdir component? Rejects empty, `.`/`..`, -/// anything containing a path separator, and any rooted/prefixed path — so a -/// removal can never climb out of the cache dir. -#[must_use] -pub fn is_safe_component(name: &str) -> bool { - if name.is_empty() || name == "." || name == ".." { - return false; - } - if name.contains('/') || name.contains('\\') { - return false; - } - let mut parts = Path::new(name).components(); - matches!(parts.next(), Some(Component::Normal(_))) && parts.next().is_none() -} - -/// Resolve the on-disk path to remove for cache subdir `name`, or an error if -/// `name` is unsafe or absent. The returned path is always a direct child of -/// `models_dir`. -pub fn resolve_removal(models_dir: &Path, name: &str) -> Result { - if !is_safe_component(name) { - return Err(format!("unsafe model name: {name:?}")); - } - let target = models_dir.join(name); - if !target.is_dir() { - return Err(format!("no cached model named {name:?}")); - } - debug_assert!( - target.starts_with(models_dir), - "removal must stay inside the cache" - ); - Ok(target) -} - -#[cfg(test)] -mod tests { - use super::*; - use std::fs; - use std::path::PathBuf; - - /// A unique temp dir for one test, wiped fresh. Clock-free (uses the - /// test's own tag), so parallel tests never collide. - fn scratch(tag: &str) -> PathBuf { - let dir = std::env::temp_dir().join(format!("mummu_manage_test_{tag}")); - let _ = fs::remove_dir_all(&dir); - fs::create_dir_all(&dir).unwrap(); - dir - } - - fn write(path: &Path, bytes: usize) { - fs::create_dir_all(path.parent().unwrap()).unwrap(); - fs::write(path, vec![b'x'; bytes]).unwrap(); - } - - #[test] - fn dir_size_sums_nested_files() { - let root = scratch("dirsize"); - write(&root.join("a.bin"), 100); - write(&root.join("sub/b.bin"), 250); - write(&root.join("sub/deep/c.bin"), 50); - assert_eq!(dir_size(&root), 400); - // A missing dir is zero, not a panic. - assert_eq!(dir_size(&root.join("nope")), 0); - fs::remove_dir_all(&root).unwrap(); - } - - #[test] - fn cached_models_lists_subdirs_largest_first_ignoring_loose_files() { - let root = scratch("cached"); - write(&root.join("qwen2.5-1.5b/model.safetensors"), 3000); - write(&root.join("all-minilm/model.safetensors"), 90); - write(&root.join("lfm2.5/model.safetensors"), 500); - write(&root.join("loose.txt"), 10_000); // not a dir → ignored - let got = cached_models(&root); - assert_eq!( - got, - vec![ - ModelDisk { - name: "qwen2.5-1.5b".into(), - bytes: 3000 - }, - ModelDisk { - name: "lfm2.5".into(), - bytes: 500 - }, - ModelDisk { - name: "all-minilm".into(), - bytes: 90 - }, - ], - ); - assert_eq!(report(&root).total_bytes, 3590); - fs::remove_dir_all(&root).unwrap(); - } - - #[test] - fn cached_models_of_missing_dir_is_empty() { - let missing = std::env::temp_dir().join("mummu_manage_test_absent_xyz"); - let _ = fs::remove_dir_all(&missing); - assert!(cached_models(&missing).is_empty()); - assert_eq!(report(&missing).total_bytes, 0); - } - - #[test] - fn is_safe_component_accepts_names_and_rejects_traversal() { - assert!(is_safe_component("qwen2.5-1.5b")); - assert!(is_safe_component("all-minilm")); - for bad in [ - "", - ".", - "..", - "a/b", - "a\\b", - "/etc", - "..\\..\\x", - "C:\\x", - "sub/", - ] { - assert!(!is_safe_component(bad), "should reject {bad:?}"); - } - } - - #[test] - fn resolve_removal_rejects_unsafe_and_absent_but_finds_present() { - let root = scratch("removal"); - write(&root.join("qwen2.5-0.5b/model.safetensors"), 10); - assert_eq!( - resolve_removal(&root, "qwen2.5-0.5b").unwrap(), - root.join("qwen2.5-0.5b") - ); - assert!(resolve_removal(&root, "not-there").is_err()); - assert!(resolve_removal(&root, "..").is_err()); - assert!(resolve_removal(&root, "../models").is_err()); - fs::remove_dir_all(&root).unwrap(); - } -} - -/// The settings-UI-facing management surface (P8): one object owning the -/// models root that composes the catalog ([`crate::registry`]), downloads -/// ([`crate::hub`], with per-chunk progress), disk accounting, and safe -/// removal. Active-model *switching* is the consumer's `ModelSlot` keyed by -/// [`ModelManager::model_dir`]. -pub struct ModelManager { - root: PathBuf, - catalog: Vec, -} - -impl ModelManager { - /// Manage `root` with the built-in catalog. - #[must_use] - pub fn new(root: PathBuf) -> Self { - Self::with_catalog(root, crate::registry::catalog()) - } - - /// Manage `root` with an app-supplied catalog (all specs must validate). - #[must_use] - pub fn with_catalog(root: PathBuf, catalog: Vec) -> Self { - assert!( - !root.as_os_str().is_empty(), - "models root must be non-empty" - ); - assert!(!catalog.is_empty(), "catalog must not be empty"); - for spec in &catalog { - if let Err(e) = spec.validate() { - panic!("invalid catalog spec: {e}"); - } - } - Self { root, catalog } - } - - #[must_use] - pub fn catalog(&self) -> &[crate::registry::ModelSpec] { - &self.catalog - } - - /// The dir a catalog model lives in (whether or not it's installed yet). - pub fn model_dir(&self, name: &str) -> Result { - self.spec(name).map(|s| s.dir(&self.root)) - } - - /// Is every required artifact of `name` on disk? (config + tokenizer + - /// single-file weights or a shard index.) - pub fn is_installed(&self, name: &str) -> Result { - let dir = self.model_dir(name)?; - let weights = dir.join("model.safetensors").is_file() - || dir.join("model.safetensors.index.json").is_file(); - Ok(weights && dir.join("config.json").is_file() && dir.join("tokenizer.json").is_file()) - } - - /// Download `name` from its spec (resumable, cache-first), reporting - /// progress per chunk. Returns the model dir ready for `load_from_dir`. - pub fn install( - &self, - name: &str, - on_progress: impl FnMut(crate::hub::Progress<'_>), - ) -> Result { - let spec = self.spec(name)?; - spec.fetch(&self.root, on_progress) - .map_err(|e| e.to_string()) - } - - /// Remove `name`'s files from disk (traversal-safe). The caller drops any - /// live `ModelSlot` first — removal only touches the disk. - pub fn remove(&self, name: &str) -> Result<(), String> { - let target = resolve_removal(&self.root, name)?; - std::fs::remove_dir_all(&target).map_err(|e| format!("remove {name:?}: {e}")) - } - - /// Disk usage for everything under the root, largest first. - #[must_use] - pub fn disk_report(&self) -> DiskReport { - report(&self.root) - } - - fn spec(&self, name: &str) -> Result<&crate::registry::ModelSpec, String> { - self.catalog.iter().find(|s| s.name == name).ok_or_else(|| { - let known: Vec<&str> = self.catalog.iter().map(|s| s.name.as_str()).collect(); - format!("unknown model {name:?}; catalog has {known:?}") - }) - } -} - -#[cfg(test)] -mod manager_tests { - use super::*; - - fn temp_root(tag: &str) -> PathBuf { - let root = std::env::temp_dir().join(format!("mummu-manager-{tag}")); - let _ = std::fs::remove_dir_all(&root); - std::fs::create_dir_all(&root).unwrap(); - root - } - - fn fake_install(root: &Path, name: &str) { - let dir = root.join(name); - std::fs::create_dir_all(&dir).unwrap(); - for f in ["config.json", "tokenizer.json", "model.safetensors"] { - std::fs::write(dir.join(f), b"{}").unwrap(); - } - } - - #[test] - fn install_state_and_report_reflect_disk() { - let root = temp_root("state"); - let mgr = ModelManager::new(root.clone()); - assert_eq!(mgr.is_installed("all-minilm-l6-v2"), Ok(false)); - - fake_install(&root, "all-minilm-l6-v2"); - assert_eq!(mgr.is_installed("all-minilm-l6-v2"), Ok(true)); - - let rep = mgr.disk_report(); - assert_eq!(rep.models.len(), 1); - assert_eq!(rep.models[0].name, "all-minilm-l6-v2"); - assert!(rep.total_bytes > 0); - std::fs::remove_dir_all(&root).ok(); - } - - #[test] - fn remove_deletes_only_the_named_model() { - let root = temp_root("remove"); - let mgr = ModelManager::new(root.clone()); - fake_install(&root, "all-minilm-l6-v2"); - fake_install(&root, "qwen2.5-0.5b-instruct"); - - mgr.remove("all-minilm-l6-v2").unwrap(); - assert_eq!(mgr.is_installed("all-minilm-l6-v2"), Ok(false)); - assert_eq!(mgr.is_installed("qwen2.5-0.5b-instruct"), Ok(true)); - std::fs::remove_dir_all(&root).ok(); - } - - #[test] - fn unknown_names_fail_loudly_with_the_catalog() { - let root = temp_root("unknown"); - let mgr = ModelManager::new(root.clone()); - let err = mgr.model_dir("nope").unwrap_err(); - assert!(err.contains("nope") && err.contains("all-minilm-l6-v2")); - assert!(mgr.remove("nope").is_err()); - std::fs::remove_dir_all(&root).ok(); - } - - #[test] - #[should_panic(expected = "catalog must not be empty")] - fn empty_catalog_is_rejected() { - let _ = ModelManager::with_catalog(PathBuf::from("x"), vec![]); - } -} diff --git a/crates/mummu/examples/src/models/lfm2.rs b/crates/mummu/examples/src/models/lfm2.rs deleted file mode 100644 index 5425307..0000000 --- a/crates/mummu/examples/src/models/lfm2.rs +++ /dev/null @@ -1,830 +0,0 @@ -//! LFM2 / LFM2.5 (Liquid Foundation Model) hybrid decoder on the shared `nn` -//! blocks: Embedding → N×{operator_norm, (gated ShortConv "LIV" | GQA attention -//! with per-head q/k RMSNorm), ffn_norm, SwiGLU} → embedding_norm → tied -//! lm-head. `layer_types[i]` in the checkpoint's `config.json` selects the -//! operator per layer (10 conv + 6 attention for the 1.2B). -//! -//! Ported from laurelane's implementation (greedy temperature-0 token match vs -//! a local Ollama `lfm2.5` reference). Mummu's parity gate (P7) re-verifies -//! here before the port is marked trusted. - -use std::path::Path; - -use burn::module::Module; -use burn::nn::{Embedding, EmbeddingConfig, RmsNorm, RmsNormConfig}; -use burn::store::{ModuleAdapter, PyTorchToBurnAdapter, SafetensorsStore}; -use burn::tensor::{Device, Int, Tensor, TensorData}; - -use crate::attn_config::{RopeScaling, check_sliding_window, sliding_window_from_gguf}; -use crate::gguf::{GgufFile, GgufMap, GgufTensorInfo, GgufValue}; -use crate::import::{ - CastFloatAdapter, DequantSink, ImportError, gguf_store, load_checked, required_file, -}; -use crate::models::CausalLm; -use crate::models::qwen2::EosIds; -use crate::nn::{ - ConvState, GqaAttention, GqaAttentionConfig, LayerKv, ShortConv, ShortConvConfig, SwiGluMlp, - SwiGluMlpConfig, causal_mask, rope_tables, -}; - -/// LFM2 architecture hyperparameters, read from the checkpoint's `config.json`. -#[derive(Debug, Clone, serde::Deserialize)] -pub struct Lfm2Config { - pub vocab_size: usize, - pub hidden_size: usize, - pub num_hidden_layers: usize, - pub num_attention_heads: usize, - pub num_key_value_heads: usize, - pub norm_eps: f64, - /// Top-level in older checkpoints (LFM2.5-1.2B); newer ones (LFM2.5-230M) - /// nest it under `rope_parameters` instead — resolved (and required) by - /// [`Self::from_json_bytes`] / [`Self::validate`]. - #[serde(default)] - pub rope_theta: f32, - /// The newer transformers convention: `{"rope_theta": …, "rope_type": …}`. - #[serde(default, skip_serializing)] - rope_parameters: Option, - /// The third spelling: a top-level `rope_scaling` object, as Qwen and the - /// Llama family write it. Refused unless plain ([`crate::attn_config`]). - #[serde(default)] - pub rope_scaling: Option, - /// The trained context length (LFM2.5-1.2B: 128 000), used to tell an - /// inert sliding window from a clipping one. - #[serde(default)] - pub max_position_embeddings: Option, - /// LFM2 has no `use_sliding_window` flag: its attention layers are full, - /// and its short-conv layers are a different mechanism entirely — so a - /// declared window would be live, and is refused. - #[serde(default)] - pub sliding_window: Option, - /// Conv kernel length `K` (the rolling decode state keeps `K-1`). - #[serde(rename = "conv_L_cache")] - pub conv_l_cache: usize, - pub block_ff_dim: usize, - pub block_multiple_of: usize, - #[serde(default)] - pub block_auto_adjust_ff_dim: bool, - /// `"full_attention"` or `"conv"` per layer. - pub layer_types: Vec, - /// `<|im_end|>` (id 7) for the instruct checkpoints. - #[serde(default)] - pub eos_token_id: EosIds, -} - -impl Lfm2Config { - #[must_use] - pub fn head_dim(&self) -> usize { - self.hidden_size / self.num_attention_heads - } - - /// SwiGLU hidden dim: `block_auto_adjust_ff_dim` shrinks to 2/3 then - /// rounds up to `block_multiple_of` (8192 for the 1.2B). - #[must_use] - pub fn ff_dim(&self) -> usize { - if self.block_auto_adjust_ff_dim { - let d = (2 * self.block_ff_dim) / 3; - let m = self.block_multiple_of.max(1); - d.div_ceil(m) * m - } else { - self.block_ff_dim - } - } - - fn is_attention(&self, layer: usize) -> bool { - self.layer_types - .get(layer) - .is_some_and(|s| s == "full_attention") - } - - /// Parse `config.json` bytes. `rope_theta` may arrive top-level (older - /// checkpoints) or nested under `rope_parameters` (the newer transformers - /// convention, LFM2.5-230M) — only plain rotary (`rope_type: "default"`) - /// is implemented, anything else fails loudly. - pub fn from_json_bytes(bytes: &[u8]) -> Result { - let mut cfg: Self = serde_json::from_slice(bytes).map_err(|e| e.to_string())?; - if let Some(rp) = cfg.rope_parameters.take() { - // Same rule, same message as every other family's `rope_scaling`: - // only plain rotary is computed, anything else is named and - // refused rather than silently approximated. - RopeScaling { - rope_type: Some(rp.rope_type.clone()), - ..RopeScaling::default() - } - .check("lfm2 config.json rope_parameters")?; - if cfg.rope_theta != 0.0 && cfg.rope_theta != rp.rope_theta { - return Err(format!( - "rope_theta given twice and disagreeing: {} (top-level) vs {} (rope_parameters)", - cfg.rope_theta, rp.rope_theta - )); - } - cfg.rope_theta = rp.rope_theta; - } - cfg.validate("lfm2 config.json")?; - Ok(cfg) - } - - /// Hyperparameters from a GGUF header's `lfm2.*` metadata. Layer kinds - /// come from llama.cpp's per-layer `attention.head_count_kv` array: - /// `0` marks a shortconv layer, nonzero an attention layer. - /// `feed_forward_length` in a GGUF is the already-adjusted SwiGLU dim, - /// so auto-adjust is off. - pub fn from_gguf(f: &GgufFile) -> Result { - let arch = f.architecture().unwrap_or(""); - if arch != "lfm2" { - return Err(format!("GGUF architecture '{arch}' is not lfm2")); - } - let meta_usize = |key: &str| -> Result { - f.get(key) - .and_then(GgufValue::as_u64) - .and_then(|v| usize::try_from(v).ok()) - .ok_or_else(|| format!("missing or non-integer GGUF metadata '{key}'")) - }; - let hidden_size = meta_usize("lfm2.embedding_length")?; - let num_hidden_layers = meta_usize("lfm2.block_count")?; - - // Per-layer kv-head counts: 0 = conv, nonzero = attention (all - // nonzero entries must agree — one KV geometry per model). - let kv_per_layer = f - .get("lfm2.attention.head_count_kv") - .and_then(GgufValue::as_array) - .ok_or("missing per-layer array 'lfm2.attention.head_count_kv'")?; - if kv_per_layer.len() != num_hidden_layers { - return Err(format!( - "head_count_kv has {} entries for {num_hidden_layers} layers", - kv_per_layer.len() - )); - } - let mut layer_types = Vec::with_capacity(num_hidden_layers); - let mut kv_heads: Option = None; - for (i, v) in kv_per_layer.iter().enumerate() { - let n = v - .as_i64() - .and_then(|n| u64::try_from(n).ok()) // i32 array in real files - .ok_or_else(|| format!("head_count_kv[{i}] is not a non-negative integer"))?; - if n == 0 { - layer_types.push("conv".to_string()); - } else { - if kv_heads.is_some_and(|k| k != n) { - return Err(format!("head_count_kv mixes {kv_heads:?} and {n}")); - } - kv_heads = Some(n); - layer_types.push("full_attention".to_string()); - } - } - let Some(kv_heads) = kv_heads else { - return Err("no attention layers in head_count_kv".into()); - }; - - let embd = f - .tensor("token_embd.weight") - .ok_or("GGUF has no token_embd.weight tensor")?; - if embd.dims.len() != 2 || embd.dims[0] != hidden_size as u64 { - return Err(format!( - "token_embd.weight dims {:?} do not match embedding_length {hidden_size}", - embd.dims - )); - } - if f.tensor("output.weight").is_some() { - return Err("LFM2 GGUF carries an untied output.weight — unsupported".into()); - } - let eps = f - .get("lfm2.attention.layer_norm_rms_epsilon") - .and_then(GgufValue::as_f32) - .ok_or("missing 'lfm2.attention.layer_norm_rms_epsilon'")?; - let theta = f - .get("lfm2.rope.freq_base") - .and_then(GgufValue::as_f32) - .ok_or("missing 'lfm2.rope.freq_base'")?; - let cfg = Self { - vocab_size: usize::try_from(embd.dims[1]).map_err(|_| "vocab too large")?, - hidden_size, - num_hidden_layers, - num_attention_heads: meta_usize("lfm2.attention.head_count")?, - num_key_value_heads: usize::try_from(kv_heads).map_err(|_| "kv heads too large")?, - norm_eps: f64::from(eps), - rope_theta: theta, - rope_parameters: None, - rope_scaling: RopeScaling::from_gguf(f, "lfm2"), - max_position_embeddings: f - .get("lfm2.context_length") - .and_then(GgufValue::as_u64) - .and_then(|v| usize::try_from(v).ok()), - sliding_window: sliding_window_from_gguf(f, "lfm2"), - conv_l_cache: meta_usize("lfm2.shortconv.l_cache")?, - block_ff_dim: meta_usize("lfm2.feed_forward_length")?, - block_multiple_of: 1, - block_auto_adjust_ff_dim: false, - layer_types, - eos_token_id: f - .get("tokenizer.ggml.eos_token_id") - .and_then(GgufValue::as_u64) - .and_then(|v| u32::try_from(v).ok()) - .map_or(EosIds::None, EosIds::One), - }; - cfg.validate("lfm2 GGUF header")?; - Ok(cfg) - } - - fn validate(&self, whose: &str) -> Result<(), String> { - if let Some(scaling) = &self.rope_scaling { - scaling.check(whose)?; - } - check_sliding_window( - self.sliding_window, - self.sliding_window.is_some(), - self.max_position_embeddings, - whose, - )?; - if self.layer_types.len() != self.num_hidden_layers { - return Err(format!( - "layer_types has {} entries but num_hidden_layers is {}", - self.layer_types.len(), - self.num_hidden_layers - )); - } - if self.num_key_value_heads == 0 - || !self - .num_attention_heads - .is_multiple_of(self.num_key_value_heads) - { - return Err(format!( - "num_attention_heads ({}) must be a positive multiple of num_key_value_heads ({})", - self.num_attention_heads, self.num_key_value_heads - )); - } - if self.conv_l_cache < 2 { - return Err(format!( - "conv_L_cache must be >= 2, got {}", - self.conv_l_cache - )); - } - if !(self.rope_theta.is_finite() && self.rope_theta > 0.0) { - return Err(format!( - "rope_theta must be a positive float (top-level or in rope_parameters), got {}", - self.rope_theta - )); - } - Ok(()) - } -} - -/// The nested rope block newer checkpoints carry in `config.json`. -#[derive(Debug, Clone, serde::Deserialize)] -struct RopeParameters { - rope_theta: f32, - #[serde(default = "RopeParameters::default_type")] - rope_type: String, -} - -impl RopeParameters { - fn default_type() -> String { - "default".into() - } -} - -/// One hybrid layer: exactly one of `conv` / `self_attn` is present, per -/// `layer_types[i]`. Field names mirror the HF checkpoint. -#[derive(Module, Debug)] -pub struct HybridLayer { - pub conv: Option, - pub self_attn: Option, - pub feed_forward: SwiGluMlp, - pub operator_norm: RmsNorm, - pub ffn_norm: RmsNorm, -} - -/// The LFM2 decoder stack (tied lm-head). -#[derive(Module, Debug)] -pub struct Lfm2 { - pub embed_tokens: Embedding, - pub layers: Vec, - pub embedding_norm: RmsNorm, -} - -/// Per-layer decode cache: conv layers roll the last `K-1` gated inputs, -/// attention layers keep the running k/v. -pub enum HybridKv { - Conv(ConvState), - Attn(LayerKv), -} - -/// A weight-loaded LFM2 plus its config. -pub struct LoadedLfm2 { - pub model: Lfm2, - pub config: Lfm2Config, - /// The parsed sibling `tokenizer_config.json`, when one was present and - /// well-formed beside a safetensors checkpoint (the load-time gate has - /// already cross-checked its EOS against `config.json`). A consumer reads - /// config-driven EOS/BOS/PAD ids from it (`eos_id()`, `bos_id()`, …). `None` - /// for a GGUF load (self-contained; no sibling file) or a dir without one. - pub tokenizer_config: Option, -} - -fn build(cfg: &Lfm2Config, device: &Device) -> Lfm2 { - let attn_cfg = GqaAttentionConfig { - hidden_size: cfg.hidden_size, - num_heads: cfg.num_attention_heads, - num_kv_heads: cfg.num_key_value_heads, - head_dim: cfg.head_dim(), - bias: false, // LFM2 projections are bias-free - qk_norm_eps: Some(cfg.norm_eps), // per-head q/k RMSNorm - qk_norm_projection: false, - }; - let conv_cfg = ShortConvConfig { - hidden_size: cfg.hidden_size, - kernel_len: cfg.conv_l_cache, - }; - let mlp_cfg = SwiGluMlpConfig { - hidden_size: cfg.hidden_size, - intermediate_size: cfg.ff_dim(), - }; - let norm = |dev: &Device| { - RmsNormConfig::new(cfg.hidden_size) - .with_epsilon(cfg.norm_eps) - .init(dev) - }; - let layers = (0..cfg.num_hidden_layers) - .map(|i| { - let attn = cfg.is_attention(i); - HybridLayer { - conv: (!attn).then(|| conv_cfg.init(device)), - self_attn: attn.then(|| attn_cfg.init(device)), - feed_forward: mlp_cfg.init(device), - operator_norm: norm(device), - ffn_norm: norm(device), - } - }) - .collect(); - Lfm2 { - embed_tokens: EmbeddingConfig::new(cfg.vocab_size, cfg.hidden_size).init(device), - layers, - embedding_norm: norm(device), - } -} - -/// Build the architecture from `dir/config.json` and load -/// `dir/model.safetensors`, checked. Key remap: strip `model.`, RmsNorm -/// `weight` → `gamma`, and LFM2's `out_proj`/`q_layernorm`/`k_layernorm` onto -/// the shared block's `o_proj`/`q_norm`/`k_norm`. -pub fn load_from_dir(dir: &Path, device: &Device) -> Result { - let cfg_path = required_file(dir, "config.json")?; - let weights = required_file(dir, "model.safetensors")?; - let cfg_bytes = std::fs::read(&cfg_path).map_err(|e| ImportError::Parse { - file: cfg_path.clone(), - reason: e.to_string(), - })?; - let config = Lfm2Config::from_json_bytes(&cfg_bytes).map_err(|reason| ImportError::Parse { - file: cfg_path, - reason, - })?; - - // Cross-check the sibling metadata (when present) before touching weights: - // tokenizer_config.json EOS agreement with config.json, a chat-template that - // speaks LFM2.5's bracket-notation tool-call convention, and added-token ids - // that match the real tokenizer.json — a repackaging mismatch fails loudly. - let tokenizer_config = crate::tokenizer::validate_checkpoint_dir( - dir, - &config.eos_token_id.to_vec(), - Some(crate::tok_config::ToolCallConvention::Lfm), - )?; - debug_assert!( - tokenizer_config - .as_ref() - .and_then(crate::tok_config::TokenizerConfig::eos_id) - .is_none_or(|id| config.eos_token_id.contains(id)), - "validate_checkpoint_dir returned a config whose EOS disagrees with config.json" - ); - - let mut model = build(&config, device); - // The float dtype comes from the DEVICE — burn 0.22 keeps the element - // type there as a runtime setting, not on a backend type. Creation sites - // still name it explicitly rather than riding the unspecified default. - let target_float = crate::backend::float_dtype(device); - let mut store = SafetensorsStore::from_file(weights.clone()) - .with_from_adapter(PyTorchToBurnAdapter.chain(CastFloatAdapter::new(target_float))) - .allow_partial(true) - .with_key_remapping(r"^model\.", "") - .with_key_remapping(r"(self_attn)\.out_proj\.", "$1.o_proj.") - .with_key_remapping(r"(self_attn)\.q_layernorm\.weight$", "$1.q_norm.gamma") - .with_key_remapping(r"(self_attn)\.k_layernorm\.weight$", "$1.k_norm.gamma") - .with_key_remapping(r"(feed_forward)\.w1\.", "$1.gate_proj.") - .with_key_remapping(r"(feed_forward)\.w2\.", "$1.down_proj.") - .with_key_remapping(r"(feed_forward)\.w3\.", "$1.up_proj.") - .with_key_remapping( - r"(operator_norm|ffn_norm|embedding_norm)\.weight$", - "$1.gamma", - ); - load_checked(&mut model, &mut store, &weights)?; - Ok(LoadedLfm2 { - model, - config, - tokenizer_config, - }) -} - -/// GGUF (llama.cpp `lfm2` arch) tensor names → the HF checkpoint names the -/// safetensors remap chain already handles. The depthwise conv kernel is the -/// one shape special-case: llama.cpp stores it squeezed (`[k, channels]` in -/// ggml dims), the checkpoint as `[channels, 1, k]` — same bytes. -fn gguf_tensor_to_hf(info: &GgufTensorInfo) -> Option { - match info.name.as_str() { - "token_embd.weight" => { - return Some(GgufMap::Rename("model.embed_tokens.weight".into())); - } - "token_embd_norm.weight" => { - return Some(GgufMap::Rename("model.embedding_norm.weight".into())); - } - _ => {} - } - let rest = info.name.strip_prefix("blk.")?; - let (layer, field) = rest.split_once('.')?; - let layer: usize = layer.parse().ok()?; - if field == "shortconv.conv.weight" { - // ggml dims [k, channels] → row-major [channels, k] → the - // checkpoint's depthwise-Conv1d shape [channels, 1, k]. - let (&k, &channels) = (info.dims.first()?, info.dims.get(1)?); - return Some(GgufMap::Reshape( - format!("model.layers.{layer}.conv.conv.weight"), - vec![channels, 1, k], - )); - } - let mapped = match field { - "attn_norm.weight" => "operator_norm.weight", - "ffn_norm.weight" => "ffn_norm.weight", - "attn_q.weight" => "self_attn.q_proj.weight", - "attn_k.weight" => "self_attn.k_proj.weight", - "attn_v.weight" => "self_attn.v_proj.weight", - "attn_output.weight" => "self_attn.out_proj.weight", - "attn_q_norm.weight" => "self_attn.q_layernorm.weight", - "attn_k_norm.weight" => "self_attn.k_layernorm.weight", - "ffn_gate.weight" => "feed_forward.w1.weight", - "ffn_down.weight" => "feed_forward.w2.weight", - "ffn_up.weight" => "feed_forward.w3.weight", - "shortconv.in_proj.weight" => "conv.in_proj.weight", - "shortconv.out_proj.weight" => "conv.out_proj.weight", - _ => return None, - }; - Some(GgufMap::Rename(format!("model.layers.{layer}.{mapped}"))) -} - -/// Load an LFM2/LFM2.5 model straight from a **GGUF** file: hyperparameters -/// from the `lfm2.*` metadata (layer kinds from the per-layer kv-head -/// array), weights dequantized and driven through the exact store pipeline -/// the safetensors path uses. -pub fn load_from_gguf(path: &Path, device: &Device) -> Result { - let parse = |reason: String| ImportError::Parse { - file: path.to_path_buf(), - reason, - }; - let f = GgufFile::open(path).map_err(|e| parse(e.to_string()))?; - let config = Lfm2Config::from_gguf(&f).map_err(parse)?; - // The scratch guard (Some only when the payload went to disk) must - // outlive `load_checked`: the store reads that file lazily. - let (base, _scratch) = gguf_store(&f, &gguf_tensor_to_hf, DequantSink::Auto, device)?; - - let mut model = build(&config, device); - let mut store = base - .with_key_remapping(r"^model\.", "") - .with_key_remapping(r"(self_attn)\.out_proj\.", "$1.o_proj.") - .with_key_remapping(r"(self_attn)\.q_layernorm\.weight$", "$1.q_norm.gamma") - .with_key_remapping(r"(self_attn)\.k_layernorm\.weight$", "$1.k_norm.gamma") - .with_key_remapping(r"(feed_forward)\.w1\.", "$1.gate_proj.") - .with_key_remapping(r"(feed_forward)\.w2\.", "$1.down_proj.") - .with_key_remapping(r"(feed_forward)\.w3\.", "$1.up_proj.") - .with_key_remapping( - r"(operator_norm|ffn_norm|embedding_norm)\.weight$", - "$1.gamma", - ); - load_checked(&mut model, &mut store, path)?; - // A GGUF is self-contained — no sibling tokenizer_config.json in this path. - Ok(LoadedLfm2 { - model, - config, - tokenizer_config: None, - }) -} - -impl CausalLm for LoadedLfm2 { - type Cache = Vec; - - fn is_eos(&self, id: u32) -> bool { - self.config.eos_token_id.contains(id) - } - - /// A fresh per-layer cache matching `layer_types`. - fn new_cache(&self) -> Self::Cache { - (0..self.config.num_hidden_layers) - .map(|i| { - if self.config.is_attention(i) { - HybridKv::Attn(None) - } else { - HybridKv::Conv(None) - } - }) - .collect() - } - - /// Forward `new_ids` (the whole prompt when `past == 0`, else one decode - /// token), updating `cache`; returns logits for the last position `[1, vocab]`. - fn forward( - &self, - new_ids: &[u32], - past: usize, - cache: &mut Self::Cache, - device: &Device, - ) -> Tensor<2> { - let t = new_ids.len(); - assert!(t >= 1, "LFM2 forward: need at least one token"); - assert!( - cache.len() == self.config.num_hidden_layers, - "LFM2 forward: cache has {} layers, model has {}", - cache.len(), - self.config.num_hidden_layers - ); - let cfg = &self.config; - - // Dtype pinned to the backend TYPE, never the per-device policy. - let ids32: Vec = new_ids.iter().map(|&i| i as i32).collect(); - let input = Tensor::<1, Int>::from_data( - TensorData::new(ids32, [t]), - (device, crate::backend::int_dtype(device)), - ) - .reshape([1, t]); - let mut x = self.model.embed_tokens.forward(input); - - let (cos, sin) = rope_tables(t, past, cfg.head_dim(), cfg.rope_theta, device); - let mask = (t > 1).then(|| causal_mask(t, past, device)); - let kk = cfg.conv_l_cache; - - for (layer, kv) in self.model.layers.iter().zip(cache.iter_mut()) { - let h = layer.operator_norm.forward(x.clone()); - let h = match (&layer.conv, &layer.self_attn, kv) { - (Some(conv), None, HybridKv::Conv(state)) => conv.forward(h, kk, state), - (None, Some(attn), HybridKv::Attn(kv_state)) => attn.forward( - h, - cfg.num_attention_heads, - cfg.num_key_value_heads, - cfg.head_dim(), - &cos, - &sin, - mask.as_ref(), - kv_state, - ), - // Layer kind and cache kind disagree — a caller bug. - _ => unreachable!("LFM2 forward: layer/cache kind mismatch"), - }; - x = x.add(h); - let h2 = layer.ffn_norm.forward(x.clone()); - x = x.add(layer.feed_forward.forward(h2)); - } - let x = self.model.embedding_norm.forward(x); - let last = x.narrow(1, t - 1, 1).reshape([1, cfg.hidden_size]); - let w = self.model.embed_tokens.weight.val(); // tied lm-head - last.matmul(w.swap_dims(0, 1)) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - /// A 3-layer toy hybrid: conv, attention, conv. - fn toy_config() -> Lfm2Config { - Lfm2Config { - vocab_size: 48, - hidden_size: 16, - num_hidden_layers: 3, - num_attention_heads: 4, - num_key_value_heads: 2, - norm_eps: 1e-5, - rope_theta: 1e4, - rope_parameters: None, - rope_scaling: None, - max_position_embeddings: Some(512), - sliding_window: None, - conv_l_cache: 3, - block_ff_dim: 48, - block_multiple_of: 8, - block_auto_adjust_ff_dim: true, - layer_types: vec!["conv".into(), "full_attention".into(), "conv".into()], - eos_token_id: EosIds::One(7), - } - } - - #[test] - fn ff_dim_auto_adjust_matches_reference_formula() { - let cfg = toy_config(); - // 2*48/3 = 32, already a multiple of 8. - assert_eq!(cfg.ff_dim(), 32); - // The 1.2B's real numbers: 2/3 of 12288 = 8192, multiple_of 8192 → 8192. - let mut big = toy_config(); - big.block_ff_dim = 12288; - big.block_multiple_of = 8192; - assert_eq!(big.ff_dim(), 8192); - } - - #[test] - fn config_rejects_layer_type_count_mismatch() { - let mut cfg = toy_config(); - cfg.layer_types.pop(); - assert!(cfg.validate("test").is_err()); - } - - /// Both `rope_theta` spellings parse; anything but plain rotary is loud. - #[test] - fn config_reads_rope_theta_top_level_or_nested() { - let base = r#"{ - "vocab_size": 48, "hidden_size": 16, "num_hidden_layers": 1, - "num_attention_heads": 4, "num_key_value_heads": 2, - "norm_eps": 1e-5, "conv_L_cache": 3, - "block_ff_dim": 48, "block_multiple_of": 8, - "layer_types": ["full_attention"]"#; - let top = format!(r#"{base}, "rope_theta": 10000.0}}"#); - assert_eq!( - Lfm2Config::from_json_bytes(top.as_bytes()) - .unwrap() - .rope_theta, - 1e4 - ); - // The 230M spelling: nested, with an explicit default rope_type. - let nested = format!( - r#"{base}, "rope_parameters": {{"rope_theta": 1000000.0, "rope_type": "default"}}}}"# - ); - assert_eq!( - Lfm2Config::from_json_bytes(nested.as_bytes()) - .unwrap() - .rope_theta, - 1e6 - ); - // Negative space: a rope scheme we don't implement must not load, - // and a config with neither spelling must not default silently. - let yarn = format!( - r#"{base}, "rope_parameters": {{"rope_theta": 1000000.0, "rope_type": "yarn"}}}}"# - ); - assert!(Lfm2Config::from_json_bytes(yarn.as_bytes()).is_err()); - let missing = format!("{base}}}"); - assert!(Lfm2Config::from_json_bytes(missing.as_bytes()).is_err()); - } - - #[test] - fn build_places_operators_by_layer_type() { - let device = crate::backend::cpu_device(); - let cfg = toy_config(); - let model = build(&cfg, &device); - assert!(model.layers[0].conv.is_some() && model.layers[0].self_attn.is_none()); - assert!(model.layers[1].conv.is_none() && model.layers[1].self_attn.is_some()); - assert!(model.layers[2].conv.is_some() && model.layers[2].self_attn.is_none()); - } - - /// Whole-stack cache equivalence for the hybrid: prefill + cached decode - /// equals one full forward, across BOTH cache kinds at once. - #[test] - fn toy_hybrid_cached_decode_matches_full_forward() { - let device = crate::backend::cpu_device(); - let cfg = toy_config(); - let loaded = LoadedLfm2 { - model: build(&cfg, &device), - config: cfg, - tokenizer_config: None, - }; - - let prompt: Vec = vec![5, 11, 2, 30]; - let mut cache = loaded.new_cache(); - let _ = loaded.forward(&prompt, 0, &mut cache, &device); - let step = loaded - .forward(&[19], prompt.len(), &mut cache, &device) - .into_data() - .to_vec::() - .unwrap(); - - let mut full_cache = loaded.new_cache(); - let all: Vec = prompt.iter().copied().chain([19]).collect(); - let full = loaded - .forward(&all, 0, &mut full_cache, &device) - .into_data() - .to_vec::() - .unwrap(); - - for (i, (c, f)) in step.iter().zip(&full).enumerate() { - assert!((c - f).abs() < 1e-4, "logit {i}: cached {c} vs full {f}"); - } - } - - #[tokio::test] - async fn greedy_generate_respects_max_tokens_bound() { - let device = crate::backend::cpu_device(); - let cfg = toy_config(); - let loaded = LoadedLfm2 { - model: build(&cfg, &device), - config: cfg, - tokenizer_config: None, - }; - let out = loaded.greedy_generate(&[1, 2], 3, &device).await.unwrap(); - assert!(out.len() <= 3); - } - - /// A synthetic GGUF header shaped like the LFM2.5-1.2B file. - fn toy_gguf() -> GgufFile { - use crate::gguf::GgmlType; - let meta = |k: &str, v: GgufValue| (k.to_string(), v); - let kv_array = GgufValue::Array( - [0u32, 2, 0] - .iter() - .map(|&v| GgufValue::U32(v)) - .collect::>(), - ); - GgufFile { - path: std::path::PathBuf::new(), - version: 3, - metadata: vec![ - meta("general.architecture", GgufValue::Str("lfm2".into())), - meta("lfm2.embedding_length", GgufValue::U32(16)), - meta("lfm2.block_count", GgufValue::U32(3)), - meta("lfm2.feed_forward_length", GgufValue::U32(32)), - meta("lfm2.attention.head_count", GgufValue::U32(4)), - meta("lfm2.attention.head_count_kv", kv_array), - meta( - "lfm2.attention.layer_norm_rms_epsilon", - GgufValue::F32(1e-5), - ), - meta("lfm2.rope.freq_base", GgufValue::F32(1e6)), - meta("lfm2.shortconv.l_cache", GgufValue::U32(3)), - meta("tokenizer.ggml.eos_token_id", GgufValue::U32(7)), - ], - tensors: vec![GgufTensorInfo { - name: "token_embd.weight".into(), - dims: vec![16, 48], - dtype: GgmlType::F32, - offset: 0, - }], - alignment: 32, - data_offset: 0, - } - } - - #[test] - fn config_from_gguf_derives_layer_types_from_kv_array() { - let cfg = Lfm2Config::from_gguf(&toy_gguf()).expect("parses"); - assert_eq!( - cfg.layer_types, - vec!["conv", "full_attention", "conv"], - "0 = conv, nonzero = attention" - ); - assert_eq!(cfg.num_key_value_heads, 2); - assert_eq!(cfg.vocab_size, 48); // from token_embd dims[1] - assert_eq!(cfg.ff_dim(), 32); // GGUF value is pre-adjusted - assert!(cfg.eos_token_id.contains(7)); - } - - #[test] - fn config_from_gguf_rejects_mixed_kv_and_missing_array() { - let mut f = toy_gguf(); - f.metadata[5].1 = GgufValue::Array(vec![ - GgufValue::U32(4), - GgufValue::U32(2), - GgufValue::U32(0), - ]); - assert!(Lfm2Config::from_gguf(&f).unwrap_err().contains("mixes")); - - let mut f = toy_gguf(); - f.metadata - .retain(|(k, _)| k != "lfm2.attention.head_count_kv"); - assert!(Lfm2Config::from_gguf(&f).is_err()); - } - - #[test] - fn gguf_names_map_onto_hf_checkpoint_names_incl_conv_reshape() { - use crate::gguf::GgmlType; - let info = |name: &str, dims: &[u64]| GgufTensorInfo { - name: name.into(), - dims: dims.to_vec(), - dtype: GgmlType::F32, - offset: 0, - }; - let name_of = |i: &GgufTensorInfo| match gguf_tensor_to_hf(i) { - Some(GgufMap::Rename(n)) => n, - other => panic!("expected Rename, got {other:?}"), - }; - assert_eq!( - name_of(&info("blk.2.attn_q_norm.weight", &[64])), - "model.layers.2.self_attn.q_layernorm.weight" - ); - assert_eq!( - name_of(&info("blk.0.shortconv.in_proj.weight", &[2048, 6144])), - "model.layers.0.conv.in_proj.weight" - ); - assert_eq!( - name_of(&info("token_embd_norm.weight", &[2048])), - "model.embedding_norm.weight" - ); - // The conv kernel un-squeezes to the checkpoint's [channels, 1, k]. - match gguf_tensor_to_hf(&info("blk.0.shortconv.conv.weight", &[3, 2048])) { - Some(GgufMap::Reshape(n, shape)) => { - assert_eq!(n, "model.layers.0.conv.conv.weight"); - assert_eq!(shape, vec![2048, 1, 3]); - } - other => panic!("expected Reshape, got {other:?}"), - } - assert!(gguf_tensor_to_hf(&info("rope_freqs.weight", &[64])).is_none()); - } -} diff --git a/crates/mummu/examples/src/models/minilm.rs b/crates/mummu/examples/src/models/minilm.rs deleted file mode 100644 index 04af96f..0000000 --- a/crates/mummu/examples/src/models/minilm.rs +++ /dev/null @@ -1,391 +0,0 @@ -//! all-MiniLM-class sentence embedder: a 6-layer post-LayerNorm BERT with -//! absolute position embeddings, full bidirectional attention over an additive -//! padding mask, exact (erf) GeLU, masked-mean pooling + L2 normalize (the -//! sentence-transformers recipe — the checkpoint's `pooler.*` weights are -//! intentionally unused). -//! -//! Ported from laurelane's implementation (cosine ~1.0 vs the Candle BERT on -//! identical weights). BERT's bidirectional attention and LayerNorm differ -//! from the causal GQA blocks in `nn`, so this file is self-contained. -//! -//! The library takes token ids + attention mask and returns embeddings; the -//! caller owns tokenization (HF `tokenizers` in the integration tests). - -use std::path::Path; - -use burn::module::Module; -use burn::nn::{Embedding, EmbeddingConfig, LayerNorm, LayerNormConfig, Linear, LinearConfig}; -use burn::store::{ModuleAdapter, PyTorchToBurnAdapter, PytorchStore, SafetensorsStore}; -use burn::tensor::{Device, Int, Tensor, TensorData, activation}; - -use crate::import::{ - CastFloatAdapter, ImportError, WeightsFile, load_checked, required_file, weights_file, -}; - -/// BERT hyperparameters, read from the checkpoint's `config.json`. -#[derive(Debug, Clone, serde::Deserialize)] -pub struct BertConfig { - pub vocab_size: usize, - pub hidden_size: usize, - pub num_hidden_layers: usize, - pub num_attention_heads: usize, - pub intermediate_size: usize, - pub max_position_embeddings: usize, - pub type_vocab_size: usize, - #[serde(default = "default_eps")] - pub layer_norm_eps: f64, -} - -fn default_eps() -> f64 { - 1e-12 -} - -impl BertConfig { - /// Parse `config.json` bytes. - pub fn from_json_bytes(bytes: &[u8]) -> Result { - let cfg: Self = serde_json::from_slice(bytes).map_err(|e| e.to_string())?; - if cfg.num_attention_heads == 0 || !cfg.hidden_size.is_multiple_of(cfg.num_attention_heads) - { - return Err(format!( - "hidden_size ({}) must be a positive multiple of num_attention_heads ({})", - cfg.hidden_size, cfg.num_attention_heads - )); - } - Ok(cfg) - } -} - -/// Word + position + token-type embeddings, post-LN. -#[derive(Module, Debug)] -pub struct Embeddings { - pub word_embeddings: Embedding, - pub position_embeddings: Embedding, - pub token_type_embeddings: Embedding, - pub layer_norm: LayerNorm, -} - -/// One post-LN BERT encoder layer (flat field names; HF's nested keys remap -/// onto these at load). -#[derive(Module, Debug)] -pub struct EncoderLayer { - pub query: Linear, - pub key: Linear, - pub value: Linear, - pub attn_output: Linear, - pub attn_layer_norm: LayerNorm, - pub intermediate: Linear, - pub output: Linear, - pub output_layer_norm: LayerNorm, -} - -/// The BERT encoder stack. -#[derive(Module, Debug)] -pub struct Bert { - pub embeddings: Embeddings, - pub layers: Vec, -} - -/// A weight-loaded MiniLM/BERT plus its config. -pub struct LoadedMiniLm { - pub model: Bert, - pub config: BertConfig, -} - -fn build(cfg: &BertConfig, device: &Device) -> Bert { - let h = cfg.hidden_size; - let eps = cfg.layer_norm_eps; - let lin = |i: usize, o: usize| LinearConfig::new(i, o).init(device); - let ln = || LayerNormConfig::new(h).with_epsilon(eps).init(device); - let layers = (0..cfg.num_hidden_layers) - .map(|_| EncoderLayer { - query: lin(h, h), - key: lin(h, h), - value: lin(h, h), - attn_output: lin(h, h), - attn_layer_norm: ln(), - intermediate: lin(h, cfg.intermediate_size), - output: lin(cfg.intermediate_size, h), - output_layer_norm: ln(), - }) - .collect(); - Bert { - embeddings: Embeddings { - word_embeddings: EmbeddingConfig::new(cfg.vocab_size, h).init(device), - position_embeddings: EmbeddingConfig::new(cfg.max_position_embeddings, h).init(device), - token_type_embeddings: EmbeddingConfig::new(cfg.type_vocab_size, h).init(device), - layer_norm: ln(), - }, - layers, - } -} - -/// HF BERT checkpoint names → our field paths, shared by every weight format. -const KEY_REMAPS: &[(&str, &str)] = &[ - (r"^bert\.", ""), // some checkpoints prefix `bert.` - (r"^embeddings\.LayerNorm\.", "embeddings.layer_norm."), - ( - r"^encoder\.layer\.(\d+)\.attention\.self\.(query|key|value)\.", - "layers.$1.$2.", - ), - ( - r"^encoder\.layer\.(\d+)\.attention\.output\.dense\.", - "layers.$1.attn_output.", - ), - ( - r"^encoder\.layer\.(\d+)\.attention\.output\.LayerNorm\.", - "layers.$1.attn_layer_norm.", - ), - ( - r"^encoder\.layer\.(\d+)\.intermediate\.dense\.", - "layers.$1.intermediate.", - ), - ( - r"^encoder\.layer\.(\d+)\.output\.dense\.", - "layers.$1.output.", - ), - ( - r"^encoder\.layer\.(\d+)\.output\.LayerNorm\.", - "layers.$1.output_layer_norm.", - ), -]; - -/// Build from `dir/config.json` and load the checkpoint, checked — -/// `model.safetensors` preferred, `pytorch_model.bin` (the state dict this -/// model family originally shipped) as the fallback. -/// NOTE: LayerNorm keys keep their `.weight`/`.bias` suffixes — the -/// `PyTorchToBurnAdapter` renames those to `gamma`/`beta` itself (unlike -/// RmsNorm in the decoder models, where the rename is manual — the asymmetry -/// is intentional; `PytorchStore` applies that adapter internally). -/// `pooler.*` stays unused (masked-mean pooling instead). -pub fn load_from_dir(dir: &Path, device: &Device) -> Result { - let cfg_path = required_file(dir, "config.json")?; - let cfg_bytes = std::fs::read(&cfg_path).map_err(|e| ImportError::Parse { - file: cfg_path.clone(), - reason: e.to_string(), - })?; - let config = BertConfig::from_json_bytes(&cfg_bytes).map_err(|reason| ImportError::Parse { - file: cfg_path, - reason, - })?; - - let mut model = build(&config, device); - // The float dtype comes from the DEVICE — burn 0.22 keeps the element - // type there as a runtime setting, not on a backend type. Creation sites - // still name it explicitly rather than riding the unspecified default. - let target_float = crate::backend::float_dtype(device); - match weights_file(dir)? { - WeightsFile::Safetensors(weights) => { - let mut store = SafetensorsStore::from_file(weights.clone()) - .with_from_adapter(PyTorchToBurnAdapter.chain(CastFloatAdapter::new(target_float))) - .allow_partial(true); - for (pattern, replacement) in KEY_REMAPS { - store = store.with_key_remapping(*pattern, *replacement); - } - load_checked(&mut model, &mut store, &weights)?; - } - WeightsFile::PytorchBin(weights) => { - // No cast adapter on this path (PytorchStore has no adapter - // chaining) — .bin-era checkpoints are f32, which every backend - // ingests directly. - let mut store = PytorchStore::from_file(weights.clone()).allow_partial(true); - for (pattern, replacement) in KEY_REMAPS { - store = store.with_key_remapping(*pattern, *replacement); - } - load_checked(&mut model, &mut store, &weights)?; - } - } - Ok(LoadedMiniLm { model, config }) -} - -fn embeddings_forward(e: &Embeddings, ids: &Tensor<2, Int>, device: &Device) -> Tensor<3> { - let [b, n] = ids.dims(); - let w = e.word_embeddings.forward(ids.clone()); // [b, n, h] - let pos = Tensor::<1, Int>::arange(0..n as i64, device).reshape([1, n]); - let p = e.position_embeddings.forward(pos); // [1, n, h] - let tt = Tensor::<2, Int>::zeros([b, n], device); - let t = e.token_type_embeddings.forward(tt); // [b, n, h] - e.layer_norm.forward(w.add(p).add(t)) -} - -fn layer_forward( - l: &EncoderLayer, - x: Tensor<3>, - add_mask: &Tensor<4>, - num_heads: usize, -) -> Tensor<3> { - let [b, n, h] = x.dims(); - debug_assert!(h.is_multiple_of(num_heads), "hidden not divisible by heads"); - let hd = h / num_heads; - - let q = l - .query - .forward(x.clone()) - .reshape([b, n, num_heads, hd]) - .swap_dims(1, 2); - let k = l - .key - .forward(x.clone()) - .reshape([b, n, num_heads, hd]) - .swap_dims(1, 2); - let v = l - .value - .forward(x.clone()) - .reshape([b, n, num_heads, hd]) - .swap_dims(1, 2); - - let scale = 1.0 / (hd as f32).sqrt(); - let scores = q - .matmul(k.swap_dims(2, 3)) - .mul_scalar(scale) - .add(add_mask.clone()); - let probs = activation::softmax(scores, 3); - let ctx = probs.matmul(v).swap_dims(1, 2).reshape([b, n, h]); - - // Self-attention output: dense → LayerNorm(+ residual). - let x = l.attn_layer_norm.forward(l.attn_output.forward(ctx).add(x)); - // Feed-forward: dense → exact GeLU → dense → LayerNorm(+ residual). - let inter = activation::gelu(l.intermediate.forward(x.clone())); - l.output_layer_norm.forward(l.output.forward(inter).add(x)) -} - -impl LoadedMiniLm { - /// Encode one tokenized string (`ids` + `mask`, 1.0 = real token) into a - /// masked-mean-pooled, **L2-normalized** sentence embedding of - /// `hidden_size` floats. - pub fn embed_ids( - &self, - ids: &[u32], - mask: &[f32], - device: &Device, - ) -> Result, String> { - let n = ids.len(); - assert!(n >= 1, "embed_ids: empty token sequence"); - assert!( - mask.len() == n, - "embed_ids: mask length {} != ids length {n}", - mask.len() - ); - debug_assert!( - n <= self.config.max_position_embeddings, - "embed_ids: sequence longer than max_position_embeddings" - ); - - let ids32: Vec = ids.iter().map(|&i| i as i32).collect(); - // Dtypes pinned to the backend TYPE, never the per-device policy. - let id_t = Tensor::<1, Int>::from_data( - TensorData::new(ids32, [n]), - (device, crate::backend::int_dtype(device)), - ) - .reshape([1, n]); - let mask_t = Tensor::<1>::from_data( - TensorData::new(mask.to_vec(), [n]), - (device, crate::backend::float_dtype(device)), - ) - .reshape([1, n]); - - // Additive padding mask [1, 1, 1, n]: 0 for real tokens, large-negative - // for padding, broadcast across heads and query positions. - let add_mask = mask_t - .clone() - .reshape([1, 1, 1, n]) - .neg() - .add_scalar(1.0) - .mul_scalar(-1e30); - - let mut x = embeddings_forward(&self.model.embeddings, &id_t, device); - for layer in &self.model.layers { - x = layer_forward(layer, x, &add_mask, self.config.num_attention_heads); - } - - // Masked mean pool → L2 normalize. - let m3 = mask_t.reshape([1, n, 1]); - let summed = x.mul(m3.clone()).sum_dim(1); // [1, 1, h] - let counts = m3.sum_dim(1); // [1, 1, 1] - let pooled = summed.div(counts).reshape([1, self.config.hidden_size]); - let norm = pooled.clone().powf_scalar(2.0).sum_dim(1).sqrt(); // [1, 1] - let normalized = pooled.div(norm); - - normalized - .into_data() - .convert::() - .to_vec::() - .map_err(|e| format!("embedding readback: {e:?}")) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - fn toy_config() -> BertConfig { - BertConfig { - vocab_size: 50, - hidden_size: 16, - num_hidden_layers: 2, - num_attention_heads: 4, - intermediate_size: 32, - max_position_embeddings: 64, - type_vocab_size: 2, - layer_norm_eps: 1e-12, - } - } - - #[test] - fn config_rejects_indivisible_heads() { - let json = br#"{ - "vocab_size": 50, "hidden_size": 15, "num_hidden_layers": 1, - "num_attention_heads": 4, "intermediate_size": 32, - "max_position_embeddings": 64, "type_vocab_size": 2 - }"#; - assert!(BertConfig::from_json_bytes(json).is_err()); - } - - #[test] - fn embeddings_are_unit_norm() { - let device = crate::backend::cpu_device(); - let cfg = toy_config(); - let loaded = LoadedMiniLm { - model: build(&cfg, &device), - config: cfg, - }; - let ids = [3u32, 7, 12, 4]; - let mask = [1.0f32; 4]; - let e = loaded.embed_ids(&ids, &mask, &device).unwrap(); - assert_eq!(e.len(), 16); - let norm: f32 = e.iter().map(|v| v * v).sum::().sqrt(); - assert!((norm - 1.0).abs() < 1e-4, "L2 norm should be 1, got {norm}"); - } - - /// Padding must not change the embedding: the padded positions are masked - /// out of attention AND the mean pool. - #[test] - fn padding_is_invisible_to_the_embedding() { - let device = crate::backend::cpu_device(); - let cfg = toy_config(); - let loaded = LoadedMiniLm { - model: build(&cfg, &device), - config: cfg, - }; - let bare = loaded - .embed_ids(&[3, 7, 12], &[1.0, 1.0, 1.0], &device) - .unwrap(); - let padded = loaded - .embed_ids(&[3, 7, 12, 0, 0], &[1.0, 1.0, 1.0, 0.0, 0.0], &device) - .unwrap(); - for (i, (a, b)) in bare.iter().zip(&padded).enumerate() { - assert!((a - b).abs() < 1e-4, "elem {i}: bare {a} vs padded {b}"); - } - } - - #[test] - #[should_panic(expected = "mask length")] - fn embed_rejects_mismatched_mask() { - let device = crate::backend::cpu_device(); - let cfg = toy_config(); - let loaded = LoadedMiniLm { - model: build(&cfg, &device), - config: cfg, - }; - let _ = loaded.embed_ids(&[1, 2, 3], &[1.0, 1.0], &device); - } -} diff --git a/crates/mummu/examples/src/models/mod.rs b/crates/mummu/examples/src/models/mod.rs deleted file mode 100644 index 99ad1c8..0000000 --- a/crates/mummu/examples/src/models/mod.rs +++ /dev/null @@ -1,204 +0,0 @@ -//! The model zoo: from-scratch architectures on the shared `nn` blocks, all -//! generic over `B: Backend`, all config-driven (hyperparameters come from the -//! checkpoint's `config.json`, never hardcoded). - -use burn::tensor::{Device, Tensor}; - -use crate::decode::{SamplerOptions, argmax_id, generate_loop, top_k_ids}; - -pub mod lfm2; -pub mod minilm; -pub mod olmoe; -pub mod qwen2; -pub mod qwen3; -pub mod qwen35; - -/// Upper bound on one [`CausalLm::warm_up`] call. A warm-up is a fixed, -/// bounded cost paid off the user's critical path — not a place to spend -/// unbounded GPU time — and the measured curve flattens after ~32 steps -/// (`mummu-bench/tests/warmup_f16.rs`), so this ceiling is 8x the useful -/// depth, not a tuning knob. -pub const MAX_WARM_UP_STEPS: usize = 256; - -/// The contract every causal LM in the zoo implements. A new architecture -/// (Hermes-class function-caller, Gemma, Qwen3, …) provides its cache type, -/// its forward pass, and its EOS check — decoding (greedy, sampled, streamed, -/// cancellable) comes for free from the shared driver. -pub trait CausalLm { - /// Per-generation decode state (KV cache, conv state, …). - type Cache; - - /// A fresh (empty) cache for one generation. - fn new_cache(&self) -> Self::Cache; - - /// Forward `new_ids` (the whole prompt when `past == 0`, else one decode - /// token), updating `cache`; returns logits for the **last** position, - /// `[1, vocab]`. - fn forward( - &self, - new_ids: &[u32], - past: usize, - cache: &mut Self::Cache, - device: &Device, - ) -> Tensor<2>; - - /// Is `id` an end-of-sequence token for this model? - fn is_eos(&self, id: u32) -> bool; - - /// Full decode: prefill once, then one token per step, stopping at EOS, - /// `max_tokens`, or a `Break` from `on_token` (streaming + cooperative - /// cancellation). Greedy (`temperature == 0`) keeps the argmax on-device. - fn generate( - &self, - prompt_ids: &[u32], - max_tokens: usize, - opts: &SamplerOptions, - device: &Device, - on_token: impl FnMut(u32) -> std::ops::ControlFlow<()>, - ) -> impl std::future::Future, String>> { - async move { - let mut cache = self.new_cache(); - generate_loop( - |ids, past| self.forward(ids, past, &mut cache, device), - prompt_ids, - max_tokens, - opts, - |id| self.is_eos(id), - on_token, - ) - .await - } - } - - /// Greedy decode (the parity-gate path): [`Self::generate`] at - /// temperature 0 with no streaming. - fn greedy_generate( - &self, - prompt_ids: &[u32], - max_tokens: usize, - device: &Device, - ) -> impl std::future::Future, String>> { - // The options must outlive the future, so own them here rather than - // passing a temporary that dies at the end of this statement. - async move { - let opts = SamplerOptions::greedy(); - self.generate(prompt_ids, max_tokens, &opts, device, |_| { - std::ops::ControlFlow::Continue(()) - }) - .await - } - } - - /// Parity probe: top-k next-token ids for a single prefill. - fn first_token( - &self, - prompt_ids: &[u32], - k: usize, - device: &Device, - ) -> impl std::future::Future, String>> { - assert!(!prompt_ids.is_empty(), "first_token: empty prompt"); - assert!(k >= 1, "first_token: k must be >= 1"); - async move { - let mut cache = self.new_cache(); - let logits = self.forward(prompt_ids, 0, &mut cache, device); - let v = logits - .into_data_async() - .await - .map_err(|e| format!("logits readback: {e:?}"))? - .convert::() - .to_vec::() - .map_err(|e| format!("logits readback: {e:?}"))?; - Ok(top_k_ids(&v, k)) - } - } - - /// Post-import **sanity smoke**: one forward on `probe_ids` must yield - /// finite, correctly-sized (`expected_vocab`-wide), non-degenerate logits. - /// The liveness gate an app calls right after `install` to catch a - /// silently-broken import — corrupt weights (NaN), a config/tokenizer vocab - /// mismatch, or a dead/zero-init forward — none of which a checked *load* - /// can see. This is not parity (an arbitrary import has no reference); it - /// proves the model actually computes. See [`crate::import::logit_sanity`]. - fn sanity_check( - &self, - probe_ids: &[u32], - expected_vocab: usize, - device: &Device, - ) -> impl std::future::Future> { - assert!(!probe_ids.is_empty(), "sanity_check: empty probe prompt"); - assert!( - expected_vocab > 0, - "sanity_check: expected_vocab must be positive" - ); - async move { - let mut cache = self.new_cache(); - let logits = self.forward(probe_ids, 0, &mut cache, device); - let v = logits - .into_data_async() - .await - .map_err(|e| format!("logits readback: {e:?}"))? - .convert::() - .to_vec::() - .map_err(|e| format!("logits readback: {e:?}"))?; - crate::import::logit_sanity(&v, expected_vocab).map_err(|e| e.to_string()) - } - } - - /// Pay the **cold-start tax off the user's critical path**: one prefill - /// plus `steps` greedy decode steps on a throwaway cache, discarded. - /// - /// A freshly-started process decodes its first tokens far slower than its - /// steady state — measured on Qwen2.5-1.5B at f16, the first 32 tokens run - /// at 12.5 tok/s against a steady 37.6, and the curve is *flat* from token - /// 33 on (`mummu-bench/tests/warmup_f16.rs`). CubeCL already persists its - /// **autotune** choices to disk across processes, so what is left is - /// per-process kernel compilation and pipeline creation, which the wgpu - /// runtime does not cache anywhere (`CompilationCache` is wired for CUDA - /// and HIP only in cubecl 0.10) — no configuration can carry it, only - /// spending it earlier can. A consumer that opens short agent turns should - /// call this once after `install`/load, beside - /// [`Self::sanity_check`]. - /// - /// Warms the **decode** step, whose kernels are shape-stable (`t == 1`). - /// Prefill kernels are keyed by prompt length, so a caller who cares about - /// TTFT should pass a `probe_ids` of its own typical prompt length rather - /// than a token or two. - /// - /// Returns the number of forwards executed (`steps + 1`). Every step reads - /// its argmax back, exactly as real decoding does — an unsynchronized - /// warm-up would queue work and return before the GPU had run any of it. - fn warm_up( - &self, - probe_ids: &[u32], - steps: usize, - device: &Device, - ) -> impl std::future::Future> { - // Validate EAGERLY, outside the future: an argument bound that only - // fires when the future is awaited is a contract the caller can hold - // wrong indefinitely (and a `should_panic` test never sees). - assert!(!probe_ids.is_empty(), "warm_up: empty probe prompt"); - assert!(steps >= 1, "warm_up: steps must be >= 1"); - assert!( - steps <= MAX_WARM_UP_STEPS, - "warm_up: {steps} steps exceeds the {MAX_WARM_UP_STEPS} bound" - ); - async move { - let mut cache = self.new_cache(); - let logits = self.forward(probe_ids, 0, &mut cache, device); - let mut next = argmax_id(logits).await?; - let mut forwards = 1usize; - for past in (probe_ids.len()..).take(steps) { - let logits = self.forward(&[next], past, &mut cache, device); - next = argmax_id(logits).await?; - forwards += 1; - } - - debug_assert_eq!( - forwards, - steps + 1, - "warm_up must run exactly one prefill plus `steps` decode forwards" - ); - Ok(forwards) - } - } -} diff --git a/crates/mummu/examples/src/models/olmoe.rs b/crates/mummu/examples/src/models/olmoe.rs deleted file mode 100644 index de7fa98..0000000 --- a/crates/mummu/examples/src/models/olmoe.rs +++ /dev/null @@ -1,1592 +0,0 @@ -//! OLMoE sparse mixture-of-experts decoder (allenai OLMoE-1B-7B), from -//! scratch on the shared `nn` blocks — the zoo's first MoE architecture. -//! Structure per layer is pre-norm like Qwen, with two deltas: -//! * the FFN is a [`SparseMoe`] — a softmax top-k router (`k = 8` of 64) -//! over narrow SwiGLU experts, `norm_topk_prob = false`; -//! * q/k RMSNorm applies to the **whole projection** (width -//! `num_heads * head_dim`) before the head split — `GqaAttention`'s -//! projection placement (OLMoE is MHA: 16 query heads, 16 KV heads). -//! -//! This is the **resident-everything** first cut: all 64 experts' weights -//! live in memory and every expert computes every token (the router mask -//! zeroes the unrouted ones) — ~7B params in f32 is ~28 GB, which targets the -//! CPU backend on the reference 128 GB machine. Expert streaming / offload is -//! the P6 placement item; keep-quantized is P9. -//! -//! Import covers **both** sources. A GGUF ships the experts already fused -//! (`ffn_*_exps`) in exactly the layout `MoeExperts` holds; the HF safetensors -//! checkpoint stores each expert separately -//! (`mlp.experts.{i}.{gate,up,down}_proj.weight`) and is sharded, so -//! [`load_from_dir`] runs it through -//! [`crate::safetensors::fuse_checkpoint_to_file`], which reads every shard and -//! stacks each 64-member expert group into one `[experts, out, in]` tensor -//! before the ordinary checked-load pipeline. It fuses to a temp file rather -//! than to RAM deliberately: the in-memory twin would need the whole payload -//! resident (13.8 GB for the 1B-7B) on top of the ~28 GB f32 model the load -//! then builds. [`load_from_gguf`] makes the same trade for the same reason: -//! its dequant streams to a temp file (~28 GB of f32) that `burn-store` then -//! mmaps back, so the payload is file-backed rather than charged to commit. - -use std::path::Path; - -use burn::module::Module; -use burn::nn::{Embedding, EmbeddingConfig, Linear, LinearConfig, RmsNorm, RmsNormConfig}; -use burn::store::{ModuleAdapter, PyTorchToBurnAdapter, SafetensorsStore}; -use burn::tensor::{Device, Int, Tensor, TensorData}; - -use crate::attn_config::{RopeScaling, check_sliding_window, sliding_window_from_gguf}; -use crate::gguf::{GgufFile, GgufMap, GgufTensorInfo, GgufValue}; -use crate::import::{ - CastFloatAdapter, DequantSink, ImportError, ScratchFile, gguf_store, load_checked, - required_file, -}; -use crate::models::CausalLm; -use crate::models::qwen2::{EosIds, gguf_f32, gguf_usize}; -use crate::nn::{ - GqaAttention, GqaAttentionConfig, LayerKv, SparseMoe, SparseMoeConfig, causal_mask, rope_tables, -}; -use crate::safetensors::{Fuse, fuse_checkpoint_to_file}; - -/// OLMoE architecture hyperparameters (HF `config.json` field names). -#[derive(Debug, Clone, serde::Deserialize)] -pub struct OlmoeConfig { - pub vocab_size: usize, - pub hidden_size: usize, - /// Per-expert SwiGLU intermediate width (1B-7B: 1024). - pub intermediate_size: usize, - pub num_hidden_layers: usize, - pub num_attention_heads: usize, - pub num_key_value_heads: usize, - pub num_experts: usize, - pub num_experts_per_tok: usize, - /// Renormalize the top-k routing weights to sum 1 (OLMoE ships `false`: - /// the raw softmax probabilities weight the mixture). - #[serde(default)] - pub norm_topk_prob: bool, - pub rms_norm_eps: f64, - pub rope_theta: f32, - /// Frequency scaling (YaRN / linear / …). `null` on OLMoE-1B-7B; a scaled - /// checkpoint is refused at load ([`crate::attn_config`]). - /// `rope_parameters` is the same object under the name newer transformers - /// writes; reading only `rope_scaling` would let a freshly-serialized - /// scaled checkpoint through as unscaled. - #[serde(default, alias = "rope_parameters")] - pub rope_scaling: Option, - /// The trained context length, used to tell an inert sliding window from - /// a clipping one. - #[serde(default)] - pub max_position_embeddings: Option, - /// OLMoE has no `use_sliding_window` flag, so a declared window is live — - /// which is why `validate` passes `sliding_window.is_some()` as *enabled*. - #[serde(default)] - pub sliding_window: Option, - #[serde(default)] - pub tie_word_embeddings: bool, - #[serde(default)] - pub eos_token_id: EosIds, -} - -impl OlmoeConfig { - /// Parse `config.json` bytes. - pub fn from_json_bytes(bytes: &[u8]) -> Result { - let cfg: Self = serde_json::from_slice(bytes).map_err(|e| e.to_string())?; - cfg.validate("olmoe config.json")?; - Ok(cfg) - } - - /// Hyperparameters from a GGUF header's `olmoe.*` metadata. - pub fn from_gguf(f: &GgufFile) -> Result { - let arch = f.architecture().unwrap_or(""); - if arch != "olmoe" { - return Err(format!("GGUF architecture '{arch}' is not olmoe")); - } - let hidden_size = gguf_usize(f, "olmoe.embedding_length")?; - let embd = f - .tensor("token_embd.weight") - .ok_or("GGUF has no token_embd.weight tensor")?; - if embd.dims.len() != 2 || embd.dims[0] != hidden_size as u64 { - return Err(format!( - "token_embd.weight dims {:?} do not match embedding_length {hidden_size}", - embd.dims - )); - } - let vocab_size = usize::try_from(embd.dims[1]).map_err(|_| "vocab too large")?; - if let Some(tokens) = f.get("tokenizer.ggml.tokens").and_then(GgufValue::as_array) - && tokens.len() > vocab_size - { - return Err(format!( - "tokenizer vocab {} exceeds embedding rows {vocab_size}", - tokens.len() - )); - } - let eos_token_id = f - .get("tokenizer.ggml.eos_token_id") - .and_then(GgufValue::as_u64) - .and_then(|v| u32::try_from(v).ok()) - .map_or(EosIds::None, EosIds::One); - // llama.cpp writes the per-expert width as expert_feed_forward_length - // when it differs from feed_forward_length; accept either spelling. - let intermediate_size = gguf_usize(f, "olmoe.expert_feed_forward_length") - .or_else(|_| gguf_usize(f, "olmoe.feed_forward_length"))?; - // OLMoE does not renormalize top-k weights; honor the metadata key - // when a file carries one, default to the architecture's `false`. - let norm_topk_prob = f - .get("olmoe.expert_weights_norm") - .and_then(GgufValue::as_bool) - .unwrap_or(false); - let cfg = Self { - vocab_size, - hidden_size, - intermediate_size, - num_hidden_layers: gguf_usize(f, "olmoe.block_count")?, - num_attention_heads: gguf_usize(f, "olmoe.attention.head_count")?, - num_key_value_heads: gguf_usize(f, "olmoe.attention.head_count_kv")?, - num_experts: gguf_usize(f, "olmoe.expert_count")?, - num_experts_per_tok: gguf_usize(f, "olmoe.expert_used_count")?, - norm_topk_prob, - rms_norm_eps: f64::from(gguf_f32(f, "olmoe.attention.layer_norm_rms_epsilon")?), - rope_theta: gguf_f32(f, "olmoe.rope.freq_base")?, - rope_scaling: RopeScaling::from_gguf(f, "olmoe"), - max_position_embeddings: f - .get("olmoe.context_length") - .and_then(GgufValue::as_u64) - .and_then(|v| usize::try_from(v).ok()), - sliding_window: sliding_window_from_gguf(f, "olmoe"), - // No separate output.weight tensor means the lm-head is tied. - tie_word_embeddings: f.tensor("output.weight").is_none(), - eos_token_id, - }; - cfg.validate("olmoe GGUF header")?; - Ok(cfg) - } - - /// `hidden_size / num_attention_heads` — OLMoE's head_dim is not decoupled. - #[must_use] - pub fn head_dim(&self) -> usize { - self.hidden_size / self.num_attention_heads.max(1) - } - - fn validate(&self, whose: &str) -> Result<(), String> { - if let Some(scaling) = &self.rope_scaling { - scaling.check(whose)?; - } - check_sliding_window( - self.sliding_window, - // No `use_sliding_window` gate in this family: a declared window - // is a live one. - self.sliding_window.is_some(), - self.max_position_embeddings, - whose, - )?; - if self.num_key_value_heads == 0 - || !self - .num_attention_heads - .is_multiple_of(self.num_key_value_heads) - { - return Err(format!( - "num_attention_heads ({}) must be a positive multiple of num_key_value_heads ({})", - self.num_attention_heads, self.num_key_value_heads - )); - } - if self.num_hidden_layers == 0 || self.vocab_size == 0 { - return Err("num_hidden_layers and vocab_size must be positive".into()); - } - if self.num_experts < 2 || !(1..=self.num_experts).contains(&self.num_experts_per_tok) { - return Err(format!( - "num_experts ({}) must be >= 2 with num_experts_per_tok ({}) in 1..=num_experts", - self.num_experts, self.num_experts_per_tok - )); - } - let hd = self.head_dim(); - if hd < 2 || !hd.is_multiple_of(2) || hd * self.num_attention_heads != self.hidden_size { - return Err(format!( - "hidden_size ({}) must split evenly into num_attention_heads ({}) even-sized heads", - self.hidden_size, self.num_attention_heads - )); - } - Ok(()) - } -} - -/// One OLMoE decoder layer. Field names mirror the HF checkpoint layout. -#[derive(Module, Debug)] -pub struct DecoderLayer { - pub self_attn: GqaAttention, - pub mlp: SparseMoe, - pub input_layernorm: RmsNorm, - pub post_attention_layernorm: RmsNorm, -} - -/// The OLMoE decoder stack (HF's `model.*` subtree). The 1B-7B ships untied. -#[derive(Module, Debug)] -pub struct Olmoe { - pub embed_tokens: Embedding, - pub layers: Vec, - pub norm: RmsNorm, - pub lm_head: Option, -} - -/// A weight-loaded OLMoE plus its config — everything a forward needs. -pub struct LoadedOlmoe { - pub model: Olmoe, - pub config: OlmoeConfig, - /// The sibling `tokenizer_config.json`, when the checkpoint dir ships one - /// — config-driven EOS/BOS/PAD for a consumer to read. `None` for a GGUF - /// load (self-contained: EOS rides the GGUF metadata). - pub tokenizer_config: Option, -} - -fn build(cfg: &OlmoeConfig, device: &Device) -> Olmoe { - let attn_cfg = GqaAttentionConfig { - hidden_size: cfg.hidden_size, - num_heads: cfg.num_attention_heads, - num_kv_heads: cfg.num_key_value_heads, - head_dim: cfg.head_dim(), - bias: false, // OLMoE projections are bias-free - qk_norm_eps: Some(cfg.rms_norm_eps), // q/k RMSNorm over the whole - qk_norm_projection: true, // projection, pre head-split - }; - let moe_cfg = SparseMoeConfig { - hidden_size: cfg.hidden_size, - expert_intermediate_size: cfg.intermediate_size, - num_experts: cfg.num_experts, - num_experts_per_tok: cfg.num_experts_per_tok, - }; - let norm = |dev: &Device| { - RmsNormConfig::new(cfg.hidden_size) - .with_epsilon(cfg.rms_norm_eps) - .init(dev) - }; - let layers = (0..cfg.num_hidden_layers) - .map(|_| DecoderLayer { - self_attn: attn_cfg.init(device), - mlp: moe_cfg.init(device), - input_layernorm: norm(device), - post_attention_layernorm: norm(device), - }) - .collect(); - let lm_head = (!cfg.tie_word_embeddings).then(|| { - LinearConfig::new(cfg.hidden_size, cfg.vocab_size) - .with_bias(false) - .init(device) - }); - Olmoe { - embed_tokens: EmbeddingConfig::new(cfg.vocab_size, cfg.hidden_size).init(device), - layers, - norm: norm(device), - lm_head, - } -} - -/// The key remap: strip `model.`, rename every RmsNorm `weight` → Burn's -/// `gamma` (incl. the projection-wide `self_attn.{q,k}_norm`). The fused -/// expert params (`mlp.experts.{gate,up,down}`) carry no `.weight` suffix — -/// they are raw `Param` fields, named by the GGUF map directly. -fn install_remaps(store: SafetensorsStore) -> SafetensorsStore { - store - .with_key_remapping(r"^model\.", "") - .with_key_remapping(r"(input_layernorm)\.weight$", "$1.gamma") - .with_key_remapping(r"(post_attention_layernorm)\.weight$", "$1.gamma") - .with_key_remapping(r"(self_attn\.q_norm)\.weight$", "$1.gamma") - .with_key_remapping(r"(self_attn\.k_norm)\.weight$", "$1.gamma") - .with_key_remapping(r"^norm\.weight$", "norm.gamma") -} - -/// GGUF (llama.cpp `olmoe` arch) tensor names → the HF-shaped names the remap -/// chain handles. `None` for anything unrecognized (a loud load error). -fn gguf_tensor_to_hf(info: &GgufTensorInfo) -> Option { - olmoe_gguf_name(&info.name).map(GgufMap::Rename) -} - -fn olmoe_gguf_name(name: &str) -> Option { - match name { - "token_embd.weight" => return Some("model.embed_tokens.weight".into()), - "output_norm.weight" => return Some("model.norm.weight".into()), - "output.weight" => return Some("lm_head.weight".into()), - _ => {} - } - let rest = name.strip_prefix("blk.")?; - let (layer, field) = rest.split_once('.')?; - let layer: usize = layer.parse().ok()?; - let mapped = match field { - "attn_norm.weight" => "input_layernorm.weight", - "ffn_norm.weight" => "post_attention_layernorm.weight", - "attn_q.weight" => "self_attn.q_proj.weight", - "attn_k.weight" => "self_attn.k_proj.weight", - "attn_v.weight" => "self_attn.v_proj.weight", - "attn_q_norm.weight" => "self_attn.q_norm.weight", - "attn_k_norm.weight" => "self_attn.k_norm.weight", - "attn_output.weight" => "self_attn.o_proj.weight", - // The router Linear — a plain 2-D weight, transposed by the adapter. - "ffn_gate_inp.weight" => "mlp.gate.weight", - // The fused 3-D expert banks — raw params, no `.weight` suffix in the - // module path. ggml dims reverse to [experts, out, in], exactly the - // `MoeExperts` layout. - "ffn_gate_exps.weight" => "mlp.experts.gate", - "ffn_up_exps.weight" => "mlp.experts.up", - "ffn_down_exps.weight" => "mlp.experts.down", - _ => return None, - }; - Some(format!("model.layers.{layer}.{mapped}")) -} - -/// Load an OLMoE model straight from a **GGUF** file: hyperparameters from -/// the `olmoe.*` metadata, weights dequantized to f32 and driven through the -/// same checked-load pipeline every other port uses. Budget note: the 1B-7B's -/// ~7B params dequantize to ~28 GB of f32 — size the target device (the -/// reference machine runs it on the 128 GB CPU backend). -/// -/// The dequant goes to a TEMP FILE, not to RAM, for the same reason -/// `load_from_dir`'s fuse does: holding ~28 GB of f32 payload *and* building -/// the ~28 GB model from it is the sum a 128 GB box with other tenants -/// actually fails to satisfy. Streaming keeps the peak at the model plus one -/// tensor. The file is this process's to delete, on success or failure. -pub fn load_from_gguf(path: &Path, device: &Device) -> Result { - let parse = |reason: String| ImportError::Parse { - file: path.to_path_buf(), - reason, - }; - let f = GgufFile::open(path).map_err(|e| parse(e.to_string()))?; - let config = OlmoeConfig::from_gguf(&f).map_err(parse)?; - // The scratch guard (Some only when the payload went to disk, which at - // OLMoE's size it always does) must outlive `load_checked`: the store - // reads that file lazily. - let (base, _scratch) = gguf_store(&f, &gguf_tensor_to_hf, DequantSink::Auto, device)?; - - let mut model = build(&config, device); - let mut store = install_remaps(base); - load_checked(&mut model, &mut store, path)?; - Ok(LoadedOlmoe { - model, - config, - tokenizer_config: None, - }) -} - -/// How a source tensor of an HF OLMoE checkpoint reaches the module. -/// -/// Everything but the expert bank passes through under its own name (the -/// `install_remaps` chain does the HF→module renaming downstream, exactly as -/// on the single-file safetensors path). The per-expert projections are the -/// N:1 case: `model.layers.3.mlp.experts.7.gate_proj.weight` is member 7 of -/// the fused `model.layers.3.mlp.experts.gate`. -/// -/// Deliberately *not* a strict allow-list: unrecognized tensors are kept, and -/// a checkpoint that renames the expert projections then fails loudly one -/// stage later in `load_checked` — which names the missing `experts.gate` -/// param in its report — rather than here with a less specific message. -fn olmoe_hf_fuse(name: &str, num_experts: usize) -> Fuse { - let Some(target) = fused_expert_target(name, num_experts) else { - return Fuse::Keep(name.to_string()); - }; - target -} - -/// `model.layers.{L}.mlp.experts.{I}.{gate,up,down}_proj.weight` → its slot in -/// the fused bank, or `None` when the name is not a per-expert projection. -fn fused_expert_target(name: &str, num_experts: usize) -> Option { - let rest = name.strip_prefix("model.layers.")?; - let (layer, rest) = rest.split_once('.')?; - layer.parse::().ok()?; - let rest = rest.strip_prefix("mlp.experts.")?; - let (index, rest) = rest.split_once('.')?; - let index: usize = index.parse().ok()?; - let projection = match rest { - "gate_proj.weight" => "gate", - "up_proj.weight" => "up", - "down_proj.weight" => "down", - _ => return None, - }; - Some(Fuse::Stack { - target: format!("model.layers.{layer}.mlp.experts.{projection}"), - index, - count: num_experts, - }) -} - -/// Load an OLMoE model from an **HF safetensors checkpoint dir**. -/// -/// The checkpoint is sharded (`model-0000N-of-0000M.safetensors` + -/// `model.safetensors.index.json`) and stores every expert separately, so the -/// shards are read and the expert groups fused into `[experts, out, in]` -/// tensors first; the fused blob then rides the SAME adapter chain and -/// `load_checked` as every other import path. Source dtype is preserved by -/// the fuse (HF ships bf16) and cast to the backend float on load. -/// -/// Budget note: the fused blob is the checkpoint's own size (~13.8 GB in bf16 -/// for the 1B-7B) and the loaded f32 model is ~28 GB — size the target device. -pub fn load_from_dir(dir: &Path, device: &Device) -> Result { - let cfg_path = required_file(dir, "config.json")?; - let cfg_bytes = std::fs::read(&cfg_path).map_err(|e| ImportError::Parse { - file: cfg_path.clone(), - reason: e.to_string(), - })?; - let config = OlmoeConfig::from_json_bytes(&cfg_bytes).map_err(|reason| ImportError::Parse { - file: cfg_path, - reason, - })?; - - // Cross-check the sibling metadata before touching weights, same as the - // dense loaders. `None` for the expected tool-call convention: Mummu ships - // no hardcoded OLMoE `chat` renderer to contradict, so only the EOS and - // added-token-id checks apply. - let tokenizer_config = - crate::tokenizer::validate_checkpoint_dir(dir, &config.eos_token_id.to_vec(), None)?; - - let num_experts = config.num_experts; - // Fuse to a TEMP FILE, not to RAM. The in-memory fuse needs the whole - // payload resident (13.8 GB for the 1B-7B) on top of the ~28 GB f32 model - // the load then builds; streaming it to disk keeps the peak at the model - // alone. The file is this process's to delete, and it is deleted whether - // the load succeeds or fails. - let fused = ScratchFile::new(dir)?; - let bytes = fuse_checkpoint_to_file( - dir, - &|name| Some(olmoe_hf_fuse(name, num_experts)), - fused.path(), - ) - .map_err(|e| ImportError::Parse { - file: dir.to_path_buf(), - reason: e.to_string(), - })?; - assert!(bytes > 8, "a fused checkpoint yields a non-empty payload"); - - let mut model = build(&config, device); - // The float dtype comes from the DEVICE — burn 0.22 keeps the element - // type there as a runtime setting, not on a backend type. Creation sites - // still name it explicitly rather than riding the unspecified default. - let target_float = crate::backend::float_dtype(device); - let mut store = install_remaps( - SafetensorsStore::from_file(fused.path().to_path_buf()) - .with_from_adapter(PyTorchToBurnAdapter.chain(CastFloatAdapter::new(target_float))) - .allow_partial(true), - ); - load_checked(&mut model, &mut store, dir)?; - Ok(LoadedOlmoe { - model, - config, - tokenizer_config, - }) -} - -impl CausalLm for LoadedOlmoe { - type Cache = Vec; - - fn new_cache(&self) -> Self::Cache { - (0..self.config.num_hidden_layers).map(|_| None).collect() - } - - fn is_eos(&self, id: u32) -> bool { - self.config.eos_token_id.contains(id) - } - - fn forward( - &self, - new_ids: &[u32], - past: usize, - cache: &mut Self::Cache, - device: &Device, - ) -> Tensor<2> { - let t = new_ids.len(); - assert!(t >= 1, "OLMoE forward: need at least one token"); - assert!( - cache.len() == self.config.num_hidden_layers, - "OLMoE forward: cache has {} layers, model has {}", - cache.len(), - self.config.num_hidden_layers - ); - let cfg = &self.config; - let hd = cfg.head_dim(); - - // Dtype pinned to the backend TYPE, never the per-device policy. - let ids32: Vec = new_ids.iter().map(|&i| i as i32).collect(); - let input = Tensor::<1, Int>::from_data( - TensorData::new(ids32, [t]), - (device, crate::backend::int_dtype(device)), - ) - .reshape([1, t]); - let mut x = self.model.embed_tokens.forward(input); // [1, t, hidden] - - let (cos, sin) = rope_tables(t, past, hd, cfg.rope_theta, device); - let mask = (t > 1).then(|| causal_mask(t, past, device)); - - for (layer, kv) in self.model.layers.iter().zip(cache.iter_mut()) { - let h = layer.input_layernorm.forward(x.clone()); - let h = layer.self_attn.forward( - h, - cfg.num_attention_heads, - cfg.num_key_value_heads, - hd, - &cos, - &sin, - mask.as_ref(), - kv, - ); - x = x.add(h); - let h2 = layer.post_attention_layernorm.forward(x.clone()); - x = x.add( - layer - .mlp - .forward(h2, cfg.num_experts_per_tok, cfg.norm_topk_prob), - ); - } - let x = self.model.norm.forward(x); - - let last = x.narrow(1, t - 1, 1).reshape([1, cfg.hidden_size]); - debug_assert!( - self.model.lm_head.is_some() != cfg.tie_word_embeddings, - "lm_head presence must match the config's tie flag" - ); - match &self.model.lm_head { - Some(head) => head.forward(last), // [1, vocab] - None => { - let w = self.model.embed_tokens.weight.val(); // [vocab, hidden] - last.matmul(w.swap_dims(0, 1)) // [1, vocab] - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::gguf::{GgmlType, GgufTensorInfo}; - - /// A synthetic toy MoE config: 4 experts, top-2, MHA, untied head. - fn toy_config() -> OlmoeConfig { - OlmoeConfig { - vocab_size: 64, - hidden_size: 16, - intermediate_size: 8, - num_hidden_layers: 2, - num_attention_heads: 4, - num_key_value_heads: 4, - num_experts: 4, - num_experts_per_tok: 2, - norm_topk_prob: false, - rms_norm_eps: 1e-5, - rope_theta: 1e4, - rope_scaling: None, - max_position_embeddings: Some(512), - sliding_window: None, - tie_word_embeddings: false, - eos_token_id: EosIds::One(2), - } - } - - #[test] - fn config_parses_the_real_1b_7b_shape() { - // The real OLMoE-1B-7B-0125-Instruct config.json shape. - let json = br#"{ - "vocab_size": 50304, "hidden_size": 2048, "intermediate_size": 1024, - "num_hidden_layers": 16, "num_attention_heads": 16, "num_key_value_heads": 16, - "num_experts": 64, "num_experts_per_tok": 8, "norm_topk_prob": false, - "rms_norm_eps": 1e-05, "rope_theta": 10000.0, - "tie_word_embeddings": false, "eos_token_id": 50279 - }"#; - let cfg = OlmoeConfig::from_json_bytes(json).unwrap(); - assert_eq!(cfg.head_dim(), 128); - assert_eq!(cfg.num_experts, 64); - assert_eq!(cfg.num_experts_per_tok, 8); - assert!(!cfg.norm_topk_prob); - assert!(!cfg.tie_word_embeddings); - assert!(cfg.eos_token_id.contains(50_279)); - assert!(cfg.rope_scaling.is_none() && cfg.sliding_window.is_none()); - } - - /// OLMoE has no `use_sliding_window` gate, so a declared window is a live - /// one — the opposite of Qwen2.5's inert field, and the reason each family - /// decides *enabled* for itself rather than sharing one guess. - #[test] - fn a_declared_window_is_live_in_a_family_without_the_gate() { - let json = br#"{ - "vocab_size": 50304, "hidden_size": 2048, "intermediate_size": 1024, - "num_hidden_layers": 16, "num_attention_heads": 16, "num_key_value_heads": 16, - "num_experts": 64, "num_experts_per_tok": 8, "rms_norm_eps": 1e-05, - "rope_theta": 10000.0, "max_position_embeddings": 4096, "sliding_window": 1024 - }"#; - let err = OlmoeConfig::from_json_bytes(json).expect_err("a live window must refuse"); - assert!(err.contains("1024"), "{err}"); - } - - #[test] - fn config_rejects_bad_expert_counts() { - let mut cfg = toy_config(); - cfg.num_experts = 1; - assert!(cfg.validate("test").is_err(), "one expert is not a mixture"); - let mut cfg = toy_config(); - cfg.num_experts_per_tok = 5; - assert!(cfg.validate("test").is_err(), "top-k above expert count"); - let mut cfg = toy_config(); - cfg.num_experts_per_tok = 0; - assert!(cfg.validate("test").is_err(), "zero top-k"); - } - - /// The load-bearing invariant: cached prefill+decode == one full forward, - /// through the MoE layers and the projection-wide q/k norm. - #[test] - fn toy_model_cached_decode_matches_full_forward() { - let device = crate::backend::cpu_device(); - let cfg = toy_config(); - let loaded = LoadedOlmoe { - model: build(&cfg, &device), - config: cfg, - tokenizer_config: None, - }; - - let prompt: Vec = vec![3, 14, 15, 9, 26]; - let mut cache = loaded.new_cache(); - let _ = loaded.forward(&prompt, 0, &mut cache, &device); - let step = loaded - .forward(&[42], prompt.len(), &mut cache, &device) - .into_data() - .to_vec::() - .unwrap(); - - let mut full_cache = loaded.new_cache(); - let all: Vec = prompt.iter().copied().chain([42]).collect(); - let full = loaded - .forward(&all, 0, &mut full_cache, &device) - .into_data() - .to_vec::() - .unwrap(); - - assert_eq!(step.len(), full.len()); - for (i, (c, f)) in step.iter().zip(&full).enumerate() { - assert!((c - f).abs() < 1e-4, "logit {i}: cached {c} vs full {f}"); - } - } - - #[test] - fn projection_qk_norm_spans_the_whole_projection() { - let device = crate::backend::cpu_device(); - let cfg = toy_config(); - let model = build(&cfg, &device); - let q_dim = cfg.num_attention_heads * cfg.head_dim(); - assert_eq!( - model.layers[0] - .self_attn - .q_norm - .as_ref() - .unwrap() - .gamma - .dims(), - [q_dim], - "OLMoE q_norm must span num_heads * head_dim, not head_dim" - ); - } - - /// A synthetic in-memory GGUF header shaped like a small `olmoe` file. - fn toy_gguf() -> GgufFile { - let meta = |k: &str, v: GgufValue| (k.to_string(), v); - GgufFile { - path: std::path::PathBuf::new(), - version: 3, - metadata: vec![ - meta("general.architecture", GgufValue::Str("olmoe".into())), - meta("olmoe.embedding_length", GgufValue::U32(16)), - meta("olmoe.block_count", GgufValue::U32(2)), - meta("olmoe.feed_forward_length", GgufValue::U32(8)), - meta("olmoe.attention.head_count", GgufValue::U32(4)), - meta("olmoe.attention.head_count_kv", GgufValue::U32(4)), - meta("olmoe.expert_count", GgufValue::U32(4)), - meta("olmoe.expert_used_count", GgufValue::U32(2)), - meta( - "olmoe.attention.layer_norm_rms_epsilon", - GgufValue::F32(1e-5), - ), - meta("olmoe.rope.freq_base", GgufValue::F32(1e4)), - meta("tokenizer.ggml.eos_token_id", GgufValue::U32(2)), - ], - tensors: vec![ - GgufTensorInfo { - name: "token_embd.weight".into(), - dims: vec![16, 64], // ggml order: [hidden, vocab] - dtype: GgmlType::F32, - offset: 0, - }, - GgufTensorInfo { - name: "output.weight".into(), - dims: vec![16, 64], - dtype: GgmlType::F32, - offset: 4096, - }, - ], - alignment: 32, - data_offset: 0, - } - } - - #[test] - fn config_from_gguf_reads_expert_metadata() { - let cfg = OlmoeConfig::from_gguf(&toy_gguf()).expect("parses"); - assert_eq!(cfg.num_experts, 4); - assert_eq!(cfg.num_experts_per_tok, 2); - assert_eq!(cfg.intermediate_size, 8); - assert!(!cfg.norm_topk_prob); - assert!(!cfg.tie_word_embeddings); // output.weight present → untied - assert!(cfg.eos_token_id.contains(2)); - } - - #[test] - fn config_from_gguf_prefers_expert_feed_forward_length() { - let mut f = toy_gguf(); - f.metadata.push(( - "olmoe.expert_feed_forward_length".into(), - GgufValue::U32(12), - )); - assert_eq!(OlmoeConfig::from_gguf(&f).unwrap().intermediate_size, 12); - } - - #[test] - fn config_from_gguf_fails_loudly_on_missing_keys_and_wrong_arch() { - let mut f = toy_gguf(); - f.metadata.retain(|(k, _)| k != "olmoe.expert_count"); - assert!( - OlmoeConfig::from_gguf(&f) - .unwrap_err() - .contains("expert_count") - ); - - let mut f = toy_gguf(); - f.metadata[0].1 = GgufValue::Str("qwen3".into()); - assert!(OlmoeConfig::from_gguf(&f).is_err()); - } - - #[test] - fn gguf_names_map_router_and_fused_expert_banks() { - assert_eq!( - olmoe_gguf_name("blk.0.ffn_gate_inp.weight").as_deref(), - Some("model.layers.0.mlp.gate.weight") - ); - // Fused expert banks are raw params — no `.weight` suffix. - assert_eq!( - olmoe_gguf_name("blk.3.ffn_gate_exps.weight").as_deref(), - Some("model.layers.3.mlp.experts.gate") - ); - assert_eq!( - olmoe_gguf_name("blk.15.ffn_down_exps.weight").as_deref(), - Some("model.layers.15.mlp.experts.down") - ); - assert_eq!( - olmoe_gguf_name("blk.1.attn_q_norm.weight").as_deref(), - Some("model.layers.1.self_attn.q_norm.weight") - ); - assert_eq!( - olmoe_gguf_name("output.weight").as_deref(), - Some("lm_head.weight") - ); - // A dense FFN tensor is not part of this architecture — loud None. - assert_eq!(olmoe_gguf_name("blk.0.ffn_gate.weight"), None); - assert_eq!(olmoe_gguf_name("rope_freqs.weight"), None); - } - - #[tokio::test] - async fn greedy_generate_respects_max_tokens_bound() { - let device = crate::backend::cpu_device(); - let cfg = toy_config(); - let loaded = LoadedOlmoe { - model: build(&cfg, &device), - config: cfg, - tokenizer_config: None, - }; - let out = loaded - .greedy_generate(&[1, 2, 3], 4, &device) - .await - .unwrap(); - assert!(out.len() <= 4); - } - - /// The HF per-expert projections fuse onto EXACTLY the module names the - /// GGUF path renames its pre-fused banks to. Both import paths must land - /// on the same params, or one of them is loading a different model. - #[test] - fn hf_expert_fusion_targets_match_the_gguf_names() { - for (projection, ggml) in [ - ("gate", "ffn_gate_exps"), - ("up", "ffn_up_exps"), - ("down", "ffn_down_exps"), - ] { - let hf = format!("model.layers.3.mlp.experts.7.{projection}_proj.weight"); - let expected = olmoe_gguf_name(&format!("blk.3.{ggml}.weight")).unwrap(); - assert_eq!( - olmoe_hf_fuse(&hf, 64), - Fuse::Stack { - target: expected, - index: 7, - count: 64, - }, - "{projection}: safetensors and GGUF must fuse to the same param" - ); - } - } - - /// Non-expert tensors pass through untouched — the `install_remaps` chain - /// does the HF→module renaming downstream, same as the dense loaders. - #[test] - fn non_expert_tensors_pass_through_by_name() { - for name in [ - "model.embed_tokens.weight", - "model.norm.weight", - "lm_head.weight", - "model.layers.0.self_attn.q_proj.weight", - "model.layers.0.self_attn.q_norm.weight", - "model.layers.0.input_layernorm.weight", - // The ROUTER is a plain Linear, not part of the expert bank — - // fusing it would be a silent disaster. - "model.layers.0.mlp.gate.weight", - ] { - assert_eq!( - olmoe_hf_fuse(name, 64), - Fuse::Keep(name.to_string()), - "{name} must pass through" - ); - } - } - - /// The expert index is parsed as a NUMBER, so the fuse plan can order the - /// bank numerically; `experts.10` must not be read as expert 1. - #[test] - fn expert_index_is_parsed_numerically() { - let at = |i: usize| match olmoe_hf_fuse( - &format!("model.layers.0.mlp.experts.{i}.gate_proj.weight"), - 64, - ) { - Fuse::Stack { index, .. } => index, - other => panic!("expected a Stack, got {other:?}"), - }; - assert_eq!(at(0), 0); - assert_eq!(at(1), 1); - assert_eq!(at(10), 10); - assert_eq!(at(63), 63); - // A non-numeric member is not an expert projection at all. - assert_eq!( - olmoe_hf_fuse("model.layers.0.mlp.experts.x.gate_proj.weight", 64), - Fuse::Keep("model.layers.0.mlp.experts.x.gate_proj.weight".to_string()) - ); - } -} - -// =========================================================================== -// P9 — the per-expert-quantized OLMoE variant. A separate model struct so -// the parity-proven fused-bank path above stays byte-for-byte untouched: -// experts live as independent (quantized) weight triples and compute -// through `SparseMoePerExpert`'s routed forward — per token only the -// top-k experts are in service. Attention, router, norms, embedding and -// head stay float (GqaAttention's Linears would hit burn 0.21's -// packed-weight reshape bug, and they are a small fraction of the bytes). -// =========================================================================== - -/// One decoder layer of the quantized variant. -#[derive(Module, Debug)] -pub struct QDecoderLayer { - pub self_attn: GqaAttention, - pub mlp: crate::nn::SparseMoePerExpert, - pub input_layernorm: RmsNorm, - pub post_attention_layernorm: RmsNorm, -} - -/// The per-expert-quantized OLMoE stack. -#[derive(Module, Debug)] -pub struct OlmoeQ { - pub embed_tokens: Embedding, - pub layers: Vec, - pub norm: RmsNorm, - pub lm_head: Option, -} - -/// A weight-loaded quantized OLMoE plus its config. With a `pool`, the -/// experts in `model` are placeholders and every layer's expert compute -/// goes through the tiered [`crate::nn::ExpertPool`] (P9 stage 3b); the -/// router, attention, norms, embedding and head stay on `B`. -pub struct LoadedOlmoeQ { - pub model: OlmoeQ, - pub config: OlmoeConfig, - pub pool: Option>, -} - -/// **Streaming** GGUF import with per-expert keep-quantized experts: each -/// fused bank is read once, split into its `num_experts` contiguous -/// members, and every member is **re-quantized independently** (its own -/// block scales) per `policy`. Peak memory = the finished model plus one -/// f32 bank. -pub fn load_from_gguf_quantized( - path: &Path, - device: &Device, - policy: crate::quant::QuantPolicy, -) -> Result { - use crate::nn::{ExpertWeights, SparseMoePerExpert}; - use burn::module::Param; - - let parse = |reason: String| ImportError::Parse { - file: path.to_path_buf(), - reason, - }; - let f = GgufFile::open(path).map_err(|e| parse(e.to_string()))?; - let config = OlmoeConfig::from_gguf(&f).map_err(parse)?; - let untied = f.tensor("output.weight").is_some(); - - let dtype = crate::backend::float_dtype(device); - let dev_tensor2 = |values: Vec, shape: [usize; 2]| { - Tensor::<2>::from_data(TensorData::new(values, shape), (device, dtype)) - }; - let dev_tensor1 = |values: Vec, n: usize| { - Tensor::<1>::from_data(TensorData::new(values, [n]), (device, dtype)) - }; - // Tiny placeholders; the completeness count below guarantees every one - // is replaced before the model is returned. - let placeholder2 = || Param::from_tensor(Tensor::<2>::zeros([1, 1], device)); - - let norm = |dev: &Device| { - RmsNormConfig::new(config.hidden_size) - .with_epsilon(config.rms_norm_eps) - .init(dev) - }; - let attn_cfg = GqaAttentionConfig { - hidden_size: config.hidden_size, - num_heads: config.num_attention_heads, - num_kv_heads: config.num_key_value_heads, - head_dim: config.head_dim(), - bias: false, - qk_norm_eps: Some(config.rms_norm_eps), - qk_norm_projection: true, - }; - let mut model = OlmoeQ { - embed_tokens: EmbeddingConfig::new(config.vocab_size, config.hidden_size).init(device), - layers: (0..config.num_hidden_layers) - .map(|_| QDecoderLayer { - self_attn: attn_cfg.init(device), - mlp: SparseMoePerExpert { - gate: LinearConfig::new(config.hidden_size, config.num_experts) - .with_bias(false) - .init(device), - experts: (0..config.num_experts) - .map(|_| ExpertWeights { - gate: placeholder2(), - up: placeholder2(), - down: placeholder2(), - }) - .collect(), - }, - input_layernorm: norm(device), - post_attention_layernorm: norm(device), - }) - .collect(), - norm: norm(device), - lm_head: untied.then(|| { - LinearConfig::new(config.hidden_size, config.vocab_size) - .with_bias(false) - .init(device) - }), - }; - - // A float linear weight from GGUF's [out, in] into Linear's [in, out]. - let linear_f32 = |values: Vec, dims_rev: &[usize]| -> Result, String> { - let &[out, inp] = dims_rev else { - return Err(format!("linear weight must be 2-D, got {dims_rev:?}")); - }; - Ok(dev_tensor2(values, [out, inp]).swap_dims(0, 1)) - }; - - let mut assigned = 0usize; - for info in &f.tensors { - let dims_rev: Vec = info.dims.iter().rev().map(|&d| d as usize).collect(); - let name = &info.name; - - // The fused banks: split into per-expert members, quantize each. - let bank_field = name - .strip_prefix("blk.") - .and_then(|r| r.split_once('.')) - .and_then(|(_, field)| match field { - "ffn_gate_exps.weight" => Some("gate"), - "ffn_up_exps.weight" => Some("up"), - "ffn_down_exps.weight" => Some("down"), - _ => None, - }); - if let Some(bank_field) = bank_field { - let layer: usize = name - .strip_prefix("blk.") - .and_then(|r| r.split_once('.')) - .and_then(|(l, _)| l.parse().ok()) - .ok_or_else(|| parse(format!("bad bank layer in '{name}'")))?; - let &[e, out, inp] = dims_rev.as_slice() else { - return Err(parse(format!("expert bank must be 3-D, got {dims_rev:?}"))); - }; - if e != config.num_experts { - return Err(parse(format!( - "bank '{name}' has {e} experts, config says {}", - config.num_experts - ))); - } - let values = f.read_tensor_f32(name).map_err(|e| parse(e.to_string()))?; - let stride = out * inp; - let mlp = &mut model - .layers - .get_mut(layer) - .ok_or_else(|| parse(format!("layer {layer} out of range")))? - .mlp; - for expert in 0..e { - let member = values[expert * stride..(expert + 1) * stride].to_vec(); - let w = dev_tensor2(member, [out, inp]).swap_dims(0, 1); // [in, out] - let w = if policy.eligible(&[inp, out]) { - crate::quant::quantize_weight(policy, w) - } else { - w - }; - let slot = &mut mlp.experts[expert]; - match bank_field { - "gate" => slot.gate = Param::from_tensor(w), - "up" => slot.up = Param::from_tensor(w), - _ => slot.down = Param::from_tensor(w), - } - } - assigned += 1; - continue; - } - - // Everything else is float, routed by the proven name table. - let mapped = - olmoe_gguf_name(name).ok_or_else(|| parse(format!("unmapped tensor name '{name}'")))?; - let values = f.read_tensor_f32(name).map_err(|e| parse(e.to_string()))?; - match mapped.as_str() { - "model.embed_tokens.weight" => { - let &[v, h] = dims_rev.as_slice() else { - return Err(parse("embedding must be 2-D".into())); - }; - model.embed_tokens.weight = Param::from_tensor(dev_tensor2(values, [v, h])); - } - "model.norm.weight" => { - model.norm.gamma = Param::from_tensor(dev_tensor1(values, dims_rev[0])); - } - "lm_head.weight" => { - let head = model - .lm_head - .as_mut() - .ok_or_else(|| parse("output.weight on a tied model".into()))?; - head.weight = Param::from_tensor(linear_f32(values, &dims_rev).map_err(parse)?); - } - other => { - let rest = other - .strip_prefix("model.layers.") - .ok_or_else(|| parse(format!("unknown path '{other}'")))?; - let (layer, field) = rest - .split_once('.') - .ok_or_else(|| parse(format!("bad layer path '{other}'")))?; - let layer: usize = layer - .parse() - .map_err(|_| parse(format!("bad layer '{other}'")))?; - let l = model - .layers - .get_mut(layer) - .ok_or_else(|| parse(format!("layer {layer} out of range")))?; - match field { - "input_layernorm.weight" => { - l.input_layernorm.gamma = - Param::from_tensor(dev_tensor1(values, dims_rev[0])); - } - "post_attention_layernorm.weight" => { - l.post_attention_layernorm.gamma = - Param::from_tensor(dev_tensor1(values, dims_rev[0])); - } - "self_attn.q_proj.weight" => { - l.self_attn.q_proj.weight = - Param::from_tensor(linear_f32(values, &dims_rev).map_err(parse)?); - } - "self_attn.k_proj.weight" => { - l.self_attn.k_proj.weight = - Param::from_tensor(linear_f32(values, &dims_rev).map_err(parse)?); - } - "self_attn.v_proj.weight" => { - l.self_attn.v_proj.weight = - Param::from_tensor(linear_f32(values, &dims_rev).map_err(parse)?); - } - "self_attn.o_proj.weight" => { - l.self_attn.o_proj.weight = - Param::from_tensor(linear_f32(values, &dims_rev).map_err(parse)?); - } - "self_attn.q_norm.weight" => { - let n = l.self_attn.q_norm.as_mut().expect("qk norm built"); - n.gamma = Param::from_tensor(dev_tensor1(values, dims_rev[0])); - } - "self_attn.k_norm.weight" => { - let n = l.self_attn.k_norm.as_mut().expect("qk norm built"); - n.gamma = Param::from_tensor(dev_tensor1(values, dims_rev[0])); - } - "mlp.gate.weight" => { - l.mlp.gate.weight = - Param::from_tensor(linear_f32(values, &dims_rev).map_err(parse)?); - } - unknown => { - return Err(parse(format!("unknown layer field '{unknown}'"))); - } - } - } - } - assigned += 1; - } - - // 6 attention + 2 norms + 1 router + 3 banks per layer, plus embedding, - // final norm, and the untied head. - let expected = config.num_hidden_layers * 12 + 2 + usize::from(untied); - if assigned != expected { - return Err(parse(format!( - "GGUF supplied {assigned} tensors, the architecture needs {expected}" - ))); - } - Ok(LoadedOlmoeQ { - model, - config, - pool: None, - }) -} - -/// How each GGUF tensor enters a `.mummu` pack: expert banks split per -/// member, attention/router/head as linears, norms as vectors. -pub fn pack_actions(info: &GgufTensorInfo) -> Option { - use crate::pack::ImportAction as A; - match info.name.as_str() { - "token_embd.weight" => return Some(A::Embedding), - "output_norm.weight" => return Some(A::Vector), - "output.weight" => return Some(A::Linear), - _ => {} - } - let rest = info.name.strip_prefix("blk.")?; - let (layer, field) = rest.split_once('.')?; - let layer: usize = layer.parse().ok()?; - Some(match field { - "ffn_gate_exps.weight" => A::ExpertBank { - layer, - proj: "gate".into(), - }, - "ffn_up_exps.weight" => A::ExpertBank { - layer, - proj: "up".into(), - }, - "ffn_down_exps.weight" => A::ExpertBank { - layer, - proj: "down".into(), - }, - "attn_norm.weight" | "ffn_norm.weight" | "attn_q_norm.weight" | "attn_k_norm.weight" => { - A::Vector - } - "attn_q.weight" - | "attn_k.weight" - | "attn_v.weight" - | "attn_output.weight" - | "ffn_gate_inp.weight" => A::Linear, - _ => return None, - }) -} - -/// Load the per-expert-quantized OLMoE from a `.mummu` pack; `choose` picks -/// each tensor's precision (experts are separate entries — the tiering hook). -pub fn load_from_pack( - dir: &Path, - device: &Device, - choose: &dyn Fn(&crate::pack::TensorEntry) -> crate::pack::Precision, -) -> Result { - load_from_pack_inner(dir, device, choose, true) -} - -fn load_from_pack_inner( - dir: &Path, - device: &Device, - choose: &dyn Fn(&crate::pack::TensorEntry) -> crate::pack::Precision, - with_experts: bool, -) -> Result { - use crate::nn::{ExpertWeights, SparseMoePerExpert}; - use crate::pack::{Pack, Role}; - use burn::module::Param; - - let parse = |reason: String| ImportError::Parse { - file: dir.to_path_buf(), - reason, - }; - let pack = Pack::open(dir).map_err(parse)?; - let header = pack.header().map_err(parse)?; - let config = OlmoeConfig::from_gguf(&header).map_err(parse)?; - let untied = pack.entry("output.weight").is_some(); - - let dtype = crate::backend::float_dtype(device); - let placeholder2 = || Param::from_tensor(Tensor::<2>::zeros([1, 1], device)); - let norm = |dev: &Device| { - RmsNormConfig::new(config.hidden_size) - .with_epsilon(config.rms_norm_eps) - .init(dev) - }; - let attn_cfg = GqaAttentionConfig { - hidden_size: config.hidden_size, - num_heads: config.num_attention_heads, - num_kv_heads: config.num_key_value_heads, - head_dim: config.head_dim(), - bias: false, - qk_norm_eps: Some(config.rms_norm_eps), - qk_norm_projection: true, - }; - let mut model = OlmoeQ { - embed_tokens: EmbeddingConfig::new(config.vocab_size, config.hidden_size).init(device), - layers: (0..config.num_hidden_layers) - .map(|_| QDecoderLayer { - self_attn: attn_cfg.init(device), - mlp: SparseMoePerExpert { - gate: LinearConfig::new(config.hidden_size, config.num_experts) - .with_bias(false) - .init(device), - experts: (0..config.num_experts) - .map(|_| ExpertWeights { - gate: placeholder2(), - up: placeholder2(), - down: placeholder2(), - }) - .collect(), - }, - input_layernorm: norm(device), - post_attention_layernorm: norm(device), - }) - .collect(), - norm: norm(device), - lm_head: untied.then(|| { - LinearConfig::new(config.hidden_size, config.vocab_size) - .with_bias(false) - .init(device) - }), - }; - - let pick = |entry: &crate::pack::TensorEntry| -> crate::pack::Precision { - let p = choose(entry); - if entry.precisions.contains_key(&p) { - p - } else { - *entry - .precisions - .keys() - .max() - .expect("pack entries store a level") - } - }; - let vec1 = |values: Vec, n: usize| { - Tensor::<1>::from_data(TensorData::new(values, [n]), (device, dtype)) - }; - - let mut assigned = 0usize; - for entry in &pack.manifest.tensors { - // Experts: addressed by role, not by name. - if let Role::Expert { layer, index, proj } = &entry.role { - if !with_experts { - continue; // a pool serves them - } - // Experts are float-or-quantized 2-D [in, out] at the chosen level. - let t = pack - .tensor::<2>(entry, pick(entry), device) - .map_err(parse)?; - let slot = &mut model - .layers - .get_mut(*layer) - .ok_or_else(|| parse(format!("layer {layer} out of range")))? - .mlp - .experts[*index]; - match proj.as_str() { - "gate" => slot.gate = Param::from_tensor(t), - "up" => slot.up = Param::from_tensor(t), - "down" => slot.down = Param::from_tensor(t), - other => return Err(parse(format!("unknown expert proj '{other}'"))), - } - assigned += 1; - continue; - } - let mapped = olmoe_gguf_name(&entry.name) - .ok_or_else(|| parse(format!("unmapped pack tensor '{}'", entry.name)))?; - // Attention/router/head stay float here; linears are already [in, out]. - let lin = |entry: &crate::pack::TensorEntry| -> Result, ImportError> { - pack.tensor::<2>(entry, crate::pack::Precision::F32, device) - .or_else(|_| pack.tensor::<2>(entry, crate::pack::Precision::F16, device)) - .map_err(parse) - }; - match mapped.as_str() { - "model.embed_tokens.weight" => { - model.embed_tokens.weight = Param::from_tensor(lin(entry)?); - } - "model.norm.weight" => { - model.norm.gamma = - Param::from_tensor(vec1(pack.read_f32(entry).map_err(parse)?, entry.shape[0])); - } - "lm_head.weight" => { - let head = model - .lm_head - .as_mut() - .ok_or_else(|| parse("output.weight on a tied model".into()))?; - head.weight = Param::from_tensor(lin(entry)?); - } - other => { - let rest = other - .strip_prefix("model.layers.") - .ok_or_else(|| parse(format!("unknown path '{other}'")))?; - let (layer, field) = rest - .split_once('.') - .ok_or_else(|| parse(format!("bad layer path '{other}'")))?; - let layer: usize = layer - .parse() - .map_err(|_| parse(format!("bad layer '{other}'")))?; - let l = model - .layers - .get_mut(layer) - .ok_or_else(|| parse(format!("layer {layer} out of range")))?; - let n = entry.shape[0]; - match field { - "input_layernorm.weight" => { - l.input_layernorm.gamma = - Param::from_tensor(vec1(pack.read_f32(entry).map_err(parse)?, n)); - } - "post_attention_layernorm.weight" => { - l.post_attention_layernorm.gamma = - Param::from_tensor(vec1(pack.read_f32(entry).map_err(parse)?, n)); - } - "self_attn.q_proj.weight" => { - l.self_attn.q_proj.weight = Param::from_tensor(lin(entry)?) - } - "self_attn.k_proj.weight" => { - l.self_attn.k_proj.weight = Param::from_tensor(lin(entry)?) - } - "self_attn.v_proj.weight" => { - l.self_attn.v_proj.weight = Param::from_tensor(lin(entry)?) - } - "self_attn.o_proj.weight" => { - l.self_attn.o_proj.weight = Param::from_tensor(lin(entry)?) - } - "self_attn.q_norm.weight" => { - let qn = l.self_attn.q_norm.as_mut().expect("qk norm built"); - qn.gamma = - Param::from_tensor(vec1(pack.read_f32(entry).map_err(parse)?, n)); - } - "self_attn.k_norm.weight" => { - let kn = l.self_attn.k_norm.as_mut().expect("qk norm built"); - kn.gamma = - Param::from_tensor(vec1(pack.read_f32(entry).map_err(parse)?, n)); - } - "mlp.gate.weight" => l.mlp.gate.weight = Param::from_tensor(lin(entry)?), - unknown => return Err(parse(format!("unknown layer field '{unknown}'"))), - } - } - } - assigned += 1; - } - // 9 non-bank tensors per layer + 3 banks × experts, plus embed, norm, head. - let per_layer = if with_experts { - 9 + 3 * config.num_experts - } else { - 9 - }; - let expected = config.num_hidden_layers * per_layer + 2 + usize::from(untied); - if assigned != expected { - return Err(parse(format!( - "pack supplied {assigned} tensors, the architecture needs {expected}" - ))); - } - Ok(LoadedOlmoeQ { - model, - config, - pool: None, - }) -} - -/// The experts of a pack, in the tier planner's flat order -/// (`layer * num_experts + index`): each entry's three projections and the -/// resident bytes per stored level (values + scales; float levels as the -/// f32 a float backend holds). -pub fn pack_expert_costs(pack: &crate::pack::Pack) -> Result, String> { - use crate::pack::{Precision, Role}; - let header = pack.header()?; - let config = OlmoeConfig::from_gguf(&header)?; - let n = config.num_hidden_layers * config.num_experts; - let mut costs = vec![crate::tier::ExpertCost::default(); n]; - for entry in &pack.manifest.tensors { - let Role::Expert { layer, index, .. } = &entry.role else { - continue; - }; - let flat = layer * config.num_experts + index; - let numel: u64 = entry.shape.iter().product::() as u64; - let cost = costs.get_mut(flat).ok_or_else(|| { - format!("expert ({layer}, {index}) outside the {n} the config declares") - })?; - for (&p, blob) in &entry.precisions { - let bytes = match p { - Precision::Q4 | Precision::Q8 => blob.values_len + blob.scales_len, - Precision::F16 | Precision::F32 => numel * 4, - }; - *cost.bytes.entry(p).or_insert(0) += bytes; - } - } - Ok(costs) -} - -/// One expert's three projections from a pack at `precision`, on `device`, -/// as a tier-tagged [`crate::nn::DeviceExpert`]. The planner's hot-swap -/// path: load the replacement, then swap it into the pool. -pub fn load_expert_from_pack( - pack: &crate::pack::Pack, - layer: usize, - index: usize, - tier: crate::tier::Tier, - device: &Device, -) -> Result { - use crate::nn::ExpertWeights; - use crate::pack::{Precision, Role}; - use burn::module::Param; - let mut gate = None; - let mut up = None; - let mut down = None; - let mut bytes = 0u64; - for entry in &pack.manifest.tensors { - let Role::Expert { - layer: l, - index: i, - proj, - } = &entry.role - else { - continue; - }; - if *l != layer || *i != index { - continue; - } - let precision = if entry.precisions.contains_key(&tier.precision) { - tier.precision - } else { - *entry - .precisions - .keys() - .max() - .expect("pack entries store a level") - }; - let blob = entry.precisions[&precision]; - bytes += match precision { - Precision::Q4 | Precision::Q8 => blob.values_len + blob.scales_len, - Precision::F16 | Precision::F32 => entry.shape.iter().product::() as u64 * 4, - }; - let t = Param::from_tensor(pack.tensor::<2>(entry, precision, device)?); - match proj.as_str() { - "gate" => gate = Some(t), - "up" => up = Some(t), - "down" => down = Some(t), - other => return Err(format!("unknown expert proj '{other}'")), - } - } - let (Some(gate), Some(up), Some(down)) = (gate, up, down) else { - return Err(format!( - "pack is missing projections of expert ({layer}, {index})" - )); - }; - Ok(crate::nn::DeviceExpert { - native_ok: std::sync::atomic::AtomicBool::new(true), - weights: ExpertWeights { gate, up, down }, - device: device.clone(), - tier, - bytes, - }) -} - -/// Load everything **but** the experts from a pack (they stay `[1, 1]` -/// placeholders) — the trunk of a pooled model. Attach the pool with -/// [`LoadedOlmoeQ::with_pool`]. -pub fn load_trunk_from_pack(dir: &Path, device: &Device) -> Result { - load_from_pack_inner(dir, device, &|_| crate::pack::Precision::F32, false) -} - -impl LoadedOlmoeQ { - /// Route every layer's expert compute through `pool`. - #[must_use] - pub fn with_pool(mut self, pool: std::sync::Arc) -> Self { - assert_eq!( - pool.num_layers(), - self.config.num_hidden_layers, - "expert pool layer count must match the model" - ); - assert_eq!( - pool.experts_per_layer(), - self.config.num_experts, - "expert pool width must match the model" - ); - self.pool = Some(pool); - self - } -} - -impl CausalLm for LoadedOlmoeQ { - type Cache = Vec; - - fn new_cache(&self) -> Self::Cache { - (0..self.config.num_hidden_layers).map(|_| None).collect() - } - - fn is_eos(&self, id: u32) -> bool { - self.config.eos_token_id.contains(id) - } - - fn forward( - &self, - new_ids: &[u32], - past: usize, - cache: &mut Self::Cache, - device: &Device, - ) -> Tensor<2> { - let t = new_ids.len(); - assert!(t >= 1, "OLMoE-Q forward: need at least one token"); - assert!( - cache.len() == self.config.num_hidden_layers, - "OLMoE-Q forward: cache has {} layers, model has {}", - cache.len(), - self.config.num_hidden_layers - ); - let cfg = &self.config; - let hd = cfg.head_dim(); - - let ids32: Vec = new_ids.iter().map(|&i| i as i32).collect(); - let input = Tensor::<1, Int>::from_data( - TensorData::new(ids32, [t]), - (device, crate::backend::int_dtype(device)), - ) - .reshape([1, t]); - let mut x = self.model.embed_tokens.forward(input); - - let (cos, sin) = rope_tables(t, past, hd, cfg.rope_theta, device); - let mask = (t > 1).then(|| causal_mask(t, past, device)); - - for (li, (layer, kv)) in self.model.layers.iter().zip(cache.iter_mut()).enumerate() { - let h = layer.input_layernorm.forward(x.clone()); - let h = layer.self_attn.forward( - h, - cfg.num_attention_heads, - cfg.num_key_value_heads, - hd, - &cos, - &sin, - mask.as_ref(), - kv, - ); - x = x.add(h); - let h2 = layer.post_attention_layernorm.forward(x.clone()); - let moe = match &self.pool { - Some(pool) => layer.mlp.forward_pooled( - h2, - cfg.num_experts_per_tok, - cfg.norm_topk_prob, - pool, - li, - ), - None => layer - .mlp - .forward(h2, cfg.num_experts_per_tok, cfg.norm_topk_prob), - }; - x = x.add(moe); - } - let x = self.model.norm.forward(x); - - let last = x.narrow(1, t - 1, 1).reshape([1, cfg.hidden_size]); - match &self.model.lm_head { - Some(head) => head.forward(last), - None => { - let w = self.model.embed_tokens.weight.val(); - last.matmul(w.swap_dims(0, 1)) - } - } - } -} diff --git a/crates/mummu/examples/src/models/qwen2.rs b/crates/mummu/examples/src/models/qwen2.rs deleted file mode 100644 index a730106..0000000 --- a/crates/mummu/examples/src/models/qwen2.rs +++ /dev/null @@ -1,831 +0,0 @@ -//! Qwen2 / Qwen2.5 decoder, from scratch on the shared `nn` blocks: -//! Embedding → N×{RmsNorm, GQA attention (RoPE, KV cache), RmsNorm, SwiGLU} -//! → RmsNorm → tied lm-head. Config-driven — the 0.5B and 1.5B tiers (and any -//! other single-file Qwen2 checkpoint) load through the same code. -//! -//! Ported from laurelane's parity-validated implementation (single-forward -//! next-token logits byte-identical to Candle's `qwen2` on identical f32 -//! weights, CPU and wgpu GPU). Mummu's own parity gate (P7) re-verifies here -//! before the port is marked trusted. - -use std::path::Path; - -use burn::module::Module; -use burn::nn::{Embedding, EmbeddingConfig, Linear, LinearConfig, RmsNorm, RmsNormConfig}; -use burn::store::{ModuleAdapter, PyTorchToBurnAdapter, SafetensorsStore}; -use burn::tensor::{Device, Int, Tensor, TensorData}; - -use crate::attn_config::{RopeScaling, check_sliding_window, sliding_window_from_gguf}; -use crate::gguf::{GgufFile, GgufMap, GgufTensorInfo, GgufValue}; -use crate::import::{ - CastFloatAdapter, DequantSink, ImportError, gguf_store, load_checked, required_file, -}; -use crate::models::CausalLm; -use crate::nn::{ - GqaAttention, GqaAttentionConfig, LayerKv, SwiGluMlp, SwiGluMlpConfig, causal_mask, rope_tables, -}; - -/// Qwen2 architecture hyperparameters, read from the checkpoint's `config.json`. -#[derive(Debug, Clone, serde::Deserialize)] -pub struct Qwen2Config { - pub vocab_size: usize, - pub hidden_size: usize, - pub intermediate_size: usize, - pub num_hidden_layers: usize, - pub num_attention_heads: usize, - pub num_key_value_heads: usize, - #[serde(default)] - pub head_dim: usize, - pub rms_norm_eps: f64, - pub rope_theta: f32, - /// Frequency scaling (YaRN / linear / …). `null` on every Qwen2.5 - /// checkpoint at its native context; a scaled one is refused at load - /// rather than answered wrong ([`crate::attn_config`]). - /// `rope_parameters` is the same object under the name newer transformers - /// writes; reading only `rope_scaling` would let a freshly-serialized - /// scaled checkpoint through as unscaled. - #[serde(default, alias = "rope_parameters")] - pub rope_scaling: Option, - /// The trained context length, used to tell an inert sliding window (one - /// that spans the whole context) from a clipping one. - #[serde(default)] - pub max_position_embeddings: Option, - /// Declared window span. Qwen2.5 ships `32768` on every tier — **inert**, - /// because `use_sliding_window` is `false`. - #[serde(default)] - pub sliding_window: Option, - /// Whether `sliding_window` is live. Qwen2's own gate; `false` everywhere - /// in the zoo, which is why the full causal mask is correct for it. - #[serde(default)] - pub use_sliding_window: bool, - #[serde(default)] - pub tie_word_embeddings: bool, - /// EOS token id(s) — `<|im_end|>` first for the instruct checkpoints. - #[serde(default)] - pub eos_token_id: EosIds, -} - -/// `eos_token_id` is a bare int in some checkpoints, a list in others. -#[derive(Debug, Clone, Default, serde::Deserialize)] -#[serde(untagged)] -pub enum EosIds { - #[default] - None, - One(u32), - Many(Vec), -} - -impl EosIds { - /// Is `id` one of the EOS ids? - #[must_use] - pub fn contains(&self, id: u32) -> bool { - match self { - Self::None => false, - Self::One(e) => *e == id, - Self::Many(v) => v.contains(&id), - } - } - - /// The EOS ids as an owned list (empty when `None`) — the shape the - /// `tokenizer_config.json` cross-check ([`crate::tok_config`]) consumes. - #[must_use] - pub fn to_vec(&self) -> Vec { - match self { - Self::None => Vec::new(), - Self::One(e) => vec![*e], - Self::Many(v) => v.clone(), - } - } -} - -/// A required GGUF metadata integer, as usize. Shared with the Qwen3 loader. -pub(crate) fn gguf_usize(f: &GgufFile, key: &str) -> Result { - f.get(key) - .and_then(GgufValue::as_u64) - .and_then(|v| usize::try_from(v).ok()) - .ok_or_else(|| format!("missing or non-integer GGUF metadata '{key}'")) -} - -/// A required GGUF metadata f32. Shared with the Qwen3 loader. -pub(crate) fn gguf_f32(f: &GgufFile, key: &str) -> Result { - f.get(key) - .and_then(GgufValue::as_f32) - .ok_or_else(|| format!("missing or non-f32 GGUF metadata '{key}'")) -} - -impl Qwen2Config { - /// Parse `config.json` bytes; derives `head_dim` when absent. - pub fn from_json_bytes(bytes: &[u8]) -> Result { - let mut cfg: Self = serde_json::from_slice(bytes).map_err(|e| e.to_string())?; - if cfg.head_dim == 0 { - cfg.head_dim = cfg.hidden_size / cfg.num_attention_heads; - } - cfg.validate("qwen2 config.json")?; - Ok(cfg) - } - - /// Hyperparameters from a GGUF header's `qwen2.*` metadata — a GGUF file - /// is self-contained, no `config.json` beside it. `vocab_size` comes from - /// the embedding tensor (llama.cpp may pad it past the tokenizer vocab). - pub fn from_gguf(f: &GgufFile) -> Result { - let arch = f.architecture().unwrap_or(""); - if arch != "qwen2" { - return Err(format!("GGUF architecture '{arch}' is not qwen2")); - } - let hidden_size = gguf_usize(f, "qwen2.embedding_length")?; - let num_attention_heads = gguf_usize(f, "qwen2.attention.head_count")?; - let embd = f - .tensor("token_embd.weight") - .ok_or("GGUF has no token_embd.weight tensor")?; - if embd.dims.len() != 2 || embd.dims[0] != hidden_size as u64 { - return Err(format!( - "token_embd.weight dims {:?} do not match embedding_length {hidden_size}", - embd.dims - )); - } - let vocab_size = usize::try_from(embd.dims[1]).map_err(|_| "vocab too large")?; - if let Some(tokens) = f.get("tokenizer.ggml.tokens").and_then(GgufValue::as_array) - && tokens.len() > vocab_size - { - return Err(format!( - "tokenizer vocab {} exceeds embedding rows {vocab_size}", - tokens.len() - )); - } - let eos_token_id = f - .get("tokenizer.ggml.eos_token_id") - .and_then(GgufValue::as_u64) - .and_then(|v| u32::try_from(v).ok()) - .map_or(EosIds::None, EosIds::One); - let head_dim = f - .get("qwen2.attention.key_length") - .and_then(GgufValue::as_u64) - .and_then(|v| usize::try_from(v).ok()) - .unwrap_or(hidden_size / num_attention_heads); - let cfg = Self { - vocab_size, - hidden_size, - intermediate_size: gguf_usize(f, "qwen2.feed_forward_length")?, - num_hidden_layers: gguf_usize(f, "qwen2.block_count")?, - num_attention_heads, - num_key_value_heads: gguf_usize(f, "qwen2.attention.head_count_kv")?, - head_dim, - rms_norm_eps: f64::from(gguf_f32(f, "qwen2.attention.layer_norm_rms_epsilon")?), - rope_theta: gguf_f32(f, "qwen2.rope.freq_base")?, - rope_scaling: RopeScaling::from_gguf(f, "qwen2"), - max_position_embeddings: f - .get("qwen2.context_length") - .and_then(GgufValue::as_u64) - .and_then(|v| usize::try_from(v).ok()), - sliding_window: sliding_window_from_gguf(f, "qwen2"), - // llama.cpp writes the key only for architectures that window, so - // its presence IS the enable — there is no `use_sliding_window` - // twin in a GGUF header the way `config.json` has one. - use_sliding_window: true, - // No separate output.weight tensor means the lm-head is tied. - tie_word_embeddings: f.tensor("output.weight").is_none(), - eos_token_id, - }; - cfg.validate("qwen2 GGUF header")?; - Ok(cfg) - } - - fn validate(&self, whose: &str) -> Result<(), String> { - if let Some(scaling) = &self.rope_scaling { - scaling.check(whose)?; - } - check_sliding_window( - self.sliding_window, - self.use_sliding_window, - self.max_position_embeddings, - whose, - )?; - if self.num_key_value_heads == 0 - || !self - .num_attention_heads - .is_multiple_of(self.num_key_value_heads) - { - return Err(format!( - "num_attention_heads ({}) must be a positive multiple of num_key_value_heads ({})", - self.num_attention_heads, self.num_key_value_heads - )); - } - if self.num_hidden_layers == 0 || self.vocab_size == 0 { - return Err("num_hidden_layers and vocab_size must be positive".into()); - } - Ok(()) - } -} - -/// One decoder layer. Field names mirror the HF checkpoint so the key remap -/// stays trivial. -#[derive(Module, Debug)] -pub struct DecoderLayer { - pub self_attn: GqaAttention, - pub mlp: SwiGluMlp, - pub input_layernorm: RmsNorm, - pub post_attention_layernorm: RmsNorm, -} - -/// The Qwen2 decoder stack (HF's `model.*` subtree). The lm-head is tied on -/// the small tiers (0.5B/1.5B safetensors); untied checkpoints — the 7B, and -/// every llama.cpp GGUF (which materializes the head as `output.weight`, at -/// higher precision than the embedding) — carry it explicitly. -#[derive(Module, Debug)] -pub struct Qwen2 { - pub embed_tokens: Embedding, - pub layers: Vec, - pub norm: RmsNorm, - pub lm_head: Option, -} - -/// A weight-loaded Qwen2 plus its config — everything a forward needs. -pub struct LoadedQwen2 { - pub model: Qwen2, - pub config: Qwen2Config, - /// The parsed sibling `tokenizer_config.json`, when one was present and - /// well-formed beside a safetensors checkpoint (the load-time gate has - /// already cross-checked its EOS against `config.json`). A consumer reads - /// config-driven EOS/BOS/PAD ids from it (`eos_id()`, `bos_id()`, …). `None` - /// for a GGUF load (self-contained; no sibling file) or a dir without one. - pub tokenizer_config: Option, -} - -fn build(cfg: &Qwen2Config, device: &Device) -> Qwen2 { - let attn_cfg = GqaAttentionConfig { - hidden_size: cfg.hidden_size, - num_heads: cfg.num_attention_heads, - num_kv_heads: cfg.num_key_value_heads, - head_dim: cfg.head_dim, - bias: true, // Qwen2 has q/k/v projection bias - qk_norm_eps: None, // and no per-head q/k norm - qk_norm_projection: false, - }; - let mlp_cfg = SwiGluMlpConfig { - hidden_size: cfg.hidden_size, - intermediate_size: cfg.intermediate_size, - }; - let norm = |dev: &Device| { - RmsNormConfig::new(cfg.hidden_size) - .with_epsilon(cfg.rms_norm_eps) - .init(dev) - }; - let layers = (0..cfg.num_hidden_layers) - .map(|_| DecoderLayer { - self_attn: attn_cfg.init(device), - mlp: mlp_cfg.init(device), - input_layernorm: norm(device), - post_attention_layernorm: norm(device), - }) - .collect(); - let lm_head = (!cfg.tie_word_embeddings).then(|| { - LinearConfig::new(cfg.hidden_size, cfg.vocab_size) - .with_bias(false) - .init(device) - }); - Qwen2 { - embed_tokens: EmbeddingConfig::new(cfg.vocab_size, cfg.hidden_size).init(device), - layers, - norm: norm(device), - lm_head, - } -} - -/// Build the architecture from `dir/config.json` and load -/// `dir/model.safetensors` into it, checked (no silent partial loads). -/// -/// The key remap maps HF names onto our HF-mirroring field paths: strip the -/// `model.` prefix and rename RmsNorm `weight` → Burn's `gamma` -/// (`PyTorchToBurnAdapter` renames Layer/Batch/Group norm params but NOT -/// RmsNorm; it does transpose the Linear weights). -pub fn load_from_dir(dir: &Path, device: &Device) -> Result { - let cfg_path = required_file(dir, "config.json")?; - let weights = required_file(dir, "model.safetensors")?; - let cfg_bytes = std::fs::read(&cfg_path).map_err(|e| ImportError::Parse { - file: cfg_path.clone(), - reason: e.to_string(), - })?; - let config = Qwen2Config::from_json_bytes(&cfg_bytes).map_err(|reason| ImportError::Parse { - file: cfg_path, - reason, - })?; - - // Cross-check the sibling metadata (when present) before touching weights: - // tokenizer_config.json EOS agreement with config.json, a chat-template that - // speaks Qwen2's Hermes/ChatML tool-call convention, and added-token ids that - // match the real tokenizer.json — a repackaging mismatch fails loudly at load. - let tokenizer_config = crate::tokenizer::validate_checkpoint_dir( - dir, - &config.eos_token_id.to_vec(), - Some(crate::tok_config::ToolCallConvention::Hermes), - )?; - debug_assert!( - tokenizer_config - .as_ref() - .and_then(crate::tok_config::TokenizerConfig::eos_id) - .is_none_or(|id| config.eos_token_id.contains(id)), - "validate_checkpoint_dir returned a config whose EOS disagrees with config.json" - ); - - let mut model = build(&config, device); - // The float dtype comes from the DEVICE — burn 0.22 keeps the element - // type there as a runtime setting, not on a backend type. Creation sites - // still name it explicitly rather than riding the unspecified default. - let target_float = crate::backend::float_dtype(device); - let mut store = SafetensorsStore::from_file(weights.clone()) - .with_from_adapter(PyTorchToBurnAdapter.chain(CastFloatAdapter::new(target_float))) - .allow_partial(true) - .with_key_remapping(r"^model\.", "") - .with_key_remapping(r"(input_layernorm)\.weight$", "$1.gamma") - .with_key_remapping(r"(post_attention_layernorm)\.weight$", "$1.gamma") - .with_key_remapping(r"^norm\.weight$", "norm.gamma"); - load_checked(&mut model, &mut store, &weights)?; - Ok(LoadedQwen2 { - model, - config, - tokenizer_config, - }) -} - -/// GGUF (llama.cpp) tensor names → the HF checkpoint names the safetensors -/// remap chain already handles. `None` for anything unrecognized — the blob -/// writer turns that into a loud error rather than dropping weights. -fn gguf_tensor_to_hf(info: &GgufTensorInfo) -> Option { - qwen2_gguf_name(&info.name).map(GgufMap::Rename) -} - -fn qwen2_gguf_name(name: &str) -> Option { - match name { - "token_embd.weight" => return Some("model.embed_tokens.weight".into()), - "output_norm.weight" => return Some("model.norm.weight".into()), - "output.weight" => return Some("lm_head.weight".into()), - _ => {} - } - let rest = name.strip_prefix("blk.")?; - let (layer, field) = rest.split_once('.')?; - let layer: usize = layer.parse().ok()?; - let mapped = match field { - "attn_norm.weight" => "input_layernorm.weight", - "ffn_norm.weight" => "post_attention_layernorm.weight", - "attn_q.weight" => "self_attn.q_proj.weight", - "attn_q.bias" => "self_attn.q_proj.bias", - "attn_k.weight" => "self_attn.k_proj.weight", - "attn_k.bias" => "self_attn.k_proj.bias", - "attn_v.weight" => "self_attn.v_proj.weight", - "attn_v.bias" => "self_attn.v_proj.bias", - "attn_output.weight" => "self_attn.o_proj.weight", - "ffn_gate.weight" => "mlp.gate_proj.weight", - "ffn_up.weight" => "mlp.up_proj.weight", - "ffn_down.weight" => "mlp.down_proj.weight", - _ => return None, - }; - Some(format!("model.layers.{layer}.{mapped}")) -} - -/// Load a Qwen2 model straight from a **GGUF** file (any dtype the dequant -/// suite covers — Q4_K_M, Q8_0, F16, …): hyperparameters from the GGUF -/// metadata, weights dequantized to f32 and driven through the exact store -/// pipeline (adapters + remaps + checked load) the safetensors path uses. -pub fn load_from_gguf(path: &Path, device: &Device) -> Result { - let parse = |reason: String| ImportError::Parse { - file: path.to_path_buf(), - reason, - }; - let f = GgufFile::open(path).map_err(|e| parse(e.to_string()))?; - let config = Qwen2Config::from_gguf(&f).map_err(parse)?; - // The scratch guard (Some only when the payload went to disk) must - // outlive `load_checked`: the store reads that file lazily. - let (base, _scratch) = gguf_store(&f, &gguf_tensor_to_hf, DequantSink::Auto, device)?; - - let mut model = build(&config, device); - let mut store = base - .with_key_remapping(r"^model\.", "") - .with_key_remapping(r"(input_layernorm)\.weight$", "$1.gamma") - .with_key_remapping(r"(post_attention_layernorm)\.weight$", "$1.gamma") - .with_key_remapping(r"^norm\.weight$", "norm.gamma"); - load_checked(&mut model, &mut store, path)?; - // A GGUF is self-contained — no sibling tokenizer_config.json in this path. - Ok(LoadedQwen2 { - model, - config, - tokenizer_config: None, - }) -} - -impl CausalLm for LoadedQwen2 { - type Cache = Vec; - - fn new_cache(&self) -> Self::Cache { - (0..self.config.num_hidden_layers).map(|_| None).collect() - } - - fn is_eos(&self, id: u32) -> bool { - self.config.eos_token_id.contains(id) - } - - fn forward( - &self, - new_ids: &[u32], - past: usize, - cache: &mut Self::Cache, - device: &Device, - ) -> Tensor<2> { - let t = new_ids.len(); - assert!(t >= 1, "Qwen2 forward: need at least one token"); - assert!( - cache.len() == self.config.num_hidden_layers, - "Qwen2 forward: cache has {} layers, model has {}", - cache.len(), - self.config.num_hidden_layers - ); - let cfg = &self.config; - - // i32 token ids: native for wgpu and the flex CPU backend alike. - // Dtype pinned to the backend TYPE, never the per-device policy. - let ids32: Vec = new_ids.iter().map(|&i| i as i32).collect(); - let input = Tensor::<1, Int>::from_data( - TensorData::new(ids32, [t]), - (device, crate::backend::int_dtype(device)), - ) - .reshape([1, t]); - let mut x = self.model.embed_tokens.forward(input); // [1, t, hidden] - - let (cos, sin) = rope_tables(t, past, cfg.head_dim, cfg.rope_theta, device); - // A single-token decode step needs no mask (the one query attends to - // all cached keys); only a multi-token prefill needs the triangle. - let mask = (t > 1).then(|| causal_mask(t, past, device)); - - for (layer, kv) in self.model.layers.iter().zip(cache.iter_mut()) { - let h = layer.input_layernorm.forward(x.clone()); - let h = layer.self_attn.forward( - h, - cfg.num_attention_heads, - cfg.num_key_value_heads, - cfg.head_dim, - &cos, - &sin, - mask.as_ref(), - kv, - ); - x = x.add(h); - let h2 = layer.post_attention_layernorm.forward(x.clone()); - x = x.add(layer.mlp.forward(h2)); - } - let x = self.model.norm.forward(x); - - let last = x.narrow(1, t - 1, 1).reshape([1, cfg.hidden_size]); - debug_assert!( - self.model.lm_head.is_some() != cfg.tie_word_embeddings, - "lm_head presence must match the config's tie flag" - ); - match &self.model.lm_head { - // Untied: the checkpoint's own head projection. - Some(head) => head.forward(last), // [1, vocab] - // Tied lm-head: logits = last_hidden @ embed_weight^T. - None => { - let w = self.model.embed_tokens.weight.val(); // [vocab, hidden] - last.matmul(w.swap_dims(0, 1)) // [1, vocab] - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - /// A synthetic 2-layer toy config: everything runs without real weights. - fn toy_config() -> Qwen2Config { - Qwen2Config { - vocab_size: 64, - hidden_size: 16, - intermediate_size: 32, - num_hidden_layers: 2, - num_attention_heads: 4, - num_key_value_heads: 2, - head_dim: 4, - rms_norm_eps: 1e-6, - rope_theta: 1e4, - rope_scaling: None, - max_position_embeddings: Some(512), - sliding_window: None, - use_sliding_window: false, - tie_word_embeddings: true, - eos_token_id: EosIds::One(2), - } - } - - #[test] - fn config_parses_hf_shape_and_derives_head_dim() { - let json = br#"{ - "vocab_size": 151936, "hidden_size": 1536, "intermediate_size": 8960, - "num_hidden_layers": 28, "num_attention_heads": 12, "num_key_value_heads": 2, - "rms_norm_eps": 1e-6, "rope_theta": 1000000.0, "tie_word_embeddings": true, - "eos_token_id": 151645 - }"#; - let cfg = Qwen2Config::from_json_bytes(json).unwrap(); - assert_eq!(cfg.head_dim, 128); // derived 1536/12 - assert!(cfg.eos_token_id.contains(151_645)); - assert!(!cfg.eos_token_id.contains(151_644)); - } - - /// The real Qwen2.5 shape carries `sliding_window: 32768` alongside - /// `use_sliding_window: false`. That window is INERT, and refusing on the - /// field's presence would reject two checkpoints Mummu has parity-verified - /// — so this is the load that must keep working. - #[test] - fn the_inert_qwen25_sliding_window_still_loads() { - let json = br#"{ - "vocab_size": 151936, "hidden_size": 1536, "intermediate_size": 8960, - "num_hidden_layers": 28, "num_attention_heads": 12, "num_key_value_heads": 2, - "rms_norm_eps": 1e-6, "rope_theta": 1000000.0, "tie_word_embeddings": true, - "max_position_embeddings": 32768, "sliding_window": 32768, - "use_sliding_window": false, "max_window_layers": 21, "rope_scaling": null - }"#; - let cfg = Qwen2Config::from_json_bytes(json).expect("the zoo's own shape must load"); - assert_eq!(cfg.sliding_window, Some(32768)); - assert!(!cfg.use_sliding_window); - assert!(cfg.rope_scaling.is_none()); - } - - /// Negative space for the same field: flip the gate and the load must fail - /// naming the span, because the full causal mask would silently let the - /// model attend past its trained window. - #[test] - fn an_enabled_sliding_window_is_refused_by_name() { - let json = br#"{ - "vocab_size": 151936, "hidden_size": 1536, "intermediate_size": 8960, - "num_hidden_layers": 28, "num_attention_heads": 12, "num_key_value_heads": 2, - "rms_norm_eps": 1e-6, "rope_theta": 1000000.0, - "max_position_embeddings": 131072, "sliding_window": 4096, - "use_sliding_window": true - }"#; - let err = Qwen2Config::from_json_bytes(json).expect_err("an enabled window must refuse"); - assert!(err.contains("4096"), "{err}"); - assert!(err.contains("qwen2 config.json"), "{err}"); - } - - /// A YaRN-scaled checkpoint would load clean and degrade only far out in - /// the context — where no short-prompt gate looks. Refuse at load instead. - #[test] - fn a_yarn_scaled_checkpoint_is_refused_at_load() { - let json = br#"{ - "vocab_size": 151936, "hidden_size": 1536, "intermediate_size": 8960, - "num_hidden_layers": 28, "num_attention_heads": 12, "num_key_value_heads": 2, - "rms_norm_eps": 1e-6, "rope_theta": 1000000.0, - "rope_scaling": {"rope_type": "yarn", "factor": 4.0, - "original_max_position_embeddings": 32768} - }"#; - let err = Qwen2Config::from_json_bytes(json).expect_err("yarn must refuse"); - assert!(err.contains("yarn"), "{err}"); - } - - #[test] - fn config_rejects_indivisible_heads() { - let json = br#"{ - "vocab_size": 100, "hidden_size": 16, "intermediate_size": 32, - "num_hidden_layers": 1, "num_attention_heads": 5, "num_key_value_heads": 2, - "rms_norm_eps": 1e-6, "rope_theta": 10000.0 - }"#; - assert!(Qwen2Config::from_json_bytes(json).is_err()); - } - - #[test] - fn eos_ids_parse_as_int_or_list() { - let one: EosIds = serde_json::from_str("7").unwrap(); - let many: EosIds = serde_json::from_str("[7, 9]").unwrap(); - assert!(one.contains(7) && !one.contains(9)); - assert!(many.contains(7) && many.contains(9) && !many.contains(8)); - } - - /// The toy model decodes through the cache identically to full re-forwards - /// — the whole-stack version of the nn-level equivalence tests. - #[test] - fn toy_model_cached_decode_matches_full_forward() { - let device = crate::backend::cpu_device(); - let cfg = toy_config(); - let loaded = LoadedQwen2 { - model: build(&cfg, &device), - config: cfg, - tokenizer_config: None, - }; - - let prompt: Vec = vec![3, 14, 15, 9, 26]; - - // Cached: prefill then one decode step for token at position 5. - let mut cache = loaded.new_cache(); - let _ = loaded.forward(&prompt, 0, &mut cache, &device); - let step = loaded - .forward(&[42], prompt.len(), &mut cache, &device) - .into_data() - .to_vec::() - .unwrap(); - - // Full: all six tokens in one forward, no cache reuse. - let mut full_cache = loaded.new_cache(); - let all: Vec = prompt.iter().copied().chain([42]).collect(); - let full = loaded - .forward(&all, 0, &mut full_cache, &device) - .into_data() - .to_vec::() - .unwrap(); - - assert_eq!(step.len(), full.len()); - for (i, (c, f)) in step.iter().zip(&full).enumerate() { - assert!((c - f).abs() < 1e-4, "logit {i}: cached {c} vs full {f}"); - } - } - - /// A synthetic in-memory GGUF header shaped like the Qwen2.5-1.5B file. - fn toy_gguf() -> GgufFile { - use crate::gguf::{GgmlType, GgufTensorInfo}; - let meta = |k: &str, v: GgufValue| (k.to_string(), v); - GgufFile { - path: std::path::PathBuf::new(), - version: 3, - metadata: vec![ - meta("general.architecture", GgufValue::Str("qwen2".into())), - meta("qwen2.embedding_length", GgufValue::U32(16)), - meta("qwen2.block_count", GgufValue::U32(2)), - meta("qwen2.feed_forward_length", GgufValue::U32(32)), - meta("qwen2.attention.head_count", GgufValue::U32(4)), - meta("qwen2.attention.head_count_kv", GgufValue::U32(2)), - meta( - "qwen2.attention.layer_norm_rms_epsilon", - GgufValue::F32(1e-6), - ), - meta("qwen2.rope.freq_base", GgufValue::F32(1e4)), - meta("tokenizer.ggml.eos_token_id", GgufValue::U32(2)), - ], - tensors: vec![GgufTensorInfo { - name: "token_embd.weight".into(), - dims: vec![16, 64], // ggml order: [hidden, vocab] - dtype: GgmlType::F32, - offset: 0, - }], - alignment: 32, - data_offset: 0, - } - } - - #[test] - fn config_from_gguf_reads_metadata_and_embedding_dims() { - let cfg = Qwen2Config::from_gguf(&toy_gguf()).expect("parses"); - assert_eq!(cfg.vocab_size, 64); // from token_embd dims[1] - assert_eq!(cfg.hidden_size, 16); - assert_eq!(cfg.num_hidden_layers, 2); - assert_eq!(cfg.head_dim, 4); // derived: hidden / heads - assert!(cfg.eos_token_id.contains(2)); - assert!(cfg.tie_word_embeddings); // no output.weight tensor - } - - #[test] - fn config_from_gguf_fails_loudly_on_missing_keys_and_wrong_arch() { - let mut f = toy_gguf(); - f.metadata.retain(|(k, _)| k != "qwen2.block_count"); - let err = Qwen2Config::from_gguf(&f).unwrap_err(); - assert!(err.contains("qwen2.block_count"), "{err}"); - - let mut f = toy_gguf(); - f.metadata[0].1 = GgufValue::Str("llama".into()); - assert!(Qwen2Config::from_gguf(&f).is_err()); - - // Embedding dims that contradict the metadata are rejected. - let mut f = toy_gguf(); - f.tensors[0].dims = vec![8, 64]; - assert!(Qwen2Config::from_gguf(&f).is_err()); - } - - #[test] - fn untied_toy_config_builds_and_uses_an_lm_head() { - let device = crate::backend::cpu_device(); - let mut cfg = toy_config(); - cfg.tie_word_embeddings = false; - let vocab = cfg.vocab_size; - let loaded = LoadedQwen2 { - model: build(&cfg, &device), - config: cfg, - tokenizer_config: None, - }; - assert!(loaded.model.lm_head.is_some()); - let mut cache = loaded.new_cache(); - let logits = loaded.forward(&[1, 2], 0, &mut cache, &device); - assert_eq!(logits.dims(), [1, vocab]); - } - - #[test] - fn config_from_gguf_detects_an_untied_head() { - use crate::gguf::{GgmlType, GgufTensorInfo}; - let mut f = toy_gguf(); - f.tensors.push(GgufTensorInfo { - name: "output.weight".into(), - dims: vec![16, 64], - dtype: GgmlType::F32, - offset: 4096, - }); - let cfg = Qwen2Config::from_gguf(&f).expect("parses"); - assert!(!cfg.tie_word_embeddings); - } - - #[test] - fn gguf_names_map_onto_hf_checkpoint_names() { - assert_eq!( - qwen2_gguf_name("token_embd.weight").as_deref(), - Some("model.embed_tokens.weight") - ); - assert_eq!( - qwen2_gguf_name("blk.27.attn_q.bias").as_deref(), - Some("model.layers.27.self_attn.q_proj.bias") - ); - assert_eq!( - qwen2_gguf_name("blk.0.ffn_down.weight").as_deref(), - Some("model.layers.0.mlp.down_proj.weight") - ); - assert_eq!( - qwen2_gguf_name("output_norm.weight").as_deref(), - Some("model.norm.weight") - ); - // Unknown names must map to None (the writer errors loudly). - assert_eq!(qwen2_gguf_name("rope_freqs.weight"), None); - assert_eq!(qwen2_gguf_name("blk.x.attn_q.weight"), None); - } - - #[tokio::test] - async fn sanity_check_passes_on_a_live_toy_model_and_flags_a_vocab_mismatch() { - let device = crate::backend::cpu_device(); - let cfg = toy_config(); - let vocab = cfg.vocab_size; - let loaded = LoadedQwen2 { - model: build(&cfg, &device), - config: cfg, - tokenizer_config: None, - }; - // A built (random-weight) model computes a live, finite, non-degenerate - // distribution — the smoke passes and reports a valid top id. - let smoke = loaded - .sanity_check(&[1, 2, 3], vocab, &device) - .await - .expect("live toy model passes the smoke"); - assert!((smoke.top_id as usize) < vocab, "top id in vocab range"); - assert!(smoke.spread > 0.0, "a live forward has positive spread"); - // The wrong expected vocab is caught as a mismatch, not a silent pass. - assert!( - loaded - .sanity_check(&[1, 2, 3], vocab + 1, &device) - .await - .is_err() - ); - } - - #[tokio::test] - async fn warm_up_runs_one_prefill_plus_its_steps_and_leaves_the_model_usable() { - let device = crate::backend::cpu_device(); - let cfg = toy_config(); - let loaded = LoadedQwen2 { - model: build(&cfg, &device), - config: cfg, - tokenizer_config: None, - }; - let forwards = loaded - .warm_up(&[1, 2, 3], 4, &device) - .await - .expect("warm-up runs on a live toy model"); - assert_eq!(forwards, 5, "one prefill plus four decode steps"); - // The warm-up cache is a throwaway: a generation after it starts from - // an empty cache and still decodes (nothing leaked into the model). - let out = loaded - .greedy_generate(&[1, 2, 3], 2, &device) - .await - .expect("decodes after a warm-up"); - assert!(!out.is_empty(), "generation after warm-up produces tokens"); - } - - #[test] - #[should_panic(expected = "exceeds the 256 bound")] - fn warm_up_rejects_an_unbounded_step_count() { - let device = crate::backend::cpu_device(); - let cfg = toy_config(); - let loaded = LoadedQwen2 { - model: build(&cfg, &device), - config: cfg, - tokenizer_config: None, - }; - let _ = loaded.warm_up(&[1, 2, 3], crate::models::MAX_WARM_UP_STEPS + 1, &device); - } - - #[tokio::test] - async fn greedy_generate_respects_max_tokens_bound() { - let device = crate::backend::cpu_device(); - let cfg = toy_config(); - let loaded = LoadedQwen2 { - model: build(&cfg, &device), - config: cfg, - tokenizer_config: None, - }; - let out = loaded - .greedy_generate(&[1, 2, 3], 4, &device) - .await - .unwrap(); - assert!(out.len() <= 4); - } -} diff --git a/crates/mummu/examples/src/models/qwen3.rs b/crates/mummu/examples/src/models/qwen3.rs deleted file mode 100644 index 1c9374a..0000000 --- a/crates/mummu/examples/src/models/qwen3.rs +++ /dev/null @@ -1,738 +0,0 @@ -//! Qwen3 dense decoder, from scratch on the shared `nn` blocks. Structurally -//! Qwen2 with three deltas the shared blocks already cover: -//! * **per-head q/k RMSNorm** over `head_dim`, applied post-projection before -//! RoPE — `GqaAttention`'s `qk_norm_eps` path (the same code the LFM2 port -//! validated against Ollama; HF Qwen3 orders it identically: -//! `q_norm(q_proj(x).view(b,t,nh,hd)).transpose(1,2)`); -//! * **no q/k/v projection bias** (`attention_bias: false`); -//! * a **decoupled `head_dim`** — `num_heads * head_dim` need not equal -//! `hidden_size` (Qwen3-4B: 32·128 = 4096 vs hidden 2560), which -//! `GqaAttentionConfig` already treats as independent. -//! -//! Everything else — RmsNorm, GQA + KV cache, SwiGLU, tied/untied lm-head — is -//! the shared stack, so this file is config + weight-key remaps only. The port -//! stays `[ ]` in the roadmap until Mummu's parity gate (P7) re-verifies it -//! against a same-weights reference; loading and decoding are proven here. - -use std::path::Path; - -use burn::module::Module; -use burn::nn::{Embedding, EmbeddingConfig, Linear, LinearConfig, RmsNorm, RmsNormConfig}; -use burn::store::{ModuleAdapter, PyTorchToBurnAdapter, SafetensorsStore}; -use burn::tensor::{Device, Int, Tensor, TensorData}; - -use crate::attn_config::{RopeScaling, check_sliding_window, sliding_window_from_gguf}; -use crate::gguf::{GgufFile, GgufMap, GgufTensorInfo, GgufValue}; -use crate::import::{ - CastFloatAdapter, DequantSink, ImportError, gguf_store, load_checked, required_file, -}; -use crate::models::CausalLm; -use crate::models::qwen2::{EosIds, gguf_f32, gguf_usize}; -use crate::nn::{ - GqaAttention, GqaAttentionConfig, LayerKv, SwiGluMlp, SwiGluMlpConfig, causal_mask, rope_tables, -}; - -/// Qwen3 architecture hyperparameters, read from the checkpoint's `config.json`. -#[derive(Debug, Clone, serde::Deserialize)] -pub struct Qwen3Config { - pub vocab_size: usize, - pub hidden_size: usize, - pub intermediate_size: usize, - pub num_hidden_layers: usize, - pub num_attention_heads: usize, - pub num_key_value_heads: usize, - /// Qwen3 always ships `head_dim` explicitly (it is decoupled from - /// `hidden_size / num_attention_heads`); we still derive it if absent. - #[serde(default)] - pub head_dim: usize, - pub rms_norm_eps: f64, - pub rope_theta: f32, - /// Frequency scaling (YaRN / linear / …). `null` on the Qwen3 checkpoints - /// in the zoo; a scaled one is refused at load rather than answered wrong - /// ([`crate::attn_config`]). - /// `rope_parameters` is the same object under the name newer transformers - /// writes; reading only `rope_scaling` would let a freshly-serialized - /// scaled checkpoint through as unscaled. - #[serde(default, alias = "rope_parameters")] - pub rope_scaling: Option, - /// The trained context length (Qwen3-0.6B: 40 960), used to tell an inert - /// sliding window from a clipping one. - #[serde(default)] - pub max_position_embeddings: Option, - /// Declared window span — `null` on Qwen3-0.6B, and gated by - /// `use_sliding_window` regardless. - #[serde(default)] - pub sliding_window: Option, - #[serde(default)] - pub use_sliding_window: bool, - #[serde(default)] - pub tie_word_embeddings: bool, - /// EOS token id(s) — `<|im_end|>` first for the instruct checkpoints. - #[serde(default)] - pub eos_token_id: EosIds, -} - -impl Qwen3Config { - /// Parse `config.json` bytes; derives `head_dim` when absent. - pub fn from_json_bytes(bytes: &[u8]) -> Result { - let mut cfg: Self = serde_json::from_slice(bytes).map_err(|e| e.to_string())?; - if cfg.head_dim == 0 { - cfg.head_dim = cfg.hidden_size / cfg.num_attention_heads; - } - cfg.validate("qwen3 config.json")?; - Ok(cfg) - } - - /// Hyperparameters from a GGUF header's `qwen3.*` metadata. Unlike Qwen2, - /// `head_dim` is **required** metadata (`qwen3.attention.key_length`), - /// because Qwen3's head_dim is decoupled — deriving it from - /// `hidden / heads` is wrong for these checkpoints (4B: 80 ≠ 128). - pub fn from_gguf(f: &GgufFile) -> Result { - let arch = f.architecture().unwrap_or(""); - if arch != "qwen3" { - return Err(format!("GGUF architecture '{arch}' is not qwen3")); - } - let hidden_size = gguf_usize(f, "qwen3.embedding_length")?; - let num_attention_heads = gguf_usize(f, "qwen3.attention.head_count")?; - let embd = f - .tensor("token_embd.weight") - .ok_or("GGUF has no token_embd.weight tensor")?; - if embd.dims.len() != 2 || embd.dims[0] != hidden_size as u64 { - return Err(format!( - "token_embd.weight dims {:?} do not match embedding_length {hidden_size}", - embd.dims - )); - } - let vocab_size = usize::try_from(embd.dims[1]).map_err(|_| "vocab too large")?; - if let Some(tokens) = f.get("tokenizer.ggml.tokens").and_then(GgufValue::as_array) - && tokens.len() > vocab_size - { - return Err(format!( - "tokenizer vocab {} exceeds embedding rows {vocab_size}", - tokens.len() - )); - } - let eos_token_id = f - .get("tokenizer.ggml.eos_token_id") - .and_then(GgufValue::as_u64) - .and_then(|v| u32::try_from(v).ok()) - .map_or(EosIds::None, EosIds::One); - // key_length is Qwen3's real head_dim; only fall back for a malformed - // file, and let validate() catch an impossible result. - let head_dim = gguf_usize(f, "qwen3.attention.key_length") - .unwrap_or(hidden_size / num_attention_heads.max(1)); - let cfg = Self { - vocab_size, - hidden_size, - intermediate_size: gguf_usize(f, "qwen3.feed_forward_length")?, - num_hidden_layers: gguf_usize(f, "qwen3.block_count")?, - num_attention_heads, - num_key_value_heads: gguf_usize(f, "qwen3.attention.head_count_kv")?, - head_dim, - rms_norm_eps: f64::from(gguf_f32(f, "qwen3.attention.layer_norm_rms_epsilon")?), - rope_theta: gguf_f32(f, "qwen3.rope.freq_base")?, - rope_scaling: RopeScaling::from_gguf(f, "qwen3"), - max_position_embeddings: f - .get("qwen3.context_length") - .and_then(GgufValue::as_u64) - .and_then(|v| usize::try_from(v).ok()), - sliding_window: sliding_window_from_gguf(f, "qwen3"), - // A GGUF header has no `use_sliding_window` twin: llama.cpp writes - // the key only for architectures that window, so presence enables. - use_sliding_window: true, - // No separate output.weight tensor means the lm-head is tied. - tie_word_embeddings: f.tensor("output.weight").is_none(), - eos_token_id, - }; - cfg.validate("qwen3 GGUF header")?; - Ok(cfg) - } - - fn validate(&self, whose: &str) -> Result<(), String> { - if let Some(scaling) = &self.rope_scaling { - scaling.check(whose)?; - } - check_sliding_window( - self.sliding_window, - self.use_sliding_window, - self.max_position_embeddings, - whose, - )?; - if self.num_key_value_heads == 0 - || !self - .num_attention_heads - .is_multiple_of(self.num_key_value_heads) - { - return Err(format!( - "num_attention_heads ({}) must be a positive multiple of num_key_value_heads ({})", - self.num_attention_heads, self.num_key_value_heads - )); - } - if self.num_hidden_layers == 0 || self.vocab_size == 0 { - return Err("num_hidden_layers and vocab_size must be positive".into()); - } - if self.head_dim < 2 || !self.head_dim.is_multiple_of(2) { - return Err(format!( - "head_dim ({}) must be even and >= 2", - self.head_dim - )); - } - Ok(()) - } -} - -/// One Qwen3 decoder layer. Field names mirror the HF checkpoint (the -/// `self_attn` submodule additionally carries `q_norm` / `k_norm`). -#[derive(Module, Debug)] -pub struct DecoderLayer { - pub self_attn: GqaAttention, - pub mlp: SwiGluMlp, - pub input_layernorm: RmsNorm, - pub post_attention_layernorm: RmsNorm, -} - -/// The Qwen3 decoder stack (HF's `model.*` subtree). Tied on the small tiers -/// (0.6B/4B safetensors, and any GGUF without a separate `output.weight`). -#[derive(Module, Debug)] -pub struct Qwen3 { - pub embed_tokens: Embedding, - pub layers: Vec, - pub norm: RmsNorm, - pub lm_head: Option, -} - -/// A weight-loaded Qwen3 plus its config — everything a forward needs. -pub struct LoadedQwen3 { - pub model: Qwen3, - pub config: Qwen3Config, - /// The parsed sibling `tokenizer_config.json`, when one was present and - /// well-formed beside a safetensors checkpoint (the load-time gate has - /// already cross-checked its EOS against `config.json`). A consumer reads - /// config-driven EOS/BOS/PAD ids from it (`eos_id()`, `bos_id()`, …). `None` - /// for a GGUF load (self-contained; no sibling file) or a dir without one. - pub tokenizer_config: Option, -} - -fn build(cfg: &Qwen3Config, device: &Device) -> Qwen3 { - let attn_cfg = GqaAttentionConfig { - hidden_size: cfg.hidden_size, - num_heads: cfg.num_attention_heads, - num_kv_heads: cfg.num_key_value_heads, - head_dim: cfg.head_dim, - bias: false, // Qwen3 dropped the q/k/v bias - qk_norm_eps: Some(cfg.rms_norm_eps), // and added per-head q/k RMSNorm - qk_norm_projection: false, - }; - let mlp_cfg = SwiGluMlpConfig { - hidden_size: cfg.hidden_size, - intermediate_size: cfg.intermediate_size, - }; - let norm = |dev: &Device| { - RmsNormConfig::new(cfg.hidden_size) - .with_epsilon(cfg.rms_norm_eps) - .init(dev) - }; - let layers = (0..cfg.num_hidden_layers) - .map(|_| DecoderLayer { - self_attn: attn_cfg.init(device), - mlp: mlp_cfg.init(device), - input_layernorm: norm(device), - post_attention_layernorm: norm(device), - }) - .collect(); - let lm_head = (!cfg.tie_word_embeddings).then(|| { - LinearConfig::new(cfg.hidden_size, cfg.vocab_size) - .with_bias(false) - .init(device) - }); - Qwen3 { - embed_tokens: EmbeddingConfig::new(cfg.vocab_size, cfg.hidden_size).init(device), - layers, - norm: norm(device), - lm_head, - } -} - -/// The safetensors key remap: strip `model.`, rename every RmsNorm `weight` → -/// Burn's `gamma`. Qwen3 adds the per-head `self_attn.q_norm` / `k_norm` to the -/// set of norms the qwen2 chain already handled. -fn install_remaps(store: SafetensorsStore) -> SafetensorsStore { - store - .with_key_remapping(r"^model\.", "") - .with_key_remapping(r"(input_layernorm)\.weight$", "$1.gamma") - .with_key_remapping(r"(post_attention_layernorm)\.weight$", "$1.gamma") - .with_key_remapping(r"(self_attn\.q_norm)\.weight$", "$1.gamma") - .with_key_remapping(r"(self_attn\.k_norm)\.weight$", "$1.gamma") - .with_key_remapping(r"^norm\.weight$", "norm.gamma") -} - -/// Build from `dir/config.json` and load `dir/model.safetensors`, checked. -pub fn load_from_dir(dir: &Path, device: &Device) -> Result { - let cfg_path = required_file(dir, "config.json")?; - let weights = required_file(dir, "model.safetensors")?; - let cfg_bytes = std::fs::read(&cfg_path).map_err(|e| ImportError::Parse { - file: cfg_path.clone(), - reason: e.to_string(), - })?; - let config = Qwen3Config::from_json_bytes(&cfg_bytes).map_err(|reason| ImportError::Parse { - file: cfg_path, - reason, - })?; - - // Cross-check the sibling metadata (when present) before touching weights: - // tokenizer_config.json's EOS must agree with config.json's, its chat-template - // must not speak a different tool-call convention than Qwen3's Hermes/ChatML - // renderer, and every added-token id it declares must match the real - // tokenizer.json. A repackaging mismatch fails loudly here rather than - // mis-stopping / mis-templating / mis-tokenizing later. - let tokenizer_config = crate::tokenizer::validate_checkpoint_dir( - dir, - &config.eos_token_id.to_vec(), - Some(crate::tok_config::ToolCallConvention::Hermes), - )?; - debug_assert!( - tokenizer_config - .as_ref() - .and_then(crate::tok_config::TokenizerConfig::eos_id) - .is_none_or(|id| config.eos_token_id.contains(id)), - "validate_checkpoint_dir returned a config whose EOS disagrees with config.json" - ); - - let mut model = build(&config, device); - // The float dtype comes from the DEVICE — burn 0.22 keeps the element - // type there as a runtime setting, not on a backend type. Creation sites - // still name it explicitly rather than riding the unspecified default. - let target_float = crate::backend::float_dtype(device); - let mut store = install_remaps( - SafetensorsStore::from_file(weights.clone()) - .with_from_adapter(PyTorchToBurnAdapter.chain(CastFloatAdapter::new(target_float))) - .allow_partial(true), - ); - load_checked(&mut model, &mut store, &weights)?; - Ok(LoadedQwen3 { - model, - config, - tokenizer_config, - }) -} - -/// GGUF (llama.cpp `qwen3` arch) tensor names → the HF checkpoint names the -/// safetensors remap chain handles. `None` for anything unrecognized. -fn gguf_tensor_to_hf(info: &GgufTensorInfo) -> Option { - qwen3_gguf_name(&info.name).map(GgufMap::Rename) -} - -fn qwen3_gguf_name(name: &str) -> Option { - match name { - "token_embd.weight" => return Some("model.embed_tokens.weight".into()), - "output_norm.weight" => return Some("model.norm.weight".into()), - "output.weight" => return Some("lm_head.weight".into()), - _ => {} - } - let rest = name.strip_prefix("blk.")?; - let (layer, field) = rest.split_once('.')?; - let layer: usize = layer.parse().ok()?; - let mapped = match field { - "attn_norm.weight" => "input_layernorm.weight", - "ffn_norm.weight" => "post_attention_layernorm.weight", - "attn_q.weight" => "self_attn.q_proj.weight", - "attn_k.weight" => "self_attn.k_proj.weight", - "attn_v.weight" => "self_attn.v_proj.weight", - "attn_q_norm.weight" => "self_attn.q_norm.weight", - "attn_k_norm.weight" => "self_attn.k_norm.weight", - "attn_output.weight" => "self_attn.o_proj.weight", - "ffn_gate.weight" => "mlp.gate_proj.weight", - "ffn_up.weight" => "mlp.up_proj.weight", - "ffn_down.weight" => "mlp.down_proj.weight", - _ => return None, - }; - Some(format!("model.layers.{layer}.{mapped}")) -} - -/// Load a Qwen3 model straight from a **GGUF** file: hyperparameters from the -/// `qwen3.*` metadata, weights dequantized to f32 and driven through the same -/// checked-load pipeline (adapters + remaps) the safetensors path uses. -pub fn load_from_gguf(path: &Path, device: &Device) -> Result { - let parse = |reason: String| ImportError::Parse { - file: path.to_path_buf(), - reason, - }; - let f = GgufFile::open(path).map_err(|e| parse(e.to_string()))?; - let config = Qwen3Config::from_gguf(&f).map_err(parse)?; - // The scratch guard (Some only when the payload went to disk) must - // outlive `load_checked`: the store reads that file lazily. - let (base, _scratch) = gguf_store(&f, &gguf_tensor_to_hf, DequantSink::Auto, device)?; - - let mut model = build(&config, device); - let mut store = install_remaps(base); - load_checked(&mut model, &mut store, path)?; - // A GGUF is self-contained — no sibling tokenizer_config.json in this path. - Ok(LoadedQwen3 { - model, - config, - tokenizer_config: None, - }) -} - -impl CausalLm for LoadedQwen3 { - type Cache = Vec; - - fn new_cache(&self) -> Self::Cache { - (0..self.config.num_hidden_layers).map(|_| None).collect() - } - - fn is_eos(&self, id: u32) -> bool { - self.config.eos_token_id.contains(id) - } - - fn forward( - &self, - new_ids: &[u32], - past: usize, - cache: &mut Self::Cache, - device: &Device, - ) -> Tensor<2> { - let t = new_ids.len(); - assert!(t >= 1, "Qwen3 forward: need at least one token"); - assert!( - cache.len() == self.config.num_hidden_layers, - "Qwen3 forward: cache has {} layers, model has {}", - cache.len(), - self.config.num_hidden_layers - ); - let cfg = &self.config; - - // Dtype pinned to the backend TYPE, never the per-device policy. - let ids32: Vec = new_ids.iter().map(|&i| i as i32).collect(); - let input = Tensor::<1, Int>::from_data( - TensorData::new(ids32, [t]), - (device, crate::backend::int_dtype(device)), - ) - .reshape([1, t]); - let mut x = self.model.embed_tokens.forward(input); // [1, t, hidden] - - let (cos, sin) = rope_tables(t, past, cfg.head_dim, cfg.rope_theta, device); - let mask = (t > 1).then(|| causal_mask(t, past, device)); - - for (layer, kv) in self.model.layers.iter().zip(cache.iter_mut()) { - let h = layer.input_layernorm.forward(x.clone()); - let h = layer.self_attn.forward( - h, - cfg.num_attention_heads, - cfg.num_key_value_heads, - cfg.head_dim, - &cos, - &sin, - mask.as_ref(), - kv, - ); - x = x.add(h); - let h2 = layer.post_attention_layernorm.forward(x.clone()); - x = x.add(layer.mlp.forward(h2)); - } - let x = self.model.norm.forward(x); - - let last = x.narrow(1, t - 1, 1).reshape([1, cfg.hidden_size]); - debug_assert!( - self.model.lm_head.is_some() != cfg.tie_word_embeddings, - "lm_head presence must match the config's tie flag" - ); - match &self.model.lm_head { - Some(head) => head.forward(last), // [1, vocab] - None => { - let w = self.model.embed_tokens.weight.val(); // [vocab, hidden] - last.matmul(w.swap_dims(0, 1)) // [1, vocab] - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::gguf::{GgmlType, GgufTensorInfo}; - - /// A synthetic toy config with a **decoupled** head_dim (num_heads·head_dim - /// = 4·6 = 24 ≠ hidden 16), exercising the Qwen3-specific shape path. - fn toy_config() -> Qwen3Config { - Qwen3Config { - vocab_size: 64, - hidden_size: 16, - intermediate_size: 32, - num_hidden_layers: 2, - num_attention_heads: 4, - num_key_value_heads: 2, - head_dim: 6, - rms_norm_eps: 1e-6, - rope_theta: 1e6, - rope_scaling: None, - max_position_embeddings: Some(512), - sliding_window: None, - use_sliding_window: false, - tie_word_embeddings: true, - eos_token_id: EosIds::One(2), - } - } - - #[test] - fn config_parses_qwen3_4b_shape() { - // The real Qwen3-4B config.json shape: head_dim is explicit and - // decoupled (32·128 = 4096 ≠ hidden 2560). - let json = br#"{ - "vocab_size": 151936, "hidden_size": 2560, "intermediate_size": 9728, - "num_hidden_layers": 36, "num_attention_heads": 32, "num_key_value_heads": 8, - "head_dim": 128, "rms_norm_eps": 1e-6, "rope_theta": 1000000.0, - "tie_word_embeddings": true, "eos_token_id": 151645 - }"#; - let cfg = Qwen3Config::from_json_bytes(json).unwrap(); - assert_eq!(cfg.head_dim, 128); // taken verbatim, NOT derived to 80 - assert_ne!(cfg.head_dim, cfg.hidden_size / cfg.num_attention_heads); - assert!(cfg.eos_token_id.contains(151_645)); - assert!(cfg.tie_word_embeddings); - } - - /// Qwen3-0.6B's own config: `rope_scaling: null`, `sliding_window: null`, - /// `use_sliding_window: false`. The load Mummu parity-verifies must keep - /// working with the new fields present. - #[test] - fn the_real_qwen3_06b_shape_with_null_scaling_still_loads() { - let json = br#"{ - "vocab_size": 151936, "hidden_size": 1024, "intermediate_size": 3072, - "num_hidden_layers": 28, "num_attention_heads": 16, "num_key_value_heads": 8, - "head_dim": 128, "rms_norm_eps": 1e-6, "rope_theta": 1000000, - "rope_scaling": null, "sliding_window": null, "use_sliding_window": false, - "max_window_layers": 28, "max_position_embeddings": 40960, - "tie_word_embeddings": true, "eos_token_id": 151645 - }"#; - let cfg = Qwen3Config::from_json_bytes(json).expect("the zoo's own shape must load"); - assert!(cfg.rope_scaling.is_none() && cfg.sliding_window.is_none()); - assert_eq!(cfg.max_position_embeddings, Some(40960)); - } - - /// The pre-4.38 spelling (`"type"` rather than `"rope_type"`) is read too, - /// so an older linear-scaled checkpoint cannot slip past. - #[test] - fn a_linear_scaled_checkpoint_is_refused_in_the_legacy_spelling() { - let json = br#"{ - "vocab_size": 151936, "hidden_size": 1024, "intermediate_size": 3072, - "num_hidden_layers": 28, "num_attention_heads": 16, "num_key_value_heads": 8, - "head_dim": 128, "rms_norm_eps": 1e-6, "rope_theta": 1000000, - "rope_scaling": {"type": "linear", "factor": 2.0} - }"#; - let err = Qwen3Config::from_json_bytes(json).expect_err("linear must refuse"); - assert!(err.contains("linear"), "{err}"); - assert!(err.contains("qwen3 config.json"), "{err}"); - } - - /// Newer transformers serializes the same object as `rope_parameters`. - /// Reading only `rope_scaling` would let a scaled checkpoint through as - /// unscaled — the exact silent failure this whole gate exists to stop. - #[test] - fn the_newer_rope_parameters_spelling_is_read_as_well() { - let json = br#"{ - "vocab_size": 151936, "hidden_size": 1024, "intermediate_size": 3072, - "num_hidden_layers": 28, "num_attention_heads": 16, "num_key_value_heads": 8, - "head_dim": 128, "rms_norm_eps": 1e-6, "rope_theta": 1000000, - "rope_parameters": {"rope_type": "yarn", "rope_theta": 1000000.0, "factor": 4.0} - }"#; - let err = Qwen3Config::from_json_bytes(json).expect_err("yarn must refuse"); - assert!(err.contains("yarn"), "{err}"); - // And the plain spelling of the same field still loads. - let plain = br#"{ - "vocab_size": 151936, "hidden_size": 1024, "intermediate_size": 3072, - "num_hidden_layers": 28, "num_attention_heads": 16, "num_key_value_heads": 8, - "head_dim": 128, "rms_norm_eps": 1e-6, "rope_theta": 1000000, - "rope_parameters": {"rope_type": "default", "rope_theta": 1000000.0} - }"#; - assert!(Qwen3Config::from_json_bytes(plain).is_ok()); - } - - #[test] - fn config_derives_head_dim_when_absent() { - let json = br#"{ - "vocab_size": 100, "hidden_size": 32, "intermediate_size": 64, - "num_hidden_layers": 1, "num_attention_heads": 4, "num_key_value_heads": 2, - "rms_norm_eps": 1e-6, "rope_theta": 1000000.0 - }"#; - let cfg = Qwen3Config::from_json_bytes(json).unwrap(); - assert_eq!(cfg.head_dim, 8); // 32 / 4 - } - - #[test] - fn config_rejects_indivisible_heads_and_odd_head_dim() { - let bad_heads = br#"{ - "vocab_size": 100, "hidden_size": 16, "intermediate_size": 32, - "num_hidden_layers": 1, "num_attention_heads": 5, "num_key_value_heads": 2, - "head_dim": 4, "rms_norm_eps": 1e-6, "rope_theta": 1e6 - }"#; - assert!(Qwen3Config::from_json_bytes(bad_heads).is_err()); - let odd_hd = br#"{ - "vocab_size": 100, "hidden_size": 16, "intermediate_size": 32, - "num_hidden_layers": 1, "num_attention_heads": 4, "num_key_value_heads": 2, - "head_dim": 5, "rms_norm_eps": 1e-6, "rope_theta": 1e6 - }"#; - assert!(Qwen3Config::from_json_bytes(odd_hd).is_err()); - } - - /// The load-bearing invariant: cached prefill+decode == one full forward. - #[test] - fn toy_model_cached_decode_matches_full_forward() { - let device = crate::backend::cpu_device(); - let cfg = toy_config(); - let loaded = LoadedQwen3 { - model: build(&cfg, &device), - config: cfg, - tokenizer_config: None, - }; - - let prompt: Vec = vec![3, 14, 15, 9, 26]; - let mut cache = loaded.new_cache(); - let _ = loaded.forward(&prompt, 0, &mut cache, &device); - let step = loaded - .forward(&[42], prompt.len(), &mut cache, &device) - .into_data() - .to_vec::() - .unwrap(); - - let mut full_cache = loaded.new_cache(); - let all: Vec = prompt.iter().copied().chain([42]).collect(); - let full = loaded - .forward(&all, 0, &mut full_cache, &device) - .into_data() - .to_vec::() - .unwrap(); - - assert_eq!(step.len(), full.len()); - for (i, (c, f)) in step.iter().zip(&full).enumerate() { - assert!((c - f).abs() < 1e-4, "logit {i}: cached {c} vs full {f}"); - } - } - - #[test] - fn untied_toy_config_builds_and_uses_an_lm_head() { - let device = crate::backend::cpu_device(); - let mut cfg = toy_config(); - cfg.tie_word_embeddings = false; - let vocab = cfg.vocab_size; - let loaded = LoadedQwen3 { - model: build(&cfg, &device), - config: cfg, - tokenizer_config: None, - }; - assert!(loaded.model.lm_head.is_some()); - let mut cache = loaded.new_cache(); - let logits = loaded.forward(&[1, 2], 0, &mut cache, &device); - assert_eq!(logits.dims(), [1, vocab]); - } - - /// A synthetic in-memory GGUF header shaped like a small `qwen3` file, - /// with the decoupled `key_length` metadata. - fn toy_gguf() -> GgufFile { - let meta = |k: &str, v: GgufValue| (k.to_string(), v); - GgufFile { - path: std::path::PathBuf::new(), - version: 3, - metadata: vec![ - meta("general.architecture", GgufValue::Str("qwen3".into())), - meta("qwen3.embedding_length", GgufValue::U32(16)), - meta("qwen3.block_count", GgufValue::U32(2)), - meta("qwen3.feed_forward_length", GgufValue::U32(32)), - meta("qwen3.attention.head_count", GgufValue::U32(4)), - meta("qwen3.attention.head_count_kv", GgufValue::U32(2)), - meta("qwen3.attention.key_length", GgufValue::U32(6)), - meta( - "qwen3.attention.layer_norm_rms_epsilon", - GgufValue::F32(1e-6), - ), - meta("qwen3.rope.freq_base", GgufValue::F32(1e6)), - meta("tokenizer.ggml.eos_token_id", GgufValue::U32(2)), - ], - tensors: vec![GgufTensorInfo { - name: "token_embd.weight".into(), - dims: vec![16, 64], // ggml order: [hidden, vocab] - dtype: GgmlType::F32, - offset: 0, - }], - alignment: 32, - data_offset: 0, - } - } - - #[test] - fn config_from_gguf_reads_decoupled_head_dim_from_key_length() { - let cfg = Qwen3Config::from_gguf(&toy_gguf()).expect("parses"); - assert_eq!(cfg.vocab_size, 64); - assert_eq!(cfg.hidden_size, 16); - assert_eq!(cfg.num_hidden_layers, 2); - // key_length (6), NOT hidden/heads (4) — the decoupled path. - assert_eq!(cfg.head_dim, 6); - assert!(cfg.eos_token_id.contains(2)); - assert!(cfg.tie_word_embeddings); // no output.weight tensor - } - - #[test] - fn config_from_gguf_fails_loudly_on_missing_keys_and_wrong_arch() { - let mut f = toy_gguf(); - f.metadata.retain(|(k, _)| k != "qwen3.block_count"); - assert!( - Qwen3Config::from_gguf(&f) - .unwrap_err() - .contains("block_count") - ); - - let mut f = toy_gguf(); - f.metadata[0].1 = GgufValue::Str("qwen2".into()); - assert!(Qwen3Config::from_gguf(&f).is_err()); - } - - #[test] - fn config_from_gguf_detects_untied_head() { - let mut f = toy_gguf(); - f.tensors.push(GgufTensorInfo { - name: "output.weight".into(), - dims: vec![16, 64], - dtype: GgmlType::F32, - offset: 4096, - }); - assert!(!Qwen3Config::from_gguf(&f).unwrap().tie_word_embeddings); - } - - #[test] - fn gguf_names_map_including_qk_norms() { - assert_eq!( - qwen3_gguf_name("blk.0.attn_q_norm.weight").as_deref(), - Some("model.layers.0.self_attn.q_norm.weight") - ); - assert_eq!( - qwen3_gguf_name("blk.35.attn_k_norm.weight").as_deref(), - Some("model.layers.35.self_attn.k_norm.weight") - ); - assert_eq!( - qwen3_gguf_name("blk.7.attn_q.weight").as_deref(), - Some("model.layers.7.self_attn.q_proj.weight") - ); - assert_eq!( - qwen3_gguf_name("output.weight").as_deref(), - Some("lm_head.weight") - ); - // Qwen3 has no q/k/v bias — a bias tensor is unrecognized (loud error). - assert_eq!(qwen3_gguf_name("blk.0.attn_q.bias"), None); - assert_eq!(qwen3_gguf_name("rope_freqs.weight"), None); - } - - #[tokio::test] - async fn greedy_generate_respects_max_tokens_bound() { - let device = crate::backend::cpu_device(); - let cfg = toy_config(); - let loaded = LoadedQwen3 { - model: build(&cfg, &device), - config: cfg, - tokenizer_config: None, - }; - let out = loaded - .greedy_generate(&[1, 2, 3], 4, &device) - .await - .unwrap(); - assert!(out.len() <= 4); - } -} diff --git a/crates/mummu/examples/src/models/qwen35.rs b/crates/mummu/examples/src/models/qwen35.rs deleted file mode 100644 index 5045d22..0000000 --- a/crates/mummu/examples/src/models/qwen35.rs +++ /dev/null @@ -1,1941 +0,0 @@ -//! Qwen3.5 / Qwen3.8 ("qwen35") hybrid decoder: Gated DeltaNet linear -//! attention on three of every four layers, gated full attention (partial -//! RoPE) on the fourth, SwiGLU MLPs, RMSNorm everywhere, tied or untied head. -//! -//! Ported from llama.cpp's reference (`src/models/qwen35.cpp` + -//! `delta-net-base.cpp`, fetched 2026-08-21), the only implementation with -//! local same-weights parity available (`llama-server` runs these GGUFs). -//! Per layer: -//! -//! - **Full attention** (`(i+1) % full_attention_interval == 0`): the q -//! projection emits query and a per-head **output gate** interleaved -//! (`[q_h | gate_h]` per head); per-head q/k RMSNorm; RoPE over only the -//! first `rope_dim` of the 256-wide heads (the metadata's MRoPE sections -//! degenerate to standard RoPE for text-only inputs); softmax attention; -//! then `out ⊙ sigmoid(gate)` before the output projection. -//! - **Gated DeltaNet** (the rest): one projection mixes q/k/v, a second -//! emits the gate `z`; the mix runs through a depthwise causal conv -//! (kernel `conv_kernel`, rolling state) + SiLU; q/k are L2-normalized -//! per head (`x / max(‖x‖, ε)`) and tiled from `n_k_heads` to -//! `n_v_heads`; the recurrence per head with state `S ∈ R^{d_k×d_v}`: -//! `S ← S·exp(g); v̂ = Sᵀk; S += k(β(v − v̂))ᵀ; o = Sᵀ(q/√d_k)` with -//! `β = σ(x·Wβ)` and `g = softplus(x·Wα + dt_bias)·a` (`a` holds -//! `-exp(A_log)`, negative); the output is gated-RMS-normed -//! (`RMS(o)·silu(z)` per head) and projected back. -//! -//! The NextN/MTP block some checkpoints append (`nextn_predict_layers = 1`) -//! is a draft head for speculative decoding, unused by the main forward — -//! its tensors are explicitly skipped on import. - -use std::path::Path; - -use burn::module::{Module, Param}; -use burn::nn::conv::{Conv1d, Conv1dConfig}; -use burn::nn::{ - Embedding, EmbeddingConfig, Linear, LinearConfig, PaddingConfig1d, RmsNorm, RmsNormConfig, -}; -use burn::tensor::{DType, Device, Int, Tensor, TensorData, activation}; - -use crate::gguf::{GgufFile, GgufMap, GgufTensorInfo, GgufValue}; -use crate::import::ImportError; -use crate::models::CausalLm; -use crate::models::qwen2::EosIds; -use crate::nn::{LayerKv, SwiGluMlp, SwiGluMlpConfig, causal_mask, repeat_kv, rope_tables}; -use crate::quant::QuantPolicy; - -/// Architecture hyperparameters, read from a GGUF header's `qwen35.*` -/// metadata (the family currently ships as GGUF; a safetensors `config.json` -/// path can join later). -#[derive(Debug, Clone)] -pub struct Qwen35Config { - pub vocab_size: usize, - pub hidden_size: usize, - /// Trunk layers only — `block_count - nextn_predict_layers`. - pub num_layers: usize, - pub num_attention_heads: usize, - pub num_key_value_heads: usize, - /// Attention head width (`key_length` == `value_length`; 256 across the - /// family — decoupled from `hidden_size / num_heads`). - pub head_dim: usize, - pub intermediate_size: usize, - pub rms_norm_eps: f64, - pub rope_theta: f32, - /// How many leading dims of each head RoPE rotates (`rope.dimension_count`). - pub rope_dim: usize, - /// Layer `i` is full attention iff `(i+1) % interval == 0`. - pub full_attention_interval: usize, - /// Depthwise conv kernel length in the DeltaNet mix path. - pub conv_kernel: usize, - /// DeltaNet value width (`ssm.inner_size` = `n_v_heads · d_state`). - pub d_inner: usize, - /// Per-head key/value width (`ssm.state_size`). - pub d_state: usize, - /// DeltaNet key/query heads (`ssm.group_count`). - pub n_k_heads: usize, - /// DeltaNet value heads (`ssm.time_step_rank` — llama.cpp's reuse). - pub n_v_heads: usize, - pub eos_token_id: EosIds, -} - -impl Qwen35Config { - #[must_use] - pub fn is_attention(&self, layer: usize) -> bool { - (layer + 1).is_multiple_of(self.full_attention_interval) - } - - /// q/k projection width in the DeltaNet mix (`n_k_heads · d_state`). - #[must_use] - pub fn key_dim(&self) -> usize { - self.n_k_heads * self.d_state - } - - /// Channels through the DeltaNet conv: q + k + v concatenated. - #[must_use] - pub fn conv_dim(&self) -> usize { - 2 * self.key_dim() + self.d_inner - } - - /// Hyperparameters from a GGUF header's `qwen35.*` metadata. - pub fn from_gguf(f: &GgufFile) -> Result { - let arch = f.architecture().unwrap_or(""); - if arch != "qwen35" { - return Err(format!("GGUF architecture '{arch}' is not qwen35")); - } - let meta_usize = |key: &str| -> Result { - f.get(key) - .and_then(GgufValue::as_u64) - .map(|v| usize::try_from(v).expect("metadata fits usize")) - .ok_or_else(|| format!("GGUF metadata missing {key}")) - }; - let meta_f32 = |key: &str| -> Result { - f.get(key) - .and_then(GgufValue::as_f32) - .ok_or_else(|| format!("GGUF metadata missing {key}")) - }; - let embd = f - .tensor("token_embd.weight") - .ok_or("GGUF has no token_embd.weight")?; - // ggml dims are fastest-varying first: [hidden, vocab]. - let vocab_size = usize::try_from(*embd.dims.get(1).ok_or("token_embd is not 2-D")?) - .expect("vocab fits usize"); - - let block_count = meta_usize("qwen35.block_count")?; - let nextn = f - .get("qwen35.nextn_predict_layers") - .and_then(GgufValue::as_u64) - .map_or(0, |v| usize::try_from(v).expect("small")); - if nextn >= block_count { - return Err(format!( - "nextn_predict_layers ({nextn}) must be below block_count ({block_count})" - )); - } - let key_length = meta_usize("qwen35.attention.key_length")?; - let value_length = meta_usize("qwen35.attention.value_length")?; - if key_length != value_length { - return Err(format!( - "key_length ({key_length}) != value_length ({value_length}) is not implemented" - )); - } - - let eos = f - .get("tokenizer.ggml.eos_token_id") - .and_then(GgufValue::as_u64) - .ok_or("GGUF metadata missing tokenizer.ggml.eos_token_id")?; - - let cfg = Self { - vocab_size, - hidden_size: meta_usize("qwen35.embedding_length")?, - num_layers: block_count - nextn, - num_attention_heads: meta_usize("qwen35.attention.head_count")?, - num_key_value_heads: meta_usize("qwen35.attention.head_count_kv")?, - head_dim: key_length, - intermediate_size: meta_usize("qwen35.feed_forward_length")?, - rms_norm_eps: f64::from(meta_f32("qwen35.attention.layer_norm_rms_epsilon")?), - rope_theta: meta_f32("qwen35.rope.freq_base")?, - rope_dim: meta_usize("qwen35.rope.dimension_count")?, - full_attention_interval: f - .get("qwen35.full_attention_interval") - .and_then(GgufValue::as_u64) - .map_or(4, |v| usize::try_from(v).expect("small")), - conv_kernel: meta_usize("qwen35.ssm.conv_kernel")?, - d_inner: meta_usize("qwen35.ssm.inner_size")?, - d_state: meta_usize("qwen35.ssm.state_size")?, - n_k_heads: meta_usize("qwen35.ssm.group_count")?, - n_v_heads: meta_usize("qwen35.ssm.time_step_rank")?, - eos_token_id: EosIds::One(u32::try_from(eos).map_err(|_| "EOS out of u32")?), - }; - cfg.validate()?; - Ok(cfg) - } - - fn validate(&self) -> Result<(), String> { - if self.d_inner != self.n_v_heads * self.d_state { - return Err(format!( - "ssm.inner_size ({}) != n_v_heads ({}) · d_state ({}) — layout not implemented", - self.d_inner, self.n_v_heads, self.d_state - )); - } - if !self.n_v_heads.is_multiple_of(self.n_k_heads) { - return Err(format!( - "n_v_heads ({}) must be a multiple of n_k_heads ({})", - self.n_v_heads, self.n_k_heads - )); - } - if self.rope_dim > self.head_dim || !self.rope_dim.is_multiple_of(2) { - return Err(format!( - "rope_dim ({}) must be even and <= head_dim ({})", - self.rope_dim, self.head_dim - )); - } - if self.full_attention_interval == 0 || self.conv_kernel < 2 { - return Err("degenerate full_attention_interval or conv_kernel".into()); - } - Ok(()) - } -} - -/// `Linear::forward` without touching the weight's shape: burn's Linear -/// unsqueezes the weight to the input rank, and a reshape on **packed** -/// quantized storage is broken in burn 0.21 (physical element count differs -/// from the logical one — measured with Q4 on flex AND wgpu, 2026-08-21). -/// Flatten the FLOAT input instead; the weight goes into the matmul as-is. -/// All qwen35 projections are bias-free. -fn qlinear(l: &Linear, x: Tensor<3>) -> Tensor<3> { - debug_assert!(l.bias.is_none(), "qwen35 projections are bias-free"); - let [b, t, d_in] = x.dims(); - let w = l.weight.val(); // [in, out] - let d_out = w.dims()[1]; - let x2 = x.reshape([b * t, d_in]); - // Decode-shape quantized weights take the packed GEMV (reads the - // stored bytes directly; on flex that is the i8 slab at 1.125 B/elem - // against the 4 B/elem an f32 slab moves). Anything else — prefill, - // float weights — keeps the plain matmul. - let y = match crate::nn::try_q4s_gemv(&x2, &w) { - Some(y) => y, - None => x2.matmul(w), - }; - y.reshape([b, t, d_out]) -} - -/// The 2-D twin of [`qlinear`] (the lm-head path). -fn qlinear2(l: &Linear, x: Tensor<2>) -> Tensor<2> { - debug_assert!(l.bias.is_none(), "qwen35 projections are bias-free"); - let w = l.weight.val(); - match crate::nn::try_q4s_gemv(&x, &w) { - Some(y) => y, - None => x.matmul(w), - } -} - -/// Gated full attention (see the module docs). Field names are this port's -/// own (the family has no HF safetensors convention to mirror yet). -#[derive(Module, Debug)] -pub struct GatedAttention { - /// Emits `[q_h | gate_h]` interleaved per head — `2·num_heads·head_dim` wide. - pub q_proj: Linear, - pub k_proj: Linear, - pub v_proj: Linear, - pub o_proj: Linear, - /// Per-head RMSNorm over `head_dim`. - pub q_norm: RmsNorm, - pub k_norm: RmsNorm, -} - -impl GatedAttention { - #[allow(clippy::too_many_arguments)] // mirrors the reference data flow - fn forward( - &self, - x: Tensor<3>, - cfg: &Qwen35Config, - cos: &Tensor<4>, - sin: &Tensor<4>, - mask: Option<&Tensor<4>>, - kv: &mut LayerKv, - ) -> Tensor<3> { - let [b, t, _] = x.dims(); - let (nh, nkv, hd) = ( - cfg.num_attention_heads, - cfg.num_key_value_heads, - cfg.head_dim, - ); - - // Split the joint projection into q and gate: per head the layout is - // [q (hd) | gate (hd)], so a [b, t, nh, 2, hd] view separates them. - let _s_qkv = crate::prof::scope("fa.qkv"); - let qg = qlinear(&self.q_proj, x.clone()).reshape([b, t, nh, 2, hd]); - let q = qg.clone().narrow(3, 0, 1).reshape([b, t, nh, hd]); - let gate = qg.narrow(3, 1, 1).reshape([b, t, nh, hd]); - - let q = self.q_norm.forward(q).swap_dims(1, 2); // [b, nh, t, hd] - let k_new = qlinear(&self.k_proj, x.clone()).reshape([b, t, nkv, hd]); - let k_new = self.k_norm.forward(k_new).swap_dims(1, 2); - let v_new = qlinear(&self.v_proj, x) - .reshape([b, t, nkv, hd]) - .swap_dims(1, 2); - - drop(_s_qkv); - let _s_rope = crate::prof::scope("fa.rope"); - // Partial RoPE: rotate the first rope_dim dims, pass the rest through. - let rope = |x: Tensor<4>| -> Tensor<4> { - let rot = x.clone().narrow(3, 0, cfg.rope_dim); - let rest = x.narrow(3, cfg.rope_dim, hd - cfg.rope_dim); - let rot = crate::nn::apply_rope(rot, cos, sin); - Tensor::cat(vec![rot, rest], 3) - }; - let q = rope(q); - let k_new = rope(k_new); - drop(_s_rope); - let _s_kv = crate::prof::scope("fa.kv"); - - let (k_all, v_all) = match kv.take() { - Some((pk, pv)) => ( - Tensor::cat(vec![pk, k_new], 2), - Tensor::cat(vec![pv, v_new], 2), - ), - None => (k_new, v_new), - }; - *kv = Some((k_all.clone(), v_all.clone())); - - let group = nh / nkv; - let k = repeat_kv(k_all, group); - let v = repeat_kv(v_all, group); - drop(_s_kv); - let _s_scores = crate::prof::scope("fa.scores"); - - // f32 island for the scores — the same overflow guard as GqaAttention. - let ambient = q.dtype(); - let scale = 1.0 / (hd as f32).sqrt(); - let mut scores = q - .cast(DType::F32) - .matmul(k.cast(DType::F32).swap_dims(2, 3)) - .mul_scalar(scale); - if let Some(m) = mask { - scores = scores.add(m.clone().cast(DType::F32)); - } - let probs = activation::softmax(scores, 3).cast(ambient); - let ctx = probs.matmul(v); // [b, nh, t, hd] - drop(_s_scores); - let _s = crate::prof::scope("fa.out"); - - // Per-head output gate: out ⊙ sigmoid(gate). - let gated = ctx - .swap_dims(1, 2) // [b, t, nh, hd] - .mul(activation::sigmoid(gate)) - .reshape([b, t, nh * hd]); - qlinear(&self.o_proj, gated) - } -} - -/// Gated DeltaNet linear attention (see the module docs). -#[derive(Module, Debug)] -pub struct GatedDeltaNet { - /// Mixes q/k/v: `hidden → 2·key_dim + d_inner`. - pub qkv_proj: Linear, - /// The gate `z`: `hidden → d_inner`. - pub z_proj: Linear, - /// Per-value-head β logits: `hidden → n_v_heads`. - pub beta_proj: Linear, - /// Per-value-head decay logits: `hidden → n_v_heads`. - pub alpha_proj: Linear, - /// Decay bias added to the α logits before softplus. - pub dt_bias: Param>, - /// `-exp(A_log)` — negative per-head decay magnitudes. - pub a: Param>, - /// Depthwise causal conv over the q/k/v mix, kernel `conv_kernel`. - pub conv1d: Conv1d, - /// Gated output RMSNorm over `d_state` (per value head). - pub norm: RmsNorm, - pub out_proj: Linear, -} - -/// DeltaNet decode state: the rolling conv window and the recurrent memory. -pub struct DeltaState { - /// Last `conv_kernel - 1` mix columns, `[b, conv_dim, k-1]`. - pub conv: Option>, - /// Per-head associative memory `[b, n_v_heads, d_state, d_state]`. - pub state: Option>, -} - -impl GatedDeltaNet { - fn forward(&self, x: Tensor<3>, cfg: &Qwen35Config, cache: &mut DeltaState) -> Tensor<3> { - let [b, t, _] = x.dims(); - let (hk, hv, ds) = (cfg.n_k_heads, cfg.n_v_heads, cfg.d_state); - let key_dim = cfg.key_dim(); - let conv_dim = cfg.conv_dim(); - let kk = cfg.conv_kernel; - let device = x.device(); - - let _s_proj = crate::prof::scope("delta.proj"); - let mixed = qlinear(&self.qkv_proj, x.clone()); // [b, t, conv_dim] - let z = qlinear(&self.z_proj, x.clone()); // [b, t, d_inner] - let beta = activation::sigmoid(qlinear(&self.beta_proj, x.clone())); // [b, t, hv] - // g = softplus(α + dt_bias) · a, with a = -exp(A_log) < 0. - let alpha = qlinear(&self.alpha_proj, x).add(self.dt_bias.val().reshape([1, 1, hv])); - let g = activation::softplus(alpha, 1.0).mul(self.a.val().reshape([1, 1, hv])); - - drop(_s_proj); - // Depthwise causal conv over the sequence, rolling the decode state - // exactly like nn::ShortConv (algebraic equivalence proven there). - let _s_conv = crate::prof::scope("delta.conv"); - let mix_cm = mixed.swap_dims(1, 2); // channel-major [b, conv_dim, t] - let conv_out = if t > 1 { - self.conv1d.forward(mix_cm.clone()).narrow(2, 0, t) - } else { - let window = match &cache.conv { - Some(prev) => Tensor::cat(vec![prev.clone(), mix_cm.clone()], 2), - None => { - let pad = Tensor::<3>::zeros([b, conv_dim, kk - 1], &device); - Tensor::cat(vec![pad, mix_cm.clone()], 2) - } - }; - let w = self.conv1d.weight.val().reshape([1, conv_dim, kk]); - window.mul(w).sum_dim(2) - }; - cache.conv = Some({ - let combined = match cache.conv.take() { - Some(prev) => Tensor::cat(vec![prev, mix_cm], 2), - None => mix_cm, - }; - let len = combined.dims()[2]; - if len >= kk - 1 { - combined.narrow(2, len - (kk - 1), kk - 1) - } else { - let pad = Tensor::<3>::zeros([b, conv_dim, (kk - 1) - len], &device); - Tensor::cat(vec![pad, combined], 2) - } - }); - let conv_out = activation::silu(conv_out.swap_dims(1, 2)); // [b, t, conv_dim] - drop(_s_conv); - let _s_split = crate::prof::scope("delta.split"); - - // Split into q/k/v and L2-normalize q/k per head: x / max(‖x‖, ε). - let eps = cfg.rms_norm_eps as f32; - let l2 = |x: Tensor<4>| -> Tensor<4> { - let norm = x.clone().powi_scalar(2).sum_dim(3).sqrt().clamp_min(eps); - x.div(norm) - }; - let q = l2(conv_out - .clone() - .narrow(2, 0, key_dim) - .reshape([b, t, hk, ds])); - let k = l2(conv_out - .clone() - .narrow(2, key_dim, key_dim) - .reshape([b, t, hk, ds])); - let v = conv_out - .narrow(2, 2 * key_dim, cfg.d_inner) - .reshape([b, t, hv, ds]); - - // Tile k-heads across the value heads (llama.cpp's ggml_repeat: - // head h_v reads k-head h_v % n_k_heads). - let tile = hv / hk; - let expand = |x: Tensor<4>| -> Tensor<4> { - if tile == 1 { - x - } else { - // [b, t, hk, ds] → [b, t, tile·hk, ds] tiling whole blocks. - x.repeat_dim(2, tile) - } - }; - let q = expand(q).swap_dims(1, 2); // [b, hv, t, ds] - let k = expand(k).swap_dims(1, 2); - let v = v.swap_dims(1, 2); - drop(_s_split); - let _s_recur = crate::prof::scope("delta.recur"); - - // The recurrence, one token at a time. State S[b, h, i, j]: - // i indexes the key dim, j the value dim. - let scale = 1.0 / (ds as f32).sqrt(); - let mut s = cache - .state - .take() - .unwrap_or_else(|| Tensor::<4>::zeros([b, hv, ds, ds], &device)); - let mut outs: Vec> = Vec::with_capacity(t); - for tau in 0..t { - let q_t = q.clone().narrow(2, tau, 1).mul_scalar(scale); // [b, hv, 1, ds] - let k_t = k.clone().narrow(2, tau, 1); // [b, hv, 1, ds] - let v_t = v.clone().narrow(2, tau, 1); // [b, hv, 1, ds] - let g_t = g.clone().narrow(1, tau, 1).reshape([b, hv, 1, 1]); // per-head decay logit - let b_t = beta.clone().narrow(1, tau, 1).reshape([b, hv, 1, 1]); - - s = s.mul(g_t.exp()); - // v̂[j] = Σ_i S[i, j]·k[i] — k over the key axis. - let v_hat = s.clone().mul(k_t.clone().swap_dims(2, 3)).sum_dim(2); // [b, hv, 1, ds] - let d = v_t.sub(v_hat).mul(b_t); // [b, hv, 1, ds] - // S += k ⊗ d (outer product over [key, value]). - s = s.add(k_t.swap_dims(2, 3).matmul(d.clone())); - // o[j] = Σ_i S[i, j]·q[i]. - let o = s.clone().mul(q_t.swap_dims(2, 3)).sum_dim(2); // [b, hv, 1, ds] - outs.push(o); - } - cache.state = Some(s); - let o = Tensor::cat(outs, 2); // [b, hv, t, ds] - drop(_s_recur); - let _s = crate::prof::scope("delta.out"); - - // Gated RMSNorm per value head, then flatten and project out. - let o = self.norm.forward(o.swap_dims(1, 2)); // [b, t, hv, ds] - let z = z.reshape([b, t, hv, ds]); - let gated = o.mul(activation::silu(z)).reshape([b, t, cfg.d_inner]); - qlinear(&self.out_proj, gated) - } -} - -/// One trunk layer: exactly one of `self_attn` / `linear_attn`. -#[derive(Module, Debug)] -pub struct Qwen35Layer { - pub input_norm: RmsNorm, - pub post_attn_norm: RmsNorm, - pub self_attn: Option, - pub linear_attn: Option, - pub mlp: SwiGluMlp, -} - -/// The qwen35 decoder stack. -#[derive(Module, Debug)] -pub struct Qwen35 { - pub embed_tokens: Embedding, - pub layers: Vec, - pub norm: RmsNorm, - /// Untied head when the checkpoint carries `output.weight`; tied to the - /// embedding otherwise. - pub lm_head: Option, -} - -/// Per-layer decode cache. -pub enum Qwen35Kv { - Attn(LayerKv), - Delta(DeltaState), -} - -/// A weight-loaded qwen35 plus its config. -pub struct LoadedQwen35 { - pub model: Qwen35, - pub config: Qwen35Config, - /// `None` for GGUF loads (self-contained; no sibling file). - pub tokenizer_config: Option, - /// P9 stage 3(c): remote FFN clusters of a partitioned pack — each - /// layer's `mlp` then holds only the *local* clusters and the pool adds - /// the rest (exact when every cluster runs). `None` = plain dense. - pub ffn_pool: Option>, - /// Opt-in skipping: clusters whose gate energy is below `tau` × the - /// row's total energy are not computed (lossy — only from a measured - /// skip table). `0.0` = exact. - pub ffn_skip_tau: f32, - /// P9 stage 4: the working-set schedule this model runs under, one entry - /// per trunk layer. `None` = every cluster is permanently resident (the - /// tier design), so there is nothing to stage. - pub ffn_plan: Option>, -} - -fn build(cfg: &Qwen35Config, device: &Device, untied_head: bool) -> Qwen35 { - let norm = |dim: usize, dev: &Device| { - RmsNormConfig::new(dim) - .with_epsilon(cfg.rms_norm_eps) - .init(dev) - }; - let linear = |inp: usize, out: usize, dev: &Device| { - LinearConfig::new(inp, out).with_bias(false).init(dev) - }; - let mlp_cfg = SwiGluMlpConfig { - hidden_size: cfg.hidden_size, - intermediate_size: cfg.intermediate_size, - }; - let layers = (0..cfg.num_layers) - .map(|i| { - let attn = cfg.is_attention(i); - Qwen35Layer { - input_norm: norm(cfg.hidden_size, device), - post_attn_norm: norm(cfg.hidden_size, device), - self_attn: attn.then(|| GatedAttention { - q_proj: linear( - cfg.hidden_size, - 2 * cfg.num_attention_heads * cfg.head_dim, - device, - ), - k_proj: linear( - cfg.hidden_size, - cfg.num_key_value_heads * cfg.head_dim, - device, - ), - v_proj: linear( - cfg.hidden_size, - cfg.num_key_value_heads * cfg.head_dim, - device, - ), - o_proj: linear( - cfg.num_attention_heads * cfg.head_dim, - cfg.hidden_size, - device, - ), - q_norm: norm(cfg.head_dim, device), - k_norm: norm(cfg.head_dim, device), - }), - linear_attn: (!attn).then(|| GatedDeltaNet { - qkv_proj: linear(cfg.hidden_size, cfg.conv_dim(), device), - z_proj: linear(cfg.hidden_size, cfg.d_inner, device), - beta_proj: linear(cfg.hidden_size, cfg.n_v_heads, device), - alpha_proj: linear(cfg.hidden_size, cfg.n_v_heads, device), - dt_bias: Param::from_tensor(Tensor::zeros([cfg.n_v_heads], device)), - a: Param::from_tensor(Tensor::zeros([cfg.n_v_heads], device)), - conv1d: Conv1dConfig::new(cfg.conv_dim(), cfg.conv_dim(), cfg.conv_kernel) - .with_groups(cfg.conv_dim()) - .with_padding(PaddingConfig1d::Explicit( - cfg.conv_kernel - 1, - cfg.conv_kernel - 1, - )) - .with_bias(false) - .init(device), - norm: norm(cfg.d_state, device), - out_proj: linear(cfg.d_inner, cfg.hidden_size, device), - }), - mlp: mlp_cfg.init(device), - } - }) - .collect(); - Qwen35 { - embed_tokens: EmbeddingConfig::new(cfg.vocab_size, cfg.hidden_size).init(device), - layers, - norm: norm(cfg.hidden_size, device), - lm_head: untied_head.then(|| linear(cfg.hidden_size, cfg.vocab_size, device)), - } -} - -/// GGUF (llama.cpp `qwen35` arch) names → this port's parameter paths. -/// `trunk_layers` gates the NextN/MTP block: any `blk.i` at or beyond it is -/// the draft head, deliberately skipped (unused by the main forward). -fn gguf_tensor_map(info: &GgufTensorInfo, trunk_layers: usize) -> Option { - match info.name.as_str() { - "token_embd.weight" => { - return Some(GgufMap::Rename("model.embed_tokens.weight".into())); - } - "output_norm.weight" => return Some(GgufMap::Rename("model.norm.weight".into())), - "output.weight" => return Some(GgufMap::Rename("lm_head.weight".into())), - _ => {} - } - let rest = info.name.strip_prefix("blk.")?; - let (layer, field) = rest.split_once('.')?; - let layer: usize = layer.parse().ok()?; - if layer >= trunk_layers { - return Some(GgufMap::Skip); // the NextN/MTP draft block - } - if field == "ssm_conv1d.weight" { - // ggml stores the depthwise kernel squeezed [k, channels]; the - // Conv1d module wants [channels, 1, k] — same bytes. - let (&k, &channels) = (info.dims.first()?, info.dims.get(1)?); - return Some(GgufMap::Reshape( - format!("model.layers.{layer}.linear_attn.conv1d.weight"), - vec![channels, 1, k], - )); - } - let mapped = qwen35_field(field)?; - Some(GgufMap::Rename(format!("model.layers.{layer}.{mapped}"))) -} - -/// GGUF per-layer field → this port's parameter path (minus the layer -/// prefix). Shared by the GGUF map and the pack loader. -fn qwen35_field(field: &str) -> Option<&'static str> { - Some(match field { - "attn_norm.weight" => "input_norm.weight", - "post_attention_norm.weight" => "post_attn_norm.weight", - // Full-attention layers. - "attn_q.weight" => "self_attn.q_proj.weight", - "attn_k.weight" => "self_attn.k_proj.weight", - "attn_v.weight" => "self_attn.v_proj.weight", - "attn_output.weight" => "self_attn.o_proj.weight", - "attn_q_norm.weight" => "self_attn.q_norm.weight", - "attn_k_norm.weight" => "self_attn.k_norm.weight", - // Gated DeltaNet layers. - "attn_qkv.weight" => "linear_attn.qkv_proj.weight", - "attn_gate.weight" => "linear_attn.z_proj.weight", - "ssm_beta.weight" => "linear_attn.beta_proj.weight", - "ssm_alpha.weight" => "linear_attn.alpha_proj.weight", - "ssm_dt.bias" => "linear_attn.dt_bias", - "ssm_a" => "linear_attn.a", - "ssm_norm.weight" => "linear_attn.norm.weight", - "ssm_out.weight" => "linear_attn.out_proj.weight", - // FFN. - "ffn_gate.weight" => "mlp.gate_proj.weight", - "ffn_up.weight" => "mlp.up_proj.weight", - "ffn_down.weight" => "mlp.down_proj.weight", - _ => return None, - }) -} - -/// Load a qwen35 model straight from a GGUF file — the classic f32 path, -/// which is [`load_from_gguf_quantized`] with quantization off. One import -/// path serves every precision (P9's "single path" rule). -pub fn load_from_gguf(path: &Path, device: &Device) -> Result { - load_from_gguf_quantized(path, device, QuantPolicy::Off) -} - -/// **Streaming** GGUF import with optional keep-quantized weights: one -/// tensor at a time is dequantized to f32 (whatever the source stored — -/// BF16, K-quants, IQ quants), moved to the device, **re-quantized** per -/// `policy` when eligible, and assigned. Peak memory is the finished model -/// plus a single f32 tensor — never the whole model at f32, which is what -/// makes the 27B tier loadable at all (its f32 form is ~109 GB). -pub fn load_from_gguf_quantized( - path: &Path, - device: &Device, - policy: QuantPolicy, -) -> Result { - let parse = |reason: String| ImportError::Parse { - file: path.to_path_buf(), - reason, - }; - let f = GgufFile::open(path).map_err(|e| parse(e.to_string()))?; - let config = Qwen35Config::from_gguf(&f).map_err(parse)?; - let untied = f.tensor("output.weight").is_some(); - let trunk = config.num_layers; - let mut model = build(&config, device, untied); - - let mut assigned = 0usize; - for info in &f.tensors { - let mapped = gguf_tensor_map(info, trunk) - .ok_or_else(|| parse(format!("unmapped tensor name '{}'", info.name)))?; - let (name, shape) = match mapped { - GgufMap::Skip => continue, - GgufMap::Rename(name) => ( - name, - info.dims - .iter() - .rev() - .map(|&d| d as usize) - .collect::>(), - ), - GgufMap::Reshape(name, shape) => (name, shape.iter().map(|&d| d as usize).collect()), - }; - let values = f - .read_tensor_f32(&info.name) - .map_err(|e| parse(e.to_string()))?; - assign_param( - &mut model, - &name, - ParamSrc::F32 { values, shape }, - policy, - device, - ) - .map_err(parse)?; - assigned += 1; - } - - // Both directions of completeness, loudly: every mapped tensor landed - // (assign_param errors otherwise) and the count matches what the - // architecture requires. - let expected = expected_tensor_count(&config, untied); - if assigned != expected { - return Err(parse(format!( - "GGUF supplied {assigned} trunk tensors, the architecture needs {expected}" - ))); - } - Ok(LoadedQwen35 { - model, - config, - tokenizer_config: None, - ffn_pool: None, - ffn_skip_tau: 0.0, - ffn_plan: None, - }) -} - -/// How many trunk tensors a checkpoint must supply (the completeness gate's -/// other half). Per attention layer 11 (2 norms + q/k/v/o + q/k norm + -/// 3 FFN), per DeltaNet layer 14 (2 norms + qkv/z + β/α/dt/a + conv + -/// ssm-norm + out + 3 FFN), plus embedding, final norm, and the untied head -/// when present. -fn expected_tensor_count(cfg: &Qwen35Config, untied: bool) -> usize { - let per_layer: usize = (0..cfg.num_layers) - .map(|i| if cfg.is_attention(i) { 11 } else { 14 }) - .sum(); - per_layer + 2 + usize::from(untied) -} - -/// Build a device tensor from row-major f32 `values` of `shape`, cast to the -/// backend float dtype. -fn device_tensor( - values: Vec, - shape: [usize; D], - device: &Device, -) -> Tensor { - let dtype = crate::backend::float_dtype(device); - Tensor::from_data(TensorData::new(values, shape), (device, dtype)) -} - -/// A 2-D **linear weight**: GGUF row-major is `[out, in]`, burn's `Linear` -/// wants `[in, out]` — transpose on device, then quantize when the policy -/// takes it. -fn linear_weight( - values: Vec, - shape: &[usize], - policy: QuantPolicy, - device: &Device, -) -> Result, String> { - let &[out, inp] = shape else { - return Err(format!("linear weight must be 2-D, got {shape:?}")); - }; - let w = device_tensor::<2>(values, [out, inp], device).swap_dims(0, 1); - Ok(if policy.eligible(&[inp, out]) { - crate::quant::quantize_weight(policy, w) - } else { - w - }) -} - -/// Where a parameter's data comes from: raw f32 (the GGUF streaming path — -/// transposed/quantized here) or a ready 2-D tensor (the pack path — already -/// `[in, out]` at its chosen precision). -pub enum ParamSrc { - F32 { values: Vec, shape: Vec }, - Ready2(Tensor<2>), -} - -/// A linear weight from either source. -fn take_linear( - ready: Option>, - values: Vec, - shape: &[usize], - policy: QuantPolicy, - device: &Device, -) -> Result, String> { - match ready { - Some(t) => Ok(t), - None => linear_weight(values, shape, policy, device), - } -} - -/// Route one mapped tensor into its module field. An unknown path is a loud -/// error — silence here would mean silently dropped weights. -fn assign_param( - model: &mut Qwen35, - name: &str, - src: ParamSrc, - policy: QuantPolicy, - device: &Device, -) -> Result<(), String> { - use burn::module::Param; - - let (values, shape_vec, ready): (Vec, Vec, Option>) = match src { - ParamSrc::F32 { values, shape } => (values, shape, None), - ParamSrc::Ready2(t) => (Vec::new(), t.dims().to_vec(), Some(t)), - }; - let shape = shape_vec.as_slice(); - - let expect_1d = |values: Vec, shape: &[usize]| -> Result, String> { - let &[n] = shape else { - return Err(format!("expected 1-D, got {shape:?}")); - }; - Ok(device_tensor::<1>(values, [n], device)) - }; - - match name { - "model.embed_tokens.weight" => { - // Embeddings stay float — token gather has no quantized kernel. - let t = match ready { - Some(t) => t, - None => { - let &[v, h] = shape else { - return Err(format!("embedding must be 2-D, got {shape:?}")); - }; - device_tensor::<2>(values, [v, h], device) - } - }; - model.embed_tokens.weight = Param::from_tensor(t); - return Ok(()); - } - "model.norm.weight" => { - model.norm.gamma = Param::from_tensor(expect_1d(values, shape)?); - return Ok(()); - } - "lm_head.weight" => { - let head = model - .lm_head - .as_mut() - .ok_or("checkpoint has output.weight but the model built a tied head")?; - head.weight = Param::from_tensor(take_linear(ready, values, shape, policy, device)?); - return Ok(()); - } - _ => {} - } - - let rest = name - .strip_prefix("model.layers.") - .ok_or_else(|| format!("unknown parameter path '{name}'"))?; - let (layer, field) = rest - .split_once('.') - .ok_or_else(|| format!("bad layer path '{name}'"))?; - let layer: usize = layer - .parse() - .map_err(|_| format!("bad layer in '{name}'"))?; - let l = model - .layers - .get_mut(layer) - .ok_or_else(|| format!("layer {layer} out of range"))?; - - let missing_attn = || format!("'{name}' targets an attention block on a DeltaNet layer"); - let missing_delta = || format!("'{name}' targets a DeltaNet block on an attention layer"); - - match field { - "input_norm.weight" => l.input_norm.gamma = Param::from_tensor(expect_1d(values, shape)?), - "post_attn_norm.weight" => { - l.post_attn_norm.gamma = Param::from_tensor(expect_1d(values, shape)?); - } - "mlp.gate_proj.weight" => { - l.mlp.gate_proj.weight = - Param::from_tensor(take_linear(ready, values, shape, policy, device)?); - } - "mlp.up_proj.weight" => { - l.mlp.up_proj.weight = - Param::from_tensor(take_linear(ready, values, shape, policy, device)?); - } - "mlp.down_proj.weight" => { - l.mlp.down_proj.weight = - Param::from_tensor(take_linear(ready, values, shape, policy, device)?); - } - _ => { - if let Some(attn_field) = field.strip_prefix("self_attn.") { - let attn = l.self_attn.as_mut().ok_or_else(missing_attn)?; - match attn_field { - "q_proj.weight" => { - attn.q_proj.weight = - Param::from_tensor(take_linear(ready, values, shape, policy, device)?); - } - "k_proj.weight" => { - attn.k_proj.weight = - Param::from_tensor(take_linear(ready, values, shape, policy, device)?); - } - "v_proj.weight" => { - attn.v_proj.weight = - Param::from_tensor(take_linear(ready, values, shape, policy, device)?); - } - "o_proj.weight" => { - attn.o_proj.weight = - Param::from_tensor(take_linear(ready, values, shape, policy, device)?); - } - "q_norm.weight" => { - attn.q_norm.gamma = Param::from_tensor(expect_1d(values, shape)?); - } - "k_norm.weight" => { - attn.k_norm.gamma = Param::from_tensor(expect_1d(values, shape)?); - } - other => return Err(format!("unknown attention field '{other}'")), - } - return Ok(()); - } - if let Some(delta_field) = field.strip_prefix("linear_attn.") { - let delta = l.linear_attn.as_mut().ok_or_else(missing_delta)?; - match delta_field { - "qkv_proj.weight" => { - delta.qkv_proj.weight = - Param::from_tensor(take_linear(ready, values, shape, policy, device)?); - } - "z_proj.weight" => { - delta.z_proj.weight = - Param::from_tensor(take_linear(ready, values, shape, policy, device)?); - } - "beta_proj.weight" => { - delta.beta_proj.weight = - Param::from_tensor(take_linear(ready, values, shape, policy, device)?); - } - "alpha_proj.weight" => { - delta.alpha_proj.weight = - Param::from_tensor(take_linear(ready, values, shape, policy, device)?); - } - "out_proj.weight" => { - delta.out_proj.weight = - Param::from_tensor(take_linear(ready, values, shape, policy, device)?); - } - "dt_bias" => delta.dt_bias = Param::from_tensor(expect_1d(values, shape)?), - "a" => delta.a = Param::from_tensor(expect_1d(values, shape)?), - "norm.weight" => { - delta.norm.gamma = Param::from_tensor(expect_1d(values, shape)?); - } - "conv1d.weight" => { - let &[ch, one, k] = shape else { - return Err(format!("conv kernel must be 3-D, got {shape:?}")); - }; - if one != 1 { - return Err(format!("conv kernel middle dim must be 1, got {one}")); - } - delta.conv1d.weight = - Param::from_tensor(device_tensor::<3>(values, [ch, 1, k], device)); - } - other => return Err(format!("unknown DeltaNet field '{other}'")), - } - return Ok(()); - } - return Err(format!("unknown layer field '{field}'")); - } - } - Ok(()) -} - -/// How each GGUF tensor enters a `.mummu` pack (see `crate::pack`). -pub fn pack_actions( - info: &GgufTensorInfo, - trunk_layers: usize, -) -> Option { - use crate::pack::ImportAction as A; - match info.name.as_str() { - "token_embd.weight" => return Some(A::Embedding), - "output_norm.weight" => return Some(A::Vector), - "output.weight" => return Some(A::Linear), - _ => {} - } - let rest = info.name.strip_prefix("blk.")?; - let (layer, field) = rest.split_once('.')?; - let layer: usize = layer.parse().ok()?; - if layer >= trunk_layers { - return Some(A::Skip); - } - Some(match field { - "ssm_conv1d.weight" => A::Conv, - "attn_norm.weight" - | "post_attention_norm.weight" - | "attn_q_norm.weight" - | "attn_k_norm.weight" - | "ssm_dt.bias" - | "ssm_a" - | "ssm_norm.weight" => A::Vector, - _ => { - qwen35_field(field)?; // known linear fields only - A::Linear - } - }) -} - -/// The FFN entry names of every trunk layer. -/// -/// Re-exported from [`crate::partition::ffn_names`], which is where it -/// belongs: every dense decoder in the zoo stores the same GGUF triple, so -/// this is not a qwen35 fact. Kept as a path so existing callers do not move. -pub use crate::partition::ffn_names; - -/// Verify-mode radial lookahead on? (`MUMMU_LOOKAHEAD=verify`). One env -/// read per process; anything but `verify` is off — there is no commit mode -/// until verify-mode acceptance earns it. -fn lookahead_verify() -> bool { - static ON: std::sync::OnceLock = std::sync::OnceLock::new(); - *ON.get_or_init(|| { - std::env::var("MUMMU_LOOKAHEAD").is_ok_and(|v| v.eq_ignore_ascii_case("verify")) - }) -} - -/// Residual-geometry probe on? (`MUMMU_RESIDUAL_PROBE=1`). -fn residual_probe() -> bool { - static ON: std::sync::OnceLock = std::sync::OnceLock::new(); - *ON.get_or_init(|| std::env::var("MUMMU_RESIDUAL_PROBE").is_ok()) -} - -/// A scratch copy of one layer's decode state for speculation: tensor -/// clones are refcounted handles, and burn ops never mutate buffers in -/// place, so the speculative forward can freely reassign the scratch -/// struct's fields while the real cache entry stays untouched. Cheap by -/// construction — this is the "never pollute S" rule made structural. -fn snapshot_kv(kv: &Qwen35Kv) -> Qwen35Kv { - match kv { - Qwen35Kv::Attn(k) => Qwen35Kv::Attn(k.clone()), - Qwen35Kv::Delta(d) => Qwen35Kv::Delta(DeltaState { - conv: d.conv.clone(), - state: d.state.clone(), - }), - } -} - -/// Pack tensor name → parameter path (the pack keeps GGUF names). -fn pack_param_path(name: &str, trunk_layers: usize) -> Option { - match name { - "token_embd.weight" => return Some("model.embed_tokens.weight".into()), - "output_norm.weight" => return Some("model.norm.weight".into()), - "output.weight" => return Some("lm_head.weight".into()), - _ => {} - } - let rest = name.strip_prefix("blk.")?; - let (layer, field) = rest.split_once('.')?; - let layer: usize = layer.parse().ok()?; - if layer >= trunk_layers { - return None; - } - if field == "ssm_conv1d.weight" { - return Some(format!("model.layers.{layer}.linear_attn.conv1d.weight")); - } - Some(format!("model.layers.{layer}.{}", qwen35_field(field)?)) -} - -/// Load from a `.mummu` pack, choosing each tensor's precision through -/// `choose` (the planner's tiering hook — per tensor, so a policy can mix -/// levels). Quantized levels arrive pre-packed (no re-quantization). -pub fn load_from_pack( - dir: &Path, - device: &Device, - choose: &dyn Fn(&crate::pack::TensorEntry) -> crate::pack::Precision, -) -> Result { - load_from_pack_inner(dir, device, choose, None, None) -} - -/// Load a **partitioned** pack with only `local(layer)` FFN clusters in each -/// layer's `mlp` (at the level `choose` picks for that entry); the caller -/// attaches the remote clusters as an `ExpertPool` via [`LoadedQwen35::with_ffn_pool`]. -/// Every layer must keep at least one local cluster. -pub fn load_from_pack_partitioned( - dir: &Path, - device: &Device, - choose: &dyn Fn(&crate::pack::TensorEntry) -> crate::pack::Precision, - local: &dyn Fn(usize) -> Vec, -) -> Result { - load_from_pack_inner(dir, device, choose, Some(local), None) -} - -/// [`load_from_pack_partitioned`], with the **embedding table pinned to -/// `embed_device`** instead of following the rest of the model. -/// -/// An embedding is a gather, not a matmul, so it is the one large tensor a -/// pack never quantizes — on the 27B it is 5.09 GB, a quarter of the model, -/// and it is read once per token while a layer's weights are read 65 times. -/// Holding it on the host frees that VRAM for weights that actually compute, -/// at the cost of one `[1, hidden]` transfer per token. -pub fn load_from_pack_partitioned_split( - dir: &Path, - device: &Device, - embed_device: &Device, - choose: &dyn Fn(&crate::pack::TensorEntry) -> crate::pack::Precision, - local: &dyn Fn(usize) -> Vec, -) -> Result { - load_from_pack_inner(dir, device, choose, Some(local), Some(embed_device)) -} - -/// Load a pack with **each layer on its own device** — the dense-model -/// placement (llama.cpp calls the knob `n_gpu_layers`). -/// -/// `layer_device(l)` says where layer `l` lives; every tensor of that layer, -/// trunk and FFN alike, is loaded there, so a layer never crosses a device -/// boundary mid-computation. Activations cross once, where the assignment -/// changes. -/// -/// This exists because the cluster-granular path is the wrong shape for a -/// dense model: it splits each layer's FFN across devices, and since every -/// cluster runs on every token there is no selectivity to pay for the -/// crossing — measured at 24.7 s/tok against 4.8 for keeping layers whole. -/// Cluster granularity stays right for a routed MoE, where only top-k -/// experts are touched. -pub fn load_from_pack_layered( - dir: &Path, - layer_device: &dyn Fn(usize) -> Device, - embed_device: &Device, - head_device: &Device, - choose: &dyn Fn(&crate::pack::TensorEntry) -> crate::pack::Precision, -) -> Result { - use crate::pack::{Pack, Role}; - let parse = |reason: String| ImportError::Parse { - file: dir.to_path_buf(), - reason, - }; - let pack = Pack::open(dir).map_err(parse)?; - let header = pack.header().map_err(parse)?; - let config = Qwen35Config::from_gguf(&header).map_err(parse)?; - let untied = pack.entry("output.weight").is_some(); - let trunk = config.num_layers; - - // Build the skeleton on the HOST, never on an accelerator. - // - // `build` materializes every parameter as f32. Building it on layer 0's - // device asked ONE device for the whole model at f32 — measured from the - // manifest, 100.20 GiB against a card with ~13 GiB usable, a 7.7x - // overshoot — so allocation failed roughly a thousand times and the card - // ended up holding nothing while the planner reported 44 of 64 layers - // placed. None of it was ever needed: the loop below replaces every - // parameter through `assign_param`, which takes the destination device - // per tensor, and the `assigned != expected` check proves none is - // missed. These are dead allocations — every one is overwritten before - // any reader sees it — so dropping them cannot change results. - // - // The per-layer `to_device` pre-move goes for the same reason: it moved - // f32 skeleton weights onto the card purely to overwrite them. - let build_host = crate::backend::cpu_device(); - let mut model = build(&config, &build_host, untied); - - let mut assigned = 0usize; - for entry in &pack.manifest.tensors { - let Some(path) = pack_param_path(&entry.name, trunk) else { - continue; - }; - // Which device this tensor belongs on: its layer's, or the special - // homes for the embedding and the head. - let device = match entry.role { - Role::Embedding => embed_device.clone(), - _ => layer_of_path(&path).map_or_else(|| head_device.clone(), &layer_device), - }; - let precision = { - let p = choose(entry); - if entry.precisions.contains_key(&p) { - p - } else { - *entry - .precisions - .keys() - .max() - .ok_or_else(|| parse(format!("'{}' has no stored precision", entry.name)))? - } - }; - let src = match entry.role { - Role::Linear | Role::Expert { .. } | Role::Embedding => { - ParamSrc::Ready2(pack.tensor::<2>(entry, precision, &device).map_err(parse)?) - } - Role::Vector | Role::Conv => ParamSrc::F32 { - values: pack.read_f32(entry).map_err(parse)?, - shape: entry.shape.clone(), - }, - }; - assign_param(&mut model, &path, src, QuantPolicy::Off, &device).map_err(parse)?; - assigned += 1; - } - // Every parameter above was assigned onto its own device. Pin the three - // specials, which the pack may not cover on every path; `to_device` is a - // no-op for a tensor already home. - model.embed_tokens = model.embed_tokens.clone().to_device(embed_device); - model.norm = model.norm.clone().to_device(head_device); - if let Some(h) = model.lm_head.take() { - model.lm_head = Some(h.to_device(head_device)); - } - let expected = expected_tensor_count(&config, untied); - if assigned != expected { - return Err(parse(format!( - "pack supplied {assigned} trunk tensors, the architecture needs {expected}" - ))); - } - Ok(LoadedQwen35 { - model, - config, - tokenizer_config: None, - ffn_pool: None, - ffn_skip_tau: 0.0, - ffn_plan: None, - }) -} - -/// The layer index a parameter path belongs to, if any -/// (`model.layers.7.mlp.gate_proj.weight` -> 7). -fn layer_of_path(path: &str) -> Option { - path.strip_prefix("model.layers.")? - .split('.') - .next()? - .parse() - .ok() -} - -/// One layer's FFN restricted to `clusters` of a partitioned pack, at -/// `precision`, in Linear layout — the local slab or a remote executor's -/// weights. Columns of gate/up and rows of down are sliced straight from -/// the stored bytes (no re-quantization). -pub fn load_ffn_clusters( - pack: &crate::pack::Pack, - layer: usize, - clusters: &[usize], - precision: crate::pack::Precision, - device: &Device, -) -> Result { - use burn::module::Param; - let part = pack - .manifest - .ffn_partition - .as_ref() - .ok_or("pack has no FFN partition")?; - let spans = part.layers.get(layer).ok_or("layer out of range")?; - let names = &part.names[layer]; - let ranges: Vec<(usize, usize)> = clusters - .iter() - .map(|&c| { - spans - .get(c) - .map(|s| (s.start, s.len)) - .ok_or_else(|| format!("cluster {c} out of range")) - }) - .collect::>()?; - if ranges.is_empty() { - return Err(format!("layer {layer}: empty cluster set")); - } - let entry = |name: &str| pack.entry(name).ok_or_else(|| format!("missing {name}")); - let pick = |e: &crate::pack::TensorEntry| { - if e.precisions.contains_key(&precision) { - precision - } else { - *e.precisions.keys().max().expect("stored level") - } - }; - let g = entry(&names[0])?; - let u = entry(&names[1])?; - let d = entry(&names[2])?; - Ok(crate::nn::ExpertWeights { - gate: Param::from_tensor(pack.tensor_cols(g, pick(g), &ranges, device)?), - up: Param::from_tensor(pack.tensor_cols(u, pick(u), &ranges, device)?), - down: Param::from_tensor(pack.tensor_rows(d, pick(d), &ranges, device)?), - }) -} - -impl LoadedQwen35 { - /// Attach the remote FFN clusters (one pool row per layer, ragged). - #[must_use] - pub fn with_ffn_pool(mut self, pool: std::sync::Arc) -> Self { - assert_eq!( - pool.num_layers(), - self.config.num_layers, - "FFN pool must have one row per layer" - ); - self.ffn_pool = Some(pool); - self - } - - /// Opt-in cluster skipping at energy threshold `tau` (see `ffn_skip_tau`). - #[must_use] - pub fn with_ffn_skip(mut self, tau: f32) -> Self { - self.ffn_skip_tau = tau.max(0.0); - self - } - - /// Run the FFN clusters as a **working set** under `plan`: each layer - /// stages what the next needs while it computes, and evicts behind - /// itself (P9 stage 4). Without a plan every cluster stays permanently - /// resident, which is the tier design. - #[must_use] - pub fn with_ffn_plan(mut self, plan: std::sync::Arc) -> Self { - self.ffn_plan = Some(plan); - self - } -} - -fn load_from_pack_inner( - dir: &Path, - device: &Device, - choose: &dyn Fn(&crate::pack::TensorEntry) -> crate::pack::Precision, - local: Option<&dyn Fn(usize) -> Vec>, - embed_device: Option<&Device>, -) -> Result { - use crate::pack::{Pack, Role}; - let parse = |reason: String| ImportError::Parse { - file: dir.to_path_buf(), - reason, - }; - let pack = Pack::open(dir).map_err(parse)?; - let header = pack.header().map_err(parse)?; - let config = Qwen35Config::from_gguf(&header).map_err(parse)?; - let untied = pack.entry("output.weight").is_some(); - let trunk = config.num_layers; - let mut model = build(&config, device, untied); - - // Partitioned FFN entries → (layer, proj index) for the local-cluster path. - let ffn_index: std::collections::HashMap<&str, (usize, usize)> = - match (&local, &pack.manifest.ffn_partition) { - (Some(_), Some(part)) => part - .names - .iter() - .enumerate() - .flat_map(|(l, n)| { - n.iter() - .enumerate() - .map(move |(i, name)| (name.as_str(), (l, i))) - }) - .collect(), - (Some(_), None) => { - return Err(parse( - "partitioned load requested but the pack has no FFN partition".into(), - )); - } - _ => std::collections::HashMap::new(), - }; - let mut assigned = 0usize; - for entry in &pack.manifest.tensors { - let Some(path) = pack_param_path(&entry.name, trunk) else { - continue; // NextN block members, if a pack kept any - }; - if let (Some(local), Some(&(layer, proj))) = (local, ffn_index.get(entry.name.as_str())) { - let clusters = local(layer); - let part = pack.manifest.ffn_partition.as_ref().expect("checked above"); - let spans = &part.layers[layer]; - let ranges: Vec<(usize, usize)> = clusters - .iter() - .map(|&c| { - spans - .get(c) - .map(|s| (s.start, s.len)) - .ok_or_else(|| parse(format!("layer {layer}: cluster {c} out of range"))) - }) - .collect::>()?; - if ranges.is_empty() { - return Err(parse(format!( - "layer {layer}: no local FFN cluster (every layer needs one)" - ))); - } - let precision = { - let p = choose(entry); - if entry.precisions.contains_key(&p) { - p - } else { - *entry.precisions.keys().max().expect("stored level") - } - }; - let t = if proj == 2 { - pack.tensor_rows(entry, precision, &ranges, device) - } else { - pack.tensor_cols(entry, precision, &ranges, device) - } - .map_err(parse)?; - assign_param( - &mut model, - &path, - ParamSrc::Ready2(t), - QuantPolicy::Off, - device, - ) - .map_err(parse)?; - assigned += 1; - continue; - } - let precision = { - let p = choose(entry); - if entry.precisions.contains_key(&p) { - p - } else { - // Fall back to the best float level the pack stored. - *entry - .precisions - .keys() - .max() - .ok_or_else(|| parse(format!("'{}' has no stored precision", entry.name)))? - } - }; - let src = match entry.role { - // The embedding may live somewhere else entirely — see - // `load_from_pack_partitioned_split`. - Role::Embedding => ParamSrc::Ready2( - pack.tensor::<2>(entry, precision, embed_device.unwrap_or(device)) - .map_err(parse)?, - ), - Role::Linear | Role::Expert { .. } => { - ParamSrc::Ready2(pack.tensor::<2>(entry, precision, device).map_err(parse)?) - } - Role::Vector | Role::Conv => ParamSrc::F32 { - values: pack.read_f32(entry).map_err(parse)?, - shape: entry.shape.clone(), - }, - }; - assign_param(&mut model, &path, src, QuantPolicy::Off, device).map_err(parse)?; - assigned += 1; - } - let expected = expected_tensor_count(&config, untied); - if assigned != expected { - return Err(parse(format!( - "pack supplied {assigned} trunk tensors, the architecture needs {expected}" - ))); - } - Ok(LoadedQwen35 { - model, - config, - tokenizer_config: None, - ffn_pool: None, - ffn_skip_tau: 0.0, - ffn_plan: None, - }) -} - -impl CausalLm for LoadedQwen35 { - type Cache = Vec; - - fn is_eos(&self, id: u32) -> bool { - self.config.eos_token_id.contains(id) - } - - fn new_cache(&self) -> Self::Cache { - (0..self.config.num_layers) - .map(|i| { - if self.config.is_attention(i) { - Qwen35Kv::Attn(None) - } else { - Qwen35Kv::Delta(DeltaState { - conv: None, - state: None, - }) - } - }) - .collect() - } - - fn forward( - &self, - new_ids: &[u32], - past: usize, - cache: &mut Self::Cache, - device: &Device, - ) -> Tensor<2> { - let t = new_ids.len(); - assert!(t >= 1, "qwen35 forward: need at least one token"); - assert!( - cache.len() == self.config.num_layers, - "qwen35 forward: cache has {} layers, model has {}", - cache.len(), - self.config.num_layers - ); - let cfg = &self.config; - - // The embedding may live on a different device from the rest of the - // model (it is a gather, so it is often left on the host to keep VRAM - // for weights that compute — see `load_from_pack_partitioned_split`). - // A gather needs its indices on the SAME device as the table, so the - // indices are built there and only the small `[1, t, hidden]` result - // crosses over. - let embed_device = self.model.embed_tokens.weight.val().device(); - let ids32: Vec = new_ids.iter().map(|&i| i as i32).collect(); - let input = Tensor::<1, Int>::from_data( - TensorData::new(ids32, [t]), - (&embed_device, crate::backend::int_dtype(&embed_device)), - ) - .reshape([1, t]); - let _prof_forward = crate::prof::scope("forward"); - let mut x = { - let _s = crate::prof::scope("embed"); - self.model.embed_tokens.forward(input).to_device(device) - }; - - let (cos, sin) = rope_tables(t, past, cfg.rope_dim, cfg.rope_theta, device); - let mask = (t > 1).then(|| causal_mask(t, past, device)); - - // Stage attribution lives in `crate::prof`: the scope guards below feed - // a flame graph (serve: POST /api/chat with {"profile": true}, then - // GET /api/profile). This forward is synchronous on the flex trunk, so - // guards never cross an await and wall time attributes cleanly; - // device-queued work lands wherever the next readback syncs, so read - // GPU bars as "where the sync happened", not as kernel time. - let n_layers = self.model.layers.len(); - // Radial-lookahead carry: layer li+1's speculative post-attn h2, - // produced during layer li's drain window on SCRATCH state, and - // verified against the exact h2 once li+1 computes it. Verify mode - // never uses the speculative value — it only measures how good it - // would have been, which is the data that decides whether a commit - // mode is ever legal. - let mut spec_carry: Option<(usize, Tensor<3>)> = None; - let (mut la_n, mut la_max) = (0u32, 0f32); - let (mut la_a1, mut la_a2, mut la_a3) = (0u32, 0u32, 0u32); - for li in 0..n_layers { - let layer = &self.model.layers[li]; - // Layers may live on different devices (the dense placement puts - // as many whole layers on the GPU as VRAM holds, the rest on the - // host). Moving `x` here is a no-op while the device does not - // change, so a same-device model pays nothing, and a split model - // crosses ONCE — where the assignment changes — instead of twice - // per layer. - let layer_device = layer.input_norm.gamma.val().device(); - if x.device() != layer_device { - x = x.to_device(&layer_device); - } - let _prof_layer = crate::prof::scope("layer"); - let h = { - let _s = crate::prof::scope("norm1"); - layer.input_norm.forward(x.clone()) - }; - // The rope tables and mask were built once on the entry device; - // an attention layer elsewhere needs them there too. - let _s_glue_rope = crate::prof::scope("glue.rope"); - let (cos_l, sin_l) = if cos.device() == layer_device { - (cos.clone(), sin.clone()) - } else { - ( - cos.clone().to_device(&layer_device), - sin.clone().to_device(&layer_device), - ) - }; - let mask_l = mask.as_ref().map(|m| { - if m.device() == layer_device { - m.clone() - } else { - m.clone().to_device(&layer_device) - } - }); - drop(_s_glue_rope); - let kv = &mut cache[li]; - let h = match (&layer.self_attn, &layer.linear_attn, kv) { - (Some(attn), None, Qwen35Kv::Attn(kv_state)) => { - let _s = crate::prof::scope("attn.full"); - attn.forward(h, cfg, &cos_l, &sin_l, mask_l.as_ref(), kv_state) - } - (None, Some(delta), Qwen35Kv::Delta(state)) => { - let _s = crate::prof::scope("attn.delta"); - delta.forward(h, cfg, state) - } - _ => unreachable!("qwen35 forward: layer/cache kind mismatch"), - }; - { - let _s = crate::prof::scope("glue.resid1"); - x = x.add(h); - } - let h2 = { - let _s = crate::prof::scope("norm2"); - layer.post_attn_norm.forward(x.clone()) - }; - // The exact h2 exists: score the speculative one from the - // previous layer's window. Pure measurement — the exact value - // is what flows onward, unconditionally. - if let Some((idx, spec_h2)) = spec_carry.take() - && idx == li - { - let _s = crate::prof::scope("spec.verify"); - let read = |t: Tensor<3>| -> f32 { - t.abs() - .max() - .into_data() - .convert::() - .to_vec::() - .map(|v| v[0]) - .unwrap_or(f32::NAN) - }; - let diff = read(spec_h2.sub(h2.clone())); - let scale = read(h2.clone()).max(1e-6); - let rel = diff / scale; - la_n += 1; - la_max = la_max.max(rel); - if rel < 1e-3 { - la_a1 += 1; - } - if rel < 1e-2 { - la_a2 += 1; - } - if rel < 5e-2 { - la_a3 += 1; - } - } - // ENQUEUE-FIRST: hand the remote FFN to the accelerators before - // the local slab runs, so they work while the host does. The - // profile that motivated this: the local mlp (0.67 s/token) ran - // serially in FRONT of a 1.59 s/token device wait, card idle. - // Exact-mode only — the skip path needs host energies up front - // and keeps the sequential call below. - let pending = match &self.ffn_pool { - Some(pool) if self.ffn_skip_tau <= 0.0 => { - if let Some(plan) = &self.ffn_plan - && let Some(sched) = plan.layers.get(li) - { - let _s = crate::prof::scope("glue.sched"); - pool.apply_schedule(li, sched, device); - } - let [b, tt, hd] = h2.dims(); - let _s = crate::prof::scope("ffn.enqueue"); - if crate::nn::trace_layer() == Some(li) && tt == 1 { - eprintln!("[tl] enqueue {}", crate::nn::trace_us()); - } - pool.run_dense_pending(li, h2.clone().reshape([b * tt, hd])) - } - _ => None, - }; - // SwiGLU spelled through qlinear (the mlp's own forward would - // reshape a packed quantized weight — see qlinear). - // Three separate scopes: gate/up multiply [1,h]x[h,inter] while - // down multiplies [1,inter]x[inter,h] — if one shape hits a slow - // kernel path, the graph should say which. - let gate = { - let _s = crate::prof::scope("mlp.gate"); - activation::silu(qlinear(&layer.mlp.gate_proj, h2.clone())) - }; - let up = { - let _s = crate::prof::scope("mlp.up"); - qlinear(&layer.mlp.up_proj, h2.clone()) - }; - let mut ffn = { - let _s = crate::prof::scope("mlp.down"); - qlinear(&layer.mlp.down_proj, gate.clone().mul(up)) - }; - // RADIAL LOOKAHEAD (MUMMU_LOOKAHEAD=verify): the dGPU is still - // draining this layer's remote FFN on its worker; the main - // thread's wait is the window. Run layer li+1's trunk on the - // known prefix a = x + local_ffn now, on SCRATCH state (tensor - // clones are refcounts; burn ops never mutate in place, and the - // real cache entry is untouched). The radial identity makes the - // parallel component of the late remote piece exact under a - // scalar; what this measures is how much the rest matters. - if lookahead_verify() && pending.is_some() && li + 1 < n_layers { - let [_, tt, _] = ffn.dims(); - if tt == 1 { - let _s = crate::prof::scope("spec.trunk"); - let a3 = x.clone().add(ffn.clone()); - let nxt = &self.model.layers[li + 1]; - let mut scratch = snapshot_kv(&cache[li + 1]); - let sh = nxt.input_norm.forward(a3.clone()); - let sh = match (&nxt.self_attn, &nxt.linear_attn, &mut scratch) { - (Some(attn), None, Qwen35Kv::Attn(kv_state)) => { - attn.forward(sh, cfg, &cos, &sin, None, kv_state) - } - (None, Some(delta), Qwen35Kv::Delta(state)) => { - delta.forward(sh, cfg, state) - } - _ => unreachable!("qwen35 lookahead: layer/cache kind mismatch"), - }; - let spec_x = a3.add(sh); - spec_carry = Some((li + 1, nxt.post_attn_norm.forward(spec_x))); - } - } - if let Some(pending) = pending { - let _s = crate::prof::scope("glue.merge"); - let tl = crate::nn::trace_layer() == Some(li) && ffn.dims()[1] == 1; - if tl { - eprintln!("[tl] join-start {}", crate::nn::trace_us()); - } - let resolved = { - let _s = crate::prof::scope("merge.resolve_call"); - pending.resolve() - }; - if tl { - eprintln!("[tl] join-done {}", crate::nn::trace_us()); - } - if let Some(remote) = resolved { - let [b, tt, hd] = ffn.dims(); - let remote = { - let _s = crate::prof::scope("merge.reshape"); - remote.reshape([b, tt, hd]) - }; - // Residual-geometry probe (MUMMU_RESIDUAL_PROBE=1): the - // radial split N(a+b) = alpha*N(a) + rstd(a+b)*(g.*b_perp) - // is exact, so lookahead's commit-mode viability is the - // size of b_perp against a — measured, not argued. - if residual_probe() && tt == 1 { - let a = x.clone().add(ffn.clone()); - let read = |t: Tensor<3>| -> f32 { - t.sum() - .into_data() - .convert::() - .to_vec::() - .map(|v| v[0]) - .unwrap_or(f32::NAN) - }; - let dot = read(a.clone().mul(remote.clone())); - let na2 = read(a.clone().mul(a)); - let nb2 = read(remote.clone().mul(remote.clone())); - if na2 > 0.0 { - let sigma = 1.0 + dot / na2; - let bperp2 = (nb2 - dot * dot / na2).max(0.0); - let alpha = sigma * (na2 / (na2 + nb2 + 2.0 * dot).max(1e-12)).sqrt(); - eprintln!( - "[residual-probe] layer={li} b_over_a={:.4} bperp_over_a={:.4} alpha_minus_1={:+.5}", - (nb2 / na2).sqrt(), - (bperp2 / na2).sqrt(), - alpha - 1.0, - ); - } - } - // Which operand carries the ~28 ms that lands on the - // FIRST op after the worker cycle? Three scoped probes: - // an op touching only the local ffn (ambient/first-op - // effects), an op touching only the remote partial (its - // first-use materialization), then the real add. The 28ms - // lands in exactly one of these and names its class. - { - let _s = crate::prof::scope("merge.warm_local"); - let _ = ffn.clone().add_scalar(0.0); - } - { - let _s = crate::prof::scope("merge.warm_remote"); - let _ = remote.clone().add_scalar(0.0); - } - { - let _s = crate::prof::scope("merge.add"); - ffn = ffn.add(remote); - } - } - } else if let Some(pool) = &self.ffn_pool { - // Working set: issue THIS layer's staging decisions before - // its FFN runs, so the transfers for the next layer overlap - // this layer's compute instead of stalling in front of it. - // Nothing here blocks — a cluster that has not landed by the - // time its layer runs simply computes on the host. - if let Some(plan) = &self.ffn_plan - && let Some(sched) = plan.layers.get(li) - { - let _s = crate::prof::scope("glue.sched"); - pool.apply_schedule(li, sched, device); - } - // Remote clusters of a partitioned FFN (exact sum; skip only - // when a measured tau was chosen). - let [b, tt, hd] = ffn.dims(); - let local_energy: Vec = if self.ffn_skip_tau > 0.0 { - gate.powf_scalar(2.0) - .sum_dim(2) - .into_data() - .convert::() - .to_vec::() - .expect("local gate energy") - } else { - Vec::new() - }; - let skip = (self.ffn_skip_tau > 0.0) - .then_some((self.ffn_skip_tau, local_energy.as_slice())); - let remote = { - let _s = crate::prof::scope("ffn.remote"); - pool.run_dense(li, h2.reshape([b * tt, hd]), skip) - }; - if let Some(remote) = remote { - let _s = crate::prof::scope("glue.merge"); - ffn = ffn.add(remote.reshape([b, tt, hd])); - } - } - { - let _s = crate::prof::scope("glue.resid2"); - x = x.add(ffn); - } - } - if la_n > 0 { - let pct = |k: u32| f64::from(k) * 100.0 / f64::from(la_n); - eprintln!( - "[lookahead] verified {la_n} layers: accept@1e-3 {:.0}% | @1e-2 {:.0}% | @5e-2 {:.0}% | worst rel {la_max:.4}", - pct(la_a1), - pct(la_a2), - pct(la_a3), - ); - } - // The final norm and head may live elsewhere than the last layer. - let head_device = self.model.norm.gamma.val().device(); - let x = if x.device() == head_device { - x - } else { - x.to_device(&head_device) - }; - let x = { - let _s = crate::prof::scope("final_norm"); - self.model.norm.forward(x) - }; - - // Last position only → logits [1, vocab]. - let last = x.narrow(1, t - 1, 1).reshape([1, cfg.hidden_size]); - // Suspect number one for unattributed time: the tied head is a - // [1, 5120] x [5120, 248320] f32 matmul, and it runs on whichever - // device holds the embedding table — the HOST, for a split model. - let _s = crate::prof::scope("lm_head"); - match &self.model.lm_head { - Some(head) => qlinear2(head, last), - None => { - // Tied head: logits = h · Eᵀ. The embedding may be on another - // device (host-resident gather table), and unlike the gather - // this IS a matmul — so run it where the big tensor lives and - // move only the `[1, vocab]` result, rather than dragging a - // multi-GB table across the bus every token. - let e = self.model.embed_tokens.weight.val(); // [vocab, hidden] - let out_device = last.device(); - let e_device = e.device(); - last.to_device(&e_device) - .matmul(e.swap_dims(0, 1)) - .to_device(&out_device) - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - /// A toy config exercising both layer kinds: layer 1 is full attention - /// (`(1+1) % 2 == 0`), layers 0 and 2 are DeltaNet. - fn toy_config() -> Qwen35Config { - Qwen35Config { - vocab_size: 64, - hidden_size: 16, - num_layers: 3, - num_attention_heads: 2, - num_key_value_heads: 1, - head_dim: 8, - intermediate_size: 24, - rms_norm_eps: 1e-6, - rope_theta: 1e4, - rope_dim: 4, - full_attention_interval: 2, - conv_kernel: 3, - d_inner: 12, // 3 v-heads × d_state 4 - d_state: 4, - n_k_heads: 1, - n_v_heads: 3, - eos_token_id: EosIds::One(0), - } - } - - fn toy_model() -> LoadedQwen35 { - let cfg = toy_config(); - cfg.validate().expect("toy config validates"); - let device = crate::backend::cpu_device(); - LoadedQwen35 { - model: build(&cfg, &device, false), - config: cfg, - tokenizer_config: None, - ffn_pool: None, - ffn_skip_tau: 0.0, - ffn_plan: None, - } - } - - /// The load-bearing invariant for BOTH caches (attention KV and the - /// DeltaNet conv window + recurrent state): prefill + one-token decode - /// steps must produce exactly the logits of a single full prefill. - #[test] - fn cached_decode_matches_full_prefill() { - let m = toy_model(); - let device = crate::backend::cpu_device(); - let ids: Vec = vec![3, 17, 42, 9, 60, 11]; - - let mut full_cache = m.new_cache(); - let full = m - .forward(&ids, 0, &mut full_cache, &device) - .into_data() - .to_vec::() - .unwrap(); - - let mut cache = m.new_cache(); - let _ = m.forward(&ids[..3], 0, &mut cache, &device); - let mut last = Vec::new(); - for (i, &id) in ids.iter().enumerate().skip(3) { - last = m - .forward(&[id], i, &mut cache, &device) - .into_data() - .to_vec::() - .unwrap(); - } - assert_eq!(full.len(), last.len()); - for (i, (f, s)) in full.iter().zip(&last).enumerate() { - assert!((f - s).abs() < 1e-4, "logit {i}: full {f} vs stepped {s}"); - } - } - - /// DeltaNet state actually carries information: the same final token - /// after different prefixes must produce different logits. - #[test] - fn recurrent_state_carries_the_prefix() { - let m = toy_model(); - let device = crate::backend::cpu_device(); - let mut c1 = m.new_cache(); - let mut c2 = m.new_cache(); - let _ = m.forward(&[1, 2, 3], 0, &mut c1, &device); - let _ = m.forward(&[9, 8, 7], 0, &mut c2, &device); - let a = m - .forward(&[5], 3, &mut c1, &device) - .into_data() - .to_vec::() - .unwrap(); - let b = m - .forward(&[5], 3, &mut c2, &device) - .into_data() - .to_vec::() - .unwrap(); - let max_diff = a - .iter() - .zip(&b) - .map(|(x, y)| (x - y).abs()) - .fold(0.0f32, f32::max); - assert!(max_diff > 1e-6, "different prefixes must change the logits"); - } - - #[test] - fn config_layer_kinds_follow_the_interval() { - let cfg = toy_config(); - assert!(!cfg.is_attention(0)); - assert!(cfg.is_attention(1)); - assert!(!cfg.is_attention(2)); - // The 27B pattern: every 4th layer. - let mut real = cfg; - real.full_attention_interval = 4; - let attn: Vec = (0..8).filter(|&i| real.is_attention(i)).collect(); - assert_eq!(attn, vec![3, 7]); - } -} diff --git a/crates/mummu/examples/src/nn/attention.rs b/crates/mummu/examples/src/nn/attention.rs deleted file mode 100644 index 8a3f16f..0000000 --- a/crates/mummu/examples/src/nn/attention.rs +++ /dev/null @@ -1,458 +0,0 @@ -//! Cache-aware grouped-query attention (GQA), shared by every decoder in the -//! zoo. Covers both proven variants: plain GQA with projection bias (Qwen2) -//! and per-head q/k RMSNorm without bias (LFM2). - -use burn::module::Module; -use burn::nn::{Linear, LinearConfig, RmsNorm, RmsNormConfig}; -use burn::tensor::{DType, Device, Tensor, TensorData, activation}; - -use super::MAX_CONTEXT_TOKENS; -use super::rope::apply_rope; - -/// Per-layer KV cache entry: cached keys and values, each `[b, nkv, seq, hd]`. -/// `None` until the layer's first forward (the prompt prefill seeds it). -pub type LayerKv = Option<(Tensor<4>, Tensor<4>)>; - -/// Additive causal mask `[1, 1, t, past+t]`: query row `i` (absolute position -/// `past+i`) may attend to key columns `0..=past+i`; future columns get a -/// large negative. `-1e4` (not `-inf`) so the f16 GPU path (max ~65504) never -/// overflows — `exp(-1e4)` is still exactly 0 after softmax. -pub fn causal_mask(t: usize, past: usize, device: &Device) -> Tensor<4> { - assert!(t >= 1, "causal_mask: need at least one query row, got t=0"); - assert!( - past + t <= MAX_CONTEXT_TOKENS, - "causal_mask: position {past}+{t} exceeds MAX_CONTEXT_TOKENS ({MAX_CONTEXT_TOKENS})" - ); - let kcols = past + t; - let mut m = vec![0f32; t * kcols]; - for i in 0..t { - let qpos = past + i; - for j in (qpos + 1)..kcols { - m[i * kcols + j] = -1e4; - } - } - // The float dtype comes from the DEVICE — burn 0.22 keeps the element - // type there as a runtime setting, not on a backend type. Creation sites - // still name it explicitly rather than riding the unspecified default. - let dtype = crate::backend::float_dtype(device); - Tensor::<2>::from_data(TensorData::new(m, [t, kcols]), (device, dtype)) - .reshape([1, 1, t, kcols]) -} - -/// GQA expand: `[b, nkv, s, hd]` → `[b, nkv*group, s, hd]`, each KV head -/// repeated `group` times contiguously (HF `repeat_kv`). -pub fn repeat_kv(x: Tensor<4>, group: usize) -> Tensor<4> { - assert!(group >= 1, "repeat_kv: group must be >= 1"); - if group == 1 { - return x; - } - let [b, nkv, s, hd] = x.dims(); - debug_assert!(nkv >= 1 && hd >= 1, "repeat_kv: degenerate kv shape"); - x.reshape([b, nkv, 1, s, hd]) - .repeat_dim(2, group) - .reshape([b, nkv * group, s, hd]) -} - -/// Grouped-query attention with a per-layer KV cache. Field names mirror the -/// HF checkpoint layout (`q_proj`/`k_proj`/`v_proj`/`o_proj`); architectures -/// whose checkpoints differ (LFM2's `out_proj`, `q_layernorm`) remap keys at -/// load time instead of renaming fields. -#[derive(Module, Debug)] -pub struct GqaAttention { - pub q_proj: Linear, - pub k_proj: Linear, - pub v_proj: Linear, - pub o_proj: Linear, - /// Per-head RMSNorm on q, applied post-projection at `[b, t, nh, hd]` - /// (LFM2-style). `None` for architectures without it (Qwen2). - pub q_norm: Option, - /// Per-head RMSNorm on k — present iff `q_norm` is. - pub k_norm: Option, -} - -/// Shape/behavior config for [`GqaAttention`]. -#[derive(Debug, Clone)] -pub struct GqaAttentionConfig { - pub hidden_size: usize, - pub num_heads: usize, - pub num_kv_heads: usize, - pub head_dim: usize, - /// Projection bias on q/k/v (Qwen2: true; LFM2: false). `o_proj` never - /// has bias in either. - pub bias: bool, - /// q/k RMSNorm epsilon (LFM2/Qwen3/OLMoE: eps of the model; Qwen2: `None`). - pub qk_norm_eps: Option, - /// Where the q/k norm applies: `false` = per-head over `head_dim` - /// (LFM2/Qwen3), `true` = over the **whole projection** before the head - /// split (OLMoE — its `q_norm`/`k_norm` span `num_heads * head_dim`). - pub qk_norm_projection: bool, -} - -impl GqaAttentionConfig { - /// Initialize the module (random weights; real weights come from import). - pub fn init(&self, device: &Device) -> GqaAttention { - assert!( - self.num_kv_heads >= 1 && self.num_heads.is_multiple_of(self.num_kv_heads), - "GQA: num_heads ({}) must be a positive multiple of num_kv_heads ({})", - self.num_heads, - self.num_kv_heads - ); - assert!( - self.head_dim >= 2 && self.head_dim.is_multiple_of(2), - "GQA: head_dim must be even and >= 2 for RoPE, got {}", - self.head_dim - ); - let q_dim = self.num_heads * self.head_dim; - let kv_dim = self.num_kv_heads * self.head_dim; - let norm = |eps: f64, dim: usize| RmsNormConfig::new(dim).with_epsilon(eps).init(device); - let q_norm_dim = if self.qk_norm_projection { - q_dim - } else { - self.head_dim - }; - let k_norm_dim = if self.qk_norm_projection { - kv_dim - } else { - self.head_dim - }; - GqaAttention { - q_proj: LinearConfig::new(self.hidden_size, q_dim) - .with_bias(self.bias) - .init(device), - k_proj: LinearConfig::new(self.hidden_size, kv_dim) - .with_bias(self.bias) - .init(device), - v_proj: LinearConfig::new(self.hidden_size, kv_dim) - .with_bias(self.bias) - .init(device), - o_proj: LinearConfig::new(q_dim, self.hidden_size) - .with_bias(false) - .init(device), - q_norm: self.qk_norm_eps.map(|eps| norm(eps, q_norm_dim)), - k_norm: self.qk_norm_eps.map(|eps| norm(eps, k_norm_dim)), - } - } -} - -/// Apply a q/k RMSNorm at the placement its gamma width implies: `head_dim` → -/// per-head at `[b, t, n, hd]` (LFM2/Qwen3), `n * head_dim` → over the whole -/// projection **before** the head split (OLMoE). The two coincide at `n == 1`. -/// Inferring from the loaded gamma keeps the module shape identical across -/// families — a checkpoint's own norm width picks its semantics. -fn qk_norm_forward( - norm: &RmsNorm, - x: Tensor<3>, // [b, t, n*hd] - n: usize, - hd: usize, -) -> Tensor<4> { - let [b, t, width] = x.dims(); - debug_assert!(width == n * hd, "q/k projection width must be n * head_dim"); - let gamma = norm.gamma.dims()[0]; - assert!( - gamma == hd || gamma == n * hd, - "q/k norm width {gamma} matches neither head_dim ({hd}) nor the projection width ({})", - n * hd - ); - if gamma == n * hd && n > 1 { - norm.forward(x).reshape([b, t, n, hd]) - } else { - norm.forward(x.reshape([b, t, n, hd])) - } -} - -impl GqaAttention { - /// Cache-aware forward: RoPE the new q/k at the offset positions, append - /// the new k/v to this layer's cache, attend over the full cached range. - /// - /// `x` is `[b, t, hidden]` — the prompt at prefill (`kv == None`), a - /// single new token per decode step after. Returns `[b, t, hidden]`. - #[allow(clippy::too_many_arguments)] // mirrors the proven reference signature - pub fn forward( - &self, - x: Tensor<3>, - num_heads: usize, - num_kv_heads: usize, - head_dim: usize, - cos: &Tensor<4>, - sin: &Tensor<4>, - mask: Option<&Tensor<4>>, - kv: &mut LayerKv, - ) -> Tensor<3> { - let [b, t, _h] = x.dims(); - let (nh, nkv, hd) = (num_heads, num_kv_heads, head_dim); - assert!( - nkv >= 1 && nh.is_multiple_of(nkv), - "GQA forward: num_heads ({nh}) must be a positive multiple of num_kv_heads ({nkv})" - ); - debug_assert!( - self.q_norm.is_some() == self.k_norm.is_some(), - "GQA forward: q_norm and k_norm must be both present or both absent" - ); - - // q/k RMSNorm (when present) applies post-projection, before - // transpose + RoPE — the LFM2 ordering, validated against Ollama. - // Placement (per-head vs whole-projection) follows the loaded norm's - // own width; see `qk_norm_forward`. - let q = self.q_proj.forward(x.clone()); - let q = match &self.q_norm { - Some(norm) => qk_norm_forward(norm, q, nh, hd), - None => q.reshape([b, t, nh, hd]), - } - .swap_dims(1, 2); - let k_new = self.k_proj.forward(x.clone()); - let k_new = match &self.k_norm { - Some(norm) => qk_norm_forward(norm, k_new, nkv, hd), - None => k_new.reshape([b, t, nkv, hd]), - } - .swap_dims(1, 2); - let v_new = self - .v_proj - .forward(x) - .reshape([b, t, nkv, hd]) - .swap_dims(1, 2); - - let q = apply_rope(q, cos, sin); - let k_new = apply_rope(k_new, cos, sin); - - // Append to (or seed) the cache, then attend over everything so far. - let (k_all, v_all) = match kv.take() { - Some((pk, pv)) => ( - Tensor::cat(vec![pk, k_new], 2), - Tensor::cat(vec![pv, v_new], 2), - ), - None => (k_new, v_new), - }; - *kv = Some((k_all.clone(), v_all.clone())); - - let group = nh / nkv; - let k = repeat_kv(k_all, group); - let v = repeat_kv(v_all, group); - - // f32 island: Qwen-class attention logits overflow f16 (max 65504) - // in the q·kᵀ scores, collapsing softmax to NaN — llama.cpp pins this - // same matmul to f32 precision for the same reason. Scores + mask + - // softmax run in f32, the probabilities (all in [0, 1]) return to the - // ambient dtype for the value matmul. Every cast is a no-op on f32 - // backends. - let ambient = q.dtype(); - let scale = 1.0 / (hd as f32).sqrt(); - let mut scores = q - .cast(DType::F32) - .matmul(k.cast(DType::F32).swap_dims(2, 3)) - .mul_scalar(scale); - if let Some(m) = mask { - scores = scores.add(m.clone().cast(DType::F32)); - } - let probs = activation::softmax(scores, 3).cast(ambient); - let ctx = probs.matmul(v).swap_dims(1, 2).reshape([b, t, nh * hd]); - self.o_proj.forward(ctx) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::nn::rope::rope_tables; - - type Dev = burn::tensor::Device; - - const HIDDEN: usize = 16; - const HEADS: usize = 4; - const KV_HEADS: usize = 2; - const HEAD_DIM: usize = 4; - const THETA: f32 = 1e4; - - fn attn(qk_norm: bool, projection: bool, device: &Dev) -> GqaAttention { - GqaAttentionConfig { - hidden_size: HIDDEN, - num_heads: HEADS, - num_kv_heads: KV_HEADS, - head_dim: HEAD_DIM, - bias: true, - qk_norm_eps: qk_norm.then_some(1e-5), - qk_norm_projection: projection, - } - .init(device) - } - - /// Deterministic pseudo-random input `[1, t, HIDDEN]`. - fn input(t: usize, seed: f32, device: &Dev) -> Tensor<3> { - let data: Vec = (0..t * HIDDEN) - .map(|i| ((i as f32 + seed) * 0.7).sin()) - .collect(); - Tensor::<2>::from_data(TensorData::new(data, [t, HIDDEN]), device).reshape([1, t, HIDDEN]) - } - - /// Full forward over `t` positions in one call (prefill-style). - fn full_forward(a: &GqaAttention, x: Tensor<3>, device: &Dev) -> Vec { - let t = x.dims()[1]; - let (cos, sin) = rope_tables(t, 0, HEAD_DIM, THETA, device); - let mask = causal_mask(t, 0, device); - let mut kv: LayerKv = None; - a.forward( - x, - HEADS, - KV_HEADS, - HEAD_DIM, - &cos, - &sin, - Some(&mask), - &mut kv, - ) - .into_data() - .to_vec::() - .unwrap() - } - - #[test] - fn causal_mask_is_strictly_causal() { - let device = crate::backend::cpu_device(); - let m = causal_mask(3, 2, &device); - assert_eq!(m.dims(), [1, 1, 3, 5]); - let v = m.into_data().to_vec::().unwrap(); - for i in 0..3 { - for j in 0..5 { - let expect = if j > 2 + i { -1e4 } else { 0.0 }; - assert_eq!(v[i * 5 + j], expect, "row {i} col {j}"); - } - } - // A single decode step attends to everything cached: all zeros. - let one = causal_mask(1, 4, &device); - assert!( - one.into_data() - .to_vec::() - .unwrap() - .iter() - .all(|&x| x == 0.0) - ); - } - - #[test] - fn repeat_kv_repeats_each_head_contiguously() { - let device = crate::backend::cpu_device(); - let x = Tensor::<1>::from_floats([1.0, 2.0, 3.0, 4.0], &device).reshape([1, 2, 1, 2]); // 2 kv heads, hd=2 - let y = repeat_kv(x, 2); // -> 4 heads - assert_eq!(y.dims(), [1, 4, 1, 2]); - assert_eq!( - y.into_data().to_vec::().unwrap(), - vec![1.0, 2.0, 1.0, 2.0, 3.0, 4.0, 3.0, 4.0] - ); - } - - /// The load-bearing invariant: prefill + one-token-at-a-time decode - /// through the KV cache must equal one full forward over the same tokens. - #[test] - fn kv_cache_decode_matches_full_forward() { - let device = crate::backend::cpu_device(); - for (qk_norm, projection) in [(false, false), (true, false), (true, true)] { - let a = attn(qk_norm, projection, &device); - let x = input(6, 3.0, &device); - - // Reference: all 6 positions in one causal forward; keep the last row. - let full = full_forward(&a, x.clone(), &device); - let last_full = &full[5 * HIDDEN..]; - - // Cached: prefill 5, then decode position 5 alone. - let mut kv: LayerKv = None; - let prefill = x.clone().narrow(1, 0, 5); - let (cos, sin) = rope_tables(5, 0, HEAD_DIM, THETA, &device); - let mask = causal_mask(5, 0, &device); - let _ = a.forward( - prefill, - HEADS, - KV_HEADS, - HEAD_DIM, - &cos, - &sin, - Some(&mask), - &mut kv, - ); - - let step = x.narrow(1, 5, 1); - let (cos1, sin1) = rope_tables(1, 5, HEAD_DIM, THETA, &device); - let out = a - .forward(step, HEADS, KV_HEADS, HEAD_DIM, &cos1, &sin1, None, &mut kv) - .into_data() - .to_vec::() - .unwrap(); - - for (i, (c, f)) in out.iter().zip(last_full).enumerate() { - assert!( - (c - f).abs() < 1e-4, - "qk_norm={qk_norm} projection={projection} elem {i}: cached {c} vs full {f}" - ); - } - } - } - - /// Causality: changing a future token must not change an earlier output row. - #[test] - fn future_tokens_cannot_affect_past_outputs() { - let device = crate::backend::cpu_device(); - let a = attn(false, false, &device); - let x1 = input(4, 1.0, &device); - // Same first 3 tokens, different 4th. - let x2 = Tensor::cat(vec![x1.clone().narrow(1, 0, 3), input(1, 99.0, &device)], 1); - let (o1, o2) = (full_forward(&a, x1, &device), full_forward(&a, x2, &device)); - // Rows 0..3 identical; row 3 differs. - for i in 0..3 * HIDDEN { - assert!((o1[i] - o2[i]).abs() < 1e-6, "past row changed at {i}"); - } - let last_differs = o1[3 * HIDDEN..] - .iter() - .zip(&o2[3 * HIDDEN..]) - .any(|(a, b)| (a - b).abs() > 1e-6); - assert!(last_differs, "the changed token's own row should differ"); - } - - #[test] - #[should_panic(expected = "must be a positive multiple")] - fn init_rejects_indivisible_head_grouping() { - let device = crate::backend::cpu_device(); - let _ = GqaAttentionConfig { - hidden_size: HIDDEN, - num_heads: 3, - num_kv_heads: 2, - head_dim: HEAD_DIM, - bias: false, - qk_norm_eps: None, - qk_norm_projection: false, - } - .init(&device); - } - - /// Negative space: the projection-wide norm is a different function than - /// the per-head norm (RMS over 16 values vs over 4) — same weights, same - /// input, different outputs. Guards against the placement silently - /// collapsing to one branch. - #[test] - fn projection_norm_differs_from_per_head_norm() { - let device = crate::backend::cpu_device(); - let per_head = attn(true, false, &device); - // Same module, but re-shaped norms: reuse per_head's projections and - // swap in projection-wide norms with unit gamma? Simpler: two configs - // share no weights, so instead check the norm widths took effect. - let projection = attn(true, true, &device); - let q_dim = HEADS * HEAD_DIM; - assert_eq!(per_head.q_norm.as_ref().unwrap().gamma.dims(), [HEAD_DIM]); - assert_eq!(projection.q_norm.as_ref().unwrap().gamma.dims(), [q_dim]); - // And the projection-placement forward is exercised end to end by the - // cache-equivalence loop above; here pin that a projection-normed - // forward actually runs (no panic) and returns the right shape. - let x = input(3, 5.0, &device); - let (cos, sin) = rope_tables(3, 0, HEAD_DIM, THETA, &device); - let mask = causal_mask(3, 0, &device); - let mut kv: LayerKv = None; - let out = projection.forward( - x, - HEADS, - KV_HEADS, - HEAD_DIM, - &cos, - &sin, - Some(&mask), - &mut kv, - ); - assert_eq!(out.dims(), [1, 3, HIDDEN]); - } -} diff --git a/crates/mummu/examples/src/nn/conv.rs b/crates/mummu/examples/src/nn/conv.rs deleted file mode 100644 index 7544954..0000000 --- a/crates/mummu/examples/src/nn/conv.rs +++ /dev/null @@ -1,247 +0,0 @@ -//! LFM2's double-gated causal short-convolution ("LIV") operator: -//! `in_proj` → split B/C/x → `B*x` → causal depthwise conv1d → `C*conv` → -//! `out_proj`, with a rolling `K-1` state so decode steps cost O(K) instead -//! of re-forwarding the sequence. The decode path was verified algebraically -//! equivalent to the padded conv in laurelane (greedy parity vs Ollama). - -use burn::module::Module; -use burn::nn::conv::{Conv1d, Conv1dConfig}; -use burn::nn::{Linear, LinearConfig, PaddingConfig1d}; -use burn::tensor::{Device, Tensor}; - -/// Rolling decode state: the last `K-1` gated inputs `[b, d, K-1]`, `None` -/// until the first forward. -pub type ConvState = Option>; - -/// Double-gated causal short conv (LFM2 "LIV" block). Field names mirror the -/// HF checkpoint (`in_proj`/`conv`/`out_proj`). -#[derive(Module, Debug)] -pub struct ShortConv { - pub in_proj: Linear, - pub conv: Conv1d, - pub out_proj: Linear, -} - -/// Shape config for [`ShortConv`]. -#[derive(Debug, Clone)] -pub struct ShortConvConfig { - pub hidden_size: usize, - /// Conv kernel length `K` (LFM2's `conv_L_cache`). - pub kernel_len: usize, -} - -impl ShortConvConfig { - /// Initialize the module (random weights; real weights come from import). - pub fn init(&self, device: &Device) -> ShortConv { - let (d, k) = (self.hidden_size, self.kernel_len); - assert!(d >= 1, "ShortConv: hidden_size must be >= 1"); - assert!( - k >= 2, - "ShortConv: kernel_len must be >= 2 (a 1-tap conv is a no-op gate)" - ); - ShortConv { - in_proj: LinearConfig::new(d, 3 * d).with_bias(false).init(device), - conv: Conv1dConfig::new(d, d, k) - .with_groups(d) - .with_padding(PaddingConfig1d::Explicit(k - 1, k - 1)) - .with_bias(false) - .init(device), - out_proj: LinearConfig::new(d, d).with_bias(false).init(device), - } - } -} - -impl ShortConv { - /// Cache-aware forward. `x` is `[b, t, d]`: the whole prompt at prefill, - /// one token per decode step after. Rolls `state` forward either way. - pub fn forward( - &self, - x: Tensor<3>, - kernel_len: usize, - state: &mut ConvState, - ) -> Tensor<3> { - let [b, t, d] = x.dims(); - let kk = kernel_len; - assert!(kk >= 2, "ShortConv forward: kernel_len must be >= 2"); - assert!(t >= 1, "ShortConv forward: need at least one position"); - debug_assert!( - state.as_ref().is_none_or(|s| s.dims() == [b, d, kk - 1]), - "ShortConv forward: stale state shape" - ); - - let bcx = self.in_proj.forward(x); // [b, t, 3d] - let bb = bcx.clone().narrow(2, 0, d); - let cc = bcx.clone().narrow(2, d, d); - let xx = bcx.narrow(2, 2 * d, d); - let bx = bb.mul(xx).swap_dims(1, 2); // input gate, channel-major [b, d, t] - - let conv_out = if t > 1 { - // Prefill: the padded depthwise Conv1d gives the full causal output. - self.conv.forward(bx.clone()).narrow(2, 0, t) // [b, d, t] - } else { - // Decode: weighted sum over the last K inputs = [cached (K-1), new (1)]. - // Equivalent to the padded conv at the new position. - let window = match state { - Some(prev) => Tensor::cat(vec![prev.clone(), bx.clone()], 2), // [b, d, K] - None => { - let pad = Tensor::<3>::zeros([b, d, kk - 1], &bx.device()); - Tensor::cat(vec![pad, bx.clone()], 2) - } - }; - let w = self.conv.weight.val().reshape([1, d, kk]); // depthwise kernel - window.mul(w).sum_dim(2) // [b, d, 1] - }; - - // Roll the state forward: keep the last (K-1) gated inputs. - let new_state = { - let combined = match state { - Some(prev) => Tensor::cat(vec![prev.clone(), bx.clone()], 2), - None => bx.clone(), - }; - let len = combined.dims()[2]; - if len >= kk - 1 { - combined.narrow(2, len - (kk - 1), kk - 1) - } else { - let pad = Tensor::<3>::zeros([b, d, (kk - 1) - len], &combined.device()); - Tensor::cat(vec![pad, combined], 2) - } - }; - *state = Some(new_state); - - let conv_out = conv_out.swap_dims(1, 2); // [b, t, d] - self.out_proj.forward(cc.mul(conv_out)) // output gate + proj - } -} - -#[cfg(test)] -mod tests { - use super::*; - use burn::tensor::TensorData; - - type Dev = burn::tensor::Device; - - const D: usize = 6; - const K: usize = 3; - - fn conv(device: &Dev) -> ShortConv { - ShortConvConfig { - hidden_size: D, - kernel_len: K, - } - .init(device) - } - - fn input(t: usize, seed: f32, device: &Dev) -> Tensor<3> { - let data: Vec = (0..t * D) - .map(|i| ((i as f32 + seed) * 0.9).cos()) - .collect(); - Tensor::<2>::from_data(TensorData::new(data, [t, D]), device).reshape([1, t, D]) - } - - /// The load-bearing invariant: prefill + one-token-at-a-time decode via - /// the rolling state must equal one full prefill over the same tokens. - #[test] - fn rolling_state_decode_matches_full_prefill() { - let device = crate::backend::cpu_device(); - let c = conv(&device); - let x = input(7, 5.0, &device); - - // Reference: all 7 positions through the padded conv. - let mut ref_state: ConvState = None; - let full = c - .forward(x.clone(), K, &mut ref_state) - .into_data() - .to_vec::() - .unwrap(); - - // Cached: prefill 4, then decode 5th..7th one at a time. - let mut state: ConvState = None; - let _ = c.forward(x.clone().narrow(1, 0, 4), K, &mut state); - for pos in 4..7 { - let out = c - .forward(x.clone().narrow(1, pos, 1), K, &mut state) - .into_data() - .to_vec::() - .unwrap(); - let expect = &full[pos * D..(pos + 1) * D]; - for (i, (got, want)) in out.iter().zip(expect).enumerate() { - assert!( - (got - want).abs() < 1e-5, - "pos {pos} elem {i}: cached {got} vs full {want}" - ); - } - } - } - - /// The conv is causal: a future token cannot change an earlier output. - #[test] - fn conv_is_causal() { - let device = crate::backend::cpu_device(); - let c = conv(&device); - let x1 = input(5, 1.0, &device); - let x2 = Tensor::cat(vec![x1.clone().narrow(1, 0, 4), input(1, 77.0, &device)], 1); - let mut s1: ConvState = None; - let mut s2: ConvState = None; - let o1 = c - .forward(x1, K, &mut s1) - .into_data() - .to_vec::() - .unwrap(); - let o2 = c - .forward(x2, K, &mut s2) - .into_data() - .to_vec::() - .unwrap(); - for i in 0..4 * D { - assert!((o1[i] - o2[i]).abs() < 1e-6, "past row changed at {i}"); - } - } - - /// Decode from a fresh (None) state equals a length-1 prefill: the - /// zero-pad seeding path. - #[test] - fn first_decode_step_seeds_state_like_prefill() { - let device = crate::backend::cpu_device(); - let c = conv(&device); - let x = input(1, 2.0, &device); - - let mut s_decode: ConvState = None; - let via_decode = c - .forward(x.clone(), K, &mut s_decode) - .into_data() - .to_vec::() - .unwrap(); - - // Same single token inside a longer prefill whose first position it is. - let longer = Tensor::cat(vec![x, input(2, 50.0, &device)], 1); - let mut s_pre: ConvState = None; - let via_prefill = c - .forward(longer, K, &mut s_pre) - .into_data() - .to_vec::() - .unwrap(); - - for i in 0..D { - assert!( - (via_decode[i] - via_prefill[i]).abs() < 1e-5, - "elem {i}: decode {} vs prefill {}", - via_decode[i], - via_prefill[i] - ); - } - // State must exist and have the rolling shape after either path. - assert_eq!(s_decode.unwrap().dims(), [1, D, K - 1]); - assert_eq!(s_pre.unwrap().dims(), [1, D, K - 1]); - } - - #[test] - #[should_panic(expected = "kernel_len must be >= 2")] - fn init_rejects_degenerate_kernel() { - let device = crate::backend::cpu_device(); - let _ = ShortConvConfig { - hidden_size: D, - kernel_len: 1, - } - .init(&device); - } -} diff --git a/crates/mummu/examples/src/nn/mlp.rs b/crates/mummu/examples/src/nn/mlp.rs deleted file mode 100644 index 28ebe85..0000000 --- a/crates/mummu/examples/src/nn/mlp.rs +++ /dev/null @@ -1,112 +0,0 @@ -//! SwiGLU feed-forward block: `down(silu(gate(x)) * up(x))`. Field names -//! mirror the HF Qwen2 layout (`gate_proj`/`up_proj`/`down_proj`); LFM2's -//! `w1`/`w3`/`w2` remap onto these at load time. - -use burn::module::Module; -use burn::nn::{Linear, LinearConfig}; -use burn::tensor::{Device, Tensor, activation}; - -/// SwiGLU MLP, no biases (both proven architectures ship it bias-free). -#[derive(Module, Debug)] -pub struct SwiGluMlp { - /// SiLU branch (LFM2: `w1`). - pub gate_proj: Linear, - /// Multiplicative branch (LFM2: `w3`). - pub up_proj: Linear, - /// Projection back to the model width (LFM2: `w2`). - pub down_proj: Linear, -} - -/// Shape config for [`SwiGluMlp`]. -#[derive(Debug, Clone)] -pub struct SwiGluMlpConfig { - pub hidden_size: usize, - pub intermediate_size: usize, -} - -impl SwiGluMlpConfig { - /// Initialize the module (random weights; real weights come from import). - pub fn init(&self, device: &Device) -> SwiGluMlp { - assert!(self.hidden_size >= 1, "SwiGLU: hidden_size must be >= 1"); - assert!( - self.intermediate_size >= 1, - "SwiGLU: intermediate_size must be >= 1" - ); - SwiGluMlp { - gate_proj: LinearConfig::new(self.hidden_size, self.intermediate_size) - .with_bias(false) - .init(device), - up_proj: LinearConfig::new(self.hidden_size, self.intermediate_size) - .with_bias(false) - .init(device), - down_proj: LinearConfig::new(self.intermediate_size, self.hidden_size) - .with_bias(false) - .init(device), - } - } -} - -impl SwiGluMlp { - /// `[b, t, hidden]` → `[b, t, hidden]`. - pub fn forward(&self, x: Tensor<3>) -> Tensor<3> { - let gate = activation::silu(self.gate_proj.forward(x.clone())); - let up = self.up_proj.forward(x); - self.down_proj.forward(gate.mul(up)) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use burn::tensor::TensorData; - - #[test] - fn forward_preserves_shape() { - let device = crate::backend::cpu_device(); - let mlp = SwiGluMlpConfig { - hidden_size: 8, - intermediate_size: 20, - } - .init(&device); - let x = Tensor::<3>::zeros([2, 3, 8], &device); - assert_eq!(mlp.forward(x).dims(), [2, 3, 8]); - } - - #[test] - fn zero_input_gives_zero_output_without_biases() { - let device = crate::backend::cpu_device(); - let mlp = SwiGluMlpConfig { - hidden_size: 4, - intermediate_size: 8, - } - .init(&device); - let x = Tensor::<3>::zeros([1, 2, 4], &device); - let out = mlp.forward(x).into_data().to_vec::().unwrap(); - assert!( - out.iter().all(|&v| v == 0.0), - "bias-free SwiGLU must map 0 to 0" - ); - } - - #[test] - fn forward_is_position_independent() { - // An MLP acts per-position: the same row through a [1,1,h] and a - // [1,2,h] batch must give identical outputs. - let device = crate::backend::cpu_device(); - let mlp = SwiGluMlpConfig { - hidden_size: 4, - intermediate_size: 8, - } - .init(&device); - let row: Vec = vec![0.3, -1.2, 0.8, 2.0]; - let single = - Tensor::<1>::from_data(TensorData::new(row.clone(), [4]), &device).reshape([1, 1, 4]); - let double = - Tensor::<1>::from_data(TensorData::new([row.clone(), row].concat(), [8]), &device) - .reshape([1, 2, 4]); - let s = mlp.forward(single).into_data().to_vec::().unwrap(); - let d = mlp.forward(double).into_data().to_vec::().unwrap(); - assert_eq!(s.as_slice(), &d[..4]); - assert_eq!(s.as_slice(), &d[4..]); - } -} diff --git a/crates/mummu/examples/src/nn/mod.rs b/crates/mummu/examples/src/nn/mod.rs deleted file mode 100644 index 3f55151..0000000 --- a/crates/mummu/examples/src/nn/mod.rs +++ /dev/null @@ -1,29 +0,0 @@ -//! Shared transformer blocks, generic over `B: Backend`. -//! -//! Extracted from laurelane's parity-validated model ports (Qwen2 and the -//! LFM2.5 hybrid — byte-identical single-forward logits and greedy sequences -//! vs Candle / Ollama references). Every model in the zoo composes these; the -//! per-model files own only architecture wiring and weight-key remaps. - -mod attention; -mod conv; -mod mlp; -mod moe; -mod rope; - -pub use attention::{GqaAttention, GqaAttentionConfig, LayerKv, causal_mask, repeat_kv}; -pub use conv::{ConvState, ShortConv, ShortConvConfig}; -pub use mlp::{SwiGluMlp, SwiGluMlpConfig}; -pub mod packed_gemv; -pub use packed_gemv::{Q4GemvOps, packed_gemv_enabled, try_q4s_gemv}; -pub use moe::{ - DeviceExpert, ExpertExec, ExpertPool, ExpertWeights, MoeExperts, Routing, SparseMoe, SparseMoeConfig, - SparseMoePerExpert, StagedExpert, trace_layer, trace_us, -}; -pub use rope::{apply_rope, rope_tables, rotate_half}; - -/// Hard ceiling on `past + t` everywhere a sequence position is materialized. -/// Nothing in the zoo has a longer trained context; a position beyond this is -/// a caller bug (e.g. a decode loop that forgot its stop condition), not a -/// workload. -pub const MAX_CONTEXT_TOKENS: usize = 131_072; diff --git a/crates/mummu/examples/src/nn/moe.rs b/crates/mummu/examples/src/nn/moe.rs deleted file mode 100644 index 96aa094..0000000 --- a/crates/mummu/examples/src/nn/moe.rs +++ /dev/null @@ -1,2204 +0,0 @@ -//! Sparse mixture-of-experts SwiGLU feed-forward (OLMoE-style): a softmax -//! top-k router over a bank of SwiGLU experts stored as **fused 3-D tensors** -//! — exactly the GGUF `ffn_{gate,up,down}_exps` layout, so a checkpoint's -//! expert bank loads as one tensor per projection instead of `num_experts` -//! separate matrices. -//! -//! First cut is **dense-mask compute**: every expert processes every token and -//! the router's sparse weight row (zero for the unrouted experts) scales the -//! results away. That wastes `1 - k/E` of the FLOPs but keeps the whole -//! forward on-device (no data-dependent gather, no host readback of routing -//! decisions) and is numerically identical to the sparse formulation. A -//! gather-based path is a perf follow-up, not a correctness one. - -use burn::module::{Module, Param}; -use burn::nn::{Linear, LinearConfig}; -use burn::tensor::{DType, Device, Distribution, Int, Tensor, TensorData, activation}; - -/// The expert bank: `num_experts` SwiGLU MLPs as three fused params in -/// `[experts, out, in]` layout (the row-major twin of ggml's -/// `ffn_*_exps.weight`). Forward transposes lazily; no per-expert modules. -#[derive(Module, Debug)] -pub struct MoeExperts { - /// `[num_experts, intermediate, hidden]` — SiLU branch. - pub gate: Param>, - /// `[num_experts, intermediate, hidden]` — multiplicative branch. - pub up: Param>, - /// `[num_experts, hidden, intermediate]` — back to the model width. - pub down: Param>, -} - -/// Router + expert bank. Field names follow the HF `Olmoe` checkpoint layout -/// (`mlp.gate` is the router Linear, `mlp.experts` the bank). -#[derive(Module, Debug)] -pub struct SparseMoe { - /// The routing projection: `hidden -> num_experts`, no bias. - pub gate: Linear, - pub experts: MoeExperts, -} - -/// Shape config for [`SparseMoe`]. -#[derive(Debug, Clone)] -pub struct SparseMoeConfig { - pub hidden_size: usize, - /// Per-expert SwiGLU intermediate width (OLMoE: 1024 — each expert is - /// narrow; capacity comes from the count). - pub expert_intermediate_size: usize, - pub num_experts: usize, - pub num_experts_per_tok: usize, -} - -impl SparseMoeConfig { - /// Initialize the module (random weights; real weights come from import). - pub fn init(&self, device: &Device) -> SparseMoe { - assert!( - self.num_experts >= 2, - "MoE: num_experts must be >= 2 (got {}); use SwiGluMlp for a dense FFN", - self.num_experts - ); - assert!( - (1..=self.num_experts).contains(&self.num_experts_per_tok), - "MoE: num_experts_per_tok ({}) must be in 1..=num_experts ({})", - self.num_experts_per_tok, - self.num_experts - ); - assert!( - self.hidden_size >= 1 && self.expert_intermediate_size >= 1, - "MoE: hidden_size and expert_intermediate_size must be >= 1" - ); - let (e, h, inter) = ( - self.num_experts, - self.hidden_size, - self.expert_intermediate_size, - ); - // Linear-style uniform init, bound by each projection's fan-in. - let init = |out: usize, inp: usize, dev: &Device| { - let bound = 1.0 / (inp as f64).sqrt(); - Param::from_tensor(Tensor::random( - [e, out, inp], - Distribution::Uniform(-bound, bound), - dev, - )) - }; - SparseMoe { - gate: LinearConfig::new(h, e).with_bias(false).init(device), - experts: MoeExperts { - gate: init(inter, h, device), - up: init(inter, h, device), - down: init(h, inter, device), - }, - } - } -} - -impl SparseMoe { - /// `[b, t, hidden]` → `[b, t, hidden]`. - /// - /// Router math mirrors HF `OlmoeSparseMoeBlock`: softmax over **all** - /// experts in f32, keep the top-`top_k` probabilities as the mixture - /// weights (renormalized to sum 1 iff `norm_topk_prob` — OLMoE ships - /// `false`). The f32 island matters on f16 backends (softmax of wide - /// logits); every cast is a no-op on f32. - pub fn forward(&self, x: Tensor<3>, top_k: usize, norm_topk_prob: bool) -> Tensor<3> { - let [b, t, h] = x.dims(); - let [e, _inter, h_in] = self.experts.gate.dims(); - assert!( - (1..=e).contains(&top_k), - "MoE forward: top_k ({top_k}) must be in 1..=num_experts ({e})" - ); - assert!( - h == h_in, - "MoE forward: input hidden {h} does not match expert hidden {h_in}" - ); - debug_assert!( - self.experts.down.dims()[1] == h, - "MoE forward: down projection must return to the model width" - ); - let ambient = x.dtype(); - let bt = b * t; - let xt = x.reshape([bt, h]); - - // Router → dense per-token weight rows [bt, e]: softmax probabilities - // where an expert is in the token's top-k, exact zero elsewhere. The - // scatter is an on-device arange-compare — burn's `one_hot` reads the - // indices back to the host, which would sync every layer. - let logits = self.gate.forward(xt.clone()); // [bt, e] - debug_assert!(logits.dims() == [bt, e], "router width must be num_experts"); - let probs = activation::softmax(logits.cast(DType::F32), 1); - let (vals, idx) = probs.topk_with_indices(top_k, 1); // both [bt, k] - let vals = if norm_topk_prob { - vals.clone().div(vals.sum_dim(1)) // [bt, k] / [bt, 1] - } else { - vals - }; - let classes = Tensor::<1, Int>::arange(0..e as i64, &xt.device()).reshape([1, 1, e as i32]); - let hit = idx - .reshape([bt, top_k, 1]) - .equal(classes.expand([bt, top_k, e])); // [bt, k, e] - let weights = hit - .float() - .cast(DType::F32) - .mul(vals.reshape([bt, top_k, 1])) - .sum_dim(1) // [bt, 1, e] - .reshape([bt, e]) - .cast(ambient); - - // Dense expert compute: one batched matmul per projection across the - // whole bank ([1, bt, h] broadcast against [e, h, *]), then the weight - // rows zero out the unrouted experts in the reduction. - let xb = xt.reshape([1, bt, h]); - let gate = xb.clone().matmul(self.experts.gate.val().swap_dims(1, 2)); // [e, bt, inter] - let up = xb.matmul(self.experts.up.val().swap_dims(1, 2)); - let acts = activation::silu(gate).mul(up); - let out = acts.matmul(self.experts.down.val().swap_dims(1, 2)); // [e, bt, h] - let w_per_expert = weights.swap_dims(0, 1).reshape([e, bt, 1]); - out.mul(w_per_expert) - .sum_dim(0) // [1, bt, h] - .reshape([b, t, h]) - } -} - -/// One expert's SwiGLU weights stored **separately** in Linear layout -/// (`[in, out]`), so each expert quantizes independently (its own block -/// scales) and only routed experts are touched at all. -#[derive(Module, Debug)] -pub struct ExpertWeights { - /// `[hidden, intermediate]` — SiLU branch. - pub gate: Param>, - /// `[hidden, intermediate]` — multiplicative branch. - pub up: Param>, - /// `[intermediate, hidden]` — back to the model width. - pub down: Param>, -} - -/// The P9 MoE variant of [`SparseMoe`]: the same router, but experts as -/// separate (typically quantized) weight triples and **routed** compute — -/// per token only its top-k experts run, so exactly `n` experts are in -/// service at a time instead of the dense-mask path's all-of-them. Routing -/// indices are read back to the host (small: `[tokens, k]` ints); the -/// dense path's no-readback rationale trades away here for the k/E FLOPs -/// and the per-expert weight independence quantization needs. -#[derive(Module, Debug)] -pub struct SparseMoePerExpert { - /// The routing projection: `hidden -> num_experts`, no bias, float. - pub gate: Linear, - pub experts: Vec, -} - -impl SparseMoePerExpert { - /// `[b, t, hidden]` → `[b, t, hidden]`. Router math identical to - /// [`SparseMoe::forward`]; expert compute is gather → three 2-D - /// matmuls (never reshaping the possibly-packed weights) → - /// scatter-add of the weighted outputs. - pub fn forward(&self, x: Tensor<3>, top_k: usize, norm_topk_prob: bool) -> Tensor<3> { - let [b, t, h] = x.dims(); - let xt = x.reshape([b * t, h]); - let routing = self.route(xt.clone(), top_k, norm_topk_prob); - self.run_local(xt, &routing).reshape([b, t, h]) - } - - /// The same layer with the experts executed by an [`ExpertPool`] — each - /// expert wherever and at whatever precision the tier plan put it (P9 - /// stage 3b). Router math stays here on `B`; `layer` indexes the pool. - pub fn forward_pooled( - &self, - x: Tensor<3>, - top_k: usize, - norm_topk_prob: bool, - pool: &ExpertPool, - layer: usize, - ) -> Tensor<3> { - let [b, t, h] = x.dims(); - let xt = x.reshape([b * t, h]); - let routing = self.route(xt.clone(), top_k, norm_topk_prob); - pool.run_layer(layer, xt, &routing).reshape([b, t, h]) - } - - /// Router: softmax → top-k → (optionally renormalized) weights, read back - /// to the host as per-expert (token rows, weights) lists. - pub fn route(&self, xt: Tensor<2>, top_k: usize, norm_topk_prob: bool) -> Routing { - let [bt, _h] = xt.dims(); - let e = self.experts.len(); - assert!( - (1..=e).contains(&top_k), - "MoE forward: top_k ({top_k}) must be in 1..=num_experts ({e})" - ); - - let logits = self.gate.forward(xt); // [bt, e] - let probs = activation::softmax(logits.cast(DType::F32), 1); - let (vals, idx) = probs.topk_with_indices(top_k, 1); // both [bt, k] - let vals = if norm_topk_prob { - vals.clone().div(vals.sum_dim(1)) - } else { - vals - }; - - // Host routing: which tokens each expert serves, with what weight. - let idx_host: Vec = idx - .into_data() - .convert::() - .to_vec::() - .expect("routing indices read back"); - let vals_host: Vec = vals - .into_data() - .convert::() - .to_vec::() - .expect("routing weights read back"); - let mut routed: Vec<(Vec, Vec)> = vec![(Vec::new(), Vec::new()); e]; - for token in 0..bt { - for slot in 0..top_k { - let expert = usize::try_from(idx_host[token * top_k + slot]) - .expect("router indices are in 0..e"); - routed[expert].0.push(token as i32); - routed[expert].1.push(vals_host[token * top_k + slot]); - } - } - Routing { - tokens: bt, - per_expert: routed, - } - } - - /// Expert compute on this module's own (same-backend) experts: gather → - /// three 2-D matmuls (never reshaping the possibly-packed weights) → - /// scatter-add of the weighted outputs. `[bt, h]` → `[bt, h]`. - pub fn run_local(&self, xt: Tensor<2>, routing: &Routing) -> Tensor<2> { - let [bt, h] = xt.dims(); - let ambient = xt.dtype(); - let device = xt.device(); - let mut out = Tensor::<2>::zeros([bt, h], &device); - for (expert, (rows, weights)) in routing.per_expert.iter().enumerate() { - if rows.is_empty() { - continue; // not in service this batch - } - let n = rows.len(); - let rows = rows.clone(); - let weights = weights.clone(); - let rows_t = Tensor::<1, Int>::from_data( - burn::tensor::TensorData::new(rows, [n]), - (&device, crate::backend::int_dtype(&device)), - ); - let x_e = xt.clone().select(0, rows_t.clone()); // [n, h] - let w = &self.experts[expert]; - let acts = - activation::silu(x_e.clone().matmul(w.gate.val())).mul(x_e.matmul(w.up.val())); - let y = acts.matmul(w.down.val()); // [n, h] - let scale = Tensor::<1>::from_data( - burn::tensor::TensorData::new(weights, [n]), - (&device, crate::backend::float_dtype(&device)), - ) - .reshape([n, 1]) - .cast(ambient); - // select_assign accumulates (scatter-add) — a token served by - // several experts sums their weighted outputs. - out = out.select_assign(0, rows_t, y.mul(scale), burn::tensor::IndexingUpdateOp::Add); - } - out - } -} - -/// One batch's routing decision, host-side: for every expert, the token -/// rows (into the flattened `[b·t, h]` input) it serves and their weights. -#[derive(Debug, Clone)] -pub struct Routing { - pub tokens: usize, - pub per_expert: Vec<(Vec, Vec)>, -} - -// =========================================================================== -// P9 stage 3(b): tiered expert execution. An expert lives on *some* device -// at *some* stored precision behind `ExpertExec`; the pool holds one per -// (layer, expert), swaps them at runtime, and counts routing hits for the -// tier planner (`crate::tier`). Data crosses devices through host f32 — -// per step only the routed rows (decode: one row per active expert). -// =========================================================================== - -/// One expert, resident somewhere, runnable from the host. -pub trait ExpertExec: Send + Sync { - /// Where/how this expert is resident. - fn tier(&self) -> crate::tier::Tier; - /// Bytes it holds on its device. - fn resident_bytes(&self) -> u64; - /// SwiGLU on `rows × hidden` f32 (row-major) → `rows × hidden`. - /// - /// The host-buffer form. Prefer [`Self::run_tensor`], which keeps the - /// data on-device when the caller is already on this expert's device. - fn run(&self, x: &[f32], rows: usize, hidden: usize) -> Vec; - - /// Like [`Self::run_tensor`], but the result stays wherever this expert - /// computed it — the caller moves it later, batched with every other - /// partial, so the enqueue returns without waiting on the device. - /// - /// Default: `run_tensor`, whose result is already caller-resident — the - /// later move is then a no-op. Device-pinned executors override to skip - /// their trailing move. - fn run_tensor_resident(&self, x: Tensor<2>) -> Tensor<2> { - self.run_tensor(x) - } - - /// Stop handing this executor's packed weights to the native quantized - /// matmul — dequantize before every multiply instead. Called after a - /// caught kernel-gap panic; default is a no-op for executors that never - /// take the native path. - fn disable_native(&self) {} - - /// Bring this expert's weights onto `device` (the working set's - /// prefetch). Default: nothing — an executor pinned to one device is - /// already where it will run, so staging it is a no-op rather than an - /// error. - fn stage(&self, _device: &Device) {} - - /// Release a staged device copy (the working set's eviction). Default: - /// nothing, for the same reason. - fn evict(&self) {} - - /// Is a device copy currently held? Pinned executors answer `true` — - /// they are always "resident" on their own device. - fn is_staged(&self) -> bool { - true - } - - /// SwiGLU on `[rows, hidden]` **as a tensor**: burn 0.22 has one tensor - /// type across devices, so this moves data only when the caller's device - /// differs from the expert's — and not at all when they match, which is - /// the difference between a per-layer host round trip and none. - /// - /// The default keeps the old behavior for executors that only implement - /// the host form. - fn run_tensor(&self, x: Tensor<2>) -> Tensor<2> { - let [rows, hidden] = x.dims(); - let device = x.device(); - let host = x - .into_data() - .convert::() - .to_vec::() - .expect("expert input read back"); - let out = self.run(&host, rows, hidden); - Tensor::<2>::from_data( - burn::tensor::TensorData::new(out, [rows, hidden]), - (&device, crate::backend::float_dtype(&device)), - ) - } - /// Per-row energy of this expert's gate activations, `Σ silu(x·g)²` — - /// the training-free router signal for skipping (P9 stage 3c). Costs - /// the gate matmul only. The default never skips. - fn gate_energy(&self, _x: &[f32], rows: usize, _hidden: usize) -> Vec { - vec![f32::INFINITY; rows] - } -} - -/// [`ExpertWeights`] on a concrete backend's device, exposed as an -/// [`ExpertExec`]. Generic over `B`, so a pool can mix CPU, wgpu and CUDA -/// experts — each at its own precision — behind one trait object. -pub struct DeviceExpert { - pub weights: ExpertWeights, - pub device: Device, - pub tier: crate::tier::Tier, - pub bytes: u64, - /// Cleared the first time this group's native quantized matmul panics - /// (the width-dependent cubecl kernel gap); afterwards every multiply - /// dequantizes first. Per GROUP, because the gap is per shape. - pub native_ok: std::sync::atomic::AtomicBool, -} - -impl ExpertExec for DeviceExpert -where - ExpertWeights: Send + Sync, - Device: Send + Sync, -{ - fn tier(&self) -> crate::tier::Tier { - self.tier - } - - fn resident_bytes(&self) -> u64 { - self.bytes - } - - fn run(&self, x: &[f32], rows: usize, hidden: usize) -> Vec { - debug_assert_eq!(x.len(), rows * hidden); - let xt = Tensor::<2>::from_data( - burn::tensor::TensorData::new(x.to_vec(), [rows, hidden]), - (&self.device, crate::backend::float_dtype(&self.device)), - ); - let w = &self.weights; - let acts = activation::silu(xt.clone().matmul(compute_weight(&w.gate))) - .mul(xt.matmul(compute_weight(&w.up))); - acts.matmul(compute_weight(&w.down)) - .into_data() - .convert::() - .to_vec::() - .expect("expert output read back") - } - - fn run_tensor(&self, x: Tensor<2>) -> Tensor<2> { - // `to_device` is a no-op when the tensor is already here, so a - // cluster group living on the caller's device costs zero transfers. - let caller = x.device(); - crate::backend::move_to(self.run_tensor_resident(x), &caller) - } - - fn run_tensor_resident(&self, x: Tensor<2>) -> Tensor<2> { - let xt = crate::backend::move_to(x, &self.device); - let w = &self.weights; - // After a caught kernel-gap panic, `native_ok` is false and packed - // weights are dequantized up front — the path that is correct at - // every width, at measured-identical speed for a single group. - let native = self.native_ok.load(std::sync::atomic::Ordering::Relaxed); - // At decode shape with a Q4S weight, the packed GEMV reads the - // stored nibbles directly — no dequant transient, no f32 weight - // traffic (nn/packed_gemv.rs). Its launch errors surface at the - // caller's readback, inside the same catch that guards q_matmul, - // so the downgrade contract is unchanged. - let mm = |xin: &Tensor<2>, p: &burn::module::Param>| -> Tensor<2> { - if native { - let wv = p.val(); - if let Some(y) = crate::nn::packed_gemv::try_q4s_gemv(xin, &wv) { - return y; - } - xin.clone().matmul(compute_weight(p)) - } else { - let t = p.val(); - let t = match t.dtype() { - DType::QFloat(_) => t.dequantize(), - _ => t, - }; - xin.clone().matmul(t) - } - }; - let acts = activation::silu(mm(&xt, &w.gate)).mul(mm(&xt, &w.up)); - mm(&acts, &w.down) - } - - fn disable_native(&self) { - self.native_ok - .store(false, std::sync::atomic::Ordering::Relaxed); - } - - fn gate_energy(&self, x: &[f32], rows: usize, hidden: usize) -> Vec { - let xt = Tensor::<2>::from_data( - burn::tensor::TensorData::new(x.to_vec(), [rows, hidden]), - (&self.device, crate::backend::float_dtype(&self.device)), - ); - activation::silu(xt.matmul(compute_weight(&self.weights.gate))) - .powf_scalar(2.0) - .sum_dim(1) - .into_data() - .convert::() - .to_vec::() - .expect("gate energy read back") - } -} - -/// A pooled expert's weight, ready to multiply. -/// -/// A quantized weight is handed to `matmul` **as-is** when this device's -/// backend multiplies it natively — that reads 4–8x fewer weight bytes and -/// skips materializing an f32 copy. Where the native path is broken it is -/// dequantized first instead (transient; the `Param` stays quantized). -/// -/// Which backends work is a property of the burn version and the device, so -/// it is **probed**, not hardcoded — see [`native_qmatmul_ok`]. -fn compute_weight(w: &Param>) -> Tensor<2> { - let t = w.val(); - match t.dtype() { - // wgpu: dequantize the 8-BIT family before the multiply — and only - // that family. This is not the retracted "wgpu Q4 is broken" claim - // of 2026-08-23 (a probe artifact); it has a production backtrace: - // burn 0.22's wgpu q_matmul panics at cubecl-std quant/view.rs:223 - // ("quantized view float vector size 1 must be a positive multiple - // of num_quants 4") when the kernel chosen for an m=1 decode step - // vectorizes the float side at 1. `num_quants 4` is four values per - // u32 — the 8-bit schemes. Q4S packs eight and has never produced - // this panic in any log; a blanket wgpu-dequantize guard was tried - // and traded the panic for something worse: dequantizing every Q4 - // group churned ~90-210 MB f32 transients through a pool that - // allocates ~1 GiB chunks, and OOM-killed generations once other - // apps held part of the card. Narrow beats broad here: - // Q8-family -> dequantize (small groups, tiny transients, covers - // the entire observed panic family); - // Q4S -> native, zero transients (speed measured identical - // either way: 1.61 ms both, bit-identical results). - // CUDA keeps the probed native path for everything. - DType::QFloat(scheme) - if is_wgpu(&t.device()) - && matches!( - scheme.value, - burn::tensor::quantization::QuantValue::Q8S - | burn::tensor::quantization::QuantValue::Q8F - | burn::tensor::quantization::QuantValue::E4M3 - | burn::tensor::quantization::QuantValue::E5M2 - ) => - { - t.dequantize() - } - DType::QFloat(_) if native_qmatmul_ok(&t.device(), t.dtype()) => t, - DType::QFloat(_) => t.dequantize(), - _ => t, - } -} - -/// Is this a wgpu device? burn 0.22 selects backends by runtime `Device` -/// value and exposes no kind accessor, so the debug form is the handle -/// available. Covers the discrete and integrated adapters alike. -fn is_wgpu(device: &Device) -> bool { - format!("{device:?}").contains("Wgpu") -} - -/// Does this device multiply a quantized weight natively — without panicking, -/// and with the right answer? -/// -/// Probed once per (device, scheme) and cached. Burn 0.21's CUDA `q_matmul` -/// panicked in kernel expansion; 0.22 fixed CUDA but wgpu still panics on Q4, -/// and upstream's own autotune candidates are documented as panicking "on the -/// level rather than declining it". A capability that varies by version, -/// backend and dtype is exactly the kind that should be measured on the -/// machine in front of us rather than asserted from a table. -/// -/// Conservative by construction: any panic, or an answer that disagrees with -/// the dequantized path, means "no". -fn native_qmatmul_ok(device: &Device, dtype: DType) -> bool { - use std::collections::HashMap; - use std::sync::{Mutex, OnceLock}; - - static CACHE: OnceLock>> = OnceLock::new(); - let key = format!("{device:?}/{dtype:?}"); - let cache = CACHE.get_or_init(|| Mutex::new(HashMap::new())); - if let Some(&hit) = cache.lock().unwrap_or_else(|e| e.into_inner()).get(&key) { - return hit; - } - - // Small enough to be free, wide enough to cross a quantization block. - let (k, n) = (64usize, 64usize); - let ok = std::panic::catch_unwind(|| { - let x = Tensor::<2>::from_data( - burn::tensor::TensorData::new(vec![0.5f32; k], [1, k]), - (device, crate::backend::float_dtype(device)), - ); - let w = Tensor::<2>::from_data( - burn::tensor::TensorData::new( - (0..k * n) - .map(|i| ((i % 17) as f32 - 8.0) * 0.1) - .collect::>(), - [k, n], - ), - (device, crate::backend::float_dtype(device)), - ); - let DType::QFloat(scheme) = dtype else { - return false; - }; - let qw = w.clone().quantize_dynamic(&scheme); - let native = x - .clone() - .matmul(qw.clone()) - .into_data() - .convert::() - .to_vec::(); - let deq = x - .matmul(qw.dequantize()) - .into_data() - .convert::() - .to_vec::(); - match (native, deq) { - (Ok(a), Ok(b)) => { - // Agreement, not just absence of a panic: a native path that - // silently computes something else is worse than one that fails. - let scale = b.iter().map(|v| v.abs()).fold(1e-3, f32::max); - a.iter().zip(&b).all(|(x, y)| (x - y).abs() <= 0.05 * scale) - } - _ => false, - } - }) - .unwrap_or(false); - - if !ok { - // Worth saying once: it silently costs bandwidth on every matmul. - eprintln!( - "[mummu] {key}: no usable native quantized matmul — dequantizing before each matmul" - ); - } - cache - .lock() - .unwrap_or_else(|e| e.into_inner()) - .insert(key, ok); - ok -} - -/// An expert whose weights live in **host RAM**, staged onto a device only -/// while it computes (P9 stage 4 — see [`crate::workingset`]). -/// -/// [`DeviceExpert`] pins its weights to one device for the process's life, -/// which is what caps how much of a model can ever run on the fast one. -/// This type inverts that: the host copy is authoritative, the device copy -/// is a cache entry the scheduler creates and drops. That is what lets a -/// model larger than VRAM still execute every layer on the GPU. -/// -/// `resident` is the staged device copy. `None` means "not staged": `run` -/// then computes on the host, which is the overflow path — never a stall, -/// because the host already holds the bytes. -pub struct StagedExpert { - /// The authoritative copy, always present, on the host device. - host: ExpertWeights, - /// The host device the weights live on (where overflow computes). - host_device: Device, - /// The staged device copy, when the scheduler has brought it in. - resident: std::sync::RwLock>, - tier: crate::tier::Tier, - bytes: u64, -} - -impl StagedExpert { - /// Hold `weights` in host RAM, unstaged. - #[must_use] - pub fn new( - weights: ExpertWeights, - host_device: Device, - tier: crate::tier::Tier, - bytes: u64, - ) -> Self { - Self { - host: weights, - host_device, - resident: std::sync::RwLock::new(None), - tier, - bytes, - } - } - - /// Stage onto `device` (the scheduler's prefetch). Idempotent: staging - /// onto the device it already sits on does nothing, so a redundant - /// prefetch costs a comparison rather than a transfer. - fn stage_on(&self, device: &Device) { - { - let held = self.resident.read().unwrap_or_else(|e| e.into_inner()); - if held.as_ref().is_some_and(|(d, _)| d == device) { - return; - } - } - let staged = ExpertWeights { - gate: burn::module::Param::from_tensor(self.host.gate.val().to_device(device)), - up: burn::module::Param::from_tensor(self.host.up.val().to_device(device)), - down: burn::module::Param::from_tensor(self.host.down.val().to_device(device)), - }; - *self.resident.write().unwrap_or_else(|e| e.into_inner()) = Some((device.clone(), staged)); - } - - /// Drop the device copy (the scheduler's eviction), freeing its memory. - /// The host copy is untouched, so the expert stays runnable. - fn evict_copy(&self) { - *self.resident.write().unwrap_or_else(|e| e.into_inner()) = None; - } - - fn staged(&self) -> bool { - self.resident - .read() - .unwrap_or_else(|e| e.into_inner()) - .is_some() - } -} - -impl ExpertExec for StagedExpert { - fn tier(&self) -> crate::tier::Tier { - self.tier - } - - fn resident_bytes(&self) -> u64 { - // Device bytes only: the host copy is the backing store, not part of - // the working set the planner budgets. - if self.staged() { self.bytes } else { 0 } - } - - fn stage(&self, device: &Device) { - self.stage_on(device); - } - - fn evict(&self) { - self.evict_copy(); - } - - fn is_staged(&self) -> bool { - self.staged() - } - - fn run(&self, x: &[f32], rows: usize, hidden: usize) -> Vec { - let device = self - .resident - .read() - .unwrap_or_else(|e| e.into_inner()) - .as_ref() - .map_or_else(|| self.host_device.clone(), |(d, _)| d.clone()); - let xt = Tensor::<2>::from_data( - burn::tensor::TensorData::new(x.to_vec(), [rows, hidden]), - (&device, crate::backend::float_dtype(&device)), - ); - self.run_tensor(xt) - .into_data() - .convert::() - .to_vec::() - .expect("expert output read back") - } - - fn run_tensor(&self, x: Tensor<2>) -> Tensor<2> { - let caller = x.device(); - let held = self.resident.read().unwrap_or_else(|e| e.into_inner()); - let (device, w) = match held.as_ref() { - // Staged: compute on the device the scheduler put it on. - Some((d, w)) => (d.clone(), w), - // Not staged — the overflow path. Compute on the host rather - // than stall waiting for a transfer that was never issued. - None => (self.host_device.clone(), &self.host), - }; - let xt = crate::backend::move_to(x, &device); - let acts = activation::silu(xt.clone().matmul(compute_weight(&w.gate))) - .mul(xt.matmul(compute_weight(&w.up))); - crate::backend::move_to(acts.matmul(compute_weight(&w.down)), &caller) - } -} - -/// The tiered expert bank of one model: `[layer][expert]` executors, -/// hot-swappable, with routing hit counters. Shared (`Arc`) between the -/// model that runs it and the planner that re-tiers it. -pub struct ExpertPool { - slots: Vec>>>, - hits: Vec>, - /// Output energy per executor (milli-units), for calibration hotness. - energy: Vec>, - /// Dense-path rows offered / computed (skip accounting). - dense_rows: [std::sync::atomic::AtomicU64; 2], -} - -impl ExpertPool { - /// Build from `[layer][expert]` executors (every layer the same width). - #[must_use] - pub fn new(slots: Vec>>) -> Self { - let counters = || -> Vec> { - slots - .iter() - .map(|l| { - (0..l.len()) - .map(|_| std::sync::atomic::AtomicU64::new(0)) - .collect() - }) - .collect() - }; - let hits = counters(); - let energy = counters(); - Self { - slots: slots - .into_iter() - .map(|l| l.into_iter().map(std::sync::RwLock::new).collect()) - .collect(), - hits, - energy, - dense_rows: [ - std::sync::atomic::AtomicU64::new(0), - std::sync::atomic::AtomicU64::new(0), - ], - } - } - - /// Output energy accumulated per executor since the last call (flat, - /// layer-major, ragged rows concatenated); resets. - pub fn take_energy(&self) -> Vec { - self.energy - .iter() - .flat_map(|l| { - l.iter() - .map(|e| e.swap(0, std::sync::atomic::Ordering::Relaxed) as f64 / 1e3) - }) - .collect() - } - - /// Dense-path skip accounting since the last call: (rows computed, rows - /// offered) summed over executors; resets. - pub fn take_dense_rows(&self) -> (u64, u64) { - ( - self.dense_rows[0].swap(0, std::sync::atomic::Ordering::Relaxed), - self.dense_rows[1].swap(0, std::sync::atomic::Ordering::Relaxed), - ) - } - - #[must_use] - pub fn num_layers(&self) -> usize { - self.slots.len() - } - - #[must_use] - pub fn experts_per_layer(&self) -> usize { - self.slots.first().map_or(0, Vec::len) - } - - /// Flat index of `(layer, expert)` — the tier planner's expert order. - #[must_use] - pub fn flat(&self, layer: usize, expert: usize) -> usize { - layer * self.experts_per_layer() + expert - } - - pub fn get(&self, layer: usize, expert: usize) -> std::sync::Arc { - self.slots[layer][expert] - .read() - .unwrap_or_else(|e| e.into_inner()) - .clone() - } - - /// Replace one expert's executor (the hot-swap). The old one is - /// returned so the caller controls when its device memory is freed. - /// Apply one layer's working-set decisions: evict what the schedule - /// gave up, stage what it prefetched. Both are cheap in-place calls on - /// the executors — no slot swap, because a [`StagedExpert`] owns its - /// host copy and its device copy at once. - /// - /// Called for layer `L` *while layer `L` computes*, so the transfers - /// overlap compute rather than sitting on the critical path. Nothing - /// here blocks: a stage that has not landed by the time its layer runs - /// simply computes on the host (see [`StagedExpert::run_tensor`]). - pub fn apply_schedule( - &self, - layer: usize, - sched: &crate::workingset::LayerSchedule, - device: &Device, - ) { - for &u in &sched.evict { - if let Some(e) = self.unit(layer, u) { - e.evict(); - } - } - for &u in &sched.prefetch { - if let Some(e) = self.unit(layer, u) { - e.stage(device); - } - } - } - - /// The executor for a flat unit id, if this pool holds it. Unit ids are - /// layer-major (`layer * experts_per_layer + index`), matching the - /// scheduler's numbering. - #[must_use] - pub fn unit( - &self, - layer: usize, - unit: crate::workingset::UnitId, - ) -> Option> { - let per = self.experts_per_layer(); - let (l, i) = unit - .checked_div(per) - .map_or((layer, unit), |l| (l, unit % per)); - // A unit id addresses its own layer; fall back to the caller's layer - // for pools whose rows are ragged (dense FFN groups). - let (l, i) = if self.slots.get(l).is_some_and(|r| i < r.len()) { - (l, i) - } else if self.slots.get(layer).is_some_and(|r| unit < r.len()) { - (layer, unit) - } else { - return None; - }; - Some( - self.slots[l][i] - .read() - .unwrap_or_else(|e| e.into_inner()) - .clone(), - ) - } - - /// Device bytes the working set currently holds — what the scheduler - /// budgets against, counting only staged copies. - #[must_use] - pub fn staged_bytes(&self) -> u64 { - (0..self.num_layers()) - .flat_map(|l| (0..self.row_len(l)).map(move |e| (l, e))) - .map(|(l, e)| { - let x = self.get(l, e); - if x.is_staged() { x.resident_bytes() } else { 0 } - }) - .sum() - } - - pub fn swap( - &self, - layer: usize, - expert: usize, - next: std::sync::Arc, - ) -> std::sync::Arc { - let mut slot = self.slots[layer][expert] - .write() - .unwrap_or_else(|e| e.into_inner()); - std::mem::replace(&mut *slot, next) - } - - /// Experts in `layer` (rows may be ragged — a dense model's remote - /// clusters differ per layer). - #[must_use] - pub fn row_len(&self, layer: usize) -> usize { - self.slots.get(layer).map_or(0, Vec::len) - } - - /// Every expert's tier, flat layer-major (ragged rows concatenated). - #[must_use] - pub fn tiers(&self) -> Vec { - (0..self.num_layers()) - .flat_map(|l| (0..self.row_len(l)).map(move |e| (l, e))) - .map(|(l, e)| self.get(l, e).tier()) - .collect() - } - - /// Resident bytes per device index. - #[must_use] - pub fn used_bytes(&self, num_devices: usize) -> Vec { - let mut used = vec![0u64; num_devices]; - for l in 0..self.num_layers() { - for e in 0..self.row_len(l) { - let x = self.get(l, e); - if let Some(u) = used.get_mut(x.tier().device) { - *u += x.resident_bytes(); - } - } - } - used - } -} - -/// Wall-clock in µs since first use, for the layer timeline trace. -pub fn trace_us() -> u128 { - static EPOCH: std::sync::OnceLock = std::sync::OnceLock::new(); - EPOCH - .get_or_init(std::time::Instant::now) - .elapsed() - .as_micros() -} - -/// Which layer (if any) the timeline trace follows (`MUMMU_TRACE_LAYER`). -pub fn trace_layer() -> Option { - static ON: std::sync::OnceLock> = std::sync::OnceLock::new(); - *ON.get_or_init(|| { - std::env::var("MUMMU_TRACE_LAYER") - .ok() - .and_then(|v| v.parse().ok()) - }) -} - -/// Raise the calling worker thread above the trunk's gemm pool. The pool -/// saturates every core through the local slab, and a default-priority -/// worker measurably could not win a core even to SUBMIT its GPU work -/// until the caller reached the join — zero overlap, the full device time -/// exposed (merge.join 27.1 ms/layer with enqueue-first ordering in -/// place). The worker needs microseconds of CPU to submit, then blocks on -/// the fence; above-normal priority preempts one pool thread for exactly -/// that sliver. -fn boost_worker_priority() { - #[cfg(windows)] - { - #[link(name = "kernel32.dll", kind = "raw-dylib", modifiers = "+verbatim")] - unsafe extern "system" { - fn GetCurrentThread() -> isize; - fn SetThreadPriority(handle: isize, priority: i32) -> i32; - } - // SAFETY: plain kernel32 calls on the current thread's pseudo - // handle; 1 = THREAD_PRIORITY_ABOVE_NORMAL. - unsafe { - SetThreadPriority(GetCurrentThread(), 1); - } - } -} - -/// Run one dense executor resident and read its result back as plain -/// bytes, with the same catch-once-downgrade-retry contract as -/// [`run_with_native_fallback`]. The catch MUST span the readback: on the -/// wgpu path the native quantized matmul only enqueues at the op call — -/// kernel expansion happens at the blocking read, where the device server -/// re-raises its panic on the reading thread — so a catch around the -/// compute alone can never see the width-dependent kernel-gap panic -/// (`num_quants 4`/`8`) this fallback exists for. Any second panic — a -/// genuine failure, OOM included — propagates loudly as ever. -fn run_readback_with_fallback(exec: &std::sync::Arc, xt: &Tensor<2>) -> TensorData { - let attempt = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - exec.run_tensor_resident(xt.clone()).into_data() - })); - match attempt { - Ok(data) => data, - Err(_) => { - eprintln!( - "[mummu] native quantized matmul panicked for one expert group \ - (cubecl kernel gap; width-dependent) — group switched to \ - dequantize-first and retried" - ); - exec.disable_native(); - exec.run_tensor_resident(xt.clone()).into_data() - } - } -} - -/// Run one dense executor with the adaptive native fallback: a panic from -/// the width-dependent cubecl kernel gap (quant/view vector-size assert, -/// `num_quants 4`/`8`) is caught ONCE, the group is switched to -/// dequantize-first, and the same input is retried. Any second panic — a -/// genuine failure, OOM included — propagates loudly as ever. -/// -/// The catch is sound here for the same reason `native_qmatmul_ok` could -/// probe this panic: it fires at kernel expand on the calling thread, before -/// submission, and the device server measurably survives it. -fn run_with_native_fallback(exec: &std::sync::Arc, xt: &Tensor<2>) -> Tensor<2> { - let attempt = - std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| exec.run_tensor(xt.clone()))); - match attempt { - Ok(y) => y, - Err(_) => { - eprintln!( - "[mummu] native quantized matmul panicked for one expert group \ - (cubecl kernel gap; width-dependent) — group switched to \ - dequantize-first and retried" - ); - exec.disable_native(); - exec.run_tensor(xt.clone()) - } - } -} - -/// One layer's remote FFN, in flight: each device's worker thread is -/// computing AND draining its own device — concurrently, exactly as the old -/// synchronous path did — and only the JOIN is deferred, so the caller's -/// local slab runs while the devices chew. -/// -/// This is the third design. The first deferred the device sync itself -/// into [`Self::resolve`], and adversarial review proved against the -/// burn-fusion sources that nothing executes at enqueue — the worker -/// streams' queued IR first ran inside resolve's drain, and resolving -/// partials one at a time serialized the devices where the old workers -/// drained them in parallel (wall `local + T_a + T_b` instead of -/// `local + max(T_a, T_b)`). The second deferred only the join, with each -/// worker building the caller-device result tensor from its readback — and -/// every layer still paid ~27 ms SOMEWHERE, migrating between scopes as -/// the code moved (worker build, join, a main-thread touch), because -/// wgpu's `into_data` returns deferred-mapped bytes: it comes back in -/// low ms while the FIRST CPU TOUCH of the bytes blocks on the GPU fence -/// (`examples/mapped-wait-probe.rs` reproduces it standalone: into_data -/// 1-4 ms, first touch 27.0-27.5 ms behind a queued GPU chain). The cost -/// was never queues, scheduling, or allocation — it is the remote FFN's -/// real GPU time surfacing at first byte access. So the third design -/// makes the WORKER touch the bytes (`to_vec` in the accumulate): the -/// fence wait lands here, concurrent with whatever the caller still has -/// to do. The layer timeline (MUMMU_TRACE_LAYER) then showed how little -/// that is: remote-heavy layers keep ~1 local cluster, so the caller -/// reaches the join ~1 ms after the enqueue and ~26 ms of GPU time is -/// exposed with nothing to overlap against. Measured per cluster at m=1: -/// dGPU 0.91 ms vs CPU 0.40 ms — the GPU is 2.3x SLOWER than the host it -/// is meant to relieve, because the quantized matmul dequantizes to f32. -/// Shrinking that is a kernel problem (packed m=1 GEMV), not a -/// scheduling one; every scheduling fix here is already in place. -pub struct PendingRemote { - handles: Vec>>, - /// The caller's device, captured at enqueue: partials come home to - /// wherever the trunk lives (wgpu-trunk placements exist), never to a - /// hardcoded CPU device. - home: Device, -} - -impl PendingRemote { - /// Join the workers, then build and sum their partials on the CALLER's - /// thread, on the caller's own device, with each readback's own dtype. - /// The workers hand back plain FULLY-MATERIALIZED bytes: wgpu readback - /// bytes are deferred-mapped, and their first CPU touch blocks on the - /// GPU fence (~27 ms measured; see `mapped-wait-probe`), so the worker - /// touches them in its accumulate — concurrent with the local slab — - /// and everything on this thread is microseconds. - #[must_use] - pub fn resolve(self) -> Option> { - let mut out: Option> = None; - for handle in self.handles { - let partial = { - let _s = crate::prof::scope("merge.join"); - match handle.join() { - Ok(p) => p, - // Same rule as everywhere in this pool: a dead worker is - // a dead generation, never a silently smaller sum. - Err(payload) => std::panic::resume_unwind(payload), - } - }; - if let Some(data) = partial { - let _s = crate::prof::scope("merge.sum"); - let dtype = data.dtype; - let p = Tensor::from_data(data, (&self.home, dtype)); - out = Some(match out { - Some(acc) => acc.add(p), - None => p, - }); - } - } - out - } -} - -impl ExpertPool { - /// The exact dense remote FFN with a deferred join: spawn one worker per - /// device (the bounded-stream rule), each running its members and moving - /// its accumulated partial to the caller — the drain — on its own - /// thread, then return without joining. Skip-mode (tau > 0) stays on - /// [`Self::run_dense`]: it needs host energies up front. - /// - /// Plain `std::thread::spawn`, not a scope: the whole point is that the - /// threads outlive this call. Everything moved in is `Arc`s and owned - /// tensors. ~2 spawns per layer is microseconds against multi-ms drains; - /// a persistent pool is the upgrade if a profile ever says otherwise. - pub fn run_dense_pending(&self, layer: usize, xt: Tensor<2>) -> Option { - let n = self.row_len(layer); - if n == 0 { - return None; - } - let [bt, _h] = xt.dims(); - let execs: Vec> = - (0..n).map(|e| self.get(layer, e)).collect(); - - let mut by_device: Vec<(String, Vec>)> = Vec::new(); - for exec in &execs { - let key = format!("{:?}", exec.tier().device); - match by_device.iter_mut().find(|(k, _)| *k == key) { - Some((_, list)) => list.push(exec.clone()), - None => by_device.push((key, vec![exec.clone()])), - } - } - for (e, _) in execs.iter().enumerate() { - self.hits[layer][e].fetch_add(bt as u64, std::sync::atomic::Ordering::Relaxed); - } - self.dense_rows[0].fetch_add(bt as u64, std::sync::atomic::Ordering::Relaxed); - self.dense_rows[1].fetch_add(bt as u64, std::sync::atomic::Ordering::Relaxed); - - let handles: Vec>> = by_device - .into_iter() - .map(|(key, members)| { - let xt = xt.clone(); - let traced = trace_layer() == Some(layer); - std::thread::spawn(move || { - boost_worker_priority(); - if traced { - eprintln!("[tl] worker-run {}", trace_us()); - } - let _w = crate::prof::scope("ffn_worker"); - let _d = crate::prof::scope(key); - // Bytes only on this thread: per executor, one - // compute-and-readback (the catch spans both — the - // kernel-gap panic surfaces at the READ, see - // run_readback_with_fallback) and a plain-f32 - // accumulate that needs no backend at all. The - // accumulate's `to_vec` is deliberate: wgpu readback - // bytes are deferred-mapped and their first CPU touch - // blocks on the GPU fence (~27 ms measured), so THIS - // thread absorbs that wait concurrent with the - // caller's local slab instead of leaking it into the - // merge. The caller rebuilds the tensor in resolve(). - let mut acc: Option<(Vec, burn::tensor::Shape)> = None; - for exec in &members { - let data = { - let _s = crate::prof::scope("compute"); - run_readback_with_fallback(exec, &xt) - }; - if traced { - eprintln!("[tl] readback-done {}", trace_us()); - } - let shape = data.shape.clone(); - let vals = data - .convert::() - .to_vec::() - .expect("remote FFN partial reads back as f32"); - match &mut acc { - Some((a, s)) => { - debug_assert_eq!(*s, shape); - for (d, v) in a.iter_mut().zip(&vals) { - *d += v; - } - } - None => acc = Some((vals, shape)), - } - } - if traced { - eprintln!("[tl] touch-done {}", trace_us()); - } - acc.map(|(vals, shape)| TensorData::new(vals, shape)) - }) - }) - .collect(); - Some(PendingRemote { - handles, - home: xt.device(), - }) - } -} - -impl ExpertPool { - /// Dense-model FFN path (P9 stage 3c): run **every** executor of - /// `layer` on every row and sum — the remote clusters' share of a - /// partitioned SwiGLU (the local slab runs on the model's own device). - /// `None` when the layer has no remote clusters. With `skip` - /// (`Some((tau, local_energy))`), a cluster is skipped for a row when its - /// gate energy is below `tau` × the row's total energy (local + all - /// remote) — the opt-in lossy mode; hit counters accumulate energy so - /// the re-tier planner sees hot clusters. - pub fn run_dense( - &self, - layer: usize, - xt: Tensor<2>, - skip: Option<(f32, &[f32])>, - ) -> Option> { - let n = self.row_len(layer); - if n == 0 { - return None; - } - let [bt, h] = xt.dims(); - let device = xt.device(); - let execs: Vec> = - (0..n).map(|e| self.get(layer, e)).collect(); - - // Exact mode (no skipping): every cluster runs on every row, so there - // is nothing to gather — sum the groups as tensors and never touch - // the host. This is the path a dense model takes, and it removes the - // per-layer round trip that dominated the 27B's decode. - if skip.is_none() { - // Run the devices CONCURRENTLY, one thread per device, and sum - // what comes back. Sequentially the layer costs the SUM of every - // device's share, so adding a second GPU bought nothing: 885 - // clusters moved from the CPU to the integrated GPU and decode - // measured 4.72 s/tok against 4.32 before, because an iGPU - // cluster (14.15 ms) is no faster than the CPU cluster it - // replaced (13.82 ms) and the move added a transfer. Run in - // parallel the layer costs the MAX instead, which is the entire - // reason to spread work across devices at all. - // - // One thread per DEVICE, never per executor. Thread-per-executor - // is what made cubecl-cuda open a stream per thread until a - // 64-layer forward exhausted VRAM (`CUDA_ERROR_OUT_OF_MEMORY`, - // "Can create a new stream"). Devices are bounded — three on this - // box — so the stream count is bounded with them. - let mut by_device: Vec<(String, Vec)> = Vec::new(); - for (e, exec) in execs.iter().enumerate() { - let key = format!("{:?}", exec.tier().device); - match by_device.iter_mut().find(|(k, _)| *k == key) { - Some((_, list)) => list.push(e), - None => by_device.push((key, vec![e])), - } - } - for (e, _) in execs.iter().enumerate() { - self.hits[layer][e].fetch_add(bt as u64, std::sync::atomic::Ordering::Relaxed); - } - self.dense_rows[0].fetch_add(bt as u64, std::sync::atomic::Ordering::Relaxed); - self.dense_rows[1].fetch_add(bt as u64, std::sync::atomic::Ordering::Relaxed); - - // One device: no threads, no join, exactly the old path. - if by_device.len() < 2 { - let mut out: Option> = None; - for exec in &execs { - let y = run_with_native_fallback(exec, &xt); - out = Some(match out { - Some(acc) => acc.add(y), - None => y, - }); - } - return out; - } - - let partials: Vec> = std::thread::scope(|scope| { - let handles: Vec<_> = by_device - .iter() - .map(|(key, members)| { - let execs = &execs; - let xt = xt.clone(); - scope.spawn(move || { - // A worker thread's stack starts empty, so this - // is a new flame-graph ROOT beside the forward's - // — read the widths as parallel wall time. - let _w = crate::prof::scope("ffn_worker"); - let _d = crate::prof::scope(key.clone()); - let mut acc: Option> = None; - for &e in members { - let y = execs[e].run_tensor(xt.clone()); - acc = Some(match acc { - Some(a) => a.add(y), - None => y, - }); - } - acc - }) - }) - .collect(); - handles - .into_iter() - .filter_map(|h| match h.join() { - Ok(partial) => partial, - // A worker panic must never become a silently missing - // partial: dropping one device's clusters from the - // FFN sum produces a WRONG answer that still reads - // fluently — observed in production when the iGPU's - // workers OOM'd and `.ok()` discarded their share. - // Re-raise on the caller so the generation fails - // loudly instead of lying. - Err(payload) => std::panic::resume_unwind(payload), - }) - .collect() - }); - - // Sum the per-device partials on the caller's device. - let mut out: Option> = None; - for p in partials { - let p = crate::backend::move_to(p, &device); - out = Some(match out { - Some(acc) => acc.add(p), - None => p, - }); - } - return out; - } - - let host: Vec = xt - .into_data() - .convert::() - .to_vec::() - .expect("FFN input read back"); - // Which rows each executor runs: all, or the rows where it matters. - let rows_per_exec: Vec> = match skip { - None => vec![(0..bt as i32).collect(); n], - Some((tau, local_energy)) => { - // Sequential, not one thread per executor: a fresh OS thread - // makes cubecl-cuda open a new CUDA stream, and a 64-layer - // forward would exhaust VRAM creating them (CUDA_ERROR_OUT_OF_MEMORY - // "Can create a new stream"). The calling thread already owns a - // stream; reuse it. - let energies: Vec> = execs - .iter() - .map(|exec| exec.gate_energy(&host, bt, h)) - .collect(); - let mut total: Vec = local_energy.to_vec(); - total.resize(bt, 0.0); - for e in &energies { - for (t, &v) in total.iter_mut().zip(e) { - *t += v; - } - } - energies - .iter() - .map(|e| { - (0..bt) - .filter(|&r| e[r] >= tau * total[r]) - .map(|r| r as i32) - .collect() - }) - .collect() - } - }; - // Sequential per executor (see the energy path above): threads here - // would each open a CUDA stream and OOM the device over 64 layers. - let outputs: Vec> = execs - .iter() - .zip(&rows_per_exec) - .map(|(exec, rows)| { - if rows.is_empty() { - return Vec::new(); - } - let mut x = Vec::with_capacity(rows.len() * h); - for &r in rows { - let r = r as usize; - x.extend_from_slice(&host[r * h..(r + 1) * h]); - } - exec.run(&x, rows.len(), h) - }) - .collect(); - let mut out = vec![0f32; bt * h]; - for ((e, rows), y) in rows_per_exec.iter().enumerate().zip(&outputs) { - let mut energy = 0f32; - for (i, &r) in rows.iter().enumerate() { - let dst = &mut out[r as usize * h..(r as usize + 1) * h]; - for (d, &v) in dst.iter_mut().zip(&y[i * h..(i + 1) * h]) { - *d += v; - energy += v * v; - } - } - // Energy-weighted "hits": what the planner treats as hotness. - self.hits[layer][e].fetch_add( - (energy * 1e3) as u64 + rows.len() as u64, - std::sync::atomic::Ordering::Relaxed, - ); - self.energy[layer][e] - .fetch_add((energy * 1e3) as u64, std::sync::atomic::Ordering::Relaxed); - self.dense_rows[0].fetch_add(rows.len() as u64, std::sync::atomic::Ordering::Relaxed); - self.dense_rows[1].fetch_add(bt as u64, std::sync::atomic::Ordering::Relaxed); - } - Some(Tensor::<2>::from_data( - burn::tensor::TensorData::new(out, [bt, h]), - (&device, crate::backend::float_dtype(&device)), - )) - } - - /// Routing hits since the last call (token-rows served per expert), - /// flat layer-major; resets the counters. - pub fn take_hits(&self) -> Vec { - self.hits - .iter() - .flat_map(|l| { - l.iter() - .map(|h| h.swap(0, std::sync::atomic::Ordering::Relaxed)) - }) - .collect() - } - - /// Run one layer's routed experts: gather the routed rows on the host, - /// execute every in-service expert concurrently on its own device, - /// scatter-add the weighted outputs, upload once. `[bt, h]` → `[bt, h]`. - pub fn run_layer(&self, layer: usize, xt: Tensor<2>, routing: &Routing) -> Tensor<2> { - let [bt, h] = xt.dims(); - let device = xt.device(); - let host: Vec = xt - .into_data() - .convert::() - .to_vec::() - .expect("MoE input read back"); - type Routed<'a> = (usize, &'a (Vec, Vec)); - let active: Vec> = routing - .per_expert - .iter() - .enumerate() - .filter(|(_, (rows, _))| !rows.is_empty()) - .collect(); - let execs: Vec> = - active.iter().map(|(e, _)| self.get(layer, *e)).collect(); - for (e, (rows, _)) in &active { - self.hits[layer][*e].fetch_add(rows.len() as u64, std::sync::atomic::Ordering::Relaxed); - } - let outputs: Vec> = std::thread::scope(|s| { - let handles: Vec<_> = active - .iter() - .zip(&execs) - .map(|((_, (rows, _)), exec)| { - let host = &host; - s.spawn(move || { - let mut x = Vec::with_capacity(rows.len() * h); - for &r in rows { - let r = r as usize; - x.extend_from_slice(&host[r * h..(r + 1) * h]); - } - exec.run(&x, rows.len(), h) - }) - }) - .collect(); - handles - .into_iter() - .map(|hd| hd.join().expect("expert thread")) - .collect() - }); - let mut out = vec![0f32; bt * h]; - for ((_, (rows, weights)), y) in active.iter().zip(&outputs) { - for (i, &r) in rows.iter().enumerate() { - let w = weights[i]; - let dst = &mut out[r as usize * h..(r as usize + 1) * h]; - for (d, &v) in dst.iter_mut().zip(&y[i * h..(i + 1) * h]) { - *d += w * v; - } - } - } - Tensor::<2>::from_data( - burn::tensor::TensorData::new(out, [bt, h]), - (&device, crate::backend::float_dtype(&device)), - ) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use burn::tensor::TensorData; - - type Dev = burn::tensor::Device; - - const HIDDEN: usize = 4; - const INTER: usize = 3; - const EXPERTS: usize = 4; - const TOP_K: usize = 2; - - /// Deterministic weights: expert `e`'s matrices are small distinct - /// sinusoids so every expert computes something different. - fn moe(device: &Dev) -> SparseMoe { - let fill = |seed: f32, dims: [usize; 3]| { - let n = dims[0] * dims[1] * dims[2]; - let data: Vec = (0..n) - .map(|i| ((i as f32) * 0.37 + seed).sin() * 0.5) - .collect(); - Param::from_tensor(Tensor::<3>::from_data(TensorData::new(data, dims), device)) - }; - let router_data: Vec = (0..HIDDEN * EXPERTS) - .map(|i| ((i as f32) * 0.61 + 1.0).cos() * 0.5) - .collect(); - let mut m = SparseMoeConfig { - hidden_size: HIDDEN, - expert_intermediate_size: INTER, - num_experts: EXPERTS, - num_experts_per_tok: TOP_K, - } - .init(device); - // Burn Linear stores weight as [in, out]. - m.gate.weight = Param::from_tensor(Tensor::<2>::from_data( - TensorData::new(router_data, [HIDDEN, EXPERTS]), - device, - )); - m.experts.gate = fill(0.1, [EXPERTS, INTER, HIDDEN]); - m.experts.up = fill(1.7, [EXPERTS, INTER, HIDDEN]); - m.experts.down = fill(3.3, [EXPERTS, HIDDEN, INTER]); - m - } - - /// The per-expert routed path must equal the dense-mask path exactly - /// (same math, different data movement) — the P9 MoE gate. - #[test] - fn per_expert_routed_matches_dense_mask() { - let device = crate::backend::cpu_device(); - let dense = moe(&device); - // Split the fused banks into per-expert Linear-layout triples. - let experts: Vec = (0..EXPERTS) - .map(|e| { - let slice = |bank: &Param>| -> Tensor<2> { - let [_, out, inp] = bank.val().dims(); - bank.val() - .narrow(0, e, 1) - .reshape([out, inp]) - .swap_dims(0, 1) // [in, out] Linear layout - }; - ExpertWeights { - gate: Param::from_tensor(slice(&dense.experts.gate)), - up: Param::from_tensor(slice(&dense.experts.up)), - down: Param::from_tensor(slice(&dense.experts.down)), - } - }) - .collect(); - let per_expert = SparseMoePerExpert { - gate: dense.gate.clone(), - experts, - }; - - for (t, seed, norm) in [(1usize, 2.0f32, false), (5, 7.0, false), (5, 7.0, true)] { - let x = input(t, seed, &device); - let a = dense - .forward(x.clone(), TOP_K, norm) - .into_data() - .to_vec::() - .unwrap(); - let b = per_expert - .forward(x, TOP_K, norm) - .into_data() - .to_vec::() - .unwrap(); - for (i, (da, db)) in a.iter().zip(&b).enumerate() { - assert!( - (da - db).abs() < 1e-5, - "t={t} norm={norm} elem {i}: dense {da} vs routed {db}" - ); - } - } - } - - /// The working set end to end: a scheduler plan drives real staging and - /// eviction in the pool, the device budget is respected, and the layer - /// still computes the right answer whether its experts were staged or - /// overflowed to the host. - #[test] - fn a_schedule_drives_staging_and_the_answer_is_unchanged() { - use crate::tier::{Precision, Tier}; - use crate::workingset::{Budget, LayerDemand, schedule}; - use std::sync::Arc; - let device = crate::backend::cpu_device(); - let dense = moe(&device); - let split = |e: usize| -> ExpertWeights { - let slice = |bank: &Param>| -> Tensor<2> { - let [_, out, inp] = bank.val().dims(); - bank.val() - .narrow(0, e, 1) - .reshape([out, inp]) - .swap_dims(0, 1) - }; - ExpertWeights { - gate: Param::from_tensor(slice(&dense.experts.gate)), - up: Param::from_tensor(slice(&dense.experts.up)), - down: Param::from_tensor(slice(&dense.experts.down)), - } - }; - let tier = Tier { - device: 0, - precision: Precision::F32, - }; - const UNIT_BYTES: u64 = 1_000; - - // One layer holding every expert, each staged-capable. - let row: Vec> = (0..EXPERTS) - .map(|e| { - Arc::new(StagedExpert::new( - split(e), - device.clone(), - tier, - UNIT_BYTES, - )) as Arc - }) - .collect(); - let pool = ExpertPool::new(vec![row]); - - // Nothing staged yet: the working set costs no device memory. - assert_eq!( - pool.staged_bytes(), - 0, - "an unstaged pool holds no device bytes" - ); - - // A schedule over two passes of this layer, with room for half the - // experts — so it must both stage and evict. - let demands: Vec = (0..2) - .map(|l| LayerDemand { - layer: l, - units: (0..EXPERTS).collect(), - }) - .collect(); - let budget = Budget { - device_bytes: UNIT_BYTES * (EXPERTS as u64 / 2), - unit_bytes: UNIT_BYTES, - stage_bytes_per_sec: (UNIT_BYTES as f64) * 2.0 / 0.010, - layer_compute_secs: 0.010, - }; - let plan = schedule(&demands, &budget); - - // Apply the first layer's decisions, as the runtime would. - pool.apply_schedule(0, &plan.layers[0], &device); - let staged = pool.staged_bytes(); - assert!( - staged <= budget.device_bytes, - "the working set must respect the device budget: {staged} > {}", - budget.device_bytes - ); - assert!(staged > 0, "a prefetching schedule must stage something"); - - // The layer's answer must match the unpooled reference regardless of - // which experts happened to be staged. - let local = SparseMoePerExpert { - gate: dense.gate.clone(), - experts: (0..EXPERTS).map(split).collect(), - }; - let x = input(3, 5.0, &device); - let want = local - .forward(x.clone(), TOP_K, true) - .into_data() - .to_vec::() - .unwrap(); - let got = local - .forward_pooled(x, TOP_K, true, &pool, 0) - .into_data() - .to_vec::() - .unwrap(); - for (i, (a, b)) in want.iter().zip(&got).enumerate() { - assert!( - (a - b).abs() < 1e-5, - "elem {i}: {a} vs {b} (staging changed the answer)" - ); - } - } - - /// P9 stage 4: staging must move WHERE an expert computes without - /// changing WHAT it computes. Same weights, same input, same answer — - /// staged, evicted, and re-staged. - #[test] - fn staging_and_eviction_never_change_the_result() { - use crate::tier::{Precision, Tier}; - let device = crate::backend::cpu_device(); - let dense = moe(&device); - let split = |e: usize| -> ExpertWeights { - let slice = |bank: &Param>| -> Tensor<2> { - let [_, out, inp] = bank.val().dims(); - bank.val() - .narrow(0, e, 1) - .reshape([out, inp]) - .swap_dims(0, 1) - }; - ExpertWeights { - gate: Param::from_tensor(slice(&dense.experts.gate)), - up: Param::from_tensor(slice(&dense.experts.up)), - down: Param::from_tensor(slice(&dense.experts.down)), - } - }; - let tier = Tier { - device: 0, - precision: Precision::F32, - }; - let staged = StagedExpert::new(split(0), device.clone(), tier, 0); - // A pinned reference on the same device, for comparison. - let pinned = DeviceExpert { - weights: split(0), - device: device.clone(), - tier, - bytes: 0, - native_ok: std::sync::atomic::AtomicBool::new(true), - }; - - let x = Tensor::<2>::from_data( - TensorData::new( - (0..HIDDEN) - .map(|i| (i as f32) * 0.25 - 0.5) - .collect::>(), - [1, HIDDEN], - ), - (&device, crate::backend::float_dtype(&device)), - ); - let want = pinned - .run_tensor(x.clone()) - .into_data() - .to_vec::() - .unwrap(); - - // Unstaged (the overflow path: computes on the host). - assert!( - !staged.is_staged(), - "a fresh StagedExpert holds no device copy" - ); - assert_eq!( - staged.resident_bytes(), - 0, - "unstaged costs no device memory" - ); - let overflow = staged - .run_tensor(x.clone()) - .into_data() - .to_vec::() - .unwrap(); - - // Staged, then evicted, then staged again. - staged.stage(&device); - assert!(staged.is_staged()); - let hot = staged - .run_tensor(x.clone()) - .into_data() - .to_vec::() - .unwrap(); - staged.stage(&device); // idempotent: no second transfer, same answer - let again = staged - .run_tensor(x.clone()) - .into_data() - .to_vec::() - .unwrap(); - staged.evict(); - assert!(!staged.is_staged(), "eviction drops the device copy"); - let after_evict = staged.run_tensor(x).into_data().to_vec::().unwrap(); - - for (i, w) in want.iter().enumerate() { - for (label, got) in [ - ("overflow", &overflow), - ("staged", &hot), - ("re-staged", &again), - ("after evict", &after_evict), - ] { - assert!( - (w - got[i]).abs() < 1e-5, - "{label} elem {i}: {} vs pinned {w}", - got[i] - ); - } - } - } - - /// P9 stage 3b: the pooled path (experts behind `ExpertExec`, host - /// round trip, concurrent execution, host scatter-add) equals the - /// same-backend routed path — here with the pool holding Q8 experts - /// on the "GPU" tier and f32 ones on the "CPU" tier side by side, and - /// a hot-swap mid-way. The f32-tier experts must match exactly; the - /// Q8 ones within quantization noise. - #[test] - fn pooled_experts_match_local_and_hot_swap() { - use crate::tier::{Precision, Tier}; - use std::sync::Arc; - let device = crate::backend::cpu_device(); - let dense = moe(&device); - let split = |e: usize| -> ExpertWeights { - let slice = |bank: &Param>| -> Tensor<2> { - let [_, out, inp] = bank.val().dims(); - bank.val() - .narrow(0, e, 1) - .reshape([out, inp]) - .swap_dims(0, 1) - }; - ExpertWeights { - gate: Param::from_tensor(slice(&dense.experts.gate)), - up: Param::from_tensor(slice(&dense.experts.up)), - down: Param::from_tensor(slice(&dense.experts.down)), - } - }; - let local = SparseMoePerExpert { - gate: dense.gate.clone(), - experts: (0..EXPERTS).map(split).collect(), - }; - // `Device` is a clonable runtime value in burn 0.22 (not a `Copy` - // marker type), so the closure clones per expert instead of copying. - let exec = |e: usize, tier: Tier| -> Arc { - Arc::new(DeviceExpert { - native_ok: std::sync::atomic::AtomicBool::new(true), - weights: split(e), - device: device.clone(), - tier, - bytes: 0, - }) - }; - let f32_tier = Tier { - device: 0, - precision: Precision::F32, - }; - let pool = ExpertPool::new(vec![(0..EXPERTS).map(|e| exec(e, f32_tier)).collect()]); - assert_eq!((pool.num_layers(), pool.experts_per_layer()), (1, EXPERTS)); - - for (t, seed, norm) in [(1usize, 2.0f32, false), (5, 7.0, true)] { - let x = input(t, seed, &device); - let a = local - .forward(x.clone(), TOP_K, norm) - .into_data() - .to_vec::() - .unwrap(); - let b = local - .forward_pooled(x, TOP_K, norm, &pool, 0) - .into_data() - .to_vec::() - .unwrap(); - for (i, (da, db)) in a.iter().zip(&b).enumerate() { - assert!( - (da - db).abs() < 1e-5, - "t={t} elem {i}: local {da} vs pooled {db}" - ); - } - } - // Hits were counted: t=1 and t=5 with top-2 → 12 routed rows total. - let hits = pool.take_hits(); - assert_eq!(hits.iter().sum::(), 12, "{hits:?}"); - assert_eq!(pool.take_hits().iter().sum::(), 0); - - // Hot-swap expert 1 onto a different tier (same weights): output - // unchanged, tier bookkeeping updated, old executor handed back. - let q_tier = Tier { - device: 1, - precision: Precision::Q8, - }; - let old = pool.swap(0, 1, exec(1, q_tier)); - assert_eq!(old.tier(), f32_tier); - assert_eq!(pool.tiers()[1], q_tier); - let x = input(5, 7.0, &device); - let a = local - .forward(x.clone(), TOP_K, true) - .into_data() - .to_vec::() - .unwrap(); - let b = local - .forward_pooled(x, TOP_K, true, &pool, 0) - .into_data() - .to_vec::() - .unwrap(); - for (i, (da, db)) in a.iter().zip(&b).enumerate() { - assert!((da - db).abs() < 1e-5, "after swap elem {i}: {da} vs {db}"); - } - } - - /// `compute_weight` hands `matmul` something that produces the RIGHT - /// ANSWER — either the quantized weight itself (where the backend - /// multiplies it natively) or a dequantized copy (where it does not). - /// - /// The test asserts the invariant, not the mechanism: which branch runs - /// depends on the backend, the burn version and the dtype, and is probed - /// at runtime. Asserting "it always dequantizes" would have to be - /// rewritten every time a backend gains a working kernel — and would - /// have failed the moment burn 0.22 fixed CUDA. - #[test] - fn compute_weight_yields_a_matmul_ready_weight() { - use crate::quant::{QuantPolicy, quantize_weight}; - use burn::tensor::TensorData; - let device = crate::backend::cpu_device(); - let vals: Vec = (0..32 * 64).map(|i| ((i as f32) * 0.05).sin()).collect(); - let t = Tensor::<2>::from_data(TensorData::new(vals.clone(), [32, 64]), &device); - - // Float param: returned as-is (still float, never re-quantized). - let out_f32 = compute_weight(&Param::from_tensor(t.clone())); - assert!( - !matches!(out_f32.dtype(), DType::QFloat(_)), - "float weight must stay float" - ); - - // Quantized param: whatever comes back must MULTIPLY correctly. - let q = quantize_weight(QuantPolicy::Q8, t.clone()); - assert!( - matches!(q.dtype(), DType::QFloat(_)), - "weight should be quantized" - ); - let ready = compute_weight(&Param::from_tensor(q)); - - let x = Tensor::<2>::from_data( - TensorData::new(vec![0.25f32; 32], [1, 32]), - (&device, crate::backend::float_dtype(&device)), - ); - let want = x.clone().matmul(t).into_data().to_vec::().unwrap(); - let got = x.matmul(ready).into_data().to_vec::().unwrap(); - let scale = want.iter().map(|v| v.abs()).fold(1e-3, f32::max); - for (i, (a, b)) in want.iter().zip(&got).enumerate() { - assert!( - (a - b).abs() <= 0.05 * scale, - "elem {i}: quantized path gave {b}, f32 gave {a}" - ); - } - } - - /// On wgpu a quantized weight must come back float from - /// `compute_weight`: the native q_matmul panics on shapes the tier - /// planner routinely produces (m=1 decode x some group widths), so the - /// packed tensor must never reach `matmul`. Storage stays quantized; - /// only the multiply dequantizes. Skips without a wgpu device. - #[test] - fn wgpu_weights_are_dequantized_before_matmul() { - if !crate::backend::inventory().has_gpu() { - return; - } - let device = crate::backend::gpu_device(); - if !is_wgpu(&device) { - return; - } - let vals: Vec = (0..64 * 64).map(|i| ((i as f32) * 0.03).sin()).collect(); - let t = Tensor::<2>::from_data( - burn::tensor::TensorData::new(vals, [64, 64]), - (&device, crate::backend::float_dtype(&device)), - ); - let q = crate::quant::quantize_weight(crate::quant::QuantPolicy::Q8, t); - assert!( - matches!(q.dtype(), DType::QFloat(_)), - "setup: weight quantized" - ); - let ready = compute_weight(&Param::from_tensor(q)); - assert!( - !matches!(ready.dtype(), DType::QFloat(_)), - "a packed weight must never reach matmul on wgpu" - ); - } - - fn input(t: usize, seed: f32, device: &Dev) -> Tensor<3> { - let data: Vec = (0..t * HIDDEN) - .map(|i| ((i as f32 + seed) * 0.9).sin()) - .collect(); - Tensor::<1>::from_data(TensorData::new(data, [t * HIDDEN]), device).reshape([1, t, HIDDEN]) - } - - fn silu(x: f32) -> f32 { - x / (1.0 + (-x).exp()) - } - - /// Hand-rolled f32 reference of the whole block (per token: router - /// softmax, top-k, sparse weighted sum of per-expert SwiGLUs). - fn reference(m: &SparseMoe, x: &[f32], t: usize, norm: bool) -> Vec { - let rw = m.gate.weight.val().into_data().to_vec::().unwrap(); // [h, e] - let gw = m.experts.gate.val().into_data().to_vec::().unwrap(); // [e, inter, h] - let uw = m.experts.up.val().into_data().to_vec::().unwrap(); - let dw = m.experts.down.val().into_data().to_vec::().unwrap(); // [e, h, inter] - let mut out = vec![0f32; t * HIDDEN]; - for tok in 0..t { - let xrow = &x[tok * HIDDEN..][..HIDDEN]; - // Router logits then softmax over all experts. - let mut logits = [0f32; EXPERTS]; - for (e, logit) in logits.iter_mut().enumerate() { - *logit = (0..HIDDEN).map(|i| xrow[i] * rw[i * EXPERTS + e]).sum(); - } - let max = logits.iter().cloned().fold(f32::MIN, f32::max); - let exps: Vec = logits.iter().map(|l| (l - max).exp()).collect(); - let z: f32 = exps.iter().sum(); - let probs: Vec = exps.iter().map(|v| v / z).collect(); - // Top-k expert ids by probability. - let mut order: Vec = (0..EXPERTS).collect(); - order.sort_by(|&a, &b| probs[b].partial_cmp(&probs[a]).unwrap()); - let picked = &order[..TOP_K]; - let denom: f32 = if norm { - picked.iter().map(|&e| probs[e]).sum() - } else { - 1.0 - }; - for &e in picked { - let w = probs[e] / denom; - // SwiGLU of expert e. - let mut act = [0f32; INTER]; - for (j, a) in act.iter_mut().enumerate() { - let g: f32 = (0..HIDDEN) - .map(|i| xrow[i] * gw[(e * INTER + j) * HIDDEN + i]) - .sum(); - let u: f32 = (0..HIDDEN) - .map(|i| xrow[i] * uw[(e * INTER + j) * HIDDEN + i]) - .sum(); - *a = silu(g) * u; - } - for i in 0..HIDDEN { - let d: f32 = (0..INTER) - .map(|j| act[j] * dw[(e * HIDDEN + i) * INTER + j]) - .sum(); - out[tok * HIDDEN + i] += w * d; - } - } - } - out - } - - #[test] - fn forward_matches_the_hand_rolled_sparse_reference() { - let device = crate::backend::cpu_device(); - let m = moe(&device); - // Both shapes the decoder actually runs: a multi-token prefill and - // the single-token decode step. - for t in [5, 1] { - for norm in [false, true] { - let x = input(t, 2.0, &device); - let xv = x.clone().into_data().to_vec::().unwrap(); - let got = m - .forward(x, TOP_K, norm) - .into_data() - .to_vec::() - .unwrap(); - let want = reference(&m, &xv, t, norm); - assert_eq!(got.len(), want.len()); - for (i, (g, w)) in got.iter().zip(&want).enumerate() { - assert!( - (g - w).abs() < 1e-5, - "t={t} norm={norm} elem {i}: got {g} vs reference {w}" - ); - } - } - } - } - - #[test] - fn top_k_equal_to_num_experts_uses_every_expert() { - // With k == E and renorm the block degenerates to a full softmax - // mixture — the reference covers it; this pins the k=E edge. - let device = crate::backend::cpu_device(); - let m = moe(&device); - let x = input(3, 0.5, &device); - let xv = x.clone().into_data().to_vec::().unwrap(); - let got = m - .forward(x, EXPERTS, false) - .into_data() - .to_vec::() - .unwrap(); - // Reference with TOP_K replaced by all experts: weights are the full - // softmax row, every expert contributes. - let mut want = vec![0f32; 3 * HIDDEN]; - { - let full = reference_all_experts(&m, &xv, 3); - want.copy_from_slice(&full); - } - for (i, (g, w)) in got.iter().zip(&want).enumerate() { - assert!( - (g - w).abs() < 1e-5, - "elem {i}: got {g} vs full-mixture {w}" - ); - } - } - - /// Full-mixture reference (every expert, softmax-weighted) for the k=E edge. - fn reference_all_experts(m: &SparseMoe, x: &[f32], t: usize) -> Vec { - let rw = m.gate.weight.val().into_data().to_vec::().unwrap(); - let gw = m.experts.gate.val().into_data().to_vec::().unwrap(); - let uw = m.experts.up.val().into_data().to_vec::().unwrap(); - let dw = m.experts.down.val().into_data().to_vec::().unwrap(); - let mut out = vec![0f32; t * HIDDEN]; - for tok in 0..t { - let xrow = &x[tok * HIDDEN..][..HIDDEN]; - let mut logits = [0f32; EXPERTS]; - for (e, logit) in logits.iter_mut().enumerate() { - *logit = (0..HIDDEN).map(|i| xrow[i] * rw[i * EXPERTS + e]).sum(); - } - let max = logits.iter().cloned().fold(f32::MIN, f32::max); - let exps: Vec = logits.iter().map(|l| (l - max).exp()).collect(); - let z: f32 = exps.iter().sum(); - for e in 0..EXPERTS { - let w = exps[e] / z; - let mut act = [0f32; INTER]; - for (j, a) in act.iter_mut().enumerate() { - let g: f32 = (0..HIDDEN) - .map(|i| xrow[i] * gw[(e * INTER + j) * HIDDEN + i]) - .sum(); - let u: f32 = (0..HIDDEN) - .map(|i| xrow[i] * uw[(e * INTER + j) * HIDDEN + i]) - .sum(); - *a = silu(g) * u; - } - for i in 0..HIDDEN { - let d: f32 = (0..INTER) - .map(|j| act[j] * dw[(e * HIDDEN + i) * INTER + j]) - .sum(); - out[tok * HIDDEN + i] += w * d; - } - } - } - out - } - - #[test] - fn norm_topk_weights_change_the_mixture() { - // norm_topk_prob renormalizes the k weights to sum 1 — unless the - // top-k already captured all the mass, outputs must differ. - let device = crate::backend::cpu_device(); - let m = moe(&device); - let x = input(4, 7.0, &device); - let a = m - .forward(x.clone(), TOP_K, false) - .into_data() - .to_vec::() - .unwrap(); - let b = m - .forward(x, TOP_K, true) - .into_data() - .to_vec::() - .unwrap(); - let differs = a.iter().zip(&b).any(|(x, y)| (x - y).abs() > 1e-7); - assert!(differs, "renormalized weights should scale the output"); - } - - #[test] - fn forward_is_position_independent() { - // MoE acts per-token: the same row in different positions/batches - // routes and computes identically. - let device = crate::backend::cpu_device(); - let m = moe(&device); - let row = input(1, 11.0, &device); - let double = Tensor::cat(vec![row.clone(), row.clone()], 1); - let s = m - .forward(row, TOP_K, false) - .into_data() - .to_vec::() - .unwrap(); - let d = m - .forward(double, TOP_K, false) - .into_data() - .to_vec::() - .unwrap(); - assert_eq!(s.as_slice(), &d[..HIDDEN]); - assert_eq!(s.as_slice(), &d[HIDDEN..]); - } - - #[test] - fn zero_input_gives_zero_output() { - // Bias-free SwiGLU experts map 0 to 0 regardless of routing. - let device = crate::backend::cpu_device(); - let m = moe(&device); - let x = Tensor::<3>::zeros([1, 2, HIDDEN], &device); - let out = m - .forward(x, TOP_K, false) - .into_data() - .to_vec::() - .unwrap(); - assert!(out.iter().all(|&v| v == 0.0)); - } - - #[test] - #[should_panic(expected = "top_k")] - fn forward_rejects_top_k_above_num_experts() { - let device = crate::backend::cpu_device(); - let m = moe(&device); - let x = input(1, 0.0, &device); - let _ = m.forward(x, EXPERTS + 1, false); - } - - #[test] - #[should_panic(expected = "num_experts_per_tok")] - fn config_rejects_zero_top_k() { - let device = crate::backend::cpu_device(); - let _ = SparseMoeConfig { - hidden_size: 4, - expert_intermediate_size: 3, - num_experts: 4, - num_experts_per_tok: 0, - } - .init(&device); - } -} diff --git a/crates/mummu/examples/src/nn/packed_gemv.rs b/crates/mummu/examples/src/nn/packed_gemv.rs deleted file mode 100644 index 55c773d..0000000 --- a/crates/mummu/examples/src/nn/packed_gemv.rs +++ /dev/null @@ -1,588 +0,0 @@ -//! Packed m=1 GEMV for block-quantized weights — the kernel the whole -//! scheduler hunt pointed at. -//! -//! Burn's own quantized matmul at m=1 lands on a documented "extremely -//! hacky fix" (burn-cubecl `kernel/matmul/base.rs`): dequantize the ENTIRE -//! weight to f32, then float-matmul — 4x the traffic of the packed bytes, -//! a transient f32 weight allocation per call (the VRAM pool churn), and -//! the measured 0.91 ms/cluster against the host slab's 0.40. This module -//! reads the packed representation directly on every backend that holds -//! clusters: -//! -//! - **wgpu/CUDA** (`CubeBackend`): a `#[cube]` kernel where each unit -//! owns one u32 word (8 Q4S nibbles) per k-step — one packed word, one -//! shared f32 scale, one x[k] broadcast, eight fused mul-adds, no -//! cross-unit reduction. Weight traffic is the packed bytes, nothing -//! else. -//! - **flex** (host): a threaded i8 GEMV over the backend's i8-unpacked -//! storage (flex stores Q4 as one i8 per element) — 1.125 B/elem of -//! traffic instead of the 4 B/elem of the f32 slab it can replace. -//! -//! The op is exact with respect to the stored quantization: it computes -//! `y[n] = Σ_k x[k] · scale(k, n/32) · q(k, n)` in f32, the same math the -//! dequantize-then-matmul reference performs, so parity against that -//! reference is a summation-order question only (tested below). -//! -//! Layout contract (asserted): `QuantValue::Q4S`, `QuantLevel::block([32])`, -//! `QuantParam::F32`, `QuantStore::PackedU32(0)` — the only format mummu -//! puts on a device (pack.rs `quantized_tensor_data`). Weights are -//! `[K, N]` row-major with blocks along N; blocks never straddle rows. - -use burn::backend::backend_extension; -use burn::backend::tensor::{FloatTensor, QuantizedTensor}; -use burn::backend::{Backend, Flex, Wgpu}; -#[cfg(feature = "cuda")] -use burn::backend::Cuda; -#[cfg(feature = "vulkan-spirv")] -use burn::backend::Vulkan; -use burn::tensor::{DType, Tensor}; -use burn::tensor::quantization::{QuantLevel, QuantParam, QuantScheme, QuantStore, QuantValue}; - -/// Is a weight tensor in the one packed format this module reads? -fn scheme_supported(scheme: &QuantScheme) -> bool { - matches!(scheme.value, QuantValue::Q4S | QuantValue::Q8S) - && matches!(scheme.level, QuantLevel::Block(b) if b.to_dim_vec(1) == [32]) - && matches!(scheme.param, QuantParam::F32) - // PackedU32(0) is what accelerators hold; flex re-tags Native after - // unpacking to i8 — both are exactly what the per-backend impls read. - && matches!(scheme.store, QuantStore::PackedU32(0) | QuantStore::Native) -} - -/// Whether the packed path is enabled (`MUMMU_PACKED_GEMV`, default on — -/// `0`/`off`/`false` falls back to burn's dequantize-first matmul -/// everywhere, matching the repo's other default-on switches). -pub fn packed_gemv_enabled() -> bool { - static ON: std::sync::OnceLock = std::sync::OnceLock::new(); - *ON.get_or_init(|| { - std::env::var("MUMMU_PACKED_GEMV").map_or(true, |v| { - !(v == "0" || v.eq_ignore_ascii_case("off") || v.eq_ignore_ascii_case("false")) - }) - }) -} - -/// Route one decode-shape matmul through the packed GEMV when everything -/// lines up (m=1, Q4S block-32, path enabled); `None` means the caller -/// should use its existing matmul. -pub fn try_q4s_gemv(x: &Tensor<2>, w: &Tensor<2>) -> Option> { - if !packed_gemv_enabled() { - return None; - } - let DType::QFloat(scheme) = w.dtype() else { - return None; - }; - if !scheme_supported(&scheme) { - return None; - } - let m = x.dims()[0]; - if m == 1 { - let y = ::q4s_gemv( - x.clone().into_dispatch(), - w.clone().into_dispatch(), - ); - return Some(Tensor::from_dispatch(y)); - } - // Prefill: row-by-row through the same exact op. Slower per element - // than a real GEMM, but it never materializes the f32 weight — the - // dequantize-first fallback's ~260 MB transient per matmul was the - // VRAM-pool churn during prefill. Capped so pathological batch shapes - // keep the old path. - if m <= 64 { - let rows: Vec> = (0..m) - .map(|r| { - let xr = x.clone().slice([r..r + 1, 0..x.dims()[1]]); - let y = ::q4s_gemv( - xr.into_dispatch(), - w.clone().into_dispatch(), - ); - Tensor::from_dispatch(y) - }) - .collect(); - return Some(Tensor::cat(rows, 0)); - } - None -} - -#[backend_extension( - Flex, - Wgpu, - Vulkan: cfg(feature = "vulkan-spirv"), - Cuda: cfg(feature = "cuda"), -)] -/// The packed-GEMV extension op. `x` is `[1, K]` f32, `w` is `[K, N]` -/// QFloat (Q4S, block-32, f32 scales, PackedU32(0)); returns `[1, N]` f32. -pub trait Q4GemvOps: Backend { - /// y = x · w, reading w's packed values and scales directly. - fn q4s_gemv(x: FloatTensor, w: QuantizedTensor) -> FloatTensor; -} - -// --------------------------------------------------------------------------- -// CubeCL backends (wgpu / vulkan / cuda), non-fusion primitive level. -// --------------------------------------------------------------------------- - -mod cube_impl { - use super::Q4GemvOps; - use burn::backend::tensor::{FloatTensor, QuantizedTensor}; - use burn::tensor::Shape; - use burn_cubecl::{ - CubeBackend, CubeRuntime, kernel::into_contiguous, ops::numeric::empty_device, - tensor::CubeTensor, - }; - use cubecl::prelude::*; - - /// Split-K packed GEMV. - /// - /// A workgroup is 32 word-columns wide (UNIT_POS_X; coalesced — at any - /// k the 32 lanes read 32 consecutive u32 words) by `split` k-slices - /// deep (UNIT_POS_Y). Each thread accumulates its word's `per_word` - /// outputs over k_len/split steps in registers; partials meet in shared - /// memory (32 * split * per_word f32 — 16 KiB at split 16, Q4) and the - /// slice-0 threads reduce and store. - /// - /// Why: the split-1 shape gave gate/up 2176 threads on a card with 8448 - /// cores, each walking 5120 dependent FMAs — measured 0.6-0.7x against - /// burn's dequantize path. `split` multiplies resident threads and - /// divides the dependent chain by the same factor; what it buys costs - /// one barrier plus a strided sum, orders of magnitude below the chain - /// it removes. - /// - /// The barrier sits OUTSIDE the validity guard: in a partial workgroup - /// (n_words not a multiple of 32 — test shapes, never the 27B's) every - /// thread must still reach sync_cube, so invalid lanes contribute zero - /// partials and skip only the final store. - #[cube(launch)] - fn packed_gemv_kernel( - w_packed: &Tensor, - scales: &Tensor, - x: &Tensor, - out: &mut Tensor, - #[comptime] per_word: usize, - #[comptime] split: usize, - ) { - let lane = usize::cast_from(UNIT_POS_X); - let slice = usize::cast_from(UNIT_POS_Y); - let wc = usize::cast_from(CUBE_POS_X) * 32 + lane; - let n_words = w_packed.shape(1); - let valid = wc < n_words; - - let k_len = x.shape(1); - let bits = comptime!(32u32 / per_word as u32); - let mask = comptime!((1u32 << (32u32 / per_word as u32)) - 1); - let sign = comptime!(1u32 << (32u32 / per_word as u32 - 1)); - let span = comptime!(1i32 << (32u32 / per_word as u32)); - - let mut acc = Array::::new(per_word); - #[unroll] - for j in 0..per_word { - acc[j] = 0.0f32; - } - if valid { - let stride_w = w_packed.stride(0); - let stride_s = scales.stride(0); - // 32 values per scale block, `per_word` per u32 word, and words - // never straddle a block — every value this thread decodes at a - // given k shares one scale. - let sc_col = wc / comptime!(32 / per_word); - let k_per = k_len / split; - let k0 = slice * k_per; - for kk in 0..k_per { - let k = k0 + kk; - let word = w_packed[k * stride_w + wc]; - let xs = x[k] * scales[k * stride_s + sc_col]; - #[unroll] - for j in 0..per_word { - let raw = (word >> (u32::cast_from(j) * bits)) & mask; - let mut q = i32::cast_from(raw); - if raw >= sign { - q -= span; - } - acc[j] += f32::cast_from(q) * xs; - } - } - } - if comptime!(split == 1) { - if valid { - let base = wc * per_word; - #[unroll] - for j in 0..per_word { - out[base + j] = acc[j]; - } - } - } else { - let mut partials = Shared::<[f32]>::new_slice(comptime!(32 * split * per_word)); - let slot = (lane * split + slice) * per_word; - #[unroll] - for j in 0..per_word { - partials[slot + j] = acc[j]; - } - sync_cube(); - if valid && slice == 0 { - let base = wc * per_word; - #[unroll] - for j in 0..per_word { - let mut total = 0.0f32; - for ss in 0..split { - total += partials[(lane * split + ss) * per_word + j]; - } - out[base + j] = total; - } - } - } - } - - /// Split-K factor (`MUMMU_GEMV_SPLIT`, default 16, max 32: the shared - /// partial buffer is 32 * split * per_word f32 — 16 KiB at (16, Q4)). - fn gemv_split_override() -> Option { - static S: std::sync::OnceLock> = std::sync::OnceLock::new(); - *S.get_or_init(|| { - std::env::var("MUMMU_GEMV_SPLIT") - .ok() - .and_then(|v| v.parse::().ok()) - .filter(|&v| (1..=32).contains(&v)) - }) - } - - /// Candidate split factors, ordered by a CAPACITY prior. - /// - /// The prior is residency, not occupancy: a weight that fits the card's - /// L2 is not starved for memory parallelism, so splitting it only adds - /// partial-sum traffic and shared-memory pressure. A weight that misses - /// L2 streams from DRAM, where extra outstanding transactions do buy - /// bandwidth. `l2_bytes` is the usable fraction (~0.75) of the device's - /// L2; when it is unknown the prior degrades to "try everything", which - /// is exactly what the measurement then resolves. - /// - /// This ORDERS the search. It never decides. The previous revision - /// decided — it shipped `split = 16` derived from a Little's-law - /// occupancy argument, and the card answered that `gate/up` is flat - /// across every split while `down` gains 2.3x. Two independent peer - /// reviews made the same derivation. A rule three parties derive and - /// the hardware refuses is not a rule. - fn split_candidates(weight_bytes: u64, l2_bytes: Option) -> Vec { - let resident = l2_bytes.is_some_and(|l2| weight_bytes <= l2); - if resident { - // L2-resident: 1 first, and only modest splits are worth a try. - vec![1, 2, 4, 8] - } else { - vec![1, 4, 8, 16, 32] - } - } - - /// The measured best split for one (device, shape, packing), cached for - /// the process. - /// - /// Autotune rather than arithmetic, because the arithmetic was wrong and - /// because the published numbers for these exact shapes disagree across - /// GPU generations. Each entry costs a handful of warm launches, once, - /// and is keyed so a different card or a different projection shape gets - /// its own answer. - fn gemv_split_for( - client: &ComputeClient, - device: &R::Device, - w_vals: &CubeTensor, - w_scales: &CubeTensor, - x: &CubeTensor, - n: usize, - n_words: usize, - k_len: usize, - per_word: usize, - ) -> usize { - if let Some(forced) = gemv_split_override() { - return forced; - } - static CACHE: std::sync::OnceLock< - std::sync::Mutex>, - > = std::sync::OnceLock::new(); - let cache = CACHE.get_or_init(|| std::sync::Mutex::new(std::collections::HashMap::new())); - let key = (format!("{device:?}"), n, k_len, per_word); - if let Some(&hit) = cache.lock().unwrap_or_else(|e| e.into_inner()).get(&key) { - return hit; - } - - let props = client.properties(); - // Usable L2 for a streaming weight; 3/4 is the common working - // fraction once the activation and output tiles are accounted for. - let l2 = u64::from(props.hardware.max_shared_memory_size as u32) - .checked_mul(0) - .and(None::) - .or_else(|| std::env::var("MUMMU_L2_MIB").ok().and_then(|v| v.parse::().ok()).map(|m| m << 20)) - .map(|b| b / 4 * 3); - let weight_bytes = (n_words as u64 * k_len as u64 * 4) + (w_scales.meta.num_elements() as u64 * 4); - - let mut best = (1usize, f64::INFINITY); - for cand in split_candidates(weight_bytes, l2) { - if k_len % cand != 0 { - continue; - } - let run = || { - let out = empty_device::( - client.clone(), - x.device.clone(), - Shape::new([1, n]), - ); - packed_gemv_kernel::launch::( - client, - CubeCount::Static((n_words as u32).div_ceil(32), 1, 1), - CubeDim { x: 32, y: cand as u32, z: 1 }, - w_vals.clone().into_tensor_arg(), - w_scales.clone().into_tensor_arg(), - x.clone().into_tensor_arg(), - out.clone().into_tensor_arg(), - per_word, - cand, - ); - out - }; - // Warm (kernel compile, autotune, pool), then time. `sync` - // blocks until the device drains, which is what makes these - // numbers comparable — a launch alone returns immediately. - run(); - let _ = cubecl::future::block_on(client.sync()); - let t0 = std::time::Instant::now(); - for _ in 0..3 { - run(); - } - let _ = cubecl::future::block_on(client.sync()); - let ms = t0.elapsed().as_secs_f64() * 1e3 / 3.0; - if ms < best.1 { - best = (cand, ms); - } - } - eprintln!( - "[mummu] gemv split for [{k_len} x {n}] (per_word {per_word}): {} at {:.3} ms", - best.0, best.1 - ); - cache - .lock() - .unwrap_or_else(|e| e.into_inner()) - .insert(key, best.0); - best.0 - } - - pub(super) fn q4s_gemv_cube( - x: CubeTensor, - w: CubeTensor, - ) -> CubeTensor { - // Values per u32 word, straight off the scheme: 8 nibbles for Q4S, - // 4 bytes for Q8S. Derived, not assumed — the values view's own - // shape below must agree with it. - let per_word: usize = match w.dtype { - burn::tensor::DType::QFloat(scheme) => 32 / scheme.value.size_bits() as usize, - other => unreachable!("packed gemv on a non-quantized weight: {other:?}"), - }; - let x = into_contiguous(x); - let (w_vals, w_scales) = w - .quantized_handles() - .expect("q4s_gemv: weight must be a quantized CubeTensor"); - let n = w.meta.shape()[1]; - let n_words = w_vals.meta.shape()[1]; - let client = x.client.clone(); - let device = x.device.clone(); - let out = empty_device::(client.clone(), device, Shape::new([1, n])); - // Split-K factor: `MUMMU_GEMV_SPLIT`, default 16 — near the - // occupancy knee for both production shapes. Forced to 1 when it - // does not divide k, so odd shapes keep the exact split-1 path. - let k_len = x.meta.shape()[1]; - let mut split = gemv_split_for( - &client, - &x.device, - &w_vals, - &w_scales, - &x, - n, - n_words, - k_len, - per_word, - ); - if split == 0 || k_len % split != 0 { - split = 1; - } - let cube_dim = CubeDim { - x: 32, - y: split as u32, - z: 1, - }; - let cubes = (n_words as u32).div_ceil(32); - debug_assert_eq!( - n_words * per_word, - n, - "packed values view must cover exactly the logical width" - ); - packed_gemv_kernel::launch::( - &client, - CubeCount::Static(cubes, 1, 1), - cube_dim, - w_vals.into_tensor_arg(), - w_scales.into_tensor_arg(), - x.into_tensor_arg(), - out.clone().into_tensor_arg(), - per_word, - split, - ); - out - } - - impl Q4GemvOps for CubeBackend { - fn q4s_gemv(x: FloatTensor, w: QuantizedTensor) -> FloatTensor { - q4s_gemv_cube::(x, w) - } - } -} - -// --------------------------------------------------------------------------- -// Flex (host): threaded i8 GEMV over the backend's unpacked storage. -// --------------------------------------------------------------------------- - -mod flex_impl { - use super::Q4GemvOps; - use burn::backend::tensor::{FloatTensor, QuantizedTensor}; - use burn::backend::{Flex, TensorMetadata}; - use burn::tensor::TensorData; - use burn_flex::FlexTensor; - - impl Q4GemvOps for Flex { - fn q4s_gemv(x: FloatTensor, w: QuantizedTensor) -> FloatTensor { - let shape = w.shape(); - let [k_len, n] = shape.dims::<2>(); - let xs_owned; - let xs: &[f32] = match x.as_slice::() { - Some(s) => s, - None => { - xs_owned = x.into_data().to_vec::().expect("f32 activations"); - &xs_owned - } - }; - let wq: &[i8] = w - .tensor() - .as_slice::() - .expect("flex quantized weight is contiguous i8"); - let scales: &[f32] = w.scales(); - let blocks = n / 32; - let mut out = vec![0f32; n]; - // Rayon par-chunks: persistent pool threads (a scope-spawn per - // call measured its cost — mlp stayed at f32-slab speed), each - // owning a disjoint 32-aligned output range and walking every - // k. The k-inner loop is written so LLVM vectorizes the - // i8-widen + FMA. - use rayon::prelude::*; - let chunk_blocks = blocks.div_ceil(rayon::current_num_threads().max(1)).max(4); - out.par_chunks_mut(chunk_blocks * 32) - .enumerate() - .for_each(|(t, chunk)| { - let b0 = t * chunk_blocks; - let nb = chunk.len() / 32; - for k in 0..k_len { - let xk = xs[k]; - let row = &wq[k * n + b0 * 32..k * n + b0 * 32 + nb * 32]; - let srow = &scales[k * blocks + b0..k * blocks + b0 + nb]; - for b in 0..nb { - let xs_s = xk * srow[b]; - let src = &row[b * 32..b * 32 + 32]; - let dst = &mut chunk[b * 32..b * 32 + 32]; - for j in 0..32 { - dst[j] += xs_s * f32::from(src[j]); - } - } - } - }); - FlexTensor::from_data(TensorData::new(out, [1, n])) - } - } -} - -// --------------------------------------------------------------------------- -// Fusion wrapper: re-enter the stream with a custom op so `Fusion` -// backends (the default `Wgpu`) hand the inner CubeBackend real tensors. -// --------------------------------------------------------------------------- - -#[cfg(feature = "fusion")] -mod fusion_impl { - use super::Q4GemvOps; - use burn::backend::tensor::{FloatTensor, QuantizedTensor}; - use burn::tensor::{DType, Shape}; - use burn_fusion::{ - Fusion, FusionBackend, FusionRuntime, - stream::{Operation, StreamId}, - }; - use burn_ir::{CustomOpIr, HandleContainer, OperationIr, OperationOutput, TensorIr}; - - impl Q4GemvOps for Fusion { - fn q4s_gemv(x: FloatTensor, w: QuantizedTensor) -> FloatTensor { - let client = x.client.clone(); - let shape_out = Shape::new([x.shape[0], w.shape[1]]); - - #[derive(derive_new::new, Clone, Debug)] - struct Gemv { - desc: CustomOpIr, - _b: core::marker::PhantomData, - } - impl Operation for Gemv { - fn execute( - &self, - handles: &mut HandleContainer< - ::FusionHandle, - >, - ) { - let ([x_ir, w_ir], [out_ir]) = self.desc.as_fixed(); - let xt = handles.get_float_tensor::(x_ir); - let wt = handles.get_quantized_tensor::(w_ir); - let y = B1::q4s_gemv(xt, wt); - handles.register_float_tensor::(&out_ir.id, y); - } - } - - let stream = StreamId::current(); - let out = TensorIr::uninit(client.create_empty_handle(), shape_out, DType::F32); - let desc = CustomOpIr::new("mummu_q4s_gemv", &[x.into_ir(), w.into_ir()], &[out]); - client - .register(stream, OperationIr::Custom(desc.clone()), Gemv::::new(desc)) - .output() - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use burn::tensor::{Distribution, Tensor}; - - /// The packed GEMV against the dequantize-then-matmul reference on the - /// host backend — same math, different summation order, so the bound - /// is tight. - #[test] - fn q4s_gemv_matches_dequant_matmul_on_flex() { - let device = crate::backend::cpu_device(); - let (k, n) = (192, 160); - let x = Tensor::<2>::random([1, k], Distribution::Uniform(-1.0, 1.0), &device); - let w = Tensor::<2>::random([k, n], Distribution::Uniform(-1.0, 1.0), &device); - let wq = crate::quant::quantize_weight(crate::quant::QuantPolicy::Q4, w); - assert!(matches!(wq.dtype(), burn::tensor::DType::QFloat(_))); - - let got = try_q4s_gemv(&x, &wq).expect("packed path must engage"); - let want = x.matmul(wq.dequantize()); - - let diff = got - .sub(want.clone()) - .abs() - .max() - .into_data() - .to_vec::() - .unwrap()[0]; - let scale = want.abs().max().into_data().to_vec::().unwrap()[0].max(1e-6); - assert!( - diff / scale < 1e-4, - "packed vs reference rel err {} (abs {diff})", - diff / scale - ); - } - - /// m != 1 and non-Q4S weights must decline, not compute. - #[test] - fn q4s_gemv_declines_out_of_contract() { - let device = crate::backend::cpu_device(); - let x2 = Tensor::<2>::random([2, 64], Distribution::Uniform(-1.0, 1.0), &device); - let wf = Tensor::<2>::random([64, 64], Distribution::Uniform(-1.0, 1.0), &device); - assert!(try_q4s_gemv(&x2, &wf).is_none(), "m=2 must decline"); - let x1 = Tensor::<2>::random([1, 64], Distribution::Uniform(-1.0, 1.0), &device); - assert!(try_q4s_gemv(&x1, &wf).is_none(), "float weight must decline"); - } -} diff --git a/crates/mummu/examples/src/nn/rope.rs b/crates/mummu/examples/src/nn/rope.rs deleted file mode 100644 index 9fdfc7f..0000000 --- a/crates/mummu/examples/src/nn/rope.rs +++ /dev/null @@ -1,150 +0,0 @@ -//! Rotary position embeddings (RoPE), HF duplicated-half layout, computed -//! manually so the same tables serve every architecture (Qwen theta 1e6 vs -//! LFM2 theta 1e6 vs others) and every backend. - -use burn::tensor::{Device, Tensor, TensorData}; - -use super::MAX_CONTEXT_TOKENS; - -/// RoPE cos/sin tables `[1, 1, t, head_dim]` for absolute positions -/// `past..past+t`, HF duplicated-half layout (each frequency written to both -/// halves of the head dim, so [`apply_rope`]'s rotate-half math lines up). -pub fn rope_tables( - t: usize, - past: usize, - head_dim: usize, - theta: f32, - device: &Device, -) -> (Tensor<4>, Tensor<4>) { - assert!(t >= 1, "rope_tables: need at least one position, got t=0"); - assert!( - head_dim >= 2 && head_dim.is_multiple_of(2), - "rope_tables: head_dim must be even and >= 2, got {head_dim}" - ); - assert!( - past + t <= MAX_CONTEXT_TOKENS, - "rope_tables: position {past}+{t} exceeds MAX_CONTEXT_TOKENS ({MAX_CONTEXT_TOKENS})" - ); - debug_assert!(theta > 0.0, "rope_tables: theta must be positive"); - - let half = head_dim / 2; - let mut cos = vec![0f32; t * head_dim]; - let mut sin = vec![0f32; t * head_dim]; - for i in 0..t { - let pos = (past + i) as f32; - for k in 0..half { - let inv = 1.0f32 / theta.powf(2.0 * k as f32 / head_dim as f32); - let (s, c) = (pos * inv).sin_cos(); - cos[i * head_dim + k] = c; - cos[i * head_dim + k + half] = c; - sin[i * head_dim + k] = s; - sin[i * head_dim + k + half] = s; - } - } - // The float dtype comes from the DEVICE — burn 0.22 keeps the element - // type there as a runtime setting, not on a backend type. Creation sites - // still name it explicitly rather than riding the unspecified default. - let dtype = crate::backend::float_dtype(device); - let cos = Tensor::<2>::from_data(TensorData::new(cos, [t, head_dim]), (device, dtype)) - .reshape([1, 1, t, head_dim]); - let sin = Tensor::<2>::from_data(TensorData::new(sin, [t, head_dim]), (device, dtype)) - .reshape([1, 1, t, head_dim]); - (cos, sin) -} - -/// `rotate_half`: split the last dim in two, return `cat([-x2, x1])`. -pub fn rotate_half(x: Tensor<4>) -> Tensor<4> { - let dims = x.dims(); - let half = dims[3] / 2; - assert!( - dims[3].is_multiple_of(2) && half >= 1, - "rotate_half: last dim must be even and >= 2, got {}", - dims[3] - ); - let x1 = x.clone().narrow(3, 0, half); - let x2 = x.narrow(3, half, half); - Tensor::cat(vec![x2.neg(), x1], 3) -} - -/// Apply RoPE: `x*cos + rotate_half(x)*sin`. `cos`/`sin` come from -/// [`rope_tables`] and must cover the same `t` and `head_dim` as `x`. -pub fn apply_rope(x: Tensor<4>, cos: &Tensor<4>, sin: &Tensor<4>) -> Tensor<4> { - let (xd, cd) = (x.dims(), cos.dims()); - assert!( - xd[2] == cd[2] && xd[3] == cd[3], - "apply_rope: x {xd:?} and tables {cd:?} disagree on [t, head_dim]" - ); - debug_assert!( - cos.dims() == sin.dims(), - "apply_rope: cos/sin shape mismatch" - ); - let rh = rotate_half(x.clone()); - x.mul(cos.clone()).add(rh.mul(sin.clone())) -} - -#[cfg(test)] -mod tests { - use super::*; - - fn to_vec(t: Tensor<4>) -> Vec { - t.into_data().to_vec::().unwrap() - } - - #[test] - fn rope_tables_position_zero_is_identity_rotation() { - let device = crate::backend::cpu_device(); - let (cos, sin) = rope_tables(1, 0, 8, 1e6, &device); - assert!(to_vec(cos).iter().all(|&c| (c - 1.0).abs() < 1e-7)); - assert!(to_vec(sin).iter().all(|&s| s.abs() < 1e-7)); - } - - #[test] - fn rope_tables_are_unit_norm_and_offset_consistent() { - let device = crate::backend::cpu_device(); - // cos^2 + sin^2 == 1 everywhere. - let (cos, sin) = rope_tables(3, 2, 16, 1e4, &device); - let (c, s) = (to_vec(cos), to_vec(sin)); - for (ci, si) in c.iter().zip(&s) { - assert!((ci * ci + si * si - 1.0).abs() < 1e-5); - } - // Row for absolute position 4 must match whether reached via past=2+i=2 - // (prefill) or past=4,t=1 (decode) — the KV-cache offset invariant. - let (cos_dec, sin_dec) = rope_tables(1, 4, 16, 1e4, &device); - let (cd, sd) = (to_vec(cos_dec), to_vec(sin_dec)); - assert_eq!(&c[2 * 16..3 * 16], &cd[..]); - assert_eq!(&s[2 * 16..3 * 16], &sd[..]); - } - - #[test] - fn rotate_half_swaps_and_negates() { - let device = crate::backend::cpu_device(); - let x = Tensor::<1>::from_floats([1.0, 2.0, 3.0, 4.0], &device).reshape([1, 1, 1, 4]); - assert_eq!(to_vec(rotate_half(x)), vec![-3.0, -4.0, 1.0, 2.0]); - } - - #[test] - fn apply_rope_at_position_zero_is_identity() { - let device = crate::backend::cpu_device(); - let x = Tensor::<1>::from_floats([0.5, -1.5, 2.0, 3.5], &device).reshape([1, 1, 1, 4]); - let (cos, sin) = rope_tables(1, 0, 4, 1e6, &device); - let y = apply_rope(x.clone(), &cos, &sin); - let (xv, yv) = (to_vec(x), to_vec(y)); - for (a, b) in xv.iter().zip(&yv) { - assert!((a - b).abs() < 1e-6); - } - } - - #[test] - #[should_panic(expected = "head_dim must be even")] - fn rope_tables_rejects_odd_head_dim() { - let device = crate::backend::cpu_device(); - let _ = rope_tables(1, 0, 7, 1e6, &device); - } - - #[test] - #[should_panic(expected = "MAX_CONTEXT_TOKENS")] - fn rope_tables_rejects_runaway_positions() { - let device = crate::backend::cpu_device(); - let _ = rope_tables(1, MAX_CONTEXT_TOKENS, 8, 1e6, &device); - } -} diff --git a/crates/mummu/examples/src/pack.rs b/crates/mummu/examples/src/pack.rs deleted file mode 100644 index 20e0d70..0000000 --- a/crates/mummu/examples/src/pack.rs +++ /dev/null @@ -1,1152 +0,0 @@ -//! The `.mummu` pack — mummu's own multi-precision model artifact (P9 -//! stage 3). Import converts a source checkpoint **once** into a directory -//! holding every tensor at every stored precision, so any device can pull -//! exactly the tensors and the precision it runs best — per MoE **expert** — -//! without re-importing, and the planner can re-tier at will. -//! -//! Layout of `.mummu/`: -//! -//! - `manifest.json` — version, source, per-tensor entries: pack name, -//! role, stored shape, and per-precision `{values, scales}` byte ranges. -//! - `header.gguf` — the source GGUF's header bytes (metadata + tensor table, -//! no payload): `GgufFile::open` parses it, so every config / tokenizer -//! reader the GGUF path already has works on a pack unchanged. -//! - `q4.bin`, `q8.bin`, `f16.bin`, `f32.bin` — one blob per stored -//! precision; each tensor's bytes are contiguous within it. -//! -//! Stored precisions stop at f32: weights born bf16/Q4 carry no more -//! information, so f64 is a *compute* option derived from f32 where a -//! backend supports it, not a storage level. Quantized levels use block-32 -//! symmetric scales (burn's `Q8S`/`Q4S` semantics) in **burn's canonical -//! quantized `TensorData` layout** (`TensorData::quantized`), which every -//! backend ingests via `q_from_data` — loading is a copy, never a -//! re-quantization. Linear weights are stored already transposed to burn's -//! `[in, out]`; expert banks are split into per-expert members on import. - -use std::collections::BTreeMap; -use std::io::{Read, Seek, SeekFrom, Write}; -use std::path::{Path, PathBuf}; - -use burn::tensor::quantization::QuantScheme; -use burn::tensor::{Device, Tensor, TensorData}; - -use crate::gguf::{GgufFile, GgufTensorInfo}; -use crate::quant::QuantPolicy; -// `scheme()` is an extension trait: the ladder lives in `mummu-mix`, which -// has no burn dependency, so the burn binding is bolted on here. -use crate::quant::SchemeExt; - -/// Pack format version (bump on any incompatible manifest/blob change). -pub const PACK_VERSION: u32 = 1; -/// Block width of the quantized levels (must match `QuantPolicy::scheme`). -pub const BLOCK: usize = 32; - -/// A stored precision level. -#[derive( - Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, serde::Serialize, serde::Deserialize, -)] -#[serde(rename_all = "lowercase")] -pub enum Precision { - Q4, - Q8, - F16, - F32, -} - -impl Precision { - pub const ALL: [Self; 4] = [Self::Q4, Self::Q8, Self::F16, Self::F32]; - - /// Blob file name inside the pack. - #[must_use] - pub fn blob_name(self) -> &'static str { - match self { - Self::Q4 => "q4.bin", - Self::Q8 => "q8.bin", - Self::F16 => "f16.bin", - Self::F32 => "f32.bin", - } - } - - /// The quant policy this level corresponds to (`Off` for floats). - #[must_use] - pub fn policy(self) -> QuantPolicy { - match self { - Self::Q4 => QuantPolicy::Q4, - Self::Q8 => QuantPolicy::Q8, - Self::F16 | Self::F32 => QuantPolicy::Off, - } - } - - /// Parse `q4,q8,f16,f32` lists (the `MUMMU_PACK_PRECISIONS` convention). - pub fn parse_list(s: &str) -> Result, String> { - let mut out = Vec::new(); - for item in s.split(',').map(str::trim).filter(|s| !s.is_empty()) { - out.push(match item.to_ascii_lowercase().as_str() { - "q4" | "int4" => Self::Q4, - "q8" | "int8" => Self::Q8, - "f16" | "half" => Self::F16, - "f32" | "float" => Self::F32, - other => return Err(format!("unknown precision {other:?}")), - }); - } - if out.is_empty() { - return Err("no precisions given".into()); - } - out.sort(); - out.dedup(); - Ok(out) - } -} - -/// What a tensor is to the model — what the loaders need to place it. -#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] -#[serde(tag = "kind", rename_all = "lowercase")] -pub enum Role { - /// A 2-D projection weight, stored `[in, out]` (burn Linear layout). - Linear, - /// The token embedding, stored `[vocab, hidden]`; never quantized. - Embedding, - /// A 1-D vector (norm gamma, bias, per-head scalars). - Vector, - /// A depthwise conv kernel `[channels, 1, k]`. - Conv, - /// One member of a split MoE expert bank, stored `[in, out]`. - Expert { - layer: usize, - index: usize, - proj: String, - }, -} - -/// Where one precision of one tensor lives inside its blob. -#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] -pub struct Blob { - pub values_offset: u64, - pub values_len: u64, - /// Zero-length for float levels. - pub scales_offset: u64, - pub scales_len: u64, -} - -/// One tensor in the pack. -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] -pub struct TensorEntry { - /// The pack name — the source tensor name, with `/e{index}` appended for - /// split expert members. Loaders map it to their module paths. - pub name: String, - pub role: Role, - /// Row-major shape AS STORED (linears already `[in, out]`). - pub shape: Vec, - pub precisions: BTreeMap, -} - -/// The pack manifest. -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] -pub struct Manifest { - pub version: u32, - pub source_file: String, - pub source_bytes: u64, - pub architecture: String, - pub precisions: Vec, - pub tensors: Vec, - /// P9 stage 3(c): the dense FFNs partitioned into neuron clusters (see - /// `crate::partition`). Absent on packs imported before partitioning. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub ffn_partition: Option, -} - -/// One contiguous cluster of a (permuted) FFN intermediate dimension. -#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] -pub struct ClusterSpan { - pub start: usize, - pub len: usize, -} - -/// One measured point of the skip trade-off: at energy threshold `tau`, -/// how far the skipped model strays from the exact one. -#[derive(Debug, Clone, Copy, PartialEq, serde::Serialize, serde::Deserialize)] -pub struct SkipPoint { - pub tau: f32, - /// Max |Δ log-prob| over the vocabulary at the measured positions. - pub max_delta_logprob: f32, - /// Fraction of measured positions whose argmax was unchanged. - pub argmax_agreement: f32, - /// Mean fraction of clusters actually computed. - pub kept_fraction: f32, -} - -/// The FFN partition of a dense model: per layer, the cluster spans of the -/// permuted intermediate dim and the three entry names; plus, once -/// calibrated, a hotness prior per cluster and the skip table. -#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)] -pub struct FfnPartition { - pub layers: Vec>, - /// `[gate, up, down]` pack names per layer. - pub names: Vec<[String; 3]>, - /// Per-layer, per-cluster activation energy share from calibration - /// (empty until `pack-calibrate` ran). - #[serde(default)] - pub hotness: Vec>, - /// Measured skip trade-off (empty until calibrated) — the planner may - /// only pick a `tau` that appears here. - #[serde(default)] - pub skip_table: Vec, -} - -/// How the importer should treat one source tensor. -#[derive(Debug, Clone)] -pub enum ImportAction { - /// Drop it (e.g. NextN/MTP blocks). - Skip, - /// A 2-D linear weight: transpose GGUF's `[out, in]` to `[in, out]`. - Linear, - /// The embedding table: keep `[vocab, hidden]`, float only. - Embedding, - /// A 1-D vector: float only. - Vector, - /// A squeezed depthwise conv kernel `[k, channels]` → `[channels, 1, k]`. - Conv, - /// A fused expert bank `[experts, out, in]`: split into members, each - /// stored `[in, out]` as `Role::Expert`. - ExpertBank { layer: usize, proj: String }, -} - -// --------------------------------------------------------------------------- -// Quantizer: block-32 symmetric, the same semantics as burn's Q8S/Q4S -// min-max calibration (scale = max|block| / range_max). These values + scales -// ARE the model at that level; burn reconstructs tensors from them verbatim. -// --------------------------------------------------------------------------- - -/// Quantize a row-major tensor whose LAST dim is a multiple of [`BLOCK`]: -/// returns (i8 values, f32 scale per block). Blocks run along rows — a -/// block never straddles two rows. -pub fn quantize_blocks( - values: &[f32], - last_dim: usize, - precision: Precision, -) -> (Vec, Vec) { - assert!( - last_dim.is_multiple_of(BLOCK) && values.len().is_multiple_of(last_dim), - "quantize_blocks: last dim {last_dim} must divide by {BLOCK} and the length" - ); - let range_max: f32 = match precision { - Precision::Q8 => 127.0, - Precision::Q4 => 7.0, - _ => panic!("quantize_blocks: {precision:?} is not a quantized level"), - }; - let mut q = Vec::with_capacity(values.len()); - let mut scales = Vec::with_capacity(values.len() / BLOCK); - for block in values.as_chunks::().0 { - let alpha = block.iter().fold(0.0f32, |m, &x| m.max(x.abs())); - let scale = if alpha > 0.0 { alpha / range_max } else { 1.0 }; - scales.push(scale); - let inv = 1.0 / scale; - q.extend( - block - .iter() - .map(|&x| (x * inv).round().clamp(-range_max, range_max) as i8), - ); - } - (q, scales) -} - -/// Pack i8 values holding 4-bit range into nibbles (two per byte, low -/// nibble first) for the on-disk Q4 blob. -fn pack_nibbles(values: &[i8]) -> Vec { - let mut out = Vec::with_capacity(values.len().div_ceil(2)); - for pair in values.chunks(2) { - let lo = (pair[0] as u8) & 0x0F; - let hi = pair.get(1).map_or(0, |&v| (v as u8) & 0x0F); - out.push(lo | (hi << 4)); - } - out -} - -/// Inverse of [`pack_nibbles`]: sign-extend each nibble back to i8. -fn unpack_nibbles(bytes: &[u8], n: usize) -> Vec { - let mut out = Vec::with_capacity(n); - for &b in bytes { - for nib in [b & 0x0F, b >> 4] { - if out.len() == n { - break; - } - // 4-bit two's complement sign extension. - out.push(if nib & 0x8 != 0 { - (nib | 0xF0) as i8 - } else { - nib as i8 - }); - } - } - out -} - -// --------------------------------------------------------------------------- -// Writer / importer -// --------------------------------------------------------------------------- - -struct BlobWriter { - file: std::io::BufWriter, - len: u64, - /// Bytes written since the last `sync_data` — a multi-hundred-GB import - /// must not leave that much dirty page cache behind (a Docker VM's OOM - /// killer counts it against the process long before writeback catches - /// up on a bind mount), so blobs are synced every [`SYNC_EVERY`] bytes. - unsynced: u64, -} - -/// Dirty-bytes bound per blob between `sync_data` calls. -const SYNC_EVERY: u64 = 1 << 30; - -impl BlobWriter { - fn append(&mut self, bytes: &[u8]) -> std::io::Result<(u64, u64)> { - let off = self.len; - self.file.write_all(bytes)?; - self.len += bytes.len() as u64; - self.unsynced += bytes.len() as u64; - if self.unsynced >= SYNC_EVERY { - self.file.flush()?; - self.file.get_ref().sync_data()?; - self.unsynced = 0; - } - Ok((off, bytes.len() as u64)) - } -} - -/// Import a GGUF into a pack at `out_dir` with the given stored precisions. -/// `map` classifies every source tensor; the importer reads each tensor -/// once (dequantizing whatever the source stored), lays it out for the -/// loaders, and writes every requested level. Float-only roles (embedding, -/// vectors, convs) get only the float levels requested (f16/f32; if neither -/// was requested, f32 is added for them). Quantized levels apply only where -/// `QuantPolicy::eligible` would. -pub fn import_gguf( - gguf_path: &Path, - out_dir: &Path, - precisions: &[Precision], - map: &dyn Fn(&GgufTensorInfo) -> Option, - mut on_progress: impl FnMut(usize, usize, &str), -) -> Result { - let f = GgufFile::open(gguf_path).map_err(|e| e.to_string())?; - std::fs::create_dir_all(out_dir).map_err(|e| format!("create {}: {e}", out_dir.display()))?; - - // Header copy: the first `data_offset` bytes are exactly metadata + - // tensor table — a valid payload-less GGUF for every header reader. - { - let mut src = std::fs::File::open(gguf_path).map_err(|e| e.to_string())?; - let mut header = vec![0u8; usize::try_from(f.data_offset).expect("header fits")]; - src.read_exact(&mut header) - .map_err(|e| format!("read header: {e}"))?; - std::fs::write(out_dir.join("header.gguf"), &header).map_err(|e| e.to_string())?; - } - - let mut precisions: Vec = precisions.to_vec(); - precisions.sort(); - precisions.dedup(); - let float_levels: Vec = { - let mut v: Vec = precisions - .iter() - .copied() - .filter(|p| matches!(p, Precision::F16 | Precision::F32)) - .collect(); - if v.is_empty() { - v.push(Precision::F32); - } - v - }; - let mut all_levels = precisions.clone(); - for p in &float_levels { - if !all_levels.contains(p) { - all_levels.push(*p); - } - } - all_levels.sort(); - - let mut writers: BTreeMap = BTreeMap::new(); - for &p in &all_levels { - let file = std::fs::File::create(out_dir.join(p.blob_name())).map_err(|e| e.to_string())?; - writers.insert( - p, - BlobWriter { - file: std::io::BufWriter::with_capacity(8 << 20, file), - len: 0, - unsynced: 0, - }, - ); - } - - let mut entries: Vec = Vec::new(); - let total = f.tensors.len(); - for (i, info) in f.tensors.iter().enumerate() { - on_progress(i, total, &info.name); - let action = map(info).ok_or_else(|| format!("unmapped tensor name '{}'", info.name))?; - if matches!(action, ImportAction::Skip) { - continue; - } - let dims_rev: Vec = info.dims.iter().rev().map(|&d| d as usize).collect(); - let values = f.read_tensor_f32(&info.name).map_err(|e| e.to_string())?; - - // Materialize the stored layout(s): one or many (expert bank) tensors. - let mut items: Vec<(String, Role, Vec, Vec)> = Vec::new(); - match action { - ImportAction::Skip => unreachable!(), - ImportAction::Linear => { - let &[out, inp] = dims_rev.as_slice() else { - return Err(format!( - "'{}' linear must be 2-D, got {dims_rev:?}", - info.name - )); - }; - items.push(( - info.name.clone(), - Role::Linear, - vec![inp, out], - transpose(&values, out, inp), - )); - } - ImportAction::Embedding => { - items.push((info.name.clone(), Role::Embedding, dims_rev.clone(), values)); - } - ImportAction::Vector => { - items.push((info.name.clone(), Role::Vector, dims_rev.clone(), values)); - } - ImportAction::Conv => { - // ggml ne = [k, ch] ⇒ row-major [ch, k] == checkpoint [ch, 1, k] bytes. - let &[ch, k] = dims_rev.as_slice() else { - return Err(format!( - "'{}' conv must be 2-D squeezed, got {dims_rev:?}", - info.name - )); - }; - items.push((info.name.clone(), Role::Conv, vec![ch, 1, k], values)); - } - ImportAction::ExpertBank { layer, proj } => { - let &[e, out, inp] = dims_rev.as_slice() else { - return Err(format!( - "'{}' expert bank must be 3-D, got {dims_rev:?}", - info.name - )); - }; - let stride = out * inp; - for expert in 0..e { - let member = &values[expert * stride..(expert + 1) * stride]; - items.push(( - format!("{}/e{expert}", info.name), - Role::Expert { - layer, - index: expert, - proj: proj.clone(), - }, - vec![inp, out], - transpose(member, out, inp), - )); - } - } - } - - for (name, role, shape, data) in items { - let quantizable = matches!(role, Role::Linear | Role::Expert { .. }) - && QuantPolicy::Q8.eligible(&shape); - let mut per: BTreeMap = BTreeMap::new(); - for &p in &all_levels { - let stored = match p { - Precision::F32 | Precision::F16 => { - float_levels.contains(&p) || !quantizable && p == Precision::F32 - } - Precision::Q4 | Precision::Q8 => quantizable && precisions.contains(&p), - }; - if !stored { - continue; - } - let w = writers.get_mut(&p).expect("writer exists"); - let blob = match p { - Precision::F32 => { - let bytes: Vec = data.iter().flat_map(|v| v.to_le_bytes()).collect(); - let (o, l) = w.append(&bytes).map_err(|e| e.to_string())?; - Blob { - values_offset: o, - values_len: l, - scales_offset: 0, - scales_len: 0, - } - } - Precision::F16 => { - let bytes: Vec = data - .iter() - .flat_map(|&v| half::f16::from_f32(v).to_le_bytes()) - .collect(); - let (o, l) = w.append(&bytes).map_err(|e| e.to_string())?; - Blob { - values_offset: o, - values_len: l, - scales_offset: 0, - scales_len: 0, - } - } - Precision::Q8 | Precision::Q4 => { - let last = *shape.last().expect("non-empty shape"); - let (q, scales) = quantize_blocks(&data, last, p); - let vbytes: Vec = if p == Precision::Q4 { - pack_nibbles(&q) - } else { - q.iter().map(|&v| v as u8).collect() - }; - let (vo, vl) = w.append(&vbytes).map_err(|e| e.to_string())?; - let sbytes: Vec = scales.iter().flat_map(|s| s.to_le_bytes()).collect(); - let (so, sl) = w.append(&sbytes).map_err(|e| e.to_string())?; - Blob { - values_offset: vo, - values_len: vl, - scales_offset: so, - scales_len: sl, - } - } - }; - per.insert(p, blob); - } - entries.push(TensorEntry { - name, - role, - shape, - precisions: per, - }); - } - } - for w in writers.values_mut() { - w.file.flush().map_err(|e| e.to_string())?; - w.file.get_ref().sync_all().map_err(|e| e.to_string())?; - } - - let manifest = Manifest { - version: PACK_VERSION, - source_file: gguf_path - .file_name() - .map_or_else(String::new, |n| n.to_string_lossy().into_owned()), - source_bytes: std::fs::metadata(gguf_path).map(|m| m.len()).unwrap_or(0), - architecture: f.architecture().unwrap_or("").to_string(), - precisions: all_levels, - tensors: entries, - ffn_partition: None, - }; - let json = serde_json::to_string_pretty(&manifest).map_err(|e| e.to_string())?; - std::fs::write(out_dir.join("manifest.json"), json).map_err(|e| e.to_string())?; - Ok(manifest) -} - -/// Row-major `[out, in]` → `[in, out]`. -fn transpose(values: &[f32], out: usize, inp: usize) -> Vec { - debug_assert_eq!(values.len(), out * inp); - let mut t = vec![0.0f32; values.len()]; - for o in 0..out { - for i in 0..inp { - t[i * out + o] = values[o * inp + i]; - } - } - t -} - -// --------------------------------------------------------------------------- -// Reader -// --------------------------------------------------------------------------- - -/// An opened pack: manifest + lazily-read blobs. -/// Burn's canonical quantized `TensorData` bytes for block-quantized i8/i4 -/// values: the values packed little-endian into u32 words (8 nibbles or 4 -/// bytes per word, element `j` at bits `j·bits`), then the f32 block scales -/// appended — exactly what `Tensor::quantize(..).into_data()` yields and -/// what every backend's `q_from_data` consumes. (`TensorData::quantized` -/// itself only handles 8-bit values under the default `PackedU32` store: -/// its Q4 reader unpacks nibbles the constructor never packed.) -pub fn quantized_tensor_data( - values: &[i8], - scales: &[f32], - shape: impl Into, - scheme: QuantScheme, -) -> TensorData { - use burn::tensor::quantization::QuantValue; - let shape: burn::tensor::Shape = shape.into(); - let bits = match scheme.value { - QuantValue::Q8S | QuantValue::Q8F => 8, - QuantValue::Q4S | QuantValue::Q4F => 4, - other => panic!("quantized_tensor_data: unsupported value type {other:?}"), - }; - let per_word = 32 / bits; - let mask = (1u32 << bits) - 1; - let mut words: Vec = Vec::with_capacity(values.len().div_ceil(per_word) + scales.len()); - for chunk in values.chunks(per_word) { - let mut w = 0u32; - for (j, &v) in chunk.iter().enumerate() { - w |= ((v as u8 as u32) & mask) << (j * bits); - } - words.push(w); - } - words.extend(scales.iter().map(|s| s.to_bits())); - TensorData::from_bytes( - burn::tensor::Bytes::from_elems(words), - shape, - burn::tensor::DType::QFloat(scheme), - ) -} - -pub struct Pack { - pub dir: PathBuf, - pub manifest: Manifest, -} - -impl Pack { - /// Is `dir` a pack (has a readable manifest)? - #[must_use] - pub fn is_pack(dir: &Path) -> bool { - dir.join("manifest.json").is_file() && dir.join("header.gguf").is_file() - } - - pub fn open(dir: &Path) -> Result { - let json = std::fs::read_to_string(dir.join("manifest.json")) - .map_err(|e| format!("read manifest: {e}"))?; - let manifest: Manifest = - serde_json::from_str(&json).map_err(|e| format!("parse manifest: {e}"))?; - if manifest.version != PACK_VERSION { - return Err(format!( - "pack version {} is not the supported {PACK_VERSION}", - manifest.version - )); - } - Ok(Self { - dir: dir.to_path_buf(), - manifest, - }) - } - - /// Write the manifest back (after partitioning / calibration). - pub fn save_manifest(&self) -> Result<(), String> { - let json = serde_json::to_string_pretty(&self.manifest).map_err(|e| e.to_string())?; - let tmp = self.dir.join("manifest.json.tmp"); - std::fs::write(&tmp, json).map_err(|e| e.to_string())?; - std::fs::rename(&tmp, self.dir.join("manifest.json")).map_err(|e| e.to_string()) - } - - /// Overwrite every stored level of `entry` from new f32 values of the - /// same shape, **in place** (same byte sizes — the quantized levels are - /// re-quantized along the same last dim). Used by the FFN partitioner, - /// whose permutation keeps shapes. - pub fn rewrite_entry(&self, entry: &TensorEntry, values: &[f32]) -> Result<(), String> { - let numel: usize = entry.shape.iter().product(); - if values.len() != numel { - return Err(format!( - "rewrite '{}': {} values for shape {:?}", - entry.name, - values.len(), - entry.shape - )); - } - let last = *entry.shape.last().expect("non-empty shape"); - for (&p, blob) in &entry.precisions { - let (vbytes, sbytes): (Vec, Vec) = match p { - Precision::F32 => ( - values.iter().flat_map(|v| v.to_le_bytes()).collect(), - Vec::new(), - ), - Precision::F16 => ( - values - .iter() - .flat_map(|&v| half::f16::from_f32(v).to_le_bytes()) - .collect(), - Vec::new(), - ), - Precision::Q8 | Precision::Q4 => { - let (q, scales) = quantize_blocks(values, last, p); - let v = if p == Precision::Q4 { - pack_nibbles(&q) - } else { - q.iter().map(|&x| x as u8).collect() - }; - (v, scales.iter().flat_map(|s| s.to_le_bytes()).collect()) - } - }; - if vbytes.len() as u64 != blob.values_len || sbytes.len() as u64 != blob.scales_len { - return Err(format!( - "rewrite '{}' {p:?}: size changed ({} / {} vs {} / {})", - entry.name, - vbytes.len(), - sbytes.len(), - blob.values_len, - blob.scales_len - )); - } - let mut file = std::fs::OpenOptions::new() - .write(true) - .open(self.dir.join(p.blob_name())) - .map_err(|e| format!("open {} for write: {e}", p.blob_name()))?; - file.seek(SeekFrom::Start(blob.values_offset)) - .map_err(|e| e.to_string())?; - file.write_all(&vbytes).map_err(|e| e.to_string())?; - if !sbytes.is_empty() { - file.seek(SeekFrom::Start(blob.scales_offset)) - .map_err(|e| e.to_string())?; - file.write_all(&sbytes).map_err(|e| e.to_string())?; - } - file.sync_data().map_err(|e| e.to_string())?; - } - Ok(()) - } - - /// A 2-D entry's **columns** `ranges` (each `(start, len)`, concatenated - /// in order) as a tensor at `precision`: `[rows, Σ len]`. Quantized - /// levels are sliced at block granularity (ranges must be block-aligned) - /// straight from the stored bytes — no re-quantization. - pub fn tensor_cols( - &self, - entry: &TensorEntry, - precision: Precision, - ranges: &[(usize, usize)], - device: &Device, - ) -> Result, String> { - let &[rows, cols] = entry.shape.as_slice() else { - return Err(format!("'{}' is not 2-D", entry.name)); - }; - let width: usize = ranges.iter().map(|r| r.1).sum(); - match precision { - Precision::Q4 | Precision::Q8 => { - if ranges - .iter() - .any(|&(s, l)| !s.is_multiple_of(BLOCK) || !l.is_multiple_of(BLOCK)) - { - return Err("column ranges must be block-aligned for quantized levels".into()); - } - let (values, scales) = self.read_quant(entry, precision)?; - let bpr = cols / BLOCK; // blocks per row - let mut v = Vec::with_capacity(rows * width); - let mut s = Vec::with_capacity(rows * width / BLOCK); - for r in 0..rows { - for &(start, len) in ranges { - v.extend_from_slice(&values[r * cols + start..r * cols + start + len]); - s.extend_from_slice( - &scales[r * bpr + start / BLOCK..r * bpr + (start + len) / BLOCK], - ); - } - } - let scheme = precision.policy().scheme().expect("quantized level"); - Ok(Tensor::from_data( - quantized_tensor_data(&v, &s, [rows, width], scheme), - device, - )) - } - Precision::F16 | Precision::F32 => { - let all = self.read_floats(entry, precision)?; - let mut v = Vec::with_capacity(rows * width); - for r in 0..rows { - for &(start, len) in ranges { - v.extend_from_slice(&all[r * cols + start..r * cols + start + len]); - } - } - let dtype = crate::backend::float_dtype(device); - Ok(Tensor::from_data( - TensorData::new(v, [rows, width]), - (device, dtype), - )) - } - } - } - - /// A 2-D entry's **rows** `ranges` (concatenated) as a tensor at - /// `precision`: `[Σ len, cols]`. - pub fn tensor_rows( - &self, - entry: &TensorEntry, - precision: Precision, - ranges: &[(usize, usize)], - device: &Device, - ) -> Result, String> { - let &[_rows, cols] = entry.shape.as_slice() else { - return Err(format!("'{}' is not 2-D", entry.name)); - }; - let height: usize = ranges.iter().map(|r| r.1).sum(); - match precision { - Precision::Q4 | Precision::Q8 => { - let (values, scales) = self.read_quant(entry, precision)?; - let bpr = cols / BLOCK; - let mut v = Vec::with_capacity(height * cols); - let mut s = Vec::with_capacity(height * bpr); - for &(start, len) in ranges { - v.extend_from_slice(&values[start * cols..(start + len) * cols]); - s.extend_from_slice(&scales[start * bpr..(start + len) * bpr]); - } - let scheme = precision.policy().scheme().expect("quantized level"); - Ok(Tensor::from_data( - quantized_tensor_data(&v, &s, [height, cols], scheme), - device, - )) - } - Precision::F16 | Precision::F32 => { - let all = self.read_floats(entry, precision)?; - let mut v = Vec::with_capacity(height * cols); - for &(start, len) in ranges { - v.extend_from_slice(&all[start * cols..(start + len) * cols]); - } - let dtype = crate::backend::float_dtype(device); - Ok(Tensor::from_data( - TensorData::new(v, [height, cols]), - (device, dtype), - )) - } - } - } - - /// The source header as a payload-less GGUF — config + tokenizer readers - /// take it as-is. - pub fn header(&self) -> Result { - GgufFile::open(&self.dir.join("header.gguf")).map_err(|e| e.to_string()) - } - - pub fn entry(&self, name: &str) -> Option<&TensorEntry> { - self.manifest.tensors.iter().find(|t| t.name == name) - } - - fn read_range(&self, precision: Precision, offset: u64, len: u64) -> Result, String> { - let mut file = std::fs::File::open(self.dir.join(precision.blob_name())) - .map_err(|e| format!("open {}: {e}", precision.blob_name()))?; - file.seek(SeekFrom::Start(offset)) - .map_err(|e| e.to_string())?; - let mut buf = vec![0u8; usize::try_from(len).expect("blob fits")]; - file.read_exact(&mut buf) - .map_err(|e| format!("read {}: {e}", precision.blob_name()))?; - Ok(buf) - } - - /// The f32 values of a tensor read from the REQUESTED float level when - /// stored: `F16` reads the half-width blob and widens — same tensor in - /// RAM, half the bytes off the disk, which matters because the pack - /// lives on an HDD RAID that runs at 100% for the whole load. Falls back - /// to [`Self::read_f32`]'s best-available order when the requested level - /// is absent. - /// - /// Exists because the float arms of `tensor`/`tensor_cols`/`tensor_rows` - /// used to funnel through `read_f32`, which prefers f32 — so every - /// "F16" load silently read `f32.bin`, including the probe that once - /// "measured" F16 speed on flex. - pub fn read_floats(&self, entry: &TensorEntry, prefer: Precision) -> Result, String> { - if prefer == Precision::F16 - && let Some(b) = entry.precisions.get(&Precision::F16) - { - let bytes = self.read_range(Precision::F16, b.values_offset, b.values_len)?; - return Ok(bytes - .as_chunks::<2>() - .0 - .iter() - .map(|c| half::f16::from_le_bytes(*c).to_f32()) - .collect()); - } - self.read_f32(entry) - } - - /// The f32 values of a tensor at its best float level (f32, else f16 - /// widened, else a quantized level dequantized). - pub fn read_f32(&self, entry: &TensorEntry) -> Result, String> { - if let Some(b) = entry.precisions.get(&Precision::F32) { - let bytes = self.read_range(Precision::F32, b.values_offset, b.values_len)?; - return Ok(bytes - .as_chunks::<4>() - .0 - .iter() - .map(|c| f32::from_le_bytes(*c)) - .collect()); - } - if let Some(b) = entry.precisions.get(&Precision::F16) { - let bytes = self.read_range(Precision::F16, b.values_offset, b.values_len)?; - return Ok(bytes - .as_chunks::<2>() - .0 - .iter() - .map(|c| half::f16::from_le_bytes(*c).to_f32()) - .collect()); - } - for p in [Precision::Q8, Precision::Q4] { - if entry.precisions.contains_key(&p) { - let (q, scales) = self.read_quant(entry, p)?; - let mut out = Vec::with_capacity(q.len()); - for (i, &v) in q.iter().enumerate() { - out.push(f32::from(v) * scales[i / BLOCK]); - } - return Ok(out); - } - } - Err(format!("'{}' has no stored precision", entry.name)) - } - - /// The (i8 values, f32 block scales) of a quantized level. - pub fn read_quant( - &self, - entry: &TensorEntry, - precision: Precision, - ) -> Result<(Vec, Vec), String> { - let b = entry - .precisions - .get(&precision) - .ok_or_else(|| format!("'{}' has no {precision:?} level", entry.name))?; - let n: usize = entry.shape.iter().product(); - let vbytes = self.read_range(precision, b.values_offset, b.values_len)?; - let values: Vec = match precision { - Precision::Q4 => unpack_nibbles(&vbytes, n), - Precision::Q8 => vbytes.iter().map(|&b| b as i8).collect(), - _ => return Err(format!("{precision:?} is not a quantized level")), - }; - let sbytes = self.read_range(precision, b.scales_offset, b.scales_len)?; - let scales: Vec = sbytes - .as_chunks::<4>() - .0 - .iter() - .map(|c| f32::from_le_bytes(*c)) - .collect(); - if values.len() != n || scales.len() != n / BLOCK { - return Err(format!("'{}' {precision:?} blob size mismatch", entry.name)); - } - Ok((values, scales)) - } - - /// Build a device tensor for `entry` at `precision`. Quantized levels - /// arrive through burn's canonical quantized `TensorData` (no - /// re-quantization); float levels through the backend's float dtype. - pub fn tensor( - &self, - entry: &TensorEntry, - precision: Precision, - device: &Device, - ) -> Result, String> { - let shape: [usize; D] = entry.shape.clone().try_into().map_err(|_| { - format!( - "'{}' is rank {}, asked for {D}", - entry.name, - entry.shape.len() - ) - })?; - match precision { - Precision::Q4 | Precision::Q8 => { - let (values, scales) = self.read_quant(entry, precision)?; - let scheme: QuantScheme = precision - .policy() - .scheme() - .expect("quantized level has a scheme"); - let data = quantized_tensor_data(&values, &scales, shape, scheme); - Ok(Tensor::from_data(data, device)) - } - Precision::F16 | Precision::F32 => { - let values = self.read_floats(entry, precision)?; - let dtype = crate::backend::float_dtype(device); - Ok(Tensor::from_data( - TensorData::new(values, shape), - (device, dtype), - )) - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn block_quantizer_roundtrips_within_half_scale() { - let vals: Vec = (0..256).map(|i| ((i as f32) * 0.37).sin() * 3.0).collect(); - for p in [Precision::Q8, Precision::Q4] { - let (q, scales) = quantize_blocks(&vals, 64, p); - assert_eq!(scales.len(), 8); - let range_max = if p == Precision::Q8 { 127.0 } else { 7.0 }; - for (i, (&v, &qq)) in vals.iter().zip(&q).enumerate() { - let back = f32::from(qq) * scales[i / BLOCK]; - assert!( - (v - back).abs() <= scales[i / BLOCK] / 2.0 + 1e-6, - "{p:?} elem {i}: {v} vs {back} (scale {})", - scales[i / BLOCK] - ); - assert!(qq.abs() as f32 <= range_max); - } - } - } - - #[test] - fn nibble_packing_roundtrips() { - let vals: Vec = (-7..=7).chain([0, 7, -7, 3]).collect(); - let packed = pack_nibbles(&vals); - assert_eq!(packed.len(), vals.len().div_ceil(2)); - assert_eq!(unpack_nibbles(&packed, vals.len()), vals); - } - - #[test] - fn precision_list_parses_and_dedups() { - assert_eq!( - Precision::parse_list("f32,q4,q8,q4").unwrap(), - vec![Precision::Q4, Precision::Q8, Precision::F32] - ); - assert!(Precision::parse_list("q3").is_err()); - } - - /// A quantized level written by the pack quantizer reconstructs to a - /// burn tensor whose dequantized values match our own dequant exactly. - #[test] - fn canonical_quantized_tensor_data_roundtrip() { - let device = crate::backend::cpu_device(); - for (p, rows, cols) in [ - (Precision::Q8, 32, 64), - (Precision::Q4, 32, 64), - // The real shape the 2B gate trips on — exercises burn's Q4 packing at scale. - (Precision::Q4, 2048, 6144), - ] { - let n = rows * cols; - let vals: Vec = (0..n).map(|i| ((i as f32) * 0.11).cos()).collect(); - let (q, scales) = quantize_blocks(&vals, cols, p); - let scheme = p.policy().scheme().unwrap(); - let data = quantized_tensor_data(&q, &scales, [rows, cols], scheme); - let t = Tensor::<2>::from_data(data, &device); - let back = t.dequantize().into_data().to_vec::().unwrap(); - assert_eq!(back.len(), n, "{p:?} [{rows}, {cols}]"); - for (i, (&qq, &b)) in q.iter().zip(&back).enumerate().step_by(97) { - let ours = f32::from(qq) * scales[i / BLOCK]; - assert!((ours - b).abs() < 1e-5, "{p:?} elem {i}: {ours} vs {b}"); - } - } - } - - #[test] - fn column_and_row_slices_round_trip_through_a_written_pack() { - // Build a tiny two-tensor pack by hand, then slice cluster ranges - // back out at f32 and Q8 and compare to the source rows/columns. - let dir = std::env::temp_dir().join(format!("mummu-pack-slice-{}", std::process::id())); - let _ = std::fs::remove_dir_all(&dir); - std::fs::create_dir_all(&dir).unwrap(); - let (rows, cols) = (64usize, 128usize); // both dims multiples of BLOCK - let gate: Vec = (0..rows * cols) - .map(|i| ((i as f32) * 0.017).sin()) - .collect(); - let down: Vec = (0..cols * rows) - .map(|i| ((i as f32) * 0.023).cos()) - .collect(); - // Write f32 and Q8 blobs with two entries. - let mut f32w = BlobWriter { - file: std::io::BufWriter::new(std::fs::File::create(dir.join("f32.bin")).unwrap()), - len: 0, - unsynced: 0, - }; - let mut q8w = BlobWriter { - file: std::io::BufWriter::new(std::fs::File::create(dir.join("q8.bin")).unwrap()), - len: 0, - unsynced: 0, - }; - let mut entry = |name: &str, vals: &[f32], shape: Vec| -> TensorEntry { - let last = *shape.last().unwrap(); - let fbytes: Vec = vals.iter().flat_map(|v| v.to_le_bytes()).collect(); - let (fo, fl) = f32w.append(&fbytes).unwrap(); - let (q, sc) = quantize_blocks(vals, last, Precision::Q8); - let qb: Vec = q.iter().map(|&x| x as u8).collect(); - let sb: Vec = sc.iter().flat_map(|s| s.to_le_bytes()).collect(); - let (qo, ql) = q8w.append(&qb).unwrap(); - let (so, sl) = q8w.append(&sb).unwrap(); - TensorEntry { - name: name.into(), - role: Role::Linear, - shape, - precisions: [ - ( - Precision::F32, - Blob { - values_offset: fo, - values_len: fl, - scales_offset: 0, - scales_len: 0, - }, - ), - ( - Precision::Q8, - Blob { - values_offset: qo, - values_len: ql, - scales_offset: so, - scales_len: sl, - }, - ), - ] - .into_iter() - .collect(), - } - }; - let ge = entry("gate", &gate, vec![rows, cols]); - let de = entry("down", &down, vec![cols, rows]); - f32w.file.flush().unwrap(); - q8w.file.flush().unwrap(); - let manifest = Manifest { - version: PACK_VERSION, - source_file: String::new(), - source_bytes: 0, - architecture: "test".into(), - precisions: vec![Precision::F32, Precision::Q8], - tensors: vec![ge.clone(), de.clone()], - ffn_partition: None, - }; - std::fs::write( - dir.join("manifest.json"), - serde_json::to_string(&manifest).unwrap(), - ) - .unwrap(); - // No header.gguf here — construct Pack directly. - let pack = Pack { - dir: dir.clone(), - manifest, - }; - let device = crate::backend::cpu_device(); - // Columns [32,32) and [96,32) of gate → [rows, 64]. - let ranges = [(32usize, 32usize), (96, 32)]; - let cslab = pack - .tensor_cols(&ge, Precision::F32, &ranges, &device) - .unwrap(); - assert_eq!(cslab.dims(), [rows, 64]); - let got = cslab.into_data().to_vec::().unwrap(); - for r in 0..rows { - for (k, &(start, len)) in ranges.iter().enumerate() { - let base: usize = ranges[..k].iter().map(|x| x.1).sum(); - for j in 0..len { - let want = gate[r * cols + start + j]; - assert!( - (got[r * 64 + base + j] - want).abs() < 1e-6, - "col f32 r{r} j{j}" - ); - } - } - } - // Same ranges as rows of down → [64, rows]. - let rslab = pack - .tensor_rows(&de, Precision::F32, &ranges, &device) - .unwrap(); - assert_eq!(rslab.dims(), [64, rows]); - let gotr = rslab.into_data().to_vec::().unwrap(); - for (k, &(start, len)) in ranges.iter().enumerate() { - let base: usize = ranges[..k].iter().map(|x| x.1).sum(); - for i in 0..len { - for c in 0..rows { - let want = down[(start + i) * rows + c]; - assert!( - (gotr[(base + i) * rows + c] - want).abs() < 1e-6, - "row f32 i{i} c{c}" - ); - } - } - } - // Q8 column slice dequantizes close to the source. - let cq = pack - .tensor_cols(&ge, Precision::Q8, &ranges, &device) - .unwrap(); - let dq = cq.dequantize().into_data().to_vec::().unwrap(); - for r in 0..rows { - for (k, &(start, len)) in ranges.iter().enumerate() { - let base: usize = ranges[..k].iter().map(|x| x.1).sum(); - for j in 0..len { - let want = gate[r * cols + start + j]; - assert!( - (dq[r * 64 + base + j] - want).abs() < 0.05, - "col q8 r{r} j{j}" - ); - } - } - } - let _ = std::fs::remove_dir_all(&dir); - } -} diff --git a/crates/mummu/examples/src/partition.rs b/crates/mummu/examples/src/partition.rs deleted file mode 100644 index f6cbd1c..0000000 --- a/crates/mummu/examples/src/partition.rs +++ /dev/null @@ -1,480 +0,0 @@ -//! **FFN partitioning** — P9 stage 3(c): turn a dense model's SwiGLU FFNs -//! into neuron clusters so the tier machinery built for MoE experts applies -//! to dense models too, **without changing the model**. -//! -//! SwiGLU is a sum over intermediate neurons: `down(silu(gate(x)) * up(x))` -//! = Σ_j silu(x·g_j)(x·u_j) d_j. Any partition of the neurons into clusters -//! therefore computes the *same* function when every cluster runs — so the -//! importer may reorder the intermediate dimension so clusters are -//! contiguous, record the cluster spans, and the runtime can hold different -//! clusters on different devices at different precisions (exact), or skip -//! low-energy clusters per token (opt-in, measured — see the skip table). -//! -//! The partition is stored **in place**: the FFN entries keep their names, -//! shapes and byte sizes, only the neuron order changes (columns of -//! `gate`/`up` in their `[hidden, inter]` Linear layout, rows of `down`), -//! every stored precision rewritten from the permuted f32. Loaders that -//! know nothing about partitions keep working unchanged. -//! -//! Clustering is MoEfication-style *parameter* clustering — balanced -//! k-means on each neuron's `gate ‖ up` weight vector — which needs no -//! calibration data (so it runs on import, for any size of model). A -//! Johnson–Lindenstrauss projection to [`PROJ_DIMS`] keeps it cheap. For -//! the exact path the cluster quality is irrelevant; it matters for -//! skipping, and the skip table measures that honestly. - -use std::collections::BTreeMap; - -use crate::pack::{ClusterSpan, FfnPartition, Pack, Precision, quantize_blocks}; - -/// Default clusters per layer (reduced to the largest divisor that keeps -/// every cluster a whole number of quantization blocks). -pub const DEFAULT_CLUSTERS: usize = 32; -/// JL projection width for the clustering features. -pub const PROJ_DIMS: usize = 192; -/// k-means iterations. -const KMEANS_ITERS: usize = 10; - -/// The three FFN entries of one layer, by pack name. -#[derive(Debug, Clone)] -pub struct FfnNames { - pub gate: String, - pub up: String, - pub down: String, -} - -/// Pick the cluster count: the largest `c <= want` with `inter % (c * block) == 0`. -#[must_use] -pub fn cluster_count(inter: usize, want: usize, block: usize) -> usize { - (1..=want) - .rev() - .find(|c| inter.is_multiple_of(c * block)) - .unwrap_or(1) -} - -/// Deterministic ±1 projection of `dims`-long vectors to [`PROJ_DIMS`] -/// (splitmix64 stream, so every import of a model yields the same clusters). -fn projection(dims: usize, seed: u64) -> Vec { - let mut state = seed ^ 0x9E37_79B9_7F4A_7C15; - let mut next = move || { - state = state.wrapping_add(0x9E37_79B9_7F4A_7C15); - let mut z = state; - z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); - z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); - z ^ (z >> 31) - }; - (0..dims * PROJ_DIMS) - .map(|_| if next() & 1 == 0 { 1.0 } else { -1.0 }) - .collect() -} - -/// Balanced k-means: `n` points of `d` features into `k` equal clusters. -/// Returns each point's cluster. Assignment is greedy by distance with a -/// capacity per cluster (the classic balanced heuristic) — deterministic. -fn balanced_kmeans(features: &[f32], n: usize, d: usize, k: usize) -> Vec { - assert!(n.is_multiple_of(k), "balanced k-means: n must divide by k"); - let cap = n / k; - // Init: evenly spaced points. - let mut centroids: Vec = (0..k) - .flat_map(|c| { - let i = c * n / k; - features[i * d..(i + 1) * d].iter().copied() - }) - .collect(); - let mut assign = vec![0usize; n]; - let mut dist = vec![0f32; n * k]; - for _ in 0..KMEANS_ITERS { - // Distances, parallel over points. - let threads = std::thread::available_parallelism().map_or(4, |p| p.get()).min(32); - let chunk = n.div_ceil(threads).max(1); - std::thread::scope(|s| { - for (ti, slab) in dist.chunks_mut(chunk * k).enumerate() { - let centroids = ¢roids; - s.spawn(move || { - let start = ti * chunk; - for (li, row) in slab.chunks_mut(k).enumerate() { - let p = &features[(start + li) * d..(start + li + 1) * d]; - for (c, out) in row.iter_mut().enumerate() { - let cen = ¢roids[c * d..(c + 1) * d]; - *out = p.iter().zip(cen).map(|(a, b)| (a - b) * (a - b)).sum(); - } - } - }); - } - }); - // Greedy balanced assignment: every (point, cluster) pair by distance. - let mut pairs: Vec<(f32, u32, u32)> = Vec::with_capacity(n * k); - for p in 0..n { - for c in 0..k { - pairs.push((dist[p * k + c], p as u32, c as u32)); - } - } - pairs.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal)); - let mut fill = vec![0usize; k]; - let mut done = vec![false; n]; - let mut left = n; - for (_, p, c) in pairs { - let (p, c) = (p as usize, c as usize); - if done[p] || fill[c] == cap { - continue; - } - assign[p] = c; - done[p] = true; - fill[c] += 1; - left -= 1; - if left == 0 { - break; - } - } - // Update centroids. - centroids.iter_mut().for_each(|v| *v = 0.0); - for p in 0..n { - let c = assign[p]; - for (acc, v) in centroids[c * d..(c + 1) * d].iter_mut().zip(&features[p * d..(p + 1) * d]) { - *acc += v; - } - } - let inv = 1.0 / cap as f32; - centroids.iter_mut().for_each(|v| *v *= inv); - } - assign -} - -/// Cluster the neurons of one layer from its `gate` and `up` weights (both -/// `[hidden, inter]` row-major): returns the permutation (new position → -/// old neuron index) with clusters contiguous, and the cluster spans. -pub fn cluster_neurons( - gate: &[f32], - up: &[f32], - hidden: usize, - inter: usize, - clusters: usize, - seed: u64, -) -> (Vec, Vec) { - assert_eq!(gate.len(), hidden * inter); - assert_eq!(up.len(), hidden * inter); - assert!(inter.is_multiple_of(clusters)); - // Features: JL projection of the 2·hidden-long neuron vector (gate col ‖ up col). - let proj = projection(2 * hidden, seed); - let mut features = vec![0f32; inter * PROJ_DIMS]; - let threads = std::thread::available_parallelism().map_or(4, |p| p.get()).min(32); - let chunk = inter.div_ceil(threads).max(1); - std::thread::scope(|s| { - for (ti, slab) in features.chunks_mut(chunk * PROJ_DIMS).enumerate() { - let proj = &proj; - s.spawn(move || { - let start = ti * chunk; - for (li, out) in slab.chunks_mut(PROJ_DIMS).enumerate() { - let j = start + li; - // Column j of gate/up: stride `inter`. - for (f, o) in out.iter_mut().enumerate() { - let mut acc = 0f32; - for h in 0..hidden { - let g = gate[h * inter + j]; - let u = up[h * inter + j]; - acc += g * proj[h * PROJ_DIMS + f] + u * proj[(hidden + h) * PROJ_DIMS + f]; - } - *o = acc; - } - } - }); - } - }); - let assign = balanced_kmeans(&features, inter, PROJ_DIMS, clusters); - let cap = inter / clusters; - let mut perm = Vec::with_capacity(inter); - let mut spans = Vec::with_capacity(clusters); - for c in 0..clusters { - let start = perm.len(); - perm.extend((0..inter).filter(|&j| assign[j] == c)); - debug_assert_eq!(perm.len() - start, cap); - spans.push(ClusterSpan { start, len: cap }); - } - (perm, spans) -} - -/// Permute the columns of a row-major `[rows, cols]` matrix: new column `p` -/// takes old column `perm[p]`. -fn permute_cols(values: &[f32], rows: usize, cols: usize, perm: &[usize]) -> Vec { - let mut out = vec![0f32; values.len()]; - for r in 0..rows { - let src = &values[r * cols..(r + 1) * cols]; - let dst = &mut out[r * cols..(r + 1) * cols]; - for (p, &j) in perm.iter().enumerate() { - dst[p] = src[j]; - } - } - out -} - -/// The FFN entry names of every trunk layer — the partitioner's input. -/// -/// Not architecture-specific: Qwen2, Qwen3, LFM2 and Qwen3.5 all store the -/// standard GGUF triple, so this lives here rather than in any one model. -/// A model whose FFN is named differently, or which has none, supplies its -/// own list (or an empty one). -#[must_use] -pub fn ffn_names(trunk_layers: usize) -> Vec { - (0..trunk_layers) - .map(|l| FfnNames { - gate: format!("blk.{l}.ffn_gate.weight"), - up: format!("blk.{l}.ffn_up.weight"), - down: format!("blk.{l}.ffn_down.weight"), - }) - .collect() -} - -/// Permute the rows of a row-major `[rows, cols]` matrix. -fn permute_rows(values: &[f32], rows: usize, cols: usize, perm: &[usize]) -> Vec { - let mut out = Vec::with_capacity(values.len()); - for &j in perm { - out.extend_from_slice(&values[j * cols..(j + 1) * cols]); - } - debug_assert_eq!(out.len(), rows * cols); - out -} - -/// Partition every layer's FFN of a pack in place. `layers` names each -/// layer's three entries; `want` is the requested cluster count. Writes -/// the partition into the manifest (replacing any previous one — which -/// must not exist, since the entries would already be permuted). -pub fn partition_pack( - pack: &mut Pack, - layers: &[FfnNames], - want: usize, - mut on_progress: impl FnMut(usize, usize), -) -> Result<(), String> { - if pack.manifest.ffn_partition.is_some() { - return Err("pack is already partitioned".into()); - } - // Crash safety: a layer's three entries are rewritten one after the - // other, so a crash in between leaves gate permuted and down not — a - // corrupt layer. The journal records, per layer, the permutation and a - // fingerprint of each entry's f32 prefix BEFORE rewriting; on a rerun a - // journaled layer is repaired by permuting only the entries whose - // fingerprint still matches the pre-state. Journals of finished layers - // are kept until the manifest is written, then removed. - let journal_dir = pack.dir.join("partition.journal"); - std::fs::create_dir_all(&journal_dir).map_err(|e| e.to_string())?; - let mut spans_per_layer = Vec::with_capacity(layers.len()); - let mut names = Vec::with_capacity(layers.len()); - for (li, n) in layers.iter().enumerate() { - on_progress(li, layers.len()); - let journal = journal_dir.join(format!("layer-{li}.json")); - if let Some(spans) = repair_layer(pack, n, &journal)? { - spans_per_layer.push(spans); - names.push([n.gate.clone(), n.up.clone(), n.down.clone()]); - continue; - } - let gate_e = pack.entry(&n.gate).ok_or_else(|| format!("missing {}", n.gate))?.clone(); - let up_e = pack.entry(&n.up).ok_or_else(|| format!("missing {}", n.up))?.clone(); - let down_e = pack.entry(&n.down).ok_or_else(|| format!("missing {}", n.down))?.clone(); - let &[hidden, inter] = gate_e.shape.as_slice() else { - return Err(format!("{} is not 2-D", n.gate)); - }; - if up_e.shape != vec![hidden, inter] || down_e.shape != vec![inter, hidden] { - return Err(format!("layer {li}: FFN shapes disagree ({:?} / {:?} / {:?})", gate_e.shape, up_e.shape, down_e.shape)); - } - let clusters = cluster_count(inter, want, crate::pack::BLOCK); - let gate = pack.read_f32(&gate_e)?; - let up = pack.read_f32(&up_e)?; - let down = pack.read_f32(&down_e)?; - let (perm, spans) = cluster_neurons(&gate, &up, hidden, inter, clusters, li as u64); - let entry = LayerJournal { - perm: perm.clone(), - spans: spans.clone(), - before: [fingerprint(&gate), fingerprint(&up), fingerprint(&down)], - done: false, - }; - std::fs::write(&journal, serde_json::to_string(&entry).map_err(|e| e.to_string())?) - .map_err(|e| e.to_string())?; - pack.rewrite_entry(&gate_e, &permute_cols(&gate, hidden, inter, &perm))?; - pack.rewrite_entry(&up_e, &permute_cols(&up, hidden, inter, &perm))?; - pack.rewrite_entry(&down_e, &permute_rows(&down, inter, hidden, &perm))?; - let entry = LayerJournal { done: true, ..entry }; - std::fs::write(&journal, serde_json::to_string(&entry).map_err(|e| e.to_string())?) - .map_err(|e| e.to_string())?; - spans_per_layer.push(spans); - names.push([n.gate.clone(), n.up.clone(), n.down.clone()]); - } - on_progress(layers.len(), layers.len()); - pack.manifest.ffn_partition = Some(FfnPartition { - layers: spans_per_layer, - names, - hotness: Vec::new(), - skip_table: Vec::new(), - }); - pack.save_manifest()?; - let _ = std::fs::remove_dir_all(&journal_dir); - Ok(()) -} - -/// Per-layer crash journal (see [`partition_pack`]). -#[derive(serde::Serialize, serde::Deserialize)] -struct LayerJournal { - perm: Vec, - spans: Vec, - /// Fingerprints of gate / up / down f32 before the rewrite. - before: [u64; 3], - done: bool, -} - -/// FNV-1a over the first 4096 values (enough to tell permuted from not). -fn fingerprint(values: &[f32]) -> u64 { - let mut h: u64 = 0xcbf2_9ce4_8422_2325; - for v in values.iter().take(4096) { - for b in v.to_le_bytes() { - h ^= u64::from(b); - h = h.wrapping_mul(0x0000_0100_0000_01b3); - } - } - h -} - -/// If `journal` exists, finish that layer: entries still matching their -/// pre-rewrite fingerprint get permuted, the rest are already done. -/// Returns the layer's spans, or `None` when there was no journal. -fn repair_layer(pack: &Pack, n: &FfnNames, journal: &std::path::Path) -> Result>, String> { - let Ok(text) = std::fs::read_to_string(journal) else { - return Ok(None); - }; - let j: LayerJournal = serde_json::from_str(&text).map_err(|e| format!("journal {}: {e}", journal.display()))?; - if j.done { - return Ok(Some(j.spans)); - } - let gate_e = pack.entry(&n.gate).ok_or_else(|| format!("missing {}", n.gate))?.clone(); - let up_e = pack.entry(&n.up).ok_or_else(|| format!("missing {}", n.up))?.clone(); - let down_e = pack.entry(&n.down).ok_or_else(|| format!("missing {}", n.down))?.clone(); - let (hidden, inter) = (gate_e.shape[0], gate_e.shape[1]); - for (i, e) in [&gate_e, &up_e, &down_e].into_iter().enumerate() { - let vals = pack.read_f32(e)?; - if fingerprint(&vals) == j.before[i] { - let permuted = if i == 2 { - permute_rows(&vals, inter, hidden, &j.perm) - } else { - permute_cols(&vals, hidden, inter, &j.perm) - }; - pack.rewrite_entry(e, &permuted)?; - } - } - let done = LayerJournal { done: true, ..j }; - std::fs::write(journal, serde_json::to_string(&done).map_err(|e| e.to_string())?).map_err(|e| e.to_string())?; - Ok(Some(done.spans)) -} - -/// Bytes one cluster of a layer's FFN costs at each stored level (gate + -/// up + down slices), for the tier planner. -pub fn cluster_costs(pack: &Pack, layer: usize) -> Result, String> { - let part = pack - .manifest - .ffn_partition - .as_ref() - .ok_or("pack has no FFN partition")?; - let names = part.names.get(layer).ok_or("layer out of range")?; - let spans = &part.layers[layer]; - let mut per_neuron: BTreeMap = BTreeMap::new(); - for name in names { - let e = pack.entry(name).ok_or_else(|| format!("missing {name}"))?; - let inter = if e.shape[0] > e.shape[1] { e.shape[0] } else { e.shape[1] }; - let numel = e.shape.iter().product::() as u64; - for (&p, blob) in &e.precisions { - let bytes = match p { - Precision::Q4 | Precision::Q8 => blob.values_len + blob.scales_len, - Precision::F32 => numel * 4, - // f16 is TWO bytes. Costing it as four made it look exactly - // as expensive as f32 to the tier planner, so it could never - // win a placement — which is why an f16 rung never appeared - // in a plan despite being the cheapest lossless option for a - // quantized source. - Precision::F16 => numel * 2, - }; - // Bytes per neuron (column/row) — the entry is [hidden, inter] or [inter, hidden]. - *per_neuron.entry(p).or_insert(0) += bytes / inter as u64; - } - } - Ok(spans - .iter() - .map(|s| crate::tier::ExpertCost { - bytes: per_neuron.iter().map(|(&p, &b)| (p, b * s.len as u64)).collect(), - }) - .collect()) -} - -/// Quantize-and-pack helper re-exported for the rewrite path's tests. -#[doc(hidden)] -pub fn requantize(values: &[f32], last_dim: usize, p: Precision) -> (Vec, Vec) { - quantize_blocks(values, last_dim, p) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn cluster_count_keeps_whole_blocks() { - assert_eq!(cluster_count(6144, 32, 32), 32); - assert_eq!(cluster_count(17408, 32, 32), 32); // 544 blocks = 32 × 17 - assert_eq!(cluster_count(17408, 64, 32), 34); // 544 = 34 × 16 - assert_eq!(cluster_count(96, 32, 32), 3); - assert_eq!(cluster_count(32, 32, 32), 1); - } - - #[test] - fn balanced_kmeans_is_balanced_and_groups_obvious_clusters() { - // 4 clear groups of 8 points in 2-D. - let mut f = Vec::new(); - for g in 0..4 { - for i in 0..8 { - f.push(g as f32 * 10.0 + (i as f32) * 0.01); - f.push(-(g as f32) * 10.0 + (i as f32) * 0.02); - } - } - let a = balanced_kmeans(&f, 32, 2, 4); - for g in 0..4 { - let first = a[g * 8]; - assert!((0..8).all(|i| a[g * 8 + i] == first), "group {g} split: {a:?}"); - } - let mut counts = [0; 4]; - a.iter().for_each(|&c| counts[c] += 1); - assert_eq!(counts, [8; 4]); - } - - #[test] - fn permutation_covers_every_neuron_once_in_contiguous_spans() { - let (hidden, inter) = (6, 64); - let gate: Vec = (0..hidden * inter).map(|i| ((i as f32) * 0.3).sin()).collect(); - let up: Vec = (0..hidden * inter).map(|i| ((i as f32) * 0.7).cos()).collect(); - let (perm, spans) = cluster_neurons(&gate, &up, hidden, inter, 4, 1); - let mut sorted = perm.clone(); - sorted.sort_unstable(); - assert_eq!(sorted, (0..inter).collect::>()); - assert_eq!(spans.len(), 4); - assert!(spans.iter().all(|s| s.len == 16)); - assert_eq!(spans.iter().map(|s| s.start).collect::>(), vec![0, 16, 32, 48]); - // Permuting columns then rows keeps the SwiGLU sum: check one input. - let x: Vec = (0..hidden).map(|h| (h as f32 + 1.0) * 0.1).collect(); - let down: Vec = (0..inter * hidden).map(|i| ((i as f32) * 0.11).sin()).collect(); - let silu = |v: f32| v / (1.0 + (-v).exp()); - let dense = |g: &[f32], u: &[f32], d: &[f32]| -> Vec { - let mut out = vec![0f32; hidden]; - for j in 0..inter { - let gj: f32 = (0..hidden).map(|h| x[h] * g[h * inter + j]).sum(); - let uj: f32 = (0..hidden).map(|h| x[h] * u[h * inter + j]).sum(); - let a = silu(gj) * uj; - for h in 0..hidden { - out[h] += a * d[j * hidden + h]; - } - } - out - }; - let ref_out = dense(&gate, &up, &down); - let p_out = dense( - &permute_cols(&gate, hidden, inter, &perm), - &permute_cols(&up, hidden, inter, &perm), - &permute_rows(&down, inter, hidden, &perm), - ); - for (a, b) in ref_out.iter().zip(&p_out) { - assert!((a - b).abs() < 1e-4, "{a} vs {b}"); - } - } -} diff --git a/crates/mummu/examples/src/plan.rs b/crates/mummu/examples/src/plan.rs deleted file mode 100644 index bc7c325..0000000 --- a/crates/mummu/examples/src/plan.rs +++ /dev/null @@ -1,333 +0,0 @@ -//! **Precision selection** — the first piece of the P6 hardware planner: -//! given a model's shape and one adapter's real capabilities, which float -//! precision is the *highest* one that still fits? -//! -//! Deliberately narrow. This is arithmetic over numbers the rest of the crate -//! already produces ([`crate::backend::inventory`] for VRAM and `SHADER_F16`, -//! a checkpoint's `config.json` for parameter count and cache geometry), and -//! it decides exactly one thing: `Gpu` (f32) or `GpuF16`. Layer placement, -//! multi-GPU sharding and CPU spill are separate ROADMAP items; quantized -//! tiers arrive with P9 and will extend [`Precision`] rather than reshape this. -//! -//! The model is calibrated against measurements, not first principles. -//! Qwen2.5-1.5B on the reference card measures ~8.0 GiB of runner VRAM in f32 -//! and ~3.6 GiB in f16 (`bench/BASELINE.md`) — roughly weights + KV cache plus -//! a fixed overhead for activations, workspaces and CubeCL's memory pools. -//! [`OVERHEAD_BYTES`] is that fixed term, and [`Fit::projected_bytes`] is the -//! whole model; both are honest approximations whose job is to keep a plan on -//! the right side of a cliff, not to predict allocator behaviour to the byte. - -use crate::backend::GpuAdapter; - -/// Fixed VRAM a live runner needs beyond weights and KV cache: activations, -/// matmul workspaces, and CubeCL's memory pools. Derived from the reference -/// measurements — Qwen2.5-1.5B (1.54 G params) reads ~8.0 GiB runner VRAM in -/// f32 against 6.2 GiB of weights, and ~3.6 GiB in f16 against 3.1 GiB of -/// weights, so the residual is ~0.5-1.8 GiB depending on dtype. 1 GiB is the -/// middle of that band and errs toward *not* promising a fit. -pub const OVERHEAD_BYTES: u64 = 1 << 30; - -/// Fraction of an adapter's VRAM a plan may claim. The rest is the display -/// server's: the reference box runs 3.5-6.5 GiB of desktop ambient on the same -/// card, and a plan that ignores it produces an allocation failure at load -/// rather than a slow model. -pub const USABLE_VRAM_FRACTION: f64 = 0.75; - -/// Float precisions the planner can pick today. Ordered highest-quality -/// first; P9's int8/int4 tiers extend this enum downward. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] -pub enum Precision { - /// `backend::Gpu` — f32 weights and KV cache. - F32, - /// `backend::GpuF16` — f16 weights and KV cache, f32 attention-score - /// island. Needs an adapter advertising `SHADER_F16`. - F16, -} - -impl Precision { - /// Bytes per stored float. - #[must_use] - pub fn bytes_per_float(self) -> u64 { - match self { - Self::F32 => 4, - Self::F16 => 2, - } - } - - /// Highest-first, the order the planner tries them in. - #[must_use] - pub fn descending() -> [Self; 2] { - [Self::F32, Self::F16] - } -} - -/// What the planner needs to know about a model to size it. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct ModelShape { - /// Total parameters (weights + embeddings), as the checkpoint reports. - pub params: u64, - /// KV-cache floats stored **per token**, summed over layers: - /// `2 * layers * num_kv_heads * head_dim`. - pub kv_floats_per_token: u64, - /// Context length the plan must hold. - pub context_tokens: usize, -} - -impl ModelShape { - /// Sizing from a decoder's hyperparameters, so a caller passes the - /// `config.json` numbers rather than pre-computing cache geometry. - #[must_use] - pub fn from_decoder( - params: u64, - layers: usize, - num_kv_heads: usize, - head_dim: usize, - context_tokens: usize, - ) -> Self { - assert!(params > 0, "ModelShape: a model has parameters"); - assert!( - layers > 0 && num_kv_heads > 0 && head_dim > 0, - "ModelShape: degenerate decoder geometry \ - (layers {layers}, kv heads {num_kv_heads}, head_dim {head_dim})" - ); - let kv_floats_per_token = 2 * layers as u64 * num_kv_heads as u64 * head_dim as u64; - Self { - params, - kv_floats_per_token, - context_tokens, - } - } - - /// Projected resident bytes at `precision`: weights + KV cache + the fixed - /// runner overhead. - #[must_use] - pub fn projected_bytes(&self, precision: Precision) -> u64 { - let per_float = precision.bytes_per_float(); - let weights = self.params.saturating_mul(per_float); - let kv = self - .kv_floats_per_token - .saturating_mul(self.context_tokens as u64) - .saturating_mul(per_float); - weights.saturating_add(kv).saturating_add(OVERHEAD_BYTES) - } -} - -/// One adapter's budget, as the planner sees it. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct DeviceBudget { - /// Total VRAM the adapter reports. - pub vram_bytes: u64, - /// Does it advertise `SHADER_F16`? Without it, f16 is not a candidate at - /// all — the dev box's own DX12 rows are exactly this case while its - /// Vulkan rows are not. - pub shader_f16: bool, -} - -impl DeviceBudget { - /// Read a budget off an enumerated adapter. `None` when VRAM is unknown - /// (wgpu exposes no portable query; only the Windows DXGI walk fills it - /// in today) — the planner refuses to guess rather than promise a fit it - /// cannot size. - #[must_use] - pub fn from_adapter(adapter: &GpuAdapter) -> Option { - Some(Self { - vram_bytes: adapter.vram_bytes?, - shader_f16: adapter.shader_f16, - }) - } - - /// Bytes a plan may claim, after leaving the display its share. - #[must_use] - pub fn usable_bytes(&self) -> u64 { - let usable = (self.vram_bytes as f64 * USABLE_VRAM_FRACTION) as u64; - debug_assert!(usable <= self.vram_bytes, "usable VRAM cannot exceed total"); - usable - } -} - -/// A precision decision, with the numbers behind it — the shape the -/// `plan`/`doctor` introspection item will render. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct Fit { - /// The chosen precision. - pub precision: Precision, - /// Projected resident bytes at that precision. - pub projected_bytes: u64, - /// Bytes the adapter allows a plan to claim. - pub usable_bytes: u64, -} - -impl Fit { - /// Headroom left over — the slack a longer context or a second model - /// would eat into. - #[must_use] - pub fn headroom_bytes(&self) -> u64 { - self.usable_bytes.saturating_sub(self.projected_bytes) - } -} - -/// Pick the **highest** precision that fits `budget`, or `None` when even f16 -/// does not — the signal that this model needs quantization (P9) or a -/// multi-device plan, not a smaller float. -/// -/// Never silently ships a worse tier than the hardware can hold, and never -/// picks f16 on an adapter that does not advertise `SHADER_F16`. -#[must_use] -pub fn pick_precision(shape: &ModelShape, budget: &DeviceBudget) -> Option { - assert!(shape.params > 0, "pick_precision: a model has parameters"); - let usable = budget.usable_bytes(); - for precision in Precision::descending() { - if precision == Precision::F16 && !budget.shader_f16 { - continue; - } - let projected = shape.projected_bytes(precision); - if projected <= usable { - return Some(Fit { - precision, - projected_bytes: projected, - usable_bytes: usable, - }); - } - } - None -} - -#[cfg(test)] -mod tests { - use super::*; - - const GIB: u64 = 1 << 30; - - /// Qwen2.5-1.5B: 1.54 G params, 28 layers, 2 kv heads, head_dim 128. - fn qwen2_1_5b(context: usize) -> ModelShape { - ModelShape::from_decoder(1_543_714_304, 28, 2, 128, context) - } - - #[test] - fn projected_size_tracks_the_measured_reference_numbers() { - let shape = qwen2_1_5b(4096); - let f32_bytes = shape.projected_bytes(Precision::F32); - let f16_bytes = shape.projected_bytes(Precision::F16); - // Measured on the reference card: ~8.0 GiB runner f32, ~3.6 GiB f16 - // (bench/BASELINE.md). The projection must land in the same - // neighbourhood, not merely be monotone — it reads 7.0 / 3.9 GiB - // here, i.e. slightly under f32's measurement and slightly over - // f16's, which is the accuracy a single fixed overhead term buys. - assert!( - (6 * GIB..=9 * GIB).contains(&f32_bytes), - "f32 projection {f32_bytes} outside the measured ~8 GiB band" - ); - assert!( - (3 * GIB..=5 * GIB).contains(&f16_bytes), - "f16 projection {f16_bytes} outside the measured ~3.6 GiB band" - ); - assert!(f16_bytes < f32_bytes, "f16 must project smaller than f32"); - } - - #[test] - fn the_reference_card_gets_f32_and_a_small_card_gets_f16() { - let shape = qwen2_1_5b(4096); - // RTX 4070 Ti SUPER as the inventory reports it: 15.7 GiB, f16 on Vulkan. - let big = DeviceBudget { - vram_bytes: 16_852_000_000, - shader_f16: true, - }; - let fit = pick_precision(&shape, &big).expect("1.5B fits a 16 GB card"); - assert_eq!(fit.precision, Precision::F32, "highest that fits wins"); - assert!(fit.headroom_bytes() > 0, "a fit leaves headroom"); - - // An 8 GB card cannot hold the f32 build but holds f16 comfortably. - let small = DeviceBudget { - vram_bytes: 8 * GIB, - shader_f16: true, - }; - let fit = pick_precision(&shape, &small).expect("f16 fits 8 GB"); - assert_eq!(fit.precision, Precision::F16); - assert!(fit.projected_bytes <= fit.usable_bytes, "a fit must fit"); - } - - #[test] - fn an_adapter_without_shader_f16_never_gets_an_f16_plan() { - let shape = qwen2_1_5b(4096); - // The dev box's own DX12 rows: same card, no SHADER_F16. - let dx12 = DeviceBudget { - vram_bytes: 8 * GIB, - shader_f16: false, - }; - assert!( - pick_precision(&shape, &dx12).is_none(), - "without SHADER_F16 the only candidate is f32, which does not fit" - ); - // With room for f32, the same adapter plans fine. - let roomy = DeviceBudget { - vram_bytes: 24 * GIB, - shader_f16: false, - }; - assert_eq!( - pick_precision(&shape, &roomy).map(|f| f.precision), - Some(Precision::F32) - ); - } - - #[test] - fn a_model_too_big_for_f16_reports_no_fit_rather_than_a_bad_plan() { - // OLMoE-1B-7B's ~7 G params: ~14 GiB in f16, past a 16 GiB card's - // usable share — exactly the case bench/BASELINE.md records as - // "GPU is out of reach until keep-quantized VRAM (P9)". - let moe = ModelShape::from_decoder(6_919_000_000, 16, 16, 64, 4096); - let card = DeviceBudget { - vram_bytes: 16_852_000_000, - shader_f16: true, - }; - assert!( - pick_precision(&moe, &card).is_none(), - "no float precision fits; the answer is quantization, not a guess" - ); - } - - #[test] - fn context_length_moves_the_decision() { - // The KV cache is the term that grows with context, so a long enough - // context must be able to push a model off f32 onto f16. - let card = DeviceBudget { - vram_bytes: 12 * GIB, - shader_f16: true, - }; - let short = pick_precision(&qwen2_1_5b(1024), &card).expect("short context fits"); - // 64k of KV adds ~3.8 GiB in f32 — past a 12 GiB card's usable share, - // but comfortable at half the width. - let long = pick_precision(&qwen2_1_5b(65_536), &card).expect("long context still fits"); - assert_eq!(short.precision, Precision::F32); - assert_eq!(long.precision, Precision::F16, "KV growth forces the drop"); - assert!( - long.projected_bytes > qwen2_1_5b(1024).projected_bytes(Precision::F16), - "a longer context must project larger at the same precision" - ); - } - - #[test] - fn a_budget_needs_known_vram_before_it_will_plan() { - // `vram_bytes: None` (every non-Windows adapter today) must yield no - // budget at all rather than a guessed one. - let unknown = GpuAdapter { - name: "test adapter".into(), - backend: wgpu::Backend::Vulkan, - device_type: wgpu::DeviceType::DiscreteGpu, - shader_f16: true, - max_buffer_bytes: 4 * GIB, - vram_bytes: None, - }; - assert!(DeviceBudget::from_adapter(&unknown).is_none()); - } - - #[test] - fn usable_vram_leaves_the_display_its_share() { - let budget = DeviceBudget { - vram_bytes: 16 * GIB, - shader_f16: true, - }; - let usable = budget.usable_bytes(); - assert!(usable < budget.vram_bytes, "never plan the whole card"); - assert_eq!(usable, 12 * GIB, "75% of 16 GiB"); - } -} diff --git a/crates/mummu/examples/src/prof.rs b/crates/mummu/examples/src/prof.rs deleted file mode 100644 index cac02c1..0000000 --- a/crates/mummu/examples/src/prof.rs +++ /dev/null @@ -1,268 +0,0 @@ -//! A scope-based wall-time profiler that renders flame graphs. -//! -//! Exists because this codebase spent a day mis-attributing a 27B's decode -//! time: three successive "the slow part is X" claims (device round trips, -//! small-op overhead, cluster dispatch) were each falsified by measurement, -//! while ~2.2 s of every 3.9 s token stayed unexplained. Sums of stage -//! timers answer "how much"; only a call-chain profile answers "where". -//! -//! Design: -//! * [`scope`] returns an RAII guard. Guards nest via a thread-local stack, -//! so a scope's identity is its full path from the thread's root — -//! `forward;layer;mlp.down` — which is exactly the folded-stack format -//! flame-graph tooling consumes. -//! * Aggregation is by path: 64 layers all fold into one `forward;layer;…` -//! bar, which is what you want — the question is "which *stage* costs", -//! not "which layer". -//! * [`folded`] emits **self time** per path (inclusive minus direct -//! children), so parent bars are never wider than their children's sum -//! and the graph does not double-count. -//! * Off by default. When disabled, [`scope`] is one atomic load and no -//! allocation; when enabled it costs one `String` join per scope — -//! hundreds of nanoseconds against stages measured in milliseconds. -//! -//! Two hard rules for instrumenting with this: -//! * **Never hold a guard across an `.await`.** The stack is thread-local; -//! a work-stealing runtime may resume the future — and drop the guard — -//! on another thread, corrupting both threads' stacks. For a timed span -//! that must cross an await, measure with `Instant` and call [`record`], -//! which writes the aggregate directly without touching any stack. -//! * Worker threads start their own root. Name the root scope after the -//! worker (`ffn_worker;`), and read those bars as parallel wall -//! time beside the main thread's — the roots of a flame graph are -//! per-thread, so their widths can legitimately sum past 100%. - -use std::cell::RefCell; -use std::collections::BTreeMap; -use std::sync::Mutex; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::time::{Duration, Instant}; - -static ENABLED: AtomicBool = AtomicBool::new(false); - -/// Path → (inclusive nanoseconds, times entered). `BTreeMap` for a -/// deterministic report; `const`-constructible, unlike `HashMap`. -static TOTALS: Mutex> = Mutex::new(BTreeMap::new()); - -thread_local! { - static STACK: RefCell> = const { RefCell::new(Vec::new()) }; -} - -/// Turn collection on or off. Serving flips this per profiled request; -/// `MUMMU_PROFILE` in the environment forces it on for every request. -pub fn set_enabled(on: bool) { - ENABLED.store(on, Ordering::Relaxed); -} - -#[must_use] -pub fn enabled() -> bool { - ENABLED.load(Ordering::Relaxed) -} - -/// Drop everything collected so far — the start of a profiled request, so -/// one flame graph describes one generation rather than a process lifetime. -pub fn reset() { - TOTALS.lock().unwrap_or_else(|e| e.into_inner()).clear(); -} - -/// Enter a scope. The returned guard records on drop; hold it for exactly -/// the region being measured, and never across an `.await` (see module docs). -#[must_use] -pub fn scope(name: impl Into) -> ScopeGuard { - if !enabled() { - return ScopeGuard { start: None }; - } - STACK.with(|s| s.borrow_mut().push(name.into())); - ScopeGuard { - start: Some(Instant::now()), - } -} - -/// Record a span directly under `path` (semicolon-separated), bypassing the -/// thread-local stack. The async-safe escape hatch: measure with `Instant` -/// across the `.await`, then attribute it here. -pub fn record(path: &str, elapsed: Duration) { - if !enabled() { - return; - } - let mut totals = TOTALS.lock().unwrap_or_else(|e| e.into_inner()); - let entry = totals.entry(path.to_string()).or_insert((0, 0)); - entry.0 += elapsed.as_nanos() as u64; - entry.1 += 1; -} - -pub struct ScopeGuard { - /// `None` when profiling was off at creation — then the guard also never - /// pushed, so it must not pop. Enabled-state changes mid-scope are - /// handled by trusting the guard's own record, not the global flag. - start: Option, -} - -impl Drop for ScopeGuard { - fn drop(&mut self) { - let Some(start) = self.start else { return }; - let elapsed = start.elapsed().as_nanos() as u64; - STACK.with(|s| { - let mut stack = s.borrow_mut(); - let path = stack.join(";"); - stack.pop(); - let mut totals = TOTALS.lock().unwrap_or_else(|e| e.into_inner()); - let entry = totals.entry(path).or_insert((0, 0)); - entry.0 += elapsed; - entry.1 += 1; - }); - } -} - -/// The collected profile in folded-stack form: one `path count` line per -/// path, where the count is **self time in microseconds** — inclusive time -/// minus the inclusive time of direct children. Standard input for any -/// flame-graph renderer, and readable enough to eyeball sorted. -#[must_use] -pub fn folded() -> String { - let totals = TOTALS.lock().unwrap_or_else(|e| e.into_inner()); - // Sum each path's direct children so self time can be derived. A direct - // child of `a;b` is `a;b;c` — one more segment, no deeper. - let mut child_sum: BTreeMap<&str, u64> = BTreeMap::new(); - for (path, &(ns, _)) in totals.iter() { - if let Some(cut) = path.rfind(';') { - *child_sum.entry(&path[..cut]).or_insert(0) += ns; - } - } - let mut out = String::new(); - for (path, &(ns, count)) in totals.iter() { - let children = child_sum.get(path.as_str()).copied().unwrap_or(0); - // A parent can measure marginally less than its children (clock - // granularity); clamp rather than emit negative bars. - let self_us = ns.saturating_sub(children) / 1_000; - if self_us == 0 { - continue; - } - // The entry count rides along as a suffix on the leaf name so the - // rendered frame reads `mlp.down (64x)` — dispatch-count questions - // ("is this called once or 2048 times?") answer themselves. - out.push_str(path); - out.push_str(&format!(" ({count}x) {self_us}\n")); - } - out -} - -/// Render folded lines to a self-contained flame-graph SVG. -#[cfg(feature = "flamegraph")] -pub fn flamegraph_svg(folded: &str) -> Result { - let mut opts = inferno::flamegraph::Options::default(); - opts.title = "mummu decode".to_string(); - opts.count_name = "µs".to_string(); - let mut svg = Vec::new(); - inferno::flamegraph::from_lines(&mut opts, folded.lines(), &mut svg) - .map_err(|e| format!("flamegraph render: {e}"))?; - String::from_utf8(svg).map_err(|e| format!("flamegraph svg not utf-8: {e}")) -} - -#[cfg(test)] -mod tests { - use super::*; - - /// Serialize: every test toggles the one global profiler. - static GATE: Mutex<()> = Mutex::new(()); - - fn run_isolated(f: impl FnOnce()) { - let _g = GATE.lock().unwrap_or_else(|e| e.into_inner()); - reset(); - set_enabled(true); - f(); - set_enabled(false); - reset(); - } - - /// Nested guards must produce nested paths, and the folded output must - /// carry SELF time: the parent's bar shrinks by its child's share, or - /// the flame graph double-counts every nanosecond. - /// - /// Asserted structurally, not by wall-clock bounds: outer does no work - /// of its own, so a correct fold gives it (near) zero self time while a - /// double-counting fold would credit it the child's entire sleep. The - /// first version bounded outer at 45 ms and flaked — Windows' ~16 ms - /// timer granularity plus parallel-test scheduling pushed a 20 ms sleep - /// to 59 ms observed. - #[test] - fn nested_scopes_fold_into_self_time() { - run_isolated(|| { - { - let _outer = scope("outer"); - { - let _inner = scope("inner"); - std::thread::sleep(Duration::from_millis(60)); - } - } - let folded = folded(); - let get = |needle: &str| -> Option { - folded - .lines() - .find(|l| l.starts_with(needle)) - .and_then(|l| l.rsplit(' ').next()) - .and_then(|v| v.parse().ok()) - }; - let inner = get("outer;inner (").expect("inner line present"); - // Sleeps never return early; only overshoot is possible. - assert!(inner >= 50_000, "inner slept 60ms, got {inner}µs"); - // A sub-µs outer self is dropped from the report entirely — that - // is the ideal outcome, not an error. - let outer = get("outer (").unwrap_or(0); - assert!( - outer < inner, - "outer ({outer}µs) credited its child's time ({inner}µs) — double-counted" - ); - }); - } - - /// Disabled must mean disabled: no entries, and — the part that broke a - /// naive design — a guard created while off must not pop a stack it - /// never pushed, even if profiling turns on before it drops. - #[test] - fn a_guard_created_while_off_never_touches_the_stack() { - run_isolated(|| { - set_enabled(false); - let dead = scope("phantom"); - set_enabled(true); - let _live = scope("real"); - drop(dead); // must NOT pop "real" - drop(_live); - let folded = folded(); - assert!(!folded.contains("phantom"), "{folded:?}"); - assert!(folded.lines().count() <= 1, "only 'real': {folded:?}"); - }); - } - - /// `record` attributes across threads and awaits without a stack; its - /// counts accumulate. - #[test] - fn record_is_stackless_and_accumulates() { - run_isolated(|| { - record("token;readback", Duration::from_micros(500)); - record("token;readback", Duration::from_micros(500)); - let folded = folded(); - assert!( - folded.contains("token;readback (2x) 1000"), - "{folded:?}" - ); - }); - } - - /// Repeated siblings aggregate into one line — 64 layers, one bar. - #[test] - fn repeated_scopes_aggregate() { - run_isolated(|| { - for _ in 0..64 { - let _l = scope("layer"); - std::thread::sleep(Duration::from_micros(300)); - } - let folded = folded(); - let line = folded - .lines() - .find(|l| l.starts_with("layer (")) - .expect("layer line"); - assert!(line.contains("(64x)"), "{line}"); - }); - } -} diff --git a/crates/mummu/examples/src/quant.rs b/crates/mummu/examples/src/quant.rs deleted file mode 100644 index e4aa776..0000000 --- a/crates/mummu/examples/src/quant.rs +++ /dev/null @@ -1,158 +0,0 @@ -//! P9 — the keep-quantized runtime policy (stage 1). -//! -//! One model path: the same module structs and the same forward code serve -//! float and quantized weights — a `Param` holding a quantized tensor -//! executes through the backend's `q_matmul` (burn-cubecl runs the mixed -//! float×quantized matmul natively; burn-flex falls back to per-op -//! dequantize, slower but with the same memory-resident win). Import -//! **re-quantizes**: whatever the source stored (BF16, Q4_K, IQ4_XS, …) is -//! dequantized per tensor and re-quantized into this one scheme. -//! -//! What stays float, deliberately: -//! - **Embeddings** — token gather (`q_select`) has no kernel on the GPU -//! backends, and a wrong silent fallback is worse than the ~2 bytes/param -//! this table costs. -//! - Norm gammas, biases, conv kernels, per-head vectors — tiny, and their -//! precision anchors the numerics. -//! -//! The **ladder itself** — which rungs exist, how wide each is, which way -//! demotion runs, and how far the source precision lets it climb — lives in -//! the `mummu-mix` crate, together with the planner that chooses among them. -//! None of that needs a tensor library. What is left here is the one part -//! that does: turning a rung into a burn `QuantScheme`. -//! -//! Scheme facts measured on this machine along the path production takes -//! (packed bytes from a pack onto the device, then multiplied — NOT a -//! synthetic tensor quantized on-device, which measures the probe rather -//! than the runtime): Q8S block-32 is correct on flex CPU, wgpu and CUDA -//! (0.58% matmul error against the pack's own f32); Q4S likewise (9.97%). -//! The long-standing "wgpu's Q4 kernel returns garbage" note is **withdrawn** -//! — 2026-08-23, `examples/pack-precision-probe.rs`. - -use burn::tensor::quantization::{ - Calibration, QuantLevel, QuantParam, QuantScheme, QuantValue, compute_q_params, compute_range, -}; -use burn::tensor::Tensor; - -pub use mummu_mix::QuantPolicy; - -/// The burn `QuantScheme` a rung denotes. -/// -/// A trait rather than an inherent method because [`QuantPolicy`] belongs to -/// `mummu-mix`, which has no burn dependency and should not gain one — this -/// is the seam between the ladder (pure data) and the tensor library. -pub trait SchemeExt { - /// `None` for the float rungs ([`QuantPolicy::Off`], [`QuantPolicy::F16`]), - /// which are not quantizations at all. - fn scheme(self) -> Option; -} - -impl SchemeExt for QuantPolicy { - fn scheme(self) -> Option { - // Block width must stay in step with `QuantPolicy::eligible`, which - // rejects rows that do not divide it. - let value = match self { - QuantPolicy::Off | QuantPolicy::F16 => return None, - QuantPolicy::Q8 => QuantValue::Q8S, - QuantPolicy::Q4 => QuantValue::Q4S, - QuantPolicy::Q2 => QuantValue::Q2S, - }; - Some( - QuantScheme::default() - .with_value(value) - .with_level(QuantLevel::block([32])) - .with_param(QuantParam::F32), - ) - } -} - -/// Quantize one weight tensor per `policy` (min-max calibration — weights -/// are static, so calibration is exact). The caller has already decided -/// eligibility; `Off` is a caller bug. -pub fn quantize_weight( - policy: QuantPolicy, - tensor: Tensor, -) -> Tensor { - let scheme = policy - .scheme() - .expect("quantize_weight called with QuantPolicy::Off"); - let range = compute_range(&scheme, &tensor, &Calibration::MinMax); - let qparams = compute_q_params(&scheme, range); - tensor.quantize(&scheme, qparams) -} - -#[cfg(test)] -mod tests { - use super::*; - - /// `LADDER`, `bits`, `demote` and `promote` must describe the SAME order. - /// They are four separate matches over one enum and drifted once already. - #[test] - fn the_ladder_is_consistent_in_every_direction() { - let ladder = QuantPolicy::LADDER; - for pair in ladder.windows(2) { - let (hi, lo) = (pair[0], pair[1]); - assert!(hi.bits() > lo.bits(), "{hi:?} should be wider than {lo:?}"); - assert_eq!(hi.demote(), Some(lo), "{hi:?} demotes to {lo:?}"); - assert_eq!(lo.promote(), Some(hi), "{lo:?} promotes to {hi:?}"); - } - assert_eq!(ladder[0].promote(), None, "nothing above the top rung"); - assert_eq!( - ladder[ladder.len() - 1].demote(), - None, - "nothing below the bottom rung" - ); - } - - /// A quantized checkpoint never earns f32, and nothing ever earns more - /// than f32 — the cap is the source. - #[test] - fn the_ceiling_follows_the_source_precision() { - // Q4_K_S: 15.36 GB for ~27 G params. - assert_eq!(QuantPolicy::ceiling_for_source(4.55), QuantPolicy::F16); - assert_eq!(QuantPolicy::ceiling_for_source(8.0), QuantPolicy::F16); - assert_eq!(QuantPolicy::ceiling_for_source(16.0), QuantPolicy::F16); - // A genuinely f32 source is the only thing that earns f32. - assert_eq!(QuantPolicy::ceiling_for_source(32.0), QuantPolicy::Off); - // ...and the ladder has no rung above it to climb to. - assert_eq!(QuantPolicy::Off.promote(), None); - } - - use super::*; - use burn::tensor::Distribution; - - #[test] - fn policy_env_parsing() { - // from_env reads the process env — test the string logic via the - // parse arms directly instead of mutating global state. - assert_eq!(QuantPolicy::Off.scheme(), None); - assert!(QuantPolicy::Q8.scheme().is_some()); - assert!(QuantPolicy::Q8.eligible(&[512, 512])); - assert!(!QuantPolicy::Q8.eligible(&[16, 16])); - assert!(!QuantPolicy::Q8.eligible(&[1024])); - assert!(!QuantPolicy::Off.eligible(&[4096, 4096])); - // Non-block-divisible rows (the 27B's [5120, 48] β/α) stay float. - assert!(!QuantPolicy::Q4.eligible(&[5120, 48])); - assert!(QuantPolicy::Q4.eligible(&[5120, 10240])); - } - - #[test] - fn q8_roundtrip_error_is_small() { - let device = crate::backend::cpu_device(); - let w = Tensor::<2>::random([64, 64], Distribution::Default, &device); - let host = w.clone().into_data().to_vec::().unwrap(); - let max_abs = host.iter().map(|v| v.abs()).fold(0.0f32, f32::max); - let q = quantize_weight(QuantPolicy::Q8, w); - let back = q.dequantize().into_data().to_vec::().unwrap(); - let max_err = host - .iter() - .zip(&back) - .map(|(a, b)| (a - b).abs()) - .fold(0.0f32, f32::max); - // Block-32 int8: worst-case error is scale/2 = max|block|/254. - assert!( - max_err <= max_abs / 100.0, - "Q8 round-trip error too large: {max_err} vs max |w| {max_abs}" - ); - } -} diff --git a/crates/mummu/examples/src/registry.rs b/crates/mummu/examples/src/registry.rs deleted file mode 100644 index 618a47f..0000000 --- a/crates/mummu/examples/src/registry.rs +++ /dev/null @@ -1,398 +0,0 @@ -//! The model registry: declarative [`ModelSpec`]s and a small built-in -//! catalog of known-good models. Adding a model to Mummu is a manifest entry -//! here (or an app-supplied spec), not new code — the spec names the source -//! repo, the architecture that loads it, and the files it needs; `fetch` -//! hands it to the P3 downloader. - -use std::path::{Path, PathBuf}; - -use crate::hub::{self, HubError, Progress}; - -/// Which from-scratch implementation loads this checkpoint. -#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] -pub enum Architecture { - /// `models::qwen2` — Qwen2 / Qwen2.5 decoder tiers. - Qwen2, - /// `models::qwen3` — Qwen3 dense decoder (per-head q/k norm, no qkv bias, - /// decoupled head_dim); the function-calling tier (4B / 9B). - Qwen3, - /// `models::lfm2` — LFM2 / LFM2.5 hybrid conv+attention. - Lfm2, - /// `models::minilm` — all-MiniLM BERT sentence embedder. - MiniLm, - /// `models::olmoe` — OLMoE sparse mixture-of-experts decoder (the zoo's - /// first MoE). Imports from GGUF (experts pre-fused) or from HF - /// safetensors (experts fused on import). - Olmoe, - /// `models::qwen35` — Qwen3.5/3.8 hybrid: Gated DeltaNet linear - /// attention + gated full attention every 4th layer. GGUF import only. - Qwen35, -} - -/// How the checkpoint's weights are stored — which fetch + load path a spec -/// takes. -#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)] -pub enum WeightFormat { - /// `config.json` + `tokenizer.json` + `model.safetensors` (or a shard - /// index); fetched by [`hub::fetch_model`], loaded by `load_from_dir`. - #[default] - Safetensors, - /// One self-contained `.gguf` file in the repo (config + tokenizer + - /// weights in the metadata); loaded by the architecture's - /// `load_from_gguf` + [`crate::tokenizer::tokenizer_from_gguf`]. - Gguf { - /// The file name inside the repo, e.g. `qwen2.5-1.5b-instruct-q4_k_m.gguf`. - file: String, - }, -} - -/// A declarative model manifest entry. -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] -pub struct ModelSpec { - /// Short cache-dir-safe name, e.g. `qwen2.5-1.5b-instruct`. - pub name: String, - /// HuggingFace repo id (`owner/name`). - pub repo: String, - /// Git revision (tag, branch, or commit) — pin for reproducibility. - pub revision: String, - pub architecture: Architecture, - /// Weight storage (absent in older manifests = safetensors). - #[serde(default)] - pub format: WeightFormat, - /// Rough on-disk size, for settings UIs and fit checks (0 = unknown). - pub disk_bytes_estimate: u64, -} - -impl ModelSpec { - /// Sanity for manifest entries (also the deserialization gate). - pub fn validate(&self) -> Result<(), String> { - if self.name.is_empty() - || !self - .name - .chars() - .all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '_' | '.')) - { - return Err(format!("bad spec name {:?}", self.name)); - } - if !self.repo.contains('/') { - return Err(format!("repo must be owner/name, got {:?}", self.repo)); - } - if self.revision.is_empty() { - return Err("revision must be non-empty (pin something)".into()); - } - if let WeightFormat::Gguf { file } = &self.format { - let safe = !file.is_empty() - && !file.contains("..") - && !file.starts_with('/') - && file.ends_with(".gguf"); - if !safe { - return Err(format!("bad gguf file name {file:?}")); - } - } - Ok(()) - } - - /// The cache directory this model lives in under `models_root`. - #[must_use] - pub fn dir(&self, models_root: &Path) -> PathBuf { - models_root.join(&self.name) - } - - /// For a GGUF spec: the local path of the model file after [`Self::fetch`]. - #[must_use] - pub fn gguf_path(&self, models_root: &Path) -> Option { - match &self.format { - WeightFormat::Gguf { file } => Some(self.dir(models_root).join(file)), - WeightFormat::Safetensors => None, - } - } - - /// Download this model into `models_root` (resumable, cache-first; see - /// [`hub::fetch_model`] / [`hub::fetch_file`]) and return its directory, - /// ready for the architecture's `load_from_dir` / `load_from_gguf`. - pub fn fetch( - &self, - models_root: &Path, - on_progress: impl FnMut(Progress<'_>), - ) -> Result { - assert!(self.validate().is_ok(), "fetch of an invalid spec"); - let dir = self.dir(models_root); - match &self.format { - WeightFormat::Safetensors => { - hub::fetch_model(&self.repo, &self.revision, &dir, on_progress) - } - WeightFormat::Gguf { file } => { - let url = hub::hub_file_url(&self.repo, &self.revision, file); - std::fs::create_dir_all(&dir).map_err(|e| HubError::Io { - path: dir.clone(), - reason: e.to_string(), - })?; - hub::fetch_file(&url, &dir.join(file), on_progress)?; - Ok(dir) - } - } - } -} - -/// The built-in catalog: the models Mummu has ported and parity-verified (or -/// is actively gating — see the ROADMAP P2 checklist for each one's status). -#[must_use] -pub fn catalog() -> Vec { - let entries = vec![ - ModelSpec { - name: "qwen2.5-1.5b-instruct".into(), - repo: "Qwen/Qwen2.5-1.5B-Instruct".into(), - revision: "main".into(), - architecture: Architecture::Qwen2, - format: WeightFormat::Safetensors, - disk_bytes_estimate: 3_100_000_000, - }, - ModelSpec { - name: "qwen2.5-0.5b-instruct".into(), - repo: "Qwen/Qwen2.5-0.5B-Instruct".into(), - revision: "main".into(), - architecture: Architecture::Qwen2, - format: WeightFormat::Safetensors, - disk_bytes_estimate: 1_000_000_000, - }, - ModelSpec { - name: "lfm2.5-1.2b".into(), - repo: "LiquidAI/LFM2.5-1.2B-Instruct".into(), - revision: "main".into(), - architecture: Architecture::Lfm2, - format: WeightFormat::Safetensors, - disk_bytes_estimate: 2_400_000_000, - }, - ModelSpec { - name: "lfm2.5-230m".into(), - repo: "LiquidAI/LFM2.5-230M".into(), - revision: "main".into(), - architecture: Architecture::Lfm2, - format: WeightFormat::Safetensors, - disk_bytes_estimate: 500_000_000, - }, - ModelSpec { - name: "all-minilm-l6-v2".into(), - repo: "sentence-transformers/all-MiniLM-L6-v2".into(), - revision: "main".into(), - architecture: Architecture::MiniLm, - format: WeightFormat::Safetensors, - disk_bytes_estimate: 91_000_000, - }, - // Single-file GGUF variants — quarter the download, same model - // (proven vs the bf16 safetensors builds in tests/real_gguf.rs). - ModelSpec { - name: "qwen2.5-1.5b-instruct-q4km".into(), - repo: "Qwen/Qwen2.5-1.5B-Instruct-GGUF".into(), - revision: "main".into(), - architecture: Architecture::Qwen2, - format: WeightFormat::Gguf { - file: "qwen2.5-1.5b-instruct-q4_k_m.gguf".into(), - }, - disk_bytes_estimate: 1_120_000_000, - }, - ModelSpec { - name: "lfm2.5-1.2b-q4km".into(), - repo: "LiquidAI/LFM2.5-1.2B-Instruct-GGUF".into(), - revision: "main".into(), - architecture: Architecture::Lfm2, - format: WeightFormat::Gguf { - file: "LFM2.5-1.2B-Instruct-Q4_K_M.gguf".into(), - }, - disk_bytes_estimate: 731_000_000, - }, - // Qwen3 dense — the local function-calling tier. 0.6B is the fast - // parity-validation / CPU tier; 4B is the BFCL sweet spot. - ModelSpec { - name: "qwen3-0.6b".into(), - repo: "Qwen/Qwen3-0.6B".into(), - revision: "main".into(), - architecture: Architecture::Qwen3, - format: WeightFormat::Safetensors, - disk_bytes_estimate: 1_500_000_000, - }, - ModelSpec { - name: "qwen3-0.6b-q4km".into(), - repo: "unsloth/Qwen3-0.6B-GGUF".into(), - revision: "main".into(), - architecture: Architecture::Qwen3, - format: WeightFormat::Gguf { - file: "Qwen3-0.6B-Q4_K_M.gguf".into(), - }, - disk_bytes_estimate: 484_000_000, - }, - ModelSpec { - name: "qwen3-4b".into(), - repo: "Qwen/Qwen3-4B".into(), - revision: "main".into(), - architecture: Architecture::Qwen3, - format: WeightFormat::Safetensors, - disk_bytes_estimate: 8_100_000_000, - }, - ModelSpec { - name: "qwen3-4b-q4km".into(), - repo: "Qwen/Qwen3-4B-GGUF".into(), - revision: "main".into(), - architecture: Architecture::Qwen3, - format: WeightFormat::Gguf { - file: "Qwen3-4B-Q4_K_M.gguf".into(), - }, - disk_bytes_estimate: 2_500_000_000, - }, - // The zoo's first MoE: 64 experts, 8 active per token (1B active / - // 7B total). Resident-everything first cut — ~28 GB dequantized to - // f32, sized for the CPU backend (128 GB reference machine). - ModelSpec { - name: "olmoe-1b-7b-0125-instruct-q4km".into(), - repo: "allenai/OLMoE-1B-7B-0125-Instruct-GGUF".into(), - revision: "main".into(), - architecture: Architecture::Olmoe, - format: WeightFormat::Gguf { - file: "OLMoE-1B-7B-0125-Instruct-Q4_K_M.gguf".into(), - }, - disk_bytes_estimate: 4_210_000_000, - }, - // The same MoE from its HF source: 3 bf16 safetensors shards + an - // index, with the 64 experts stored separately. `load_from_dir` fuses - // them into the `[experts, out, in]` banks the module holds. - ModelSpec { - name: "olmoe-1b-7b-0125-instruct".into(), - repo: "allenai/OLMoE-1B-7B-0125-Instruct".into(), - revision: "main".into(), - architecture: Architecture::Olmoe, - format: WeightFormat::Safetensors, - disk_bytes_estimate: 13_800_000_000, - }, - // Qwen3.5 hybrid (Gated DeltaNet + interval attention) — the zoo's - // first linear-attention family. BF16 is the parity-reference build; - // Q8_0 is the practical download (identical f32 footprint once - // dequantized — mummu has no keep-quantized runtime yet, see P9). - ModelSpec { - name: "qwen3.5-2b".into(), - repo: "unsloth/Qwen3.5-2B-GGUF".into(), - revision: "main".into(), - architecture: Architecture::Qwen35, - format: WeightFormat::Gguf { - file: "Qwen3.5-2B-BF16.gguf".into(), - }, - disk_bytes_estimate: 4_500_000_000, - }, - ModelSpec { - name: "qwen3.5-2b-q8".into(), - repo: "unsloth/Qwen3.5-2B-GGUF".into(), - revision: "main".into(), - architecture: Architecture::Qwen35, - format: WeightFormat::Gguf { - file: "Qwen3.5-2B-Q8_0.gguf".into(), - }, - disk_bytes_estimate: 2_400_000_000, - }, - // Qwen3.8-27B: the header parses and every tensor dequantizes (the - // full IQ family shipped 2026-08-21), but a 27B at f32 is ~109 GB — - // loading it needs the P9 keep-quantized runtime. The entry exists - // so big-RAM hosts can try the import and everyone else gets the - // loud size error instead of silence. - ModelSpec { - name: "qwen3.8-27b-ud-q4ks".into(), - repo: "unsloth/Qwen3.8-27B-GGUF".into(), - revision: "main".into(), - architecture: Architecture::Qwen35, - format: WeightFormat::Gguf { - file: "Qwen3.8-27B-UD-Q4_K_S.gguf".into(), - }, - disk_bytes_estimate: 16_400_000_000, - }, - ]; - debug_assert!( - entries.iter().all(|s| s.validate().is_ok()), - "built-in catalog must validate" - ); - entries -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn builtin_catalog_validates_and_names_are_unique() { - let cat = catalog(); - assert!(cat.len() >= 3); - for spec in &cat { - spec.validate().unwrap_or_else(|e| panic!("{e}")); - } - let mut names: Vec<&str> = cat.iter().map(|s| s.name.as_str()).collect(); - names.sort_unstable(); - names.dedup(); - assert_eq!(names.len(), cat.len(), "duplicate names in the catalog"); - } - - #[test] - fn spec_dir_is_rooted_and_named() { - let spec = &catalog()[0]; - let dir = spec.dir(Path::new("root/models")); - assert_eq!(dir, Path::new("root/models").join(&spec.name)); - } - - #[test] - fn traversal_names_are_rejected() { - let mut spec = catalog()[0].clone(); - spec.name = "../escape".into(); - assert!(spec.validate().is_err(), "path traversal must not validate"); - } - - #[test] - fn bare_repo_is_rejected() { - let mut spec = catalog()[0].clone(); - spec.repo = "qwen".into(); - assert!(spec.validate().is_err()); - } - - #[test] - fn specs_round_trip_through_json() { - let spec = &catalog()[0]; - let json = serde_json::to_string(spec).unwrap(); - let back: ModelSpec = serde_json::from_str(&json).unwrap(); - assert_eq!(back.name, spec.name); - assert_eq!(back.architecture, spec.architecture); - } - - #[test] - fn gguf_specs_round_trip_and_old_manifests_default_to_safetensors() { - let gguf = catalog() - .into_iter() - .find(|s| matches!(s.format, WeightFormat::Gguf { .. })) - .expect("catalog has a gguf entry"); - let json = serde_json::to_string(&gguf).unwrap(); - let back: ModelSpec = serde_json::from_str(&json).unwrap(); - assert_eq!(back.format, gguf.format); - let file = match &gguf.format { - WeightFormat::Gguf { file } => file.clone(), - WeightFormat::Safetensors => unreachable!(), - }; - assert_eq!( - gguf.gguf_path(Path::new("root")), - Some(Path::new("root").join(&gguf.name).join(file)) - ); - - // A manifest written before `format` existed still deserializes. - let old = r#"{"name":"m","repo":"a/b","revision":"main", - "architecture":"Qwen2","disk_bytes_estimate":1}"#; - let back: ModelSpec = serde_json::from_str(old).unwrap(); - assert_eq!(back.format, WeightFormat::Safetensors); - assert_eq!(back.gguf_path(Path::new("root")), None); - } - - #[test] - fn bad_gguf_file_names_are_rejected() { - let mut spec = catalog()[0].clone(); - for bad in ["", "../up.gguf", "/abs.gguf", "weights.bin"] { - spec.format = WeightFormat::Gguf { file: bad.into() }; - assert!(spec.validate().is_err(), "{bad:?} must not validate"); - } - spec.format = WeightFormat::Gguf { - file: "ok-model.q4_k_m.gguf".into(), - }; - assert!(spec.validate().is_ok()); - } -} diff --git a/crates/mummu/examples/src/safetensors.rs b/crates/mummu/examples/src/safetensors.rs deleted file mode 100644 index 66ca655..0000000 --- a/crates/mummu/examples/src/safetensors.rs +++ /dev/null @@ -1,1192 +0,0 @@ -//! Safetensors **reader** + a fusing rewriter for checkpoints whose on-disk -//! layout does not match the module layout. -//! -//! `burn-store` reads safetensors for us on the ordinary path, so this module -//! exists for the two things it cannot do: -//! -//! 1. **Sharded checkpoints.** Anything past ~5 GB ships as -//! `model-0000N-of-0000M.safetensors` + a `model.safetensors.index.json` -//! weight map. `import::weights_file` only ever finds a single -//! `model.safetensors`. -//! 2. **N:1 tensor fusion.** `burn-store`'s remapping is 1:1 (rename); an MoE -//! checkpoint stores every expert separately -//! (`mlp.experts.{0..63}.gate_proj.weight`) while the module holds ONE -//! fused `[experts, out, in]` param — exactly the ggml `ffn_*_exps` layout -//! a GGUF ships pre-fused. -//! -//! The fuse reads every shard's header, plans the output layout, -//! validates it, and only then copies payload bytes — so a checkpoint missing -//! expert 37 fails before a single weight byte is read, rather than loading -//! clean and computing wrong. The result is an ordinary safetensors blob that -//! goes through the SAME `SafetensorsStore` + adapter-chain + `load_checked` -//! pipeline as every other import path (the GGUF path's -//! `dequant_to_safetensors` is the precedent). Two ways out, byte-identical -//! and unit-tested as such: [`fuse_checkpoint`] returns the blob in memory, -//! and [`fuse_checkpoint_to_file`] streams it to disk for checkpoints too big -//! to hold twice. -//! -//! Source dtypes are preserved verbatim — the bf16→backend-float cast stays -//! where it already lives, in `CastFloatAdapter` on the load pipeline. - -use std::collections::HashMap; -use std::fs::File; -use std::io::{BufReader, Read, Seek, SeekFrom, Write}; -use std::path::{Path, PathBuf}; - -/// Largest safetensors JSON header accepted, per shard. Real headers are a -/// few hundred KB (one entry per tensor); 64 MiB is a corrupt or hostile file. -const MAX_HEADER_BYTES: u64 = 64 << 20; - -/// Largest tensor count accepted across a whole checkpoint. A 64-expert MoE -/// at 16 layers already declares ~3 000; 1M is a runaway index. -const MAX_TENSORS: usize = 1 << 20; - -/// Largest fused payload either fuse will produce. Matches the GGUF path's -/// ceiling (the reference machine has 128 GB) — note the payload keeps the -/// SOURCE dtype, so a bf16 checkpoint costs half its f32 footprint. Only -/// [`fuse_checkpoint`] holds this much at once; [`fuse_checkpoint_to_file`] -/// streams it and never buffers more than one tensor. -const MAX_FUSED_BYTES: u64 = 48 << 30; - -/// Largest SINGLE tensor (or fused member) the streaming fuse will buffer. -/// The widest real one we carry is a 64-expert projection member at ~4 MiB; -/// 4 GiB is a corrupt header claiming an absurd tensor. -const MAX_PART_BYTES: u64 = 4 << 30; - -/// Largest number of shards in an index (mirrors `hub::MAX_SHARDS`). -const MAX_SHARDS: usize = 256; - -/// What went wrong reading or fusing a safetensors checkpoint. -#[derive(Debug, thiserror::Error)] -pub enum SafetensorsError { - #[error("safetensors {path}: {source}")] - Io { - path: String, - source: std::io::Error, - }, - #[error("safetensors {path}: header is not valid JSON: {reason}")] - BadHeader { path: String, reason: String }, - #[error("safetensors {path}: {what} {count} exceeds the {bound} bound")] - OverBound { - path: String, - what: &'static str, - count: u64, - bound: u64, - }, - #[error("safetensors {path}: tensor '{name}': {reason}")] - BadTensor { - path: String, - name: String, - reason: String, - }, - /// A fused group is not exactly `count` distinct members `0..count`. This - /// is the load-bearing check: a silently short expert bank would load - /// clean and compute wrong. - #[error("fused tensor '{target}': {reason}")] - BadGroup { target: String, reason: String }, - #[error("no safetensors checkpoint in {0} (looked for model.safetensors and *.index.json)")] - NoCheckpoint(PathBuf), - /// The index names a shard that is not on disk. A sibling `.part` means an - /// interrupted download, which is worth saying out loud: the alternative is - /// this surfacing as a bare `os error 2` from deep inside the load. - #[error("incomplete checkpoint {dir}: index names {missing}, which is not present{hint}")] - MissingShard { - dir: PathBuf, - missing: String, - hint: &'static str, - }, -} - -/// `u64` -> `usize` as an error rather than a panic. -/// -/// Every caller has already bounded `value` against a file length or one of -/// the `MAX_*` ceilings, so on a 64-bit target this cannot fail — but an -/// import path is exactly where "cannot fail" should still be an `Err` -/// instead of an `.expect()`, so a 32-bit build degrades to a clean error. -fn to_usize(value: u64, path: &Path, what: &'static str) -> Result { - debug_assert!( - value <= usize::MAX as u64, - "{what} fits usize on this target" - ); - usize::try_from(value).map_err(|_| SafetensorsError::OverBound { - path: path.display().to_string(), - what, - count: value, - bound: usize::MAX as u64, - }) -} - -/// One tensor as the on-disk header describes it. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct TensorEntry { - /// Safetensors dtype token, verbatim (`BF16`, `F32`, `I64`, …). - pub dtype: String, - /// Row-major shape. - pub shape: Vec, - /// `[start, end)` within the shard's payload region. - pub offsets: (u64, u64), -} - -impl TensorEntry { - /// Payload length in bytes. - #[must_use] - pub fn byte_len(&self) -> u64 { - debug_assert!( - self.offsets.1 >= self.offsets.0, - "offsets validated on read" - ); - self.offsets.1 - self.offsets.0 - } - - /// Elements implied by the shape (1 for a scalar — an empty shape). - #[must_use] - pub fn element_count(&self) -> u64 { - self.shape.iter().product() - } -} - -/// A parsed safetensors header: the tensor table plus where its payload starts. -#[derive(Debug)] -pub struct SafetensorsHeader { - pub path: PathBuf, - /// Entries in **file order** (the order payload bytes appear), which is - /// what makes a sequential read cheap. - pub tensors: Vec<(String, TensorEntry)>, - /// Absolute file offset where the payload region begins. - pub data_offset: u64, -} - -impl SafetensorsHeader { - /// Read and validate one shard's header. No payload bytes are read. - pub fn open(path: &Path) -> Result { - let io = |source: std::io::Error| SafetensorsError::Io { - path: path.display().to_string(), - source, - }; - let mut file = BufReader::new(File::open(path).map_err(io)?); - let file_len = file.get_ref().metadata().map_err(io)?.len(); - - let mut len_bytes = [0u8; 8]; - file.read_exact(&mut len_bytes).map_err(io)?; - let header_len = u64::from_le_bytes(len_bytes); - if header_len > MAX_HEADER_BYTES || header_len.saturating_add(8) > file_len { - return Err(SafetensorsError::OverBound { - path: path.display().to_string(), - what: "header bytes", - count: header_len, - bound: MAX_HEADER_BYTES.min(file_len), - }); - } - let mut header = vec![0u8; to_usize(header_len, path, "header bytes")?]; - file.read_exact(&mut header).map_err(io)?; - - let data_offset = 8 + header_len; - assert!(data_offset <= file_len, "payload starts inside the file"); - let payload_len = file_len - data_offset; - - let parsed = Self::parse_header(&header, path, payload_len)?; - Ok(Self { - path: path.to_path_buf(), - tensors: parsed, - data_offset, - }) - } - - /// Parse the JSON header into a file-ordered tensor table, validating every - /// entry against the payload region it claims. - fn parse_header( - header: &[u8], - path: &Path, - payload_len: u64, - ) -> Result, SafetensorsError> { - let bad_header = |reason: String| SafetensorsError::BadHeader { - path: path.display().to_string(), - reason, - }; - let json: serde_json::Value = - serde_json::from_slice(header).map_err(|e| bad_header(e.to_string()))?; - let object = json - .as_object() - .ok_or_else(|| bad_header("header is not a JSON object".into()))?; - if object.len() > MAX_TENSORS { - return Err(SafetensorsError::OverBound { - path: path.display().to_string(), - what: "tensor entries", - count: object.len() as u64, - bound: MAX_TENSORS as u64, - }); - } - - let mut tensors = Vec::with_capacity(object.len()); - for (name, value) in object { - // `__metadata__` is a free-form string map, not a tensor. - if name == "__metadata__" { - continue; - } - let bad = |reason: String| SafetensorsError::BadTensor { - path: path.display().to_string(), - name: name.clone(), - reason, - }; - let entry = value - .as_object() - .ok_or_else(|| bad("entry is not an object".into()))?; - let dtype = entry - .get("dtype") - .and_then(serde_json::Value::as_str) - .ok_or_else(|| bad("missing 'dtype'".into()))? - .to_string(); - let element_bytes = - dtype_bytes(&dtype).ok_or_else(|| bad(format!("unsupported dtype '{dtype}'")))?; - let shape: Vec = entry - .get("shape") - .and_then(serde_json::Value::as_array) - .ok_or_else(|| bad("missing 'shape'".into()))? - .iter() - .map(|d| d.as_u64().ok_or_else(|| bad("non-integer dim".into()))) - .collect::>()?; - let offsets = entry - .get("data_offsets") - .and_then(serde_json::Value::as_array) - .ok_or_else(|| bad("missing 'data_offsets'".into()))?; - if offsets.len() != 2 { - return Err(bad(format!("data_offsets has {} entries", offsets.len()))); - } - let start = offsets[0] - .as_u64() - .ok_or_else(|| bad("non-integer data_offset".into()))?; - let end = offsets[1] - .as_u64() - .ok_or_else(|| bad("non-integer data_offset".into()))?; - if end < start || end > payload_len { - return Err(bad(format!( - "data_offsets [{start}, {end}) outside the {payload_len}-byte payload" - ))); - } - // Positive AND negative space: the declared shape must be exactly - // the bytes claimed — a mismatch is a corrupt or mis-declared - // tensor, never something to load and hope about. - let expected = shape - .iter() - .try_fold(element_bytes, |acc: u64, d| acc.checked_mul(*d)) - .ok_or_else(|| bad("shape overflows u64 bytes".into()))?; - if expected != end - start { - return Err(bad(format!( - "shape {shape:?} of {dtype} implies {expected} bytes, header claims {}", - end - start - ))); - } - tensors.push(( - name.clone(), - TensorEntry { - dtype, - shape, - offsets: (start, end), - }, - )); - } - // File order (by payload offset) makes the copy pass a sequential read. - tensors.sort_by_key(|(_, t)| t.offsets.0); - Ok(tensors) - } - - /// Read one tensor's payload bytes. - fn read_payload(&self, entry: &TensorEntry, into: &mut [u8]) -> Result<(), SafetensorsError> { - assert_eq!( - into.len() as u64, - entry.byte_len(), - "read_payload: destination sized to the entry" - ); - let io = |source: std::io::Error| SafetensorsError::Io { - path: self.path.display().to_string(), - source, - }; - let mut file = File::open(&self.path).map_err(io)?; - file.seek(SeekFrom::Start(self.data_offset + entry.offsets.0)) - .map_err(io)?; - file.read_exact(into).map_err(io) - } -} - -/// Bytes per element for a safetensors dtype token, or `None` if unsupported. -#[must_use] -pub fn dtype_bytes(dtype: &str) -> Option { - Some(match dtype { - "BOOL" | "U8" | "I8" | "F8_E4M3" | "F8_E5M2" => 1, - "U16" | "I16" | "F16" | "BF16" => 2, - "U32" | "I32" | "F32" => 4, - "U64" | "I64" | "F64" => 8, - _ => return None, - }) -} - -/// What [`fuse_checkpoint`] should do with a source tensor. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum Fuse { - /// Copy through under this (possibly renamed) target name. - Keep(String), - /// Contribute to `target` as member `index` of `count`, stacked along a - /// NEW leading axis — the N:1 case (`mlp.experts.7.gate_proj.weight` is - /// member 7 of the fused `mlp.experts.gate`). - Stack { - target: String, - index: usize, - count: usize, - }, - /// Drop this tensor (present in the checkpoint, unused by the module). - Drop, -} - -/// Every safetensors shard of a checkpoint dir, in index order. -/// -/// A single `model.safetensors` is the one-shard case; otherwise -/// `model.safetensors.index.json`'s weight map names the shards. -pub fn checkpoint_shards(dir: &Path) -> Result, SafetensorsError> { - let single = dir.join("model.safetensors"); - if single.is_file() { - return Ok(vec![single]); - } - let index_path = dir.join("model.safetensors.index.json"); - if !index_path.is_file() { - return Err(SafetensorsError::NoCheckpoint(dir.to_path_buf())); - } - let bytes = std::fs::read(&index_path).map_err(|source| SafetensorsError::Io { - path: index_path.display().to_string(), - source, - })?; - let names = crate::hub::shards_from_index(&bytes, &index_path).map_err(|e| { - SafetensorsError::BadHeader { - path: index_path.display().to_string(), - reason: e.to_string(), - } - })?; - assert!( - !names.is_empty() && names.len() <= MAX_SHARDS, - "shards_from_index bounds the count" - ); - - // The index is a manifest, not evidence. Check every shard is actually on - // disk HERE, so a half-fetched checkpoint is one clear error at planning - // time rather than an `os error 2` raised after the first shards have - // already been opened — and so callers can use this function as the - // "is the checkpoint complete?" question it looks like. - let paths: Vec = names.into_iter().map(|n| dir.join(n)).collect(); - for path in &paths { - if path.is_file() { - continue; - } - let name = path.file_name().unwrap_or_default().to_string_lossy(); - let interrupted = path.with_file_name(format!("{name}.part")).is_file(); - return Err(SafetensorsError::MissingShard { - dir: dir.to_path_buf(), - missing: name.into_owned(), - hint: if interrupted { - " — a .part sibling is present, so the download was interrupted" - } else { - "" - }, - }); - } - debug_assert!( - paths.iter().all(|p| p.is_file()), - "every shard verified present before returning" - ); - Ok(paths) -} - -/// Read every shard of the checkpoint in `dir` and build ONE in-memory -/// safetensors blob, applying `map` to each source tensor name. -/// -/// Fused (`Fuse::Stack`) groups are concatenated in **numeric member order**, -/// gaining a leading `count` axis — `count` tensors of shape `[out, in]` -/// become one `[count, out, in]`. Every group must be exactly complete; -/// a missing or duplicate member is a loud [`SafetensorsError::BadGroup`], -/// raised during planning, before any payload byte is read. -pub fn fuse_checkpoint( - dir: &Path, - map: &dyn Fn(&str) -> Option, -) -> Result, SafetensorsError> { - let mut blob = Vec::new(); - fuse_into(dir, map, &mut blob)?; - Ok(blob) -} - -/// [`fuse_checkpoint`] straight to a file, never holding the payload in RAM. -/// -/// This is the variant a real checkpoint wants. The in-memory form needs the -/// whole fused payload resident — 13.8 GB for OLMoE-1B-7B — *on top of* the -/// model the load then builds (~28 GB at f32), and that sum is what a 128 GB -/// box with other tenants actually fails to satisfy. Writing to disk trades -/// the spike for temp space and lets `SafetensorsStore::from_file` page the -/// weights in as it needs them. -pub fn fuse_checkpoint_to_file( - dir: &Path, - map: &dyn Fn(&str) -> Option, - out: &Path, -) -> Result { - let io = |source: std::io::Error| SafetensorsError::Io { - path: out.display().to_string(), - source, - }; - let file = File::create(out).map_err(io)?; - let mut sink = std::io::BufWriter::with_capacity(1 << 20, file); - let written = fuse_into(dir, map, &mut sink)?; - sink.flush().map_err(io)?; - sink.into_inner() - .map_err(|e| SafetensorsError::Io { - path: out.display().to_string(), - source: e.into_error(), - })? - .sync_all() - .map_err(io)?; - Ok(written) -} - -/// The shared fuse: plan, then stream header + payload into `sink` in output -/// order, one part at a time. -/// -/// Writes are strictly ascending because `plan` carries tensors in output -/// order and each tensor's `parts` are already in destination order — so this -/// never needs the payload addressable at once, only ONE part at a time. -fn fuse_into( - dir: &Path, - map: &dyn Fn(&str) -> Option, - sink: &mut W, -) -> Result { - let shard_paths = checkpoint_shards(dir)?; - let headers = shard_paths - .iter() - .map(|p| SafetensorsHeader::open(p)) - .collect::, _>>()?; - - let plan = plan_output(&headers, map)?; - let total: u64 = plan.iter().map(|p| p.len).sum(); - if total > MAX_FUSED_BYTES { - return Err(SafetensorsError::OverBound { - path: dir.display().to_string(), - what: "fused payload bytes", - count: total, - bound: MAX_FUSED_BYTES, - }); - } - - // Header first: names, dtypes, shapes, and the contiguous offsets the - // copy pass will fill. - let mut header = String::from("{"); - for (i, p) in plan.iter().enumerate() { - if i > 0 { - header.push(','); - } - let json_name = - serde_json::to_string(&p.name).map_err(|e| SafetensorsError::BadTensor { - path: dir.display().to_string(), - name: p.name.clone(), - reason: format!("name is not encodable as JSON: {e}"), - })?; - header.push_str(&format!( - "{json_name}:{{\"dtype\":\"{}\",\"shape\":{:?},\"data_offsets\":[{},{}]}}", - p.dtype, - p.shape, - p.start, - p.start + p.len, - )); - } - header.push('}'); - - let io = |source: std::io::Error| SafetensorsError::Io { - path: dir.display().to_string(), - source, - }; - sink.write_all(&(header.len() as u64).to_le_bytes()) - .map_err(io)?; - sink.write_all(header.as_bytes()).map_err(io)?; - - // One reusable buffer, sized to the largest single part rather than to the - // payload: for a 64-expert bank that is one expert projection (~4 MB), not - // the 13.8 GB whole. - let widest = plan - .iter() - .flat_map(|p| p.parts.iter()) - .map(|(_, entry, _)| entry.byte_len()) - .max() - .unwrap_or(0); - if widest > MAX_PART_BYTES { - return Err(SafetensorsError::OverBound { - path: dir.display().to_string(), - what: "single tensor bytes", - count: widest, - bound: MAX_PART_BYTES, - }); - } - let mut buf = vec![0u8; to_usize(widest, dir, "single tensor bytes")?]; - - let mut written = 0u64; - for p in &plan { - debug_assert_eq!(written, p.start, "tensors are written in output order"); - for (shard, entry, dst_within) in &p.parts { - debug_assert_eq!( - written, - p.start + dst_within, - "parts are written in destination order" - ); - let len = to_usize(entry.byte_len(), dir, "member byte length")?; - let slot = &mut buf[..len]; - headers[*shard].read_payload(entry, slot)?; - sink.write_all(slot).map_err(io)?; - written += len as u64; - } - } - assert_eq!( - written, total, - "every planned byte was written exactly once" - ); - Ok(written) -} - -/// First pass: walk every shard and record, in first-seen order, what each -/// output target is built from. -/// -/// This is where a checkpoint's *claims* are checked — an unmapped name, a -/// member index outside `0..count`, two members disagreeing about the group -/// size, or the same member twice. What it does NOT check is completeness; -/// that is [`plan_output`]'s job, once every shard has been seen. -fn collect_groups( - headers: &[SafetensorsHeader], - map: &dyn Fn(&str) -> Option, -) -> Result<(Vec, HashMap), SafetensorsError> { - debug_assert!(!headers.is_empty(), "collecting needs a shard header"); - let mut order: Vec = Vec::new(); - let mut groups: HashMap = HashMap::new(); - - for (shard, header) in headers.iter().enumerate() { - for (name, entry) in &header.tensors { - let action = map(name).ok_or_else(|| SafetensorsError::BadTensor { - path: header.path.display().to_string(), - name: name.clone(), - reason: "unmapped tensor name".into(), - })?; - let (target, member, count) = match action { - Fuse::Drop => continue, - Fuse::Keep(target) => (target, 0usize, 1usize), - Fuse::Stack { - target, - index, - count, - } => { - if index >= count { - return Err(SafetensorsError::BadGroup { - target, - reason: format!("member index {index} outside 0..{count}"), - }); - } - (target, index, count) - } - }; - let group = groups.entry(target.clone()).or_insert_with(|| { - order.push(target.clone()); - Group { - count, - stacked: count > 1, - members: Vec::new(), - } - }); - if group.count != count { - return Err(SafetensorsError::BadGroup { - target, - reason: format!("member count disagrees ({} vs {count})", group.count), - }); - } - if group.members.iter().any(|(i, _, _)| *i == member) { - return Err(SafetensorsError::BadGroup { - target, - reason: format!("duplicate member {member}"), - }); - } - group.members.push((member, shard, entry.clone())); - } - } - - debug_assert_eq!( - order.len(), - groups.len(), - "every ordered target has exactly one group" - ); - Ok((order, groups)) -} - -/// Second pass: lay the collected groups out into the output, and apply THE -/// completeness check. -/// -/// Nothing here reads payload bytes — a malformed checkpoint fails before the -/// expensive pass. Completeness can only be judged after [`collect_groups`] -/// has seen every shard, because a group's members are free to be split -/// across shards in any order. -fn plan_output( - headers: &[SafetensorsHeader], - map: &dyn Fn(&str) -> Option, -) -> Result, SafetensorsError> { - assert!( - !headers.is_empty(), - "planning needs at least one shard header" - ); - assert!( - headers.len() <= MAX_SHARDS, - "shard count is bounded before planning" - ); - - let (order, mut groups) = collect_groups(headers, map)?; - let mut plan = Vec::with_capacity(order.len()); - let mut cursor = 0u64; - for target in order { - // `order` is pushed only from `or_insert_with`, so every target in it - // has a group and none repeats — an `Err` here would mean that - // invariant broke, not that the checkpoint is bad. - let Some(mut group) = groups.remove(&target) else { - return Err(SafetensorsError::BadGroup { - target, - reason: "planned target has no collected members".into(), - }); - }; - // THE load-bearing check: exactly `count` members, ids 0..count. - if group.members.len() != group.count { - return Err(SafetensorsError::BadGroup { - target, - reason: format!( - "{} of {} members present — a short group would load clean and compute wrong", - group.members.len(), - group.count - ), - }); - } - // Numeric member order, NOT the lexicographic order the names imply - // (…experts.10 sorts before …experts.2 as text). - group.members.sort_by_key(|(i, _, _)| *i); - debug_assert!( - group - .members - .iter() - .enumerate() - .all(|(i, (m, _, _))| i == *m), - "a complete, duplicate-free group is exactly 0..count once sorted" - ); - - let first = &group.members[0].2; - for (_, _, entry) in &group.members { - if entry.dtype != first.dtype || entry.shape != first.shape { - return Err(SafetensorsError::BadGroup { - target, - reason: format!( - "members disagree on layout ({} {:?} vs {} {:?})", - first.dtype, first.shape, entry.dtype, entry.shape - ), - }); - } - } - - let shape = if group.stacked { - std::iter::once(group.count as u64) - .chain(first.shape.iter().copied()) - .collect() - } else { - first.shape.clone() - }; - let member_bytes = first.byte_len(); - let len = member_bytes * group.count as u64; - let parts = group - .members - .iter() - .enumerate() - .map(|(slot, (_, shard, entry))| (*shard, entry.clone(), slot as u64 * member_bytes)) - .collect(); - plan.push(PlannedNamed { - name: target, - dtype: first.dtype.clone(), - shape, - start: cursor, - len, - parts, - }); - cursor += len; - } - assert!( - !plan.is_empty(), - "a checkpoint dir with shards yields at least one output tensor" - ); - Ok(plan) -} - -/// Accumulator for one output tensor while planning. -struct Group { - count: usize, - stacked: bool, - members: Vec<(usize, usize, TensorEntry)>, -} - -/// A planned output tensor: where it lands in the blob and what fills it. -struct PlannedNamed { - name: String, - dtype: String, - shape: Vec, - start: u64, - len: u64, - /// Sources, in destination order: (shard index, entry, byte offset within - /// this tensor). One part for `Keep`, `count` parts for `Stack`. - parts: Vec<(usize, TensorEntry, u64)>, -} - -#[cfg(test)] -mod tests { - use super::*; - - /// Build a synthetic safetensors file: `(name, dtype, shape, bytes)`. - fn write_st(path: &Path, tensors: &[(&str, &str, Vec, Vec)]) { - let mut header = String::from("{"); - let mut data: Vec = Vec::new(); - for (i, (name, dtype, shape, bytes)) in tensors.iter().enumerate() { - if i > 0 { - header.push(','); - } - let start = data.len(); - data.extend_from_slice(bytes); - header.push_str(&format!( - "{:?}:{{\"dtype\":\"{dtype}\",\"shape\":{shape:?},\"data_offsets\":[{start},{}]}}", - name, - data.len() - )); - } - header.push('}'); - let mut blob = (header.len() as u64).to_le_bytes().to_vec(); - blob.extend_from_slice(header.as_bytes()); - blob.extend_from_slice(&data); - std::fs::write(path, blob).unwrap(); - } - - /// A fresh empty scratch dir under the OS temp root. - fn scratch(tag: &str) -> PathBuf { - let dir = std::env::temp_dir().join(format!("mummu_st_{tag}")); - let _ = std::fs::remove_dir_all(&dir); - std::fs::create_dir_all(&dir).unwrap(); - dir - } - - /// One f32 tensor's little-endian bytes. - fn f32s(values: &[f32]) -> Vec { - values.iter().flat_map(|v| v.to_le_bytes()).collect() - } - - /// Read a tensor back out of a fused blob by name. - fn read_back(blob: &[u8], name: &str) -> (String, Vec, Vec) { - let header_len = u64::from_le_bytes(blob[..8].try_into().unwrap()) as usize; - let json: serde_json::Value = serde_json::from_slice(&blob[8..8 + header_len]).unwrap(); - let entry = &json[name]; - let dtype = entry["dtype"].as_str().unwrap().to_string(); - let shape: Vec = entry["shape"] - .as_array() - .unwrap() - .iter() - .map(|d| d.as_u64().unwrap()) - .collect(); - let start = entry["data_offsets"][0].as_u64().unwrap() as usize; - let end = entry["data_offsets"][1].as_u64().unwrap() as usize; - let base = 8 + header_len; - let (words, rest) = blob[base + start..base + end].as_chunks::<4>(); - assert!(rest.is_empty(), "f32 payload is a whole number of words"); - let values = words.iter().copied().map(f32::from_le_bytes).collect(); - (dtype, shape, values) - } - - #[test] - fn header_parses_shape_dtype_and_offsets() { - let dir = scratch("header"); - let path = dir.join("model.safetensors"); - write_st( - &path, - &[ - ("a", "F32", vec![2, 2], f32s(&[1.0, 2.0, 3.0, 4.0])), - ("b", "F32", vec![3], f32s(&[5.0, 6.0, 7.0])), - ], - ); - let h = SafetensorsHeader::open(&path).unwrap(); - assert_eq!(h.tensors.len(), 2); - assert_eq!(h.tensors[0].0, "a"); - assert_eq!(h.tensors[0].1.shape, vec![2, 2]); - assert_eq!(h.tensors[0].1.byte_len(), 16); - assert_eq!(h.tensors[0].1.element_count(), 4); - assert_eq!(h.tensors[1].1.offsets, (16, 28)); - std::fs::remove_dir_all(&dir).unwrap(); - } - - #[test] - fn header_rejects_a_shape_that_disagrees_with_its_bytes() { - let dir = scratch("badshape"); - let path = dir.join("model.safetensors"); - // Declares [4] f32 (16 B) but only 8 B of payload are claimed. - let mut blob = Vec::new(); - let header = r#"{"a":{"dtype":"F32","shape":[4],"data_offsets":[0,8]}}"#.to_string(); - blob.extend_from_slice(&(header.len() as u64).to_le_bytes()); - blob.extend_from_slice(header.as_bytes()); - blob.extend_from_slice(&[0u8; 8]); - std::fs::write(&path, blob).unwrap(); - - let err = SafetensorsHeader::open(&path).unwrap_err(); - assert!( - matches!(&err, SafetensorsError::BadTensor { reason, .. } - if reason.contains("implies 16 bytes")), - "got {err}" - ); - std::fs::remove_dir_all(&dir).unwrap(); - } - - #[test] - fn header_rejects_an_unsupported_dtype() { - let dir = scratch("baddtype"); - let path = dir.join("model.safetensors"); - write_st(&path, &[("a", "COMPLEX128", vec![1], vec![0u8; 16])]); - let err = SafetensorsHeader::open(&path).unwrap_err(); - assert!( - matches!(&err, SafetensorsError::BadTensor { reason, .. } - if reason.contains("unsupported dtype")), - "got {err}" - ); - std::fs::remove_dir_all(&dir).unwrap(); - } - - /// THE ordering test. Members are stacked in NUMERIC order; the names sort - /// lexicographically as 0, 1, 10, 2, … so a text sort silently permutes - /// the expert bank — a model that loads clean and computes wrong. - #[test] - fn stack_orders_members_numerically_not_lexicographically() { - let dir = scratch("order"); - let count = 12usize; - let tensors: Vec<(String, &str, Vec, Vec)> = (0..count) - .map(|i| { - ( - format!("mlp.experts.{i}.gate_proj.weight"), - "F32", - vec![1], - // Expert i holds exactly the value i. - f32s(&[i as f32]), - ) - }) - .collect(); - let refs: Vec<(&str, &str, Vec, Vec)> = tensors - .iter() - .map(|(n, d, s, b)| (n.as_str(), *d, s.clone(), b.clone())) - .collect(); - write_st(&dir.join("model.safetensors"), &refs); - - let blob = fuse_checkpoint(&dir, &|name| { - let idx: usize = name - .strip_prefix("mlp.experts.")? - .strip_suffix(".gate_proj.weight")? - .parse() - .ok()?; - Some(Fuse::Stack { - target: "mlp.experts.gate".into(), - index: idx, - count, - }) - }) - .unwrap(); - - let (dtype, shape, values) = read_back(&blob, "mlp.experts.gate"); - assert_eq!(dtype, "F32"); - assert_eq!(shape, vec![count as u64, 1], "gains a leading expert axis"); - // Slot i must hold expert i — the whole point. - assert_eq!( - values, - (0..count).map(|i| i as f32).collect::>(), - "experts must stack in numeric order (lexicographic would give 0,1,10,11,2,…)" - ); - std::fs::remove_dir_all(&dir).unwrap(); - } - - #[test] - fn stack_rejects_a_short_group() { - let dir = scratch("short"); - // Declare 4 members but ship only 3. - write_st( - &dir.join("model.safetensors"), - &[ - ("e.0", "F32", vec![1], f32s(&[0.0])), - ("e.1", "F32", vec![1], f32s(&[1.0])), - ("e.2", "F32", vec![1], f32s(&[2.0])), - ], - ); - let err = fuse_checkpoint(&dir, &|name| { - let idx: usize = name.strip_prefix("e.")?.parse().ok()?; - Some(Fuse::Stack { - target: "e".into(), - index: idx, - count: 4, - }) - }) - .unwrap_err(); - assert!( - matches!(&err, SafetensorsError::BadGroup { reason, .. } - if reason.contains("3 of 4 members")), - "got {err}" - ); - std::fs::remove_dir_all(&dir).unwrap(); - } - - #[test] - fn stack_rejects_duplicate_members_and_layout_disagreement() { - let dir = scratch("dup"); - write_st( - &dir.join("model.safetensors"), - &[ - ("a", "F32", vec![1], f32s(&[0.0])), - ("b", "F32", vec![1], f32s(&[1.0])), - ], - ); - // Both map to member 0 of the same target. - let err = fuse_checkpoint(&dir, &|_| { - Some(Fuse::Stack { - target: "t".into(), - index: 0, - count: 2, - }) - }) - .unwrap_err(); - assert!( - matches!(&err, SafetensorsError::BadGroup { reason, .. } - if reason.contains("duplicate member 0")), - "got {err}" - ); - - // Same group, different shapes. - let dir2 = scratch("layout"); - write_st( - &dir2.join("model.safetensors"), - &[ - ("a", "F32", vec![1], f32s(&[0.0])), - ("b", "F32", vec![2], f32s(&[1.0, 2.0])), - ], - ); - let err = fuse_checkpoint(&dir2, &|name| { - Some(Fuse::Stack { - target: "t".into(), - index: usize::from(name == "b"), - count: 2, - }) - }) - .unwrap_err(); - assert!( - matches!(&err, SafetensorsError::BadGroup { reason, .. } - if reason.contains("disagree on layout")), - "got {err}" - ); - std::fs::remove_dir_all(&dir).unwrap(); - std::fs::remove_dir_all(&dir2).unwrap(); - } - - #[test] - fn unmapped_tensor_is_a_loud_error_and_drop_is_explicit() { - let dir = scratch("unmapped"); - write_st( - &dir.join("model.safetensors"), - &[ - ("keep", "F32", vec![1], f32s(&[1.0])), - ("junk", "F32", vec![1], f32s(&[2.0])), - ], - ); - // No mapping for "junk" -> loud. - let err = fuse_checkpoint(&dir, &|name| { - (name == "keep").then(|| Fuse::Keep("keep".into())) - }) - .unwrap_err(); - assert!( - matches!(&err, SafetensorsError::BadTensor { reason, .. } - if reason.contains("unmapped tensor name")), - "got {err}" - ); - // Explicitly dropping it is fine, and it leaves the blob. - let blob = fuse_checkpoint(&dir, &|name| { - Some(if name == "keep" { - Fuse::Keep("keep".into()) - } else { - Fuse::Drop - }) - }) - .unwrap(); - let (_, shape, values) = read_back(&blob, "keep"); - assert_eq!((shape, values), (vec![1], vec![1.0])); - let header_len = u64::from_le_bytes(blob[..8].try_into().unwrap()) as usize; - let header = std::str::from_utf8(&blob[8..8 + header_len]).unwrap(); - assert!(!header.contains("junk"), "dropped tensor is absent"); - std::fs::remove_dir_all(&dir).unwrap(); - } - - /// A group whose members live in DIFFERENT shards still fuses correctly — - /// a real 3-shard checkpoint can split a layer across a boundary. - #[test] - fn shards_are_discovered_and_fused_across_boundaries() { - let dir = scratch("shards"); - write_st( - &dir.join("model-00001-of-00002.safetensors"), - &[("e.1", "F32", vec![1], f32s(&[11.0]))], - ); - write_st( - &dir.join("model-00002-of-00002.safetensors"), - &[("e.0", "F32", vec![1], f32s(&[10.0]))], - ); - std::fs::write( - dir.join("model.safetensors.index.json"), - br#"{"weight_map":{"e.1":"model-00001-of-00002.safetensors", - "e.0":"model-00002-of-00002.safetensors"}}"#, - ) - .unwrap(); - - let shards = checkpoint_shards(&dir).unwrap(); - assert_eq!(shards.len(), 2, "both shards discovered from the index"); - - let blob = fuse_checkpoint(&dir, &|name| { - let idx: usize = name.strip_prefix("e.")?.parse().ok()?; - Some(Fuse::Stack { - target: "e".into(), - index: idx, - count: 2, - }) - }) - .unwrap(); - let (_, shape, values) = read_back(&blob, "e"); - assert_eq!(shape, vec![2, 1]); - // Member 0 came from shard 2, member 1 from shard 1 — order is by - // member index, never by shard order. - assert_eq!(values, vec![10.0, 11.0]); - std::fs::remove_dir_all(&dir).unwrap(); - } - - #[test] - fn a_dir_with_no_checkpoint_is_a_loud_error() { - let dir = scratch("empty"); - assert!(matches!( - checkpoint_shards(&dir), - Err(SafetensorsError::NoCheckpoint(_)) - )); - std::fs::remove_dir_all(&dir).unwrap(); - } - - /// An index is a manifest, not evidence. A half-fetched checkpoint must - /// fail HERE — before any shard is opened — or it surfaces as a bare - /// `os error 2` from inside the load, and any caller using - /// `checkpoint_shards(..).is_ok()` as "is this checkpoint complete?" - /// silently believes an interrupted download is ready to use. - #[test] - fn an_index_naming_a_missing_shard_is_a_loud_error_not_a_late_os_error() { - let dir = scratch("missing_shard"); - write_st( - &dir.join("model-00001-of-00002.safetensors"), - &[("e.0", "F32", vec![1], f32s(&[10.0]))], - ); - std::fs::write( - dir.join("model.safetensors.index.json"), - br#"{"weight_map":{"e.0":"model-00001-of-00002.safetensors", - "e.1":"model-00002-of-00002.safetensors"}}"#, - ) - .unwrap(); - - // Shard 2 is absent entirely: named, no hint. - match checkpoint_shards(&dir) { - Err(SafetensorsError::MissingShard { missing, hint, .. }) => { - assert_eq!(missing, "model-00002-of-00002.safetensors"); - assert!( - hint.is_empty(), - "no .part sibling, so no interrupted-download hint" - ); - } - other => panic!("expected MissingShard, got {other:?}"), - } - - // The same shard as an interrupted download: the error says so, which - // is the difference between "re-fetch me" and "this repo is broken". - std::fs::write( - dir.join("model-00002-of-00002.safetensors.part"), - b"partial", - ) - .unwrap(); - match checkpoint_shards(&dir) { - Err(SafetensorsError::MissingShard { hint, .. }) => { - assert!( - hint.contains("interrupted"), - "hint names the .part sibling: {hint:?}" - ); - } - other => panic!("expected MissingShard, got {other:?}"), - } - - std::fs::remove_dir_all(&dir).unwrap(); - } - - /// The streaming fuse and the in-memory fuse must be the SAME bytes — - /// otherwise "load a big model from a file" and "load a small one from - /// RAM" are two different importers, and only one of them is tested. - #[test] - fn fusing_to_a_file_is_byte_identical_to_fusing_in_memory() { - let dir = scratch("fuse_to_file"); - write_st( - &dir.join("model-00001-of-00002.safetensors"), - &[ - ("e.1", "F32", vec![2], f32s(&[11.0, 12.0])), - ("keep", "F32", vec![1], f32s(&[7.0])), - ], - ); - write_st( - &dir.join("model-00002-of-00002.safetensors"), - &[("e.0", "F32", vec![2], f32s(&[10.0, 9.0]))], - ); - std::fs::write( - dir.join("model.safetensors.index.json"), - br#"{"weight_map":{"e.1":"model-00001-of-00002.safetensors", - "keep":"model-00001-of-00002.safetensors", - "e.0":"model-00002-of-00002.safetensors"}}"#, - ) - .unwrap(); - - let map = |name: &str| -> Option { - if name == "keep" { - return Some(Fuse::Keep("keep".into())); - } - let idx: usize = name.strip_prefix("e.")?.parse().ok()?; - Some(Fuse::Stack { - target: "e".into(), - index: idx, - count: 2, - }) - }; - - let in_memory = fuse_checkpoint(&dir, &map).unwrap(); - let out = dir.join("fused.safetensors"); - let written = fuse_checkpoint_to_file(&dir, &map, &out).unwrap(); - let on_disk = std::fs::read(&out).unwrap(); - - assert_eq!( - in_memory, on_disk, - "the two fuse paths must agree byte for byte" - ); - // The return value counts PAYLOAD bytes, so the file is that plus the - // 8-byte length prefix and the JSON header. - assert!( - written < on_disk.len() as u64, - "payload {written} B sits inside a {} B file", - on_disk.len() - ); - - // And the fused content is still correct, not merely self-consistent. - let (_, shape, values) = read_back(&on_disk, "e"); - assert_eq!(shape, vec![2, 2]); - assert_eq!(values, vec![10.0, 9.0, 11.0, 12.0], "numeric member order"); - - std::fs::remove_dir_all(&dir).unwrap(); - } - - #[test] - fn dtype_bytes_covers_the_float_and_int_tokens() { - assert_eq!(dtype_bytes("BF16"), Some(2)); - assert_eq!(dtype_bytes("F32"), Some(4)); - assert_eq!(dtype_bytes("I64"), Some(8)); - assert_eq!(dtype_bytes("BOOL"), Some(1)); - assert_eq!(dtype_bytes("NOPE"), None); - } -} diff --git a/crates/mummu/examples/src/template.rs b/crates/mummu/examples/src/template.rs deleted file mode 100644 index 9e04d25..0000000 --- a/crates/mummu/examples/src/template.rs +++ /dev/null @@ -1,413 +0,0 @@ -//! Render a checkpoint's **own** imported chat template — the general -//! fallback for models whose family has no hardcoded [`crate::chat`] -//! renderer. Behind the non-default `jinja-template` feature. -//! -//! The zoo's prompt wrapping is from-scratch and byte-verified: one -//! [`ChatMl`] constructor per family, each proven byte-identical to -//! `transformers.apply_chat_template` on the real checkpoint's template by -//! `tests/template_gate.rs`. That is the right shape for a model Mummu has -//! ported — the bytes are pinned by a test, not by a template file that can -//! change under us. It is no shape at all for a model Mummu has *not* ported, -//! which is exactly what the import suite (P3) is for: a checkpoint arrives -//! with a `chat_template` nobody has written a renderer for. -//! -//! So the rule this module encodes, and the one a consumer should follow: -//! -//! - **A family renderer exists → use it.** [`ChatMl::qwen2`], -//! [`ChatMl::qwen3`], [`ChatMl::lfm2`]. Byte-pinned, no Jinja at runtime. -//! - **No family renderer → use [`ImportedTemplate`].** The checkpoint's own -//! template is the authority on its own prompt format, and rendering it is -//! strictly better than guessing ChatML. -//! -//! [`Renderer`] is that rule as a value, for a consumer that holds one -//! renderer and does not want to branch at every call site. -//! -//! What this module does NOT do is replace the family renderers. The gate -//! proved the two agree today on Qwen2.5/Qwen3/LFM2.5; that agreement is a -//! *result*, and the from-scratch path stays the shipping one. - -use std::path::Path; - -use hf_chat_template::{ - ChatTemplate, ChatTemplateField, Message, RenderInput, TokenField, - TokenizerConfig as HfTokenizerConfig, -}; - -use crate::chat::{ChatMl, MAX_TOOLS, MAX_TURNS, Role, ToolSpec, Turn}; -use crate::tok_config::{SpecialToken, TokenizerConfig}; - -/// Largest prompt a render will return. A chat template is a Jinja program -/// from an untrusted checkpoint; a loop over a long history can produce far -/// more than it was handed. Real prompts are kilobytes — a 8 MiB result is a -/// runaway template, not a conversation. -const MAX_RENDERED_BYTES: usize = 8 * 1024 * 1024; - -/// What went wrong rendering a checkpoint's imported template. -#[derive(Debug, thiserror::Error)] -pub enum TemplateError { - /// The checkpoint declares no chat template at all — neither the - /// `chat_template` key of `tokenizer_config.json` nor a standalone - /// `chat_template.jinja` beside it. There is nothing to render with. - #[error("checkpoint declares no chat_template")] - Absent, - /// The template is not valid Jinja, or failed while rendering (a - /// `raise_exception` in the template lands here too, with its message). - #[error("chat template: {0}")] - Jinja(String), - /// The render produced more than [`MAX_RENDERED_BYTES`]. - #[error("chat template rendered {got} bytes, over the {MAX_RENDERED_BYTES} byte bound")] - TooLarge { got: usize }, - /// A tool signature could not be serialized to JSON for the template. - #[error("tool {name:?}: {reason}")] - BadTool { name: String, reason: String }, -} - -impl From for TemplateError { - fn from(e: hf_chat_template::Error) -> Self { - TemplateError::Jinja(e.to_string()) - } -} - -/// A compiled chat template imported from a checkpoint. -/// -/// Compiling is the expensive half (Jinja parse); hold one per model and -/// render many prompts from it. -pub struct ImportedTemplate { - inner: ChatTemplate, -} - -impl std::fmt::Debug for ImportedTemplate { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.write_str("ImportedTemplate") - } -} - -/// Our resolved special-token slot in the shape the Jinja context wants. -fn token_field(slot: &Option) -> Option { - slot.as_ref().map(|t| TokenField::Str(t.content.clone())) -} - -impl ImportedTemplate { - /// Compile the template a parsed [`TokenizerConfig`] carries. - /// - /// The config's BOS/EOS/PAD/UNK slots go into the render context under - /// the names templates use (`bos_token`, …) — many families' templates - /// end a turn with `{{ eos_token }}` rather than a literal, so a template - /// compiled without them renders a prompt the model never saw in - /// training. - pub fn from_config(config: &TokenizerConfig) -> Result { - let source = config - .chat_template - .as_deref() - .ok_or(TemplateError::Absent)?; - assert!( - !source.trim().is_empty(), - "TokenizerConfig never stores a blank chat_template" - ); - let hf = HfTokenizerConfig { - chat_template: Some(ChatTemplateField::Single(source.to_string())), - bos_token: token_field(&config.bos_token), - eos_token: token_field(&config.eos_token), - pad_token: token_field(&config.pad_token), - unk_token: token_field(&config.unk_token), - extra: Default::default(), - }; - let inner = ChatTemplate::from_tokenizer_config(&hf)?; - Ok(Self { inner }) - } - - /// Read a checkpoint directory's `tokenizer_config.json` (falling back to - /// a standalone `chat_template.jinja`, per [`TokenizerConfig::from_dir`]) - /// and compile what it declares. - pub fn from_dir(dir: &Path) -> Result { - assert!(!dir.as_os_str().is_empty(), "from_dir: empty dir"); - let config = - TokenizerConfig::from_dir(dir).map_err(|e| TemplateError::Jinja(e.to_string()))?; - Self::from_config(&config) - } - - /// Render a conversation, with the assistant generation prefix appended — - /// the same contract as [`ChatMl::render`]. - pub fn render(&self, turns: &[Turn]) -> Result { - self.render_with_tools(&[], turns) - } - - /// Render a conversation that advertises `tools`, with the assistant - /// generation prefix appended — the same contract as - /// [`ChatMl::render_with_tools`]. The template decides *where* and *how* - /// the signatures appear; that is the whole point of using it. - /// - /// Each tool is handed over in the shape `transformers` itself produces - /// (`get_json_schema`): `{"type": "function", "function": {name, - /// description, parameters}}`. That is what the mainstream templates - /// (Hermes/Qwen and everything modelled on them) unpack — but the key is - /// *open*: a template runs `tool | tojson` on whatever it is given, and - /// some families want the signature bare (LFM2.5 does; its own renderer - /// covers it). Use [`Self::render_with_tools_json`] when a checkpoint's - /// template wants a different shape. - pub fn render_with_tools( - &self, - tools: &[ToolSpec], - turns: &[Turn], - ) -> Result { - let json: Vec = tools.iter().map(tool_json).collect::>()?; - self.render_with_tools_json(&json, turns) - } - - /// Render with tool signatures given as raw JSON, for a template whose - /// `tools` shape is not the `transformers` default (see - /// [`Self::render_with_tools`]). - pub fn render_with_tools_json( - &self, - tools: &[serde_json::Value], - turns: &[Turn], - ) -> Result { - assert!( - turns.len() <= MAX_TURNS, - "imported render: {} turns exceeds the {MAX_TURNS} bound", - turns.len() - ); - assert!( - tools.len() <= MAX_TOOLS, - "imported render: {} tools exceeds the {MAX_TOOLS} bound", - tools.len() - ); - let input = RenderInput { - messages: turns.iter().map(message_from).collect(), - tools: tools.to_vec(), - add_generation_prompt: true, - ..RenderInput::default() - }; - let out = self.inner.render(&input)?; - if out.len() > MAX_RENDERED_BYTES { - return Err(TemplateError::TooLarge { got: out.len() }); - } - Ok(out) - } -} - -/// One [`Turn`] as the message shape `transformers` hands a template. -/// -/// An assistant turn carrying structured `tool_calls` passes them as data and -/// drops its `content` — the calls' rendered markers in [`Turn::content`] are -/// the *family* renderer's wire format, and re-emitting them here would -/// double-wrap under a template that writes its own. A turn is either a -/// tool-call turn or a text turn, never both, which is what the constructors -/// in [`crate::chat`] build. -fn message_from(turn: &Turn) -> Message { - let role = match turn.role { - Role::System => "system", - Role::User => "user", - Role::Assistant => "assistant", - Role::Tool => "tool", - }; - if turn.role == Role::Assistant && !turn.tool_calls.is_empty() { - let mut m = Message::new(role, ""); - m.content = None; - m.tool_calls = turn - .tool_calls - .iter() - .map(|c| serde_json::to_value(c).unwrap_or(serde_json::Value::Null)) - .collect(); - debug_assert_eq!(m.tool_calls.len(), turn.tool_calls.len()); - return m; - } - Message::new(role, turn.content.clone()) -} - -/// One [`ToolSpec`] in the Hermes wire shape templates expect under `tools`: -/// `{"type": "function", "function": {name, description, parameters}}`. -fn tool_json(spec: &ToolSpec) -> Result { - let function = serde_json::to_value(spec).map_err(|e| TemplateError::BadTool { - name: spec.name.clone(), - reason: e.to_string(), - })?; - let mut wire = serde_json::Map::new(); - wire.insert("type".into(), serde_json::Value::String("function".into())); - wire.insert("function".into(), function); - Ok(serde_json::Value::Object(wire)) -} - -/// Which renderer a model uses, as a value. -/// -/// Construct [`Renderer::Family`] whenever the architecture is one Mummu has -/// ported (its bytes are pinned by the template gate); fall back to -/// [`Renderer::Imported`] for anything else. [`Renderer::for_checkpoint`] -/// applies exactly that rule. -#[derive(Debug)] -pub enum Renderer { - /// A byte-verified from-scratch renderer. - Family(ChatMl), - /// The checkpoint's own Jinja template. - Imported(ImportedTemplate), -} - -impl Renderer { - /// Take `family` when the caller has one for this architecture, else - /// compile the checkpoint's own template from `dir`. - /// - /// The family renderer is not second-guessed: passing `Some` never reads - /// the template. That is deliberate — the gate proves the family - /// renderers match, and a checkpoint repackaged with a foreign template - /// is caught at load by the consistency gate in `tokenizer.rs`, not - /// silently obeyed here. - pub fn for_checkpoint(family: Option, dir: &Path) -> Result { - match family { - Some(chat_ml) => Ok(Self::Family(chat_ml)), - None => ImportedTemplate::from_dir(dir).map(Self::Imported), - } - } - - /// Render a conversation with the generation prefix appended. - pub fn render(&self, turns: &[Turn]) -> Result { - match self { - Self::Family(c) => Ok(c.render(turns)), - Self::Imported(t) => t.render(turns), - } - } - - /// Render a tool-advertising conversation with the generation prefix. - pub fn render_with_tools( - &self, - tools: &[ToolSpec], - turns: &[Turn], - ) -> Result { - match self { - Self::Family(c) => Ok(c.render_with_tools(tools, turns)), - Self::Imported(t) => t.render_with_tools(tools, turns), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::chat::ToolCall; - - /// A minimal ChatML-shaped template, in the same Jinja dialect a real - /// checkpoint ships: enough to exercise roles, tools and tool calls - /// without needing a multi-GB checkpoint on disk. - const TOY_TEMPLATE: &str = concat!( - "{%- if tools %}{{ tools | tojson }}\n{%- endif %}", - "{%- for m in messages %}", - "<|{{ m.role }}|>", - "{%- if m.tool_calls %}", - "{%- for c in m.tool_calls %}[call {{ c.name }} {{ c.arguments | tojson }}]{%- endfor %}", - "{%- else %}{{ m.content }}{%- endif %}", - "{{ eos_token }}", - "{%- endfor %}", - "{%- if add_generation_prompt %}<|assistant|>{%- endif %}" - ); - - fn config_with(template: Option<&str>) -> TokenizerConfig { - TokenizerConfig { - chat_template: template.map(str::to_string), - eos_token: Some(SpecialToken { - content: "".into(), - id: Some(7), - special: true, - }), - ..TokenizerConfig::default() - } - } - - #[test] - fn renders_roles_and_the_generation_prompt_from_the_imported_template() { - let t = ImportedTemplate::from_config(&config_with(Some(TOY_TEMPLATE))) - .expect("toy template compiles"); - let out = t - .render(&[Turn::system("be brief"), Turn::user("hi")]) - .expect("renders"); - assert_eq!( - out, "<|system|>be brief<|user|>hi<|assistant|>", - "roles, the config's eos_token and the generation prefix all reach the template" - ); - } - - /// Tools reach the template as structured JSON in the Hermes wire shape, - /// and the template — not us — decides where they land. - #[test] - fn tools_reach_the_template_as_hermes_wire_json() { - let t = ImportedTemplate::from_config(&config_with(Some(TOY_TEMPLATE))).expect("compiles"); - let spec = ToolSpec { - name: "get_weather".into(), - description: "weather".into(), - parameters: serde_json::json!({"type": "object"}), - }; - let out = t - .render_with_tools(&[spec], &[Turn::user("weather?")]) - .expect("renders"); - assert!(out.starts_with("["), "tools block leads: {out}"); - // `tojson` spells separators the Python way (`": "`, `", "`) — the - // same spelling `chat::python_json` pins for the from-scratch path. - assert!( - out.contains(r#""type": "function""#) && out.contains(r#""name": "get_weather""#), - "the tool arrives wrapped as a function spec: {out}" - ); - } - - /// The load-bearing difference from the family renderers: an assistant - /// tool-call turn passes its calls as DATA, so the template writes its - /// own markers instead of inheriting Hermes' `` wrapping. - #[test] - fn assistant_tool_calls_pass_structurally_not_as_hermes_markers() { - let t = ImportedTemplate::from_config(&config_with(Some(TOY_TEMPLATE))).expect("compiles"); - let calls = [ToolCall { - name: "get_weather".into(), - arguments: serde_json::json!({"city": "Paris"}), - }]; - let out = t - .render(&[ - Turn::user("weather?"), - Turn::assistant_tool_calls(&calls), - Turn::tool_response("{\"c\": 21}"), - ]) - .expect("renders"); - assert!( - out.contains("[call get_weather "), - "the template wrote its own call markers: {out}" - ); - assert!( - !out.contains(""), - "Hermes markers must NOT leak through as content: {out}" - ); - assert!(out.contains("<|tool|>{\"c\": 21}"), "tool role turn: {out}"); - } - - #[test] - fn a_checkpoint_without_a_template_is_a_loud_absent() { - let err = ImportedTemplate::from_config(&config_with(None)).unwrap_err(); - assert!(matches!(err, TemplateError::Absent), "got {err:?}"); - } - - #[test] - fn a_broken_template_is_a_loud_jinja_error_not_a_panic() { - let err = ImportedTemplate::from_config(&config_with(Some("{% for x in %}"))) - .expect_err("unbalanced Jinja must not compile"); - assert!(matches!(err, TemplateError::Jinja(_)), "got {err:?}"); - } - - /// A runaway template (a loop that multiplies its input) must hit the - /// byte bound rather than return a prompt nothing can tokenize. - #[test] - fn a_runaway_render_trips_the_byte_bound() { - let bomb = "{%- for _ in range(4000) %}{{ messages[0].content }}{%- endfor %}"; - let t = ImportedTemplate::from_config(&config_with(Some(bomb))).expect("compiles"); - let big = "x".repeat(4096); - let err = t - .render(&[Turn::user(big)]) - .expect_err("must trip the bound"); - assert!(matches!(err, TemplateError::TooLarge { .. }), "got {err:?}"); - } - - /// `Renderer` never second-guesses a family renderer: given one, it does - /// not read the checkpoint dir at all (here: a dir that does not exist). - #[test] - fn renderer_prefers_the_family_renderer_without_touching_the_checkpoint() { - let r = Renderer::for_checkpoint(Some(ChatMl::qwen3()), Path::new("no/such/dir")) - .expect("a family renderer needs no files"); - assert!(matches!(r, Renderer::Family(_))); - let ours = r.render(&[Turn::user("hi")]).expect("renders"); - assert_eq!(ours, ChatMl::qwen3().render(&[Turn::user("hi")])); - } -} diff --git a/crates/mummu/examples/src/tier.rs b/crates/mummu/examples/src/tier.rs deleted file mode 100644 index d08c7b5..0000000 --- a/crates/mummu/examples/src/tier.rs +++ /dev/null @@ -1,528 +0,0 @@ -//! **Tier planning** — P9 stage 3(b): which device runs which MoE expert at -//! which stored precision, all at once. -//! -//! A `.mummu` pack holds every expert at every level ([`crate::pack`]), so -//! the runtime can keep *different* experts resident on *different* -//! devices at *different* precisions simultaneously — int4 on the CPU, -//! int8 on an integrated GPU, f32 on the discrete card — and move them as -//! routing statistics shift ("hot-swapping"). This module decides the -//! assignment; [`crate::nn::ExpertPool`] executes it. -//! -//! The plan is two-phase and deterministic: -//! -//! 1. **Admission** — every expert gets the *cheapest* slot that exists -//! (smallest bytes on the slowest device), so a plan either exists for -//! all experts or fails loudly with the shortfall. Nothing is silently -//! dropped. -//! 2. **Promotion** — experts are visited hottest-first and each moves to -//! the most desirable tier that still fits: faster device first, then -//! higher precision within it. Capacity is checked against the bytes -//! already promised, so the result never oversubscribes a device. -//! -//! Hotness is the caller's (routing hit counts smoothed over requests); -//! uniform hotness degrades to "fill the best device in expert order". -//! Re-planning with new hotness and diffing against the live plan gives -//! the swap list — the pool applies exactly those moves. - -use std::collections::BTreeMap; - -pub use crate::pack::Precision; - -/// What kind of device a tier lives on — the planner's speed ordering -/// follows `speed`, not this tag; it exists for labels and defaults. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, serde::Serialize, serde::Deserialize)] -pub enum DeviceClass { - Cpu, - IntegratedGpu, - DiscreteGpu, -} - -/// One device the planner may place experts on. -#[derive(Debug, Clone, PartialEq)] -pub struct TierDevice { - pub name: String, - pub class: DeviceClass, - /// The precision ladder this device runs experts at, **best first** - /// (e.g. a discrete GPU `[F32, Q8]`, a CPU `[Q8, Q4]`). A level absent - /// here is never placed on this device, whatever the pack stores. - pub ladder: Vec, - /// Work this device is already committed to every step, in expert - /// equivalents. The trunk lives on exactly one device and runs on every - /// token; without it that device looks idle to the scheduler. - pub preload_units: usize, - /// Bytes of expert weights this device may hold (after whatever else - /// — the trunk, caches, the desktop — already lives there). - pub budget_bytes: u64, - /// Relative throughput rank; higher runs hotter experts. Ties break - /// toward the earlier device in the slice. - pub speed: u32, -} - -/// One expert's placement: device index into the planner's device slice -/// and the stored level it is loaded at. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)] -pub struct Tier { - pub device: usize, - pub precision: Precision, -} - -/// Resident bytes of one expert at each level the pack stores for it -/// (values + scales), as a device would hold them. Levels a device's -/// backend widens (f16 on an f32 backend) are the caller's to cost. -#[derive(Debug, Clone, Default, PartialEq, Eq)] -pub struct ExpertCost { - pub bytes: BTreeMap, -} - -/// A complete assignment: one tier per expert (planner input order) and the -/// bytes each device ends up holding. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct TierPlan { - pub tiers: Vec, - pub used_bytes: Vec, -} - -impl TierPlan { - /// Experts whose tier differs between `self` (live) and `next`: the - /// swap list, in expert order. - #[must_use] - pub fn diff(&self, next: &TierPlan) -> Vec<(usize, Tier)> { - self.tiers - .iter() - .zip(&next.tiers) - .enumerate() - .filter(|(_, (a, b))| a != b) - .map(|(i, (_, b))| (i, *b)) - .collect() - } - - /// Experts per (device, precision), for logs and `/api/ps`-style summaries. - #[must_use] - pub fn histogram(&self) -> BTreeMap<(usize, Precision), usize> { - let mut h = BTreeMap::new(); - for t in &self.tiers { - *h.entry((t.device, t.precision)).or_insert(0) += 1; - } - h - } -} - -/// Plan tiers for `costs.len()` experts over `devices`, hottest-first. -/// `hotness` is per expert (any non-negative scale; empty = uniform). -/// -/// Errors when some expert fits nowhere even at its cheapest level — the -/// message names the expert and the shortfall, never a silent drop. -pub fn plan_tiers( - devices: &[TierDevice], - costs: &[ExpertCost], - hotness: &[f64], -) -> Result { - if devices.is_empty() { - return Err("tier plan: no devices".into()); - } - if !hotness.is_empty() && hotness.len() != costs.len() { - return Err(format!( - "tier plan: {} hotness values for {} experts", - hotness.len(), - costs.len() - )); - } - // Device visiting order: fastest first (stable on ties). - let mut by_speed: Vec = (0..devices.len()).collect(); - by_speed.sort_by(|&a, &b| devices[b].speed.cmp(&devices[a].speed)); - - // Desirability order of every (device, precision) slot: faster device - // first, then that device's ladder best-first. - let slots: Vec = by_speed - .iter() - .flat_map(|&d| { - devices[d].ladder.iter().map(move |&p| Tier { - device: d, - precision: p, - }) - }) - .collect(); - // The admission order is the reverse: cheapest bytes on the slowest device. - // Host residency differs from stored size for Q4: flex unpacks nibbles - // to one i8 per element at load (`q_from_data`), so a Q4 blob occupies - // 1.125 B/elem resident against 0.625 stored — exactly 9/5. Without - // this the planner overpacks the host by 1.8x and the commit charge, - // not the budget, becomes the limit. - let cost_of = |tier: Tier, e: usize| -> Option { - let stored = costs[e].bytes.get(&tier.precision).copied()?; - let host_q4 = devices[tier.device].class == DeviceClass::Cpu - && tier.precision == Precision::Q4; - Some(if host_q4 { stored * 9 / 5 } else { stored }) - }; - - let mut used = vec![0u64; devices.len()]; - let mut tiers: Vec = Vec::with_capacity(costs.len()); - - // 1. Admission: every expert at its cheapest available slot. - for (e, cost) in costs.iter().enumerate() { - let mut best: Option<(u64, Tier)> = None; - for &slot in slots.iter().rev() { - let Some(bytes) = cost_of(slot, e) else { continue }; - if used[slot.device] + bytes <= devices[slot.device].budget_bytes - && best.is_none_or(|(b, _)| bytes < b) - { - best = Some((bytes, slot)); - } - } - let Some((bytes, slot)) = best else { - let cheapest = cost.bytes.values().copied().min().unwrap_or(0); - let free: u64 = devices - .iter() - .zip(&used) - .map(|(d, &u)| d.budget_bytes.saturating_sub(u)) - .max() - .unwrap_or(0); - return Err(format!( - "tier plan: expert {e} ({cheapest} bytes at its cheapest level) fits no device — \ - largest free budget is {free} bytes after admitting experts 0..{e}" - )); - }; - used[slot.device] += bytes; - tiers.push(slot); - } - - // 2a. Scheduler A: how many experts each device SHOULD hold. - // - // Admission below put everything on the slowest device and promotion - // pulls it up, fastest-slot-first — which fills the quick device and - // dumps the remainder on the slow ones. That is the right shape when - // devices run one after another, and the wrong one now that they run - // concurrently: the layer then costs the slowest device's share, so the - // objective is for every device to FINISH TOGETHER, not for the fast one - // to be busiest. Measured, fill-first put 996 experts on a device that - // takes 1.59 ms each and 1052 on devices taking ~14 ms — the slow side - // ran ~9x longer and decided the layer. - // - // `schedule::divide` gives the makespan-minimizing split; promotion - // treats it as a quota rather than a target, so a device is never filled - // past its share while a slower one still has work it could have taken. - let quota = { - let sched: Vec = devices - .iter() - .map(|dev| { - // Cheapest an expert can be on this device, over the whole - // set — what its budget divides into. - let cheapest = costs - .iter() - .filter_map(|c| { - dev.ladder.iter().filter_map(|p| c.bytes.get(p).copied()).min() - }) - .max() - .unwrap_or(u64::MAX); - crate::schedule::Device { - name: dev.name.clone(), - throughput: f64::from(dev.speed), - capacity_units: if cheapest == 0 || cheapest == u64::MAX { - 0 - } else { - (dev.budget_bytes / cheapest) as usize - }, - // Work this device already owes every token, in cluster - // equivalents — the trunk, for whichever device holds it. - // Counting it is what stops a host that is already - // saturated from being handed a "fair share" on top. - preload_units: dev.preload_units, - } - }) - .collect(); - crate::schedule::divide(&sched, costs.len()).units - }; - let mut held = vec![0usize; devices.len()]; - for t in &tiers { - held[t.device] += 1; - } - - // 2b. Promotion, hottest first. - let mut order: Vec = (0..costs.len()).collect(); - if !hotness.is_empty() { - order.sort_by(|&a, &b| { - hotness[b] - .partial_cmp(&hotness[a]) - .unwrap_or(std::cmp::Ordering::Equal) - .then(a.cmp(&b)) - }); - } - // Two phases, because placement and precision optimize different things - // and one pass let the wrong one win. - // - // Moving an expert from a 14.15 ms device to a 1.59 ms one saves 12.6 ms - // EVERY token. Upgrading the precision of an expert already on the fast - // device saves nothing — for a quantized checkpoint the extra bytes buy - // no accuracy at all (measured 0.0000 relative error for f16 against a - // 4.55 bits/param source) — and those bytes are capacity another expert - // could have used. Interleaved, precision won: the fast device ended up - // holding 610 experts at mixed rungs while the slow one held 1374. - // - // So: fill the fast devices first, at their CHEAPEST rung, and only then - // spend whatever is left over on precision. - - // Phase 1 — placement. Hottest first, onto the fastest device with room. - for &e in &order { - let cur = tiers[e]; - let cur_bytes = cost_of(cur, e).expect("admitted tier has a cost"); - for &d in &by_speed { - if d == cur.device { - continue; - } - if held[d] >= quota[d] { - continue; // scheduler A's balanced share for this device - } - // Move when the destination is faster — or when the CURRENT - // holder is past its own balanced share, even at equal-or-lower - // speed. The quota is the arbiter, not raw speed: it already - // prices in preloaded work, so a device level with the CPU on - // throughput is still a win while the CPU carries the trunk. - // The old strict "only strictly faster" gate made an - // equal-speed device UNREACHABLE — admission parks everything - // on the cheapest host slot, and with the integrated GPU rated - // 71 against the host's 72 it received zero clusters in every - // default plan (found by adversarial review, verified against - // this function line by line). - if devices[d].speed <= devices[cur.device].speed && held[cur.device] <= quota[cur.device] - { - continue; // no makespan gain from this move - } - // Cheapest rung this device accepts — capacity beats precision. - let Some((bytes, slot)) = devices[d] - .ladder - .iter() - .filter_map(|&p| { - let slot = Tier { device: d, precision: p }; - cost_of(slot, e).map(|b| (b, slot)) - }) - .min_by_key(|&(b, _)| b) - else { - continue; - }; - if used[d] + bytes <= devices[d].budget_bytes { - used[cur.device] -= cur_bytes; - used[d] += bytes; - held[cur.device] -= 1; - held[d] += 1; - tiers[e] = slot; - break; - } - } - } - - // Phase 2 — precision, with what is left, and without moving anything. - for &e in &order { - let cur = tiers[e]; - let cur_bytes = cost_of(cur, e).expect("admitted tier has a cost"); - for &p in &devices[cur.device].ladder { - let slot = Tier { device: cur.device, precision: p }; - if slot == cur { - break; // the ladder is best-first: nothing finer remains - } - let Some(bytes) = cost_of(slot, e) else { continue }; - if used[cur.device] - cur_bytes + bytes <= devices[cur.device].budget_bytes { - used[cur.device] = used[cur.device] - cur_bytes + bytes; - tiers[e] = slot; - break; - } - } - } - - Ok(TierPlan { - tiers, - used_bytes: used, - }) -} - -/// Hotness smoothing across requests: an exponential moving average of -/// per-expert routing hits, normalized per request so a long prompt does -/// not drown a short one. `alpha` in (0, 1]: 1 = this request only. -pub fn smooth_hotness(prev: &mut Vec, hits: &[u64], alpha: f64) { - if prev.len() != hits.len() { - *prev = vec![0.0; hits.len()]; - } - let total: u64 = hits.iter().sum(); - if total == 0 { - return; - } - let inv = 1.0 / total as f64; - for (p, &h) in prev.iter_mut().zip(hits) { - *p = (1.0 - alpha) * *p + alpha * (h as f64 * inv); - } -} - -#[cfg(test)] -mod tests { - use super::*; - - fn cost(q4: u64, q8: u64, f32: u64) -> ExpertCost { - ExpertCost { - bytes: [(Precision::Q4, q4), (Precision::Q8, q8), (Precision::F32, f32)] - .into_iter() - .collect(), - } - } - - fn devices(cpu_budget: u64, gpu_budget: u64) -> Vec { - vec![ - TierDevice { - name: "cpu".into(), - class: DeviceClass::Cpu, - ladder: vec![Precision::Q8, Precision::Q4], - budget_bytes: cpu_budget, - speed: 1, - preload_units: 0, - }, - TierDevice { - name: "gpu".into(), - class: DeviceClass::DiscreteGpu, - ladder: vec![Precision::F32, Precision::Q8], - budget_bytes: gpu_budget, - speed: 10, - preload_units: 0, - }, - ] - } - - /// The fast device is filled with EXPERTS before any of its bytes go on - /// precision. - /// - /// This reverses the earlier policy, deliberately and on measurement. Here - /// the GPU's 8 bytes hold either one f32 expert or four q8 ones, and it is - /// 10x the CPU's speed: four experts on it beats one on it and three on - /// the CPU, by a wide margin. The old ordering interleaved placement and - /// precision, so precision won — on the real 27B that left the 1.59 ms - /// device holding 610 clusters while the 14.15 ms device held 1374. - /// - /// Precision above the source buys nothing anyway (measured 0.0000 - /// relative error for f16 against a 4.55 bits/param checkpoint), while a - /// byte spent on it is capacity another expert could have used. - #[test] - fn the_fast_device_is_filled_with_experts_before_precision() { - let costs = vec![cost(1, 2, 8); 4]; - let plan = plan_tiers(&devices(100, 8), &costs, &[0.1, 0.5, 0.3, 0.1]).unwrap(); - assert!( - plan.tiers.iter().all(|t| t.device == 1), - "every expert should be on the fast device: {:?}", - plan.tiers - ); - assert!( - plan.tiers.iter().all(|t| t.precision == Precision::Q8), - "at its cheapest rung, not its best: {:?}", - plan.tiers - ); - assert_eq!(plan.used_bytes, vec![0, 8]); - } - - /// An equal-speed device must still receive work when the current - /// holder is over its balanced share. The trunk-preloaded host and the - /// integrated GPU are level on measured throughput; the old strictly- - /// faster gate made the iGPU unreachable — every cluster admitted to - /// the cheap host slot and stayed there, while scheduler A's quota - /// assumed the iGPU would carry ~a fifth of the slow-tier work. - #[test] - fn an_equal_speed_idle_device_relieves_an_overloaded_one() { - let devices = vec![ - TierDevice { - name: "cpu".into(), - class: DeviceClass::Cpu, - ladder: vec![Precision::Q8, Precision::Q4], - budget_bytes: 1_000, - speed: 72, - // Busy with the trunk: its balanced share of extra work is - // near zero, so held > quota from the first admission. - preload_units: 1_000, - }, - TierDevice { - name: "igpu".into(), - class: DeviceClass::IntegratedGpu, - // The expensive rung only — the kernel-safety shape of the - // real integrated GPU after the source cap. - ladder: vec![Precision::F32], - budget_bytes: 1_000, - speed: 71, - preload_units: 0, - }, - ]; - let costs = vec![cost(1, 2, 8); 8]; - let plan = plan_tiers(&devices, &costs, &[]).unwrap(); - let on_igpu = plan.tiers.iter().filter(|t| t.device == 1).count(); - assert!( - on_igpu > 0, - "the idle equal-speed device took nothing: {:?}", - plan.tiers - ); - assert!( - plan.tiers.iter().filter(|t| t.device == 1).all(|t| t.precision == Precision::F32), - "only its own ladder's rung may be used: {:?}", - plan.tiers - ); - } - - /// ...and when the fast device is genuinely full, hotness still decides - /// who gets in. - #[test] - fn the_hottest_expert_wins_the_last_slot_on_the_fast_device() { - let costs = vec![cost(1, 2, 8); 4]; - // GPU budget 2 = exactly one expert at its cheapest rung (q8). - let plan = plan_tiers(&devices(100, 2), &costs, &[0.1, 0.5, 0.3, 0.1]).unwrap(); - assert_eq!(plan.tiers[1], Tier { device: 1, precision: Precision::Q8 }); - assert!( - plan.tiers.iter().enumerate().all(|(e, t)| e == 1 || t.device == 0), - "only the hottest gets the fast device: {:?}", - plan.tiers - ); - } - - #[test] - fn admission_spills_to_int4_when_budgets_are_tight() { - // CPU holds 4 bytes: four q4 experts exactly; GPU off (0 budget). - let costs = vec![cost(1, 2, 8); 4]; - let plan = plan_tiers(&devices(4, 0), &costs, &[]).unwrap(); - assert!(plan.tiers.iter().all(|t| *t == Tier { device: 0, precision: Precision::Q4 })); - assert_eq!(plan.used_bytes, vec![4, 0]); - } - - #[test] - fn promotion_upgrades_precision_within_a_device_when_room_allows() { - // CPU holds 5: admission puts four q4 (4 bytes); promotion lifts the - // hottest to q8 (frees 1, costs 2 → 5). The next can't (6 > 5). - let costs = vec![cost(1, 2, 8); 4]; - let plan = plan_tiers(&devices(5, 0), &costs, &[0.0, 0.0, 1.0, 0.0]).unwrap(); - assert_eq!(plan.tiers[2].precision, Precision::Q8); - assert_eq!(plan.tiers.iter().filter(|t| t.precision == Precision::Q4).count(), 3); - assert_eq!(plan.used_bytes[0], 5); - } - - #[test] - fn no_fit_is_a_loud_error_naming_the_expert() { - let costs = vec![cost(1, 2, 8); 4]; - let err = plan_tiers(&devices(3, 0), &costs, &[]).unwrap_err(); - assert!(err.contains("expert 3"), "{err}"); - } - - #[test] - fn diff_lists_only_moved_experts() { - let costs = vec![cost(1, 2, 8); 3]; - // A GPU budget of 2 holds exactly ONE expert at its cheapest rung, so - // hotness decides which — without that contention every expert lands - // on the fast device in both plans and there is no diff to observe. - let a = plan_tiers(&devices(100, 2), &costs, &[1.0, 0.0, 0.0]).unwrap(); - let b = plan_tiers(&devices(100, 2), &costs, &[0.0, 1.0, 0.0]).unwrap(); - let moves = a.diff(&b); - assert_eq!(moves.len(), 2, "{moves:?}"); - assert!(moves.contains(&(1, Tier { device: 1, precision: Precision::Q8 })), "{moves:?}"); - assert!(moves.contains(&(0, Tier { device: 0, precision: Precision::Q8 })), "{moves:?}"); - } - - #[test] - fn smoothing_normalizes_per_request() { - let mut h = Vec::new(); - smooth_hotness(&mut h, &[3, 1], 1.0); - assert_eq!(h, vec![0.75, 0.25]); - smooth_hotness(&mut h, &[0, 4], 0.5); - assert_eq!(h, vec![0.375, 0.625]); - smooth_hotness(&mut h, &[0, 0], 0.5); // empty request changes nothing - assert_eq!(h, vec![0.375, 0.625]); - } -} diff --git a/crates/mummu/examples/src/tok_config.rs b/crates/mummu/examples/src/tok_config.rs deleted file mode 100644 index 7aa4381..0000000 --- a/crates/mummu/examples/src/tok_config.rs +++ /dev/null @@ -1,934 +0,0 @@ -//! `tokenizer_config.json` import — the **special-tokens map + chat template** -//! half of P3's "tokenizer + chat-template import". -//! -//! HF checkpoints carry the tokenizer *pipeline* in `tokenizer.json` (loaded -//! directly by the `tokenizers` crate) but keep the *conventions* a runner -//! needs — which tokens are BOS/EOS/PAD/UNK, whether to prepend BOS, the full -//! id→token map of added specials, and the Jinja **chat template** — in a -//! sibling `tokenizer_config.json`. This module parses that file into a -//! [`TokenizerConfig`] so an app can discover a model's declared special -//! tokens (and their ids) and its embedded chat template without hardcoding. -//! -//! It does **not** render the Jinja template — Mummu's prompt wrapping stays -//! the byte-verified code in [`crate::chat`]. The imported template is the -//! source of truth to check those renderers against and to detect a model's -//! tool-call convention; the imported special-token *ids* are a cross-check -//! that the config and the tokenizer agree (verified on real weights in -//! `tests/real_tokenizer_config.rs`). -//! -//! Union-typed fields are handled faithfully: a special-token field is `null`, -//! a bare string, or an `AddedToken` object `{ "content": … }`; `chat_template` -//! is a string or a list of `{ "name", "template" }` (the `default` entry, else -//! the first, is chosen). Every parse is total and bounded — malformed input is -//! a loud [`ImportError::Parse`], never a panic. - -use std::collections::HashMap; -use std::path::Path; - -use serde_json::Value; - -use crate::import::ImportError; - -/// The file this module reads. -pub const FILE_NAME: &str = "tokenizer_config.json"; - -/// A standalone chat-template file recent `transformers` `save_pretrained` -/// writes *instead of* the `chat_template` key of [`FILE_NAME`] (some -/// checkpoints — e.g. Gemma4 — ship it only here). [`TokenizerConfig::from_dir`] -/// falls back to it when the JSON key is absent, so the template-dependent -/// checks keep working for those checkpoints. -pub const CHAT_TEMPLATE_FILE: &str = "chat_template.jinja"; - -/// Hard cap on the config file size. Chat templates (esp. tool-calling ones) -/// are the large part — a few tens of KiB in practice; 4 MiB is generous -/// headroom while still refusing a pathological or wrong file outright. -const MAX_CONFIG_BYTES: u64 = 4 * 1024 * 1024; - -/// Hard cap on `added_tokens_decoder` entries — bounds the id space we ingest -/// (real vocabs top out in the low hundreds of thousands of *tokens*, of which -/// only the added specials appear here). -const MAX_ADDED_TOKENS: usize = 1_048_576; - -/// One entry of `added_tokens_decoder`: a token added on top of the base BPE -/// vocab, at a known id. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct AddedTokenInfo { - /// The token id (the JSON map key). - pub id: u32, - /// The token's literal text. - pub content: String, - /// Whether it is a special (control) token, never split out of text. - pub special: bool, -} - -/// A resolved special-token slot (BOS/EOS/PAD/UNK). `id` is `Some` when the -/// token's content was found in `added_tokens_decoder` — i.e. the config's -/// own tables agree on the id. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct SpecialToken { - /// The token's literal text. - pub content: String, - /// The id `added_tokens_decoder` assigns this content, if present. - pub id: Option, - /// Whether it is flagged special (authoritatively from the added-token - /// entry when found, else the field's own flag, else `true`). - pub special: bool, -} - -/// The tool-call convention a chat template speaks, detected from its marker -/// tokens. Lets an app pick the matching [`crate::chat`] render style -/// (`render_with_tools` Hermes vs LFM) from the checkpoint's own template -/// instead of hardcoding it per model. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum ToolCallConvention { - /// Hermes / Qwen (`{json}` with a `` block) — - /// [`crate::chat`]'s Hermes-style renderer + `parse_tool_calls`. - Hermes, - /// LFM2.5 (a Pythonic call list between `<|tool_call_start|>` / - /// `<|tool_call_end|>`) — [`crate::chat`]'s LFM-style renderer + - /// `parse_tool_calls_lfm`. - Lfm, -} - -/// A tokenizer id disagreement found by [`TokenizerConfig::check_ids_against`] -/// or [`TokenizerConfig::check_eos_agrees`]: the config recorded `found` for -/// `content`, but the authoritative source says `expected`. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct IdMismatch { - /// Which slot disagreed, for the message (e.g. `"eos_token_id"`). - pub what: &'static str, - /// The token text at issue. - pub content: String, - /// The authoritative id (from the tokenizer or `config.json`); `None` when - /// the token is absent there. - pub expected: Option, - /// The id `tokenizer_config.json` recorded. - pub found: Option, -} - -/// The conventions imported from a `tokenizer_config.json`. -#[derive(Debug, Clone, Default, PartialEq, Eq)] -pub struct TokenizerConfig { - /// Prepend BOS on encode (`add_bos_token`, default `false`). - pub add_bos_token: bool, - /// Append EOS on encode (`add_eos_token`, default `false`). - pub add_eos_token: bool, - pub bos_token: Option, - pub eos_token: Option, - pub pad_token: Option, - pub unk_token: Option, - /// `model_max_length`, when present and an integer. - pub model_max_length: Option, - /// The raw Jinja chat template (not rendered here). From the JSON - /// `chat_template` key, or — via [`Self::from_dir`] when that key is absent — - /// a standalone sibling [`CHAT_TEMPLATE_FILE`]. - pub chat_template: Option, - /// The full `added_tokens_decoder` map, sorted by ascending id. - pub added_tokens: Vec, -} - -/// Build an [`ImportError::Parse`] for `file` with `reason`. -fn parse_err(file: &Path, reason: impl Into) -> ImportError { - ImportError::Parse { - file: file.to_path_buf(), - reason: reason.into(), - } -} - -/// Read `dir/chat_template.jinja` if present — the standalone template file -/// recent `transformers` writes instead of the `chat_template` JSON key. Absent, -/// a non-file, or an empty/whitespace-only body → `Ok(None)` (as good as no -/// template); a present but oversized or non-UTF-8 file is a loud -/// [`ImportError::Parse`], never a panic. Bounded by the same [`MAX_CONFIG_BYTES`] -/// cap as the JSON file (a real template is tens of KiB). -fn read_chat_template_file(dir: &Path) -> Result, ImportError> { - assert!( - !dir.as_os_str().is_empty(), - "read_chat_template_file: empty dir" - ); - let path = dir.join(CHAT_TEMPLATE_FILE); - let Ok(meta) = std::fs::metadata(&path) else { - return Ok(None); // no standalone template beside the checkpoint - }; - if !meta.is_file() { - return Ok(None); - } - if meta.len() > MAX_CONFIG_BYTES { - return Err(parse_err( - &path, - format!( - "{CHAT_TEMPLATE_FILE} is {} bytes (> {MAX_CONFIG_BYTES} cap)", - meta.len() - ), - )); - } - let bytes = std::fs::read(&path).map_err(|e| parse_err(&path, format!("read: {e}")))?; - let text = String::from_utf8(bytes).map_err(|e| parse_err(&path, format!("not utf-8: {e}")))?; - if text.trim().is_empty() { - return Ok(None); - } - debug_assert!( - !text.trim().is_empty(), - "a returned template has non-whitespace content" - ); - Ok(Some(text)) -} - -impl TokenizerConfig { - /// Read + parse `dir/tokenizer_config.json`. - /// - /// [`ImportError::MissingFile`] if absent, [`ImportError::Parse`] if it is - /// oversized, unreadable, or malformed. - pub fn from_dir(dir: &Path) -> Result { - assert!(!dir.as_os_str().is_empty(), "from_dir: empty dir"); - let path = dir.join(FILE_NAME); - let meta = std::fs::metadata(&path).map_err(|_| ImportError::MissingFile(path.clone()))?; - if !meta.is_file() { - return Err(ImportError::MissingFile(path)); - } - if meta.len() > MAX_CONFIG_BYTES { - return Err(parse_err( - &path, - format!( - "{FILE_NAME} is {} bytes (> {MAX_CONFIG_BYTES} cap)", - meta.len() - ), - )); - } - let bytes = std::fs::read(&path).map_err(|e| parse_err(&path, format!("read: {e}")))?; - debug_assert!( - bytes.len() as u64 <= MAX_CONFIG_BYTES, - "size checked before read" - ); - let mut cfg = Self::from_json(&bytes, &path)?; - // Fall back to a standalone `chat_template.jinja` only when the JSON key - // was absent — a checkpoint that ships the template in the file (Gemma4) - // otherwise reads as having no template, silently disabling the - // template-dependent checks. A present JSON key wins (never overridden). - if cfg.chat_template.is_none() { - cfg.chat_template = read_chat_template_file(dir)?; - } - Ok(cfg) - } - - /// Parse `tokenizer_config.json` bytes (`file` only labels errors). - pub fn from_json(bytes: &[u8], file: &Path) -> Result { - assert!(!file.as_os_str().is_empty(), "from_json: empty file label"); - let root: Value = - serde_json::from_slice(bytes).map_err(|e| parse_err(file, format!("json: {e}")))?; - let obj = root - .as_object() - .ok_or_else(|| parse_err(file, "top-level is not a JSON object"))?; - - let added = parse_added_tokens(obj.get("added_tokens_decoder"), file)?; - assert!( - added.len() <= MAX_ADDED_TOKENS, - "added tokens bounded by parser" - ); - let index: HashMap<&str, (u32, bool)> = added - .iter() - .map(|a| (a.content.as_str(), (a.id, a.special))) - .collect(); - - let cfg = Self { - add_bos_token: bool_field(obj, "add_bos_token"), - add_eos_token: bool_field(obj, "add_eos_token"), - bos_token: resolve_special(obj.get("bos_token"), &index), - eos_token: resolve_special(obj.get("eos_token"), &index), - pad_token: resolve_special(obj.get("pad_token"), &index), - unk_token: resolve_special(obj.get("unk_token"), &index), - model_max_length: obj.get("model_max_length").and_then(Value::as_u64), - chat_template: obj.get("chat_template").and_then(parse_chat_template), - added_tokens: added, - }; - debug_assert!( - cfg.added_tokens.windows(2).all(|w| w[0].id < w[1].id), - "added tokens are sorted and unique after parse" - ); - Ok(cfg) - } - - /// The EOS token id, if the config resolved one. - #[must_use] - pub fn eos_id(&self) -> Option { - self.eos_token.as_ref().and_then(|t| t.id) - } - - /// The BOS token id, if the config resolved one. - #[must_use] - pub fn bos_id(&self) -> Option { - self.bos_token.as_ref().and_then(|t| t.id) - } - - /// The PAD token id, if the config resolved one. - #[must_use] - pub fn pad_id(&self) -> Option { - self.pad_token.as_ref().and_then(|t| t.id) - } - - /// Whether an embedded chat template was imported. - #[must_use] - pub fn has_chat_template(&self) -> bool { - self.chat_template.is_some() - } - - /// Detect the tool-call convention the imported `chat_template` implies, - /// by its marker tokens. `None` when there is no template or it names no - /// tool markers we recognize (a chat-only model). LFM's - /// `<|tool_call_start|>` / `<|tool_list_start|>` are checked first because - /// they are unambiguous; a Hermes template is identified by ``. - #[must_use] - pub fn tool_call_convention(&self) -> Option { - let template = self.chat_template.as_deref()?; - if template.contains("tool_call_start") || template.contains("tool_list_start") { - return Some(ToolCallConvention::Lfm); - } - if template.contains("") { - return Some(ToolCallConvention::Hermes); - } - None - } - - /// Cross-check every added-token id against an authoritative `token_to_id` - /// (the loaded HF [`tokenizers::Tokenizer`]): each `added_tokens_decoder` - /// entry must resolve, in the real tokenizer, to exactly the id the config - /// recorded. Catches a checkpoint whose `tokenizer_config.json` disagrees - /// with its `tokenizer.json` (a repackaging bug that would silently corrupt - /// special-token handling). The resolved BOS/EOS/PAD/UNK slots need no - /// separate check — each carries an id only because its content was found - /// in `added_tokens_decoder`, so it is already covered here. - /// - /// Takes a closure, not a `Tokenizer`, so this module stays free of a - /// tokenizer dependency and any id source can be validated. `Ok` when all - /// agree; `Err` lists every mismatch (bounded by the added-token count). - pub fn check_ids_against( - &self, - token_to_id: impl Fn(&str) -> Option, - ) -> Result<(), Vec> { - assert!( - self.added_tokens.len() <= MAX_ADDED_TOKENS, - "added tokens bounded by the parser" - ); - let mut bad = Vec::new(); - for a in &self.added_tokens { - let found = token_to_id(&a.content); - if found != Some(a.id) { - bad.push(IdMismatch { - what: "added token", - content: a.content.clone(), - expected: found, - found: Some(a.id), - }); - } - } - if bad.is_empty() { - return Ok(()); - } - debug_assert!( - !bad.is_empty(), - "the Err branch reports at least one mismatch" - ); - Err(bad) - } - - /// The fail-loud consistency gate a safetensors loader runs after parsing - /// `config.json`: the imported `tokenizer_config.json` must not contradict - /// (a) `config.json`'s EOS set, nor (b) this family's byte-verified - /// [`crate::chat`] renderer, whose tool-call convention is - /// `expected_convention` (`None` = the family has no tool renderer, so the - /// template's convention is not checked). Returns a human-readable reason on - /// the first disagreement; the loader wraps it as - /// [`ImportError::Inconsistent`]. - /// - /// The template check only fires when the imported template *declares* a - /// convention (`tool_call_convention()` is `Some`): a base/chat-only - /// checkpoint whose template names no tool markers is not forced to match. - /// This catches the real bug — a checkpoint packaged with a *different* - /// tool-call style than the loader's renderer emits (e.g. an LFM template - /// dropped into a Qwen dir) — without rejecting tool-less templates. - pub fn check_consistency( - &self, - config_eos_ids: &[u32], - expected_convention: Option, - ) -> Result<(), String> { - if let Err(m) = self.check_eos_agrees(config_eos_ids) { - return Err(format!( - "tokenizer_config EOS {:?} ({:?}) is not in config.json eos_token_id {config_eos_ids:?}", - m.found, m.content - )); - } - if let (Some(expected), Some(found)) = (expected_convention, self.tool_call_convention()) - && found != expected - { - return Err(format!( - "chat_template tool-call convention {found:?} contradicts this model's {expected:?} renderer", - )); - } - Ok(()) - } - - /// Cross-check the config's declared EOS against `config.json`'s - /// `eos_token_id` set: if `tokenizer_config.json` resolved an EOS id, it - /// must be one of the ids `config.json` names. Catches the packaging bug - /// where the two files disagree on which token ends a turn (the model then - /// never stops, or stops on the wrong id). `Ok` when they agree or when no - /// EOS id was resolved (nothing to check); `Err` names the disagreement. - pub fn check_eos_agrees(&self, config_eos_ids: &[u32]) -> Result<(), IdMismatch> { - assert!( - config_eos_ids.len() <= 256, - "config.json eos_token_id sets are small; got {}", - config_eos_ids.len() - ); - let Some(eos) = self.eos_token.as_ref() else { - return Ok(()); - }; - let Some(id) = eos.id else { - return Ok(()); - }; - assert!( - !eos.content.is_empty(), - "a resolved EOS has non-empty content" - ); - if config_eos_ids.contains(&id) { - Ok(()) - } else { - Err(IdMismatch { - what: "eos_token_id", - content: eos.content.clone(), - expected: config_eos_ids.first().copied(), - found: Some(id), - }) - } - } -} - -/// Run the [`TokenizerConfig::check_consistency`] gate against the -/// `tokenizer_config.json` beside a checkpoint's other files, if one is present. -/// -/// `tokenizer_config.json` is **optional** — a GGUF-derived dir or a minimal -/// checkpoint may not ship one — so a missing file is `Ok(None)` (nothing to -/// validate). A file that is *present but malformed* propagates its parse error, -/// and a present, well-formed file that *disagrees* with `config_eos_ids` or the -/// family renderer becomes [`ImportError::Inconsistent`]. On success the parsed -/// config is returned (`Some`) so a loader can reuse it (e.g. future BOS wiring). -pub fn validate_dir( - dir: &Path, - config_eos_ids: &[u32], - expected_convention: Option, -) -> Result, ImportError> { - assert!(!dir.as_os_str().is_empty(), "validate_dir: empty dir"); - let cfg = match TokenizerConfig::from_dir(dir) { - Ok(cfg) => cfg, - // The file is optional: absence is not an error, it is "nothing to check". - Err(ImportError::MissingFile(_)) => return Ok(None), - Err(e) => return Err(e), - }; - cfg.check_consistency(config_eos_ids, expected_convention) - .map_err(|reason| ImportError::Inconsistent { - file: dir.join(FILE_NAME), - reason, - })?; - Ok(Some(cfg)) -} - -/// A boolean field, defaulting to `false` when absent or non-boolean. -fn bool_field(obj: &serde_json::Map, key: &str) -> bool { - obj.get(key).and_then(Value::as_bool).unwrap_or(false) -} - -/// Parse `added_tokens_decoder` (a `"id" -> { content, special, … }` map) into -/// an id-sorted [`AddedTokenInfo`] list. Absent → empty. Malformed keys, -/// missing content, an oversize map, or a duplicate id are loud errors. -fn parse_added_tokens(v: Option<&Value>, file: &Path) -> Result, ImportError> { - let Some(v) = v else { - return Ok(Vec::new()); - }; - let map = v - .as_object() - .ok_or_else(|| parse_err(file, "added_tokens_decoder is not an object"))?; - if map.len() > MAX_ADDED_TOKENS { - return Err(parse_err( - file, - format!( - "added_tokens_decoder has {} entries (> {MAX_ADDED_TOKENS} cap)", - map.len() - ), - )); - } - let mut out = Vec::with_capacity(map.len()); - for (key, entry) in map { - let id = key.parse::().map_err(|_| { - parse_err( - file, - format!("added_tokens_decoder key {key:?} is not a token id"), - ) - })?; - let content = entry - .as_object() - .and_then(|o| o.get("content")) - .and_then(Value::as_str) - .ok_or_else(|| parse_err(file, format!("added token {id} has no string content")))?; - let special = entry - .as_object() - .and_then(|o| o.get("special")) - .and_then(Value::as_bool) - .unwrap_or(true); - out.push(AddedTokenInfo { - id, - content: content.to_string(), - special, - }); - } - out.sort_by_key(|a| a.id); - // Distinct JSON keys can still collide numerically ("1" vs "01"); reject - // rather than silently keep an ambiguous id. - for w in out.windows(2) { - if w[0].id == w[1].id { - return Err(parse_err( - file, - format!("duplicate added-token id {}", w[0].id), - )); - } - } - Ok(out) -} - -/// A special-token field's `(content, own_special_flag)`. `null`/absent/other -/// shapes → `None`; a bare string assumes special; an object reads `content` -/// (+ optional `special`). -fn special_content(v: &Value) -> Option<(String, bool)> { - match v { - Value::String(s) => Some((s.clone(), true)), - Value::Object(o) => { - let content = o.get("content").and_then(Value::as_str)?; - let special = o.get("special").and_then(Value::as_bool).unwrap_or(true); - Some((content.to_string(), special)) - } - _ => None, - } -} - -/// Resolve a special-token field to a [`SpecialToken`], looking its id up in -/// the added-token index (which is authoritative for the id and the special -/// flag when the content is found there). -fn resolve_special(v: Option<&Value>, index: &HashMap<&str, (u32, bool)>) -> Option { - let (content, own_special) = special_content(v?)?; - let (id, special) = match index.get(content.as_str()) { - Some(&(id, sp)) => (Some(id), sp), - None => (None, own_special), - }; - Some(SpecialToken { - content, - id, - special, - }) -} - -/// Extract the chat template: a bare string, or the `default` (else first) -/// entry of a `[{ "name", "template" }]` list. Other shapes → `None`. -fn parse_chat_template(v: &Value) -> Option { - match v { - Value::String(s) => Some(s.clone()), - Value::Array(arr) => { - let mut first: Option = None; - for item in arr { - let Some(tmpl) = item.get("template").and_then(Value::as_str) else { - continue; - }; - if item.get("name").and_then(Value::as_str) == Some("default") { - return Some(tmpl.to_string()); - } - if first.is_none() { - first = Some(tmpl.to_string()); - } - } - first - } - _ => None, - } -} - -#[cfg(test)] -mod tests { - use super::*; - - fn parse(json: &str) -> TokenizerConfig { - TokenizerConfig::from_json(json.as_bytes(), Path::new("tokenizer_config.json")) - .expect("valid config parses") - } - - #[test] - fn string_form_specials_resolve_ids_from_added_tokens() { - let cfg = parse( - r#"{ - "add_bos_token": false, - "eos_token": "<|im_end|>", - "pad_token": "<|endoftext|>", - "unk_token": null, - "added_tokens_decoder": { - "151643": {"content": "<|endoftext|>", "special": true}, - "151645": {"content": "<|im_end|>", "special": true} - } - }"#, - ); - assert_eq!(cfg.eos_id(), Some(151645)); - assert_eq!(cfg.pad_id(), Some(151643)); - assert!(cfg.eos_token.as_ref().unwrap().special); - assert_eq!(cfg.unk_token, None, "null token is absent"); - assert!(!cfg.add_bos_token); - assert_eq!(cfg.added_tokens.len(), 2); - // added_tokens are sorted by id. - assert_eq!(cfg.added_tokens[0].id, 151643); - assert_eq!(cfg.added_tokens[1].id, 151645); - } - - #[test] - fn object_form_special_reads_content_and_flag() { - let cfg = parse( - r#"{ - "bos_token": {"content": "", "special": true, "lstrip": false}, - "added_tokens_decoder": {"1": {"content": "", "special": true}} - }"#, - ); - let bos = cfg.bos_token.expect("bos present"); - assert_eq!(bos.content, ""); - assert_eq!(bos.id, Some(1)); - assert!(bos.special); - } - - #[test] - fn unresolved_special_has_no_id_but_keeps_content() { - // eos content is not in added_tokens_decoder → id None, content kept. - let cfg = parse(r#"{ "eos_token": "" }"#); - let eos = cfg.eos_token.expect("eos present"); - assert_eq!(eos.content, ""); - assert_eq!(eos.id, None); - assert!(eos.special, "bare-string tokens default special"); - } - - #[test] - fn chat_template_string_and_list_forms() { - let s = parse(r#"{ "chat_template": "{{ hi }}" }"#); - assert_eq!(s.chat_template.as_deref(), Some("{{ hi }}")); - assert!(s.has_chat_template()); - - let list = parse( - r#"{ "chat_template": [ - {"name": "tool_use", "template": "TOOLS"}, - {"name": "default", "template": "DEFAULT"} - ] }"#, - ); - assert_eq!( - list.chat_template.as_deref(), - Some("DEFAULT"), - "default entry wins" - ); - - let first = parse( - r#"{ "chat_template": [ - {"name": "a", "template": "FIRST"}, - {"name": "b", "template": "SECOND"} - ] }"#, - ); - assert_eq!( - first.chat_template.as_deref(), - Some("FIRST"), - "no default → first" - ); - } - - #[test] - fn tool_call_convention_is_detected_from_template_markers() { - // Hermes: + block (Qwen-style). - let hermes = parse( - r#"{ "chat_template": "…{{sig}}…\n{json}\n…" }"#, - ); - assert_eq!( - hermes.tool_call_convention(), - Some(ToolCallConvention::Hermes) - ); - - // LFM: the distinctive <|tool_call_start|> markers win even if the - // template also mentions a generic tool word. - let lfm = parse( - r#"{ "chat_template": "…List of tools: […]…<|tool_call_start|>[f(x=1)]<|tool_call_end|>…" }"#, - ); - assert_eq!(lfm.tool_call_convention(), Some(ToolCallConvention::Lfm)); - - // A chat-only template names no tool markers. - let chat_only = parse(r#"{ "chat_template": "<|im_start|>{{content}}<|im_end|>" }"#); - assert_eq!(chat_only.tool_call_convention(), None); - - // No template at all, and an empty template, are both None (no panic). - assert_eq!(parse("{}").tool_call_convention(), None); - assert_eq!( - parse(r#"{ "chat_template": "" }"#).tool_call_convention(), - None - ); - } - - #[test] - fn check_ids_against_passes_when_agreeing_and_lists_mismatches() { - let cfg = parse( - r#"{ - "eos_token": "<|im_end|>", - "added_tokens_decoder": { - "100": {"content": "<|im_end|>", "special": true}, - "200": {"content": "<|extra|>", "special": true} - } - }"#, - ); - // Agreeing tokenizer: both ids match. - let agree = |t: &str| match t { - "<|im_end|>" => Some(100), - "<|extra|>" => Some(200), - _ => None, - }; - assert!(cfg.check_ids_against(agree).is_ok()); - - // A tokenizer that disagrees on one id and is missing the other. - let disagree = |t: &str| match t { - "<|im_end|>" => Some(999), - _ => None, - }; - let bad = cfg.check_ids_against(disagree).unwrap_err(); - assert_eq!(bad.len(), 2, "both added tokens mismatch"); - let end = bad.iter().find(|m| m.content == "<|im_end|>").unwrap(); - assert_eq!(end.found, Some(100), "config recorded 100"); - assert_eq!(end.expected, Some(999), "tokenizer says 999"); - let extra = bad.iter().find(|m| m.content == "<|extra|>").unwrap(); - assert_eq!(extra.expected, None, "tokenizer is missing it"); - } - - #[test] - fn check_eos_agrees_matches_config_json_eos_set() { - let cfg = parse( - r#"{ - "eos_token": "<|im_end|>", - "added_tokens_decoder": {"151645": {"content": "<|im_end|>", "special": true}} - }"#, - ); - // config.json eos_token_id lists the same id → agree. - assert!(cfg.check_eos_agrees(&[151_645]).is_ok()); - // A list that includes it (Qwen lists multiple) → agree. - assert!(cfg.check_eos_agrees(&[151_643, 151_645]).is_ok()); - // A disagreeing config.json → loud mismatch naming both sides. - let m = cfg.check_eos_agrees(&[151_643]).unwrap_err(); - assert_eq!(m.what, "eos_token_id"); - assert_eq!(m.found, Some(151_645), "tokenizer_config's eos id"); - assert_eq!(m.expected, Some(151_643), "config.json's first eos id"); - - // No resolved EOS id → nothing to check, always Ok. - let no_eos = parse(r#"{ "eos_token": "" }"#); // unresolved (no added_tokens) - assert!(no_eos.check_eos_agrees(&[7]).is_ok()); - assert!(parse("{}").check_eos_agrees(&[]).is_ok()); - } - - #[test] - fn check_consistency_passes_agreeing_and_flags_each_disagreement() { - // A well-formed Hermes checkpoint: EOS agrees, template is Hermes. - let cfg = parse( - r#"{ - "eos_token": "<|im_end|>", - "chat_template": "…{{s}}…\n{j}\n…", - "added_tokens_decoder": {"151645": {"content": "<|im_end|>", "special": true}} - }"#, - ); - assert!( - cfg.check_consistency(&[151_645], Some(ToolCallConvention::Hermes)) - .is_ok() - ); - // EOS not in config.json's set → loud reason mentioning the id. - let eos_err = cfg - .check_consistency(&[151_643], Some(ToolCallConvention::Hermes)) - .unwrap_err(); - assert!(eos_err.contains("151645"), "reason names the stray id"); - // Right EOS, but the loader expected the LFM convention → contradiction. - let conv_err = cfg - .check_consistency(&[151_645], Some(ToolCallConvention::Lfm)) - .unwrap_err(); - assert!(conv_err.contains("Hermes") && conv_err.contains("Lfm")); - } - - #[test] - fn check_consistency_skips_template_check_when_none_expected_or_declared() { - // A tool-less (base) template declares no convention: any expectation is - // fine, only the EOS is checked. - let base = parse( - r#"{ - "eos_token": "<|end|>", - "chat_template": "<|im_start|>{{content}}<|im_end|>", - "added_tokens_decoder": {"9": {"content": "<|end|>", "special": true}} - }"#, - ); - assert!( - base.check_consistency(&[9], Some(ToolCallConvention::Hermes)) - .is_ok(), - "no declared convention → template not forced to match" - ); - // expected None → the template convention is never checked even if set. - let lfm = parse(r#"{ "chat_template": "…<|tool_call_start|>[f()]<|tool_call_end|>…" }"#); - assert!(lfm.check_consistency(&[], None).is_ok()); - } - - #[test] - fn validate_dir_is_ok_when_file_absent_and_loud_on_mismatch() { - use crate::import::ImportError; - let dir = std::env::temp_dir().join("mummu_tokcfg_validate_test"); - std::fs::create_dir_all(&dir).unwrap(); - let _ = std::fs::remove_file(dir.join(FILE_NAME)); - // Absent file → Ok(None): the file is optional. - assert!(matches!( - validate_dir(&dir, &[1, 2], Some(ToolCallConvention::Hermes)), - Ok(None) - )); - // Present + agreeing → Ok(Some(cfg)). - std::fs::write( - dir.join(FILE_NAME), - br#"{"eos_token":"<|im_end|>","added_tokens_decoder":{"5":{"content":"<|im_end|>","special":true}}}"#, - ) - .unwrap(); - assert!(matches!( - validate_dir(&dir, &[5], Some(ToolCallConvention::Hermes)), - Ok(Some(_)) - )); - // Present + disagreeing EOS → loud Inconsistent (not a silent pass). - assert!(matches!( - validate_dir(&dir, &[7], Some(ToolCallConvention::Hermes)), - Err(ImportError::Inconsistent { .. }) - )); - std::fs::remove_dir_all(&dir).unwrap(); - } - - #[test] - fn empty_config_defaults_everything() { - let cfg = parse("{}"); - assert!(!cfg.add_bos_token && !cfg.add_eos_token); - assert_eq!(cfg.eos_token, None); - assert_eq!(cfg.model_max_length, None); - assert!(!cfg.has_chat_template()); - assert!(cfg.added_tokens.is_empty()); - } - - #[test] - fn model_max_length_and_add_eos_are_read() { - let cfg = parse(r#"{ "add_eos_token": true, "model_max_length": 131072 }"#); - assert!(cfg.add_eos_token); - assert_eq!(cfg.model_max_length, Some(131072)); - } - - #[test] - fn malformed_inputs_fail_loudly() { - let file = Path::new("tokenizer_config.json"); - // Non-object top level. - assert!(TokenizerConfig::from_json(b"[1,2,3]", file).is_err()); - // Invalid JSON. - assert!(TokenizerConfig::from_json(b"{ not json", file).is_err()); - // Non-numeric added-token key. - assert!( - TokenizerConfig::from_json( - br#"{"added_tokens_decoder": {"oops": {"content": "x"}}}"#, - file - ) - .is_err() - ); - // Added token without content. - assert!( - TokenizerConfig::from_json( - br#"{"added_tokens_decoder": {"1": {"special": true}}}"#, - file - ) - .is_err() - ); - // Duplicate numeric id ("1" vs "01"). - assert!( - TokenizerConfig::from_json( - br#"{"added_tokens_decoder": {"1": {"content": "a"}, "01": {"content": "b"}}}"#, - file - ) - .is_err() - ); - } - - #[test] - fn from_dir_reads_and_reports_missing() { - let dir = std::env::temp_dir().join("mummu_tokcfg_test"); - std::fs::create_dir_all(&dir).unwrap(); - let _ = std::fs::remove_file(dir.join(FILE_NAME)); - assert!(matches!( - TokenizerConfig::from_dir(&dir), - Err(ImportError::MissingFile(_)) - )); - std::fs::write(dir.join(FILE_NAME), br#"{"eos_token": "<|end|>"}"#).unwrap(); - let cfg = TokenizerConfig::from_dir(&dir).expect("present config parses"); - assert_eq!(cfg.eos_token.unwrap().content, "<|end|>"); - std::fs::remove_dir_all(&dir).unwrap(); - } - - #[test] - fn from_dir_falls_back_to_standalone_chat_template_jinja() { - let dir = std::env::temp_dir().join("mummu_tokcfg_jinja_fallback"); - let _ = std::fs::remove_dir_all(&dir); - std::fs::create_dir_all(&dir).unwrap(); - // tokenizer_config.json has NO chat_template key; the template lives in a - // sibling chat_template.jinja (the Gemma4-style layout). - std::fs::write(dir.join(FILE_NAME), br#"{"eos_token": "<|im_end|>"}"#).unwrap(); - std::fs::write( - dir.join(CHAT_TEMPLATE_FILE), - b"<|im_start|>system\n{{x}}<|im_end|>\n{{s}}\n\n{j}\n", - ) - .unwrap(); - let cfg = TokenizerConfig::from_dir(&dir).expect("config + jinja fallback parses"); - assert!(cfg.has_chat_template(), "the .jinja template was picked up"); - assert!( - cfg.chat_template - .as_deref() - .unwrap() - .contains("<|im_start|>") - ); - // The convention gate now works for this checkpoint (it would have been - // None — silently unchecked — without the fallback). - assert_eq!( - cfg.tool_call_convention(), - Some(ToolCallConvention::Hermes), - "the standalone template's tool-call convention is detected" - ); - std::fs::remove_dir_all(&dir).unwrap(); - } - - #[test] - fn from_dir_json_chat_template_wins_over_the_jinja_file() { - let dir = std::env::temp_dir().join("mummu_tokcfg_jinja_precedence"); - let _ = std::fs::remove_dir_all(&dir); - std::fs::create_dir_all(&dir).unwrap(); - // Both present: the JSON key is authoritative and must not be overridden. - std::fs::write( - dir.join(FILE_NAME), - br#"{"chat_template": "FROM_JSON_KEY"}"#, - ) - .unwrap(); - std::fs::write(dir.join(CHAT_TEMPLATE_FILE), b"FROM_JINJA_FILE").unwrap(); - let cfg = TokenizerConfig::from_dir(&dir).expect("parses"); - assert_eq!( - cfg.chat_template.as_deref(), - Some("FROM_JSON_KEY"), - "a present JSON chat_template wins over the standalone file" - ); - std::fs::remove_dir_all(&dir).unwrap(); - } - - #[test] - fn from_dir_treats_an_empty_jinja_file_as_absent() { - let dir = std::env::temp_dir().join("mummu_tokcfg_jinja_empty"); - let _ = std::fs::remove_dir_all(&dir); - std::fs::create_dir_all(&dir).unwrap(); - std::fs::write(dir.join(FILE_NAME), b"{}").unwrap(); - std::fs::write(dir.join(CHAT_TEMPLATE_FILE), b" \n\t ").unwrap(); - let cfg = TokenizerConfig::from_dir(&dir).expect("parses"); - assert!( - !cfg.has_chat_template(), - "a whitespace-only .jinja is as good as no template" - ); - std::fs::remove_dir_all(&dir).unwrap(); - } -} diff --git a/crates/mummu/examples/src/tokenizer.rs b/crates/mummu/examples/src/tokenizer.rs deleted file mode 100644 index 20cd00b..0000000 --- a/crates/mummu/examples/src/tokenizer.rs +++ /dev/null @@ -1,1033 +0,0 @@ -//! Tokenizer construction from non-`tokenizer.json` sources. -//! -//! **GGUF metadata** ([`tokenizer_from_gguf`]) — llama.cpp stores the -//! tokenizer as `tokenizer.ggml.*` metadata: the vocab (`tokens`, index = -//! token id), per-token types, BPE `merges`, and a `pre` identifier naming -//! the pre-tokenizer regex (the ecosystem hardcodes the regex per model -//! family, exactly as llama.cpp's `llama_vocab` does). Rebuilt as: -//! NFC → Split(pre regex) → ByteLevel → BPE, with control/user-defined -//! tokens re-added as special/non-special added tokens. Verified -//! byte-identical vs the checkpoint's `tokenizer.json` in `tests/real_gguf.rs`. -//! -//! **SentencePiece `tokenizer.model`** ([`tokenizer_from_spm`]) — the SPM -//! proto the Llama/Gemma/T5 families ship. A bounded hand-rolled protobuf -//! reader (the `gguf.rs` approach — the schema is tiny and frozen) feeds the -//! same pipeline HF's `convert_slow_tokenizer` assembles: Precompiled -//! charsmap (+ multi-space collapse) → Metaspace → **Unigram**. Verified -//! byte-identical vs the checkpoint's `tokenizer.json` in `tests/real_spm.rs`. - -use std::path::Path; - -use tokenizers::models::bpe::{BPE, Merges, Vocab}; -use tokenizers::normalizers::unicode::NFC; -use tokenizers::pre_tokenizers::byte_level::ByteLevel; -use tokenizers::pre_tokenizers::sequence::Sequence; -use tokenizers::pre_tokenizers::split::{Split, SplitPattern}; -use tokenizers::processors::template::TemplateProcessing; -use tokenizers::{AddedToken, SplitDelimiterBehavior, Tokenizer}; - -use crate::gguf::{GgufFile, GgufValue}; -use crate::import::ImportError; -use crate::tok_config::{self, TokenizerConfig, ToolCallConvention}; - -/// llama.cpp token types (`llama_token_type`). -const TOKEN_TYPE_NORMAL: i64 = 1; -const TOKEN_TYPE_CONTROL: i64 = 3; -const TOKEN_TYPE_USER_DEFINED: i64 = 4; -const TOKEN_TYPE_UNUSED: i64 = 5; - -/// How a model family's byte-level BPE pipeline is configured — the part a -/// GGUF names by `tokenizer.ggml.pre` id instead of carrying explicitly. -struct PreSpec { - /// The pre-tokenizer split regex (from the family's `tokenizer.json`). - regex: &'static str, - /// Whether the pipeline NFC-normalizes first. - nfc: bool, -} - -/// The per-family registry, keyed by `tokenizer.ggml.pre` — the same registry -/// llama.cpp keeps in its vocab loader. Only families we actually run are -/// listed; unknown ids are a loud error (a wrong regex silently produces -/// wrong token ids). -fn pre_spec(pre: &str) -> Option { - match pre { - // Qwen2/2.5 (matches the checkpoint's tokenizer.json byte for byte). - "qwen2" => Some(PreSpec { - regex: r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+", - nfc: true, - }), - // Qwen3.5/3.8: qwen2's pattern with combining marks (\p{M}) folded - // into the letter class (llama.cpp PRE_TYPE_QWEN35, from the - // family's tokenizer.json). - "qwen35" => Some(PreSpec { - regex: r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?[\p{L}\p{M}]+|\p{N}| ?[^\s\p{L}\p{M}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+", - nfc: true, - }), - // LFM2/LFM2.5: digits split in groups of ≤3, no normalizer. - "lfm2" => Some(PreSpec { - regex: r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+", - nfc: false, - }), - // OLMo/OLMoE (GPT-NeoX lineage): the stock GPT-2 regex, contractions - // case-SENSITIVE — the family's tokenizer.json is ByteLevel with - // `use_regex: true` (exactly this pattern) behind an NFC normalizer. - "olmo" => Some(PreSpec { - regex: r"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+", - nfc: true, - }), - _ => None, - } -} - -/// A required `tokenizer.ggml.*` metadata array. -fn required_array<'f>(f: &'f GgufFile, key: &str) -> Result<&'f [GgufValue], String> { - f.get(key) - .and_then(GgufValue::as_array) - .ok_or_else(|| format!("missing or non-array GGUF metadata '{key}'")) -} - -/// Build an HF [`Tokenizer`] from a GGUF file's `tokenizer.ggml.*` metadata. -/// -/// Supports the `gpt2` (byte-level BPE) tokenizer model with a known `pre` -/// regex. Token ids are the `tokens` array indexes; control (type 3) tokens -/// become special added tokens, user-defined (type 4) become non-special -/// added tokens, unused (type 5) padding entries are skipped. Every added -/// token's id is verified after construction — a drifted id would silently -/// corrupt every prompt, so it fails loudly instead. -pub fn tokenizer_from_gguf(f: &GgufFile) -> Result { - let model = f - .get("tokenizer.ggml.model") - .and_then(GgufValue::as_str) - .ok_or("missing GGUF metadata 'tokenizer.ggml.model'")?; - if model != "gpt2" { - return Err(format!( - "tokenizer model '{model}' is not supported yet (only gpt2 byte-level BPE)" - )); - } - let pre = f - .get("tokenizer.ggml.pre") - .and_then(GgufValue::as_str) - .ok_or("missing GGUF metadata 'tokenizer.ggml.pre'")?; - let Some(spec) = pre_spec(pre) else { - return Err(format!("unknown pre-tokenizer id '{pre}'")); - }; - - let tokens = required_array(f, "tokenizer.ggml.tokens")?; - let types = required_array(f, "tokenizer.ggml.token_type")?; - if tokens.len() != types.len() { - return Err(format!( - "tokens ({}) and token_type ({}) lengths differ", - tokens.len(), - types.len() - )); - } - - // The BPE vocab: every non-padding token at id = array index. Control/ - // user-defined tokens go in TOO — `add_tokens` below then reuses the - // model id it finds, which is what lets added tokens live at LOW ids - // (LFM2 puts its 500+ specials at 0..) instead of only after the vocab. - let mut vocab = Vocab::default(); - // (index, content, is_special) for everything re-added post-BPE. - let mut added: Vec<(usize, String, bool)> = Vec::new(); - for (index, (token, ty)) in tokens.iter().zip(types).enumerate() { - let text = token - .as_str() - .ok_or_else(|| format!("token {index} is not a string"))?; - let ty = ty - .as_i64() - .ok_or_else(|| format!("token_type {index} is not an integer"))?; - #[allow(clippy::cast_possible_truncation)] // bounded by MAX_ARRAY_LEN - match ty { - TOKEN_TYPE_NORMAL => { - vocab.insert(text.to_string(), index as u32); - } - TOKEN_TYPE_CONTROL | TOKEN_TYPE_USER_DEFINED => { - vocab.insert(text.to_string(), index as u32); - added.push((index, text.to_string(), ty == TOKEN_TYPE_CONTROL)); - } - TOKEN_TYPE_UNUSED => {} // vocab-padding entries ([PADn]) - other => return Err(format!("token {index} has unsupported type {other}")), - } - } - assert!(!vocab.is_empty(), "a tokenizer must have normal tokens"); - - let mut merges = Merges::with_capacity(required_array(f, "tokenizer.ggml.merges")?.len()); - for (i, m) in required_array(f, "tokenizer.ggml.merges")? - .iter() - .enumerate() - { - let m = m - .as_str() - .ok_or_else(|| format!("merge {i} is not a string"))?; - let (a, b) = m - .split_once(' ') - .ok_or_else(|| format!("merge {i} ('{m}') is not 'left right'"))?; - merges.push((a.to_string(), b.to_string())); - } - - let bpe = BPE::builder() - .vocab_and_merges(vocab, merges) - .build() - .map_err(|e| format!("BPE build: {e}"))?; - let mut tok = Tokenizer::new(bpe); - if spec.nfc { - tok.with_normalizer(Some(NFC)) - .map_err(|e| format!("NFC normalizer: {e}"))?; - } - let split = Split::new( - SplitPattern::Regex(spec.regex.to_string()), - SplitDelimiterBehavior::Isolated, - false, - ) - .map_err(|e| format!("pre-tokenizer regex: {e}"))?; - // ByteLevel exactly as the HF checkpoints configure it: no prefix space, - // no offset trimming, regex handled by the Split stage above. - let byte_level = ByteLevel::new(false, false, false); - tok.with_pre_tokenizer(Some(Sequence::new(vec![split.into(), byte_level.into()]))); - tok.with_decoder(Some(byte_level)); - - // `tokenizer.ggml.add_bos_token` → a BOS-prepending template processor - // (what the family's tokenizer.json does); otherwise the offsets-only - // ByteLevel processor. - let add_bos = f - .get("tokenizer.ggml.add_bos_token") - .and_then(|v| match *v { - GgufValue::Bool(b) => Some(b), - _ => None, - }) - .unwrap_or(false); - if add_bos { - let bos_id = f - .get("tokenizer.ggml.bos_token_id") - .and_then(GgufValue::as_u64) - .and_then(|v| u32::try_from(v).ok()) - .ok_or("add_bos_token is set but tokenizer.ggml.bos_token_id is missing")?; - let bos = tokens - .get(bos_id as usize) - .and_then(GgufValue::as_str) - .ok_or_else(|| format!("bos_token_id {bos_id} is out of vocab range"))?; - let template = TemplateProcessing::builder() - .try_single(format!("{bos} $A")) - .map_err(|e| format!("BOS template: {e}"))? - .special_tokens(vec![(bos.to_string(), bos_id)]) - .build() - .map_err(|e| format!("BOS template: {e}"))?; - tok.with_post_processor(Some(template)); - } else { - tok.with_post_processor(Some(byte_level)); - } - - // Re-add control/user-defined tokens in id order, then verify every id - // landed where the GGUF says it lives. - for (_, text, special) in &added { - let t = AddedToken::from(text.clone(), *special); - if *special { - tok.add_special_tokens([t]) - .map_err(|e| format!("add special token '{text}': {e}"))?; - } else { - tok.add_tokens([t]) - .map_err(|e| format!("add token '{text}': {e}"))?; - } - } - for (index, text, _) in &added { - let got = tok.token_to_id(text); - #[allow(clippy::cast_possible_truncation)] // bounded by MAX_ARRAY_LEN - if got != Some(*index as u32) { - return Err(format!( - "added token '{text}' resolved to id {got:?}, GGUF says {index} — \ - non-contiguous added-token ids are not supported" - )); - } - } - Ok(tok) -} - -/// The HF fast-tokenizer file a checkpoint ships beside its weights. -pub const TOKENIZER_JSON: &str = "tokenizer.json"; - -/// Cap on the number of individual mismatches spelled out in an -/// [`ImportError::Inconsistent`] message — the rest are summarized as a count so -/// a pathologically-broken checkpoint can't produce an unbounded error string. -const MAX_LISTED_MISMATCHES: usize = 8; - -/// The full checkpoint-metadata gate a safetensors loader runs after parsing -/// `config.json` and **before** reading any weight bytes. It layers the -/// tokenizer-opening id cross-check on top of the tokenizer-free -/// [`tok_config::validate_dir`] gate (EOS agreement + tool-call convention), -/// giving the loaders one call that fails loudly on *any* metadata -/// disagreement at load time rather than at generate time. -/// -/// Both sibling files are optional, and each absence is "nothing to check", not -/// an error: -/// * no `tokenizer_config.json` (a GGUF-derived or minimal dir) → `Ok(None)`; -/// * a `tokenizer_config.json` present but no `tokenizer.json` beside it → the -/// EOS/convention checks still run, the id cross-check is skipped. -/// -/// A present, well-formed `tokenizer_config.json` whose declared added-token ids -/// disagree with the real `tokenizer.json` is an [`ImportError::Inconsistent`] -/// (a repackaging bug a checked *weight* load cannot see). On success the parsed -/// config is returned so a loader can reuse it (e.g. future config-driven BOS). -pub fn validate_checkpoint_dir( - dir: &Path, - config_eos_ids: &[u32], - expected_convention: Option, -) -> Result, ImportError> { - assert!( - !dir.as_os_str().is_empty(), - "validate_checkpoint_dir: empty dir" - ); - assert!( - config_eos_ids.len() <= 256, - "config.json eos_token_id sets are small; got {}", - config_eos_ids.len() - ); - let cfg = tok_config::validate_dir(dir, config_eos_ids, expected_convention)?; - if let Some(cfg) = &cfg { - check_added_token_ids(dir, cfg)?; - } - Ok(cfg) -} - -/// Cross-check every added-token id `cfg` declares against the id the sibling -/// `dir/tokenizer.json` assigns that content. `Ok` when the tokenizer file is -/// absent (nothing to cross-check) or every id agrees; an -/// [`ImportError::Inconsistent`] naming the disagreements otherwise. A -/// `tokenizer.json` that is present but unreadable/malformed is a loud -/// [`ImportError::Parse`], never a panic. -fn check_added_token_ids(dir: &Path, cfg: &TokenizerConfig) -> Result<(), ImportError> { - assert!( - !dir.as_os_str().is_empty(), - "check_added_token_ids: empty dir" - ); - let path = dir.join(TOKENIZER_JSON); - if !path.is_file() { - return Ok(()); // no fast tokenizer beside the checkpoint — nothing to check - } - let tok = Tokenizer::from_file(&path).map_err(|e| ImportError::Parse { - file: path.clone(), - reason: format!("load {TOKENIZER_JSON}: {e}"), - })?; - let mismatches = match cfg.check_ids_against(|t| tok.token_to_id(t)) { - Ok(()) => return Ok(()), - Err(m) => m, - }; - assert!( - !mismatches.is_empty(), - "the Err branch lists at least one mismatch" - ); - debug_assert!( - mismatches.len() <= cfg.added_tokens.len(), - "at most one mismatch per declared added token" - ); - let shown = mismatches.len().min(MAX_LISTED_MISMATCHES); - let mut reason = format!( - "{} added-token id(s) in {} disagree with {}:", - mismatches.len(), - tok_config::FILE_NAME, - TOKENIZER_JSON, - ); - for m in mismatches.iter().take(shown) { - reason.push_str(&format!( - " {:?} (config={:?}, tokenizer={:?});", - m.content, m.found, m.expected - )); - } - if mismatches.len() > shown { - reason.push_str(&format!(" (+{} more)", mismatches.len() - shown)); - } - Err(ImportError::Inconsistent { file: path, reason }) -} - -// --------------------------------------------------------------------------- -// SentencePiece `tokenizer.model` import (Llama/Gemma/T5-family checkpoints -// that ship the SPM proto instead of — or beside — a `tokenizer.json`). -// --------------------------------------------------------------------------- - -/// Largest `tokenizer.model` file the reader will load — an order of magnitude -/// past the largest real proto (Gemma's 256k-piece model is ~4 MiB). -const MAX_SPM_FILE_BYTES: u64 = 64 * 1024 * 1024; -/// Most pieces a proto may declare (Gemma: 256k; bound leaves headroom). -const MAX_SPM_PIECES: usize = 1_048_576; - -/// SentencePiece `ModelProto.SentencePiece.Type` values. -const SPM_TYPE_NORMAL: i64 = 1; -const SPM_TYPE_UNKNOWN: i64 = 2; -const SPM_TYPE_CONTROL: i64 = 3; -const SPM_TYPE_USER_DEFINED: i64 = 4; -const SPM_TYPE_UNUSED: i64 = 5; -const SPM_TYPE_BYTE: i64 = 6; - -/// `TrainerSpec.model_type` values. -const SPM_MODEL_UNIGRAM: i64 = 1; -const SPM_MODEL_BPE: i64 = 2; - -/// One vocab piece out of the proto: `(text, score, type)` at index = id. -struct SpmPiece { - text: String, - score: f64, - kind: i64, -} - -/// The slice of a SentencePiece `ModelProto` that tokenizer assembly needs. -struct SpmProto { - pieces: Vec, - model_type: i64, - unk_id: i64, - precompiled_charsmap: Vec, - add_dummy_prefix: bool, - remove_extra_whitespaces: bool, - escape_whitespaces: bool, -} - -/// A bounded protobuf wire-format reader (the same hand-rolled-parser approach -/// as `gguf.rs` — the proto schema is tiny and frozen, a protobuf codegen -/// dependency would be heavier than the format). -struct ProtoReader<'a> { - buf: &'a [u8], - pos: usize, -} - -impl<'a> ProtoReader<'a> { - fn new(buf: &'a [u8]) -> Self { - Self { buf, pos: 0 } - } - - fn done(&self) -> bool { - debug_assert!(self.pos <= self.buf.len(), "pos never overshoots"); - self.pos >= self.buf.len() - } - - /// A base-128 varint, bounded to 10 bytes (the u64 maximum). - fn varint(&mut self) -> Result { - let mut out: u64 = 0; - for shift in 0..10u32 { - let Some(&b) = self.buf.get(self.pos) else { - return Err("varint runs past the end of the buffer".into()); - }; - self.pos += 1; - out |= u64::from(b & 0x7f) << (7 * shift).min(63); - if b & 0x80 == 0 { - return Ok(out); - } - } - Err("varint longer than 10 bytes".into()) - } - - /// The next field key as `(field_number, wire_type)`. - fn key(&mut self) -> Result<(u64, u8), String> { - let key = self.varint()?; - #[allow(clippy::cast_possible_truncation)] // wire type is 3 bits - Ok((key >> 3, (key & 0x7) as u8)) - } - - /// A length-delimited payload (wire type 2). - fn bytes(&mut self) -> Result<&'a [u8], String> { - let len = usize::try_from(self.varint()?).map_err(|_| "length overflows usize")?; - let end = self - .pos - .checked_add(len) - .filter(|&e| e <= self.buf.len()) - .ok_or_else(|| format!("length-delimited field of {len} B runs past the end"))?; - let out = &self.buf[self.pos..end]; - self.pos = end; - Ok(out) - } - - /// Skip one field of the given wire type (unknown/uninteresting fields). - fn skip(&mut self, wire_type: u8) -> Result<(), String> { - match wire_type { - 0 => self.varint().map(|_| ()), - 1 => self.advance(8), - 2 => self.bytes().map(|_| ()), - 5 => self.advance(4), - other => Err(format!("unsupported protobuf wire type {other}")), - } - } - - fn advance(&mut self, n: usize) -> Result<(), String> { - let end = self - .pos - .checked_add(n) - .filter(|&e| e <= self.buf.len()) - .ok_or("fixed-width field runs past the end")?; - self.pos = end; - Ok(()) - } - - /// A fixed32 float (wire type 5). - fn float32(&mut self) -> Result { - let end = self.pos + 4; - let bytes: [u8; 4] = self - .buf - .get(self.pos..end) - .and_then(|s| s.try_into().ok()) - .ok_or("float32 runs past the end")?; - self.pos = end; - Ok(f32::from_le_bytes(bytes)) - } -} - -/// Parse one `ModelProto.SentencePiece` message: `piece`(1), `score`(2), -/// `type`(3, default NORMAL). -fn parse_spm_piece(buf: &[u8]) -> Result { - let mut r = ProtoReader::new(buf); - let mut piece = SpmPiece { - text: String::new(), - score: 0.0, - kind: SPM_TYPE_NORMAL, - }; - while !r.done() { - let (field, wire) = r.key()?; - match (field, wire) { - (1, 2) => { - piece.text = String::from_utf8(r.bytes()?.to_vec()) - .map_err(|_| "piece text is not UTF-8".to_string())?; - } - (2, 5) => piece.score = f64::from(r.float32()?), - (3, 0) => { - piece.kind = i64::try_from(r.varint()?).map_err(|_| "piece type overflows")?; - } - (_, w) => r.skip(w)?, - } - } - Ok(piece) -} - -/// Parse a SentencePiece `ModelProto`: `pieces`(1, repeated), -/// `trainer_spec`(2) for `model_type`(3)/`unk_id`(40), `normalizer_spec`(3) -/// for `precompiled_charsmap`(2)/`add_dummy_prefix`(3)/ -/// `remove_extra_whitespaces`(4). Unknown fields are skipped, truncation is a -/// loud error, and the piece count is bounded. -fn parse_model_proto(buf: &[u8]) -> Result { - let mut proto = SpmProto { - pieces: Vec::new(), - model_type: SPM_MODEL_UNIGRAM, - // The sentencepiece defaults (overridden by every real trainer_spec). - unk_id: 0, - precompiled_charsmap: Vec::new(), - add_dummy_prefix: true, - remove_extra_whitespaces: true, - escape_whitespaces: true, - }; - let mut r = ProtoReader::new(buf); - while !r.done() { - let (field, wire) = r.key()?; - match (field, wire) { - (1, 2) => { - if proto.pieces.len() >= MAX_SPM_PIECES { - return Err(format!("more than {MAX_SPM_PIECES} pieces")); - } - proto.pieces.push(parse_spm_piece(r.bytes()?)?); - } - (2, 2) => { - let mut t = ProtoReader::new(r.bytes()?); - while !t.done() { - let (f, w) = t.key()?; - match (f, w) { - (3, 0) => { - proto.model_type = - i64::try_from(t.varint()?).map_err(|_| "model_type overflows")?; - } - (40, 0) => { - proto.unk_id = - i64::try_from(t.varint()?).map_err(|_| "unk_id overflows")?; - } - (_, w) => t.skip(w)?, - } - } - } - (3, 2) => { - let mut n = ProtoReader::new(r.bytes()?); - while !n.done() { - let (f, w) = n.key()?; - match (f, w) { - (2, 2) => proto.precompiled_charsmap = n.bytes()?.to_vec(), - (3, 0) => proto.add_dummy_prefix = n.varint()? != 0, - (4, 0) => proto.remove_extra_whitespaces = n.varint()? != 0, - (5, 0) => proto.escape_whitespaces = n.varint()? != 0, - (_, w) => n.skip(w)?, - } - } - } - (_, w) => r.skip(w)?, - } - } - if proto.pieces.is_empty() { - return Err("proto declares no pieces".into()); - } - Ok(proto) -} - -/// Build an HF [`Tokenizer`] from a SentencePiece `tokenizer.model` proto -/// (the Llama/Gemma/T5-family format) — the same pipelines HF's own -/// `convert_slow_tokenizer` assembles, dispatched on the proto's -/// `model_type`: -/// -/// - **UNIGRAM** (T5/ALBERT/Gemma): `Precompiled` charsmap normalizer (+ a -/// `{2,}`-space collapse when `remove_extra_whitespaces`), a Metaspace -/// pre-tokenizer/decoder driven by `add_dummy_prefix`, and a `Unigram` -/// model over the proto's pieces (index = token id). -/// - **BPE** (Llama-2 family): merges reconstructed from the vocab + scores -/// (HF's `SentencePieceExtractor` algorithm, literally), a `▁`-prepend + -/// space→`▁` normalizer, no pre-tokenizer, and the -/// Replace/ByteFallback/Fuse/Strip decoder chain. -/// -/// In both, CONTROL/UNKNOWN pieces become special added tokens, USER_DEFINED -/// become plain added tokens, and every added id is verified post-build -/// exactly like the GGUF path. Tokens a checkpoint adds *beyond* the proto -/// (T5's ``, chat specials) live in sibling metadata -/// (`tokenizer_config.json`), not the proto — add them via -/// [`Tokenizer::add_special_tokens`] after this returns. -/// -/// Faithfulness is verified against the same checkpoints' `tokenizer.json` -/// (byte-identical ids over a battery of prompts) in `tests/real_spm.rs`. -pub fn tokenizer_from_spm(path: &Path) -> Result { - let meta = std::fs::metadata(path).map_err(|e| format!("stat {}: {e}", path.display()))?; - if meta.len() > MAX_SPM_FILE_BYTES { - return Err(format!( - "{} is {} B — over the {MAX_SPM_FILE_BYTES} B tokenizer.model bound", - path.display(), - meta.len() - )); - } - let buf = std::fs::read(path).map_err(|e| format!("read {}: {e}", path.display()))?; - let proto = parse_model_proto(&buf).map_err(|e| format!("{}: {e}", path.display()))?; - let mut tok = match proto.model_type { - SPM_MODEL_UNIGRAM => assemble_spm_unigram(&proto)?, - SPM_MODEL_BPE => assemble_spm_bpe(&proto)?, - other => { - return Err(format!( - "SentencePiece model_type {other} is not supported (UNIGRAM and BPE are)" - )); - } - }; - add_and_verify_spm_specials(&mut tok, &proto)?; - Ok(tok) -} - -/// The proto's declared `unk_id`, range-checked against the pieces. -fn spm_unk_index(proto: &SpmProto) -> Result { - debug_assert!(!proto.pieces.is_empty(), "parse rejects empty protos"); - usize::try_from(proto.unk_id) - .ok() - .filter(|&u| u < proto.pieces.len()) - .ok_or_else(|| format!("unk_id {} is out of vocab range", proto.unk_id)) -} - -/// The UNIGRAM assembly (T5/ALBERT/Gemma family) — see [`tokenizer_from_spm`]. -fn assemble_spm_unigram(proto: &SpmProto) -> Result { - use tokenizers::decoders::metaspace::{Metaspace, PrependScheme}; - use tokenizers::models::unigram::Unigram; - use tokenizers::normalizers::replace::ReplacePattern; - use tokenizers::normalizers::{Precompiled, Replace, Sequence as NormSequence}; - - let unk = spm_unk_index(proto)?; - let byte_fallback = proto.pieces.iter().any(|p| p.kind == SPM_TYPE_BYTE); - let vocab: Vec<(String, f64)> = proto - .pieces - .iter() - .map(|p| (p.text.clone(), p.score)) - .collect(); - debug_assert!(vocab.len() == proto.pieces.len(), "vocab keeps every index"); - let model = - Unigram::from(vocab, Some(unk), byte_fallback).map_err(|e| format!("unigram: {e}"))?; - let mut tok = Tokenizer::new(model); - - let mut normalizers: Vec = Vec::with_capacity(2); - if !proto.precompiled_charsmap.is_empty() { - let pre = Precompiled::from(&proto.precompiled_charsmap) - .map_err(|e| format!("precompiled charsmap: {e}"))?; - normalizers.push(pre.into()); - } - if proto.remove_extra_whitespaces { - // A REGEX pattern (multi-space collapse), exactly as the reference - // tokenizer.json spells it — a string pattern would match literally. - let collapse = Replace::new(ReplacePattern::Regex(" {2,}".into()), " ") - .map_err(|e| format!("replace: {e}"))?; - normalizers.push(collapse.into()); - } - if !normalizers.is_empty() { - tok.with_normalizer(Some(NormSequence::new(normalizers))) - .map_err(|e| format!("normalizer: {e}"))?; - } - - let prepend = if proto.add_dummy_prefix { - PrependScheme::Always - } else { - PrependScheme::Never - }; - let metaspace = Metaspace::new('\u{2581}', prepend, true); - tok.with_pre_tokenizer(Some(metaspace.clone())); - tok.with_decoder(Some(metaspace)); - Ok(tok) -} - -/// The BPE assembly (Llama-2 family) — see [`tokenizer_from_spm`]. The -/// pipeline shape is pinned by the family's own `tokenizer.json`: -/// `Prepend(▁)` + `Replace(" "→"▁")` normalizers, NO pre-tokenizer, and the -/// `Replace("▁"→" ")` / `ByteFallback` / `Fuse` / `Strip(" ")` decoder chain. -fn assemble_spm_bpe(proto: &SpmProto) -> Result { - use tokenizers::decoders::byte_fallback::ByteFallback; - use tokenizers::decoders::fuse::Fuse; - use tokenizers::decoders::sequence::Sequence as DecoderSequence; - use tokenizers::decoders::strip::Strip; - use tokenizers::normalizers::replace::ReplacePattern; - use tokenizers::normalizers::{Precompiled, Prepend, Replace, Sequence as NormSequence}; - - let unk = spm_unk_index(proto)?; - let byte_fallback = proto.pieces.iter().any(|p| p.kind == SPM_TYPE_BYTE); - let mut vocab = Vocab::default(); - for (index, p) in proto.pieces.iter().enumerate() { - #[allow(clippy::cast_possible_truncation)] // bounded by MAX_SPM_PIECES - vocab.insert(p.text.clone(), index as u32); - } - if vocab.len() != proto.pieces.len() { - return Err("duplicate piece text in the proto".into()); - } - let merges = extract_spm_merges(&proto.pieces, &vocab); - let bpe = BPE::builder() - .vocab_and_merges(vocab, merges) - .unk_token(proto.pieces[unk].text.clone()) - .fuse_unk(true) - .byte_fallback(byte_fallback) - .build() - .map_err(|e| format!("BPE build: {e}"))?; - let mut tok = Tokenizer::new(bpe); - - let mut normalizers: Vec = Vec::with_capacity(3); - if !proto.precompiled_charsmap.is_empty() { - let pre = Precompiled::from(&proto.precompiled_charsmap) - .map_err(|e| format!("precompiled charsmap: {e}"))?; - normalizers.push(pre.into()); - } - if proto.add_dummy_prefix { - normalizers.push(Prepend::new("\u{2581}".into()).into()); - } - if proto.escape_whitespaces { - let escape = Replace::new(ReplacePattern::String(" ".into()), "\u{2581}") - .map_err(|e| format!("replace: {e}"))?; - normalizers.push(escape.into()); - } - if !normalizers.is_empty() { - tok.with_normalizer(Some(NormSequence::new(normalizers))) - .map_err(|e| format!("normalizer: {e}"))?; - } - - let unescape = Replace::new(ReplacePattern::String("\u{2581}".into()), " ") - .map_err(|e| format!("replace: {e}"))?; - let mut decoders: Vec = vec![unescape.into()]; - if byte_fallback { - decoders.push(ByteFallback::new().into()); - } - decoders.push(Fuse::new().into()); - if proto.add_dummy_prefix { - // Strip the ONE leading space the dummy prefix injected. - decoders.push(Strip::new(' ', 1, 0).into()); - } - tok.with_decoder(Some(DecoderSequence::new(decoders))); - Ok(tok) -} - -/// Reconstruct BPE merges from a SentencePiece BPE proto's vocab + scores — -/// HF's `SentencePieceExtractor.extract(vocab_scores)`, literally: every -/// piece contributes each split `(l, r)` whose halves are both pieces, local -/// candidates ordered by `(id(l), id(r))`; the whole list is then -/// STABLE-sorted by score **descending** (a higher score is an earlier -/// merge; ties keep piece-id order, exactly like Python's stable sort over -/// an insertion-ordered dict). -fn extract_spm_merges(pieces: &[SpmPiece], vocab: &Vocab) -> Merges { - debug_assert!(!pieces.is_empty(), "parse rejects empty protos"); - debug_assert!(vocab.len() == pieces.len(), "one vocab entry per piece"); - let mut scored: Vec<(String, String, f64)> = Vec::new(); - for p in pieces { - let mut local: Vec<(&str, &str)> = Vec::new(); - for (split, _) in p.text.char_indices().skip(1) { - let (l, r) = p.text.split_at(split); - if vocab.contains_key(l) && vocab.contains_key(r) { - local.push((l, r)); - } - } - local.sort_by_key(|&(l, r)| (vocab[l], vocab[r])); - scored.extend( - local - .into_iter() - .map(|(l, r)| (l.to_string(), r.to_string(), p.score)), - ); - } - scored.sort_by(|a, b| b.2.total_cmp(&a.2)); - scored.into_iter().map(|(l, r, _)| (l, r)).collect() -} - -/// CONTROL/UNKNOWN pieces are the proto's specials (``, ``, ``, -/// …); USER_DEFINED are plain added tokens. Re-add + verify every id like -/// the GGUF path — a drifted id would silently corrupt every prompt. -fn add_and_verify_spm_specials(tok: &mut Tokenizer, proto: &SpmProto) -> Result<(), String> { - let mut added: Vec<(usize, &str, bool)> = Vec::new(); - for (index, p) in proto.pieces.iter().enumerate() { - match p.kind { - SPM_TYPE_CONTROL | SPM_TYPE_UNKNOWN => added.push((index, &p.text, true)), - SPM_TYPE_USER_DEFINED => added.push((index, &p.text, false)), - SPM_TYPE_NORMAL | SPM_TYPE_UNUSED | SPM_TYPE_BYTE => {} - other => return Err(format!("piece {index} has unsupported type {other}")), - } - } - for &(_, text, special) in &added { - let t = AddedToken::from(text.to_string(), special); - if special { - tok.add_special_tokens([t]) - .map_err(|e| format!("add special token '{text}': {e}"))?; - } else { - tok.add_tokens([t]) - .map_err(|e| format!("add token '{text}': {e}"))?; - } - } - for &(index, text, _) in &added { - let got = tok.token_to_id(text); - #[allow(clippy::cast_possible_truncation)] // bounded by MAX_SPM_PIECES - if got != Some(index as u32) { - return Err(format!( - "added token '{text}' resolved to id {got:?}, the proto says {index}" - )); - } - } - Ok(()) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::gguf::GgufValue; - - /// A minimal synthetic GGUF header carrying a 4-token BPE tokenizer. - fn toy_gguf() -> GgufFile { - let strs = |v: &[&str]| { - GgufValue::Array(v.iter().map(|s| GgufValue::Str((*s).to_string())).collect()) - }; - let ints = |v: &[i32]| GgufValue::Array(v.iter().map(|&i| GgufValue::I32(i)).collect()); - GgufFile { - path: std::path::PathBuf::new(), - version: 3, - metadata: vec![ - ("tokenizer.ggml.model".into(), GgufValue::Str("gpt2".into())), - ("tokenizer.ggml.pre".into(), GgufValue::Str("qwen2".into())), - ( - "tokenizer.ggml.tokens".into(), - strs(&["a", "b", "ab", "<|stop|>", "[PAD4]"]), - ), - ("tokenizer.ggml.token_type".into(), ints(&[1, 1, 1, 3, 5])), - ("tokenizer.ggml.merges".into(), strs(&["a b"])), - ], - tensors: vec![], - alignment: 32, - data_offset: 0, - } - } - - #[test] - fn toy_gguf_tokenizer_encodes_merges_and_specials() { - let tok = tokenizer_from_gguf(&toy_gguf()).expect("builds"); - let ids = tok.encode("ab", true).expect("encodes"); - assert_eq!(ids.get_ids(), &[2], "the merge applies"); - assert_eq!(tok.token_to_id("<|stop|>"), Some(3), "control token id"); - let ids = tok.encode("<|stop|>ab", true).expect("encodes"); - assert_eq!(ids.get_ids(), &[3, 2], "special token is never split"); - // Unused padding entries are skipped, not part of the vocab. - assert_eq!(tok.token_to_id("[PAD4]"), None); - } - - #[test] - fn unknown_model_pre_and_bad_merges_fail_loudly() { - let mut f = toy_gguf(); - f.metadata[0].1 = GgufValue::Str("sentencepiece".into()); - assert!(tokenizer_from_gguf(&f).unwrap_err().contains("gpt2")); - - let mut f = toy_gguf(); - f.metadata[1].1 = GgufValue::Str("llama-bpe".into()); - assert!( - tokenizer_from_gguf(&f) - .unwrap_err() - .contains("pre-tokenizer") - ); - - let mut f = toy_gguf(); - f.metadata[4].1 = GgufValue::Array(vec![GgufValue::Str("nospace".into())]); - assert!(tokenizer_from_gguf(&f).unwrap_err().contains("merge")); - } - - // ---- SentencePiece proto reader + assembly ---------------------------- - - fn pb_varint(mut v: u64) -> Vec { - let mut out = Vec::new(); - loop { - let byte = (v & 0x7f) as u8; - v >>= 7; - if v == 0 { - out.push(byte); - return out; - } - out.push(byte | 0x80); - } - } - - fn pb_field(field: u64, wire: u8, payload: &[u8]) -> Vec { - let mut out = pb_varint((field << 3) | u64::from(wire)); - if wire == 2 { - out.extend(pb_varint(payload.len() as u64)); - } - out.extend_from_slice(payload); - out - } - - /// One `SentencePiece` message: piece(1)=text, score(2)=f32, type(3). - fn pb_piece(text: &str, score: f32, kind: Option) -> Vec { - let mut msg = pb_field(1, 2, text.as_bytes()); - msg.extend(pb_field(2, 5, &score.to_le_bytes())); - if let Some(k) = kind { - msg.extend(pb_field(3, 0, &pb_varint(k as u64))); - } - msg - } - - /// A tiny UNIGRAM proto: ``(UNKNOWN) ``(CONTROL) then normal - /// pieces, `unk_id = 0`, `add_dummy_prefix` + `remove_extra_whitespaces`. - fn toy_spm(model_type: i64) -> Vec { - let mut buf = Vec::new(); - for msg in [ - pb_piece("", 0.0, Some(SPM_TYPE_UNKNOWN)), - pb_piece("", 0.0, Some(SPM_TYPE_CONTROL)), - pb_piece("\u{2581}", -2.0, None), - pb_piece("\u{2581}hello", -1.0, None), - pb_piece("hello", -3.0, None), - ] { - buf.extend(pb_field(1, 2, &msg)); - } - let mut trainer = pb_field(3, 0, &pb_varint(model_type as u64)); - trainer.extend(pb_field(40, 0, &pb_varint(0))); // unk_id - buf.extend(pb_field(2, 2, &trainer)); - let mut norm = pb_field(3, 0, &pb_varint(1)); // add_dummy_prefix - norm.extend(pb_field(4, 0, &pb_varint(1))); // remove_extra_whitespaces - buf.extend(pb_field(3, 2, &norm)); - buf - } - - fn spm_file(bytes: &[u8]) -> std::path::PathBuf { - use std::sync::atomic::{AtomicU64, Ordering}; - // Unique per call: parallel tests in this process must never share a - // file (same-length protos would otherwise collide and race). - static NEXT: AtomicU64 = AtomicU64::new(0); - let n = NEXT.fetch_add(1, Ordering::Relaxed); - let path = - std::env::temp_dir().join(format!("mummu-spm-test-{}-{n}.model", std::process::id())); - std::fs::write(&path, bytes).expect("temp proto writes"); - path - } - - #[test] - fn toy_spm_tokenizer_encodes_with_metaspace_and_specials() { - let path = spm_file(&toy_spm(SPM_MODEL_UNIGRAM)); - let tok = tokenizer_from_spm(&path).expect("builds"); - let _ = std::fs::remove_file(&path); - - // add_dummy_prefix: "hello" → "▁hello" → piece 3. - let ids = tok.encode("hello", false).expect("encodes"); - assert_eq!(ids.get_ids(), &[3], "metaspace prefix + unigram pick"); - // remove_extra_whitespaces collapses the run before metaspace. - let ids = tok.encode("hello hello", false).expect("encodes"); - assert_eq!(ids.get_ids(), &[3, 3], "space run collapses to one"); - // The CONTROL piece is a special added token at its proto index. - assert_eq!(tok.token_to_id(""), Some(1)); - let ids = tok.encode("hello", false).expect("encodes"); - assert_eq!(ids.get_ids()[0], 1, "special token is never split"); - } - - #[test] - fn spm_rejects_unknown_type_truncation_and_bad_unk() { - // model_type 3 = WORD — neither UNIGRAM nor BPE. - let path = spm_file(&toy_spm(3)); - let err = tokenizer_from_spm(&path).unwrap_err(); - let _ = std::fs::remove_file(&path); - assert!( - err.contains("not supported"), - "WORD protos are a loud error: {err}" - ); - - let full = toy_spm(SPM_MODEL_UNIGRAM); - let path = spm_file(&full[..full.len() - 3]); - let err = tokenizer_from_spm(&path).unwrap_err(); - let _ = std::fs::remove_file(&path); - assert!( - err.contains("end") || err.contains("varint"), - "truncation is a loud error, not a panic: {err}" - ); - - // unk_id past the vocab is rejected. - let mut buf = Vec::new(); - buf.extend(pb_field(1, 2, &pb_piece("x", 0.0, None))); - let mut trainer = pb_field(3, 0, &pb_varint(SPM_MODEL_UNIGRAM as u64)); - trainer.extend(pb_field(40, 0, &pb_varint(7))); - buf.extend(pb_field(2, 2, &trainer)); - let path = spm_file(&buf); - let err = tokenizer_from_spm(&path).unwrap_err(); - let _ = std::fs::remove_file(&path); - assert!(err.contains("unk_id"), "out-of-range unk_id: {err}"); - } - - /// A BPE-shaped toy proto: single chars + the merged pieces, scores - /// encoding merge order (higher = earlier merge, the SPM convention). - fn toy_spm_bpe() -> Vec { - let mut buf = Vec::new(); - for msg in [ - pb_piece("", 0.0, Some(SPM_TYPE_UNKNOWN)), - pb_piece("", 0.0, Some(SPM_TYPE_CONTROL)), - pb_piece("\u{2581}", 0.0, None), - pb_piece("h", 0.0, None), - pb_piece("e", 0.0, None), - pb_piece("l", 0.0, None), - pb_piece("o", 0.0, None), - pb_piece("he", -1.0, None), - pb_piece("ll", -2.0, None), - pb_piece("llo", -3.0, None), - pb_piece("hello", -4.0, None), - pb_piece("\u{2581}hello", -5.0, None), - ] { - buf.extend(pb_field(1, 2, &msg)); - } - let mut trainer = pb_field(3, 0, &pb_varint(SPM_MODEL_BPE as u64)); - trainer.extend(pb_field(40, 0, &pb_varint(0))); - buf.extend(pb_field(2, 2, &trainer)); - let mut norm = pb_field(3, 0, &pb_varint(1)); // add_dummy_prefix - norm.extend(pb_field(5, 0, &pb_varint(1))); // escape_whitespaces - buf.extend(pb_field(3, 2, &norm)); - buf - } - - #[test] - fn toy_spm_bpe_merges_chain_and_decode_round_trips() { - let path = spm_file(&toy_spm_bpe()); - let tok = tokenizer_from_spm(&path).expect("BPE proto builds"); - let _ = std::fs::remove_file(&path); - - // "hello" → prepend+escape "▁hello" → chars merge all the way up the - // score-ordered chain (h+e, l+l, ll+o, he+llo, ▁+hello) → one piece. - let ids = tok.encode("hello", false).expect("encodes"); - assert_eq!(ids.get_ids(), &[11], "the merge chain reaches ▁hello"); - // Decode chain strips the dummy prefix back off. - let text = tok.decode(&[11], true).expect("decodes"); - assert_eq!(text, "hello", "Replace/Fuse/Strip decode chain"); - // The CONTROL piece is a special added token at its proto index. - assert_eq!(tok.token_to_id(""), Some(1)); - } - - #[test] - fn spm_proto_reader_skips_unknown_fields() { - // An unknown length-delimited field (99) + an unknown varint field - // (98) must be skipped, leaving the known fields intact. - let mut buf = pb_field(99, 2, b"ignored"); - buf.extend(pb_field(98, 0, &pb_varint(12345))); - buf.extend(toy_spm(SPM_MODEL_UNIGRAM)); - let proto = parse_model_proto(&buf).expect("parses around unknown fields"); - assert_eq!(proto.pieces.len(), 5); - assert_eq!(proto.unk_id, 0); - assert!(proto.add_dummy_prefix); - } -} diff --git a/crates/mummu/examples/src/tune.rs b/crates/mummu/examples/src/tune.rs deleted file mode 100644 index 0c28036..0000000 --- a/crates/mummu/examples/src/tune.rs +++ /dev/null @@ -1,253 +0,0 @@ -//! The **autotune cache**: where CubeCL persists its kernel picks, and how to -//! throw them away. -//! -//! CubeCL benchmarks several implementations of each kernel the first time it -//! sees one and writes the winner to disk, keyed by (device, kernel, -//! checksum). Later processes load those picks instead of re-tuning, which is -//! what makes a cold start bearable — but the cache has **no invalidation and -//! no re-tune trigger**, so a pick made while the machine was busy is -//! indistinguishable from a good one and is believed forever. Measured on -//! 2026-08-09 (see `bench/BASELINE.md`): a tune that happened during a -//! contended moment cost **21–27 % of f16 decode throughput** in every -//! subsequent process, while the f32 picks from the same moment were -//! unaffected — so the symptom is silent, partial, and permanent. -//! -//! This module is the repair: report where the cache lives and delete it, so a -//! consumer can offer a "re-tune GPU kernels" action instead of shipping a bad -//! tune to a user forever. It reads the same configuration CubeCL reads -//! (`[cubecl.autotune] cache` from the `cubecl.toml` / `burn.toml` discovered -//! by walking up from the process CWD), so the path is right by construction -//! rather than by convention. - -use std::path::{Path, PathBuf}; - - -/// File-name stem of the environment database cubecl 0.11 persists autotune -/// picks into, under the configured cache root. Load-bearing for safety: -/// [`clear_autotune_cache`] removes only files starting with this stem, never -/// the root itself (which defaults to the Cargo `target/` tree). -const ENVIRONMENT_DB_STEM: &str = "environment"; - -/// Bound on a cache walk. The layout is -/// `/autotune///.json.log` — a few hundred -/// files on a normal machine. A runaway count means the root is pointing -/// somewhere it should not, and is an error rather than a long walk. -const MAX_CACHE_FILES: usize = 65_536; -/// Bound on recursion depth for the same reason (the real layout is 3 deep). -const MAX_CACHE_DEPTH: usize = 8; - -/// What went wrong inspecting or clearing the cache. -#[derive(Debug, thiserror::Error)] -pub enum TuneError { - /// The cache directory could not be read or removed. - #[error("autotune cache i/o at {path}: {message}")] - Io { - /// The path being read or removed. - path: PathBuf, - /// The underlying OS error. - message: String, - }, - /// The cache tree is larger or deeper than any real autotune cache, which - /// means the configured root is not what we think it is. Refused rather - /// than walked (or deleted). - #[error("autotune cache at {path} is implausible ({what}) — refusing to touch it")] - Implausible { - /// The configured cache directory. - path: PathBuf, - /// Which bound was exceeded. - what: String, - }, -} - -/// Where the autotune cache lives, and how much of it there is. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct TuneCacheReport { - /// The `/autotune` directory CubeCL writes to. - pub dir: PathBuf, - /// Number of cache files found (0 when the cache does not exist yet). - pub files: usize, - /// Total bytes of those files. - pub bytes: u64, -} - -impl TuneCacheReport { - /// Has anything been tuned and persisted yet? - #[must_use] - pub fn is_empty(&self) -> bool { - self.files == 0 - } -} - -/// The directory CubeCL persists autotune picks to, per the configuration it -/// would itself discover. -/// -/// **Reads the global config**, which initializes it if no one has yet — the -/// same one-shot singleton `RuntimeConfig::set` writes to. A consumer that -/// wants to `set` a custom config must do so *before* calling this (and before -/// building any backend), exactly as CubeCL requires. -#[must_use] -pub fn autotune_cache_dir() -> PathBuf { - // cubecl 0.11 moved autotune persistence out of a per-file directory and - // into an environment database under the configured cache root (the - // `AutotuneConfig` no longer carries a path at all — only `disable_cache`). - // The root is still the thing a consumer wants to report and clear, so - // this returns it directly. - cubecl_runtime::config::cache::CacheConfig::default().root() -} - -/// Measure the persisted cache without changing it. A missing directory is -/// not an error — it means nothing has been tuned yet. -pub fn autotune_cache_report() -> Result { - let dir = autotune_cache_dir(); - let (files, bytes) = measure(&dir, 0)?; - Ok(TuneCacheReport { dir, files, bytes }) -} - -/// Delete the persisted autotune cache, returning what was removed. -/// -/// The next process to run will re-tune from scratch and write fresh picks — -/// **the next one**, not this one: a running process has already loaded the -/// cache into memory and will keep using and re-writing it, so a consumer -/// should treat this as "re-tune on next launch" (or call it before building a -/// backend). Idempotent: clearing an absent cache reports zero and succeeds. -pub fn clear_autotune_cache() -> Result { - let report = autotune_cache_report()?; - if !report.dir.exists() { - debug_assert!(report.is_empty(), "an absent cache cannot hold files"); - return Ok(report); - } - // Remove the environment database files, never the root itself: under - // cubecl 0.11's default (`CacheConfig::Target`) that root is the Cargo - // `target/` tree, so a recursive delete here would blow away the build. - let mut removed = false; - for entry in std::fs::read_dir(&report.dir).map_err(|e| TuneError::Io { - path: report.dir.clone(), - message: e.to_string(), - })? { - let entry = entry.map_err(|e| TuneError::Io { - path: report.dir.clone(), - message: e.to_string(), - })?; - let path = entry.path(); - let is_db = path - .file_name() - .and_then(|n| n.to_str()) - .is_some_and(|n| n.starts_with(ENVIRONMENT_DB_STEM)); - if is_db && path.is_file() { - std::fs::remove_file(&path).map_err(|e| TuneError::Io { - path: path.clone(), - message: e.to_string(), - })?; - removed = true; - } - } - debug_assert!( - removed || report.is_empty(), - "a non-empty cache should have had a database to remove" - ); - Ok(report) -} - -/// Bounded recursive walk: `(file count, total bytes)` under `dir`. -fn measure(dir: &Path, depth: usize) -> Result<(usize, u64), TuneError> { - if depth > MAX_CACHE_DEPTH { - return Err(TuneError::Implausible { - path: dir.to_path_buf(), - what: format!("deeper than {MAX_CACHE_DEPTH} levels"), - }); - } - if !dir.is_dir() { - return Ok((0, 0)); - } - let entries = std::fs::read_dir(dir).map_err(|e| TuneError::Io { - path: dir.to_path_buf(), - message: e.to_string(), - })?; - - let (mut files, mut bytes) = (0usize, 0u64); - for entry in entries { - let entry = entry.map_err(|e| TuneError::Io { - path: dir.to_path_buf(), - message: e.to_string(), - })?; - let path = entry.path(); - if path.is_dir() { - let (f, b) = measure(&path, depth + 1)?; - files += f; - bytes += b; - } else { - files += 1; - bytes += entry.metadata().map(|m| m.len()).unwrap_or(0); - } - if files > MAX_CACHE_FILES { - return Err(TuneError::Implausible { - path: dir.to_path_buf(), - what: format!("more than {MAX_CACHE_FILES} files"), - }); - } - } - Ok((files, bytes)) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn the_cache_dir_is_an_absolute_root() { - let dir = autotune_cache_dir(); - // And it must be an absolute path — the roots CubeCL can resolve to - // (CWD, the project target dir, the user config dir, or an explicit - // file path) are all absolute in practice, and a relative one would - // make "clear the cache" depend on the caller's CWD. - assert!(dir.is_absolute(), "cache dir {dir:?} must be absolute"); - } - - #[test] - fn measuring_an_absent_directory_reports_empty() { - let missing = std::env::temp_dir().join("mummu-no-such-autotune-dir-9e3f"); - assert!(!missing.exists(), "fixture path must not exist"); - assert_eq!( - measure(&missing, 0).expect("absent is not an error"), - (0, 0) - ); - } - - #[test] - fn measuring_counts_files_and_bytes_across_nested_dirs() { - let root = std::env::temp_dir().join("mummu-tune-measure-a41c"); - let nested = root.join("0.10.0").join("device-4-0"); - std::fs::create_dir_all(&nested).expect("fixture dirs"); - std::fs::write(nested.join("matmul.json.log"), b"12345").expect("fixture file"); - std::fs::write(nested.join("reduce.json.log"), b"678").expect("fixture file"); - - let (files, bytes) = measure(&root, 0).expect("walks"); - assert_eq!(files, 2, "both nested files counted"); - assert_eq!(bytes, 8, "byte totals summed across directories"); - - std::fs::remove_dir_all(&root).expect("fixture cleanup"); - } - - #[test] - fn measuring_refuses_an_implausibly_deep_tree() { - // Depth is checked before the directory is read, so a synthetic path - // is enough — no need to build a 9-level fixture. - let deep = std::env::temp_dir().join("mummu-tune-depth"); - let err = measure(&deep, MAX_CACHE_DEPTH + 1).expect_err("too deep is an error"); - assert!( - matches!(err, TuneError::Implausible { .. }), - "expected Implausible, got {err:?}" - ); - } - - #[test] - fn clearing_an_absent_cache_is_a_successful_no_op() { - // `clear` on a machine that has never tuned must not error; the real - // dir may or may not exist here, so assert on the shape of the result. - let before = autotune_cache_report().expect("report"); - if !before.dir.exists() { - let cleared = clear_autotune_cache().expect("clearing nothing succeeds"); - assert!(cleared.is_empty(), "nothing to clear reports empty"); - } - } -} diff --git a/crates/mummu/examples/src/vram.rs b/crates/mummu/examples/src/vram.rs deleted file mode 100644 index 1ce2498..0000000 --- a/crates/mummu/examples/src/vram.rs +++ /dev/null @@ -1,218 +0,0 @@ -//! How much video memory is *actually* free right now, across every process. -//! -//! Placement needs a number that moves when someone else takes VRAM. Two -//! sources, and the difference between them matters: -//! -//! * **DXGI's budget** ([`crate::backend::video_memory`]) is what the OS says -//! *this process* may use. Windows permits oversubscription and pages VRAM -//! behind your back, so it happily reports ~15 GiB of a 16 GiB card while -//! another process holds 9 GiB of it. Measured on this box, 2026-08-23. -//! Useful as a ceiling, useless as "what is free". -//! * **NVML** reports the card's global `total`/`used`/`free` — the same -//! numbers `nvidia-smi` prints, because that is what nvidia-smi calls. This -//! is the honest answer, and it is what a rebalance needs. -//! -//! NVML is loaded at runtime rather than linked, because it ships with the -//! NVIDIA driver and a machine without one must still run: a missing DLL -//! degrades to `None`, never to a failed process start. - -/// A snapshot of one adapter's global memory use. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct Memory { - pub total: u64, - /// Held by every process on the machine, this one included. - pub used: u64, - pub free: u64, -} - -impl Memory { - /// What a model may take without pushing the card into paging, leaving - /// `reserve` for the desktop and for allocations that are not weights - /// (activations, KV state, kernel workspaces). - /// - /// Saturating on purpose: when the card is already fuller than the - /// reserve, the answer is zero, not a wrapped enormous number. - #[must_use] - pub fn headroom(self, reserve: u64) -> u64 { - self.free.saturating_sub(reserve) - } -} - -/// Global VRAM use for the primary GPU, or `None` when nothing on this -/// machine will say. -/// -/// Callers must treat `None` as "no information" and hold their current -/// placement — assuming plenty risks an OOM mid-generation, and assuming -/// pressure needlessly demotes a model that was running fine. -#[must_use] -pub fn memory() -> Option { - nvml::memory() -} - -/// NVML, loaded by hand so its absence is a `None` and not a link error. -#[cfg(windows)] -mod nvml { - use super::Memory; - use core::ffi::{c_char, c_void}; - use std::sync::OnceLock; - - /// `nvmlMemory_t`, verbatim layout. - #[repr(C)] - #[derive(Default, Clone, Copy)] - struct NvmlMemory { - total: u64, - free: u64, - used: u64, - } - - type Init = unsafe extern "C" fn() -> i32; - type HandleByIndex = unsafe extern "C" fn(u32, *mut *mut c_void) -> i32; - type GetMemoryInfo = unsafe extern "C" fn(*mut c_void, *mut NvmlMemory) -> i32; - - #[link(name = "kernel32", kind = "raw-dylib")] - unsafe extern "system" { - fn LoadLibraryA(name: *const c_char) -> *mut c_void; - fn GetProcAddress(module: *mut c_void, name: *const c_char) -> *mut c_void; - } - - /// The three entry points we need, resolved once. - struct Api { - handle_by_index: HandleByIndex, - get_memory_info: GetMemoryInfo, - } - - // SAFETY: the fields are function pointers into a DLL that is never - // unloaded (no FreeLibrary anywhere), so they stay valid for the process - // lifetime and are safe to call from any thread — NVML is thread-safe. - unsafe impl Send for Api {} - unsafe impl Sync for Api {} - - /// Resolve NVML once, but only CACHE A SUCCESS. - /// - /// Caching a failure was a real bug: one transient miss made this return - /// `None` for the life of the process, `backend_budget` then fell back to - /// its configured ceiling unreduced, and the planner put 14.53 GiB of - /// weights on a 16 GiB card — every generation dying with `wgpu error: - /// Out of Memory` while a standalone probe read the card fine. A reading - /// this load-bearing must be allowed to recover. - fn api() -> Option<&'static Api> { - static API: OnceLock = OnceLock::new(); - if let Some(api) = API.get() { - return Some(api); - } - let resolved = resolve()?; - // A race just means two threads resolved it; both are equivalent. - let _ = API.set(resolved); - API.get() - } - - /// One attempt at loading NVML and finding its entry points. - fn resolve() -> Option { - (|| { - // SAFETY: literal, NUL-terminated names; every returned pointer - // is null-checked before it is transmuted to a function pointer. - unsafe { - let module = LoadLibraryA(c"nvml.dll".as_ptr()); - if module.is_null() { - return None; - } - let symbol = |name: &core::ffi::CStr| { - let p = GetProcAddress(module, name.as_ptr()); - (!p.is_null()).then_some(p) - }; - // `_v2` where NVML versioned the ABI; the unsuffixed names - // are the older, incompatible signatures. - let init: Init = core::mem::transmute(symbol(c"nvmlInit_v2")?); - let handle_by_index: HandleByIndex = - core::mem::transmute(symbol(c"nvmlDeviceGetHandleByIndex_v2")?); - let get_memory_info: GetMemoryInfo = - core::mem::transmute(symbol(c"nvmlDeviceGetMemoryInfo")?); - // NVML_SUCCESS is 0. Init is idempotent and refcounted; we - // never shut down, matching the never-unloaded module above. - if init() != 0 { - return None; - } - Some(Api { - handle_by_index, - get_memory_info, - }) - } - })() - } - - pub fn memory() -> Option { - let api = api()?; - // SAFETY: `api` resolved successfully, so NVML is initialised. Both - // calls write through out-pointers to stack locals and are checked - // against NVML_SUCCESS before the values are read. - unsafe { - let mut device: *mut c_void = core::ptr::null_mut(); - // Device 0: the primary GPU. Multi-GPU placement picks its own - // devices and is a separate concern from this global reading. - if (api.handle_by_index)(0, &mut device) != 0 || device.is_null() { - return None; - } - let mut mem = NvmlMemory::default(); - if (api.get_memory_info)(device, &mut mem) != 0 { - return None; - } - Some(Memory { - total: mem.total, - used: mem.used, - free: mem.free, - }) - } - } -} - -#[cfg(not(windows))] -mod nvml { - use super::Memory; - - /// NVML exists on Linux as `libnvidia-ml.so.1`; wiring it up is the same - /// shape as the Windows path and a follow-up. - pub fn memory() -> Option { - None - } -} - -#[cfg(test)] -mod tests { - use super::*; - - /// Whatever NVML reports has to be internally consistent and match the - /// card. This is the guard against a wrong struct layout or a mis-resolved - /// symbol, both of which would return plausible-looking nonsense. - #[test] - fn reported_memory_is_self_consistent() { - let Some(m) = memory() else { - return; // no NVIDIA driver here; nothing to check - }; - assert!(m.total > 0, "a card with no memory is a bad reading"); - assert!( - m.used + m.free <= m.total + (64 << 20), - "used {} + free {} overshoots total {}", - m.used, - m.free, - m.total - ); - assert!(m.free <= m.total); - // Anything under 256 MiB or over 256 GiB is not a GPU we can believe. - assert!( - (256 << 20..=256u64 << 30).contains(&m.total), - "implausible total {}", - m.total - ); - } - - /// Headroom never wraps, however full the card is. - #[test] - fn headroom_saturates_when_the_card_is_full() { - let m = Memory { - total: 16 << 30, - used: 16 << 30, - free: 0, - }; - assert_eq!(m.headroom(2 << 30), 0); - } -} diff --git a/crates/mummu/examples/src/workingset.rs b/crates/mummu/examples/src/workingset.rs deleted file mode 100644 index ad7c6e6..0000000 --- a/crates/mummu/examples/src/workingset.rs +++ /dev/null @@ -1,426 +0,0 @@ -//! **VRAM as a working set** — P9 stage 4: host RAM is the backing store, -//! device memory is a cache, and the CPU (or an iGPU) computes whatever the -//! cache could not supply in time. -//! -//! The tiering in [`crate::tier`] gives every unit a *permanent* home, which -//! caps how much of a model can ever run on the fast device: a 15 GB model -//! and a 9 GiB budget means most units live — and therefore compute — on the -//! CPU forever. The working-set design breaks that cap. Every unit lives in -//! host RAM; the device holds only what the *current* layers need, staged in -//! ahead of use and evicted behind it, so every layer can execute on the -//! fast device without any unit owning a permanent slot. -//! -//! # What makes this schedulable rather than a cache heuristic -//! -//! The access sequence is **known in advance and identical every token**: -//! layer 0, 1, … L-1, then the next token repeats it. Two consequences the -//! scheduler exploits, neither available to a general-purpose cache: -//! -//! - **Eviction is optimal, not heuristic.** Bélády's rule — evict the entry -//! whose next use is furthest away — is normally unimplementable because -//! it needs the future. Here the future is a `for` loop. [`Plan::victim`] -//! applies it exactly. -//! - **Prefetch distance is a decision, not a guess.** Staging for layer -//! `L + d` is issued while layer `L` computes; `d` is chosen from measured -//! staging bandwidth and per-layer compute time, not from access history. -//! -//! # The bound that shapes the policy -//! -//! Streaming is not free, and for a *dense* model it cannot beat the CPU on -//! bandwidth alone: every unit is needed every layer, so streaming the whole -//! model each token moves the whole model across the bus each token — at -//! which point the CPU reading the same bytes from its own DDR is no worse. -//! Streaming wins in exactly two situations, and the policy is built around -//! them: -//! -//! 1. **Selectivity.** A routed MoE touches `top_k` of `E` experts per token, -//! so it stages a `k/E` fraction. This is why the MoE conversion matters -//! beyond fit: it turns "move everything" into "move what was routed". -//! 2. **Overlap.** Staging that happens *while the previous layer computes* -//! costs nothing on the critical path until it exceeds compute time. -//! -//! So the policy is: **pin what fits, stream what is selective, overlap -//! always, and fall back to host compute on a miss** — never stall waiting -//! for a transfer, because a stall is strictly worse than computing in place -//! (the CPU already holds the bytes). - -use std::collections::HashMap; - -/// A unit of placement: one MoE expert, or one FFN neuron cluster. -pub type UnitId = usize; - -/// Where a unit's compute happened — what the scheduler reports back so a -/// caller can see whether the cache is earning its keep. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum Served { - /// Resident on the device when the layer ran (the fast path). - Cached, - /// Staged in time by the prefetcher. - Prefetched, - /// Not resident and not staged in time — computed on the host instead. - /// Not an error: computing in place beats stalling on a transfer. - Overflow, -} - -/// One layer's demand: the units it needs, in the order it needs them. -#[derive(Debug, Clone)] -pub struct LayerDemand { - pub layer: usize, - pub units: Vec, -} - -/// What the scheduler decided for one layer. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct LayerSchedule { - pub layer: usize, - /// Units to compute on the device (resident by the time the layer runs). - pub on_device: Vec, - /// Units to compute on the host — the cache could not supply them in - /// time and stalling would be worse. - pub overflow: Vec, - /// Staging to *issue* now, for a later layer, while this layer computes. - pub prefetch: Vec, - /// Units evicted to make room for `prefetch`. - pub evict: Vec, -} - -/// Sizes and speeds the scheduler plans against. All measured, never assumed -/// — see `examples/stage-probe.rs`. -#[derive(Debug, Clone)] -pub struct Budget { - /// Device memory the working set may use. - pub device_bytes: u64, - /// Bytes per unit (uniform in practice: clusters are equal-sized by - /// construction, experts by architecture). - pub unit_bytes: u64, - /// Host→device staging bandwidth, bytes/sec. - pub stage_bytes_per_sec: f64, - /// How long one layer's compute takes on the device, seconds. Staging - /// issued during a layer is free until it exceeds this. - pub layer_compute_secs: f64, -} - -impl Budget { - /// How many units fit in the device budget. - #[must_use] - pub fn capacity(&self) -> usize { - if self.unit_bytes == 0 { - return 0; - } - usize::try_from(self.device_bytes / self.unit_bytes).unwrap_or(usize::MAX) - } - - /// How many units can be staged inside one layer's compute window — - /// the prefetcher's per-layer budget. Staging beyond this lands late and - /// would stall, so the scheduler sends the remainder to the host instead. - #[must_use] - pub fn stageable_per_layer(&self) -> usize { - if self.unit_bytes == 0 || self.stage_bytes_per_sec <= 0.0 { - return 0; - } - let per_unit_secs = self.unit_bytes as f64 / self.stage_bytes_per_sec; - if per_unit_secs <= 0.0 { - return usize::MAX; - } - (self.layer_compute_secs / per_unit_secs).floor().max(0.0) as usize - } -} - -/// The schedule for one full pass (one token) over a known layer sequence. -#[derive(Debug, Clone)] -pub struct Plan { - pub layers: Vec, - /// Units pinned for the whole pass: they fit, and they are used often - /// enough that streaming them would be pure waste. - pub pinned: Vec, - /// Fraction of unit-uses served from the device. - pub hit_rate: f64, -} - -impl Plan { - /// Bélády's optimal victim: of `resident`, the unit whose next use is - /// furthest in the future (or never). Implementable here only because - /// the access sequence is known — see the module header. - /// - /// `next_use[u]` is the index of `u`'s next use, `usize::MAX` if unused - /// again this pass. - #[must_use] - pub fn victim(resident: &[UnitId], next_use: &HashMap) -> Option { - resident - .iter() - .copied() - .max_by_key(|u| next_use.get(u).copied().unwrap_or(usize::MAX)) - } -} - -/// Build a working-set schedule for one pass over `demands`. -/// -/// The policy, in order: -/// -/// 1. **Pin the hot core.** Units used in every layer (a dense model's local -/// slab, an MoE's always-hot experts) never leave the device: streaming -/// something needed every layer is pure overhead. Pinning is capped so -/// the stream always has room to work. -/// 2. **Prefetch ahead.** While layer `L` computes, issue staging for the -/// units layer `L + 1` needs and does not have, bounded by -/// [`Budget::stageable_per_layer`] — what actually fits in the compute -/// window. -/// 3. **Evict optimally.** Make room with Bélády's rule. -/// 4. **Overflow to the host.** Anything not resident and not stageable in -/// time is computed on the host. Never stall: the host already has the -/// bytes, so waiting is strictly worse than computing. -pub fn schedule(demands: &[LayerDemand], budget: &Budget) -> Plan { - let capacity = budget.capacity(); - let per_layer_stage = budget.stageable_per_layer(); - - // --- 1. pin the units every layer needs ------------------------------ - let mut uses: HashMap = HashMap::new(); - for d in demands { - for &u in &d.units { - *uses.entry(u).or_insert(0) += 1; - } - } - // Reserve room for streaming: pinning the entire cache would leave the - // prefetcher nowhere to land, turning every non-pinned unit into - // overflow. - let pin_cap = capacity.saturating_sub(per_layer_stage.max(1)).min(capacity); - let mut hot: Vec<(UnitId, usize)> = uses.iter().map(|(&u, &n)| (u, n)).collect(); - // Hottest first; ties by id so a plan is reproducible. - hot.sort_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(&b.0))); - let pinned: Vec = hot - .iter() - .filter(|(_, n)| *n > 1) // used more than once — worth keeping - .take(pin_cap) - .map(|(u, _)| *u) - .collect(); - - // --- flat access sequence, for optimal eviction ---------------------- - // position -> unit, so "next use of u after position i" is a lookup. - let mut seq: Vec = Vec::new(); - let mut layer_span: Vec<(usize, usize)> = Vec::with_capacity(demands.len()); - for d in demands { - let start = seq.len(); - seq.extend(d.units.iter().copied()); - layer_span.push((start, seq.len())); - } - - let mut resident: Vec = pinned.clone(); - let mut layers: Vec = Vec::with_capacity(demands.len()); - let (mut hits, mut total) = (0usize, 0usize); - - for (li, d) in demands.iter().enumerate() { - let mut on_device = Vec::new(); - let mut overflow = Vec::new(); - for &u in &d.units { - total += 1; - if resident.contains(&u) { - on_device.push(u); - hits += 1; - } else { - // Not staged in time: compute on the host rather than stall. - overflow.push(u); - } - } - - // --- 2/3. prefetch for the next layer, evicting optimally -------- - let mut prefetch = Vec::new(); - let mut evict = Vec::new(); - if let Some(next) = demands.get(li + 1) { - // Next use of each unit, measured from the end of this layer — - // the horizon eviction decisions are made against. - let from = layer_span[li].1; - let mut next_use: HashMap = HashMap::new(); - for (pos, &u) in seq.iter().enumerate().skip(from) { - next_use.entry(u).or_insert(pos); - } - for &u in &next.units { - if prefetch.len() >= per_layer_stage { - break; // the rest cannot land in time; they will overflow - } - if resident.contains(&u) || prefetch.contains(&u) { - continue; - } - if resident.len() >= capacity { - // Evict, but never a pinned unit and never something this - // prefetch round just brought in. - let candidates: Vec = resident - .iter() - .copied() - .filter(|r| !pinned.contains(r) && !prefetch.contains(r)) - .collect(); - let Some(v) = Plan::victim(&candidates, &next_use) else { - break; // nothing evictable — the rest overflows - }; - resident.retain(|r| *r != v); - evict.push(v); - } - resident.push(u); - prefetch.push(u); - } - } - - layers.push(LayerSchedule { - layer: d.layer, - on_device, - overflow, - prefetch, - evict, - }); - } - - let hit_rate = if total == 0 { - 0.0 - } else { - hits as f64 / total as f64 - }; - Plan { - layers, - pinned, - hit_rate, - } -} - -#[cfg(test)] -mod tests { - use super::*; - - fn dense(layers: usize, units_per_layer: usize) -> Vec { - // A dense partitioned model: every layer needs its OWN units, and - // needs all of them (exact mode). - (0..layers) - .map(|l| LayerDemand { - layer: l, - units: (0..units_per_layer) - .map(|u| l * units_per_layer + u) - .collect(), - }) - .collect() - } - - fn budget(capacity_units: usize, stage_per_layer: usize) -> Budget { - // unit = 1 MB; pick bandwidth so exactly `stage_per_layer` units fit - // in one layer's compute window. - let unit = 1_000_000u64; - Budget { - device_bytes: unit * capacity_units as u64, - unit_bytes: unit, - stage_bytes_per_sec: (unit as f64) * stage_per_layer as f64 / 0.010, - layer_compute_secs: 0.010, - } - } - - #[test] - fn capacity_and_stage_window_come_from_measurements() { - let b = budget(8, 3); - assert_eq!(b.capacity(), 8); - assert_eq!(b.stageable_per_layer(), 3); - // A zero-bandwidth device can stage nothing (everything overflows). - let none = Budget { - stage_bytes_per_sec: 0.0, - ..budget(8, 3) - }; - assert_eq!(none.stageable_per_layer(), 0); - } - - #[test] - fn beladys_rule_evicts_the_furthest_next_use() { - let resident = [1usize, 2, 3]; - let next: HashMap = [(1, 10), (2, 99), (3, 20)].into_iter().collect(); - assert_eq!(Plan::victim(&resident, &next), Some(2)); - // A unit with no next use at all is the best possible victim. - let next: HashMap = [(1, 10), (3, 20)].into_iter().collect(); - assert_eq!(Plan::victim(&resident, &next), Some(2)); - assert_eq!(Plan::victim(&[], &next), None); - } - - #[test] - fn a_model_that_fits_is_pinned_and_never_streams() { - // 4 layers x 2 units = 8 units, capacity 16: everything fits. - // Nothing should be evicted, and after the first pass every use is a - // hit. (Units are used once per pass here, so pinning is driven by - // capacity, not by reuse — what matters is that nothing thrashes.) - let demands = dense(4, 2); - let plan = schedule(&demands, &budget(16, 4)); - assert!( - plan.layers.iter().all(|l| l.evict.is_empty()), - "a model that fits must never evict: {:?}", - plan.layers - ); - } - - #[test] - fn streaming_prefetches_the_next_layer_while_this_one_computes() { - // Capacity 4, can stage 2 units per layer window; each layer needs 2. - let demands = dense(6, 2); - let plan = schedule(&demands, &budget(4, 2)); - // Layer 0 issues staging for layer 1's units. - assert_eq!(plan.layers[0].prefetch, vec![2, 3], "{:?}", plan.layers[0]); - // And layer 1 then finds them resident — the point of the pipeline. - assert_eq!(plan.layers[1].on_device, vec![2, 3]); - assert!(plan.layers[1].overflow.is_empty()); - } - - #[test] - fn what_cannot_be_staged_in_time_overflows_to_the_host_never_stalls() { - // Each layer needs 4 units but only 1 can be staged per window. - let demands = dense(4, 4); - let plan = schedule(&demands, &budget(8, 1)); - let overflowed: usize = plan.layers.iter().map(|l| l.overflow.len()).sum(); - assert!( - overflowed > 0, - "a starved stage window must produce host overflow, not a stall" - ); - // Every unit is still accounted for: computed somewhere, every layer. - for (l, d) in plan.layers.iter().zip(&demands) { - assert_eq!( - l.on_device.len() + l.overflow.len(), - d.units.len(), - "every unit must be computed somewhere" - ); - } - } - - #[test] - fn a_routed_moe_reuses_its_hot_experts() { - // The selective case the design exists for: 8 layers routing over a - // small hot set, so the cache actually pays off. - let demands: Vec = (0..8) - .map(|l| LayerDemand { - layer: l, - units: vec![0, 1, (l % 3) + 2], // two always-hot + one rotating - }) - .collect(); - let plan = schedule(&demands, &budget(6, 2)); - assert!( - plan.pinned.contains(&0) && plan.pinned.contains(&1), - "always-hot experts must be pinned, not streamed: {:?}", - plan.pinned - ); - assert!( - plan.hit_rate > 0.6, - "a routed workload with a hot core should mostly hit: {}", - plan.hit_rate - ); - } - - #[test] - fn pinning_never_starves_the_stream() { - // Even when everything looks hot, the scheduler must leave room for - // staging — a fully pinned cache turns every miss into permanent - // overflow. - let demands: Vec = (0..10) - .map(|l| LayerDemand { - layer: l, - units: vec![0, 1, 2, l + 3], - }) - .collect(); - let b = budget(4, 2); - let plan = schedule(&demands, &b); - assert!( - plan.pinned.len() < b.capacity(), - "pinning must leave staging room: pinned {} of {}", - plan.pinned.len(), - b.capacity() - ); - } -}