From 35da9ea3a5e1faf676368bd5a995a20aaac49765 Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Thu, 6 Aug 2026 06:26:01 -0500 Subject: [PATCH 1/5] deps: cargo update (8 crates) - budgets and gates hold `cargo update` moved 8 transitive crates to their newest Rust-1.99-nightly-compatible versions: aho-corasick 1.1.4 -> 1.1.5, android_system_properties 0.1.5 -> 0.1.6, macerator 0.3.3 -> 0.3.4, macerator-macros 0.1.5 -> 0.2.0, minijinja + minijinja-contrib 2.21.0 -> 2.22.0, regex-automata 0.4.16 -> 0.4.18, zlib-rs 0.6.6 -> 0.6.7. macerator is burn-flex's SIMD layer and minijinja backs the hf-chat-template dev-dependency behind the template byte gate, so both touch load-bearing paths. `cargo upgrade --incompatible` offers exactly one bump - wgpu 29 -> 30 - which stays held for the documented reason: wgpu is not ours to pick, burn 0.21 resolves 29 transitively and our direct handle exists only so the startup adapter probe speaks the same wgpu burn does. Taking 30 alone would compile two wgpu copies and probe a different one than the backend uses. It unblocks with a burn bump, not a `cargo upgrade`. Pin watch: burn's newest tag is still 0.22.0-pre.1 (0.21.0 remains the latest stable), so the P0 migration item stays gated; tokenizers 0.23.1 is current. Verified green: cargo fmt --check clean, cargo clippy --all-targets with no warnings from our crates, cargo build, and 202 library unit tests passing by name. Budget gates on an idle card (2.2 GiB ambient, 0% util): GPU 99.2 ms TTFT / 12.9 tok/s (budgets 150 ms / 10 tok/s), CPU 15.86 tok/s (budget 6). The GPU gate's first post-update run read 9.5 tok/s - the documented autotune-cache transient after a dependency change, steady on re-run. Co-Authored-By: Claude Opus 4.8 --- Cargo.lock | 32 ++++++++++++++++---------------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 474205f..966d6a3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -55,9 +55,9 @@ dependencies = [ [[package]] name = "aho-corasick" -version = "1.1.4" +version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" dependencies = [ "memchr", ] @@ -97,9 +97,9 @@ checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" [[package]] name = "android_system_properties" -version = "0.1.5" +version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +checksum = "ae221649c9976a6f6c56ae1facf410f3ddb33cc661c4b7b61020a912d4237fbc" dependencies = [ "libc", ] @@ -3444,9 +3444,9 @@ dependencies = [ [[package]] name = "macerator" -version = "0.3.3" +version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "508a2f720538bb7e3ea3cb6d098615b107cb82ae387d77d906276905e9ec894b" +checksum = "2ddeaadb76e307b1d2b4836adee77910d405b9f6c1ca031bc81a7a398e2b22f2" dependencies = [ "bytemuck", "cfg_aliases", @@ -3460,9 +3460,9 @@ dependencies = [ [[package]] name = "macerator-macros" -version = "0.1.5" +version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed5ce85961d618ce9794bdf822bfe96fe9dd341aa5b033b454f7a8d96e79b9b1" +checksum = "d2a397d020af346358830d6f28579f26283003ce87af3dd79173ac9ee2d3ac28" dependencies = [ "darling 0.20.11", "proc-macro2", @@ -3536,9 +3536,9 @@ checksum = "38d1115007560874e373613744c6fba374c17688327a71c1476d1a5954cc857b" [[package]] name = "minijinja" -version = "2.21.0" +version = "2.22.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cb3d648e68cea56d9858d535ee28f9538404e2dd8cb08ed0bd05dca379477f39" +checksum = "ef84a52be188a1d4124bd717903fdde96ca4705f2b56adfe2d91fcc57fcb6987" dependencies = [ "indexmap", "memo-map", @@ -3548,9 +3548,9 @@ dependencies = [ [[package]] name = "minijinja-contrib" -version = "2.21.0" +version = "2.22.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "85342f6fac0be8ccd5bd00d9066be538f34f393f577b75d81b17c8398a6b43bb" +checksum = "fd6e279dc925840c2d9ebc1e9f85410611d01292561297463ba3e516c3ad94dd" dependencies = [ "minijinja", "serde", @@ -4663,9 +4663,9 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.4.16" +version = "0.4.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fcfdb36bda0c880c5931cdc7a2bcdc8ba4556847b9d912bca70bc94708711ad" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" dependencies = [ "aho-corasick", "memchr", @@ -6719,9 +6719,9 @@ dependencies = [ [[package]] name = "zlib-rs" -version = "0.6.6" +version = "0.6.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b142a20ec14a91d5bc708c1dc21b080c550113d8aa77afa29635673a65dd02c5" +checksum = "34b31d188d9d685a4f9c7b46d6e36631b07058d2cfe190267adce54dc230bf12" [[package]] name = "zmij" From 7d06519204eed31f9cfb908156096cd1cb8e45c2 Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Thu, 6 Aug 2026 07:03:41 -0500 Subject: [PATCH 2/5] bench: flash attention measured and rejected; long-prefill row lands Closes the ROADMAP perf item "Evaluate CubeCL's now-complete flash-attention kernel". It reaches Mummu as burn 0.21's `tensor::module::attention`, which the wgpu backend routes through an autotune set of cubek flash variants, and it is a genuine drop-in for the explicit q.kT -> scale -> mask -> f32-softmax -> .v chain: * `scale: None` asks for the op's own 1/sqrt(head_dim) default - the same factor we applied, and passing it explicitly would silently disqualify the flash kernel (burn-cubecl routes any custom scale to the unfused fallback); * `is_causal: true` reproduces our additive mask exactly - the causal boundary aligns bottom-right, col > row + (seq_k - seq_q), which is the KV cache's rule at every `past`, decode step included; * the f32 island survives *inside* the kernel (AccumulatorPrecision::Strict(F32)), so the reason f16 attention does not NaN is preserved rather than thrown away. It was implemented, proven equivalent by a unit test against the explicit formulation written out longhand, then A/B'd on an idle card (criterion, two runs per arm) - and reverted, because the measurement says it costs more than it buys here: f32 TTFT 90.9 -> 91.9 ms | f32 prefill@2048 593 -> 629 ms (+6.0%) f32 decode 60.0 -> 61.1 ms/tok | f16 TTFT 20.4 -> 16.7 ms (-18%) f16 prefill@2048 210 -> 164 ms (-22%) | f16 decode 20.5 -> 22.9 (+11%) The one real win is f16 prefill, where the accelerated plane matmuls have tiles to fill. Decode is seq_q = 1 - a matvec with no tile reuse, where flash is overhead and (leading hypothesis) an opaque node the Fusion backend cannot absorb the way it absorbs the explicit chain. Taking only the winning quadrant would fork the hottest leaf function on dtype, on a path no strict parity gate covers; that is now its own ROADMAP item with the f16 parity leg it needs listed as a prerequisite. Kept from the work, because it is what made the evaluation decisive: a permanent `ttft_prefill_2048` row in the criterion bench AND in the budget gate. The ~36-token bench prompt materializes a 62 KiB scores tensor; 2048 tokens materializes 201 MiB - only the second length can see an attention change at all. Recorded 593 ms f32 / 210 ms f16, budget 900 ms; the gate passes at 592 ms. The control runs also re-baselined both dtypes, and found drift nobody claimed: f16 decode is 2.7x faster than the 2026-07-12 record (54.5 -> 20.5 ms/token) while f32 decode is 10% slower (54.3 -> 60.0). Both are pre-existing at HEAD. BASELINE.md now records today's numbers and retires its "f16 buys VRAM, not speed" reading; bisecting the f32 drift is a new ROADMAP item, as is the observation that a 10 tok/s ceiling four ms/token below the recorded number can catch a collapse but never a drift. Verified green: fmt, clippy --all-targets (no warnings from our crates), 202 lib unit tests by name, build, and the extended budget gate at 99.2 ms TTFT / 13.2 tok/s / 592 ms prefill@2048. Co-Authored-By: Claude Opus 4.8 --- ROADMAP.md | 46 +++++++++++++++++++++- bench/BASELINE.md | 58 +++++++++++++++++++++++----- crates/mummu-bench/benches/runner.rs | 26 +++++++++++++ crates/mummu-bench/tests/budget.rs | 31 ++++++++++++++- 4 files changed, 150 insertions(+), 11 deletions(-) diff --git a/ROADMAP.md b/ROADMAP.md index afff3e9..10aea4c 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -123,7 +123,7 @@ a benchmark holds/improves its budget; README perf claims link an artifact. with validation armed (no OOB found — clean bill), and the budget gates hold their numbers from the opted-out bench crate (12.4 tok/s with the files ≈ 12.1 without; first run after a config change can read ~30% low while autotune re-tunes — re-run before believing a regression). -- [ ] Evaluate **CubeCL's now-complete flash-attention kernel** for the decode/prefill attention step — +- [x] Evaluate **CubeCL's now-complete flash-attention kernel** for the decode/prefill attention step — the releases page reports a full implementation (causal **masking**, partitions, row-wise reductions, multi-plane ops). Mummu currently materializes attention explicitly (q·kᵀ → f32 softmax island → ·v); a fused flash-attention kernel collapses those into one dispatch, which is squarely @@ -131,6 +131,50 @@ a benchmark holds/improves its budget; README perf claims link an artifact. the O(t²) scores tensor at prefill). Gate strictly: the f32-softmax island is the whole reason the f16 parity holds, so any flash path must re-pass the parity harness (both legs) AND hold/beat `bench/BASELINE.md` before adoption. — https://github.com/tracel-ai/cubecl/releases *(2026-07-17 research)* + *(2026-08-06) **Evaluated end to end, and rejected on measurement — the numbers are in + `bench/BASELINE.md`.** It reaches Mummu as burn 0.21's `tensor::module::attention`, which the + wgpu backend routes to `cubek`'s flash kernel through an autotune set (flash-blackbox-accelerated + variants vs an unfused fallback). It is a genuine drop-in: `scale: None` asks for the op's own + `1/sqrt(head_dim)` default — the same factor, and passing it explicitly would silently + *disqualify* the flash kernel (burn-cubecl routes any custom scale to the fallback) — `is_causal: + true` reproduces our mask exactly (its causal boundary aligns bottom-right, `col > row + (seq_k − + seq_q)`, which IS the KV cache's rule at every `past`), and **the f32 island survives inside the + kernel** (`AccumulatorPrecision::Strict(F32)`), so the whole reason f16 attention doesn't NaN is + preserved rather than discarded. Implemented, proven equivalent by a new unit test against the + explicit formulation written out longhand, then A/B'd (criterion, idle card, two runs per arm) — + and reverted: f16 prefill @2048 −22 % and f16 TTFT −18 % (real wins, the accelerated plane + matmuls have tiles to fill), but f16 decode **+11 %**, f32 decode +1.8 %, f32 prefill @2048 + **+6.0 %**. Decode is `seq_q = 1`, a matvec with no tile reuse where flash is pure overhead and + (leading hypothesis) an opaque node the Fusion backend can't absorb the way it absorbs the + explicit chain's scale/mask/softmax/cast. Adopting only the winning quadrant means a + dtype-conditional fork in the hottest leaf function on a path **no strict parity gate covers** + (the gates run f32 and GGUF-dequant-to-f32) — see the split item below. Kept from the work: a + permanent **`ttft_prefill_2048` row** in the criterion bench and the budget gate (593 ms f32 / + 210 ms f16 recorded, ≤ 900 ms budget), the row where the attention formulation is visible at all — + the ~36-token bench prompt makes a 62 KiB scores tensor, 2048 tokens makes 201 MiB.* +- [ ] **Adopt flash attention for f16 prefill only, once f16 has parity coverage** — the winning + quadrant of the 2026-08-06 evaluation above: `t > 1` (prefill) on an f16 ambient dtype is + −22 % prefill @2048 and −18 % TTFT, and it drops the O(t²) scores tensor (201 MiB at 2048 × 12 + heads today, and it is the term that ends long-context prefill on a 16 GB card — a P6 fit lever, + not only a latency one). Two prerequisites, both deliberate: (a) an **f16 parity leg** — every + strict gate today runs f32 or GGUF-dequant-to-f32, so an f16-only numeric fork would ship + unverified; the cheap shape is an in-process f16-vs-f32 first-forward agreement assert (the + dtype-pinning work makes both aliases coexist — `real_mixed_dtype.rs` already does exactly this + for one token) and the honest shape is llama.cpp at f16 on the `llama_ref` harness; (b) accept a + dtype- **and** length-conditional branch in `GqaAttention::forward`, or find a formulation that + isn't conditional. Re-measure first: the numbers are burn-0.21/wgpu-29-specific. +- [ ] **Bisect the f32 decode drift: 54.3 → 60.0 ms/token since 2026-07-12** — the 2026-08-06 + re-measure found the f32 decode row 10 % slower than recorded while f16 got **2.7× faster** + (54.5 → 20.5 ms/token) over the same period. Neither move was caused by a Mummu change that + claimed them, and the budget gate never noticed because its ceiling (10 tok/s) sits four + ms/token below the recorded number — a ceiling that loose cannot catch drift, only collapse. + Suspects, cheapest first: the 2026-07-30 dtype pinning (explicit `(device, dtype)` at every + creation site may have moved f32 off a fast path while putting f16 on one), a dependency bump in + the CubeCL/wgpu stack, an autotune cache re-tuned against a different machine state, or the + GPU driver. Route: check out the 2026-07-12 tree and re-bench it on today's machine — if it + also reads 60 ms/token the cause is under the repo, if it reads 54 the cause is in it, and + `git bisect` over the bench closes it either way. Worth doing before any further f32 perf work + builds on a number that moved. ## Phases diff --git a/bench/BASELINE.md b/bench/BASELINE.md index 3582b1d..696d88d 100644 --- a/bench/BASELINE.md +++ b/bench/BASELINE.md @@ -8,25 +8,39 @@ recorded numbers (and this file's date) only on a legitimate improvement. ## Qwen2.5-1.5B-Instruct · single GPU · f32 -| Metric | Recorded (2026-07-12, SPIR-V) | Budget | +| Metric | Recorded (2026-08-06) | Budget | | --- | --- | --- | -| TTFT (fresh cache: full prefill + first token) | 96.7 ms | ≤ 150 ms | -| Decode latency (32 greedy tokens, warm KV cache) | 1.737 s → **54.3 ms/token ≈ 18.4 tok/s** | ≥ 10 tok/s | +| TTFT (fresh cache: full prefill + first token) | 90.9 ms | ≤ 150 ms | +| Prefill @ 2048 tokens (`ttft_prefill_2048`) | **593 ms** | ≤ 900 ms | +| Decode latency (32 greedy tokens, warm KV cache) | 1.921 s → **60.0 ms/token ≈ 16.7 tok/s** | ≥ 10 tok/s | | Peak GPU memory during the real-inference suite (whole card, ~3.5 GiB desktop ambient → ~8.0 GiB runner) | 11.5 GiB | ≤ 13 GiB whole-card | ## Qwen2.5-1.5B-Instruct · single GPU · **f16** (weights + KV; f32 attention-score island) -| Metric | Recorded (2026-07-12, SPIR-V) | Budget | +| Metric | Recorded (2026-08-06) | Budget | | --- | --- | --- | -| TTFT (fresh cache: full prefill + first token) | 97.2 ms | ≤ 150 ms | -| Decode latency (32 greedy tokens, warm KV cache) | 1.744 s → **54.5 ms/token ≈ 18.4 tok/s** | ≥ 10 tok/s | +| TTFT (fresh cache: full prefill + first token) | 20.4 ms | ≤ 150 ms | +| Prefill @ 2048 tokens (`ttft_prefill_2048`) | **210 ms** | ≤ 900 ms | +| Decode latency (32 greedy tokens, warm KV cache) | 0.656 s → **20.5 ms/token ≈ 48.8 tok/s** | ≥ 10 tok/s | | Peak GPU memory during decode (whole card, 3.1 GiB ambient → **~3.6 GiB runner**, 2026-07-11 measure) | 6.75 GiB | ≤ 8 GiB whole-card | +**2026-08-06 re-measure — the f16 path is 2.7× faster than 2026-07-12 recorded, f32 10% slower.** +Both tables above were re-run on an idle card (criterion, three runs each, unchanged shipping code) +because the flash-attention evaluation below needed an honest control. f16 decode moved +54.5 → 20.5 ms/token and f16 TTFT 97.2 → 20.4 ms — an improvement nothing in this repo claimed, so +it came from underneath (driver, or a dependency in the CubeCL/wgpu stack, or the 2026-07-30 dtype +pinning removing a per-call policy lookup). f32 decode moved the other way, 54.3 → 60.0 ms/token; +the gate never tripped because its ceiling is 10 tok/s, four ms/token of slack away. Both drifts are +pre-existing at HEAD, not caused by anything this run shipped — bisecting the f32 one is a ROADMAP +item. + The SPIR-V compiler (burn `vulkan` feature) cut decode latency **23%** on both dtypes (70.7 → 54.3 ms/token) at the cost of ~9 ms TTFT (88.4 → 96.7 ms, still ⅔ under its ceiling); parity -held byte-identically (max |Δlogit| 2.670e-5, Ollama greedy leg exact). Decode remains -dispatch-bound — f16 still buys VRAM (~7.9 → ~3.6 GiB runner), not speed; ~54 ms/token streams f32 -weights at only ~114 GB/s vs the card's ~672 GB/s, so kernel/dispatch overhead is still the ceiling. +held byte-identically (max |Δlogit| 2.670e-5, Ollama greedy leg exact). **The "f16 buys VRAM, not +speed" reading from that run no longer holds**: as of the 2026-08-06 re-measure f16 decodes at +20.5 ms/token against f32's 60.0 — 2.9× — so f16 is now the fast path as well as the small one, and +the f32 path is the one that looks dispatch-bound (60 ms/token streams ~6.2 GB of f32 weights at +~103 GB/s against the card's ~672 GB/s). ## Qwen2.5-0.5B-Instruct · CPU (burn-flex) · f32 @@ -52,6 +66,32 @@ state. Use the 0.76 s/token row for regression comparisons; both are recorded so mistaken for drift. Notes +- 2026-08-06: **burn 0.21's fused `attention` op (CubeCL flash attention) was measured and rejected.** + Swapping the explicit q·kᵀ → scale → mask → f32-softmax → ·v chain for one + `tensor::module::attention(…, is_causal: true)` dispatch is a drop-in — same default `1/sqrt(hd)` + scale, same bottom-right causal alignment, and the kernel keeps the f32 island itself + (`AccumulatorPrecision::Strict(F32)`) — and it passed a new equivalence unit test against the + explicit formulation. It still does not ship, because the A/B (criterion, same session, idle card, + two runs per arm) says it costs more than it buys on this hardware: + + | metric | explicit (shipping) | fused SDPA | delta | + | --- | --- | --- | --- | + | f32 TTFT (36 tok) | 90.9 ms | 91.9 ms | +1.1 % | + | f32 prefill @ 2048 | 593 ms | 629 ms | **+6.0 %** | + | f32 decode | 60.0 ms/token | 61.1 ms/token | +1.8 % | + | f16 TTFT (36 tok) | 20.4 ms | 16.7 ms | **−18 %** | + | f16 prefill @ 2048 | 210 ms | 164 ms | **−22 %** | + | f16 decode | 20.5 ms/token | 22.9 ms/token | **+11 %** | + + The one real win is f16 prefill, where the accelerated flash kernel's plane matmuls have tiles to + fill. Decode is `seq_q = 1` — a matvec with no tile reuse, where flash's machinery is pure overhead + and (the leading hypothesis) an opaque op the Fusion backend cannot absorb into the surrounding + stream the way it absorbs the explicit chain's scale/mask/softmax/cast. f32 loses everywhere for + want of an accelerated path to exploit. Adopting only the winning quadrant would mean a + dtype-conditional fork in the hottest leaf function, on a path no strict parity gate covers (the + gates run f32 and GGUF-dequant-to-f32), so it waits for a deliberate decision with f16 parity + coverage behind it. Re-measure after the burn 0.22 / wgpu 30 bump: wgpu 30 lifts `SHADER_F16` to + WGSL, which changes which kernels are even candidates here. - 2026-07-11: the f32 attention-score island (NaN fix for f16) coincided with an f32 *improvement* (TTFT 100.5 → 88.4 ms, decode 13.3 → 14.1 tok/s) — softmax now always runs in f32 with fusion re-tuning around it; both budget gates re-passed (`budget.rs` 96.8 ms / 10.2 tok/s, `budget_cpu.rs` diff --git a/crates/mummu-bench/benches/runner.rs b/crates/mummu-bench/benches/runner.rs index 75ef992..3a3389f 100644 --- a/crates/mummu-bench/benches/runner.rs +++ b/crates/mummu-bench/benches/runner.rs @@ -25,6 +25,12 @@ use tokenizers::Tokenizer; /// jitter, short enough that the KV cache stays near its steady-state length. const DECODE_STEPS_PER_SAMPLE: usize = 32; +/// Prefill length for the long-context TTFT row. Chosen so the explicit +/// attention path's `[1, heads, t, t]` scores tensor is hundreds of MiB — +/// large enough that avoiding it is a measurable effect, small enough that +/// the 16 GB reference card still holds the f32 model beside it. +const LONG_PREFILL_TOKENS: usize = 2048; + fn qwen2_dir() -> Option { let dir = PathBuf::from(std::env::var_os("MUMMU_QWEN2_DIR")?); dir.is_dir().then_some(dir) @@ -67,6 +73,26 @@ where }); }); + // Long-context prefill: the same work, but at a sequence length where the + // attention formulation actually matters. Explicit attention materializes + // a `[1, heads, t, t]` scores tensor — 201 MiB of f32 at t = 2048 for this + // model's 12 heads, versus 62 KiB at the 36-token prompt above — so this + // is the row that moves when the attention step changes shape. + let long_ids: Vec = ids + .iter() + .cycle() + .take(LONG_PREFILL_TOKENS) + .copied() + .collect(); + assert_eq!(long_ids.len(), LONG_PREFILL_TOKENS); + group.bench_function("ttft_prefill_2048", |b| { + b.iter(|| { + let mut cache = loaded.new_cache(); + let logits = loaded.forward(&long_ids, 0, &mut cache, &device); + black_box(argmax_id(logits).expect("argmax")) + }); + }); + // Decode: per sample, prefill once (untimed), then time N greedy decode // steps through the warm KV cache. Per-token latency = measured / N. group.bench_function("decode_32_tokens", |b| { diff --git a/crates/mummu-bench/tests/budget.rs b/crates/mummu-bench/tests/budget.rs index 44dd996..15ffa71 100644 --- a/crates/mummu-bench/tests/budget.rs +++ b/crates/mummu-bench/tests/budget.rs @@ -25,6 +25,13 @@ const TTFT_BUDGET_MS: f64 = 150.0; const DECODE_BUDGET_TOKENS_PER_S: f64 = 10.0; const DECODE_STEPS: usize = 32; +/// Long-context prefill: the row where the attention formulation shows up at +/// all (explicit attention materializes a `[1, heads, t, t]` scores tensor — +/// 201 MiB of f32 here, against 62 KiB at the ~36-token prompt). Budget is +/// ~1.5x the recorded number, matching the looseness of the rows above. +const LONG_PREFILL_TOKENS: usize = 2048; +const LONG_PREFILL_BUDGET_MS: f64 = 900.0; + fn qwen2_dir() -> Option { let dir = PathBuf::from(std::env::var_os("MUMMU_QWEN2_DIR")?); dir.is_dir().then_some(dir) @@ -64,9 +71,27 @@ fn qwen2_stays_inside_its_perf_budgets() { } let tok_per_s = DECODE_STEPS as f64 / start.elapsed().as_secs_f64(); + // Long-context prefill over the same weights, prompt tiled to length. + let long_ids: Vec = ids + .iter() + .cycle() + .take(LONG_PREFILL_TOKENS) + .copied() + .collect(); + assert_eq!(long_ids.len(), LONG_PREFILL_TOKENS); + let mut long_cache = loaded.new_cache(); + let _ = argmax_id(loaded.forward(&long_ids, 0, &mut long_cache, &device)).expect("warm-up"); + let start = Instant::now(); + let mut long_cache = loaded.new_cache(); + let logits = loaded.forward(&long_ids, 0, &mut long_cache, &device); + let _ = argmax_id(logits).expect("argmax"); + let long_prefill_ms = start.elapsed().as_secs_f64() * 1e3; + eprintln!( "[budget] TTFT {ttft_ms:.1} ms (budget {TTFT_BUDGET_MS} ms), \ - decode {tok_per_s:.1} tok/s (budget {DECODE_BUDGET_TOKENS_PER_S} tok/s)" + decode {tok_per_s:.1} tok/s (budget {DECODE_BUDGET_TOKENS_PER_S} tok/s), \ + prefill@{LONG_PREFILL_TOKENS} {long_prefill_ms:.0} ms \ + (budget {LONG_PREFILL_BUDGET_MS} ms)" ); assert!( ttft_ms <= TTFT_BUDGET_MS, @@ -76,4 +101,8 @@ fn qwen2_stays_inside_its_perf_budgets() { tok_per_s >= DECODE_BUDGET_TOKENS_PER_S, "decode regression: {tok_per_s:.1} tok/s < {DECODE_BUDGET_TOKENS_PER_S} tok/s budget" ); + assert!( + long_prefill_ms <= LONG_PREFILL_BUDGET_MS, + "long-prefill regression: {long_prefill_ms:.0} ms > {LONG_PREFILL_BUDGET_MS} ms budget" + ); } From f75c2ec10ed3567ef1435808dd226ddcad7eac97 Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Thu, 6 Aug 2026 07:18:05 -0500 Subject: [PATCH 3/5] template: render a checkpoint's own chat_template (feature jinja-template) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes the P3 item "General fallback chat renderer via hf-chat-template" - payoff (2) of the 2026-07-23 evaluation, whose payoff (1) (the byte gate) proved the crate reproduces `transformers.apply_chat_template` byte for byte on three families. A model whose family has no hardcoded renderer is now promptable from the authority on its own prompt format: its own template. Both open questions the item left get decided here. (a) The dependency is promoted from dev-only to an *optional* runtime dependency - the same crate at the same 0.2.1 the gate already trusts - behind the non-default feature `jinja-template`. A default build still carries no Jinja engine, so the from-scratch ethos holds for the zoo, and a consumer that must run an un-ported checkpoint opts in. (b) The selection rule ships as a value rather than a convention: `Renderer::for_checkpoint(family: Option, dir)` takes the family renderer when the caller has one and falls back to `ImportedTemplate` otherwise. It deliberately does NOT second-guess a family renderer by reading the template - the gate pins those bytes, and a checkpoint repackaged with a foreign template is caught at LOAD by the tokenizer.rs consistency gate, not silently obeyed at render. The API mirrors ChatMl (`render` / `render_with_tools`) plus `render_with_tools_json`, because the `tools` shape is genuinely open: mainstream templates unpack the transformers `{"type":"function","function":{…}}` wrapper, LFM2.5's wants the signature bare. Everything is bounded and fail-loud - `Absent` (no template declared), `Jinja` (bad syntax, or a template `raise_exception`), `TooLarge` at 8 MiB, `BadTool` - so a runaway template trips the byte bound instead of returning a prompt nothing can tokenize. One model change, additive: `chat::Turn` gains `tool_calls: Vec`, populated by `assistant_tool_calls{,_lfm}` BESIDE the rendered content. The family renderers never read it, so prompt bytes are unchanged by construction and by the gate; the imported path passes the calls as data so an arbitrary template writes its OWN call markers rather than inheriting Hermes' `` wrapping, which would double-wrap under any other convention. Proof: 7 unit tests over a toy Jinja template (roles, eos_token from the config, tool wire shape, structural calls, absent/broken template, and a render bomb hitting the byte bound) - no fixture needed - plus tests/imported_render.rs on real checkpoints. Byte-identical to ChatMl::qwen3() on plain 142 B, tools 748 B and FC history 324 B, and to ChatMl::lfm2() on plain 157 B and tools 379 B; the LFM leg also exercises the standalone chat_template.jinja fallback and the bos_token injection the gate previously had to hand-inject. Verified green: fmt, clippy --all-targets in BOTH feature configurations, 209 lib unit tests with the feature / 202 without, all 10 template-gate legs re-passed byte-identically, and Qwen3-0.6B still greedy-emits a parseable `` on the 4070 Ti SUPER. Co-Authored-By: Claude Opus 4.8 --- Cargo.toml | 7 + README.md | 14 + ROADMAP.md | 26 +- crates/mummu/Cargo.toml | 15 +- crates/mummu/src/chat.rs | 19 +- crates/mummu/src/lib.rs | 3 + crates/mummu/src/template.rs | 413 ++++++++++++++++++++++++++ crates/mummu/tests/imported_render.rs | 179 +++++++++++ 8 files changed, 672 insertions(+), 4 deletions(-) create mode 100644 crates/mummu/src/template.rs create mode 100644 crates/mummu/tests/imported_render.rs diff --git a/Cargo.toml b/Cargo.toml index 86052cc..c31d156 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -83,6 +83,13 @@ once_cell = "1" criterion = "0.8" +# Renders an HF `chat_template` Jinja byte-identically to +# `transformers.apply_chat_template` (minijinja + a transformers compat +# layer). Two roles, one version: the test-only reference behind the template +# BYTE gate, and — behind mummu's non-default `jinja-template` feature — the +# fallback renderer for checkpoints with no hardcoded family renderer. +hf-chat-template = "0.2.1" + [profile.release] # Thin LTO: inference is GPU-kernel-bound, so fat LTO's extra minutes per link # buy almost nothing here while tripling the verify-loop cost. Revisit against diff --git a/README.md b/README.md index 4801b9d..2d61194 100644 --- a/README.md +++ b/README.md @@ -51,6 +51,20 @@ It exists because two local-first apps — **[laurelane](https://github.com/phys neither → no behavior change). On a successful safetensors load the parsed `TokenizerConfig` is surfaced on the returned `Loaded{Qwen2,Qwen3,Lfm2}` struct (`tokenizer_config`), so a consumer reads config-driven EOS/BOS/PAD ids straight off the model; a GGUF load surfaces `None` (self-contained). +- **Fallback chat renderer for un-ported models** (optional feature `jinja-template`) — + `mummu::template::ImportedTemplate` renders a checkpoint's **own** imported `chat_template` Jinja, so a + model whose family has no hardcoded renderer is still promptable from the authority on its prompt + format: its own template. The selection rule ships as a value — `Renderer::for_checkpoint(family, dir)` + takes a byte-verified family renderer when one exists and falls back to the template otherwise, never + second-guessing the family renderer. Bounded and fail-loud (no template → `Absent`, bad Jinja → + `Jinja`, a runaway render → `TooLarge` at 8 MiB), with the config's BOS/EOS/PAD/UNK reaching the render + context and assistant tool calls passed **structurally** so the template writes its own call markers + instead of inheriting Hermes'. Verified against the from-scratch path on real checkpoints + (`tests/imported_render.rs`): byte-identical to `ChatMl::qwen3()` on plain (142 B), tools (748 B) and + full FC history (324 B), and to `ChatMl::lfm2()` on plain (157 B) and tools (379 B) — the LFM leg also + proving the standalone `chat_template.jinja` fallback and the `bos_token` injection. The feature is + **off by default**: the zoo's from-scratch renderers cover it byte-for-byte, and a default build carries + no Jinja engine. - **Import validation** — a two-stage error taxonomy: `ImportError` for the file→module stage (missing file, parse, load, and an `Incomplete` per-tensor missing/errored diff) and `SanityError` for the runtime liveness a checked load can't see — NaN/Inf logits, a vocab-width mismatch, or a diff --git a/ROADMAP.md b/ROADMAP.md index 10aea4c..3cc9d28 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -830,11 +830,35 @@ The subsystem that turns "a model on HuggingFace or on disk" into a loaded, pari the local cache). Known family divergences are PINNED to their exact deltas so any other drift still fails: Qwen2.5's no-system branding preamble ("You are Qwen, …") vs our neutral one (with tools AND the plain injected default turn), Qwen3's no-system no-preamble, Qwen3's history think-stripping. -- [ ] **General fallback chat renderer via `hf-chat-template`** — payoff (2) of the evaluation above: for a +- [x] **General fallback chat renderer via `hf-chat-template`** — payoff (2) of the evaluation above: for a checkpoint whose family has no hardcoded `chat` renderer, render prompts from its own imported `chat_template` (the byte gate proved fidelity on Qwen3). Weigh promoting the dep from dev to optional runtime feature vs the from-scratch ethos; needs the P8/consumer-facing API decision of when to prefer the imported template over a family renderer. *(2026-07-23, split from the evaluation.)* + *(2026-08-06) **Shipped as `mummu::template`, behind the non-default feature `jinja-template`.** Both + open questions decided: **(a) the dep** is promoted from dev-only to an *optional* runtime dependency + — same crate, same 0.2.1 the byte gate already trusts as the transformers-equivalent reference — so a + default build still carries no Jinja engine and the from-scratch ethos holds for the zoo, while a + consumer that must run an un-ported checkpoint opts in. **(b) the selection rule** is a value, not a + convention: `Renderer::for_checkpoint(family: Option, dir)` takes the family renderer when the + caller has one and falls back to `ImportedTemplate` otherwise, and it deliberately does **not** + second-guess a family renderer by reading the template (the gate pins those bytes; a checkpoint + repackaged with a foreign template is caught at *load* by the `tokenizer.rs` consistency gate, not + silently obeyed at render). API mirrors `ChatMl`: `render` / `render_with_tools`, plus + `render_with_tools_json` because the `tools` shape is genuinely open — the mainstream templates unpack + the `transformers` `{"type":"function","function":{…}}` wrapper, LFM2.5's wants the signature bare. + Bounded + fail-loud throughout (`Absent` / `Jinja` / `TooLarge` at 8 MiB / `BadTool`; a runaway + template trips the byte bound rather than returning an untokenizable prompt — unit-tested with a + render bomb). The one model change: `chat::Turn` gained an additive `tool_calls: Vec` field + that `assistant_tool_calls{,_lfm}` now populate **beside** the rendered content — the family renderers + never read it (prompt bytes unchanged by construction *and* by the gate), but the imported path passes + calls as data so the template writes its own markers instead of inheriting Hermes' `` + wrapping. Proof: 7 unit tests over a toy Jinja template (no fixture needed) + `tests/imported_render.rs` + on real checkpoints — byte-identical to `ChatMl::qwen3()` on plain 142 B / tools 748 B / FC history + 324 B and to `ChatMl::lfm2()` on plain 157 B / tools 379 B (that leg also exercising the standalone + `chat_template.jinja` fallback and `bos_token` injection, which the gate had to hand-inject); all 10 + template-gate legs re-passed unchanged, 209 unit tests with the feature (202 without), clippy clean in + both configurations, and Qwen3-0.6B still greedy-emits a parseable `` on the 4070 Ti SUPER.* - [x] **`ChatMl::qwen3()` with history think-stripping** — the byte gate documented that Qwen3's template strips `…` reasoning from assistant turns at/before the last user query while our shared `ChatMl::qwen2()` renderer re-renders history verbatim (fine for fresh prompts + tool loops, diff --git a/crates/mummu/Cargo.toml b/crates/mummu/Cargo.toml index 3d14925..ad53fe2 100644 --- a/crates/mummu/Cargo.toml +++ b/crates/mummu/Cargo.toml @@ -6,8 +6,19 @@ edition.workspace = true license.workspace = true repository.workspace = true +[features] +default = [] +# Render a checkpoint's OWN imported `chat_template` Jinja (crate::template), +# for models whose family has no hardcoded `chat` renderer. Off by default: +# the from-scratch renderers cover the zoo byte-for-byte and a default build +# should not carry a Jinja engine to wrap a prompt. Turning it on pulls +# hf-chat-template (minijinja) — the same crate, same version, that the +# template BYTE gate already trusts as the transformers-equivalent reference. +jinja-template = ["dep:hf-chat-template"] + [dependencies] burn = { workspace = true } +hf-chat-template = { workspace = true, optional = true } wgpu = { workspace = true } pollster = { workspace = true } half = { workspace = true } @@ -28,7 +39,9 @@ burn-store = { workspace = true } # wrapping stays the hardcoded, byte-verified `chat` renderers. # Default features stay ON: they carry minijinja's Python-compat layer # (str.startswith etc.), which real HF templates (Qwen3's included) call. -hf-chat-template = "0.2.1" +# Also the dependency behind the optional `jinja-template` feature above; +# listing it here too keeps the gate running in a plain `cargo test`. +hf-chat-template = { workspace = true } [target.'cfg(windows)'.dependencies] # One syscall (GlobalMemoryStatusEx) for the device inventory's total-RAM diff --git a/crates/mummu/src/chat.rs b/crates/mummu/src/chat.rs index ae79c1d..4fdff6c 100644 --- a/crates/mummu/src/chat.rs +++ b/crates/mummu/src/chat.rs @@ -75,6 +75,15 @@ impl Role { 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 { @@ -83,6 +92,7 @@ impl Turn { Self { role: Role::System, content: content.into(), + tool_calls: Vec::new(), } } @@ -91,6 +101,7 @@ impl Turn { Self { role: Role::User, content: content.into(), + tool_calls: Vec::new(), } } @@ -99,6 +110,7 @@ impl Turn { Self { role: Role::Assistant, content: content.into(), + tool_calls: Vec::new(), } } @@ -124,6 +136,7 @@ impl Turn { Self { role: Role::Assistant, content: blocks.join("\n"), + tool_calls: calls.to_vec(), } } @@ -145,6 +158,7 @@ impl Turn { "<|tool_call_start|>{}<|tool_call_end|>", pythonic_calls(calls) ), + tool_calls: calls.to_vec(), } } @@ -156,6 +170,7 @@ impl Turn { Self { role: Role::Tool, content: content.into(), + tool_calls: Vec::new(), } } } @@ -337,7 +352,7 @@ pub struct ToolCall { pub const MAX_TOOL_CALLS: usize = 64; /// Most tools one render will advertise. -const MAX_TOOLS: usize = 128; +pub(crate) const MAX_TOOLS: usize = 128; /// What went wrong extracting tool calls from a model response. #[derive(Debug, thiserror::Error)] @@ -690,7 +705,7 @@ impl<'a> PythonicParser<'a> { /// Longest conversation a single render will wrap — a generous bound that /// still catches an unbounded history being passed by mistake. -const MAX_TURNS: usize = 1024; +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 diff --git a/crates/mummu/src/lib.rs b/crates/mummu/src/lib.rs index 364f76e..7983ede 100644 --- a/crates/mummu/src/lib.rs +++ b/crates/mummu/src/lib.rs @@ -20,5 +20,8 @@ pub mod manage; pub mod models; pub mod nn; pub mod registry; +/// Render a checkpoint's own imported chat template (feature `jinja-template`). +#[cfg(feature = "jinja-template")] +pub mod template; pub mod tok_config; pub mod tokenizer; diff --git a/crates/mummu/src/template.rs b/crates/mummu/src/template.rs new file mode 100644 index 0000000..9e04d25 --- /dev/null +++ b/crates/mummu/src/template.rs @@ -0,0 +1,413 @@ +//! 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/tests/imported_render.rs b/crates/mummu/tests/imported_render.rs new file mode 100644 index 0000000..2503e15 --- /dev/null +++ b/crates/mummu/tests/imported_render.rs @@ -0,0 +1,179 @@ +//! REAL-CHECKPOINT proof for the general fallback renderer (`crate::template`, +//! feature `jinja-template`): rendering a checkpoint's OWN imported chat +//! template through Mummu's public API must byte-match the from-scratch family +//! renderer that the template byte gate already pins. +//! +//! `tests/template_gate.rs` proves the same equality by driving +//! `hf-chat-template` directly from the test; this proves the *module* wires +//! it correctly — special tokens into the context, tools in the Hermes wire +//! shape, assistant tool calls passed structurally rather than as pre-wrapped +//! markers, and the generation prompt on. A checkpoint outside the zoo has no +//! family renderer to compare against, so this is the only place the fallback +//! path's fidelity can be measured at all. Ignored by default; run with: +//! +//! ```text +//! MUMMU_QWEN3_DIR=path/to/qwen3-0.6b \ +//! cargo test -p mummu --features jinja-template --test imported_render -- --ignored --nocapture +//! ``` +#![cfg(feature = "jinja-template")] + +use std::path::PathBuf; + +use mummu::chat::{ChatMl, ToolCall, ToolSpec, Turn}; +use mummu::template::{ImportedTemplate, Renderer}; + +fn qwen3_dir() -> Option { + let dir = PathBuf::from(std::env::var_os("MUMMU_QWEN3_DIR")?); + dir.is_dir().then_some(dir) +} + +fn weather_spec() -> ToolSpec { + ToolSpec { + name: "get_weather".into(), + description: "Get the current weather in a city.".into(), + parameters: serde_json::json!({ + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"} + }, + "required": ["city"] + }), + } +} + +/// Report where two renders diverge, for a readable failure. +fn diff(label: &str, ours: &str, reference: &str) -> String { + match ours + .bytes() + .zip(reference.bytes()) + .position(|(a, b)| a != b) + .or_else(|| (ours.len() != reference.len()).then(|| ours.len().min(reference.len()))) + { + None => format!("{label}: byte-identical ({} B)", ours.len()), + Some(at) => { + let lo = at.saturating_sub(60); + format!( + "{label}: DIVERGES at byte {at}\n imported…{:?}\n family …{:?}", + &ours[lo..(at + 60).min(ours.len())], + &reference[lo..(at + 60).min(reference.len())], + ) + } + } +} + +/// Plain conversation, tools, and a full function-calling history: the +/// imported template rendered through `ImportedTemplate` must equal +/// `ChatMl::qwen3()` byte for byte on all three. +#[test] +#[ignore = "needs the local Qwen3 checkpoint dir (MUMMU_QWEN3_DIR)"] +fn imported_qwen3_template_byte_matches_the_family_renderer() { + let dir = qwen3_dir().expect("set MUMMU_QWEN3_DIR to a Qwen3 checkpoint dir"); + let imported = ImportedTemplate::from_dir(&dir).expect("the checkpoint's template compiles"); + let family = ChatMl::qwen3(); + + let plain = [ + Turn::system("You are a helpful assistant."), + Turn::user("List the first five prime numbers."), + ]; + let ours = imported.render(&plain).expect("renders"); + let reference = family.render(&plain); + println!("{}", diff("plain", &ours, &reference)); + assert_eq!(ours, reference, "plain render must byte-match"); + + let tools = [weather_spec()]; + let turns = [ + Turn::system("You are a helpful assistant."), + Turn::user("What's the weather in Paris?"), + ]; + let ours = imported.render_with_tools(&tools, &turns).expect("renders"); + let reference = family.render_with_tools(&tools, &turns); + println!("{}", diff("tools", &ours, &reference)); + assert_eq!(ours, reference, "tools render must byte-match"); + + let calls = [ToolCall { + name: "get_weather".into(), + arguments: serde_json::json!({"city": "Paris"}), + }]; + let history = [ + Turn::system("You are a helpful assistant."), + Turn::user("What's the weather in Paris?"), + Turn::assistant_tool_calls(&calls), + Turn::tool_response("{\"temp_c\": 21}"), + ]; + let ours = imported.render(&history).expect("renders"); + let reference = family.render(&history); + println!("{}", diff("fc-history", &ours, &reference)); + assert_eq!( + ours, reference, + "FC history must byte-match — the structural tool_calls path" + ); +} + +/// LFM2.5 is the other half of the wiring proof: its checkpoint ships NO +/// `chat_template` JSON key (only a standalone `chat_template.jinja`), and its +/// template opens with `{{- bos_token -}}` — so this leg exercises both the +/// file fallback in `from_dir` AND the special tokens reaching the render +/// context from the parsed config. Get either wrong and the prompt loses its +/// BOS silently. +#[test] +#[ignore = "needs the local LFM2.5 checkpoint dir (MUMMU_LFM2_DIR) with chat_template.jinja"] +fn imported_lfm2_template_byte_matches_the_family_renderer() { + let Some(dir) = std::env::var_os("MUMMU_LFM2_DIR") + .map(PathBuf::from) + .filter(|d| d.is_dir()) + else { + panic!("set MUMMU_LFM2_DIR to an LFM2.5 checkpoint dir"); + }; + let imported = ImportedTemplate::from_dir(&dir).expect("standalone chat_template.jinja loads"); + let family = ChatMl::lfm2(); + + let turns = [ + Turn::system("You are a helpful assistant."), + Turn::user("List the first five prime numbers."), + ]; + let ours = imported.render(&turns).expect("renders"); + let reference = family.render(&turns); + println!("{}", diff("lfm plain", &ours, &reference)); + assert!( + ours.starts_with("<|startoftext|>"), + "the config's bos_token reached the template: {:?}", + &ours[..ours.len().min(40)] + ); + assert_eq!(ours, reference, "LFM plain render must byte-match"); + + // LFM's template wants BARE tool JSON, not the transformers wrapper — + // the shape `render_with_tools_json` exists for. + let bare = serde_json::json!({ + "name": "get_weather", + "description": "Get the current weather in a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"} + }, + "required": ["city"] + } + }); + let ours = imported + .render_with_tools_json(&[bare], &turns) + .expect("renders"); + let reference = family.render_with_tools(&[weather_spec()], &turns); + println!("{}", diff("lfm tools", &ours, &reference)); + assert_eq!(ours, reference, "LFM tools render must byte-match"); +} + +/// The selection rule as a value: `Renderer` with no family renderer falls +/// back to the checkpoint's own template and renders the same bytes. +#[test] +#[ignore = "needs the local Qwen3 checkpoint dir (MUMMU_QWEN3_DIR)"] +fn renderer_falls_back_to_the_checkpoint_template_when_no_family_renderer() { + let dir = qwen3_dir().expect("set MUMMU_QWEN3_DIR to a Qwen3 checkpoint dir"); + let fallback = Renderer::for_checkpoint(None, &dir).expect("falls back to the template"); + assert!(matches!(fallback, Renderer::Imported(_))); + + let turns = [Turn::user("Say hi.")]; + let ours = fallback.render(&turns).expect("renders"); + let reference = ChatMl::qwen3().render(&turns); + println!("{}", diff("fallback", &ours, &reference)); + assert_eq!(ours, reference); +} From 28c5d7ea1196a06b5412ec9f36ad768a2f07f6cc Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Thu, 6 Aug 2026 07:20:50 -0500 Subject: [PATCH 4/5] docs: fold 2026-08-06 research - graph capture is the dispatch lever Research-and-fold pass, four surgical ROADMAP additions, each with sources: P0 / burn 0.22 - still pre-release (0.22.0-pre.1 newest, 0.21.0 latest stable), so the migration stays gated, but reading the notes properly reframes it from a cost into the most valuable pending item: 0.22 ships **graph capture, explicitly to cut CPU-side launch overhead**. Also arriving with it: LoRA/QLoRA in-framework (P10 becomes wiring), the remote backend gaining multi-device + client-side operation-graph caching + async reads (P6), and dequant->op->quant fallbacks for slice/gather/select/expand plus BitNet b1.58 calibration (P9). Perf / dispatch-bound decode - the next lever is now named. A kernel launch costs ~5-10 us of CPU time, a forward dispatches hundreds, and on batch-1 decode that sequencing is 20-40% of inference time; capture the decode step's launch sequence once and replay it. That is a far better bet than shaving kernels one at a time - this run's flash-attention A/B is the evidence per-kernel substitution does not move the number. Also noted: the dispatch-bound premise is now f32-ONLY, since f16 decodes 2.9x faster after the re-measure. P5 / speculative decoding - calibration that changes the item's expected value. On consumer GPUs a small-draft speculation is frequently a net LOSS: 7B target + 0.5B draft measured 0.27x on an RTX 5060 Ti (flipping to 1.4x only at a 14B target), and a public 19-configuration llama.cpp study on Qwen3.6-35B-A3B found no variant achieving net speedup on an RTX 3090 while vLLM got +27.5% on the same hardware. So: build the MTP-head route, gate the draft-model route on a per-(target, draft, hardware) measurement, and expect the batched verify to be where winners differ from losers. P9 / keep-quantized kernels - LlamaWeb's split is corroborated from CUDA (Marlin-class kernels unpack int4 to f16 in the register file, no intermediate), so it is the general answer, not a WebGPU workaround; its *interface* (a scheme = an unpacker + a dequant routine, all downstream kernels format-agnostic) is the shape to copy for our GGUF reader. Caveat sharpened by this run: measure the win against f16, not f32, or it will look better than it is. Co-Authored-By: Claude Opus 4.8 --- ROADMAP.md | 53 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/ROADMAP.md b/ROADMAP.md index 3cc9d28..0396820 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -89,6 +89,18 @@ a benchmark holds/improves its budget; README perf claims link an artifact. on an otherwise-idle GPU (5% util). Decode throughput tracking CPU availability is what a dispatch-bound path looks like. Operationally: **run the budget gates on a quiet machine** or they report contention as a regression.* + *(2026-08-06)* Two updates. **(1) The premise is now f32-only.** This run's re-measure has f16 + decoding at 20.5 ms/token against f32's 60.0 — f16 is no longer "same speed, half the VRAM", it is + 2.9× faster, so the dispatch-bound reading applies to the f32 path and the f16 path has evidently + found an accelerated one. Any further work here should target f32 or, better, ask whether f32 is + still the right default at all. **(2) The next lever is named: graph capture.** The industry + framing matches our numbers exactly — a kernel launch costs ~5–10 µs of CPU time, an LLM forward + dispatches hundreds of them, and on batch-1 decode that CPU-side sequencing is 20–40 % of total + inference time; the standard fix is to capture the decode step's whole launch sequence once and + replay it, which keeps per-token overhead flat instead of paying it per kernel. burn **0.22 ships + graph capture** for this purpose (P0 item), which is a far better bet than shaving kernels one at a + time — this run's flash-attention A/B is the evidence that per-kernel substitution does not move + this number. — https://gigagpu.com/cuda-graph-optimization-inference/ - [x] Evaluate Burn 0.21's `burn.toml` project config — per-subsystem tuning + a CubeCL kernel-validation layer without recompiling; useful as a debug switch for kernel-level parity hunts — https://burn.dev/blog/release-0.21.0/ *(2026-07-17 research)* Concretely, a `burn.toml` dropped at @@ -191,6 +203,18 @@ a benchmark holds/improves its budget; README perf claims link an artifact. Do NOT adopt a pre-release; when 0.22.0 stabilizes: migrate on a branch, re-run every parity gate + budget, and expect the backend aliases + dtype helpers + all loaders' `target_float` derivation to change shape. *(2026-07-30 research)* — https://github.com/Tracel-AI/burn/releases + *(2026-08-06 research)* Still pre-release (0.22.0-pre.1 is the newest tag; 0.21.0 the latest stable), + so this stays gated — but reading the pre-release notes properly turns it from a migration *cost* + into the run's most valuable pending item, because **0.22 ships graph capture, explicitly to cut + CPU-side launch overhead**. That is the exact bottleneck the dispatch-bound decode item has been + chasing since 2026-07-11 (an f16-matches-f32 measurement, then SPIR-V, then this run's flash-attention + A/B all pointing at per-dispatch cost rather than bandwidth), and it is the one lever that attacks it + *generically* rather than one kernel at a time. Three more items arrive with it: **LoRA/QLoRA** land + in-framework (P10 becomes wiring, not implementation), the **remote backend** gains multi-device + + client-side operation-graph caching + async reads (P6 multi-GPU), and quantization gains + dequant→op→quant fallbacks for slice/gather/select/expand plus BitNet b1.58 calibration (P9). Plan the + migration around measuring graph capture on the decode loop first — if it lands the dispatch win, it + reorders everything below it in the perf section. - [x] Silence the pre-existing `LNK4098` (LIBCMT defaultlib conflict) the 2026-07 nightly toolchain's new `linker_messages` lint now surfaces when linking the `mummu` lib-test binary — find which native dep object embeds the static-CRT directive (tokenizers' C++ deps are the suspects) and @@ -979,6 +1003,23 @@ The subsystem that turns "a model on HuggingFace or on disk" into a loaded, pari greedy; (b) a small same-tokenizer zoo model as drafter (0.5B drafts for 9B). Needs: batched-verify forward through the existing KV cache (rollback on reject), driver support in `generate_loop`. *(2026-07-30 research)* — https://github.com/JustVugg/colibri + *(2026-08-06 research)* **Calibration that changes this item's expected value: on consumer GPUs a + small-draft-model speculation is frequently a net LOSS, not a 1.5–2× win.** The 2026 measurements to + plan against: a 7B target + 0.5B draft on an RTX 5060 Ti ran at **0.27×** (i.e. ~4× slower) and only + flipped to 1.4× when the target grew to 14B — the draft's cost only amortizes when the target is + expensive enough per token; and a public 19-configuration llama.cpp study on Qwen3.6-35B-A3B with a + vocab-matched Qwen3.5-0.8B drafter found **no variant achieving a net speedup** on a single RTX 3090 + (ngram-cache, ngram-mod and classic draft all lost), while vLLM on the same hardware got +27.5 % — + i.e. the engine's verify-batching quality, not the idea, decides it. Consequences for Mummu: (a) + route (a) **native MTP heads** is the one to build — no second model to pay for; (b) route (b) small- + model drafting must be gated on an end-to-end measurement per (target, draft, hardware) triple, never + shipped on the literature's headline; (c) the batched-verify forward has to be genuinely batched + through the KV cache, since that is where the engines that win differ from the ones that lose. Our + own dispatch-bound decode makes (c) harder and the win smaller: a k-token verify is one forward + either way, so speculation trades dispatches for compute — which is the right direction here, and + worth re-checking after graph capture (P0) moves the dispatch baseline. — + https://github.com/thc1006/qwen3.6-speculative-decoding-rtx3090 · + https://inventivehq.com/blog/llama-cpp-speculative-decoding-consumer-gpu - [ ] **Grammar-constrained decoding** *(colibri parity)* — colibri forces structured output via `.gbnf` grammars (llama.cpp's GBNF convention) and even uses grammar-forced *drafts* to speed structured generation. Mummu's tool-calling currently *trusts* the model to emit parseable `` JSON @@ -1207,6 +1248,18 @@ The VRAM lever the P6 planner pulls to make the largest useful model fit the use here may be VRAM-only for us until the dispatch gap closes; (b) their gains lean on subgroup matrix ops where available, with portable fallbacks. — https://arxiv.org/html/2605.20706v1 *(2026-07-16 research)* + *(2026-08-06 research)* Corroborated from the CUDA side, which matters because it means the design is + the *general* answer and not a WebGPU workaround: production int4 kernels (Marlin and descendants) + fuse dequantization into the matmul so the int4 values unpack to f16 **in the register file**, never + materializing a full-size intermediate — the same "straight into registers" split LlamaWeb measured + for the decode matvec. Also worth copying from LlamaWeb is its *interface*: a new quantization scheme + is an **unpacker + a dequant routine**, and every downstream kernel (matmul, attention) stays + format-agnostic — the shape that let one representation carry 21 formats, and the natural fit for our + GGUF reader, which already parses every K-quant block layout structurally. One caveat sharpens with + this run's numbers: LlamaWeb's headline decode win came from being bandwidth-bound, and our f32 path + is not — but our **f16** path now decodes 2.9× faster than f32, so measure the keep-quantized win + against f16, not f32, or it will look better than it is. — + https://www.tensortonic.com/llm-internals/quantization - [ ] Evaluate **CubeCL's quantization primitives** for the keep-quantized matmul: recent CubeCL ships block-scaled MMA, global quantization for matmul, quantized tensor views, and FP4/FP2 formats — the kernel substrate a Q4-weights × f16-activations decode path would ride (vs hand-writing a From aeaee7a31a050a6693e5afb7a82d135ec6136741 Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Thu, 6 Aug 2026 07:40:47 -0500 Subject: [PATCH 5/5] bench: the old f16 rows were f32 - bisected, retracted, gated Closes the f32-decode-drift item opened earlier this run, and the answer is not the one the item expected. Two old trees were checked out and benched on the same idle machine the same day as HEAD: e44debf (2026-07-12, recorded 54.3 f32 / 54.5 f16): f32 60.1 ms/token, f16 60.2 ms/token c1826e7 (just before the 2026-07-30 dtype pinning): f32 60.0 ms/token, f16 PANICS with DTypeMismatch HEAD (2026-08-06): f32 60.0 ms/token, f16 20.5 ms/token (1) f32 never regressed. The same 2026-07-12 code reads 60.1 ms/token today, so the 54.3 recorded then was machine state - driver, OS, background load - and not a Mummu change. f32 decode has been ~60 ms/token on this card the whole time. (2) The f16 rows recorded on 2026-07-11 and 2026-07-12 were f32 runs wearing an f16 label. The criterion bench builds `Gpu` and then `GpuF16` in ONE process; before the dtype pinning the f32 leg locked the per-device default dtype policy and the "f16" model ran in f32 - which is exactly why those rows matched f32 to a tenth of a millisecond (70.9 vs 70.7, then 54.5 vs 54.3) and were believed. Once the loaders took `target_float` from the TYPE (2026-07-23) the mismatch turned loud rather than silent, which is the DTypeMismatch at c1826e7; the 2026-07-30 pinning of every runtime creation site fixed it. Real f16 decode is 2.9x faster than f32 and always was - the harness could not see it. So the standing "f16 buys VRAM, not speed" reading is withdrawn, and with it the f16 half of "decode is dispatch-bound" and of SPIR-V's "+30% on BOTH dtypes" (only their f32 halves were ever measured). The f32 dispatch-bound reading stands on its own evidence - the ~30% swing from host CPU load on an idle GPU. The f16 VRAM figures are unaffected: they come from real_f16.rs, one alias per process, always genuinely f16. Guard so this cannot recur silently: `mummu-bench/tests/budget_f16.rs`, an f16 budget gate in its OWN test binary, which asserts `logits.dtype() == F16` before believing a single number. A gate that cannot tell which precision it measured is not a gate. Its budgets are set against what IT measures (21 ms TTFT / 16.9 tok/s, budgets 60 ms / 12 tok/s), not against criterion's steady-state row - and that gap is itself a new finding: a cold f16 session's first 32 tokens run at roughly f32 speed because a single burst never leaves autotune. Both numbers are recorded together, and closing the gap (a persisted CubeCL autotune cache) is a new ROADMAP item. Verified green: fmt, clippy --all-targets, 202 lib unit tests, and all three GPU/CPU budget gates - f32 99.1 ms / 12.7 tok/s / 592 ms prefill@2048, f16 21 ms / 16.9 tok/s, CPU 16.6 tok/s. Co-Authored-By: Claude Opus 4.8 --- ROADMAP.md | 67 +++++++++++---- bench/BASELINE.md | 54 +++++++++--- crates/mummu-bench/tests/budget_f16.rs | 111 +++++++++++++++++++++++++ 3 files changed, 202 insertions(+), 30 deletions(-) create mode 100644 crates/mummu-bench/tests/budget_f16.rs diff --git a/ROADMAP.md b/ROADMAP.md index 0396820..2f860c5 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -89,11 +89,15 @@ a benchmark holds/improves its budget; README perf claims link an artifact. on an otherwise-idle GPU (5% util). Decode throughput tracking CPU availability is what a dispatch-bound path looks like. Operationally: **run the budget gates on a quiet machine** or they report contention as a regression.* - *(2026-08-06)* Two updates. **(1) The premise is now f32-only.** This run's re-measure has f16 - decoding at 20.5 ms/token against f32's 60.0 — f16 is no longer "same speed, half the VRAM", it is - 2.9× faster, so the dispatch-bound reading applies to the f32 path and the f16 path has evidently - found an accelerated one. Any further work here should target f32 or, better, ask whether f32 is - still the right default at all. **(2) The next lever is named: graph capture.** The industry + *(2026-08-06)* Two updates. **(1) The 2026-07-11 f16 leg of this argument is RETRACTED — it was + never an f16 measurement.** "f16 (half the weight traffic) decodes at exactly f32's speed" came + from a bench that builds `Gpu` then `GpuF16` in one process; the f32 leg locked the per-device + default dtype policy, so the f16 model ran in f32 and of course matched to a tenth of a + millisecond. Bisected this run (see the closed item above): real f16 decodes at **20.5 ms/token + against f32's 60.0**. So the dispatch-bound reading holds for the **f32** path — which the + independent evidence still supports (the ~30 % swing from *host CPU* load on an idle GPU; SPIR-V's + +30 %) — and the open question becomes whether f32 is the right default at all rather than why f16 + is not faster. **(2) The next lever is named: graph capture.** The industry framing matches our numbers exactly — a kernel launch costs ~5–10 µs of CPU time, an LLM forward dispatches hundreds of them, and on batch-1 decode that CPU-side sequencing is 20–40 % of total inference time; the standard fix is to capture the decode step's whole launch sequence once and @@ -175,18 +179,41 @@ a benchmark holds/improves its budget; README perf claims link an artifact. for one token) and the honest shape is llama.cpp at f16 on the `llama_ref` harness; (b) accept a dtype- **and** length-conditional branch in `GqaAttention::forward`, or find a formulation that isn't conditional. Re-measure first: the numbers are burn-0.21/wgpu-29-specific. -- [ ] **Bisect the f32 decode drift: 54.3 → 60.0 ms/token since 2026-07-12** — the 2026-08-06 - re-measure found the f32 decode row 10 % slower than recorded while f16 got **2.7× faster** - (54.5 → 20.5 ms/token) over the same period. Neither move was caused by a Mummu change that - claimed them, and the budget gate never noticed because its ceiling (10 tok/s) sits four - ms/token below the recorded number — a ceiling that loose cannot catch drift, only collapse. - Suspects, cheapest first: the 2026-07-30 dtype pinning (explicit `(device, dtype)` at every - creation site may have moved f32 off a fast path while putting f16 on one), a dependency bump in - the CubeCL/wgpu stack, an autotune cache re-tuned against a different machine state, or the - GPU driver. Route: check out the 2026-07-12 tree and re-bench it on today's machine — if it - also reads 60 ms/token the cause is under the repo, if it reads 54 the cause is in it, and - `git bisect` over the bench closes it either way. Worth doing before any further f32 perf work - builds on a number that moved. +- [x] **Bisect the f32 decode drift: 54.3 → 60.0 ms/token since 2026-07-12** — the 2026-08-06 + re-measure found the f32 decode row 10 % slower than recorded while f16 read **2.7× faster**. + *(2026-08-06, closed the same run — and it turned up something bigger than the drift.)* The + 2026-07-12 tree (`e44debf`) and the pre-dtype-pinning tree (`c1826e7`) were checked out and + benched on the same idle machine: f32 reads **60.1 / 60.0 / 60.0 ms/token** across all three + trees, so **f32 never regressed** — the 54.3 recorded on 2026-07-12 is not reproducible from + that same code today, i.e. it was machine state (driver/OS/background), never a Mummu change. + The f16 finding is the real one: `c1826e7` **panics `DTypeMismatch`** on the f16 bench leg, and + `e44debf` reports f16 at 60.2 ms/token — a tenth of a millisecond from its own f32 row. That is + the one-alias-per-process hazard (root-caused 2026-07-23, fixed 2026-07-30): this bench builds + `Gpu` then `GpuF16` in ONE process, the f32 leg locked the per-device default dtype policy, and + **the "f16" rows recorded on 2026-07-11 and 2026-07-12 were f32 runs wearing an f16 label** — + which is exactly why they matched f32 so implausibly closely (70.9 vs 70.7, then 54.5 vs 54.3). + Real f16 decode is **20.5 ms/token, 2.9× faster than f32**, and always was; the harness could + not see it. Consequences folded above: the dispatch-bound item's f16 leg is retracted, and + `bench/BASELINE.md` carries the three-tree table. The f16 *VRAM* figures are unaffected — they + come from `real_f16.rs`, one alias per process, always genuinely f16. Standing lesson for the + perf suite: **a measurement that agrees with its control to a tenth of a percent is evidence of + a wiring bug, not of a null result.** + *(2026-08-06, same run)* Guard added so this class of bug cannot recur silently: + **`mummu-bench/tests/budget_f16.rs`**, an f16 budget gate in its OWN test binary (one dtype alias + per process) that asserts `logits.dtype() == F16` before it believes a single number. Recorded + 21 ms TTFT / 16.9 tok/s, budgets 60 ms / 12 tok/s. +- [ ] **Close the f16 autotune warm-up gap: 16.9 tok/s cold vs 48.8 steady** — building the f16 gate + surfaced it. The gate prefills twice, then times 32 decode steps once, and gets 16.9 tok/s + (~59 ms/token); criterion, which runs the same 32 steps across many samples, gets 48.8 (20.5 + ms/token). So an f16 session's **first ~32 tokens run at roughly f32 speed** and only then does the + 2.9× appear — the f16 path has many more autotune variants to try than f32 (whose same-harness gap + is only 13.2 vs 16.7 tok/s). That is a real user-facing cost for a runner whose consumers open + short-lived agent turns, and it is not a perf mystery so much as a caching question: CubeCL's + autotune cache is configurable (`[cubecl.autotune] cache = "local"|"target"|"global"|{file=…}`, + already in the repo-root `burn.toml`'s vocabulary) and a persisted, per-machine cache should let a + cold process start warm. Route: measure whether a `global`/file-backed autotune cache carries the + tuning across processes, and if so ship it as the default for consumers; if not, consider a + warm-up prefill at model install. Gate on the budget rows like everything else. ## Phases @@ -1071,6 +1098,12 @@ that fits the model AND uses every device to the fullest. gate re-passed both legs (max |Δlogit| 2.670e-5, unchanged; Ollama greedy byte-identical), and f32 perf *improved* (TTFT 100.5 → 88.4 ms, decode 13.3 → 14.1 tok/s). f16 benches recorded in `bench/BASELINE.md`: 88.0 ms TTFT, 14.1 tok/s — speed parity with f32, VRAM halved.* + *(2026-08-06 correction)* The islands themselves are unaffected — they were validated by + `real_f16.rs`, one dtype alias per process, genuinely f16 (no NaN, 6.75 GiB, coherent output). But + the **"speed parity with f32" bench line above is withdrawn**: that row came from the two-alias + bench process and was an f32 run mislabelled f16 (bisected this run — see the closed drift item in + the perf section). Real f16 decode is 20.5 ms/token vs f32's 60.0. The same withdrawal applies to + the SPIR-V item's "+30 % on BOTH dtypes" — only its f32 half was ever measured. - [x] Evaluate burn-wgpu's **`spirv` compiler feature** on Vulkan (CubeCL SPIR-V backend instead of WGSL/naga): claims significantly faster matmul incl. TensorCores at f16 — could be the cheapest decode-tok/s lever on the dev GPU; gate on the parity harness + `bench/BASELINE.md` — diff --git a/bench/BASELINE.md b/bench/BASELINE.md index 696d88d..6377861 100644 --- a/bench/BASELINE.md +++ b/bench/BASELINE.md @@ -24,23 +24,51 @@ recorded numbers (and this file's date) only on a legitimate improvement. | Decode latency (32 greedy tokens, warm KV cache) | 0.656 s → **20.5 ms/token ≈ 48.8 tok/s** | ≥ 10 tok/s | | Peak GPU memory during decode (whole card, 3.1 GiB ambient → **~3.6 GiB runner**, 2026-07-11 measure) | 6.75 GiB | ≤ 8 GiB whole-card | -**2026-08-06 re-measure — the f16 path is 2.7× faster than 2026-07-12 recorded, f32 10% slower.** -Both tables above were re-run on an idle card (criterion, three runs each, unchanged shipping code) -because the flash-attention evaluation below needed an honest control. f16 decode moved -54.5 → 20.5 ms/token and f16 TTFT 97.2 → 20.4 ms — an improvement nothing in this repo claimed, so -it came from underneath (driver, or a dependency in the CubeCL/wgpu stack, or the 2026-07-30 dtype -pinning removing a per-call policy lookup). f32 decode moved the other way, 54.3 → 60.0 ms/token; -the gate never tripped because its ceiling is 10 tok/s, four ms/token of slack away. Both drifts are -pre-existing at HEAD, not caused by anything this run shipped — bisecting the f32 one is a ROADMAP -item. +Gated by `mummu-bench/tests/budget_f16.rs`, which lives in its **own test binary** — one dtype alias +per process is what makes the row above real, and the gate asserts `logits.dtype() == F16` so a +policy-poisoned run fails loudly instead of quietly reporting f32 numbers. Its budgets are set +against what *it* measures, not the criterion row: **21 ms TTFT / 16.9 tok/s** (budgets 60 ms / +12 tok/s). The ~3× decode gap against criterion's 48.8 tok/s is warm-up, not contradiction — a single +32-step burst never leaves f16 autotune, so the gate number is roughly what the first 32 tokens of a +cold f16 session cost and criterion's is steady state. Two honest numbers for the same path, recorded +together so the gap is not mistaken for drift (the OLMoE rows below do the same). + +**2026-08-06 re-measure, then bisected — and the old f16 rows were never f16.** +Both tables were re-run on an idle card (criterion, three runs each, unchanged shipping code) because +the flash-attention evaluation below needed an honest control, and they came back far from the +2026-07-12 record: f16 decode 54.5 → **20.5 ms/token**, f16 TTFT 97.2 → **20.4 ms**, while f32 decode +went the *other* way, 54.3 → 60.0. Neither move was a mystery for long — the 2026-07-12 tree +(`e44debf`) and the pre-2026-07-30 tree (`c1826e7`) were checked out and benched on the same idle +machine the same day: + +| tree | f32 decode | f16 decode | +| --- | --- | --- | +| `e44debf` (2026-07-12, recorded 54.3 / 54.5) | 60.1 ms/token | 60.2 ms/token | +| `c1826e7` (just before the 2026-07-30 dtype pinning) | 60.0 ms/token | **panics `DTypeMismatch`** | +| HEAD (2026-08-06) | 60.0 ms/token | **20.5 ms/token** | + +Two conclusions, and the second one retires a standing belief: + +1. **f32 never regressed.** The same 2026-07-12 code reads 60.1 ms/token today. The 54.3 recorded then + is not reproducible from that tree now, so it was machine state (driver, OS, background load), + never a Mummu change. f32 decode has been ~60 ms/token on this card the whole time. +2. **The old f16 bench rows were f32 runs wearing an f16 label.** This bench instantiates `Gpu` and + then `GpuF16` in one process; before the dtype pinning, the f32 leg locked the per-device default + dtype policy and the "f16" model ran in f32 — which is exactly why 2026-07-11 and 2026-07-12 + recorded f16 as *matching f32 to a tenth of a millisecond* (70.9 vs 70.7, then 54.5 vs 54.3). Once + the loaders started taking `target_float` from the TYPE (2026-07-23) the mismatch became loud + instead of silent — hence the `DTypeMismatch` panic at `c1826e7` — and the 2026-07-30 pinning of + every runtime creation site fixed it. **Real f16 decode is 2.9× faster than f32**, and always was; + the bench simply could not see it. (The f16 VRAM figures below are unaffected: they were measured + by `real_f16.rs`, which runs one alias per process and so was always genuinely f16.) The SPIR-V compiler (burn `vulkan` feature) cut decode latency **23%** on both dtypes (70.7 → 54.3 ms/token) at the cost of ~9 ms TTFT (88.4 → 96.7 ms, still ⅔ under its ceiling); parity held byte-identically (max |Δlogit| 2.670e-5, Ollama greedy leg exact). **The "f16 buys VRAM, not -speed" reading from that run no longer holds**: as of the 2026-08-06 re-measure f16 decodes at -20.5 ms/token against f32's 60.0 — 2.9× — so f16 is now the fast path as well as the small one, and -the f32 path is the one that looks dispatch-bound (60 ms/token streams ~6.2 GB of f32 weights at -~103 GB/s against the card's ~672 GB/s). +speed" reading from that run was wrong, and the bisect above says why** — the f16 row it rested on +was an f32 run. f16 decodes at 20.5 ms/token against f32's 60.0, so f16 is the fast path as well as +the small one. What survives is the *f32* reading: 60 ms/token streams ~6.2 GB of f32 weights at +~103 GB/s against the card's ~672 GB/s, so that path is dispatch-bound, not bandwidth-bound. ## Qwen2.5-0.5B-Instruct · CPU (burn-flex) · f32 diff --git a/crates/mummu-bench/tests/budget_f16.rs b/crates/mummu-bench/tests/budget_f16.rs new file mode 100644 index 0000000..fd88a17 --- /dev/null +++ b/crates/mummu-bench/tests/budget_f16.rs @@ -0,0 +1,111 @@ +//! The **f16** perf-budget gate from `bench/BASELINE.md`: Qwen2.5-1.5B TTFT +//! and greedy decode on `GpuF16`. Ignored by default (multi-GB weights + a +//! SHADER_F16 GPU); run with +//! +//! ```text +//! MUMMU_QWEN2_DIR=path/to/qwen2.5-1.5b cargo test -p mummu-bench --release --test budget_f16 -- --ignored --nocapture +//! ``` +//! +//! This lives in its OWN test binary on purpose, and that is the whole +//! reason it exists. Burn resolves an unspecified tensor dtype against a +//! per-DEVICE policy that the first alias to touch the device locks; a +//! process that builds `Gpu` before `GpuF16` therefore ran "f16" work in f32 +//! for weeks, and the criterion bench — which benches both aliases in one +//! process — recorded f16 rows that matched f32 to a tenth of a millisecond +//! and were believed (see the 2026-08-06 bisect in `bench/BASELINE.md`). One +//! alias per process makes that impossible, and the dtype assert below makes +//! it loud if it ever becomes possible again: a gate that cannot tell which +//! precision it measured is not a gate. + +use std::path::PathBuf; +use std::time::Instant; + +use burn::tensor::DType; +use mummu::backend::{GpuF16, inventory}; +use mummu::decode::argmax_id; +use mummu::models::CausalLm; +use mummu::models::qwen2; +use tokenizers::Tokenizer; + +/// Budgets are set against what THIS harness measures (25 ms / 16.4 tok/s on +/// 2026-08-06), not against criterion's steady-state f16 row (20.4 ms / +/// 48.8 tok/s) — the two differ by ~3x on decode because a single 32-step +/// burst never leaves autotune warm-up, and the f16 path has far more +/// variants to tune than f32. Both numbers are honest about different +/// things: criterion's is steady-state throughput, this one is roughly what +/// the first 32 tokens of a cold session cost. Budgets sit ~0.75x / 2.4x off +/// the measurement, matching the f32 gate's looseness. +const TTFT_BUDGET_MS: f64 = 60.0; +const DECODE_BUDGET_TOKENS_PER_S: f64 = 12.0; +const DECODE_STEPS: usize = 32; + +fn qwen2_dir() -> Option { + let dir = PathBuf::from(std::env::var_os("MUMMU_QWEN2_DIR")?); + dir.is_dir().then_some(dir) +} + +#[test] +#[ignore = "needs multi-GB local weights (MUMMU_QWEN2_DIR) + a SHADER_F16 GPU"] +fn qwen2_f16_stays_inside_its_perf_budgets() { + let Some(dir) = qwen2_dir() else { + panic!("set MUMMU_QWEN2_DIR to a dir with config.json/tokenizer.json/model.safetensors"); + }; + if !inventory().any_shader_f16() { + eprintln!("[budget/f16] no SHADER_F16 adapter — skipping"); + return; + } + let text = "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\nExplain, in three sentences, why the sky is blue.<|im_end|>\n<|im_start|>assistant\n"; + let tok = Tokenizer::from_file(dir.join("tokenizer.json")).expect("tokenizer.json loads"); + let ids = tok.encode(text, false).expect("encodes").get_ids().to_vec(); + assert!(ids.len() >= 16, "budget prompt suspiciously short"); + + let device = burn::tensor::Device::::default(); + let loaded = qwen2::load_from_dir::(&dir, &device).expect("weights load checked"); + + // Warm-up: first-run autotune + pipeline compilation must not count. + // TWO passes, not one — a single f16 prefill leaves the prefill kernels + // still re-tuning and the TTFT reading swings 24 -> 110 ms between runs. + // The first pass also produces the tensor whose dtype proves this really + // is the f16 path — the assert this whole binary exists for. + for pass in 0..2 { + let mut cache = loaded.new_cache(); + let logits = loaded.forward(&ids, 0, &mut cache, &device); + if pass == 0 { + assert_eq!( + logits.dtype(), + DType::F16, + "this gate must measure f16: a device policy locked by another alias would \ + silently make these f32 numbers" + ); + } + let _ = argmax_id(logits).expect("warm-up argmax"); + } + + // TTFT: fresh cache, full prefill, first token (argmax readback = sync). + let start = Instant::now(); + let mut cache = loaded.new_cache(); + let logits = loaded.forward(&ids, 0, &mut cache, &device); + let mut next = argmax_id(logits).expect("argmax"); + let ttft_ms = start.elapsed().as_secs_f64() * 1e3; + + // Decode throughput over a warm cache. + let start = Instant::now(); + for past in (ids.len()..).take(DECODE_STEPS) { + let logits = loaded.forward(&[next], past, &mut cache, &device); + next = argmax_id(logits).expect("argmax"); + } + let tok_per_s = DECODE_STEPS as f64 / start.elapsed().as_secs_f64(); + + eprintln!( + "[budget/f16] TTFT {ttft_ms:.1} ms (budget {TTFT_BUDGET_MS} ms), \ + decode {tok_per_s:.1} tok/s (budget {DECODE_BUDGET_TOKENS_PER_S} tok/s)" + ); + assert!( + ttft_ms <= TTFT_BUDGET_MS, + "f16 TTFT regression: {ttft_ms:.1} ms > {TTFT_BUDGET_MS} ms budget" + ); + assert!( + tok_per_s >= DECODE_BUDGET_TOKENS_PER_S, + "f16 decode regression: {tok_per_s:.1} tok/s < {DECODE_BUDGET_TOKENS_PER_S} tok/s budget" + ); +}