fb103ac5ae
Every registered string now lives in one table (codecs/registry.rs) with its kind, status, media class, registering SPEC section, edition and registration date. `wai registry export` writes registry/capabilities.json, its JSON Schema, the generated table in REGISTRY.md and vocab v2; --check fails on drift. vocab v1 stays stable. A candidate capability (a registered string whose payload format no revision has pinned) must not be emitted. The C packers, the browser writers and `wai wrap` refuse one, and `codecs::is_dispatchable` is the single dispatch and selection predicate (SPEC §4): candidates, companion-only capabilities, records and derivations are not dispatchable, and an unregistered string is decoded only where the sink advertises it. The codec glossary is linked to the registry through a reviewed table and held to its schema by a library test; `org` becomes `body`, from a closed allowlist. The SPEC gains §5 sections on kinds and candidates, a §8 edition rule, corrected citations and rewritten appendices. The quantum-ops vectors are regenerated with neutral device ids. CI declares the feature-gated test targets with required-features, checks the registry export, builds wai-web for wasm32, and runs two public-text checks (comparative prose, and private name patterns) with a self-test. The specification, extensions, research notes, READMEs, glossary and source comments are rewritten to name no company or product and to make no comparisons. Co-Authored-By: Claude Opus 5.5 <[email protected]>
1153 lines
48 KiB
Rust
1153 lines
48 KiB
Rust
//! Integer-exact transform primitives — the missing piece for a
|
||
//! `NeuralIntegerExact` *attention-based* decode (the attention and MetaFormer
|
||
//! transform class, whose synthesis is a 600-op attention + normalization net rather than
|
||
//! a plain conv stack). Pure i64, no float, no `ort`: every op is byte-identical
|
||
//! on every machine, the same discipline as `int_synth` (conv) and `int_entropy`
|
||
//! (entropy). Each primitive is a fixed-point implementation with a **stated
|
||
//! error bound** versus the float reference — the honest contract for replacing
|
||
//! a float op with one that can go in a receipt.
|
||
//!
|
||
//! Scope: the hard, reusable cores — softmax, sigmoid, gelu, layernorm, matmul,
|
||
//! linear, single- and multi-head attention — and a pre-norm [`TransformerBlock`]
|
||
//! that composes them into the repeating MetaFormer unit. Assembling a *specific*
|
||
//! net (e.g. JPEG AI's 609-op synthesis) by stacking blocks with trained weights
|
||
//! on top of these is a separate, larger build.
|
||
|
||
/// Probability output scale of [`softmax`]: results are at `2^PROB_FRAC`
|
||
/// (they sum to ~`2^PROB_FRAC`).
|
||
pub const PROB_FRAC: u32 = 16;
|
||
const EXP_FRAC: u32 = 30; // internal exp() fixed-point scale
|
||
|
||
// Integer-rational constants (value × 10^15) so the fixed-point constants are
|
||
// computed by integer arithmetic — identical on every platform, no float drift.
|
||
const LN2_E15: i64 = 693_147_180_559_945; // ln 2
|
||
// 2nd-order fit of exp(-t) on t∈[0, ln2): c0 + c1·t + c2·t² (max |err| ≈ 2e-3).
|
||
const C0_E15: i64 = 1_000_000_000_000_000;
|
||
const C1_E15: i64 = -968_890_000_000_000;
|
||
const C2_E15: i64 = 357_000_000_000_000;
|
||
|
||
#[inline]
|
||
fn from_e15(v_e15: i64, shift: u32) -> i64 {
|
||
((v_e15 as i128 * (1i128 << shift)) / 1_000_000_000_000_000i128) as i64
|
||
}
|
||
|
||
/// Integer-exact softmax over `x` (fixed-point, `frac` fractional bits).
|
||
/// Returns probabilities at scale `2^PROB_FRAC`. Deterministic by construction:
|
||
/// subtract max, `exp(d) = 2^-q · poly(exp(-r))` with `d = -(q·ln2 + r)`, then
|
||
/// integer-normalize. No float anywhere; `frac ≤ 16` keeps the i128 intermediates
|
||
/// in range.
|
||
pub fn softmax(x: &[i64], frac: u32) -> Vec<i64> {
|
||
let n = x.len();
|
||
if n == 0 {
|
||
return Vec::new();
|
||
}
|
||
let xmax = *x.iter().max().unwrap();
|
||
let ln2 = from_e15(LN2_E15, frac).max(1);
|
||
let (c0, c1, c2) = (
|
||
from_e15(C0_E15, EXP_FRAC),
|
||
from_e15(C1_E15, EXP_FRAC),
|
||
from_e15(C2_E15, EXP_FRAC),
|
||
);
|
||
|
||
let mut exps = vec![0i64; n];
|
||
for (i, &xi) in x.iter().enumerate() {
|
||
let neg = xmax - xi; // ≥ 0, scale 2^frac
|
||
let q = (neg / ln2) as u32; // power of two: exp(d) = 2^-q · exp(-r)
|
||
let r = neg - q as i64 * ln2; // ∈ [0, ln2), scale 2^frac
|
||
// exp(-r) at 2^EXP_FRAC = c0 + c1·t + c2·t², t = r/2^frac
|
||
let t1 = ((c1 as i128 * r as i128) >> frac) as i64;
|
||
let t2 = ((c2 as i128 * r as i128 * r as i128) >> (2 * frac)) as i64;
|
||
let e = (c0 + t1 + t2).max(0); // exp(-r), scale 2^EXP_FRAC
|
||
exps[i] = if q >= EXP_FRAC + 2 { 0 } else { e >> q };
|
||
}
|
||
|
||
let sum: i64 = exps.iter().sum();
|
||
if sum == 0 {
|
||
// UNREACHABLE, and load-bearingly so — do not make it reachable without
|
||
// fixing it. The model checker marks this branch dead (see
|
||
// `kani_softmax_sums_to_exactly_one`): the argmax element always has
|
||
// `neg = xmax - xmax = 0`, hence `q = 0`, `r = 0`, `e = c0 > 0`, so
|
||
// `exps[argmax] > 0` and `sum > 0` for every input.
|
||
//
|
||
// That matters because this line would BREAK the exact-normalisation
|
||
// contract asserted below. `n * ((1<<PROB_FRAC)/n)` is `1<<PROB_FRAC`
|
||
// only when `n` divides it: for n = 3 it is 65_535, not 65_536. A
|
||
// denormalised distribution here becomes a denormalised CDF, and the
|
||
// entropy decode then selects the wrong interval and desyncs the rest
|
||
// of the stream — the exact failure SPEC §7 exists to prevent.
|
||
//
|
||
// So the invariant currently holds by unreachability rather than by
|
||
// construction. Any change to the exp reduction that lets this fire
|
||
// must distribute the remainder here the way the tail of this function
|
||
// does.
|
||
return vec![(1i64 << PROB_FRAC) / n as i64; n];
|
||
}
|
||
let mut probs: Vec<i64> = exps
|
||
.iter()
|
||
.map(|&e| (((e as i128) << PROB_FRAC) / sum as i128) as i64)
|
||
.collect();
|
||
// Integer division truncates each prob down by < 1 LSB; give the whole
|
||
// remainder to the largest element so the distribution sums to EXACTLY
|
||
// 2^PROB_FRAC, deterministically (argmax of exps = argmax of probs).
|
||
let total: i64 = probs.iter().sum();
|
||
let argmax = exps.iter().enumerate().max_by_key(|&(_, &e)| e).map(|(i, _)| i).unwrap();
|
||
probs[argmax] += (1i64 << PROB_FRAC) - total;
|
||
probs
|
||
}
|
||
|
||
// ─────────────────────────────────────────────────────────────────────
|
||
// Kani proof harnesses — the normalisation invariant, checked
|
||
// ─────────────────────────────────────────────────────────────────────
|
||
//
|
||
// `softmax` states its own contract above: the remainder is handed to the
|
||
// largest element "so the distribution sums to EXACTLY 2^PROB_FRAC". That is
|
||
// load-bearing rather than cosmetic — these probabilities become a CDF, and a
|
||
// CDF that does not normalise makes the entropy decode select the wrong
|
||
// interval and desync the rest of the stream, exactly the failure SPEC §7 is
|
||
// written to prevent. A comment is not a guarantee; this is.
|
||
//
|
||
// The logits are bounded to a range a real hyper-decoder produces. That bound
|
||
// is a genuine precondition, not a convenience: `xmax - xi` is an i64
|
||
// subtraction, so unbounded logits can overflow it — which panics in debug and
|
||
// wraps in release, i.e. two different softmaxes depending on the build. The
|
||
// bound is stated here so the precondition is explicit rather than assumed.
|
||
|
||
/// The distribution sums to exactly one, and no probability is negative.
|
||
///
|
||
/// **Arity and bound are what the solver can close, and are stated rather than
|
||
/// tuned quietly.** Each element costs an i128 polynomial evaluation and an
|
||
/// i128 division; at three symbolic logits over ±2^40 the instance did not
|
||
/// close in twenty minutes of propositional reduction. This is the two-logit
|
||
/// case over ±2^20 — a real instance of the invariant, not a universal claim.
|
||
/// [`kani_softmax_singleton_is_certain`] covers the one-symbol case.
|
||
#[cfg(kani)]
|
||
#[kani::proof]
|
||
#[kani::unwind(5)]
|
||
#[kani::solver(kissat)]
|
||
fn kani_softmax_sums_to_exactly_one() {
|
||
let a: i64 = kani::any();
|
||
let b: i64 = kani::any();
|
||
// Logits at scale 2^frac from an integer hyper-decoder; bounded well inside
|
||
// the range where the reduction arithmetic is exact.
|
||
const LIM: i64 = 1i64 << 20;
|
||
kani::assume(a > -LIM && a < LIM);
|
||
kani::assume(b > -LIM && b < LIM);
|
||
|
||
let p = softmax(&[a, b], PROB_FRAC);
|
||
let total: i64 = p.iter().sum();
|
||
kani::assert(total == 1i64 << PROB_FRAC, "the distribution sums to exactly one");
|
||
for &v in &p {
|
||
kani::assert(v >= 0, "no probability is negative");
|
||
}
|
||
kani::cover!(true, "softmax is reachable under these bounds");
|
||
}
|
||
|
||
/// The degenerate single-element case: all the mass on one symbol.
|
||
#[cfg(kani)]
|
||
#[kani::proof]
|
||
#[kani::unwind(4)]
|
||
#[kani::solver(kissat)]
|
||
fn kani_softmax_singleton_is_certain() {
|
||
let a: i64 = kani::any();
|
||
kani::assume(a > -(1i64 << 40) && a < (1i64 << 40));
|
||
let p = softmax(&[a], PROB_FRAC);
|
||
kani::assert(p.len() == 1, "one logit yields one probability");
|
||
kani::assert(p[0] == 1i64 << PROB_FRAC, "a lone symbol carries all the mass");
|
||
}
|
||
|
||
/// `exp(-a)` for `a ≥ 0` (fixed-point, `frac` bits), result at `2^frac`
|
||
/// (so in `[0, 2^frac]`). Same reduce-by-`ln2` + 2nd-order polynomial as
|
||
/// [`softmax`]'s inner exp, factored out so the activations can share it.
|
||
fn exp_neg(a: i64, frac: u32) -> i64 {
|
||
let one = 1i64 << frac;
|
||
if a <= 0 {
|
||
return one;
|
||
}
|
||
let ln2 = from_e15(LN2_E15, frac).max(1);
|
||
let (c0, c1, c2) = (
|
||
from_e15(C0_E15, EXP_FRAC),
|
||
from_e15(C1_E15, EXP_FRAC),
|
||
from_e15(C2_E15, EXP_FRAC),
|
||
);
|
||
let q = (a / ln2) as u32; // exp(-a) = 2^-q · exp(-r)
|
||
let r = a - q as i64 * ln2; // ∈ [0, ln2), scale 2^frac
|
||
let t1 = ((c1 as i128 * r as i128) >> frac) as i64;
|
||
let t2 = ((c2 as i128 * r as i128 * r as i128) >> (2 * frac)) as i64;
|
||
let e = (c0 + t1 + t2).max(0); // exp(-r), scale 2^EXP_FRAC
|
||
if q >= EXP_FRAC + 2 {
|
||
return 0;
|
||
}
|
||
(e >> (EXP_FRAC - frac)) >> q // → exp(-a) at 2^frac (frac ≤ EXP_FRAC)
|
||
}
|
||
|
||
/// Integer-exact logistic sigmoid `1/(1+e^-z)` (fixed-point, `frac ≤ 16` bits),
|
||
/// result at `2^frac` ∈ `[0, 2^frac]`. Built on [`exp_neg`], reflecting through
|
||
/// `z<0` so the polynomial is only ever evaluated on a non-negative argument
|
||
/// (where it is accurate and never overflows). Byte-identical on every machine.
|
||
pub fn sigmoid(z: i64, frac: u32) -> i64 {
|
||
let one = 1i64 << frac;
|
||
if z >= 0 {
|
||
let e = exp_neg(z, frac); // e^-z ∈ [0,1] at 2^frac
|
||
((one as i128 * one as i128) / (one + e) as i128) as i64
|
||
} else {
|
||
let e = exp_neg(-z, frac); // e^z ∈ [0,1] at 2^frac
|
||
((e as i128 * one as i128) / (one + e) as i128) as i64
|
||
}
|
||
}
|
||
|
||
// GELU's sigmoid-approximation coefficient 1.702 (value × 10^15).
|
||
const GELU_A_E15: i64 = 1_702_000_000_000_000;
|
||
|
||
/// Integer-exact GELU via the sigmoid approximation `x·σ(1.702·x)`
|
||
/// (fixed-point, `frac ≤ 16` bits), output at `2^frac`. The activation a
|
||
/// MetaFormer/JPEG-AI FFN needs; composes [`sigmoid`], so it is byte-identical
|
||
/// on every machine. The error bound is stated versus the **float of the same
|
||
/// formula** — i.e. it measures the fixed-point error, not the σ-approximation
|
||
/// (that choice of GELU is the op being made integer-exact).
|
||
pub fn gelu(x: &[i64], frac: u32) -> Vec<i64> {
|
||
let a = from_e15(GELU_A_E15, frac); // 1.702 at 2^frac
|
||
x.iter()
|
||
.map(|&xi| {
|
||
let z = ((a as i128 * xi as i128) >> frac) as i64; // 1.702·x at 2^frac
|
||
let s = sigmoid(z, frac); // σ(1.702·x) at 2^frac
|
||
((xi as i128 * s as i128) >> frac) as i64
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
/// Integer-exact layer normalization over `x` (fixed-point, `frac` bits):
|
||
/// `y = (x − mean) / sqrt(var + eps) · gamma + beta`. `gamma`/`beta` are at the
|
||
/// same `2^frac` scale; `eps` is at `2^(2·frac)` (added to the variance). The
|
||
/// reciprocal-std is an exact integer `isqrt` — no float, no Newton drift — so
|
||
/// the result is byte-identical on every machine. Output at `2^frac`.
|
||
pub fn layernorm(x: &[i64], frac: u32, gamma: &[i64], beta: &[i64], eps: i64) -> Vec<i64> {
|
||
let n = x.len();
|
||
if n == 0 {
|
||
return Vec::new();
|
||
}
|
||
let ni = n as i64;
|
||
// mean (round-to-nearest), scale 2^frac
|
||
let sum: i128 = x.iter().map(|&v| v as i128).sum();
|
||
let mean = ((sum + (ni as i128) / 2) / ni as i128) as i64;
|
||
// variance = Σ(x−mean)² / n, scale 2^(2·frac)
|
||
let var: i128 = x.iter().map(|&v| { let d = (v - mean) as i128; d * d }).sum::<i128>() / ni as i128;
|
||
// std = isqrt(var + eps), scale 2^frac (since var is at 2^(2·frac))
|
||
let std = ((var + eps as i128).max(1) as u128).isqrt() as i64;
|
||
let std = std.max(1);
|
||
(0..n)
|
||
.map(|i| {
|
||
let d = x[i] - mean; // 2^frac
|
||
let norm = (((d as i128) << frac) / std as i128) as i64; // (d/std) at 2^frac
|
||
((norm as i128 * gamma[i] as i128) >> frac) as i64 + beta[i]
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
/// Integer matmul `C = A·B`, requantized by an arithmetic `>> shift`. `A` is
|
||
/// `m×k`, `B` is `k×n`, both i64 fixed-point, row-major; the accumulation is
|
||
/// i64. Deterministic — the integer core of every linear / attention projection.
|
||
pub fn matmul(a: &[i64], b: &[i64], m: usize, k: usize, n: usize, shift: u32) -> Vec<i64> {
|
||
let mut c = vec![0i64; m * n];
|
||
for i in 0..m {
|
||
for p in 0..k {
|
||
let aip = a[i * k + p] as i128;
|
||
if aip == 0 {
|
||
continue;
|
||
}
|
||
let brow = p * n;
|
||
let crow = i * n;
|
||
for j in 0..n {
|
||
c[crow + j] += ((aip * b[brow + j] as i128) >> shift) as i64;
|
||
}
|
||
}
|
||
}
|
||
c
|
||
}
|
||
|
||
/// [`matmul`]'s accumulation is total for fixed-point operands of realistic
|
||
/// magnitude — and the bound is the point of the harness.
|
||
///
|
||
/// The product is formed in i128 and then **truncated to i64**, after which the
|
||
/// accumulator is a plain `i64 +=`. That `+=` is the hazard: for unbounded
|
||
/// operands it overflows, panicking in debug and wrapping in release, i.e. two
|
||
/// different matmuls depending on the build. Nothing in the signature enforces a
|
||
/// magnitude, so the safe operating range is a precondition callers inherit
|
||
/// silently. Writing it down here makes it checkable instead.
|
||
///
|
||
/// `±2^31` at Q16 is a real activation range: each product is at most
|
||
/// `2^62 >> 16 = 2^46`, and `k` of them accumulate well inside i64.
|
||
#[cfg(kani)]
|
||
#[kani::proof]
|
||
#[kani::unwind(4)]
|
||
#[kani::solver(kissat)]
|
||
fn kani_matmul_accumulation_is_total() {
|
||
const LIM: i64 = 1i64 << 31;
|
||
let bounded = || -> i64 {
|
||
let v: i64 = kani::any();
|
||
kani::assume(v > -LIM && v < LIM);
|
||
v
|
||
};
|
||
// Written out rather than looped: an assume-loop of its own would need a
|
||
// larger unwind bound than the code under test, and a harness that cannot
|
||
// unwind reports nothing about the function.
|
||
let a: [i64; 4] = [bounded(), bounded(), bounded(), bounded()]; // 2x2
|
||
let b: [i64; 4] = [bounded(), bounded(), bounded(), bounded()]; // 2x2
|
||
let c = matmul(&a, &b, 2, 2, 2, PROB_FRAC);
|
||
kani::assert(c.len() == 4, "the result has m*n entries");
|
||
kani::cover!(true, "the accumulation is reachable");
|
||
}
|
||
|
||
|
||
/// [`layernorm`] is total: the divide is guarded, the arithmetic stays in range.
|
||
///
|
||
/// The division by `std` is already defended twice (`.max(1)` on the radicand
|
||
/// and again on the root), so the question is whether the surrounding
|
||
/// fixed-point arithmetic holds and whether `isqrt` over u128 is something a
|
||
/// bounded model checker can actually close. Both are worth knowing.
|
||
///
|
||
/// `gamma`/`beta` are indexed by the length of `x` with no length check in the
|
||
/// signature, so they are supplied at matching length here — that mismatch is a
|
||
/// caller-contract violation, not a property of this function.
|
||
///
|
||
/// **Read the bound before trusting the coverage.** This closes at ±2^10, which
|
||
/// is *below* the magnitudes a real Q16 activation carries; at ±2^20 the
|
||
/// instance did not close in 25 minutes. So this proves the guarded divide and
|
||
/// the surrounding arithmetic hold for small operands — it does not establish
|
||
/// `layernorm` over its working range, and it should not be read as doing so.
|
||
/// The obstacle is `isqrt` over a symbolic u128, a limit of bounded model
|
||
/// checking rather than a property of the code.
|
||
#[cfg(kani)]
|
||
#[kani::proof]
|
||
#[kani::unwind(4)]
|
||
#[kani::solver(kissat)]
|
||
fn kani_layernorm_is_total() {
|
||
// Deliberately small. `std = isqrt(var + eps)` over a symbolic u128 is the
|
||
// cost centre: at ±2^20 the instance did not close in 25 minutes. A tighter
|
||
// magnitude shrinks the radicand and with it the solver's search.
|
||
const LIM: i64 = 1i64 << 10;
|
||
let bounded = || -> i64 {
|
||
let v: i64 = kani::any();
|
||
kani::assume(v > -LIM && v < LIM);
|
||
v
|
||
};
|
||
let x = [bounded(), bounded()];
|
||
let gamma = [bounded(), bounded()];
|
||
let beta = [bounded(), bounded()];
|
||
let out = layernorm(&x, PROB_FRAC, &gamma, &beta, 1);
|
||
kani::assert(out.len() == 2, "one output per input");
|
||
kani::cover!(true, "layernorm is reachable");
|
||
}
|
||
|
||
/// Single-head scaled dot-product attention, fully integer:
|
||
/// `out = softmax(Q·Kᵀ / √d) · V`. `q`/`k`/`v` are `seq×d` row-major at `2^frac`;
|
||
/// output is `seq×d` at `2^frac`. Composes [`matmul`] + [`softmax`]; every step
|
||
/// is integer, so the whole block is byte-identical on every machine — a
|
||
/// `NeuralIntegerExact` attention block.
|
||
pub fn attention(q: &[i64], k: &[i64], v: &[i64], seq: usize, d: usize, frac: u32) -> Vec<i64> {
|
||
const SC: u32 = 16;
|
||
// 1/√d at 2^SC, by exact integer isqrt.
|
||
let sqrt_d = ((d as u128) << (2 * SC)).isqrt() as i64; // √d at 2^SC
|
||
let inv_sqrt_d = (1i128 << (2 * SC)) / sqrt_d.max(1) as i128; // 1/√d at 2^SC
|
||
|
||
let mut out = vec![0i64; seq * d];
|
||
for i in 0..seq {
|
||
// scores_j = (Qᵢ·Kⱼ / √d) at 2^frac, then softmax over j.
|
||
let mut scores = vec![0i64; seq];
|
||
for j in 0..seq {
|
||
let mut raw = 0i128; // Σ q·k at 2^(2·frac)
|
||
for t in 0..d {
|
||
raw += q[i * d + t] as i128 * k[j * d + t] as i128;
|
||
}
|
||
// → at 2^frac: raw·(1/√d) / 2^(frac+SC)
|
||
scores[j] = ((raw * inv_sqrt_d) >> (frac + SC)) as i64;
|
||
}
|
||
let attn = softmax(&scores, frac); // at 2^PROB_FRAC
|
||
for t in 0..d {
|
||
let mut acc = 0i128; // Σ attn·v at 2^(PROB_FRAC+frac)
|
||
for j in 0..seq {
|
||
acc += attn[j] as i128 * v[j * d + t] as i128;
|
||
}
|
||
out[i * d + t] = (acc >> PROB_FRAC) as i64; // → 2^frac
|
||
}
|
||
}
|
||
out
|
||
}
|
||
|
||
/// Integer-exact multi-head attention: split the `d`-wide model dimension into
|
||
/// `heads` contiguous heads of width `d/heads`, run [`attention`] independently
|
||
/// per head (each scaled by `√(d/heads)`, as in a real transformer), and concat
|
||
/// the heads back into a `seq×d` output at `2^frac`. Pure composition of the
|
||
/// tested single-head block, so the whole thing stays byte-identical on every
|
||
/// machine. `heads` must divide `d`.
|
||
pub fn multi_head_attention(
|
||
q: &[i64],
|
||
k: &[i64],
|
||
v: &[i64],
|
||
seq: usize,
|
||
d: usize,
|
||
heads: usize,
|
||
frac: u32,
|
||
) -> Vec<i64> {
|
||
assert!(heads > 0 && d % heads == 0, "heads must divide d");
|
||
let dh = d / heads;
|
||
let mut out = vec![0i64; seq * d];
|
||
// Gather head `h`'s columns [h·dh, h·dh+dh) into a dense seq×dh buffer.
|
||
let head_slice = |m: &[i64], h: usize| -> Vec<i64> {
|
||
let mut s = vec![0i64; seq * dh];
|
||
for i in 0..seq {
|
||
for t in 0..dh {
|
||
s[i * dh + t] = m[i * d + h * dh + t];
|
||
}
|
||
}
|
||
s
|
||
};
|
||
for h in 0..heads {
|
||
let oh = attention(&head_slice(q, h), &head_slice(k, h), &head_slice(v, h), seq, dh, frac);
|
||
for i in 0..seq {
|
||
for t in 0..dh {
|
||
out[i * d + h * dh + t] = oh[i * dh + t];
|
||
}
|
||
}
|
||
}
|
||
out
|
||
}
|
||
|
||
/// Integer-exact inverse GDN (IGDN) — the divisive-normalization activation of
|
||
/// the CNN learned-codec synthesis (bmshj2018 / mbt2018 / cheng2020 `g_s`):
|
||
/// `y_i = x_i · sqrt(beta_i + Σ_j gamma_ij · x_j²)`, per spatial position across
|
||
/// the `c` channels. `x` is `[c, hw]` channel-major at `2^frac`; `gamma` is
|
||
/// `c·c` (row-major, `[out_i][in_j]`) at `2^gf`; `beta` is `c` at `2^(2·frac)`.
|
||
/// The sqrt is exact integer `isqrt` — no float, no Newton — so the op is
|
||
/// byte-identical on every machine (the missing piece to make a real CNN learned
|
||
/// codec's synthesis `NeuralIntegerExact`).
|
||
///
|
||
/// Runs the fast kernel ([`crate::int_kernels`]) unless the current
|
||
/// [`crate::int_kernels::Policy`] selects [`reference::igdn`]; both give the same
|
||
/// integer for every output element.
|
||
pub fn igdn(x: &[i64], c: usize, hw: usize, frac: u32, gamma: &[i64], beta: &[i64], gf: u32) -> Vec<i64> {
|
||
if crate::int_kernels::policy().path == crate::int_kernels::KernelPath::Fast
|
||
&& let Some(out) = crate::int_kernels::igdn(x, c, hw, frac, gamma, beta, gf)
|
||
{
|
||
return out;
|
||
}
|
||
reference::igdn(x, c, hw, frac, gamma, beta, gf)
|
||
}
|
||
|
||
/// The scalar reference kernel the fast IGDN is checked against, kept verbatim.
|
||
#[doc(hidden)]
|
||
pub mod reference {
|
||
/// Integer-exact inverse GDN (IGDN) — the divisive-normalization activation of
|
||
/// the CNN learned-codec synthesis (bmshj2018 / mbt2018 / cheng2020 `g_s`):
|
||
/// `y_i = x_i · sqrt(beta_i + Σ_j gamma_ij · x_j²)`, per spatial position across
|
||
/// the `c` channels. `x` is `[c, hw]` channel-major at `2^frac`; `gamma` is
|
||
/// `c·c` (row-major, `[out_i][in_j]`) at `2^gf`; `beta` is `c` at `2^(2·frac)`.
|
||
/// The sqrt is exact integer `isqrt` — no float, no Newton — so the op is
|
||
/// byte-identical on every machine (the missing piece to make a real CNN learned
|
||
/// codec's synthesis `NeuralIntegerExact`).
|
||
pub fn igdn(x: &[i64], c: usize, hw: usize, frac: u32, gamma: &[i64], beta: &[i64], gf: u32) -> Vec<i64> {
|
||
let mut out = vec![0i64; c * hw];
|
||
for p in 0..hw {
|
||
// x_j² at 2^(2·frac)
|
||
let mut x2 = vec![0i128; c];
|
||
for j in 0..c {
|
||
let xj = x[j * hw + p] as i128;
|
||
x2[j] = xj * xj;
|
||
}
|
||
for i in 0..c {
|
||
// norm_i = beta_i + Σ_j gamma_ij·x_j² at 2^(2·frac)
|
||
let mut norm = beta[i] as i128;
|
||
let grow = i * c;
|
||
for j in 0..c {
|
||
norm += (gamma[grow + j] as i128 * x2[j]) >> gf;
|
||
}
|
||
// sqrt(norm@2^(2frac)) = isqrt(norm) at 2^frac
|
||
let s = (norm.max(0) as u128).isqrt() as i64;
|
||
out[i * hw + p] = ((x[i * hw + p] as i128 * s as i128) >> frac) as i64;
|
||
}
|
||
}
|
||
out
|
||
}
|
||
}
|
||
|
||
/// Per-token layer norm of a `seq×d` row-major matrix: [`layernorm`] applied
|
||
/// independently to each row over its `d` features. `gamma`/`beta` (length `d`)
|
||
/// and `eps` are shared across rows, as in a transformer. Output `seq×d` at `2^frac`.
|
||
pub fn layernorm_rows(
|
||
x: &[i64],
|
||
seq: usize,
|
||
d: usize,
|
||
frac: u32,
|
||
gamma: &[i64],
|
||
beta: &[i64],
|
||
eps: i64,
|
||
) -> Vec<i64> {
|
||
let mut out = vec![0i64; seq * d];
|
||
for i in 0..seq {
|
||
let y = layernorm(&x[i * d..(i + 1) * d], frac, gamma, beta, eps);
|
||
out[i * d..(i + 1) * d].copy_from_slice(&y);
|
||
}
|
||
out
|
||
}
|
||
|
||
/// Affine linear layer `Y = X·W + b` (fixed-point): `X` is `seq×d_in`, `W` is
|
||
/// `d_in×d_out` row-major, `b` length `d_out` (added per output column); all at
|
||
/// `2^frac`, output `seq×d_out` at `2^frac`. The integer core of every
|
||
/// projection / FFN layer.
|
||
pub fn linear(
|
||
x: &[i64],
|
||
seq: usize,
|
||
d_in: usize,
|
||
w: &[i64],
|
||
d_out: usize,
|
||
b: &[i64],
|
||
frac: u32,
|
||
) -> Vec<i64> {
|
||
let mut y = matmul(x, w, seq, d_in, d_out, frac);
|
||
for i in 0..seq {
|
||
for j in 0..d_out {
|
||
y[i * d_out + j] += b[j];
|
||
}
|
||
}
|
||
y
|
||
}
|
||
|
||
/// Elementwise residual add of two equal-length fixed-point buffers.
|
||
pub fn add(a: &[i64], b: &[i64]) -> Vec<i64> {
|
||
a.iter().zip(b).map(|(&x, &y)| x + y).collect()
|
||
}
|
||
|
||
/// A pre-norm transformer (MetaFormer) encoder block, fully integer-exact:
|
||
///
|
||
/// ```text
|
||
/// a = x + Wo·MHA( LN1(x)·{Wq,Wk,Wv} )
|
||
/// y = a + ( gelu( LN2(a)·W1 + b1 )·W2 + b2 )
|
||
/// ```
|
||
///
|
||
/// Q/K/V/O projections are bias-free (`d×d`); the FFN is `d → hidden → d` with
|
||
/// bias and a [`gelu`] nonlinearity. All weights/activations share the `2^frac`
|
||
/// fixed-point scale; `eps` is at `2^(2·frac)`. Every step composes a primitive
|
||
/// in this module, so a whole stack of these is byte-identical on every machine
|
||
/// — the repeating unit of a `NeuralIntegerExact` MetaFormer synthesis net.
|
||
pub struct TransformerBlock {
|
||
/// Model width.
|
||
pub d: usize,
|
||
/// Attention heads (must divide `d`).
|
||
pub heads: usize,
|
||
/// FFN inner width.
|
||
pub hidden: usize,
|
||
/// Fixed-point fractional bits (`≤ 16`).
|
||
pub frac: u32,
|
||
/// LayerNorm epsilon at `2^(2·frac)`.
|
||
pub eps: i64,
|
||
/// Query projection, `d×d` row-major.
|
||
pub wq: Vec<i64>,
|
||
/// Key projection, `d×d` row-major.
|
||
pub wk: Vec<i64>,
|
||
/// Value projection, `d×d` row-major.
|
||
pub wv: Vec<i64>,
|
||
/// Output projection, `d×d` row-major.
|
||
pub wo: Vec<i64>,
|
||
/// Pre-attention LayerNorm scale (length `d`).
|
||
pub ln1_g: Vec<i64>,
|
||
/// Pre-attention LayerNorm shift (length `d`).
|
||
pub ln1_b: Vec<i64>,
|
||
/// Pre-FFN LayerNorm scale (length `d`).
|
||
pub ln2_g: Vec<i64>,
|
||
/// Pre-FFN LayerNorm shift (length `d`).
|
||
pub ln2_b: Vec<i64>,
|
||
/// FFN up-projection, `d×hidden` row-major.
|
||
pub w1: Vec<i64>,
|
||
/// FFN up-projection bias (length `hidden`).
|
||
pub b1: Vec<i64>,
|
||
/// FFN down-projection, `hidden×d` row-major.
|
||
pub w2: Vec<i64>,
|
||
/// FFN down-projection bias (length `d`).
|
||
pub b2: Vec<i64>,
|
||
}
|
||
|
||
impl TransformerBlock {
|
||
/// Run the block on a `seq×d` input (row-major, `2^frac`). Output `seq×d`.
|
||
pub fn forward(&self, x: &[i64], seq: usize) -> Vec<i64> {
|
||
let (d, f) = (self.d, self.frac);
|
||
// attention sublayer: pre-norm → Q/K/V → MHA → Wo → residual
|
||
let ln1 = layernorm_rows(x, seq, d, f, &self.ln1_g, &self.ln1_b, self.eps);
|
||
let q = matmul(&ln1, &self.wq, seq, d, d, f);
|
||
let k = matmul(&ln1, &self.wk, seq, d, d, f);
|
||
let v = matmul(&ln1, &self.wv, seq, d, d, f);
|
||
let attn = multi_head_attention(&q, &k, &v, seq, d, self.heads, f);
|
||
let attn_o = matmul(&attn, &self.wo, seq, d, d, f);
|
||
let a = add(x, &attn_o);
|
||
// FFN sublayer: pre-norm → up → gelu → down → residual
|
||
let ln2 = layernorm_rows(&a, seq, d, f, &self.ln2_g, &self.ln2_b, self.eps);
|
||
let h = gelu(&linear(&ln2, seq, d, &self.w1, self.hidden, &self.b1, f), f);
|
||
let ffn = linear(&h, seq, self.hidden, &self.w2, d, &self.b2, f);
|
||
add(&a, &ffn)
|
||
}
|
||
}
|
||
|
||
/// A stack of [`TransformerBlock`]s — the MetaFormer synthesis skeleton. Running
|
||
/// a `seq×d` input through every block in turn; with trained per-block weights
|
||
/// this *is* the attention-based synthesis net (e.g. JPEG AI's), and because
|
||
/// every block is integer-exact the whole stack is byte-identical on every
|
||
/// machine — its output can go in a receipt.
|
||
pub struct TransformerEncoder {
|
||
/// The blocks, applied in order. All must share the same `d`/`frac`.
|
||
pub blocks: Vec<TransformerBlock>,
|
||
}
|
||
|
||
impl TransformerEncoder {
|
||
/// New encoder from an ordered list of blocks.
|
||
pub fn new(blocks: Vec<TransformerBlock>) -> Self {
|
||
Self { blocks }
|
||
}
|
||
|
||
/// Number of stacked blocks.
|
||
pub fn depth(&self) -> usize {
|
||
self.blocks.len()
|
||
}
|
||
|
||
/// Run `x` (`seq×d`, `2^frac`) through every block in order. An empty stack
|
||
/// is the identity.
|
||
pub fn forward(&self, x: &[i64], seq: usize) -> Vec<i64> {
|
||
let mut h = x.to_vec();
|
||
for b in &self.blocks {
|
||
h = b.forward(&h, seq);
|
||
}
|
||
h
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn igdn_matches_float_and_is_deterministic() {
|
||
const FRAC: u32 = 14;
|
||
const GF: u32 = 16;
|
||
let (c, hw) = (6usize, 4usize);
|
||
let q = |v: f64, f: u32| (v * (1i64 << f) as f64).round() as i64;
|
||
// deterministic pseudo-random x∈[-1,1], gamma∈[0,0.3], beta∈[0.1,0.4]
|
||
let r = |i: usize, s: usize| {
|
||
let h = (i.wrapping_mul(2654435761).wrapping_add(s.wrapping_mul(40503))) & 0xffff;
|
||
h as f64 / 65535.0
|
||
};
|
||
let xf: Vec<f64> = (0..c * hw).map(|i| 2.0 * r(i, 1) - 1.0).collect();
|
||
let gf_: Vec<f64> = (0..c * c).map(|i| 0.3 * r(i, 2)).collect();
|
||
let bf: Vec<f64> = (0..c).map(|i| 0.1 + 0.3 * r(i, 3)).collect();
|
||
let xi: Vec<i64> = xf.iter().map(|&v| q(v, FRAC)).collect();
|
||
let gi: Vec<i64> = gf_.iter().map(|&v| q(v, GF)).collect();
|
||
let bi: Vec<i64> = bf.iter().map(|&v| q(v, 2 * FRAC)).collect();
|
||
|
||
let yi = igdn(&xi, c, hw, FRAC, &gi, &bi, GF);
|
||
assert_eq!(yi, igdn(&xi, c, hw, FRAC, &gi, &bi, GF)); // deterministic
|
||
|
||
// float reference: y_i = x_i·sqrt(beta_i + Σ_j gamma_ij·x_j²)
|
||
let mut worst = 0.0f64;
|
||
for p in 0..hw {
|
||
for i in 0..c {
|
||
let mut norm = bf[i];
|
||
for j in 0..c {
|
||
norm += gf_[i * c + j] * xf[j * hw + p].powi(2);
|
||
}
|
||
let yf = xf[i * hw + p] * norm.max(0.0).sqrt();
|
||
worst = worst.max((yi[i * hw + p] as f64 / (1i64 << FRAC) as f64 - yf).abs());
|
||
}
|
||
}
|
||
assert!(worst < 5e-3, "igdn max abs error {worst} exceeds the 5e-3 bound");
|
||
}
|
||
|
||
fn float_softmax(x: &[f64]) -> Vec<f64> {
|
||
let m = x.iter().cloned().fold(f64::MIN, f64::max);
|
||
let e: Vec<f64> = x.iter().map(|&v| (v - m).exp()).collect();
|
||
let s: f64 = e.iter().sum();
|
||
e.iter().map(|&v| v / s).collect()
|
||
}
|
||
|
||
#[test]
|
||
fn softmax_matches_float_within_bound() {
|
||
const FRAC: u32 = 12;
|
||
let cases: &[&[f64]] = &[
|
||
&[0.0, 1.0, 2.0, 3.0],
|
||
&[-5.0, -1.0, 0.0, 0.5, 4.0],
|
||
&[2.2, 2.2, 2.2],
|
||
&[-8.0, 8.0],
|
||
&[0.1, -0.2, 0.3, -0.4, 0.5, -0.6, 0.7],
|
||
];
|
||
let mut worst = 0.0f64;
|
||
for &c in cases {
|
||
let xi: Vec<i64> = c.iter().map(|&v| (v * (1i64 << FRAC) as f64).round() as i64).collect();
|
||
let pi = softmax(&xi, FRAC);
|
||
let pf = float_softmax(c);
|
||
for (a, b) in pi.iter().zip(&pf) {
|
||
let af = *a as f64 / (1i64 << PROB_FRAC) as f64;
|
||
worst = worst.max((af - b).abs());
|
||
}
|
||
// probabilities sum to EXACTLY 1.0 (the remainder rides argmax).
|
||
let tot: i64 = pi.iter().sum();
|
||
assert_eq!(tot, 1i64 << PROB_FRAC, "softmax must sum to exactly 2^{PROB_FRAC}");
|
||
}
|
||
assert!(worst < 5e-3, "softmax max abs error {worst} exceeds the 5e-3 bound");
|
||
}
|
||
|
||
fn float_layernorm(x: &[f64], g: &[f64], b: &[f64], eps: f64) -> Vec<f64> {
|
||
let n = x.len() as f64;
|
||
let mean = x.iter().sum::<f64>() / n;
|
||
let var = x.iter().map(|&v| (v - mean).powi(2)).sum::<f64>() / n;
|
||
let s = (var + eps).sqrt();
|
||
(0..x.len()).map(|i| (x[i] - mean) / s * g[i] + b[i]).collect()
|
||
}
|
||
|
||
#[test]
|
||
fn layernorm_matches_float_within_bound() {
|
||
const FRAC: u32 = 14;
|
||
let eps_f = 1e-5;
|
||
let eps = (eps_f * (1i64 << (2 * FRAC)) as f64).round() as i64;
|
||
let cases: &[&[f64]] = &[
|
||
&[1.0, 2.0, 3.0, 4.0, 5.0],
|
||
&[-2.0, 0.5, 0.5, 3.0],
|
||
&[0.1, -0.1, 0.2, -0.2, 0.05, 0.3],
|
||
&[7.0, 7.0, 7.0, 7.0],
|
||
];
|
||
let q = |v: f64| (v * (1i64 << FRAC) as f64).round() as i64;
|
||
let mut worst = 0.0f64;
|
||
for &c in cases {
|
||
let g: Vec<f64> = c.iter().enumerate().map(|(i, _)| 1.0 + 0.1 * i as f64).collect();
|
||
let b: Vec<f64> = c.iter().enumerate().map(|(i, _)| 0.05 * i as f64 - 0.1).collect();
|
||
let xi: Vec<i64> = c.iter().map(|&v| q(v)).collect();
|
||
let gi: Vec<i64> = g.iter().map(|&v| q(v)).collect();
|
||
let bi: Vec<i64> = b.iter().map(|&v| q(v)).collect();
|
||
let yi = layernorm(&xi, FRAC, &gi, &bi, eps);
|
||
let yf = float_layernorm(c, &g, &b, eps_f);
|
||
for (a, e) in yi.iter().zip(&yf) {
|
||
worst = worst.max((*a as f64 / (1i64 << FRAC) as f64 - e).abs());
|
||
}
|
||
}
|
||
assert!(worst < 5e-3, "layernorm max abs error {worst} exceeds the 5e-3 bound");
|
||
}
|
||
|
||
#[test]
|
||
fn layernorm_is_deterministic() {
|
||
let x = vec![100i64, -300, 4096, 9000, -50];
|
||
let g = vec![1 << 14; 5];
|
||
let b = vec![0i64; 5];
|
||
assert_eq!(layernorm(&x, 14, &g, &b, 3), layernorm(&x, 14, &g, &b, 3));
|
||
}
|
||
|
||
#[test]
|
||
fn matmul_matches_float() {
|
||
const FRAC: u32 = 14;
|
||
let q = |v: f64| (v * (1i64 << FRAC) as f64).round() as i64;
|
||
// A 2×3, B 3×2
|
||
let af = [0.5, -1.0, 2.0, 1.5, 0.25, -0.5];
|
||
let bf = [1.0, 0.5, -2.0, 0.0, 0.75, 1.0];
|
||
let a: Vec<i64> = af.iter().map(|&v| q(v)).collect();
|
||
let b: Vec<i64> = bf.iter().map(|&v| q(v)).collect();
|
||
let c = matmul(&a, &b, 2, 3, 2, FRAC); // result at 2^FRAC
|
||
// float reference
|
||
let mut cf = [0.0f64; 4];
|
||
for i in 0..2 { for j in 0..2 { for p in 0..3 { cf[i*2+j] += af[i*3+p]*bf[p*2+j]; } } }
|
||
let mut worst = 0.0f64;
|
||
for (x, e) in c.iter().zip(&cf) {
|
||
worst = worst.max((*x as f64 / (1i64 << FRAC) as f64 - e).abs());
|
||
}
|
||
assert!(worst < 2e-3, "matmul max abs error {worst}");
|
||
}
|
||
|
||
#[test]
|
||
fn attention_block_matches_float_and_is_deterministic() {
|
||
const FRAC: u32 = 13;
|
||
let (seq, d) = (4usize, 8usize);
|
||
// deterministic pseudo-random Q,K,V in [-1,1]
|
||
let genv = |salt: usize| -> Vec<f64> {
|
||
(0..seq * d).map(|i| {
|
||
let h = (i.wrapping_mul(2654435761).wrapping_add(salt.wrapping_mul(40503))) & 0xffff;
|
||
h as f64 / 32768.0 - 1.0
|
||
}).collect()
|
||
};
|
||
let (qf, kf, vf) = (genv(1), genv(2), genv(3));
|
||
let q = |v: &f64| (v * (1i64 << FRAC) as f64).round() as i64;
|
||
let (qi, ki, vi): (Vec<i64>, Vec<i64>, Vec<i64>) =
|
||
(qf.iter().map(q).collect(), kf.iter().map(q).collect(), vf.iter().map(q).collect());
|
||
|
||
let out = attention(&qi, &ki, &vi, seq, d, FRAC);
|
||
assert_eq!(out, attention(&qi, &ki, &vi, seq, d, FRAC)); // deterministic
|
||
|
||
// float reference: softmax(QKᵀ/√d)·V
|
||
let sd = (d as f64).sqrt();
|
||
let mut worst = 0.0f64;
|
||
for i in 0..seq {
|
||
let mut sc = vec![0.0f64; seq];
|
||
for j in 0..seq {
|
||
let mut dot = 0.0; for t in 0..d { dot += qf[i*d+t]*kf[j*d+t]; }
|
||
sc[j] = dot / sd;
|
||
}
|
||
let a = float_softmax(&sc);
|
||
for t in 0..d {
|
||
let mut o = 0.0; for j in 0..seq { o += a[j]*vf[j*d+t]; }
|
||
worst = worst.max((out[i*d+t] as f64 / (1i64 << FRAC) as f64 - o).abs());
|
||
}
|
||
}
|
||
assert!(worst < 1e-2, "attention max abs error {worst} exceeds 1e-2");
|
||
}
|
||
|
||
#[test]
|
||
fn softmax_is_deterministic_and_float_free() {
|
||
// Same integer input → identical bytes, every call, every machine.
|
||
let x = vec![-3000i64, 100, 4096, 4096, -100, 9000];
|
||
assert_eq!(softmax(&x, 12), softmax(&x, 12));
|
||
// Monotone: a larger logit never gets a smaller probability.
|
||
let p = softmax(&[0i64, 1 << 12, 2 << 12, 3 << 12], 12);
|
||
assert!(p.windows(2).all(|w| w[0] <= w[1]));
|
||
}
|
||
|
||
fn float_gelu(x: f64) -> f64 {
|
||
x / (1.0 + (-1.702 * x).exp()) // x·σ(1.702·x) — the op being integerized
|
||
}
|
||
|
||
#[test]
|
||
fn gelu_matches_float_within_bound() {
|
||
const FRAC: u32 = 14;
|
||
let q = |v: f64| (v * (1i64 << FRAC) as f64).round() as i64;
|
||
let xs: Vec<f64> = (-60..=60).map(|i| i as f64 / 10.0).collect(); // [-6, 6]
|
||
let xi: Vec<i64> = xs.iter().map(|&v| q(v)).collect();
|
||
let yi = gelu(&xi, FRAC);
|
||
let mut worst = 0.0f64;
|
||
for (a, &xf) in yi.iter().zip(&xs) {
|
||
worst = worst.max((*a as f64 / (1i64 << FRAC) as f64 - float_gelu(xf)).abs());
|
||
}
|
||
assert!(worst < 5e-3, "gelu max abs error {worst} exceeds the 5e-3 bound");
|
||
// Shape sanity: gelu(0)=0, saturates to ~x for large x, ~0 for very negative.
|
||
assert_eq!(gelu(&[0], FRAC), vec![0]);
|
||
let big = gelu(&[q(6.0)], FRAC)[0] as f64 / (1i64 << FRAC) as f64;
|
||
assert!((big - 6.0).abs() < 1e-2);
|
||
}
|
||
|
||
#[test]
|
||
fn sigmoid_matches_float_and_is_deterministic() {
|
||
const FRAC: u32 = 16;
|
||
let q = |v: f64| (v * (1i64 << FRAC) as f64).round() as i64;
|
||
let mut worst = 0.0f64;
|
||
for i in -80..=80 {
|
||
let z = i as f64 / 10.0;
|
||
let s = sigmoid(q(z), FRAC);
|
||
assert_eq!(s, sigmoid(q(z), FRAC)); // deterministic
|
||
let sf = 1.0 / (1.0 + (-z).exp());
|
||
worst = worst.max((s as f64 / (1i64 << FRAC) as f64 - sf).abs());
|
||
}
|
||
assert!(worst < 3e-3, "sigmoid max abs error {worst} exceeds 3e-3");
|
||
}
|
||
|
||
#[test]
|
||
fn multi_head_one_head_equals_single_head() {
|
||
const FRAC: u32 = 13;
|
||
let (seq, d) = (4usize, 8usize);
|
||
let mk = |salt: usize| -> Vec<i64> {
|
||
(0..seq * d)
|
||
.map(|i| {
|
||
let h = (i.wrapping_mul(2654435761).wrapping_add(salt.wrapping_mul(40503))) & 0xffff;
|
||
(h as i64 - 32768) >> 2
|
||
})
|
||
.collect()
|
||
};
|
||
let (q, k, v) = (mk(1), mk(2), mk(3));
|
||
assert_eq!(
|
||
multi_head_attention(&q, &k, &v, seq, d, 1, FRAC),
|
||
attention(&q, &k, &v, seq, d, FRAC),
|
||
"1 head must reduce to single-head attention"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn multi_head_matches_float_and_is_deterministic() {
|
||
const FRAC: u32 = 13;
|
||
let (seq, d, heads) = (4usize, 8usize, 2usize);
|
||
let dh = d / heads;
|
||
let genf = |salt: usize| -> Vec<f64> {
|
||
(0..seq * d)
|
||
.map(|i| {
|
||
let h = (i.wrapping_mul(2654435761).wrapping_add(salt.wrapping_mul(40503))) & 0xffff;
|
||
h as f64 / 32768.0 - 1.0
|
||
})
|
||
.collect()
|
||
};
|
||
let (qf, kf, vf) = (genf(1), genf(2), genf(3));
|
||
let q = |v: &f64| (v * (1i64 << FRAC) as f64).round() as i64;
|
||
let (qi, ki, vi): (Vec<i64>, Vec<i64>, Vec<i64>) =
|
||
(qf.iter().map(q).collect(), kf.iter().map(q).collect(), vf.iter().map(q).collect());
|
||
|
||
let out = multi_head_attention(&qi, &ki, &vi, seq, d, heads, FRAC);
|
||
assert_eq!(out, multi_head_attention(&qi, &ki, &vi, seq, d, heads, FRAC)); // deterministic
|
||
|
||
// Float reference: per head, softmax(QKᵀ/√dh)·V over that head's columns.
|
||
let sd = (dh as f64).sqrt();
|
||
let mut worst = 0.0f64;
|
||
for h in 0..heads {
|
||
let col = |m: &[f64], i: usize, t: usize| m[i * d + h * dh + t];
|
||
for i in 0..seq {
|
||
let mut sc = vec![0.0f64; seq];
|
||
for j in 0..seq {
|
||
let mut dot = 0.0;
|
||
for t in 0..dh {
|
||
dot += col(&qf, i, t) * col(&kf, j, t);
|
||
}
|
||
sc[j] = dot / sd;
|
||
}
|
||
let a = float_softmax(&sc);
|
||
for t in 0..dh {
|
||
let mut o = 0.0;
|
||
for j in 0..seq {
|
||
o += a[j] * col(&vf, j, t);
|
||
}
|
||
let got = out[i * d + h * dh + t] as f64 / (1i64 << FRAC) as f64;
|
||
worst = worst.max((got - o).abs());
|
||
}
|
||
}
|
||
}
|
||
assert!(worst < 1e-2, "multi-head attention max abs error {worst} exceeds 1e-2");
|
||
}
|
||
|
||
#[test]
|
||
fn transformer_block_matches_float_and_is_deterministic() {
|
||
const FRAC: u32 = 13;
|
||
let (seq, d, heads, hidden) = (3usize, 8usize, 2usize, 16usize);
|
||
let scale = (1i64 << FRAC) as f64;
|
||
let qf = |v: f64| (v * scale).round() as i64;
|
||
let qv = |f: &[f64]| f.iter().map(|&v| qf(v)).collect::<Vec<i64>>();
|
||
// deterministic pseudo-random weights in [-0.5, 0.5]
|
||
let rnd = |salt: usize, n: usize| -> Vec<f64> {
|
||
(0..n)
|
||
.map(|i| {
|
||
let h = (i.wrapping_mul(2654435761).wrapping_add(salt.wrapping_mul(40503))) & 0xffff;
|
||
(h as f64 / 32768.0 - 1.0) * 0.5
|
||
})
|
||
.collect()
|
||
};
|
||
let xf = rnd(0, seq * d);
|
||
let (wqf, wkf, wvf, wof) = (rnd(1, d * d), rnd(2, d * d), rnd(3, d * d), rnd(4, d * d));
|
||
let (w1f, b1f) = (rnd(5, d * hidden), rnd(6, hidden));
|
||
let (w2f, b2f) = (rnd(7, hidden * d), rnd(8, d));
|
||
let ln1g: Vec<f64> = (0..d).map(|i| 1.0 + 0.05 * i as f64).collect();
|
||
let ln1b: Vec<f64> = (0..d).map(|i| 0.02 * i as f64 - 0.05).collect();
|
||
let eps_f = 1e-5;
|
||
|
||
let block = TransformerBlock {
|
||
d,
|
||
heads,
|
||
hidden,
|
||
frac: FRAC,
|
||
eps: (eps_f * (1i64 << (2 * FRAC)) as f64).round() as i64,
|
||
wq: qv(&wqf),
|
||
wk: qv(&wkf),
|
||
wv: qv(&wvf),
|
||
wo: qv(&wof),
|
||
ln1_g: qv(&ln1g),
|
||
ln1_b: qv(&ln1b),
|
||
ln2_g: qv(&ln1g),
|
||
ln2_b: qv(&ln1b),
|
||
w1: qv(&w1f),
|
||
b1: qv(&b1f),
|
||
w2: qv(&w2f),
|
||
b2: qv(&b2f),
|
||
};
|
||
let xi = qv(&xf);
|
||
let out = block.forward(&xi, seq);
|
||
assert_eq!(out, block.forward(&xi, seq), "block must be deterministic");
|
||
|
||
// Float reference mirroring forward() exactly.
|
||
let fmatmul = |x: &[f64], w: &[f64], m: usize, ki: usize, n: usize| -> Vec<f64> {
|
||
let mut c = vec![0.0; m * n];
|
||
for i in 0..m {
|
||
for p in 0..ki {
|
||
for j in 0..n {
|
||
c[i * n + j] += x[i * ki + p] * w[p * n + j];
|
||
}
|
||
}
|
||
}
|
||
c
|
||
};
|
||
let fln = |x: &[f64], g: &[f64], b: &[f64]| -> Vec<f64> {
|
||
let rows = x.len() / d;
|
||
let mut o = vec![0.0; x.len()];
|
||
for i in 0..rows {
|
||
let r = &x[i * d..(i + 1) * d];
|
||
let mean = r.iter().sum::<f64>() / d as f64;
|
||
let var = r.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / d as f64;
|
||
let s = (var + eps_f).sqrt();
|
||
for t in 0..d {
|
||
o[i * d + t] = (r[t] - mean) / s * g[t] + b[t];
|
||
}
|
||
}
|
||
o
|
||
};
|
||
let fmha = |q: &[f64], k: &[f64], v: &[f64]| -> Vec<f64> {
|
||
let dh = d / heads;
|
||
let sd = (dh as f64).sqrt();
|
||
let mut out = vec![0.0; seq * d];
|
||
for hh in 0..heads {
|
||
for i in 0..seq {
|
||
let mut sc = vec![0.0; seq];
|
||
for j in 0..seq {
|
||
let mut dot = 0.0;
|
||
for t in 0..dh {
|
||
dot += q[i * d + hh * dh + t] * k[j * d + hh * dh + t];
|
||
}
|
||
sc[j] = dot / sd;
|
||
}
|
||
let a = float_softmax(&sc);
|
||
for t in 0..dh {
|
||
let mut o = 0.0;
|
||
for j in 0..seq {
|
||
o += a[j] * v[j * d + hh * dh + t];
|
||
}
|
||
out[i * d + hh * dh + t] = o;
|
||
}
|
||
}
|
||
}
|
||
out
|
||
};
|
||
let addb = |y: &mut [f64], b: &[f64], n: usize| {
|
||
for i in 0..(y.len() / n) {
|
||
for j in 0..n {
|
||
y[i * n + j] += b[j];
|
||
}
|
||
}
|
||
};
|
||
|
||
let ln1 = fln(&xf, &ln1g, &ln1b);
|
||
let qm = fmatmul(&ln1, &wqf, seq, d, d);
|
||
let km = fmatmul(&ln1, &wkf, seq, d, d);
|
||
let vm = fmatmul(&ln1, &wvf, seq, d, d);
|
||
let attn = fmha(&qm, &km, &vm);
|
||
let attn_o = fmatmul(&attn, &wof, seq, d, d);
|
||
let a: Vec<f64> = xf.iter().zip(&attn_o).map(|(&x, &y)| x + y).collect();
|
||
let ln2 = fln(&a, &ln1g, &ln1b);
|
||
let mut h1 = fmatmul(&ln2, &w1f, seq, d, hidden);
|
||
addb(&mut h1, &b1f, hidden);
|
||
let h: Vec<f64> = h1.iter().map(|&v| float_gelu(v)).collect();
|
||
let mut f2 = fmatmul(&h, &w2f, seq, hidden, d);
|
||
addb(&mut f2, &b2f, d);
|
||
let yf: Vec<f64> = a.iter().zip(&f2).map(|(&x, &y)| x + y).collect();
|
||
|
||
let mut worst = 0.0f64;
|
||
for (o, e) in out.iter().zip(&yf) {
|
||
worst = worst.max((*o as f64 / scale - e).abs());
|
||
}
|
||
assert!(worst < 3e-2, "transformer block max abs error {worst} exceeds 3e-2");
|
||
}
|
||
|
||
#[test]
|
||
fn transformer_encoder_stacks_and_is_deterministic() {
|
||
const FRAC: u32 = 13;
|
||
let (seq, d, heads, hidden) = (3usize, 8usize, 2usize, 16usize);
|
||
let scale = (1i64 << FRAC) as f64;
|
||
let qf = |v: f64| (v * scale).round() as i64;
|
||
let qv = |f: &[f64]| f.iter().map(|&v| qf(v)).collect::<Vec<i64>>();
|
||
let rnd = |salt: usize, n: usize| -> Vec<f64> {
|
||
(0..n)
|
||
.map(|i| {
|
||
let h = (i.wrapping_mul(2654435761).wrapping_add(salt.wrapping_mul(40503))) & 0xffff;
|
||
(h as f64 / 32768.0 - 1.0) * 0.5
|
||
})
|
||
.collect()
|
||
};
|
||
let ln_g: Vec<f64> = (0..d).map(|i| 1.0 + 0.05 * i as f64).collect();
|
||
let ln_b: Vec<f64> = (0..d).map(|i| 0.02 * i as f64 - 0.05).collect();
|
||
let eps = (1e-5 * (1i64 << (2 * FRAC)) as f64).round() as i64;
|
||
let mk_block = |s: usize| TransformerBlock {
|
||
d,
|
||
heads,
|
||
hidden,
|
||
frac: FRAC,
|
||
eps,
|
||
wq: qv(&rnd(s + 1, d * d)),
|
||
wk: qv(&rnd(s + 2, d * d)),
|
||
wv: qv(&rnd(s + 3, d * d)),
|
||
wo: qv(&rnd(s + 4, d * d)),
|
||
ln1_g: qv(&ln_g),
|
||
ln1_b: qv(&ln_b),
|
||
ln2_g: qv(&ln_g),
|
||
ln2_b: qv(&ln_b),
|
||
w1: qv(&rnd(s + 5, d * hidden)),
|
||
b1: qv(&rnd(s + 6, hidden)),
|
||
w2: qv(&rnd(s + 7, hidden * d)),
|
||
b2: qv(&rnd(s + 8, d)),
|
||
};
|
||
let xi = qv(&rnd(0, seq * d));
|
||
|
||
// Empty stack is the identity.
|
||
let empty = TransformerEncoder::new(vec![]);
|
||
assert_eq!(empty.depth(), 0);
|
||
assert_eq!(empty.forward(&xi, seq), xi);
|
||
|
||
// A 2-block stack equals applying the two blocks in sequence...
|
||
let manual = mk_block(20).forward(&mk_block(10).forward(&xi, seq), seq);
|
||
let enc = TransformerEncoder::new(vec![mk_block(10), mk_block(20)]);
|
||
assert_eq!(enc.depth(), 2);
|
||
let out = enc.forward(&xi, seq);
|
||
assert_eq!(out, manual, "encoder must equal sequential block application");
|
||
// ...and is deterministic to the byte.
|
||
assert_eq!(out, enc.forward(&xi, seq), "encoder must be deterministic");
|
||
}
|
||
}
|
||
|
||
/// The fast IGDN against the scalar reference, element by element, in each of its
|
||
/// arithmetic modes: 32×32→64-bit terms, `i64` terms, the `i128` fallback, and the
|
||
/// `i64`-sum / `i128`-output-product combination.
|
||
#[cfg(test)]
|
||
mod igdn_equivalence {
|
||
use super::*;
|
||
use crate::int_kernels::testkit::{Rng, FAST};
|
||
use crate::int_kernels::with_policy;
|
||
|
||
#[allow(clippy::too_many_arguments)]
|
||
fn check(g: &mut Rng, cases: usize, xr: i64, gr: i64, br: i64, gf_lo: u32, signed_gamma: bool, label: &str) {
|
||
for _ in 0..cases {
|
||
let c = 1 + g.below(12);
|
||
let hw = 1 + g.below(150);
|
||
let frac = 8 + g.below(8) as u32;
|
||
let gf = gf_lo + g.below(8) as u32;
|
||
let x = g.vec(c * hw, xr, 10);
|
||
let gamma: Vec<i64> =
|
||
g.vec(c * c, gr, 20).into_iter().map(|v| if signed_gamma { v } else { v.abs() }).collect();
|
||
let beta: Vec<i64> = g.vec(c, br, 0).into_iter().map(i64::abs).collect();
|
||
let want = reference::igdn(&x, c, hw, frac, &gamma, &beta, gf);
|
||
for p in FAST {
|
||
let got = with_policy(p, || igdn(&x, c, hw, frac, &gamma, &beta, gf));
|
||
assert_eq!(got, want, "igdn[{label}] {p:?} c={c} hw={hw} frac={frac} gf={gf}");
|
||
}
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn igdn_equals_reference_in_every_mode() {
|
||
let mut g = Rng(0x16D_0001);
|
||
// Q12 activations, the committed models' magnitudes (32×32→64 terms)
|
||
check(&mut g, 120, 1 << 15, 1 << 18, 1 << 26, 16, false, "narrow");
|
||
// x² beyond i32, terms still in i64; signed gamma (floor shifts of negatives)
|
||
check(&mut g, 120, 1 << 20, 1 << 16, 1 << 30, 16, true, "i64");
|
||
// gamma·x² beyond i64: the i128 path
|
||
check(&mut g, 120, 1 << 30, 1 << 24, 1 << 40, 16, true, "i128");
|
||
// every operand fits i32 but the SUM of terms does not fit i64: the sum
|
||
// bound, not the operand widths, must send this to the i128 path
|
||
check(&mut g, 120, 46340, i32::MAX as i64, 1 << 62, 0, false, "sum-bound");
|
||
// |x| ≥ 2^31 with gamma ≤ 1: i64 sums, i128 output product
|
||
check(&mut g, 60, (1i64 << 31) + (1 << 29), 1, 1 << 20, 30, false, "wide-x");
|
||
// an all-zero gamma says nothing about x: x² beyond i64 must still not be
|
||
// formed in i64 (the reference squares in i128)
|
||
check(&mut g, 40, 1i64 << 40, 0, 1 << 20, 16, false, "zero-gamma-huge-x");
|
||
}
|
||
|
||
/// A position whose norm is negative (negative beta) exercises the `max(0)` clamp.
|
||
#[test]
|
||
fn igdn_negative_norm_clamps_like_reference() {
|
||
let (c, hw) = (3usize, 5usize);
|
||
let x: Vec<i64> = (0..c * hw).map(|i| (i as i64 - 7) * 311).collect();
|
||
let gamma = vec![1i64; c * c];
|
||
let beta = vec![-(1i64 << 40), 5, -3];
|
||
let want = reference::igdn(&x, c, hw, 12, &gamma, &beta, 16);
|
||
for p in FAST {
|
||
assert_eq!(with_policy(p, || igdn(&x, c, hw, 12, &gamma, &beta, 16)), want, "{p:?}");
|
||
}
|
||
}
|
||
}
|