DaisyChain-Train / web /public /verified_core.js
Quazim0t0's picture
Upload web/public/verified_core.js with huggingface_hub
7fcb6bd verified
Raw
History Blame Contribute Delete
25.2 kB
// Verified INT8 compute — the emulated GPU logic, in the browser.
// A layer's forward runs THROUGH the units: quantize -> LUT multiply -> requant
// -> optional ReLU -> dequant. Backward is a straight-through estimator (the
// integer path has no gradient), so ordinary float weights still learn.
// Same units as the Python/Docker DaisyChain; here they're lookup tables.
(function (root) {
"use strict";
let TC; // TrainCore (matmul/transpose) — resolved per environment at the end
function quantize(X) {
let mx = 0; for (let i = 0; i < X.length; i++) { const a = Math.abs(X[i]); if (a > mx) mx = a; }
const scale = Math.max(mx / 127, 1e-8);
const q = new Int8Array(X.length);
for (let i = 0; i < X.length; i++) { let v = Math.round(X[i] / scale); q[i] = v < -128 ? -128 : v > 127 ? 127 : v; }
return { q, scale };
}
// int8 matmul via the verified multiply LUT: acc(m×n) = sum_k mulLUT[Xq,Wq]
function lutMatmulJS(Xq, Wq, m, k, n, L) {
const C = new Int32Array(m * n), mul = L.mul;
for (let i = 0; i < m; i++) {
for (let p = 0; p < k; p++) {
const au = (Xq[i * k + p] & 0xFF) * 256, wo = p * n, co = i * n;
for (let j = 0; j < n; j++) C[co + j] += mul[au + (Wq[wo + j] & 0xFF)];
}
}
return C;
}
// ---- 3xINT8 fast-accurate GEMM --------------------------------------------
// The CUTLASS example-27 "3xTF32" scheme, ported to the verified units:
// split each float into a coarse int8 part plus an int8-quantized residual,
// run three EXACT LUT GEMMs (hi·hi, hi·lo, lo·hi), drop the negligible
// lo·lo, and recombine. Same big/small decomposition NVIDIA uses to recover
// near-fp32 accuracy from TF32 tensor cores — here it recovers ~14-bit
// accuracy from the 8-bit units, at 3× the unit ops. Every product still
// goes through the verified mul8 LUT.
function quantize2(X) {
const hi = quantize(X);
const r = new Float32Array(X.length);
for (let i = 0; i < X.length; i++) r[i] = X[i] - hi.q[i] * hi.scale;
const lo = quantize(r);
return { hi, lo };
}
function combine3(hh, hl, lh, x, w, len) {
const out = new Float32Array(len);
const shh = x.hi.scale * w.hi.scale, shl = x.hi.scale * w.lo.scale, slh = x.lo.scale * w.hi.scale;
for (let i = 0; i < len; i++) out[i] = hh[i] * shh + hl[i] * shl + lh[i] * slh;
return out;
}
function lutMatmul3JS(Xf, Wf, m, k, n, L) { // sync, CPU LUT path
const x = quantize2(Xf), w = quantize2(Wf);
return combine3(lutMatmulJS(x.hi.q, w.hi.q, m, k, n, L),
lutMatmulJS(x.hi.q, w.lo.q, m, k, n, L),
lutMatmulJS(x.lo.q, w.hi.q, m, k, n, L), x, w, m * n);
}
async function lutMatmul3(Xf, Wf, m, k, n, L, matmulInt8) { // any backend
const x = quantize2(Xf), w = quantize2(Wf);
const mm = matmulInt8 || lutMatmulJS;
const [hh, hl, lh] = await Promise.all([
mm(x.hi.q, w.hi.q, m, k, n, L),
mm(x.hi.q, w.lo.q, m, k, n, L),
mm(x.lo.q, w.hi.q, m, k, n, L),
]);
return combine3(hh, hl, lh, x, w, m * n);
}
// ---- block-scaled verified GEMM (CUTLASS ex. 67/81 blockwise scaling) ------
// Per-ROW scales for the activations and per-COLUMN scales for the weights:
// the integer math through the mul8 LUT is completely unchanged — only the
// dequant uses rs[row]·cs[col] instead of one tensor-wide product, so a single
// outlier no longer crushes the quantization resolution of every other
// row/column. One LUT pass at this granularity beats the per-tensor 3-pass.
function quantizeRows(X, rows, cols) {
const q = new Int8Array(rows * cols), s = new Float32Array(rows);
for (let r = 0; r < rows; r++) {
let mx = 0;
for (let c = 0; c < cols; c++) { const a = Math.abs(X[r * cols + c]); if (a > mx) mx = a; }
const sc = Math.max(mx / 127, 1e-8); s[r] = sc;
for (let c = 0; c < cols; c++) {
const v = Math.round(X[r * cols + c] / sc);
q[r * cols + c] = v < -128 ? -128 : v > 127 ? 127 : v;
}
}
return { q, s };
}
function quantizeCols(W, rows, cols) {
const q = new Int8Array(rows * cols), s = new Float32Array(cols);
for (let c = 0; c < cols; c++) {
let mx = 0;
for (let r = 0; r < rows; r++) { const a = Math.abs(W[r * cols + c]); if (a > mx) mx = a; }
s[c] = Math.max(mx / 127, 1e-8);
}
for (let r = 0; r < rows; r++)
for (let c = 0; c < cols; c++) {
const v = Math.round(W[r * cols + c] / s[c]);
q[r * cols + c] = v < -128 ? -128 : v > 127 ? 127 : v;
}
return { q, s };
}
// ---- B2B MLP chain (CUTLASS ex. 13 two-GEMM fusion + ex. 23 epilogue
// reduction), respecced for cross-device exactness ---------------------------
// The MLP is the one back-to-back GEMM pair with no layernorm/softmax between
// (ReLU is already fused in the epilogue), so the intermediate h1 can be
// quantized ON the GPU and fed straight to the second GEMM. Two rules make
// that fleet-safe:
// 1. The per-row |max| (ex. 23) uses only comparisons — exact on any
// hardware, order-independent — and comes back to JS as ~1KB.
// 2. Scale DERIVATION (two divisions) stays in JS f64, which IEEE requires
// to be exactly rounded and is therefore identical on every device.
// WGSL division is only 2.5 ULP — a fork waiting to happen — but WGSL
// multiply/add are correctly rounded and floor/clamp are exact. So the
// quantize step is respecced from round(x / scale) to
// floor(f32(x * invScale) + 0.5) — floor(x+0.5) IS Math.round's tie
// rule — and the fround-stepped mirror below is bit-identical to the
// GPU kernel, which is exact-gated against it at init.
// NOTE this changes which int8 a value on a rounding boundary lands on
// (≤1 step) vs quantizeRows, so old and new builds cannot co-train — the
// divergence guard stops such mixed groups by design.
function rowAbsMax(X, rows, cols) {
const mx = new Float32Array(rows);
for (let r = 0; r < rows; r++) {
let m = 0;
for (let c = 0; c < cols; c++) { const a = Math.abs(X[r * cols + c]); if (a > m) m = a; }
mx[r] = m;
}
return mx;
}
function scalesFromAbsMax(mx) { // f64 divisions: exactly rounded, device-identical
const scale = new Float32Array(mx.length), inv = new Float32Array(mx.length);
for (let i = 0; i < mx.length; i++) {
scale[i] = Math.max(mx[i] / 127, 1e-8);
inv[i] = 1 / scale[i]; // recip of the STORED f32 scale
}
return { scale, inv };
}
function quantizeRowsInv(X, rows, cols, inv) { // bit-exact mirror of the GPU quantize kernel
const q = new Int8Array(rows * cols);
for (let r = 0; r < rows; r++) {
const iv = inv[r];
for (let c = 0; c < cols; c++) {
const n = Math.floor(f32(f32(X[r * cols + c] * iv) + 0.5));
q[r * cols + c] = n < -128 ? -128 : n > 127 ? 127 : n;
}
}
return q;
}
// the chained MLP: X @ W1 -> ReLU (fused) -> absmax -> quantize -> @ W2.
// d = { m, k, h, n }; gpuMlp (from webgpu.js) runs both GEMMs + the
// on-GPU quantize with one tiny absmax readback between; without it the CPU
// mirror chain runs — SAME math, so mixed GPU/CPU fleets stay bit-identical.
async function vmlpBlock(Xf, W1f, W2f, d, L, gpuMlp, audit) {
const x = quantizeRows(Xf, d.m, d.k);
const w1 = quantizeCols(W1f, d.k, d.h);
const w2 = quantizeCols(W2f, d.h, d.n);
if (gpuMlp) {
const r = await gpuMlp(x.q, w1.q, w2.q, x.s, w1.s, w2.s, d);
if (audit && audit.due()) {
// audit BOTH live GEMMs: gemm1 against the units directly; gemm2 by
// reconstructing its exact operand through the proven quantize mirror
const bad1 = auditTile(x.q, w1.q, x.s, w1.s, { m: d.m, k: d.k, n: d.h, relu: true }, r.h1, L, audit.cells);
if (bad1) { audit.fail("mlp gemm1: " + bad1); return r; }
const sc = scalesFromAbsMax(rowAbsMax(r.h1, d.m, d.h));
const hq = quantizeRowsInv(r.h1, d.m, d.h, sc.inv);
const bad2 = auditTile(hq, w2.q, sc.scale, w2.s, { m: d.m, k: d.h, n: d.n }, r.out, L, audit.cells);
if (bad2) audit.fail("mlp gemm2: " + bad2);
}
return r;
}
const h1 = bgemmJS(x.q, w1.q, x.s, w1.s, { m: d.m, k: d.k, n: d.h, batch: 1, relu: true }, L);
const sc = scalesFromAbsMax(rowAbsMax(h1, d.m, d.h));
const hq = quantizeRowsInv(h1, d.m, d.h, sc.inv);
const out = bgemmJS(hq, w2.q, sc.scale, w2.s, { m: d.m, k: d.h, n: d.n, batch: 1 }, L);
return { h1, out };
}
// ---- epilogue mirror -------------------------------------------------------
// BIT-EXACT mirror of the WGSL epilogue `f32(s) * a * b`. WGSL rounds to f32
// after the int->float conversion and after EACH multiply; plain JS would do
// the whole chain in f64 and round once, which differs in the last ulp. That
// last-ulp gap is what used to force a tolerance into the kernel gates —
// mirroring the rounding exactly is what lets the gates compare with `!==`.
const f32 = Math.fround;
function epi(s, a, b) { return f32(f32(f32(s) * a) * b); }
// f32 equality at the BIT level: `!==` says -0 === 0, but replicas are
// compared by hashing raw bytes, so an audit that can't see the sign of zero
// could pass a device that later forks the fleet. (Real ISAs have non-IEEE
// modes that flush -0 to +0 — e.g. RDNA2 output modifiers / legacy muls.)
const _fb = new Float32Array(1), _ub = new Uint32Array(_fb.buffer);
function bitDiff(a, b) { _fb[0] = a; const u = _ub[0]; _fb[0] = b; return u !== _ub[0]; }
// CPU mirror of the fused GPU kernel: batched int8 GEMM through the LUT with
// the epilogue (block dequant + optional ReLU) applied before returning —
// exactly what the WGSL kernel does on-device. d.acc=true returns the raw
// int32 accumulator instead (the exact oracle the fused kernel normally hides).
function bgemmJS(Xq, Wq, rs, cs, d, L) {
const { m, k, n } = d, batch = d.batch || 1, relu = !!d.relu, mul = L.mul;
const raw = !!d.acc;
const out = raw ? new Int32Array(batch * m * n) : new Float32Array(batch * m * n);
const acc = new Int32Array(n);
for (let bz = 0; bz < batch; bz++) {
const xo = bz * m * k, wo = bz * k * n, oo = bz * m * n, co = bz * n;
for (let i = 0; i < m; i++) {
acc.fill(0);
const xrow = xo + i * k;
for (let p = 0; p < k; p++) {
const au = (Xq[xrow + p] & 0xFF) * 256, wrow = wo + p * n;
for (let j = 0; j < n; j++) acc[j] += mul[au + (Wq[wrow + j] & 0xFF)];
}
const orow = oo + i * n;
if (raw) { for (let j = 0; j < n; j++) out[orow + j] = acc[j]; continue; }
const rscale = rs[bz * m + i];
for (let j = 0; j < n; j++) {
const v = epi(acc[j], rscale, cs[co + j]);
out[orow + j] = relu && v < 0 ? 0 : v;
}
}
}
return out;
}
// Recompute a handful of RANDOM output cells of a live GEMM through the LUT
// mirror and compare against what the kernel produced. Sampling cells instead
// of whole matrices makes this cheap enough to run continuously, at the real
// shapes training uses — not once at boot on toy inputs.
// STRATIFIED sampling. Uniformly random cells are the wrong instrument for
// the bugs that actually occur here: a bounds-guard off-by-one or a pack-tail
// padding bug lives on the LAST row/column, and uniform sampling finds that
// with probability ~1/n per cell — at the 16512-wide logits GEMM, never. So
// the first cells are the structurally dangerous ones (corners, last row,
// last column, last batch) chosen deterministically, and the remainder are
// random interior cells that catch diffuse bugs. Same principle as poisoning
// the buffer pool: construct the dangerous case, don't wait to land on it.
function auditTile(Xq, Wq, rs, cs, d, got, L, nCells) {
const { m, k, n } = d, batch = d.batch || 1, relu = !!d.relu, mul = L.mul;
const N = nCells || 8;
const edges = [[0, m - 1, n - 1], [0, 0, n - 1], [0, m - 1, 0], [0, 0, 0],
[batch - 1, m - 1, n - 1], [batch - 1, 0, 0]];
for (let t = 0; t < N; t++) {
let bz, i, j;
if (t < edges.length) { bz = edges[t][0]; i = edges[t][1]; j = edges[t][2]; }
else { bz = (Math.random() * batch) | 0; i = (Math.random() * m) | 0; j = (Math.random() * n) | 0; }
let acc = 0;
const xrow = bz * m * k + i * k, wo = bz * k * n;
for (let p = 0; p < k; p++) acc += mul[(Xq[xrow + p] & 0xFF) * 256 + (Wq[wo + p * n + j] & 0xFF)];
let v = epi(acc, rs[bz * m + i], cs[bz * n + j]);
if (relu && v < 0) v = 0;
const idx = (bz * m + i) * n + j;
if (bitDiff(got[idx], v))
return `GEMM audit failed at [b${bz},${i},${j}] shape ${m}x${k}x${n}: kernel ${Object.is(got[idx], -0) ? "-0" : got[idx]} vs units ${Object.is(v, -0) ? "-0" : v}`;
}
return null;
}
// ---- exact mirror of the split-K f32 GEMM ----------------------------------
// The f32 backward GEMM was the last kernel gated by a TOLERANCE (allclose at
// 1e-3) — and this project's own gate mutation test shows allclose waving
// through real bugs. The reason was real though: split-K accumulates in a
// different ORDER than a naive reference, so bit-equality against the naive
// one is impossible. The fix is the same as the epilogue mirror: reproduce
// the kernel's order exactly, then compare with `!==`.
// partials: for z in 0..S-1, sum p in [z*ks, min(k,(z+1)*ks)) in order
// reduce: sum the S partials in ascending z
// `fma` selects the rounding schedule for `s + a*b`: WGSL PERMITS a compiler
// to contract that into a fused multiply-add (one rounding) instead of two.
// Which one the device does is a fact about the device, so the gate tries
// both and reports which matches rather than assuming.
function fgemmMirror(A, Bm, d, fma) {
const { m, k, n } = d, transA = !!d.transA;
const S = k > 4096 ? Math.min(16, Math.ceil(k / 2048)) : 1;
const ks = Math.ceil(k / S);
const out = new Float32Array(m * n);
for (let row = 0; row < m; row++)
for (let col = 0; col < n; col++) {
let acc = 0; // reduce pass, ascending z
for (let z = 0; z < S; z++) {
const p0 = z * ks, p1 = Math.min(k, p0 + ks);
let s = 0; // one partial, in order
for (let p = p0; p < p1; p++) {
const a = transA ? A[p * m + row] : A[row * k + p];
s = fma ? f32(s + a * Bm[p * n + col]) // single rounding
: f32(s + f32(a * Bm[p * n + col]));
}
acc = f32(acc + s);
}
out[row * n + col] = acc;
}
return out;
}
// ---- live audits for the fused attention kernels ---------------------------
// The attention kernels had exact INIT gates but nothing at live shapes —
// the exact gap the GEMM audit exists to close, left open on the kernels with
// the trickiest indexing (head-strided gather, scatter write-back). These
// recompute individual output cells from the units, stratified like
// auditTile: last/first token pair, last head, last channel first, then
// random. Cost is hd (or T) multiply-adds per cell.
function auditAttScores(qq, kq, qs, ks, d, got, L, nCells) {
const { B, T, heads, hd } = d, C = heads * hd, mul = L.mul, raw = !!d.acc;
const N = nCells || 8;
const edges = [[B - 1, heads - 1, T - 1, T - 1], [0, 0, 0, 0],
[0, heads - 1, T - 1, 0], [B - 1, 0, 0, T - 1]];
for (let t = 0; t < N; t++) {
let bi, h, ti, tj;
if (t < edges.length) { bi = edges[t][0]; h = edges[t][1]; ti = edges[t][2]; tj = edges[t][3]; }
else { bi = (Math.random() * B) | 0; h = (Math.random() * heads) | 0;
ti = (Math.random() * T) | 0; tj = (Math.random() * T) | 0; }
const bz = bi * heads + h;
const qo = (bi * T + ti) * C + h * hd, ko = (bi * T + tj) * C + h * hd;
let acc = 0;
for (let p = 0; p < hd; p++) acc += mul[(qq[qo + p] & 0xFF) * 256 + (kq[ko + p] & 0xFF)];
const v = raw ? acc : epi(acc, qs[(bi * T + ti) * heads + h], ks[(bi * T + tj) * heads + h]);
const idx = (bz * T + ti) * T + tj;
if (raw ? got[idx] !== v : bitDiff(got[idx], v))
return `att.scores audit failed at [b${bi},h${h},${ti},${tj}] B${B}T${T}H${heads}d${hd}: kernel ${got[idx]} vs units ${v}`;
}
return null;
}
function auditAttCtx(aq, vq, as, vs, d, got, L, nCells) {
const { B, T, heads, hd } = d, C = heads * hd, mul = L.mul, raw = !!d.acc;
const N = nCells || 8;
const edges = [[B - 1, heads - 1, T - 1, hd - 1], [0, 0, 0, 0],
[0, heads - 1, T - 1, 0], [B - 1, 0, 0, hd - 1]];
for (let t = 0; t < N; t++) {
let bi, h, ti, j;
if (t < edges.length) { bi = edges[t][0]; h = edges[t][1]; ti = edges[t][2]; j = edges[t][3]; }
else { bi = (Math.random() * B) | 0; h = (Math.random() * heads) | 0;
ti = (Math.random() * T) | 0; j = (Math.random() * hd) | 0; }
const bz = bi * heads + h, ao = (bz * T + ti) * T;
let acc = 0;
for (let tj = 0; tj < T; tj++)
acc += mul[(aq[ao + tj] & 0xFF) * 256 + (vq[(bi * T + tj) * C + h * hd + j] & 0xFF)];
const v = raw ? acc : epi(acc, as[bz * T + ti], vs[(bi * heads + h) * hd + j]);
const idx = (bi * T + ti) * C + h * hd + j;
if (raw ? got[idx] !== v : bitDiff(got[idx], v))
return `att.ctx audit failed at [b${bi},h${h},${ti},${j}] B${B}T${T}H${heads}d${hd}: kernel ${got[idx]} vs units ${v}`;
}
return null;
}
// block-scaled verified GEMM, float in → float out.
// d = { m, k, n, batch=1, relu=false }; X is (batch·m)×k, W is batch×(k×n)
// gpuBgemm (from webgpu.js) runs the batched kernel with the fused epilogue;
// without it the CPU LUT mirror runs. Every product goes through the units.
async function vgemmBlock(Xf, Wf, d, L, gpuBgemm, audit) {
const { m, k, n } = d, batch = d.batch || 1;
const x = quantizeRows(Xf, batch * m, k);
let wq, ws;
if (batch === 1) {
const w = quantizeCols(Wf, k, n); wq = w.q; ws = w.s;
} else {
wq = new Int8Array(batch * k * n); ws = new Float32Array(batch * n);
for (let bz = 0; bz < batch; bz++) {
const w = quantizeCols(Wf.subarray(bz * k * n, (bz + 1) * k * n), k, n);
wq.set(w.q, bz * k * n); ws.set(w.s, bz * n);
}
}
if (gpuBgemm) {
const out = await gpuBgemm(x.q, wq, x.s, ws, d);
// continuous re-verification at LIVE shapes: the boot gate only ever saw
// toy inputs, so sample a few real cells against the units as we go
if (audit && audit.due()) {
const bad = auditTile(x.q, wq, x.s, ws, d, out, L, audit.cells);
if (bad) audit.fail(bad);
}
return out;
}
return bgemmJS(x.q, wq, x.s, ws, d, L);
}
// ---- gather-fused attention through the units (CUTLASS ex. 36/52) ----------
// The kernels read q/k/v/ctx directly in their natural BT×C layout with
// head-strided indexing — no JS gather copies, no kᵀ transpose, and the
// context write scatters straight back into BT×C. Quantization stays
// block-scaled: q/k/a per (token,head) row, v per (head,channel) column.
// The (BT·heads)×hd row view of q/k IS the contiguous buffer, so
// quantizeRows(q, BT·heads, hd) gives per-(token,head) scales for free.
function quantizeHeadCols(v, B, T, heads, hd) { // per (batch,head,channel) column
const C = heads * hd;
const q = new Int8Array(B * T * C), s = new Float32Array(B * heads * hd);
for (let bi = 0; bi < B; bi++)
for (let h = 0; h < heads; h++)
for (let j = 0; j < hd; j++) {
let mx = 0;
for (let ti = 0; ti < T; ti++) {
const a = Math.abs(v[(bi * T + ti) * C + h * hd + j]);
if (a > mx) mx = a;
}
const sc = Math.max(mx / 127, 1e-8);
s[(bi * heads + h) * hd + j] = sc;
for (let ti = 0; ti < T; ti++) {
const idx = (bi * T + ti) * C + h * hd + j;
const w = Math.round(v[idx] / sc);
q[idx] = w < -128 ? -128 : w > 127 ? 127 : w;
}
}
return { q, s };
}
// scores S[bz,ti,tj] = q_row(bi,ti,h) · k_row(bi,tj,h), every product via the LUT
// d.acc=true returns the raw int32 accumulator (exact oracle for the kernel gate)
function attScoresJS(qq, kq, qs, ks, d, L) {
const { B, T, heads, hd } = d, C = heads * hd, mul = L.mul, raw = !!d.acc;
const out = raw ? new Int32Array(B * heads * T * T) : new Float32Array(B * heads * T * T);
for (let bi = 0; bi < B; bi++) for (let h = 0; h < heads; h++) {
const bz = bi * heads + h;
for (let ti = 0; ti < T; ti++) {
const qo = (bi * T + ti) * C + h * hd, rscale = qs[(bi * T + ti) * heads + h];
for (let tj = 0; tj < T; tj++) {
const ko = (bi * T + tj) * C + h * hd;
let acc = 0;
for (let p = 0; p < hd; p++) acc += mul[(qq[qo + p] & 0xFF) * 256 + (kq[ko + p] & 0xFF)];
out[(bz * T + ti) * T + tj] = raw ? acc : epi(acc, rscale, ks[(bi * T + tj) * heads + h]);
}
}
}
return out;
}
// ctx[(bi,ti),(h,j)] = Σ_tj a[bz,ti,tj]·v[(bi,tj),(h,j)] — scatter fused into BT×C
function attCtxJS(aq, vq, as, vs, d, L) {
const { B, T, heads, hd } = d, C = heads * hd, mul = L.mul, raw = !!d.acc;
const out = raw ? new Int32Array(B * T * C) : new Float32Array(B * T * C);
for (let bi = 0; bi < B; bi++) for (let h = 0; h < heads; h++) {
const bz = bi * heads + h;
for (let ti = 0; ti < T; ti++) {
const ao = (bz * T + ti) * T, rscale = as[bz * T + ti];
for (let j = 0; j < hd; j++) {
let acc = 0;
for (let tj = 0; tj < T; tj++)
acc += mul[(aq[ao + tj] & 0xFF) * 256 + (vq[(bi * T + tj) * C + h * hd + j] & 0xFF)];
out[(bi * T + ti) * C + h * hd + j] = raw ? acc : epi(acc, rscale, vs[(bi * heads + h) * hd + j]);
}
}
}
return out;
}
// one verified layer forward; returns float out (+ cache for STE backward).
// Every product goes through the verified INT8 multiply (mul8 LUT) with exact
// int32 accumulation — i.e. an emulated INT8 tensor-core GEMM — then dequant.
async function linearFwd(X, W, m, k, n, L, useRelu, matmulInt8) {
const xq = quantize(X), wq = quantize(W);
const acc = await (matmulInt8 || lutMatmulJS)(xq.q, wq.q, m, k, n, L); // verified multiply
const dq = xq.scale * wq.scale;
const out = new Float32Array(m * n);
const mask = useRelu ? new Uint8Array(m * n) : null;
for (let i = 0; i < m * n; i++) {
let v = acc[i] * dq;
if (useRelu) { if (v > 0) mask[i] = 1; else v = 0; }
out[i] = v;
}
return { out, mask };
}
// 2-layer MLP: X→H (relu) →dout. Forward through verified units, MSE loss.
async function forward(X, y, W1, W2, D, L, matmulInt8) {
const { n, din, h, dout } = D;
const l1 = await linearFwd(X, W1, n, din, h, L, true, matmulInt8);
const l2 = await linearFwd(l1.out, W2, n, h, dout, L, false, matmulInt8);
const resid = new Float32Array(n * dout); let loss = 0;
for (let i = 0; i < resid.length; i++) { const r = l2.out[i] - y[i]; resid[i] = r; loss += r * r; }
loss /= resid.length;
return { loss, resid, z1: l1.out, mask1: l1.mask };
}
// STE backward (verified matmul treated as float X@W). Returns flat [gW1, gW2].
function backward(X, W1, W2, fwd, D) {
const { n, din, h, dout } = D;
const { resid, z1, mask1 } = fwd;
const s = 2 / n;
const dout_ = new Float32Array(resid.length);
for (let i = 0; i < resid.length; i++) dout_[i] = resid[i] * s;
const mm = TC.matmul, tr = TC.transpose;
const gW2 = mm(tr(z1, n, h), dout_, h, n, dout); // z1ᵀ @ dout
const dz1 = mm(dout_, tr(W2, h, dout), n, dout, h); // dout @ W2ᵀ
for (let i = 0; i < dz1.length; i++) if (!mask1[i]) dz1[i] = 0; // relu grad
const gW1 = mm(tr(X, n, din), dz1, din, n, h); // Xᵀ @ dz1
const g = new Float32Array(gW1.length + gW2.length);
g.set(gW1, 0); g.set(gW2, gW1.length);
return g;
}
function splitApply(W1, W2, gAvg, lr) {
for (let i = 0; i < W1.length; i++) W1[i] -= lr * gAvg[i];
for (let j = 0; j < W2.length; j++) W2[j] -= lr * gAvg[W1.length + j];
}
const api = { quantize, quantize2, quantizeRows, quantizeCols, quantizeHeadCols, lutMatmulJS, lutMatmul3JS, lutMatmul3,
bgemmJS, vgemmBlock, auditTile, epi, attScoresJS, attCtxJS, linearFwd, forward, backward, splitApply,
rowAbsMax, scalesFromAbsMax, quantizeRowsInv, vmlpBlock, bitDiff,
auditAttScores, auditAttCtx, fgemmMirror };
if (typeof module !== "undefined" && module.exports) { TC = require("./traincore.js"); module.exports = api; }
else { TC = root.TrainCore; root.Verified = api; }
})(typeof self !== "undefined" ? self : this);