diff --git a/Cargo.lock b/Cargo.lock index 4649f66..2e828e6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -136,7 +136,7 @@ checksum = "0ae92a5119aa49cdbcf6b9f893fe4e1d98b04ccbf82ee0584ad948a44a734dea" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -296,7 +296,7 @@ dependencies = [ "regex", "rustc-hash 2.1.3", "shlex 1.3.0", - "syn", + "syn 2.0.118", ] [[package]] @@ -388,7 +388,7 @@ dependencies = [ "proc-macro2", "quote", "rustversion", - "syn", + "syn 2.0.118", ] [[package]] @@ -580,7 +580,7 @@ dependencies = [ "derive-new", "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -842,7 +842,7 @@ checksum = "f65693059b6b9c588b9f62fed1cedbf0a8b805631457ea162d68f0de186f3de5" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -895,6 +895,26 @@ dependencies = [ "pkg-config", ] +[[package]] +name = "c2rust-bitfields" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "46dc7d2bffa0d0b3d47eb2dc69973466858281446c2ac9f6d8a10e92ab1017df" +dependencies = [ + "c2rust-bitfields-derive", +] + +[[package]] +name = "c2rust-bitfields-derive" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebe1117afa5937ce280034e31fa1e84ed1824a252f75380327eed438535333f8" +dependencies = [ + "proc-macro2", + "quote", + "syn 1.0.109", +] + [[package]] name = "candle-core" version = "0.10.2" @@ -1558,7 +1578,7 @@ dependencies = [ "prettyplease", "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -1570,7 +1590,7 @@ dependencies = [ "darling 0.23.0", "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -1622,6 +1642,23 @@ dependencies = [ "web-time", ] +[[package]] +name = "cubecl-spirv" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd12ded588654255e94e3d80a58e0988fdf50dd6b80365795072b04b99884f88" +dependencies = [ + "bitflags", + "cubecl-common", + "cubecl-core", + "cubecl-opt", + "cubecl-runtime", + "half", + "hashbrown 0.16.1", + "serde", + "tracel-rspirv", +] + [[package]] name = "cubecl-std" version = "0.10.0" @@ -1645,6 +1682,7 @@ version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4d3663c6e1a86187172b2cee495b6b533e7d91a2cb37e26f421649e315fe4c3a" dependencies = [ + "ash", "async-channel", "bytemuck", "cfg-if", @@ -1653,12 +1691,14 @@ dependencies = [ "cubecl-core", "cubecl-ir", "cubecl-runtime", + "cubecl-spirv", "derive-new", "derive_more", "half", "hashbrown 0.16.1", "log", "sanitize-filename", + "tracel-ash", "tracing", "wasm-bindgen-futures", "wgpu", @@ -1848,7 +1888,7 @@ dependencies = [ "proc-macro2", "quote", "strsim", - "syn", + "syn 2.0.118", ] [[package]] @@ -1861,7 +1901,7 @@ dependencies = [ "ident_case", "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -1874,7 +1914,7 @@ dependencies = [ "proc-macro2", "quote", "strsim", - "syn", + "syn 2.0.118", ] [[package]] @@ -1885,7 +1925,7 @@ checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" dependencies = [ "darling_core 0.20.11", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -1896,7 +1936,7 @@ checksum = "d38308df82d1080de0afee5d069fa14b0326a88c14f15c5ccda35b4a6c414c81" dependencies = [ "darling_core 0.21.3", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -1907,7 +1947,7 @@ checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" dependencies = [ "darling_core 0.23.0", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -1945,7 +1985,7 @@ checksum = "2cdc8d50f426189eef89dac62fabfa0abb27d5cc008f25bf4156a0203325becc" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -1966,7 +2006,7 @@ dependencies = [ "darling 0.20.11", "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -1976,7 +2016,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c" dependencies = [ "derive_builder_core", - "syn", + "syn 2.0.118", ] [[package]] @@ -1998,7 +2038,7 @@ dependencies = [ "proc-macro2", "quote", "rustc_version", - "syn", + "syn 2.0.118", "unicode-xid", ] @@ -2071,7 +2111,7 @@ checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -2191,7 +2231,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -2213,7 +2253,7 @@ dependencies = [ "darling 0.21.3", "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -2233,7 +2273,7 @@ checksum = "44f23cf4b44bfce11a86ace86f8a73ffdec849c9fd00a386a53d278bd9e81fb3" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -3142,7 +3182,7 @@ checksum = "c34819042dc3d3971c46c2190835914dfbe0c3c13f61449b2997f4e9722dfa60" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -3200,7 +3240,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "38c0b942f458fe50cdac086d2f946512305e5631e720728f2a61aabcd47a6264" dependencies = [ "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -3413,7 +3453,7 @@ dependencies = [ "darling 0.20.11", "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -3526,7 +3566,7 @@ checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -3551,7 +3591,7 @@ dependencies = [ "pollster", "serde", "serde_json", - "sha2 0.10.9", + "sha2 0.11.0", "thiserror 2.0.18", "tokenizers", "ureq 3.3.0", @@ -3746,7 +3786,7 @@ checksum = "ed3955f1a9c7c0c15e092f9c887db08b1fc683305fdf6eb6684f22555355e202" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -4157,7 +4197,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" dependencies = [ "proc-macro2", - "syn", + "syn 2.0.118", ] [[package]] @@ -4185,7 +4225,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4488a4a36b9a4ba6b9334a32a39971f77c1436ec82c38707bce707699cc3bbcb" dependencies = [ "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -4857,7 +4897,7 @@ checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -5070,6 +5110,17 @@ version = "2.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" +[[package]] +name = "syn" +version = "1.0.109" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "syn" version = "2.0.118" @@ -5098,7 +5149,7 @@ checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -5209,7 +5260,7 @@ checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -5220,7 +5271,7 @@ checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -5471,6 +5522,16 @@ version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" +[[package]] +name = "tracel-ash" +version = "0.39.3+sdk1.4.350" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d97fcf7e656528db871e679727140819a221d8c2aa5b90bfb5eb492a0ecc16e6" +dependencies = [ + "ash", + "c2rust-bitfields", +] + [[package]] name = "tracel-llvm" version = "20.1.4-7" @@ -5522,7 +5583,7 @@ dependencies = [ "proc-macro2", "quote", "regex", - "syn", + "syn 2.0.118", "tracel-llvm-bundler", "tracel-tblgen-rs", "unindent", @@ -5537,6 +5598,17 @@ dependencies = [ "tracel-llvm-bundler", ] +[[package]] +name = "tracel-rspirv" +version = "0.12.1+sdk-1.4.341.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1653aa21b867351f48c51f1063a2f872f8e82931951cae469d8a53aa4d7d72e8" +dependencies = [ + "bitflags", + "rustc-hash 2.1.3", + "serde", +] + [[package]] name = "tracel-tblgen-rs" version = "20.1.4-7" @@ -5569,7 +5641,7 @@ checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -5791,7 +5863,7 @@ checksum = "41b6d82be61465f97d42bd1d15bf20f3b0a3a0905018f38f9d6f6962055b0b5c" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -5882,7 +5954,7 @@ dependencies = [ "bumpalo", "proc-macro2", "quote", - "syn", + "syn 2.0.118", "wasm-bindgen-shared", ] @@ -6203,7 +6275,7 @@ checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -6214,7 +6286,7 @@ checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -6407,7 +6479,7 @@ checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.118", "synstructure", ] @@ -6428,7 +6500,7 @@ checksum = "e2e817b7b52d0c7358d3246da9d69935ebb18116b2b102b4230dac079b4862f5" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] @@ -6448,7 +6520,7 @@ checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.118", "synstructure", ] @@ -6488,7 +6560,7 @@ checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.118", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 9767c26..ba97a81 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,6 +17,11 @@ repository = "https://github.com/physics515/Mummu" burn = { version = "0.21", default-features = false, features = [ "std", "wgpu", + # `vulkan` switches CubeCL's kernel compiler on Vulkan adapters from + # WGSL/naga to direct SPIR-V; other APIs keep WGSL in the same binary. + # Adopted 2026-07-12: decode 70.7 -> 54.3 ms/token with parity + # byte-identical (see bench/BASELINE.md). + "vulkan", "fusion", "autotune", "store", @@ -49,7 +54,7 @@ tokenizers = "0.22" ureq = "3" # Stream-verifies downloads against the Hub's announced LFS sha256 # (X-Linked-ETag); SHA-NI accelerated, ~GB/s — negligible next to the network. -sha2 = "0.10" +sha2 = "0.11" serde = { version = "1", features = ["derive"] } serde_json = "1" diff --git a/README.md b/README.md index af78ea0..1e6d1a0 100644 --- a/README.md +++ b/README.md @@ -41,22 +41,35 @@ It exists because two local-first apps — **[laurelane](https://github.com/phys - **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. -- **Function calling (Hermes-style)** — advertise `ToolSpec`s through `render_with_tools` (the exact - `# Tools`/`` template Qwen2.5/Qwen3 are trained on), feed results back as merged - `` turns, and extract calls with a bounded `parse_tool_calls`; proven end-to-end on - the real GPU (Qwen2.5-1.5B emitted a parseable `get_weather({"city": "Paris"})` call, - `tests/real_toolcall.rs`). +- **Function calling (both zoo conventions)** — advertise `ToolSpec`s through `render_with_tools` in + the convention the model was trained on: **Hermes** for Qwen2.5/Qwen3 (the exact + `# Tools`/`` JSON template, results as merged `` turns) and **LFM** for + LFM2.5 (bare tool JSON on a `List of tools:` system line, Pythonic calls in + `<|tool_call_start|>` tokens, results as real `tool` turns, past-turn `` stripping); both + parsers are bounded with a loud error taxonomy. Proven end-to-end on the real GPU: Qwen2.5-1.5B + emitted a parseable Hermes call and LFM2.5-1.2B emitted exactly + `<|tool_call_start|>[get_weather(city="Paris")]<|tool_call_end|>` (`tests/real_toolcall.rs`, + `tests/real_toolcall_lfm.rs`). - **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** (14.1 tok/s / 88 ms TTFT); the parity gate re-passes - unchanged on f32, where the island casts are no-ops ([bench/BASELINE.md](bench/BASELINE.md)). -- **Benchmarked** — Qwen2.5-1.5B on the reference GPU: **TTFT 88.4 ms, decode 14.1 tok/s** (f32, 11.9 GiB - whole-card peak ≈ 7.9 GiB runner; f16: 88.0 ms, 14.1 tok/s, 6.75 GiB ≈ 3.6 GiB runner) — recorded with + 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)). +- **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`). - **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`. +- **GGUF import (container + K-quant dequant)** — `mummu::gguf` parses the llama.cpp container (typed, + bounded metadata; fully validated tensor table) and dequantizes F32/F16/BF16, Q8_0, and the + **Q4_K/Q6_K superblocks** to f32 — proven against the model's true weights: on the real Qwen2.5-1.5B + Q4_K_M file the F32 norms come back **bit-exact** vs the bf16 safetensors of the same checkpoint and + Q4_K embedding rows dequantize at cosine 0.9975 (`tests/real_gguf.rs`). Next: the remaining quant + types and GGUF→model load (tracked in P3/P9). - **Hub downloads** — streaming HuggingFace fetches into the model cache: resumable (`.part` + HTTP Range, proven byte-identical after an interrupted transfer), length-verified, shard-index aware, with a per-chunk progress callback; verified end-to-end by downloading all-MiniLM and embedding with it. diff --git a/ROADMAP.md b/ROADMAP.md index 77cf327..889bef9 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -46,7 +46,11 @@ a benchmark holds/improves its budget; README perf claims link an artifact. ~83 GB/s vs the 4070 Ti SUPER's ~672 GB/s — the SPIR-V compiler feature (P6 item) and f16 are the levers to chase; re-baseline after each. *(2026-07-11) Confirmed empirically: f16 (half the weight traffic) decodes at exactly f32's speed — 70.9 vs 70.7 ms/token — so bandwidth isn't the limiter; - SPIR-V (TensorCores at f16) is the remaining lever.* + SPIR-V (TensorCores at f16) is the remaining lever.* *(2026-07-12) SPIR-V pulled: decode + 70.7 → 54.3 ms/token (+30% tok/s) on BOTH dtypes — and f16 still exactly matches f32, so the path + stays dispatch-bound (~114 GB/s effective vs ~672). Next levers: whatever closes the remaining + per-token dispatch gap (fewer kernels per step — deeper fusion of the decode step, or CubeCL + graph/megakernel work) rather than bandwidth.* - [ ] Evaluate Burn 0.21's `burn.toml` project config — per-subsystem tuning + a CubeCL kernel-validation layer without recompiling; useful as a debug switch for kernel-level parity hunts — https://burn.dev/blog/release-0.21.0/ @@ -111,6 +115,12 @@ a benchmark holds/improves its budget; README perf claims link an artifact. semantic test (`tests/real_minilm.rs`).* *(2026-07-10) **Parity PASSED**: embedding matches the Candle f32 reference (`minilm-probe` fixture) at cosine 0.99999994 with max |Δcomponent| 1.2e-7 (bound 1e-4); semantic sanity re-verified on real weights (paraphrase 0.556 vs cross-topic ≈ 0).* +- [ ] **Qwen3.5 small tier** as the next zoo port: released Feb 2026 in 0.8B / 2B / 4B / 9B, with + 2026 GGUF re-releases specifically improving tool-calling (chat-template fixes) — the 4B/9B are + the function-calling sweet spot BFCL identified (9B 66.1%), and unsloth ships ready GGUFs for the + P3 import path to chew on; parity reference = Ollama `qwen3.5:9b` (already pulled locally) — + https://unsloth.ai/docs/models/qwen3.5 · https://huggingface.co/unsloth/Qwen3.5-9B-GGUF + *(2026-07-12 research)* - [x] A `Model` trait so new architectures (Hermes-class function-callers, Gemma, Qwen3, …) slot in. *(2026-07-10) `models::CausalLm` — associated `Cache` type; a port supplies `new_cache` / `forward` / `is_eos` and inherits `generate` / `greedy_generate` / `first_token` from the shared @@ -165,6 +175,21 @@ The subsystem that turns "a model on HuggingFace or on disk" into a loaded, pari (Q2K–Q8K) AND IQ-quants, SIMD-optimized, importance-matrix support, HF-compatible config generation — evaluate as dependency-or-reference before hand-porting ggml-quants — https://docs.rs/pmetal-gguf/latest/pmetal_gguf/ + *(2026-07-12) **Container reader shipped** (`mummu::gguf`, no new deps): magic/version (v2/v3 LE), + typed+bounded metadata KVs (strings ≤ 1 MiB, arrays ≤ 4M, nesting ≤ 2), tensor table with + per-entry validation (known dtype, aligned offset, whole blocks, unique names), K-quant block + layouts recorded (`block_size`/`bytes_per_block` for Q4_0…Q8_K), fail-loud error taxonomy. + REAL-FILE proof (`tests/real_gguf.rs`): the local Qwen2.5-1.5B **Q4_K_M** parses — v3, 26 kvs, + 339 tensors, Q4_K `token_embd [1536, 151936]`, 198 K-quant tensors, ~1.04 GiB payload located + (3B file cross-checked: 435 tensors). 7 unit tests over a synthetic-bytes builder.* + *(2026-07-12, same run) **Dequant shipped** for the Q4_K_M set — F32/F16/BF16, Q8_0, and the + K-quant superblocks **Q4_K/Q6_K** (exact ports of ggml-quants' reference dequantizers, incl. the + packed 6-bit scale/min encoding) + `GgufFile::read_tensor_f32`. Proof against the model's TRUE + weights (`real_gguf.rs` — the same checkpoint exists locally as bf16 safetensors AND Q4_K_M GGUF): + the GGUF's F32 `output_norm.weight` is **bit-exact** vs the bf16 originals, and dequantized Q4_K + embedding rows hit cosine **0.9975** vs truth (garbage layout ⇒ ≈ 0) — 5 hand-computed-block unit + tests (121 total). NEXT slice: remaining dequants (Q4_0/Q5/Q2_K/Q3_K/Q5_K), GGUF→model load + (name remap + ggml dim-order transpose), tokenizer-from-GGUF-metadata.* - [ ] **GPTQ / AWQ** (HF safetensors) — import the calibration-quantized int4/int8 layouts most "quantized on the Hub" models ship as (a `.safetensors` payload + a quant config), dequant or keep-quant into Burn. - [ ] **ONNX** (optional) — `burn-import` ONNX→Burn for models distributed as ONNX graphs. @@ -206,10 +231,21 @@ The subsystem that turns "a model on HuggingFace or on disk" into a loaded, pari (calls + prose, loud error taxonomy). REAL-GPU proof (`tests/real_toolcall.rs`): Qwen2.5-1.5B greedy-emitted `{"name": "get_weather", "arguments": {"city": "Paris"}}` from a rendered prompt and the parser round-tripped it. 10 new unit tests.* -- [ ] LFM2.5 bracket-notation tool-call template + parser (`<|tool_list_start|>` special tokens, +- [x] LFM2.5 bracket-notation tool-call template + parser (`<|tool_list_start|>` special tokens, Python-ish call syntax) — with Hermes/Qwen2.5 (0.880 agent score) done, LFM2.5-1.2B (same score, fastest at ~1.5 s on 2026's 21-model local tool-calling benchmark) is the other target — https://mikeveerman.be/blog/github-2026-02-06-tool-calling-benchmark/ + *(2026-07-12) Shipped to LFM**2.5**'s actual wire format (its `chat_template.jinja` + model card — + NOT the gen-1 `<|tool_list_start|>` wrapping, which 2.5 dropped): tools as bare JSON on a + `List of tools: […]` system line (no default preamble), tool results as real `tool` role turns, + `` reasoning stripped from all but the last assistant history turn, calls emitted as a + Pythonic list in `<|tool_call_start|>…<|tool_call_end|>`. `chat`: style-split `render_with_tools`, + `Turn::assistant_tool_calls_lfm`, and a bounded recursive-descent parser (`parse_tool_calls_lfm`, + depth ≤ 8, ≤ 64 calls, Python AND JSON literal spellings, byte-offset error taxonomy). REAL-GPU + proof (`tests/real_toolcall_lfm.rs`): LFM2.5-1.2B greedy-emitted exactly + `<|tool_call_start|>[get_weather(city="Paris")]<|tool_call_end|>` from our rendered prompt and the + parser round-tripped it; the Qwen2 parity gate re-passed both legs after the template refactor + (max |Δlogit| 2.670e-5, Ollama greedy byte-identical). 16 new unit tests (109 total).* *(2026-07-10 research)* 2026 community numbers back the plan: Qwen3-8B keeps tool-calling score through Q4_K_M (0.919 quantized vs 0.933 full — quant does NOT cost tool reliability, good news for P9); BFCL shows a capability cliff below ~7B (Qwen3.5-9B 66.1% vs 4B 50.3%), so the zoo's @@ -263,10 +299,16 @@ that fits the model AND uses every device to the fullest. gate re-passed both legs (max |Δlogit| 2.670e-5, unchanged; Ollama greedy byte-identical), and f32 perf *improved* (TTFT 100.5 → 88.4 ms, decode 13.3 → 14.1 tok/s). f16 benches recorded in `bench/BASELINE.md`: 88.0 ms TTFT, 14.1 tok/s — speed parity with f32, VRAM halved.* -- [ ] Evaluate burn-wgpu's **`spirv` compiler feature** on Vulkan (CubeCL SPIR-V backend instead of +- [x] Evaluate burn-wgpu's **`spirv` compiler feature** on Vulkan (CubeCL SPIR-V backend instead of WGSL/naga): claims significantly faster matmul incl. TensorCores at f16 — could be the cheapest decode-tok/s lever on the dev GPU; gate on the parity harness + `bench/BASELINE.md` — https://github.com/tracel-ai/burn/blob/main/crates/burn-wgpu/README.md + *(2026-07-12) **Adopted** (burn `vulkan` feature; runtime reports `fusion>>`, + auto-selected on Vulkan adapters only — other APIs keep WGSL, no code changes): decode + **70.7 → 54.3 ms/token (14.1 → 18.4 tok/s, +30%)** on f32 AND f16; TTFT 88.4 → 96.7 ms (well + under the 150 ms ceiling); VRAM peak unchanged (11.5 GiB whole-card). Parity gate byte-identical + (max |Δlogit| 2.670e-5, Ollama greedy exact); f16 island + CPU budget gates re-passed + (108.4 ms / 11.7 tok/s GPU gate, 13.2 tok/s CPU). `bench/BASELINE.md` re-baselined.* - [ ] **Placement plan** — given model size + KV-cache + display headroom and the device set, choose a **fit-and-fill** plan: single GPU when it fits; **shard layers across multiple GPUs** (pipeline/ layer-parallel over Burn's multi-device tensors — Burn gives the multi-device *primitives*, not automatic @@ -299,6 +341,11 @@ that fits the model AND uses every device to the fullest. (temperature 0 for the greedy leg) — that plus an fp16 GGUF of the same revision is a workable logits leg without Python — https://docs.liquid.ai/deployment/on-device/llama-cpp · https://github.com/ggml-org/llama.cpp/blob/master/tools/server/README.md + *(2026-07-12 research)* Caution for that route: llama.cpp's own chat/tool layer mishandles + LFM2.5 (issue #23838 — its parser rejects the documented `<|tool_call_start|>[…]` format), so + drive `llama-server` in RAW completion mode (`/completion`, no chat template) and render prompts + with our byte-verified `ChatMl::lfm2()` — never through llama.cpp's template stack — + https://github.com/ggml-org/llama.cpp/issues/23838 - [ ] Wire the perf suite (above) into the parity harness so a correctness *or* budget regression fails CI. ### P8 — Model management API diff --git a/bench/BASELINE.md b/bench/BASELINE.md index 8079d4a..b5075aa 100644 --- a/bench/BASELINE.md +++ b/bench/BASELINE.md @@ -1,29 +1,32 @@ # Benchmark baseline & budgets -Reference machine: Ryzen 9 7950X3D · 128 GB · **RTX 4070 Ti SUPER 16 GB** (wgpu/Vulkan, `Fusion`). +Reference machine: Ryzen 9 7950X3D · 128 GB · **RTX 4070 Ti SUPER 16 GB** (wgpu/Vulkan, +`Fusion`, **SPIR-V compiler** — `fusion>>` since 2026-07-12). Bench: `MUMMU_QWEN2_DIR= cargo bench -p mummu-bench` (criterion, `benches/runner.rs`; fixed ~36-token ChatML prompt). A change that pushes a budget over its ceiling does not ship; update the recorded numbers (and this file's date) only on a legitimate improvement. ## Qwen2.5-1.5B-Instruct · single GPU · f32 -| Metric | Recorded (2026-07-11) | Budget | +| Metric | Recorded (2026-07-12, SPIR-V) | Budget | | --- | --- | --- | -| TTFT (fresh cache: full prefill + first token) | 88.4 ms | ≤ 150 ms | -| Decode latency (32 greedy tokens, warm KV cache) | 2.263 s → **70.7 ms/token ≈ 14.1 tok/s** | ≥ 10 tok/s | -| Peak GPU memory during decode (whole card, ~4.0 GiB desktop ambient → ~7.9 GiB runner) | 11.9 GiB | ≤ 13 GiB whole-card | +| TTFT (fresh cache: full prefill + first token) | 96.7 ms | ≤ 150 ms | +| Decode latency (32 greedy tokens, warm KV cache) | 1.737 s → **54.3 ms/token ≈ 18.4 tok/s** | ≥ 10 tok/s | +| Peak GPU memory during the real-inference suite (whole card, ~3.5 GiB desktop ambient → ~8.0 GiB runner) | 11.5 GiB | ≤ 13 GiB whole-card | ## Qwen2.5-1.5B-Instruct · single GPU · **f16** (weights + KV; f32 attention-score island) -| Metric | Recorded (2026-07-11) | Budget | +| Metric | Recorded (2026-07-12, SPIR-V) | Budget | | --- | --- | --- | -| TTFT (fresh cache: full prefill + first token) | 88.0 ms | ≤ 150 ms | -| Decode latency (32 greedy tokens, warm KV cache) | 2.270 s → **70.9 ms/token ≈ 14.1 tok/s** | ≥ 10 tok/s | -| Peak GPU memory during decode (whole card, 3.1 GiB ambient → **~3.6 GiB runner**) | 6.75 GiB | ≤ 8 GiB whole-card | - -f16 decode speed matches f32 — the WGSL decode path is dispatch-bound, not bandwidth-bound (see Notes), -so halved weight traffic buys nothing yet; the win is VRAM (**~7.9 → ~3.6 GiB runner**, room for larger -models/contexts). The SPIR-V compiler feature (ROADMAP P6) is the identified speed lever for both dtypes. +| TTFT (fresh cache: full prefill + first token) | 97.2 ms | ≤ 150 ms | +| Decode latency (32 greedy tokens, warm KV cache) | 1.744 s → **54.5 ms/token ≈ 18.4 tok/s** | ≥ 10 tok/s | +| Peak GPU memory during decode (whole card, 3.1 GiB ambient → **~3.6 GiB runner**, 2026-07-11 measure) | 6.75 GiB | ≤ 8 GiB whole-card | + +The SPIR-V compiler (burn `vulkan` feature) cut decode latency **23%** on both dtypes +(70.7 → 54.3 ms/token) at the cost of ~9 ms TTFT (88.4 → 96.7 ms, still ⅔ under its ceiling); parity +held byte-identically (max |Δlogit| 2.670e-5, Ollama greedy leg exact). Decode remains +dispatch-bound — f16 still buys VRAM (~7.9 → ~3.6 GiB runner), not speed; ~54 ms/token streams f32 +weights at only ~114 GB/s vs the card's ~672 GB/s, so kernel/dispatch overhead is still the ceiling. ## Qwen2.5-0.5B-Instruct · CPU (burn-flex) · f32 diff --git a/crates/mummu/src/chat.rs b/crates/mummu/src/chat.rs index fdbce37..30ffdd8 100644 --- a/crates/mummu/src/chat.rs +++ b/crates/mummu/src/chat.rs @@ -5,11 +5,18 @@ //! the exact prompt committed in the Candle logits fixture). //! //! Both zoo LLMs speak ChatML; LFM2.5 additionally prefixes `<|startoftext|>`. -//! Tool use follows the Hermes convention Qwen2.5/Qwen3 ship in their chat -//! template: tool signatures in a `` block of the system turn, calls -//! emitted as `{json}`, results returned inside -//! `` blocks of a user turn. LFM2.5's bracket notation is a -//! P4 follow-up. +//! Tool use comes in two conventions, selected by the per-model constructor: +//! +//! - **Hermes** (Qwen2.5/Qwen3): tool signatures in a `` block of the +//! system turn, calls emitted as `{json}`, results +//! returned inside `` blocks of a user turn. +//! - **LFM** (LFM2.5, per its `chat_template.jinja` + model card): tool +//! signatures as bare JSON in a `List of tools: […]` line of the system +//! turn, calls emitted as a *Pythonic call list* between the +//! `<|tool_call_start|>`/`<|tool_call_end|>` special tokens — e.g. +//! `[get_weather(city="Paris")]` — results returned in a dedicated `tool` +//! role turn, and ``-prefixed reasoning stripped from every +//! assistant history turn but the last. /// Who is speaking in a [`Turn`]. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -18,17 +25,27 @@ pub enum Role { User, Assistant, /// A tool result going back to the model. Hermes-style templates render - /// these inside a *user* turn as `` blocks. + /// these inside a *user* turn as `` blocks; LFM-style + /// templates give each one its own `tool` role turn. Tool, } +/// Which tool-use convention a template speaks (see the module docs). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ToolCallStyle { + Hermes, + Lfm, +} + impl Role { - fn tag(self) -> &'static str { - match self { - Self::System => "system", - Self::User => "user", - Self::Assistant => "assistant", - Self::Tool => "user", // Hermes: tool results ride in a user turn + fn tag(self, style: ToolCallStyle) -> &'static str { + match (self, style) { + (Self::System, _) => "system", + (Self::User, _) => "user", + (Self::Assistant, _) => "assistant", + // Hermes: tool results ride in a user turn; LFM: a real tool turn. + (Self::Tool, ToolCallStyle::Hermes) => "user", + (Self::Tool, ToolCallStyle::Lfm) => "tool", } } } @@ -90,8 +107,30 @@ impl Turn { } } - /// A tool's result going back to the model; renders as a - /// `` block (consecutive ones merge into one user turn). + /// An assistant turn that invokes tools in LFM's Pythonic convention: + /// the calls render as one bracketed call list between the + /// `<|tool_call_start|>`/`<|tool_call_end|>` special tokens — exactly + /// what an LFM2.5 model emits — so histories re-render faithfully. + #[must_use] + pub fn assistant_tool_calls_lfm(calls: &[ToolCall]) -> Self { + assert!(!calls.is_empty(), "assistant_tool_calls_lfm: no calls"); + assert!( + calls.len() <= MAX_TOOL_CALLS, + "assistant_tool_calls_lfm: {} calls exceeds the {MAX_TOOL_CALLS} bound", + calls.len() + ); + Self { + role: Role::Assistant, + content: format!( + "<|tool_call_start|>{}<|tool_call_end|>", + pythonic_calls(calls) + ), + } + } + + /// A tool's result going back to the model. Hermes templates render it as + /// a `` block (consecutive ones merge into one user turn); + /// LFM templates give it its own `tool` role turn. #[must_use] pub fn tool_response(content: impl Into) -> Self { Self { @@ -101,6 +140,100 @@ impl Turn { } } +/// Deepest literal nesting the Pythonic renderer/parser will follow — far +/// past any real argument payload, and the recursion bound for both. +const MAX_VALUE_DEPTH: usize = 8; + +/// Render tool calls as LFM's Pythonic call list: `[name(k=v, …), …]`. +/// JSON scalars map to Python spellings (`true`→`True`, `null`→`None`); +/// strings/lists/objects render as Python literals. +fn pythonic_calls(calls: &[ToolCall]) -> String { + assert!(!calls.is_empty(), "pythonic_calls: no calls"); + assert!(calls.len() <= MAX_TOOL_CALLS, "pythonic_calls: over bound"); + let mut out = String::from("["); + for (i, call) in calls.iter().enumerate() { + assert!(!call.name.is_empty(), "pythonic_calls: unnamed call"); + if i > 0 { + out.push_str(", "); + } + out.push_str(&call.name); + out.push('('); + match &call.arguments { + serde_json::Value::Object(args) => { + for (j, (key, value)) in args.iter().enumerate() { + if j > 0 { + out.push_str(", "); + } + out.push_str(key); + out.push('='); + python_literal(value, &mut out, 0); + } + } + serde_json::Value::Null => {} + other => panic!("pythonic_calls: arguments must be an object or null, got {other}"), + } + out.push(')'); + } + out.push(']'); + debug_assert!(out.starts_with('[') && out.ends_with(']'), "list shape"); + out +} + +/// Append one JSON value as a Python literal. Strings render double-quoted +/// with JSON-style escapes (valid Python), scalars as Python spellings. +fn python_literal(value: &serde_json::Value, out: &mut String, depth: usize) { + assert!(depth <= MAX_VALUE_DEPTH, "python_literal: value too deep"); + match value { + serde_json::Value::Null => out.push_str("None"), + serde_json::Value::Bool(true) => out.push_str("True"), + serde_json::Value::Bool(false) => out.push_str("False"), + serde_json::Value::Number(n) => out.push_str(&n.to_string()), + serde_json::Value::String(s) => python_string_literal(s, out), + serde_json::Value::Array(items) => { + out.push('['); + for (i, item) in items.iter().enumerate() { + if i > 0 { + out.push_str(", "); + } + python_literal(item, out, depth + 1); + } + out.push(']'); + } + serde_json::Value::Object(map) => { + out.push('{'); + for (i, (key, val)) in map.iter().enumerate() { + if i > 0 { + out.push_str(", "); + } + python_string_literal(key, out); + out.push_str(": "); + python_literal(val, out, depth + 1); + } + out.push('}'); + } + } +} + +/// Double-quoted Python string literal with JSON-compatible escapes. +fn python_string_literal(s: &str, out: &mut String) { + out.push('"'); + for c in s.chars() { + match c { + '"' => out.push_str("\\\""), + '\\' => out.push_str("\\\\"), + '\n' => out.push_str("\\n"), + '\t' => out.push_str("\\t"), + '\r' => out.push_str("\\r"), + c if (c as u32) < 0x20 => { + out.push_str(&format!("\\u{:04x}", c as u32)); + } + c => out.push(c), + } + } + out.push('"'); + debug_assert!(out.ends_with('"'), "string literal closes"); +} + /// A callable tool signature, serialized into the system prompt exactly as /// the Hermes-style templates expect: `{"type": "function", "function": /// {"name": …, "description": …, "parameters": }}`. @@ -138,10 +271,20 @@ const MAX_TOOLS: usize = 128; /// What went wrong extracting tool calls from a model response. #[derive(Debug, thiserror::Error)] pub enum ToolCallError { - #[error("tool call {index}: unclosed tag")] + #[error("tool call {index}: unclosed tool-call tag")] Unclosed { index: usize }, #[error("tool call {index}: {reason}")] BadJson { index: usize, reason: String }, + #[error("tool call block {index}, byte {offset}: {reason}")] + Syntax { + /// Which `<|tool_call_start|>` block (0-based) failed to parse. + index: usize, + /// Byte offset *inside the block* where parsing stopped. + offset: usize, + reason: String, + }, + #[error("tool call arguments nested deeper than {MAX_VALUE_DEPTH} levels")] + TooDeep, #[error("more than {MAX_TOOL_CALLS} tool calls in one response")] TooMany, } @@ -178,30 +321,333 @@ pub fn parse_tool_calls(text: &str) -> Result<(Vec, String), ToolCallE Ok((calls, prose.trim().to_string())) } +/// Extract LFM-style tool calls from a model response: every +/// `<|tool_call_start|>…<|tool_call_end|>` block parses as a *Pythonic call +/// list* (`[name(k=v, …), …]`); the text outside the blocks (the model's +/// prose, trimmed) comes back alongside. Text with no blocks is simply +/// `(vec![], text)` — not an error. +pub fn parse_tool_calls_lfm(text: &str) -> Result<(Vec, String), ToolCallError> { + const OPEN: &str = "<|tool_call_start|>"; + const CLOSE: &str = "<|tool_call_end|>"; + let mut calls = Vec::new(); + let mut prose = String::new(); + let mut rest = text; + let mut block = 0; + while let Some(start) = rest.find(OPEN) { + prose.push_str(&rest[..start]); + let after_open = &rest[start + OPEN.len()..]; + let Some(end) = after_open.find(CLOSE) else { + return Err(ToolCallError::Unclosed { index: block }); + }; + let parsed = PythonicParser::new(after_open[..end].trim(), block).parse_call_list()?; + if calls.len() + parsed.len() > MAX_TOOL_CALLS { + return Err(ToolCallError::TooMany); + } + calls.extend(parsed); + rest = &after_open[end + CLOSE.len()..]; + block += 1; + } + prose.push_str(rest); + debug_assert!(calls.len() <= MAX_TOOL_CALLS, "bound enforced per block"); + Ok((calls, prose.trim().to_string())) +} + +/// A bounded recursive-descent parser for LFM's Pythonic call list. +/// +/// Grammar (whitespace-tolerant, trailing commas allowed): +/// ```text +/// calls := '[' [ call (',' call)* [','] ] ']' +/// call := ident '(' [ kwarg (',' kwarg)* [','] ] ')' +/// kwarg := ident '=' value +/// value := 'True' | 'False' | 'None' | 'true' | 'false' | 'null' +/// | number | string | '[' … ']' | '{' string ':' value, … '}' +/// ``` +/// Lowercase JSON spellings are accepted because LFM2.5 documents a JSON +/// fallback mode and small models mix the two. +struct PythonicParser<'a> { + src: &'a str, + pos: usize, + block: usize, +} + +impl<'a> PythonicParser<'a> { + fn new(src: &'a str, block: usize) -> Self { + Self { src, pos: 0, block } + } + + fn fail(&self, reason: impl Into) -> ToolCallError { + ToolCallError::Syntax { + index: self.block, + offset: self.pos, + reason: reason.into(), + } + } + + fn skip_ws(&mut self) { + while self.src[self.pos..].starts_with(|c: char| c.is_ascii_whitespace()) { + self.pos += 1; + } + } + + fn peek(&self) -> Option { + self.src[self.pos..].chars().next() + } + + fn eat(&mut self, c: char) -> bool { + if self.peek() == Some(c) { + self.pos += c.len_utf8(); + true + } else { + false + } + } + + fn expect(&mut self, c: char, what: &str) -> Result<(), ToolCallError> { + self.skip_ws(); + if self.eat(c) { + Ok(()) + } else { + Err(self.fail(format!("expected '{c}' {what}"))) + } + } + + /// The whole block: a bracketed list of calls, then end of input. + fn parse_call_list(mut self) -> Result, ToolCallError> { + assert!(self.pos == 0, "parse_call_list: parser already consumed"); + self.expect('[', "to open the call list")?; + let mut calls = Vec::new(); + loop { + self.skip_ws(); + if self.eat(']') { + break; + } + if calls.len() == MAX_TOOL_CALLS { + return Err(ToolCallError::TooMany); + } + calls.push(self.parse_call()?); + self.skip_ws(); + if !self.eat(',') && self.peek() != Some(']') { + return Err(self.fail("expected ',' or ']' after a call")); + } + } + self.skip_ws(); + if self.pos != self.src.len() { + return Err(self.fail("trailing text after the call list")); + } + debug_assert!(calls.len() <= MAX_TOOL_CALLS, "bound enforced in loop"); + Ok(calls) + } + + fn parse_call(&mut self) -> Result { + let name = self.parse_ident("a function name")?; + self.expect('(', "to open the arguments")?; + let mut args = serde_json::Map::new(); + loop { + self.skip_ws(); + if self.eat(')') { + break; + } + let key = self.parse_ident("an argument name")?; + self.expect('=', "between argument name and value")?; + self.skip_ws(); + let value = self.parse_value(0)?; + args.insert(key, value); + self.skip_ws(); + if !self.eat(',') && self.peek() != Some(')') { + return Err(self.fail("expected ',' or ')' after an argument")); + } + } + debug_assert!(!name.is_empty(), "parse_ident never returns empty"); + Ok(ToolCall { + name, + arguments: serde_json::Value::Object(args), + }) + } + + fn parse_ident(&mut self, what: &str) -> Result { + self.skip_ws(); + let start = self.pos; + if self + .peek() + .is_some_and(|c| c.is_ascii_alphabetic() || c == '_') + { + self.pos += 1; + while self + .peek() + .is_some_and(|c| c.is_ascii_alphanumeric() || c == '_') + { + self.pos += 1; + } + } + if self.pos == start { + return Err(self.fail(format!("expected {what}"))); + } + Ok(self.src[start..self.pos].to_string()) + } + + fn parse_value(&mut self, depth: usize) -> Result { + if depth > MAX_VALUE_DEPTH { + return Err(ToolCallError::TooDeep); + } + self.skip_ws(); + match self.peek() { + Some('"') | Some('\'') => Ok(serde_json::Value::String(self.parse_string()?)), + Some('[') => self.parse_list(depth), + Some('{') => self.parse_dict(depth), + Some(c) if c == '-' || c.is_ascii_digit() => self.parse_number(), + Some(c) if c.is_ascii_alphabetic() => { + let word = self.parse_ident("a literal")?; + match word.as_str() { + "True" | "true" => Ok(serde_json::Value::Bool(true)), + "False" | "false" => Ok(serde_json::Value::Bool(false)), + "None" | "null" => Ok(serde_json::Value::Null), + other => Err(self.fail(format!("unknown literal '{other}'"))), + } + } + _ => Err(self.fail("expected a value")), + } + } + + fn parse_list(&mut self, depth: usize) -> Result { + assert!(self.peek() == Some('['), "parse_list: caller checked '['"); + self.pos += 1; + let mut items = Vec::new(); + loop { + self.skip_ws(); + if self.eat(']') { + break; + } + items.push(self.parse_value(depth + 1)?); + self.skip_ws(); + if !self.eat(',') && self.peek() != Some(']') { + return Err(self.fail("expected ',' or ']' in a list")); + } + } + Ok(serde_json::Value::Array(items)) + } + + fn parse_dict(&mut self, depth: usize) -> Result { + assert!(self.peek() == Some('{'), "parse_dict: caller checked brace"); + self.pos += 1; + let mut map = serde_json::Map::new(); + loop { + self.skip_ws(); + if self.eat('}') { + break; + } + self.skip_ws(); + if !matches!(self.peek(), Some('"') | Some('\'')) { + return Err(self.fail("dict keys must be strings")); + } + let key = self.parse_string()?; + self.expect(':', "between dict key and value")?; + let value = self.parse_value(depth + 1)?; + map.insert(key, value); + self.skip_ws(); + if !self.eat(',') && self.peek() != Some('}') { + return Err(self.fail("expected ',' or '}' in a dict")); + } + } + Ok(serde_json::Value::Object(map)) + } + + /// A single- or double-quoted string with Python/JSON escapes. + fn parse_string(&mut self) -> Result { + let quote = self.peek().expect("parse_string: caller checked a quote"); + assert!(quote == '"' || quote == '\'', "caller checked the quote"); + self.pos += 1; + let mut out = String::new(); + loop { + let Some(c) = self.peek() else { + return Err(self.fail("unterminated string")); + }; + self.pos += c.len_utf8(); + match c { + c if c == quote => break, + '\\' => out.push(self.parse_escape()?), + c => out.push(c), + } + } + Ok(out) + } + + fn parse_escape(&mut self) -> Result { + let Some(c) = self.peek() else { + return Err(self.fail("dangling escape at end of string")); + }; + self.pos += c.len_utf8(); + match c { + '"' | '\'' | '\\' | '/' => Ok(c), + 'n' => Ok('\n'), + 't' => Ok('\t'), + 'r' => Ok('\r'), + 'u' => { + let hex = self + .src + .get(self.pos..self.pos + 4) + .ok_or_else(|| self.fail("truncated \\u escape"))?; + let code = u32::from_str_radix(hex, 16).map_err(|_| self.fail("bad \\u escape"))?; + self.pos += 4; + char::from_u32(code).ok_or_else(|| self.fail("\\u escape is not a scalar")) + } + other => Err(self.fail(format!("unsupported escape '\\{other}'"))), + } + } + + fn parse_number(&mut self) -> Result { + let start = self.pos; + self.eat('-'); + while self + .peek() + .is_some_and(|c| c.is_ascii_digit() || matches!(c, '.' | 'e' | 'E' | '+' | '-')) + { + self.pos += 1; + } + let text = &self.src[start..self.pos]; + debug_assert!(!text.is_empty(), "caller checked a digit or '-'"); + if let Ok(i) = text.parse::() { + return Ok(serde_json::Value::Number(i.into())); + } + let f = text + .parse::() + .map_err(|_| self.fail(format!("bad number '{text}'")))?; + serde_json::Number::from_f64(f) + .map(serde_json::Value::Number) + .ok_or_else(|| self.fail(format!("non-finite number '{text}'"))) + } +} + /// Longest conversation a single render will wrap — a generous bound that /// still catches an unbounded history being passed by mistake. const MAX_TURNS: usize = 1024; /// The ChatML template family: `<|im_start|>role\ncontent<|im_end|>\n` per /// turn, then an open assistant turn for the model to complete. `bos` is -/// prepended once when a model requires a start-of-text token. +/// prepended once when a model requires a start-of-text token; `tool_style` +/// picks the tool-use convention (see the module docs). #[derive(Debug, Clone)] pub struct ChatMl { bos: Option<&'static str>, + tool_style: ToolCallStyle, } impl ChatMl { - /// Qwen2 / Qwen2.5-Instruct: plain ChatML, no BOS. + /// Qwen2 / Qwen2.5-Instruct: plain ChatML, no BOS, Hermes tool use. #[must_use] pub fn qwen2() -> Self { - Self { bos: None } + Self { + bos: None, + tool_style: ToolCallStyle::Hermes, + } } - /// LFM2 / LFM2.5-Instruct: ChatML behind `<|startoftext|>`. + /// LFM2 / LFM2.5-Instruct: ChatML behind `<|startoftext|>`, LFM + /// (Pythonic) tool use. #[must_use] pub fn lfm2() -> Self { Self { bos: Some("<|startoftext|>"), + tool_style: ToolCallStyle::Lfm, } } @@ -221,10 +667,11 @@ impl ChatMl { "chat render: the template opens the assistant turn itself; \ a trailing assistant turn would double it" ); + let last_assistant = turns.iter().rposition(|t| t.role == Role::Assistant); let mut out = String::from(self.bos.unwrap_or("")); let mut i = 0; while i < turns.len() { - if turns[i].role == Role::Tool { + if self.tool_style == ToolCallStyle::Hermes && turns[i].role == Role::Tool { // Hermes: consecutive tool results merge into ONE user turn, // each wrapped in its own block. out.push_str("<|im_start|>user"); @@ -237,9 +684,9 @@ impl ChatMl { out.push_str("<|im_end|>\n"); } else { out.push_str("<|im_start|>"); - out.push_str(turns[i].role.tag()); + out.push_str(turns[i].role.tag(self.tool_style)); out.push('\n'); - out.push_str(&turns[i].content); + out.push_str(self.turn_content(&turns[i], i, last_assistant)); out.push_str("<|im_end|>\n"); i += 1; } @@ -249,11 +696,43 @@ impl ChatMl { out } - /// [`render`](Self::render) with Hermes-style function calling: the tool - /// signatures are advertised in a `# Tools` section of the system turn — - /// the exact wording and tag structure Qwen2.5/Qwen3 ship in their chat - /// template. An existing leading system turn provides the preamble; a - /// conversation without one gets a neutral "You are a helpful assistant." + /// What a turn's body renders as. LFM templates strip ``-prefixed + /// reasoning from every assistant history turn but the last (the + /// `keep_past_thinking=false` default of LFM2.5's `chat_template.jinja`); + /// everything else passes through. + fn turn_content<'a>( + &self, + turn: &'a Turn, + index: usize, + last_assistant: Option, + ) -> &'a str { + debug_assert!(index <= MAX_TURNS, "index bounded by the render assert"); + debug_assert!( + turn.role != Role::Assistant || last_assistant.is_some(), + "an assistant turn implies a last-assistant index" + ); + let is_past_assistant = + turn.role == Role::Assistant && last_assistant.is_some_and(|l| index != l); + if self.tool_style == ToolCallStyle::Lfm + && is_past_assistant + && let Some(end) = turn.content.rfind("") + { + return turn.content[end + "".len()..].trim(); + } + &turn.content + } + + /// [`render`](Self::render) with function calling in this template's + /// convention. + /// + /// - **Hermes** (Qwen2.5/Qwen3): tool signatures advertised in a + /// `# Tools` section of the system turn — the exact wording and tag + /// structure those models ship in their chat template. A conversation + /// without a system turn gets a neutral "You are a helpful assistant." + /// - **LFM** (LFM2.5): bare tool JSON on a `List of tools: […]` line + /// appended to the system turn — the exact shape of LFM2.5's + /// `chat_template.jinja` + model card, which injects *no* default + /// preamble: without a system turn the tools line stands alone. #[must_use] pub fn render_with_tools(&self, tools: &[ToolSpec], turns: &[Turn]) -> String { assert!( @@ -269,7 +748,15 @@ impl ChatMl { tools.iter().all(|t| !t.name.is_empty()), "render_with_tools: every tool needs a name" ); + match self.tool_style { + ToolCallStyle::Hermes => self.render_with_tools_hermes(tools, turns), + ToolCallStyle::Lfm => self.render_with_tools_lfm(tools, turns), + } + } + fn render_with_tools_hermes(&self, tools: &[ToolSpec], turns: &[Turn]) -> String { + debug_assert!(!tools.is_empty(), "checked by render_with_tools"); + debug_assert!(self.tool_style == ToolCallStyle::Hermes, "hermes only"); let (preamble, rest) = match turns.first() { Some(t) if t.role == Role::System => (t.content.as_str(), &turns[1..]), _ => ("You are a helpful assistant.", turns), @@ -300,6 +787,34 @@ impl ChatMl { wrapped.extend_from_slice(rest); self.render(&wrapped) } + + fn render_with_tools_lfm(&self, tools: &[ToolSpec], turns: &[Turn]) -> String { + debug_assert!(!tools.is_empty(), "checked by render_with_tools"); + debug_assert!(self.tool_style == ToolCallStyle::Lfm, "lfm only"); + let (preamble, rest) = match turns.first() { + Some(t) if t.role == Role::System => (t.content.as_str(), &turns[1..]), + _ => ("", turns), + }; + let mut system = String::from(preamble); + if !system.is_empty() { + system.push('\n'); + } + system.push_str("List of tools: ["); + for (i, tool) in tools.iter().enumerate() { + if i > 0 { + system.push_str(", "); + } + let json = serde_json::to_string(tool).unwrap_or_default(); + debug_assert!(!json.is_empty(), "a ToolSpec always serializes"); + system.push_str(&json); + } + system.push(']'); + + let mut wrapped = Vec::with_capacity(rest.len() + 1); + wrapped.push(Turn::system(system)); + wrapped.extend_from_slice(rest); + self.render(&wrapped) + } } #[cfg(test)] @@ -478,4 +993,251 @@ mod tests { fn trailing_assistant_turn_is_rejected() { let _ = ChatMl::qwen2().render(&[Turn::user("q"), Turn::assistant("half-done")]); } + + // ---- LFM (Pythonic) tool use ---------------------------------------- + + /// The tools line must match LFM2.5's `chat_template.jinja` byte shape: + /// `List of tools: [{bare tool json}, …]` appended to the system turn + /// with a `\n`, tools comma-joined, bare (no Hermes `"type":"function"` + /// wrapper — the model card's examples show `{"name": …}` directly). + #[test] + fn lfm_tools_render_matches_the_lfm25_template_shape() { + let raw = ChatMl::lfm2().render_with_tools( + &[weather_tool()], + &[ + Turn::system("You are a helpful assistant."), + Turn::user("Weather in Paris?"), + ], + ); + let expected = "<|startoftext|><|im_start|>system\nYou are a helpful assistant.\n\ + List of tools: [{\"name\":\"get_weather\",\"description\":\"Get the current weather for a city.\",\"parameters\":{\"properties\":{\"city\":{\"type\":\"string\"}},\"required\":[\"city\"],\"type\":\"object\"}}]<|im_end|>\n\ + <|im_start|>user\nWeather in Paris?<|im_end|>\n\ + <|im_start|>assistant\n"; + assert_eq!(raw, expected); + } + + /// LFM2.5's template injects NO default preamble: without a system turn + /// the tools line stands alone as the whole system prompt. + #[test] + fn lfm_tools_render_without_a_system_turn_has_no_preamble() { + let raw = ChatMl::lfm2().render_with_tools(&[weather_tool()], &[Turn::user("hi")]); + assert!(raw.starts_with("<|startoftext|><|im_start|>system\nList of tools: [")); + assert!(raw.contains("<|im_start|>user\nhi<|im_end|>")); + } + + #[test] + fn lfm_tools_comma_join_in_one_list() { + let mut second = weather_tool(); + second.name = "get_time".into(); + let raw = ChatMl::lfm2().render_with_tools(&[weather_tool(), second], &[Turn::user("hi")]); + assert!(raw.contains("\"name\":\"get_weather\"")); + assert!(raw.contains("}, {\"name\":\"get_time\"")); + assert_eq!(raw.matches("List of tools: [").count(), 1); + } + + /// LFM tool results are real `tool` role turns — one each, no Hermes + /// merging, no `` wrapper. + #[test] + fn lfm_tool_responses_render_as_tool_turns() { + let calls = [ToolCall { + name: "get_weather".into(), + arguments: serde_json::json!({"city": "Paris"}), + }]; + let raw = ChatMl::lfm2().render(&[ + Turn::user("Weather in Paris and Lyon?"), + Turn::assistant_tool_calls_lfm(&calls), + Turn::tool_response("{\"temp_c\": 21}"), + Turn::tool_response("{\"temp_c\": 24}"), + ]); + assert!(raw.contains( + "<|im_start|>assistant\n<|tool_call_start|>[get_weather(city=\"Paris\")]\ + <|tool_call_end|><|im_end|>\n" + )); + assert!(raw.contains("<|im_start|>tool\n{\"temp_c\": 21}<|im_end|>\n")); + assert!(raw.contains("<|im_start|>tool\n{\"temp_c\": 24}<|im_end|>\n")); + assert!(!raw.contains("")); + assert_eq!(raw.matches("<|im_start|>tool\n").count(), 2); + } + + /// Hermes rendering is untouched by the LFM additions: tool turns still + /// merge into a user turn. + #[test] + fn hermes_tool_turns_still_merge_after_the_style_split() { + let raw = ChatMl::qwen2().render(&[ + Turn::user("q"), + Turn::tool_response("r1"), + Turn::tool_response("r2"), + ]); + assert_eq!(raw.matches("<|im_start|>user").count(), 2); + assert!(!raw.contains("<|im_start|>tool")); + } + + /// The LFM2.5 template strips `` reasoning from every assistant + /// history turn but the LAST (keep_past_thinking=false default). + #[test] + fn lfm_strips_past_thinking_but_keeps_the_last() { + let raw = ChatMl::lfm2().render(&[ + Turn::user("a?"), + Turn::assistant("hmm\n\nAlpha."), + Turn::user("b?"), + Turn::assistant("later thoughts\n\nBeta."), + Turn::user("c?"), + ]); + assert!(raw.contains("<|im_start|>assistant\nAlpha.<|im_end|>")); + assert!(raw.contains("later thoughts")); + assert!(!raw.contains("hmm")); + } + + /// Qwen2 (Hermes) does no thinking-stripping — not part of its template. + #[test] + fn hermes_keeps_past_thinking_verbatim() { + let raw = ChatMl::qwen2().render(&[ + Turn::user("a?"), + Turn::assistant("hmmAlpha."), + Turn::user("b?"), + ]); + assert!(raw.contains("hmmAlpha.")); + } + + #[test] + fn pythonic_rendering_covers_the_scalar_spellings() { + let calls = [ToolCall { + name: "f".into(), + arguments: serde_json::json!({ + "s": "he said \"hi\"\n", + "i": -3, + "x": 1.5, + "yes": true, + "no": false, + "nothing": null, + "list": [1, "two"], + "map": {"k": true} + }), + }]; + let turn = Turn::assistant_tool_calls_lfm(&calls); + // serde_json object keys iterate in sorted order. + assert_eq!( + turn.content, + "<|tool_call_start|>[f(i=-3, list=[1, \"two\"], map={\"k\": True}, \ + no=False, nothing=None, s=\"he said \\\"hi\\\"\\n\", x=1.5, \ + yes=True)]<|tool_call_end|>" + ); + } + + /// The model card's own example parses to the exact call. + #[test] + fn lfm_parse_handles_the_model_card_example() { + let text = "<|tool_call_start|>[get_candidate_status(candidate_id=\"12345\")]\ + <|tool_call_end|>"; + let (calls, prose) = parse_tool_calls_lfm(text).unwrap(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].name, "get_candidate_status"); + assert_eq!(calls[0].arguments["candidate_id"], "12345"); + assert!(prose.is_empty()); + } + + #[test] + fn lfm_parse_extracts_multiple_calls_and_prose() { + let text = "Checking both.\n<|tool_call_start|>[get_weather(city=\"Paris\"), \ + get_weather(city='Lyon', units=None)]<|tool_call_end|> done"; + let (calls, prose) = parse_tool_calls_lfm(text).unwrap(); + assert_eq!(calls.len(), 2); + assert_eq!(calls[0].arguments["city"], "Paris"); + assert_eq!(calls[1].arguments["city"], "Lyon"); + assert_eq!(calls[1].arguments["units"], serde_json::Value::Null); + assert_eq!(prose, "Checking both.\n done"); + } + + #[test] + fn lfm_parse_of_plain_text_is_empty_not_an_error() { + let (calls, prose) = parse_tool_calls_lfm("The answer is 4.").unwrap(); + assert!(calls.is_empty()); + assert_eq!(prose, "The answer is 4."); + } + + #[test] + fn lfm_parse_accepts_python_and_json_literal_spellings() { + let text = "<|tool_call_start|>[f(a=True, b=false, c=null, d=None, \ + e=[1, 2.5, -3], g={\"k\": \"v\", 'k2': True,})]<|tool_call_end|>"; + let (calls, _) = parse_tool_calls_lfm(text).unwrap(); + let args = &calls[0].arguments; + assert_eq!(args["a"], true); + assert_eq!(args["b"], false); + assert_eq!(args["c"], serde_json::Value::Null); + assert_eq!(args["d"], serde_json::Value::Null); + assert_eq!(args["e"], serde_json::json!([1, 2.5, -3])); + assert_eq!(args["g"], serde_json::json!({"k": "v", "k2": true})); + } + + #[test] + fn lfm_parse_handles_string_escapes() { + let text = "<|tool_call_start|>[f(s='it\\'s \\\"q\\\" \\u00e9\\n')]<|tool_call_end|>"; + let (calls, _) = parse_tool_calls_lfm(text).unwrap(); + assert_eq!(calls[0].arguments["s"], "it's \"q\" \u{e9}\n"); + } + + #[test] + fn lfm_parse_rejects_malformed_blocks_loudly() { + // Unclosed special token. + assert!(matches!( + parse_tool_calls_lfm("<|tool_call_start|>[f()]"), + Err(ToolCallError::Unclosed { index: 0 }) + )); + // Not a call list. + assert!(matches!( + parse_tool_calls_lfm("<|tool_call_start|>f()<|tool_call_end|>"), + Err(ToolCallError::Syntax { index: 0, .. }) + )); + // Positional args are not in the grammar. + assert!(matches!( + parse_tool_calls_lfm("<|tool_call_start|>[f(\"paris\")]<|tool_call_end|>"), + Err(ToolCallError::Syntax { .. }) + )); + // Unterminated string. + assert!(matches!( + parse_tool_calls_lfm("<|tool_call_start|>[f(a=\"oops)]<|tool_call_end|>"), + Err(ToolCallError::Syntax { .. }) + )); + // Trailing junk after the list. + assert!(matches!( + parse_tool_calls_lfm("<|tool_call_start|>[f()] junk<|tool_call_end|>"), + Err(ToolCallError::Syntax { .. }) + )); + } + + #[test] + fn lfm_parse_bounds_depth_and_call_count() { + // 9 levels of list nesting exceeds MAX_VALUE_DEPTH = 8. + let deep = format!( + "<|tool_call_start|>[f(a={}1{})]<|tool_call_end|>", + "[".repeat(9), + "]".repeat(9) + ); + assert!(matches!( + parse_tool_calls_lfm(&deep), + Err(ToolCallError::TooDeep) + )); + let many = format!( + "<|tool_call_start|>[{}]<|tool_call_end|>", + vec!["f()"; MAX_TOOL_CALLS + 1].join(", ") + ); + assert!(matches!( + parse_tool_calls_lfm(&many), + Err(ToolCallError::TooMany) + )); + } + + /// The whole LFM loop: a rendered history turn re-parses to the same + /// calls, through the Pythonic spelling and back. + #[test] + fn lfm_tool_calls_round_trip_through_render_and_parse() { + let calls = vec![ToolCall { + name: "lookup".into(), + arguments: serde_json::json!({"q": "primes", "k": 5, "deep": {"a": [true, null]}}), + }]; + let turn = Turn::assistant_tool_calls_lfm(&calls); + let (parsed, prose) = parse_tool_calls_lfm(&turn.content).unwrap(); + assert_eq!(parsed, calls); + assert!(prose.is_empty()); + } } diff --git a/crates/mummu/src/gguf.rs b/crates/mummu/src/gguf.rs new file mode 100644 index 0000000..bad834a --- /dev/null +++ b/crates/mummu/src/gguf.rs @@ -0,0 +1,987 @@ +//! GGUF container reader — the first slice of P3's "run what the ecosystem +//! ships" import path. GGUF (llama.cpp's format) is one file: a small header +//! of typed metadata key-values, a tensor table (name, shape, quantized +//! dtype, offset), then an aligned blob of tensor payloads. +//! +//! This module reads the *header*: every metadata value typed and bounded, +//! every tensor located and size-checked — and fails loudly on anything +//! malformed, oversized, or unknown. Tensor payloads are *located*, never +//! loaded here; dequantizing them into Burn tensors is the next slice. + +use std::fs::File; +use std::io::{BufReader, Read, Seek}; +use std::path::Path; + +/// GGUF file magic, little-endian `"GGUF"`. +const MAGIC: [u8; 4] = *b"GGUF"; + +/// Versions this reader understands (v2 moved counts to u64; v3 is v2 plus a +/// big-endian variant this reader rejects by magic). +const SUPPORTED_VERSIONS: [u32; 2] = [2, 3]; + +/// Default payload alignment when `general.alignment` is absent. +const DEFAULT_ALIGNMENT: u64 = 32; + +/// Most metadata key-values a sane model file carries (real models: ~20-40). +const MAX_KVS: u64 = 4096; + +/// Most tensors a supported model carries (Qwen2.5-1.5B: 339). +const MAX_TENSORS: u64 = 65_536; + +/// Longest metadata string (chat templates run ~10 KiB; 1 MiB is generous). +const MAX_STRING_BYTES: u64 = 1 << 20; + +/// Longest metadata array (tokenizer vocab/merges run ~152k entries). +const MAX_ARRAY_LEN: u64 = 1 << 22; + +/// GGML allows at most 4 tensor dimensions. +const MAX_DIMS: u32 = 4; + +/// What went wrong reading a GGUF header. +#[derive(Debug, thiserror::Error)] +pub enum GgufError { + #[error("gguf {path}: {source}")] + Io { + path: String, + source: std::io::Error, + }, + #[error("gguf {path}: bad magic {found:02x?} (big-endian GGUF is unsupported)")] + BadMagic { path: String, found: [u8; 4] }, + #[error("gguf {path}: unsupported version {version} (supported: {SUPPORTED_VERSIONS:?})")] + UnsupportedVersion { path: String, version: u32 }, + #[error("gguf {path}: {what} count {count} exceeds the {bound} bound")] + OverBound { + path: String, + what: &'static str, + count: u64, + bound: u64, + }, + #[error("gguf {path}: metadata '{key}': {reason}")] + BadValue { + path: String, + key: String, + reason: String, + }, + #[error("gguf {path}: tensor {index}: {reason}")] + BadTensor { + path: String, + index: usize, + reason: String, + }, +} + +/// One typed metadata value. Arrays are homogeneous per the spec; nested +/// arrays are legal but bounded to one level of nesting in practice. +#[derive(Debug, Clone, PartialEq)] +pub enum GgufValue { + U8(u8), + I8(i8), + U16(u16), + I16(i16), + U32(u32), + I32(i32), + F32(f32), + Bool(bool), + Str(String), + Array(Vec), + U64(u64), + I64(i64), + F64(f64), +} + +impl GgufValue { + /// The value as a string, if it is one. + #[must_use] + pub fn as_str(&self) -> Option<&str> { + match self { + Self::Str(s) => Some(s), + _ => None, + } + } + + /// The value widened to u64, if it is any unsigned integer. + #[must_use] + pub fn as_u64(&self) -> Option { + match *self { + Self::U8(v) => Some(u64::from(v)), + Self::U16(v) => Some(u64::from(v)), + Self::U32(v) => Some(u64::from(v)), + Self::U64(v) => Some(v), + _ => None, + } + } + + /// The value as f32, if it is one. + #[must_use] + pub fn as_f32(&self) -> Option { + match *self { + Self::F32(v) => Some(v), + _ => None, + } + } + + /// The value as an array slice, if it is one. + #[must_use] + pub fn as_array(&self) -> Option<&[GgufValue]> { + match self { + Self::Array(items) => Some(items), + _ => None, + } + } +} + +/// A GGML tensor dtype, as stored on disk. Quantized types pack fixed-size +/// blocks; `block_size`/`bytes_per_block` give the layout the dequant slice +/// (and size validation here) needs. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[allow(non_camel_case_types)] // the ecosystem's canonical spellings +pub enum GgmlType { + F32, + F16, + Q4_0, + Q4_1, + Q5_0, + Q5_1, + Q8_0, + Q8_1, + Q2_K, + Q3_K, + Q4_K, + Q5_K, + Q6_K, + Q8_K, + BF16, +} + +impl GgmlType { + /// Decode the on-disk type id; unknown ids are a loud error, never a guess. + fn from_id(id: u32) -> Option { + let ty = match id { + 0 => Self::F32, + 1 => Self::F16, + 2 => Self::Q4_0, + 3 => Self::Q4_1, + 6 => Self::Q5_0, + 7 => Self::Q5_1, + 8 => Self::Q8_0, + 9 => Self::Q8_1, + 10 => Self::Q2_K, + 11 => Self::Q3_K, + 12 => Self::Q4_K, + 13 => Self::Q5_K, + 14 => Self::Q6_K, + 15 => Self::Q8_K, + 30 => Self::BF16, + _ => return None, + }; + Some(ty) + } + + /// Elements per quantization block (1 for plain float types). + #[must_use] + pub fn block_size(self) -> u64 { + match self { + Self::F32 | Self::F16 | Self::BF16 => 1, + Self::Q4_0 | Self::Q4_1 | Self::Q5_0 | Self::Q5_1 | Self::Q8_0 | Self::Q8_1 => 32, + Self::Q2_K | Self::Q3_K | Self::Q4_K | Self::Q5_K | Self::Q6_K | Self::Q8_K => 256, + } + } + + /// Bytes one block occupies on disk (ggml's type sizes). + #[must_use] + pub fn bytes_per_block(self) -> u64 { + match self { + Self::F32 => 4, + Self::F16 | Self::BF16 => 2, + Self::Q4_0 => 18, // f16 d + 16 B qs + Self::Q4_1 => 20, // f16 d + f16 m + 16 B qs + Self::Q5_0 => 22, // f16 d + 4 B qh + 16 B qs + Self::Q5_1 => 24, // f16 d + f16 m + 4 B qh + 16 B qs + Self::Q8_0 => 34, // f16 d + 32 i8 + Self::Q8_1 => 36, // f16 d + f16 s + 32 i8 + Self::Q2_K => 84, // 16 B scales + 64 B qs + f16 d + f16 dmin + Self::Q3_K => 110, // 32 B hmask + 64 B qs + 12 B scales + f16 d + Self::Q4_K => 144, // f16 d + f16 dmin + 12 B scales + 128 B qs + Self::Q5_K => 176, // Q4_K + 32 B qh + Self::Q6_K => 210, // 128 B ql + 64 B qh + 16 i8 scales + f16 d + Self::Q8_K => 292, // f32 d + 256 i8 + 16 i16 bsums + } + } +} + +/// One entry of the tensor table: where a tensor lives and what shape/dtype +/// it has. `offset` is relative to [`GgufFile::data_offset`] and is a +/// multiple of the file's alignment (validated on read). +#[derive(Debug, Clone)] +pub struct GgufTensorInfo { + pub name: String, + /// Dimensions in ggml order (fastest-varying first — the *reverse* of + /// the row-major order safetensors/PyTorch shapes use). + pub dims: Vec, + pub dtype: GgmlType, + pub offset: u64, +} + +impl GgufTensorInfo { + /// Total element count. + #[must_use] + pub fn element_count(&self) -> u64 { + self.dims.iter().product() + } + + /// Exact on-disk payload size. Element counts are validated to be whole + /// blocks at parse time, so this is always exact for parsed tensors. + #[must_use] + pub fn byte_len(&self) -> u64 { + let elements = self.element_count(); + debug_assert!( + elements.is_multiple_of(self.dtype.block_size()), + "parse validated whole blocks" + ); + elements / self.dtype.block_size() * self.dtype.bytes_per_block() + } +} + +/// A parsed GGUF header: typed metadata + the located tensor table. +#[derive(Debug)] +pub struct GgufFile { + /// Where this header was read from (payload reads re-open it). + pub path: std::path::PathBuf, + pub version: u32, + /// Metadata in file order (keys are unique per the spec). + pub metadata: Vec<(String, GgufValue)>, + pub tensors: Vec, + /// Payload alignment (`general.alignment`, default 32). + pub alignment: u64, + /// Absolute file offset where the aligned tensor payload blob begins. + pub data_offset: u64, +} + +impl GgufFile { + /// Read and validate a GGUF header (metadata + tensor table only — no + /// tensor payloads are loaded). + pub fn open(path: &Path) -> Result { + let file = File::open(path).map_err(|source| GgufError::Io { + path: path.display().to_string(), + source, + })?; + let mut r = Reader { + inner: BufReader::new(file), + path: path.display().to_string(), + }; + let parsed = r.read_file(path)?; + // Positive space: the header must end before the payload it locates. + assert!( + parsed.data_offset.is_multiple_of(parsed.alignment), + "data offset is aligned by construction" + ); + Ok(parsed) + } + + /// Read one tensor's payload and dequantize it to f32, in the on-disk + /// (ggml fastest-varying-first) element order. + pub fn read_tensor_f32(&self, name: &str) -> Result, GgufError> { + use std::io::SeekFrom; + let info = self.tensor(name).ok_or_else(|| GgufError::BadValue { + path: self.path.display().to_string(), + key: name.to_string(), + reason: "no such tensor".into(), + })?; + let mut file = File::open(&self.path).map_err(|source| GgufError::Io { + path: self.path.display().to_string(), + source, + })?; + file.seek(SeekFrom::Start(self.data_offset + info.offset)) + .map_err(|source| GgufError::Io { + path: self.path.display().to_string(), + source, + })?; + let byte_len = usize::try_from(info.byte_len()).map_err(|_| GgufError::OverBound { + path: self.path.display().to_string(), + what: "tensor payload bytes", + count: info.byte_len(), + bound: usize::MAX as u64, + })?; + let mut bytes = vec![0u8; byte_len]; + file.read_exact(&mut bytes) + .map_err(|source| GgufError::Io { + path: self.path.display().to_string(), + source, + })?; + let out = dequantize(info.dtype, &bytes).map_err(|reason| GgufError::BadTensor { + path: self.path.display().to_string(), + index: 0, + reason: format!("{name}: {reason}"), + })?; + assert_eq!( + out.len() as u64, + info.element_count(), + "dequant must yield exactly the tensor's elements" + ); + Ok(out) + } + + /// Look up a metadata value by exact key. + #[must_use] + pub fn get(&self, key: &str) -> Option<&GgufValue> { + self.metadata.iter().find(|(k, _)| k == key).map(|(_, v)| v) + } + + /// The model architecture (`general.architecture`), when present. + #[must_use] + pub fn architecture(&self) -> Option<&str> { + self.get("general.architecture").and_then(GgufValue::as_str) + } + + /// Look up a tensor by exact name. + #[must_use] + pub fn tensor(&self, name: &str) -> Option<&GgufTensorInfo> { + self.tensors.iter().find(|t| t.name == name) + } +} + +/// Sequential little-endian reader over the header bytes. +struct Reader { + inner: BufReader, + path: String, +} + +impl Reader { + fn io_err(&self, source: std::io::Error) -> GgufError { + GgufError::Io { + path: self.path.clone(), + source, + } + } + + fn bytes(&mut self) -> Result<[u8; N], GgufError> { + let mut buf = [0u8; N]; + self.inner + .read_exact(&mut buf) + .map_err(|e| self.io_err(e))?; + Ok(buf) + } + + fn u32(&mut self) -> Result { + Ok(u32::from_le_bytes(self.bytes()?)) + } + + fn u64(&mut self) -> Result { + Ok(u64::from_le_bytes(self.bytes()?)) + } + + /// A length-prefixed UTF-8 string, bounded by [`MAX_STRING_BYTES`]. + fn string(&mut self, what: &'static str) -> Result { + let len = self.u64()?; + if len > MAX_STRING_BYTES { + return Err(GgufError::OverBound { + path: self.path.clone(), + what, + count: len, + bound: MAX_STRING_BYTES, + }); + } + let mut buf = vec![0u8; usize::try_from(len).expect("bounded above")]; + self.inner + .read_exact(&mut buf) + .map_err(|e| self.io_err(e))?; + String::from_utf8(buf).map_err(|e| GgufError::BadValue { + path: self.path.clone(), + key: what.to_string(), + reason: format!("invalid UTF-8: {e}"), + }) + } + + /// One typed metadata value. `depth` bounds array nesting. + fn value(&mut self, key: &str, type_id: u32, depth: u32) -> Result { + let bad = |reason: String, path: &str| GgufError::BadValue { + path: path.to_string(), + key: key.to_string(), + reason, + }; + let v = match type_id { + 0 => GgufValue::U8(self.bytes::<1>()?[0]), + #[allow(clippy::cast_possible_wrap)] // bit-exact reinterpret is the format + 1 => GgufValue::I8(self.bytes::<1>()?[0] as i8), + 2 => GgufValue::U16(u16::from_le_bytes(self.bytes()?)), + 3 => GgufValue::I16(i16::from_le_bytes(self.bytes()?)), + 4 => GgufValue::U32(self.u32()?), + 5 => GgufValue::I32(i32::from_le_bytes(self.bytes()?)), + 6 => GgufValue::F32(f32::from_le_bytes(self.bytes()?)), + 7 => match self.bytes::<1>()?[0] { + 0 => GgufValue::Bool(false), + 1 => GgufValue::Bool(true), + other => return Err(bad(format!("bool byte {other}"), &self.path)), + }, + 8 => GgufValue::Str(self.string("metadata string")?), + 9 => { + if depth >= 2 { + return Err(bad("arrays nested deeper than 2".into(), &self.path)); + } + let elem_type = self.u32()?; + let len = self.u64()?; + if len > MAX_ARRAY_LEN { + return Err(GgufError::OverBound { + path: self.path.clone(), + what: "metadata array", + count: len, + bound: MAX_ARRAY_LEN, + }); + } + let mut items = Vec::with_capacity(usize::try_from(len).expect("bounded above")); + for _ in 0..len { + items.push(self.value(key, elem_type, depth + 1)?); + } + GgufValue::Array(items) + } + 10 => GgufValue::U64(self.u64()?), + 11 => GgufValue::I64(i64::from_le_bytes(self.bytes()?)), + 12 => GgufValue::F64(f64::from_le_bytes(self.bytes()?)), + other => return Err(bad(format!("unknown value type {other}"), &self.path)), + }; + Ok(v) + } + + /// The whole header: magic, version, metadata, tensor table, alignment. + fn read_file(&mut self, source_path: &Path) -> Result { + let magic = self.bytes::<4>()?; + if magic != MAGIC { + return Err(GgufError::BadMagic { + path: self.path.clone(), + found: magic, + }); + } + let version = self.u32()?; + if !SUPPORTED_VERSIONS.contains(&version) { + return Err(GgufError::UnsupportedVersion { + path: self.path.clone(), + version, + }); + } + let tensor_count = self.u64()?; + let kv_count = self.u64()?; + for (what, count, bound) in [ + ("tensor", tensor_count, MAX_TENSORS), + ("metadata kv", kv_count, MAX_KVS), + ] { + if count > bound { + return Err(GgufError::OverBound { + path: self.path.clone(), + what, + count, + bound, + }); + } + } + + let mut metadata = Vec::with_capacity(usize::try_from(kv_count).expect("bounded above")); + for _ in 0..kv_count { + let key = self.string("metadata key")?; + let type_id = self.u32()?; + let value = self.value(&key, type_id, 0)?; + metadata.push((key, value)); + } + + let alignment = metadata + .iter() + .find(|(k, _)| k == "general.alignment") + .and_then(|(_, v)| v.as_u64()) + .unwrap_or(DEFAULT_ALIGNMENT); + if alignment == 0 || !alignment.is_power_of_two() { + return Err(GgufError::BadValue { + path: self.path.clone(), + key: "general.alignment".into(), + reason: format!("{alignment} is not a power of two"), + }); + } + + let tensors = self.tensor_table(tensor_count, alignment)?; + let header_end = self.inner.stream_position().map_err(|e| self.io_err(e))?; + let data_offset = header_end.div_ceil(alignment) * alignment; + debug_assert!(data_offset >= header_end, "padding never rewinds"); + Ok(GgufFile { + path: source_path.to_path_buf(), + version, + metadata, + tensors, + alignment, + data_offset, + }) + } + + /// The tensor table, with every entry's shape/dtype/offset validated. + fn tensor_table( + &mut self, + count: u64, + alignment: u64, + ) -> Result, GgufError> { + assert!(count <= MAX_TENSORS, "caller bounded the count"); + assert!(alignment.is_power_of_two(), "caller validated alignment"); + let mut tensors: Vec = + Vec::with_capacity(usize::try_from(count).expect("bounded above")); + for index in 0..usize::try_from(count).expect("bounded above") { + let bad = |reason: String, path: &str| GgufError::BadTensor { + path: path.to_string(), + index, + reason, + }; + let name = self.string("tensor name")?; + let n_dims = self.u32()?; + if n_dims == 0 || n_dims > MAX_DIMS { + return Err(bad(format!("{n_dims} dims (1..={MAX_DIMS})"), &self.path)); + } + let mut dims = Vec::with_capacity(n_dims as usize); + for _ in 0..n_dims { + dims.push(self.u64()?); + } + let type_id = self.u32()?; + let Some(dtype) = GgmlType::from_id(type_id) else { + return Err(bad(format!("unknown ggml type id {type_id}"), &self.path)); + }; + let offset = self.u64()?; + if !offset.is_multiple_of(alignment) { + return Err(bad( + format!("offset {offset} not {alignment}-aligned"), + &self.path, + )); + } + let elements: u64 = dims.iter().product(); + if elements == 0 || !elements.is_multiple_of(dtype.block_size()) { + return Err(bad( + format!("{elements} elements is not whole {dtype:?} blocks"), + &self.path, + )); + } + if tensors.iter().any(|t| t.name == name) { + return Err(bad(format!("duplicate tensor name '{name}'"), &self.path)); + } + tensors.push(GgufTensorInfo { + name, + dims, + dtype, + offset, + }); + } + Ok(tensors) + } +} + +// ---- Dequantization ------------------------------------------------------ +// +// Exact ports of ggml's reference dequantizers (ggml-quants.c) for the types +// a Q4_K_M file actually carries: plain floats, Q8_0, and the K-quant +// superblocks Q4_K / Q6_K. Layouts follow `GgmlType::bytes_per_block`. + +/// Dequantize a whole tensor payload to f32. `bytes` must be whole blocks of +/// `dtype` (guaranteed for payload slices sized by [`GgufTensorInfo::byte_len`]). +pub fn dequantize(dtype: GgmlType, bytes: &[u8]) -> Result, String> { + let bpb = usize::try_from(dtype.bytes_per_block()).expect("small"); + if bytes.is_empty() || !bytes.len().is_multiple_of(bpb) { + return Err(format!( + "{} bytes is not whole {dtype:?} blocks of {bpb}", + bytes.len() + )); + } + let blocks = bytes.len() / bpb; + let block_elems = usize::try_from(dtype.block_size()).expect("small"); + let mut out = Vec::with_capacity(blocks * block_elems); + for block in bytes.chunks_exact(bpb) { + match dtype { + GgmlType::F32 => out.push(f32::from_le_bytes(block.try_into().expect("4 bytes"))), + GgmlType::F16 => out.push(f16_to_f32(u16::from_le_bytes([block[0], block[1]]))), + GgmlType::BF16 => { + out.push(f32::from_bits( + u32::from(u16::from_le_bytes([block[0], block[1]])) << 16, + )); + } + GgmlType::Q8_0 => dequant_q8_0(block, &mut out), + GgmlType::Q4_K => dequant_q4_k(block, &mut out), + GgmlType::Q6_K => dequant_q6_k(block, &mut out), + other => return Err(format!("dequant for {other:?} is not implemented yet")), + } + } + assert_eq!(out.len(), blocks * block_elems, "whole blocks out"); + Ok(out) +} + +/// IEEE 754 half → f32 (no `half` dep on this path; exhaustive over u16 in +/// tests against the `half` crate the workspace already carries). +fn f16_to_f32(bits: u16) -> f32 { + f32::from(half::f16::from_bits(bits)) +} + +/// Q8_0: f16 scale + 32 signed bytes; `x = d * q`. +fn dequant_q8_0(block: &[u8], out: &mut Vec) { + assert_eq!(block.len(), 34, "Q8_0 block is 34 bytes"); + let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]])); + #[allow(clippy::cast_possible_wrap)] // bit-exact reinterpret is the format + out.extend(block[2..34].iter().map(|&q| d * f32::from(q as i8))); +} + +/// The Q4_K/Q5_K 6-bit (scale, min) pair for sub-block `j` — ggml's +/// `get_scale_min_k4`. +fn scale_min_k4(scales: &[u8], j: usize) -> (f32, f32) { + assert_eq!(scales.len(), 12, "K-quant scale block is 12 bytes"); + assert!(j < 8, "8 sub-blocks per superblock"); + let (sc, m) = if j < 4 { + (scales[j] & 63, scales[j + 4] & 63) + } else { + ( + (scales[j + 4] & 0x0F) | ((scales[j - 4] >> 6) << 4), + (scales[j + 4] >> 4) | ((scales[j] >> 6) << 4), + ) + }; + (f32::from(sc), f32::from(m)) +} + +/// Q4_K: 256-element superblock — f16 d + f16 dmin + 12 B packed 6-bit +/// (scale, min) pairs + 128 B of 4-bit quants; `x = d·sc·q − dmin·m`. +fn dequant_q4_k(block: &[u8], out: &mut Vec) { + assert_eq!(block.len(), 144, "Q4_K superblock is 144 bytes"); + let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]])); + let dmin = f16_to_f32(u16::from_le_bytes([block[2], block[3]])); + let scales = &block[4..16]; + let qs = &block[16..144]; + // 4 chunks of 64 values; each chunk reads 32 bytes — low nibbles first. + for chunk in 0..4 { + let (sc1, m1) = scale_min_k4(scales, chunk * 2); + let (sc2, m2) = scale_min_k4(scales, chunk * 2 + 1); + let q = &qs[chunk * 32..chunk * 32 + 32]; + out.extend(q.iter().map(|&b| d * sc1 * f32::from(b & 0x0F) - dmin * m1)); + out.extend(q.iter().map(|&b| d * sc2 * f32::from(b >> 4) - dmin * m2)); + } +} + +/// Q6_K: 256-element superblock — 128 B low-4 + 64 B high-2 + 16 i8 +/// sub-scales + f16 d; `x = d·sc·(q − 32)`. +fn dequant_q6_k(block: &[u8], out: &mut Vec) { + assert_eq!(block.len(), 210, "Q6_K superblock is 210 bytes"); + let (ql_all, rest) = block.split_at(128); + let (qh_all, rest) = rest.split_at(64); + let (scales, d_bytes) = rest.split_at(16); + let d = f16_to_f32(u16::from_le_bytes([d_bytes[0], d_bytes[1]])); + let start = out.len(); + out.resize(start + 256, 0.0); + let y = &mut out[start..]; + // Two halves of 128 values, each consuming 64 ql / 32 qh / 8 scales. + for half_idx in 0..2 { + let ql = &ql_all[half_idx * 64..half_idx * 64 + 64]; + let qh = &qh_all[half_idx * 32..half_idx * 32 + 32]; + let sc = &scales[half_idx * 8..half_idx * 8 + 8]; + let base = half_idx * 128; + for l in 0..32 { + let is = l / 16; + let q1 = i16::from((ql[l] & 0x0F) | ((qh[l] & 3) << 4)) - 32; + let q2 = i16::from((ql[l + 32] & 0x0F) | (((qh[l] >> 2) & 3) << 4)) - 32; + let q3 = i16::from((ql[l] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32; + let q4 = i16::from((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32; + #[allow(clippy::cast_possible_wrap)] // i8 sub-scales are the format + let s = |i: usize| f32::from(sc[i] as i8); + y[base + l] = d * s(is) * f32::from(q1); + y[base + l + 32] = d * s(is + 2) * f32::from(q2); + y[base + l + 64] = d * s(is + 4) * f32::from(q3); + y[base + l + 96] = d * s(is + 6) * f32::from(q4); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Write; + + /// Minimal in-memory GGUF builder for tests. + struct TestGguf { + buf: Vec, + tensor_count: u64, + kv_count: u64, + kvs: Vec, + tensors: Vec, + } + + impl TestGguf { + fn new() -> Self { + Self { + buf: Vec::new(), + tensor_count: 0, + kv_count: 0, + kvs: Vec::new(), + tensors: Vec::new(), + } + } + + fn push_str(out: &mut Vec, s: &str) { + out.extend_from_slice(&(s.len() as u64).to_le_bytes()); + out.extend_from_slice(s.as_bytes()); + } + + fn kv_str(mut self, key: &str, value: &str) -> Self { + Self::push_str(&mut self.kvs, key); + self.kvs.extend_from_slice(&8u32.to_le_bytes()); + Self::push_str(&mut self.kvs, value); + self.kv_count += 1; + self + } + + fn kv_u32(mut self, key: &str, value: u32) -> Self { + Self::push_str(&mut self.kvs, key); + self.kvs.extend_from_slice(&4u32.to_le_bytes()); + self.kvs.extend_from_slice(&value.to_le_bytes()); + self.kv_count += 1; + self + } + + fn kv_str_array(mut self, key: &str, values: &[&str]) -> Self { + Self::push_str(&mut self.kvs, key); + self.kvs.extend_from_slice(&9u32.to_le_bytes()); + self.kvs.extend_from_slice(&8u32.to_le_bytes()); + self.kvs + .extend_from_slice(&(values.len() as u64).to_le_bytes()); + for v in values { + Self::push_str(&mut self.kvs, v); + } + self.kv_count += 1; + self + } + + fn tensor(mut self, name: &str, dims: &[u64], type_id: u32, offset: u64) -> Self { + Self::push_str(&mut self.tensors, name); + self.tensors + .extend_from_slice(&(dims.len() as u32).to_le_bytes()); + for d in dims { + self.tensors.extend_from_slice(&d.to_le_bytes()); + } + self.tensors.extend_from_slice(&type_id.to_le_bytes()); + self.tensors.extend_from_slice(&offset.to_le_bytes()); + self.tensor_count += 1; + self + } + + fn build(mut self) -> Vec { + self.buf.extend_from_slice(&MAGIC); + self.buf.extend_from_slice(&3u32.to_le_bytes()); + self.buf.extend_from_slice(&self.tensor_count.to_le_bytes()); + self.buf.extend_from_slice(&self.kv_count.to_le_bytes()); + self.buf.extend_from_slice(&self.kvs); + self.buf.extend_from_slice(&self.tensors); + self.buf + } + } + + fn open_bytes(bytes: &[u8]) -> Result { + use std::sync::atomic::{AtomicU64, Ordering}; + // Parallel tests in one process must never share a temp file. + static NEXT: AtomicU64 = AtomicU64::new(0); + let dir = std::env::temp_dir().join("mummu-gguf-tests"); + std::fs::create_dir_all(&dir).expect("temp dir"); + let path = dir.join(format!( + "t-{}-{}.gguf", + std::process::id(), + NEXT.fetch_add(1, Ordering::Relaxed) + )); + let mut f = File::create(&path).expect("temp file"); + f.write_all(bytes).expect("write"); + drop(f); + let result = GgufFile::open(&path); + let _ = std::fs::remove_file(&path); + result + } + + #[test] + fn minimal_file_round_trips() { + let bytes = TestGguf::new() + .kv_str("general.architecture", "qwen2") + .kv_u32("qwen2.block_count", 28) + .kv_str_array("tokenizer.ggml.tokens", &["a", "b", "c"]) + .tensor("token_embd.weight", &[64, 2], 0, 0) + .tensor("blk.0.attn_q.weight", &[256], 12, 512) + .build(); + let f = open_bytes(&bytes).expect("parses"); + assert_eq!(f.version, 3); + assert_eq!(f.architecture(), Some("qwen2")); + assert_eq!( + f.get("qwen2.block_count").and_then(GgufValue::as_u64), + Some(28) + ); + assert_eq!( + f.get("tokenizer.ggml.tokens") + .and_then(GgufValue::as_array) + .map(<[GgufValue]>::len), + Some(3) + ); + let embd = f.tensor("token_embd.weight").expect("present"); + assert_eq!(embd.dims, vec![64, 2]); + assert_eq!(embd.dtype, GgmlType::F32); + assert_eq!(embd.byte_len(), 64 * 2 * 4); + let q = f.tensor("blk.0.attn_q.weight").expect("present"); + assert_eq!(q.dtype, GgmlType::Q4_K); + assert_eq!(q.byte_len(), 144); // one 256-element Q4_K superblock + assert_eq!(f.alignment, DEFAULT_ALIGNMENT); + assert!(f.data_offset.is_multiple_of(f.alignment)); + assert!(f.data_offset >= (bytes.len() as u64)); + } + + #[test] + fn bad_magic_is_rejected() { + let mut bytes = TestGguf::new().build(); + bytes[..4].copy_from_slice(b"FUGG"); + assert!(matches!( + open_bytes(&bytes), + Err(GgufError::BadMagic { .. }) + )); + } + + #[test] + fn unsupported_version_is_rejected() { + let mut bytes = TestGguf::new().build(); + bytes[4..8].copy_from_slice(&1u32.to_le_bytes()); + assert!(matches!( + open_bytes(&bytes), + Err(GgufError::UnsupportedVersion { version: 1, .. }) + )); + } + + #[test] + fn truncated_file_is_an_io_error_not_a_hang() { + let bytes = TestGguf::new() + .kv_str("general.architecture", "qwen2") + .build(); + assert!(matches!( + open_bytes(&bytes[..bytes.len() - 3]), + Err(GgufError::Io { .. }) + )); + } + + #[test] + fn oversized_counts_are_rejected() { + let mut bytes = TestGguf::new().build(); + // tensor_count lives at bytes 8..16. + bytes[8..16].copy_from_slice(&(MAX_TENSORS + 1).to_le_bytes()); + assert!(matches!( + open_bytes(&bytes), + Err(GgufError::OverBound { what: "tensor", .. }) + )); + } + + #[test] + fn unknown_value_type_and_ggml_type_are_rejected() { + let mut with_kv = TestGguf::new().kv_u32("some.key", 1).build(); + // The kv's type id (4 = u32) sits right after the 8-byte key string + // prefix + 8 bytes of key: magic(4)+ver(4)+counts(16)+len(8)+key(8). + with_kv[40..44].copy_from_slice(&99u32.to_le_bytes()); + assert!(matches!( + open_bytes(&with_kv), + Err(GgufError::BadValue { .. }) + )); + + let with_tensor = TestGguf::new().tensor("t", &[32], 63, 0).build(); + assert!(matches!( + open_bytes(&with_tensor), + Err(GgufError::BadTensor { .. }) + )); + } + + // ---- dequant --------------------------------------------------------- + + fn f16_bytes(v: f32) -> [u8; 2] { + half::f16::from_f32(v).to_bits().to_le_bytes() + } + + #[test] + fn float_widths_dequantize_exactly() { + let f32_bytes = 1.5f32.to_le_bytes(); + assert_eq!(dequantize(GgmlType::F32, &f32_bytes).unwrap(), vec![1.5]); + assert_eq!( + dequantize(GgmlType::F16, &f16_bytes(-0.25)).unwrap(), + vec![-0.25] + ); + // bf16 is the top half of the f32 bit pattern. + let bf16 = (2.0f32.to_bits() >> 16) as u16; + assert_eq!( + dequantize(GgmlType::BF16, &bf16.to_le_bytes()).unwrap(), + vec![2.0] + ); + } + + #[test] + fn q8_0_block_matches_hand_computation() { + let mut block = Vec::new(); + block.extend_from_slice(&f16_bytes(0.5)); + #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] + block.extend((0..32).map(|i| (i - 16) as i8 as u8)); + let out = dequantize(GgmlType::Q8_0, &block).unwrap(); + assert_eq!(out.len(), 32); + assert_eq!(out[0], 0.5 * -16.0); + assert_eq!(out[16], 0.0); + assert_eq!(out[31], 0.5 * 15.0); + } + + #[test] + fn q4_k_superblock_matches_hand_computation() { + let mut block = vec![0u8; 144]; + block[0..2].copy_from_slice(&f16_bytes(1.0)); // d + block[2..4].copy_from_slice(&f16_bytes(0.5)); // dmin + // Sub-block 0: sc=2, m=1 · sub-block 1: sc=3, m=0 (direct 6-bit slots). + block[4] = 2; + block[5] = 3; + block[8] = 1; + // Sub-block 4: packed slot — sc = scales[8] & 0xF = 1, m = scales[8] >> 4 = 2. + block[12] = 0x21; + // First quant byte of chunk 0: low nibble 1 (sub 0), high nibble 5 (sub 1). + block[16] = 0x51; + // First quant byte of chunk 2 (sub-blocks 4/5): low nibble 4. + block[16 + 64] = 0x04; + let out = dequantize(GgmlType::Q4_K, &block).unwrap(); + assert_eq!(out.len(), 256); + assert_eq!(out[0], 2.0 * 1.0 - 0.5 * 1.0); // d·sc0·q − dmin·m0 = 1.5 + assert_eq!(out[32], 3.0 * 5.0); // sub 1: m=0 + assert_eq!(out[128], 1.0 * 4.0 - 0.5 * 2.0); // sub 4 via packed scales + // A zero quant in sub-block 0 still subtracts the min. + assert_eq!(out[1], -0.5); + } + + #[test] + fn q6_k_superblock_matches_hand_computation() { + let mut block = vec![0u8; 210]; + block[0] = 0x0F; // ql[0]: low 4 bits = 15 + block[128] = 0b0000_0011; // qh[0]: high 2 bits = 3 for q1 + block[192] = 2; // scales[0] = 2 + block[194] = 1; // scales[2] = 1 + block[208..210].copy_from_slice(&f16_bytes(1.0)); // d + let out = dequantize(GgmlType::Q6_K, &block).unwrap(); + assert_eq!(out.len(), 256); + // q1 = (15 | 3<<4) − 32 = 31, scale 2 → 62. + assert_eq!(out[0], 62.0); + // q2 = (0 | 0) − 32 = −32, scale sc[2]=1 → −32. + assert_eq!(out[32], -32.0); + // Zero scale zeroes the value even though q3 = −32. + assert_eq!(out[64], 0.0); + } + + #[test] + fn dequant_rejects_partial_blocks_and_unimplemented_types() { + assert!(dequantize(GgmlType::Q8_0, &[0u8; 33]).is_err()); + assert!(dequantize(GgmlType::Q8_0, &[]).is_err()); + assert!(dequantize(GgmlType::Q2_K, &[0u8; 84]).is_err()); + } + + #[test] + fn misaligned_offsets_partial_blocks_and_duplicates_are_rejected() { + let misaligned = TestGguf::new().tensor("t", &[32], 8, 7).build(); + assert!(matches!( + open_bytes(&misaligned), + Err(GgufError::BadTensor { .. }) + )); + // 100 elements is not whole 256-element Q4_K superblocks. + let partial = TestGguf::new().tensor("t", &[100], 12, 0).build(); + assert!(matches!( + open_bytes(&partial), + Err(GgufError::BadTensor { .. }) + )); + let dup = TestGguf::new() + .tensor("t", &[32], 8, 0) + .tensor("t", &[32], 8, 64) + .build(); + assert!(matches!(open_bytes(&dup), Err(GgufError::BadTensor { .. }))); + } +} diff --git a/crates/mummu/src/lib.rs b/crates/mummu/src/lib.rs index cc6b7e8..8239a95 100644 --- a/crates/mummu/src/lib.rs +++ b/crates/mummu/src/lib.rs @@ -13,6 +13,7 @@ pub mod backend; pub mod cache; pub mod chat; pub mod decode; +pub mod gguf; pub mod hub; pub mod import; pub mod manage; diff --git a/crates/mummu/tests/real_gguf.rs b/crates/mummu/tests/real_gguf.rs new file mode 100644 index 0000000..5ed947d --- /dev/null +++ b/crates/mummu/tests/real_gguf.rs @@ -0,0 +1,161 @@ +//! Real-file GGUF header proof: parse an actual llama.cpp-quantized model +//! (Qwen2.5-1.5B-Instruct Q4_K_M) and check the header describes the model +//! we know. Ignored by default; run with +//! +//! ```text +//! MUMMU_GGUF_PATH=path/to/qwen2.5-1.5b-instruct-q4_k_m.gguf cargo test -p mummu --test real_gguf -- --ignored --nocapture +//! ``` + +use std::path::PathBuf; + +use mummu::gguf::{GgmlType, GgufFile, GgufValue}; + +fn gguf_path() -> Option { + let path = PathBuf::from(std::env::var_os("MUMMU_GGUF_PATH")?); + path.is_file().then_some(path) +} + +#[test] +#[ignore = "needs a local GGUF model file (MUMMU_GGUF_PATH)"] +fn real_qwen2_gguf_header_parses_and_describes_the_model() { + let Some(path) = gguf_path() else { + panic!("set MUMMU_GGUF_PATH to a local .gguf model file"); + }; + let f = GgufFile::open(&path).expect("header parses"); + + eprintln!( + "[real_gguf] v{} · {} kvs · {} tensors · align {} · data at {}", + f.version, + f.metadata.len(), + f.tensors.len(), + f.alignment, + f.data_offset + ); + assert_eq!(f.architecture(), Some("qwen2")); + + // The tokenizer vocab must be a real array of strings. + let tokens = f + .get("tokenizer.ggml.tokens") + .and_then(GgufValue::as_array) + .expect("vocab array present"); + assert!( + tokens.len() > 100_000, + "Qwen vocab is ~152k: {}", + tokens.len() + ); + + // The embedding tensor exists with the config's dims (ggml order: + // fastest-varying first — [hidden, vocab]). + let embd = f.tensor("token_embd.weight").expect("embedding present"); + assert_eq!(embd.dims[0], 1536, "Qwen2.5-1.5B hidden size"); + assert_eq!(embd.dims[1] as usize, tokens.len(), "vocab rows"); + + // Every tensor: known dtype (the parser guarantees it), aligned offset, + // whole blocks; and at least one K-quant tensor is actually present. + for t in &f.tensors { + assert!(t.offset.is_multiple_of(f.alignment), "{}", t.name); + assert!(t.byte_len() > 0, "{}", t.name); + } + let kquants = f + .tensors + .iter() + .filter(|t| matches!(t.dtype, GgmlType::Q4_K | GgmlType::Q6_K)) + .count(); + assert!(kquants > 0, "a q4_k_m file carries K-quant tensors"); + let payload_bytes: u64 = f.tensors.iter().map(|t| t.byte_len()).sum(); + eprintln!( + "[real_gguf] token_embd {:?} {:?} · {} K-quant tensors · payload ~{:.2} GiB", + embd.dtype, + embd.dims, + kquants, + payload_bytes as f64 / f64::from(1u32 << 30) + ); +} + +/// Minimal safetensors reader for the cross-check: header JSON + raw bf16 +/// tensor bytes widened to f32. +fn safetensors_bf16_f32(path: &std::path::Path, name: &str) -> Vec { + use std::io::{Read, Seek, SeekFrom}; + let mut file = std::fs::File::open(path).expect("safetensors opens"); + let mut len_bytes = [0u8; 8]; + file.read_exact(&mut len_bytes).expect("header length"); + let header_len = u64::from_le_bytes(len_bytes); + let mut header = vec![0u8; usize::try_from(header_len).expect("sane header")]; + file.read_exact(&mut header).expect("header json"); + let header: serde_json::Value = serde_json::from_slice(&header).expect("header parses"); + let entry = header.get(name).unwrap_or_else(|| panic!("{name} present")); + assert_eq!(entry["dtype"], "BF16", "cross-check expects bf16 weights"); + let start = entry["data_offsets"][0].as_u64().expect("start offset"); + let end = entry["data_offsets"][1].as_u64().expect("end offset"); + file.seek(SeekFrom::Start(8 + header_len + start)) + .expect("seek to tensor"); + let mut raw = vec![0u8; usize::try_from(end - start).expect("sane tensor")]; + file.read_exact(&mut raw).expect("tensor bytes"); + raw.chunks_exact(2) + .map(|b| f32::from_bits(u32::from(u16::from_le_bytes([b[0], b[1]])) << 16)) + .collect() +} + +fn cosine(a: &[f32], b: &[f32]) -> f64 { + assert_eq!(a.len(), b.len()); + let dot: f64 = a + .iter() + .zip(b) + .map(|(x, y)| f64::from(*x) * f64::from(*y)) + .sum(); + let na: f64 = a.iter().map(|x| f64::from(*x).powi(2)).sum::().sqrt(); + let nb: f64 = b.iter().map(|x| f64::from(*x).powi(2)).sum::().sqrt(); + dot / (na * nb) +} + +/// Dequantization proof against the model's TRUE weights: the same +/// checkpoint exists here as bf16 safetensors and as a Q4_K_M GGUF, so +/// - the GGUF's F32 norm tensors must equal the bf16 originals EXACTLY +/// (bf16 → f32 widening is lossless and that is how llama.cpp converts); +/// - a dequantized Q4_K embedding row must land within quantization error +/// of the original (cosine ≈ 1; garbage layout decode would be ≈ 0). +#[test] +#[ignore = "needs the local GGUF (MUMMU_GGUF_PATH) + safetensors (MUMMU_QWEN2_DIR) of the same model"] +fn real_qwen2_gguf_dequant_matches_the_true_weights() { + let Some(path) = gguf_path() else { + panic!("set MUMMU_GGUF_PATH to the qwen2.5-1.5b-instruct q4_k_m gguf"); + }; + let st_dir = std::env::var_os("MUMMU_QWEN2_DIR").map(PathBuf::from); + let Some(st) = st_dir.filter(|d| d.join("model.safetensors").is_file()) else { + panic!("set MUMMU_QWEN2_DIR to the same model's safetensors dir"); + }; + let st = st.join("model.safetensors"); + let f = GgufFile::open(&path).expect("header parses"); + + // Leg 1: F32 norm — exact. + let ours = f + .read_tensor_f32("output_norm.weight") + .expect("norm dequantizes"); + let reference = safetensors_bf16_f32(&st, "model.norm.weight"); + assert_eq!(ours.len(), reference.len(), "same norm size"); + let exact = ours + .iter() + .zip(&reference) + .all(|(a, b)| a.to_bits() == b.to_bits()); + assert!(exact, "GGUF F32 norm must be the bf16 weights, bit-exact"); + eprintln!( + "[real_gguf] output_norm.weight: {} f32 values bit-exact vs safetensors", + ours.len() + ); + + // Leg 2: Q4_K embedding rows — within quantization error of the truth. + let embd = f.read_tensor_f32("token_embd.weight").expect("dequantizes"); + let truth = safetensors_bf16_f32(&st, "model.embed_tokens.weight"); + assert_eq!(embd.len(), truth.len(), "same embedding size"); + let hidden = 1536; + for row in [9707usize, 100_000] { + let ours_row = &embd[row * hidden..(row + 1) * hidden]; + let true_row = &truth[row * hidden..(row + 1) * hidden]; + let cos = cosine(ours_row, true_row); + eprintln!("[real_gguf] Q4_K embd row {row}: cosine {cos:.5} vs bf16 truth"); + assert!( + cos > 0.97, + "row {row}: cosine {cos} — layout decode is wrong" + ); + } +} diff --git a/crates/mummu/tests/real_toolcall_lfm.rs b/crates/mummu/tests/real_toolcall_lfm.rs new file mode 100644 index 0000000..3e5b105 --- /dev/null +++ b/crates/mummu/tests/real_toolcall_lfm.rs @@ -0,0 +1,78 @@ +//! Real-model tool-use proof for the LFM convention: render a `List of +//! tools:` prompt, let LFM2.5-1.2B decode on the GPU, and parse the Pythonic +//! `<|tool_call_start|>` block it actually emits. Ignored by default; run with +//! +//! ```text +//! MUMMU_LFM2_DIR=path/to/lfm2.5-1.2b cargo test -p mummu --release --test real_toolcall_lfm -- --ignored --nocapture +//! ``` + +use std::path::PathBuf; + +use mummu::backend::Gpu; +use mummu::chat::{ChatMl, ToolSpec, Turn, parse_tool_calls_lfm}; +use mummu::models::CausalLm; +use mummu::models::lfm2; +use tokenizers::Tokenizer; + +fn lfm2_dir() -> Option { + let dir = PathBuf::from(std::env::var_os("MUMMU_LFM2_DIR")?); + dir.is_dir().then_some(dir) +} + +#[test] +#[ignore = "needs multi-GB local weights (MUMMU_LFM2_DIR) + the reference GPU"] +fn lfm2_emits_a_parseable_pythonic_tool_call() { + let Some(dir) = lfm2_dir() else { + panic!("set MUMMU_LFM2_DIR to a dir with config.json/tokenizer.json/model.safetensors"); + }; + + let tools = [ToolSpec { + name: "get_weather".into(), + description: "Get the current weather for a city.".into(), + parameters: serde_json::json!({ + "type": "object", + "properties": { + "city": { "type": "string", "description": "City name" } + }, + "required": ["city"] + }), + }]; + let raw = ChatMl::lfm2().render_with_tools( + &tools, + &[Turn::user("What is the weather in Paris right now?")], + ); + + let tok = Tokenizer::from_file(dir.join("tokenizer.json")).expect("tokenizer.json loads"); + let prompt = tok + .encode(raw.as_str(), false) + .expect("encodes") + .get_ids() + .to_vec(); + + let device = burn::tensor::Device::::default(); + let loaded = lfm2::load_from_dir::(&dir, &device).expect("weights load checked"); + let ids = loaded + .greedy_generate(&prompt, 128, &device) + .expect("greedy decode"); + // Special tokens stay in: the <|tool_call_start|> markers ARE the format. + let text = tok.decode(&ids, false).expect("decode"); + eprintln!("[real_toolcall_lfm] model emitted: {text:?}"); + + let (calls, prose) = parse_tool_calls_lfm(&text).expect("emitted tool call parses"); + assert_eq!( + calls.len(), + 1, + "expected exactly one call, prose: {prose:?}" + ); + assert_eq!(calls[0].name, "get_weather"); + assert_eq!( + calls[0].arguments["city"].as_str(), + Some("Paris"), + "arguments: {:?}", + calls[0].arguments + ); + eprintln!( + "[real_toolcall_lfm] parsed: {} with {}", + calls[0].name, calls[0].arguments + ); +}