jouleclaw: North decode 3-6x and prefill 6-11x on the CPU, bit-identical #649
10 changed files with 3284 additions and 58 deletions
|
|
@ -1439,6 +1439,144 @@ uses `Digits` (isolate each digit) + `ByteLevel
|
|||
instead of explicit `Split` entries. Both now map onto the
|
||||
chained-split machinery from wave 60.
|
||||
|
||||
## North Mini Code on the CPU: decode 3.3–5.2× and prefill 5.9–10.6× against base at its default thread count, bit-identical (2026-11)
|
||||
|
||||
The pure-Rust CPU path (`--features rayon`) of North Mini Code 1.0
|
||||
(Cohere2-MoE, 30B total / 3B active, BF16 safetensors), made
|
||||
faster without changing a bit of any logit. Measured with
|
||||
`examples/north_bench.rs` (prints every step's token id and a hash
|
||||
of all 262144 logits' bits), base = `main` before this change. The
|
||||
ratios in the title are median against median in the table under
|
||||
**Speed**, with base at rayon's default thread count (every logical
|
||||
CPU) and the new code at its own default; that section also gives
|
||||
base at 12 threads and says what was not compared.
|
||||
|
||||
**Where the time went** (`JOULECLAW_PROF=1`, base, M3 Ultra under
|
||||
load, ~600-token prompt): prefill 45.3 s = attention 36% (every
|
||||
head on one thread, a scalar add chain per score), experts 50%
|
||||
(one expert after another), q/k/v + o 11%; decode 212 ms/token =
|
||||
attention 26%, q/k/v 20%, o_proj 20%, experts 28%, lm_head 3%.
|
||||
lm_head streams its 1 GB in ~7 ms (~140 GB/s) while each layer's
|
||||
16–50 MB projections moved 12–40 GB/s: decode was bound by its ~7
|
||||
dispatches per layer (three of them single-threaded GEMVs) and a
|
||||
GEMV kernel held to ~16–20 GB/s per core by its two accumulator
|
||||
chains, not by memory bandwidth. Base decode under load was also
|
||||
fastest at 8 threads (94 ms) and slowest at the default 32
|
||||
(174 ms).
|
||||
|
||||
**What changed** (each output keeps its exact reduction order):
|
||||
|
||||
- the BF16 kernel computes blocks of 4 activation rows × 2 weight
|
||||
rows, 2 × 4 or 1 × 4 (up to 16 accumulator chains instead of 2)
|
||||
with a 1 KiB prefetch;
|
||||
- decode runs three dispatches per layer (q/k/v + router;
|
||||
attention heads + expert gate/up; o_proj + expert down) of
|
||||
~256 KiB row chunks on a pool whose threads stay awake between
|
||||
dispatches (`spin_pool`);
|
||||
- attention runs per head (decode) and per (head, 16 query rows)
|
||||
(prefill) in parallel, with a bit-exact NEON score kernel;
|
||||
- prefill runs the experts in parallel and the projections as
|
||||
64 × 64 output tiles; RoPE angles are computed once per position;
|
||||
- `forward_extend` batches tokens into a state that already holds
|
||||
a context, bit-identical to feeding them one decode step at a
|
||||
time;
|
||||
- threads: one per performance core, at most 16
|
||||
(`JOULECLAW_THREADS` overrides).
|
||||
|
||||
**Identity.** Five prompts (a raw completion, Python, OpenQASM, a
|
||||
596-token and a 2698-token Rust prompt), 72 greedy tokens each:
|
||||
the same token ids and the same hash of every step's logits as
|
||||
base, on an M3 Ultra and on an M4 Max, and the same on both
|
||||
machines; every repeat of the timing runs below gave that one hash
|
||||
sequence too. Unit tests pin the kernels and both forward passes
|
||||
against the old code bit for bit.
|
||||
|
||||
**What is re-checked without the checkpoint.**
|
||||
`src/cohere2_moe_reference.txt` holds logit hashes of three synthetic
|
||||
Cohere2-MoE checkpoints, generated by base's `forward_prompt` and
|
||||
`forward_token` (with its one-row BF16 kernel) on macOS arm64, Linux
|
||||
arm64 and Linux x86_64. The three platforms give different bits, so
|
||||
each has its own lines. One checkpoint is wide enough that decode
|
||||
splits its projections into several row chunks and prefill runs the
|
||||
2-D tiles. The test `stored_reference::logits_match_the_stored_reference`
|
||||
recomputes the hashes with the reference passes, the CPU passes and
|
||||
`forward_extend`, at the default thread count and at 1 and 3 threads.
|
||||
CI runs it in the Linux arm64 workspace job. That job's container has
|
||||
glibc 2.35, and the Linux arm64 lines were made with 2.39: with 2.35's
|
||||
libm loaded in its place on the host that made them, the test gives
|
||||
the stored lines too, and CI's own first run of it was pending when
|
||||
this was written. The macOS lines are checked by running the tests on
|
||||
a Mac. `examples/cohere2_reference.rs` generates the lines, at commit
|
||||
30c4cf1ee, for a platform not yet in the file. **What still needs the checkpoint:** identity at North's
|
||||
own sizes (hidden 2048, 128 experts with 8 routed, a 262144-row
|
||||
lm_head, contexts past the 4096-token window) and on the five prompts
|
||||
above, whose text and hash sequences are not stored, and every speed
|
||||
figure here. The per-run logs are not in the repository.
|
||||
|
||||
**Speed.** Medians of 4 runs (the 72-token identity run and three
|
||||
48-token timing runs; base at 12 threads had only the three), base
|
||||
and new interleaved in the same conditions, each at its default
|
||||
thread count unless the column says otherwise: base at rayon's
|
||||
default, every logical CPU (32 on the M3 Ultra, 16 on the M4 Max),
|
||||
the new code at one thread per performance core (16, its cap, on the
|
||||
M3 Ultra; 12 on the M4 Max). M3 Ultra (24P+8E): shared, load average
|
||||
30–130, with another North benchmark and CI running. M4 Max
|
||||
(12P+4E): lightly loaded (load average 7–29).
|
||||
|
||||
| | M3 Ultra base, 32 thr | M3 Ultra base, 12 thr | M3 Ultra new, 16 thr | M4 Max base, 16 thr | M4 Max base, 12 thr | M4 Max new, 12 thr |
|
||||
|---|---|---|---|---|---|---|
|
||||
| decode at ~600 context | 267 ms | 168 ms | **54 ms** | 116 ms | 106 ms | **35 ms** |
|
||||
| decode at ~2700 context | 405 ms | 364 ms | **80 ms** | 246 ms | 224 ms | **48 ms** |
|
||||
| prefill, 596 tokens | 9.9 tok/s | 12.0 tok/s | **91 tok/s** | 17.8 tok/s | 19.8 tok/s | **105 tok/s** |
|
||||
| prefill, 2698 tokens | 5.9 tok/s | 4.5 tok/s | **63 tok/s** | 9.0 tok/s | 9.2 tok/s | **79 tok/s** |
|
||||
|
||||
Paired run-for-run ratios against base at its default thread count
|
||||
(base / new, all four pairs of each cell). M3 Ultra decode 4.3–4.8×
|
||||
at ~600 and 5.0–6.6× at ~2700, apart from one pair in each column,
|
||||
20.2× and 21.2×, where base's run took 3.6× and 5.8× its median
|
||||
step (967 ms at load 71, 2358 ms at load 131); those two runs are
|
||||
inside the table's medians.
|
||||
Prefill 5.4–10.3× and 9.7–11.8×. M4 Max decode 3.3–3.5× at ~600
|
||||
apart from the identity run's pair, 2.0× (the new code's median step
|
||||
was 59.7 ms in that run, 34–37 ms in the others), and 5.1–6.3× at
|
||||
~2700; prefill 5.2–6.8× and 7.0–9.0×.
|
||||
|
||||
Against base at 12 threads (the three timing pairs): M3 Ultra decode
|
||||
2.9–3.8× and 3.3–5.8×, prefill 5.9–10.6× and 11.9–14.1×; M4 Max
|
||||
decode 3.1× and 4.5–5.2×, prefill 5.8–6.1× and 8.4–9.5×. Base at 8
|
||||
threads, its fastest in the thread sweep above (94 ms), was not run
|
||||
in pairs with the new code: against this table's 54 ms that is about
|
||||
1.7×, but measured under other load, so it is not a ratio this
|
||||
benchmark establishes.
|
||||
|
||||
A 69-token follow-up turn fed into a ~630-token context
|
||||
(`north_bench --then` and `--then-seq`): `forward_extend` 1.23 s
|
||||
(M3 Ultra) / 0.79–0.80 s (M4 Max), against 3.3–4.1 s / 2.4 s one
|
||||
`forward_token` at a time on the new code, and an estimated ~18 s at
|
||||
base's decode speed (69 × its 267 ms median step; not measured);
|
||||
same logits. A caller gets the first figure only by
|
||||
feeding its turn through `forward_extend`; one that feeds
|
||||
`forward_token` per token gets the second.
|
||||
|
||||
**Did not help, or not enough:** more threads (base and new were
|
||||
both slower above 16 under load); for decode, running each step
|
||||
inside a rayon pool rather than handing each of its loops to the
|
||||
pool (no clear win once measured in-process; decode's loops now run
|
||||
on the spin pool, while prefill, whose projections and experts are
|
||||
rayon loops, does run inside its pool, `in_cpu_pool`); a rayon pool
|
||||
of any size against the spin pool for decode (2.5–4.4× slower under
|
||||
load: 164–465 ms against 65–106 on an M3 Ultra; on a quiet M4 Max
|
||||
40–42 ms against 34 with the same 12 threads, 18–24% slower);
|
||||
prefetch distances other than 1 KiB (256 B–4 KiB swept).
|
||||
|
||||
**Next levers:** the sliding layers trim their KV cache by moving
|
||||
the whole window each token past 4096 positions (a ring buffer
|
||||
removes ~0.6 GB of copies per token there); decode now moves ~6.5
|
||||
GB of weights per token near what the CPU side reads from memory,
|
||||
so the next step is fewer bytes (narrower weights, which would
|
||||
change logits) or several sequences sharing each weight read; and
|
||||
prefill experts are bound by the 4 × 2 NEON block's arithmetic.
|
||||
|
||||
## Cohere2-MoE / North Mini Code (wave 60)
|
||||
|
||||
Fourth topology, first attention-only one:
|
||||
|
|
|
|||
|
|
@ -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<f32> {
|
||||
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<u64>, Vec<f32>)> = Vec::new();
|
||||
let mut push = |name: String, shape: Vec<u64>| {
|
||||
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<u8> = 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: `<model> <output> <hash>`.
|
||||
fn scenario(name: &str, model: &Cohere2MoeModel, shape: &Shape) -> Vec<String> {
|
||||
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<u32> = (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}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,184 @@
|
|||
//! Speed and identity harness for Cohere2-MoE (North Mini Code)
|
||||
//! decode: one prompt from a file, batched prefill, greedy
|
||||
//! decode, and per step the token id, a hash of every logit's
|
||||
//! bits (so two builds can be compared bit for bit, not only
|
||||
//! token for token) and the step's wall-clock.
|
||||
//!
|
||||
//! ```bash
|
||||
//! cargo run --release --features rayon --example north_bench -- \
|
||||
//! /path/to/snapshot-dir prompt.txt [n_tokens] [--chat]
|
||||
//! ```
|
||||
//!
|
||||
//! `--chat` wraps the file's text as one user turn of the
|
||||
//! model's chat template, the shape the program-writing loop
|
||||
//! feeds. `--then <file>` afterwards feeds that file's text as a
|
||||
//! further user turn into the same context with
|
||||
//! `forward_extend` (`--then-seq <file>`: one `forward_token` per
|
||||
//! token, the same logits) and greedily decodes 8 more tokens.
|
||||
//! With `JOULECLAW_PROF=1` it also prints where the time went, op
|
||||
//! by op, for the prefill and per decoded token.
|
||||
|
||||
use jouleclaw_loader_gguf::bpe_tokenizer::BpeTokenizer;
|
||||
use jouleclaw_loader_gguf::cohere2_moe::{prof, Cohere2MoeConfig, Cohere2MoeModel};
|
||||
use jouleclaw_loader_gguf::safetensors_reader::read_safetensors_file;
|
||||
use std::path::PathBuf;
|
||||
use std::time::Instant;
|
||||
|
||||
/// FNV-1a over the bit patterns of every logit.
|
||||
fn logits_hash(logits: &[f32]) -> u64 {
|
||||
let mut h = 0xcbf2_9ce4_8422_2325u64;
|
||||
for v in logits {
|
||||
for b in v.to_bits().to_le_bytes() {
|
||||
h = (h ^ b as u64).wrapping_mul(0x0000_0100_0000_01b3);
|
||||
}
|
||||
}
|
||||
h
|
||||
}
|
||||
|
||||
fn print_prof(label: &str, per: f64) {
|
||||
if !prof::enabled() {
|
||||
return;
|
||||
}
|
||||
let rows = prof::take();
|
||||
let total: f64 = rows.iter().map(|r| r.1).sum();
|
||||
println!("profile {label} (ms{}):", if per > 1.0 { " per token" } else { "" });
|
||||
for (name, ms) in rows.iter().filter(|r| r.1 > 0.0) {
|
||||
println!(" {:<34} {:>10.2} {:>5.1}%", name, ms / per, 100.0 * ms / total.max(1e-9));
|
||||
}
|
||||
println!(" {:<34} {:>10.2}", "total", total / per);
|
||||
}
|
||||
|
||||
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
if args.len() < 3 {
|
||||
eprintln!("usage: north_bench <snapshot-dir> <prompt-file> [n_tokens] [--chat]");
|
||||
std::process::exit(1);
|
||||
}
|
||||
let dir = PathBuf::from(&args[1]);
|
||||
let text = std::fs::read_to_string(&args[2])?;
|
||||
let n_tokens: usize = args.get(3).map(|s| s.parse()).transpose()?.unwrap_or(64);
|
||||
let chat = args.iter().any(|a| a == "--chat");
|
||||
let then = |flag: &str| args.iter().position(|a| a == flag).and_then(|i| args.get(i + 1)).cloned();
|
||||
let (then_file, then_seq) = match (then("--then"), then("--then-seq")) {
|
||||
(Some(f), _) => (Some(f), false),
|
||||
(None, Some(f)) => (Some(f), true),
|
||||
(None, None) => (None, false),
|
||||
};
|
||||
|
||||
let config: Cohere2MoeConfig = serde_json::from_slice(&std::fs::read(dir.join("config.json"))?)?;
|
||||
let mut shard_paths: Vec<PathBuf> = 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::<Result<Vec<_>, _>>()?;
|
||||
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::<f64>() / 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::<Vec<_>>().join(","));
|
||||
println!("hashes {}", hashes.iter().map(|h| format!("{h:016x}")).collect::<Vec<_>>().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::<Vec<_>>().join(","), logits_hash(&logits));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
|
@ -161,6 +161,210 @@ pub fn attention_forward_f32(
|
|||
}
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
// One query row, bit for bit, fast (2026-11)
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
/// One query row of [`attention_forward_f32`] against keys and
|
||||
/// values `0..kv_len` (`k`, `v` hold at least `kv_len` rows),
|
||||
/// with exactly its arithmetic: the same scores, max, softmax
|
||||
/// and weighted sum of V, every value rounded at the same step
|
||||
/// in the same order, so `out_row` is bit-identical. A decode
|
||||
/// step calls it once per head (the heads run in parallel); a
|
||||
/// prefill once per (head, query row), where `kv_len = t + 1`
|
||||
/// is the causal mask.
|
||||
///
|
||||
/// What makes it fast without changing a bit: the reference
|
||||
/// adds each score's 128 products one at a time on one scalar
|
||||
/// chain (an f32 sum may not be reordered), so it waits on the
|
||||
/// add's latency for every product. Here the products of four
|
||||
/// keys sit in the four lanes of a NEON register (a 4×4
|
||||
/// transpose of four key rows) and four such registers run at
|
||||
/// once, so sixteen scores advance together, each lane still
|
||||
/// adding its key's products in `d` order: multiply, round,
|
||||
/// add, round, as the scalar code does (Rust never fuses them).
|
||||
/// The weighted sum of V keeps its per-`d` running sums in
|
||||
/// registers across keys, again multiply then add in key order.
|
||||
/// `scores` is scratch of at least `kv_len`.
|
||||
pub fn attention_row_f32(
|
||||
q_row: &[f32],
|
||||
k: &[f32],
|
||||
v: &[f32],
|
||||
kv_len: usize,
|
||||
d_head: usize,
|
||||
scores: &mut [f32],
|
||||
out_row: &mut [f32],
|
||||
) {
|
||||
assert_eq!(q_row.len(), d_head);
|
||||
assert!(k.len() >= kv_len * d_head);
|
||||
assert!(v.len() >= kv_len * d_head);
|
||||
assert!(scores.len() >= kv_len);
|
||||
assert_eq!(out_row.len(), d_head);
|
||||
assert!(d_head > 0);
|
||||
let inv_sqrt_d = 1.0_f32 / (d_head as f32).sqrt();
|
||||
let scores = &mut scores[..kv_len];
|
||||
qk_scores(q_row, k, kv_len, d_head, inv_sqrt_d, scores);
|
||||
let mut max_score = f32::NEG_INFINITY;
|
||||
for &score in scores.iter() {
|
||||
if score > max_score {
|
||||
max_score = score;
|
||||
}
|
||||
}
|
||||
let mut sum = 0.0f32;
|
||||
for s in scores.iter_mut() {
|
||||
let e = (*s - max_score).exp();
|
||||
*s = e;
|
||||
sum += e;
|
||||
}
|
||||
let inv_sum = if sum > 0.0 { 1.0 / sum } else { 0.0 };
|
||||
for s in scores.iter_mut() {
|
||||
*s *= inv_sum;
|
||||
}
|
||||
weighted_v(scores, v, d_head, out_row);
|
||||
}
|
||||
|
||||
/// `scores[s] = (Σ_d q[d]·k[s][d], summed in d order) · inv_sqrt_d`.
|
||||
fn qk_scores(q: &[f32], k: &[f32], kv_len: usize, d_head: usize, inv_sqrt_d: f32, scores: &mut [f32]) {
|
||||
#[cfg(target_arch = "aarch64")]
|
||||
let s0 = if d_head.is_multiple_of(4) { qk_scores_neon(q, k, kv_len, d_head, inv_sqrt_d, scores) } else { 0 };
|
||||
#[cfg(not(target_arch = "aarch64"))]
|
||||
let s0 = 0;
|
||||
for s in s0..kv_len {
|
||||
let k_row = &k[s * d_head..(s + 1) * d_head];
|
||||
let mut acc = 0.0f32;
|
||||
for d in 0..d_head {
|
||||
acc += q[d] * k_row[d];
|
||||
}
|
||||
scores[s] = acc * inv_sqrt_d;
|
||||
}
|
||||
}
|
||||
|
||||
/// The NEON part of [`qk_scores`]: keys in blocks of 16, then 4.
|
||||
/// Returns how many leading keys it scored (`d_head % 4 == 0`).
|
||||
#[cfg(target_arch = "aarch64")]
|
||||
fn qk_scores_neon(q: &[f32], k: &[f32], kv_len: usize, d_head: usize, inv_sqrt_d: f32, scores: &mut [f32]) -> usize {
|
||||
let mut s0 = 0;
|
||||
{
|
||||
// SAFETY: rows s0..s0+16 (or +4) are below kv_len, so
|
||||
// every load stays inside k[..kv_len * d_head]; q has
|
||||
// d_head elements.
|
||||
unsafe {
|
||||
while s0 + 16 <= kv_len {
|
||||
let acc = qk_block::<4>(q.as_ptr(), k.as_ptr().add(s0 * d_head), d_head);
|
||||
for (b, a) in acc.iter().enumerate() {
|
||||
let mut lanes = [0f32; 4];
|
||||
std::arch::aarch64::vst1q_f32(lanes.as_mut_ptr(), *a);
|
||||
for (l, &x) in lanes.iter().enumerate() {
|
||||
scores[s0 + 4 * b + l] = x * inv_sqrt_d;
|
||||
}
|
||||
}
|
||||
s0 += 16;
|
||||
}
|
||||
while s0 + 4 <= kv_len {
|
||||
let acc = qk_block::<1>(q.as_ptr(), k.as_ptr().add(s0 * d_head), d_head);
|
||||
let mut lanes = [0f32; 4];
|
||||
std::arch::aarch64::vst1q_f32(lanes.as_mut_ptr(), acc[0]);
|
||||
for (l, &x) in lanes.iter().enumerate() {
|
||||
scores[s0 + l] = x * inv_sqrt_d;
|
||||
}
|
||||
s0 += 4;
|
||||
}
|
||||
}
|
||||
}
|
||||
s0
|
||||
}
|
||||
|
||||
/// `B` blocks of four consecutive key rows starting at `k`:
|
||||
/// lane `l` of block `b` is key `4b + l`'s dot product with `q`,
|
||||
/// accumulated `acc = acc + q[d]·k[d]` (separately rounded) in
|
||||
/// `d` order.
|
||||
#[cfg(target_arch = "aarch64")]
|
||||
#[inline(always)]
|
||||
unsafe fn qk_block<const B: usize>(
|
||||
q: *const f32,
|
||||
k: *const f32,
|
||||
d_head: usize,
|
||||
) -> [std::arch::aarch64::float32x4_t; B] {
|
||||
use std::arch::aarch64::*;
|
||||
unsafe {
|
||||
let mut acc = [vdupq_n_f32(0.0); B];
|
||||
let mut d = 0;
|
||||
while d < d_head {
|
||||
let qv = vld1q_f32(q.add(d));
|
||||
for b in 0..B {
|
||||
let base = k.add(4 * b * d_head + d);
|
||||
let r0 = vld1q_f32(base);
|
||||
let r1 = vld1q_f32(base.add(d_head));
|
||||
let r2 = vld1q_f32(base.add(2 * d_head));
|
||||
let r3 = vld1q_f32(base.add(3 * d_head));
|
||||
// 4×4 transpose: c_j lane l = key l's element d + j.
|
||||
let t0 = vtrn1q_f32(r0, r1);
|
||||
let t1 = vtrn2q_f32(r0, r1);
|
||||
let t2 = vtrn1q_f32(r2, r3);
|
||||
let t3 = vtrn2q_f32(r2, r3);
|
||||
let c0 = vreinterpretq_f32_f64(vtrn1q_f64(vreinterpretq_f64_f32(t0), vreinterpretq_f64_f32(t2)));
|
||||
let c1 = vreinterpretq_f32_f64(vtrn1q_f64(vreinterpretq_f64_f32(t1), vreinterpretq_f64_f32(t3)));
|
||||
let c2 = vreinterpretq_f32_f64(vtrn2q_f64(vreinterpretq_f64_f32(t0), vreinterpretq_f64_f32(t2)));
|
||||
let c3 = vreinterpretq_f32_f64(vtrn2q_f64(vreinterpretq_f64_f32(t1), vreinterpretq_f64_f32(t3)));
|
||||
let mut a = acc[b];
|
||||
a = vaddq_f32(a, vmulq_laneq_f32::<0>(c0, qv));
|
||||
a = vaddq_f32(a, vmulq_laneq_f32::<1>(c1, qv));
|
||||
a = vaddq_f32(a, vmulq_laneq_f32::<2>(c2, qv));
|
||||
a = vaddq_f32(a, vmulq_laneq_f32::<3>(c3, qv));
|
||||
acc[b] = a;
|
||||
}
|
||||
d += 4;
|
||||
}
|
||||
acc
|
||||
}
|
||||
}
|
||||
|
||||
/// `out[d] = Σ_s w[s]·v[s][d]`, in key order, skipping zero
|
||||
/// weights — the reference's `out[d] += w·v[d]` loop with the
|
||||
/// running sums held in registers, 64 lanes at a time.
|
||||
fn weighted_v(w: &[f32], v: &[f32], d_head: usize, out: &mut [f32]) {
|
||||
#[cfg(target_arch = "aarch64")]
|
||||
if d_head.is_multiple_of(4) {
|
||||
use std::arch::aarch64::*;
|
||||
let mut d0 = 0;
|
||||
while d0 < d_head {
|
||||
let lanes = (d_head - d0).min(64);
|
||||
// SAFETY: d0 + lanes ≤ d_head, and rows s < w.len()
|
||||
// lie inside v (asserted by the caller).
|
||||
unsafe {
|
||||
let mut acc = [vdupq_n_f32(0.0); 16];
|
||||
let nv = lanes / 4;
|
||||
for (s, &weight) in w.iter().enumerate() {
|
||||
if weight == 0.0 {
|
||||
continue;
|
||||
}
|
||||
let row = v.as_ptr().add(s * d_head + d0);
|
||||
for (i, a) in acc.iter_mut().enumerate().take(nv) {
|
||||
*a = vaddq_f32(*a, vmulq_n_f32(vld1q_f32(row.add(4 * i)), weight));
|
||||
}
|
||||
}
|
||||
for (i, a) in acc.iter().enumerate().take(nv) {
|
||||
vst1q_f32(out.as_mut_ptr().add(d0 + 4 * i), *a);
|
||||
}
|
||||
}
|
||||
d0 += lanes;
|
||||
}
|
||||
return;
|
||||
}
|
||||
for o in out.iter_mut() {
|
||||
*o = 0.0;
|
||||
}
|
||||
for (s, &weight) in w.iter().enumerate() {
|
||||
if weight == 0.0 {
|
||||
continue;
|
||||
}
|
||||
let v_row = &v[s * d_head..(s + 1) * d_head];
|
||||
for d in 0..d_head {
|
||||
out[d] += weight * v_row[d];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
// Flash-attention chunked variant (wave 29)
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
|
|
@ -297,6 +501,53 @@ fn as_u8(s: &[f32]) -> &[u8] {
|
|||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn wave(n: usize, seed: usize, scale: f32) -> Vec<f32> {
|
||||
(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
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -0,0 +1,101 @@
|
|||
# The stored reference for the Cohere2-MoE decoder (North Mini Code's
|
||||
# architecture), checked without the checkpoint by the test
|
||||
# cohere2_moe::tests::stored_reference::logits_match_the_stored_reference.
|
||||
#
|
||||
# A line: <os>-<arch> <checkpoint> <output> <FNV-1a 64 of the logits' f32 bits,
|
||||
# little-endian>. The checkpoints are that test's synthetic tiny, medium and
|
||||
# wide configs; the scenario prefills a prompt, decodes, feeds a chunk into
|
||||
# the live state and decodes twice more.
|
||||
#
|
||||
# Generated by examples/cohere2_reference.rs run at commit 30c4cf1ee (main
|
||||
# before PR #649): its forward_prompt and forward_token, with and without the
|
||||
# rayon feature (the lines were identical), on 2026-11-06. That file says how
|
||||
# to add a platform.
|
||||
#
|
||||
# The bits depend on the platform's math library as well as on the code:
|
||||
# x86_64 runs scalar kernels where arm64 runs NEON, and the two arm64 systems,
|
||||
# with the same kernels, still differ on the medium and wide checkpoints. So
|
||||
# each platform's lines name, on their `math-library` line, the library they
|
||||
# were made with. On Linux arm64, glibc 2.35's libm (that of CI's container)
|
||||
# loaded in place of 2.39's gives the same lines.
|
||||
math-library macos-aarch64 the system math library (libsystem_m) of macOS 27.2 (26B123)
|
||||
math-library linux-aarch64 the libm of Ubuntu GLIBC 2.39-0ubuntu8.9
|
||||
math-library linux-x86_64 the libm of Ubuntu GLIBC 2.43-2ubuntu2.4
|
||||
macos-aarch64 tiny prefill 6b3c5457df8a475b
|
||||
macos-aarch64 tiny decode0 6b734576ec270b6a
|
||||
macos-aarch64 tiny decode1 daa682bb61855607
|
||||
macos-aarch64 tiny decode2 a96a525073b3b793
|
||||
macos-aarch64 tiny decode3 9bb8aa8e5d425aba
|
||||
macos-aarch64 tiny extend c3fda182453d886a
|
||||
macos-aarch64 tiny after0 66a97258116e7597
|
||||
macos-aarch64 tiny after1 9d7b37bbb2325b5f
|
||||
macos-aarch64 medium prefill b82cf3338dfa66f0
|
||||
macos-aarch64 medium decode0 c66f9c5495e8969b
|
||||
macos-aarch64 medium decode1 8fe95e4149a712e2
|
||||
macos-aarch64 medium decode2 68f81c6d7139db90
|
||||
macos-aarch64 medium decode3 058a2083f6058017
|
||||
macos-aarch64 medium decode4 b5957c64fb853c5a
|
||||
macos-aarch64 medium decode5 c90b8d2ab8fe7230
|
||||
macos-aarch64 medium extend 6448cf5f1255e448
|
||||
macos-aarch64 medium after0 d491c6a5b890724a
|
||||
macos-aarch64 medium after1 0284c81caf0cbe33
|
||||
macos-aarch64 wide prefill 5756c9cbfe56f1ac
|
||||
macos-aarch64 wide decode0 a61a0a894569af2a
|
||||
macos-aarch64 wide decode1 ae4763855789a897
|
||||
macos-aarch64 wide decode2 86cda1dc39403823
|
||||
macos-aarch64 wide decode3 159cdb1461c22868
|
||||
macos-aarch64 wide extend 772006dffd4eae2f
|
||||
macos-aarch64 wide after0 0dd7d76f2077f8c5
|
||||
macos-aarch64 wide after1 f43a082fc856a8d1
|
||||
linux-aarch64 tiny prefill 6b3c5457df8a475b
|
||||
linux-aarch64 tiny decode0 6b734576ec270b6a
|
||||
linux-aarch64 tiny decode1 daa682bb61855607
|
||||
linux-aarch64 tiny decode2 a96a525073b3b793
|
||||
linux-aarch64 tiny decode3 9bb8aa8e5d425aba
|
||||
linux-aarch64 tiny extend c3fda182453d886a
|
||||
linux-aarch64 tiny after0 66a97258116e7597
|
||||
linux-aarch64 tiny after1 9d7b37bbb2325b5f
|
||||
linux-aarch64 medium prefill e73f43bb8f743349
|
||||
linux-aarch64 medium decode0 efd7dfbbbc45d7a6
|
||||
linux-aarch64 medium decode1 2367b4995fca5e82
|
||||
linux-aarch64 medium decode2 36e603a23c2e410d
|
||||
linux-aarch64 medium decode3 aa3d37da0de5837b
|
||||
linux-aarch64 medium decode4 830cf607dfafe891
|
||||
linux-aarch64 medium decode5 414613d72da1cdb5
|
||||
linux-aarch64 medium extend 0509a4b9f8d8925f
|
||||
linux-aarch64 medium after0 750ebc98f4384676
|
||||
linux-aarch64 medium after1 9161ec76f213720f
|
||||
linux-aarch64 wide prefill 39650607a70ba636
|
||||
linux-aarch64 wide decode0 3396985a061dbd5f
|
||||
linux-aarch64 wide decode1 cab2298070b6e8a5
|
||||
linux-aarch64 wide decode2 c9bccfd56c20432f
|
||||
linux-aarch64 wide decode3 9ffbfed9d2dc8cce
|
||||
linux-aarch64 wide extend a017b0616c1771ab
|
||||
linux-aarch64 wide after0 5929322a051cd849
|
||||
linux-aarch64 wide after1 2d1b2ab4c0a12842
|
||||
linux-x86_64 tiny prefill f256b8a75cf96d05
|
||||
linux-x86_64 tiny decode0 fb05b8bf618b395b
|
||||
linux-x86_64 tiny decode1 0c7e7b75513c1a3b
|
||||
linux-x86_64 tiny decode2 231bbcdb71db9c65
|
||||
linux-x86_64 tiny decode3 0df0412449659098
|
||||
linux-x86_64 tiny extend ae95b2ca568c2802
|
||||
linux-x86_64 tiny after0 93df2260f050c0e7
|
||||
linux-x86_64 tiny after1 3bf419607737244f
|
||||
linux-x86_64 medium prefill 4652623cd80b24e1
|
||||
linux-x86_64 medium decode0 dc02242d91e91dbd
|
||||
linux-x86_64 medium decode1 fcb0e3e903d187d2
|
||||
linux-x86_64 medium decode2 c1c4c18cbb9298f8
|
||||
linux-x86_64 medium decode3 370686ea9a04b164
|
||||
linux-x86_64 medium decode4 800ba78022c1d084
|
||||
linux-x86_64 medium decode5 4a6b5efb48a72994
|
||||
linux-x86_64 medium extend 4efe2d7fd73de655
|
||||
linux-x86_64 medium after0 d4553a9d2cef7e4e
|
||||
linux-x86_64 medium after1 cf0729eb66a67a7f
|
||||
linux-x86_64 wide prefill 3c21f84d4dd6d2e5
|
||||
linux-x86_64 wide decode0 c533e37794392f5d
|
||||
linux-x86_64 wide decode1 929641e6f2ab1eaf
|
||||
linux-x86_64 wide decode2 9677d2965fd35073
|
||||
linux-x86_64 wide decode3 15fb690a596a3419
|
||||
linux-x86_64 wide extend 6851feba7437339f
|
||||
linux-x86_64 wide after0 a0c0eb86889d1563
|
||||
linux-x86_64 wide after1 a2eaa676e2afe1c5
|
||||
|
|
@ -62,6 +62,22 @@ mod enabled {
|
|||
.as_ref()
|
||||
}
|
||||
|
||||
/// True when the operator opted in (`JOULECLAW_FABRIC=1`), so
|
||||
/// decode projections may leave the CPU.
|
||||
pub fn fabric_engaged() -> bool {
|
||||
fabric().is_some()
|
||||
}
|
||||
|
||||
/// True when a batched (prefill) projection stays on the CPU's
|
||||
/// NEON kernels: the fabric is off. Every route off the CPU
|
||||
/// needs the fabric, the Metal-4 tensor path too
|
||||
/// ([`matmul_bf16_mpp_bt`] refuses without it), so
|
||||
/// `JOULECLAW_MPP_PREFILL` plays no part here: it only chooses
|
||||
/// among the fabric's paths once the fabric is on.
|
||||
pub fn prefill_on_cpu() -> bool {
|
||||
fabric().is_none()
|
||||
}
|
||||
|
||||
/// Widen a BF16 row-major block to f32 row-major with
|
||||
/// NEON (`vshll_n_u16` — the bf16→f32 decode IS a 16-bit
|
||||
/// shift), sequential loads AND stores. The scalar
|
||||
|
|
@ -1334,7 +1350,7 @@ mod enabled {
|
|||
|
||||
#[cfg(feature = "fabric")]
|
||||
pub use enabled::{
|
||||
ledger_reset, ledger_snapshot, matmul_bf16_amx, matmul_bf16_auto, matmul_bf16_auto_multi,
|
||||
fabric_engaged, prefill_on_cpu, ledger_reset, ledger_snapshot, matmul_bf16_amx, matmul_bf16_auto, matmul_bf16_auto_multi,
|
||||
matmul_bf16_attn_block, matmul_bf16_gpu, matmul_bf16_mlp_window, matmul_bf16_moe_window,
|
||||
matmul_bf16_mpp_bt, matmul_bf16_norm_qkv_rope_window, matmul_bf16_norm_window, AutoJob,
|
||||
LedgerEntry,
|
||||
|
|
@ -1345,6 +1361,19 @@ pub use enabled::{
|
|||
#[cfg(feature = "fabric-meter")]
|
||||
pub use enabled::{tune_shape_live, tune_shape_live_process};
|
||||
|
||||
/// Without the `fabric` feature nothing leaves the CPU.
|
||||
#[cfg(not(feature = "fabric"))]
|
||||
pub fn fabric_engaged() -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
/// Without the `fabric` feature every prefill projection runs on
|
||||
/// the CPU's NEON kernels.
|
||||
#[cfg(not(feature = "fabric"))]
|
||||
pub fn prefill_on_cpu() -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
/// One op in a fused dispatch window (no-fabric stub).
|
||||
#[cfg(not(feature = "fabric"))]
|
||||
pub struct AutoJob<'a> {
|
||||
|
|
@ -1456,6 +1485,82 @@ pub fn matmul_bf16_auto(
|
|||
}
|
||||
}
|
||||
|
||||
/// Tests whose answer depends on the fabric switch
|
||||
/// (`JOULECLAW_FABRIC`), which a process reads once. Such a test
|
||||
/// runs where the switch is what it needs: in this process if it
|
||||
/// already is, else alone in a child of the test binary. No test sets
|
||||
/// the variable in its own process, so one test's need never decides
|
||||
/// the switch for the others in the binary.
|
||||
#[cfg(test)]
|
||||
pub(crate) mod fabric_switch {
|
||||
/// Set in a child that [`run_alone`] started.
|
||||
const CHILD: &str = "JOULECLAW_TEST_CHILD";
|
||||
|
||||
/// True when the fabric's state in this process is `on`, so the
|
||||
/// calling test goes on here. Otherwise the test runs alone in a
|
||||
/// child with the switch set accordingly and must pass there, and
|
||||
/// this returns false: the caller returns, its check made.
|
||||
/// `module` is the caller's `module_path!()`, `test` its name.
|
||||
pub(crate) fn here(on: bool, module: &str, test: &str) -> bool {
|
||||
if super::fabric_engaged() == on {
|
||||
return true;
|
||||
}
|
||||
run_alone(module, test, &[("JOULECLAW_FABRIC", on.then_some("1"))]);
|
||||
false
|
||||
}
|
||||
|
||||
/// Whether this process is a child that [`run_alone`] started.
|
||||
pub(crate) fn in_child() -> bool {
|
||||
std::env::var_os(CHILD).is_some()
|
||||
}
|
||||
|
||||
/// Run the test `module::test` alone in a child of this test
|
||||
/// binary with `env` set (`None`: removed), fail unless the child
|
||||
/// reports that exactly that one test ran and passed, and return
|
||||
/// what it printed.
|
||||
pub(crate) fn run_alone(module: &str, test: &str, env: &[(&str, Option<&str>)]) -> String {
|
||||
assert!(!in_child(), "{module}::{test}: a child test process would start another, with {env:?}");
|
||||
spawn(module, test, env, true)
|
||||
}
|
||||
|
||||
/// [`run_alone`], but the child is a top-level test process: it
|
||||
/// may start children of its own.
|
||||
#[cfg_attr(not(feature = "fabric"), allow(dead_code))]
|
||||
pub(crate) fn run_as_top_level(module: &str, test: &str, env: &[(&str, Option<&str>)]) -> String {
|
||||
spawn(module, test, env, false)
|
||||
}
|
||||
|
||||
fn spawn(module: &str, test: &str, env: &[(&str, Option<&str>)], child: bool) -> String {
|
||||
// libtest names a test by its module path without the crate.
|
||||
let path = match module.split_once("::") {
|
||||
Some((_, m)) => format!("{m}::{test}"),
|
||||
None => test.to_string(),
|
||||
};
|
||||
let mut cmd = std::process::Command::new(std::env::current_exe().expect("the test binary's path"));
|
||||
cmd.args([path.as_str(), "--exact", "--nocapture", "--test-threads=1"]);
|
||||
if child {
|
||||
cmd.env(CHILD, "1");
|
||||
} else {
|
||||
cmd.env_remove(CHILD);
|
||||
}
|
||||
for &(k, v) in env {
|
||||
match v {
|
||||
Some(v) => cmd.env(k, v),
|
||||
None => cmd.env_remove(k),
|
||||
};
|
||||
}
|
||||
let out = cmd.output().expect("start the test binary");
|
||||
let stdout = String::from_utf8_lossy(&out.stdout).into_owned();
|
||||
assert!(
|
||||
out.status.success() && stdout.contains("test result: ok. 1 passed"),
|
||||
"{path} in a child process with {env:?}: {}\n{stdout}\n{}",
|
||||
out.status,
|
||||
String::from_utf8_lossy(&out.stderr)
|
||||
);
|
||||
stdout
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(all(test, feature = "fabric", target_os = "macos", target_arch = "aarch64"))]
|
||||
mod tests {
|
||||
|
||||
|
|
@ -1579,9 +1684,12 @@ mod tests {
|
|||
/// off / GPU uncertified / pre-Metal-4 (wrapper → false).
|
||||
#[test]
|
||||
fn mpp_bt_matches_neon_when_available() {
|
||||
// SAFETY: single-threaded test setup, before any
|
||||
// fabric() init in this test.
|
||||
unsafe { std::env::set_var("JOULECLAW_FABRIC", "1") };
|
||||
// The switch is read once per process; setting it here
|
||||
// would turn the fabric on for every other test in this
|
||||
// binary that had not read it yet.
|
||||
if !super::fabric_switch::here(true, module_path!(), "mpp_bt_matches_neon_when_available") {
|
||||
return;
|
||||
}
|
||||
// SHAPE SWEEP — square, non-square, and non-tile-aligned
|
||||
// M/N (the square-tiling wrapper must handle them all).
|
||||
let shapes = [
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -132,6 +132,26 @@ pub fn matmul_bf16_scalar(
|
|||
/// replaces eight shift-and-cast scalar ops, then two
|
||||
/// `vfmaq_f32` accumulate.
|
||||
///
|
||||
/// ## Every output is the one-row dot product, bit for bit
|
||||
///
|
||||
/// Each `out[i, j]` is accumulated exactly as
|
||||
/// [`matmul_bf16_neon_one_row`] does it: lane `l` of `acc0`
|
||||
/// holds the fused multiply-adds of elements `8c + l` and lane
|
||||
/// `l` of `acc1` those of `8c + 4 + l`, in chunk order; then
|
||||
/// `vaddvq(acc0) + vaddvq(acc1)`, then the `k % 8` tail added
|
||||
/// one product at a time, then the bias. What changed (2026-11)
|
||||
/// is only how many of those dot products run at once: a block
|
||||
/// of 4 activation rows × 2 weight rows, 2 × 4 (2 or 3 rows
|
||||
/// left) or 1 × 4 (1 row left), with 1 weight row at a ragged
|
||||
/// edge, shares its loads and keeps up to 16 independent
|
||||
/// accumulator chains (8 outputs, two chains each) in flight.
|
||||
/// The one-row loop has two chains, so it was bound by the
|
||||
/// fused multiply-add's latency (about 16 GB/s of weight per
|
||||
/// core on an M3 Ultra) rather than by memory. Weight rows are
|
||||
/// walked in tiles of at most 64 KiB, so a tile stays in L1
|
||||
/// while every activation row passes over it, and each row is
|
||||
/// prefetched 1 KiB ahead.
|
||||
///
|
||||
/// Falls back to [`matmul_bf16_scalar`] off aarch64.
|
||||
#[cfg(target_arch = "aarch64")]
|
||||
pub fn matmul_bf16_neon(
|
||||
|
|
@ -140,6 +160,214 @@ pub fn matmul_bf16_neon(
|
|||
bias: &[f32],
|
||||
m: usize, n: usize, k: usize,
|
||||
out: &mut [f32],
|
||||
) {
|
||||
assert_eq!(a.len(), m * k);
|
||||
assert_eq!(w.len(), 2 * n * k);
|
||||
assert_eq!(out.len(), m * n);
|
||||
assert!(bias.is_empty() || bias.len() == n);
|
||||
// SAFETY: the asserts above bound every access the blocked
|
||||
// kernel makes: rows `0..m` of `a` (`k` f32 each), rows
|
||||
// `0..n` of `w` (`2k` bytes each) and `m * n` outputs at
|
||||
// stride `n`.
|
||||
unsafe { neon_blocked::matmul(a.as_ptr(), w.as_ptr(), bias, m, n, k, out.as_mut_ptr(), n) }
|
||||
}
|
||||
|
||||
/// The blocked BF16 kernel behind [`matmul_bf16_neon`] and the
|
||||
/// threaded tiles: register blocks of activation rows × weight
|
||||
/// rows whose every dot product keeps the one-row loop's exact
|
||||
/// accumulation order (see [`matmul_bf16_neon`]).
|
||||
#[cfg(target_arch = "aarch64")]
|
||||
mod neon_blocked {
|
||||
use super::bf16_to_f32;
|
||||
use std::arch::aarch64::*;
|
||||
|
||||
/// Software prefetch distance in each weight row. Measured
|
||||
/// (2026-11, one core streaming 1 GB of BF16 from DRAM): with
|
||||
/// no prefetch the 4-row block reads 14–15 GB/s on an M3
|
||||
/// Ultra and 19 on an M4 Max, against 16–17 and 26 at 1 KiB
|
||||
/// ahead; 256 B to 4 KiB were swept.
|
||||
const PREFETCH_BYTES: usize = 1024;
|
||||
|
||||
/// The chunked (8-lane) part of `MR × NR` dot products:
|
||||
/// activation rows `a + i*lda`, weight rows `w + j*ldw`
|
||||
/// bytes, `chunks` 8-element chunks. Per output, lanes and
|
||||
/// order are the one-row loop's.
|
||||
#[inline(always)]
|
||||
unsafe fn block<const MR: usize, const NR: usize>(
|
||||
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<const MR: usize, const NR: usize>(
|
||||
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::<MR, NR>(a0, k, w0, 2 * k, k / 8);
|
||||
for i in 0..MR {
|
||||
for j in 0..NR {
|
||||
let v = finish(r[i][j], a0.add(i * k), w0.add(j * 2 * k), k, bias, j0 + j);
|
||||
*out.add((i0 + i) * ldo + j0 + j) = v;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Weight rows `0..n`, activation rows `0..m`; output
|
||||
/// `out[i * ldo + j]`. `bias` is indexed by `j` (empty for
|
||||
/// none).
|
||||
///
|
||||
/// # Safety
|
||||
/// `a` holds `m * k` f32, `w` holds `2 * n * k` bytes, and
|
||||
/// `out` is writable at every `i * ldo + j`.
|
||||
pub unsafe fn matmul(
|
||||
a: *const f32,
|
||||
w: *const u8,
|
||||
bias: &[f32],
|
||||
m: usize,
|
||||
n: usize,
|
||||
k: usize,
|
||||
out: *mut f32,
|
||||
ldo: usize,
|
||||
) {
|
||||
unsafe {
|
||||
// Weight rows per L1 tile: at most 64 KiB, a multiple
|
||||
// of the 4-row block.
|
||||
let tile = ((64 * 1024) / (2 * k).max(1)).clamp(4, 256) / 4 * 4;
|
||||
let mut j0 = 0;
|
||||
while j0 < n {
|
||||
let jt = (n - j0).min(tile);
|
||||
let mut i = 0;
|
||||
while i < m {
|
||||
let mr = (m - i).min(4);
|
||||
let mut j = j0;
|
||||
let jend = j0 + jt;
|
||||
match mr {
|
||||
4 => {
|
||||
while j + 2 <= jend {
|
||||
run::<4, 2>(a, w, bias, k, i, j, out, ldo);
|
||||
j += 2;
|
||||
}
|
||||
while j < jend {
|
||||
run::<4, 1>(a, w, bias, k, i, j, out, ldo);
|
||||
j += 1;
|
||||
}
|
||||
}
|
||||
2 | 3 => {
|
||||
while j + 4 <= jend {
|
||||
run::<2, 4>(a, w, bias, k, i, j, out, ldo);
|
||||
j += 4;
|
||||
}
|
||||
while j < jend {
|
||||
run::<2, 1>(a, w, bias, k, i, j, out, ldo);
|
||||
j += 1;
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
while j + 4 <= jend {
|
||||
run::<1, 4>(a, w, bias, k, i, j, out, ldo);
|
||||
j += 4;
|
||||
}
|
||||
while j < jend {
|
||||
run::<1, 1>(a, w, bias, k, i, j, out, ldo);
|
||||
j += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
i += if mr == 3 { 2 } else { mr };
|
||||
}
|
||||
j0 += jt;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The pre-2026-11 one-row NEON loop, kept as the reference the
|
||||
/// blocked [`matmul_bf16_neon`] is tested against bit for bit.
|
||||
#[cfg(target_arch = "aarch64")]
|
||||
pub fn matmul_bf16_neon_one_row(
|
||||
a: &[f32],
|
||||
w: &[u8],
|
||||
bias: &[f32],
|
||||
m: usize, n: usize, k: usize,
|
||||
out: &mut [f32],
|
||||
) {
|
||||
assert_eq!(a.len(), m * k);
|
||||
assert_eq!(w.len(), 2 * n * k);
|
||||
|
|
@ -205,6 +433,216 @@ pub fn matmul_bf16_neon(
|
|||
matmul_bf16_scalar(a, w, bias, m, n, k, out)
|
||||
}
|
||||
|
||||
/// Non-aarch64 fallback: delegates to the scalar kernel.
|
||||
#[cfg(not(target_arch = "aarch64"))]
|
||||
pub fn matmul_bf16_neon_one_row(
|
||||
a: &[f32],
|
||||
w: &[u8],
|
||||
bias: &[f32],
|
||||
m: usize, n: usize, k: usize,
|
||||
out: &mut [f32],
|
||||
) {
|
||||
matmul_bf16_scalar(a, w, bias, m, n, k, out)
|
||||
}
|
||||
|
||||
/// `[m, n]` output tile `rows × cols` written in place at
|
||||
/// stride `ldo`: the per-tile body of the threaded kernels.
|
||||
///
|
||||
/// # Safety
|
||||
/// `a` holds `m * k` f32 (only rows `i0..i0+rows` are read),
|
||||
/// `w` holds the tile's `cols` weight rows (`2 * cols * k`
|
||||
/// bytes), and `out` is writable at `(i0 + i) * ldo + j` for
|
||||
/// every `i < rows`, `j < cols`, with no other writer.
|
||||
#[cfg(feature = "rayon")]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
unsafe fn matmul_bf16_tile(
|
||||
a: &[f32],
|
||||
w: &[u8],
|
||||
bias: &[f32],
|
||||
i0: usize,
|
||||
rows: usize,
|
||||
cols: usize,
|
||||
k: usize,
|
||||
out: *mut f32,
|
||||
ldo: usize,
|
||||
) {
|
||||
#[cfg(target_arch = "aarch64")]
|
||||
unsafe {
|
||||
neon_blocked::matmul(a.as_ptr().add(i0 * k), w.as_ptr(), bias, rows, cols, k, out.add(i0 * ldo), ldo);
|
||||
}
|
||||
#[cfg(not(target_arch = "aarch64"))]
|
||||
{
|
||||
let mut buf = vec![0f32; rows * cols];
|
||||
matmul_bf16_scalar(&a[i0 * k..(i0 + rows) * k], w, bias, rows, cols, k, &mut buf);
|
||||
for i in 0..rows {
|
||||
for j in 0..cols {
|
||||
unsafe { *out.add((i0 + i) * ldo + j) = buf[i * cols + j] };
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A raw output pointer handed to the tiles of one matmul; each
|
||||
/// tile writes a disjoint `(row, column)` block.
|
||||
#[cfg(feature = "rayon")]
|
||||
#[derive(Clone, Copy)]
|
||||
struct OutPtr(*mut f32);
|
||||
// SAFETY: every tile of one call writes only its own block, and
|
||||
// the call returns after all of them finish.
|
||||
#[cfg(feature = "rayon")]
|
||||
unsafe impl Send for OutPtr {}
|
||||
#[cfg(feature = "rayon")]
|
||||
unsafe impl Sync for OutPtr {}
|
||||
#[cfg(feature = "rayon")]
|
||||
impl OutPtr {
|
||||
/// The pointer. A tile's closure reaches it through this
|
||||
/// method, never through the field: naming the field would
|
||||
/// capture the raw pointer alone, which is neither `Send`
|
||||
/// nor `Sync`, where a method call captures the `OutPtr`.
|
||||
fn get(self) -> *mut f32 {
|
||||
self.0
|
||||
}
|
||||
}
|
||||
|
||||
/// The thread pool the CPU forward passes run in: its size is
|
||||
/// `JOULECLAW_THREADS`, else `RAYON_NUM_THREADS`, else
|
||||
/// [`default_cpu_threads`].
|
||||
#[cfg(feature = "rayon")]
|
||||
pub fn cpu_pool() -> &'static rayon::ThreadPool {
|
||||
static POOL: std::sync::OnceLock<rayon::ThreadPool> = 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::<usize>().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<usize> {
|
||||
let mut v: libc::c_int = 0;
|
||||
let mut len = std::mem::size_of::<libc::c_int>();
|
||||
// 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<usize> {
|
||||
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<R: Send>(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<usize> = std::sync::OnceLock::new();
|
||||
let kb = *KB.get_or_init(|| {
|
||||
std::env::var("JOULECLAW_CHUNK_KB")
|
||||
.ok()
|
||||
.and_then(|v| v.parse::<usize>().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<GemvChunk<'a>>,
|
||||
a: &'a [f32],
|
||||
w: &'a [u8],
|
||||
n: usize,
|
||||
k: usize,
|
||||
out: &'a mut [f32],
|
||||
) {
|
||||
assert_eq!(a.len(), k);
|
||||
assert_eq!(w.len(), 2 * n * k);
|
||||
assert_eq!(out.len(), n);
|
||||
let rows = gemv_rows_per_chunk(k);
|
||||
for (c, o) in out.chunks_mut(rows).enumerate() {
|
||||
let j0 = c * rows;
|
||||
let nr = o.len();
|
||||
into.push(GemvChunk { a, w: &w[j0 * 2 * k..(j0 + nr) * 2 * k], k, out: o });
|
||||
}
|
||||
}
|
||||
|
||||
/// BF16 matmul with rayon multi-thread fan-out — wave 25
|
||||
/// dispatch shape: `m == 1` parallel over `n` (independent
|
||||
/// dot products), `m > 1` parallel over `m` (contiguous
|
||||
|
|
@ -257,8 +695,12 @@ pub fn matmul_bf16_threaded(
|
|||
return Ok(());
|
||||
}
|
||||
if m == 1 {
|
||||
// 2026-11: many small chunks instead of one fixed share
|
||||
// per thread (see [`gemv_rows_per_chunk`]), at least
|
||||
// four per thread. Each output is the same dot product
|
||||
// whatever the split.
|
||||
let threads = rayon::current_num_threads().max(1);
|
||||
let chunk_n = n.div_ceil(threads).max(1);
|
||||
let chunk_n = gemv_rows_per_chunk(k).min(n.div_ceil(4 * threads).div_ceil(4) * 4).max(4);
|
||||
out.par_chunks_mut(chunk_n).enumerate().for_each(
|
||||
|(idx, out_chunk)| {
|
||||
let start_j = idx * chunk_n;
|
||||
|
|
@ -273,59 +715,31 @@ pub fn matmul_bf16_threaded(
|
|||
},
|
||||
);
|
||||
} else {
|
||||
let threads = rayon::current_num_threads().max(1);
|
||||
// Axis choice is a bandwidth question (wave 54).
|
||||
// Parallel-over-m makes EVERY thread stream the whole
|
||||
// [n,k] weight matrix — traffic = threads × weights,
|
||||
// which is exactly wrong for the small-m prefill shape
|
||||
// where weights dominate. Parallel-over-n reads the
|
||||
// weights once in total but every thread reads all of
|
||||
// A [m,k]. Crossover: m·k·4 (A per thread) vs
|
||||
// n·k·2/threads (weights per thread) → prefer the n
|
||||
// axis while m < n / (2·threads).
|
||||
if m < n / (2 * threads).max(1) {
|
||||
// Parallel over n. Each thread computes [m, chunk_n]
|
||||
// into its own buffer (the kernel's natural output
|
||||
// shape for a weight-row slice), then scatters rows
|
||||
// into the strided columns of out — the scatter is
|
||||
// O(m·n) f32 moves, noise next to the matmul.
|
||||
let chunk_n = n.div_ceil(threads).max(1);
|
||||
let pieces: Vec<(usize, usize, Vec<f32>)> = (0..n.div_ceil(chunk_n))
|
||||
.into_par_iter()
|
||||
.map(|idx| {
|
||||
let start_j = idx * chunk_n;
|
||||
let this_n = chunk_n.min(n - start_j);
|
||||
let w_chunk =
|
||||
&w[start_j * row_bytes..(start_j + this_n) * row_bytes];
|
||||
let bias_chunk: &[f32] = if bias.is_empty() {
|
||||
&[]
|
||||
} else {
|
||||
&bias[start_j..start_j + this_n]
|
||||
};
|
||||
let mut buf = vec![0f32; m * this_n];
|
||||
matmul_bf16_neon(a, w_chunk, bias_chunk, m, this_n, k, &mut buf);
|
||||
(start_j, this_n, buf)
|
||||
})
|
||||
.collect();
|
||||
for (start_j, this_n, buf) in pieces {
|
||||
for i in 0..m {
|
||||
out[i * n + start_j..i * n + start_j + this_n]
|
||||
.copy_from_slice(&buf[i * this_n..(i + 1) * this_n]);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Large m: A traffic dominates — parallel over m,
|
||||
// each thread a contiguous row chunk.
|
||||
let chunk_m = m.div_ceil(threads).max(1);
|
||||
out.par_chunks_mut(chunk_m * n).enumerate().for_each(
|
||||
|(idx, out_chunk)| {
|
||||
let start_i = idx * chunk_m;
|
||||
let this_m = out_chunk.len() / n;
|
||||
let a_chunk = &a[start_i * k..(start_i + this_m) * k];
|
||||
matmul_bf16_neon(a_chunk, w, bias, this_m, n, k, out_chunk);
|
||||
},
|
||||
);
|
||||
}
|
||||
// 2026-11: 2-D tiles. Each task owns a block of up to 64
|
||||
// activation rows × 64 weight rows and writes it in
|
||||
// place, so the weights are read about m/64 times in
|
||||
// total instead of once per activation row (the old
|
||||
// parallel-over-m split streamed every weight per row,
|
||||
// and parallel-over-n gave each thread all of A). Inside
|
||||
// a tile the blocked kernel keeps every output's
|
||||
// one-row accumulation order.
|
||||
const TM: usize = 64;
|
||||
let tn = 64usize;
|
||||
let ptr = OutPtr(out.as_mut_ptr());
|
||||
let n_tiles = n.div_ceil(tn);
|
||||
(0..m.div_ceil(TM) * n_tiles).into_par_iter().for_each(|t| {
|
||||
let (bi, bj) = (t / n_tiles, t % n_tiles);
|
||||
let i0 = bi * TM;
|
||||
let rows = TM.min(m - i0);
|
||||
let j0 = bj * tn;
|
||||
let cols = tn.min(n - j0);
|
||||
let w_tile = &w[j0 * row_bytes..(j0 + cols) * row_bytes];
|
||||
let b_tile: &[f32] = if bias.is_empty() { &[] } else { &bias[j0..j0 + cols] };
|
||||
// SAFETY: tile (bi, bj) writes rows i0..i0+rows of
|
||||
// columns j0..j0+cols of `out` ([m, n], stride n),
|
||||
// which no other tile touches; `a` is [m, k].
|
||||
unsafe { matmul_bf16_tile(a, w_tile, b_tile, i0, rows, cols, k, ptr.get().add(j0), n) };
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
|
@ -550,6 +964,80 @@ mod tests {
|
|||
}
|
||||
}
|
||||
|
||||
/// The blocked kernel's every output is the one-row loop's,
|
||||
/// bit for bit: block shapes 1/2/3/4/5+ rows, weight rows
|
||||
/// that do not fill a 4-row or 2-row block, `k` with and
|
||||
/// without a `k % 8` tail, `k < 8`, with and without bias.
|
||||
#[test]
|
||||
fn bf16_blocked_is_bit_identical_to_one_row() {
|
||||
for &m in &[1usize, 2, 3, 4, 5, 6, 7, 9] {
|
||||
for &n in &[1usize, 2, 3, 4, 5, 7, 9, 17, 70] {
|
||||
for &k in &[1usize, 7, 8, 9, 16, 203, 1030] {
|
||||
let a = synth_f32(m * k, m + 3 * k);
|
||||
let w = to_bf16_bytes(&synth_f32(n * k, 7 * n + k));
|
||||
for with_bias in [false, true] {
|
||||
let bias: Vec<f32> = 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<f32> = (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
|
||||
|
|
|
|||
|
|
@ -0,0 +1,481 @@
|
|||
//! A fork-join pool for decode-sized parallel loops: threads that
|
||||
//! stay awake between loops.
|
||||
//!
|
||||
//! A decode step of a mixture-of-experts model runs about 150
|
||||
//! small parallel loops one after another (each layer's q/k/v,
|
||||
//! attention + experts, o_proj + experts), each a fraction of a
|
||||
//! millisecond. A pool whose threads go to sleep between loops pays
|
||||
//! an OS wake-up per thread per loop, and on a busy host a woken
|
||||
//! thread can wait a long time for a core. Here the threads spin
|
||||
//! (`spin_loop`, no syscall) for a short while after a loop ends,
|
||||
//! so the next loop finds them running, and park only when no loop
|
||||
//! comes. The caller's thread takes units too.
|
||||
//!
|
||||
//! Units are claimed one at a time from a shared counter, so a slow
|
||||
//! or descheduled thread holds at most the unit it is running.
|
||||
//! Which thread runs which unit never changes a result: every unit
|
||||
//! writes only its own outputs.
|
||||
|
||||
use std::any::Any;
|
||||
use std::sync::atomic::{AtomicBool, AtomicPtr, AtomicU64, AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, Condvar, Mutex, OnceLock, PoisonError};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
/// How long an idle thread spins before it parks.
|
||||
const SPIN: Duration = Duration::from_micros(300);
|
||||
|
||||
/// Bits of the claim word that hold the next unit's index; the
|
||||
/// epoch has the other 40.
|
||||
const IDX_BITS: u32 = 24;
|
||||
|
||||
/// The unit index of a finished loop's claim word, and the most units
|
||||
/// one loop may have (`run` asserts it): no loop has more, so a thread
|
||||
/// that still holds a finished loop's epoch fails the bound and claims
|
||||
/// nothing. 2^24 - 1 is far above any loop the decoder makes (a
|
||||
/// 262144-row lm_head in 4-row chunks is 65536 units; prefill
|
||||
/// attention for a 500k-token prompt about 1M).
|
||||
const RETIRED: u64 = (1 << IDX_BITS) - 1;
|
||||
|
||||
/// The epoch's part of a claim word.
|
||||
fn epoch_of(word: u64) -> u64 {
|
||||
word >> IDX_BITS
|
||||
}
|
||||
|
||||
/// One loop: run `f(i)` for `i < len`.
|
||||
struct Job<'a> {
|
||||
f: &'a (dyn Fn(usize) + Sync),
|
||||
}
|
||||
|
||||
struct Shared {
|
||||
/// `(epoch << IDX_BITS) | next unit` of the current loop, or
|
||||
/// `(epoch << IDX_BITS) | RETIRED` once a new loop is being
|
||||
/// published.
|
||||
claim: AtomicU64,
|
||||
/// The current loop's unit count. It carries no epoch, so a
|
||||
/// thread that read an older loop's claim word may read a newer
|
||||
/// count; `run` retires the old word before it stores the count,
|
||||
/// which makes that thread's claim fail (see `run`).
|
||||
len: AtomicUsize,
|
||||
/// The current loop's job, valid from the claim of one of its
|
||||
/// units until that unit is counted done.
|
||||
job: AtomicPtr<Job<'static>>,
|
||||
/// Units finished in the current loop.
|
||||
done: AtomicUsize,
|
||||
/// A unit panicked: its payload is in `payloads`.
|
||||
panicked: AtomicBool,
|
||||
/// The payloads of the current loop's panicked units, in the order
|
||||
/// they were caught. The caller re-raises the first and drops the
|
||||
/// rest once the loop has finished and the pool is released, so no
|
||||
/// payload's `Drop` runs while another thread may be in the loop.
|
||||
payloads: Mutex<Vec<Box<dyn Any + Send>>>,
|
||||
/// Parked threads wait here for a new epoch.
|
||||
sleep: Mutex<u64>,
|
||||
wake: Condvar,
|
||||
}
|
||||
|
||||
/// The pool. Built once per process ([`spin_pool`]).
|
||||
pub struct SpinPool {
|
||||
shared: Arc<Shared>,
|
||||
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<Window>,
|
||||
}
|
||||
|
||||
/// See [`SpinPool::window`].
|
||||
#[cfg(test)]
|
||||
type Window = Box<dyn Fn(&Shared) + Send + Sync>;
|
||||
|
||||
impl Shared {
|
||||
/// Claim and run units of loop `epoch` until none are left.
|
||||
/// Returns how many this thread ran.
|
||||
fn work(&self, epoch: u64) -> usize {
|
||||
let mut ran = 0;
|
||||
loop {
|
||||
let cur = self.claim.load(Ordering::Acquire);
|
||||
if epoch_of(cur) != epoch {
|
||||
return ran;
|
||||
}
|
||||
let idx = (cur & RETIRED) as usize;
|
||||
// If this is a newer loop's count, `cur` is out of date:
|
||||
// the count was stored (Release) after the word was
|
||||
// retired, so the exchange below sees the retired word
|
||||
// and fails.
|
||||
if idx >= self.len.load(Ordering::Acquire) {
|
||||
return ran;
|
||||
}
|
||||
if self
|
||||
.claim
|
||||
.compare_exchange_weak(cur, cur + 1, Ordering::AcqRel, Ordering::Relaxed)
|
||||
.is_err()
|
||||
{
|
||||
continue;
|
||||
}
|
||||
// Unit `idx` of `epoch` is ours alone. The claim
|
||||
// succeeded while `epoch` was unfinished, so the loop that
|
||||
// owns the job cannot return (and drop it) before we count
|
||||
// this unit done, and no newer loop has replaced it.
|
||||
//
|
||||
// One theoretical exception: the exchange compares the
|
||||
// whole word, so a thread descheduled between loading `cur`
|
||||
// and the exchange for exactly 2^40 - 1 loops (the epoch's
|
||||
// period; about 7.7 years of loops at 4500 a second, the
|
||||
// decoder's rate at 30 tokens a second) would find the same
|
||||
// word again, in a different loop. With 32 epoch bits that
|
||||
// was about 11 days of uninterrupted decoding.
|
||||
let job = self.job.load(Ordering::Acquire);
|
||||
// SAFETY: as above; the job was published (Release) before
|
||||
// `epoch`'s claim counter, which we read with Acquire.
|
||||
let f = unsafe { (*job).f };
|
||||
if let Err(payload) = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| f(idx))) {
|
||||
// Kept, not dropped: dropping it would run its `Drop`
|
||||
// here, inside the loop.
|
||||
self.payloads.lock().unwrap_or_else(PoisonError::into_inner).push(payload);
|
||||
self.panicked.store(true, Ordering::Relaxed);
|
||||
}
|
||||
ran += 1;
|
||||
self.done.fetch_add(1, Ordering::Release);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The epoch after `e`: 40 bits, wrapping, never 0 (a new thread
|
||||
/// has seen epoch 0).
|
||||
fn next_epoch(e: u64) -> u64 {
|
||||
match (e + 1) & ((1 << (64 - IDX_BITS)) - 1) {
|
||||
0 => 1,
|
||||
n => n,
|
||||
}
|
||||
}
|
||||
|
||||
fn worker(shared: Arc<Shared>) {
|
||||
let mut seen = 0u64;
|
||||
loop {
|
||||
let start = Instant::now();
|
||||
let mut epoch;
|
||||
loop {
|
||||
epoch = epoch_of(shared.claim.load(Ordering::Acquire));
|
||||
if epoch != seen || start.elapsed() > SPIN {
|
||||
break;
|
||||
}
|
||||
for _ in 0..64 {
|
||||
std::hint::spin_loop();
|
||||
}
|
||||
}
|
||||
if epoch == seen {
|
||||
let mut s = shared.sleep.lock().unwrap();
|
||||
while *s == seen {
|
||||
s = shared.wake.wait(s).unwrap();
|
||||
}
|
||||
epoch = *s;
|
||||
}
|
||||
seen = epoch;
|
||||
shared.work(epoch);
|
||||
}
|
||||
}
|
||||
|
||||
impl SpinPool {
|
||||
/// A pool of `threads` threads, the caller's included (so
|
||||
/// `threads - 1` are spawned). They live for the process.
|
||||
pub fn new(threads: usize) -> Self {
|
||||
let pool = Self::unstarted(threads);
|
||||
for i in 1..pool.threads {
|
||||
let s = pool.shared.clone();
|
||||
std::thread::Builder::new()
|
||||
.name(format!("jouleclaw-spin-{i}"))
|
||||
.spawn(move || worker(s))
|
||||
.expect("spawn a pool thread");
|
||||
}
|
||||
pool
|
||||
}
|
||||
|
||||
/// The pool's state for `threads` threads, none of them started.
|
||||
fn unstarted(threads: usize) -> Self {
|
||||
let shared = Arc::new(Shared {
|
||||
claim: AtomicU64::new(0),
|
||||
len: AtomicUsize::new(0),
|
||||
job: AtomicPtr::new(std::ptr::null_mut()),
|
||||
done: AtomicUsize::new(0),
|
||||
panicked: AtomicBool::new(false),
|
||||
payloads: Mutex::new(Vec::new()),
|
||||
sleep: Mutex::new(0),
|
||||
wake: Condvar::new(),
|
||||
});
|
||||
SpinPool {
|
||||
shared,
|
||||
threads: threads.max(1),
|
||||
busy: Mutex::new(()),
|
||||
#[cfg(test)]
|
||||
window: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Threads that run units, the caller's included.
|
||||
pub fn threads(&self) -> usize {
|
||||
self.threads
|
||||
}
|
||||
|
||||
/// Run `f(i)` for every `i < len` on the pool and the calling
|
||||
/// thread; returns when all have run. A unit must not call `run`
|
||||
/// on the same pool: the loop holds the pool until every unit has
|
||||
/// run, so the inner call would wait forever.
|
||||
pub fn run(&self, len: usize, f: &(dyn Fn(usize) + Sync)) {
|
||||
if len == 0 {
|
||||
return;
|
||||
}
|
||||
if self.threads == 1 || len == 1 {
|
||||
for i in 0..len {
|
||||
f(i);
|
||||
}
|
||||
return;
|
||||
}
|
||||
assert!(len as u64 <= RETIRED, "too many units in one loop");
|
||||
// `busy` cannot be poisoned: nothing panics while it is held.
|
||||
// From here to `drop(one)` the code is atomics; the `sleep`
|
||||
// lock, never held across a panic (its holders only set or
|
||||
// wait on the epoch); `Shared::work`, which catches a unit's
|
||||
// panic and keeps its payload instead of dropping it; moving
|
||||
// the payloads out of their list, whose lock is held only to
|
||||
// push or take; and, in tests only, the `window` hook. A unit's panic is
|
||||
// re-raised after the guard is released. Were the lock ever
|
||||
// poisoned anyway, it would be by a panic in the test hook,
|
||||
// before this loop's epoch is published, when no other thread
|
||||
// can reach this loop's job: the pool's state is still that
|
||||
// of the last finished loop, so the lock is taken as it is.
|
||||
let one = self.busy.lock().unwrap_or_else(PoisonError::into_inner);
|
||||
let sh = &*self.shared;
|
||||
let job = Job { f };
|
||||
// SAFETY: the job lives until this function returns, which is
|
||||
// after every unit has been counted done; no thread touches it
|
||||
// after its unit (see `Shared::work`). The lifetime is erased
|
||||
// only to store the pointer.
|
||||
let ptr = &job as *const Job<'_> as *mut Job<'static>;
|
||||
// Publication, in this order. A thread can still be inside
|
||||
// `work` for the previous loop, holding its epoch: woken late
|
||||
// from park, or descheduled after its last unit. That loop's
|
||||
// claim word reads (epoch, its unit count), so were this
|
||||
// loop's count visible first and larger, the thread would pass
|
||||
// the bound, claim a unit the old loop does not have, and run
|
||||
// a unit of this loop that is then claimed again, counted
|
||||
// twice, and lets this call return while it still runs.
|
||||
// 1. Retire the old word: its index passes no bound.
|
||||
// 2. The count, Release: a thread that reads it then sees the
|
||||
// retired word in its exchange, which fails.
|
||||
// 3. The job, then the new epoch, Release: a thread that reads
|
||||
// the epoch sees this loop's count and job.
|
||||
let old = epoch_of(sh.claim.load(Ordering::Relaxed));
|
||||
sh.claim.store((old << IDX_BITS) | RETIRED, Ordering::Relaxed);
|
||||
sh.done.store(0, Ordering::Relaxed);
|
||||
sh.len.store(len, Ordering::Release);
|
||||
sh.job.store(ptr, Ordering::Release);
|
||||
#[cfg(test)]
|
||||
if let Some(w) = &self.window {
|
||||
w(sh);
|
||||
}
|
||||
let epoch = next_epoch(old);
|
||||
sh.claim.store(epoch << IDX_BITS, Ordering::Release);
|
||||
{
|
||||
let mut s = sh.sleep.lock().unwrap();
|
||||
*s = epoch;
|
||||
sh.wake.notify_all();
|
||||
}
|
||||
sh.work(epoch);
|
||||
while sh.done.load(Ordering::Acquire) < len {
|
||||
std::hint::spin_loop();
|
||||
}
|
||||
let payloads = if sh.panicked.swap(false, Ordering::Relaxed) {
|
||||
std::mem::take(&mut *sh.payloads.lock().unwrap_or_else(PoisonError::into_inner))
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
drop(one);
|
||||
// The loop has finished and the pool is free: payloads may run
|
||||
// their `Drop` now. The first unit to panic is re-raised as it
|
||||
// was, payload and all; the others are dropped.
|
||||
let mut payloads = payloads.into_iter();
|
||||
if let Some(first) = payloads.next() {
|
||||
drop(payloads);
|
||||
std::panic::resume_unwind(first);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The process's decode pool, sized like
|
||||
/// [`crate::matmul_bf16::cpu_pool`].
|
||||
pub fn spin_pool() -> &'static SpinPool {
|
||||
static POOL: OnceLock<SpinPool> = 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<AtomicUsize> = (0..len).map(|_| AtomicUsize::new(0)).collect();
|
||||
pool.run(len, &|i| {
|
||||
hits[i].fetch_add(1, Ordering::Relaxed);
|
||||
});
|
||||
assert!(hits.iter().all(|h| h.load(Ordering::Relaxed) == 1), "len {len}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_epoch_wraps_without_reaching_zero() {
|
||||
assert_eq!(next_epoch(1), 2);
|
||||
assert_eq!(next_epoch((1 << 40) - 2), (1 << 40) - 1);
|
||||
assert_eq!(next_epoch((1 << 40) - 1), 1);
|
||||
// A pool whose counter is just below the wrap keeps running.
|
||||
let pool = SpinPool::new(4);
|
||||
pool.shared.claim.store(((1 << 40) - 3) << IDX_BITS, Ordering::Relaxed);
|
||||
let n = AtomicUsize::new(0);
|
||||
for _ in 0..6 {
|
||||
pool.run(50, &|_| {
|
||||
n.fetch_add(1, Ordering::Relaxed);
|
||||
});
|
||||
}
|
||||
assert_eq!(n.load(Ordering::Relaxed), 300);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn loops_after_the_threads_parked_still_run() {
|
||||
let pool = SpinPool::new(4);
|
||||
let n = AtomicUsize::new(0);
|
||||
pool.run(100, &|_| {
|
||||
n.fetch_add(1, Ordering::Relaxed);
|
||||
});
|
||||
std::thread::sleep(SPIN * 4);
|
||||
pool.run(100, &|_| {
|
||||
n.fetch_add(1, Ordering::Relaxed);
|
||||
});
|
||||
assert_eq!(n.load(Ordering::Relaxed), 200);
|
||||
}
|
||||
|
||||
/// A thread still inside `work` for a loop that has finished (a
|
||||
/// worker woken late from park holds that loop's epoch, and so
|
||||
/// does one descheduled after its last unit) claims nothing of
|
||||
/// the next loop, also when the next loop has more units.
|
||||
/// Deterministic: the pool has no threads of its own, and the
|
||||
/// late thread runs in the window in which the new loop's unit
|
||||
/// count and job are published and its epoch is not.
|
||||
#[test]
|
||||
fn a_late_thread_of_a_finished_loop_claims_nothing_of_the_next() {
|
||||
let late = Arc::new(AtomicUsize::new(0));
|
||||
let armed = Arc::new(AtomicBool::new(false));
|
||||
let mut pool = SpinPool::unstarted(2);
|
||||
let (l, a) = (late.clone(), armed.clone());
|
||||
pool.window = Some(Box::new(move |sh: &Shared| {
|
||||
if a.load(Ordering::Relaxed) {
|
||||
let finished = epoch_of(sh.claim.load(Ordering::Relaxed));
|
||||
let ran = std::thread::scope(|s| s.spawn(|| sh.work(finished)).join().unwrap());
|
||||
l.fetch_add(ran, Ordering::Relaxed);
|
||||
}
|
||||
}));
|
||||
let hits: Vec<AtomicUsize> = (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<usize> = 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<usize> = HITS[..len].iter().map(|h| h.swap(0, Ordering::Relaxed)).collect();
|
||||
assert!(counts.iter().all(|&c| c == 1), "round {round}, loop of {len}: {counts:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The caller gets the unit's own panic, its payload, not a
|
||||
/// message of the pool's.
|
||||
#[test]
|
||||
fn a_units_panic_reaches_the_caller_with_its_payload() {
|
||||
#[derive(Debug, PartialEq)]
|
||||
struct UnitFailed(usize);
|
||||
let pool = SpinPool::new(4);
|
||||
let failed = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
|
||||
pool.run(16, &|i| {
|
||||
if i == 5 {
|
||||
std::panic::panic_any(UnitFailed(5));
|
||||
}
|
||||
});
|
||||
}));
|
||||
let payload = failed.expect_err("the unit's panic reaches the caller");
|
||||
assert_eq!(payload.downcast_ref::<UnitFailed>(), Some(&UnitFailed(5)));
|
||||
}
|
||||
|
||||
/// The epoch outlasts 32 bits, so a thread would have to sleep
|
||||
/// between reading the claim word and claiming through 2^40 - 1
|
||||
/// loops for a stale word to match again; the unit index has 24
|
||||
/// bits, and a loop with more units than that is refused before it
|
||||
/// takes the pool.
|
||||
#[test]
|
||||
fn the_epoch_has_40_bits_and_a_loop_at_most_2_pow_24_minus_1_units() {
|
||||
assert_eq!(next_epoch(0xFFFF_FFFF), 0x1_0000_0000);
|
||||
assert_eq!(next_epoch((1 << 40) - 2), (1 << 40) - 1);
|
||||
assert_eq!(next_epoch((1 << 40) - 1), 1);
|
||||
let pool = SpinPool::new(4);
|
||||
let refused = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| pool.run(1 << 24, &|_| {})));
|
||||
assert!(refused.is_err(), "a loop of 2^24 units is refused");
|
||||
let n = AtomicUsize::new(0);
|
||||
pool.run(50, &|_| {
|
||||
n.fetch_add(1, Ordering::Relaxed);
|
||||
});
|
||||
assert_eq!(n.load(Ordering::Relaxed), 50, "and the pool runs the next loop");
|
||||
}
|
||||
|
||||
/// A unit that panics fails its own loop and nothing after it:
|
||||
/// the panic reaches the caller, and the pool runs every later
|
||||
/// loop.
|
||||
#[test]
|
||||
fn a_panicking_unit_leaves_the_pool_usable() {
|
||||
let pool = SpinPool::new(4);
|
||||
let failed = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
|
||||
pool.run(16, &|i| assert_ne!(i, 5, "unit 5 fails"));
|
||||
}));
|
||||
assert!(failed.is_err(), "the unit's panic reaches the caller");
|
||||
let n = AtomicUsize::new(0);
|
||||
for _ in 0..3 {
|
||||
pool.run(16, &|_| {
|
||||
n.fetch_add(1, Ordering::Relaxed);
|
||||
});
|
||||
}
|
||||
assert_eq!(n.load(Ordering::Relaxed), 48);
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue