From 6db018cd04fda93fff18c577ffcf2668ced72e82 Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Thu, 5 Nov 2026 07:45:42 -0500 Subject: [PATCH 01/14] jouleclaw: a per-op profile and a bit-identity bench for North decode JOULECLAW_PROF=1 times every op of the Cohere2-MoE forward passes (rmsnorm, q/k/v, rope and the KV append, attention, o_proj, router, the expert stages, the residual, lm_head). It reads clocks only. north_bench runs one prompt from a file and prints, per decoded token, its id and a hash of every logit's bits, so two builds can be compared bit for bit, with the median ms per token and the profile. Co-Authored-By: Claude Opus 5.5 --- .../examples/north_bench.rs | 134 ++++++++++++++++++ .../jouleclaw-loader-gguf/src/cohere2_moe.rs | 118 +++++++++++++++ 2 files changed, 252 insertions(+) create mode 100644 jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/examples/north_bench.rs 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..9d40b3bdb --- /dev/null +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/examples/north_bench.rs @@ -0,0 +1,134 @@ +//! 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. 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 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)); + Ok(()) +} 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..f02c84b87 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,93 @@ 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; + const N: usize = 15; + 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", + ]; + + static NS: [AtomicU64; N] = [const { AtomicU64::new(0) }; N]; + + 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 { + 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 // ───────────────────────────────────────────────────────────────────── @@ -400,12 +487,15 @@ impl Cohere2MoeModel { pub fn forward_token(&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 +580,7 @@ impl Cohere2MoeModel { } } attn_out.copy_from_slice(&o_out); + lap.mark(prof::ATTN_BLOCK_GPU); break 'attn; } lstate.gpu_dead = true; @@ -501,6 +592,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 +609,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 +625,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 +644,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 +656,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 +678,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 +706,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 +720,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 +732,7 @@ impl Cohere2MoeModel { *l *= cfg.logit_scale; } } + lap.mark(prof::LM_HEAD); state.position += 1; logits } @@ -648,15 +751,18 @@ impl Cohere2MoeModel { 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 +776,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 +788,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 +833,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 +851,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 +871,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 +887,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 +911,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,6 +931,7 @@ impl Cohere2MoeModel { *l *= cfg.logit_scale; } } + lap.mark(prof::LM_HEAD); state.position = t_len; logits } -- 2.54.0 (Apple Git-157) From 9d928f493ed4c08c6afa5b66085326656d5ae8ee Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Thu, 5 Nov 2026 08:47:44 -0500 Subject: [PATCH 02/14] jouleclaw: a blocked BF16 kernel that keeps every dot product's order matmul_bf16_neon computed one weight row at a time on two accumulator chains, so a core waited on the fused multiply-add's latency: about 16-20 GB/s of weight per core, in cache or not. It now computes a block of up to 4 activation rows x 4 weight rows at once (up to 16 chains, shared loads), walks weight rows in 64 KiB L1 tiles, and prefetches each row 1 KiB ahead. Every output is still accumulated exactly as the one-row loop does it (the same lanes, the same chunk order, the same final reduction, tail and bias), so results do not change by a bit; the old loop stays as matmul_bf16_neon_one_row and the tests compare the two bit for bit over block shapes, tails and biases. One core, measured: 61 against 20 GB/s on cache-resident weights (M4 Max), and from DRAM 26 against 19 (M4 Max) and 16-17 against 14-15 without the prefetch (M3 Ultra). The threaded kernel splits an m = 1 GEMV into small row chunks (about 256 KiB, at least four per thread) instead of one fixed share per thread, so a slow or descheduled thread no longer holds up the rest, and splits m > 1 into 64 x 64 output tiles written in place, so a prefill reads each weight about m/64 times instead of once per activation row. GemvChunk / push_gemv_chunks expose the chunks, and cpu_pool / in_cpu_pool give the CPU passes a pool sized by JOULECLAW_THREADS (else RAYON_NUM_THREADS, else one thread per performance core, at most 16). Co-Authored-By: Claude Opus 5.5 --- .../jouleclaw-loader-gguf/src/matmul_bf16.rs | 585 ++++++++++++++++-- 1 file changed, 531 insertions(+), 54 deletions(-) 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..9994a162f 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,24 @@ 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 up to 4 activation rows × 4 weight rows shares its loads +/// and keeps up to 16 independent accumulator chains 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 +158,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 +431,206 @@ 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 {} + +/// 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 +683,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 +703,32 @@ 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 ptr = ptr; + 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.0.add(j0), n) }; + }); } Ok(()) } @@ -550,6 +953,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 -- 2.54.0 (Apple Git-157) From 3fed234fd238351db2978c4680a8c93316f1fc0e Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Thu, 5 Nov 2026 08:47:44 -0500 Subject: [PATCH 03/14] jouleclaw: one attention row, bit-identical to attention_forward_f32, fast attention_row_f32 computes one query row against keys 0..kv_len with attention_forward_f32's exact arithmetic. The reference adds each score's products on one scalar chain (an f32 sum may not be reordered), so it waits on the add's latency for every product. Here four keys share a NEON register (a 4x4 transpose of four key rows) and four registers run at once, each lane still multiplying, rounding, adding and rounding in d order; the weighted sum of V keeps its running sums in registers across keys, multiply then add in key order, skipping zero weights as the reference does. The test compares it with the reference bit for bit for decode and causal rows, head widths that are and are not multiples of 4, key counts around the blocks, and weights that underflow to zero. Co-Authored-By: Claude Opus 5.5 --- .../jouleclaw-loader-gguf/src/attention.rs | 251 ++++++++++++++++++ 1 file changed, 251 insertions(+) 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..596c84b62 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 % 4 == 0 { 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 % 4 == 0 { + 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 -- 2.54.0 (Apple Git-157) From 624b8389df774bdec552b6a56f6255bb86735379 Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Thu, 5 Nov 2026 10:17:04 -0500 Subject: [PATCH 04/14] jouleclaw: North decode 3-6x and prefill 6-11x on the CPU, bit-identical The profile (JOULECLAW_PROF=1, the M3 Ultra build host under load, a ~600-token prompt) showed North Mini Code's CPU path waiting, not streaming: decode spent 212 ms/token in ~7 dispatches per layer, three of them single-threaded GEMVs, with all 32 attention heads on one thread; its 16-50 MB projections moved 12-40 GB/s while the 1 GB lm_head, one dispatch, moved ~140 GB/s. Prefill (13 tok/s) spent 36% in attention on one thread and 50% running the experts one after another. Base decode was also fastest at 8 threads and slowest at the default 32. Cohere2MoeModel's forward passes now have CPU versions, used when nothing may leave the CPU (always without the fabric feature): - decode: three dispatches per layer (q/k/v + router; attention heads + expert gate/up; o_proj + expert down) of ~256 KiB row chunks, on spin_pool, a pool whose threads spin briefly between dispatches instead of sleeping and claim units one at a time; - prefill: attention per (head, 16 query rows) in parallel, the experts in parallel, projections as 64 x 64 tiles; - RoPE angles once per position; - forward_extend: batched feeding of a state that already holds a context, equal to one forward_token per token (the program loop's repair rounds feed their turns that way today); - threads: one per performance core, at most 16 (JOULECLAW_THREADS overrides). Every output keeps its reduction order, so nothing changes by a bit. The previous passes stay as forward_*_reference (the fabric path), and tests compare the CPU passes with them bit for bit, prefill, decode and extend, across the sliding window. On the real checkpoint, five prompts (6 to 2698 tokens) x 72 greedy tokens gave base's token ids and the same hash of every step's logits, on the M3 Ultra and on an M4 Max. Medians, base vs new, interleaved (M3 Ultra shared with another North benchmark and CI, load 30-130 / M4 Max lightly loaded): decode at ~600 context 267 -> 54 ms / 116 -> 35 ms; at ~2700 context 405 -> 80 ms / 246 -> 48 ms; prefill of 596 tokens 9.9 -> 91 / 17.8 -> 105 tok/s, of 2698 tokens 5.9 -> 63 / 9.0 -> 79 tok/s. BENCHMARKS.md has every run and the paired ratios. north_bench gains --then / --then-seq to time and check forward_extend. Co-Authored-By: Claude Opus 5.5 --- .../jouleclaw-loader-gguf/BENCHMARKS.md | 84 ++ .../examples/north_bench.rs | 54 +- .../jouleclaw-loader-gguf/src/cohere2_moe.rs | 746 +++++++++++++++++- .../src/fabric_dispatch.rs | 30 +- .../crates/jouleclaw-loader-gguf/src/lib.rs | 2 + .../jouleclaw-loader-gguf/src/spin_pool.rs | 255 ++++++ 6 files changed, 1167 insertions(+), 4 deletions(-) create mode 100644 jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/spin_pool.rs diff --git a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/BENCHMARKS.md b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/BENCHMARKS.md index 8efb073e9..7c0bc4b58 100644 --- a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/BENCHMARKS.md +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/BENCHMARKS.md @@ -1439,6 +1439,90 @@ 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: 3–6× decode, 6–11× prefill, 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. + +**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 up to 4 activation rows × 4 + weight rows (16 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. + +**Speed.** Medians of 4 runs (3 for base at 12 threads), base and +new interleaved in the same conditions. 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 | M3 Ultra base, 12 thr | M3 Ultra new | M4 Max base | M4 Max new | +|---|---|---|---|---|---| +| decode at ~600 context | 267 ms | 168 ms | **54 ms** | 116 ms | **35 ms** | +| decode at ~2700 context | 405 ms | 364 ms | **80 ms** | 246 ms | **48 ms** | +| prefill, 596 tokens | 9.9 tok/s | 12.0 tok/s | **91 tok/s** | 17.8 tok/s | **105 tok/s** | +| prefill, 2698 tokens | 5.9 tok/s | 4.5 tok/s | **63 tok/s** | 9.0 tok/s | **79 tok/s** | + +Paired run-for-run ratios (base / new): M3 Ultra decode 4.3–4.8× +at ~600 and 5.0–6.6× at ~2700 (one base run at load 130–131 was +20×, left out), prefill 5.4–10.3× and 9.8–11.8×; M4 Max decode +3.3–3.5× and 5.1–5.3×, prefill 5.2–6.8× and 7.0–9.0×. + +A 69-token follow-up turn fed into a ~630-token context: +`forward_extend` 1.23 s (M3 Ultra) / 0.79 s (M4 Max), against +3.3–4.1 s / 2.4 s one `forward_token` at a time on the new code +and ~18 s at base's decode speed; same logits. + +**Did not help, or not enough:** more threads (base and new were +both slower above 16 under load); running the steps inside a +rayon pool rather than handing each loop to it (no clear win once +measured in-process); a rayon pool of any size against the spin +pool (2.5–4.4× slower under load, 16–24% on a quiet machine); +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/north_bench.rs b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/examples/north_bench.rs index 9d40b3bdb..c7ed8d7c6 100644 --- a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/examples/north_bench.rs +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/examples/north_bench.rs @@ -11,8 +11,12 @@ //! //! `--chat` wraps the file's text as one user turn of the //! model's chat template, the shape the program-writing loop -//! feeds. With `JOULECLAW_PROF=1` it also prints where the time -//! went, op by op, for the prefill and per decoded token. +//! 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}; @@ -54,6 +58,12 @@ fn main() -> Result<(), Box> { 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)? @@ -130,5 +140,45 @@ fn main() -> Result<(), Box> { 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/cohere2_moe.rs b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/cohere2_moe.rs index f02c84b87..37150967d 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 @@ -70,7 +70,11 @@ pub mod prof { pub const EXPERT_COMBINE: usize = 12; pub const RESIDUAL: usize = 13; pub const LM_HEAD: usize = 14; - const N: usize = 15; + 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", @@ -87,10 +91,15 @@ pub mod prof { "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")) @@ -100,6 +109,7 @@ pub mod prof { 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)) } @@ -484,7 +494,24 @@ 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(); @@ -745,7 +772,24 @@ impl Cohere2MoeModel { /// Sliding layers attend per-row windows; the cache is /// trimmed to the window afterwards, exactly matching /// rolling decode. Requires a fresh state. + /// + /// On the CPU this is [`Self::forward_prompt_cpu`]; when a + /// batched projection may leave the CPU (the fabric, or the + /// Metal-4 tensor path of a `fabric` build) 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; @@ -937,6 +981,570 @@ impl Cohere2MoeModel { } } +// ───────────────────────────────────────────────────────────────────── +// 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_*_matches_reference` tests pin it. + +/// `(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). + /// A conversation that adds a turn to a cached context feeds + /// it this way. Where a batched projection may leave the CPU + /// (see [`Self::forward_prompt`]) it runs the tokens one at a + /// time instead, 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 { @@ -1142,6 +1750,142 @@ 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 weights larger than a model + /// this size has; matmul_bf16's tests pin that.) + 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 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) { + 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]) { + 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"); + } + } + + #[test] + fn extend_matches_sequential_decode_tiny() { + // 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() { + // 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() { + let model = tiny_model(); + cpu_matches_reference(&model, &[1, 5, 9, 2, 7], 8); + } + + #[test] + fn cpu_passes_match_reference_medium() { + 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); + } + #[test] fn generate_through_the_trait() { let model = tiny_model(); 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..d20271f97 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,21 @@ 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 would stay on the + /// CPU's NEON kernels: the fabric is off and the Metal-4 + /// tensor prefill path is switched off + /// (`JOULECLAW_MPP_PREFILL=0`; it is on by default in a + /// `fabric` build). + pub fn prefill_on_cpu() -> bool { + fabric().is_none() && !mpp_prefill() + } + /// 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 +1349,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 +1360,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> { 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/spin_pool.rs b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/spin_pool.rs new file mode 100644 index 000000000..5300f514a --- /dev/null +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/spin_pool.rs @@ -0,0 +1,255 @@ +//! 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::sync::atomic::{AtomicBool, AtomicPtr, AtomicU64, AtomicUsize, Ordering}; +use std::sync::{Arc, Condvar, Mutex, OnceLock}; +use std::time::{Duration, Instant}; + +/// How long an idle thread spins before it parks. +const SPIN: Duration = Duration::from_micros(300); + +/// One loop: run `f(i)` for `i < len`. +struct Job<'a> { + f: &'a (dyn Fn(usize) + Sync), +} + +struct Shared { + /// `(epoch << 32) | next unit` of the current loop. + claim: AtomicU64, + /// The current loop's unit count. A plain number: a thread may + /// read a newer loop's count, which only makes its claim fail. + 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; the caller re-raises it. + panicked: AtomicBool, + /// 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<()>, +} + +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 cur >> 32 != epoch { + return ran; + } + let idx = (cur & 0xFFFF_FFFF) as usize; + 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. + 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 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| f(idx))).is_err() { + self.panicked.store(true, Ordering::Relaxed); + } + ran += 1; + self.done.fetch_add(1, Ordering::Release); + } + } +} + +/// The epoch after `e`: 32 bits, wrapping, never 0 (a new thread +/// has seen epoch 0). +fn next_epoch(e: u64) -> u64 { + match (e + 1) & 0xFFFF_FFFF { + 0 => 1, + n => n, + } +} + +fn worker(shared: Arc) { + let mut seen = 0u64; + loop { + let start = Instant::now(); + let mut epoch; + loop { + epoch = shared.claim.load(Ordering::Acquire) >> 32; + 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 threads = threads.max(1); + 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), + sleep: Mutex::new(0), + wake: Condvar::new(), + }); + for i in 1..threads { + let s = shared.clone(); + std::thread::Builder::new() + .name(format!("jouleclaw-spin-{i}")) + .spawn(move || worker(s)) + .expect("spawn a pool thread"); + } + SpinPool { shared, threads, busy: Mutex::new(()) } + } + + /// 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. + 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; + } + let _one = self.busy.lock().unwrap(); + let sh = &*self.shared; + assert!(len < 1 << 32, "too many units in one loop"); + 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>; + sh.done.store(0, Ordering::Relaxed); + sh.len.store(len, Ordering::Relaxed); + sh.job.store(ptr, Ordering::Release); + let epoch = next_epoch(sh.claim.load(Ordering::Relaxed) >> 32); + sh.claim.store(epoch << 32, 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(); + } + if sh.panicked.swap(false, Ordering::Relaxed) { + panic!("a unit of a spin-pool loop panicked"); + } + } +} + +/// 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(0xFFFF_FFFE), 0xFFFF_FFFF); + assert_eq!(next_epoch(0xFFFF_FFFF), 1); + // A pool whose counter is just below the wrap keeps running. + let pool = SpinPool::new(4); + pool.shared.claim.store(0xFFFF_FFFD << 32, 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); + } +} -- 2.54.0 (Apple Git-157) From a947e2ed12b809cde2d3e1fd8f83aa56abde6065 Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Thu, 5 Nov 2026 17:08:18 -0500 Subject: [PATCH 05/14] jouleclaw-loader-gguf: the two lints the macOS job denies The macOS energy job ends by linting this crate with warnings denied, in two configurations. It failed in the first and so never reached the second: - attention.rs tested `d_head % 4 == 0` in two places (manual_is_multiple_of). They read `d_head.is_multiple_of(4)` now, as the rest of the crate does; the value is the same for every d_head. - behind it, with `rayon` on, matmul_bf16_threaded's tile closure rebound `ptr` to itself (redundant_locals). The rebinding was not idle: it made the closure capture the OutPtr, which is Sync, where naming `ptr.0` alone captures the raw pointer, which is not. OutPtr::get does the same by being a method, so the closure calls that and the rebinding goes. Neither changes a result: the same branch is taken and the same pointer is written through. The tests that pin both paths bit for bit pass on macOS arm64, Linux aarch64 and Linux x86_64: attention_row_is_bit_identical_to_reference, over head widths that are and are not multiples of four, and bf16_threaded_above_the_floor_is_bit_identical, over the tiled shapes. Co-Authored-By: Claude Fable 5.1 --- .../crates/jouleclaw-loader-gguf/src/attention.rs | 4 ++-- .../crates/jouleclaw-loader-gguf/src/matmul_bf16.rs | 13 +++++++++++-- 2 files changed, 13 insertions(+), 4 deletions(-) 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 596c84b62..25b13958c 100644 --- a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/attention.rs +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/attention.rs @@ -226,7 +226,7 @@ pub fn attention_row_f32( /// `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 % 4 == 0 { qk_scores_neon(q, k, kv_len, d_head, inv_sqrt_d, scores) } else { 0 }; + 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 { @@ -324,7 +324,7 @@ unsafe fn qk_block( /// 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 % 4 == 0 { + if d_head.is_multiple_of(4) { use std::arch::aarch64::*; let mut d0 = 0; while d0 < d_head { 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 9994a162f..401f37066 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 @@ -491,6 +491,16 @@ struct OutPtr(*mut f32); 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 @@ -716,7 +726,6 @@ pub fn matmul_bf16_threaded( 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 ptr = ptr; let (bi, bj) = (t / n_tiles, t % n_tiles); let i0 = bi * TM; let rows = TM.min(m - i0); @@ -727,7 +736,7 @@ pub fn matmul_bf16_threaded( // 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.0.add(j0), n) }; + unsafe { matmul_bf16_tile(a, w_tile, b_tile, i0, rows, cols, k, ptr.get().add(j0), n) }; }); } Ok(()) -- 2.54.0 (Apple Git-157) From 02231036763551691d8180bd20d34686f3d608e4 Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Thu, 5 Nov 2026 17:08:18 -0500 Subject: [PATCH 06/14] jouleclaw-loader-gguf: the CPU-against-reference tests say when the reference has left the CPU In a `fabric` build one of fabric_dispatch's tests sets JOULECLAW_FABRIC=1, and that switch is one per process. Every test that reads it afterwards runs with the fabric engaged, and the reference passes then hand their projections to the fabric's members, which round differently. cpu_passes_match_reference_medium and extend_matches_sequential_decode_medium compared the CPU passes with those and failed whenever the two ran in one process (five runs of five on the M3 Ultra). Alone they pass, five of five; with the switch exported they fail the same way. The tiny ones passed only because their shapes are below the fabric's floors. No identity is claimed with the fabric engaged: forward_token and forward_prompt take the reference passes then precisely because those may leave the CPU. So both helpers ask first, say on stderr that they are not checking, and return. Asking settles the switch for the process, so a test that goes on runs with the fabric off to its end. Without the `fabric` feature the answer is always "on the CPU", so the default and `rayon` runs check everything they did. No CI job runs these tests in a `fabric` build. This was found running the crate's lib tests under `rayon fabric-meter` by hand, and that run is now 321 passed, 2 ignored. Co-Authored-By: Claude Fable 5.1 --- .../jouleclaw-loader-gguf/src/cohere2_moe.rs | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) 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 37150967d..a8f1b8256 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 @@ -1797,11 +1797,31 @@ pub(crate) mod tests { v.iter().map(|x| x.to_bits()).collect() } + /// True when the reference passes are CPU arithmetic, which + /// is when they are the CPU passes' oracle. In a `fabric` + /// build one of `fabric_dispatch`'s tests sets + /// `JOULECLAW_FABRIC=1` for the whole test process; the + /// reference passes then hand their projections to the + /// fabric's members, which round differently (it is what + /// those passes are kept for), and no identity with them is + /// claimed. The switch is read once per process, so the + /// answer given here holds for the rest of the test. + fn reference_is_on_the_cpu(what: &str) -> bool { + let on_cpu = !crate::fabric_dispatch::fabric_engaged(); + if !on_cpu { + eprintln!("{what}: not checked. The fabric is engaged in this process (JOULECLAW_FABRIC=1), so the reference passes are not CPU arithmetic."); + } + on_cpu + } + /// 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) { + if !reference_is_on_the_cpu("cpu_matches_reference") { + return; + } 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); @@ -1832,6 +1852,9 @@ pub(crate) mod tests { /// 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]) { + if !reference_is_on_the_cpu("extend_matches_sequential") { + return; + } let mut s_ext = model.new_state(); let mut s_seq = model.new_state(); let mut at = 0; -- 2.54.0 (Apple Git-157) From 16c1a6f252fee40f109bf950298288170393fc3b Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Fri, 6 Nov 2026 03:03:31 -0500 Subject: [PATCH 07/14] jouleclaw-loader-gguf: a late thread of a finished loop claims nothing of the next, and a unit's panic leaves the pool usable The spin pool published a new loop's unit count before it replaced the claim word, which still read (the finished loop's epoch, its unit count). A thread still inside `work` for that loop (woken late from park, or descheduled after its last unit) passed the bound against the larger count and claimed units the finished loop does not have. It ran units of the new loop that were then claimed and run again, `done` was counted past `len`, and `run` could return while a unit was still writing buffers its caller then read and freed. `run` now retires the old word first (its index passes no bound), then stores the count with Release, so a thread that reads the new count sees the retired word in its exchange, which fails; then the job, then the new epoch. Three tests, each failing before this change: - a late thread in the window between the count and the epoch, with no pool threads (deterministic): it ran units 3..7 of the next loop, and every unit of that loop ran twice; - the pool's own threads woken late from park, with the publication held open 200 us: in the first round, units 2..63 of a 64-unit loop ran twice; - a panicking unit: `run` re-raised the panic while it held the `busy` guard, which poisoned it, and the next loop panicked on the PoisonError. The guard is now dropped before the panic is re-raised, and a poisoned lock is taken as it is: the last loop has finished, so the pool's state is whole. All six pool tests passed in 21 runs on a loaded M3 Ultra. Co-Authored-By: Claude Opus 5.5 --- .../jouleclaw-loader-gguf/src/spin_pool.rs | 172 ++++++++++++++++-- 1 file changed, 155 insertions(+), 17 deletions(-) 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 index 5300f514a..2a9799c55 100644 --- a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/spin_pool.rs +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/spin_pool.rs @@ -17,22 +17,30 @@ //! writes only its own outputs. use std::sync::atomic::{AtomicBool, AtomicPtr, AtomicU64, AtomicUsize, Ordering}; -use std::sync::{Arc, Condvar, Mutex, OnceLock}; +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); +/// The unit index of a finished loop's claim word: no loop has this +/// many units (`run` asserts fewer than 2^32), so a thread that still +/// holds that loop's epoch fails the bound and claims nothing. +const RETIRED: u64 = 0xFFFF_FFFF; + /// One loop: run `f(i)` for `i < len`. struct Job<'a> { f: &'a (dyn Fn(usize) + Sync), } struct Shared { - /// `(epoch << 32) | next unit` of the current loop. + /// `(epoch << 32) | next unit` of the current loop, or + /// `(epoch << 32) | RETIRED` once a new loop is being published. claim: AtomicU64, - /// The current loop's unit count. A plain number: a thread may - /// read a newer loop's count, which only makes its claim fail. + /// 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. @@ -52,8 +60,16 @@ pub struct SpinPool { 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. @@ -65,6 +81,10 @@ impl Shared { return ran; } let idx = (cur & 0xFFFF_FFFF) 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; } @@ -131,7 +151,19 @@ 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 threads = threads.max(1); + 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), @@ -141,14 +173,13 @@ impl SpinPool { sleep: Mutex::new(0), wake: Condvar::new(), }); - for i in 1..threads { - let s = shared.clone(); - std::thread::Builder::new() - .name(format!("jouleclaw-spin-{i}")) - .spawn(move || worker(s)) - .expect("spawn a pool thread"); + SpinPool { + shared, + threads: threads.max(1), + busy: Mutex::new(()), + #[cfg(test)] + window: None, } - SpinPool { shared, threads, busy: Mutex::new(()) } } /// Threads that run units, the caller's included. @@ -168,19 +199,42 @@ impl SpinPool { } return; } - let _one = self.busy.lock().unwrap(); - let sh = &*self.shared; assert!(len < 1 << 32, "too many units in one loop"); + // Nothing panics while this is held: a unit's panic is caught + // and re-raised after the guard is gone. A poisoned lock (from + // a panic this code does not expect) still guards a pool whose + // last loop finished, so it 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 = sh.claim.load(Ordering::Relaxed) >> 32; + sh.claim.store((old << 32) | RETIRED, Ordering::Relaxed); sh.done.store(0, Ordering::Relaxed); - sh.len.store(len, Ordering::Relaxed); + sh.len.store(len, Ordering::Release); sh.job.store(ptr, Ordering::Release); - let epoch = next_epoch(sh.claim.load(Ordering::Relaxed) >> 32); + #[cfg(test)] + if let Some(w) = &self.window { + w(sh); + } + let epoch = next_epoch(old); sh.claim.store(epoch << 32, Ordering::Release); { let mut s = sh.sleep.lock().unwrap(); @@ -191,7 +245,9 @@ impl SpinPool { while sh.done.load(Ordering::Acquire) < len { std::hint::spin_loop(); } - if sh.panicked.swap(false, Ordering::Relaxed) { + let panicked = sh.panicked.swap(false, Ordering::Relaxed); + drop(one); + if panicked { panic!("a unit of a spin-pool loop panicked"); } } @@ -252,4 +308,86 @@ mod tests { }); 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 = sh.claim.load(Ordering::Relaxed) >> 32; + 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:?}"); + } + } + } + + /// 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); + } } -- 2.54.0 (Apple Git-157) From 0e45132e209a2c7362bb0aa3b5aaf5eee1dce135 Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Fri, 6 Nov 2026 03:03:41 -0500 Subject: [PATCH 08/14] jouleclaw-loader-gguf: prefill and extend take the CPU passes whenever the fabric is off, and the identity tests run where the switch is off `prefill_on_cpu()` was `fabric().is_none() && !mpp_prefill()`. A `fabric` build with the fabric off (the default) and JOULECLAW_MPP_PREFILL unset therefore took the reference prefill (attention one head at a time, experts one after another) and fed `forward_extend` one token at a time, though nothing can leave the CPU then: the Metal-4 path refuses without the fabric, like every other route off the CPU. It is now `fabric().is_none()`. Logits were never affected; the batched CPU prefill and extend were. The identity tests compare the CPU passes with the reference passes, which are CPU arithmetic only with the fabric off. The switch is read once per process, and `mpp_bt_matches_neon_when_available` set JOULECLAW_FABRIC=1 in the test process, so what those tests compared depended on which test read the switch first; since 022310367 they passed with nothing checked when it was on. Now no test sets the variable in its own process. A test that needs the switch one way runs here if it already is, otherwise alone in a child of the test binary with the switch set, and the child must report one passed test (`fabric_switch`). The helpers assert that the reference is CPU arithmetic and that `forward_extend` batches, instead of returning. Before this change, in a `rayon fabric-meter` build on macOS: with the switch unset, the new predicate test and both extend tests failed (prefill_on_cpu() false with the fabric off); with JOULECLAW_FABRIC=1, all four identity tests failed the CPU-arithmetic assert. After it, all pass both ways: with the switch on the four run in children, and with it off the MPP test runs in a child with it on. Co-Authored-By: Claude Opus 5.5 --- .../jouleclaw-loader-gguf/src/cohere2_moe.rs | 73 +++++++++++------ .../src/fabric_dispatch.rs | 80 ++++++++++++++++--- 2 files changed, 120 insertions(+), 33 deletions(-) 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 a8f1b8256..2baaf81d3 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 @@ -773,9 +773,10 @@ impl Cohere2MoeModel { /// trimmed to the window afterwards, exactly matching /// rolling decode. Requires a fresh state. /// - /// On the CPU this is [`Self::forward_prompt_cpu`]; when a - /// batched projection may leave the CPU (the fabric, or the - /// Metal-4 tensor path of a `fabric` build) it is + /// 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() { @@ -1568,6 +1569,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 @@ -1797,21 +1799,15 @@ pub(crate) mod tests { v.iter().map(|x| x.to_bits()).collect() } - /// True when the reference passes are CPU arithmetic, which - /// is when they are the CPU passes' oracle. In a `fabric` - /// build one of `fabric_dispatch`'s tests sets - /// `JOULECLAW_FABRIC=1` for the whole test process; the - /// reference passes then hand their projections to the - /// fabric's members, which round differently (it is what - /// those passes are kept for), and no identity with them is - /// claimed. The switch is read once per process, so the - /// answer given here holds for the rest of the test. - fn reference_is_on_the_cpu(what: &str) -> bool { - let on_cpu = !crate::fabric_dispatch::fabric_engaged(); - if !on_cpu { - eprintln!("{what}: not checked. The fabric is engaged in this process (JOULECLAW_FABRIC=1), so the reference passes are not CPU arithmetic."); - } - on_cpu + /// 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 @@ -1819,9 +1815,7 @@ pub(crate) mod tests { /// 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) { - if !reference_is_on_the_cpu("cpu_matches_reference") { - return; - } + 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); @@ -1852,9 +1846,11 @@ pub(crate) mod tests { /// 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]) { - if !reference_is_on_the_cpu("extend_matches_sequential") { - return; - } + 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; @@ -1881,8 +1877,15 @@ pub(crate) mod tests { } } + // 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]); @@ -1890,6 +1893,9 @@ pub(crate) mod tests { #[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(); @@ -1898,17 +1904,36 @@ pub(crate) mod tests { #[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(); 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 d20271f97..56fb0dcef 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 @@ -68,13 +68,14 @@ mod enabled { fabric().is_some() } - /// True when a batched (prefill) projection would stay on the - /// CPU's NEON kernels: the fabric is off and the Metal-4 - /// tensor prefill path is switched off - /// (`JOULECLAW_MPP_PREFILL=0`; it is on by default in a - /// `fabric` build). + /// 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() && !mpp_prefill() + fabric().is_none() } /// Widen a BF16 row-major block to f32 row-major with @@ -1484,6 +1485,64 @@ 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), and fail unless the + /// child reports that exactly that one test ran and passed. + pub(crate) fn run_alone(module: &str, test: &str, env: &[(&str, Option<&str>)]) { + // 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(), + }; + assert!(!in_child(), "{path}: a child test process would start another, with {env:?}"); + 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"]).env(CHILD, "1"); + 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); + 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) + ); + } +} + #[cfg(all(test, feature = "fabric", target_os = "macos", target_arch = "aarch64"))] mod tests { @@ -1607,9 +1666,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 = [ -- 2.54.0 (Apple Git-157) From 97c3accb809c5b4334f2d7e01577c7411fc35abe Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Fri, 6 Nov 2026 03:03:50 -0500 Subject: [PATCH 09/14] jouleclaw-loader-gguf: a stored reference for Cohere2-MoE logits that CI checks without the checkpoint The five-prompt, two-Mac identity result has no stored reference, and nothing in CI could re-check it. src/cohere2_moe_reference.txt now holds FNV-1a hashes of every output's logits for three synthetic checkpoints: tiny, medium, and a wide one whose decode projections are several row chunks and whose prefill reaches the 2-D tiles. The scenario is a prefill, decode steps, a chunk fed into the live state, and two more steps. The hashes come from the code before the CPU passes: commit 30c4cf1ee (main before this pull request), its forward_prompt and forward_token and its one-row BF16 kernel, with and without rayon (identical lines), on macOS arm64, Linux arm64 and Linux x86_64. The platforms give different bits, so each has its own 26 lines. `stored_reference::logits_match_the_stored_reference` recomputes them with the reference passes, the CPU passes and the public entry points, then, with rayon, in child processes at 1 thread and at 3 threads with 8 KiB GEMV chunks. It passes on all three platforms. On Linux arm64 it also passes with the glibc 2.35 libm of CI's container bound in place of the host's 2.39; CI's linux-arm64 workspace job runs it. BENCHMARKS.md says what this re-checks and what still needs the real checkpoint: identity at North's own sizes and on the five prompts, whose text and hashes are not stored, and every speed figure. Co-Authored-By: Claude Opus 5.5 --- .../jouleclaw-loader-gguf/BENCHMARKS.md | 19 ++ .../jouleclaw-loader-gguf/src/cohere2_moe.rs | 180 +++++++++++++++++- .../src/cohere2_moe_reference.txt | 94 +++++++++ 3 files changed, 290 insertions(+), 3 deletions(-) create mode 100644 jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/cohere2_moe_reference.txt diff --git a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/BENCHMARKS.md b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/BENCHMARKS.md index 7c0bc4b58..4214f5030 100644 --- a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/BENCHMARKS.md +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/BENCHMARKS.md @@ -1486,6 +1486,25 @@ 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 (its container's glibc 2.35 +math library gives the stored lines too); the macOS lines are checked +by running the tests on a Mac. **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 tables give medians and paired ranges; the per-run +logs are not kept. + **Speed.** Medians of 4 runs (3 for base at 12 threads), base and new interleaved in the same conditions. M3 Ultra (24P+8E): shared, load average 30–130, with another North benchmark and CI running. 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 2baaf81d3..bd82d9e63 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 @@ -1013,7 +1013,9 @@ impl Cohere2MoeModel { // 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_*_matches_reference` tests pin it. +// 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 @@ -1755,8 +1757,9 @@ pub(crate) mod tests { /// 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 weights larger than a model - /// this size has; matmul_bf16's tests pin that.) + /// 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, @@ -1945,4 +1948,175 @@ 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 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): at the default thread count here and, with + /// `rayon`, in child processes at 1 thread and at 3 threads + /// with 8 KiB GEMV chunks (both are read once per process). + #[test] + fn logits_match_the_stored_reference() { + if !fabric_switch::here(false, module_path!(), "logits_match_the_stored_reference") { + return; + } + let platform = format!("{}-{}", std::env::consts::OS, std::env::consts::ARCH); + let want: Vec<&str> = include_str!("cohere2_moe_reference.txt") + .lines() + .filter_map(|l| l.strip_prefix(platform.as_str())?.strip_prefix(' ')) + .collect(); + assert_eq!(want.len(), 26, "stored lines for {platform}"); + 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), + )); + } + for (passes, lines) in &got { + assert_eq!(lines, &want, "the {passes} passes on {platform}"); + } + #[cfg(feature = "rayon")] + if !fabric_switch::in_child() { + let test = "logits_match_the_stored_reference"; + let off = ("JOULECLAW_FABRIC", None); + fabric_switch::run_alone(module_path!(), test, &[off, ("JOULECLAW_THREADS", Some("1"))]); + fabric_switch::run_alone( + module_path!(), + test, + &[off, ("JOULECLAW_THREADS", Some("3")), ("JOULECLAW_CHUNK_KB", Some("8"))], + ); + } + } + } } 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..6a74207b0 --- /dev/null +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/cohere2_moe_reference.txt @@ -0,0 +1,94 @@ +# 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 the code before the CPU passes: commit 30c4cf1ee (main before +# PR #649), its forward_prompt and forward_token running the same scenario, +# with and without the rayon feature (the lines were identical), on macOS +# arm64, Linux arm64 and Linux x86_64 on 2026-11-06. The platforms differ: +# 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, which +# leaves their math libraries (exp, powf, sin, cos) as the likely cause. +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 -- 2.54.0 (Apple Git-157) From 3e6a60bae5b63e57516cc26c31a9602df07475b8 Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Fri, 6 Nov 2026 03:03:56 -0500 Subject: [PATCH 10/14] jouleclaw-loader-gguf: forward_extend's doc says what it is for, not who calls it It said that a conversation adding a turn to a cached context feeds it this way. Nothing in this repository does: its only caller is examples/north_bench.rs, and a conversation that feeds its turn one forward_token per token gets the same logits, more slowly. The doc now says what the function is for and what the alternative gives. BENCHMARKS.md now says where the follow-up-turn figures come from (north_bench --then and --then-seq), and which of them a caller gets by which route. Co-Authored-By: Claude Opus 5.5 --- .../crates/jouleclaw-loader-gguf/BENCHMARKS.md | 11 +++++++---- .../crates/jouleclaw-loader-gguf/src/cohere2_moe.rs | 10 ++++++---- 2 files changed, 13 insertions(+), 8 deletions(-) diff --git a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/BENCHMARKS.md b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/BENCHMARKS.md index 4214f5030..385419fcc 100644 --- a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/BENCHMARKS.md +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/BENCHMARKS.md @@ -1522,10 +1522,13 @@ at ~600 and 5.0–6.6× at ~2700 (one base run at load 130–131 was 20×, left out), prefill 5.4–10.3× and 9.8–11.8×; M4 Max decode 3.3–3.5× and 5.1–5.3×, prefill 5.2–6.8× and 7.0–9.0×. -A 69-token follow-up turn fed into a ~630-token context: -`forward_extend` 1.23 s (M3 Ultra) / 0.79 s (M4 Max), against -3.3–4.1 s / 2.4 s one `forward_token` at a time on the new code -and ~18 s at base's decode speed; same logits. +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 s (M4 Max), against 3.3–4.1 s / 2.4 s one +`forward_token` at a time on the new code and ~18 s at base's +decode speed; 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); running the steps inside a 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 bd82d9e63..07a42746d 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 @@ -1329,10 +1329,12 @@ impl Cohere2MoeModel { /// [`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). - /// A conversation that adds a turn to a cached context feeds - /// it this way. Where a batched projection may leave the CPU - /// (see [`Self::forward_prompt`]) it runs the tokens one at a - /// time instead, which is the same thing. + /// 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() { -- 2.54.0 (Apple Git-157) From b38c067054824ed3c0bd2e330f88800addc7a273 Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Fri, 6 Nov 2026 03:06:33 -0500 Subject: [PATCH 11/14] jouleclaw-loader-gguf: the spin pool's run says that a unit must not start a loop on it A loop holds the pool's `busy` lock until every unit has run. A unit that calls `run` on the same pool with more than one unit therefore waits forever: on the caller's thread for a lock it already holds, on a pool thread for a lock held by a caller that waits for that very unit. No unit in this crate does so (decode's units are GEMV chunks and attention rows); the precondition is now written where `run` is. Co-Authored-By: Claude Opus 5.5 --- .../crates/jouleclaw-loader-gguf/src/spin_pool.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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 index 2a9799c55..610950fcf 100644 --- a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/spin_pool.rs +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/spin_pool.rs @@ -188,7 +188,9 @@ impl SpinPool { } /// Run `f(i)` for every `i < len` on the pool and the calling - /// thread; returns when all have run. + /// 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; -- 2.54.0 (Apple Git-157) From 2d6fe0f3df151b3bfeb3e56cb3a952df09cd2efb Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Fri, 6 Nov 2026 07:18:22 -0500 Subject: [PATCH 12/14] jouleclaw-loader-gguf: the spin pool re-raises a unit's own panic, and its epoch has 40 bits A unit's panic reached the caller as a new panic with the pool's message, so its payload was lost. Each panicked unit's payload is now kept in a list rather than dropped: dropping it would run its Drop inside the loop, where a panicking Drop would leave the unit uncounted and the caller waiting. Once the loop has finished and the pool is released, the caller drops the others and re-raises the first with resume_unwind, payload and all. The claim word was a 32-bit epoch and a 32-bit unit index. A thread descheduled between loading the word and its exchange across exactly 2^32 - 1 loops would find the same word again in a different loop: about 11 days of decoding at 4500 loops a second. The word is now a 40-bit epoch and a 24-bit index (about 7.7 years). A loop may have at most 2^24 - 1 units, which `run` asserts before it takes the pool; the decoder's largest are about 65 thousand (lm_head in 4-row chunks) and about a million (prefill attention for a 500k-token prompt). The safety comment states the remaining case. The comment on `busy` now says exactly what it relies on: nothing between taking it and releasing it panics, and why each step cannot. Tests failing first: the unit's payload reached the caller as None (a string from the pool); next_epoch(0xFFFF_FFFF) was 1. Both pass now, as does the reviewer's loom model ported to the new layout (6 of 6, the retire step still found necessary by its mutation test) and its stress harness at 2 to 24 threads on macOS arm64, Linux arm64 and Linux x86_64. Co-Authored-By: Claude Opus 5.5 --- .../jouleclaw-loader-gguf/src/spin_pool.rs | 142 ++++++++++++++---- 1 file changed, 114 insertions(+), 28 deletions(-) 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 index 610950fcf..96ff0195c 100644 --- a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/spin_pool.rs +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/spin_pool.rs @@ -16,6 +16,7 @@ //! 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}; @@ -23,10 +24,22 @@ use std::time::{Duration, Instant}; /// How long an idle thread spins before it parks. const SPIN: Duration = Duration::from_micros(300); -/// The unit index of a finished loop's claim word: no loop has this -/// many units (`run` asserts fewer than 2^32), so a thread that still -/// holds that loop's epoch fails the bound and claims nothing. -const RETIRED: u64 = 0xFFFF_FFFF; +/// 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> { @@ -34,8 +47,9 @@ struct Job<'a> { } struct Shared { - /// `(epoch << 32) | next unit` of the current loop, or - /// `(epoch << 32) | RETIRED` once a new loop is being published. + /// `(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 @@ -47,8 +61,13 @@ struct Shared { job: AtomicPtr>, /// Units finished in the current loop. done: AtomicUsize, - /// A unit panicked; the caller re-raises it. + /// 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, @@ -77,10 +96,10 @@ impl Shared { let mut ran = 0; loop { let cur = self.claim.load(Ordering::Acquire); - if cur >> 32 != epoch { + if epoch_of(cur) != epoch { return ran; } - let idx = (cur & 0xFFFF_FFFF) as usize; + 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 @@ -99,11 +118,22 @@ impl Shared { // 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 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| f(idx))).is_err() { + 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; @@ -112,10 +142,10 @@ impl Shared { } } -/// The epoch after `e`: 32 bits, wrapping, never 0 (a new thread +/// 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) & 0xFFFF_FFFF { + match (e + 1) & ((1 << (64 - IDX_BITS)) - 1) { 0 => 1, n => n, } @@ -127,7 +157,7 @@ fn worker(shared: Arc) { let start = Instant::now(); let mut epoch; loop { - epoch = shared.claim.load(Ordering::Acquire) >> 32; + epoch = epoch_of(shared.claim.load(Ordering::Acquire)); if epoch != seen || start.elapsed() > SPIN { break; } @@ -170,6 +200,7 @@ impl SpinPool { 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(), }); @@ -201,11 +232,19 @@ impl SpinPool { } return; } - assert!(len < 1 << 32, "too many units in one loop"); - // Nothing panics while this is held: a unit's panic is caught - // and re-raised after the guard is gone. A poisoned lock (from - // a panic this code does not expect) still guards a pool whose - // last loop finished, so it is taken as it is. + 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 }; @@ -227,8 +266,8 @@ impl SpinPool { // 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 = sh.claim.load(Ordering::Relaxed) >> 32; - sh.claim.store((old << 32) | RETIRED, Ordering::Relaxed); + 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); @@ -237,7 +276,7 @@ impl SpinPool { w(sh); } let epoch = next_epoch(old); - sh.claim.store(epoch << 32, Ordering::Release); + sh.claim.store(epoch << IDX_BITS, Ordering::Release); { let mut s = sh.sleep.lock().unwrap(); *s = epoch; @@ -247,10 +286,19 @@ impl SpinPool { while sh.done.load(Ordering::Acquire) < len { std::hint::spin_loop(); } - let panicked = sh.panicked.swap(false, Ordering::Relaxed); + 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); - if panicked { - panic!("a unit of a spin-pool loop panicked"); + // 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); } } } @@ -283,11 +331,11 @@ mod tests { #[test] fn the_epoch_wraps_without_reaching_zero() { assert_eq!(next_epoch(1), 2); - assert_eq!(next_epoch(0xFFFF_FFFE), 0xFFFF_FFFF); - assert_eq!(next_epoch(0xFFFF_FFFF), 1); + 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(0xFFFF_FFFD << 32, Ordering::Relaxed); + pool.shared.claim.store(((1 << 40) - 3) << IDX_BITS, Ordering::Relaxed); let n = AtomicUsize::new(0); for _ in 0..6 { pool.run(50, &|_| { @@ -326,7 +374,7 @@ mod tests { let (l, a) = (late.clone(), armed.clone()); pool.window = Some(Box::new(move |sh: &Shared| { if a.load(Ordering::Relaxed) { - let finished = sh.claim.load(Ordering::Relaxed) >> 32; + 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); } @@ -374,6 +422,44 @@ mod tests { } } + /// 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. -- 2.54.0 (Apple Git-157) From 5e03008268d11357bd0432039fc2490291de1082 Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Fri, 6 Nov 2026 07:18:22 -0500 Subject: [PATCH 13/14] jouleclaw-loader-gguf: the stored reference names its math library, has its generator, and runs every configuration whatever the switch The stored logit hashes depend on the platform's math library as well as on the code, but a mismatch blamed only the code. Each platform's lines now carry a `math-library` line naming the library they were made with (macOS 27.2's libsystem_m, Ubuntu GLIBC 2.39 on Linux arm64, 2.43 on Linux x86_64). A mismatch names it, says whether the ways of running the model agree with each other (if they do, this environment's library may round differently), and says how to tell. examples/cohere2_reference.rs is the generator: the test's checkpoints, scenario and hash, through forward_prompt and forward_token only, so it runs unchanged at 30c4cf1ee. Run there, with and without rayon, it reproduces all 78 stored lines on all three platforms, and it prints the `math-library` line for a new one. With JOULECLAW_FABRIC=1 in the environment, the main check ran in a child process and the 1- and 3-thread checks never ran at all. The test process now checks nothing itself: it starts one child per configuration, each with the fabric off, and each child must report its check. Tests failing first: the mismatch message did not name the library; a test run with JOULECLAW_FABRIC=1 reported 0 of its 3 checks. Both pass now, with the switch unset and set. Co-Authored-By: Claude Opus 5.5 --- .../examples/cohere2_reference.rs | 280 ++++++++++++++++++ .../jouleclaw-loader-gguf/src/cohere2_moe.rs | 151 ++++++++-- .../src/cohere2_moe_reference.txt | 19 +- .../src/fabric_dispatch.rs | 30 +- 4 files changed, 446 insertions(+), 34 deletions(-) create mode 100644 jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/examples/cohere2_reference.rs 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/src/cohere2_moe.rs b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/src/cohere2_moe.rs index 07a42746d..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 @@ -2053,23 +2053,135 @@ pub(crate) mod tests { 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): at the default thread count here and, with - /// `rayon`, in child processes at 1 thread and at 3 threads - /// with 8 KiB GEMV chunks (both are read once per process). + /// 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::here(false, module_path!(), "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 want: Vec<&str> = include_str!("cohere2_moe_reference.txt") - .lines() - .filter_map(|l| l.strip_prefix(platform.as_str())?.strip_prefix(' ')) - .collect(); - assert_eq!(want.len(), 26, "stored lines for {platform}"); let mut got: [(&str, Vec); 3] = [("reference", Vec::new()), ("cpu", Vec::new()), ("public", Vec::new())]; for (name, cfg, shape) in stored_models() { @@ -2105,20 +2217,15 @@ pub(crate) mod tests { |ts, s| m.forward_extend(ts, s), )); } - for (passes, lines) in &got { - assert_eq!(lines, &want, "the {passes} passes on {platform}"); - } - #[cfg(feature = "rayon")] - if !fabric_switch::in_child() { - let test = "logits_match_the_stored_reference"; - let off = ("JOULECLAW_FABRIC", None); - fabric_switch::run_alone(module_path!(), test, &[off, ("JOULECLAW_THREADS", Some("1"))]); - fabric_switch::run_alone( - module_path!(), - test, - &[off, ("JOULECLAW_THREADS", Some("3")), ("JOULECLAW_CHUNK_KB", Some("8"))], - ); + 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 index 6a74207b0..93652a62f 100644 --- 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 @@ -7,13 +7,20 @@ # wide configs; the scenario prefills a prompt, decodes, feeds a chunk into # the live state and decodes twice more. # -# Generated by the code before the CPU passes: commit 30c4cf1ee (main before -# PR #649), its forward_prompt and forward_token running the same scenario, -# with and without the rayon feature (the lines were identical), on macOS -# arm64, Linux arm64 and Linux x86_64 on 2026-11-06. The platforms differ: +# 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, which -# leaves their math libraries (exp, powf, sin, cos) as the likely cause. +# 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 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 56fb0dcef..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 @@ -1515,17 +1515,34 @@ pub(crate) mod fabric_switch { } /// Run the test `module::test` alone in a child of this test - /// binary with `env` set (`None`: removed), and fail unless the - /// child reports that exactly that one test ran and passed. - pub(crate) fn run_alone(module: &str, test: &str, env: &[(&str, Option<&str>)]) { + /// 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(), }; - assert!(!in_child(), "{path}: a child test process would start another, with {env:?}"); 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"]).env(CHILD, "1"); + 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), @@ -1533,13 +1550,14 @@ pub(crate) mod fabric_switch { }; } let out = cmd.output().expect("start the test binary"); - let stdout = String::from_utf8_lossy(&out.stdout); + 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 } } -- 2.54.0 (Apple Git-157) From 895778bbb3aa004cad265b74c966f340435d4073 Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Fri, 6 Nov 2026 07:18:42 -0500 Subject: [PATCH 14/14] jouleclaw-loader-gguf: BENCHMARKS.md says what each North figure compares, from the run logs Corrected against the per-run logs of the benchmark (final.log on both Macs): - The headline compared with base at its default thread count, which it did not say. It now gives median-against-median ratios, decode 3.3-5.2x and prefill 5.9-10.6x, and names that baseline. - The M3 Ultra decode median at ~600 context, 267/54 ms = 4.96x, lay outside the stated paired range 4.3-4.8x. The range left out a 20.2x pair in that very column, and a second, 21.2x, at ~2700. The text named one at load 130-131 (it was 131, at ~2700). Both are now named with their loads, and the text says they are inside the medians. - The M4 Max decode ranges silently left out the identity run's pairs (2.0x and 6.3x) while the prefill ranges kept them. They are now stated; the M3 prefill floor at ~2700 was 9.7, not 9.8. - Base at 12 threads is now in the table for the M4 Max too, with paired ratios against it on both Macs. Base at 8 threads, its fastest in the earlier sweep, was never run in pairs; the text says so instead of implying the ratio. - ~18 s for a turn at base's speed is an estimate (69 x 267 ms) and says so. A rayon pool against the spin pool was 18-24% slower on a quiet M4 Max (40-42 against 34 ms), not 16-24%. "Running the steps inside a rayon pool did not help" is about decode; prefill does run inside its pool, and the text says so. - CI's glibc 2.35 giving the stored lines was shown on the host that made them, with 2.35's libm loaded, not by CI; the text says CI's run was pending. The BF16 kernel runs 4 x 2, 2 x 4 and 1 x 4 blocks, not 4 x 4; its doc comment and the benchmark's bullet now say so. Co-Authored-By: Claude Opus 5.5 --- .../jouleclaw-loader-gguf/BENCHMARKS.md | 92 +++++++++++++------ .../jouleclaw-loader-gguf/src/matmul_bf16.rs | 6 +- 2 files changed, 66 insertions(+), 32 deletions(-) diff --git a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/BENCHMARKS.md b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/BENCHMARKS.md index 385419fcc..577566336 100644 --- a/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/BENCHMARKS.md +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/BENCHMARKS.md @@ -1439,13 +1439,17 @@ 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: 3–6× decode, 6–11× prefill, bit-identical (2026-11) +## 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. +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 @@ -1462,8 +1466,9 @@ fastest at 8 threads (94 ms) and slowest at the default 32 **What changed** (each output keeps its exact reduction order): -- the BF16 kernel computes blocks of up to 4 activation rows × 4 - weight rows (16 chains instead of 2) with a 1 KiB prefetch; +- 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 @@ -1496,45 +1501,72 @@ 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 (its container's glibc 2.35 -math library gives the stored lines too); the macOS lines are checked -by running the tests on a Mac. **What still needs the checkpoint:** identity at North's +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 tables give medians and paired ranges; the per-run -logs are not kept. +figure here. The per-run logs are not in the repository. -**Speed.** Medians of 4 runs (3 for base at 12 threads), base and -new interleaved in the same conditions. 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). +**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 | M3 Ultra base, 12 thr | M3 Ultra new | M4 Max base | M4 Max new | -|---|---|---|---|---|---| -| decode at ~600 context | 267 ms | 168 ms | **54 ms** | 116 ms | **35 ms** | -| decode at ~2700 context | 405 ms | 364 ms | **80 ms** | 246 ms | **48 ms** | -| prefill, 596 tokens | 9.9 tok/s | 12.0 tok/s | **91 tok/s** | 17.8 tok/s | **105 tok/s** | -| prefill, 2698 tokens | 5.9 tok/s | 4.5 tok/s | **63 tok/s** | 9.0 tok/s | **79 tok/s** | +| | 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 (base / new): M3 Ultra decode 4.3–4.8× -at ~600 and 5.0–6.6× at ~2700 (one base run at load 130–131 was -20×, left out), prefill 5.4–10.3× and 9.8–11.8×; M4 Max decode -3.3–3.5× and 5.1–5.3×, prefill 5.2–6.8× and 7.0–9.0×. +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 s (M4 Max), against 3.3–4.1 s / 2.4 s one -`forward_token` at a time on the new code and ~18 s at base's -decode speed; same logits. A caller gets the first figure only by +(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); running the steps inside a -rayon pool rather than handing each loop to it (no clear win once -measured in-process); a rayon pool of any size against the spin -pool (2.5–4.4× slower under load, 16–24% on a quiet machine); +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 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 401f37066..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 @@ -141,8 +141,10 @@ pub fn matmul_bf16_scalar( /// `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 up to 4 activation rows × 4 weight rows shares its loads -/// and keeps up to 16 independent accumulator chains in flight. +/// 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 -- 2.54.0 (Apple Git-157)