diff --git a/Cargo.lock b/Cargo.lock index 6a37b7e..ec336e5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -585,6 +585,7 @@ version = "0.21.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0098bf6555c151517b4c98ca0e8ed1ab6a976e40f5f4e967c117b00eee21e192" dependencies = [ + "aligned-vec", "burn-backend", "burn-ir", "burn-std", @@ -592,7 +593,9 @@ dependencies = [ "gemm", "half", "libm", + "macerator", "num-traits", + "rayon", ] [[package]] @@ -3393,6 +3396,7 @@ name = "mummu" version = "0.1.0" dependencies = [ "burn", + "burn-flex", "half", "once_cell", "pollster", @@ -3400,6 +3404,7 @@ dependencies = [ "serde_json", "thiserror 2.0.18", "tokenizers", + "ureq 3.3.0", "wgpu", ] @@ -3407,8 +3412,10 @@ dependencies = [ name = "mummu-bench" version = "0.1.0" dependencies = [ + "burn", "criterion", "mummu", + "tokenizers", ] [[package]] @@ -3924,9 +3931,9 @@ dependencies = [ [[package]] name = "pollster" -version = "0.4.0" +version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f3a9f18d041e6d0e102a0a46750538147e5e8992d3b4873aaafee2520b00ce3" +checksum = "2560ec4d09f0d4db1fb45518d099d74954dd3da5855a0fcb84f2779b5b41a59e" [[package]] name = "portable-atomic" @@ -5225,7 +5232,7 @@ dependencies = [ "libc", "serde", "serde_json", - "ureq", + "ureq 2.12.1", "zip 0.6.6", ] @@ -5516,6 +5523,35 @@ dependencies = [ "webpki-roots 0.26.11", ] +[[package]] +name = "ureq" +version = "3.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dea7109cdcd5864d4eeb1b58a1648dc9bf520360d7af16ec26d0a9354bafcfc0" +dependencies = [ + "base64 0.22.1", + "flate2", + "log", + "percent-encoding", + "rustls", + "rustls-pki-types", + "ureq-proto", + "utf8-zero", + "webpki-roots 1.0.8", +] + +[[package]] +name = "ureq-proto" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e994ba84b0bd1b1b0cf92878b7ef898a5c1760108fe7b6010327e274917a808c" +dependencies = [ + "base64 0.22.1", + "http", + "httparse", + "log", +] + [[package]] name = "url" version = "2.5.8" @@ -5528,6 +5564,12 @@ dependencies = [ "serde", ] +[[package]] +name = "utf8-zero" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8c0a043c9540bae7c578c88f91dda8bd82e59ae27c21baca69c8b191aaf5a6e" + [[package]] name = "utf8_iter" version = "1.0.4" diff --git a/Cargo.toml b/Cargo.toml index 81bbd06..242574f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,21 +16,30 @@ repository = "https://github.com/physics515/Mummu" # `Wgpu` transparently becomes `Fusion`. burn = { version = "0.21", default-features = false, features = [ "std", - "ndarray", "wgpu", "fusion", "autotune", "store", ] } +# The CPU backend: pure-Rust SIMD + gemm, burn-ndarray's designated successor +# (0.21 release notes); swapped in 2026-07-10 after the MiniLM Candle-parity +# gate passed on it bit-for-bit-equivalent (cosine 0.99999994, max component +# diff 1.3e-7) with all 80 unit tests green. +burn-flex = "0.21" + # Direct handle on wgpu (the same version Burn resolves) for the cheap startup # adapter probe: enumerate GPUs to decide GPU vs CPU before building a backend. wgpu = "29" # Block on wgpu 29's async adapter enumeration from the sync startup probe. -pollster = "0.4" +pollster = "1.0" # f16 element type for the GPU backend; the same `half` version Burn pulls. half = "2" tokenizers = "0.22" +# Blocking HTTP (rustls) for the P3 model downloader; small, sync like the rest +# of the library surface. +ureq = "3" + serde = { version = "1", features = ["derive"] } serde_json = "1" thiserror = "2" diff --git a/README.md b/README.md index 6d30e7a..f8df458 100644 --- a/README.md +++ b/README.md @@ -2,13 +2,13 @@ **A from-scratch, single-binary local model runner in Rust — import any model, quantize it to fit, and run it on all the hardware you have: CPU, one GPU, or several. No Ollama, no cloud, no CUDA toolchain.** -Mummu imports models from HuggingFace (or disk), **(auto-)quantizes** them to fit, and runs them **natively in Rust on [Burn](https://burn.dev)** across every device you have. One binary — a runtime probe inventories your GPUs (Vulkan / DX12 / Metal via `wgpu`) and your CPU (`ndarray`) and places the model to use them **to the fullest**: several GPUs together, with CPU offload when VRAM is short, no feature-split builds and no CUDA toolchain. Models are **reimplemented from scratch**, generic over the Burn backend, and **parity-tested byte-for-byte** against a reference so the reimplementations can be trusted. +Mummu imports models from HuggingFace (or disk), **(auto-)quantizes** them to fit, and runs them **natively in Rust on [Burn](https://burn.dev)** across every device you have. One binary — a runtime probe inventories your GPUs (Vulkan / DX12 / Metal via `wgpu`) and your CPU (`burn-flex`) and places the model to use them **to the fullest**: several GPUs together, with CPU offload when VRAM is short, no feature-split builds and no CUDA toolchain. Models are **reimplemented from scratch**, generic over the Burn backend, and **parity-tested byte-for-byte** against a reference so the reimplementations can be trusted. It exists because two local-first apps — **[laurelane](https://github.com/physics515/laurelane)** (a private budgeting cockpit) and **[Nanna](https://github.com/physics515/Nanna)** (an always-on local AI presence) — were building the *same* runner twice. Mummu is that runner, extracted and generalized: each app consumes it as a dependency and keeps only its own domain glue. Laurelane proved the blueprint (Qwen2.5 / LFM2.5 / all-MiniLM ported to Burn, byte-identical parity vs Candle, validated on an RTX 4070 Ti SUPER 16 GB); Mummu is where it lives, hardens, and grows. ## What it is -- **One binary, every device** — compile both `Wgpu` (Vulkan/DX12/Metal, no CUDA toolchain) and `NdArray` (CPU); a runtime probe enumerates **all** adapters + the CPU and places the model across them — a single GPU, **several GPUs together**, or GPU + CPU hybrid. No feature-split builds, no per-vendor path. +- **One binary, every device** — compile both `Wgpu` (Vulkan/DX12/Metal, no CUDA toolchain) and `burn-flex` (CPU); a runtime probe enumerates **all** adapters + the CPU and places the model across them — a single GPU, **several GPUs together**, or GPU + CPU hybrid. No feature-split builds, no per-vendor path. - **Models from scratch, generic over `B: Backend`** — a growing zoo (Qwen2/2.5, LFM2/2.5 hybrid conv+attention, all-MiniLM embedder) built on shared blocks (RmsNorm · GQA · RoPE · SwiGLU · tied lm-head · depthwise causal conv), with a clean trait to add more. - **Trustworthy reimplementations** — every port must pass a **parity gate**: single-forward top-k logits *and* a short greedy sequence match a reference (Candle, or a local Ollama of the same model) exactly. - **Fast** — per-layer KV cache (+ conv-state cache for hybrids), on-GPU argmax (sync only the winning index), sampling, **token streaming**, cooperative cancellation; kernel `fusion` + `autotune`; an opt-in **f16** path to roughly halve VRAM. @@ -19,7 +19,7 @@ It exists because two local-first apps — **[laurelane](https://github.com/phys ## Status — what runs today - **Workspace + backends** — `crates/mummu` (library) + `crates/mummu-bench` (criterion); one binary - compiles both `Wgpu` (with `fusion` + `autotune`) and `NdArray`, with a cached runtime GPU probe and a + compiles both `Wgpu` (with `fusion` + `autotune`) and `burn-flex` (CPU), with a cached runtime GPU probe and a device inventory that records per-adapter/per-API `SHADER_F16`. - **Shared blocks, generic over `B: Backend`** — cache-aware GQA attention (optional per-head q/k RMSNorm), manual RoPE, SwiGLU, and LFM2's double-gated causal short-conv with rolling decode state; @@ -28,9 +28,27 @@ It exists because two local-first apps — **[laurelane](https://github.com/phys fail-loud load (never silently zero-init); `config.json`-driven hyperparameters. - **Three models ported and running on real weights** — Qwen2/2.5, the LFM2/2.5 hybrid, and the all-MiniLM sentence embedder; Qwen2.5-1.5B and LFM2.5-1.2B load and greedy-decode correctly on the - reference GPU (wgpu/Vulkan). The P7 parity harness will gate them "trusted" against a reference before - the roadmap ticks them done. -- **Model-cache disk accounting** — per-model disk usage + traversal-safe removal validation (`manage`). + reference GPU (wgpu/Vulkan). +- **Qwen2.5 and MiniLM are parity-verified** — the two-leg P7 gate passes for Qwen2.5-1.5B on the + reference GPU: single-forward top-5 logits match a Candle f32 reference (max |Δlogit| 2.7e-5, + `tests/parity_qwen2.rs` + the committed `tools/candle-probe` fixture) and a 24-token greedy sequence + matches `ollama qwen2.5:1.5b-instruct-fp16` byte-for-byte. The MiniLM embedder matches its Candle + reference at cosine 0.99999994 (max |Δcomponent| 1.2e-7, `tests/real_minilm.rs`). LFM2.5 still awaits + a same-weights reference (tracked in P7). +- **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. +- **Benchmarked** — Qwen2.5-1.5B f32 on the reference GPU: **TTFT 100.5 ms, decode 13.3 tok/s, 11.9 GiB + whole-card peak** (~7.9 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`. +- **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. +- **Process-lifetime model cache** — `ModelSlot` loads a checkpoint once per process, switches models by + key, and `clear()`s to free VRAM; Burn's `Param` isn't `Sync`, so access serializes behind its mutex. ## Design principles diff --git a/ROADMAP.md b/ROADMAP.md index 1b31496..23bf7bd 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -6,7 +6,7 @@ > README (perf claims link a benchmark artifact); everything not-done / discovered / next is a `[ ]` > here; git history + PRs are the record. Edit surgically; never rewrite wholesale. -**Stack:** Rust 2024 · **Burn 0.21** (`wgpu` 29 + `ndarray`, `fusion` + `autotune`, multi-device) · +**Stack:** Rust 2024 · **Burn 0.21** (`wgpu` 29 + `burn-flex` CPU, `fusion` + `autotune`, multi-device) · `burn-store` · HF `tokenizers` · runs on **any hardware** — CPU, one GPU, or several (multi-GPU + CPU offload). Reference dev machine: Ryzen 9 7950X3D · 128 GB · RTX 4070 Ti SUPER 16 GB. @@ -34,9 +34,17 @@ Performance is a **gate**, not a phase. Governing metric: **task throughput @ bu TTFT while fitting the *target device set's* VRAM/RAM (one GPU, several, or CPU): the biggest useful model that fits, run as fast as the hardware allows, every device busy. A change ships only when parity holds and a benchmark holds/improves its budget; README perf claims link an artifact. -- [ ] `mummu-bench` (criterion) — TTFT, decode tok/s, VRAM per model × device-set (1 GPU / N GPUs / CPU / hybrid). - *(2026-07-09) Crate + harness wired (smoke bench); real model benches land with P5.* -- [ ] `bench/BASELINE.md` budgets (VRAM ceiling, min decode tok/s, max TTFT) + a `cargo test` gate that fails a regression. +- [x] `mummu-bench` (criterion) — TTFT, decode tok/s, VRAM per model × device-set (1 GPU / N GPUs / CPU / hybrid). + *(2026-07-09) Crate + harness wired (smoke bench); real model benches land with P5.* *(2026-07-10) + Real Qwen2.5-1.5B GPU benches live: TTFT 100.5 ms, decode 13.3 tok/s (f32, criterion). More + model × device-set combos accrete as those paths land (CPU, f16, multi-GPU).* +- [x] `bench/BASELINE.md` budgets (VRAM ceiling, min decode tok/s, max TTFT) + a `cargo test` gate that fails a regression. + *(2026-07-10) Baseline recorded (TTFT ≤ 150 ms, ≥ 10 tok/s, ≤ 13 GiB whole-card; measured 100.5 ms / + 13.3 tok/s / 11.9 GiB peak incl. ~4 GiB ambient); gate = `mummu-bench/tests/budget.rs` (ignored, + weights + GPU) passing at 110.4 ms / 11.8 tok/s.* +- [ ] Decode is dispatch-bound, not bandwidth-bound: 75 ms/token streams ~6.2 GB of f32 weights at only + ~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. - [ ] 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/ @@ -60,10 +68,20 @@ a benchmark holds/improves its budget; README perf claims link an artifact. f16=true on Vulkan, false on DX12, + an integrated AMD GPU (a real second adapter for multi-GPU).* - [x] `fusion` + `autotune` on (`Wgpu` becomes `Fusion`; needs `recursion_limit = 512`). *(2026-07-09) Workspace features + crate-level `recursion_limit`.* -- [ ] Evaluate **burn-flex** (Burn 0.21's new pure-Rust CPU backend; `burn-ndarray` is now on a +- [x] Evaluate **burn-flex** (Burn 0.21's new pure-Rust CPU backend; `burn-ndarray` is now on a deprecation path) as the `Cpu` alias replacement — SIMD + gemm, no_std, and built-in per-tensor/ per-block quantization (~40 quantized ops) that P9 could ride on. Gate on parity + a CPU decode bench — https://github.com/antimora/burn-flex · https://burn.dev/blog/release-0.21.0/ + *(2026-07-10) **Swapped in** (`Cpu = burn_flex::Flex`, ndarray feature dropped): the + MiniLM Candle-parity gate passes on Flex (cosine 0.99999994, max |Δcomponent| 1.3e-7 — equivalent + to ndarray's 1.2e-7) and all 80 unit tests are green, incl. the cache-equivalence proofs that run + on the CPU backend. A dedicated CPU decode tok/s bench still wants a CPU-tier model (0.5B) — next + item.* +- [x] CPU decode bench for `bench/BASELINE.md`: pull Qwen2.5-0.5B (catalog entry exists) and record + decode tok/s on the Flex backend, so CPU-only machines get a budget row and Flex regressions are + caught like GPU ones. *(2026-07-10) 0.5B fetched through the registry/hub path (988 MB in ~13 s), + decodes coherently on Flex at **11.7 tok/s** (7950X3D, f32); budget ≥ 6 tok/s gated by + `mummu-bench/tests/budget_cpu.rs`.* ### P2 — Model zoo (from scratch, generic over `B`) *(ex-laurelane)* - [x] Shared blocks: RmsNorm, GQA attention, RoPE (manual rotate-half), SwiGLU MLP, tied lm-head, @@ -71,10 +89,12 @@ a benchmark holds/improves its budget; README perf claims link an artifact. (optional per-head q/k RMSNorm covers Qwen2 AND LFM2), `SwiGluMlp`, `ShortConv` (LIV) with rolling state, RoPE + causal mask; 18 unit tests incl. the prefill+decode ≡ full-forward equivalence for both the KV cache and the conv state. RmsNorm/tied-head come from burn::nn / the model files.* -- [ ] **Qwen2 / Qwen2.5** decoder (1.5B / 0.5B tiers). *(2026-07-09) Ported (`models::qwen2`, +- [x] **Qwen2 / Qwen2.5** decoder (1.5B / 0.5B tiers). *(2026-07-09) Ported (`models::qwen2`, config-driven, checked safetensors load); REAL GPU inference verified — Qwen2.5-1.5B on the 4070 Ti SUPER greedy-decoded "2+2 equals 4.", top-5 probe led by id 9707 "Hello"; toy-model cache-equivalence - unit test. Stays `[ ]` until the P7 parity harness passes it against a logits reference.* + unit test.* *(2026-07-10) **Parity gate PASSED** (`tests/parity_qwen2.rs`): top-5 logits match the + Candle f32 reference exactly by id with max |Δlogit| 2.7e-5 (bound 1e-3), and a 24-token greedy + sequence matches `ollama qwen2.5:1.5b-instruct-fp16` byte-for-byte on the 4070 Ti SUPER.* - [ ] **LFM2.5-1.2B** hybrid (6 GQA-attention w/ per-head q/k RMSNorm + 10 double-gated short-conv "LIV" blocks, SwiGLU, tied head, conv-state cache; ChatML, EOS `<|im_end|>`). *(2026-07-09) Ported (`models::lfm2`, hybrid cache, LFM2→shared-block key remap); toy hybrid cache-equivalence test; @@ -83,26 +103,46 @@ a benchmark holds/improves its budget; README perf claims link an artifact. gate passes: the local `ollama lfm2.5:latest` tag now resolves to the 8.5B **MoE** Q4 w/ thinking (verified via `ollama show` 2026-07-09) — not the same weights, so no valid local reference exists; see the P7 reference item.* -- [ ] **all-MiniLM** BERT sentence-embedder (6-layer post-LN bidirectional attention + GeLU FFN, +- [x] **all-MiniLM** BERT sentence-embedder (6-layer post-LN bidirectional attention + GeLU FFN, masked-mean-pool + L2-normalize). *(2026-07-09) Ported (`models::minilm`, ids+mask in → L2-normalized embedding out; tokenization stays caller-side); unit tests incl. padding-invisibility; real-weights - semantic test (`tests/real_minilm.rs`). Stays `[ ]` until the P7 cosine-vs-Candle parity re-runs here.* -- [ ] A `Model` trait so new architectures (Hermes-class function-callers, Gemma, Qwen3, …) slot in. + 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).* +- [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 + driver (static dispatch, single code path). Both LLMs now implement it; the parity gate re-ran + green through the trait, and the real-inference suite shares one `ModelSlot` (4 GPU tests, one + 3.1 GB load — found and fixed a 2-models-in-VRAM blowup in the old per-test loads).* ### P3 — Model import suite (any source / any format → a running model) The subsystem that turns "a model on HuggingFace or on disk" into a loaded, parity-checked Mummu model. **Data-driven** — adding a model is a manifest entry, not new code. All import is Burn-native (`burn-store` / `burn-import`). -- [ ] **Sources** — HuggingFace Hub (repo id + revision), local paths, and a bundled resources dir (checked +- [x] **Sources** — HuggingFace Hub (repo id + revision), local paths, and a bundled resources dir (checked first). Streaming download into a per-user cache: resumable (`.part`), integrity-checked, and - **sharded-checkpoint aware** (read `*.index.json`, fetch + merge shards). + **sharded-checkpoint aware** (read `*.index.json`, fetch + merge shards). *(2026-07-10) `mummu::hub` + (ureq 3, https-only): `fetch_file` streams through `.part` with HTTP-Range resume + Content-Length + verification + cache-first; `fetch_model` pulls config/tokenizer/weights with the + `model.safetensors.index.json` shard fallback; `Progress` callback per chunk (feeds P8). Real-network + proof: all-MiniLM (90.8 MB) downloaded → checked-load → unit-norm embedding; a half-seeded `.part` + resumed at byte 249,507/466,247 and finished byte-identical. Local paths are already first-class + (`load_from_dir`); bundled-resources-dir precedence is app wiring.* +- [ ] Stronger download integrity: verify the Hub's LFS sha256 (`X-Linked-ETag`) instead of length-only; + re-verify on cache hits behind a flag. - [x] **safetensors** *(ex-laurelane)* — `burn-store` `SafetensorsStore` + `PyTorchToBurnAdapter`; the primary path. *(2026-07-09) `import::{CastFloatAdapter, load_checked}`: bf16→backend-float cast + fail-loud load; proven by loading the real 3.1 GB Qwen2.5-1.5B and 2.3 GB LFM2.5 checkpoints with zero missing keys.* - [ ] **PyTorch state dicts** (`.pth` / `pytorch_model*.bin`) — for models not shipped as safetensors. - [ ] **GGUF** (llama.cpp) — parse the GGUF container (metadata KV + tensor table), map tensors to modules, and **dequantize** Q4/Q5/Q8/K-quant blocks into Burn tensors (or hand keep-quantized to P9). GGUF is how - most small models are distributed — this makes the whole ecosystem importable. + most small models are distributed — this makes the whole ecosystem importable. *(2026-07-10 research)* + K-quant superblocks are 256 values: Q4_K = fp16 d + fp16 dmin + 12 B of 6-bit sub-scales/mins + 128 B + of 4-bit q (144 B total), `x = d·scale_i·q − dmin·min_i`; Q6_K = 128 B low-4 + 64 B high-2 + 16×i8 + sub-scales + fp16 d (210 B). Rust references: llama.cpp ggml-quants + the `rage-quant` crate + (Q8_0/Q4_K/Q6_K dequant + SIMD dot) — https://haroldbenoit.com/notes/ml/llms/quantization/llama.cpp/k-quants-implementation · + https://crates.io/crates/rage-quant - [ ] **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. @@ -119,19 +159,34 @@ The subsystem that turns "a model on HuggingFace or on disk" into a loaded, pari `layer_types` + auto-adjusted ff_dim); both real checkpoints parse and drive the build.* - [ ] **Tokenizer + chat-template import** — HF `tokenizer.json` (fast), SentencePiece `tokenizer.model`, BPE merges/vocab; special-tokens map + the chat template from `tokenizer_config.json`. -- [ ] **Model registry / manifest** — a declarative `ModelSpec` (repo, architecture, weight format, dtype, +- [x] **Model registry / manifest** — a declarative `ModelSpec` (repo, architecture, weight format, dtype, tokenizer, chat template, size tier) + a small built-in catalog of known-good models (Qwen2.5, LFM2.5, - MiniLM, …); adding a model = a manifest entry. + MiniLM, …); adding a model = a manifest entry. *(2026-07-10) `mummu::registry`: `ModelSpec` + (name/repo/revision/architecture/size, validated incl. traversal-safe names, serde round-trip) + + `spec.fetch(models_root, progress)` onto the hub downloader; built-in catalog: Qwen2.5-1.5B/0.5B, + LFM2.5-1.2B, all-MiniLM (repo ids match laurelane's validated constants); the network proof now + fetches spec-driven. Weight-format/dtype/chat-template fields accrete as those import paths land.* - [ ] **Import validation** — checked load + a first-token parity smoke against a reference before a model is marked trusted; a clear error taxonomy (missing file, bad shard, key mismatch, unsupported dtype). ### P4 — Tokenizer & chat templates *(ex-laurelane)* -- [ ] HF `tokenizers` (pinned); explicit chat templates (ChatML + per-model), correct special/EOS tokens. +- [x] HF `tokenizers` (pinned); explicit chat templates (ChatML + per-model), correct special/EOS tokens. + *(2026-07-10) `mummu::chat`: `Turn`/`Role` + a `ChatMl` renderer with per-model constructors + (`qwen2` plain, `lfm2` with `<|startoftext|>` BOS); byte-verified — the Qwen2 parity gate now + renders its prompt through the template and still matches the Candle fixture and the Ollama fp16 + greedy leg exactly. EOS stays config-driven (`EosIds`); tool-use templates are the next item.* - [ ] Hermes-style tool-use chat template (Qwen3 ships it in `tokenizer_config.json`) + LFM2.5's bracket-notation tool-call output — the two top scorers (0.880 agent score; LFM2.5-1.2B also the fastest at ~1.5 s) on 2026's 21-model local tool-calling benchmark; function calling is why the apps want a local runner — https://mikeveerman.be/blog/github-2026-02-06-tool-calling-benchmark/ · https://qwen.readthedocs.io/en/latest/framework/function_call.html + *(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 + function-calling tier should target the 7–9B class once quant lands; Hermes 4 (Qwen3-14B fine-tune) + emits `` tags after an explicit reasoning step — easy to parse with the same template + machinery — https://www.promptquorum.com/power-local-llm/best-local-models-tool-calling-2026 · + https://localaimaster.com/blog/best-ollama-models-for-agents ### P5 — Decode engine *(ex-laurelane)* - [x] Per-layer KV cache (+ conv-state cache for hybrids); prompt prefilled once, then one token/step. @@ -139,9 +194,17 @@ The subsystem that turns "a model on HuggingFace or on disk" into a loaded, pari full-forward proven by unit tests at block AND whole-model level, then by real GPU decode.* - [x] On-GPU argmax (sync only the winning index); single-token decode skips the causal mask. *(2026-07-09) `decode::argmax_id`; `t == 1` builds no mask.* -- [ ] Sampling beyond greedy (temperature / top-p); **token streaming** via a callback/channel; - cooperative interrupt/cancellation between tokens. -- [ ] Process-lifetime model + tokenizer cache (per backend; behind a `Mutex` since Burn `Param` isn't `Sync`). +- [x] Sampling beyond greedy (temperature / top-p); **token streaming** via a callback/channel; + cooperative interrupt/cancellation between tokens. *(2026-07-10) `decode::{SamplerOptions, + sample_id, Pcg32, generate_loop}` + per-model `generate(…, on_token)`: temperature/top-k/top-p over + an O(vocab) partial select, deterministic in-house PCG32 (no rand dep), `ControlFlow` callback = + streaming AND cancellation; greedy stays on-GPU argmax and re-passed the Qwen2 parity gate; real-GPU + proof: seeded sampled stream replays identically and cancels at 8 tokens (10 new unit tests).* +- [x] Process-lifetime model + tokenizer cache (per backend; behind a `Mutex` since Burn `Param` isn't `Sync`). + *(2026-07-10) `cache::ModelSlot` — one static slot per (model, backend); `with(key, load, f)` + loads once, reuses on key match, drop-then-swap on a different checkpoint dir (the P8 + active-model-switch primitive), `clear()` frees VRAM; real-GPU proof: two decodes through a static + slot performed exactly one 3.1 GB load (6 unit tests + threaded-static test).* ### P6 — Hardware planner: precision, placement & full utilization The "use all the hardware" phase — inventory the machine, then pick the precision and the device placement @@ -151,7 +214,19 @@ that fits the model AND uses every device to the fullest. - [ ] **Precision selection** — pick a per-device dtype (f32 / **f16** / int8 / int4) that fits: f16 via `Wgpu` (needs wgpu ≥ 27 `SHADER_F16` polyfill — *laurelane compiles it + a startup `SHADER_F16` diagnostic; finish on-GPU validation here*: no naga crash, ~halved VRAM, coherent output); - drop to int8/int4 (P9) when f16 still won't fit. + drop to int8/int4 (P9) when f16 still won't fit. *(2026-07-10) On-GPU validation ran + (`tests/real_f16.rs`, the standing gate): **2 of 3 claims hold** — shaders compile + run on + Vulkan/SHADER_F16, VRAM drops 11.9 → 8.7 GiB whole-card (~7.9 → ~4.7 GiB runner), but logits + collapse to NaN (the GPU argmax returns the out-of-vocab sentinel 151936 = vocab_size; now caught + loudly by a decode guard). Coherent-output remains open below.* +- [ ] **f16 mixed-precision islands** — Qwen2.5-1.5B in pure f16 NaNs out (overflow in the + softmax/RmsNorm/logit reductions; f16 max is 65 504). Keep weights + matmuls f16 but compute the + numerically hot reductions (attention softmax, RmsNorm accumulation, final logits) in f32, then + re-run the f16 gate and the parity harness. +- [ ] 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 - [ ] **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 @@ -165,20 +240,32 @@ that fits the model AND uses every device to the fullest. counts, precision, CPU-offload cap) for power users. ### P7 — Parity & performance harness *(ex-laurelane)* -- [ ] Parity gate: single-forward top-k logits + a short greedy sequence must match a reference (Candle, +- [x] Parity gate: single-forward top-k logits + a short greedy sequence must match a reference (Candle, or a local Ollama of the same model) — the trust gate every port passes. *(2026-07-09) Greedy leg - exists (`tests/parity_lfm2.rs`, Ollama raw-mode temperature-0 via curl); blocked on a same-weights - reference: `ollama lfm2.5:latest` is now the 8.5B MoE and no 1.2B tag exists.* -- [ ] Stand up same-weights references: a Candle-based logits probe (dev-dependency or a small side + exists (`tests/parity_lfm2.rs`, Ollama raw-mode temperature-0 via curl).* *(2026-07-10) Both legs + live and passing for Qwen2.5-1.5B (`tests/parity_qwen2.rs`): logits leg vs a committed + `tools/candle-probe` fixture, greedy leg vs Ollama fp16. LFM2.5 still lacks a same-weights + reference (see the P2 item); candle-transformers has no LFM2, so its logits leg needs another + route (llama.cpp logprobs, or an HF transformers dump).* +- [x] Stand up same-weights references: a Candle-based logits probe (dev-dependency or a small side harness, as laurelane's Qwen2 validation did) + pull `qwen2.5:1.5b-instruct-fp16` in Ollama for the - Qwen greedy leg. + Qwen greedy leg. *(2026-07-10) `tools/candle-probe` (out-of-workspace bin, Candle =0.9.1 CPU f32) + prints top-k (id, logit) JSON for the fixed prompt; fixture committed under + `crates/mummu/tests/fixtures/`; fp16 Ollama tag pulled and validated.* +- [ ] LFM2.5 same-weights reference for the parity gate: no Candle port exists — candidate routes are + llama.cpp `logprobs` on the fp16 GGUF, or a one-shot HF `transformers` logits dump matched to the + safetensors revision. - [ ] Wire the perf suite (above) into the parity harness so a correctness *or* budget regression fails CI. ### P8 — Model management API -- [ ] Download progress · disk usage · switch/remove models — an app-agnostic API the consumers' settings +- [x] Download progress · disk usage · switch/remove models — an app-agnostic API the consumers' settings UIs call. *(laurelane has disk-usage + remove; add progress + active-model switch.)* - *(2026-07-09) Disk usage + traversal-safe removal validation shipped as `manage` (5 unit tests); - download progress + active-model switch still open.* + *(2026-07-09) Disk usage + traversal-safe removal validation shipped as `manage` (5 unit tests).* + *(2026-07-10) Composed into `manage::ModelManager`: catalog listing, `is_installed`, + `install(name, on_progress)` (resumable hub fetch with per-chunk progress), traversal-safe + `remove`, `disk_report`; active-model switch = a consumer `ModelSlot` keyed by + `manager.model_dir(name)` (drop-then-swap + `clear()` already proven on the real GPU). 4 new + unit tests.* ### P9 — Quantization (fit any model to the hardware) The VRAM lever the P6 planner pulls to make the largest useful model fit the user's actual devices. diff --git a/bench/BASELINE.md b/bench/BASELINE.md new file mode 100644 index 0000000..8440f32 --- /dev/null +++ b/bench/BASELINE.md @@ -0,0 +1,30 @@ +# Benchmark baseline & budgets + +Reference machine: Ryzen 9 7950X3D · 128 GB · **RTX 4070 Ti SUPER 16 GB** (wgpu/Vulkan, `Fusion`, +f32). 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-10) | Budget | +| --- | --- | --- | +| TTFT (fresh cache: full prefill + first token) | 100.5 ms | ≤ 150 ms | +| Decode latency (32 greedy tokens, warm KV cache) | 2.414 s → **75.4 ms/token ≈ 13.3 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 | + +## Qwen2.5-0.5B-Instruct · CPU (burn-flex) · f32 + +| Metric | Recorded (2026-07-10) | Budget | +| --- | --- | --- | +| Decode (8 greedy tokens, warm KV cache; `mummu-bench/tests/budget_cpu.rs`) | **11.7 tok/s** | ≥ 6 tok/s | + +Datapoint (not yet a budget): the same model on `GpuF16` peaks at **8.7 GiB whole-card (~4.7 GiB +runner)** — VRAM roughly halves as expected — but decodes NaN today (see the ROADMAP P6 +mixed-precision-islands item), so no f16 perf row exists yet. + +Notes +- Effective weight-streaming bandwidth at 75 ms/token over ~6.2 GB of f32 weights is ~83 GB/s vs the + card's ~672 GB/s — the decode path is kernel/dispatch-bound, not bandwidth-bound. The CubeCL SPIR-V + compiler feature and the f16 path (ROADMAP P6) are the identified levers. +- `harness_smoke` (sub-ns) exists only to keep `cargo bench` green without the multi-GB weights. diff --git a/crates/mummu-bench/Cargo.toml b/crates/mummu-bench/Cargo.toml index 1f6b75d..2ccff67 100644 --- a/crates/mummu-bench/Cargo.toml +++ b/crates/mummu-bench/Cargo.toml @@ -9,9 +9,11 @@ publish = false [dependencies] mummu = { path = "../mummu" } +burn = { workspace = true } [dev-dependencies] criterion = { workspace = true } +tokenizers = { workspace = true } [[bench]] name = "runner" diff --git a/crates/mummu-bench/benches/runner.rs b/crates/mummu-bench/benches/runner.rs index bd295ba..fd53000 100644 --- a/crates/mummu-bench/benches/runner.rs +++ b/crates/mummu-bench/benches/runner.rs @@ -1,9 +1,42 @@ -//! Runner benchmark stub. Real benches (TTFT, decode tok/s, VRAM per model × -//! device-set) land with the decode engine (P5); until then this measures a -//! trivial baseline so the criterion harness itself is wired and green. +//! Real-model benchmarks: TTFT and decode tok/s for Qwen2.5-1.5B on the +//! machine's default GPU. Budgets and the last recorded numbers live in +//! `bench/BASELINE.md`. Run with +//! +//! ```text +//! MUMMU_QWEN2_DIR=path/to/qwen2.5-1.5b cargo bench -p mummu-bench +//! ``` +//! +//! Without `MUMMU_QWEN2_DIR` (or weights on disk) only the harness smoke runs, +//! so `cargo bench` stays green on machines without the multi-GB checkpoint. -use criterion::{Criterion, criterion_group, criterion_main}; use std::hint::black_box; +use std::path::PathBuf; +use std::time::{Duration, Instant}; + +use criterion::{Criterion, criterion_group, criterion_main}; +use mummu::backend::Gpu; +use mummu::decode::argmax_id; +use mummu::models::CausalLm; +use mummu::models::qwen2::{self, LoadedQwen2}; +use tokenizers::Tokenizer; + +/// Decode steps timed per criterion sample: long enough to amortize per-step +/// jitter, short enough that the KV cache stays near its steady-state length. +const DECODE_STEPS_PER_SAMPLE: usize = 32; + +fn qwen2_dir() -> Option { + let dir = PathBuf::from(std::env::var_os("MUMMU_QWEN2_DIR")?); + dir.is_dir().then_some(dir) +} + +/// The benchmark prompt — fixed so numbers are comparable across runs. +fn prompt_ids(dir: &std::path::Path) -> Vec { + let text = "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\nExplain, in three sentences, why the sky is blue.<|im_end|>\n<|im_start|>assistant\n"; + let tok = Tokenizer::from_file(dir.join("tokenizer.json")).expect("tokenizer.json loads"); + let ids = tok.encode(text, false).expect("encodes").get_ids().to_vec(); + assert!(ids.len() >= 16, "benchmark prompt suspiciously short"); + ids +} fn bench_harness_smoke(c: &mut Criterion) { c.bench_function("harness_smoke", |b| { @@ -11,5 +44,50 @@ fn bench_harness_smoke(c: &mut Criterion) { }); } -criterion_group!(benches, bench_harness_smoke); +fn bench_qwen2_real(c: &mut Criterion) { + let Some(dir) = qwen2_dir() else { + eprintln!("[mummu-bench] MUMMU_QWEN2_DIR not set — skipping real-model benches"); + return; + }; + let device = burn::tensor::Device::::default(); + let loaded: LoadedQwen2 = + qwen2::load_from_dir(&dir, &device).expect("weights load checked"); + let ids = prompt_ids(&dir); + + // TTFT: fresh cache, full prefill, first token argmax (the id readback is + // the GPU sync point, so the measured span covers real work end-to-end). + let mut group = c.benchmark_group("qwen2.5-1.5b/gpu"); + group.sample_size(10); + group.bench_function("ttft_prefill_first_token", |b| { + b.iter(|| { + let mut cache = loaded.new_cache(); + let logits = loaded.forward(&ids, 0, &mut cache, &device); + black_box(argmax_id(logits).expect("argmax")) + }); + }); + + // Decode: per sample, prefill once (untimed), then time N greedy decode + // steps through the warm KV cache. Per-token latency = measured / N. + group.bench_function("decode_32_tokens", |b| { + b.iter_custom(|iters| { + let mut total = Duration::ZERO; + for _ in 0..iters { + let mut cache = loaded.new_cache(); + let logits = loaded.forward(&ids, 0, &mut cache, &device); + let mut next = argmax_id(logits).expect("argmax"); + let start = Instant::now(); + for (step, past) in (ids.len()..).take(DECODE_STEPS_PER_SAMPLE).enumerate() { + let logits = loaded.forward(&[next], past, &mut cache, &device); + next = argmax_id(logits).expect("argmax"); + black_box((step, next)); + } + total += start.elapsed(); + } + total + }); + }); + group.finish(); +} + +criterion_group!(benches, bench_harness_smoke, bench_qwen2_real); criterion_main!(benches); diff --git a/crates/mummu-bench/tests/budget.rs b/crates/mummu-bench/tests/budget.rs new file mode 100644 index 0000000..44dd996 --- /dev/null +++ b/crates/mummu-bench/tests/budget.rs @@ -0,0 +1,79 @@ +//! The perf-budget gate from `bench/BASELINE.md`: fails when TTFT or decode +//! throughput for Qwen2.5-1.5B regresses past its budget on the reference +//! GPU. Ignored by default (multi-GB weights + a real GPU); run with +//! +//! ```text +//! MUMMU_QWEN2_DIR=path/to/qwen2.5-1.5b cargo test -p mummu-bench --release -- --ignored --nocapture +//! ``` +//! +//! Budgets are deliberately looser than the recorded numbers (see +//! BASELINE.md) so machine jitter doesn't flake the gate while a real +//! regression (a kernel falling off a fast path, an accidental sync per +//! layer) still trips it. + +use std::path::PathBuf; +use std::time::Instant; + +use mummu::backend::Gpu; +use mummu::decode::argmax_id; +use mummu::models::CausalLm; +use mummu::models::qwen2; +use tokenizers::Tokenizer; + +/// From bench/BASELINE.md (recorded 100.5 ms / 13.3 tok/s on 2026-07-10). +const TTFT_BUDGET_MS: f64 = 150.0; +const DECODE_BUDGET_TOKENS_PER_S: f64 = 10.0; +const DECODE_STEPS: usize = 32; + +fn qwen2_dir() -> Option { + let dir = PathBuf::from(std::env::var_os("MUMMU_QWEN2_DIR")?); + dir.is_dir().then_some(dir) +} + +#[test] +#[ignore = "needs multi-GB local weights (MUMMU_QWEN2_DIR) + the reference GPU"] +fn qwen2_stays_inside_its_perf_budgets() { + let Some(dir) = qwen2_dir() else { + panic!("set MUMMU_QWEN2_DIR to a dir with config.json/tokenizer.json/model.safetensors"); + }; + let text = "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\nExplain, in three sentences, why the sky is blue.<|im_end|>\n<|im_start|>assistant\n"; + let tok = Tokenizer::from_file(dir.join("tokenizer.json")).expect("tokenizer.json loads"); + let ids = tok.encode(text, false).expect("encodes").get_ids().to_vec(); + assert!(ids.len() >= 16, "budget prompt suspiciously short"); + + let device = burn::tensor::Device::::default(); + let loaded = qwen2::load_from_dir::(&dir, &device).expect("weights load checked"); + + // Warm-up: first-run autotune + pipeline compilation must not count. + let mut cache = loaded.new_cache(); + let logits = loaded.forward(&ids, 0, &mut cache, &device); + let _ = argmax_id(logits).expect("warm-up argmax"); + + // TTFT: fresh cache, full prefill, first token (argmax readback = sync). + let start = Instant::now(); + let mut cache = loaded.new_cache(); + let logits = loaded.forward(&ids, 0, &mut cache, &device); + let mut next = argmax_id(logits).expect("argmax"); + let ttft_ms = start.elapsed().as_secs_f64() * 1e3; + + // Decode throughput over a warm cache. + let start = Instant::now(); + for past in (ids.len()..).take(DECODE_STEPS) { + let logits = loaded.forward(&[next], past, &mut cache, &device); + next = argmax_id(logits).expect("argmax"); + } + let tok_per_s = DECODE_STEPS as f64 / start.elapsed().as_secs_f64(); + + eprintln!( + "[budget] TTFT {ttft_ms:.1} ms (budget {TTFT_BUDGET_MS} ms), \ + decode {tok_per_s:.1} tok/s (budget {DECODE_BUDGET_TOKENS_PER_S} tok/s)" + ); + assert!( + ttft_ms <= TTFT_BUDGET_MS, + "TTFT regression: {ttft_ms:.1} ms > {TTFT_BUDGET_MS} ms budget" + ); + assert!( + tok_per_s >= DECODE_BUDGET_TOKENS_PER_S, + "decode regression: {tok_per_s:.1} tok/s < {DECODE_BUDGET_TOKENS_PER_S} tok/s budget" + ); +} diff --git a/crates/mummu-bench/tests/budget_cpu.rs b/crates/mummu-bench/tests/budget_cpu.rs new file mode 100644 index 0000000..95241ad --- /dev/null +++ b/crates/mummu-bench/tests/budget_cpu.rs @@ -0,0 +1,75 @@ +//! The CPU-tier perf-budget gate from `bench/BASELINE.md`: Qwen2.5-0.5B +//! greedy decode on the `Cpu` (burn-flex) backend. Ignored by default +//! (weights on disk); run with +//! +//! ```text +//! MUMMU_QWEN2_05B_DIR=path/to/qwen2.5-0.5b cargo test -p mummu-bench --release --test budget_cpu -- --ignored --nocapture +//! ``` + +use std::path::PathBuf; +use std::time::Instant; + +use mummu::backend::Cpu; +use mummu::decode::argmax_id; +use mummu::models::CausalLm; +use mummu::models::qwen2; +use tokenizers::Tokenizer; + +/// From bench/BASELINE.md (recorded 11.7 tok/s on 2026-07-10; ~2x headroom). +const DECODE_BUDGET_TOKENS_PER_S: f64 = 6.0; +const DECODE_STEPS: usize = 8; + +fn qwen2_05b_dir() -> Option { + let dir = PathBuf::from(std::env::var_os("MUMMU_QWEN2_05B_DIR")?); + dir.is_dir().then_some(dir) +} + +#[test] +#[ignore = "needs local Qwen2.5-0.5B weights (MUMMU_QWEN2_05B_DIR)"] +fn qwen2_05b_cpu_decode_stays_inside_its_budget() { + let Some(dir) = qwen2_05b_dir() else { + panic!( + "set MUMMU_QWEN2_05B_DIR to a dir with config.json/tokenizer.json/model.safetensors" + ); + }; + let raw = mummu::chat::ChatMl::qwen2().render(&[ + mummu::chat::Turn::system("You are a concise assistant."), + mummu::chat::Turn::user("What is 2+2? Answer in one short sentence."), + ]); + let tok = Tokenizer::from_file(dir.join("tokenizer.json")).expect("tokenizer.json loads"); + let ids = tok + .encode(raw.as_str(), false) + .expect("encodes") + .get_ids() + .to_vec(); + + let device = burn::tensor::Device::::default(); + let loaded = qwen2::load_from_dir::(&dir, &device).expect("weights load checked"); + + // Prefill (uncounted warm-up: first-run dispatch paths), then time decode. + let mut cache = loaded.new_cache(); + let logits = loaded.forward(&ids, 0, &mut cache, &device); + let mut next = argmax_id(logits).expect("argmax"); + + let start = Instant::now(); + let mut out = Vec::with_capacity(DECODE_STEPS); + for past in (ids.len()..).take(DECODE_STEPS) { + let logits = loaded.forward(&[next], past, &mut cache, &device); + next = argmax_id(logits).expect("argmax"); + out.push(next); + } + let tok_per_s = DECODE_STEPS as f64 / start.elapsed().as_secs_f64(); + let text = tok.decode(&out, true).expect("decode"); + + eprintln!( + "[budget/cpu] 0.5B decode {tok_per_s:.2} tok/s on flex (budget {DECODE_BUDGET_TOKENS_PER_S}); text: {text:?}" + ); + assert!( + text.contains('4'), + "CPU decode must stay coherent, got: {text:?}" + ); + assert!( + tok_per_s >= DECODE_BUDGET_TOKENS_PER_S, + "CPU decode regression: {tok_per_s:.2} tok/s < {DECODE_BUDGET_TOKENS_PER_S} tok/s budget" + ); +} diff --git a/crates/mummu/Cargo.toml b/crates/mummu/Cargo.toml index 62af7ea..36b3b16 100644 --- a/crates/mummu/Cargo.toml +++ b/crates/mummu/Cargo.toml @@ -15,4 +15,6 @@ tokenizers = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } thiserror = { workspace = true } +ureq = { workspace = true } once_cell = { workspace = true } +burn-flex = { workspace = true } diff --git a/crates/mummu/src/backend.rs b/crates/mummu/src/backend.rs index 8893c3a..5b7da37 100644 --- a/crates/mummu/src/backend.rs +++ b/crates/mummu/src/backend.rs @@ -4,7 +4,7 @@ //! choice on a consumer. What this module owns is the *default* policy proven //! in laurelane: enumerate GPU adapters once with a cheap `wgpu` probe (no //! device creation), run on **wgpu (GPU via Vulkan/DX12/Metal — no CUDA -//! toolchain)** when a hardware adapter is present, else **ndarray (CPU)**. +//! toolchain)** when a hardware adapter is present, else **burn-flex (CPU)**. //! No feature-split builds. //! //! The probe also records which adapters advertise `SHADER_F16` — the input @@ -22,8 +22,8 @@ pub type Gpu = burn::backend::Wgpu; /// where `SHADER_F16` is available (check [`DeviceInventory::any_shader_f16`]). pub type GpuF16 = burn::backend::Wgpu; -/// CPU backend (pure-Rust ndarray). -pub type Cpu = burn::backend::NdArray; +/// CPU backend (burn-flex: pure-Rust SIMD + gemm; burn-ndarray's successor). +pub type Cpu = burn_flex::Flex; /// One enumerated GPU adapter, as reported by wgpu. #[derive(Debug, Clone)] @@ -118,7 +118,7 @@ pub fn device_label() -> &'static str { if use_gpu() { "GPU (wgpu)" } else { - "CPU (ndarray)" + "CPU (flex)" } } @@ -151,7 +151,7 @@ mod tests { if use_gpu() { assert_eq!(label, "GPU (wgpu)"); } else { - assert_eq!(label, "CPU (ndarray)"); + assert_eq!(label, "CPU (flex)"); } } diff --git a/crates/mummu/src/cache.rs b/crates/mummu/src/cache.rs new file mode 100644 index 0000000..3d12fb4 --- /dev/null +++ b/crates/mummu/src/cache.rs @@ -0,0 +1,170 @@ +//! Process-lifetime model caching. Loading a checkpoint costs seconds and +//! gigabytes, so consumers keep one [`ModelSlot`] static per (model, backend) +//! and pay the load once. Burn's `Param` is not `Sync`, so the loaded value +//! lives behind a `Mutex` and is only reachable inside [`ModelSlot::with`] — +//! which also serializes inference, the right default for a single GPU. +//! +//! Switching to a different checkpoint dir through the same slot drops the +//! old model (freeing its VRAM/RAM) and loads the new one — this is the +//! active-model-switch primitive the P8 management API builds on. + +use std::path::{Path, PathBuf}; +use std::sync::Mutex; + +struct Entry { + key: PathBuf, + value: T, +} + +/// A one-model cache slot, keyed by checkpoint directory. +pub struct ModelSlot { + inner: Mutex>>, +} + +impl Default for ModelSlot { + fn default() -> Self { + Self::new() + } +} + +impl ModelSlot { + #[must_use] + pub const fn new() -> Self { + Self { + inner: Mutex::new(None), + } + } + + /// Run `f` with the model for `key`, loading it first if the slot is + /// empty or holds a different checkpoint (the old model is dropped + /// before `load` runs, so peak memory stays one model per slot). + pub fn with( + &self, + key: &Path, + load: impl FnOnce(&Path) -> Result, + f: impl FnOnce(&T) -> R, + ) -> Result { + assert!(!key.as_os_str().is_empty(), "model cache: empty key"); + let mut guard = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + let hit = guard.as_ref().is_some_and(|e| e.key == key); + if !hit { + *guard = None; // free the old model before loading the new one + let value = load(key)?; + *guard = Some(Entry { + key: key.to_path_buf(), + value, + }); + } + let entry = guard.as_ref().expect("slot was just filled"); + debug_assert!(entry.key == key, "slot must hold the requested model"); + Ok(f(&entry.value)) + } + + /// Drop the cached model (freeing its VRAM/RAM). No-op when empty. + pub fn clear(&self) { + *self.inner.lock().unwrap_or_else(|e| e.into_inner()) = None; + } + + /// The checkpoint dir currently loaded, if any — for settings UIs. + #[must_use] + pub fn loaded_key(&self) -> Option { + self.inner + .lock() + .unwrap_or_else(|e| e.into_inner()) + .as_ref() + .map(|e| e.key.clone()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::convert::Infallible; + + #[test] + fn loads_once_and_reuses_for_the_same_key() { + let slot: ModelSlot = ModelSlot::new(); + let mut loads = 0; + for _ in 0..3 { + let got = slot + .with::<_, Infallible>( + Path::new("model-a"), + |k| { + loads += 1; + Ok(k.display().to_string()) + }, + |m| m.clone(), + ) + .unwrap(); + assert_eq!(got, "model-a"); + } + assert_eq!(loads, 1, "same key must load exactly once"); + assert_eq!(slot.loaded_key().as_deref(), Some(Path::new("model-a"))); + } + + #[test] + fn switching_key_reloads_and_replaces() { + let slot: ModelSlot = ModelSlot::new(); + let mut loads = 0; + let mut run = |key: &str| { + slot.with::<_, Infallible>( + Path::new(key), + |k| { + loads += 1; + Ok(k.display().to_string()) + }, + |m| m.clone(), + ) + .unwrap() + }; + assert_eq!(run("model-a"), "model-a"); + assert_eq!(run("model-b"), "model-b"); // switch: drop a, load b + assert_eq!(run("model-b"), "model-b"); // hit + assert_eq!(loads, 2); + assert_eq!(slot.loaded_key().as_deref(), Some(Path::new("model-b"))); + } + + #[test] + fn failed_load_leaves_the_slot_empty() { + let slot: ModelSlot = ModelSlot::new(); + let err = slot.with(Path::new("bad"), |_| Err("boom"), |m: &String| m.clone()); + assert_eq!(err, Err("boom")); + assert_eq!(slot.loaded_key(), None, "a failed load must not cache"); + } + + #[test] + fn clear_unloads() { + let slot: ModelSlot = ModelSlot::new(); + slot.with::<_, Infallible>(Path::new("m"), |_| Ok(7), |_| ()) + .unwrap(); + assert!(slot.loaded_key().is_some()); + slot.clear(); + assert_eq!(slot.loaded_key(), None); + } + + #[test] + #[should_panic(expected = "empty key")] + fn empty_key_is_rejected() { + let slot: ModelSlot = ModelSlot::new(); + let _ = slot.with::<_, Infallible>(Path::new(""), |_| Ok(1), |_| ()); + } + + /// The slot is usable as a `static` (the whole point). + static GLOBAL: ModelSlot = ModelSlot::new(); + + #[test] + fn works_as_a_static_across_threads() { + let handles: Vec<_> = (0..4) + .map(|_| { + std::thread::spawn(|| { + GLOBAL + .with::<_, Infallible>(Path::new("shared"), |_| Ok(41), |v| v + 1) + .unwrap() + }) + }) + .collect(); + for h in handles { + assert_eq!(h.join().unwrap(), 42); + } + } +} diff --git a/crates/mummu/src/chat.rs b/crates/mummu/src/chat.rs new file mode 100644 index 0000000..10fbfe6 --- /dev/null +++ b/crates/mummu/src/chat.rs @@ -0,0 +1,175 @@ +//! Explicit chat templates. Prompt wrapping is part of a model's contract — +//! an implicit or slightly-wrong template silently ruins output quality — so +//! templates are code here, never guessed: each per-model constructor is +//! byte-verified against the parity references (the Qwen2 template renders +//! the exact prompt committed in the Candle logits fixture). +//! +//! Both zoo LLMs speak ChatML; LFM2.5 additionally prefixes `<|startoftext|>`. +//! Tool-use templates (Hermes-style, LFM2.5 bracket notation) are P4 follow-ups. + +/// Who is speaking in a [`Turn`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Role { + System, + User, + Assistant, +} + +impl Role { + fn tag(self) -> &'static str { + match self { + Self::System => "system", + Self::User => "user", + Self::Assistant => "assistant", + } + } +} + +/// One message in a conversation. +#[derive(Debug, Clone)] +pub struct Turn { + pub role: Role, + pub content: String, +} + +impl Turn { + #[must_use] + pub fn system(content: impl Into) -> Self { + Self { + role: Role::System, + content: content.into(), + } + } + + #[must_use] + pub fn user(content: impl Into) -> Self { + Self { + role: Role::User, + content: content.into(), + } + } + + #[must_use] + pub fn assistant(content: impl Into) -> Self { + Self { + role: Role::Assistant, + content: content.into(), + } + } +} + +/// 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. +#[derive(Debug, Clone)] +pub struct ChatMl { + bos: Option<&'static str>, +} + +impl ChatMl { + /// Qwen2 / Qwen2.5-Instruct: plain ChatML, no BOS. + #[must_use] + pub fn qwen2() -> Self { + Self { bos: None } + } + + /// LFM2 / LFM2.5-Instruct: ChatML behind `<|startoftext|>`. + #[must_use] + pub fn lfm2() -> Self { + Self { + bos: Some("<|startoftext|>"), + } + } + + /// Render a conversation into the raw prompt string, ending with the open + /// assistant turn the model completes. The caller tokenizes the result + /// with special tokens enabled by the tokenizer itself, not re-added. + #[must_use] + pub fn render(&self, turns: &[Turn]) -> String { + assert!(!turns.is_empty(), "chat render: no turns"); + assert!( + turns.len() <= MAX_TURNS, + "chat render: {} turns exceeds the {MAX_TURNS} bound", + turns.len() + ); + assert!( + turns.last().map(|t| t.role) != Some(Role::Assistant), + "chat render: the template opens the assistant turn itself; \ + a trailing assistant turn would double it" + ); + let mut out = String::from(self.bos.unwrap_or("")); + for turn in turns { + out.push_str("<|im_start|>"); + out.push_str(turn.role.tag()); + out.push('\n'); + out.push_str(&turn.content); + out.push_str("<|im_end|>\n"); + } + out.push_str("<|im_start|>assistant\n"); + debug_assert!(out.ends_with("assistant\n"), "render must open a turn"); + out + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn qwen2_render_matches_the_parity_verified_shape() { + // This exact string (with this system prompt + user text) is what the + // Candle fixture and the Ollama fp16 greedy leg were verified against. + let raw = ChatMl::qwen2().render(&[ + Turn::system("You are a helpful assistant."), + Turn::user("List the first five prime numbers."), + ]); + assert_eq!( + raw, + "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n\ + <|im_start|>user\nList the first five prime numbers.<|im_end|>\n\ + <|im_start|>assistant\n" + ); + } + + #[test] + fn lfm2_render_prefixes_bos_and_matches_the_parity_shape() { + let raw = ChatMl::lfm2().render(&[Turn::user("List the first five prime numbers.")]); + assert_eq!( + raw, + "<|startoftext|><|im_start|>user\nList the first five prime numbers.<|im_end|>\n\ + <|im_start|>assistant\n" + ); + } + + #[test] + fn multi_turn_history_renders_in_order() { + let raw = ChatMl::qwen2().render(&[ + Turn::user("Hi."), + Turn::assistant("Hello!"), + Turn::user("Bye."), + ]); + assert_eq!( + raw, + "<|im_start|>user\nHi.<|im_end|>\n\ + <|im_start|>assistant\nHello!<|im_end|>\n\ + <|im_start|>user\nBye.<|im_end|>\n\ + <|im_start|>assistant\n" + ); + } + + #[test] + #[should_panic(expected = "no turns")] + fn empty_conversation_is_rejected() { + let _ = ChatMl::qwen2().render(&[]); + } + + #[test] + #[should_panic(expected = "double it")] + fn trailing_assistant_turn_is_rejected() { + let _ = ChatMl::qwen2().render(&[Turn::user("q"), Turn::assistant("half-done")]); + } +} diff --git a/crates/mummu/src/decode.rs b/crates/mummu/src/decode.rs index 3ea9a3e..c9161ab 100644 --- a/crates/mummu/src/decode.rs +++ b/crates/mummu/src/decode.rs @@ -1,8 +1,20 @@ -//! Decode-loop primitives shared by every causal model: on-device argmax and -//! a top-k probe. The full engine (sampling, streaming, interrupts) is P5. +//! Decode-loop primitives shared by every causal model: on-device argmax, a +//! top-k probe, temperature/top-p sampling with a deterministic in-house RNG, +//! and the streaming `generate_loop` driver with cooperative cancellation. + +use std::ops::ControlFlow; use burn::tensor::{Tensor, backend::Backend}; +/// Hard ceiling on the vocab a sampled step will read back to the CPU +/// (~4 MB of f32 at the bound); anything larger is a wiring bug, not a model. +const VOCAB_READBACK_BOUND: usize = 1 << 20; + +/// Candidate-set cap when sampling: top-p truncation happens *within* the +/// `top_k` highest-logit tokens, so the post-softmax walk is O(k log k), not +/// O(vocab log vocab). 1024 keeps >99.9% of realistic nucleus mass. +const DEFAULT_TOP_K: usize = 1024; + /// Greedy next-token id from `[1, vocab]` logits. The argmax runs /// **on-device** and only the single winning index is synced back — vs. /// copying a whole ~150k-logit vector to the CPU every decode step. @@ -29,6 +41,214 @@ pub fn top_k_ids(v: &[f32], k: usize) -> Vec { idx.into_iter().take(k).map(|i| i as u32).collect() } +/// Sampling knobs for one generation. `temperature == 0` means exact greedy +/// (argmax stays on-device; nothing else is consulted). +#[derive(Debug, Clone)] +pub struct SamplerOptions { + /// 0 = greedy; higher flattens the distribution. Must be finite and >= 0. + pub temperature: f32, + /// Nucleus mass in (0, 1]: sample only from the smallest prefix of + /// probability-sorted candidates whose mass reaches this. + pub top_p: f32, + /// Candidate-set cap applied before top-p (>= 1). + pub top_k: usize, + /// RNG seed: the same (options, logits, seed) always picks the same token. + pub seed: u64, +} + +impl Default for SamplerOptions { + fn default() -> Self { + Self { + temperature: 0.0, + top_p: 1.0, + top_k: DEFAULT_TOP_K, + seed: 0, + } + } +} + +impl SamplerOptions { + /// Greedy decoding (temperature 0) — the parity-gate configuration. + #[must_use] + pub fn greedy() -> Self { + Self::default() + } + + fn validate(&self) { + assert!( + self.temperature.is_finite() && self.temperature >= 0.0, + "sampler: temperature must be finite and >= 0, got {}", + self.temperature + ); + assert!( + self.top_p > 0.0 && self.top_p <= 1.0, + "sampler: top_p must be in (0, 1], got {}", + self.top_p + ); + assert!(self.top_k >= 1, "sampler: top_k must be >= 1"); + } +} + +/// PCG-XSH-RR 32 (O'Neill): a tiny deterministic RNG so sampling is +/// reproducible from a seed without pulling in a rand dependency. +pub struct Pcg32 { + state: u64, + inc: u64, +} + +impl Pcg32 { + const MULT: u64 = 6_364_136_223_846_793_005; + + #[must_use] + pub fn new(seed: u64) -> Self { + // Fixed stream; the standard seeding dance (advance, add, advance). + let mut rng = Self { + state: 0, + inc: (54 << 1) | 1, + }; + rng.next_u32(); + rng.state = rng.state.wrapping_add(seed); + rng.next_u32(); + rng + } + + pub fn next_u32(&mut self) -> u32 { + let old = self.state; + self.state = old.wrapping_mul(Self::MULT).wrapping_add(self.inc); + let xorshifted = (((old >> 18) ^ old) >> 27) as u32; + let rot = (old >> 59) as u32; + xorshifted.rotate_right(rot) + } + + /// Uniform in [0, 1) with 24 bits of mantissa. + pub fn next_f32(&mut self) -> f32 { + (self.next_u32() >> 8) as f32 * (1.0 / (1 << 24) as f32) + } +} + +/// Sample a token id from raw logits with temperature + top-k + top-p. +/// Pure and deterministic given (logits, opts, rng state). `temperature == 0` +/// callers should use [`argmax_id`] instead (asserted here). +#[must_use] +pub fn sample_id(logits: &[f32], opts: &SamplerOptions, rng: &mut Pcg32) -> u32 { + opts.validate(); + assert!(!logits.is_empty(), "sample_id: empty logits"); + assert!( + logits.len() <= VOCAB_READBACK_BOUND, + "sample_id: vocab {} exceeds the readback bound", + logits.len() + ); + assert!( + opts.temperature > 0.0, + "sample_id: temperature 0 is the argmax path" + ); + + // Top-k prefilter: O(vocab) partial select, then sort just the candidates. + let k = opts.top_k.min(logits.len()); + let mut idx: Vec = (0..logits.len() as u32).collect(); + let by_logit_desc = |&a: &u32, &b: &u32| logits[b as usize].total_cmp(&logits[a as usize]); + if k < idx.len() { + idx.select_nth_unstable_by(k - 1, by_logit_desc); + idx.truncate(k); + } + idx.sort_unstable_by(by_logit_desc); + + // Temperature softmax over the candidates (max-subtracted: never overflows). + let max_logit = logits[idx[0] as usize]; + let mut probs: Vec = idx + .iter() + .map(|&i| ((logits[i as usize] - max_logit) / opts.temperature).exp()) + .collect(); + let total: f32 = probs.iter().sum(); + debug_assert!(total > 0.0, "softmax mass must be positive"); + for p in &mut probs { + *p /= total; + } + + // Nucleus: keep the smallest probability-sorted prefix with mass >= top_p + // (probs are already descending because idx is logit-sorted). + let mut cut = probs.len(); + let mut mass = 0.0_f32; + for (i, &p) in probs.iter().enumerate() { + mass += p; + if mass >= opts.top_p { + cut = i + 1; + break; + } + } + debug_assert!(cut >= 1, "nucleus must keep at least the top token"); + + // Draw within the (renormalized) nucleus by cumulative walk. + let nucleus_mass: f32 = probs[..cut].iter().sum(); + let mut u = rng.next_f32() * nucleus_mass; + let mut chosen = idx[cut - 1]; // fallback: rounding can leave u > 0 at the end + for (i, &p) in probs[..cut].iter().enumerate() { + if u < p { + chosen = idx[i]; + break; + } + u -= p; + } + assert!( + (chosen as usize) < logits.len(), + "sampled id out of the vocab" + ); + chosen +} + +/// The shared decode driver: prefill once via `step`, then one token per +/// iteration. Emits each accepted token through `on_token`; a `Break` return +/// cancels cooperatively *before* the next forward. EOS is never emitted. +/// `step(new_ids, past)` returns `[1, vocab]` logits for the last position. +pub fn generate_loop( + mut step: impl FnMut(&[u32], usize) -> Tensor, + prompt_ids: &[u32], + max_tokens: usize, + opts: &SamplerOptions, + is_eos: impl Fn(u32) -> bool, + mut on_token: impl FnMut(u32) -> ControlFlow<()>, +) -> Result, String> { + opts.validate(); + assert!(!prompt_ids.is_empty(), "generate_loop: empty prompt"); + assert!(max_tokens >= 1, "generate_loop: max_tokens must be >= 1"); + + let mut rng = Pcg32::new(opts.seed); + let greedy = opts.temperature == 0.0; + let mut logits = step(prompt_ids, 0); + let mut out: Vec = Vec::with_capacity(max_tokens); + for past in (prompt_ids.len()..).take(max_tokens) { + let vocab = logits.dims()[1] as u32; + let next = if greedy { + argmax_id(logits)? + } else { + let v = logits + .into_data() + .convert::() + .to_vec::() + .map_err(|e| format!("logits readback: {e:?}"))?; + sample_id(&v, opts, &mut rng) + }; + // A GPU argmax over NaN logits can return an out-of-range sentinel + // (observed: exactly `vocab` on f16 numeric collapse) — fail loudly + // instead of emitting garbage ids the tokenizer silently drops. + if next >= vocab { + return Err(format!( + "decode step {past}: id {next} is outside the {vocab}-token vocab — NaN logits / numeric collapse on this backend?" + )); + } + if is_eos(next) { + break; + } + out.push(next); + if on_token(next).is_break() { + break; + } + logits = step(&[next], past); + } + debug_assert!(out.len() <= max_tokens); + Ok(out) +} + #[cfg(test)] mod tests { use super::*; @@ -42,6 +262,155 @@ mod tests { assert_eq!(argmax_id(logits).unwrap(), 2); } + #[test] + fn pcg32_is_deterministic_per_seed_and_in_unit_range() { + let (mut a, mut b) = (Pcg32::new(7), Pcg32::new(7)); + let seq_a: Vec = (0..8).map(|_| a.next_u32()).collect(); + let seq_b: Vec = (0..8).map(|_| b.next_u32()).collect(); + assert_eq!(seq_a, seq_b, "same seed must replay the same stream"); + + let mut c = Pcg32::new(8); + let seq_c: Vec = (0..8).map(|_| c.next_u32()).collect(); + assert_ne!(seq_a, seq_c, "different seeds must diverge"); + + let mut r = Pcg32::new(99); + for _ in 0..1000 { + let f = r.next_f32(); + assert!((0.0..1.0).contains(&f), "next_f32 out of [0,1): {f}"); + } + } + + #[test] + fn sample_id_peaked_logits_always_pick_the_peak() { + let logits = [0.0f32, 30.0, -5.0, 1.0]; + let opts = SamplerOptions { + temperature: 0.8, + top_p: 0.95, + ..SamplerOptions::default() + }; + for seed in 0..32 { + let mut rng = Pcg32::new(seed); + assert_eq!(sample_id(&logits, &opts, &mut rng), 1); + } + } + + #[test] + fn sample_id_top_k_one_is_argmax_at_any_temperature() { + let logits = [1.0f32, 3.0, 2.0, 2.9]; + let opts = SamplerOptions { + temperature: 10.0, + top_p: 1.0, + top_k: 1, + seed: 0, + }; + for seed in 0..16 { + let mut rng = Pcg32::new(seed); + assert_eq!( + sample_id( + &logits, + &SamplerOptions { + seed, + ..opts.clone() + }, + &mut rng + ), + 1 + ); + } + } + + #[test] + fn sample_id_tiny_top_p_degenerates_to_argmax() { + let logits = [1.0f32, 1.1, 0.9, 1.05]; + let opts = SamplerOptions { + temperature: 5.0, + top_p: 0.01, + ..SamplerOptions::default() + }; + for seed in 0..16 { + let mut rng = Pcg32::new(seed); + assert_eq!(sample_id(&logits, &opts, &mut rng), 1); + } + } + + #[test] + fn sample_id_high_temperature_spreads_over_candidates() { + let logits = [2.0f32, 2.0, 2.0, 2.0]; + let opts = SamplerOptions { + temperature: 1.0, + top_p: 1.0, + ..SamplerOptions::default() + }; + let picks: std::collections::HashSet = (0..64) + .map(|seed| sample_id(&logits, &opts, &mut Pcg32::new(seed))) + .collect(); + assert!( + picks.len() >= 3, + "uniform logits over 64 seeds should hit >= 3 of 4 ids, got {picks:?}" + ); + for &p in &picks { + assert!(p < 4); + } + } + + #[test] + #[should_panic(expected = "argmax path")] + fn sample_id_rejects_temperature_zero() { + let mut rng = Pcg32::new(0); + let _ = sample_id(&[1.0, 2.0], &SamplerOptions::greedy(), &mut rng); + } + + /// A fixed toy vocab where the "model" always prefers id 2, then id 3 + /// after seeing 2 — enough to drive the loop without weights. + fn toy_step(device: &burn::tensor::Device) -> impl FnMut(&[u32], usize) -> Tensor { + let device = *device; + move |new_ids: &[u32], _past: usize| { + let peak = if new_ids.last() == Some(&2) { 3 } else { 2 }; + let mut v = vec![0.0f32; 8]; + v[peak] = 9.0; + Tensor::::from_floats(v.as_slice(), &device).reshape([1, 8]) + } + } + + #[test] + fn generate_loop_greedy_follows_argmax_and_stops_at_eos() { + let device = burn::tensor::Device::::default(); + let out = generate_loop( + toy_step(&device), + &[1], + 6, + &SamplerOptions::greedy(), + |id| id == 3, // treat the follow-up token as EOS + |_| std::ops::ControlFlow::Continue(()), + ) + .unwrap(); + assert_eq!(out, vec![2], "one token, then EOS never emitted"); + } + + #[test] + fn generate_loop_cancels_cooperatively_between_tokens() { + let device = burn::tensor::Device::::default(); + let mut streamed = Vec::new(); + let out = generate_loop( + toy_step(&device), + &[1], + 100, + &SamplerOptions::greedy(), + |_| false, // no EOS: only the callback can stop this + |id| { + streamed.push(id); + if streamed.len() == 2 { + std::ops::ControlFlow::Break(()) + } else { + std::ops::ControlFlow::Continue(()) + } + }, + ) + .unwrap(); + assert_eq!(out.len(), 2, "break after the 2nd token stops the loop"); + assert_eq!(streamed, out, "every emitted token was streamed"); + } + #[test] fn top_k_ids_orders_descending() { let v = [0.1f32, 5.0, -2.0, 3.0]; diff --git a/crates/mummu/src/hub.rs b/crates/mummu/src/hub.rs new file mode 100644 index 0000000..3424c40 --- /dev/null +++ b/crates/mummu/src/hub.rs @@ -0,0 +1,318 @@ +//! Model downloads: HuggingFace Hub (or any HTTP host) → the local model +//! cache. Streaming, **resumable** (a `.part` picks up where a killed +//! download stopped, via HTTP `Range`), length-verified, and +//! **sharded-checkpoint aware** (`model.safetensors.index.json` → fetch every +//! shard). Progress surfaces through a callback so app settings UIs can show +//! it (P8). Completed files are cache-first: an existing destination is never +//! re-fetched. + +use std::io::{Read, Write}; +use std::path::{Path, PathBuf}; + +/// Streaming copy granularity: big enough to amortize syscalls, small enough +/// to keep progress callbacks responsive. +const CHUNK_BYTES: usize = 64 * 1024; + +/// Hard per-file ceiling — larger than any model shard we'd fetch (shards are +/// conventionally ≤ ~10 GB); anything past this is a wiring bug or a hostile +/// server, not a model. +const MAX_FILE_BYTES: u64 = 64 << 30; + +/// Ceiling on shards in an index — the largest public checkpoints ship tens. +const MAX_SHARDS: usize = 512; + +/// Everything that can go wrong fetching a model. +#[derive(Debug, thiserror::Error)] +pub enum HubError { + #[error("http {url}: {reason}")] + Http { url: String, reason: String }, + #[error("io {path}: {reason}")] + Io { path: PathBuf, reason: String }, + #[error("{url}: expected {expected} bytes, received {received}")] + Incomplete { + url: String, + expected: u64, + received: u64, + }, + #[error("shard index {path}: {reason}")] + BadIndex { path: PathBuf, reason: String }, +} + +/// Download progress for one file, reported after every chunk. +#[derive(Debug, Clone)] +pub struct Progress<'a> { + pub file: &'a str, + pub received_bytes: u64, + /// Total including any resumed prefix; `None` when the server omits it. + pub total_bytes: Option, +} + +/// `https://huggingface.co/{repo}/resolve/{revision}/{file}` — the Hub's +/// stable raw-file endpoint. +#[must_use] +pub fn hub_file_url(repo: &str, revision: &str, file: &str) -> String { + assert!( + !repo.is_empty() && repo.contains('/'), + "repo must be owner/name, got {repo:?}" + ); + assert!(!revision.is_empty(), "revision must be non-empty"); + format!("https://huggingface.co/{repo}/resolve/{revision}/{file}") +} + +/// The unique shard files referenced by a `*.index.json` (weight_map values, +/// deduped, sorted for a deterministic fetch order). +pub fn shards_from_index(index_json: &[u8], index_path: &Path) -> Result, HubError> { + let v: serde_json::Value = + serde_json::from_slice(index_json).map_err(|e| HubError::BadIndex { + path: index_path.to_path_buf(), + reason: e.to_string(), + })?; + let map = v["weight_map"] + .as_object() + .ok_or_else(|| HubError::BadIndex { + path: index_path.to_path_buf(), + reason: "no weight_map object".into(), + })?; + let mut shards: Vec = map + .values() + .filter_map(|s| s.as_str().map(str::to_string)) + .collect(); + shards.sort_unstable(); + shards.dedup(); + if shards.is_empty() || shards.len() > MAX_SHARDS { + return Err(HubError::BadIndex { + path: index_path.to_path_buf(), + reason: format!("{} shards (expected 1..={MAX_SHARDS})", shards.len()), + }); + } + Ok(shards) +} + +/// The in-flight twin of `dest` (`.part`). +fn part_path(dest: &Path) -> PathBuf { + let mut p = dest.as_os_str().to_owned(); + p.push(".part"); + PathBuf::from(p) +} + +/// Fetch `url` into `dest`, streaming through `.part` and resuming any +/// earlier partial download. No-op when `dest` already exists (cache-first). +/// `on_progress` fires after every chunk with cumulative counts. +pub fn fetch_file( + url: &str, + dest: &Path, + mut on_progress: impl FnMut(Progress<'_>), +) -> Result<(), HubError> { + assert!(url.starts_with("https://"), "refusing non-https url: {url}"); + if dest.exists() { + return Ok(()); // cache hit — never re-fetch a completed file + } + let file_label = dest + .file_name() + .map(|n| n.to_string_lossy().into_owned()) + .unwrap_or_default(); + assert!(!file_label.is_empty(), "dest must name a file: {dest:?}"); + if let Some(parent) = dest.parent() { + std::fs::create_dir_all(parent).map_err(|e| HubError::Io { + path: parent.to_path_buf(), + reason: e.to_string(), + })?; + } + + let part = part_path(dest); + let resume_from = std::fs::metadata(&part).map(|m| m.len()).unwrap_or(0); + + let mut req = ureq::get(url); + if resume_from > 0 { + req = req.header("Range", format!("bytes={resume_from}-")); + } + let mut resp = req.call().map_err(|e| HubError::Http { + url: url.into(), + reason: e.to_string(), + })?; + // A server that ignores Range (200 instead of 206) restarts the body from + // byte 0 — truncate our part file to match, never splice mismatched halves. + let resumed = resp.status() == 206 && resume_from > 0; + let body_len: Option = resp + .headers() + .get("content-length") + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.parse().ok()); + let already = if resumed { resume_from } else { 0 }; + let total = body_len.map(|l| l + already); + if let Some(t) = total { + assert!( + t <= MAX_FILE_BYTES, + "{url}: {t} bytes exceeds the file bound" + ); + } + + let mut out = std::fs::OpenOptions::new() + .create(true) + .append(resumed) + .write(true) + .truncate(!resumed) + .open(&part) + .map_err(|e| HubError::Io { + path: part.clone(), + reason: e.to_string(), + })?; + + let mut reader = resp.body_mut().as_reader(); + let mut received = already; + let mut buf = vec![0u8; CHUNK_BYTES]; + loop { + let n = reader.read(&mut buf).map_err(|e| HubError::Http { + url: url.into(), + reason: e.to_string(), + })?; + if n == 0 { + break; + } + out.write_all(&buf[..n]).map_err(|e| HubError::Io { + path: part.clone(), + reason: e.to_string(), + })?; + received += n as u64; + assert!( + received <= MAX_FILE_BYTES, + "{url}: stream exceeded the file bound" + ); + on_progress(Progress { + file: &file_label, + received_bytes: received, + total_bytes: total, + }); + } + drop(out); + + if let Some(expected) = total + && received != expected + { + // Keep the .part for a future resume; report loudly. + return Err(HubError::Incomplete { + url: url.into(), + expected, + received, + }); + } + std::fs::rename(&part, dest).map_err(|e| HubError::Io { + path: dest.to_path_buf(), + reason: e.to_string(), + }) +} + +/// Fetch a whole model from the Hub into `dest_dir`: `config.json`, +/// `tokenizer.json`, and the weights — `model.safetensors` when the repo is +/// single-file, else every shard listed by `model.safetensors.index.json`. +/// Returns `dest_dir` ready for the per-model `load_from_dir`. +pub fn fetch_model( + repo: &str, + revision: &str, + dest_dir: &Path, + mut on_progress: impl FnMut(Progress<'_>), +) -> Result { + for file in ["config.json", "tokenizer.json"] { + fetch_file( + &hub_file_url(repo, revision, file), + &dest_dir.join(file), + &mut on_progress, + )?; + } + // Single-file first (the common case for the small-model tiers we target). + let single = fetch_file( + &hub_file_url(repo, revision, "model.safetensors"), + &dest_dir.join("model.safetensors"), + &mut on_progress, + ); + if single.is_ok() { + return Ok(dest_dir.to_path_buf()); + } + // Fall back to a sharded checkpoint; if there's no index either, report + // the original single-file error (the more useful signal). + let index_name = "model.safetensors.index.json"; + let index_dest = dest_dir.join(index_name); + if fetch_file( + &hub_file_url(repo, revision, index_name), + &index_dest, + &mut on_progress, + ) + .is_err() + { + return single.map(|()| dest_dir.to_path_buf()); + } + let index_bytes = std::fs::read(&index_dest).map_err(|e| HubError::Io { + path: index_dest.clone(), + reason: e.to_string(), + })?; + for shard in shards_from_index(&index_bytes, &index_dest)? { + fetch_file( + &hub_file_url(repo, revision, &shard), + &dest_dir.join(&shard), + &mut on_progress, + )?; + } + Ok(dest_dir.to_path_buf()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn hub_url_has_the_resolve_shape() { + assert_eq!( + hub_file_url("Qwen/Qwen2.5-1.5B-Instruct", "main", "config.json"), + "https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct/resolve/main/config.json" + ); + } + + #[test] + #[should_panic(expected = "owner/name")] + fn bare_repo_names_are_rejected() { + let _ = hub_file_url("qwen", "main", "config.json"); + } + + #[test] + fn shard_index_dedupes_and_sorts() { + let idx = br#"{"metadata":{},"weight_map":{ + "a.weight":"model-00002-of-00002.safetensors", + "b.weight":"model-00001-of-00002.safetensors", + "c.weight":"model-00001-of-00002.safetensors"}}"#; + let shards = shards_from_index(idx, Path::new("x.index.json")).unwrap(); + assert_eq!( + shards, + vec![ + "model-00001-of-00002.safetensors", + "model-00002-of-00002.safetensors" + ] + ); + } + + #[test] + fn shard_index_without_weight_map_is_rejected() { + let err = shards_from_index(b"{}", Path::new("x.index.json")); + assert!(matches!(err, Err(HubError::BadIndex { .. }))); + } + + #[test] + fn part_path_appends_suffix() { + assert_eq!( + part_path(Path::new("m/model.safetensors")), + Path::new("m/model.safetensors.part") + ); + } + + #[test] + fn existing_dest_is_a_cache_hit_without_any_http() { + // A bogus URL proves no request is made when the file already exists. + let dir = std::env::temp_dir().join("mummu-hub-test-cache-hit"); + std::fs::create_dir_all(&dir).unwrap(); + let dest = dir.join("present.bin"); + std::fs::write(&dest, b"already here").unwrap(); + let mut calls = 0; + fetch_file("https://invalid.invalid/x", &dest, |_| calls += 1).unwrap(); + assert_eq!(calls, 0, "cache hit must not stream"); + std::fs::remove_dir_all(&dir).ok(); + } +} diff --git a/crates/mummu/src/lib.rs b/crates/mummu/src/lib.rs index 712073f..cc6b7e8 100644 --- a/crates/mummu/src/lib.rs +++ b/crates/mummu/src/lib.rs @@ -10,8 +10,12 @@ #![recursion_limit = "512"] pub mod backend; +pub mod cache; +pub mod chat; pub mod decode; +pub mod hub; pub mod import; pub mod manage; pub mod models; pub mod nn; +pub mod registry; diff --git a/crates/mummu/src/manage.rs b/crates/mummu/src/manage.rs index 2c3428e..5968e81 100644 --- a/crates/mummu/src/manage.rs +++ b/crates/mummu/src/manage.rs @@ -220,3 +220,153 @@ mod tests { fs::remove_dir_all(&root).unwrap(); } } + +/// The settings-UI-facing management surface (P8): one object owning the +/// models root that composes the catalog ([`crate::registry`]), downloads +/// ([`crate::hub`], with per-chunk progress), disk accounting, and safe +/// removal. Active-model *switching* is the consumer's `ModelSlot` keyed by +/// [`ModelManager::model_dir`]. +pub struct ModelManager { + root: PathBuf, + catalog: Vec, +} + +impl ModelManager { + /// Manage `root` with the built-in catalog. + #[must_use] + pub fn new(root: PathBuf) -> Self { + Self::with_catalog(root, crate::registry::catalog()) + } + + /// Manage `root` with an app-supplied catalog (all specs must validate). + #[must_use] + pub fn with_catalog(root: PathBuf, catalog: Vec) -> Self { + assert!( + !root.as_os_str().is_empty(), + "models root must be non-empty" + ); + assert!(!catalog.is_empty(), "catalog must not be empty"); + for spec in &catalog { + if let Err(e) = spec.validate() { + panic!("invalid catalog spec: {e}"); + } + } + Self { root, catalog } + } + + #[must_use] + pub fn catalog(&self) -> &[crate::registry::ModelSpec] { + &self.catalog + } + + /// The dir a catalog model lives in (whether or not it's installed yet). + pub fn model_dir(&self, name: &str) -> Result { + self.spec(name).map(|s| s.dir(&self.root)) + } + + /// Is every required artifact of `name` on disk? (config + tokenizer + + /// single-file weights or a shard index.) + pub fn is_installed(&self, name: &str) -> Result { + let dir = self.model_dir(name)?; + let weights = dir.join("model.safetensors").is_file() + || dir.join("model.safetensors.index.json").is_file(); + Ok(weights && dir.join("config.json").is_file() && dir.join("tokenizer.json").is_file()) + } + + /// Download `name` from its spec (resumable, cache-first), reporting + /// progress per chunk. Returns the model dir ready for `load_from_dir`. + pub fn install( + &self, + name: &str, + on_progress: impl FnMut(crate::hub::Progress<'_>), + ) -> Result { + let spec = self.spec(name)?; + spec.fetch(&self.root, on_progress) + .map_err(|e| e.to_string()) + } + + /// Remove `name`'s files from disk (traversal-safe). The caller drops any + /// live `ModelSlot` first — removal only touches the disk. + pub fn remove(&self, name: &str) -> Result<(), String> { + let target = resolve_removal(&self.root, name)?; + std::fs::remove_dir_all(&target).map_err(|e| format!("remove {name:?}: {e}")) + } + + /// Disk usage for everything under the root, largest first. + #[must_use] + pub fn disk_report(&self) -> DiskReport { + report(&self.root) + } + + fn spec(&self, name: &str) -> Result<&crate::registry::ModelSpec, String> { + self.catalog.iter().find(|s| s.name == name).ok_or_else(|| { + let known: Vec<&str> = self.catalog.iter().map(|s| s.name.as_str()).collect(); + format!("unknown model {name:?}; catalog has {known:?}") + }) + } +} + +#[cfg(test)] +mod manager_tests { + use super::*; + + fn temp_root(tag: &str) -> PathBuf { + let root = std::env::temp_dir().join(format!("mummu-manager-{tag}")); + let _ = std::fs::remove_dir_all(&root); + std::fs::create_dir_all(&root).unwrap(); + root + } + + fn fake_install(root: &Path, name: &str) { + let dir = root.join(name); + std::fs::create_dir_all(&dir).unwrap(); + for f in ["config.json", "tokenizer.json", "model.safetensors"] { + std::fs::write(dir.join(f), b"{}").unwrap(); + } + } + + #[test] + fn install_state_and_report_reflect_disk() { + let root = temp_root("state"); + let mgr = ModelManager::new(root.clone()); + assert_eq!(mgr.is_installed("all-minilm-l6-v2"), Ok(false)); + + fake_install(&root, "all-minilm-l6-v2"); + assert_eq!(mgr.is_installed("all-minilm-l6-v2"), Ok(true)); + + let rep = mgr.disk_report(); + assert_eq!(rep.models.len(), 1); + assert_eq!(rep.models[0].name, "all-minilm-l6-v2"); + assert!(rep.total_bytes > 0); + std::fs::remove_dir_all(&root).ok(); + } + + #[test] + fn remove_deletes_only_the_named_model() { + let root = temp_root("remove"); + let mgr = ModelManager::new(root.clone()); + fake_install(&root, "all-minilm-l6-v2"); + fake_install(&root, "qwen2.5-0.5b-instruct"); + + mgr.remove("all-minilm-l6-v2").unwrap(); + assert_eq!(mgr.is_installed("all-minilm-l6-v2"), Ok(false)); + assert_eq!(mgr.is_installed("qwen2.5-0.5b-instruct"), Ok(true)); + std::fs::remove_dir_all(&root).ok(); + } + + #[test] + fn unknown_names_fail_loudly_with_the_catalog() { + let root = temp_root("unknown"); + let mgr = ModelManager::new(root.clone()); + let err = mgr.model_dir("nope").unwrap_err(); + assert!(err.contains("nope") && err.contains("all-minilm-l6-v2")); + assert!(mgr.remove("nope").is_err()); + std::fs::remove_dir_all(&root).ok(); + } + + #[test] + #[should_panic(expected = "catalog must not be empty")] + fn empty_catalog_is_rejected() { + let _ = ModelManager::with_catalog(PathBuf::from("x"), vec![]); + } +} diff --git a/crates/mummu/src/models/lfm2.rs b/crates/mummu/src/models/lfm2.rs index 617c97f..4675fd9 100644 --- a/crates/mummu/src/models/lfm2.rs +++ b/crates/mummu/src/models/lfm2.rs @@ -15,8 +15,8 @@ use burn::nn::{Embedding, EmbeddingConfig, RmsNorm, RmsNormConfig}; use burn::store::{ModuleAdapter, PyTorchToBurnAdapter, SafetensorsStore}; use burn::tensor::{Int, Tensor, TensorData, backend::Backend}; -use crate::decode::{argmax_id, top_k_ids}; use crate::import::{CastFloatAdapter, ImportError, load_checked, required_file}; +use crate::models::CausalLm; use crate::models::qwen2::EosIds; use crate::nn::{ ConvState, GqaAttention, GqaAttentionConfig, LayerKv, ShortConv, ShortConvConfig, SwiGluMlp, @@ -219,10 +219,15 @@ pub fn load_from_dir( Ok(LoadedLfm2 { model, config }) } -impl LoadedLfm2 { +impl CausalLm for LoadedLfm2 { + type Cache = Vec>; + + fn is_eos(&self, id: u32) -> bool { + self.config.eos_token_id.contains(id) + } + /// A fresh per-layer cache matching `layer_types`. - #[must_use] - pub fn new_cache(&self) -> Vec> { + fn new_cache(&self) -> Self::Cache { (0..self.config.num_hidden_layers) .map(|i| { if self.config.is_attention(i) { @@ -236,11 +241,11 @@ impl LoadedLfm2 { /// Forward `new_ids` (the whole prompt when `past == 0`, else one decode /// token), updating `cache`; returns logits for the last position `[1, vocab]`. - pub fn forward( + fn forward( &self, new_ids: &[u32], past: usize, - cache: &mut [HybridKv], + cache: &mut Self::Cache, device: &B::Device, ) -> Tensor { let t = new_ids.len(); @@ -288,50 +293,6 @@ impl LoadedLfm2 { let w = self.model.embed_tokens.weight.val(); // tied lm-head last.matmul(w.swap_dims(0, 1)) } - - /// Greedy decode: prefill once, then one token per step, stopping at any - /// config EOS id or `max_tokens`. Argmax runs on-device. - pub fn greedy_generate( - &self, - prompt_ids: &[u32], - max_tokens: usize, - device: &B::Device, - ) -> Result, String> { - assert!(!prompt_ids.is_empty(), "greedy_generate: empty prompt"); - assert!(max_tokens >= 1, "greedy_generate: max_tokens must be >= 1"); - let mut cache = self.new_cache(); - let mut logits = self.forward(prompt_ids, 0, &mut cache, device); - let mut out: Vec = Vec::with_capacity(max_tokens); - for past in (prompt_ids.len()..).take(max_tokens) { - let next = argmax_id(logits)?; - if self.config.eos_token_id.contains(next) { - break; - } - out.push(next); - logits = self.forward(&[next], past, &mut cache, device); - } - debug_assert!(out.len() <= max_tokens); - Ok(out) - } - - /// Parity probe: top-k next-token ids for a single prefill. - pub fn first_token( - &self, - prompt_ids: &[u32], - k: usize, - device: &B::Device, - ) -> Result, String> { - assert!(!prompt_ids.is_empty(), "first_token: empty prompt"); - assert!(k >= 1, "first_token: k must be >= 1"); - let mut cache = self.new_cache(); - let logits = self.forward(prompt_ids, 0, &mut cache, device); - let v = logits - .into_data() - .convert::() - .to_vec::() - .map_err(|e| format!("logits readback: {e:?}"))?; - Ok(top_k_ids(&v, k)) - } } #[cfg(test)] diff --git a/crates/mummu/src/models/mod.rs b/crates/mummu/src/models/mod.rs index c6de96b..30dda60 100644 --- a/crates/mummu/src/models/mod.rs +++ b/crates/mummu/src/models/mod.rs @@ -2,6 +2,94 @@ //! generic over `B: Backend`, all config-driven (hyperparameters come from the //! checkpoint's `config.json`, never hardcoded). +use burn::tensor::{Tensor, backend::Backend}; + +use crate::decode::{SamplerOptions, generate_loop, top_k_ids}; + pub mod lfm2; pub mod minilm; pub mod qwen2; + +/// The contract every causal LM in the zoo implements. A new architecture +/// (Hermes-class function-caller, Gemma, Qwen3, …) provides its cache type, +/// its forward pass, and its EOS check — decoding (greedy, sampled, streamed, +/// cancellable) comes for free from the shared driver. +pub trait CausalLm { + /// Per-generation decode state (KV cache, conv state, …). + type Cache; + + /// A fresh (empty) cache for one generation. + fn new_cache(&self) -> Self::Cache; + + /// Forward `new_ids` (the whole prompt when `past == 0`, else one decode + /// token), updating `cache`; returns logits for the **last** position, + /// `[1, vocab]`. + fn forward( + &self, + new_ids: &[u32], + past: usize, + cache: &mut Self::Cache, + device: &B::Device, + ) -> Tensor; + + /// Is `id` an end-of-sequence token for this model? + fn is_eos(&self, id: u32) -> bool; + + /// Full decode: prefill once, then one token per step, stopping at EOS, + /// `max_tokens`, or a `Break` from `on_token` (streaming + cooperative + /// cancellation). Greedy (`temperature == 0`) keeps the argmax on-device. + fn generate( + &self, + prompt_ids: &[u32], + max_tokens: usize, + opts: &SamplerOptions, + device: &B::Device, + on_token: impl FnMut(u32) -> std::ops::ControlFlow<()>, + ) -> Result, String> { + let mut cache = self.new_cache(); + generate_loop( + |ids, past| self.forward(ids, past, &mut cache, device), + prompt_ids, + max_tokens, + opts, + |id| self.is_eos(id), + on_token, + ) + } + + /// Greedy decode (the parity-gate path): [`Self::generate`] at + /// temperature 0 with no streaming. + fn greedy_generate( + &self, + prompt_ids: &[u32], + max_tokens: usize, + device: &B::Device, + ) -> Result, String> { + self.generate( + prompt_ids, + max_tokens, + &SamplerOptions::greedy(), + device, + |_| std::ops::ControlFlow::Continue(()), + ) + } + + /// Parity probe: top-k next-token ids for a single prefill. + fn first_token( + &self, + prompt_ids: &[u32], + k: usize, + device: &B::Device, + ) -> Result, String> { + assert!(!prompt_ids.is_empty(), "first_token: empty prompt"); + assert!(k >= 1, "first_token: k must be >= 1"); + let mut cache = self.new_cache(); + let logits = self.forward(prompt_ids, 0, &mut cache, device); + let v = logits + .into_data() + .convert::() + .to_vec::() + .map_err(|e| format!("logits readback: {e:?}"))?; + Ok(top_k_ids(&v, k)) + } +} diff --git a/crates/mummu/src/models/qwen2.rs b/crates/mummu/src/models/qwen2.rs index 198046c..941d715 100644 --- a/crates/mummu/src/models/qwen2.rs +++ b/crates/mummu/src/models/qwen2.rs @@ -15,8 +15,8 @@ use burn::nn::{Embedding, EmbeddingConfig, RmsNorm, RmsNormConfig}; use burn::store::{ModuleAdapter, PyTorchToBurnAdapter, SafetensorsStore}; use burn::tensor::{Int, Tensor, TensorData, backend::Backend}; -use crate::decode::{argmax_id, top_k_ids}; use crate::import::{CastFloatAdapter, ImportError, load_checked, required_file}; +use crate::models::CausalLm; use crate::nn::{ GqaAttention, GqaAttentionConfig, LayerKv, SwiGluMlp, SwiGluMlpConfig, causal_mask, rope_tables, }; @@ -185,21 +185,22 @@ pub fn load_from_dir( Ok(LoadedQwen2 { model, config }) } -impl LoadedQwen2 { - /// A fresh (empty) per-layer KV cache. - #[must_use] - pub fn new_cache(&self) -> Vec> { +impl CausalLm for LoadedQwen2 { + type Cache = Vec>; + + fn new_cache(&self) -> Self::Cache { (0..self.config.num_hidden_layers).map(|_| None).collect() } - /// Forward `new_ids` (the whole prompt when `past == 0`, else one decode - /// token), updating `cache`; returns logits for the **last** position, - /// `[1, vocab]`. - pub fn forward( + fn is_eos(&self, id: u32) -> bool { + self.config.eos_token_id.contains(id) + } + + fn forward( &self, new_ids: &[u32], past: usize, - cache: &mut [LayerKv], + cache: &mut Self::Cache, device: &B::Device, ) -> Tensor { let t = new_ids.len(); @@ -212,7 +213,7 @@ impl LoadedQwen2 { ); let cfg = &self.config; - // i32 token ids: native for wgpu, upcast for ndarray — portable either way. + // i32 token ids: native for wgpu and the flex CPU backend alike. let ids32: Vec = new_ids.iter().map(|&i| i as i32).collect(); let input = Tensor::::from_data(TensorData::new(ids32, [t]), device).reshape([1, t]); @@ -246,52 +247,6 @@ impl LoadedQwen2 { let w = self.model.embed_tokens.weight.val(); // [vocab, hidden] last.matmul(w.swap_dims(0, 1)) // [1, vocab] } - - /// Greedy decode: prefill `prompt_ids` once, then one token per step - /// through the KV cache, stopping at any config EOS id or `max_tokens`. - /// The argmax runs on-device; only the winning index is synced back. - pub fn greedy_generate( - &self, - prompt_ids: &[u32], - max_tokens: usize, - device: &B::Device, - ) -> Result, String> { - assert!(!prompt_ids.is_empty(), "greedy_generate: empty prompt"); - assert!(max_tokens >= 1, "greedy_generate: max_tokens must be >= 1"); - let mut cache = self.new_cache(); - let mut logits = self.forward(prompt_ids, 0, &mut cache, device); - let mut out: Vec = Vec::with_capacity(max_tokens); - for past in (prompt_ids.len()..).take(max_tokens) { - let next = argmax_id(logits)?; - if self.config.eos_token_id.contains(next) { - break; - } - out.push(next); - logits = self.forward(&[next], past, &mut cache, device); - } - debug_assert!(out.len() <= max_tokens); - Ok(out) - } - - /// Parity probe: greedy next-token id + top-k ids for a single prefill. - pub fn first_token( - &self, - prompt_ids: &[u32], - k: usize, - device: &B::Device, - ) -> Result, String> { - assert!(!prompt_ids.is_empty(), "first_token: empty prompt"); - assert!(k >= 1, "first_token: k must be >= 1"); - let mut cache = self.new_cache(); - let logits = self.forward(prompt_ids, 0, &mut cache, device); - // `convert` so an f16 GPU tensor reads back as f32. - let v = logits - .into_data() - .convert::() - .to_vec::() - .map_err(|e| format!("logits readback: {e:?}"))?; - Ok(top_k_ids(&v, k)) - } } #[cfg(test)] diff --git a/crates/mummu/src/registry.rs b/crates/mummu/src/registry.rs new file mode 100644 index 0000000..22bf095 --- /dev/null +++ b/crates/mummu/src/registry.rs @@ -0,0 +1,167 @@ +//! The model registry: declarative [`ModelSpec`]s and a small built-in +//! catalog of known-good models. Adding a model to Mummu is a manifest entry +//! here (or an app-supplied spec), not new code — the spec names the source +//! repo, the architecture that loads it, and the files it needs; `fetch` +//! hands it to the P3 downloader. + +use std::path::{Path, PathBuf}; + +use crate::hub::{self, HubError, Progress}; + +/// Which from-scratch implementation loads this checkpoint. +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub enum Architecture { + /// `models::qwen2` — Qwen2 / Qwen2.5 decoder tiers. + Qwen2, + /// `models::lfm2` — LFM2 / LFM2.5 hybrid conv+attention. + Lfm2, + /// `models::minilm` — all-MiniLM BERT sentence embedder. + MiniLm, +} + +/// A declarative model manifest entry. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct ModelSpec { + /// Short cache-dir-safe name, e.g. `qwen2.5-1.5b-instruct`. + pub name: String, + /// HuggingFace repo id (`owner/name`). + pub repo: String, + /// Git revision (tag, branch, or commit) — pin for reproducibility. + pub revision: String, + pub architecture: Architecture, + /// Rough on-disk size, for settings UIs and fit checks (0 = unknown). + pub disk_bytes_estimate: u64, +} + +impl ModelSpec { + /// Sanity for manifest entries (also the deserialization gate). + pub fn validate(&self) -> Result<(), String> { + if self.name.is_empty() + || !self + .name + .chars() + .all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '_' | '.')) + { + return Err(format!("bad spec name {:?}", self.name)); + } + if !self.repo.contains('/') { + return Err(format!("repo must be owner/name, got {:?}", self.repo)); + } + if self.revision.is_empty() { + return Err("revision must be non-empty (pin something)".into()); + } + Ok(()) + } + + /// The cache directory this model lives in under `models_root`. + #[must_use] + pub fn dir(&self, models_root: &Path) -> PathBuf { + models_root.join(&self.name) + } + + /// Download this model into `models_root` (resumable, cache-first; see + /// [`hub::fetch_model`]) and return its directory, ready for the + /// architecture's `load_from_dir`. + pub fn fetch( + &self, + models_root: &Path, + on_progress: impl FnMut(Progress<'_>), + ) -> Result { + assert!(self.validate().is_ok(), "fetch of an invalid spec"); + hub::fetch_model( + &self.repo, + &self.revision, + &self.dir(models_root), + on_progress, + ) + } +} + +/// The built-in catalog: the models Mummu has ported and parity-verified (or +/// is actively gating — see the ROADMAP P2 checklist for each one's status). +#[must_use] +pub fn catalog() -> Vec { + let entries = vec![ + ModelSpec { + name: "qwen2.5-1.5b-instruct".into(), + repo: "Qwen/Qwen2.5-1.5B-Instruct".into(), + revision: "main".into(), + architecture: Architecture::Qwen2, + disk_bytes_estimate: 3_100_000_000, + }, + ModelSpec { + name: "qwen2.5-0.5b-instruct".into(), + repo: "Qwen/Qwen2.5-0.5B-Instruct".into(), + revision: "main".into(), + architecture: Architecture::Qwen2, + disk_bytes_estimate: 1_000_000_000, + }, + ModelSpec { + name: "lfm2.5-1.2b".into(), + repo: "LiquidAI/LFM2.5-1.2B-Instruct".into(), + revision: "main".into(), + architecture: Architecture::Lfm2, + disk_bytes_estimate: 2_400_000_000, + }, + ModelSpec { + name: "all-minilm-l6-v2".into(), + repo: "sentence-transformers/all-MiniLM-L6-v2".into(), + revision: "main".into(), + architecture: Architecture::MiniLm, + disk_bytes_estimate: 91_000_000, + }, + ]; + debug_assert!( + entries.iter().all(|s| s.validate().is_ok()), + "built-in catalog must validate" + ); + entries +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn builtin_catalog_validates_and_names_are_unique() { + let cat = catalog(); + assert!(cat.len() >= 3); + for spec in &cat { + spec.validate().unwrap_or_else(|e| panic!("{e}")); + } + let mut names: Vec<&str> = cat.iter().map(|s| s.name.as_str()).collect(); + names.sort_unstable(); + names.dedup(); + assert_eq!(names.len(), cat.len(), "duplicate names in the catalog"); + } + + #[test] + fn spec_dir_is_rooted_and_named() { + let spec = &catalog()[0]; + let dir = spec.dir(Path::new("root/models")); + assert_eq!(dir, Path::new("root/models").join(&spec.name)); + } + + #[test] + fn traversal_names_are_rejected() { + let mut spec = catalog()[0].clone(); + spec.name = "../escape".into(); + assert!(spec.validate().is_err(), "path traversal must not validate"); + } + + #[test] + fn bare_repo_is_rejected() { + let mut spec = catalog()[0].clone(); + spec.repo = "qwen".into(); + assert!(spec.validate().is_err()); + } + + #[test] + fn specs_round_trip_through_json() { + let spec = &catalog()[0]; + let json = serde_json::to_string(spec).unwrap(); + let back: ModelSpec = serde_json::from_str(&json).unwrap(); + assert_eq!(back.name, spec.name); + assert_eq!(back.architecture, spec.architecture); + } +} diff --git a/crates/mummu/tests/fixtures/minilm_embedding.json b/crates/mummu/tests/fixtures/minilm_embedding.json new file mode 100644 index 0000000..8adb79a --- /dev/null +++ b/crates/mummu/tests/fixtures/minilm_embedding.json @@ -0,0 +1 @@ +{"embedding":[0.04393353685736656,0.058934397995471954,0.048178404569625854,0.07754811644554138,0.02674442157149315,-0.03762957826256752,-0.0026050335727632046,-0.059943001717329025,-0.002496055094525218,0.022072870284318924,0.04802590236067772,0.055755265057086945,-0.03894545137882233,-0.026616763323545456,0.007693361025303602,-0.026237716898322105,-0.036416083574295044,-0.0378161258995533,0.07407807558774948,-0.049504972994327545,-0.05852169916033745,-0.06361961364746094,0.032435011118650436,0.022008514031767845,-0.07106366008520126,-0.033157795667648315,-0.06941040605306625,-0.05003735423088074,0.07462680339813232,-0.11113379895687103,-0.012306325137615204,0.037745632231235504,-0.028031358495354652,0.014535312540829182,-0.03155853971838951,-0.0805836170911789,0.05835270136594772,0.0025900632608681917,0.03928028419613838,0.025769561529159546,0.049850549548864365,-0.0017561803106218576,-0.04552976414561272,0.02926071733236313,-0.10201725363731384,0.05222867429256439,-0.07908996939659119,-0.010285761207342148,0.009202475659549236,0.013073237612843513,-0.040477730333805084,-0.02779257297515869,0.012466756626963615,0.06728328764438629,0.06812478601932526,-0.0075711957179009914,-0.006099399644881487,-0.042377714067697525,0.05178159102797508,-0.015670713037252426,0.00956356804817915,0.041239071637392044,0.021495992317795753,0.010429366491734982,0.02733498066663742,0.01870620623230934,-0.02696070820093155,-0.07005422562360764,-0.10470044612884521,-0.00189871562179178,0.017701681703329086,-0.057472553104162216,-0.014422387816011906,0.00047053731395862997,0.002332293428480625,-0.02519208937883377,0.04930036514997482,-0.05096093937754631,0.06319832801818848,0.014916544780135155,-0.02707667276263237,-0.0452875941991806,-0.049059413373470306,0.037494100630283356,0.03845794498920441,0.0015689957654103637,0.03099225088953972,0.020163051784038544,-0.012436269782483578,-0.03067200444638729,-0.027881937101483345,-0.06891822069883347,-0.05136770382523537,0.021479584276676178,0.01157471165060997,0.0012540665920823812,0.018876586109399796,-0.04423192888498306,-0.04498177021741867,-0.0034186693374067545,0.013113119639456272,0.020009955391287804,0.12109972536563873,0.023107515648007393,-0.022015923634171486,-0.032884638756513596,-0.003155076177790761,0.00011780133354477584,0.09914983808994293,0.016523899510502815,-0.004696679767221212,-0.014536595903337002,-0.0037107416428625584,0.09651360660791397,0.028590811416506767,0.021348150447010994,-0.0717645213007927,-0.024114245548844337,-0.04409407824277878,-0.10734684020280838,0.06799451261758804,0.1304667741060257,-0.07970301806926727,0.0067951250821352005,-0.023751135915517807,-0.046163663268089294,-0.02996511198580265,-3.6940997306864896e-33,0.07309695333242416,-0.022017216309905052,-0.08614645898342133,-0.07143796980381012,-0.06367415934801102,-0.07218630611896515,-0.005930451676249504,-0.023364149034023285,-0.02836580015718937,0.04774352163076401,-0.08061760663986206,-0.0015648269327357411,0.01384436059743166,-0.028623534366488457,-0.033538658171892166,-0.11377748847007751,-0.009176341816782951,-0.010810140520334244,0.03231958672404289,0.05883808061480522,0.0334208607673645,0.1079879105091095,-0.03727129101753235,-0.029677074402570724,0.05171902850270271,-0.022533860057592392,-0.0696091279387474,-0.021447502076625824,-0.023341042920947075,0.048219989985227585,-0.035876642912626266,-0.04689906910061836,-0.03978738933801651,0.1108132153749466,-0.014300765469670296,-0.11846451461315155,0.05829155072569847,-0.06258892267942429,-0.029404137283563614,0.0603238083422184,-0.0024441825225949287,0.016011586412787437,0.026723351329565048,0.02495306357741356,-0.06493190675973892,-0.010680138133466244,0.02814648114144802,0.010356350801885128,-0.0006635468453168869,0.01981854811310768,-0.030428847298026085,0.006284233182668686,0.05152682214975357,-0.04753745347261429,-0.06444211304187775,0.09550321847200394,0.07558581978082657,-0.028157508000731468,-0.03499659523367882,0.101816326379776,0.019873205572366714,-0.03680365905165672,0.0029351995326578617,-0.0500744991004467,0.15093205869197845,-0.06160794571042061,-0.08588124066591263,0.007139936555176973,-0.013306569308042526,0.07804044336080551,0.017525048926472664,0.04212796315550804,0.03579389303922653,-0.13295046985149384,0.035697005689144135,-0.020311687141656876,0.012491010129451752,-0.038035593926906586,0.04915434867143631,-0.015654122456908226,0.12141832709312439,-0.08086451888084412,-0.046878162771463394,0.04108423367142677,-0.018431849777698517,0.06696905195713043,0.004335977137088776,0.02273155190050602,-0.013642889447510242,-0.04532381147146225,-0.03928292542695999,-0.006298878695815802,0.05296089127659798,-0.03690645471215248,0.07116769254207611,2.3334323261730595e-33,0.10523135960102081,-0.04818737879395485,0.06959184259176254,0.06569761782884598,-0.04651494696736336,0.05144921690225601,-0.012447616085410118,0.032087281346321106,-0.09233565628528595,0.05009322986006737,-0.03288767486810684,0.01391384843736887,-0.000870247429702431,-0.004909032490104437,0.10394642502069473,0.00032163140713237226,0.05281103029847145,-0.011799014173448086,0.02315662056207657,0.013176782988011837,-0.052596259862184525,0.03267020359635353,0.00030862470157444477,0.06411285698413849,0.03885011374950409,0.05880080536007881,0.08297936618328094,-0.0188149306923151,-0.022637803107500076,-0.10047369450330734,-0.03837522119283676,-0.058808114379644394,0.0018241882789880037,-0.04269958287477493,0.025019565597176552,0.06400594860315323,-0.0377483107149601,-0.006839089095592499,-0.0025460575707256794,-0.09760432690382004,0.01884760521352291,-0.0008831329178065062,0.017361151054501534,0.07107905298471451,0.03303936496376991,0.006934205535799265,-0.05605221912264824,0.05146341398358345,-0.042954206466674805,0.04600774124264717,-0.00878827553242445,0.03172885999083519,0.04939659684896469,0.029518958181142807,-0.05051916465163231,-0.05431869998574257,0.0001499919599154964,-0.02766146883368492,0.034687869250774384,-0.02108902484178543,0.013806032948195934,0.029988670721650124,0.013974426314234734,-0.004264745395630598,-0.01503373309969902,-0.08760955184698105,-0.06850544363260269,-0.04281417280435562,0.07769449055194855,-0.07102848589420319,-0.007376902271062136,0.021372729912400246,0.013556227087974548,-0.07904649525880814,0.005476593505591154,0.08306638151407242,0.11414800584316254,0.001807591412216425,0.0875491127371788,-0.04160451143980026,0.015541664324700832,-0.010120641440153122,-0.007324312813580036,0.010796627961099148,-0.06628167629241943,0.03984139859676361,-0.1167115569114685,0.06429930776357651,0.040291979908943176,-0.06547412276268005,0.019505267962813377,0.08099958300590515,0.05364629253745079,0.0767969861626625,-0.013485239818692207,-1.76919012773169e-8,-0.04439351335167885,0.009206416085362434,-0.08795906603336334,0.042692217975854874,0.07313647866249084,0.01684274896979332,-0.040326233953237534,0.018513130024075508,0.08441726863384247,-0.0374477244913578,0.030299603939056396,0.029064133763313293,0.06368786096572876,0.028975049033761024,-0.01472698524594307,0.017754321917891502,-0.03368951380252838,0.017316089943051338,0.033787500113248825,0.17682604491710663,-0.017553281038999557,-0.06030776724219322,-0.014339401386678219,-0.023853661492466927,-0.04455317184329033,-0.028985051438212395,-0.0896776095032692,-0.0017593814991414547,-0.02614855021238327,0.005939929746091366,-0.0518355593085289,0.08572797477245331,-0.0818398967385292,0.00835439283400774,0.04007897898554802,0.04177645966410637,0.10457352548837662,-0.0028656665235757828,0.0196691807359457,0.005810448434203863,0.013325382024049759,0.045100051909685135,-0.02175874263048172,-0.013949304819107056,-0.06869924813508987,-0.0029411043506115675,-0.031076492741703987,-0.10585442185401917,0.06916244328022003,-0.04241151362657547,-0.0467681884765625,-0.036475103348493576,0.045039962977170944,0.060981735587120056,-0.0656561553478241,-0.005456398241221905,-0.018622752279043198,-0.06314841657876968,-0.0387437529861927,0.03467336297035217,0.05554580315947533,0.052162766456604004,0.05610648915171623,0.10206391662359238],"ids":[101,1996,4248,2829,4419,14523,2058,1996,13971,3899,1012,102,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],"reference":"candle-0.9.1 cpu f32","sentence":"The quick brown fox jumps over the lazy dog."} diff --git a/crates/mummu/tests/fixtures/qwen2_5_1_5b_first_logits.json b/crates/mummu/tests/fixtures/qwen2_5_1_5b_first_logits.json new file mode 100644 index 0000000..10ed084 --- /dev/null +++ b/crates/mummu/tests/fixtures/qwen2_5_1_5b_first_logits.json @@ -0,0 +1,54 @@ +{ + "prompt_ids": [ + 151644, + 8948, + 198, + 2610, + 525, + 264, + 10950, + 17847, + 13, + 151645, + 198, + 151644, + 872, + 198, + 852, + 279, + 1156, + 4236, + 10250, + 5109, + 13, + 151645, + 198, + 151644, + 77091, + 198 + ], + "prompt_raw": "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\nList the first five prime numbers.<|im_end|>\n<|im_start|>assistant\n", + "reference": "candle-0.9.1 cpu f32", + "top_k": [ + { + "id": 785, + "logit": 24.636566162109375 + }, + { + "id": 16, + "logit": 22.25823211669922 + }, + { + "id": 32, + "logit": 21.57818031311035 + }, + { + "id": 1249, + "logit": 21.026325225830078 + }, + { + "id": 8420, + "logit": 20.907291412353516 + } + ] +} diff --git a/crates/mummu/tests/parity_lfm2.rs b/crates/mummu/tests/parity_lfm2.rs index d5c0d53..51ea5ca 100644 --- a/crates/mummu/tests/parity_lfm2.rs +++ b/crates/mummu/tests/parity_lfm2.rs @@ -16,6 +16,7 @@ use std::path::PathBuf; use std::process::Command; use mummu::backend::Gpu; +use mummu::models::CausalLm; use mummu::models::lfm2; use tokenizers::Tokenizer; @@ -27,9 +28,10 @@ fn lfm2_dir() -> Option { dir.is_dir().then_some(dir) } -/// The LFM2.5 ChatML wrapping (BOS + user turn + assistant open). +/// The LFM2.5 ChatML wrapping (BOS + user turn + assistant open), rendered by +/// the library's own template (the same shape laurelane validated vs Ollama). fn chatml(user: &str) -> String { - format!("<|startoftext|><|im_start|>user\n{user}<|im_end|>\n<|im_start|>assistant\n") + mummu::chat::ChatMl::lfm2().render(&[mummu::chat::Turn::user(user)]) } /// Greedy reference from Ollama in raw mode (identical prompt text, diff --git a/crates/mummu/tests/parity_qwen2.rs b/crates/mummu/tests/parity_qwen2.rs new file mode 100644 index 0000000..497aec9 --- /dev/null +++ b/crates/mummu/tests/parity_qwen2.rs @@ -0,0 +1,202 @@ +//! Qwen2.5 parity: greedy sequence vs a local Ollama fp16 reference of the +//! SAME weights (`qwen2.5:1.5b-instruct-fp16` — the fp16 GGUF conversion of +//! the Qwen2.5-1.5B-Instruct checkpoint we load from safetensors). Ignored by +//! default; run with +//! +//! ```text +//! MUMMU_QWEN2_DIR=path/to/qwen2.5-1.5b cargo test -p mummu --release --test parity_qwen2 -- --ignored --nocapture +//! ``` +//! +//! This is the greedy leg of the P7 parity gate; the top-k-logits leg runs +//! against the Candle probe (`tools/candle-probe`) which exposes raw logits. + +use std::path::PathBuf; +use std::process::Command; + +use mummu::backend::Gpu; +use mummu::models::CausalLm; +use mummu::models::qwen2; +use tokenizers::Tokenizer; + +const REFERENCE_TAG: &str = "qwen2.5:1.5b-instruct-fp16"; +const PROMPT: &str = "List the first five prime numbers."; +const MAX_TOKENS: usize = 24; + +fn qwen2_dir() -> Option { + let dir = PathBuf::from(std::env::var_os("MUMMU_QWEN2_DIR")?); + dir.is_dir().then_some(dir) +} + +/// The Qwen2.5-Instruct ChatML wrapping (system + user turn + assistant open), +/// rendered by the library's own template — the fixture equality assert below +/// is what byte-verifies `chat::ChatMl::qwen2` against the Candle reference. +fn chatml(user: &str) -> String { + mummu::chat::ChatMl::qwen2().render(&[ + mummu::chat::Turn::system("You are a helpful assistant."), + mummu::chat::Turn::user(user), + ]) +} + +/// Greedy reference from Ollama in raw mode (identical prompt text, +/// temperature 0), via curl so the test needs no HTTP dependency. +fn ollama_greedy(raw_prompt: &str, max_tokens: usize) -> Result { + let body = serde_json::json!({ + "model": REFERENCE_TAG, + "prompt": raw_prompt, + "raw": true, + "stream": false, + "options": { "temperature": 0.0, "num_predict": max_tokens } + }); + let out = Command::new("curl") + .args([ + "-s", + "-X", + "POST", + "http://localhost:11434/api/generate", + "-d", + ]) + .arg(body.to_string()) + .output() + .map_err(|e| format!("curl spawn: {e}"))?; + if !out.status.success() { + return Err(format!("curl exit {:?}", out.status.code())); + } + let v: serde_json::Value = + serde_json::from_slice(&out.stdout).map_err(|e| format!("ollama response parse: {e}"))?; + v.get("response") + .and_then(|r| r.as_str()) + .map(str::to_string) + .ok_or_else(|| format!("no response field in: {v}")) +} + +/// The committed Candle reference (generated by `tools/candle-probe`; see its +/// header for the refresh command). +const LOGITS_FIXTURE: &str = include_str!("fixtures/qwen2_5_1_5b_first_logits.json"); + +/// Max |Δlogit| tolerated between Burn (wgpu f32) and Candle (CPU f32) on the +/// same weights: different backends legitimately reorder float reductions. +/// Measured on the dev GPU (4070 Ti SUPER, Vulkan): 2.7e-5 — this bound gives +/// ~40x headroom; anything past it is a real numeric bug, not noise. +const LOGIT_ABS_TOLERANCE: f32 = 1.0e-3; + +#[test] +#[ignore = "needs local Qwen2.5 weights (MUMMU_QWEN2_DIR)"] +fn qwen2_first_forward_top_k_matches_candle_reference() { + let Some(dir) = qwen2_dir() else { + panic!("set MUMMU_QWEN2_DIR to a dir with config.json/tokenizer.json/model.safetensors"); + }; + let fixture: serde_json::Value = serde_json::from_str(LOGITS_FIXTURE).expect("fixture parses"); + + // The fixture must have been generated from the *same rendered prompt*; + // fail loudly on drift instead of comparing apples to oranges. + let raw = chatml(PROMPT); + assert_eq!( + fixture["prompt_raw"].as_str().expect("prompt_raw"), + raw, + "fixture prompt drifted from the test prompt — regenerate with tools/candle-probe" + ); + + let tok = Tokenizer::from_file(dir.join("tokenizer.json")).expect("tokenizer.json loads"); + let prompt_ids = tok + .encode(raw.as_str(), false) + .expect("encodes") + .get_ids() + .to_vec(); + let fixture_ids: Vec = fixture["prompt_ids"] + .as_array() + .expect("prompt_ids") + .iter() + .map(|v| v.as_u64().expect("id") as u32) + .collect(); + assert_eq!(prompt_ids, fixture_ids, "tokenizations diverge"); + + let device = burn::tensor::Device::::default(); + let loaded = qwen2::load_from_dir::(&dir, &device).expect("weights load checked"); + let mut cache = loaded.new_cache(); + let logits = loaded + .forward(&prompt_ids, 0, &mut cache, &device) + .into_data() + .convert::() + .to_vec::() + .expect("logits readback"); + + let expected: Vec<(usize, f32)> = fixture["top_k"] + .as_array() + .expect("top_k") + .iter() + .map(|e| { + ( + e["id"].as_u64().expect("id") as usize, + e["logit"].as_f64().expect("logit") as f32, + ) + }) + .collect(); + assert!(expected.len() >= 5, "fixture needs at least top-5"); + + let mut indexed: Vec<(usize, f32)> = logits.iter().copied().enumerate().collect(); + indexed.sort_by(|a, b| b.1.total_cmp(&a.1)); + let ours = &indexed[..expected.len()]; + + let our_ids: Vec = ours.iter().map(|&(id, _)| id).collect(); + let ref_ids: Vec = expected.iter().map(|&(id, _)| id).collect(); + let max_abs_diff = ours + .iter() + .zip(expected.iter()) + .map(|(&(_, a), &(_, b))| (a - b).abs()) + .fold(0.0_f32, f32::max); + eprintln!("[parity/qwen2] top-{} ids ours: {our_ids:?}", ours.len()); + eprintln!("[parity/qwen2] top-{} ids ref : {ref_ids:?}", ours.len()); + eprintln!("[parity/qwen2] max |Δlogit| vs candle: {max_abs_diff:e}"); + + assert_eq!( + our_ids, ref_ids, + "top-k ids diverge from the Candle reference" + ); + assert!( + max_abs_diff <= LOGIT_ABS_TOLERANCE, + "logits diverge: max |Δ| = {max_abs_diff} > {LOGIT_ABS_TOLERANCE}" + ); +} + +#[test] +#[ignore = "needs local Qwen2.5 weights (MUMMU_QWEN2_DIR) + a running Ollama with qwen2.5:1.5b-instruct-fp16"] +fn qwen2_greedy_sequence_matches_ollama_reference() { + let Some(dir) = qwen2_dir() else { + panic!("set MUMMU_QWEN2_DIR to a dir with config.json/tokenizer.json/model.safetensors"); + }; + let raw = chatml(PROMPT); + + // Reference first: fail with the transport error (not a weights error) + // when Ollama isn't up. + let reference = match ollama_greedy(&raw, MAX_TOKENS) { + Ok(r) => r, + Err(e) => panic!("Ollama reference unavailable: {e}"), + }; + + let tok = Tokenizer::from_file(dir.join("tokenizer.json")).expect("tokenizer.json loads"); + let prompt_ids = tok + .encode(raw.as_str(), false) + .expect("encodes") + .get_ids() + .to_vec(); + + let device = burn::tensor::Device::::default(); + let loaded = qwen2::load_from_dir::(&dir, &device).expect("weights load checked"); + let ids = loaded + .greedy_generate(&prompt_ids, MAX_TOKENS, &device) + .expect("greedy decode"); + let ours = tok.decode(&ids, true).expect("decode"); + + eprintln!("[parity/qwen2] ours ({} tokens): {ours:?}", ids.len()); + eprintln!("[parity/qwen2] ollama : {reference:?}"); + + // Exact-match the greedy prefix over the shorter of the two (EOS/num_predict + // may truncate one side earlier than the other). + let n = ours.trim().len().min(reference.trim().len()); + assert!(n >= 8, "outputs too short to compare: {n} chars"); + assert_eq!( + &ours.trim()[..n], + &reference.trim()[..n], + "greedy sequences diverge" + ); +} diff --git a/crates/mummu/tests/real_f16.rs b/crates/mummu/tests/real_f16.rs new file mode 100644 index 0000000..79d3e03 --- /dev/null +++ b/crates/mummu/tests/real_f16.rs @@ -0,0 +1,68 @@ +//! f16 on-GPU validation (the P6 precision milestone): load Qwen2.5 on the +//! `GpuF16` backend and prove the three claims — no shader-compile crash, +//! materially lower VRAM than f32, coherent greedy output. Ignored by +//! default; run with +//! +//! ```text +//! MUMMU_QWEN2_DIR=path/to/qwen2.5-1.5b cargo test -p mummu --release --test real_f16 -- --ignored --nocapture +//! ``` + +use std::path::PathBuf; + +use mummu::backend::{GpuF16, inventory}; +use mummu::models::CausalLm; +use mummu::models::qwen2; +use tokenizers::Tokenizer; + +fn qwen2_dir() -> Option { + let dir = PathBuf::from(std::env::var_os("MUMMU_QWEN2_DIR")?); + dir.is_dir().then_some(dir) +} + +#[test] +#[ignore = "needs multi-GB local weights (MUMMU_QWEN2_DIR) + a SHADER_F16 GPU"] +fn qwen2_decodes_coherently_in_f16_on_gpu() { + let Some(dir) = qwen2_dir() else { + panic!("set MUMMU_QWEN2_DIR to a dir with config.json/tokenizer.json/model.safetensors"); + }; + assert!( + inventory().any_shader_f16(), + "no adapter advertises SHADER_F16 — cannot validate f16 here" + ); + + let raw = mummu::chat::ChatMl::qwen2().render(&[ + mummu::chat::Turn::system("You are a concise assistant."), + mummu::chat::Turn::user("What is 2+2? Answer in one short sentence."), + ]); + 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(); + + // Claim 1 (no crash): building the backend + loading casts bf16 -> f16 via + // CastFloatAdapter (load_from_dir targets the backend float dtype). + let device = burn::tensor::Device::::default(); + let loaded = qwen2::load_from_dir::(&dir, &device).expect("f16 weights load checked"); + + // Claim 3 (coherent output): greedy still answers arithmetic. + let start = std::time::Instant::now(); + let ids = loaded + .greedy_generate(&prompt, 32, &device) + .expect("f16 greedy decode"); + let secs = start.elapsed().as_secs_f64(); + assert!(!ids.is_empty(), "f16 decode produced no tokens before EOS"); + let text = tok.decode(&ids, true).expect("decode"); + eprintln!( + "[real_f16] {} tokens in {secs:.2}s (incl. prefill): {text:?} ids={:?}", + ids.len(), + &ids[..ids.len().min(12)] + ); + assert!( + text.contains('4'), + "expected the f16 answer to mention 4, got: {text:?}" + ); + // Claim 2 (VRAM) is measured outside the process (nvidia-smi peak while + // this test runs) — recorded in bench/BASELINE.md. +} diff --git a/crates/mummu/tests/real_hub.rs b/crates/mummu/tests/real_hub.rs new file mode 100644 index 0000000..f6599e1 --- /dev/null +++ b/crates/mummu/tests/real_hub.rs @@ -0,0 +1,135 @@ +//! Real-network Hub download proof: fetch a small model from HuggingFace, +//! checked-load it, and run it. Ignored by default (network + ~90 MB); run with +//! +//! ```text +//! MUMMU_HUB_DEST=some/tmp/dir cargo test -p mummu --release --test real_hub -- --ignored --nocapture +//! ``` + +use std::path::PathBuf; + +use mummu::backend::Cpu; +use mummu::hub; +use mummu::models::minilm; +use tokenizers::Tokenizer; + +/// Small enough to download in seconds, real enough to prove the pipeline. +const REPO: &str = "sentence-transformers/all-MiniLM-L6-v2"; + +/// The MiniLM entry from the built-in catalog (also pins the repo above). +fn minilm_spec() -> mummu::registry::ModelSpec { + mummu::registry::catalog() + .into_iter() + .find(|s| s.repo == REPO) + .expect("MiniLM is in the built-in catalog") +} + +#[test] +#[ignore = "needs network (MUMMU_HUB_DEST names the download dir)"] +fn hub_download_then_load_then_embed() { + let Some(dest) = std::env::var_os("MUMMU_HUB_DEST").map(PathBuf::from) else { + panic!("set MUMMU_HUB_DEST to a scratch dir for the ~90 MB download"); + }; + let mut events = 0u64; + let mut last_total = None; + // Spec-driven: the catalog entry names the repo/revision/dir. + let dir = minilm_spec() + .fetch(&dest, |p| { + events += 1; + last_total = p.total_bytes; + }) + .expect("hub fetch"); + // Either the network streamed (progress fired) or everything was already + // cached from a prior run (zero events) — both are correct; the load + // below is the real proof either way. + eprintln!( + "[real_hub] fetched into {} ({events} progress events, total {last_total:?})", + dir.display() + ); + + for f in ["config.json", "tokenizer.json", "model.safetensors"] { + assert!(dir.join(f).is_file(), "{f} must exist after fetch"); + } + + // Checked load + a real embedding: the downloaded artifacts are usable. + let device = burn::tensor::Device::::default(); + let loaded = minilm::load_from_dir::(&dir, &device).expect("checked load"); + let tok = Tokenizer::from_file(dir.join("tokenizer.json")).expect("tokenizer loads"); + let enc = tok + .encode("Downloads that verify themselves.", true) + .expect("encodes"); + let mask: Vec = enc.get_attention_mask().iter().map(|&m| m as f32).collect(); + let embedding = loaded + .embed_ids(enc.get_ids(), &mask, &device) + .expect("embeds"); + let norm: f32 = embedding.iter().map(|v| v * v).sum::().sqrt(); + eprintln!( + "[real_hub] embedding dims {}, norm {norm:.6}", + embedding.len() + ); + assert_eq!(embedding.len(), 384); + assert!((norm - 1.0).abs() < 1e-3, "L2 norm should be 1, got {norm}"); +} + +#[test] +#[ignore = "needs network (MUMMU_HUB_DEST names the download dir)"] +fn hub_resume_completes_a_partial_download_byte_identical() { + let Some(dest) = std::env::var_os("MUMMU_HUB_DEST").map(PathBuf::from) else { + panic!("set MUMMU_HUB_DEST to a scratch dir"); + }; + let dir = dest.join("resume-proof"); + let url = hub::hub_file_url(REPO, "main", "tokenizer.json"); + let full = dir.join("tokenizer.json"); + + // Reference download (or cache hit from a prior run). + hub::fetch_file(&url, &full, |_| {}).expect("reference fetch"); + let reference = std::fs::read(&full).expect("reference bytes"); + assert!(reference.len() > 200_000, "file too small to prove resume"); + + // Simulate a killed download: keep only the first half as `.part`. + std::fs::remove_file(&full).expect("drop completed file"); + let half = reference.len() / 2; + std::fs::write(dir.join("tokenizer.json.part"), &reference[..half]).expect("seed part"); + + let mut first_event_received = None; + hub::fetch_file(&url, &full, |p| { + first_event_received.get_or_insert(p.received_bytes); + }) + .expect("resumed fetch"); + + let resumed = std::fs::read(&full).expect("resumed bytes"); + assert_eq!( + resumed, reference, + "resumed download must be byte-identical to a full one" + ); + // The first progress event must already include the resumed prefix — + // proof the transfer continued rather than restarted. + let first = first_event_received.expect("progress fired"); + assert!( + first > half as u64, + "first event at {first} bytes should sit past the {half}-byte seed" + ); + eprintln!( + "[real_hub] resume: seeded {half} bytes, first event at {first}, final {} bytes identical", + resumed.len() + ); +} + +/// One-shot helper the nightly uses to pull the CPU-tier Qwen into a cache +/// dir (also a second real proof of the sharded/single-file fetch path on a +/// 1 GB checkpoint). +#[test] +#[ignore = "needs network (MUMMU_HUB_DEST names the download dir); ~1 GB"] +fn hub_fetches_the_cpu_tier_qwen() { + let Some(dest) = std::env::var_os("MUMMU_HUB_DEST").map(PathBuf::from) else { + panic!("set MUMMU_HUB_DEST to a scratch dir for the ~1 GB download"); + }; + let spec = mummu::registry::catalog() + .into_iter() + .find(|s| s.name == "qwen2.5-0.5b-instruct") + .expect("0.5B is in the catalog"); + let dir = spec.fetch(&dest, |_| {}).expect("hub fetch"); + for f in ["config.json", "tokenizer.json", "model.safetensors"] { + assert!(dir.join(f).is_file(), "{f} must exist after fetch"); + } + eprintln!("[real_hub] 0.5B fetched into {}", dir.display()); +} diff --git a/crates/mummu/tests/real_inference.rs b/crates/mummu/tests/real_inference.rs index 93b1ca0..273a4ac 100644 --- a/crates/mummu/tests/real_inference.rs +++ b/crates/mummu/tests/real_inference.rs @@ -6,18 +6,35 @@ //! ``` //! //! where the dir holds `config.json`, `tokenizer.json`, `model.safetensors`. +//! +//! All GPU tests share one [`mummu::cache::ModelSlot`] static — the suite pays +//! the multi-GB load once, and the slot's mutex serializes GPU access, so the +//! whole file stays within a single model's VRAM even with parallel test +//! threads (two concurrent 6 GB loads would blow the 16 GB reference card). use std::path::PathBuf; use mummu::backend::{Cpu, Gpu, use_gpu}; -use mummu::models::qwen2; +use mummu::models::CausalLm; +use mummu::models::qwen2::{self, LoadedQwen2}; use tokenizers::Tokenizer; +/// One model for the whole suite (see the module docs). +static QWEN2_SLOT: mummu::cache::ModelSlot> = mummu::cache::ModelSlot::new(); + fn qwen2_dir() -> Option { let dir = PathBuf::from(std::env::var_os("MUMMU_QWEN2_DIR")?); dir.is_dir().then_some(dir) } +/// Run `f` with the shared GPU model (loading it on first use). +fn with_gpu_model(dir: &std::path::Path, f: impl FnOnce(&LoadedQwen2) -> R) -> R { + let device = burn::tensor::Device::::default(); + QWEN2_SLOT + .with(dir, |d| qwen2::load_from_dir::(d, &device), f) + .expect("weights load checked") +} + /// ChatML prompt for the Qwen2.5 instruct checkpoints. fn chatml(system: &str, user: &str) -> String { format!( @@ -54,11 +71,10 @@ fn qwen2_greedy_decodes_coherent_text_on_default_device() { let (ids, label) = if use_gpu() { let device = burn::tensor::Device::::default(); - let loaded = qwen2::load_from_dir::(&dir, &device).expect("weights load checked"); ( - loaded - .greedy_generate(&prompt, 32, &device) - .expect("decode"), + with_gpu_model(&dir, |m| { + m.greedy_generate(&prompt, 32, &device).expect("decode") + }), "GPU", ) } else { @@ -94,11 +110,99 @@ fn qwen2_first_token_probe_reports_top5() { }; let prompt = encode(&dir, &chatml("You are a helpful assistant.", "Say hello.")); let device = burn::tensor::Device::::default(); - let loaded = qwen2::load_from_dir::(&dir, &device).expect("weights load checked"); - let top5 = loaded.first_token(&prompt, 5, &device).expect("probe"); + let top5 = with_gpu_model(&dir, |m| m.first_token(&prompt, 5, &device).expect("probe")); assert_eq!(top5.len(), 5); eprintln!( "[real_inference] top-5 next-token ids: {top5:?} → {:?}", decode(&dir, &top5[..1]) ); } + +#[test] +#[ignore = "needs multi-GB local weights (MUMMU_QWEN2_DIR)"] +fn qwen2_sampled_streaming_is_seeded_deterministic_and_cancellable() { + let Some(dir) = qwen2_dir() else { + panic!("set MUMMU_QWEN2_DIR to a dir with config.json/tokenizer.json/model.safetensors"); + }; + let prompt = encode( + &dir, + &chatml("You are a poet.", "Write one line about the sea."), + ); + let device = burn::tensor::Device::::default(); + + let opts = mummu::decode::SamplerOptions { + temperature: 0.7, + top_p: 0.9, + seed: 42, + ..mummu::decode::SamplerOptions::default() + }; + + let (cancelled, streamed, replay) = with_gpu_model(&dir, |loaded| { + // Leg 1: streaming + cooperative cancellation after 8 tokens. + let mut streamed = Vec::new(); + let cancelled = loaded + .generate(&prompt, 64, &opts, &device, |id| { + streamed.push(id); + if streamed.len() == 8 { + std::ops::ControlFlow::Break(()) + } else { + std::ops::ControlFlow::Continue(()) + } + }) + .expect("sampled decode"); + + // Leg 2: the same seed replays the same sampled prefix. + let replay = loaded + .generate(&prompt, 8, &opts, &device, |_| { + std::ops::ControlFlow::Continue(()) + }) + .expect("replay decode"); + (cancelled, streamed, replay) + }); + + assert_eq!(cancelled.len(), 8, "cancel after 8 streamed tokens"); + assert_eq!(cancelled, streamed, "returned ids == streamed ids"); + assert_eq!( + replay, cancelled, + "same (prompt, options, seed) must resample the same tokens" + ); + + let text = decode(&dir, &cancelled); + eprintln!("[real_inference/sampled] 8 tokens @ T=0.7 p=0.9 seed=42: {text:?}"); + assert!(!text.trim().is_empty(), "sampled text should be non-empty"); +} + +#[test] +#[ignore = "needs multi-GB local weights (MUMMU_QWEN2_DIR)"] +fn qwen2_model_slot_reuses_the_loaded_model() { + let Some(dir) = qwen2_dir() else { + panic!("set MUMMU_QWEN2_DIR to a dir with config.json/tokenizer.json/model.safetensors"); + }; + let device = burn::tensor::Device::::default(); + let prompt = encode( + &dir, + "<|im_start|>user\nSay hi.<|im_end|>\n<|im_start|>assistant\n", + ); + + // Another test may already have populated the shared slot, so the + // invariant is "no reload between consecutive same-key calls", observed + // via the load closure never firing once the slot is warm. + let mut loads_after_warm = 0; + for round in 0..2 { + let ids = QWEN2_SLOT + .with( + &dir, + |d| { + if round > 0 { + loads_after_warm += 1; + } + qwen2::load_from_dir::(d, &device) + }, + |m| m.greedy_generate(&prompt, 4, &device).expect("decode"), + ) + .expect("slot load"); + assert!(!ids.is_empty(), "cached model must still decode"); + } + assert_eq!(loads_after_warm, 0, "a warm slot must never reload"); + assert_eq!(QWEN2_SLOT.loaded_key().as_deref(), Some(dir.as_path())); +} diff --git a/crates/mummu/tests/real_minilm.rs b/crates/mummu/tests/real_minilm.rs index f394b47..9ffa3b4 100644 --- a/crates/mummu/tests/real_minilm.rs +++ b/crates/mummu/tests/real_minilm.rs @@ -58,3 +58,64 @@ fn minilm_embeds_similar_sentences_closer_than_dissimilar() { let norm: f32 = cat1.iter().map(|v| v * v).sum::().sqrt(); assert!((norm - 1.0).abs() < 1e-3, "L2 norm should be 1, got {norm}"); } + +/// The committed Candle reference embedding (generated by +/// `tools/candle-probe`'s `minilm-probe` bin; see its header for the refresh +/// command). +const EMBEDDING_FIXTURE: &str = include_str!("fixtures/minilm_embedding.json"); + +/// Max |Δcomponent| tolerated between Burn (ndarray f32) and Candle (CPU f32) +/// on the same weights — both CPU f32, so only reduction-order noise remains. +const COMPONENT_ABS_TOLERANCE: f32 = 1.0e-4; + +#[test] +#[ignore = "needs local MiniLM weights (MUMMU_MINILM_DIR)"] +fn minilm_embedding_matches_candle_reference() { + let Some(dir) = minilm_dir() else { + panic!("set MUMMU_MINILM_DIR to a dir with config.json/tokenizer.json/model.safetensors"); + }; + let fixture: serde_json::Value = + serde_json::from_str(EMBEDDING_FIXTURE).expect("fixture parses"); + let sentence = fixture["sentence"].as_str().expect("sentence"); + + let tok = Tokenizer::from_file(dir.join("tokenizer.json")).expect("tokenizer.json loads"); + let enc = tok.encode(sentence, true).expect("encodes"); + let ids = enc.get_ids().to_vec(); + let fixture_ids: Vec = fixture["ids"] + .as_array() + .expect("ids") + .iter() + .map(|v| v.as_u64().expect("id") as u32) + .collect(); + assert_eq!(ids, fixture_ids, "tokenizations diverge"); + + let mask: Vec = enc.get_attention_mask().iter().map(|&m| m as f32).collect(); + let device = burn::tensor::Device::::default(); + let loaded = minilm::load_from_dir::(&dir, &device).expect("weights load checked"); + let ours = loaded.embed_ids(&ids, &mask, &device).expect("embeds"); + + let reference: Vec = fixture["embedding"] + .as_array() + .expect("embedding") + .iter() + .map(|v| v.as_f64().expect("component") as f32) + .collect(); + assert_eq!(ours.len(), reference.len(), "embedding widths diverge"); + + let cos = cosine(&ours, &reference); + let max_abs_diff = ours + .iter() + .zip(reference.iter()) + .map(|(a, b)| (a - b).abs()) + .fold(0.0_f32, f32::max); + eprintln!("[real_minilm] cosine vs candle = {cos:.8}, max |Δ| = {max_abs_diff:e}"); + + assert!( + cos >= 0.999_99, + "embedding diverges from the Candle reference: cosine = {cos}" + ); + assert!( + max_abs_diff <= COMPONENT_ABS_TOLERANCE, + "components diverge: max |Δ| = {max_abs_diff} > {COMPONENT_ABS_TOLERANCE}" + ); +} diff --git a/tools/candle-probe/Cargo.lock b/tools/candle-probe/Cargo.lock new file mode 100644 index 0000000..ca7f3d9 --- /dev/null +++ b/tools/candle-probe/Cargo.lock @@ -0,0 +1,1845 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "getrandom", + "once_cell", + "serde", + "version_check", + "zerocopy", +] + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "arbitrary" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3d036a3c4ab069c7b410a2ce876bd74808d2d0888a82667669f8e783a898bf1" +dependencies = [ + "derive_arbitrary", +] + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "base64" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" + +[[package]] +name = "bit-set" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0700ddab506f33b20a03b13996eccd309a48e5ff77d0d95926aa0210fb4e95f1" +dependencies = [ + "bit-vec", +] + +[[package]] +name = "bit-vec" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "349f9b6a179ed607305526ca489b34ad0a41aed5f7980fa90eb03160b69598fb" + +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" + +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + +[[package]] +name = "candle-core" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9f51e2ecf6efe9737af8f993433c839f956d2b6ed4fd2dd4a7c6d8b0fa667ff" +dependencies = [ + "byteorder", + "gemm 0.17.1", + "half", + "memmap2", + "num-traits", + "num_cpus", + "rand", + "rand_distr", + "rayon", + "safetensors", + "thiserror 1.0.69", + "ug", + "yoke", + "zip", +] + +[[package]] +name = "candle-nn" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c1980d53280c8f9e2c6cbe1785855d7ff8010208b46e21252b978badf13ad69d" +dependencies = [ + "candle-core", + "half", + "num-traits", + "rayon", + "safetensors", + "serde", + "thiserror 1.0.69", +] + +[[package]] +name = "candle-probe" +version = "0.1.0" +dependencies = [ + "candle-core", + "candle-nn", + "candle-transformers", + "serde_json", + "tokenizers", +] + +[[package]] +name = "candle-transformers" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "186cb80045dbe47e0b387ea6d3e906f02fb3056297080d9922984c90e90a72b0" +dependencies = [ + "byteorder", + "candle-core", + "candle-nn", + "fancy-regex", + "num-traits", + "rand", + "rayon", + "serde", + "serde_json", + "serde_plain", + "tracing", +] + +[[package]] +name = "castaway" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a" +dependencies = [ + "rustversion", +] + +[[package]] +name = "cc" +version = "1.2.66" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f5d6cac793997bd970000024b2934968efe83b382de4fdcf4fcb46b6ee4ad996" +dependencies = [ + "find-msvc-tools", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "compact_str" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dfdd1c2274d9aa354115b09dc9a901d6c5576818cdf70d14cae2bdb47df00ab" +dependencies = [ + "castaway", + "cfg-if", + "itoa", + "rustversion", + "ryu", + "serde", + "static_assertions", +] + +[[package]] +name = "console" +version = "0.16.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fe5f465a4f6fee88fad41b85d990f84c835335e85b5d9e6e63e0d06d28cba7c" +dependencies = [ + "encode_unicode", + "libc", + "unicode-width", + "windows-sys", +] + +[[package]] +name = "crc32fast" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "darling" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" +dependencies = [ + "darling_core", + "darling_macro", +] + +[[package]] +name = "darling_core" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn", +] + +[[package]] +name = "darling_macro" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" +dependencies = [ + "darling_core", + "quote", + "syn", +] + +[[package]] +name = "dary_heap" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe" +dependencies = [ + "serde", +] + +[[package]] +name = "derive_arbitrary" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e567bd82dcff979e4b03460c307b3cdc9e96fde3d73bed1496d2bc75d9dd62a" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "derive_builder" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "507dfb09ea8b7fa618fcf76e953f4f5e192547945816d5358edffe39f6f94947" +dependencies = [ + "derive_builder_macro", +] + +[[package]] +name = "derive_builder_core" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8" +dependencies = [ + "darling", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "derive_builder_macro" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c" +dependencies = [ + "derive_builder_core", + "syn", +] + +[[package]] +name = "displaydoc" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "dyn-stack" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56e53799688f5632f364f8fb387488dd05db9fe45db7011be066fc20e7027f8b" +dependencies = [ + "bytemuck", + "reborrow", +] + +[[package]] +name = "dyn-stack" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c4713e43e2886ba72b8271aa66c93d722116acf7a75555cce11dcde84388fe8" +dependencies = [ + "bytemuck", + "dyn-stack-macros", +] + +[[package]] +name = "dyn-stack-macros" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1d926b4d407d372f141f93bb444696142c29d32962ccbd3531117cf3aa0bfa9" + +[[package]] +name = "either" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "encode_unicode" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" + +[[package]] +name = "enum-as-inner" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1e6a265c649f3f5979b601d26f1d05ada116434c87741c9493cb56218f76cbc" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "esaxx-rs" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" +dependencies = [ + "cc", +] + +[[package]] +name = "fancy-regex" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "531e46835a22af56d1e3b66f04844bed63158bc094a628bec1d321d9b4c44bf2" +dependencies = [ + "bit-set", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + +[[package]] +name = "futures-core" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" + +[[package]] +name = "futures-task" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" + +[[package]] +name = "futures-util" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" +dependencies = [ + "futures-core", + "futures-task", + "pin-project-lite", + "slab", +] + +[[package]] +name = "gemm" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ab24cc62135b40090e31a76a9b2766a501979f3070fa27f689c27ec04377d32" +dependencies = [ + "dyn-stack 0.10.0", + "gemm-c32 0.17.1", + "gemm-c64 0.17.1", + "gemm-common 0.17.1", + "gemm-f16 0.17.1", + "gemm-f32 0.17.1", + "gemm-f64 0.17.1", + "num-complex", + "num-traits", + "paste", + "raw-cpuid 10.7.0", + "seq-macro", +] + +[[package]] +name = "gemm" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab96b703d31950f1aeddded248bc95543c9efc7ac9c4a21fda8703a83ee35451" +dependencies = [ + "dyn-stack 0.13.2", + "gemm-c32 0.18.2", + "gemm-c64 0.18.2", + "gemm-common 0.18.2", + "gemm-f16 0.18.2", + "gemm-f32 0.18.2", + "gemm-f64 0.18.2", + "num-complex", + "num-traits", + "paste", + "raw-cpuid 11.6.0", + "seq-macro", +] + +[[package]] +name = "gemm-c32" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9c030d0b983d1e34a546b86e08f600c11696fde16199f971cd46c12e67512c0" +dependencies = [ + "dyn-stack 0.10.0", + "gemm-common 0.17.1", + "num-complex", + "num-traits", + "paste", + "raw-cpuid 10.7.0", + "seq-macro", +] + +[[package]] +name = "gemm-c32" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6db9fd9f40421d00eea9dd0770045a5603b8d684654816637732463f4073847" +dependencies = [ + "dyn-stack 0.13.2", + "gemm-common 0.18.2", + "num-complex", + "num-traits", + "paste", + "raw-cpuid 11.6.0", + "seq-macro", +] + +[[package]] +name = "gemm-c64" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fbb5f2e79fefb9693d18e1066a557b4546cd334b226beadc68b11a8f9431852a" +dependencies = [ + "dyn-stack 0.10.0", + "gemm-common 0.17.1", + "num-complex", + "num-traits", + "paste", + "raw-cpuid 10.7.0", + "seq-macro", +] + +[[package]] +name = "gemm-c64" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfcad8a3d35a43758330b635d02edad980c1e143dc2f21e6fd25f9e4eada8edf" +dependencies = [ + "dyn-stack 0.13.2", + "gemm-common 0.18.2", + "num-complex", + "num-traits", + "paste", + "raw-cpuid 11.6.0", + "seq-macro", +] + +[[package]] +name = "gemm-common" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2e7ea062c987abcd8db95db917b4ffb4ecdfd0668471d8dc54734fdff2354e8" +dependencies = [ + "bytemuck", + "dyn-stack 0.10.0", + "half", + "num-complex", + "num-traits", + "once_cell", + "paste", + "pulp 0.18.22", + "raw-cpuid 10.7.0", + "rayon", + "seq-macro", + "sysctl 0.5.5", +] + +[[package]] +name = "gemm-common" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a352d4a69cbe938b9e2a9cb7a3a63b7e72f9349174a2752a558a8a563510d0f3" +dependencies = [ + "bytemuck", + "dyn-stack 0.13.2", + "half", + "libm", + "num-complex", + "num-traits", + "once_cell", + "paste", + "pulp 0.21.5", + "raw-cpuid 11.6.0", + "rayon", + "seq-macro", + "sysctl 0.6.0", +] + +[[package]] +name = "gemm-f16" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ca4c06b9b11952071d317604acb332e924e817bd891bec8dfb494168c7cedd4" +dependencies = [ + "dyn-stack 0.10.0", + "gemm-common 0.17.1", + "gemm-f32 0.17.1", + "half", + "num-complex", + "num-traits", + "paste", + "raw-cpuid 10.7.0", + "rayon", + "seq-macro", +] + +[[package]] +name = "gemm-f16" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cff95ae3259432f3c3410eaa919033cd03791d81cebd18018393dc147952e109" +dependencies = [ + "dyn-stack 0.13.2", + "gemm-common 0.18.2", + "gemm-f32 0.18.2", + "half", + "num-complex", + "num-traits", + "paste", + "raw-cpuid 11.6.0", + "rayon", + "seq-macro", +] + +[[package]] +name = "gemm-f32" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9a69f51aaefbd9cf12d18faf273d3e982d9d711f60775645ed5c8047b4ae113" +dependencies = [ + "dyn-stack 0.10.0", + "gemm-common 0.17.1", + "num-complex", + "num-traits", + "paste", + "raw-cpuid 10.7.0", + "seq-macro", +] + +[[package]] +name = "gemm-f32" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc8d3d4385393304f407392f754cd2dc4b315d05063f62cf09f47b58de276864" +dependencies = [ + "dyn-stack 0.13.2", + "gemm-common 0.18.2", + "num-complex", + "num-traits", + "paste", + "raw-cpuid 11.6.0", + "seq-macro", +] + +[[package]] +name = "gemm-f64" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa397a48544fadf0b81ec8741e5c0fba0043008113f71f2034def1935645d2b0" +dependencies = [ + "dyn-stack 0.10.0", + "gemm-common 0.17.1", + "num-complex", + "num-traits", + "paste", + "raw-cpuid 10.7.0", + "seq-macro", +] + +[[package]] +name = "gemm-f64" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35b2a4f76ce4b8b16eadc11ccf2e083252d8237c1b589558a49b0183545015bd" +dependencies = [ + "dyn-stack 0.13.2", + "gemm-common 0.18.2", + "num-complex", + "num-traits", + "paste", + "raw-cpuid 11.6.0", + "seq-macro", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "wasip2", +] + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "bytemuck", + "cfg-if", + "crunchy", + "num-traits", + "rand", + "rand_distr", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "hermit-abi" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" + +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown", +] + +[[package]] +name = "indicatif" +version = "0.18.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9433806cd6b4ec1aba79c021c7e4c58fb4c3b9977c085062e611ac929998fb0c" +dependencies = [ + "console", + "portable-atomic", + "unicode-width", + "unit-prefix", + "web-time", +] + +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "js-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102" +dependencies = [ + "cfg-if", + "futures-util", + "wasm-bindgen", +] + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libm" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[package]] +name = "macro_rules_attribute" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65049d7923698040cd0b1ddcced9b0eb14dd22c5f86ae59c3740eab64a676520" +dependencies = [ + "macro_rules_attribute-proc_macro", + "paste", +] + +[[package]] +name = "macro_rules_attribute-proc_macro" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "670fdfda89751bc4a84ac13eaa63e205cf0fd22b4c9a5fbfa085b63c1f1d3a30" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "memmap2" +version = "0.9.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1219ed1b7f229ee7104d281dd01d6802fe28bb6e95d292942c4daacdeb798c0" +dependencies = [ + "libc", + "stable_deref_trait", +] + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "monostate" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67" +dependencies = [ + "monostate-impl", + "serde", + "serde_core", +] + +[[package]] +name = "monostate-impl" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "num" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35bd024e8b2ff75562e5f34e7f4905839deb4b22955ef5e73d2fea1b9813cb23" +dependencies = [ + "num-bigint", + "num-complex", + "num-integer", + "num-iter", + "num-rational", + "num-traits", +] + +[[package]] +name = "num-bigint" +version = "0.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c89e69e7e0f03bea5ef08013795c25018e101932225a656383bd384495ecc367" +dependencies = [ + "num-integer", + "num-traits", +] + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "bytemuck", + "num-traits", +] + +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-iter" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c92800bd69a1eac91786bcfe9da64a897eb72911b8dc3095decbd07429e8048b" +dependencies = [ + "num-integer", + "num-traits", +] + +[[package]] +name = "num-rational" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" +dependencies = [ + "num-bigint", + "num-integer", + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", + "libm", +] + +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + +[[package]] +name = "num_enum" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d0bca838442ec211fa11de3a8b0e0e8f3a4522575b5c4c06ed722e005036f26" +dependencies = [ + "num_enum_derive", + "rustversion", +] + +[[package]] +name = "num_enum_derive" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "680998035259dcfcafe653688bf2aa6d3e2dc05e98be6ab46afb089dc84f1df8" +dependencies = [ + "proc-macro-crate", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "onig" +version = "6.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc3cbf698f9438986c11a880c90a6d04b9de27575afd28bbf45b154b6c709e2" +dependencies = [ + "bitflags 2.13.0", + "libc", + "once_cell", + "onig_sys", +] + +[[package]] +name = "onig_sys" +version = "69.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e68317604e77e53b85896388e1a803c1d21b74c899ec9e5e1112db90735edd7" +dependencies = [ + "cc", + "pkg-config", +] + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + +[[package]] +name = "portable-atomic" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + +[[package]] +name = "proc-macro-crate" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f" +dependencies = [ + "toml_edit", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "pulp" +version = "0.18.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a0a01a0dc67cf4558d279f0c25b0962bd08fc6dec0137699eae304103e882fe6" +dependencies = [ + "bytemuck", + "libm", + "num-complex", + "reborrow", +] + +[[package]] +name = "pulp" +version = "0.21.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96b86df24f0a7ddd5e4b95c94fc9ed8a98f1ca94d3b01bdce2824097e7835907" +dependencies = [ + "bytemuck", + "cfg-if", + "libm", + "num-complex", + "reborrow", + "version_check", +] + +[[package]] +name = "quote" +version = "1.0.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "rand" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" +dependencies = [ + "rand_chacha", + "rand_core", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom", +] + +[[package]] +name = "rand_distr" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463" +dependencies = [ + "num-traits", + "rand", +] + +[[package]] +name = "raw-cpuid" +version = "10.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c297679cb867470fa8c9f67dbba74a78d78e3e98d7cf2b08d6d71540f797332" +dependencies = [ + "bitflags 1.3.2", +] + +[[package]] +name = "raw-cpuid" +version = "11.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "498cd0dc59d73224351ee52a95fee0f1a617a2eae0e7d9d720cc622c73a54186" +dependencies = [ + "bitflags 2.13.0", +] + +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-cond" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f" +dependencies = [ + "either", + "itertools", + "rayon", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "reborrow" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03251193000f4bd3b042892be858ee50e8b3719f2b08e5833ac4353724632430" + +[[package]] +name = "regex" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a0e75113e14dc5acb068cd0786884f214f1312650a3d36d269f5c4f3cdee8a2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f388202e4b80542a0921078cc23b6333bcf1409c1e3f86404cae4766a6131db" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + +[[package]] +name = "safetensors" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44560c11236a6130a46ce36c836a62936dc81ebf8c36a37947423571be0e55b6" +dependencies = [ + "serde", + "serde_json", +] + +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + +[[package]] +name = "seq-macro" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.150" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "serde_plain" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ce1fc6db65a611022b23a0dec6975d63fb80a302cb3388835ff02c097258d50" +dependencies = [ + "serde", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + +[[package]] +name = "spm_precompiled" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326" +dependencies = [ + "base64", + "nom", + "serde", + "unicode-segmentation", +] + +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + +[[package]] +name = "syn" +version = "2.0.118" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "sysctl" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec7dddc5f0fee506baf8b9fdb989e242f17e4b11c61dfbb0635b705217199eea" +dependencies = [ + "bitflags 2.13.0", + "byteorder", + "enum-as-inner", + "libc", + "thiserror 1.0.69", + "walkdir", +] + +[[package]] +name = "sysctl" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01198a2debb237c62b6826ec7081082d951f46dbb64b0e8c7649a452230d1dfc" +dependencies = [ + "bitflags 2.13.0", + "byteorder", + "enum-as-inner", + "libc", + "thiserror 1.0.69", + "walkdir", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl 2.0.18", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tokenizers" +version = "0.22.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b238e22d44a15349529690fb07bd645cf58149a1b1e44d6cb5bd1641ff1a6223" +dependencies = [ + "ahash", + "aho-corasick", + "compact_str", + "dary_heap", + "derive_builder", + "esaxx-rs", + "getrandom", + "indicatif", + "itertools", + "log", + "macro_rules_attribute", + "monostate", + "onig", + "paste", + "rand", + "rayon", + "rayon-cond", + "regex", + "regex-syntax", + "serde", + "serde_json", + "spm_precompiled", + "thiserror 2.0.18", + "unicode-normalization-alignments", + "unicode-segmentation", + "unicode_categories", +] + +[[package]] +name = "toml_datetime" +version = "1.1.1+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_edit" +version = "0.25.12+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2153edc6955a6c354fad8f5efd38b6a8769bdccf9fe50f8e1329f81b0baa5d7" +dependencies = [ + "indexmap", + "toml_datetime", + "toml_parser", + "winnow", +] + +[[package]] +name = "toml_parser" +version = "1.1.2+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2abe9b86193656635d2411dc43050282ca48aa31c2451210f4202550afb7526" +dependencies = [ + "winnow", +] + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", +] + +[[package]] +name = "ug" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90b70b37e9074642bc5f60bb23247fd072a84314ca9e71cdf8527593406a0dd3" +dependencies = [ + "gemm 0.18.2", + "half", + "libloading", + "memmap2", + "num", + "num-traits", + "num_cpus", + "rayon", + "safetensors", + "serde", + "thiserror 1.0.69", + "tracing", + "yoke", +] + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unicode-normalization-alignments" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de" +dependencies = [ + "smallvec", +] + +[[package]] +name = "unicode-segmentation" +version = "1.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" + +[[package]] +name = "unicode-width" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" + +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + +[[package]] +name = "unit-prefix" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81e544489bf3d8ef66c953931f56617f423cd4b5494be343d9b9d3dda037b9a3" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "wasm-bindgen" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "winnow" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0592e1c9d151f854e6fd382574c3a0855250e1d9b2f99d9281c6e6391af352f1" +dependencies = [ + "memchr", +] + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "yoke" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "120e6aef9aa629e3d4f52dc8cc43a015c7724194c97dfaf45180d2daf2b77f40" +dependencies = [ + "serde", + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2380878cad4ac9aac1e2435f3eb4020e8374b5f13c296cb75b4620ff8e229154" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerocopy" +version = "0.8.54" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7cbbc0a705a0fd05cc3676525980d2bf5a9bc4adac6d6475209a7887cf59d19" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.54" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2e817b7b52d0c7358d3246da9d69935ebb18116b2b102b4230dac079b4862f5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zerofrom" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zip" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9cc23c04387f4da0374be4533ad1208cbb091d5c11d070dfef13676ad6497164" +dependencies = [ + "arbitrary", + "crc32fast", + "crossbeam-utils", + "displaydoc", + "indexmap", + "num_enum", + "thiserror 1.0.69", +] + +[[package]] +name = "zmij" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" diff --git a/tools/candle-probe/Cargo.toml b/tools/candle-probe/Cargo.toml new file mode 100644 index 0000000..acb1c40 --- /dev/null +++ b/tools/candle-probe/Cargo.toml @@ -0,0 +1,26 @@ +# A standalone same-weights reference for the P7 parity gate: load a Qwen2 +# safetensors checkpoint with Candle (CPU, f32), run one forward over a fixed +# ChatML prompt, and print the top-k (id, logit) pairs as JSON. The output is +# committed as a fixture under crates/mummu/tests/fixtures/ so the mummu-side +# logits test runs without Candle installed. Deliberately OUTSIDE the cargo +# workspace: Candle must never become a mummu dependency (Burn is the one +# runtime — see the ROADMAP North Star). +[package] +name = "candle-probe" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] +# =0.9.1: the same Candle laurelane validated byte-identical parity against. +candle-core = "=0.9.1" +candle-nn = "=0.9.1" +candle-transformers = "=0.9.1" +tokenizers = "0.22" +serde_json = "1" + +[workspace] + +[[bin]] +name = "minilm-probe" +path = "src/minilm.rs" diff --git a/tools/candle-probe/src/main.rs b/tools/candle-probe/src/main.rs new file mode 100644 index 0000000..85f281c --- /dev/null +++ b/tools/candle-probe/src/main.rs @@ -0,0 +1,65 @@ +//! Same-weights Candle reference for the P7 parity gate. +//! +//! Usage: `candle-probe [k]` — loads `config.json` / +//! `tokenizer.json` / `model.safetensors` from `` (a Qwen2-family +//! checkpoint), runs ONE forward over the fixed ChatML prompt below on CPU in +//! f32, and prints a JSON object with the prompt, the top-k ids, and their +//! logits. Redirect the output into +//! `crates/mummu/tests/fixtures/_first_logits.json` to refresh the +//! committed fixture that `tests/parity_qwen2.rs` compares Burn against. + +use candle_core::{DType, Device, Tensor}; +use candle_nn::VarBuilder; +use candle_transformers::models::qwen2::{Config, ModelForCausalLM}; + +/// Must stay identical to `PROMPT`/`chatml` in `crates/mummu/tests/parity_qwen2.rs`; +/// the fixture carries the rendered prompt so the test can verify it drifted nowhere. +const PROMPT: &str = "List the first five prime numbers."; + +fn chatml(user: &str) -> String { + format!( + "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n{user}<|im_end|>\n<|im_start|>assistant\n" + ) +} + +fn main() -> Result<(), Box> { + let mut args = std::env::args().skip(1); + let dir = std::path::PathBuf::from(args.next().ok_or("usage: candle-probe [k]")?); + let k: usize = args.next().map(|s| s.parse()).transpose()?.unwrap_or(5); + assert!(k >= 1 && k <= 64, "k out of range: {k}"); + assert!(dir.is_dir(), "not a directory: {}", dir.display()); + + let device = Device::Cpu; + let config: Config = serde_json::from_slice(&std::fs::read(dir.join("config.json"))?)?; + let vb = unsafe { + VarBuilder::from_mmaped_safetensors(&[dir.join("model.safetensors")], DType::F32, &device)? + }; + let mut model = ModelForCausalLM::new(&config, vb)?; + + let raw = chatml(PROMPT); + let tokenizer = tokenizers::Tokenizer::from_file(dir.join("tokenizer.json"))?; + let ids = tokenizer.encode(raw.as_str(), false)?.get_ids().to_vec(); + assert!(!ids.is_empty(), "prompt tokenized to nothing"); + + let input = Tensor::new(ids.as_slice(), &device)?.unsqueeze(0)?; + // seqlen_offset 0: full prefill; the model returns last-position logits + // (batch and seq dims of 1) — flatten to the bare vocab vector. + let logits = model.forward(&input, 0)?.flatten_all()?.to_vec1::()?; + assert!(logits.len() > k, "vocab smaller than k"); + + let mut indexed: Vec<(usize, f32)> = logits.iter().copied().enumerate().collect(); + indexed.sort_by(|a, b| b.1.total_cmp(&a.1)); + let top: Vec = indexed[..k] + .iter() + .map(|&(id, logit)| serde_json::json!({ "id": id, "logit": logit })) + .collect(); + + let out = serde_json::json!({ + "reference": "candle-0.9.1 cpu f32", + "prompt_raw": raw, + "prompt_ids": ids, + "top_k": top, + }); + println!("{}", serde_json::to_string_pretty(&out)?); + Ok(()) +} diff --git a/tools/candle-probe/src/minilm.rs b/tools/candle-probe/src/minilm.rs new file mode 100644 index 0000000..a43f142 --- /dev/null +++ b/tools/candle-probe/src/minilm.rs @@ -0,0 +1,61 @@ +//! Same-weights Candle reference for the MiniLM leg of the P7 parity gate. +//! +//! Usage: `minilm-probe ` — loads the all-MiniLM checkpoint from +//! `` (`config.json` / `tokenizer.json` / `model.safetensors`), +//! embeds the fixed sentence below on CPU in f32 (masked-mean-pool + +//! L2-normalize, exactly Mummu's pipeline), and prints the sentence, token +//! ids, and the full embedding vector as JSON. Redirect into +//! `crates/mummu/tests/fixtures/minilm_embedding.json` to refresh the +//! committed fixture that `tests/real_minilm.rs` compares Burn against. + +use candle_core::{DType, Device, Tensor}; +use candle_nn::VarBuilder; +use candle_transformers::models::bert::{BertModel, Config}; + +/// Must stay identical to the sentence in the parity test in +/// `crates/mummu/tests/real_minilm.rs`; the fixture carries it for a drift check. +const SENTENCE: &str = "The quick brown fox jumps over the lazy dog."; + +fn main() -> Result<(), Box> { + let mut args = std::env::args().skip(1); + let dir = std::path::PathBuf::from(args.next().ok_or("usage: minilm-probe ")?); + assert!(dir.is_dir(), "not a directory: {}", dir.display()); + + let device = Device::Cpu; + let config: Config = serde_json::from_slice(&std::fs::read(dir.join("config.json"))?)?; + let vb = unsafe { + VarBuilder::from_mmaped_safetensors(&[dir.join("model.safetensors")], DType::F32, &device)? + }; + let model = BertModel::load(vb, &config)?; + + let tokenizer = tokenizers::Tokenizer::from_file(dir.join("tokenizer.json"))?; + let enc = tokenizer.encode(SENTENCE, true)?; + let ids = enc.get_ids().to_vec(); + let mask: Vec = enc.get_attention_mask().to_vec(); + assert!(!ids.is_empty(), "sentence tokenized to nothing"); + assert_eq!(ids.len(), mask.len(), "ids/mask length mismatch"); + + let t = ids.len(); + let input = Tensor::new(ids.as_slice(), &device)?.unsqueeze(0)?; + let type_ids = input.zeros_like()?; + let attention = Tensor::new(mask.as_slice(), &device)?.unsqueeze(0)?; + let hidden = model.forward(&input, &type_ids, Some(&attention))?; // [1, t, h] + + // Masked mean pool + L2 normalize — the sentence-transformers recipe and + // exactly what `LoadedMiniLm::embed_ids` does. + let mask_f = attention.to_dtype(DType::F32)?.reshape((1, t, 1))?; + let summed = hidden.broadcast_mul(&mask_f)?.sum(1)?; // [1, h] + let count = mask_f.sum(1)?; // [1, 1] + let mean = summed.broadcast_div(&count)?; + let norm = mean.sqr()?.sum_keepdim(1)?.sqrt()?; + let embedding = mean.broadcast_div(&norm)?.flatten_all()?.to_vec1::()?; + + let out = serde_json::json!({ + "reference": "candle-0.9.1 cpu f32", + "sentence": SENTENCE, + "ids": ids, + "embedding": embedding, + }); + println!("{}", serde_json::to_string(&out)?); + Ok(()) +}