diff --git a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/BENCHMARKS.md b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/BENCHMARKS.md index 8efb073e9..577566336 100644 --- a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/BENCHMARKS.md +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/BENCHMARKS.md @@ -1439,6 +1439,144 @@ uses `Digits` (isolate each digit) + `ByteLevel instead of explicit `Split` entries. Both now map onto the chained-split machinery from wave 60. +## North Mini Code on the CPU: decode 3.3–5.2× and prefill 5.9–10.6× against base at its default thread count, bit-identical (2026-11) + +The pure-Rust CPU path (`--features rayon`) of North Mini Code 1.0 +(Cohere2-MoE, 30B total / 3B active, BF16 safetensors), made +faster without changing a bit of any logit. Measured with +`examples/north_bench.rs` (prints every step's token id and a hash +of all 262144 logits' bits), base = `main` before this change. The +ratios in the title are median against median in the table under +**Speed**, with base at rayon's default thread count (every logical +CPU) and the new code at its own default; that section also gives +base at 12 threads and says what was not compared. + +**Where the time went** (`JOULECLAW_PROF=1`, base, M3 Ultra under +load, ~600-token prompt): prefill 45.3 s = attention 36% (every +head on one thread, a scalar add chain per score), experts 50% +(one expert after another), q/k/v + o 11%; decode 212 ms/token = +attention 26%, q/k/v 20%, o_proj 20%, experts 28%, lm_head 3%. +lm_head streams its 1 GB in ~7 ms (~140 GB/s) while each layer's +16–50 MB projections moved 12–40 GB/s: decode was bound by its ~7 +dispatches per layer (three of them single-threaded GEMVs) and a +GEMV kernel held to ~16–20 GB/s per core by its two accumulator +chains, not by memory bandwidth. Base decode under load was also +fastest at 8 threads (94 ms) and slowest at the default 32 +(174 ms). + +**What changed** (each output keeps its exact reduction order): + +- the BF16 kernel computes blocks of 4 activation rows × 2 weight + rows, 2 × 4 or 1 × 4 (up to 16 accumulator chains instead of 2) + with a 1 KiB prefetch; +- decode runs three dispatches per layer (q/k/v + router; + attention heads + expert gate/up; o_proj + expert down) of + ~256 KiB row chunks on a pool whose threads stay awake between + dispatches (`spin_pool`); +- attention runs per head (decode) and per (head, 16 query rows) + (prefill) in parallel, with a bit-exact NEON score kernel; +- prefill runs the experts in parallel and the projections as + 64 × 64 output tiles; RoPE angles are computed once per position; +- `forward_extend` batches tokens into a state that already holds + a context, bit-identical to feeding them one decode step at a + time; +- threads: one per performance core, at most 16 + (`JOULECLAW_THREADS` overrides). + +**Identity.** Five prompts (a raw completion, Python, OpenQASM, a +596-token and a 2698-token Rust prompt), 72 greedy tokens each: +the same token ids and the same hash of every step's logits as +base, on an M3 Ultra and on an M4 Max, and the same on both +machines; every repeat of the timing runs below gave that one hash +sequence too. Unit tests pin the kernels and both forward passes +against the old code bit for bit. + +**What is re-checked without the checkpoint.** +`src/cohere2_moe_reference.txt` holds logit hashes of three synthetic +Cohere2-MoE checkpoints, generated by base's `forward_prompt` and +`forward_token` (with its one-row BF16 kernel) on macOS arm64, Linux +arm64 and Linux x86_64. The three platforms give different bits, so +each has its own lines. One checkpoint is wide enough that decode +splits its projections into several row chunks and prefill runs the +2-D tiles. The test `stored_reference::logits_match_the_stored_reference` +recomputes the hashes with the reference passes, the CPU passes and +`forward_extend`, at the default thread count and at 1 and 3 threads. +CI runs it in the Linux arm64 workspace job. That job's container has +glibc 2.35, and the Linux arm64 lines were made with 2.39: with 2.35's +libm loaded in its place on the host that made them, the test gives +the stored lines too, and CI's own first run of it was pending when +this was written. The macOS lines are checked by running the tests on +a Mac. `examples/cohere2_reference.rs` generates the lines, at commit +30c4cf1ee, for a platform not yet in the file. **What still needs the checkpoint:** identity at North's +own sizes (hidden 2048, 128 experts with 8 routed, a 262144-row +lm_head, contexts past the 4096-token window) and on the five prompts +above, whose text and hash sequences are not stored, and every speed +figure here. The per-run logs are not in the repository. + +**Speed.** Medians of 4 runs (the 72-token identity run and three +48-token timing runs; base at 12 threads had only the three), base +and new interleaved in the same conditions, each at its default +thread count unless the column says otherwise: base at rayon's +default, every logical CPU (32 on the M3 Ultra, 16 on the M4 Max), +the new code at one thread per performance core (16, its cap, on the +M3 Ultra; 12 on the M4 Max). M3 Ultra (24P+8E): shared, load average +30–130, with another North benchmark and CI running. M4 Max +(12P+4E): lightly loaded (load average 7–29). + +| | M3 Ultra base, 32 thr | M3 Ultra base, 12 thr | M3 Ultra new, 16 thr | M4 Max base, 16 thr | M4 Max base, 12 thr | M4 Max new, 12 thr | +|---|---|---|---|---|---|---| +| decode at ~600 context | 267 ms | 168 ms | **54 ms** | 116 ms | 106 ms | **35 ms** | +| decode at ~2700 context | 405 ms | 364 ms | **80 ms** | 246 ms | 224 ms | **48 ms** | +| prefill, 596 tokens | 9.9 tok/s | 12.0 tok/s | **91 tok/s** | 17.8 tok/s | 19.8 tok/s | **105 tok/s** | +| prefill, 2698 tokens | 5.9 tok/s | 4.5 tok/s | **63 tok/s** | 9.0 tok/s | 9.2 tok/s | **79 tok/s** | + +Paired run-for-run ratios against base at its default thread count +(base / new, all four pairs of each cell). M3 Ultra decode 4.3–4.8× +at ~600 and 5.0–6.6× at ~2700, apart from one pair in each column, +20.2× and 21.2×, where base's run took 3.6× and 5.8× its median +step (967 ms at load 71, 2358 ms at load 131); those two runs are +inside the table's medians. +Prefill 5.4–10.3× and 9.7–11.8×. M4 Max decode 3.3–3.5× at ~600 +apart from the identity run's pair, 2.0× (the new code's median step +was 59.7 ms in that run, 34–37 ms in the others), and 5.1–6.3× at +~2700; prefill 5.2–6.8× and 7.0–9.0×. + +Against base at 12 threads (the three timing pairs): M3 Ultra decode +2.9–3.8× and 3.3–5.8×, prefill 5.9–10.6× and 11.9–14.1×; M4 Max +decode 3.1× and 4.5–5.2×, prefill 5.8–6.1× and 8.4–9.5×. Base at 8 +threads, its fastest in the thread sweep above (94 ms), was not run +in pairs with the new code: against this table's 54 ms that is about +1.7×, but measured under other load, so it is not a ratio this +benchmark establishes. + +A 69-token follow-up turn fed into a ~630-token context +(`north_bench --then` and `--then-seq`): `forward_extend` 1.23 s +(M3 Ultra) / 0.79–0.80 s (M4 Max), against 3.3–4.1 s / 2.4 s one +`forward_token` at a time on the new code, and an estimated ~18 s at +base's decode speed (69 × its 267 ms median step; not measured); +same logits. A caller gets the first figure only by +feeding its turn through `forward_extend`; one that feeds +`forward_token` per token gets the second. + +**Did not help, or not enough:** more threads (base and new were +both slower above 16 under load); for decode, running each step +inside a rayon pool rather than handing each of its loops to the +pool (no clear win once measured in-process; decode's loops now run +on the spin pool, while prefill, whose projections and experts are +rayon loops, does run inside its pool, `in_cpu_pool`); a rayon pool +of any size against the spin pool for decode (2.5–4.4× slower under +load: 164–465 ms against 65–106 on an M3 Ultra; on a quiet M4 Max +40–42 ms against 34 with the same 12 threads, 18–24% slower); +prefetch distances other than 1 KiB (256 B–4 KiB swept). + +**Next levers:** the sliding layers trim their KV cache by moving +the whole window each token past 4096 positions (a ring buffer +removes ~0.6 GB of copies per token there); decode now moves ~6.5 +GB of weights per token near what the CPU side reads from memory, +so the next step is fewer bytes (narrower weights, which would +change logits) or several sequences sharing each weight read; and +prefill experts are bound by the 4 × 2 NEON block's arithmetic. + ## Cohere2-MoE / North Mini Code (wave 60) Fourth topology, first attention-only one: diff --git a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/examples/cohere2_reference.rs b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/examples/cohere2_reference.rs new file mode 100644 index 000000000..9a5dcd56e --- /dev/null +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/examples/cohere2_reference.rs @@ -0,0 +1,280 @@ +//! Generator of `src/cohere2_moe_reference.txt`, the stored reference +//! the test `cohere2_moe::tests::stored_reference` checks. +//! +//! It builds the test's three synthetic Cohere2-MoE checkpoints (the +//! same weights, byte for byte), runs the test's scenario through +//! `forward_prompt` and `forward_token`, and prints one line per output +//! in the file's format, plus a `math-library` line naming this +//! system's math library, on which the bits depend as well as on the +//! code. +//! +//! The stored lines are this program's output at commit 30c4cf1ee (main +//! before the CPU passes), whose `forward_prompt` and `forward_token` +//! are the code the CPU passes must match. It uses only API that commit +//! has, so to add a platform: +//! +//! ```bash +//! git worktree add --detach /tmp/base 30c4cf1ee +//! cp examples/cohere2_reference.rs /tmp/base/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/examples/ +//! cd /tmp/base/jouleclaw/jouleclaw-rs +//! cargo run -q -p jouleclaw-loader-gguf --example cohere2_reference +//! cargo run -q -p jouleclaw-loader-gguf --features rayon --example cohere2_reference # must print the same +//! ``` +//! +//! and append its lines to the file. Run at a later commit, it prints +//! that commit's lines, which the test requires to be the same. + +use jouleclaw_loader_gguf::cohere2_moe::{Cohere2MoeConfig, Cohere2MoeModel}; +use jouleclaw_loader_gguf::safetensors_reader::read_safetensors_file; + +fn tiny_config() -> Cohere2MoeConfig { + Cohere2MoeConfig { + hidden_size: 8, + num_hidden_layers: 4, + vocab_size: 13, + layer_types: vec![ + "full_attention".into(), + "sliding_attention".into(), + "sliding_attention".into(), + "full_attention".into(), + ], + num_attention_heads: 2, + num_key_value_heads: 1, + head_dim: 4, + rope_theta: 10_000.0, + sliding_window: 3, + num_experts: 4, + num_experts_per_tok: 2, + intermediate_size: 6, + expert_selection_fn: "sigmoid".into(), + norm_topk_prob: false, + first_k_dense_replace: 1, + prefix_dense_intermediate_size: Some(10), + prefix_dense_sliding_window_pattern: 1, + rms_norm_eps: 1e-6, + logit_scale: 0.5, + } +} + +fn medium_config() -> Cohere2MoeConfig { + Cohere2MoeConfig { + hidden_size: 64, + num_hidden_layers: 4, + vocab_size: 37, + layer_types: vec![ + "full_attention".into(), + "sliding_attention".into(), + "sliding_attention".into(), + "full_attention".into(), + ], + num_attention_heads: 4, + num_key_value_heads: 2, + head_dim: 32, + rope_theta: 10_000.0, + sliding_window: 20, + num_experts: 8, + num_experts_per_tok: 3, + intermediate_size: 24, + expert_selection_fn: "sigmoid".into(), + norm_topk_prob: false, + first_k_dense_replace: 1, + prefix_dense_intermediate_size: Some(40), + prefix_dense_sliding_window_pattern: 1, + rms_norm_eps: 1e-6, + logit_scale: 0.5, + } +} + +fn wide_config() -> Cohere2MoeConfig { + Cohere2MoeConfig { + hidden_size: 256, + num_hidden_layers: 3, + vocab_size: 1031, + layer_types: vec![ + "full_attention".into(), + "sliding_attention".into(), + "full_attention".into(), + ], + num_attention_heads: 8, + num_key_value_heads: 2, + head_dim: 32, + rope_theta: 10_000.0, + sliding_window: 16, + num_experts: 8, + num_experts_per_tok: 2, + intermediate_size: 96, + expert_selection_fn: "sigmoid".into(), + norm_topk_prob: false, + first_k_dense_replace: 1, + prefix_dense_intermediate_size: Some(2048), + prefix_dense_sliding_window_pattern: 1, + rms_norm_eps: 1e-6, + logit_scale: 0.5, + } +} + +/// A scenario: `prompt` tokens prefilled, `decode` steps, `extend` +/// tokens fed into the live state, two more steps. +struct Shape { + prompt: u32, + decode: u32, + extend: u32, +} + +/// Deterministic weights for a tensor, from its name. +fn fill(name: &str, n: usize) -> Vec { + let mut h = 1469598103934665603u64; + for b in name.bytes() { + h = (h ^ b as u64).wrapping_mul(1099511628211); + } + (0..n) + .map(|i| { + let mut x = h.wrapping_add((i as u64).wrapping_mul(0x9E3779B97F4A7C15)); + x ^= x >> 33; + x = x.wrapping_mul(0xFF51AFD7ED558CCD); + x ^= x >> 33; + ((x % 2000) as f32 / 1000.0 - 1.0) * 0.4 + }) + .collect() +} + +/// The checkpoint for `cfg` as a BF16 safetensors file at `path`. +fn build_checkpoint(path: &std::path::Path, cfg: &Cohere2MoeConfig) { + let h = cfg.hidden_size as u64; + let mut tensors: Vec<(String, Vec, Vec)> = Vec::new(); + let mut push = |name: String, shape: Vec| { + let n: u64 = shape.iter().product(); + let vals = fill(&name, n as usize); + tensors.push((name, shape, vals)); + }; + push("model.embed_tokens.weight".into(), vec![cfg.vocab_size as u64, h]); + push("model.norm.weight".into(), vec![h]); + let q_dim = (cfg.num_attention_heads * cfg.head_dim) as u64; + let kv_dim = (cfg.num_key_value_heads * cfg.head_dim) as u64; + for i in 0..cfg.num_hidden_layers { + let p = format!("model.layers.{i}"); + push(format!("{p}.input_layernorm.weight"), vec![h]); + push(format!("{p}.self_attn.q_proj.weight"), vec![q_dim, h]); + push(format!("{p}.self_attn.k_proj.weight"), vec![kv_dim, h]); + push(format!("{p}.self_attn.v_proj.weight"), vec![kv_dim, h]); + push(format!("{p}.self_attn.o_proj.weight"), vec![h, q_dim]); + if i < cfg.first_k_dense_replace { + let pf = cfg.prefix_dense_intermediate_size.unwrap() as u64; + push(format!("{p}.mlp.gate_proj.weight"), vec![pf, h]); + push(format!("{p}.mlp.up_proj.weight"), vec![pf, h]); + push(format!("{p}.mlp.down_proj.weight"), vec![h, pf]); + } else { + push(format!("{p}.mlp.gate.weight"), vec![cfg.num_experts as u64, h]); + for e in 0..cfg.num_experts { + let f = cfg.intermediate_size as u64; + push(format!("{p}.mlp.experts.{e}.gate_proj.weight"), vec![f, h]); + push(format!("{p}.mlp.experts.{e}.up_proj.weight"), vec![f, h]); + push(format!("{p}.mlp.experts.{e}.down_proj.weight"), vec![h, f]); + } + } + } + let mut header = String::from("{"); + let mut data: Vec = Vec::new(); + for (i, (name, shape, vals)) in tensors.iter().enumerate() { + let start = data.len(); + for &v in vals { + // BF16 by truncation, as the test's checkpoints are. + data.extend_from_slice(&((v.to_bits() >> 16) as u16).to_le_bytes()); + } + let end = data.len(); + if i > 0 { + header.push(','); + } + header.push_str(&format!( + "\"{name}\":{{\"dtype\":\"BF16\",\"shape\":{shape:?},\"data_offsets\":[{start},{end}]}}" + )); + } + header.push('}'); + let mut out = Vec::new(); + out.extend_from_slice(&(header.len() as u64).to_le_bytes()); + out.extend_from_slice(header.as_bytes()); + out.extend_from_slice(&data); + std::fs::write(path, out).expect("write the checkpoint"); +} + +/// FNV-1a over every logit's bits, little-endian (north_bench's hash). +fn logits_hash(logits: &[f32]) -> u64 { + let mut h = 0xcbf2_9ce4_8422_2325u64; + for v in logits { + for b in v.to_bits().to_le_bytes() { + h = (h ^ b as u64).wrapping_mul(0x0000_0100_0000_01b3); + } + } + h +} + +/// The scenario's lines for one model: ` `. +fn scenario(name: &str, model: &Cohere2MoeModel, shape: &Shape) -> Vec { + let v = model.config.vocab_size as u32; + let mut out = Vec::new(); + let mut line = |what: String, l: &[f32]| out.push(format!("{name} {what} {:016x}", logits_hash(l))); + let mut state = model.new_state(); + let prompt: Vec = (0..shape.prompt).map(|i| (i * 7 + 3) % v).collect(); + line("prefill".into(), &model.forward_prompt(&prompt, &mut state)); + for i in 0..shape.decode { + line(format!("decode{i}"), &model.forward_token((i * 5 + 2) % v, &mut state)); + } + // The chunk is fed one decode step at a time (`forward_extend`, + // which batches it, gives the same logits and does not exist at + // the commit the stored lines come from). + let mut fed = Vec::new(); + for i in 0..shape.extend { + fed = model.forward_token((i * 11 + 5) % v, &mut state); + } + line("extend".into(), &fed); + for i in 0..2 { + line(format!("after{i}"), &model.forward_token((i * 3 + 1) % v, &mut state)); + } + out +} + +/// This system's math library, as `math-library` names it. +fn math_library() -> String { + let run = |cmd: &str, args: &[&str]| { + std::process::Command::new(cmd) + .args(args) + .output() + .ok() + .map(|o| String::from_utf8_lossy(&o.stdout).lines().next().unwrap_or("").trim().to_string()) + .unwrap_or_default() + }; + match std::env::consts::OS { + "macos" => format!( + "the system math library (libsystem_m) of macOS {} ({})", + run("sw_vers", &["-productVersion"]), + run("sw_vers", &["-buildVersion"]) + ), + "linux" => { + // `ldd (Ubuntu GLIBC 2.39-0ubuntu8.9) 2.39`: the C library's name and version. + let l = run("ldd", &["--version"]); + let libc = l.split_once('(').and_then(|(_, r)| r.rsplit_once(')')).map_or(l.as_str(), |(n, _)| n); + format!("the libm of {libc}") + } + os => format!("the system math library of {os}"), + } +} + +fn main() { + let platform = format!("{}-{}", std::env::consts::OS, std::env::consts::ARCH); + println!("math-library {platform} {}", math_library()); + let models = [ + ("tiny", tiny_config(), Shape { prompt: 5, decode: 4, extend: 4 }), + ("medium", medium_config(), Shape { prompt: 45, decode: 6, extend: 25 }), + ("wide", wide_config(), Shape { prompt: 24, decode: 4, extend: 20 }), + ]; + for (name, cfg, shape) in models { + let path = std::env::temp_dir().join(format!("cohere2-reference-{name}-{}.safetensors", std::process::id())); + build_checkpoint(&path, &cfg); + let st = read_safetensors_file(&path).expect("read the checkpoint"); + let model = Cohere2MoeModel::bind(cfg, st).expect("bind the checkpoint"); + let _ = std::fs::remove_file(&path); + for line in scenario(name, &model, &shape) { + println!("{platform} {line}"); + } + } +} diff --git a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/examples/north_bench.rs b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/examples/north_bench.rs new file mode 100644 index 000000000..c7ed8d7c6 --- /dev/null +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/examples/north_bench.rs @@ -0,0 +1,184 @@ +//! Speed and identity harness for Cohere2-MoE (North Mini Code) +//! decode: one prompt from a file, batched prefill, greedy +//! decode, and per step the token id, a hash of every logit's +//! bits (so two builds can be compared bit for bit, not only +//! token for token) and the step's wall-clock. +//! +//! ```bash +//! cargo run --release --features rayon --example north_bench -- \ +//! /path/to/snapshot-dir prompt.txt [n_tokens] [--chat] +//! ``` +//! +//! `--chat` wraps the file's text as one user turn of the +//! model's chat template, the shape the program-writing loop +//! feeds. `--then ` afterwards feeds that file's text as a +//! further user turn into the same context with +//! `forward_extend` (`--then-seq `: one `forward_token` per +//! token, the same logits) and greedily decodes 8 more tokens. +//! With `JOULECLAW_PROF=1` it also prints where the time went, op +//! by op, for the prefill and per decoded token. + +use jouleclaw_loader_gguf::bpe_tokenizer::BpeTokenizer; +use jouleclaw_loader_gguf::cohere2_moe::{prof, Cohere2MoeConfig, Cohere2MoeModel}; +use jouleclaw_loader_gguf::safetensors_reader::read_safetensors_file; +use std::path::PathBuf; +use std::time::Instant; + +/// FNV-1a over the bit patterns of every logit. +fn logits_hash(logits: &[f32]) -> u64 { + let mut h = 0xcbf2_9ce4_8422_2325u64; + for v in logits { + for b in v.to_bits().to_le_bytes() { + h = (h ^ b as u64).wrapping_mul(0x0000_0100_0000_01b3); + } + } + h +} + +fn print_prof(label: &str, per: f64) { + if !prof::enabled() { + return; + } + let rows = prof::take(); + let total: f64 = rows.iter().map(|r| r.1).sum(); + println!("profile {label} (ms{}):", if per > 1.0 { " per token" } else { "" }); + for (name, ms) in rows.iter().filter(|r| r.1 > 0.0) { + println!(" {:<34} {:>10.2} {:>5.1}%", name, ms / per, 100.0 * ms / total.max(1e-9)); + } + println!(" {:<34} {:>10.2}", "total", total / per); +} + +fn main() -> Result<(), Box> { + let args: Vec = std::env::args().collect(); + if args.len() < 3 { + eprintln!("usage: north_bench [n_tokens] [--chat]"); + std::process::exit(1); + } + let dir = PathBuf::from(&args[1]); + let text = std::fs::read_to_string(&args[2])?; + let n_tokens: usize = args.get(3).map(|s| s.parse()).transpose()?.unwrap_or(64); + let chat = args.iter().any(|a| a == "--chat"); + let then = |flag: &str| args.iter().position(|a| a == flag).and_then(|i| args.get(i + 1)).cloned(); + let (then_file, then_seq) = match (then("--then"), then("--then-seq")) { + (Some(f), _) => (Some(f), false), + (None, Some(f)) => (Some(f), true), + (None, None) => (None, false), + }; + + let config: Cohere2MoeConfig = serde_json::from_slice(&std::fs::read(dir.join("config.json"))?)?; + let mut shard_paths: Vec = std::fs::read_dir(&dir)? + .filter_map(|e| e.ok().map(|e| e.path())) + .filter(|p| { + p.extension().is_some_and(|e| e == "safetensors") + && p.file_name().is_some_and(|n| n.to_string_lossy().starts_with("model")) + }) + .collect(); + shard_paths.sort(); + let shards = shard_paths + .iter() + .map(|p| read_safetensors_file(p).map_err(|e| format!("{p:?}: {e:?}"))) + .collect::, _>>()?; + let model = Cohere2MoeModel::bind_sharded(config, shards)?; + let tokenizer = BpeTokenizer::from_file(dir.join("tokenizer.json"))?; + + let body = if chat { + format!( + "<|START_OF_TURN_TOKEN|><|USER_TOKEN|>{}<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>", + text.trim_end() + ) + } else { + text + }; + let mut prompt = vec![2u32]; // bos_token_id + prompt.extend(tokenizer.encode(&body)); + println!("prompt_tokens {}", prompt.len()); + + let mut state = model.new_state(); + let t_pre = Instant::now(); + let mut logits = model.forward_prompt(&prompt, &mut state); + let pre = t_pre.elapsed().as_secs_f64(); + println!( + "prefill_ms {:.0} tok_s {:.2}", + pre * 1e3, + prompt.len() as f64 / pre + ); + print_prof("prefill", 1.0); + + let mut ids = Vec::with_capacity(n_tokens); + let mut hashes = Vec::with_capacity(n_tokens); + let mut ms = Vec::with_capacity(n_tokens); + for step in 0..n_tokens { + let next = logits + .iter() + .enumerate() + .fold((0usize, f32::NEG_INFINITY), |(bi, bv), (i, &v)| if v > bv { (i, v) } else { (bi, bv) }) + .0 as u32; + let hash = logits_hash(&logits); + ids.push(next); + hashes.push(hash); + if step + 1 == n_tokens { + println!("step {step:>3} id {next:>6} hash {hash:016x}"); + break; + } + let t = Instant::now(); + logits = model.forward_token(next, &mut state); + let dt = t.elapsed().as_secs_f64() * 1e3; + ms.push(dt); + println!("step {step:>3} id {next:>6} hash {hash:016x} ms {dt:.1}"); + } + let mut sorted = ms.clone(); + sorted.sort_by(f64::total_cmp); + let median = sorted.get(sorted.len() / 2).copied().unwrap_or(0.0); + let mean = ms.iter().sum::() / ms.len().max(1) as f64; + println!( + "decode_ms median {median:.1} mean {mean:.1} over {} steps, context {}..{}", + ms.len(), + prompt.len(), + state.position + ); + print_prof("decode", ms.len().max(1) as f64); + println!("ids {}", ids.iter().map(|i| i.to_string()).collect::>().join(",")); + println!("hashes {}", hashes.iter().map(|h| format!("{h:016x}")).collect::>().join(",")); + println!("text {:?}", tokenizer.decode(&ids)); + + if let Some(file) = then_file { + let more = std::fs::read_to_string(&file)?; + let mut feed = vec![*ids.last().unwrap()]; + feed.extend(tokenizer.encode(&format!( + "<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|USER_TOKEN|>{}<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>", + more.trim_end() + ))); + let t = Instant::now(); + logits = if then_seq { + let mut l = Vec::new(); + for &x in &feed { + l = model.forward_token(x, &mut state); + } + l + } else { + model.forward_extend(&feed, &mut state) + }; + let dt = t.elapsed().as_secs_f64(); + println!( + "then_{} {} tokens in {:.0} ms ({:.2} tok/s), context {}, hash {:016x}", + if then_seq { "seq" } else { "extend" }, + feed.len(), + dt * 1e3, + feed.len() as f64 / dt, + state.position, + logits_hash(&logits) + ); + let mut more_ids = Vec::new(); + for _ in 0..8 { + let next = logits + .iter() + .enumerate() + .fold((0usize, f32::NEG_INFINITY), |(bi, bv), (i, &v)| if v > bv { (i, v) } else { (bi, bv) }) + .0 as u32; + more_ids.push(next); + logits = model.forward_token(next, &mut state); + } + println!("then_ids {} hash {:016x}", more_ids.iter().map(|i| i.to_string()).collect::>().join(","), logits_hash(&logits)); + } + Ok(()) +} diff --git a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/attention.rs b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/attention.rs index cfe6c59cd..25b13958c 100644 --- a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/attention.rs +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/attention.rs @@ -161,6 +161,210 @@ pub fn attention_forward_f32( } } +// ───────────────────────────────────────────────────────────────────── +// One query row, bit for bit, fast (2026-11) +// ───────────────────────────────────────────────────────────────────── + +/// One query row of [`attention_forward_f32`] against keys and +/// values `0..kv_len` (`k`, `v` hold at least `kv_len` rows), +/// with exactly its arithmetic: the same scores, max, softmax +/// and weighted sum of V, every value rounded at the same step +/// in the same order, so `out_row` is bit-identical. A decode +/// step calls it once per head (the heads run in parallel); a +/// prefill once per (head, query row), where `kv_len = t + 1` +/// is the causal mask. +/// +/// What makes it fast without changing a bit: the reference +/// adds each score's 128 products one at a time on one scalar +/// chain (an f32 sum may not be reordered), so it waits on the +/// add's latency for every product. Here the products of four +/// keys sit in the four lanes of a NEON register (a 4×4 +/// transpose of four key rows) and four such registers run at +/// once, so sixteen scores advance together, each lane still +/// adding its key's products in `d` order: multiply, round, +/// add, round, as the scalar code does (Rust never fuses them). +/// The weighted sum of V keeps its per-`d` running sums in +/// registers across keys, again multiply then add in key order. +/// `scores` is scratch of at least `kv_len`. +pub fn attention_row_f32( + q_row: &[f32], + k: &[f32], + v: &[f32], + kv_len: usize, + d_head: usize, + scores: &mut [f32], + out_row: &mut [f32], +) { + assert_eq!(q_row.len(), d_head); + assert!(k.len() >= kv_len * d_head); + assert!(v.len() >= kv_len * d_head); + assert!(scores.len() >= kv_len); + assert_eq!(out_row.len(), d_head); + assert!(d_head > 0); + let inv_sqrt_d = 1.0_f32 / (d_head as f32).sqrt(); + let scores = &mut scores[..kv_len]; + qk_scores(q_row, k, kv_len, d_head, inv_sqrt_d, scores); + let mut max_score = f32::NEG_INFINITY; + for &score in scores.iter() { + if score > max_score { + max_score = score; + } + } + let mut sum = 0.0f32; + for s in scores.iter_mut() { + let e = (*s - max_score).exp(); + *s = e; + sum += e; + } + let inv_sum = if sum > 0.0 { 1.0 / sum } else { 0.0 }; + for s in scores.iter_mut() { + *s *= inv_sum; + } + weighted_v(scores, v, d_head, out_row); +} + +/// `scores[s] = (Σ_d q[d]·k[s][d], summed in d order) · inv_sqrt_d`. +fn qk_scores(q: &[f32], k: &[f32], kv_len: usize, d_head: usize, inv_sqrt_d: f32, scores: &mut [f32]) { + #[cfg(target_arch = "aarch64")] + let s0 = if d_head.is_multiple_of(4) { qk_scores_neon(q, k, kv_len, d_head, inv_sqrt_d, scores) } else { 0 }; + #[cfg(not(target_arch = "aarch64"))] + let s0 = 0; + for s in s0..kv_len { + let k_row = &k[s * d_head..(s + 1) * d_head]; + let mut acc = 0.0f32; + for d in 0..d_head { + acc += q[d] * k_row[d]; + } + scores[s] = acc * inv_sqrt_d; + } +} + +/// The NEON part of [`qk_scores`]: keys in blocks of 16, then 4. +/// Returns how many leading keys it scored (`d_head % 4 == 0`). +#[cfg(target_arch = "aarch64")] +fn qk_scores_neon(q: &[f32], k: &[f32], kv_len: usize, d_head: usize, inv_sqrt_d: f32, scores: &mut [f32]) -> usize { + let mut s0 = 0; + { + // SAFETY: rows s0..s0+16 (or +4) are below kv_len, so + // every load stays inside k[..kv_len * d_head]; q has + // d_head elements. + unsafe { + while s0 + 16 <= kv_len { + let acc = qk_block::<4>(q.as_ptr(), k.as_ptr().add(s0 * d_head), d_head); + for (b, a) in acc.iter().enumerate() { + let mut lanes = [0f32; 4]; + std::arch::aarch64::vst1q_f32(lanes.as_mut_ptr(), *a); + for (l, &x) in lanes.iter().enumerate() { + scores[s0 + 4 * b + l] = x * inv_sqrt_d; + } + } + s0 += 16; + } + while s0 + 4 <= kv_len { + let acc = qk_block::<1>(q.as_ptr(), k.as_ptr().add(s0 * d_head), d_head); + let mut lanes = [0f32; 4]; + std::arch::aarch64::vst1q_f32(lanes.as_mut_ptr(), acc[0]); + for (l, &x) in lanes.iter().enumerate() { + scores[s0 + l] = x * inv_sqrt_d; + } + s0 += 4; + } + } + } + s0 +} + +/// `B` blocks of four consecutive key rows starting at `k`: +/// lane `l` of block `b` is key `4b + l`'s dot product with `q`, +/// accumulated `acc = acc + q[d]·k[d]` (separately rounded) in +/// `d` order. +#[cfg(target_arch = "aarch64")] +#[inline(always)] +unsafe fn qk_block( + q: *const f32, + k: *const f32, + d_head: usize, +) -> [std::arch::aarch64::float32x4_t; B] { + use std::arch::aarch64::*; + unsafe { + let mut acc = [vdupq_n_f32(0.0); B]; + let mut d = 0; + while d < d_head { + let qv = vld1q_f32(q.add(d)); + for b in 0..B { + let base = k.add(4 * b * d_head + d); + let r0 = vld1q_f32(base); + let r1 = vld1q_f32(base.add(d_head)); + let r2 = vld1q_f32(base.add(2 * d_head)); + let r3 = vld1q_f32(base.add(3 * d_head)); + // 4×4 transpose: c_j lane l = key l's element d + j. + let t0 = vtrn1q_f32(r0, r1); + let t1 = vtrn2q_f32(r0, r1); + let t2 = vtrn1q_f32(r2, r3); + let t3 = vtrn2q_f32(r2, r3); + let c0 = vreinterpretq_f32_f64(vtrn1q_f64(vreinterpretq_f64_f32(t0), vreinterpretq_f64_f32(t2))); + let c1 = vreinterpretq_f32_f64(vtrn1q_f64(vreinterpretq_f64_f32(t1), vreinterpretq_f64_f32(t3))); + let c2 = vreinterpretq_f32_f64(vtrn2q_f64(vreinterpretq_f64_f32(t0), vreinterpretq_f64_f32(t2))); + let c3 = vreinterpretq_f32_f64(vtrn2q_f64(vreinterpretq_f64_f32(t1), vreinterpretq_f64_f32(t3))); + let mut a = acc[b]; + a = vaddq_f32(a, vmulq_laneq_f32::<0>(c0, qv)); + a = vaddq_f32(a, vmulq_laneq_f32::<1>(c1, qv)); + a = vaddq_f32(a, vmulq_laneq_f32::<2>(c2, qv)); + a = vaddq_f32(a, vmulq_laneq_f32::<3>(c3, qv)); + acc[b] = a; + } + d += 4; + } + acc + } +} + +/// `out[d] = Σ_s w[s]·v[s][d]`, in key order, skipping zero +/// weights — the reference's `out[d] += w·v[d]` loop with the +/// running sums held in registers, 64 lanes at a time. +fn weighted_v(w: &[f32], v: &[f32], d_head: usize, out: &mut [f32]) { + #[cfg(target_arch = "aarch64")] + if d_head.is_multiple_of(4) { + use std::arch::aarch64::*; + let mut d0 = 0; + while d0 < d_head { + let lanes = (d_head - d0).min(64); + // SAFETY: d0 + lanes ≤ d_head, and rows s < w.len() + // lie inside v (asserted by the caller). + unsafe { + let mut acc = [vdupq_n_f32(0.0); 16]; + let nv = lanes / 4; + for (s, &weight) in w.iter().enumerate() { + if weight == 0.0 { + continue; + } + let row = v.as_ptr().add(s * d_head + d0); + for (i, a) in acc.iter_mut().enumerate().take(nv) { + *a = vaddq_f32(*a, vmulq_n_f32(vld1q_f32(row.add(4 * i)), weight)); + } + } + for (i, a) in acc.iter().enumerate().take(nv) { + vst1q_f32(out.as_mut_ptr().add(d0 + 4 * i), *a); + } + } + d0 += lanes; + } + return; + } + for o in out.iter_mut() { + *o = 0.0; + } + for (s, &weight) in w.iter().enumerate() { + if weight == 0.0 { + continue; + } + let v_row = &v[s * d_head..(s + 1) * d_head]; + for d in 0..d_head { + out[d] += weight * v_row[d]; + } + } +} + // ───────────────────────────────────────────────────────────────────── // Flash-attention chunked variant (wave 29) // ───────────────────────────────────────────────────────────────────── @@ -297,6 +501,53 @@ fn as_u8(s: &[f32]) -> &[u8] { mod tests { use super::*; + fn wave(n: usize, seed: usize, scale: f32) -> Vec { + (0..n) + .map(|i| { + let x = ((i * 2654435761 + seed * 40503) % 1000) as f32 / 500.0 - 1.0; + x * scale + }) + .collect() + } + + /// One row of [`attention_row_f32`] is the reference's row, + /// bit for bit: decode-shaped (one query, every key) and + /// causal prefill rows, key counts around the 4- and 16-key + /// blocks, head widths that are and are not multiples of 4, + /// and a spread large enough that some softmax weights + /// underflow to exactly zero (the skipped-key branch). + #[test] + fn attention_row_is_bit_identical_to_reference() { + for &d in &[4usize, 6, 8, 32, 128] { + for &seq in &[1usize, 2, 3, 4, 5, 15, 16, 17, 33, 100] { + for &scale in &[0.3f32, 40.0] { + let q = wave(seq * d, 1 + d, scale); + let k = wave(seq * d, 2 + seq, scale); + let v = wave(seq * d, 3, 1.0); + // Causal prefill: every row against the reference. + let mut want = vec![0f32; seq * d]; + attention_forward_f32(&q, &k, &v, seq, seq, d, true, &mut want); + let mut scores = vec![0f32; seq]; + for t in 0..seq { + let mut got = vec![0f32; d]; + attention_row_f32(&q[t * d..(t + 1) * d], &k, &v, t + 1, d, &mut scores, &mut got); + for (x, y) in want[t * d..(t + 1) * d].iter().zip(got.iter()) { + assert_eq!(x.to_bits(), y.to_bits(), "causal d={d} seq={seq} t={t} scale={scale}"); + } + } + // Decode: one query against all keys. + let mut want1 = vec![0f32; d]; + attention_forward_f32(&q[..d], &k, &v, 1, seq, d, false, &mut want1); + let mut got1 = vec![0f32; d]; + attention_row_f32(&q[..d], &k, &v, seq, d, &mut scores, &mut got1); + for (x, y) in want1.iter().zip(got1.iter()) { + assert_eq!(x.to_bits(), y.to_bits(), "decode d={d} seq={seq} scale={scale}"); + } + } + } + } + } + #[test] fn attention_identity_v_returns_v_when_q_eq_k() { // Q = K and V is the identity-style sequence → with diff --git a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/cohere2_moe.rs b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/cohere2_moe.rs index bdadbc156..ad244ad87 100644 --- a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/cohere2_moe.rs +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/cohere2_moe.rs @@ -39,6 +39,103 @@ use crate::rmsnorm::rmsnorm_f32_out; use crate::safetensors_reader::{SafetensorsDtype, SafetensorsModel, SafetensorsTensorInfo}; use crate::ssm::silu; +// ───────────────────────────────────────────────────────────────────── +// Per-op profile +// ───────────────────────────────────────────────────────────────────── + +/// Per-op wall-clock profile of [`Cohere2MoeModel::forward_token`] +/// and [`Cohere2MoeModel::forward_prompt`], on when the process +/// starts with `JOULECLAW_PROF=1`. Each forward pass keeps a +/// [`Lap`] and marks the end of every op; the elapsed time since +/// the previous mark goes to that op's bucket. Off, a mark is one +/// branch on a cached flag. It reads clocks only, never the data, +/// so it cannot change a logit. +pub mod prof { + use std::sync::atomic::{AtomicU64, Ordering}; + use std::sync::OnceLock; + use std::time::Instant; + + pub const EMBED: usize = 0; + pub const NORM: usize = 1; + pub const QKV: usize = 2; + pub const ROPE_CACHE: usize = 3; + pub const ATTN: usize = 4; + pub const O_PROJ: usize = 5; + pub const ATTN_BLOCK_GPU: usize = 6; + pub const DENSE_MLP: usize = 7; + pub const ROUTER: usize = 8; + pub const EXPERT_GATE_UP: usize = 9; + pub const EXPERT_ACT: usize = 10; + pub const EXPERT_DOWN: usize = 11; + pub const EXPERT_COMBINE: usize = 12; + pub const RESIDUAL: usize = 13; + pub const LM_HEAD: usize = 14; + pub const D1: usize = 15; + pub const D2: usize = 16; + pub const D3: usize = 17; + pub const EXPERTS_BY_EXPERT: usize = 18; + const N: usize = 19; + pub const NAMES: [&str; N] = [ + "embed", + "rmsnorm", + "q/k/v proj", + "rope + kv append", + "attention (scores, softmax, AV)", + "o_proj", + "attention block (gpu)", + "dense mlp (layer 0)", + "router + top-k", + "experts gate/up", + "experts silu*up", + "experts down", + "experts combine", + "residual", + "final norm + lm_head", + "[cpu] q/k/v + router (+dense gate/up)", + "[cpu] attention + expert gate/up", + "[cpu] o_proj + expert down", + "[cpu] experts, all stages, by expert", + ]; + + static NS: [AtomicU64; N] = [const { AtomicU64::new(0) }; N]; + + /// Whether the process started with `JOULECLAW_PROF=1`. + pub fn enabled() -> bool { + static ON: OnceLock = OnceLock::new(); + *ON.get_or_init(|| std::env::var("JOULECLAW_PROF").as_deref() == Ok("1")) + } + + /// The running clock of one forward pass. + pub struct Lap(Option); + + impl Lap { + /// Start a pass's clock (a no-op when profiling is off). + pub fn start() -> Self { + Lap(enabled().then(Instant::now)) + } + /// Charge the time since the previous mark to `op`. + #[inline] + pub fn mark(&mut self, op: usize) { + if let Some(t) = self.0.as_mut() { + let now = Instant::now(); + NS[op].fetch_add((now - *t).as_nanos() as u64, Ordering::Relaxed); + *t = now; + } + } + } + + /// The accumulated `(op, milliseconds)` since the last take, + /// and reset them. + pub fn take() -> Vec<(&'static str, f64)> { + NAMES + .iter() + .zip(NS.iter()) + .map(|(n, a)| (*n, a.swap(0, Ordering::Relaxed) as f64 / 1e6)) + .collect() + } +} +use prof::Lap; + // ───────────────────────────────────────────────────────────────────── // Config // ───────────────────────────────────────────────────────────────────── @@ -397,15 +494,35 @@ impl Cohere2MoeModel { /// One decode step: token in, next-token logits out /// (already scaled by `logit_scale`; lm_head is the tied /// embedding table). + /// + /// On the CPU this is [`Self::forward_token_cpu`]; with the + /// fabric engaged (`JOULECLAW_FABRIC=1` in a `fabric` build) + /// it is [`Self::forward_token_reference`], whose attention + /// block and expert windows may run on the GPU. pub fn forward_token(&self, token: u32, state: &mut DecodeState) -> Vec { + if crate::fabric_dispatch::fabric_engaged() { + self.forward_token_reference(token, state) + } else { + self.forward_token_cpu(token, state) + } + } + + /// The decode step as it was before 2026-11: each projection + /// its own dispatch, attention one head after another. Kept + /// as the fabric path and as the reference + /// [`Self::forward_token_cpu`] must match bit for bit. + pub fn forward_token_reference(&self, token: u32, state: &mut DecodeState) -> Vec { let cfg = &self.config; let h_dim = cfg.hidden_size; + let mut lap = Lap::start(); let mut h = vec![0f32; h_dim]; self.embed(token, &mut h); let mut normed = vec![0f32; h_dim]; + lap.mark(prof::EMBED); for (li, (layer, lstate)) in self.layers.iter().zip(state.layers.iter_mut()).enumerate() { rmsnorm_f32_out(&h, &layer.input_norm, 1, h_dim, cfg.rms_norm_eps, &mut normed); + lap.mark(prof::NORM); // ── attention branch ── let hd = cfg.head_dim; @@ -490,6 +607,7 @@ impl Cohere2MoeModel { } } attn_out.copy_from_slice(&o_out); + lap.mark(prof::ATTN_BLOCK_GPU); break 'attn; } lstate.gpu_dead = true; @@ -501,6 +619,7 @@ impl Cohere2MoeModel { self.proj(layer.q_proj, &normed, &mut q); self.proj(layer.k_proj, &normed, &mut k); self.proj(layer.v_proj, &normed, &mut v); + lap.mark(prof::QKV); if cfg.uses_rope(li) { rope_interleaved_inplace(&mut q, n_q, hd, state.position, cfg.rope_theta); rope_interleaved_inplace(&mut k, n_kv, hd, state.position, cfg.rope_theta); @@ -517,6 +636,7 @@ impl Cohere2MoeModel { } } } + lap.mark(prof::ROPE_CACHE); let seq = lstate.k_cache[0].len() / hd; let mut attn = vec![0f32; n_q * hd]; for qi in 0..n_q { @@ -532,7 +652,9 @@ impl Cohere2MoeModel { &mut attn[qi * hd..(qi + 1) * hd], ); } + lap.mark(prof::ATTN); self.proj(layer.o_proj, &attn, &mut attn_out); + lap.mark(prof::O_PROJ); break 'attn; } // 'attn @@ -549,6 +671,7 @@ impl Cohere2MoeModel { u_buf[i] *= silu(g_buf[i]); } self.proj(*down, &u_buf, &mut mlp_out); + lap.mark(prof::DENSE_MLP); } Mlp::Moe { router, experts } => { // Wave 63: one dispatch per expert stage @@ -560,6 +683,7 @@ impl Cohere2MoeModel { let f = cfg.intermediate_size; let picks = self.route(&logits); let n_p = picks.len(); + lap.mark(prof::ROUTER); let mut gu = vec![0f32; n_p * 2 * f]; { let mut jobs: Vec = @@ -581,12 +705,14 @@ impl Cohere2MoeModel { } crate::fabric_dispatch::matmul_bf16_auto_multi(&mut jobs); } + lap.mark(prof::EXPERT_GATE_UP); for chunk in gu.chunks_mut(2 * f) { let (g_half, u_half) = chunk.split_at_mut(f); for i in 0..f { u_half[i] *= silu(g_half[i]); } } + lap.mark(prof::EXPERT_ACT); let mut e_outs = vec![0f32; n_p * h_dim]; { let mut jobs: Vec = @@ -607,11 +733,13 @@ impl Cohere2MoeModel { } crate::fabric_dispatch::matmul_bf16_auto_multi(&mut jobs); } + lap.mark(prof::EXPERT_DOWN); for (out, &(_, weight)) in e_outs.chunks(h_dim).zip(picks.iter()) { for (mv, ev) in mlp_out.iter_mut().zip(out.iter()) { *mv += weight * *ev; } } + lap.mark(prof::EXPERT_COMBINE); } } @@ -619,6 +747,7 @@ impl Cohere2MoeModel { for i in 0..h_dim { h[i] += attn_out[i] + mlp_out[i]; } + lap.mark(prof::RESIDUAL); } let mut normed_f = vec![0f32; h_dim]; @@ -630,6 +759,7 @@ impl Cohere2MoeModel { *l *= cfg.logit_scale; } } + lap.mark(prof::LM_HEAD); state.position += 1; logits } @@ -642,21 +772,42 @@ impl Cohere2MoeModel { /// Sliding layers attend per-row windows; the cache is /// trimmed to the window afterwards, exactly matching /// rolling decode. Requires a fresh state. + /// + /// With the fabric off this is [`Self::forward_prompt_cpu`]. With + /// the fabric engaged (`JOULECLAW_FABRIC=1` in a `fabric` build) + /// a batched projection may leave the CPU, for AMX, the GPU or + /// the Metal-4 tensor path, and it is + /// [`Self::forward_prompt_reference`]. pub fn forward_prompt(&self, tokens: &[u32], state: &mut DecodeState) -> Vec { + if crate::fabric_dispatch::prefill_on_cpu() { + self.forward_prompt_cpu(tokens, state) + } else { + self.forward_prompt_reference(tokens, state) + } + } + + /// The prefill as it was before 2026-11 (one expert after + /// another, attention one head after another). Kept for the + /// fabric and as the reference [`Self::forward_prompt_cpu`] + /// must match bit for bit. + pub fn forward_prompt_reference(&self, tokens: &[u32], state: &mut DecodeState) -> Vec { assert!(!tokens.is_empty(), "empty prompt"); assert_eq!(state.position, 0, "forward_prompt requires a fresh state"); let cfg = &self.config; let h_dim = cfg.hidden_size; let t_len = tokens.len(); + let mut lap = Lap::start(); let mut h = vec![0f32; t_len * h_dim]; for (t, &tok) in tokens.iter().enumerate() { self.embed(tok, &mut h[t * h_dim..(t + 1) * h_dim]); } let mut normed = vec![0f32; t_len * h_dim]; + lap.mark(prof::EMBED); for (li, (layer, lstate)) in self.layers.iter().zip(state.layers.iter_mut()).enumerate() { rmsnorm_f32_out(&h, &layer.input_norm, t_len, h_dim, cfg.rms_norm_eps, &mut normed); + lap.mark(prof::NORM); // ── attention branch, batched projections ── let hd = cfg.head_dim; @@ -670,6 +821,7 @@ impl Cohere2MoeModel { self.proj_m(layer.q_proj, &normed, t_len, &mut q); self.proj_m(layer.k_proj, &normed, t_len, &mut k); self.proj_m(layer.v_proj, &normed, t_len, &mut v); + lap.mark(prof::QKV); for t in 0..t_len { if cfg.uses_rope(li) { rope_interleaved_inplace(&mut q[t * n_q * hd..(t + 1) * n_q * hd], n_q, hd, t, cfg.rope_theta); @@ -681,6 +833,7 @@ impl Cohere2MoeModel { lstate.v_cache[g].extend_from_slice(&v[o..o + hd]); } } + lap.mark(prof::ROPE_CACHE); let mut attn = vec![0f32; t_len * n_q * hd]; let mut q_head = vec![0f32; t_len * hd]; let mut out_head = vec![0f32; t_len * hd]; @@ -725,8 +878,10 @@ impl Cohere2MoeModel { lstate.v_cache[g].drain(..cut); } } + lap.mark(prof::ATTN); let mut attn_out = vec![0f32; t_len * h_dim]; self.proj_m(layer.o_proj, &attn, t_len, &mut attn_out); + lap.mark(prof::O_PROJ); // ── MLP branch ── let mut mlp_out = vec![0f32; t_len * h_dim]; @@ -741,6 +896,7 @@ impl Cohere2MoeModel { u_buf[i] *= silu(g_buf[i]); } self.proj_m(*down, &u_buf, t_len, &mut mlp_out); + lap.mark(prof::DENSE_MLP); } Mlp::Moe { router, experts } => { // Wave 56 pattern: batched router, tokens @@ -760,6 +916,7 @@ impl Cohere2MoeModel { } } let mut pick_out = vec![0f32; t_len * k_picks * h_dim]; + lap.mark(prof::ROUTER); for (e, members) in groups.iter().enumerate() { if members.is_empty() { continue; @@ -775,11 +932,14 @@ impl Cohere2MoeModel { let mut u_buf = vec![0f32; n_t * f]; self.proj_m(ge, &x_rows, n_t, &mut g_buf); self.proj_m(ue, &x_rows, n_t, &mut u_buf); + lap.mark(prof::EXPERT_GATE_UP); for i in 0..n_t * f { u_buf[i] *= silu(g_buf[i]); } + lap.mark(prof::EXPERT_ACT); let mut out_rows = vec![0f32; n_t * h_dim]; self.proj_m(de, &u_buf, n_t, &mut out_rows); + lap.mark(prof::EXPERT_DOWN); for (i, &(t, slot, _)) in members.iter().enumerate() { let dst = (t * k_picks + slot) * h_dim; pick_out[dst..dst + h_dim] @@ -796,12 +956,14 @@ impl Cohere2MoeModel { } } } + lap.mark(prof::EXPERT_COMBINE); } } for i in 0..t_len * h_dim { h[i] += attn_out[i] + mlp_out[i]; } + lap.mark(prof::RESIDUAL); } let last = &h[(t_len - 1) * h_dim..t_len * h_dim]; @@ -814,11 +976,580 @@ impl Cohere2MoeModel { *l *= cfg.logit_scale; } } + lap.mark(prof::LM_HEAD); state.position = t_len; logits } } +// ───────────────────────────────────────────────────────────────────── +// CPU forward passes (2026-11) +// ───────────────────────────────────────────────────────────────────── +// +// The same arithmetic as the reference passes, value for value, +// laid out for a many-core CPU. What changed is only which work +// runs at the same time: +// +// - decode: one parallel dispatch for everything that reads the +// normed row (q, k, v, router), one for the attention heads +// together with the experts' gate/up (both need only that +// dispatch's results), one for o_proj together with the +// experts' down; lm_head is a fourth. Every projection is split +// into ~256 KiB row chunks, so a loaded pool balances instead of +// waiting on a thread's fixed share, and the dispatches run on +// [`crate::spin_pool`], whose threads stay awake between them. +// The reference ran about seven dispatches per layer, three of +// them single-threaded, and all attention on one thread. +// - prefill: attention runs per (head, query row) in parallel, +// the experts run in parallel (one task per expert and its +// tokens), and the projections use the 2-D tiled kernel. +// - RoPE angles are computed once per position, not once per head +// and layer. +// - the batched prefill also extends a state that already holds a +// context ([`Cohere2MoeModel::forward_extend`]), which equals +// feeding those tokens one decode step at a time. +// +// Each output keeps its reduction order (the blocked BF16 kernel +// and [`crate::attention::attention_row_f32`] are bit-exact to the +// one-row loops), and every sum across tasks (the expert combine, +// the residual) still runs in the reference's order on one thread. +// The `cpu_passes_match_reference_*` tests pin it against the +// reference passes, and `stored_reference` against hashes the code +// before these passes produced. + +/// `(sin, cos)` of each interleaved RoPE pair at `position`, +/// computed exactly as [`rope_interleaved_inplace`] does (an f64 +/// angle rounded to f32), once per position. +fn rope_table(position: usize, head_dim: usize, theta: f64) -> Vec<(f32, f32)> { + (0..head_dim / 2) + .map(|i| { + let angle = theta.powf(-((2 * i) as f64) / head_dim as f64) * position as f64; + (angle.sin() as f32, angle.cos() as f32) + }) + .collect() +} + +/// [`rope_interleaved_inplace`] with the angles of [`rope_table`]. +fn rope_apply(x: &mut [f32], n_heads: usize, head_dim: usize, table: &[(f32, f32)]) { + for hd in 0..n_heads { + let row = &mut x[hd * head_dim..(hd + 1) * head_dim]; + for (i, &(sin, cos)) in table.iter().enumerate() { + let a = row[2 * i]; + let b = row[2 * i + 1]; + row[2 * i] = a * cos - b * sin; + row[2 * i + 1] = b * cos + a * sin; + } + } +} + +/// One unit of a parallel dispatch. +enum Work<'a> { + Gemv(crate::matmul_bf16::GemvChunk<'a>), + /// One decode head: `q` against cache rows `0..kv_len`. + Attn { q: &'a [f32], k: &'a [f32], v: &'a [f32], kv_len: usize, out: &'a mut [f32] }, + /// Prefill query rows `t0..t0 + out.len() / hd` of head `qi` + /// (`q` is `[t_len, n_q · hd]`). The cache holds `base` rows + /// from before this call, then the call's own; row `t` + /// attends cache rows up to `base + t`, the last `window` of + /// them when sliding. + AttnRows { + q: &'a [f32], + qi: usize, + n_q: usize, + base: usize, + t0: usize, + k: &'a [f32], + v: &'a [f32], + window: Option, + hd: usize, + out: &'a mut [f32], + }, +} + +impl Work<'_> { + fn run(&mut self) { + match self { + Work::Gemv(c) => c.run(), + Work::Attn { q, k, v, kv_len, out } => { + let mut scores = vec![0f32; *kv_len]; + crate::attention::attention_row_f32(q, k, v, *kv_len, q.len(), &mut scores, out); + } + Work::AttnRows { q, qi, n_q, base, t0, k, v, window, hd, out } => { + let hd = *hd; + let rows = out.len() / hd; + let mut scores = vec![0f32; *base + *t0 + rows]; + for (r, o) in out.chunks_mut(hd).enumerate() { + let t = *t0 + r; + // Keys 0..total are what a decode step at this + // row's position holds before its window trim; + // the window is the last `w` of them (all of + // them while it does not bind). + let total = *base + t + 1; + let start = window.map_or(0, |w| total - w.min(total)); + let len = total - start; + let q_row = &q[(t * *n_q + *qi) * hd..(t * *n_q + *qi + 1) * hd]; + crate::attention::attention_row_f32( + q_row, + &k[start * hd..(start + len) * hd], + &v[start * hd..(start + len) * hd], + len, + hd, + &mut scores, + o, + ); + } + } + } + } +} + +/// Run every unit, in parallel when `rayon` is on, on the +/// [`crate::spin_pool`]: the units are claimed one at a time, so +/// threads that are running take the remaining ones and nobody +/// waits on a slow thread's fixed share, and the pool's threads +/// stay awake between a decode step's loops. Measured on North +/// decode, interleaved step by step in one process (2026-11): on +/// an idle M4 Max 34 ms/token against 40–42 with the same 12 +/// threads in a rayon pool; on an M3 Ultra under load ~55–60, +/// 65–106 ms against 164–465 with rayon at 8–24 threads. +fn run_work(mut work: Vec>) { + #[cfg(feature = "rayon")] + { + /// The units, shared with the pool's threads. + struct Units<'a, 'b>(*mut Work<'a>, std::marker::PhantomData<&'b mut [Work<'a>]>); + // SAFETY: the pool runs each index on exactly one thread (its + // claim counter), so the `&mut` taken per index never aliases, + // and `run` returns only after every index has run. + unsafe impl Sync for Units<'_, '_> {} + let units = Units(work.as_mut_ptr(), std::marker::PhantomData); + crate::spin_pool::spin_pool().run(work.len(), &|i| { + let u = &units; + // SAFETY: as above; `i < work.len()`. + unsafe { (*u.0.add(i)).run() } + }); + } + #[cfg(not(feature = "rayon"))] + for w in work.iter_mut() { + w.run(); + } +} + +impl Cohere2MoeModel { + /// Queue `out = W · a` (m = 1) as row chunks. + fn push_proj<'a>(&'a self, work: &mut Vec>, w: MatRef, a: &'a [f32], out: &'a mut [f32]) { + let t = self.tensor(w); + let (n, k) = (t.shape[0] as usize, t.shape[1] as usize); + let mut chunks = Vec::new(); + crate::matmul_bf16::push_gemv_chunks(&mut chunks, a, self.shards[w.shard].tensor_bytes(t), n, k, out); + work.extend(chunks.into_iter().map(Work::Gemv)); + } + + /// One decode step on the CPU, bit-identical to + /// [`Self::forward_token_reference`] (see the section note). + /// Its parallel loops run on [`crate::spin_pool`]. + pub fn forward_token_cpu(&self, token: u32, state: &mut DecodeState) -> Vec { + let cfg = &self.config; + let h_dim = cfg.hidden_size; + let hd = cfg.head_dim; + let n_q = cfg.num_attention_heads; + let n_kv = cfg.num_key_value_heads; + let q_per_kv = n_q / n_kv; + let f = cfg.intermediate_size; + let mut lap = Lap::start(); + let mut h = vec![0f32; h_dim]; + self.embed(token, &mut h); + let mut normed = vec![0f32; h_dim]; + let rope = rope_table(state.position, hd, cfg.rope_theta); + lap.mark(prof::EMBED); + + for (li, (layer, lstate)) in self.layers.iter().zip(state.layers.iter_mut()).enumerate() { + rmsnorm_f32_out(&h, &layer.input_norm, 1, h_dim, cfg.rms_norm_eps, &mut normed); + lap.mark(prof::NORM); + + // ── dispatch 1: everything that reads only `normed` ── + let mut q = vec![0f32; n_q * hd]; + let mut k = vec![0f32; n_kv * hd]; + let mut v = vec![0f32; n_kv * hd]; + let mut router_logits = Vec::new(); + let (mut dg, mut du) = (Vec::new(), Vec::new()); + match &layer.mlp { + Mlp::Moe { .. } => router_logits = vec![0f32; cfg.num_experts], + Mlp::Dense { gate, .. } => { + let pf = self.tensor(*gate).shape[0] as usize; + dg = vec![0f32; pf]; + du = vec![0f32; pf]; + } + } + { + let mut work = Vec::new(); + self.push_proj(&mut work, layer.q_proj, &normed, &mut q); + self.push_proj(&mut work, layer.k_proj, &normed, &mut k); + self.push_proj(&mut work, layer.v_proj, &normed, &mut v); + match &layer.mlp { + Mlp::Moe { router, .. } => { + self.push_proj(&mut work, *router, &normed, &mut router_logits) + } + Mlp::Dense { gate, up, .. } => { + self.push_proj(&mut work, *gate, &normed, &mut dg); + self.push_proj(&mut work, *up, &normed, &mut du); + } + } + run_work(work); + } + lap.mark(prof::D1); + + if cfg.uses_rope(li) { + rope_apply(&mut q, n_q, hd, &rope); + rope_apply(&mut k, n_kv, hd, &rope); + } + for g in 0..n_kv { + lstate.k_cache[g].extend_from_slice(&k[g * hd..(g + 1) * hd]); + lstate.v_cache[g].extend_from_slice(&v[g * hd..(g + 1) * hd]); + if cfg.is_sliding(li) { + let max = cfg.sliding_window * hd; + if lstate.k_cache[g].len() > max { + let cut = lstate.k_cache[g].len() - max; + lstate.k_cache[g].drain(..cut); + lstate.v_cache[g].drain(..cut); + } + } + } + lap.mark(prof::ROPE_CACHE); + + let picks = match &layer.mlp { + Mlp::Moe { .. } => self.route(&router_logits), + Mlp::Dense { .. } => { + for i in 0..du.len() { + du[i] *= silu(dg[i]); + } + Vec::new() + } + }; + lap.mark(prof::ROUTER); + + // ── dispatch 2: attention heads + expert gate/up ── + let seq = lstate.k_cache[0].len() / hd; + let n_p = picks.len(); + let mut attn = vec![0f32; n_q * hd]; + let mut gu = vec![0f32; n_p * 2 * f]; + let mut mlp_out = vec![0f32; h_dim]; + { + let mut work = Vec::new(); + for (qi, out) in attn.chunks_mut(hd).enumerate() { + let g = qi / q_per_kv; + work.push(Work::Attn { + q: &q[qi * hd..(qi + 1) * hd], + k: &lstate.k_cache[g], + v: &lstate.v_cache[g], + kv_len: seq, + out, + }); + } + match &layer.mlp { + Mlp::Moe { experts, .. } => { + for (chunk, &(e, _)) in gu.chunks_mut(2 * f).zip(picks.iter()) { + let (g_half, u_half) = chunk.split_at_mut(f); + let (ge, ue, _) = experts[e]; + self.push_proj(&mut work, ge, &normed, g_half); + self.push_proj(&mut work, ue, &normed, u_half); + } + } + Mlp::Dense { down, .. } => self.push_proj(&mut work, *down, &du, &mut mlp_out), + } + run_work(work); + } + lap.mark(prof::D2); + for chunk in gu.chunks_mut(2 * f) { + let (g_half, u_half) = chunk.split_at_mut(f); + for i in 0..f { + u_half[i] *= silu(g_half[i]); + } + } + lap.mark(prof::EXPERT_ACT); + + // ── dispatch 3: o_proj + expert down ── + let mut attn_out = vec![0f32; h_dim]; + let mut e_outs = vec![0f32; n_p * h_dim]; + { + let mut work = Vec::new(); + self.push_proj(&mut work, layer.o_proj, &attn, &mut attn_out); + if let Mlp::Moe { experts, .. } = &layer.mlp { + for ((chunk, out), &(e, _)) in + gu.chunks(2 * f).zip(e_outs.chunks_mut(h_dim)).zip(picks.iter()) + { + let (_, _, de) = experts[e]; + self.push_proj(&mut work, de, &chunk[f..], out); + } + } + run_work(work); + } + lap.mark(prof::D3); + for (out, &(_, weight)) in e_outs.chunks(h_dim).zip(picks.iter()) { + for (mv, ev) in mlp_out.iter_mut().zip(out.iter()) { + *mv += weight * *ev; + } + } + lap.mark(prof::EXPERT_COMBINE); + + for i in 0..h_dim { + h[i] += attn_out[i] + mlp_out[i]; + } + lap.mark(prof::RESIDUAL); + } + + let mut normed_f = vec![0f32; h_dim]; + rmsnorm_f32_out(&h, &self.final_norm, 1, h_dim, cfg.rms_norm_eps, &mut normed_f); + let mut logits = vec![0f32; cfg.vocab_size]; + { + let mut work = Vec::new(); + self.push_proj(&mut work, self.embeddings, &normed_f, &mut logits); // tied lm_head + run_work(work); + } + if cfg.logit_scale != 1.0 { + for l in logits.iter_mut() { + *l *= cfg.logit_scale; + } + } + lap.mark(prof::LM_HEAD); + state.position += 1; + logits + } + + /// Batched prefill on the CPU, bit-identical to + /// [`Self::forward_prompt_reference`] (see the section + /// note). Requires a fresh state. It runs on a thread of + /// [`crate::matmul_bf16::cpu_pool`]. + pub fn forward_prompt_cpu(&self, tokens: &[u32], state: &mut DecodeState) -> Vec { + assert_eq!(state.position, 0, "forward_prompt requires a fresh state"); + crate::matmul_bf16::in_cpu_pool(|| self.extend_cpu_in_pool(tokens, state)) + } + + /// Feed `tokens` to a state at any position and return the + /// logits after the last one: bit-identical to calling + /// [`Self::forward_token`] once per token, but batched like a + /// prefill (each new row attends exactly the keys a decode + /// step at its position would, the sliding window included). + /// It is the fast way to add a turn to a state that already + /// holds a conversation: a caller that feeds the turn through + /// [`Self::forward_token`] instead gets the same logits, one + /// decode step per token. With the fabric engaged (see + /// [`Self::forward_prompt`]) it runs the tokens one at a time + /// itself, which is the same thing. + pub fn forward_extend(&self, tokens: &[u32], state: &mut DecodeState) -> Vec { + assert!(!tokens.is_empty(), "nothing to feed"); + if crate::fabric_dispatch::prefill_on_cpu() { + crate::matmul_bf16::in_cpu_pool(|| self.extend_cpu_in_pool(tokens, state)) + } else { + let mut logits = Vec::new(); + for &t in tokens { + logits = self.forward_token(t, state); + } + logits + } + } + + fn extend_cpu_in_pool(&self, tokens: &[u32], state: &mut DecodeState) -> Vec { + assert!(!tokens.is_empty(), "empty prompt"); + let p0 = state.position; + let cfg = &self.config; + let h_dim = cfg.hidden_size; + let t_len = tokens.len(); + let hd = cfg.head_dim; + let n_q = cfg.num_attention_heads; + let n_kv = cfg.num_key_value_heads; + let q_per_kv = n_q / n_kv; + let w = cfg.sliding_window; + let mut lap = Lap::start(); + + let mut h = vec![0f32; t_len * h_dim]; + for (t, &tok) in tokens.iter().enumerate() { + self.embed(tok, &mut h[t * h_dim..(t + 1) * h_dim]); + } + let mut normed = vec![0f32; t_len * h_dim]; + let ropes: Vec> = + (0..t_len).map(|t| rope_table(p0 + t, hd, cfg.rope_theta)).collect(); + lap.mark(prof::EMBED); + + for (li, (layer, lstate)) in self.layers.iter().zip(state.layers.iter_mut()).enumerate() { + rmsnorm_f32_out(&h, &layer.input_norm, t_len, h_dim, cfg.rms_norm_eps, &mut normed); + lap.mark(prof::NORM); + // Cache rows already held (a sliding layer's are at most + // its window), before this call's keys are appended. + let base = lstate.k_cache[0].len() / hd; + + let mut q = vec![0f32; t_len * n_q * hd]; + let mut k = vec![0f32; t_len * n_kv * hd]; + let mut v = vec![0f32; t_len * n_kv * hd]; + self.proj_m(layer.q_proj, &normed, t_len, &mut q); + self.proj_m(layer.k_proj, &normed, t_len, &mut k); + self.proj_m(layer.v_proj, &normed, t_len, &mut v); + lap.mark(prof::QKV); + for t in 0..t_len { + if cfg.uses_rope(li) { + rope_apply(&mut q[t * n_q * hd..(t + 1) * n_q * hd], n_q, hd, &ropes[t]); + rope_apply(&mut k[t * n_kv * hd..(t + 1) * n_kv * hd], n_kv, hd, &ropes[t]); + } + for g in 0..n_kv { + let o = (t * n_kv + g) * hd; + lstate.k_cache[g].extend_from_slice(&k[o..o + hd]); + lstate.v_cache[g].extend_from_slice(&v[o..o + hd]); + } + } + lap.mark(prof::ROPE_CACHE); + + // ── attention: (head, block of query rows) in parallel ── + let window = cfg.is_sliding(li).then_some(w); + let mut by_head = vec![0f32; n_q * t_len * hd]; + { + const ROWS: usize = 16; + let mut work = Vec::new(); + for (qi, head) in by_head.chunks_mut(t_len * hd).enumerate() { + let g = qi / q_per_kv; + for (b, out) in head.chunks_mut(ROWS * hd).enumerate() { + work.push(Work::AttnRows { + q: &q, + qi, + n_q, + base, + t0: b * ROWS, + k: &lstate.k_cache[g], + v: &lstate.v_cache[g], + window, + hd, + out, + }); + } + } + run_work(work); + } + let mut attn = vec![0f32; t_len * n_q * hd]; + for qi in 0..n_q { + for t in 0..t_len { + let dst = (t * n_q + qi) * hd; + let src = (qi * t_len + t) * hd; + attn[dst..dst + hd].copy_from_slice(&by_head[src..src + hd]); + } + } + // Trim sliding caches to the window: the state rolling + // decode would have left. + if cfg.is_sliding(li) && base + t_len > w { + for g in 0..n_kv { + let cut = (base + t_len - w) * hd; + lstate.k_cache[g].drain(..cut); + lstate.v_cache[g].drain(..cut); + } + } + lap.mark(prof::ATTN); + let mut attn_out = vec![0f32; t_len * h_dim]; + self.proj_m(layer.o_proj, &attn, t_len, &mut attn_out); + lap.mark(prof::O_PROJ); + + let mut mlp_out = vec![0f32; t_len * h_dim]; + match &layer.mlp { + Mlp::Dense { gate, up, down } => { + let pf = self.tensor(*gate).shape[0] as usize; + let mut g_buf = vec![0f32; t_len * pf]; + let mut u_buf = vec![0f32; t_len * pf]; + self.proj_m(*gate, &normed, t_len, &mut g_buf); + self.proj_m(*up, &normed, t_len, &mut u_buf); + for i in 0..t_len * pf { + u_buf[i] *= silu(g_buf[i]); + } + self.proj_m(*down, &u_buf, t_len, &mut mlp_out); + lap.mark(prof::DENSE_MLP); + } + Mlp::Moe { router, experts } => { + let n_e = cfg.num_experts; + let f = cfg.intermediate_size; + let k_picks = cfg.num_experts_per_tok; + let mut router_logits = vec![0f32; t_len * n_e]; + self.proj_m(*router, &normed, t_len, &mut router_logits); + let picks_per_token: Vec> = (0..t_len) + .map(|t| self.route(&router_logits[t * n_e..(t + 1) * n_e])) + .collect(); + // Tokens grouped by expert; `row_of[t·k + slot]` + // is that pick's row in its expert's group. + let mut groups: Vec> = vec![Vec::new(); n_e]; + let mut row_of = vec![0usize; t_len * k_picks]; + for (t, picks) in picks_per_token.iter().enumerate() { + for (slot, &(e, _)) in picks.iter().enumerate() { + row_of[t * k_picks + slot] = groups[e].len(); + groups[e].push(t); + } + } + lap.mark(prof::ROUTER); + let expert = |e: usize| -> Vec { + let members = &groups[e]; + if members.is_empty() { + return Vec::new(); + } + let n_t = members.len(); + let mut x_rows = vec![0f32; n_t * h_dim]; + for (i, &t) in members.iter().enumerate() { + x_rows[i * h_dim..(i + 1) * h_dim] + .copy_from_slice(&normed[t * h_dim..(t + 1) * h_dim]); + } + let (ge, ue, de) = experts[e]; + let bytes = |m: MatRef| self.shards[m.shard].tensor_bytes(self.tensor(m)); + let mut g_buf = vec![0f32; n_t * f]; + let mut u_buf = vec![0f32; n_t * f]; + crate::matmul_bf16::matmul_bf16_neon(&x_rows, bytes(ge), &[], n_t, f, h_dim, &mut g_buf); + crate::matmul_bf16::matmul_bf16_neon(&x_rows, bytes(ue), &[], n_t, f, h_dim, &mut u_buf); + for i in 0..n_t * f { + u_buf[i] *= silu(g_buf[i]); + } + let mut out_rows = vec![0f32; n_t * h_dim]; + crate::matmul_bf16::matmul_bf16_neon(&u_buf, bytes(de), &[], n_t, h_dim, f, &mut out_rows); + out_rows + }; + #[cfg(feature = "rayon")] + let outs: Vec> = { + use rayon::prelude::*; + (0..n_e).into_par_iter().with_max_len(1).map(expert).collect() + }; + #[cfg(not(feature = "rayon"))] + let outs: Vec> = (0..n_e).map(expert).collect(); + lap.mark(prof::EXPERTS_BY_EXPERT); + for (t, picks) in picks_per_token.iter().enumerate() { + let m_row = &mut mlp_out[t * h_dim..(t + 1) * h_dim]; + for (slot, &(e, weight)) in picks.iter().enumerate() { + let i = row_of[t * k_picks + slot]; + let src = &outs[e][i * h_dim..(i + 1) * h_dim]; + for (mv, ev) in m_row.iter_mut().zip(src.iter()) { + *mv += weight * *ev; + } + } + } + lap.mark(prof::EXPERT_COMBINE); + } + } + + for i in 0..t_len * h_dim { + h[i] += attn_out[i] + mlp_out[i]; + } + lap.mark(prof::RESIDUAL); + } + + let last = &h[(t_len - 1) * h_dim..t_len * h_dim]; + let mut normed_f = vec![0f32; h_dim]; + rmsnorm_f32_out(last, &self.final_norm, 1, h_dim, cfg.rms_norm_eps, &mut normed_f); + let mut logits = vec![0f32; cfg.vocab_size]; + { + let mut work = Vec::new(); + self.push_proj(&mut work, self.embeddings, &normed_f, &mut logits); + run_work(work); + } + if cfg.logit_scale != 1.0 { + for l in logits.iter_mut() { + *l *= cfg.logit_scale; + } + } + lap.mark(prof::LM_HEAD); + state.position = p0 + t_len; + logits + } +} + impl crate::hybrid_lm::HybridLm for Cohere2MoeModel { type State = DecodeState; fn new_state(&self) -> Self::State { @@ -842,6 +1573,7 @@ impl crate::hybrid_lm::HybridLm for Cohere2MoeModel { #[cfg(test)] pub(crate) mod tests { use super::*; + use crate::fabric_dispatch::fabric_switch; use crate::safetensors_reader::read_safetensors_file; /// Tiny config with a TINY sliding window (3) so the @@ -1024,6 +1756,189 @@ pub(crate) mod tests { assert!((picks[0].1 + picks[1].1 - 1.0).abs() > 0.1); } + /// A config that reaches the fast paths the tiny one skips: + /// 32-wide heads (the 16-key attention blocks) and a prompt + /// longer than its sliding window. (Splitting a projection + /// into several row chunks needs larger weights: the stored + /// reference's wide config has them, and matmul_bf16's tests + /// pin the kernels at North's widths.) + fn medium_config() -> Cohere2MoeConfig { + Cohere2MoeConfig { + hidden_size: 64, + num_hidden_layers: 4, + vocab_size: 37, + layer_types: vec![ + "full_attention".into(), + "sliding_attention".into(), + "sliding_attention".into(), + "full_attention".into(), + ], + num_attention_heads: 4, + num_key_value_heads: 2, + head_dim: 32, + rope_theta: 10_000.0, + sliding_window: 20, + num_experts: 8, + num_experts_per_tok: 3, + intermediate_size: 24, + expert_selection_fn: "sigmoid".into(), + norm_topk_prob: false, + first_k_dense_replace: 1, + prefix_dense_intermediate_size: Some(40), + prefix_dense_sliding_window_pattern: 1, + rms_norm_eps: 1e-6, + logit_scale: 0.5, + } + } + + fn model_of(cfg: Cohere2MoeConfig, tag: &str) -> Cohere2MoeModel { + let path = tmp(tag); + build_tiny_checkpoint(&path, &cfg); + let st = read_safetensors_file(&path).unwrap(); + let m = Cohere2MoeModel::bind(cfg, st).unwrap(); + let _ = std::fs::remove_file(&path); + m + } + + fn bits(v: &[f32]) -> Vec { + v.iter().map(|x| x.to_bits()).collect() + } + + /// The reference passes are the CPU passes' oracle only as CPU + /// arithmetic: with the fabric engaged they hand projections to + /// its members, which round differently (it is what those passes + /// are kept for). + fn assert_reference_on_the_cpu() { + assert!( + !crate::fabric_dispatch::fabric_engaged(), + "the fabric is engaged in this process, so the reference passes are not CPU arithmetic" + ); + } + + /// The CPU passes are the reference passes, bit for bit: the + /// prefill logits, then every decode step (past the sliding + /// window, so the rolling trim runs), from a CPU prefill and + /// from a reference prefill. + fn cpu_matches_reference(model: &Cohere2MoeModel, prompt: &[u32], steps: usize) { + assert_reference_on_the_cpu(); + let mut s_ref = model.new_state(); + let mut s_cpu = model.new_state(); + let l_ref = model.forward_prompt_reference(prompt, &mut s_ref); + let l_cpu = model.forward_prompt_cpu(prompt, &mut s_cpu); + assert_eq!(bits(&l_ref), bits(&l_cpu), "prefill logits"); + for step in 0..steps { + let tok = (step as u32 * 5 + 2) % model.config.vocab_size as u32; + let a = model.forward_token_reference(tok, &mut s_ref); + let b = model.forward_token_cpu(tok, &mut s_cpu); + assert_eq!(bits(&a), bits(&b), "decode step {step}"); + } + // Sequential CPU decode from scratch equals both, too. + let mut s_seq = model.new_state(); + let mut last = Vec::new(); + for &t in prompt { + last = model.forward_token_cpu(t, &mut s_seq); + } + let mut s_ref2 = model.new_state(); + let mut last_ref = Vec::new(); + for &t in prompt { + last_ref = model.forward_token_reference(t, &mut s_ref2); + } + assert_eq!(bits(&last), bits(&last_ref), "sequential decode"); + } + + /// `forward_extend` in chunks equals feeding the same tokens + /// one at a time through the reference decode, bit for bit, + /// at every chunk boundary and in the decode after it; the + /// chunks start before, straddle and pass the sliding window. + fn extend_matches_sequential(model: &Cohere2MoeModel, tokens: &[u32], cuts: &[usize]) { + assert_reference_on_the_cpu(); + assert!( + crate::fabric_dispatch::prefill_on_cpu(), + "nothing can leave the CPU with the fabric off, so forward_extend must batch the tokens" + ); + let mut s_ext = model.new_state(); + let mut s_seq = model.new_state(); + let mut at = 0; + for &cut in cuts.iter().chain(std::iter::once(&tokens.len())) { + let chunk = &tokens[at..cut]; + let got = if at == 0 { + model.forward_prompt_cpu(chunk, &mut s_ext) + } else { + model.forward_extend(chunk, &mut s_ext) + }; + let mut want = Vec::new(); + for &t in chunk { + want = model.forward_token_reference(t, &mut s_seq); + } + assert_eq!(bits(&got), bits(&want), "after feeding {cut} tokens"); + assert_eq!(s_ext.position, s_seq.position); + at = cut; + } + for step in 0..6u32 { + let t = (step * 3 + 1) % model.config.vocab_size as u32; + let a = model.forward_token_cpu(t, &mut s_ext); + let b = model.forward_token_reference(t, &mut s_seq); + assert_eq!(bits(&a), bits(&b), "decode {step} after the extends"); + } + } + + // The four tests below compare against the reference passes, so + // they run where the fabric is off: here, or in a child process + // when this one has it on (`fabric_switch`). + + #[test] + fn extend_matches_sequential_decode_tiny() { + if !fabric_switch::here(false, module_path!(), "extend_matches_sequential_decode_tiny") { + return; + } + // Window 3: chunks of 2, 1, 4 and 3 tokens. + let model = tiny_model(); + extend_matches_sequential(&model, &[1, 5, 9, 2, 7, 3, 11, 4, 6, 8], &[2, 3, 7]); + } + + #[test] + fn extend_matches_sequential_decode_medium() { + if !fabric_switch::here(false, module_path!(), "extend_matches_sequential_decode_medium") { + return; + } + // Window 20: 13 tokens, then 9 (crossing it), then 25. + let model = model_of(medium_config(), "medium-extend"); + let tokens: Vec = (0..47u32).map(|i| (i * 11 + 5) % 37).collect(); + extend_matches_sequential(&model, &tokens, &[13, 22]); + } + + #[test] + fn cpu_passes_match_reference_tiny() { + if !fabric_switch::here(false, module_path!(), "cpu_passes_match_reference_tiny") { + return; + } + let model = tiny_model(); + cpu_matches_reference(&model, &[1, 5, 9, 2, 7], 8); + } + + #[test] + fn cpu_passes_match_reference_medium() { + if !fabric_switch::here(false, module_path!(), "cpu_passes_match_reference_medium") { + return; + } + let model = model_of(medium_config(), "medium"); + let prompt: Vec = (0..45u32).map(|i| (i * 7 + 3) % 37).collect(); + cpu_matches_reference(&model, &prompt, 30); + } + + /// Every route a batched projection has off the CPU needs the + /// fabric (the Metal-4 tensor path too), so with the fabric off + /// `forward_prompt` and `forward_extend` take the CPU passes, in + /// every build and whatever `JOULECLAW_MPP_PREFILL` says. + #[test] + fn prefill_takes_the_cpu_passes_while_the_fabric_is_off() { + assert_eq!( + crate::fabric_dispatch::prefill_on_cpu(), + !crate::fabric_dispatch::fabric_engaged(), + "prefill_on_cpu() against the fabric being off" + ); + } + #[test] fn generate_through_the_trait() { let model = tiny_model(); @@ -1035,4 +1950,282 @@ pub(crate) mod tests { assert_eq!(out.len(), 4); assert!(out.iter().all(|&t| (t as usize) < model.config.vocab_size)); } + + /// The stored reference, `cohere2_moe_reference.txt`: logit hashes + /// of three synthetic checkpoints from the code before the CPU + /// passes, for the platforms it was generated on. + #[cfg(any( + all(target_os = "macos", target_arch = "aarch64"), + all(target_os = "linux", any(target_arch = "aarch64", target_arch = "x86_64")) + ))] + mod stored_reference { + use super::*; + + /// FNV-1a over every logit's bits, little-endian (north_bench's + /// hash). + fn logits_hash(logits: &[f32]) -> u64 { + let mut h = 0xcbf2_9ce4_8422_2325u64; + for v in logits { + for b in v.to_bits().to_le_bytes() { + h = (h ^ b as u64).wrapping_mul(0x0000_0100_0000_01b3); + } + } + h + } + + /// A config past the thresholds the tiny and medium ones stay + /// under. In decode the tied lm_head (1031 rows), the dense + /// prefix's gate and up (2048 rows) and its down (k = 2048) are + /// each several 256 KiB row chunks; in a prefill of 16 tokens or + /// more the dense projections clear the threaded kernel's 16 MiB + /// floor (m × weight bytes), so they run as its 2-D tiles. + fn wide_config() -> Cohere2MoeConfig { + Cohere2MoeConfig { + hidden_size: 256, + num_hidden_layers: 3, + vocab_size: 1031, + layer_types: vec![ + "full_attention".into(), + "sliding_attention".into(), + "full_attention".into(), + ], + num_attention_heads: 8, + num_key_value_heads: 2, + head_dim: 32, + rope_theta: 10_000.0, + sliding_window: 16, + num_experts: 8, + num_experts_per_tok: 2, + intermediate_size: 96, + expert_selection_fn: "sigmoid".into(), + norm_topk_prob: false, + first_k_dense_replace: 1, + prefix_dense_intermediate_size: Some(2048), + prefix_dense_sliding_window_pattern: 1, + rms_norm_eps: 1e-6, + logit_scale: 0.5, + } + } + + /// A stored scenario: `prompt` tokens prefilled, `decode` steps, + /// `extend` tokens fed into the live state at once, two more + /// steps. + struct Shape { + prompt: u32, + decode: u32, + extend: u32, + } + + /// The three checkpoints of the stored reference. Each prompt and + /// each fed chunk is longer than the model's sliding window. + fn stored_models() -> [(&'static str, Cohere2MoeConfig, Shape); 3] { + [ + ("tiny", tiny_config(), Shape { prompt: 5, decode: 4, extend: 4 }), + ("medium", medium_config(), Shape { prompt: 45, decode: 6, extend: 25 }), + ("wide", wide_config(), Shape { prompt: 24, decode: 4, extend: 20 }), + ] + } + + /// Run the scenario with the given prefill, decode step and feed, + /// and return one line per output: ` `. + fn scenario( + name: &str, + model: &Cohere2MoeModel, + shape: &Shape, + prefill: impl Fn(&[u32], &mut DecodeState) -> Vec, + step: impl Fn(u32, &mut DecodeState) -> Vec, + feed: impl Fn(&[u32], &mut DecodeState) -> Vec, + ) -> Vec { + let v = model.config.vocab_size as u32; + let mut out = Vec::new(); + let mut line = |what: String, l: &[f32]| out.push(format!("{name} {what} {:016x}", logits_hash(l))); + let mut state = model.new_state(); + let prompt: Vec = (0..shape.prompt).map(|i| (i * 7 + 3) % v).collect(); + line("prefill".into(), &prefill(&prompt, &mut state)); + for i in 0..shape.decode { + line(format!("decode{i}"), &step((i * 5 + 2) % v, &mut state)); + } + let chunk: Vec = (0..shape.extend).map(|i| (i * 11 + 5) % v).collect(); + line("extend".into(), &feed(&chunk, &mut state)); + for i in 0..2 { + line(format!("after{i}"), &step((i * 3 + 1) % v, &mut state)); + } + out + } + + /// The line a check of the stored reference prints, so the + /// process that started it can count the checks made. + const CHECKED: &str = "stored reference checked:"; + + /// `got` (each way of running the model, with its lines) + /// against `stored`'s lines for `platform`. The bits depend on + /// the math library as well as on the code, so a mismatch names + /// the library the lines were made with (the file's + /// `math-library` line) and says whether the ways of running the + /// model at least agree with each other. + fn compare_with_stored(stored: &str, platform: &str, got: &[(&str, Vec)]) -> Result<(), String> { + let want: Vec<&str> = stored + .lines() + .filter_map(|l| l.strip_prefix(platform)?.strip_prefix(' ')) + .collect(); + if want.is_empty() { + return Err(format!( + "no stored lines for {platform}: generate them with examples/cohere2_reference.rs \ + (see its header) and add them to cohere2_moe_reference.txt" + )); + } + let library = stored + .lines() + .find_map(|l| l.strip_prefix("math-library ")?.strip_prefix(platform)?.strip_prefix(' ')) + .unwrap_or("a math library the file does not name"); + let wrong: Vec<&str> = got.iter().filter(|(_, lines)| lines != &want).map(|(p, _)| *p).collect(); + if wrong.is_empty() { + return Ok(()); + } + let mut msg = format!( + "on {platform}, the {} passes differ from the stored lines, which were made with {library}.\n", + wrong.join(", ") + ); + if got.windows(2).all(|w| w[0].1 == w[1].1) { + msg.push_str( + "The ways of running the model agree with each other. Either a change moved the arithmetic \ + of all of them alike, or this environment's math library rounds exp, powf, sin or cos \ + differently from the one above. To tell which, run examples/cohere2_reference.rs at commit \ + 30c4cf1ee here: if its lines differ from the stored ones too, the math library does.\n", + ); + } else { + msg.push_str("The ways of running the model also differ from each other, which no math library explains.\n"); + } + for (passes, lines) in got { + for (g, w) in lines.iter().zip(&want) { + if g != w { + msg.push_str(&format!(" {passes}: got {g}, stored {w}\n")); + } + } + if lines.len() != want.len() { + msg.push_str(&format!(" {passes}: {} lines, {} stored\n", lines.len(), want.len())); + } + } + Err(msg) + } + + /// A mismatch names the math library the stored lines were made + /// with, since the bits depend on it as well as on the code, and + /// says whether the ways of running the model agree with each + /// other. + #[test] + fn a_mismatch_names_the_math_library_the_lines_were_made_with() { + let stored = "math-library testos-testarch libm 9.9 of TestOS\n\ + testos-testarch tiny prefill 0000000000000001\n"; + let line = |h: &str| vec![format!("tiny prefill {h}")]; + let agreeing = [("reference", line("02")), ("cpu", line("02")), ("public", line("02"))]; + let err = compare_with_stored(stored, "testos-testarch", &agreeing).expect_err("the lines differ"); + assert!(err.contains("libm 9.9 of TestOS") && err.contains("math library"), "{err}"); + assert!(err.contains("agree with each other"), "{err}"); + let disagreeing = [("reference", line("02")), ("cpu", line("03")), ("public", line("02"))]; + let err = compare_with_stored(stored, "testos-testarch", &disagreeing).expect_err("the lines differ"); + assert!(err.contains("differ from each other"), "{err}"); + let matching = [("reference", line("0000000000000001"))]; + assert_eq!(compare_with_stored(stored, "testos-testarch", &matching), Ok(())); + } + + /// With the fabric switched on for the whole test run, the stored + /// reference is still checked in every configuration (each in a + /// child process with the fabric off). + #[cfg(feature = "fabric")] + #[test] + fn with_the_fabric_on_the_stored_reference_is_checked_in_every_configuration() { + let out = fabric_switch::run_as_top_level( + module_path!(), + "logits_match_the_stored_reference", + &[("JOULECLAW_FABRIC", Some("1"))], + ); + let checked = out.lines().filter(|l| l.contains(CHECKED)).count(); + let want = if cfg!(feature = "rayon") { 3 } else { 1 }; + assert_eq!(checked, want, "checks reported by the test run with JOULECLAW_FABRIC=1:\n{out}"); + } + + /// The CPU passes, the reference passes and the public entry + /// points all give this platform's stored hashes, which the + /// code before the CPU passes produced (see the file's + /// header). The fabric switch, the thread count and the GEMV + /// chunk size are each read once per process, so the test + /// process checks nothing itself: it starts one child per + /// configuration, each with the fabric off, at the default + /// thread count and, with `rayon`, at 1 thread and at 3 threads + /// with 8 KiB chunks, and each child must report its check. + #[test] + fn logits_match_the_stored_reference() { + if fabric_switch::in_child() { + check_against_the_stored_reference(); + return; + } + #[allow(unused_mut)] + let mut configurations: Vec)>> = vec![vec![]]; + #[cfg(feature = "rayon")] + configurations.extend([ + vec![("JOULECLAW_THREADS", Some("1"))], + vec![("JOULECLAW_THREADS", Some("3")), ("JOULECLAW_CHUNK_KB", Some("8"))], + ]); + for configuration in configurations { + let env = [vec![("JOULECLAW_FABRIC", None)], configuration].concat(); + let out = fabric_switch::run_alone(module_path!(), "logits_match_the_stored_reference", &env); + // libtest may print the test's name on the same line. + let checks: Vec<&str> = out.lines().filter_map(|l| l.find(CHECKED).map(|i| &l[i..])).collect(); + assert_eq!(checks.len(), 1, "the child with {env:?} reported {} checks:\n{out}", checks.len()); + println!("{}", checks[0]); + } + } + + /// One configuration's check, in a child process the test + /// started with the fabric off. + fn check_against_the_stored_reference() { + assert!(!crate::fabric_dispatch::fabric_engaged(), "the child runs with the fabric off"); + let platform = format!("{}-{}", std::env::consts::OS, std::env::consts::ARCH); + let mut got: [(&str, Vec); 3] = + [("reference", Vec::new()), ("cpu", Vec::new()), ("public", Vec::new())]; + for (name, cfg, shape) in stored_models() { + let m = model_of(cfg, name); + got[0].1.extend(scenario( + name, + &m, + &shape, + |t, s| m.forward_prompt_reference(t, s), + |t, s| m.forward_token_reference(t, s), + |ts, s| { + let mut l = Vec::new(); + for &t in ts { + l = m.forward_token_reference(t, s); + } + l + }, + )); + got[1].1.extend(scenario( + name, + &m, + &shape, + |t, s| m.forward_prompt_cpu(t, s), + |t, s| m.forward_token_cpu(t, s), + |ts, s| crate::matmul_bf16::in_cpu_pool(|| m.extend_cpu_in_pool(ts, s)), + )); + got[2].1.extend(scenario( + name, + &m, + &shape, + |t, s| m.forward_prompt(t, s), + |t, s| m.forward_token(t, s), + |ts, s| m.forward_extend(ts, s), + )); + } + if let Err(e) = compare_with_stored(include_str!("cohere2_moe_reference.txt"), &platform, &got) { + panic!("{e}"); + } + let var = |k: &str| std::env::var(k).unwrap_or_else(|_| "unset".into()); + println!( + "{CHECKED} {platform}, JOULECLAW_THREADS={}, JOULECLAW_CHUNK_KB={}", + var("JOULECLAW_THREADS"), + var("JOULECLAW_CHUNK_KB") + ); + } + } } diff --git a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/cohere2_moe_reference.txt b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/cohere2_moe_reference.txt new file mode 100644 index 000000000..93652a62f --- /dev/null +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/cohere2_moe_reference.txt @@ -0,0 +1,101 @@ +# The stored reference for the Cohere2-MoE decoder (North Mini Code's +# architecture), checked without the checkpoint by the test +# cohere2_moe::tests::stored_reference::logits_match_the_stored_reference. +# +# A line: - . The checkpoints are that test's synthetic tiny, medium and +# wide configs; the scenario prefills a prompt, decodes, feeds a chunk into +# the live state and decodes twice more. +# +# Generated by examples/cohere2_reference.rs run at commit 30c4cf1ee (main +# before PR #649): its forward_prompt and forward_token, with and without the +# rayon feature (the lines were identical), on 2026-11-06. That file says how +# to add a platform. +# +# The bits depend on the platform's math library as well as on the code: +# x86_64 runs scalar kernels where arm64 runs NEON, and the two arm64 systems, +# with the same kernels, still differ on the medium and wide checkpoints. So +# each platform's lines name, on their `math-library` line, the library they +# were made with. On Linux arm64, glibc 2.35's libm (that of CI's container) +# loaded in place of 2.39's gives the same lines. +math-library macos-aarch64 the system math library (libsystem_m) of macOS 27.2 (26B123) +math-library linux-aarch64 the libm of Ubuntu GLIBC 2.39-0ubuntu8.9 +math-library linux-x86_64 the libm of Ubuntu GLIBC 2.43-2ubuntu2.4 +macos-aarch64 tiny prefill 6b3c5457df8a475b +macos-aarch64 tiny decode0 6b734576ec270b6a +macos-aarch64 tiny decode1 daa682bb61855607 +macos-aarch64 tiny decode2 a96a525073b3b793 +macos-aarch64 tiny decode3 9bb8aa8e5d425aba +macos-aarch64 tiny extend c3fda182453d886a +macos-aarch64 tiny after0 66a97258116e7597 +macos-aarch64 tiny after1 9d7b37bbb2325b5f +macos-aarch64 medium prefill b82cf3338dfa66f0 +macos-aarch64 medium decode0 c66f9c5495e8969b +macos-aarch64 medium decode1 8fe95e4149a712e2 +macos-aarch64 medium decode2 68f81c6d7139db90 +macos-aarch64 medium decode3 058a2083f6058017 +macos-aarch64 medium decode4 b5957c64fb853c5a +macos-aarch64 medium decode5 c90b8d2ab8fe7230 +macos-aarch64 medium extend 6448cf5f1255e448 +macos-aarch64 medium after0 d491c6a5b890724a +macos-aarch64 medium after1 0284c81caf0cbe33 +macos-aarch64 wide prefill 5756c9cbfe56f1ac +macos-aarch64 wide decode0 a61a0a894569af2a +macos-aarch64 wide decode1 ae4763855789a897 +macos-aarch64 wide decode2 86cda1dc39403823 +macos-aarch64 wide decode3 159cdb1461c22868 +macos-aarch64 wide extend 772006dffd4eae2f +macos-aarch64 wide after0 0dd7d76f2077f8c5 +macos-aarch64 wide after1 f43a082fc856a8d1 +linux-aarch64 tiny prefill 6b3c5457df8a475b +linux-aarch64 tiny decode0 6b734576ec270b6a +linux-aarch64 tiny decode1 daa682bb61855607 +linux-aarch64 tiny decode2 a96a525073b3b793 +linux-aarch64 tiny decode3 9bb8aa8e5d425aba +linux-aarch64 tiny extend c3fda182453d886a +linux-aarch64 tiny after0 66a97258116e7597 +linux-aarch64 tiny after1 9d7b37bbb2325b5f +linux-aarch64 medium prefill e73f43bb8f743349 +linux-aarch64 medium decode0 efd7dfbbbc45d7a6 +linux-aarch64 medium decode1 2367b4995fca5e82 +linux-aarch64 medium decode2 36e603a23c2e410d +linux-aarch64 medium decode3 aa3d37da0de5837b +linux-aarch64 medium decode4 830cf607dfafe891 +linux-aarch64 medium decode5 414613d72da1cdb5 +linux-aarch64 medium extend 0509a4b9f8d8925f +linux-aarch64 medium after0 750ebc98f4384676 +linux-aarch64 medium after1 9161ec76f213720f +linux-aarch64 wide prefill 39650607a70ba636 +linux-aarch64 wide decode0 3396985a061dbd5f +linux-aarch64 wide decode1 cab2298070b6e8a5 +linux-aarch64 wide decode2 c9bccfd56c20432f +linux-aarch64 wide decode3 9ffbfed9d2dc8cce +linux-aarch64 wide extend a017b0616c1771ab +linux-aarch64 wide after0 5929322a051cd849 +linux-aarch64 wide after1 2d1b2ab4c0a12842 +linux-x86_64 tiny prefill f256b8a75cf96d05 +linux-x86_64 tiny decode0 fb05b8bf618b395b +linux-x86_64 tiny decode1 0c7e7b75513c1a3b +linux-x86_64 tiny decode2 231bbcdb71db9c65 +linux-x86_64 tiny decode3 0df0412449659098 +linux-x86_64 tiny extend ae95b2ca568c2802 +linux-x86_64 tiny after0 93df2260f050c0e7 +linux-x86_64 tiny after1 3bf419607737244f +linux-x86_64 medium prefill 4652623cd80b24e1 +linux-x86_64 medium decode0 dc02242d91e91dbd +linux-x86_64 medium decode1 fcb0e3e903d187d2 +linux-x86_64 medium decode2 c1c4c18cbb9298f8 +linux-x86_64 medium decode3 370686ea9a04b164 +linux-x86_64 medium decode4 800ba78022c1d084 +linux-x86_64 medium decode5 4a6b5efb48a72994 +linux-x86_64 medium extend 4efe2d7fd73de655 +linux-x86_64 medium after0 d4553a9d2cef7e4e +linux-x86_64 medium after1 cf0729eb66a67a7f +linux-x86_64 wide prefill 3c21f84d4dd6d2e5 +linux-x86_64 wide decode0 c533e37794392f5d +linux-x86_64 wide decode1 929641e6f2ab1eaf +linux-x86_64 wide decode2 9677d2965fd35073 +linux-x86_64 wide decode3 15fb690a596a3419 +linux-x86_64 wide extend 6851feba7437339f +linux-x86_64 wide after0 a0c0eb86889d1563 +linux-x86_64 wide after1 a2eaa676e2afe1c5 diff --git a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/fabric_dispatch.rs b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/fabric_dispatch.rs index 55de7f7a7..53ad34ca6 100644 --- a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/fabric_dispatch.rs +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/fabric_dispatch.rs @@ -62,6 +62,22 @@ mod enabled { .as_ref() } + /// True when the operator opted in (`JOULECLAW_FABRIC=1`), so + /// decode projections may leave the CPU. + pub fn fabric_engaged() -> bool { + fabric().is_some() + } + + /// True when a batched (prefill) projection stays on the CPU's + /// NEON kernels: the fabric is off. Every route off the CPU + /// needs the fabric, the Metal-4 tensor path too + /// ([`matmul_bf16_mpp_bt`] refuses without it), so + /// `JOULECLAW_MPP_PREFILL` plays no part here: it only chooses + /// among the fabric's paths once the fabric is on. + pub fn prefill_on_cpu() -> bool { + fabric().is_none() + } + /// Widen a BF16 row-major block to f32 row-major with /// NEON (`vshll_n_u16` — the bf16→f32 decode IS a 16-bit /// shift), sequential loads AND stores. The scalar @@ -1334,7 +1350,7 @@ mod enabled { #[cfg(feature = "fabric")] pub use enabled::{ - ledger_reset, ledger_snapshot, matmul_bf16_amx, matmul_bf16_auto, matmul_bf16_auto_multi, + fabric_engaged, prefill_on_cpu, ledger_reset, ledger_snapshot, matmul_bf16_amx, matmul_bf16_auto, matmul_bf16_auto_multi, matmul_bf16_attn_block, matmul_bf16_gpu, matmul_bf16_mlp_window, matmul_bf16_moe_window, matmul_bf16_mpp_bt, matmul_bf16_norm_qkv_rope_window, matmul_bf16_norm_window, AutoJob, LedgerEntry, @@ -1345,6 +1361,19 @@ pub use enabled::{ #[cfg(feature = "fabric-meter")] pub use enabled::{tune_shape_live, tune_shape_live_process}; +/// Without the `fabric` feature nothing leaves the CPU. +#[cfg(not(feature = "fabric"))] +pub fn fabric_engaged() -> bool { + false +} + +/// Without the `fabric` feature every prefill projection runs on +/// the CPU's NEON kernels. +#[cfg(not(feature = "fabric"))] +pub fn prefill_on_cpu() -> bool { + true +} + /// One op in a fused dispatch window (no-fabric stub). #[cfg(not(feature = "fabric"))] pub struct AutoJob<'a> { @@ -1456,6 +1485,82 @@ pub fn matmul_bf16_auto( } } +/// Tests whose answer depends on the fabric switch +/// (`JOULECLAW_FABRIC`), which a process reads once. Such a test +/// runs where the switch is what it needs: in this process if it +/// already is, else alone in a child of the test binary. No test sets +/// the variable in its own process, so one test's need never decides +/// the switch for the others in the binary. +#[cfg(test)] +pub(crate) mod fabric_switch { + /// Set in a child that [`run_alone`] started. + const CHILD: &str = "JOULECLAW_TEST_CHILD"; + + /// True when the fabric's state in this process is `on`, so the + /// calling test goes on here. Otherwise the test runs alone in a + /// child with the switch set accordingly and must pass there, and + /// this returns false: the caller returns, its check made. + /// `module` is the caller's `module_path!()`, `test` its name. + pub(crate) fn here(on: bool, module: &str, test: &str) -> bool { + if super::fabric_engaged() == on { + return true; + } + run_alone(module, test, &[("JOULECLAW_FABRIC", on.then_some("1"))]); + false + } + + /// Whether this process is a child that [`run_alone`] started. + pub(crate) fn in_child() -> bool { + std::env::var_os(CHILD).is_some() + } + + /// Run the test `module::test` alone in a child of this test + /// binary with `env` set (`None`: removed), fail unless the child + /// reports that exactly that one test ran and passed, and return + /// what it printed. + pub(crate) fn run_alone(module: &str, test: &str, env: &[(&str, Option<&str>)]) -> String { + assert!(!in_child(), "{module}::{test}: a child test process would start another, with {env:?}"); + spawn(module, test, env, true) + } + + /// [`run_alone`], but the child is a top-level test process: it + /// may start children of its own. + #[cfg_attr(not(feature = "fabric"), allow(dead_code))] + pub(crate) fn run_as_top_level(module: &str, test: &str, env: &[(&str, Option<&str>)]) -> String { + spawn(module, test, env, false) + } + + fn spawn(module: &str, test: &str, env: &[(&str, Option<&str>)], child: bool) -> String { + // libtest names a test by its module path without the crate. + let path = match module.split_once("::") { + Some((_, m)) => format!("{m}::{test}"), + None => test.to_string(), + }; + let mut cmd = std::process::Command::new(std::env::current_exe().expect("the test binary's path")); + cmd.args([path.as_str(), "--exact", "--nocapture", "--test-threads=1"]); + if child { + cmd.env(CHILD, "1"); + } else { + cmd.env_remove(CHILD); + } + for &(k, v) in env { + match v { + Some(v) => cmd.env(k, v), + None => cmd.env_remove(k), + }; + } + let out = cmd.output().expect("start the test binary"); + let stdout = String::from_utf8_lossy(&out.stdout).into_owned(); + assert!( + out.status.success() && stdout.contains("test result: ok. 1 passed"), + "{path} in a child process with {env:?}: {}\n{stdout}\n{}", + out.status, + String::from_utf8_lossy(&out.stderr) + ); + stdout + } +} + #[cfg(all(test, feature = "fabric", target_os = "macos", target_arch = "aarch64"))] mod tests { @@ -1579,9 +1684,12 @@ mod tests { /// off / GPU uncertified / pre-Metal-4 (wrapper → false). #[test] fn mpp_bt_matches_neon_when_available() { - // SAFETY: single-threaded test setup, before any - // fabric() init in this test. - unsafe { std::env::set_var("JOULECLAW_FABRIC", "1") }; + // The switch is read once per process; setting it here + // would turn the fabric on for every other test in this + // binary that had not read it yet. + if !super::fabric_switch::here(true, module_path!(), "mpp_bt_matches_neon_when_available") { + return; + } // SHAPE SWEEP — square, non-square, and non-tile-aligned // M/N (the square-tiling wrapper must handle them all). let shapes = [ diff --git a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/lib.rs b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/lib.rs index 48ec3ae3d..752fe0bcb 100644 --- a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/lib.rs +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/lib.rs @@ -94,6 +94,8 @@ pub mod devstral_resident; pub mod hybrid_lm; pub mod falcon_h1; pub mod cohere2_moe; +#[cfg(feature = "rayon")] +pub mod spin_pool; pub mod mellum; pub mod fabric_dispatch; pub mod bpe_tokenizer; diff --git a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/matmul_bf16.rs b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/matmul_bf16.rs index 5fb483828..47d0bb914 100644 --- a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/matmul_bf16.rs +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/matmul_bf16.rs @@ -132,6 +132,26 @@ pub fn matmul_bf16_scalar( /// replaces eight shift-and-cast scalar ops, then two /// `vfmaq_f32` accumulate. /// +/// ## Every output is the one-row dot product, bit for bit +/// +/// Each `out[i, j]` is accumulated exactly as +/// [`matmul_bf16_neon_one_row`] does it: lane `l` of `acc0` +/// holds the fused multiply-adds of elements `8c + l` and lane +/// `l` of `acc1` those of `8c + 4 + l`, in chunk order; then +/// `vaddvq(acc0) + vaddvq(acc1)`, then the `k % 8` tail added +/// one product at a time, then the bias. What changed (2026-11) +/// is only how many of those dot products run at once: a block +/// of 4 activation rows × 2 weight rows, 2 × 4 (2 or 3 rows +/// left) or 1 × 4 (1 row left), with 1 weight row at a ragged +/// edge, shares its loads and keeps up to 16 independent +/// accumulator chains (8 outputs, two chains each) in flight. +/// The one-row loop has two chains, so it was bound by the +/// fused multiply-add's latency (about 16 GB/s of weight per +/// core on an M3 Ultra) rather than by memory. Weight rows are +/// walked in tiles of at most 64 KiB, so a tile stays in L1 +/// while every activation row passes over it, and each row is +/// prefetched 1 KiB ahead. +/// /// Falls back to [`matmul_bf16_scalar`] off aarch64. #[cfg(target_arch = "aarch64")] pub fn matmul_bf16_neon( @@ -140,6 +160,214 @@ pub fn matmul_bf16_neon( bias: &[f32], m: usize, n: usize, k: usize, out: &mut [f32], +) { + assert_eq!(a.len(), m * k); + assert_eq!(w.len(), 2 * n * k); + assert_eq!(out.len(), m * n); + assert!(bias.is_empty() || bias.len() == n); + // SAFETY: the asserts above bound every access the blocked + // kernel makes: rows `0..m` of `a` (`k` f32 each), rows + // `0..n` of `w` (`2k` bytes each) and `m * n` outputs at + // stride `n`. + unsafe { neon_blocked::matmul(a.as_ptr(), w.as_ptr(), bias, m, n, k, out.as_mut_ptr(), n) } +} + +/// The blocked BF16 kernel behind [`matmul_bf16_neon`] and the +/// threaded tiles: register blocks of activation rows × weight +/// rows whose every dot product keeps the one-row loop's exact +/// accumulation order (see [`matmul_bf16_neon`]). +#[cfg(target_arch = "aarch64")] +mod neon_blocked { + use super::bf16_to_f32; + use std::arch::aarch64::*; + + /// Software prefetch distance in each weight row. Measured + /// (2026-11, one core streaming 1 GB of BF16 from DRAM): with + /// no prefetch the 4-row block reads 14–15 GB/s on an M3 + /// Ultra and 19 on an M4 Max, against 16–17 and 26 at 1 KiB + /// ahead; 256 B to 4 KiB were swept. + const PREFETCH_BYTES: usize = 1024; + + /// The chunked (8-lane) part of `MR × NR` dot products: + /// activation rows `a + i*lda`, weight rows `w + j*ldw` + /// bytes, `chunks` 8-element chunks. Per output, lanes and + /// order are the one-row loop's. + #[inline(always)] + unsafe fn block( + a: *const f32, + lda: usize, + w: *const u8, + ldw: usize, + chunks: usize, + ) -> [[f32; NR]; MR] { + unsafe { + let mut lo = [[vdupq_n_f32(0.0); NR]; MR]; + let mut hi = [[vdupq_n_f32(0.0); NR]; MR]; + for c in 0..chunks { + let kk = c * 8; + if c % 4 == 0 { + // One prefetch per 64-byte line of each weight + // row, PREFETCH_BYTES ahead. A hint: it never + // faults, and past a row's end it lands in the + // next rows, which come next. + for j in 0..NR { + let p = w.wrapping_add(j * ldw + kk * 2 + PREFETCH_BYTES); + std::arch::asm!("prfm pldl1keep, [{0}]", in(reg) p, options(nostack, readonly, preserves_flags)); + } + } + let mut wl = [vdupq_n_f32(0.0); NR]; + let mut wh = [vdupq_n_f32(0.0); NR]; + for j in 0..NR { + let bits8: uint16x8_t = vld1q_u16(w.add(j * ldw + kk * 2) as *const u16); + wl[j] = vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(bits8), 16)); + wh[j] = vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(bits8), 16)); + } + for i in 0..MR { + let a_lo = vld1q_f32(a.add(i * lda + kk)); + let a_hi = vld1q_f32(a.add(i * lda + kk + 4)); + for j in 0..NR { + lo[i][j] = vfmaq_f32(lo[i][j], a_lo, wl[j]); + hi[i][j] = vfmaq_f32(hi[i][j], a_hi, wh[j]); + } + } + } + let mut r = [[0f32; NR]; MR]; + for i in 0..MR { + for j in 0..NR { + r[i][j] = vaddvq_f32(lo[i][j]) + vaddvq_f32(hi[i][j]); + } + } + r + } + } + + /// Finish one output the way the one-row loop does: the + /// scalar tail, one product at a time, then the bias. + #[inline(always)] + unsafe fn finish( + mut accum: f32, + a_row: *const f32, + w_row: *const u8, + k: usize, + bias: &[f32], + j: usize, + ) -> f32 { + unsafe { + for c in (k / 8) * 8..k { + let bits = u16::from_le_bytes([*w_row.add(c * 2), *w_row.add(c * 2 + 1)]); + accum += *a_row.add(c) * bf16_to_f32(bits); + } + } + if !bias.is_empty() { + accum += bias[j]; + } + accum + } + + /// One `MR × NR` block, finished and stored at + /// `out[i * ldo + j]`. + #[inline(always)] + unsafe fn run( + a: *const f32, + w: *const u8, + bias: &[f32], + k: usize, + i0: usize, + j0: usize, + out: *mut f32, + ldo: usize, + ) { + unsafe { + let a0 = a.add(i0 * k); + let w0 = w.add(j0 * 2 * k); + let r = block::(a0, k, w0, 2 * k, k / 8); + for i in 0..MR { + for j in 0..NR { + let v = finish(r[i][j], a0.add(i * k), w0.add(j * 2 * k), k, bias, j0 + j); + *out.add((i0 + i) * ldo + j0 + j) = v; + } + } + } + } + + /// Weight rows `0..n`, activation rows `0..m`; output + /// `out[i * ldo + j]`. `bias` is indexed by `j` (empty for + /// none). + /// + /// # Safety + /// `a` holds `m * k` f32, `w` holds `2 * n * k` bytes, and + /// `out` is writable at every `i * ldo + j`. + pub unsafe fn matmul( + a: *const f32, + w: *const u8, + bias: &[f32], + m: usize, + n: usize, + k: usize, + out: *mut f32, + ldo: usize, + ) { + unsafe { + // Weight rows per L1 tile: at most 64 KiB, a multiple + // of the 4-row block. + let tile = ((64 * 1024) / (2 * k).max(1)).clamp(4, 256) / 4 * 4; + let mut j0 = 0; + while j0 < n { + let jt = (n - j0).min(tile); + let mut i = 0; + while i < m { + let mr = (m - i).min(4); + let mut j = j0; + let jend = j0 + jt; + match mr { + 4 => { + while j + 2 <= jend { + run::<4, 2>(a, w, bias, k, i, j, out, ldo); + j += 2; + } + while j < jend { + run::<4, 1>(a, w, bias, k, i, j, out, ldo); + j += 1; + } + } + 2 | 3 => { + while j + 4 <= jend { + run::<2, 4>(a, w, bias, k, i, j, out, ldo); + j += 4; + } + while j < jend { + run::<2, 1>(a, w, bias, k, i, j, out, ldo); + j += 1; + } + } + _ => { + while j + 4 <= jend { + run::<1, 4>(a, w, bias, k, i, j, out, ldo); + j += 4; + } + while j < jend { + run::<1, 1>(a, w, bias, k, i, j, out, ldo); + j += 1; + } + } + } + i += if mr == 3 { 2 } else { mr }; + } + j0 += jt; + } + } + } +} + +/// The pre-2026-11 one-row NEON loop, kept as the reference the +/// blocked [`matmul_bf16_neon`] is tested against bit for bit. +#[cfg(target_arch = "aarch64")] +pub fn matmul_bf16_neon_one_row( + a: &[f32], + w: &[u8], + bias: &[f32], + m: usize, n: usize, k: usize, + out: &mut [f32], ) { assert_eq!(a.len(), m * k); assert_eq!(w.len(), 2 * n * k); @@ -205,6 +433,216 @@ pub fn matmul_bf16_neon( matmul_bf16_scalar(a, w, bias, m, n, k, out) } +/// Non-aarch64 fallback: delegates to the scalar kernel. +#[cfg(not(target_arch = "aarch64"))] +pub fn matmul_bf16_neon_one_row( + a: &[f32], + w: &[u8], + bias: &[f32], + m: usize, n: usize, k: usize, + out: &mut [f32], +) { + matmul_bf16_scalar(a, w, bias, m, n, k, out) +} + +/// `[m, n]` output tile `rows × cols` written in place at +/// stride `ldo`: the per-tile body of the threaded kernels. +/// +/// # Safety +/// `a` holds `m * k` f32 (only rows `i0..i0+rows` are read), +/// `w` holds the tile's `cols` weight rows (`2 * cols * k` +/// bytes), and `out` is writable at `(i0 + i) * ldo + j` for +/// every `i < rows`, `j < cols`, with no other writer. +#[cfg(feature = "rayon")] +#[allow(clippy::too_many_arguments)] +unsafe fn matmul_bf16_tile( + a: &[f32], + w: &[u8], + bias: &[f32], + i0: usize, + rows: usize, + cols: usize, + k: usize, + out: *mut f32, + ldo: usize, +) { + #[cfg(target_arch = "aarch64")] + unsafe { + neon_blocked::matmul(a.as_ptr().add(i0 * k), w.as_ptr(), bias, rows, cols, k, out.add(i0 * ldo), ldo); + } + #[cfg(not(target_arch = "aarch64"))] + { + let mut buf = vec![0f32; rows * cols]; + matmul_bf16_scalar(&a[i0 * k..(i0 + rows) * k], w, bias, rows, cols, k, &mut buf); + for i in 0..rows { + for j in 0..cols { + unsafe { *out.add((i0 + i) * ldo + j) = buf[i * cols + j] }; + } + } + } +} + +/// A raw output pointer handed to the tiles of one matmul; each +/// tile writes a disjoint `(row, column)` block. +#[cfg(feature = "rayon")] +#[derive(Clone, Copy)] +struct OutPtr(*mut f32); +// SAFETY: every tile of one call writes only its own block, and +// the call returns after all of them finish. +#[cfg(feature = "rayon")] +unsafe impl Send for OutPtr {} +#[cfg(feature = "rayon")] +unsafe impl Sync for OutPtr {} +#[cfg(feature = "rayon")] +impl OutPtr { + /// The pointer. A tile's closure reaches it through this + /// method, never through the field: naming the field would + /// capture the raw pointer alone, which is neither `Send` + /// nor `Sync`, where a method call captures the `OutPtr`. + fn get(self) -> *mut f32 { + self.0 + } +} + +/// The thread pool the CPU forward passes run in: its size is +/// `JOULECLAW_THREADS`, else `RAYON_NUM_THREADS`, else +/// [`default_cpu_threads`]. +#[cfg(feature = "rayon")] +pub fn cpu_pool() -> &'static rayon::ThreadPool { + static POOL: std::sync::OnceLock = std::sync::OnceLock::new(); + POOL.get_or_init(|| { + rayon::ThreadPoolBuilder::new() + .num_threads(cpu_threads()) + .thread_name(|i| format!("jouleclaw-cpu-{i}")) + .build() + .expect("the CPU thread pool") + }) +} + +/// The CPU passes' thread count: `JOULECLAW_THREADS`, else +/// `RAYON_NUM_THREADS`, else [`default_cpu_threads`]. +#[cfg(feature = "rayon")] +pub fn cpu_threads() -> usize { + let env = |k: &str| std::env::var(k).ok().and_then(|v| v.parse::().ok()).filter(|&n| n > 0); + env("JOULECLAW_THREADS").or_else(|| env("RAYON_NUM_THREADS")).unwrap_or_else(default_cpu_threads) +} + +/// The default size of [`cpu_pool`] and the decode +/// [`crate::spin_pool`]: one thread per performance core +/// (`hw.perflevel0.physicalcpu` on macOS, else every logical CPU), +/// at most 16. Measured on North (2026-11): a 4-row BF16 GEMV +/// streams 15–26 GB/s per core, so about a dozen threads already +/// meet what the CPU side can pull from memory; on an M3 Ultra (24 +/// performance cores) 16 to 24 threads decoded no faster, and with +/// rayon more threads prefilled slower. The E-cores are left out: +/// a fixed share on a slow core is what a fork-join waits for. +#[cfg(feature = "rayon")] +pub fn default_cpu_threads() -> usize { + performance_cores() + .or_else(|| std::thread::available_parallelism().ok().map(|n| n.get())) + .unwrap_or(1) + .min(16) +} + +#[cfg(all(feature = "rayon", target_os = "macos"))] +fn performance_cores() -> Option { + let mut v: libc::c_int = 0; + let mut len = std::mem::size_of::(); + // SAFETY: a NUL-terminated name, a c_int out-buffer and its + // size; no new value is set. + let rc = unsafe { + libc::sysctlbyname( + c"hw.perflevel0.physicalcpu".as_ptr(), + (&mut v as *mut libc::c_int).cast(), + &mut len, + std::ptr::null_mut(), + 0, + ) + }; + (rc == 0 && v > 0).then_some(v as usize) +} + +#[cfg(all(feature = "rayon", not(target_os = "macos")))] +fn performance_cores() -> Option { + None +} + +/// Run `f` on a thread of [`cpu_pool`], so the rayon loops inside +/// it start from a pool thread: the caller's thread does its share +/// of each loop instead of handing the loop to the pool and +/// sleeping until a pool thread wakes it, once per loop. The CPU +/// prefill runs in it. Without `rayon`, `f` runs here. +pub fn in_cpu_pool(f: impl FnOnce() -> R + Send) -> R { + #[cfg(feature = "rayon")] + { + if rayon::current_thread_index().is_some() { + return f(); + } + cpu_pool().install(f) + } + #[cfg(not(feature = "rayon"))] + f() +} + +/// Weight rows per decode (m = 1) work unit: about +/// `JOULECLAW_CHUNK_KB` (default 256) KiB of weight, a multiple +/// of the kernel's 4-row block. Small enough that a pool whose +/// threads are slow, late or descheduled (E-cores, a loaded +/// host) still balances: the threads that are running take the +/// next unit, and nobody waits on one thread's fixed share. +pub fn gemv_rows_per_chunk(k: usize) -> usize { + static KB: std::sync::OnceLock = std::sync::OnceLock::new(); + let kb = *KB.get_or_init(|| { + std::env::var("JOULECLAW_CHUNK_KB") + .ok() + .and_then(|v| v.parse::().ok()) + .filter(|&v| v > 0) + .unwrap_or(256) + }); + ((kb * 1024) / (2 * k).max(1)).max(4) / 4 * 4 +} + +/// One weight-row chunk of an m = 1 BF16 GEMV: `out = W · a` +/// for the chunk's rows. A decode step collects the chunks of +/// every independent projection into one list and runs them as +/// one parallel dispatch. Each output is [`matmul_bf16_neon`]'s +/// dot product, so how the rows are split never changes a bit. +pub struct GemvChunk<'a> { + pub a: &'a [f32], + pub w: &'a [u8], + pub k: usize, + pub out: &'a mut [f32], +} + +impl GemvChunk<'_> { + /// Compute this chunk's outputs (one thread). + pub fn run(&mut self) { + let n = self.out.len(); + matmul_bf16_neon(self.a, self.w, &[], 1, n, self.k, self.out); + } +} + +/// Split `out[n] = W[n, k] · a[k]` into [`GemvChunk`]s of +/// [`gemv_rows_per_chunk`] rows, appended to `into`. +pub fn push_gemv_chunks<'a>( + into: &mut Vec>, + a: &'a [f32], + w: &'a [u8], + n: usize, + k: usize, + out: &'a mut [f32], +) { + assert_eq!(a.len(), k); + assert_eq!(w.len(), 2 * n * k); + assert_eq!(out.len(), n); + let rows = gemv_rows_per_chunk(k); + for (c, o) in out.chunks_mut(rows).enumerate() { + let j0 = c * rows; + let nr = o.len(); + into.push(GemvChunk { a, w: &w[j0 * 2 * k..(j0 + nr) * 2 * k], k, out: o }); + } +} + /// BF16 matmul with rayon multi-thread fan-out — wave 25 /// dispatch shape: `m == 1` parallel over `n` (independent /// dot products), `m > 1` parallel over `m` (contiguous @@ -257,8 +695,12 @@ pub fn matmul_bf16_threaded( return Ok(()); } if m == 1 { + // 2026-11: many small chunks instead of one fixed share + // per thread (see [`gemv_rows_per_chunk`]), at least + // four per thread. Each output is the same dot product + // whatever the split. let threads = rayon::current_num_threads().max(1); - let chunk_n = n.div_ceil(threads).max(1); + let chunk_n = gemv_rows_per_chunk(k).min(n.div_ceil(4 * threads).div_ceil(4) * 4).max(4); out.par_chunks_mut(chunk_n).enumerate().for_each( |(idx, out_chunk)| { let start_j = idx * chunk_n; @@ -273,59 +715,31 @@ pub fn matmul_bf16_threaded( }, ); } else { - let threads = rayon::current_num_threads().max(1); - // Axis choice is a bandwidth question (wave 54). - // Parallel-over-m makes EVERY thread stream the whole - // [n,k] weight matrix — traffic = threads × weights, - // which is exactly wrong for the small-m prefill shape - // where weights dominate. Parallel-over-n reads the - // weights once in total but every thread reads all of - // A [m,k]. Crossover: m·k·4 (A per thread) vs - // n·k·2/threads (weights per thread) → prefer the n - // axis while m < n / (2·threads). - if m < n / (2 * threads).max(1) { - // Parallel over n. Each thread computes [m, chunk_n] - // into its own buffer (the kernel's natural output - // shape for a weight-row slice), then scatters rows - // into the strided columns of out — the scatter is - // O(m·n) f32 moves, noise next to the matmul. - let chunk_n = n.div_ceil(threads).max(1); - let pieces: Vec<(usize, usize, Vec)> = (0..n.div_ceil(chunk_n)) - .into_par_iter() - .map(|idx| { - let start_j = idx * chunk_n; - let this_n = chunk_n.min(n - start_j); - let w_chunk = - &w[start_j * row_bytes..(start_j + this_n) * row_bytes]; - let bias_chunk: &[f32] = if bias.is_empty() { - &[] - } else { - &bias[start_j..start_j + this_n] - }; - let mut buf = vec![0f32; m * this_n]; - matmul_bf16_neon(a, w_chunk, bias_chunk, m, this_n, k, &mut buf); - (start_j, this_n, buf) - }) - .collect(); - for (start_j, this_n, buf) in pieces { - for i in 0..m { - out[i * n + start_j..i * n + start_j + this_n] - .copy_from_slice(&buf[i * this_n..(i + 1) * this_n]); - } - } - } else { - // Large m: A traffic dominates — parallel over m, - // each thread a contiguous row chunk. - let chunk_m = m.div_ceil(threads).max(1); - out.par_chunks_mut(chunk_m * n).enumerate().for_each( - |(idx, out_chunk)| { - let start_i = idx * chunk_m; - let this_m = out_chunk.len() / n; - let a_chunk = &a[start_i * k..(start_i + this_m) * k]; - matmul_bf16_neon(a_chunk, w, bias, this_m, n, k, out_chunk); - }, - ); - } + // 2026-11: 2-D tiles. Each task owns a block of up to 64 + // activation rows × 64 weight rows and writes it in + // place, so the weights are read about m/64 times in + // total instead of once per activation row (the old + // parallel-over-m split streamed every weight per row, + // and parallel-over-n gave each thread all of A). Inside + // a tile the blocked kernel keeps every output's + // one-row accumulation order. + const TM: usize = 64; + let tn = 64usize; + let ptr = OutPtr(out.as_mut_ptr()); + let n_tiles = n.div_ceil(tn); + (0..m.div_ceil(TM) * n_tiles).into_par_iter().for_each(|t| { + let (bi, bj) = (t / n_tiles, t % n_tiles); + let i0 = bi * TM; + let rows = TM.min(m - i0); + let j0 = bj * tn; + let cols = tn.min(n - j0); + let w_tile = &w[j0 * row_bytes..(j0 + cols) * row_bytes]; + let b_tile: &[f32] = if bias.is_empty() { &[] } else { &bias[j0..j0 + cols] }; + // SAFETY: tile (bi, bj) writes rows i0..i0+rows of + // columns j0..j0+cols of `out` ([m, n], stride n), + // which no other tile touches; `a` is [m, k]. + unsafe { matmul_bf16_tile(a, w_tile, b_tile, i0, rows, cols, k, ptr.get().add(j0), n) }; + }); } Ok(()) } @@ -550,6 +964,80 @@ mod tests { } } + /// The blocked kernel's every output is the one-row loop's, + /// bit for bit: block shapes 1/2/3/4/5+ rows, weight rows + /// that do not fill a 4-row or 2-row block, `k` with and + /// without a `k % 8` tail, `k < 8`, with and without bias. + #[test] + fn bf16_blocked_is_bit_identical_to_one_row() { + for &m in &[1usize, 2, 3, 4, 5, 6, 7, 9] { + for &n in &[1usize, 2, 3, 4, 5, 7, 9, 17, 70] { + for &k in &[1usize, 7, 8, 9, 16, 203, 1030] { + let a = synth_f32(m * k, m + 3 * k); + let w = to_bf16_bytes(&synth_f32(n * k, 7 * n + k)); + for with_bias in [false, true] { + let bias: Vec = if with_bias { + (0..n).map(|j| j as f32 * 0.37 - 1.0).collect() + } else { + Vec::new() + }; + let mut want = vec![0f32; m * n]; + let mut got = vec![0f32; m * n]; + matmul_bf16_neon_one_row(&a, &w, &bias, m, n, k, &mut want); + matmul_bf16_neon(&a, &w, &bias, m, n, k, &mut got); + for (x, y) in want.iter().zip(got.iter()) { + assert_eq!(x.to_bits(), y.to_bits(), "m={m} n={n} k={k} bias={with_bias}"); + } + } + } + } + } + } + + /// Row chunks of a GEMV (the decode dispatch's units) give + /// the one-row loop's outputs, whatever the chunk size. + #[test] + fn gemv_chunks_are_bit_identical_to_one_row() { + let (n, k) = (300usize, 2048usize); + let a = synth_f32(k, 11); + let w = to_bf16_bytes(&synth_f32(n * k, 5)); + let mut want = vec![0f32; n]; + matmul_bf16_neon_one_row(&a, &w, &[], 1, n, k, &mut want); + let mut got = vec![0f32; n]; + let mut chunks = Vec::new(); + push_gemv_chunks(&mut chunks, &a, &w, n, k, &mut got); + assert!(chunks.len() > 1, "the test must split the rows"); + for c in chunks.iter_mut() { + c.run(); + } + drop(chunks); + for (x, y) in want.iter().zip(got.iter()) { + assert_eq!(x.to_bits(), y.to_bits()); + } + } + + /// The threaded paths above the dispatch floor (16 MiB of + /// weight × m): the m = 1 row chunks and the 2-D tiles + /// (several row tiles and a ragged column tile), against the + /// one-row loop, bit for bit. + #[cfg(feature = "rayon")] + #[test] + fn bf16_threaded_above_the_floor_is_bit_identical() { + for (m, n, k) in [(1usize, 4100usize, 2048usize), (70, 130, 1030), (3, 2100, 2048)] { + assert!(m * 2 * n * k >= 16 * 1024 * 1024, "shape must clear the floor"); + let a = synth_f32(m * k, 13); + let w = to_bf16_bytes(&synth_f32(n * k, 17)); + let bias: Vec = (0..n).map(|j| (j % 5) as f32 * 0.5).collect(); + let mut want = vec![0f32; m * n]; + let mut got = vec![0f32; m * n]; + matmul_bf16_neon_one_row(&a, &w, &bias, m, n, k, &mut want); + matmul_bf16_threaded(&a, &w, &bias, m, n, k, &mut got).expect("rayon feature on"); + for (x, y) in want.iter().zip(got.iter()) { + assert_eq!(x.to_bits(), y.to_bits(), "m={m} n={n} k={k}"); + } + } + } + #[test] fn bulk_dequant_parity() { let vals = synth_f32(257, 5); // odd length diff --git a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/spin_pool.rs b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/spin_pool.rs new file mode 100644 index 000000000..96ff0195c --- /dev/null +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/spin_pool.rs @@ -0,0 +1,481 @@ +//! A fork-join pool for decode-sized parallel loops: threads that +//! stay awake between loops. +//! +//! A decode step of a mixture-of-experts model runs about 150 +//! small parallel loops one after another (each layer's q/k/v, +//! attention + experts, o_proj + experts), each a fraction of a +//! millisecond. A pool whose threads go to sleep between loops pays +//! an OS wake-up per thread per loop, and on a busy host a woken +//! thread can wait a long time for a core. Here the threads spin +//! (`spin_loop`, no syscall) for a short while after a loop ends, +//! so the next loop finds them running, and park only when no loop +//! comes. The caller's thread takes units too. +//! +//! Units are claimed one at a time from a shared counter, so a slow +//! or descheduled thread holds at most the unit it is running. +//! Which thread runs which unit never changes a result: every unit +//! writes only its own outputs. + +use std::any::Any; +use std::sync::atomic::{AtomicBool, AtomicPtr, AtomicU64, AtomicUsize, Ordering}; +use std::sync::{Arc, Condvar, Mutex, OnceLock, PoisonError}; +use std::time::{Duration, Instant}; + +/// How long an idle thread spins before it parks. +const SPIN: Duration = Duration::from_micros(300); + +/// Bits of the claim word that hold the next unit's index; the +/// epoch has the other 40. +const IDX_BITS: u32 = 24; + +/// The unit index of a finished loop's claim word, and the most units +/// one loop may have (`run` asserts it): no loop has more, so a thread +/// that still holds a finished loop's epoch fails the bound and claims +/// nothing. 2^24 - 1 is far above any loop the decoder makes (a +/// 262144-row lm_head in 4-row chunks is 65536 units; prefill +/// attention for a 500k-token prompt about 1M). +const RETIRED: u64 = (1 << IDX_BITS) - 1; + +/// The epoch's part of a claim word. +fn epoch_of(word: u64) -> u64 { + word >> IDX_BITS +} + +/// One loop: run `f(i)` for `i < len`. +struct Job<'a> { + f: &'a (dyn Fn(usize) + Sync), +} + +struct Shared { + /// `(epoch << IDX_BITS) | next unit` of the current loop, or + /// `(epoch << IDX_BITS) | RETIRED` once a new loop is being + /// published. + claim: AtomicU64, + /// The current loop's unit count. It carries no epoch, so a + /// thread that read an older loop's claim word may read a newer + /// count; `run` retires the old word before it stores the count, + /// which makes that thread's claim fail (see `run`). + len: AtomicUsize, + /// The current loop's job, valid from the claim of one of its + /// units until that unit is counted done. + job: AtomicPtr>, + /// Units finished in the current loop. + done: AtomicUsize, + /// A unit panicked: its payload is in `payloads`. + panicked: AtomicBool, + /// The payloads of the current loop's panicked units, in the order + /// they were caught. The caller re-raises the first and drops the + /// rest once the loop has finished and the pool is released, so no + /// payload's `Drop` runs while another thread may be in the loop. + payloads: Mutex>>, + /// Parked threads wait here for a new epoch. + sleep: Mutex, + wake: Condvar, +} + +/// The pool. Built once per process ([`spin_pool`]). +pub struct SpinPool { + shared: Arc, + threads: usize, + /// One loop at a time. + busy: Mutex<()>, + /// Tests only: run inside [`SpinPool::run`] after the new loop's + /// unit count and job are published and before its epoch is. + #[cfg(test)] + window: Option, +} + +/// See [`SpinPool::window`]. +#[cfg(test)] +type Window = Box; + +impl Shared { + /// Claim and run units of loop `epoch` until none are left. + /// Returns how many this thread ran. + fn work(&self, epoch: u64) -> usize { + let mut ran = 0; + loop { + let cur = self.claim.load(Ordering::Acquire); + if epoch_of(cur) != epoch { + return ran; + } + let idx = (cur & RETIRED) as usize; + // If this is a newer loop's count, `cur` is out of date: + // the count was stored (Release) after the word was + // retired, so the exchange below sees the retired word + // and fails. + if idx >= self.len.load(Ordering::Acquire) { + return ran; + } + if self + .claim + .compare_exchange_weak(cur, cur + 1, Ordering::AcqRel, Ordering::Relaxed) + .is_err() + { + continue; + } + // Unit `idx` of `epoch` is ours alone. The claim + // succeeded while `epoch` was unfinished, so the loop that + // owns the job cannot return (and drop it) before we count + // this unit done, and no newer loop has replaced it. + // + // One theoretical exception: the exchange compares the + // whole word, so a thread descheduled between loading `cur` + // and the exchange for exactly 2^40 - 1 loops (the epoch's + // period; about 7.7 years of loops at 4500 a second, the + // decoder's rate at 30 tokens a second) would find the same + // word again, in a different loop. With 32 epoch bits that + // was about 11 days of uninterrupted decoding. + let job = self.job.load(Ordering::Acquire); + // SAFETY: as above; the job was published (Release) before + // `epoch`'s claim counter, which we read with Acquire. + let f = unsafe { (*job).f }; + if let Err(payload) = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| f(idx))) { + // Kept, not dropped: dropping it would run its `Drop` + // here, inside the loop. + self.payloads.lock().unwrap_or_else(PoisonError::into_inner).push(payload); + self.panicked.store(true, Ordering::Relaxed); + } + ran += 1; + self.done.fetch_add(1, Ordering::Release); + } + } +} + +/// The epoch after `e`: 40 bits, wrapping, never 0 (a new thread +/// has seen epoch 0). +fn next_epoch(e: u64) -> u64 { + match (e + 1) & ((1 << (64 - IDX_BITS)) - 1) { + 0 => 1, + n => n, + } +} + +fn worker(shared: Arc) { + let mut seen = 0u64; + loop { + let start = Instant::now(); + let mut epoch; + loop { + epoch = epoch_of(shared.claim.load(Ordering::Acquire)); + if epoch != seen || start.elapsed() > SPIN { + break; + } + for _ in 0..64 { + std::hint::spin_loop(); + } + } + if epoch == seen { + let mut s = shared.sleep.lock().unwrap(); + while *s == seen { + s = shared.wake.wait(s).unwrap(); + } + epoch = *s; + } + seen = epoch; + shared.work(epoch); + } +} + +impl SpinPool { + /// A pool of `threads` threads, the caller's included (so + /// `threads - 1` are spawned). They live for the process. + pub fn new(threads: usize) -> Self { + let pool = Self::unstarted(threads); + for i in 1..pool.threads { + let s = pool.shared.clone(); + std::thread::Builder::new() + .name(format!("jouleclaw-spin-{i}")) + .spawn(move || worker(s)) + .expect("spawn a pool thread"); + } + pool + } + + /// The pool's state for `threads` threads, none of them started. + fn unstarted(threads: usize) -> Self { + let shared = Arc::new(Shared { + claim: AtomicU64::new(0), + len: AtomicUsize::new(0), + job: AtomicPtr::new(std::ptr::null_mut()), + done: AtomicUsize::new(0), + panicked: AtomicBool::new(false), + payloads: Mutex::new(Vec::new()), + sleep: Mutex::new(0), + wake: Condvar::new(), + }); + SpinPool { + shared, + threads: threads.max(1), + busy: Mutex::new(()), + #[cfg(test)] + window: None, + } + } + + /// Threads that run units, the caller's included. + pub fn threads(&self) -> usize { + self.threads + } + + /// Run `f(i)` for every `i < len` on the pool and the calling + /// thread; returns when all have run. A unit must not call `run` + /// on the same pool: the loop holds the pool until every unit has + /// run, so the inner call would wait forever. + pub fn run(&self, len: usize, f: &(dyn Fn(usize) + Sync)) { + if len == 0 { + return; + } + if self.threads == 1 || len == 1 { + for i in 0..len { + f(i); + } + return; + } + assert!(len as u64 <= RETIRED, "too many units in one loop"); + // `busy` cannot be poisoned: nothing panics while it is held. + // From here to `drop(one)` the code is atomics; the `sleep` + // lock, never held across a panic (its holders only set or + // wait on the epoch); `Shared::work`, which catches a unit's + // panic and keeps its payload instead of dropping it; moving + // the payloads out of their list, whose lock is held only to + // push or take; and, in tests only, the `window` hook. A unit's panic is + // re-raised after the guard is released. Were the lock ever + // poisoned anyway, it would be by a panic in the test hook, + // before this loop's epoch is published, when no other thread + // can reach this loop's job: the pool's state is still that + // of the last finished loop, so the lock is taken as it is. + let one = self.busy.lock().unwrap_or_else(PoisonError::into_inner); + let sh = &*self.shared; + let job = Job { f }; + // SAFETY: the job lives until this function returns, which is + // after every unit has been counted done; no thread touches it + // after its unit (see `Shared::work`). The lifetime is erased + // only to store the pointer. + let ptr = &job as *const Job<'_> as *mut Job<'static>; + // Publication, in this order. A thread can still be inside + // `work` for the previous loop, holding its epoch: woken late + // from park, or descheduled after its last unit. That loop's + // claim word reads (epoch, its unit count), so were this + // loop's count visible first and larger, the thread would pass + // the bound, claim a unit the old loop does not have, and run + // a unit of this loop that is then claimed again, counted + // twice, and lets this call return while it still runs. + // 1. Retire the old word: its index passes no bound. + // 2. The count, Release: a thread that reads it then sees the + // retired word in its exchange, which fails. + // 3. The job, then the new epoch, Release: a thread that reads + // the epoch sees this loop's count and job. + let old = epoch_of(sh.claim.load(Ordering::Relaxed)); + sh.claim.store((old << IDX_BITS) | RETIRED, Ordering::Relaxed); + sh.done.store(0, Ordering::Relaxed); + sh.len.store(len, Ordering::Release); + sh.job.store(ptr, Ordering::Release); + #[cfg(test)] + if let Some(w) = &self.window { + w(sh); + } + let epoch = next_epoch(old); + sh.claim.store(epoch << IDX_BITS, Ordering::Release); + { + let mut s = sh.sleep.lock().unwrap(); + *s = epoch; + sh.wake.notify_all(); + } + sh.work(epoch); + while sh.done.load(Ordering::Acquire) < len { + std::hint::spin_loop(); + } + let payloads = if sh.panicked.swap(false, Ordering::Relaxed) { + std::mem::take(&mut *sh.payloads.lock().unwrap_or_else(PoisonError::into_inner)) + } else { + Vec::new() + }; + drop(one); + // The loop has finished and the pool is free: payloads may run + // their `Drop` now. The first unit to panic is re-raised as it + // was, payload and all; the others are dropped. + let mut payloads = payloads.into_iter(); + if let Some(first) = payloads.next() { + drop(payloads); + std::panic::resume_unwind(first); + } + } +} + +/// The process's decode pool, sized like +/// [`crate::matmul_bf16::cpu_pool`]. +pub fn spin_pool() -> &'static SpinPool { + static POOL: OnceLock = OnceLock::new(); + POOL.get_or_init(|| SpinPool::new(crate::matmul_bf16::cpu_threads())) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn every_unit_runs_exactly_once_across_many_loops() { + let pool = SpinPool::new(6); + for len in [1usize, 2, 5, 64, 1000] { + for _ in 0..50 { + let hits: Vec = (0..len).map(|_| AtomicUsize::new(0)).collect(); + pool.run(len, &|i| { + hits[i].fetch_add(1, Ordering::Relaxed); + }); + assert!(hits.iter().all(|h| h.load(Ordering::Relaxed) == 1), "len {len}"); + } + } + } + + #[test] + fn the_epoch_wraps_without_reaching_zero() { + assert_eq!(next_epoch(1), 2); + assert_eq!(next_epoch((1 << 40) - 2), (1 << 40) - 1); + assert_eq!(next_epoch((1 << 40) - 1), 1); + // A pool whose counter is just below the wrap keeps running. + let pool = SpinPool::new(4); + pool.shared.claim.store(((1 << 40) - 3) << IDX_BITS, Ordering::Relaxed); + let n = AtomicUsize::new(0); + for _ in 0..6 { + pool.run(50, &|_| { + n.fetch_add(1, Ordering::Relaxed); + }); + } + assert_eq!(n.load(Ordering::Relaxed), 300); + } + + #[test] + fn loops_after_the_threads_parked_still_run() { + let pool = SpinPool::new(4); + let n = AtomicUsize::new(0); + pool.run(100, &|_| { + n.fetch_add(1, Ordering::Relaxed); + }); + std::thread::sleep(SPIN * 4); + pool.run(100, &|_| { + n.fetch_add(1, Ordering::Relaxed); + }); + assert_eq!(n.load(Ordering::Relaxed), 200); + } + + /// A thread still inside `work` for a loop that has finished (a + /// worker woken late from park holds that loop's epoch, and so + /// does one descheduled after its last unit) claims nothing of + /// the next loop, also when the next loop has more units. + /// Deterministic: the pool has no threads of its own, and the + /// late thread runs in the window in which the new loop's unit + /// count and job are published and its epoch is not. + #[test] + fn a_late_thread_of_a_finished_loop_claims_nothing_of_the_next() { + let late = Arc::new(AtomicUsize::new(0)); + let armed = Arc::new(AtomicBool::new(false)); + let mut pool = SpinPool::unstarted(2); + let (l, a) = (late.clone(), armed.clone()); + pool.window = Some(Box::new(move |sh: &Shared| { + if a.load(Ordering::Relaxed) { + let finished = epoch_of(sh.claim.load(Ordering::Relaxed)); + let ran = std::thread::scope(|s| s.spawn(|| sh.work(finished)).join().unwrap()); + l.fetch_add(ran, Ordering::Relaxed); + } + })); + let hits: Vec = (0..8).map(|_| AtomicUsize::new(0)).collect(); + let f = |i: usize| { + hits[i].fetch_add(1, Ordering::Relaxed); + }; + pool.run(3, &f); + armed.store(true, Ordering::Relaxed); + pool.run(8, &f); + let counts: Vec = hits.iter().map(|h| h.load(Ordering::Relaxed)).collect(); + assert_eq!(late.load(Ordering::Relaxed), 0, "the finished loop's late thread ran units of the next: {counts:?}"); + assert_eq!(counts, [2, 2, 2, 1, 1, 1, 1, 1]); + } + + /// The same with the pool's own threads, woken late from park. + /// Each round lets the threads park, then runs a 2-unit loop, + /// which the caller finishes before the threads it woke are + /// running, and at once a 64-unit loop whose publication is held + /// open for 200 µs, about what a descheduled caller would hold + /// it. Every unit of every loop runs exactly once. + #[test] + fn threads_woken_late_run_no_unit_twice() { + // Static, so that a unit still running when a broken pool + // returns early has something to write to. + static HITS: [AtomicUsize; 64] = [const { AtomicUsize::new(0) }; 64]; + fn hit(i: usize) { + HITS[i].fetch_add(1, Ordering::Relaxed); + } + let mut pool = SpinPool::new(6); + pool.window = Some(Box::new(|_: &Shared| { + let t = Instant::now(); + while t.elapsed() < Duration::from_micros(200) { + std::hint::spin_loop(); + } + })); + for round in 0..100 { + std::thread::sleep(SPIN * 3); + for len in [2usize, 64] { + pool.run(len, &hit); + let counts: Vec = HITS[..len].iter().map(|h| h.swap(0, Ordering::Relaxed)).collect(); + assert!(counts.iter().all(|&c| c == 1), "round {round}, loop of {len}: {counts:?}"); + } + } + } + + /// The caller gets the unit's own panic, its payload, not a + /// message of the pool's. + #[test] + fn a_units_panic_reaches_the_caller_with_its_payload() { + #[derive(Debug, PartialEq)] + struct UnitFailed(usize); + let pool = SpinPool::new(4); + let failed = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + pool.run(16, &|i| { + if i == 5 { + std::panic::panic_any(UnitFailed(5)); + } + }); + })); + let payload = failed.expect_err("the unit's panic reaches the caller"); + assert_eq!(payload.downcast_ref::(), Some(&UnitFailed(5))); + } + + /// The epoch outlasts 32 bits, so a thread would have to sleep + /// between reading the claim word and claiming through 2^40 - 1 + /// loops for a stale word to match again; the unit index has 24 + /// bits, and a loop with more units than that is refused before it + /// takes the pool. + #[test] + fn the_epoch_has_40_bits_and_a_loop_at_most_2_pow_24_minus_1_units() { + assert_eq!(next_epoch(0xFFFF_FFFF), 0x1_0000_0000); + assert_eq!(next_epoch((1 << 40) - 2), (1 << 40) - 1); + assert_eq!(next_epoch((1 << 40) - 1), 1); + let pool = SpinPool::new(4); + let refused = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| pool.run(1 << 24, &|_| {}))); + assert!(refused.is_err(), "a loop of 2^24 units is refused"); + let n = AtomicUsize::new(0); + pool.run(50, &|_| { + n.fetch_add(1, Ordering::Relaxed); + }); + assert_eq!(n.load(Ordering::Relaxed), 50, "and the pool runs the next loop"); + } + + /// A unit that panics fails its own loop and nothing after it: + /// the panic reaches the caller, and the pool runs every later + /// loop. + #[test] + fn a_panicking_unit_leaves_the_pool_usable() { + let pool = SpinPool::new(4); + let failed = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + pool.run(16, &|i| assert_ne!(i, 5, "unit 5 fails")); + })); + assert!(failed.is_err(), "the unit's panic reaches the caller"); + let n = AtomicUsize::new(0); + for _ in 0..3 { + pool.run(16, &|_| { + n.fetch_add(1, Ordering::Relaxed); + }); + } + assert_eq!(n.load(Ordering::Relaxed), 48); + } +}