From f0ed82e78cae865142dd5a4435b31204f1c3e97a Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Sun, 9 Aug 2026 06:34:06 -0500 Subject: [PATCH 1/7] deps: refresh the lockfile - 24 patch/minor bumps, every gate green `cargo update` moved 24 transitive packages to their newest Rust-1.99-nightly-compatible versions (notably tracel-ash 0.39.3 -> 0.39.5+sdk1.4.357 on the Vulkan path, ureq 3.3 -> 3.4, minijinja 2.22 -> 2.23 under the template gate, thiserror 2.0.19 -> 2.0.20, zerocopy 0.8.55 -> 0.8.56). No manifest edit was needed: `cargo upgrade --incompatible` offers only wgpu 29 -> 30, which is the standing intentional pin (burn 0.21 resolves wgpu 29 transitively, so 30 unblocks with a burn bump, not a `cargo upgrade`), and every other direct dep is already at its newest stable. burn stays 0.21.0 - 0.22.0-pre.1 is still the only 0.22 tag, so the P0 migration item stays gated. Verified: fmt clean, clippy --all-targets no warnings, 202 unit tests green, cargo build green, and the budget gates hold - GPU 110.4 ms TTFT / 13.1 tok/s / 597 ms prefill@2048, f16 30.5 ms / 13.9 tok/s, CPU 14.2 tok/s. Both GPU gates read low on their FIRST post-bump run (9.1 and 6.2 tok/s) and recovered on re-run, the documented autotune-invalidation transient. Co-Authored-By: Claude Opus 4.8 --- Cargo.lock | 144 ++++++++++++++++++++++++++++------------------------- 1 file changed, 75 insertions(+), 69 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 966d6a3..6d25543 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -208,7 +208,7 @@ dependencies = [ "num-traits", "pastey 0.1.1", "rayon", - "thiserror 2.0.19", + "thiserror 2.0.20", "v_frame", "y4m", ] @@ -263,6 +263,12 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" + [[package]] name = "base64ct" version = "1.8.3" @@ -464,7 +470,7 @@ dependencies = [ "rand_distr 0.6.0", "serde", "spin", - "thiserror 2.0.19", + "thiserror 2.0.20", ] [[package]] @@ -505,7 +511,7 @@ dependencies = [ "serde", "serde_json", "spin", - "thiserror 2.0.19", + "thiserror 2.0.20", "uuid", ] @@ -836,13 +842,13 @@ dependencies = [ [[package]] name = "bytemuck_derive" -version = "1.11.0" +version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f65693059b6b9c588b9f62fed1cedbf0a8b805631457ea162d68f0de186f3de5" +checksum = "fc0e56a716f1e132ff6bf4bdac1c944a3fcdc1cae65f70a4a2a1ac3b401d2d1f" dependencies = [ "proc-macro2", "quote", - "syn 2.0.119", + "syn 3.0.3", ] [[package]] @@ -933,7 +939,7 @@ dependencies = [ "rand_distr 0.5.1", "rayon", "safetensors 0.7.0", - "thiserror 2.0.19", + "thiserror 2.0.20", "tokenizers 0.22.2", "yoke", "zip 7.2.0", @@ -965,9 +971,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.4.0" +version = "1.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5add81bb678e6cb321aff7fa0dc7689ad82b112dbc032cea19f91d6b8e3582b9" +checksum = "5d262e149917187838d5b42777c8253bcb64500067342904e7d429499a6f277e" dependencies = [ "find-msvc-tools", "jobserver", @@ -1067,18 +1073,18 @@ dependencies = [ [[package]] name = "clap" -version = "4.6.5" +version = "4.6.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "301b56658598e48f3648647ac6fc887be7e7108eddfa4e9b63fcf3ec58c0cadf" +checksum = "473c7e07f409a8d772161724aa8db6a765a2532a70f9667eeb7b49d3d02fbdca" dependencies = [ "clap_builder", ] [[package]] name = "clap_builder" -version = "4.6.5" +version = "4.6.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94a65403d1a1bd28f7dc68eb8506e8874808ee5eecb59298de588e2e1407a078" +checksum = "7b48fea5a88e9ae728a2dcbedbfc0e730f7d60da42e1cb049a83c9fb8b789889" dependencies = [ "anstyle", "clap_lex", @@ -1635,7 +1641,7 @@ dependencies = [ "serde", "serde_json", "spin", - "thiserror 2.0.19", + "thiserror 2.0.20", "toml", "tracing", "wasm-bindgen-futures", @@ -1824,7 +1830,7 @@ dependencies = [ "half", "num-traits", "serde", - "thiserror 2.0.19", + "thiserror 2.0.20", ] [[package]] @@ -2374,9 +2380,9 @@ dependencies = [ [[package]] name = "find-msvc-tools" -version = "0.1.9" +version = "0.1.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" +checksum = "26b73573e6edcd2af0cdf47bd6cb58f0b3839491263c314eaad1ccf24430e1de" [[package]] name = "fixedbitset" @@ -2740,7 +2746,7 @@ dependencies = [ "hashbrown 0.16.1", "log", "presser", - "thiserror 2.0.19", + "thiserror 2.0.20", "windows", ] @@ -3269,9 +3275,9 @@ dependencies = [ [[package]] name = "js-sys" -version = "0.3.103" +version = "0.3.104" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102" +checksum = "0e0c1080212aad755ea003d18543e8768dd432c48819efd73a7bf1e39b7a5a3a" dependencies = [ "cfg-if", "futures-util", @@ -3536,9 +3542,9 @@ checksum = "38d1115007560874e373613744c6fba374c17688327a71c1476d1a5954cc857b" [[package]] name = "minijinja" -version = "2.22.0" +version = "2.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef84a52be188a1d4124bd717903fdde96ca4705f2b56adfe2d91fcc57fcb6987" +checksum = "42d74234349a775546a83af0f0c0c0e3a73227dee4950a542cda81a47240b3e6" dependencies = [ "indexmap", "memo-map", @@ -3548,9 +3554,9 @@ dependencies = [ [[package]] name = "minijinja-contrib" -version = "2.22.0" +version = "2.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fd6e279dc925840c2d9ebc1e9f85410611d01292561297463ba3e516c3ad94dd" +checksum = "35c96d8fae7fa4743bbcf06f486ff43f81837f5e8826f8bf525dbc4c0e03be5d" dependencies = [ "minijinja", "serde", @@ -3635,9 +3641,9 @@ dependencies = [ "serde", "serde_json", "sha2 0.11.0", - "thiserror 2.0.19", + "thiserror 2.0.20", "tokenizers 0.23.1", - "ureq 3.3.0", + "ureq 3.4.0", "wgpu", "windows-sys 0.61.2", ] @@ -3674,7 +3680,7 @@ dependencies = [ "once_cell", "rustc-hash 1.1.0", "spirv", - "thiserror 2.0.19", + "thiserror 2.0.20", "unicode-ident", ] @@ -4335,7 +4341,7 @@ dependencies = [ "rustc-hash 2.1.3", "rustls", "socket2", - "thiserror 2.0.19", + "thiserror 2.0.20", "tokio", "tracing", "web-time", @@ -4357,7 +4363,7 @@ dependencies = [ "rustls", "rustls-pki-types", "slab", - "thiserror 2.0.19", + "thiserror 2.0.20", "tinyvec", "tracing", "web-time", @@ -4539,7 +4545,7 @@ dependencies = [ "rand 0.9.5", "rand_chacha 0.9.0", "simd_helpers", - "thiserror 2.0.19", + "thiserror 2.0.20", "v_frame", "wasm-bindgen", ] @@ -4646,7 +4652,7 @@ checksum = "a4e608c6638b9c18977b00b475ac1f28d14e84b27d8d42f70e0bf1e3dec127ac" dependencies = [ "getrandom 0.2.17", "libredox", - "thiserror 2.0.19", + "thiserror 2.0.20", ] [[package]] @@ -5306,11 +5312,11 @@ dependencies = [ [[package]] name = "thiserror" -version = "2.0.19" +version = "2.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09a43598840e33d5b0331f38c5e30d13bb11c11210a4b58f0d9b18a5a5eefcd9" +checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f" dependencies = [ - "thiserror-impl 2.0.19", + "thiserror-impl 2.0.20", ] [[package]] @@ -5326,9 +5332,9 @@ dependencies = [ [[package]] name = "thiserror-impl" -version = "2.0.19" +version = "2.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43cbfe0cf76104d42a574802844187e84a305e531ed54455f11fbde0f10541cd" +checksum = "bc04cd3e1236dd4a98afca4569f2deb3f120e5422a4023be2cb683f8486292af" dependencies = [ "proc-macro2", "quote", @@ -5440,7 +5446,7 @@ dependencies = [ "serde", "serde_json", "spm_precompiled", - "thiserror 2.0.19", + "thiserror 2.0.20", "unicode-normalization-alignments", "unicode-segmentation", "unicode_categories", @@ -5474,7 +5480,7 @@ dependencies = [ "serde", "serde_json", "spm_precompiled", - "thiserror 2.0.19", + "thiserror 2.0.20", "unicode-normalization-alignments", "unicode-segmentation", "unicode_categories", @@ -5619,9 +5625,9 @@ checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" [[package]] name = "tracel-ash" -version = "0.39.3+sdk1.4.350" +version = "0.39.5+sdk1.4.357" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d97fcf7e656528db871e679727140819a221d8c2aa5b90bfb5eb492a0ecc16e6" +checksum = "9802d90b2e66f00a395a3acf337ddd49dd1e90cb3287fc31a4148f875d7be7c5" dependencies = [ "ash", "c2rust-bitfields", @@ -5713,7 +5719,7 @@ dependencies = [ "bindgen", "cc", "paste", - "thiserror 2.0.19", + "thiserror 2.0.20", "tracel-llvm-bundler", ] @@ -5882,11 +5888,11 @@ dependencies = [ [[package]] name = "ureq" -version = "3.3.0" +version = "3.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dea7109cdcd5864d4eeb1b58a1648dc9bf520360d7af16ec26d0a9354bafcfc0" +checksum = "972d7902c8735f2695410b8aed7df6ed12a47394aa1c8d7af49f0497b731a94d" dependencies = [ - "base64 0.22.1", + "base64 0.23.1", "flate2", "log", "percent-encoding", @@ -5899,11 +5905,11 @@ dependencies = [ [[package]] name = "ureq-proto" -version = "0.6.0" +version = "0.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e994ba84b0bd1b1b0cf92878b7ef898a5c1760108fe7b6010327e274917a808c" +checksum = "da5f78b09e6941e1a0f2e30e695e4b120377b54d5e0aec11b594bb57b3971613" dependencies = [ - "base64 0.22.1", + "base64 0.23.1", "http", "httparse", "log", @@ -6009,9 +6015,9 @@ dependencies = [ [[package]] name = "wasm-bindgen" -version = "0.2.126" +version = "0.2.127" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4" +checksum = "1b70935747edd64d89de3efa29d73789b806c15798f8e7dca4d8ac356b50ce70" dependencies = [ "cfg-if", "once_cell", @@ -6022,9 +6028,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-futures" -version = "0.4.76" +version = "0.4.77" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c62df1340f32221cb9c54d6a27b030e3dba64361d4a95bed55f9aacb44da291d" +checksum = "6b7777d5cc23d0e91404e53ce2d5e8ec7acae3026b16233dba62cd3246457950" dependencies = [ "js-sys", "wasm-bindgen", @@ -6032,9 +6038,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.126" +version = "0.2.127" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1" +checksum = "77775f8f3f7217702089053b94958f8f54061a3f663417df76e19cbdcca29bc1" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -6042,9 +6048,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.126" +version = "0.2.127" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e" +checksum = "e11d33f857dc2fb11b8bc75aee111aa9cbeb12cd9f25efd3d4c2a3dd4e235284" dependencies = [ "bumpalo", "proc-macro2", @@ -6055,9 +6061,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.126" +version = "0.2.127" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24" +checksum = "7ef64dbcc55df09c7e5a46182d181c2cfa3e925f3da937ea764728b4bbb9dcbf" dependencies = [ "unicode-ident", ] @@ -6076,9 +6082,9 @@ dependencies = [ [[package]] name = "web-sys" -version = "0.3.103" +version = "0.3.104" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8622dcb61c0bcc9fffa6938bed81210af2da9a7e4a1a834b2e37a59b6dfb6141" +checksum = "c435338968042f4f59a557f690a253676d47ce13ceb55d70100e7facf6620a30" dependencies = [ "js-sys", "wasm-bindgen", @@ -6172,7 +6178,7 @@ dependencies = [ "raw-window-handle", "rustc-hash 1.1.0", "smallvec", - "thiserror 2.0.19", + "thiserror 2.0.20", "wgpu-core-deps-apple", "wgpu-core-deps-emscripten", "wgpu-core-deps-windows-linux-android", @@ -6251,7 +6257,7 @@ dependencies = [ "raw-window-metal", "renderdoc-sys", "smallvec", - "thiserror 2.0.19", + "thiserror 2.0.20", "wasm-bindgen", "wayland-sys", "web-sys", @@ -6539,9 +6545,9 @@ dependencies = [ [[package]] name = "xml-rs" -version = "0.8.28" +version = "0.8.29" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3ae8337f8a065cfc972643663ea4279e04e7256de865aa66fe25cec5fb912d3f" +checksum = "e450f9b2ed1dff33c94c12589a87338689467b9c4f5d8a5710bd09a847d2c8a7" [[package]] name = "xxhash-rust" @@ -6580,18 +6586,18 @@ dependencies = [ [[package]] name = "zerocopy" -version = "0.8.55" +version = "0.8.56" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5a105cd7b140f6eeec8acff2ea38135d3cab283ada58540f629fe51e46696eb" +checksum = "556764e583adb45a9f8d413c2a147fa7e8d821e48e12b14fd560b607998b75eb" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.55" +version = "0.8.56" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fe976fb70c78cd64cccfe3a6fc142244e8a77b70959b30faf9d0ac37ee228eb" +checksum = "f2ab42fc20575779bd240faa45f94a74256f755c0fa9e89f0ede20d91d0cdfc1" dependencies = [ "proc-macro2", "quote", @@ -6790,9 +6796,9 @@ dependencies = [ [[package]] name = "zune-core" -version = "0.5.1" +version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cb8a0807f7c01457d0379ba880ba6322660448ddebc890ce29bb64da71fb40f9" +checksum = "d56377fd46368984a170bc5aac5567e52ca5da874caa60bea39fcbca78fb658b" [[package]] name = "zune-inflate" From 08b9ed8e096b0e33041283ef799efad26f7cd442 Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Sun, 9 Aug 2026 07:49:30 -0500 Subject: [PATCH 2/7] perf: the f16 cold-start tax is one 32-token burst - measured, then paid off the critical path Closes the ROADMAP item "close the f16 autotune warm-up gap", and the route it proposed turned out to rest on a wrong premise. CubeCL ALREADY persists autotune across processes by default (`[cubecl.autotune] cache` defaults to `target`; this repo's live cache is crates/mummu-bench/target/autotune/), so a `global`/file-backed location changes where the cache lives, never whether tuning carries. What a cold process still pays is kernel compilation + pipeline creation, which the wgpu runtime caches nowhere - cubecl 0.10 wires `CompilationCache` for the CUDA and HIP runtimes only, so `[cubecl.compilation] cache` is inert on our path. No configuration can carry it; only spending it earlier can. New `mummu-bench/tests/warmup_f16.rs` measures the whole curve in one process: 8 bursts x 32 decode steps, each burst a fresh cache + untimed prefill so KV length stays constant (decoding 256 CONSECUTIVE tokens instead confounds warm-up with attention length - measured, it reads ~20% low by the last burst). The first 32 tokens run at 12.5-16.3 tok/s and the curve is FLAT from token 33 on at 37-42. So the cold tax is 2.5-3.0x, exactly one burst deep, and `budget_f16.rs`'s 16.9 tok/s is not a mystery - it IS the first burst. Shipped `CausalLm::warm_up(probe_ids, steps, device)`: one prefill plus `steps` greedy decode steps on a throwaway cache, every step's argmax read back (an unsynchronized warm-up returns before the GPU runs anything), bounded by `MAX_WARM_UP_STEPS = 256`, default-implemented so every zoo model inherits it beside `sanity_check`. REAL-GPU proof in its own binary (`warmup_api_f16.rs` - warm-up is a once-per-process effect, so a second test in one binary would prove nothing): after a 4.21 s `warm_up(&ids, 32)` a cold process's FIRST burst runs at 41.9 tok/s against the next burst's 41.0, ratio 1.02x where un-warmed it is 0.33x. Two findings folded into bench/BASELINE.md rather than acted on here: 1. A STALE AUTOTUNE CACHE silently cost 21-27% of f16 decode. This run's first budget run happened on a contended machine (9.1 tok/s f32, failing its own gate before recovering to 13.1); autotune tuned under that contention, wrote its picks, and every later process loaded them and never re-tuned. Deleting the cache: f16 decode_32_tokens 1.0279 s -> 0.8109/0.8371 s, f32 unmoved (1.9646 -> 1.9797). New ROADMAP item. 2. The gap's other half - "48.8 tok/s steady" - is not reproducible. Criterion measures 36.8 today, and the SAME numbers come back on the pre-`cargo update` lockfile, so it is machine state, not a regression. Measured the mechanism directly: +22pp of host CPU load costs f32 +3.2% and f16 +9.8%, because f16 does ~3x less GPU work per dispatch. README's retracted f16 claims are corrected while here: "identical speed" (an f32 run mislabelled f16, retracted 2026-08-06 in BASELINE.md but never fixed in the README) becomes the measured ~2.3x, and the benchmark bullet carries today's full criterion set. Verified: fmt clean, clippy --all-targets no warnings, 204 unit tests (2 new), cargo build green, warmup_f16 15.4 -> 41.4 tok/s and the f16 budget gate 25.2 ms / 15.5 tok/s. Co-Authored-By: Claude Opus 4.8 --- README.md | 22 +++- ROADMAP.md | 44 ++++++- bench/BASELINE.md | 53 ++++++++ crates/mummu-bench/tests/warmup_api_f16.rs | 103 +++++++++++++++ crates/mummu-bench/tests/warmup_f16.rs | 145 +++++++++++++++++++++ crates/mummu/src/models/mod.rs | 63 ++++++++- crates/mummu/src/models/qwen2.rs | 34 +++++ 7 files changed, 456 insertions(+), 8 deletions(-) create mode 100644 crates/mummu-bench/tests/warmup_api_f16.rs create mode 100644 crates/mummu-bench/tests/warmup_f16.rs diff --git a/README.md b/README.md index 2d61194..2748051 100644 --- a/README.md +++ b/README.md @@ -130,9 +130,18 @@ It exists because two local-first apps — **[laurelane](https://github.com/phys (serde_json `preserve_order`) — the exact bytes the reference stack renders and models emit back. - **f16 inference, validated** — Qwen2.5-1.5B runs coherently on `GpuF16` (weights + KV in f16, the q·kᵀ attention scores + softmax computed in an f32 island to stop f16 overflow): **~3.6 GiB runner - VRAM vs ~7.9 GiB f32, at identical speed**; the parity gate re-passes unchanged on f32, where the - island casts are no-ops ([bench/BASELINE.md](bench/BASELINE.md)). The same island covers the Qwen3 + VRAM vs ~7.9 GiB f32, and ~2.3× the decode throughput** (27.1 vs 61.8 ms/token, measured + 2026-08-09); the parity gate re-passes unchanged on f32, where the island casts are no-ops + ([bench/BASELINE.md](bench/BASELINE.md)). The same island covers the Qwen3 arch — Qwen3-0.6B decodes coherently in f16 (its qk-norm + decoupled head_dim ride the same f32 scores). +- **Warm-up API** — a freshly-started process decodes its first ~32 tokens at roughly a third of its + steady rate (per-process kernel compilation + pipeline creation; CubeCL persists *autotune* across + processes but the wgpu runtime caches no compiled kernels). `CausalLm::warm_up(probe_ids, steps, + device)` pays that cost off the user's critical path — one prefill plus `steps` greedy decode steps on + a throwaway cache, bounded and synchronized. Measured on the reference GPU: after a 4.2 s warm-up a + cold process's first 32-token burst runs at **41.9 tok/s vs the next burst's 41.0** (un-warmed, that + ratio is 0.33×) — `mummu-bench/tests/warmup_api_f16.rs`, curve in + [bench/BASELINE.md](bench/BASELINE.md). - **In-process mixed precision is defined behavior** — every runtime tensor-creation site pins its dtype to the backend type (`backend::{float_dtype, int_dtype}`), so an f32 (`Gpu`) and an f16 (`GpuF16`) model can share one process and one device regardless of Burn's first-touch-locked @@ -142,10 +151,11 @@ It exists because two local-first apps — **[laurelane](https://github.com/phys - **SPIR-V kernels on Vulkan** — CubeCL compiles direct SPIR-V (burn's `vulkan` feature) instead of WGSL/naga on Vulkan adapters, worth **+30% decode throughput** on the reference GPU with parity byte-identical; other APIs (DX12/Metal) transparently keep WGSL in the same binary. -- **Benchmarked** — Qwen2.5-1.5B on the reference GPU: **TTFT 96.7 ms, decode 18.4 tok/s** (f32, - 11.5 GiB whole-card peak ≈ 8.0 GiB runner; f16: 97.2 ms, 18.4 tok/s, ~3.6 GiB runner) — recorded with - budgets in [bench/BASELINE.md](bench/BASELINE.md), enforced by an opt-in regression gate - (`mummu-bench/tests/budget.rs`). +- **Benchmarked** — Qwen2.5-1.5B on the reference GPU, criterion: **f32 TTFT 98.2 ms, decode + 16.2 tok/s, prefill@2048 597 ms** (11.5 GiB whole-card peak ≈ 8.0 GiB runner); **f16 TTFT 24.9 ms, + decode 36.8 tok/s, prefill@2048 241 ms** (~3.6 GiB runner) — recorded with budgets in + [bench/BASELINE.md](bench/BASELINE.md), enforced by opt-in regression gates + (`mummu-bench/tests/budget{,_f16,_cpu,_moe}.rs`, one dtype alias per process). - **Model management** — `ModelManager` gives settings UIs the whole lifecycle over a declarative model catalog (`registry::ModelSpec`): install with per-chunk download progress, `is_installed`, per-model disk usage, and traversal-safe removal; model switching rides `ModelSlot`. diff --git a/ROADMAP.md b/ROADMAP.md index 2f860c5..7cc1766 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -202,7 +202,7 @@ a benchmark holds/improves its budget; README perf claims link an artifact. **`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 +- [x] **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 @@ -214,6 +214,48 @@ a benchmark holds/improves its budget; README perf claims link an artifact. 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. + *(2026-08-09) **Measured, and the caching premise was wrong — shipped the warm-up instead.** The + route's first branch is a non-question: CubeCL **already** persists autotune across processes by + default (`[cubecl.autotune] cache` defaults to `target`; this repo's live cache is + `crates/mummu-bench/target/autotune/`, written and re-loaded on every run), so a `global`/file + location changes *where* the cache lives, never *whether* tuning carries. What the cold process + still pays is **kernel compilation + pipeline creation**, which the wgpu runtime caches nowhere — + cubecl 0.10 wires `CompilationCache` for the CUDA and HIP runtimes only, so the + `[cubecl.compilation] cache` knob is inert on our path. No configuration can carry it; only + spending it earlier can. New harness `mummu-bench/tests/warmup_f16.rs` measures the whole curve in + one process (8 bursts × 32 decode steps, each burst a fresh cache + untimed prefill so KV length + stays constant): **the first 32 tokens run at 12.5–16.3 tok/s, burst 2 onward is flat at + 37–41**, i.e. the cold tax is 2.5–3.0× and exactly ONE burst deep — nothing beyond it is + recoverable in-process. So `budget_f16.rs`'s 16.9 tok/s is not a mystery, it IS the first burst. + Shipped: **`CausalLm::warm_up(probe_ids, steps, device)`** — one prefill plus `steps` greedy decode + steps on a throwaway cache, every step's argmax read back (an unsynchronized warm-up returns before + the GPU runs anything), bounded by `MAX_WARM_UP_STEPS = 256`, default-implemented so every zoo model + inherits it beside `sanity_check`. REAL-GPU proof (`mummu-bench/tests/warmup_api_f16.rs`, own binary + because warm-up is a once-per-process effect): after a 4.21 s `warm_up(&ids, 32)`, a cold process's + FIRST 32-token burst runs at **41.9 tok/s** vs the next burst's 41.0 — ratio 1.02× where un-warmed + it is 0.33×. Warms the decode step (shape-stable at `t == 1`); prefill kernels key on prompt length, + documented on the fn. The other half of the gap's premise — "48.8 steady" — did not survive: + criterion measures 36.8 tok/s today and the same numbers come back on the pre-`cargo update` + lockfile, so it is machine state, not a regression (see `bench/BASELINE.md`, incl. the measured + f16-vs-f32 host-CPU sensitivity, +9.8 % vs +3.2 % for the same added load).* +- [ ] **A stale autotune cache is permanent, silent, and cost 21–27 % of f16 decode** — found while + closing the item above, and it is the reason "just persist the autotune cache" is not a free win + for consumers. This run's FIRST budget run happened on a contended machine (9.1 tok/s f32, failing + its own gate before recovering to 13.1); autotune tuned under that contention, wrote its picks to + `crates/mummu-bench/target/autotune/`, and **every later process loaded them and never re-tuned**. + Deleting the directory and re-running the same code on the same machine: f16 + `decode_32_tokens` 1.0279 s → 0.8109 / 0.8371 s, while f32 was unmoved (1.9646 → 1.9797 s). CubeCL + offers no invalidation, no re-tune trigger, and no confidence signal — the cache is keyed by + (device, kernel, checksum), and a pick made under load is indistinguishable from a good one. That + is a real hazard for a runner whose consumers ship a persisted cache to end users' machines, where + the tune may happen during install (busy) and be believed forever. Routes: (a) re-tune on demand — + an API to clear the cache root (Mummu knows where it is; `CacheConfig::root()` is public) plus a + documented "re-tune" action in the consumer's settings; (b) pin the cache to a Mummu-owned location + via `RuntimeConfig::set` (cubecl exposes a one-shot programmatic setter that must run before the + first `get()`) so the runner controls invalidation rather than inheriting the CWD walk-up; (c) + validate on load — time one warm-up burst against a recorded expectation and clear the cache when + it reads far low, which the `warm_up` API already has the shape for. Gate any of them on + `bench/BASELINE.md` like everything else. *(2026-08-09, measured this run.)* ## Phases diff --git a/bench/BASELINE.md b/bench/BASELINE.md index 6377861..eac4e81 100644 --- a/bench/BASELINE.md +++ b/bench/BASELINE.md @@ -33,6 +33,59 @@ against what *it* measures, not the criterion row: **21 ms TTFT / 16.9 tok/s** ( 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-09: the warm-up gap is now a curve, not an inference — and it is exactly one burst deep.** +`mummu-bench/tests/warmup_f16.rs` measures what the two rows above could only bracket: a cold f16 +process, no warm-up pass, then 8 bursts of 32 decode steps, each burst a fresh cache + untimed +prefill (criterion's `decode_32_tokens` sample, replicated so KV length is constant across bursts — +decoding 256 *consecutive* tokens instead confounds warm-up with attention length and reads ~20 % +low by the last burst). Two runs, same session: + +| cumulative tokens | 32 | 64 | 96 | 128 | 160 | 192 | 224 | 256 | +| --- | --- | --- | --- | --- | --- | --- | --- | --- | +| run A, tok/s | 12.5 | 37.4 | 38.8 | 36.7 | 38.1 | 37.4 | 36.8 | 37.6 | +| run B, tok/s | 16.3 | 41.1 | 39.7 | 39.3 | 40.3 | 40.6 | 41.3 | 40.3 | + +The first 32 tokens cost **2.5–3.0×** steady state and the curve is flat from token 33 on. That +settles the shape of the gap: `budget_f16.rs`'s number IS the first burst (15.5–16.3 tok/s measured +the same session), and nothing beyond one burst is recoverable in-process. It also names what the +cost is *not*: CubeCL already persists **autotune** choices across processes by default +(`[cubecl.autotune] cache` defaults to `target` — this repo's live cache is +`crates/mummu-bench/target/autotune/`), so what is left is per-process kernel compilation and +pipeline creation, which the wgpu runtime caches nowhere — `CompilationCache` is wired for the CUDA +and HIP runtimes only in cubecl 0.10, and `[cubecl.compilation] cache` is inert on our path. No +configuration carries it; only spending it earlier does, which is what `CausalLm::warm_up` is for. +Proof it works (`mummu-bench/tests/warmup_api_f16.rs`): after one `warm_up(&ids, 32, …)` costing +**4.21 s**, a cold process's FIRST burst runs at **41.9 tok/s** against the next burst's 41.0 — +a ratio of 1.02× where the un-warmed ratio is 0.33×. + +**2026-08-09: a stale autotune cache silently cost 21–27 % of f16 decode.** Found while measuring the +above, and it is the reason a persisted cache is not a free win. The day's first budget run happened +on a contended machine (it measured 9.1 tok/s f32 and failed its own gate before recovering to 13.1 +on a re-run); autotune tuned under that contention, wrote its picks to +`crates/mummu-bench/target/autotune/`, and **every later process loaded those picks and never +re-tuned**. Deleting the directory and re-running the same code on the same machine: + +| criterion `decode_32_tokens` | stale cache | after deleting the cache | +| --- | --- | --- | +| f32 | 1.9646 s | 1.9797 s (unchanged) | +| f16 | 1.0279 s | **0.8109 s / 0.8371 s** (two runs) | + +f32's picks were unaffected; f16's were not. The operational rule: **delete the autotune cache before +believing an f16 regression**, the same way the 2026-07-24 note says to re-run after a config change. +The general lesson for the runner — a persistent autotune cache makes a bad tune permanent and +silent, with no invalidation and no re-tune trigger — is a ROADMAP item, not something a benchmark +file can fix. + +**2026-08-09 re-measure (context for the rows above, budgets unchanged).** Full criterion run on the +same idle-ish machine (~15 % ambient host CPU from other work), fresh autotune cache: f32 TTFT +98.2 ms / prefill@2048 597 ms / decode 61.8 ms/token, and f16 TTFT 24.9 ms / prefill@2048 241 ms / +decode 27.1 ms/token. The f32 rows reproduce the 2026-08-06 record within 3–8 %; the f16 rows read +15–32 % slower. That asymmetry is the dispatch-bound thesis showing up as *sensitivity*, and it was +measured directly rather than assumed: adding 8 spinning threads (host CPU 15 % → 37 %) costs f32 +**+3.2 %** and f16 **+9.8 %** — f16 does ~3× less GPU work per dispatch, so the same host-side +sequencing is ~3× more of its per-token time. The recorded columns are left at 2026-08-06 because +today's machine was not quieter, not because today's numbers are wrong. + **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 diff --git a/crates/mummu-bench/tests/warmup_api_f16.rs b/crates/mummu-bench/tests/warmup_api_f16.rs new file mode 100644 index 0000000..cebea16 --- /dev/null +++ b/crates/mummu-bench/tests/warmup_api_f16.rs @@ -0,0 +1,103 @@ +//! Proof that `CausalLm::warm_up` actually pays the cold-start tax +//! `warmup_f16.rs` measures: after one warm-up call, a cold process's FIRST +//! 32-token burst already runs at steady-state speed instead of ~1/3 of it. +//! +//! Own test binary for two reasons: one dtype alias per process (see +//! `budget_f16.rs`), and — the load-bearing one — warm-up is a +//! once-per-process effect, so a second test in the same binary would find +//! the GPU already warm and prove nothing. +//! +//! ```text +//! MUMMU_QWEN2_DIR=path/to/qwen2.5-1.5b cargo test -p mummu-bench --release --test warmup_api_f16 -- --ignored --nocapture +//! ``` + +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; + +/// Same burst size as `warmup_f16.rs` and `budget_f16.rs`, so the numbers +/// here sit on the same axis as the recorded rows. +const BURST_TOKENS: usize = 32; +/// Warm-up depth: the measured curve flattens after one 32-token burst, so +/// this is exactly the depth the harness says is needed — not a margin. +const WARM_UP_STEPS: usize = 32; +/// How close to the second burst the FIRST burst must land for the warm-up to +/// count as having worked. Un-warmed, the ratio is ~0.33 (12.5 vs 37.6 tok/s); +/// this floor sits well above that and below run-to-run jitter. +const WARMED_RATIO_FLOOR: f64 = 0.80; + +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 warm_up_puts_the_first_burst_at_steady_state() { + 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!("[warmup-api/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, "warm-up prompt suspiciously short"); + + let device = burn::tensor::Device::::default(); + let loaded = qwen2::load_from_dir::(&dir, &device).expect("weights load checked"); + + // The API under test — the whole cold cost is meant to land here. + let warm_start = Instant::now(); + let forwards = loaded + .warm_up(&ids, WARM_UP_STEPS, &device) + .expect("warm-up runs"); + let warm_s = warm_start.elapsed().as_secs_f64(); + assert_eq!( + forwards, + WARM_UP_STEPS + 1, + "warm_up must report one prefill plus its decode steps" + ); + + let mut rates = [0.0f64; 2]; + for rate in &mut rates { + let mut cache = loaded.new_cache(); + let logits = loaded.forward(&ids, 0, &mut cache, &device); + assert_eq!( + logits.dtype(), + DType::F16, + "this harness must measure f16: a device policy locked by another alias would \ + silently make these f32 numbers" + ); + let mut next = argmax_id(logits).expect("argmax"); + let start = Instant::now(); + for past in (ids.len()..).take(BURST_TOKENS) { + let logits = loaded.forward(&[next], past, &mut cache, &device); + next = argmax_id(logits).expect("argmax"); + } + let elapsed = start.elapsed().as_secs_f64(); + assert!(elapsed > 0.0, "a burst cannot take zero time"); + *rate = BURST_TOKENS as f64 / elapsed; + } + + let [first, second] = rates; + let ratio = first / second; + eprintln!( + "[warmup-api/f16] warm_up {warm_s:.2} s ({forwards} forwards), \ + then burst 1 {first:.1} tok/s vs burst 2 {second:.1} tok/s = {ratio:.2}x \ + (floor {WARMED_RATIO_FLOOR}x; un-warmed this ratio is ~0.33x)" + ); + assert!( + ratio >= WARMED_RATIO_FLOOR, + "warm_up did not warm the decode path: first burst {first:.1} tok/s is only \ + {ratio:.2}x the second burst's {second:.1} tok/s" + ); +} diff --git a/crates/mummu-bench/tests/warmup_f16.rs b/crates/mummu-bench/tests/warmup_f16.rs new file mode 100644 index 0000000..e057337 --- /dev/null +++ b/crates/mummu-bench/tests/warmup_f16.rs @@ -0,0 +1,145 @@ +//! The **warm-up curve** the f16 budget gate only sees one point of: how many +//! decoded tokens a *cold process* needs before f16 reaches its steady-state +//! rate, and how much of that cost the persistent autotune cache already +//! carries across processes. +//! +//! `bench/BASELINE.md` records two honest f16 decode numbers that differ ~3x — +//! 16.9 tok/s from `budget_f16.rs` (one 32-step burst) and 48.8 tok/s from +//! criterion (many 32-step samples). This harness measures the whole curve in +//! ONE process instead of one point in each, so the gap stops being an +//! inference from two harnesses and becomes a measurement. +//! +//! Own test binary, like every `GpuF16` leg: one dtype alias per process (see +//! `budget_f16.rs`). +//! +//! ```text +//! MUMMU_QWEN2_DIR=path/to/qwen2.5-1.5b cargo test -p mummu-bench --release --test warmup_f16 -- --ignored --nocapture +//! ``` + +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; + +/// Tokens per burst — the same 32 the two existing f16 rows use, so a burst +/// here is directly comparable to `budget_f16.rs`'s single measurement. +const BURST_TOKENS: usize = 32; +/// Bounded: 8 bursts = 256 decoded tokens, ~5-15 s of GPU time depending on +/// where in the curve the run sits. Long enough to reach the criterion +/// steady state, short enough that the gate stays a gate. +const BURSTS: usize = 8; +/// The steady-state floor. Measured 37.6 / 40.3 tok/s on an idle-ish machine +/// — but this row is unusually sensitive to *host* CPU contention (f16 does +/// ~3x less GPU work per dispatch than f32, so host-side sequencing is a +/// larger share of its per-token time; see `bench/BASELINE.md`), and a run +/// sharing the box with a workspace build measured 25.8. The floor is +/// therefore set just above the **f32** decode rate (~16 tok/s): what it +/// defends is the claim that steady-state f16 is meaningfully the faster path, +/// which no amount of contention should erase, rather than a specific number +/// the machine cannot promise. +const STEADY_BUDGET_TOKENS_PER_S: f64 = 20.0; + +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 f16_decode_warms_up_within_a_bounded_number_of_tokens() { + 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!("[warmup/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, "warm-up prompt suspiciously short"); + + let device = burn::tensor::Device::::default(); + let load_start = Instant::now(); + let loaded = qwen2::load_from_dir::(&dir, &device).expect("weights load checked"); + let load_s = load_start.elapsed().as_secs_f64(); + + // NO warm-up pass: the cold cost is the measurement. The first prefill is + // timed like any other, and its logits carry the dtype proof this binary + // exists for. + let start = Instant::now(); + let mut cache = loaded.new_cache(); + let logits = loaded.forward(&ids, 0, &mut cache, &device); + assert_eq!( + logits.dtype(), + DType::F16, + "this harness must measure f16: a device policy locked by another alias would \ + silently make these f32 numbers" + ); + let cold_top_id = argmax_id(logits).expect("argmax"); + let cold_ttft_ms = start.elapsed().as_secs_f64() * 1e3; + + // Each burst reproduces criterion's `decode_32_tokens` sample exactly: + // fresh cache, untimed prefill, then 32 timed steps. Decoding 256 + // CONSECUTIVE tokens instead would grow the KV cache 8x across the run and + // confound warm-up with attention length — measured, and it does: the + // consecutive form's last burst reads ~20% below its own plateau. + let mut rates = Vec::with_capacity(BURSTS); + for _ in 0..BURSTS { + let mut cache = loaded.new_cache(); + let logits = loaded.forward(&ids, 0, &mut cache, &device); + let mut next = argmax_id(logits).expect("argmax"); + assert_eq!( + next, cold_top_id, + "every burst prefills the same prompt, so its first token must not change" + ); + let start = Instant::now(); + let mut past = ids.len(); + for _ in 0..BURST_TOKENS { + let logits = loaded.forward(&[next], past, &mut cache, &device); + next = argmax_id(logits).expect("argmax"); + past += 1; + } + let elapsed = start.elapsed().as_secs_f64(); + assert!(elapsed > 0.0, "a burst cannot take zero time"); + assert_eq!( + past, + ids.len() + BURST_TOKENS, + "cache position must advance one step per decoded token" + ); + rates.push(BURST_TOKENS as f64 / elapsed); + } + assert_eq!(rates.len(), BURSTS, "every burst must record a rate"); + + let first = rates[0]; + let steady = rates[BURSTS - 1]; + let curve: Vec = rates + .iter() + .enumerate() + .map(|(i, r)| format!("{}:{r:.1}", (i + 1) * BURST_TOKENS)) + .collect(); + eprintln!( + "[warmup/f16] load {load_s:.1} s, cold TTFT {cold_ttft_ms:.1} ms, \ + burst tok/s by cumulative token — {}", + curve.join(" ") + ); + eprintln!( + "[warmup/f16] first burst {first:.1} tok/s, steady {steady:.1} tok/s, \ + warm-up cost {:.2}x (budget: steady >= {STEADY_BUDGET_TOKENS_PER_S} tok/s)", + steady / first + ); + + assert!( + steady >= STEADY_BUDGET_TOKENS_PER_S, + "f16 steady-state regression: {steady:.1} tok/s < {STEADY_BUDGET_TOKENS_PER_S} tok/s budget" + ); + assert!( + steady >= first, + "warm-up must not run backwards: steady {steady:.1} tok/s < first burst {first:.1} tok/s" + ); +} diff --git a/crates/mummu/src/models/mod.rs b/crates/mummu/src/models/mod.rs index 1a1dbcd..f8d567a 100644 --- a/crates/mummu/src/models/mod.rs +++ b/crates/mummu/src/models/mod.rs @@ -4,7 +4,7 @@ use burn::tensor::{Tensor, backend::Backend}; -use crate::decode::{SamplerOptions, generate_loop, top_k_ids}; +use crate::decode::{SamplerOptions, argmax_id, generate_loop, top_k_ids}; pub mod lfm2; pub mod minilm; @@ -12,6 +12,13 @@ pub mod olmoe; pub mod qwen2; pub mod qwen3; +/// Upper bound on one [`CausalLm::warm_up`] call. A warm-up is a fixed, +/// bounded cost paid off the user's critical path — not a place to spend +/// unbounded GPU time — and the measured curve flattens after ~32 steps +/// (`mummu-bench/tests/warmup_f16.rs`), so this ceiling is 8x the useful +/// depth, not a tuning knob. +pub const MAX_WARM_UP_STEPS: usize = 256; + /// The contract every causal LM in the zoo implements. A new architecture /// (Hermes-class function-caller, Gemma, Qwen3, …) provides its cache type, /// its forward pass, and its EOS check — decoding (greedy, sampled, streamed, @@ -122,4 +129,58 @@ pub trait CausalLm { .map_err(|e| format!("logits readback: {e:?}"))?; crate::import::logit_sanity(&v, expected_vocab).map_err(|e| e.to_string()) } + + /// Pay the **cold-start tax off the user's critical path**: one prefill + /// plus `steps` greedy decode steps on a throwaway cache, discarded. + /// + /// A freshly-started process decodes its first tokens far slower than its + /// steady state — measured on Qwen2.5-1.5B at f16, the first 32 tokens run + /// at 12.5 tok/s against a steady 37.6, and the curve is *flat* from token + /// 33 on (`mummu-bench/tests/warmup_f16.rs`). CubeCL already persists its + /// **autotune** choices to disk across processes, so what is left is + /// per-process kernel compilation and pipeline creation, which the wgpu + /// runtime does not cache anywhere (`CompilationCache` is wired for CUDA + /// and HIP only in cubecl 0.10) — no configuration can carry it, only + /// spending it earlier can. A consumer that opens short agent turns should + /// call this once after `install`/load, beside + /// [`Self::sanity_check`]. + /// + /// Warms the **decode** step, whose kernels are shape-stable (`t == 1`). + /// Prefill kernels are keyed by prompt length, so a caller who cares about + /// TTFT should pass a `probe_ids` of its own typical prompt length rather + /// than a token or two. + /// + /// Returns the number of forwards executed (`steps + 1`). Every step reads + /// its argmax back, exactly as real decoding does — an unsynchronized + /// warm-up would queue work and return before the GPU had run any of it. + fn warm_up( + &self, + probe_ids: &[u32], + steps: usize, + device: &B::Device, + ) -> Result { + assert!(!probe_ids.is_empty(), "warm_up: empty probe prompt"); + assert!(steps >= 1, "warm_up: steps must be >= 1"); + assert!( + steps <= MAX_WARM_UP_STEPS, + "warm_up: {steps} steps exceeds the {MAX_WARM_UP_STEPS} bound" + ); + + let mut cache = self.new_cache(); + let logits = self.forward(probe_ids, 0, &mut cache, device); + let mut next = argmax_id(logits)?; + let mut forwards = 1usize; + for past in (probe_ids.len()..).take(steps) { + let logits = self.forward(&[next], past, &mut cache, device); + next = argmax_id(logits)?; + forwards += 1; + } + + debug_assert_eq!( + forwards, + steps + 1, + "warm_up must run exactly one prefill plus `steps` decode forwards" + ); + Ok(forwards) + } } diff --git a/crates/mummu/src/models/qwen2.rs b/crates/mummu/src/models/qwen2.rs index 770a0b8..76908d4 100644 --- a/crates/mummu/src/models/qwen2.rs +++ b/crates/mummu/src/models/qwen2.rs @@ -691,6 +691,40 @@ mod tests { assert!(loaded.sanity_check(&[1, 2, 3], vocab + 1, &device).is_err()); } + #[test] + fn warm_up_runs_one_prefill_plus_its_steps_and_leaves_the_model_usable() { + let device = Dev::default(); + let cfg = toy_config(); + let loaded = LoadedQwen2:: { + model: build(&cfg, &device), + config: cfg, + tokenizer_config: None, + }; + let forwards = loaded + .warm_up(&[1, 2, 3], 4, &device) + .expect("warm-up runs on a live toy model"); + assert_eq!(forwards, 5, "one prefill plus four decode steps"); + // The warm-up cache is a throwaway: a generation after it starts from + // an empty cache and still decodes (nothing leaked into the model). + let out = loaded + .greedy_generate(&[1, 2, 3], 2, &device) + .expect("decodes after a warm-up"); + assert!(!out.is_empty(), "generation after warm-up produces tokens"); + } + + #[test] + #[should_panic(expected = "exceeds the 256 bound")] + fn warm_up_rejects_an_unbounded_step_count() { + let device = Dev::default(); + let cfg = toy_config(); + let loaded = LoadedQwen2:: { + model: build(&cfg, &device), + config: cfg, + tokenizer_config: None, + }; + let _ = loaded.warm_up(&[1, 2, 3], crate::models::MAX_WARM_UP_STEPS + 1, &device); + } + #[test] fn greedy_generate_respects_max_tokens_bound() { let device = Dev::default(); From 4c671f1c81734e4cebb833ed62ef64caaaa9e732 Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Sun, 9 Aug 2026 07:50:30 -0500 Subject: [PATCH 3/7] parity: f16 is now a verified path, not just a live one - both legs pass first run Every strict gate in the repo ran f32 (or GGUF-dequantized-to-f32): the Candle logits fixture, the llama.cpp GGUF legs, the Ollama greedy leg. The f16 path had liveness only - `real_f16.rs` (no NaN, right VRAM, coherent text) and a one-token f16-vs-f32 agreement in `real_mixed_dtype.rs`. That gap is why the 2026-08-06 flash-attention evaluation could not adopt its one winning quadrant (f16 prefill, -22% @2048): an f16-only numeric fork would have shipped unverified. New `tests/parity_f16.rs` closes it with the honest shape the ROADMAP asked for - llama.cpp on the SAME Q4_K_M file, our side loaded onto `GpuF16` - and both legs passed on the first run, on the 4070 Ti SUPER: qwen2-f16 top-5 ids exact in order [785, 32, 16, 1249, 8420], 24-token greedy byte-identical, max |dlogprob| 2.5284926197284596e-1 qwen3-f16 top-5 ids exact in order [151667, 151644, 151645, 99966, 131545], greedy byte-identical incl. the tokens, max |dlogprob| 3.938617118639698e-1 Both f16 numbers come in BELOW their f32 twins (2.66e-1 / 4.02e-1), so f16 adds nothing measurable on top of the reference's own Q8_K activation-quantization noise - the f32-softmax attention island is doing exactly what it was added for. Own test binary, deliberately: `GpuF16` locks Burn's per-device default dtype policy, so one alias per process is what keeps the numbers real. The only production-adjacent change is a test refactor: the comparator moved out of `parity_gguf.rs` into a shared `tests/gguf_compare/` module and gained explicit `port` + `tolerance` parameters. It sits beside `llama_ref` rather than inside it because `parity_lfm2.rs` uses only the transport half, and an unused pub fn in its binary would be a dead_code warning (no `#![allow]`). f32 legs re-ran unchanged after the move - qwen3 bit-identical at 4.015608155114805e-1, qwen2 at 2.66e-1 with top-5 exact and greedy byte-identical. Fixture note: the Qwen2.5-1.5B Q4_K_M GGUF (1.07 GB) was not in the local cache and was fetched for this run. Verified: fmt clean, clippy --all-targets no warnings, 204 unit tests, cargo build green. Co-Authored-By: Claude Opus 4.8 --- README.md | 5 + ROADMAP.md | 14 +++ crates/mummu/tests/gguf_compare/mod.rs | 138 +++++++++++++++++++++++++ crates/mummu/tests/parity_f16.rs | 106 +++++++++++++++++++ crates/mummu/tests/parity_gguf.rs | 135 +++--------------------- 5 files changed, 276 insertions(+), 122 deletions(-) create mode 100644 crates/mummu/tests/gguf_compare/mod.rs create mode 100644 crates/mummu/tests/parity_f16.rs diff --git a/README.md b/README.md index 2748051..91a31f7 100644 --- a/README.md +++ b/README.md @@ -101,6 +101,11 @@ It exists because two local-first apps — **[laurelane](https://github.com/phys exact in order, 24-token greedy byte-identical, max |Δlogprob| 3.7e-1 — so the MoE router and expert bank are verified against a reference, not just plausible. The MiniLM embedder matches its Candle reference at cosine 0.99999994 (max |Δcomponent| 1.2e-7, `tests/real_minilm.rs`). + **The f16 path is parity-verified too** (`tests/parity_f16.rs`, its own binary because `GpuF16` + locks Burn's per-device dtype policy): the same llama.cpp comparison with our side loaded onto + `GpuF16` passes for Qwen2.5-1.5B and Qwen3-0.6B — top-5 ids exact in order, 24-token greedy + byte-identical, max |Δlogprob| 2.5e-1 / 3.9e-1, *below* the f32 legs' own 2.7e-1 / 4.0e-1. So half + precision is a verified path, not merely a live one. - **Sampling, streaming, cancellation** — temperature / top-k / top-p sampling (deterministic per seed), per-token streaming through a `ControlFlow` callback, and cooperative between-token cancellation; greedy decoding keeps the argmax on-device. diff --git a/ROADMAP.md b/ROADMAP.md index 7cc1766..9dfbe59 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -179,6 +179,20 @@ 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. + *(2026-08-09)* **Prerequisite (a) is DONE — the honest shape, and it passed first run on both + architectures.** New `tests/parity_f16.rs` (own binary: `GpuF16` locks the per-device dtype + policy) runs the SAME strict comparison every other port passes, with our side loaded onto + `GpuF16`: **Qwen2.5-1.5B Q4_K_M** — top-5 ids match llama.cpp **exactly in order** + (785, 32, 16, 1249, 8420), 24-token greedy **byte-identical**, max |Δlogprob| + **2.5284926197284596e-1**; **Qwen3-0.6B Q4_K_M** — top-5 exact in order + (151667, 151644, 151645, 99966, 131545), greedy byte-identical incl. the `` tokens, max + |Δlogprob| **3.938617118639698e-1**. Both f16 numbers are *below* their f32 twins (2.66e-1 / + 4.02e-1), i.e. f16 adds nothing measurable on top of the reference's own Q8_K activation-quant + noise — the f32-softmax island is doing its job. Enabling this cost only a refactor: the + comparator moved out of `parity_gguf.rs` into a shared `tests/gguf_compare/` module (beside + `llama_ref`, which `parity_lfm2.rs` still uses for transport only) and gained explicit `port` + + `tolerance` parameters; the f32 legs re-passed unchanged (qwen3 bit-identical at + 4.015608155114805e-1). What remains for this item is (b) plus a re-measure of the A/B. - [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 diff --git a/crates/mummu/tests/gguf_compare/mod.rs b/crates/mummu/tests/gguf_compare/mod.rs new file mode 100644 index 0000000..5ab0d50 --- /dev/null +++ b/crates/mummu/tests/gguf_compare/mod.rs @@ -0,0 +1,138 @@ +//! The shared quantized-reference comparison used by every GGUF parity leg: +//! llama.cpp running the SAME .gguf file our loader loads, compared by top-k +//! first-forward ids and a byte-identical greedy sequence. +//! +//! Lives beside `llama_ref` rather than inside it because `parity_lfm2.rs` +//! drives the reference server directly (its BF16-vs-safetensors legs are +//! shaped differently) and only needs the transport half. + +use mummu::gguf::GgufFile; +use mummu::models::CausalLm; + +use crate::llama_ref::{LlamaServer, logprobs_at, server_exe}; +/// The shared GGUF prompt for every quantized-reference leg. +pub const PROMPT: &str = "List the first five prime numbers."; +/// Greedy tokens compared byte-for-byte. +pub const MAX_TOKENS: usize = 24; +/// Top-k compared on the first forward. +pub const TOP_K: usize = 5; +/// The reference's activation-quantization noise reshuffles near-ties deep in +/// the top-k (measured: LFM2.5 swaps ranks 4-5), so the strict-order assert +/// covers the top 3 and the rest is set overlap. +pub const STRICT_ORDER_K: usize = 3; +/// Minimum top-k set overlap. +pub const MIN_SET_OVERLAP: usize = 4; + +/// Next free port for a reference server; each leg gets its own so a binary's +/// tests can run concurrently. +pub fn next_port(base: u16) -> u16 { + static NEXT: std::sync::atomic::AtomicU16 = std::sync::atomic::AtomicU16::new(0); + base + NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed) +} + +/// One quantized-reference comparison: `load` builds our model from the GGUF, +/// `render` wraps the prompt in the model's chat template, and `tolerance` +/// bounds |Δlogprob| over the top-k. Generic over the backend — the dense +/// tiers compare on `Gpu`, OLMoE's ~28 GB f32 build only fits the CPU backend, +/// and the f16 leg runs `GpuF16` from its own binary (one dtype alias per +/// process). Panics (test style) on any divergence. +pub fn compare_against_llama_cpp( + tag: &str, + gguf: &std::path::Path, + port: u16, + tolerance: f64, + load: impl FnOnce(&std::path::Path, &burn::tensor::Device) -> M, + render: impl FnOnce(&str) -> String, +) where + B: burn::tensor::backend::Backend, + M: CausalLm, +{ + let exe = server_exe().expect("set MUMMU_LLAMA_SERVER to a llama.cpp llama-server binary"); + + // Tokenizer from the same single file, like everything else on our side. + let f = GgufFile::open(gguf).expect("gguf header parses"); + let tok = mummu::tokenizer::tokenizer_from_gguf(&f).expect("tokenizer from gguf metadata"); + drop(f); + let raw = render(PROMPT); + let ids = tok + .encode(raw.as_str(), false) + .expect("encodes") + .get_ids() + .to_vec(); + assert!(ids.len() >= 8, "rendered prompt suspiciously short"); + + // Reference first: fail with the transport error, not a weights error. + let server = LlamaServer::start(&exe, gguf, port).expect("llama-server starts"); + let reference = server + .greedy_completion(&ids, MAX_TOKENS, 10) + .expect("reference completion"); + assert!( + !reference.steps.is_empty(), + "reference returned no logprob steps" + ); + let ref_top: Vec<(u32, f64)> = reference.steps[0].iter().copied().take(TOP_K).collect(); + assert_eq!( + ref_top.len(), + TOP_K, + "reference returned fewer than top-{TOP_K}" + ); + + let device = burn::tensor::Device::::default(); + let loaded = load(gguf, &device); + let mut cache = loaded.new_cache(); + let logits = loaded + .forward(&ids, 0, &mut cache, &device) + .into_data() + .convert::() + .to_vec::() + .expect("logits readback"); + drop(cache); + + let mut indexed: Vec<(usize, f32)> = logits.iter().copied().enumerate().collect(); + indexed.sort_by(|a, b| b.1.total_cmp(&a.1)); + let our_ids: Vec = indexed[..TOP_K].iter().map(|&(id, _)| id as u32).collect(); + let ref_ids: Vec = ref_top.iter().map(|&(id, _)| id).collect(); + let ours_lp = logprobs_at(&logits, &our_ids); + let max_abs_diff = ours_lp + .iter() + .zip(ref_top.iter()) + .map(|(a, &(_, b))| (a - b).abs()) + .fold(0.0_f64, f64::max); + + let out_ids = loaded + .greedy_generate(&ids, MAX_TOKENS, &device) + .expect("greedy decode"); + let ours = tok.decode(&out_ids, true).expect("decode"); + + eprintln!("[parity/gguf/{tag}] top-{TOP_K} ids ours: {our_ids:?}"); + eprintln!("[parity/gguf/{tag}] top-{TOP_K} ids ref : {ref_ids:?}"); + eprintln!("[parity/gguf/{tag}] max |Δlogprob| vs llama.cpp: {max_abs_diff:e}"); + eprintln!( + "[parity/gguf/{tag}] ours ({} tokens): {ours:?}", + out_ids.len() + ); + eprintln!( + "[parity/gguf/{tag}] llama.cpp : {:?}", + reference.content + ); + + assert_eq!( + &our_ids[..STRICT_ORDER_K], + &ref_ids[..STRICT_ORDER_K], + "top-{STRICT_ORDER_K} ids diverge from llama.cpp on the same quantized file" + ); + let overlap = our_ids.iter().filter(|id| ref_ids.contains(id)).count(); + assert!( + overlap >= MIN_SET_OVERLAP, + "top-{TOP_K} sets overlap only {overlap} (need >= {MIN_SET_OVERLAP}): \ + ours {our_ids:?} vs ref {ref_ids:?}" + ); + assert!( + max_abs_diff <= tolerance, + "logprobs diverge: max |Δ| = {max_abs_diff} > {tolerance}" + ); + let (a, b) = (ours.trim(), reference.content.trim()); + let n = a.len().min(b.len()); + assert!(n >= 8, "outputs too short to compare: {n} chars"); + assert_eq!(&a[..n], &b[..n], "greedy sequences diverge"); +} diff --git a/crates/mummu/tests/parity_f16.rs b/crates/mummu/tests/parity_f16.rs new file mode 100644 index 0000000..4388aa2 --- /dev/null +++ b/crates/mummu/tests/parity_f16.rs @@ -0,0 +1,106 @@ +//! **f16 parity leg** — the trust gate the f16 path never had. +//! +//! Every strict gate in the repo runs f32 (or GGUF-dequantized-to-f32): the +//! Candle logits fixture, the llama.cpp GGUF legs, the Ollama greedy leg. The +//! f16 path was validated only for *liveness* (`real_f16.rs`: no NaN, right +//! VRAM, coherent text) and for agreeing with f32 on ONE token in one process +//! (`real_mixed_dtype.rs`). That is not parity, and it is why the 2026-08-06 +//! flash-attention evaluation could not adopt its one winning quadrant (f16 +//! prefill) — an f16-only numeric fork would have shipped unverified. +//! +//! This leg closes that: llama.cpp runs the SAME Q4_K_M file, our side loads +//! it onto **`GpuF16`** (dequantize once, cast to f16 on load, f32 attention +//! -score island), and the two are compared by the same top-k + byte-identical +//! greedy asserts as every other port. Own test binary because `GpuF16` locks +//! Burn's per-device dtype policy — one alias per process (see +//! `mummu-bench/tests/budget_f16.rs`). +//! +//! ```text +//! MUMMU_GGUF_PATH=path/to/qwen2.5-1.5b-instruct-q4_k_m.gguf \ +//! MUMMU_QWEN3_GGUF_PATH=path/to/qwen3-0.6b-q4_k_m.gguf \ +//! MUMMU_LLAMA_SERVER=path/to/llama-server.exe \ +//! cargo test -p mummu --release --test parity_f16 -- --ignored --nocapture +//! ``` + +mod gguf_compare; +mod llama_ref; + +use std::path::PathBuf; + +use gguf_compare::{compare_against_llama_cpp, next_port}; +use mummu::backend::{GpuF16, inventory}; +use mummu::models::{qwen2, qwen3}; + +/// Max |Δlogprob| for the f16 legs. The f32 legs run at 7.5e-1 against +/// measured 2.66e-1 (Qwen2) / 4.02e-1 (Qwen3); f16 adds its own rounding on +/// top of the reference's Q8_K activation quantization, so this sits one step +/// looser. The primary assert is unchanged and unrelaxed: the top-3 ids in +/// order and the 24-token greedy sequence byte-identical. +const LOGPROB_ABS_TOLERANCE: f64 = 1.5e0; + +/// Port range for this binary; distinct from `parity_gguf`'s and +/// `parity_lfm2`'s so a concurrent run never collides. +const PORT_BASE: u16 = 18501; + +fn env_path(var: &str, what: &str) -> PathBuf { + let p = std::env::var_os(var) + .map(PathBuf::from) + .unwrap_or_else(|| panic!("set {var} to {what}")); + assert!(p.is_file(), "{var} is not a file: {p:?}"); + p +} + +/// Skip (rather than fail) on a machine whose GPU cannot do f16 at all. +fn shader_f16_or_skip(tag: &str) -> bool { + if inventory().any_shader_f16() { + return true; + } + eprintln!("[parity/gguf/{tag}] no SHADER_F16 adapter — skipping"); + false +} + +#[test] +#[ignore = "needs the Qwen2.5 Q4_K_M GGUF (MUMMU_GGUF_PATH), llama-server \ + (MUMMU_LLAMA_SERVER), and a SHADER_F16 GPU"] +fn qwen2_q4_gguf_matches_llama_cpp_in_f16() { + if !shader_f16_or_skip("qwen2-f16") { + return; + } + let gguf = env_path("MUMMU_GGUF_PATH", "the qwen2.5-1.5b-instruct q4_k_m gguf"); + compare_against_llama_cpp( + "qwen2-f16", + &gguf, + next_port(PORT_BASE), + LOGPROB_ABS_TOLERANCE, + |p, d| qwen2::load_from_gguf::(p, d).expect("gguf load checked"), + |user| { + mummu::chat::ChatMl::qwen2().render(&[ + mummu::chat::Turn::system("You are a helpful assistant."), + mummu::chat::Turn::user(user), + ]) + }, + ); +} + +#[test] +#[ignore = "needs a Qwen3 Q4_K_M GGUF (MUMMU_QWEN3_GGUF_PATH), llama-server \ + (MUMMU_LLAMA_SERVER), and a SHADER_F16 GPU"] +fn qwen3_q4_gguf_matches_llama_cpp_in_f16() { + if !shader_f16_or_skip("qwen3-f16") { + return; + } + let gguf = env_path("MUMMU_QWEN3_GGUF_PATH", "a qwen3 q4_k_m gguf"); + compare_against_llama_cpp( + "qwen3-f16", + &gguf, + next_port(PORT_BASE), + LOGPROB_ABS_TOLERANCE, + |p, d| qwen3::load_from_gguf::(p, d).expect("gguf load checked"), + |user| { + mummu::chat::ChatMl::qwen2().render(&[ + mummu::chat::Turn::system("You are a helpful assistant."), + mummu::chat::Turn::user(user), + ]) + }, + ); +} diff --git a/crates/mummu/tests/parity_gguf.rs b/crates/mummu/tests/parity_gguf.rs index 1d54770..bdfafe8 100644 --- a/crates/mummu/tests/parity_gguf.rs +++ b/crates/mummu/tests/parity_gguf.rs @@ -16,20 +16,16 @@ //! cargo test -p mummu --release --test parity_gguf -- --ignored --nocapture //! ``` +mod gguf_compare; mod llama_ref; use std::path::PathBuf; -use llama_ref::{LlamaServer, logprobs_at}; +use gguf_compare::{compare_against_llama_cpp, next_port}; use mummu::backend::{Cpu, Gpu}; use mummu::gguf::GgufFile; -use mummu::models::CausalLm; use mummu::models::{lfm2, olmoe, qwen2, qwen3}; -const PROMPT: &str = "List the first five prime numbers."; -const MAX_TOKENS: usize = 24; -const TOP_K: usize = 5; - /// Max |Δlogprob| over the top-k between our load (weights dequantized to f32 /// once, wgpu compute) and llama.cpp on the SAME Q4_K_M file (CPU kernels /// that also quantize the *activations* to Q8_K per integer dot product — the @@ -41,11 +37,9 @@ const TOP_K: usize = 5; /// is the primary assert. const LOGPROB_ABS_TOLERANCE: f64 = 7.5e-1; -/// The reference's activation-quantization noise (see above) reshuffles -/// near-ties deep in the top-k (measured: LFM2.5 swaps ranks 4-5), so the -/// strict-order assert covers the top 3 and the rest is set overlap. -const STRICT_ORDER_K: usize = 3; -const MIN_SET_OVERLAP: usize = 4; +/// Port range for this binary's servers; distinct from `parity_lfm2`'s and +/// `parity_f16`'s so concurrent legs never collide. +const PORT_BASE: u16 = 18481; fn env_path(var: &str, what: &str) -> PathBuf { let p = std::env::var_os(var) @@ -55,117 +49,6 @@ fn env_path(var: &str, what: &str) -> PathBuf { p } -fn next_port() -> u16 { - // This binary's two tests may run concurrently: one port each, in a range - // distinct from parity_lfm2's. - static NEXT_PORT: std::sync::atomic::AtomicU16 = std::sync::atomic::AtomicU16::new(18481); - NEXT_PORT.fetch_add(1, std::sync::atomic::Ordering::Relaxed) -} - -/// One quantized-reference comparison: `load` builds our model from the GGUF, -/// `render` wraps the prompt in the model's chat template. Generic over the -/// backend — the dense tiers compare on `Gpu`; OLMoE's ~28 GB f32 build only -/// fits the CPU backend. Panics (test style) on any divergence. -fn compare_against_llama_cpp( - tag: &str, - gguf: &std::path::Path, - load: impl FnOnce(&std::path::Path, &burn::tensor::Device) -> M, - render: impl FnOnce(&str) -> String, -) where - B: burn::tensor::backend::Backend, - M: CausalLm, -{ - let exe = - llama_ref::server_exe().expect("set MUMMU_LLAMA_SERVER to a llama.cpp llama-server binary"); - - // Tokenizer from the same single file, like everything else on our side. - let f = GgufFile::open(gguf).expect("gguf header parses"); - let tok = mummu::tokenizer::tokenizer_from_gguf(&f).expect("tokenizer from gguf metadata"); - drop(f); - let raw = render(PROMPT); - let ids = tok - .encode(raw.as_str(), false) - .expect("encodes") - .get_ids() - .to_vec(); - assert!(ids.len() >= 8, "rendered prompt suspiciously short"); - - // Reference first: fail with the transport error, not a weights error. - let server = LlamaServer::start(&exe, gguf, next_port()).expect("llama-server starts"); - let reference = server - .greedy_completion(&ids, MAX_TOKENS, 10) - .expect("reference completion"); - assert!( - !reference.steps.is_empty(), - "reference returned no logprob steps" - ); - let ref_top: Vec<(u32, f64)> = reference.steps[0].iter().copied().take(TOP_K).collect(); - assert_eq!( - ref_top.len(), - TOP_K, - "reference returned fewer than top-{TOP_K}" - ); - - let device = burn::tensor::Device::::default(); - let loaded = load(gguf, &device); - let mut cache = loaded.new_cache(); - let logits = loaded - .forward(&ids, 0, &mut cache, &device) - .into_data() - .convert::() - .to_vec::() - .expect("logits readback"); - drop(cache); - - let mut indexed: Vec<(usize, f32)> = logits.iter().copied().enumerate().collect(); - indexed.sort_by(|a, b| b.1.total_cmp(&a.1)); - let our_ids: Vec = indexed[..TOP_K].iter().map(|&(id, _)| id as u32).collect(); - let ref_ids: Vec = ref_top.iter().map(|&(id, _)| id).collect(); - let ours_lp = logprobs_at(&logits, &our_ids); - let max_abs_diff = ours_lp - .iter() - .zip(ref_top.iter()) - .map(|(a, &(_, b))| (a - b).abs()) - .fold(0.0_f64, f64::max); - - let out_ids = loaded - .greedy_generate(&ids, MAX_TOKENS, &device) - .expect("greedy decode"); - let ours = tok.decode(&out_ids, true).expect("decode"); - - eprintln!("[parity/gguf/{tag}] top-{TOP_K} ids ours: {our_ids:?}"); - eprintln!("[parity/gguf/{tag}] top-{TOP_K} ids ref : {ref_ids:?}"); - eprintln!("[parity/gguf/{tag}] max |Δlogprob| vs llama.cpp: {max_abs_diff:e}"); - eprintln!( - "[parity/gguf/{tag}] ours ({} tokens): {ours:?}", - out_ids.len() - ); - eprintln!( - "[parity/gguf/{tag}] llama.cpp : {:?}", - reference.content - ); - - assert_eq!( - &our_ids[..STRICT_ORDER_K], - &ref_ids[..STRICT_ORDER_K], - "top-{STRICT_ORDER_K} ids diverge from llama.cpp on the same quantized file" - ); - let overlap = our_ids.iter().filter(|id| ref_ids.contains(id)).count(); - assert!( - overlap >= MIN_SET_OVERLAP, - "top-{TOP_K} sets overlap only {overlap} (need >= {MIN_SET_OVERLAP}): \ - ours {our_ids:?} vs ref {ref_ids:?}" - ); - assert!( - max_abs_diff <= LOGPROB_ABS_TOLERANCE, - "logprobs diverge: max |Δ| = {max_abs_diff} > {LOGPROB_ABS_TOLERANCE}" - ); - let (a, b) = (ours.trim(), reference.content.trim()); - let n = a.len().min(b.len()); - assert!(n >= 8, "outputs too short to compare: {n} chars"); - assert_eq!(&a[..n], &b[..n], "greedy sequences diverge"); -} - #[test] #[ignore = "needs the Qwen2.5 Q4_K_M GGUF (MUMMU_GGUF_PATH) + llama-server (MUMMU_LLAMA_SERVER)"] fn qwen2_q4_gguf_matches_llama_cpp_on_the_same_file() { @@ -173,6 +56,8 @@ fn qwen2_q4_gguf_matches_llama_cpp_on_the_same_file() { compare_against_llama_cpp( "qwen2", &gguf, + next_port(PORT_BASE), + LOGPROB_ABS_TOLERANCE, |p, d| qwen2::load_from_gguf::(p, d).expect("gguf load checked"), |user| { mummu::chat::ChatMl::qwen2().render(&[ @@ -194,6 +79,8 @@ fn qwen3_q4_gguf_matches_llama_cpp_on_the_same_file() { compare_against_llama_cpp( "qwen3", &gguf, + next_port(PORT_BASE), + LOGPROB_ABS_TOLERANCE, |p, d| qwen3::load_from_gguf::(p, d).expect("gguf load checked"), |user| { mummu::chat::ChatMl::qwen2().render(&[ @@ -214,6 +101,8 @@ fn lfm2_q4_gguf_matches_llama_cpp_on_the_same_file() { compare_against_llama_cpp( "lfm2", &gguf, + next_port(PORT_BASE), + LOGPROB_ABS_TOLERANCE, |p, d| lfm2::load_from_gguf::(p, d).expect("gguf load checked"), |user| mummu::chat::ChatMl::lfm2().render(&[mummu::chat::Turn::user(user)]), ); @@ -250,6 +139,8 @@ fn olmoe_q4_gguf_matches_llama_cpp_on_the_same_file() { compare_against_llama_cpp( "olmoe", &gguf, + next_port(PORT_BASE), + LOGPROB_ABS_TOLERANCE, |p, d| olmoe::load_from_gguf::(p, d).expect("gguf load checked"), move |user| format!("{bos}<|user|>\n{user}\n<|assistant|>\n"), ); From d8d6cabfdae1ff7cfc2e363d97b86d82836ba0a3 Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Sun, 9 Aug 2026 08:08:39 -0500 Subject: [PATCH 4/7] docs: flash attention for f16 prefill is rejected on CORRECTNESS - the new parity leg caught it Closes the ROADMAP item "adopt flash attention for f16 prefill only, once f16 has parity coverage". Prerequisite (a) shipped two commits ago; this run then did (b) and the answer changed from "wait" to "no". The conditional was built exactly as scoped - `use_fused_attention(t, ambient, masked)` gating on `t > 1 && f16 && masked` (the measured quadrant, nothing wider), picking between a fused `tensor::module::attention(.., is_causal: true)` and the existing chain, with a CPU-backend unit test holding the two formulations to each other at four (past, t) pairs. THE WIN REPRODUCED, two runs per arm, f32 as an untouched control: f16 TTFT (36 tok) 24.9 / 24.6 ms -> 21.0 / 20.8 ms -15.6% f16 prefill @ 2048 240.6 / 239.5 ms -> 224.2 / 221.3 ms -7.2% f16 decode 830.8 ms -> 806.7 ms (noise; path not taken) f32 (control) 98.2 -> 98.0 ms, 597.3 -> 598.5 ms, 1.979 -> 1.988 s Then `tests/parity_f16.rs` failed: Qwen2.5-1.5B Q4_K_M on GpuF16 returns NON-FINITE logits through the fused kernel, while the identical weights through the explicit chain are byte-identical to llama.cpp. The 2026-08-06 evaluation assumed the f32 score island survives inside the kernel (`AccumulatorPrecision::Strict(F32)`); it does not, for the very model whose q.k^T overflow motivated that island - the fused path reproduces the pre-island 2026-07-11 NaN. It is also model-dependent: Qwen3-0.6B PASSED the same fused path (top-3 exact, greedy byte-identical, max |dlogprob| 3.9386e-1 against the explicit 3.9386e-1, only the 5th tail id reshuffling). That is exactly what makes it unshippable in a shared leaf function - correct for narrow models, silently NaN for wide ones. Reverted; the tree keeps the explicit chain, and the f16 parity legs re-passed green after the revert (2.5284926197284596e-1 / 3.938617118639698e-1, both greedy byte-identical). What would reopen it: burn 0.22 / wgpu 30, or an upstream fix that makes the kernel's score accumulation genuinely f32 for f16 inputs. Standing lesson recorded in bench/BASELINE.md: run the parity leg BEFORE the A/B. The measurement was never the hard part - and an evaluation without a parity gate behind it would have shipped this. Co-Authored-By: Claude Opus 4.8 --- ROADMAP.md | 28 ++++++++++++++++++++++++++-- bench/BASELINE.md | 22 ++++++++++++++++++++++ 2 files changed, 48 insertions(+), 2 deletions(-) diff --git a/ROADMAP.md b/ROADMAP.md index 9dfbe59..b20e726 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -168,7 +168,7 @@ a benchmark holds/improves its budget; README perf claims link an artifact. 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 +- [x] **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, @@ -192,7 +192,31 @@ a benchmark holds/improves its budget; README perf claims link an artifact. comparator moved out of `parity_gguf.rs` into a shared `tests/gguf_compare/` module (beside `llama_ref`, which `parity_lfm2.rs` still uses for transport only) and gained explicit `port` + `tolerance` parameters; the f32 legs re-passed unchanged (qwen3 bit-identical at - 4.015608155114805e-1). What remains for this item is (b) plus a re-measure of the A/B. + 4.015608155114805e-1). + *(2026-08-09, same run) **Item CLOSED — rejected on correctness, and the parity leg built two + hours earlier is what caught it.** With (a) in hand, (b) was implemented exactly as scoped: a + `use_fused_attention(t, ambient, masked)` gate (`t > 1 && f16 && masked` — the measured + quadrant, nothing wider) picking between a new `attend_fused` and the existing chain, extracted + as `attend_explicit`; plus a CPU-backend unit test holding the two formulations to each other at + four `(past, t)` pairs (they agree to 1e-5, so the bottom-right causal alignment and the implicit + `1/sqrt(head_dim)` scale are right) and a gate test pinning the quadrant. **The A/B reproduced + the win** — two runs per arm, same session, idle-ish card, f32 rows as an untouched control: + f16 TTFT 24.9/24.6 → **21.0/20.8 ms (−15.6 %)**, f16 prefill@2048 240.6/239.5 → **224.2/221.3 ms + (−7.2 %)**, f16 decode unchanged within noise (830.8 → 806.7 ms, and it cannot take the path by + construction), f32 TTFT 98.2 → 98.0 ms / prefill 597.3 → 598.5 ms. **Then the f16 parity gate + failed: Qwen2.5-1.5B Q4_K_M on `GpuF16` returns NON-FINITE logits through the fused kernel** + (`logprobs_at`'s finiteness assert), while the same weights through the explicit chain are fine + and llama.cpp-identical. So the 2026-08-06 reading that "the f32 island survives inside the + kernel (`AccumulatorPrecision::Strict(F32)`)" does **not** hold in practice for the very model + whose q·kᵀ overflow motivated the island — the fused path reproduces the pre-island 2026-07-11 + NaN. It is also model-dependent: **Qwen3-0.6B passed** through the same fused path (top-3 exact, + greedy byte-identical, max |Δlogprob| 3.9386004209988457e-1 vs the explicit 3.938617118639698e-1 + — only the 5th tail id reshuffled), which is exactly what makes it unshippable as a rule in a + shared leaf function: it would be correct for narrow models and silently NaN for wide ones. + Reverted; the tree keeps the explicit chain. What would reopen this: burn 0.22 / wgpu 30 (the P0 + item — wgpu 30 lifts `SHADER_F16` to WGSL, changing which kernels are candidates at all), or an + upstream fix that makes the kernel's score accumulation genuinely f32 for f16 inputs. Re-run + `tests/parity_f16.rs` FIRST next time; the measurement was never the hard part.* - [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 diff --git a/bench/BASELINE.md b/bench/BASELINE.md index eac4e81..1e4fba9 100644 --- a/bench/BASELINE.md +++ b/bench/BASELINE.md @@ -173,6 +173,28 @@ Notes 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-08-09: **the f16-prefill quadrant was implemented, re-measured, and rejected — this time on + CORRECTNESS.** With the f16 parity legs in place (`tests/parity_f16.rs`, shipped the same run), the + conditional was built exactly as scoped (`t > 1 && f16 && masked` picks the fused kernel; every + other call keeps the explicit chain) and **the win reproduced**, two runs per arm, f32 as an + untouched control: + + | metric | explicit (control) | fused f16 prefill | delta | + | --- | --- | --- | --- | + | f16 TTFT (36 tok) | 24.9 / 24.6 ms | 21.0 / 20.8 ms | **−15.6 %** | + | f16 prefill @ 2048 | 240.6 / 239.5 ms | 224.2 / 221.3 ms | **−7.2 %** | + | f16 decode | 868.7 / 792.8 ms | 822.3 / 791.1 ms | noise (path not taken) | + | f32 TTFT / prefill@2048 / decode | 98.2 ms / 597.3 ms / 1.9787 s | 98.0 ms / 598.5 ms / 1.9877 s | control, unmoved | + + Then the parity gate killed it: **Qwen2.5-1.5B Q4_K_M on `GpuF16` returns non-finite logits through + the fused kernel**, while the identical weights through the explicit chain are llama.cpp-identical. + The 2026-08-06 note above assumed the f32 score island survives inside the kernel + (`AccumulatorPrecision::Strict(F32)`); it does not, for the very model whose q·kᵀ overflow motivated + that island — the fused path reproduces the pre-island 2026-07-11 NaN. And it is model-dependent: + **Qwen3-0.6B passed** the same fused path (top-3 exact, greedy byte-identical, max |Δlogprob| + 3.9386e-1 vs the explicit 3.9386e-1), which is precisely what makes it unshippable in a shared leaf — + correct for narrow models, silently NaN for wide ones. Reverted. The standing lesson: **run + `tests/parity_f16.rs` before the A/B, not after** — the measurement was never the hard part. - 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` From c92b75fe9a777390ac11aa221a0e184543c0be3f Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Sun, 9 Aug 2026 08:14:19 -0500 Subject: [PATCH 5/7] tune: give consumers a way to throw away a bad autotune cache Route (a) of the hazard this run measured: CubeCL persists its kernel picks to disk and reloads them forever, with no invalidation and no re-tune trigger, so a pick made while the machine was busy is indistinguishable from a good one. Measured today (bench/BASELINE.md): a tune taken during a contended moment cost 21-27% of f16 decode throughput in EVERY subsequent process, while the f32 picks from the same moment were fine - silent, partial, permanent. New `mummu::tune`: - `autotune_cache_dir()` - where CubeCL will persist picks, read out of the very config CubeCL itself discovers (`CubeClRuntimeConfig::get().autotune.cache.root()` + the `autotune` segment `CacheOption::name` adds), so the path is right by construction instead of a hardcoded copy of the discovery rule that would drift the first time cubecl changes it. - `autotune_cache_report()` - files + bytes; an absent cache is empty, not an error. - `clear_autotune_cache()` - removes it and returns what it removed: the "re-tune GPU kernels" action a settings UI needs. Bounded and fail-loud: a tree deeper than 8 levels or wider than 65 536 files is `TuneError::Implausible` rather than a long walk or a wide delete, and every path the module touches ends in the `autotune` segment by construction (asserted before any removal), so a misconfigured root cannot widen the delete. Documented honestly - it takes effect on the NEXT process, since a running one has already loaded the cache into memory. One new direct dependency, `cubecl-runtime 0.10` - the same version burn 0.21 already resolves through burn-cubecl, so it is feature-unified and costs no compile time. Same precedent as the direct `burn-store` and `wgpu` handles. Proof: 5 unit tests plus a REAL-GPU test (`tests/real_autotune_cache.rs`) that clears the cache, runs four 512-square matmuls with readbacks, finds 3 files / 8 303 bytes written to exactly the reported directory, then clears them and confirms empty. It touches `crates/mummu/target/autotune` and never the bench crate's - CubeCL's root is the walk-up from the process CWD, so the recorded benchmark numbers keep their own tuning undisturbed. Still open on the ROADMAP item: (b) pinning the cache to a Mummu-owned location via `RuntimeConfig::set`, and (c) detecting a bad tune automatically rather than only exposing the repair. Verified: fmt clean, clippy --all-targets no warnings, 209 unit tests (5 new), cargo build green. Co-Authored-By: Claude Opus 4.8 --- Cargo.lock | 1 + Cargo.toml | 7 + README.md | 9 + ROADMAP.md | 20 ++ crates/mummu/Cargo.toml | 1 + crates/mummu/src/lib.rs | 1 + crates/mummu/src/tune.rs | 243 ++++++++++++++++++++++ crates/mummu/tests/real_autotune_cache.rs | 102 +++++++++ 8 files changed, 384 insertions(+) create mode 100644 crates/mummu/src/tune.rs create mode 100644 crates/mummu/tests/real_autotune_cache.rs diff --git a/Cargo.lock b/Cargo.lock index 6d25543..4ca1880 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3634,6 +3634,7 @@ dependencies = [ "burn", "burn-flex", "burn-store", + "cubecl-runtime", "half", "hf-chat-template", "once_cell", diff --git a/Cargo.toml b/Cargo.toml index c31d156..a46b3c3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -43,6 +43,13 @@ burn-store = { version = "0.21", default-features = false, features = [ "pytorch", ] } +# Direct handle on cubecl-runtime (the same version burn 0.21 resolves through +# burn-cubecl) to read the autotune-cache location out of the very config +# CubeCL itself discovers — `crate::tune` reports and clears it. Feature- +# unified, so this adds no compile time; hardcoding the discovery rule instead +# would drift from cubecl's the first time it changes. +cubecl-runtime = "0.10" + # Direct handle on wgpu (the same version Burn resolves) for the cheap startup # adapter probe: enumerate GPUs to decide GPU vs CPU before building a backend. wgpu = "29" diff --git a/README.md b/README.md index 91a31f7..bf2b043 100644 --- a/README.md +++ b/README.md @@ -161,6 +161,15 @@ It exists because two local-first apps — **[laurelane](https://github.com/phys decode 36.8 tok/s, prefill@2048 241 ms** (~3.6 GiB runner) — recorded with budgets in [bench/BASELINE.md](bench/BASELINE.md), enforced by opt-in regression gates (`mummu-bench/tests/budget{,_f16,_cpu,_moe}.rs`, one dtype alias per process). +- **Autotune-cache control** — CubeCL benchmarks each kernel once and persists the winner to disk, so + later processes start warm; but the cache has no invalidation, so a pick made while the machine was + busy is believed forever (measured 2026-08-09: a tune taken during a contended moment cost 21–27% of + f16 decode in every subsequent process, silently). `mummu::tune` is the repair a settings UI needs: + `autotune_cache_dir()` reports where the picks live — read out of the very config CubeCL discovers, + not a hardcoded copy of the rule — `autotune_cache_report()` measures it, and + `clear_autotune_cache()` removes it so the next launch re-tunes. Bounded and fail-loud (an + implausibly deep or wide tree is an error, never a wide delete). Proven on the real GPU + (`tests/real_autotune_cache.rs`). - **Model management** — `ModelManager` gives settings UIs the whole lifecycle over a declarative model catalog (`registry::ModelSpec`): install with per-chunk download progress, `is_installed`, per-model disk usage, and traversal-safe removal; model switching rides `ModelSlot`. diff --git a/ROADMAP.md b/ROADMAP.md index b20e726..55a243f 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -294,6 +294,26 @@ a benchmark holds/improves its budget; README perf claims link an artifact. validate on load — time one warm-up burst against a recorded expectation and clear the cache when it reads far low, which the `warm_up` API already has the shape for. Gate any of them on `bench/BASELINE.md` like everything else. *(2026-08-09, measured this run.)* + *(2026-08-09, same run) **Route (a) shipped: `mummu::tune`.*** `autotune_cache_dir()` reports + where CubeCL will persist picks — read out of the very config CubeCL discovers + (`CubeClRuntimeConfig::get().autotune.cache.root()` joined with the `autotune` segment + `CacheOption::name` adds), so the path is right by construction rather than by a hardcoded copy + of the discovery rule; `autotune_cache_report()` measures it (files + bytes, absent = empty, not + an error); `clear_autotune_cache()` removes it and returns what it removed, which is the + "re-tune GPU kernels" action a consumer's settings UI needs. Bounded and fail-loud throughout — + a tree deeper than 8 levels or wider than 65 536 files is `TuneError::Implausible` rather than a + long walk or a wide delete, and every path this module touches ends in the `autotune` segment by + construction (asserted), so a misconfigured root cannot widen the delete. Documented honestly: + it takes effect on the **next** process (a running one has already loaded the cache into memory). + One new direct dep, `cubecl-runtime 0.10`, the same version burn 0.21 already resolves through + burn-cubecl — feature-unified, zero compile-time cost, the `burn-store`/`wgpu` precedent. + 5 unit tests + a REAL-GPU proof (`tests/real_autotune_cache.rs`): it clears the cache, runs four + 512² matmuls with readbacks, and finds **3 files / 8 303 bytes** written to exactly the reported + directory, then clears them again and confirms empty. That test deliberately touches + `crates/mummu/target/autotune` and never the bench crate's — CubeCL's root is the walk-up from + the process CWD, so the recorded benchmark numbers keep their own tuning. Still open here: + (b) pinning the cache to a Mummu-owned location via `RuntimeConfig::set`, and (c) detecting a + bad tune automatically rather than exposing the repair. ## Phases diff --git a/crates/mummu/Cargo.toml b/crates/mummu/Cargo.toml index ad53fe2..1efcce0 100644 --- a/crates/mummu/Cargo.toml +++ b/crates/mummu/Cargo.toml @@ -31,6 +31,7 @@ sha2 = { workspace = true } once_cell = { workspace = true } burn-flex = { workspace = true } burn-store = { workspace = true } +cubecl-runtime = { workspace = true } [dev-dependencies] # Test-only reference renderer for the imported `chat_template` Jinja: claims diff --git a/crates/mummu/src/lib.rs b/crates/mummu/src/lib.rs index 7983ede..90bab04 100644 --- a/crates/mummu/src/lib.rs +++ b/crates/mummu/src/lib.rs @@ -25,3 +25,4 @@ pub mod registry; pub mod template; pub mod tok_config; pub mod tokenizer; +pub mod tune; diff --git a/crates/mummu/src/tune.rs b/crates/mummu/src/tune.rs new file mode 100644 index 0000000..884dc18 --- /dev/null +++ b/crates/mummu/src/tune.rs @@ -0,0 +1,243 @@ +//! The **autotune cache**: where CubeCL persists its kernel picks, and how to +//! throw them away. +//! +//! CubeCL benchmarks several implementations of each kernel the first time it +//! sees one and writes the winner to disk, keyed by (device, kernel, +//! checksum). Later processes load those picks instead of re-tuning, which is +//! what makes a cold start bearable — but the cache has **no invalidation and +//! no re-tune trigger**, so a pick made while the machine was busy is +//! indistinguishable from a good one and is believed forever. Measured on +//! 2026-08-09 (see `bench/BASELINE.md`): a tune that happened during a +//! contended moment cost **21–27 % of f16 decode throughput** in every +//! subsequent process, while the f32 picks from the same moment were +//! unaffected — so the symptom is silent, partial, and permanent. +//! +//! This module is the repair: report where the cache lives and delete it, so a +//! consumer can offer a "re-tune GPU kernels" action instead of shipping a bad +//! tune to a user forever. It reads the same configuration CubeCL reads +//! (`[cubecl.autotune] cache` from the `cubecl.toml` / `burn.toml` discovered +//! by walking up from the process CWD), so the path is right by construction +//! rather than by convention. + +use std::path::{Path, PathBuf}; + +use cubecl_runtime::config::{CubeClRuntimeConfig, RuntimeConfig}; + +/// Directory segment CubeCL appends to the configured cache root for autotune +/// data (`CacheOption::name("autotune")`). Load-bearing for safety: every path +/// this module deletes ends in this segment, so a misconfigured root can never +/// turn a clear into a wider delete. +const AUTOTUNE_SEGMENT: &str = "autotune"; + +/// Bound on a cache walk. The layout is +/// `/autotune///.json.log` — a few hundred +/// files on a normal machine. A runaway count means the root is pointing +/// somewhere it should not, and is an error rather than a long walk. +const MAX_CACHE_FILES: usize = 65_536; +/// Bound on recursion depth for the same reason (the real layout is 3 deep). +const MAX_CACHE_DEPTH: usize = 8; + +/// What went wrong inspecting or clearing the cache. +#[derive(Debug, thiserror::Error)] +pub enum TuneError { + /// The cache directory could not be read or removed. + #[error("autotune cache i/o at {path}: {message}")] + Io { + /// The path being read or removed. + path: PathBuf, + /// The underlying OS error. + message: String, + }, + /// The cache tree is larger or deeper than any real autotune cache, which + /// means the configured root is not what we think it is. Refused rather + /// than walked (or deleted). + #[error("autotune cache at {path} is implausible ({what}) — refusing to touch it")] + Implausible { + /// The configured cache directory. + path: PathBuf, + /// Which bound was exceeded. + what: String, + }, +} + +/// Where the autotune cache lives, and how much of it there is. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TuneCacheReport { + /// The `/autotune` directory CubeCL writes to. + pub dir: PathBuf, + /// Number of cache files found (0 when the cache does not exist yet). + pub files: usize, + /// Total bytes of those files. + pub bytes: u64, +} + +impl TuneCacheReport { + /// Has anything been tuned and persisted yet? + #[must_use] + pub fn is_empty(&self) -> bool { + self.files == 0 + } +} + +/// The directory CubeCL persists autotune picks to, per the configuration it +/// would itself discover. +/// +/// **Reads the global config**, which initializes it if no one has yet — the +/// same one-shot singleton `RuntimeConfig::set` writes to. A consumer that +/// wants to `set` a custom config must do so *before* calling this (and before +/// building any backend), exactly as CubeCL requires. +#[must_use] +pub fn autotune_cache_dir() -> PathBuf { + let dir = CubeClRuntimeConfig::get() + .autotune + .cache + .root() + .join(AUTOTUNE_SEGMENT); + debug_assert!( + dir.ends_with(AUTOTUNE_SEGMENT), + "the cache dir must always end in the autotune segment" + ); + dir +} + +/// Measure the persisted cache without changing it. A missing directory is +/// not an error — it means nothing has been tuned yet. +pub fn autotune_cache_report() -> Result { + let dir = autotune_cache_dir(); + let (files, bytes) = measure(&dir, 0)?; + Ok(TuneCacheReport { dir, files, bytes }) +} + +/// Delete the persisted autotune cache, returning what was removed. +/// +/// The next process to run will re-tune from scratch and write fresh picks — +/// **the next one**, not this one: a running process has already loaded the +/// cache into memory and will keep using and re-writing it, so a consumer +/// should treat this as "re-tune on next launch" (or call it before building a +/// backend). Idempotent: clearing an absent cache reports zero and succeeds. +pub fn clear_autotune_cache() -> Result { + let report = autotune_cache_report()?; + assert!( + report.dir.ends_with(AUTOTUNE_SEGMENT), + "refusing to remove a directory that is not the autotune cache: {:?}", + report.dir + ); + if !report.dir.exists() { + debug_assert!(report.is_empty(), "an absent cache cannot hold files"); + return Ok(report); + } + std::fs::remove_dir_all(&report.dir).map_err(|e| TuneError::Io { + path: report.dir.clone(), + message: e.to_string(), + })?; + Ok(report) +} + +/// Bounded recursive walk: `(file count, total bytes)` under `dir`. +fn measure(dir: &Path, depth: usize) -> Result<(usize, u64), TuneError> { + if depth > MAX_CACHE_DEPTH { + return Err(TuneError::Implausible { + path: dir.to_path_buf(), + what: format!("deeper than {MAX_CACHE_DEPTH} levels"), + }); + } + if !dir.is_dir() { + return Ok((0, 0)); + } + let entries = std::fs::read_dir(dir).map_err(|e| TuneError::Io { + path: dir.to_path_buf(), + message: e.to_string(), + })?; + + let (mut files, mut bytes) = (0usize, 0u64); + for entry in entries { + let entry = entry.map_err(|e| TuneError::Io { + path: dir.to_path_buf(), + message: e.to_string(), + })?; + let path = entry.path(); + if path.is_dir() { + let (f, b) = measure(&path, depth + 1)?; + files += f; + bytes += b; + } else { + files += 1; + bytes += entry.metadata().map(|m| m.len()).unwrap_or(0); + } + if files > MAX_CACHE_FILES { + return Err(TuneError::Implausible { + path: dir.to_path_buf(), + what: format!("more than {MAX_CACHE_FILES} files"), + }); + } + } + Ok((files, bytes)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn the_cache_dir_always_ends_in_the_autotune_segment() { + let dir = autotune_cache_dir(); + assert!( + dir.ends_with(AUTOTUNE_SEGMENT), + "cache dir {dir:?} must end in {AUTOTUNE_SEGMENT}" + ); + // And it must be an absolute path — the roots CubeCL can resolve to + // (CWD, the project target dir, the user config dir, or an explicit + // file path) are all absolute in practice, and a relative one would + // make "clear the cache" depend on the caller's CWD. + assert!(dir.is_absolute(), "cache dir {dir:?} must be absolute"); + } + + #[test] + fn measuring_an_absent_directory_reports_empty() { + let missing = std::env::temp_dir().join("mummu-no-such-autotune-dir-9e3f"); + assert!(!missing.exists(), "fixture path must not exist"); + assert_eq!( + measure(&missing, 0).expect("absent is not an error"), + (0, 0) + ); + } + + #[test] + fn measuring_counts_files_and_bytes_across_nested_dirs() { + let root = std::env::temp_dir().join("mummu-tune-measure-a41c"); + let nested = root.join("0.10.0").join("device-4-0"); + std::fs::create_dir_all(&nested).expect("fixture dirs"); + std::fs::write(nested.join("matmul.json.log"), b"12345").expect("fixture file"); + std::fs::write(nested.join("reduce.json.log"), b"678").expect("fixture file"); + + let (files, bytes) = measure(&root, 0).expect("walks"); + assert_eq!(files, 2, "both nested files counted"); + assert_eq!(bytes, 8, "byte totals summed across directories"); + + std::fs::remove_dir_all(&root).expect("fixture cleanup"); + } + + #[test] + fn measuring_refuses_an_implausibly_deep_tree() { + // Depth is checked before the directory is read, so a synthetic path + // is enough — no need to build a 9-level fixture. + let deep = std::env::temp_dir().join("mummu-tune-depth"); + let err = measure(&deep, MAX_CACHE_DEPTH + 1).expect_err("too deep is an error"); + assert!( + matches!(err, TuneError::Implausible { .. }), + "expected Implausible, got {err:?}" + ); + } + + #[test] + fn clearing_an_absent_cache_is_a_successful_no_op() { + // `clear` on a machine that has never tuned must not error; the real + // dir may or may not exist here, so assert on the shape of the result. + let before = autotune_cache_report().expect("report"); + assert!(before.dir.ends_with(AUTOTUNE_SEGMENT)); + if !before.dir.exists() { + let cleared = clear_autotune_cache().expect("clearing nothing succeeds"); + assert!(cleared.is_empty(), "nothing to clear reports empty"); + } + } +} diff --git a/crates/mummu/tests/real_autotune_cache.rs b/crates/mummu/tests/real_autotune_cache.rs new file mode 100644 index 0000000..0f76bbe --- /dev/null +++ b/crates/mummu/tests/real_autotune_cache.rs @@ -0,0 +1,102 @@ +//! Real-GPU proof for `mummu::tune`: CubeCL really does persist autotune +//! picks to the directory we report, and clearing it really does remove them. +//! +//! This test deliberately deletes the cache — its own crate's, never the +//! benchmark crate's. CubeCL resolves the cache root by walking up from the +//! process CWD for a `Cargo.toml`, and cargo runs these tests with CWD = +//! `crates/mummu`, so the tree touched here is `crates/mummu/target/autotune` +//! while the recorded benchmark numbers keep tuning out of +//! `crates/mummu-bench/target/autotune`. +//! +//! ```text +//! cargo test -p mummu --release --test real_autotune_cache -- --ignored --nocapture +//! ``` + +use std::time::{Duration, Instant}; + +use burn::tensor::Tensor; +use mummu::backend::{Gpu, use_gpu}; +use mummu::tune::{autotune_cache_report, clear_autotune_cache}; + +/// Square matmul side. Big enough that CubeCL autotunes it (and small enough +/// to stay well inside any GPU's memory). +const N: usize = 512; +/// Matmuls run to provoke tuning. +const ROUNDS: usize = 4; +/// Autotune commits its winner from a worker, so the file may land shortly +/// after the work does. Bounded wait, not a sleep-and-hope. +const WRITE_TRIES: u32 = 40; +const WRITE_INTERVAL: Duration = Duration::from_millis(250); + +#[test] +#[ignore = "needs a real GPU, and deletes this crate's autotune cache"] +fn autotune_picks_are_persisted_where_we_report_and_clearing_removes_them() { + if !use_gpu() { + eprintln!("[tune] no GPU adapter — skipping"); + return; + } + + // Start from a known-empty cache so "files appeared" means this run. + let cleared = clear_autotune_cache().expect("clearing succeeds"); + eprintln!( + "[tune] cache dir {:?} — cleared {} files / {} bytes", + cleared.dir, cleared.files, cleared.bytes + ); + let empty = autotune_cache_report().expect("report after clear"); + assert!( + empty.is_empty(), + "cache must be empty right after a clear, found {} files", + empty.files + ); + assert_eq!(empty.dir, cleared.dir, "the reported dir must be stable"); + + // Provoke autotuning with real GPU work: a few matmuls, each result read + // back so the work is actually executed rather than queued. + let device = burn::tensor::Device::::default(); + for round in 0..ROUNDS { + let a = Tensor::::ones([N, N], &device).mul_scalar(1.0 + round as f32); + let b = Tensor::::ones([N, N], &device); + let sum = a + .matmul(b) + .sum() + .into_data() + .convert::() + .to_vec::() + .expect("matmul readback"); + assert_eq!(sum.len(), 1, "sum reduces to one element"); + assert!(sum[0].is_finite(), "matmul produced a non-finite sum"); + } + + // The winner is committed from a worker thread, so poll (bounded). + let start = Instant::now(); + let mut report = autotune_cache_report().expect("report after work"); + for _ in 0..WRITE_TRIES { + if !report.is_empty() { + break; + } + std::thread::sleep(WRITE_INTERVAL); + report = autotune_cache_report().expect("report while waiting"); + } + eprintln!( + "[tune] after {ROUNDS} matmuls: {} files / {} bytes in {:?} (waited {:?})", + report.files, + report.bytes, + report.dir, + start.elapsed() + ); + assert!( + !report.is_empty(), + "GPU work must leave autotune picks in {:?}", + report.dir + ); + assert!(report.bytes > 0, "persisted picks cannot be zero bytes"); + + // And clearing takes them away again — the "re-tune on next launch" action. + let cleared = clear_autotune_cache().expect("second clear succeeds"); + assert_eq!( + cleared.files, report.files, + "the clear must report exactly what it removed" + ); + let after = autotune_cache_report().expect("report after second clear"); + assert!(after.is_empty(), "cache must be empty after clearing"); +} From 71e58d82fe3c43aa85947e3f27135d9e0a923dd2 Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Sun, 9 Aug 2026 08:16:50 -0500 Subject: [PATCH 6/7] docs: fold 2026-08-09 research - upstream is already answering the autotune-cache hazard Pin watch: burn 0.22 is STILL pre-release (0.22.0-pre.1 the newest tag, 0.21.0 the latest stable), cubecl still 0.11.0-pre.1, tokenizers 0.23.1 current - the P0 migration stays gated and `cargo upgrade --incompatible` offers only the standing wgpu 29->30 pin. Three concrete folds, all tied to items this run touched: 1. Autotune-cache hazard (the item this run opened and half-shipped): cubecl 0.11 adds CUBECL_AUTOTUNE_CACHE (PR #1423, "Disable persistent tune cache option") - an env var that bypasses persistent read AND write and keeps tuning in-memory per process. That is a cleaner route (c) than anything we would build, and it is worth adopting for mummu-bench the moment it lands: this run's 21-27% f16 swing came from exactly the inheritance it disables. cubecl 0.11 also moves autotune scoring from latency to throughput (#1422, #1408), so re-measure the hazard's magnitude after the bump instead of assuming it survives. And CubeCL documents the persisted caches as a SHIPPING artifact ("ship a warm cache with your binary when you know the deployment target"), which is the mirror image of the hazard and the right shape for a consumer with a fixed build target. 2. Warm-up gap: "no configuration can carry it" is a cubecl-0.10-on-wgpu statement, not a permanent one. wgpu ships Device::create_pipeline_cache (gfx-rs/wgpu #5293) explicitly for startup time, so a cubecl-wgpu compilation cache is buildable upstream - the thing to look for on the 0.22 / wgpu-30 bump. Recorded with the verification that produced today's finding: at cubecl 0.10 CompilationCache is constructed in the cuda and hip runtimes only. 3. P9 keep-quantized kernels: a third independent corroboration of the fused-dequant-into-shared-memory/registers split, this time from a shipped WebGPU product (PrismML's 1-bit 27B) rather than a paper. The design question is settled; only the substrate choice is open. Co-Authored-By: Claude Opus 4.8 --- ROADMAP.md | 49 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 49 insertions(+) diff --git a/ROADMAP.md b/ROADMAP.md index 55a243f..ba3d60a 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -42,6 +42,15 @@ https://github.com/Tracel-AI/burn/releases* *(2026-08-03) Pin watch: burn 0.22 i CubeCL tagged **0.11.0-pre.1** the same day (2026-07-29): a frontend mega-refactor (references), a **Metal backend**, a new CPU runtime, tiled layouts, and CUDA stream priority hints — all of which arrive with the burn bump, not before — https://github.com/tracel-ai/cubecl/releases* +*(2026-08-09) Pin watch: **burn 0.22 is still pre-release** (0.22.0-pre.1 remains the newest tag, +0.21.0 the latest stable) and cubecl still 0.11.0-pre.1, so the P0 migration stays gated; tokenizers +0.23.1 and every other direct dep are already current, and `cargo upgrade --incompatible` offers only +the standing wgpu 29→30 pin. Two cubecl-0.11 changelog entries matter to this run's autotune work and +are folded into the items below: **#1423 "Disable persistent tune cache option"** (a +`CUBECL_AUTOTUNE_CACHE` env var that bypasses persistent cache read/write and keeps tuning +in-memory-per-process) and **#1422 "Feat/autotune throughput"** (throughput-based autotuning, beside +#1408 "Peak device throughput") — https://github.com/tracel-ai/cubecl/releases · +https://github.com/tracel-ai/cubecl/pull/1423* ## North Star @@ -276,6 +285,19 @@ a benchmark holds/improves its budget; README perf claims link an artifact. criterion measures 36.8 tok/s today and the same numbers come back on the pre-`cargo update` lockfile, so it is machine state, not a regression (see `bench/BASELINE.md`, incl. the measured f16-vs-f32 host-CPU sensitivity, +9.8 % vs +3.2 % for the same added load).* + *(2026-08-09 research, to re-check when burn 0.22 lands)* The "no configuration can carry it" + finding is a **cubecl-0.10-on-wgpu** statement, not a permanent one, and there are two routes to + removing the residual cost rather than paying it earlier. (1) wgpu itself ships + `Device::create_pipeline_cache` / `PipelineCacheDescriptor` (gfx-rs/wgpu #5293) — a driver-blob + pipeline cache explicitly for "reducing program startup time" — so a cubecl-wgpu compilation + cache is buildable upstream and is the thing to look for on the 0.22 / wgpu-30 bump. (2) CubeCL's + documentation already describes **shipping a warm compilation + autotune cache with the binary** + for a known deployment target, which for a consumer like Nanna (fixed Tauri build, known GPU + classes) converts the cold start into a build-time cost. Verify the wiring before believing + either: at cubecl 0.10 `CompilationCache` is constructed in the **cuda and hip** runtimes only + (read from source this run), which is exactly why `[cubecl.compilation] cache` is inert on our + path. — https://github.com/gfx-rs/wgpu/issues/5293 · + https://docs.rs/wgpu/latest/wgpu/struct.PipelineCache.html - [ ] **A stale autotune cache is permanent, silent, and cost 21–27 % of f16 decode** — found while closing the item above, and it is the reason "just persist the autotune cache" is not a free win for consumers. This run's FIRST budget run happened on a contended machine (9.1 tok/s f32, failing @@ -314,6 +336,23 @@ a benchmark holds/improves its budget; README perf claims link an artifact. the process CWD, so the recorded benchmark numbers keep their own tuning. Still open here: (b) pinning the cache to a Mummu-owned location via `RuntimeConfig::set`, and (c) detecting a bad tune automatically rather than exposing the repair. + *(2026-08-09 research)* **Upstream is moving on both halves, and it arrives with the burn 0.22 + bump — re-scope (b)/(c) then, don't hand-roll them now.** cubecl 0.11 adds + **`CUBECL_AUTOTUNE_CACHE`** (PR #1423, "Disable persistent tune cache option"): an env var that + bypasses persistent read *and* write and keeps tuning `in_memory_cache` per process. That is a + cleaner (c) than anything we would build — a consumer that suspects a bad tune can run once with + persistence off and compare, and a *test* harness can opt out entirely so a benchmark never + inherits another run's picks (worth adopting for `mummu-bench` the moment it lands: this run's + 21–27 % f16 swing came from exactly that inheritance). Also in 0.11: **#1422 "Feat/autotune + throughput"** + **#1408 "Peak device throughput"**, i.e. autotune scoring moves from latency to + throughput against a measured device peak — plausibly *more* noise-robust, so re-measure the + hazard's magnitude after the bump rather than assuming it persists. And CubeCL's own + documentation frames the persisted caches as a **shipping** artifact — "ship a warm cache with + your binary when you know the deployment target, so the cold-start cost is paid once at build + time" — which is the mirror image of this item and worth evaluating for consumers with a fixed + target (Nanna's Tauri build): a *good* tune shipped deliberately, rather than whatever the user's + first busy minute produced. — https://github.com/tracel-ai/cubecl/pull/1423 · + https://github.com/tracel-ai/cubecl ## Phases @@ -1393,6 +1432,16 @@ The VRAM lever the P6 planner pulls to make the largest useful model fit the use 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 + *(2026-08-09 research)* Third independent corroboration, this time from a shipped WebGPU product + rather than a paper: PrismML's 1-bit 27B WebGPU runner describes the same split — hand-written + WGSL matmul kernels with **fused dequantization in shared memory, weights never materialized as + fp16 in VRAM**. Three sources (LlamaWeb on WebGPU, Marlin-class CUDA int4 kernels, this) now + agree on the same shape, so the design question for our P9 kernel is settled and only the + *substrate* choice is open (hand-written CubeCL kernel vs the CubeCL quantization primitives in + the item below). Also worth watching for the KV half above: llama.cpp's TurboQuant discussion + (#20969) is the current state of extreme KV-cache quantization — + https://essamamdani.com/blog/prismml-bonsai-27b-1-bit-27b-model-runs-phone-webgpu-july-2026 · + https://github.com/ggml-org/llama.cpp/discussions/20969 - [ ] 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 8b36856a570315676ee893ca9f5f91cd4b4acaf6 Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Sun, 9 Aug 2026 08:20:55 -0500 Subject: [PATCH 7/7] plan: pick the highest float precision that actually fits one adapter The float half of P6's "precision selection". Everything it needs already existed on both sides - a checkpoint's config.json shape, and `backend::inventory()`'s per-adapter VRAM + SHADER_F16 - but a consumer still had to guess between `Gpu` and `GpuF16` by hand. `mummu::plan::pick_precision(&ModelShape, &DeviceBudget) -> Option` returns the HIGHEST precision that fits, or `None`, which is the honest "no float tier fits; this needs quantization or several devices" rather than a silently-worse tier. Supporting types keep callers out of the arithmetic: `ModelShape::from_decoder` takes the config.json numbers and derives KV geometry itself; `DeviceBudget::from_adapter` reads straight off an enumerated adapter and returns `None` when `vram_bytes` is unknown (every non-Windows adapter today) rather than guessing a budget; `Fit` carries projected/usable bytes and `headroom_bytes()` - already the shape the `plan`/`doctor` introspection item will render. Two constants carry the judgement, calibrated against bench/BASELINE.md rather than first principles: OVERHEAD_BYTES (1 GiB for activations/workspaces/CubeCL pools - the residual between measured runner VRAM and weights+KV) and USABLE_VRAM_FRACTION (0.75, because the reference box runs 3.5-6.5 GiB of desktop ambient on the same card and a plan that ignores it fails at load, not slowly). 7 unit tests pin the decisions to real hardware and real models rather than to the formula: - Qwen2.5-1.5B projects 7.0 GiB f32 / 3.9 GiB f16 against the measured 8.0 / 3.6. - The 15.7 GiB reference card gets f32; an 8 GiB card gets f16. - The dev box's own DX12 rows (same card, no SHADER_F16) never get an f16 plan - and do get f32 once the card is big enough. - A 64k context pushes a 12 GiB card from f32 down to f16. - OLMoE-1B-7B on a 16 GiB card returns None, matching what bench/BASELINE.md records as "out of reach until keep-quantized VRAM". Stays [ ] on the ROADMAP for the int8/int4 tiers, which extend `Precision` downward when P9 lands. Verified: fmt clean, clippy --all-targets no warnings, 216 unit tests (7 new), cargo build green. Co-Authored-By: Claude Opus 4.8 --- README.md | 9 ++ ROADMAP.md | 20 +++ crates/mummu/src/lib.rs | 1 + crates/mummu/src/plan.rs | 333 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 363 insertions(+) create mode 100644 crates/mummu/src/plan.rs diff --git a/README.md b/README.md index bf2b043..8b86963 100644 --- a/README.md +++ b/README.md @@ -161,6 +161,15 @@ It exists because two local-first apps — **[laurelane](https://github.com/phys decode 36.8 tok/s, prefill@2048 241 ms** (~3.6 GiB runner) — recorded with budgets in [bench/BASELINE.md](bench/BASELINE.md), enforced by opt-in regression gates (`mummu-bench/tests/budget{,_f16,_cpu,_moe}.rs`, one dtype alias per process). +- **Precision selection** — `mummu::plan::pick_precision` answers "which dtype fits this card?" from + numbers the crate already has: a model's `config.json` shape on one side, `backend::inventory()`'s + per-adapter VRAM and `SHADER_F16` on the other. It returns the **highest** precision that fits + (f32 before f16) with the projected and usable byte counts behind the decision, `None` when no float + tier fits — the honest "this needs quantization or several devices", never a silently-worse tier — + and never plans f16 on an adapter that doesn't advertise it. Its overhead and headroom constants are + calibrated against [bench/BASELINE.md](bench/BASELINE.md), and its tests pin the decisions to real + hardware: the 15.7 GiB reference card takes Qwen2.5-1.5B in f32, an 8 GiB card in f16, a 64k context + forces a 12 GiB card down to f16, and OLMoE-1B-7B on 16 GiB reports no fit. - **Autotune-cache control** — CubeCL benchmarks each kernel once and persists the winner to disk, so later processes start warm; but the cache has no invalidation, so a pick made while the machine was busy is believed forever (measured 2026-08-09: a tune taken during a contended moment cost 21–27% of diff --git a/ROADMAP.md b/ROADMAP.md index ba3d60a..2079a74 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -1225,6 +1225,26 @@ that fits the model AND uses every device to the fullest. `Wgpu`; drop to int8/int4 (P9) when f16 still won't fit. *(2026-07-11) The f16 backend itself is now **fully validated** (all 3 claims — see the islands item below); what remains here is the *picking* logic, which rides the placement-plan item + P9.* + *(2026-08-09) **The float half of the picking logic shipped: `mummu::plan`.*** + `pick_precision(&ModelShape, &DeviceBudget) -> Option` returns the **highest** precision that + fits one adapter, or `None` — which is the honest answer "no float tier fits; this needs + quantization or a multi-device plan", never a silently-worse tier. `ModelShape::from_decoder` + takes the `config.json` numbers (params, layers, kv heads, head_dim, context) and derives the KV + geometry itself; `DeviceBudget::from_adapter` reads straight off `backend::inventory()` and + returns `None` when `vram_bytes` is unknown (every non-Windows adapter today) rather than + guessing; `Fit` carries the projected and usable byte counts plus `headroom_bytes()`, which is + already the shape the `plan`/`doctor` introspection item will render. Two constants carry the + judgement and are calibrated against `bench/BASELINE.md` rather than first principles: + `OVERHEAD_BYTES` (1 GiB for activations/workspaces/CubeCL pools — the residual between measured + runner VRAM and weights+KV, ~0.5–1.8 GiB depending on dtype) and `USABLE_VRAM_FRACTION` (0.75, + because the reference box runs 3.5–6.5 GiB of desktop ambient on the same card and a plan that + ignores it fails at load). 7 unit tests pin the decisions against real hardware and real models: + Qwen2.5-1.5B projects 7.0 GiB f32 / 3.9 GiB f16 against the measured 8.0 / 3.6; the 15.7 GiB + reference card gets f32 and an 8 GiB card gets f16; the dev box's own **DX12 rows (no + `SHADER_F16`) never get an f16 plan**; a 64k context pushes a 12 GiB card from f32 down to f16; + and OLMoE-1B-7B on a 16 GiB card returns `None`, matching what `bench/BASELINE.md` records as + "GPU is out of reach until keep-quantized VRAM (P9)". Item stays `[ ]` for the int8/int4 tiers, + which extend `Precision` downward once P9 lands. - [x] **f16 mixed-precision islands** — Qwen2.5-1.5B in pure f16 NaNs out (overflow in the softmax/RmsNorm/logit reductions; f16 max is 65 504). Keep weights + matmuls f16 but compute the numerically hot reductions (attention softmax, RmsNorm accumulation, final logits) in f32, then diff --git a/crates/mummu/src/lib.rs b/crates/mummu/src/lib.rs index 90bab04..8bf0e82 100644 --- a/crates/mummu/src/lib.rs +++ b/crates/mummu/src/lib.rs @@ -19,6 +19,7 @@ pub mod import; pub mod manage; pub mod models; pub mod nn; +pub mod plan; pub mod registry; /// Render a checkpoint's own imported chat template (feature `jinja-template`). #[cfg(feature = "jinja-template")] diff --git a/crates/mummu/src/plan.rs b/crates/mummu/src/plan.rs new file mode 100644 index 0000000..bc7c325 --- /dev/null +++ b/crates/mummu/src/plan.rs @@ -0,0 +1,333 @@ +//! **Precision selection** — the first piece of the P6 hardware planner: +//! given a model's shape and one adapter's real capabilities, which float +//! precision is the *highest* one that still fits? +//! +//! Deliberately narrow. This is arithmetic over numbers the rest of the crate +//! already produces ([`crate::backend::inventory`] for VRAM and `SHADER_F16`, +//! a checkpoint's `config.json` for parameter count and cache geometry), and +//! it decides exactly one thing: `Gpu` (f32) or `GpuF16`. Layer placement, +//! multi-GPU sharding and CPU spill are separate ROADMAP items; quantized +//! tiers arrive with P9 and will extend [`Precision`] rather than reshape this. +//! +//! The model is calibrated against measurements, not first principles. +//! Qwen2.5-1.5B on the reference card measures ~8.0 GiB of runner VRAM in f32 +//! and ~3.6 GiB in f16 (`bench/BASELINE.md`) — roughly weights + KV cache plus +//! a fixed overhead for activations, workspaces and CubeCL's memory pools. +//! [`OVERHEAD_BYTES`] is that fixed term, and [`Fit::projected_bytes`] is the +//! whole model; both are honest approximations whose job is to keep a plan on +//! the right side of a cliff, not to predict allocator behaviour to the byte. + +use crate::backend::GpuAdapter; + +/// Fixed VRAM a live runner needs beyond weights and KV cache: activations, +/// matmul workspaces, and CubeCL's memory pools. Derived from the reference +/// measurements — Qwen2.5-1.5B (1.54 G params) reads ~8.0 GiB runner VRAM in +/// f32 against 6.2 GiB of weights, and ~3.6 GiB in f16 against 3.1 GiB of +/// weights, so the residual is ~0.5-1.8 GiB depending on dtype. 1 GiB is the +/// middle of that band and errs toward *not* promising a fit. +pub const OVERHEAD_BYTES: u64 = 1 << 30; + +/// Fraction of an adapter's VRAM a plan may claim. The rest is the display +/// server's: the reference box runs 3.5-6.5 GiB of desktop ambient on the same +/// card, and a plan that ignores it produces an allocation failure at load +/// rather than a slow model. +pub const USABLE_VRAM_FRACTION: f64 = 0.75; + +/// Float precisions the planner can pick today. Ordered highest-quality +/// first; P9's int8/int4 tiers extend this enum downward. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub enum Precision { + /// `backend::Gpu` — f32 weights and KV cache. + F32, + /// `backend::GpuF16` — f16 weights and KV cache, f32 attention-score + /// island. Needs an adapter advertising `SHADER_F16`. + F16, +} + +impl Precision { + /// Bytes per stored float. + #[must_use] + pub fn bytes_per_float(self) -> u64 { + match self { + Self::F32 => 4, + Self::F16 => 2, + } + } + + /// Highest-first, the order the planner tries them in. + #[must_use] + pub fn descending() -> [Self; 2] { + [Self::F32, Self::F16] + } +} + +/// What the planner needs to know about a model to size it. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ModelShape { + /// Total parameters (weights + embeddings), as the checkpoint reports. + pub params: u64, + /// KV-cache floats stored **per token**, summed over layers: + /// `2 * layers * num_kv_heads * head_dim`. + pub kv_floats_per_token: u64, + /// Context length the plan must hold. + pub context_tokens: usize, +} + +impl ModelShape { + /// Sizing from a decoder's hyperparameters, so a caller passes the + /// `config.json` numbers rather than pre-computing cache geometry. + #[must_use] + pub fn from_decoder( + params: u64, + layers: usize, + num_kv_heads: usize, + head_dim: usize, + context_tokens: usize, + ) -> Self { + assert!(params > 0, "ModelShape: a model has parameters"); + assert!( + layers > 0 && num_kv_heads > 0 && head_dim > 0, + "ModelShape: degenerate decoder geometry \ + (layers {layers}, kv heads {num_kv_heads}, head_dim {head_dim})" + ); + let kv_floats_per_token = 2 * layers as u64 * num_kv_heads as u64 * head_dim as u64; + Self { + params, + kv_floats_per_token, + context_tokens, + } + } + + /// Projected resident bytes at `precision`: weights + KV cache + the fixed + /// runner overhead. + #[must_use] + pub fn projected_bytes(&self, precision: Precision) -> u64 { + let per_float = precision.bytes_per_float(); + let weights = self.params.saturating_mul(per_float); + let kv = self + .kv_floats_per_token + .saturating_mul(self.context_tokens as u64) + .saturating_mul(per_float); + weights.saturating_add(kv).saturating_add(OVERHEAD_BYTES) + } +} + +/// One adapter's budget, as the planner sees it. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DeviceBudget { + /// Total VRAM the adapter reports. + pub vram_bytes: u64, + /// Does it advertise `SHADER_F16`? Without it, f16 is not a candidate at + /// all — the dev box's own DX12 rows are exactly this case while its + /// Vulkan rows are not. + pub shader_f16: bool, +} + +impl DeviceBudget { + /// Read a budget off an enumerated adapter. `None` when VRAM is unknown + /// (wgpu exposes no portable query; only the Windows DXGI walk fills it + /// in today) — the planner refuses to guess rather than promise a fit it + /// cannot size. + #[must_use] + pub fn from_adapter(adapter: &GpuAdapter) -> Option { + Some(Self { + vram_bytes: adapter.vram_bytes?, + shader_f16: adapter.shader_f16, + }) + } + + /// Bytes a plan may claim, after leaving the display its share. + #[must_use] + pub fn usable_bytes(&self) -> u64 { + let usable = (self.vram_bytes as f64 * USABLE_VRAM_FRACTION) as u64; + debug_assert!(usable <= self.vram_bytes, "usable VRAM cannot exceed total"); + usable + } +} + +/// A precision decision, with the numbers behind it — the shape the +/// `plan`/`doctor` introspection item will render. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Fit { + /// The chosen precision. + pub precision: Precision, + /// Projected resident bytes at that precision. + pub projected_bytes: u64, + /// Bytes the adapter allows a plan to claim. + pub usable_bytes: u64, +} + +impl Fit { + /// Headroom left over — the slack a longer context or a second model + /// would eat into. + #[must_use] + pub fn headroom_bytes(&self) -> u64 { + self.usable_bytes.saturating_sub(self.projected_bytes) + } +} + +/// Pick the **highest** precision that fits `budget`, or `None` when even f16 +/// does not — the signal that this model needs quantization (P9) or a +/// multi-device plan, not a smaller float. +/// +/// Never silently ships a worse tier than the hardware can hold, and never +/// picks f16 on an adapter that does not advertise `SHADER_F16`. +#[must_use] +pub fn pick_precision(shape: &ModelShape, budget: &DeviceBudget) -> Option { + assert!(shape.params > 0, "pick_precision: a model has parameters"); + let usable = budget.usable_bytes(); + for precision in Precision::descending() { + if precision == Precision::F16 && !budget.shader_f16 { + continue; + } + let projected = shape.projected_bytes(precision); + if projected <= usable { + return Some(Fit { + precision, + projected_bytes: projected, + usable_bytes: usable, + }); + } + } + None +} + +#[cfg(test)] +mod tests { + use super::*; + + const GIB: u64 = 1 << 30; + + /// Qwen2.5-1.5B: 1.54 G params, 28 layers, 2 kv heads, head_dim 128. + fn qwen2_1_5b(context: usize) -> ModelShape { + ModelShape::from_decoder(1_543_714_304, 28, 2, 128, context) + } + + #[test] + fn projected_size_tracks_the_measured_reference_numbers() { + let shape = qwen2_1_5b(4096); + let f32_bytes = shape.projected_bytes(Precision::F32); + let f16_bytes = shape.projected_bytes(Precision::F16); + // Measured on the reference card: ~8.0 GiB runner f32, ~3.6 GiB f16 + // (bench/BASELINE.md). The projection must land in the same + // neighbourhood, not merely be monotone — it reads 7.0 / 3.9 GiB + // here, i.e. slightly under f32's measurement and slightly over + // f16's, which is the accuracy a single fixed overhead term buys. + assert!( + (6 * GIB..=9 * GIB).contains(&f32_bytes), + "f32 projection {f32_bytes} outside the measured ~8 GiB band" + ); + assert!( + (3 * GIB..=5 * GIB).contains(&f16_bytes), + "f16 projection {f16_bytes} outside the measured ~3.6 GiB band" + ); + assert!(f16_bytes < f32_bytes, "f16 must project smaller than f32"); + } + + #[test] + fn the_reference_card_gets_f32_and_a_small_card_gets_f16() { + let shape = qwen2_1_5b(4096); + // RTX 4070 Ti SUPER as the inventory reports it: 15.7 GiB, f16 on Vulkan. + let big = DeviceBudget { + vram_bytes: 16_852_000_000, + shader_f16: true, + }; + let fit = pick_precision(&shape, &big).expect("1.5B fits a 16 GB card"); + assert_eq!(fit.precision, Precision::F32, "highest that fits wins"); + assert!(fit.headroom_bytes() > 0, "a fit leaves headroom"); + + // An 8 GB card cannot hold the f32 build but holds f16 comfortably. + let small = DeviceBudget { + vram_bytes: 8 * GIB, + shader_f16: true, + }; + let fit = pick_precision(&shape, &small).expect("f16 fits 8 GB"); + assert_eq!(fit.precision, Precision::F16); + assert!(fit.projected_bytes <= fit.usable_bytes, "a fit must fit"); + } + + #[test] + fn an_adapter_without_shader_f16_never_gets_an_f16_plan() { + let shape = qwen2_1_5b(4096); + // The dev box's own DX12 rows: same card, no SHADER_F16. + let dx12 = DeviceBudget { + vram_bytes: 8 * GIB, + shader_f16: false, + }; + assert!( + pick_precision(&shape, &dx12).is_none(), + "without SHADER_F16 the only candidate is f32, which does not fit" + ); + // With room for f32, the same adapter plans fine. + let roomy = DeviceBudget { + vram_bytes: 24 * GIB, + shader_f16: false, + }; + assert_eq!( + pick_precision(&shape, &roomy).map(|f| f.precision), + Some(Precision::F32) + ); + } + + #[test] + fn a_model_too_big_for_f16_reports_no_fit_rather_than_a_bad_plan() { + // OLMoE-1B-7B's ~7 G params: ~14 GiB in f16, past a 16 GiB card's + // usable share — exactly the case bench/BASELINE.md records as + // "GPU is out of reach until keep-quantized VRAM (P9)". + let moe = ModelShape::from_decoder(6_919_000_000, 16, 16, 64, 4096); + let card = DeviceBudget { + vram_bytes: 16_852_000_000, + shader_f16: true, + }; + assert!( + pick_precision(&moe, &card).is_none(), + "no float precision fits; the answer is quantization, not a guess" + ); + } + + #[test] + fn context_length_moves_the_decision() { + // The KV cache is the term that grows with context, so a long enough + // context must be able to push a model off f32 onto f16. + let card = DeviceBudget { + vram_bytes: 12 * GIB, + shader_f16: true, + }; + let short = pick_precision(&qwen2_1_5b(1024), &card).expect("short context fits"); + // 64k of KV adds ~3.8 GiB in f32 — past a 12 GiB card's usable share, + // but comfortable at half the width. + let long = pick_precision(&qwen2_1_5b(65_536), &card).expect("long context still fits"); + assert_eq!(short.precision, Precision::F32); + assert_eq!(long.precision, Precision::F16, "KV growth forces the drop"); + assert!( + long.projected_bytes > qwen2_1_5b(1024).projected_bytes(Precision::F16), + "a longer context must project larger at the same precision" + ); + } + + #[test] + fn a_budget_needs_known_vram_before_it_will_plan() { + // `vram_bytes: None` (every non-Windows adapter today) must yield no + // budget at all rather than a guessed one. + let unknown = GpuAdapter { + name: "test adapter".into(), + backend: wgpu::Backend::Vulkan, + device_type: wgpu::DeviceType::DiscreteGpu, + shader_f16: true, + max_buffer_bytes: 4 * GIB, + vram_bytes: None, + }; + assert!(DeviceBudget::from_adapter(&unknown).is_none()); + } + + #[test] + fn usable_vram_leaves_the_display_its_share() { + let budget = DeviceBudget { + vram_bytes: 16 * GIB, + shader_f16: true, + }; + let usable = budget.usable_bytes(); + assert!(usable < budget.vram_bytes, "never plan the whole card"); + assert_eq!(usable, 12 * GIB, "75% of 16 GiB"); + } +}