web: restore LF line endings on verified_core.js / transformer.js
Browse filesThe previous upload carried CRLF line endings on these two files (a local git checkout artifact). Content is otherwise unchanged; this restores the repo's LF convention so the files diff cleanly.
- web/public/transformer.js +583 -583
- web/public/verified_core.js +519 -519
web/public/transformer.js
CHANGED
|
@@ -1,583 +1,583 @@
|
|
| 1 |
-
// A miniature transformer language model that trains THROUGH the verified INT8
|
| 2 |
-
// units: every matrix product in the forward pass — QKV projections, attention
|
| 3 |
-
// scores, attention·values, output projection, the MLP, and the unembedding —
|
| 4 |
-
// runs through the verified multiply LUT (an emulated INT8 tensor core).
|
| 5 |
-
// Backward is a straight-through estimator in float (the integer path has no
|
| 6 |
-
// gradient), exactly like the Python VerifiedLinear.
|
| 7 |
-
//
|
| 8 |
-
// Task: next-character prediction on a deterministic, self-generated corpus
|
| 9 |
-
// (every peer builds the same text from the same seed — nothing to download).
|
| 10 |
-
(function (root) {
|
| 11 |
-
"use strict";
|
| 12 |
-
|
| 13 |
-
let TC, V; // TrainCore / Verified — resolved per environment at the end
|
| 14 |
-
|
| 15 |
-
function mulberry32(a) { return function () { a |= 0; a = a + 0x6D2B79F5 | 0; let t = Math.imul(a ^ a >>> 15, 1 | a); t = t + Math.imul(t ^ t >>> 7, 61 | t) ^ t; return ((t ^ t >>> 14) >>> 0) / 4294967296; }; }
|
| 16 |
-
function randn(n, rng) { const o = new Float32Array(n); for (let i = 0; i < n; i += 2) { let u = 0, v = 0; while (u === 0) u = rng(); while (v === 0) v = rng(); const m = Math.sqrt(-2 * Math.log(u)); o[i] = m * Math.cos(2 * Math.PI * v); if (i + 1 < n) o[i + 1] = m * Math.sin(2 * Math.PI * v); } return o; }
|
| 17 |
-
|
| 18 |
-
// ---- corpus: deterministic cottagecore prose, identical on every peer -----
|
| 19 |
-
const W_ADJ = ["mossy", "golden", "amber", "quiet", "little", "misty", "sunny", "wild", "cozy", "dusty", "merry", "brave"];
|
| 20 |
-
const W_NOUN = ["fox", "hare", "owl", "badger", "toad", "sparrow", "otter", "deer", "mushroom", "acorn", "willow", "robin", "river", "meadow", "garden", "lantern"];
|
| 21 |
-
const W_VERB = ["naps", "sings", "wanders", "hides", "dreams", "waits", "dances", "listens", "rests", "grows"];
|
| 22 |
-
const W_PREP = ["by", "under", "near", "beside", "beyond", "inside"];
|
| 23 |
-
function buildCorpus() {
|
| 24 |
-
const rng = mulberry32(20260712);
|
| 25 |
-
const pick = (a) => a[Math.floor(rng() * a.length)];
|
| 26 |
-
let s = "";
|
| 27 |
-
while (s.length < 60000)
|
| 28 |
-
s += `the ${pick(W_ADJ)} ${pick(W_NOUN)} ${pick(W_VERB)} ${pick(W_PREP)} the ${pick(W_ADJ)} ${pick(W_NOUN)}. `;
|
| 29 |
-
return s;
|
| 30 |
-
}
|
| 31 |
-
// ---- tokenizer -------------------------------------------------------------
|
| 32 |
-
// Spikewhale tokenizer (tokenizer.json): byte-level greedy longest-match
|
| 33 |
-
// ("length-max"), ~16.5k tokens. Until it loads (or if the file is missing)
|
| 34 |
-
// a 96-char byte-level vocab keeps the app working — but ALL devices in a
|
| 35 |
-
// group must use the same tokenizer (the config broadcast enforces it).
|
| 36 |
-
const FALLBACK_CHARS = [...Array(95)].map((_, i) => String.fromCharCode(32 + i)).concat(["\n"]);
|
| 37 |
-
let tok = {
|
| 38 |
-
name: "char-96 (fallback)",
|
| 39 |
-
vocab: Object.fromEntries(FALLBACK_CHARS.map((c, i) => [c, i])),
|
| 40 |
-
ids: FALLBACK_CHARS, maxLen: 1, size: FALLBACK_CHARS.length,
|
| 41 |
-
unk: 0, specials: new Set(),
|
| 42 |
-
};
|
| 43 |
-
tok.unk = tok.vocab[" "];
|
| 44 |
-
function vocabSize() { return tok.size; }
|
| 45 |
-
function tokenizerName() { return tok.name; }
|
| 46 |
-
function loadTokenizerData(d) { // plain {vocab, vocab_size, max_token_len}
|
| 47 |
-
const ids = new Array(d.vocab_size);
|
| 48 |
-
for (const [t, i] of Object.entries(d.vocab)) ids[i] = t;
|
| 49 |
-
tok = { name: `Spikewhale length-max (${d.vocab_size} tokens)`,
|
| 50 |
-
vocab: d.vocab, ids, maxLen: d.max_token_len || 24, size: d.vocab_size,
|
| 51 |
-
unk: d.vocab["<unk>"] ?? 1,
|
| 52 |
-
specials: new Set(["<pad>", "<unk>", "<bos>", "<eos>", ...(d.special_tokens || [])]) };
|
| 53 |
-
IDS = encode(CORPUS); // re-tokenize whatever corpus is loaded
|
| 54 |
-
return tok.name;
|
| 55 |
-
}
|
| 56 |
-
async function loadTokenizer(url) {
|
| 57 |
-
const r = await fetch(url || "tokenizer.json");
|
| 58 |
-
if (!r.ok) throw new Error(`tokenizer.json HTTP ${r.status}`);
|
| 59 |
-
return loadTokenizerData(await r.json());
|
| 60 |
-
}
|
| 61 |
-
function toLatin1(s) { const b = new TextEncoder().encode(s); let o = ""; for (const x of b) o += String.fromCharCode(x); return o; }
|
| 62 |
-
function encode(text) { // greedy longest match over bytes
|
| 63 |
-
const s = toLatin1(text), out = [];
|
| 64 |
-
let i = 0;
|
| 65 |
-
while (i < s.length) {
|
| 66 |
-
let m = null;
|
| 67 |
-
for (let L = Math.min(tok.maxLen, s.length - i); L > 0; L--) {
|
| 68 |
-
const sub = s.substr(i, L);
|
| 69 |
-
if (sub in tok.vocab) { m = sub; break; }
|
| 70 |
-
}
|
| 71 |
-
if (m === null) { out.push(tok.unk); i++; continue; }
|
| 72 |
-
out.push(tok.vocab[m]); i += m.length;
|
| 73 |
-
}
|
| 74 |
-
return Int32Array.from(out);
|
| 75 |
-
}
|
| 76 |
-
function decode(idArr) {
|
| 77 |
-
let s = "";
|
| 78 |
-
for (const id of idArr) {
|
| 79 |
-
const t = tok.ids[id];
|
| 80 |
-
if (t === undefined || tok.specials.has(t)) continue;
|
| 81 |
-
s += t;
|
| 82 |
-
}
|
| 83 |
-
const bytes = Uint8Array.from([...s].map(c => c.charCodeAt(0)));
|
| 84 |
-
return new TextDecoder().decode(bytes);
|
| 85 |
-
}
|
| 86 |
-
let CORPUS = buildCorpus();
|
| 87 |
-
let IDS = encode(CORPUS);
|
| 88 |
-
let DATASET = "built-in corpus";
|
| 89 |
-
|
| 90 |
-
// Training text: FineWeb-Edu (10BT sample), HARDCODED as the only dataset.
|
| 91 |
-
// The serving Space reads random slices of the parquet shards straight off
|
| 92 |
-
// the HF CDN with range requests (see server.js /data) — no dependency on
|
| 93 |
-
// the datasets-server rows API, which 503s routinely. Each device pulls its
|
| 94 |
-
// own random slice (that's data parallelism — batches were always
|
| 95 |
-
// per-device anyway). Offline or on failure the built-in corpus stays.
|
| 96 |
-
const DEFAULT_DS = "HuggingFaceFW/fineweb-edu";
|
| 97 |
-
async function streamDataset() { // dataset choice removed on purpose
|
| 98 |
-
const r = await fetch("data");
|
| 99 |
-
if (!r.ok) throw new Error(`/data HTTP ${r.status}`);
|
| 100 |
-
const text = (await r.text()).replace(/[^\x20-\x7e\n]/g, " ");
|
| 101 |
-
if (text.length < 10000) throw new Error("too little text returned");
|
| 102 |
-
CORPUS = text.slice(0, 500000);
|
| 103 |
-
IDS = encode(CORPUS);
|
| 104 |
-
DATASET = `${DEFAULT_DS} · 10BT sample (parquet via this Space)`;
|
| 105 |
-
return { name: DATASET, chars: CORPUS.length };
|
| 106 |
-
}
|
| 107 |
-
const streamFineWebEdu = () => streamDataset();
|
| 108 |
-
function datasetName() { return DATASET; }
|
| 109 |
-
|
| 110 |
-
// ---- verified matmul: block-scaled INT8 through the units ------------------
|
| 111 |
-
// CUTLASS ex. 67/81 blockwise scaling: per-row activation scales × per-column
|
| 112 |
-
// weight scales, one exact LUT/DP4A GEMM, dequant (+ optional fused ReLU) in
|
| 113 |
-
// the kernel epilogue (ex. 12). Replaces the per-tensor 3-pass: same outlier
|
| 114 |
-
// robustness at one third of the unit ops.
|
| 115 |
-
async function vmm(Xf, Wf, m, k, n, ctx, relu) {
|
| 116 |
-
// ctx.audit re-checks random cells of this LIVE GEMM against the units
|
| 117 |
-
return V.vgemmBlock(Xf, Wf, { m, k, n, batch: 1, relu: !!relu }, ctx.L, ctx.bgemm, ctx.audit);
|
| 118 |
-
}
|
| 119 |
-
// CUTLASS ex. 45 (dual GEMM): sibling GEMMs that share the same LEFT operand
|
| 120 |
-
// run as ONE batched dispatch, and the shared operand is quantized ONCE
|
| 121 |
-
// instead of once per sibling. Used for the q/k/v projections — same X
|
| 122 |
-
// (ln1.y), three weights, identical shapes. Bit-identical to three separate
|
| 123 |
-
// vmm calls: quantizeRows is deterministic (same input -> same int8+scales),
|
| 124 |
-
// the tiled copies index exactly like separate batch elements, and block
|
| 125 |
-
// scales are per-row/per-column PER BATCH ELEMENT, so concatenation changes
|
| 126 |
-
// no scale and no product. The batched kernel is the same exact-gated bgemm
|
| 127 |
-
// that training already runs, and the live-shape audit still samples it.
|
| 128 |
-
async function vmmShared3(Xf, Wa, Wb, Wc, m, k, n, ctx) {
|
| 129 |
-
const x = V.quantizeRows(Xf, m, k);
|
| 130 |
-
const xq = new Int8Array(3 * m * k), xs = new Float32Array(3 * m);
|
| 131 |
-
for (let i = 0; i < 3; i++) { xq.set(x.q, i * m * k); xs.set(x.s, i * m); }
|
| 132 |
-
const wq = new Int8Array(3 * k * n), ws = new Float32Array(3 * n);
|
| 133 |
-
[Wa, Wb, Wc].forEach((W, i) => { const w = V.quantizeCols(W, k, n); wq.set(w.q, i * k * n); ws.set(w.s, i * n); });
|
| 134 |
-
const d = { m, k, n, batch: 3 };
|
| 135 |
-
let out;
|
| 136 |
-
if (ctx.bgemm) {
|
| 137 |
-
out = await ctx.bgemm(xq, wq, xs, ws, d);
|
| 138 |
-
if (ctx.audit && ctx.audit.due()) {
|
| 139 |
-
const bad = V.auditTile(xq, wq, xs, ws, d, out, ctx.L, ctx.audit.cells);
|
| 140 |
-
if (bad) ctx.audit.fail(bad);
|
| 141 |
-
}
|
| 142 |
-
} else {
|
| 143 |
-
out = V.bgemmJS(xq, wq, xs, ws, d, ctx.L);
|
| 144 |
-
}
|
| 145 |
-
const MN = m * n;
|
| 146 |
-
return [out.subarray(0, MN), out.subarray(MN, 2 * MN), out.subarray(2 * MN, 3 * MN)];
|
| 147 |
-
}
|
| 148 |
-
|
| 149 |
-
// ---- layernorm (no affine) -------------------------------------------------
|
| 150 |
-
function lnFwd(x, rows, C) {
|
| 151 |
-
const y = new Float32Array(rows * C), sig = new Float32Array(rows);
|
| 152 |
-
for (let r = 0; r < rows; r++) {
|
| 153 |
-
let mu = 0; for (let j = 0; j < C; j++) mu += x[r * C + j]; mu /= C;
|
| 154 |
-
let v = 0; for (let j = 0; j < C; j++) { const d = x[r * C + j] - mu; v += d * d; }
|
| 155 |
-
const s = Math.sqrt(v / C + 1e-5); sig[r] = s;
|
| 156 |
-
for (let j = 0; j < C; j++) y[r * C + j] = (x[r * C + j] - mu) / s;
|
| 157 |
-
}
|
| 158 |
-
return { y, sig };
|
| 159 |
-
}
|
| 160 |
-
function lnBwd(dy, y, sig, rows, C) {
|
| 161 |
-
const dx = new Float32Array(rows * C);
|
| 162 |
-
for (let r = 0; r < rows; r++) {
|
| 163 |
-
let mdy = 0, mdyy = 0;
|
| 164 |
-
for (let j = 0; j < C; j++) { mdy += dy[r * C + j]; mdyy += dy[r * C + j] * y[r * C + j]; }
|
| 165 |
-
mdy /= C; mdyy /= C;
|
| 166 |
-
for (let j = 0; j < C; j++) dx[r * C + j] = (dy[r * C + j] - mdy - y[r * C + j] * mdyy) / sig[r];
|
| 167 |
-
}
|
| 168 |
-
return dx;
|
| 169 |
-
}
|
| 170 |
-
|
| 171 |
-
// ---- model -----------------------------------------------------------------
|
| 172 |
-
// cfg: { c: width, t: seq len, b: batch/device, layers, heads, steps, lr }
|
| 173 |
-
// engine: the Compute backend object ({bgemm} for the fused WebGPU path) or a
|
| 174 |
-
// legacy matmulInt8 function (Node tests, inference kit) -> CPU LUT mirror
|
| 175 |
-
function init(cfg, L, engine, audit) {
|
| 176 |
-
const c = cfg.c, layers = cfg.layers || 2, heads = cfg.heads || 2, hidden = 2 * c;
|
| 177 |
-
let seed = 100;
|
| 178 |
-
const mk = (nEl, scale) => { const w = randn(nEl, mulberry32(seed++)); for (let i = 0; i < nEl; i++) w[i] *= scale; return w; };
|
| 179 |
-
const params = [], names = [];
|
| 180 |
-
const add = (name, w) => { params.push(w); names.push(name); return w; };
|
| 181 |
-
const m = {
|
| 182 |
-
cfg: { ...cfg, layers, heads, hidden, vocab: vocabSize() },
|
| 183 |
-
ctx: { L, bgemm: (engine && engine.bgemm) || null,
|
| 184 |
-
att: (engine && engine.att) || null, fgemm: (engine && engine.fgemm) || null,
|
| 185 |
-
fgemm2: (engine && engine.fgemm2) || null,
|
| 186 |
-
mlp: (engine && engine.mlp) || null,
|
| 187 |
-
audit: audit || null, unitBackward: !!cfg.unitBackward },
|
| 188 |
-
emb: add("emb", mk(vocabSize() * c, 0.08)),
|
| 189 |
-
pos: add("pos", mk(cfg.t * c, 0.02)),
|
| 190 |
-
blocks: [], params, names,
|
| 191 |
-
};
|
| 192 |
-
for (let l = 0; l < layers; l++)
|
| 193 |
-
m.blocks.push({
|
| 194 |
-
Wq: add(`b${l}.Wq`, mk(c * c, 0.08)), Wk: add(`b${l}.Wk`, mk(c * c, 0.08)),
|
| 195 |
-
Wv: add(`b${l}.Wv`, mk(c * c, 0.08)), Wo: add(`b${l}.Wo`, mk(c * c, 0.08)),
|
| 196 |
-
W1: add(`b${l}.W1`, mk(c * hidden, 0.08)), W2: add(`b${l}.W2`, mk(hidden * c, 0.08)),
|
| 197 |
-
});
|
| 198 |
-
// weight-tied unembedding: logits use embᵀ (no separate Wu). Halves the
|
| 199 |
-
// vocab-sized parameters — and with a 16k vocab that's ~half of ALL
|
| 200 |
-
// parameters, so gradients over the wire shrink ~2× too.
|
| 201 |
-
m.nParams = params.reduce((a, p) => a + p.length, 0);
|
| 202 |
-
return m;
|
| 203 |
-
}
|
| 204 |
-
|
| 205 |
-
function sampleBatch(cfg) {
|
| 206 |
-
const { b, t } = cfg;
|
| 207 |
-
const X = new Int32Array(b * t), Y = new Int32Array(b * t);
|
| 208 |
-
for (let i = 0; i < b; i++) {
|
| 209 |
-
const off = Math.floor(Math.random() * (IDS.length - t - 1));
|
| 210 |
-
for (let j = 0; j < t; j++) { X[i * t + j] = IDS[off + j]; Y[i * t + j] = IDS[off + j + 1]; }
|
| 211 |
-
}
|
| 212 |
-
return { X, Y };
|
| 213 |
-
}
|
| 214 |
-
|
| 215 |
-
// ---- head layout helpers ---------------------------------------------------
|
| 216 |
-
// q/k/v live as BT×C with head h owning columns [h*hd, (h+1)*hd). The backward
|
| 217 |
-
// wants every head as its own GEMM problem, so gather once into BH×T×hd and
|
| 218 |
-
// scatter back at the end — one pass each, instead of slicing per head inside
|
| 219 |
-
// the loop and paying a GPU dispatch per tiny matmul.
|
| 220 |
-
function gatherHeads(x, B, T, C, heads, hd) { // BT×C -> BH×T×hd
|
| 221 |
-
const out = new Float32Array(B * heads * T * hd);
|
| 222 |
-
for (let bi = 0; bi < B; bi++)
|
| 223 |
-
for (let h = 0; h < heads; h++) {
|
| 224 |
-
const bz = bi * heads + h;
|
| 225 |
-
for (let ti = 0; ti < T; ti++)
|
| 226 |
-
for (let j = 0; j < hd; j++) out[(bz * T + ti) * hd + j] = x[(bi * T + ti) * C + h * hd + j];
|
| 227 |
-
}
|
| 228 |
-
return out;
|
| 229 |
-
}
|
| 230 |
-
function scatterHeadsAcc(dst, src, B, T, C, heads, hd) { // BH×T×hd -> += BT×C
|
| 231 |
-
for (let bi = 0; bi < B; bi++)
|
| 232 |
-
for (let h = 0; h < heads; h++) {
|
| 233 |
-
const bz = bi * heads + h;
|
| 234 |
-
for (let ti = 0; ti < T; ti++)
|
| 235 |
-
for (let j = 0; j < hd; j++) dst[(bi * T + ti) * C + h * hd + j] += src[(bz * T + ti) * hd + j];
|
| 236 |
-
}
|
| 237 |
-
}
|
| 238 |
-
function batchedTranspose(x, batch, rows, cols) { // per-batch rows×cols -> cols×rows
|
| 239 |
-
const out = new Float32Array(batch * rows * cols);
|
| 240 |
-
for (let b = 0; b < batch; b++) {
|
| 241 |
-
const o = b * rows * cols;
|
| 242 |
-
for (let r = 0; r < rows; r++)
|
| 243 |
-
for (let c = 0; c < cols; c++) out[o + c * rows + r] = x[o + r * cols + c];
|
| 244 |
-
}
|
| 245 |
-
return out;
|
| 246 |
-
}
|
| 247 |
-
|
| 248 |
-
// ---- forward THROUGH the verified units (caches kept for STE backward) -----
|
| 249 |
-
async function forward(m, X, Y) {
|
| 250 |
-
const { c: C, t: T, b: B, layers, heads, hidden, vocab } = m.cfg;
|
| 251 |
-
const BT = B * T, hd = C / heads, ctx = m.ctx;
|
| 252 |
-
const cache = { X, Y, blocks: [] };
|
| 253 |
-
let x = new Float32Array(BT * C);
|
| 254 |
-
for (let i = 0; i < BT; i++) {
|
| 255 |
-
const id = X[i], tpos = i % T;
|
| 256 |
-
for (let j = 0; j < C; j++) x[i * C + j] = m.emb[id * C + j] + m.pos[tpos * C + j];
|
| 257 |
-
}
|
| 258 |
-
for (let l = 0; l < layers; l++) {
|
| 259 |
-
const bl = m.blocks[l], cb = { xin: x };
|
| 260 |
-
const l1 = lnFwd(x, BT, C); cb.ln1 = l1;
|
| 261 |
-
// q/k/v share the same left operand — one batched dispatch, one quantize
|
| 262 |
-
// of ln1.y instead of three (CUTLASS ex. 45; see vmmShared3)
|
| 263 |
-
const [q, k, v] = await vmmShared3(l1.y, bl.Wq, bl.Wk, bl.Wv, BT, C, C, ctx);
|
| 264 |
-
cb.q = q; cb.k = k; cb.v = v;
|
| 265 |
-
const scale = 1 / Math.sqrt(hd);
|
| 266 |
-
// gather-FUSED attention (CUTLASS ex. 36/52): the kernels read q/k/v in
|
| 267 |
-
// their natural BT×C layout with head-strided indexing and scatter ctx
|
| 268 |
-
// straight back — no JS gather copies, no kᵀ transpose. All B×H heads in
|
| 269 |
-
// one dispatch per stage, every product through the verified units.
|
| 270 |
-
const BH = B * heads, dAtt = { B, T, heads, hd };
|
| 271 |
-
// per-(token,head) row quantization: the (BT·heads)×hd view IS the buffer
|
| 272 |
-
const qq = V.quantizeRows(q, BT * heads, hd), kq = V.quantizeRows(k, BT * heads, hd);
|
| 273 |
-
const sAll = ctx.att ? await ctx.att.scores(qq.q, kq.q, qq.s, kq.s, dAtt)
|
| 274 |
-
: V.attScoresJS(qq.q, kq.q, qq.s, kq.s, dAtt, ctx.L);
|
| 275 |
-
// live-shape audit: the init gate only ever saw four test shapes
|
| 276 |
-
if (ctx.att && ctx.audit && ctx.audit.due()) {
|
| 277 |
-
const bad = V.auditAttScores(qq.q, kq.q, qq.s, kq.s, dAtt, sAll, ctx.L, ctx.audit.cells);
|
| 278 |
-
if (bad) ctx.audit.fail(bad);
|
| 279 |
-
}
|
| 280 |
-
const aAll = new Float32Array(BH * T * T); // causal softmax
|
| 281 |
-
for (let bz = 0; bz < BH; bz++) {
|
| 282 |
-
const so = bz * T * T;
|
| 283 |
-
for (let ti = 0; ti < T; ti++) {
|
| 284 |
-
let mx = -1e30;
|
| 285 |
-
for (let tj = 0; tj <= ti; tj++) mx = Math.max(mx, sAll[so + ti * T + tj] * scale);
|
| 286 |
-
let z = 0;
|
| 287 |
-
for (let tj = 0; tj <= ti; tj++) { const e = Math.exp(sAll[so + ti * T + tj] * scale - mx); aAll[so + ti * T + tj] = e; z += e; }
|
| 288 |
-
for (let tj = 0; tj <= ti; tj++) aAll[so + ti * T + tj] /= z;
|
| 289 |
-
}
|
| 290 |
-
}
|
| 291 |
-
const aq = V.quantizeRows(aAll, BH * T, T);
|
| 292 |
-
const vq = V.quantizeHeadCols(v, B, T, heads, hd);
|
| 293 |
-
const ctxOut = ctx.att ? await ctx.att.ctx(aq.q, vq.q, aq.s, vq.s, dAtt)
|
| 294 |
-
: V.attCtxJS(aq.q, vq.q, aq.s, vq.s, dAtt, ctx.L);
|
| 295 |
-
if (ctx.att && ctx.audit && ctx.audit.due()) {
|
| 296 |
-
const bad = V.auditAttCtx(aq.q, vq.q, aq.s, vq.s, dAtt, ctxOut, ctx.L, ctx.audit.cells);
|
| 297 |
-
if (bad) ctx.audit.fail(bad);
|
| 298 |
-
}
|
| 299 |
-
cb.aAll = aAll; // backward slices heads from q/k/v/aAll
|
| 300 |
-
cb.ctxOut = ctxOut;
|
| 301 |
-
const attnOut = await vmm(ctxOut, bl.Wo, BT, C, C, ctx);
|
| 302 |
-
const x2 = new Float32Array(BT * C);
|
| 303 |
-
for (let i = 0; i < x2.length; i++) x2[i] = x[i] + attnOut[i];
|
| 304 |
-
cb.x2 = x2;
|
| 305 |
-
const l2 = lnFwd(x2, BT, C); cb.ln2 = l2;
|
| 306 |
-
// CUTLASS ex. 13 + 23: both MLP GEMMs run back-to-back on the GPU. The
|
| 307 |
-
// intermediate h1 is quantized ON-DEVICE (exact-gated respec — see
|
| 308 |
-
// vmlpBlock in verified_core.js) and only its per-row absmax (~1KB)
|
| 309 |
-
// visits JS between the GEMMs; h1 itself comes back solely because the
|
| 310 |
-
// STE backward needs it. CPU devices run the bit-identical mirror chain.
|
| 311 |
-
const { h1, out: mlpOut } = await V.vmlpBlock(l2.y, bl.W1, bl.W2,
|
| 312 |
-
{ m: BT, k: C, h: hidden, n: C }, ctx.L, ctx.mlp, ctx.audit);
|
| 313 |
-
const mask = new Uint8Array(h1.length);
|
| 314 |
-
for (let i = 0; i < h1.length; i++) if (h1[i] > 0) mask[i] = 1;
|
| 315 |
-
cb.h1 = h1; cb.mask = mask;
|
| 316 |
-
x = new Float32Array(BT * C);
|
| 317 |
-
for (let i = 0; i < x.length; i++) x[i] = x2[i] + mlpOut[i];
|
| 318 |
-
cache.blocks.push(cb);
|
| 319 |
-
}
|
| 320 |
-
const lf = lnFwd(x, BT, C); cache.lnf = lf; cache.xf = x;
|
| 321 |
-
const logits = await vmm(lf.y, TC.transpose(m.emb, vocab, C), BT, C, vocab, ctx); // tied: embᵀ
|
| 322 |
-
// cross-entropy + dlogits
|
| 323 |
-
let loss = 0;
|
| 324 |
-
const dlogits = new Float32Array(BT * vocab);
|
| 325 |
-
for (let i = 0; i < BT; i++) {
|
| 326 |
-
let mx = -1e30;
|
| 327 |
-
for (let j = 0; j < vocab; j++) mx = Math.max(mx, logits[i * vocab + j]);
|
| 328 |
-
let z = 0;
|
| 329 |
-
for (let j = 0; j < vocab; j++) z += Math.exp(logits[i * vocab + j] - mx);
|
| 330 |
-
const lz = Math.log(z) + mx;
|
| 331 |
-
loss += lz - logits[i * vocab + Y[i]];
|
| 332 |
-
for (let j = 0; j < vocab; j++)
|
| 333 |
-
dlogits[i * vocab + j] = (Math.exp(logits[i * vocab + j] - lz) - (j === Y[i] ? 1 : 0)) / BT;
|
| 334 |
-
}
|
| 335 |
-
loss /= BT;
|
| 336 |
-
cache.dlogits = dlogits;
|
| 337 |
-
return { loss, cache, logits };
|
| 338 |
-
}
|
| 339 |
-
|
| 340 |
-
// ---- STE backward (float), mirrors forward exactly --------------------------
|
| 341 |
-
// The two vocab-sized matmuls run on the split-K f32 GPU kernel when
|
| 342 |
-
// available (CUTLASS ex. 06) — same float math, off the JS thread.
|
| 343 |
-
async function backward(m, cache) {
|
| 344 |
-
const { c: C, t: T, b: B, layers, heads, hidden, vocab } = m.cfg;
|
| 345 |
-
const BT = B * T, hd = C / heads, tr = TC.transpose;
|
| 346 |
-
const g = m.params.map(p => new Float32Array(p.length));
|
| 347 |
-
const gi = Object.fromEntries(m.names.map((n, i) => [n, i]));
|
| 348 |
-
// Every matmul here goes through `bmm`. With ctx.unitBackward the STE
|
| 349 |
-
// gradient is computed BY the verified units (block-scaled int8, exact int32
|
| 350 |
-
// accumulate) instead of in float. STE is a claim about the math — pretend
|
| 351 |
-
// the quantizer was the identity — not about the datatype that evaluates it,
|
| 352 |
-
// so the two are orthogonal and this stays a correct STE.
|
| 353 |
-
const units = !!m.ctx.unitBackward;
|
| 354 |
-
const bmm = units
|
| 355 |
-
? (A, Bm, mm_, k, n) => vmm(A, Bm, mm_, k, n, m.ctx)
|
| 356 |
-
: async (A, Bm, mm_, k, n) => TC.matmul(A, Bm, mm_, k, n);
|
| 357 |
-
// batched: all `batch` problems in ONE dispatch (CUTLASS ex. 05/24). The
|
| 358 |
-
// per-head backward is 4 GEMMs x B x heads of tiny matrices; issued one at a
|
| 359 |
-
// time the GPU spends all its time on dispatch overhead rather than math.
|
| 360 |
-
const bmmB = units
|
| 361 |
-
? (A, Bm, rows, k, n, batch) =>
|
| 362 |
-
V.vgemmBlock(A, Bm, { m: rows, k, n, batch }, m.ctx.L, m.ctx.bgemm, m.ctx.audit)
|
| 363 |
-
: async (A, Bm, rows, k, n, batch) => {
|
| 364 |
-
const out = new Float32Array(batch * rows * n);
|
| 365 |
-
for (let bz = 0; bz < batch; bz++)
|
| 366 |
-
out.set(TC.matmul(A.subarray(bz * rows * k, (bz + 1) * rows * k),
|
| 367 |
-
Bm.subarray(bz * k * n, (bz + 1) * k * n), rows, k, n), bz * rows * n);
|
| 368 |
-
return out;
|
| 369 |
-
};
|
| 370 |
-
// tied unembed: logits = lnf @ embᵀ, so the unembedding gradient flows
|
| 371 |
-
// straight into emb — dlogitsᵀ @ lnf is V×C, emb's own shape
|
| 372 |
-
let dlnfIn;
|
| 373 |
-
if (m.ctx.fgemm2 && !units) {
|
| 374 |
-
// Both GEMMs consume dlogits (BT x vocab, ~17 MB at the 16512 vocab).
|
| 375 |
-
// fgemm2 uploads it ONCE and runs both on one submit — profiling had
|
| 376 |
-
// this pair at 55% of the step, over half of it re-uploading the same
|
| 377 |
-
// operand. Bit-identical to the two separate calls (gated at init).
|
| 378 |
-
[g[gi.emb], dlnfIn] = await m.ctx.fgemm2(
|
| 379 |
-
cache.dlogits,
|
| 380 |
-
cache.lnf.y, { m: vocab, k: BT, n: C, transA: true },
|
| 381 |
-
m.emb, { m: BT, k: vocab, n: C }); // split-K shape
|
| 382 |
-
} else if (m.ctx.fgemm && !units) {
|
| 383 |
-
[g[gi.emb], dlnfIn] = await Promise.all([
|
| 384 |
-
m.ctx.fgemm(cache.dlogits, cache.lnf.y, { m: vocab, k: BT, n: C, transA: true }),
|
| 385 |
-
m.ctx.fgemm(cache.dlogits, m.emb, { m: BT, k: vocab, n: C }), // split-K shape
|
| 386 |
-
]);
|
| 387 |
-
} else if (units && m.ctx.bgemm) {
|
| 388 |
-
// units + GPU: the g.emb operand is dlogitsᵀ (vocab×BT, ~4M elements), and
|
| 389 |
-
// tr() + quantizeRows() is three full passes over it in JS. Quantizing the
|
| 390 |
-
// COLUMNS of dlogits directly into transposed int8 is one pass and
|
| 391 |
-
// bit-identical: same |max| scan, same rounds, in the same order — only
|
| 392 |
-
// the write pattern changes. The GEMM itself still goes through ctx.bgemm
|
| 393 |
-
// (exact-gated), and the live-shape audit still samples it.
|
| 394 |
-
const quantizeColsAsRows = (X, rows, cols) => { // == quantizeRows(tr(X), cols, rows)
|
| 395 |
-
V.assertFinite(X, "quantizeColsAsRows");
|
| 396 |
-
const q = new Int8Array(cols * rows), s = new Float32Array(cols);
|
| 397 |
-
for (let c = 0; c < cols; c++) {
|
| 398 |
-
let mx = 0;
|
| 399 |
-
for (let r = 0; r < rows; r++) { const a = Math.abs(X[r * cols + c]); if (a > mx) mx = a; }
|
| 400 |
-
const sc = Math.max(mx / 127, 1e-8); s[c] = sc;
|
| 401 |
-
for (let r = 0; r < rows; r++) {
|
| 402 |
-
const v = Math.round(X[r * cols + c] / sc);
|
| 403 |
-
q[c * rows + r] = v < -128 ? -128 : v > 127 ? 127 : v;
|
| 404 |
-
}
|
| 405 |
-
}
|
| 406 |
-
return { q, s };
|
| 407 |
-
};
|
| 408 |
-
const dlq = quantizeColsAsRows(cache.dlogits, BT, vocab); // dlogitsᵀ quantized, one pass
|
| 409 |
-
const wq2 = V.quantizeCols(cache.lnf.y, BT, C);
|
| 410 |
-
const dEmb = { m: vocab, k: BT, n: C, batch: 1 };
|
| 411 |
-
const [gEmb, dIn] = await Promise.all([
|
| 412 |
-
m.ctx.bgemm(dlq.q, wq2.q, dlq.s, wq2.s, dEmb),
|
| 413 |
-
bmm(cache.dlogits, m.emb, BT, vocab, C),
|
| 414 |
-
]);
|
| 415 |
-
if (m.ctx.audit && m.ctx.audit.due()) {
|
| 416 |
-
const bad = V.auditTile(dlq.q, wq2.q, dlq.s, wq2.s, dEmb, gEmb, m.ctx.L, m.ctx.audit.cells);
|
| 417 |
-
if (bad) m.ctx.audit.fail(bad);
|
| 418 |
-
}
|
| 419 |
-
g[gi.emb] = gEmb; dlnfIn = dIn;
|
| 420 |
-
} else {
|
| 421 |
-
// independent GEMMs — overlap them (these are the two vocab-sized calls,
|
| 422 |
-
// the largest in the whole backward; each is its own round trip)
|
| 423 |
-
[g[gi.emb], dlnfIn] = await Promise.all([
|
| 424 |
-
bmm(tr(cache.dlogits, BT, vocab), cache.lnf.y, vocab, BT, C),
|
| 425 |
-
bmm(cache.dlogits, m.emb, BT, vocab, C),
|
| 426 |
-
]);
|
| 427 |
-
}
|
| 428 |
-
let dx = lnBwd(dlnfIn, cache.lnf.y, cache.lnf.sig, BT, C);
|
| 429 |
-
const scale = 1 / Math.sqrt(hd);
|
| 430 |
-
// concat helper for fusing sibling GEMMs into one batched dispatch
|
| 431 |
-
const cat = (...arrs) => {
|
| 432 |
-
const out = new Float32Array(arrs.reduce((a, x) => a + x.length, 0));
|
| 433 |
-
let o = 0; for (const x of arrs) { out.set(x, o); o += x.length; }
|
| 434 |
-
return out;
|
| 435 |
-
};
|
| 436 |
-
for (let l = layers - 1; l >= 0; l--) {
|
| 437 |
-
const bl = m.blocks[l], cb = cache.blocks[l];
|
| 438 |
-
// mlp: x3 = x2 + relu(ln2 @ W1) @ W2
|
| 439 |
-
// gW2 and dh1 are independent — overlap their dispatches. On GPU each bmm
|
| 440 |
-
// is a full upload/submit/readback round trip, so sequential awaits leave
|
| 441 |
-
// the GPU idle between every pair; this is pure latency, not arithmetic,
|
| 442 |
-
// and each GEMM's int32 accumulation is exact so overlap changes no bit.
|
| 443 |
-
const dmlpOut = dx; // residual passthrough handled below
|
| 444 |
-
const [gW2, dh1] = await Promise.all([
|
| 445 |
-
bmm(tr(cb.h1, BT, hidden), dmlpOut, hidden, BT, C),
|
| 446 |
-
bmm(dmlpOut, tr(bl.W2, hidden, C), BT, C, hidden),
|
| 447 |
-
]);
|
| 448 |
-
g[gi[`b${l}.W2`]] = gW2;
|
| 449 |
-
for (let i = 0; i < dh1.length; i++) if (!cb.mask[i]) dh1[i] = 0;
|
| 450 |
-
const [gW1, dln2raw] = await Promise.all([
|
| 451 |
-
bmm(tr(cb.ln2.y, BT, C), dh1, C, BT, hidden),
|
| 452 |
-
bmm(dh1, tr(bl.W1, C, hidden), BT, hidden, C),
|
| 453 |
-
]);
|
| 454 |
-
g[gi[`b${l}.W1`]] = gW1;
|
| 455 |
-
const dln2in = lnBwd(dln2raw, cb.ln2.y, cb.ln2.sig, BT, C);
|
| 456 |
-
const dx2 = new Float32Array(BT * C);
|
| 457 |
-
for (let i = 0; i < dx2.length; i++) dx2[i] = dx[i] + dln2in[i];
|
| 458 |
-
// attention: x2 = xin + (ctxOut @ Wo)
|
| 459 |
-
const [gWo, dctx] = await Promise.all([
|
| 460 |
-
bmm(tr(cb.ctxOut, BT, C), dx2, C, BT, C),
|
| 461 |
-
bmm(dx2, tr(bl.Wo, C, C), BT, C, C),
|
| 462 |
-
]);
|
| 463 |
-
g[gi[`b${l}.Wo`]] = gWo;
|
| 464 |
-
// gather every head once, then run each stage as ONE batched GEMM over all
|
| 465 |
-
// B*heads problems: 4 dispatches per layer instead of 4 per head.
|
| 466 |
-
const BH = B * heads;
|
| 467 |
-
const qb = gatherHeads(cb.q, B, T, C, heads, hd);
|
| 468 |
-
const kb = gatherHeads(cb.k, B, T, C, heads, hd);
|
| 469 |
-
const vb = gatherHeads(cb.v, B, T, C, heads, hd);
|
| 470 |
-
const dchb = gatherHeads(dctx, B, T, C, heads, hd);
|
| 471 |
-
const aT = batchedTranspose(cb.aAll, BH, T, T); // BH×T×T
|
| 472 |
-
const vT = batchedTranspose(vb, BH, T, hd); // BH×hd×T
|
| 473 |
-
const [dvAll, daAll] = await Promise.all([
|
| 474 |
-
bmmB(aT, dchb, T, T, hd, BH), // aᵀ @ dctx
|
| 475 |
-
bmmB(dchb, vT, T, hd, T, BH), // dctx @ vᵀ
|
| 476 |
-
]);
|
| 477 |
-
// softmax backward is elementwise + a causal row reduction: stays in float
|
| 478 |
-
// (no matrix math here, so nothing for the units to do)
|
| 479 |
-
const dsAll = new Float32Array(BH * T * T);
|
| 480 |
-
for (let bz = 0; bz < BH; bz++) {
|
| 481 |
-
const o = bz * T * T;
|
| 482 |
-
for (let ti = 0; ti < T; ti++) {
|
| 483 |
-
let dot = 0;
|
| 484 |
-
for (let tj = 0; tj <= ti; tj++) dot += daAll[o + ti * T + tj] * cb.aAll[o + ti * T + tj];
|
| 485 |
-
for (let tj = 0; tj <= ti; tj++)
|
| 486 |
-
dsAll[o + ti * T + tj] = cb.aAll[o + ti * T + tj] * (daAll[o + ti * T + tj] - dot) * scale;
|
| 487 |
-
}
|
| 488 |
-
}
|
| 489 |
-
const dsT = batchedTranspose(dsAll, BH, T, T);
|
| 490 |
-
const [dqAll, dkAll] = await Promise.all([
|
| 491 |
-
bmmB(dsAll, kb, T, T, hd, BH), // ds @ k
|
| 492 |
-
bmmB(dsT, qb, T, T, hd, BH), // dsᵀ @ q
|
| 493 |
-
]);
|
| 494 |
-
const dq = new Float32Array(BT * C), dk = new Float32Array(BT * C), dv = new Float32Array(BT * C);
|
| 495 |
-
scatterHeadsAcc(dq, dqAll, B, T, C, heads, hd);
|
| 496 |
-
scatterHeadsAcc(dk, dkAll, B, T, C, heads, hd);
|
| 497 |
-
scatterHeadsAcc(dv, dvAll, B, T, C, heads, hd);
|
| 498 |
-
// The QKV weight grads share the same left operand (ln1ᵀ), and the three
|
| 499 |
-
// dln1in terms share one shape — each trio fuses into ONE batched GEMM
|
| 500 |
-
// (batch=3) instead of three dispatches. Bit-identical to separate calls:
|
| 501 |
-
// block scales are per-row of X and per-column of W PER BATCH ELEMENT, so
|
| 502 |
-
// concatenation changes no scale and no product.
|
| 503 |
-
const ln1T = tr(cb.ln1.y, BT, C);
|
| 504 |
-
const [gQKV, dIn3] = await Promise.all([
|
| 505 |
-
bmmB(cat(ln1T, ln1T, ln1T), cat(dq, dk, dv), C, BT, C, 3),
|
| 506 |
-
bmmB(cat(dq, dk, dv), cat(tr(bl.Wq, C, C), tr(bl.Wk, C, C), tr(bl.Wv, C, C)), BT, C, C, 3),
|
| 507 |
-
]);
|
| 508 |
-
const CC = C * C, BTC = BT * C;
|
| 509 |
-
g[gi[`b${l}.Wq`]] = gQKV.slice(0, CC);
|
| 510 |
-
g[gi[`b${l}.Wk`]] = gQKV.slice(CC, 2 * CC);
|
| 511 |
-
g[gi[`b${l}.Wv`]] = gQKV.slice(2 * CC, 3 * CC);
|
| 512 |
-
// sum the three dln1in terms in q,k,v order with an f32 round after EACH
|
| 513 |
-
// add — the old code accumulated into a Float32Array element three times,
|
| 514 |
-
// which rounds per step; a bare q+k+v here would run in f64 and round
|
| 515 |
-
// once, a last-ulp difference that forks replicas. (Exactly the epilogue
|
| 516 |
-
// mirror lesson: match the rounding schedule, not just the values.)
|
| 517 |
-
const dln1in = new Float32Array(BTC);
|
| 518 |
-
for (let i = 0; i < BTC; i++)
|
| 519 |
-
dln1in[i] = Math.fround(Math.fround(dIn3[i] + dIn3[BTC + i]) + dIn3[2 * BTC + i]);
|
| 520 |
-
const dxin = lnBwd(dln1in, cb.ln1.y, cb.ln1.sig, BT, C);
|
| 521 |
-
dx = new Float32Array(BT * C);
|
| 522 |
-
for (let i = 0; i < dx.length; i++) dx[i] = dx2[i] + dxin[i];
|
| 523 |
-
}
|
| 524 |
-
// embedding + positional
|
| 525 |
-
const ge = g[gi.emb], gp = g[gi.pos];
|
| 526 |
-
for (let i = 0; i < BT; i++) {
|
| 527 |
-
const id = cache.X[i], tpos = i % T;
|
| 528 |
-
for (let j = 0; j < C; j++) { ge[id * C + j] += dx[i * C + j]; gp[tpos * C + j] += dx[i * C + j]; }
|
| 529 |
-
}
|
| 530 |
-
// flatten
|
| 531 |
-
const flat = new Float32Array(m.nParams);
|
| 532 |
-
let off = 0;
|
| 533 |
-
for (const t of g) { flat.set(t, off); off += t.length; }
|
| 534 |
-
return flat;
|
| 535 |
-
}
|
| 536 |
-
|
| 537 |
-
async function trainStep(m) {
|
| 538 |
-
const { X, Y } = sampleBatch(m.cfg);
|
| 539 |
-
const { loss, cache } = await forward(m, X, Y);
|
| 540 |
-
const grad = await backward(m, cache);
|
| 541 |
-
return { loss, grad };
|
| 542 |
-
}
|
| 543 |
-
|
| 544 |
-
function applyUpdate(m, upd) { // W -= upd (lr folded in by the optimizer)
|
| 545 |
-
let off = 0;
|
| 546 |
-
for (const p of m.params) { for (let i = 0; i < p.length; i++) p[i] -= upd[off + i]; off += p.length; }
|
| 547 |
-
}
|
| 548 |
-
function getFlatParams(m) {
|
| 549 |
-
const flat = new Float32Array(m.nParams);
|
| 550 |
-
let off = 0;
|
| 551 |
-
for (const p of m.params) { flat.set(p, off); off += p.length; }
|
| 552 |
-
return flat;
|
| 553 |
-
}
|
| 554 |
-
function setFlatParams(m, flat) {
|
| 555 |
-
let off = 0;
|
| 556 |
-
for (const p of m.params) { p.set(flat.subarray(off, off + p.length)); off += p.length; }
|
| 557 |
-
}
|
| 558 |
-
|
| 559 |
-
// greedy sampling — watch the model actually speak
|
| 560 |
-
async function generate(m, prompt, nChars) {
|
| 561 |
-
const { t: T } = m.cfg;
|
| 562 |
-
let ids = [...encode(prompt)];
|
| 563 |
-
for (let step = 0; step < nChars; step++) {
|
| 564 |
-
const win = ids.slice(-T);
|
| 565 |
-
const X = new Int32Array(T), Y = new Int32Array(T);
|
| 566 |
-
for (let i = 0; i < win.length; i++) X[T - win.length + i] = win[i];
|
| 567 |
-
const save = m.cfg.b; m.cfg.b = 1;
|
| 568 |
-
const { logits } = await forward(m, X, Y);
|
| 569 |
-
m.cfg.b = save;
|
| 570 |
-
const row = (T - 1) * m.cfg.vocab;
|
| 571 |
-
let best = 0, bv = -1e30;
|
| 572 |
-
for (let j = 0; j < m.cfg.vocab; j++) if (logits[row + j] > bv) { bv = logits[row + j]; best = j; }
|
| 573 |
-
ids.push(best);
|
| 574 |
-
}
|
| 575 |
-
return decode(ids);
|
| 576 |
-
}
|
| 577 |
-
|
| 578 |
-
const api = { init, trainStep, applyUpdate, getFlatParams, setFlatParams, generate,
|
| 579 |
-
streamFineWebEdu, streamDataset, datasetName, loadTokenizer, loadTokenizerData,
|
| 580 |
-
vocabSize, tokenizerName, encode, decode };
|
| 581 |
-
if (typeof module !== "undefined" && module.exports) { TC = require("./traincore.js"); V = require("./verified_core.js"); module.exports = api; }
|
| 582 |
-
else { TC = root.TrainCore; V = root.Verified; root.Transformer = api; }
|
| 583 |
-
})(typeof self !== "undefined" ? self : this);
|
|
|
|
| 1 |
+
// A miniature transformer language model that trains THROUGH the verified INT8
|
| 2 |
+
// units: every matrix product in the forward pass — QKV projections, attention
|
| 3 |
+
// scores, attention·values, output projection, the MLP, and the unembedding —
|
| 4 |
+
// runs through the verified multiply LUT (an emulated INT8 tensor core).
|
| 5 |
+
// Backward is a straight-through estimator in float (the integer path has no
|
| 6 |
+
// gradient), exactly like the Python VerifiedLinear.
|
| 7 |
+
//
|
| 8 |
+
// Task: next-character prediction on a deterministic, self-generated corpus
|
| 9 |
+
// (every peer builds the same text from the same seed — nothing to download).
|
| 10 |
+
(function (root) {
|
| 11 |
+
"use strict";
|
| 12 |
+
|
| 13 |
+
let TC, V; // TrainCore / Verified — resolved per environment at the end
|
| 14 |
+
|
| 15 |
+
function mulberry32(a) { return function () { a |= 0; a = a + 0x6D2B79F5 | 0; let t = Math.imul(a ^ a >>> 15, 1 | a); t = t + Math.imul(t ^ t >>> 7, 61 | t) ^ t; return ((t ^ t >>> 14) >>> 0) / 4294967296; }; }
|
| 16 |
+
function randn(n, rng) { const o = new Float32Array(n); for (let i = 0; i < n; i += 2) { let u = 0, v = 0; while (u === 0) u = rng(); while (v === 0) v = rng(); const m = Math.sqrt(-2 * Math.log(u)); o[i] = m * Math.cos(2 * Math.PI * v); if (i + 1 < n) o[i + 1] = m * Math.sin(2 * Math.PI * v); } return o; }
|
| 17 |
+
|
| 18 |
+
// ---- corpus: deterministic cottagecore prose, identical on every peer -----
|
| 19 |
+
const W_ADJ = ["mossy", "golden", "amber", "quiet", "little", "misty", "sunny", "wild", "cozy", "dusty", "merry", "brave"];
|
| 20 |
+
const W_NOUN = ["fox", "hare", "owl", "badger", "toad", "sparrow", "otter", "deer", "mushroom", "acorn", "willow", "robin", "river", "meadow", "garden", "lantern"];
|
| 21 |
+
const W_VERB = ["naps", "sings", "wanders", "hides", "dreams", "waits", "dances", "listens", "rests", "grows"];
|
| 22 |
+
const W_PREP = ["by", "under", "near", "beside", "beyond", "inside"];
|
| 23 |
+
function buildCorpus() {
|
| 24 |
+
const rng = mulberry32(20260712);
|
| 25 |
+
const pick = (a) => a[Math.floor(rng() * a.length)];
|
| 26 |
+
let s = "";
|
| 27 |
+
while (s.length < 60000)
|
| 28 |
+
s += `the ${pick(W_ADJ)} ${pick(W_NOUN)} ${pick(W_VERB)} ${pick(W_PREP)} the ${pick(W_ADJ)} ${pick(W_NOUN)}. `;
|
| 29 |
+
return s;
|
| 30 |
+
}
|
| 31 |
+
// ---- tokenizer -------------------------------------------------------------
|
| 32 |
+
// Spikewhale tokenizer (tokenizer.json): byte-level greedy longest-match
|
| 33 |
+
// ("length-max"), ~16.5k tokens. Until it loads (or if the file is missing)
|
| 34 |
+
// a 96-char byte-level vocab keeps the app working — but ALL devices in a
|
| 35 |
+
// group must use the same tokenizer (the config broadcast enforces it).
|
| 36 |
+
const FALLBACK_CHARS = [...Array(95)].map((_, i) => String.fromCharCode(32 + i)).concat(["\n"]);
|
| 37 |
+
let tok = {
|
| 38 |
+
name: "char-96 (fallback)",
|
| 39 |
+
vocab: Object.fromEntries(FALLBACK_CHARS.map((c, i) => [c, i])),
|
| 40 |
+
ids: FALLBACK_CHARS, maxLen: 1, size: FALLBACK_CHARS.length,
|
| 41 |
+
unk: 0, specials: new Set(),
|
| 42 |
+
};
|
| 43 |
+
tok.unk = tok.vocab[" "];
|
| 44 |
+
function vocabSize() { return tok.size; }
|
| 45 |
+
function tokenizerName() { return tok.name; }
|
| 46 |
+
function loadTokenizerData(d) { // plain {vocab, vocab_size, max_token_len}
|
| 47 |
+
const ids = new Array(d.vocab_size);
|
| 48 |
+
for (const [t, i] of Object.entries(d.vocab)) ids[i] = t;
|
| 49 |
+
tok = { name: `Spikewhale length-max (${d.vocab_size} tokens)`,
|
| 50 |
+
vocab: d.vocab, ids, maxLen: d.max_token_len || 24, size: d.vocab_size,
|
| 51 |
+
unk: d.vocab["<unk>"] ?? 1,
|
| 52 |
+
specials: new Set(["<pad>", "<unk>", "<bos>", "<eos>", ...(d.special_tokens || [])]) };
|
| 53 |
+
IDS = encode(CORPUS); // re-tokenize whatever corpus is loaded
|
| 54 |
+
return tok.name;
|
| 55 |
+
}
|
| 56 |
+
async function loadTokenizer(url) {
|
| 57 |
+
const r = await fetch(url || "tokenizer.json");
|
| 58 |
+
if (!r.ok) throw new Error(`tokenizer.json HTTP ${r.status}`);
|
| 59 |
+
return loadTokenizerData(await r.json());
|
| 60 |
+
}
|
| 61 |
+
function toLatin1(s) { const b = new TextEncoder().encode(s); let o = ""; for (const x of b) o += String.fromCharCode(x); return o; }
|
| 62 |
+
function encode(text) { // greedy longest match over bytes
|
| 63 |
+
const s = toLatin1(text), out = [];
|
| 64 |
+
let i = 0;
|
| 65 |
+
while (i < s.length) {
|
| 66 |
+
let m = null;
|
| 67 |
+
for (let L = Math.min(tok.maxLen, s.length - i); L > 0; L--) {
|
| 68 |
+
const sub = s.substr(i, L);
|
| 69 |
+
if (sub in tok.vocab) { m = sub; break; }
|
| 70 |
+
}
|
| 71 |
+
if (m === null) { out.push(tok.unk); i++; continue; }
|
| 72 |
+
out.push(tok.vocab[m]); i += m.length;
|
| 73 |
+
}
|
| 74 |
+
return Int32Array.from(out);
|
| 75 |
+
}
|
| 76 |
+
function decode(idArr) {
|
| 77 |
+
let s = "";
|
| 78 |
+
for (const id of idArr) {
|
| 79 |
+
const t = tok.ids[id];
|
| 80 |
+
if (t === undefined || tok.specials.has(t)) continue;
|
| 81 |
+
s += t;
|
| 82 |
+
}
|
| 83 |
+
const bytes = Uint8Array.from([...s].map(c => c.charCodeAt(0)));
|
| 84 |
+
return new TextDecoder().decode(bytes);
|
| 85 |
+
}
|
| 86 |
+
let CORPUS = buildCorpus();
|
| 87 |
+
let IDS = encode(CORPUS);
|
| 88 |
+
let DATASET = "built-in corpus";
|
| 89 |
+
|
| 90 |
+
// Training text: FineWeb-Edu (10BT sample), HARDCODED as the only dataset.
|
| 91 |
+
// The serving Space reads random slices of the parquet shards straight off
|
| 92 |
+
// the HF CDN with range requests (see server.js /data) — no dependency on
|
| 93 |
+
// the datasets-server rows API, which 503s routinely. Each device pulls its
|
| 94 |
+
// own random slice (that's data parallelism — batches were always
|
| 95 |
+
// per-device anyway). Offline or on failure the built-in corpus stays.
|
| 96 |
+
const DEFAULT_DS = "HuggingFaceFW/fineweb-edu";
|
| 97 |
+
async function streamDataset() { // dataset choice removed on purpose
|
| 98 |
+
const r = await fetch("data");
|
| 99 |
+
if (!r.ok) throw new Error(`/data HTTP ${r.status}`);
|
| 100 |
+
const text = (await r.text()).replace(/[^\x20-\x7e\n]/g, " ");
|
| 101 |
+
if (text.length < 10000) throw new Error("too little text returned");
|
| 102 |
+
CORPUS = text.slice(0, 500000);
|
| 103 |
+
IDS = encode(CORPUS);
|
| 104 |
+
DATASET = `${DEFAULT_DS} · 10BT sample (parquet via this Space)`;
|
| 105 |
+
return { name: DATASET, chars: CORPUS.length };
|
| 106 |
+
}
|
| 107 |
+
const streamFineWebEdu = () => streamDataset();
|
| 108 |
+
function datasetName() { return DATASET; }
|
| 109 |
+
|
| 110 |
+
// ---- verified matmul: block-scaled INT8 through the units ------------------
|
| 111 |
+
// CUTLASS ex. 67/81 blockwise scaling: per-row activation scales × per-column
|
| 112 |
+
// weight scales, one exact LUT/DP4A GEMM, dequant (+ optional fused ReLU) in
|
| 113 |
+
// the kernel epilogue (ex. 12). Replaces the per-tensor 3-pass: same outlier
|
| 114 |
+
// robustness at one third of the unit ops.
|
| 115 |
+
async function vmm(Xf, Wf, m, k, n, ctx, relu) {
|
| 116 |
+
// ctx.audit re-checks random cells of this LIVE GEMM against the units
|
| 117 |
+
return V.vgemmBlock(Xf, Wf, { m, k, n, batch: 1, relu: !!relu }, ctx.L, ctx.bgemm, ctx.audit);
|
| 118 |
+
}
|
| 119 |
+
// CUTLASS ex. 45 (dual GEMM): sibling GEMMs that share the same LEFT operand
|
| 120 |
+
// run as ONE batched dispatch, and the shared operand is quantized ONCE
|
| 121 |
+
// instead of once per sibling. Used for the q/k/v projections — same X
|
| 122 |
+
// (ln1.y), three weights, identical shapes. Bit-identical to three separate
|
| 123 |
+
// vmm calls: quantizeRows is deterministic (same input -> same int8+scales),
|
| 124 |
+
// the tiled copies index exactly like separate batch elements, and block
|
| 125 |
+
// scales are per-row/per-column PER BATCH ELEMENT, so concatenation changes
|
| 126 |
+
// no scale and no product. The batched kernel is the same exact-gated bgemm
|
| 127 |
+
// that training already runs, and the live-shape audit still samples it.
|
| 128 |
+
async function vmmShared3(Xf, Wa, Wb, Wc, m, k, n, ctx) {
|
| 129 |
+
const x = V.quantizeRows(Xf, m, k);
|
| 130 |
+
const xq = new Int8Array(3 * m * k), xs = new Float32Array(3 * m);
|
| 131 |
+
for (let i = 0; i < 3; i++) { xq.set(x.q, i * m * k); xs.set(x.s, i * m); }
|
| 132 |
+
const wq = new Int8Array(3 * k * n), ws = new Float32Array(3 * n);
|
| 133 |
+
[Wa, Wb, Wc].forEach((W, i) => { const w = V.quantizeCols(W, k, n); wq.set(w.q, i * k * n); ws.set(w.s, i * n); });
|
| 134 |
+
const d = { m, k, n, batch: 3 };
|
| 135 |
+
let out;
|
| 136 |
+
if (ctx.bgemm) {
|
| 137 |
+
out = await ctx.bgemm(xq, wq, xs, ws, d);
|
| 138 |
+
if (ctx.audit && ctx.audit.due()) {
|
| 139 |
+
const bad = V.auditTile(xq, wq, xs, ws, d, out, ctx.L, ctx.audit.cells);
|
| 140 |
+
if (bad) ctx.audit.fail(bad);
|
| 141 |
+
}
|
| 142 |
+
} else {
|
| 143 |
+
out = V.bgemmJS(xq, wq, xs, ws, d, ctx.L);
|
| 144 |
+
}
|
| 145 |
+
const MN = m * n;
|
| 146 |
+
return [out.subarray(0, MN), out.subarray(MN, 2 * MN), out.subarray(2 * MN, 3 * MN)];
|
| 147 |
+
}
|
| 148 |
+
|
| 149 |
+
// ---- layernorm (no affine) -------------------------------------------------
|
| 150 |
+
function lnFwd(x, rows, C) {
|
| 151 |
+
const y = new Float32Array(rows * C), sig = new Float32Array(rows);
|
| 152 |
+
for (let r = 0; r < rows; r++) {
|
| 153 |
+
let mu = 0; for (let j = 0; j < C; j++) mu += x[r * C + j]; mu /= C;
|
| 154 |
+
let v = 0; for (let j = 0; j < C; j++) { const d = x[r * C + j] - mu; v += d * d; }
|
| 155 |
+
const s = Math.sqrt(v / C + 1e-5); sig[r] = s;
|
| 156 |
+
for (let j = 0; j < C; j++) y[r * C + j] = (x[r * C + j] - mu) / s;
|
| 157 |
+
}
|
| 158 |
+
return { y, sig };
|
| 159 |
+
}
|
| 160 |
+
function lnBwd(dy, y, sig, rows, C) {
|
| 161 |
+
const dx = new Float32Array(rows * C);
|
| 162 |
+
for (let r = 0; r < rows; r++) {
|
| 163 |
+
let mdy = 0, mdyy = 0;
|
| 164 |
+
for (let j = 0; j < C; j++) { mdy += dy[r * C + j]; mdyy += dy[r * C + j] * y[r * C + j]; }
|
| 165 |
+
mdy /= C; mdyy /= C;
|
| 166 |
+
for (let j = 0; j < C; j++) dx[r * C + j] = (dy[r * C + j] - mdy - y[r * C + j] * mdyy) / sig[r];
|
| 167 |
+
}
|
| 168 |
+
return dx;
|
| 169 |
+
}
|
| 170 |
+
|
| 171 |
+
// ---- model -----------------------------------------------------------------
|
| 172 |
+
// cfg: { c: width, t: seq len, b: batch/device, layers, heads, steps, lr }
|
| 173 |
+
// engine: the Compute backend object ({bgemm} for the fused WebGPU path) or a
|
| 174 |
+
// legacy matmulInt8 function (Node tests, inference kit) -> CPU LUT mirror
|
| 175 |
+
function init(cfg, L, engine, audit) {
|
| 176 |
+
const c = cfg.c, layers = cfg.layers || 2, heads = cfg.heads || 2, hidden = 2 * c;
|
| 177 |
+
let seed = 100;
|
| 178 |
+
const mk = (nEl, scale) => { const w = randn(nEl, mulberry32(seed++)); for (let i = 0; i < nEl; i++) w[i] *= scale; return w; };
|
| 179 |
+
const params = [], names = [];
|
| 180 |
+
const add = (name, w) => { params.push(w); names.push(name); return w; };
|
| 181 |
+
const m = {
|
| 182 |
+
cfg: { ...cfg, layers, heads, hidden, vocab: vocabSize() },
|
| 183 |
+
ctx: { L, bgemm: (engine && engine.bgemm) || null,
|
| 184 |
+
att: (engine && engine.att) || null, fgemm: (engine && engine.fgemm) || null,
|
| 185 |
+
fgemm2: (engine && engine.fgemm2) || null,
|
| 186 |
+
mlp: (engine && engine.mlp) || null,
|
| 187 |
+
audit: audit || null, unitBackward: !!cfg.unitBackward },
|
| 188 |
+
emb: add("emb", mk(vocabSize() * c, 0.08)),
|
| 189 |
+
pos: add("pos", mk(cfg.t * c, 0.02)),
|
| 190 |
+
blocks: [], params, names,
|
| 191 |
+
};
|
| 192 |
+
for (let l = 0; l < layers; l++)
|
| 193 |
+
m.blocks.push({
|
| 194 |
+
Wq: add(`b${l}.Wq`, mk(c * c, 0.08)), Wk: add(`b${l}.Wk`, mk(c * c, 0.08)),
|
| 195 |
+
Wv: add(`b${l}.Wv`, mk(c * c, 0.08)), Wo: add(`b${l}.Wo`, mk(c * c, 0.08)),
|
| 196 |
+
W1: add(`b${l}.W1`, mk(c * hidden, 0.08)), W2: add(`b${l}.W2`, mk(hidden * c, 0.08)),
|
| 197 |
+
});
|
| 198 |
+
// weight-tied unembedding: logits use embᵀ (no separate Wu). Halves the
|
| 199 |
+
// vocab-sized parameters — and with a 16k vocab that's ~half of ALL
|
| 200 |
+
// parameters, so gradients over the wire shrink ~2× too.
|
| 201 |
+
m.nParams = params.reduce((a, p) => a + p.length, 0);
|
| 202 |
+
return m;
|
| 203 |
+
}
|
| 204 |
+
|
| 205 |
+
function sampleBatch(cfg) {
|
| 206 |
+
const { b, t } = cfg;
|
| 207 |
+
const X = new Int32Array(b * t), Y = new Int32Array(b * t);
|
| 208 |
+
for (let i = 0; i < b; i++) {
|
| 209 |
+
const off = Math.floor(Math.random() * (IDS.length - t - 1));
|
| 210 |
+
for (let j = 0; j < t; j++) { X[i * t + j] = IDS[off + j]; Y[i * t + j] = IDS[off + j + 1]; }
|
| 211 |
+
}
|
| 212 |
+
return { X, Y };
|
| 213 |
+
}
|
| 214 |
+
|
| 215 |
+
// ---- head layout helpers ---------------------------------------------------
|
| 216 |
+
// q/k/v live as BT×C with head h owning columns [h*hd, (h+1)*hd). The backward
|
| 217 |
+
// wants every head as its own GEMM problem, so gather once into BH×T×hd and
|
| 218 |
+
// scatter back at the end — one pass each, instead of slicing per head inside
|
| 219 |
+
// the loop and paying a GPU dispatch per tiny matmul.
|
| 220 |
+
function gatherHeads(x, B, T, C, heads, hd) { // BT×C -> BH×T×hd
|
| 221 |
+
const out = new Float32Array(B * heads * T * hd);
|
| 222 |
+
for (let bi = 0; bi < B; bi++)
|
| 223 |
+
for (let h = 0; h < heads; h++) {
|
| 224 |
+
const bz = bi * heads + h;
|
| 225 |
+
for (let ti = 0; ti < T; ti++)
|
| 226 |
+
for (let j = 0; j < hd; j++) out[(bz * T + ti) * hd + j] = x[(bi * T + ti) * C + h * hd + j];
|
| 227 |
+
}
|
| 228 |
+
return out;
|
| 229 |
+
}
|
| 230 |
+
function scatterHeadsAcc(dst, src, B, T, C, heads, hd) { // BH×T×hd -> += BT×C
|
| 231 |
+
for (let bi = 0; bi < B; bi++)
|
| 232 |
+
for (let h = 0; h < heads; h++) {
|
| 233 |
+
const bz = bi * heads + h;
|
| 234 |
+
for (let ti = 0; ti < T; ti++)
|
| 235 |
+
for (let j = 0; j < hd; j++) dst[(bi * T + ti) * C + h * hd + j] += src[(bz * T + ti) * hd + j];
|
| 236 |
+
}
|
| 237 |
+
}
|
| 238 |
+
function batchedTranspose(x, batch, rows, cols) { // per-batch rows×cols -> cols×rows
|
| 239 |
+
const out = new Float32Array(batch * rows * cols);
|
| 240 |
+
for (let b = 0; b < batch; b++) {
|
| 241 |
+
const o = b * rows * cols;
|
| 242 |
+
for (let r = 0; r < rows; r++)
|
| 243 |
+
for (let c = 0; c < cols; c++) out[o + c * rows + r] = x[o + r * cols + c];
|
| 244 |
+
}
|
| 245 |
+
return out;
|
| 246 |
+
}
|
| 247 |
+
|
| 248 |
+
// ---- forward THROUGH the verified units (caches kept for STE backward) -----
|
| 249 |
+
async function forward(m, X, Y) {
|
| 250 |
+
const { c: C, t: T, b: B, layers, heads, hidden, vocab } = m.cfg;
|
| 251 |
+
const BT = B * T, hd = C / heads, ctx = m.ctx;
|
| 252 |
+
const cache = { X, Y, blocks: [] };
|
| 253 |
+
let x = new Float32Array(BT * C);
|
| 254 |
+
for (let i = 0; i < BT; i++) {
|
| 255 |
+
const id = X[i], tpos = i % T;
|
| 256 |
+
for (let j = 0; j < C; j++) x[i * C + j] = m.emb[id * C + j] + m.pos[tpos * C + j];
|
| 257 |
+
}
|
| 258 |
+
for (let l = 0; l < layers; l++) {
|
| 259 |
+
const bl = m.blocks[l], cb = { xin: x };
|
| 260 |
+
const l1 = lnFwd(x, BT, C); cb.ln1 = l1;
|
| 261 |
+
// q/k/v share the same left operand — one batched dispatch, one quantize
|
| 262 |
+
// of ln1.y instead of three (CUTLASS ex. 45; see vmmShared3)
|
| 263 |
+
const [q, k, v] = await vmmShared3(l1.y, bl.Wq, bl.Wk, bl.Wv, BT, C, C, ctx);
|
| 264 |
+
cb.q = q; cb.k = k; cb.v = v;
|
| 265 |
+
const scale = 1 / Math.sqrt(hd);
|
| 266 |
+
// gather-FUSED attention (CUTLASS ex. 36/52): the kernels read q/k/v in
|
| 267 |
+
// their natural BT×C layout with head-strided indexing and scatter ctx
|
| 268 |
+
// straight back — no JS gather copies, no kᵀ transpose. All B×H heads in
|
| 269 |
+
// one dispatch per stage, every product through the verified units.
|
| 270 |
+
const BH = B * heads, dAtt = { B, T, heads, hd };
|
| 271 |
+
// per-(token,head) row quantization: the (BT·heads)×hd view IS the buffer
|
| 272 |
+
const qq = V.quantizeRows(q, BT * heads, hd), kq = V.quantizeRows(k, BT * heads, hd);
|
| 273 |
+
const sAll = ctx.att ? await ctx.att.scores(qq.q, kq.q, qq.s, kq.s, dAtt)
|
| 274 |
+
: V.attScoresJS(qq.q, kq.q, qq.s, kq.s, dAtt, ctx.L);
|
| 275 |
+
// live-shape audit: the init gate only ever saw four test shapes
|
| 276 |
+
if (ctx.att && ctx.audit && ctx.audit.due()) {
|
| 277 |
+
const bad = V.auditAttScores(qq.q, kq.q, qq.s, kq.s, dAtt, sAll, ctx.L, ctx.audit.cells);
|
| 278 |
+
if (bad) ctx.audit.fail(bad);
|
| 279 |
+
}
|
| 280 |
+
const aAll = new Float32Array(BH * T * T); // causal softmax
|
| 281 |
+
for (let bz = 0; bz < BH; bz++) {
|
| 282 |
+
const so = bz * T * T;
|
| 283 |
+
for (let ti = 0; ti < T; ti++) {
|
| 284 |
+
let mx = -1e30;
|
| 285 |
+
for (let tj = 0; tj <= ti; tj++) mx = Math.max(mx, sAll[so + ti * T + tj] * scale);
|
| 286 |
+
let z = 0;
|
| 287 |
+
for (let tj = 0; tj <= ti; tj++) { const e = Math.exp(sAll[so + ti * T + tj] * scale - mx); aAll[so + ti * T + tj] = e; z += e; }
|
| 288 |
+
for (let tj = 0; tj <= ti; tj++) aAll[so + ti * T + tj] /= z;
|
| 289 |
+
}
|
| 290 |
+
}
|
| 291 |
+
const aq = V.quantizeRows(aAll, BH * T, T);
|
| 292 |
+
const vq = V.quantizeHeadCols(v, B, T, heads, hd);
|
| 293 |
+
const ctxOut = ctx.att ? await ctx.att.ctx(aq.q, vq.q, aq.s, vq.s, dAtt)
|
| 294 |
+
: V.attCtxJS(aq.q, vq.q, aq.s, vq.s, dAtt, ctx.L);
|
| 295 |
+
if (ctx.att && ctx.audit && ctx.audit.due()) {
|
| 296 |
+
const bad = V.auditAttCtx(aq.q, vq.q, aq.s, vq.s, dAtt, ctxOut, ctx.L, ctx.audit.cells);
|
| 297 |
+
if (bad) ctx.audit.fail(bad);
|
| 298 |
+
}
|
| 299 |
+
cb.aAll = aAll; // backward slices heads from q/k/v/aAll
|
| 300 |
+
cb.ctxOut = ctxOut;
|
| 301 |
+
const attnOut = await vmm(ctxOut, bl.Wo, BT, C, C, ctx);
|
| 302 |
+
const x2 = new Float32Array(BT * C);
|
| 303 |
+
for (let i = 0; i < x2.length; i++) x2[i] = x[i] + attnOut[i];
|
| 304 |
+
cb.x2 = x2;
|
| 305 |
+
const l2 = lnFwd(x2, BT, C); cb.ln2 = l2;
|
| 306 |
+
// CUTLASS ex. 13 + 23: both MLP GEMMs run back-to-back on the GPU. The
|
| 307 |
+
// intermediate h1 is quantized ON-DEVICE (exact-gated respec — see
|
| 308 |
+
// vmlpBlock in verified_core.js) and only its per-row absmax (~1KB)
|
| 309 |
+
// visits JS between the GEMMs; h1 itself comes back solely because the
|
| 310 |
+
// STE backward needs it. CPU devices run the bit-identical mirror chain.
|
| 311 |
+
const { h1, out: mlpOut } = await V.vmlpBlock(l2.y, bl.W1, bl.W2,
|
| 312 |
+
{ m: BT, k: C, h: hidden, n: C }, ctx.L, ctx.mlp, ctx.audit);
|
| 313 |
+
const mask = new Uint8Array(h1.length);
|
| 314 |
+
for (let i = 0; i < h1.length; i++) if (h1[i] > 0) mask[i] = 1;
|
| 315 |
+
cb.h1 = h1; cb.mask = mask;
|
| 316 |
+
x = new Float32Array(BT * C);
|
| 317 |
+
for (let i = 0; i < x.length; i++) x[i] = x2[i] + mlpOut[i];
|
| 318 |
+
cache.blocks.push(cb);
|
| 319 |
+
}
|
| 320 |
+
const lf = lnFwd(x, BT, C); cache.lnf = lf; cache.xf = x;
|
| 321 |
+
const logits = await vmm(lf.y, TC.transpose(m.emb, vocab, C), BT, C, vocab, ctx); // tied: embᵀ
|
| 322 |
+
// cross-entropy + dlogits
|
| 323 |
+
let loss = 0;
|
| 324 |
+
const dlogits = new Float32Array(BT * vocab);
|
| 325 |
+
for (let i = 0; i < BT; i++) {
|
| 326 |
+
let mx = -1e30;
|
| 327 |
+
for (let j = 0; j < vocab; j++) mx = Math.max(mx, logits[i * vocab + j]);
|
| 328 |
+
let z = 0;
|
| 329 |
+
for (let j = 0; j < vocab; j++) z += Math.exp(logits[i * vocab + j] - mx);
|
| 330 |
+
const lz = Math.log(z) + mx;
|
| 331 |
+
loss += lz - logits[i * vocab + Y[i]];
|
| 332 |
+
for (let j = 0; j < vocab; j++)
|
| 333 |
+
dlogits[i * vocab + j] = (Math.exp(logits[i * vocab + j] - lz) - (j === Y[i] ? 1 : 0)) / BT;
|
| 334 |
+
}
|
| 335 |
+
loss /= BT;
|
| 336 |
+
cache.dlogits = dlogits;
|
| 337 |
+
return { loss, cache, logits };
|
| 338 |
+
}
|
| 339 |
+
|
| 340 |
+
// ---- STE backward (float), mirrors forward exactly --------------------------
|
| 341 |
+
// The two vocab-sized matmuls run on the split-K f32 GPU kernel when
|
| 342 |
+
// available (CUTLASS ex. 06) — same float math, off the JS thread.
|
| 343 |
+
async function backward(m, cache) {
|
| 344 |
+
const { c: C, t: T, b: B, layers, heads, hidden, vocab } = m.cfg;
|
| 345 |
+
const BT = B * T, hd = C / heads, tr = TC.transpose;
|
| 346 |
+
const g = m.params.map(p => new Float32Array(p.length));
|
| 347 |
+
const gi = Object.fromEntries(m.names.map((n, i) => [n, i]));
|
| 348 |
+
// Every matmul here goes through `bmm`. With ctx.unitBackward the STE
|
| 349 |
+
// gradient is computed BY the verified units (block-scaled int8, exact int32
|
| 350 |
+
// accumulate) instead of in float. STE is a claim about the math — pretend
|
| 351 |
+
// the quantizer was the identity — not about the datatype that evaluates it,
|
| 352 |
+
// so the two are orthogonal and this stays a correct STE.
|
| 353 |
+
const units = !!m.ctx.unitBackward;
|
| 354 |
+
const bmm = units
|
| 355 |
+
? (A, Bm, mm_, k, n) => vmm(A, Bm, mm_, k, n, m.ctx)
|
| 356 |
+
: async (A, Bm, mm_, k, n) => TC.matmul(A, Bm, mm_, k, n);
|
| 357 |
+
// batched: all `batch` problems in ONE dispatch (CUTLASS ex. 05/24). The
|
| 358 |
+
// per-head backward is 4 GEMMs x B x heads of tiny matrices; issued one at a
|
| 359 |
+
// time the GPU spends all its time on dispatch overhead rather than math.
|
| 360 |
+
const bmmB = units
|
| 361 |
+
? (A, Bm, rows, k, n, batch) =>
|
| 362 |
+
V.vgemmBlock(A, Bm, { m: rows, k, n, batch }, m.ctx.L, m.ctx.bgemm, m.ctx.audit)
|
| 363 |
+
: async (A, Bm, rows, k, n, batch) => {
|
| 364 |
+
const out = new Float32Array(batch * rows * n);
|
| 365 |
+
for (let bz = 0; bz < batch; bz++)
|
| 366 |
+
out.set(TC.matmul(A.subarray(bz * rows * k, (bz + 1) * rows * k),
|
| 367 |
+
Bm.subarray(bz * k * n, (bz + 1) * k * n), rows, k, n), bz * rows * n);
|
| 368 |
+
return out;
|
| 369 |
+
};
|
| 370 |
+
// tied unembed: logits = lnf @ embᵀ, so the unembedding gradient flows
|
| 371 |
+
// straight into emb — dlogitsᵀ @ lnf is V×C, emb's own shape
|
| 372 |
+
let dlnfIn;
|
| 373 |
+
if (m.ctx.fgemm2 && !units) {
|
| 374 |
+
// Both GEMMs consume dlogits (BT x vocab, ~17 MB at the 16512 vocab).
|
| 375 |
+
// fgemm2 uploads it ONCE and runs both on one submit — profiling had
|
| 376 |
+
// this pair at 55% of the step, over half of it re-uploading the same
|
| 377 |
+
// operand. Bit-identical to the two separate calls (gated at init).
|
| 378 |
+
[g[gi.emb], dlnfIn] = await m.ctx.fgemm2(
|
| 379 |
+
cache.dlogits,
|
| 380 |
+
cache.lnf.y, { m: vocab, k: BT, n: C, transA: true },
|
| 381 |
+
m.emb, { m: BT, k: vocab, n: C }); // split-K shape
|
| 382 |
+
} else if (m.ctx.fgemm && !units) {
|
| 383 |
+
[g[gi.emb], dlnfIn] = await Promise.all([
|
| 384 |
+
m.ctx.fgemm(cache.dlogits, cache.lnf.y, { m: vocab, k: BT, n: C, transA: true }),
|
| 385 |
+
m.ctx.fgemm(cache.dlogits, m.emb, { m: BT, k: vocab, n: C }), // split-K shape
|
| 386 |
+
]);
|
| 387 |
+
} else if (units && m.ctx.bgemm) {
|
| 388 |
+
// units + GPU: the g.emb operand is dlogitsᵀ (vocab×BT, ~4M elements), and
|
| 389 |
+
// tr() + quantizeRows() is three full passes over it in JS. Quantizing the
|
| 390 |
+
// COLUMNS of dlogits directly into transposed int8 is one pass and
|
| 391 |
+
// bit-identical: same |max| scan, same rounds, in the same order — only
|
| 392 |
+
// the write pattern changes. The GEMM itself still goes through ctx.bgemm
|
| 393 |
+
// (exact-gated), and the live-shape audit still samples it.
|
| 394 |
+
const quantizeColsAsRows = (X, rows, cols) => { // == quantizeRows(tr(X), cols, rows)
|
| 395 |
+
V.assertFinite(X, "quantizeColsAsRows");
|
| 396 |
+
const q = new Int8Array(cols * rows), s = new Float32Array(cols);
|
| 397 |
+
for (let c = 0; c < cols; c++) {
|
| 398 |
+
let mx = 0;
|
| 399 |
+
for (let r = 0; r < rows; r++) { const a = Math.abs(X[r * cols + c]); if (a > mx) mx = a; }
|
| 400 |
+
const sc = Math.max(mx / 127, 1e-8); s[c] = sc;
|
| 401 |
+
for (let r = 0; r < rows; r++) {
|
| 402 |
+
const v = Math.round(X[r * cols + c] / sc);
|
| 403 |
+
q[c * rows + r] = v < -128 ? -128 : v > 127 ? 127 : v;
|
| 404 |
+
}
|
| 405 |
+
}
|
| 406 |
+
return { q, s };
|
| 407 |
+
};
|
| 408 |
+
const dlq = quantizeColsAsRows(cache.dlogits, BT, vocab); // dlogitsᵀ quantized, one pass
|
| 409 |
+
const wq2 = V.quantizeCols(cache.lnf.y, BT, C);
|
| 410 |
+
const dEmb = { m: vocab, k: BT, n: C, batch: 1 };
|
| 411 |
+
const [gEmb, dIn] = await Promise.all([
|
| 412 |
+
m.ctx.bgemm(dlq.q, wq2.q, dlq.s, wq2.s, dEmb),
|
| 413 |
+
bmm(cache.dlogits, m.emb, BT, vocab, C),
|
| 414 |
+
]);
|
| 415 |
+
if (m.ctx.audit && m.ctx.audit.due()) {
|
| 416 |
+
const bad = V.auditTile(dlq.q, wq2.q, dlq.s, wq2.s, dEmb, gEmb, m.ctx.L, m.ctx.audit.cells);
|
| 417 |
+
if (bad) m.ctx.audit.fail(bad);
|
| 418 |
+
}
|
| 419 |
+
g[gi.emb] = gEmb; dlnfIn = dIn;
|
| 420 |
+
} else {
|
| 421 |
+
// independent GEMMs — overlap them (these are the two vocab-sized calls,
|
| 422 |
+
// the largest in the whole backward; each is its own round trip)
|
| 423 |
+
[g[gi.emb], dlnfIn] = await Promise.all([
|
| 424 |
+
bmm(tr(cache.dlogits, BT, vocab), cache.lnf.y, vocab, BT, C),
|
| 425 |
+
bmm(cache.dlogits, m.emb, BT, vocab, C),
|
| 426 |
+
]);
|
| 427 |
+
}
|
| 428 |
+
let dx = lnBwd(dlnfIn, cache.lnf.y, cache.lnf.sig, BT, C);
|
| 429 |
+
const scale = 1 / Math.sqrt(hd);
|
| 430 |
+
// concat helper for fusing sibling GEMMs into one batched dispatch
|
| 431 |
+
const cat = (...arrs) => {
|
| 432 |
+
const out = new Float32Array(arrs.reduce((a, x) => a + x.length, 0));
|
| 433 |
+
let o = 0; for (const x of arrs) { out.set(x, o); o += x.length; }
|
| 434 |
+
return out;
|
| 435 |
+
};
|
| 436 |
+
for (let l = layers - 1; l >= 0; l--) {
|
| 437 |
+
const bl = m.blocks[l], cb = cache.blocks[l];
|
| 438 |
+
// mlp: x3 = x2 + relu(ln2 @ W1) @ W2
|
| 439 |
+
// gW2 and dh1 are independent — overlap their dispatches. On GPU each bmm
|
| 440 |
+
// is a full upload/submit/readback round trip, so sequential awaits leave
|
| 441 |
+
// the GPU idle between every pair; this is pure latency, not arithmetic,
|
| 442 |
+
// and each GEMM's int32 accumulation is exact so overlap changes no bit.
|
| 443 |
+
const dmlpOut = dx; // residual passthrough handled below
|
| 444 |
+
const [gW2, dh1] = await Promise.all([
|
| 445 |
+
bmm(tr(cb.h1, BT, hidden), dmlpOut, hidden, BT, C),
|
| 446 |
+
bmm(dmlpOut, tr(bl.W2, hidden, C), BT, C, hidden),
|
| 447 |
+
]);
|
| 448 |
+
g[gi[`b${l}.W2`]] = gW2;
|
| 449 |
+
for (let i = 0; i < dh1.length; i++) if (!cb.mask[i]) dh1[i] = 0;
|
| 450 |
+
const [gW1, dln2raw] = await Promise.all([
|
| 451 |
+
bmm(tr(cb.ln2.y, BT, C), dh1, C, BT, hidden),
|
| 452 |
+
bmm(dh1, tr(bl.W1, C, hidden), BT, hidden, C),
|
| 453 |
+
]);
|
| 454 |
+
g[gi[`b${l}.W1`]] = gW1;
|
| 455 |
+
const dln2in = lnBwd(dln2raw, cb.ln2.y, cb.ln2.sig, BT, C);
|
| 456 |
+
const dx2 = new Float32Array(BT * C);
|
| 457 |
+
for (let i = 0; i < dx2.length; i++) dx2[i] = dx[i] + dln2in[i];
|
| 458 |
+
// attention: x2 = xin + (ctxOut @ Wo)
|
| 459 |
+
const [gWo, dctx] = await Promise.all([
|
| 460 |
+
bmm(tr(cb.ctxOut, BT, C), dx2, C, BT, C),
|
| 461 |
+
bmm(dx2, tr(bl.Wo, C, C), BT, C, C),
|
| 462 |
+
]);
|
| 463 |
+
g[gi[`b${l}.Wo`]] = gWo;
|
| 464 |
+
// gather every head once, then run each stage as ONE batched GEMM over all
|
| 465 |
+
// B*heads problems: 4 dispatches per layer instead of 4 per head.
|
| 466 |
+
const BH = B * heads;
|
| 467 |
+
const qb = gatherHeads(cb.q, B, T, C, heads, hd);
|
| 468 |
+
const kb = gatherHeads(cb.k, B, T, C, heads, hd);
|
| 469 |
+
const vb = gatherHeads(cb.v, B, T, C, heads, hd);
|
| 470 |
+
const dchb = gatherHeads(dctx, B, T, C, heads, hd);
|
| 471 |
+
const aT = batchedTranspose(cb.aAll, BH, T, T); // BH×T×T
|
| 472 |
+
const vT = batchedTranspose(vb, BH, T, hd); // BH×hd×T
|
| 473 |
+
const [dvAll, daAll] = await Promise.all([
|
| 474 |
+
bmmB(aT, dchb, T, T, hd, BH), // aᵀ @ dctx
|
| 475 |
+
bmmB(dchb, vT, T, hd, T, BH), // dctx @ vᵀ
|
| 476 |
+
]);
|
| 477 |
+
// softmax backward is elementwise + a causal row reduction: stays in float
|
| 478 |
+
// (no matrix math here, so nothing for the units to do)
|
| 479 |
+
const dsAll = new Float32Array(BH * T * T);
|
| 480 |
+
for (let bz = 0; bz < BH; bz++) {
|
| 481 |
+
const o = bz * T * T;
|
| 482 |
+
for (let ti = 0; ti < T; ti++) {
|
| 483 |
+
let dot = 0;
|
| 484 |
+
for (let tj = 0; tj <= ti; tj++) dot += daAll[o + ti * T + tj] * cb.aAll[o + ti * T + tj];
|
| 485 |
+
for (let tj = 0; tj <= ti; tj++)
|
| 486 |
+
dsAll[o + ti * T + tj] = cb.aAll[o + ti * T + tj] * (daAll[o + ti * T + tj] - dot) * scale;
|
| 487 |
+
}
|
| 488 |
+
}
|
| 489 |
+
const dsT = batchedTranspose(dsAll, BH, T, T);
|
| 490 |
+
const [dqAll, dkAll] = await Promise.all([
|
| 491 |
+
bmmB(dsAll, kb, T, T, hd, BH), // ds @ k
|
| 492 |
+
bmmB(dsT, qb, T, T, hd, BH), // dsᵀ @ q
|
| 493 |
+
]);
|
| 494 |
+
const dq = new Float32Array(BT * C), dk = new Float32Array(BT * C), dv = new Float32Array(BT * C);
|
| 495 |
+
scatterHeadsAcc(dq, dqAll, B, T, C, heads, hd);
|
| 496 |
+
scatterHeadsAcc(dk, dkAll, B, T, C, heads, hd);
|
| 497 |
+
scatterHeadsAcc(dv, dvAll, B, T, C, heads, hd);
|
| 498 |
+
// The QKV weight grads share the same left operand (ln1ᵀ), and the three
|
| 499 |
+
// dln1in terms share one shape — each trio fuses into ONE batched GEMM
|
| 500 |
+
// (batch=3) instead of three dispatches. Bit-identical to separate calls:
|
| 501 |
+
// block scales are per-row of X and per-column of W PER BATCH ELEMENT, so
|
| 502 |
+
// concatenation changes no scale and no product.
|
| 503 |
+
const ln1T = tr(cb.ln1.y, BT, C);
|
| 504 |
+
const [gQKV, dIn3] = await Promise.all([
|
| 505 |
+
bmmB(cat(ln1T, ln1T, ln1T), cat(dq, dk, dv), C, BT, C, 3),
|
| 506 |
+
bmmB(cat(dq, dk, dv), cat(tr(bl.Wq, C, C), tr(bl.Wk, C, C), tr(bl.Wv, C, C)), BT, C, C, 3),
|
| 507 |
+
]);
|
| 508 |
+
const CC = C * C, BTC = BT * C;
|
| 509 |
+
g[gi[`b${l}.Wq`]] = gQKV.slice(0, CC);
|
| 510 |
+
g[gi[`b${l}.Wk`]] = gQKV.slice(CC, 2 * CC);
|
| 511 |
+
g[gi[`b${l}.Wv`]] = gQKV.slice(2 * CC, 3 * CC);
|
| 512 |
+
// sum the three dln1in terms in q,k,v order with an f32 round after EACH
|
| 513 |
+
// add — the old code accumulated into a Float32Array element three times,
|
| 514 |
+
// which rounds per step; a bare q+k+v here would run in f64 and round
|
| 515 |
+
// once, a last-ulp difference that forks replicas. (Exactly the epilogue
|
| 516 |
+
// mirror lesson: match the rounding schedule, not just the values.)
|
| 517 |
+
const dln1in = new Float32Array(BTC);
|
| 518 |
+
for (let i = 0; i < BTC; i++)
|
| 519 |
+
dln1in[i] = Math.fround(Math.fround(dIn3[i] + dIn3[BTC + i]) + dIn3[2 * BTC + i]);
|
| 520 |
+
const dxin = lnBwd(dln1in, cb.ln1.y, cb.ln1.sig, BT, C);
|
| 521 |
+
dx = new Float32Array(BT * C);
|
| 522 |
+
for (let i = 0; i < dx.length; i++) dx[i] = dx2[i] + dxin[i];
|
| 523 |
+
}
|
| 524 |
+
// embedding + positional
|
| 525 |
+
const ge = g[gi.emb], gp = g[gi.pos];
|
| 526 |
+
for (let i = 0; i < BT; i++) {
|
| 527 |
+
const id = cache.X[i], tpos = i % T;
|
| 528 |
+
for (let j = 0; j < C; j++) { ge[id * C + j] += dx[i * C + j]; gp[tpos * C + j] += dx[i * C + j]; }
|
| 529 |
+
}
|
| 530 |
+
// flatten
|
| 531 |
+
const flat = new Float32Array(m.nParams);
|
| 532 |
+
let off = 0;
|
| 533 |
+
for (const t of g) { flat.set(t, off); off += t.length; }
|
| 534 |
+
return flat;
|
| 535 |
+
}
|
| 536 |
+
|
| 537 |
+
async function trainStep(m) {
|
| 538 |
+
const { X, Y } = sampleBatch(m.cfg);
|
| 539 |
+
const { loss, cache } = await forward(m, X, Y);
|
| 540 |
+
const grad = await backward(m, cache);
|
| 541 |
+
return { loss, grad };
|
| 542 |
+
}
|
| 543 |
+
|
| 544 |
+
function applyUpdate(m, upd) { // W -= upd (lr folded in by the optimizer)
|
| 545 |
+
let off = 0;
|
| 546 |
+
for (const p of m.params) { for (let i = 0; i < p.length; i++) p[i] -= upd[off + i]; off += p.length; }
|
| 547 |
+
}
|
| 548 |
+
function getFlatParams(m) {
|
| 549 |
+
const flat = new Float32Array(m.nParams);
|
| 550 |
+
let off = 0;
|
| 551 |
+
for (const p of m.params) { flat.set(p, off); off += p.length; }
|
| 552 |
+
return flat;
|
| 553 |
+
}
|
| 554 |
+
function setFlatParams(m, flat) {
|
| 555 |
+
let off = 0;
|
| 556 |
+
for (const p of m.params) { p.set(flat.subarray(off, off + p.length)); off += p.length; }
|
| 557 |
+
}
|
| 558 |
+
|
| 559 |
+
// greedy sampling — watch the model actually speak
|
| 560 |
+
async function generate(m, prompt, nChars) {
|
| 561 |
+
const { t: T } = m.cfg;
|
| 562 |
+
let ids = [...encode(prompt)];
|
| 563 |
+
for (let step = 0; step < nChars; step++) {
|
| 564 |
+
const win = ids.slice(-T);
|
| 565 |
+
const X = new Int32Array(T), Y = new Int32Array(T);
|
| 566 |
+
for (let i = 0; i < win.length; i++) X[T - win.length + i] = win[i];
|
| 567 |
+
const save = m.cfg.b; m.cfg.b = 1;
|
| 568 |
+
const { logits } = await forward(m, X, Y);
|
| 569 |
+
m.cfg.b = save;
|
| 570 |
+
const row = (T - 1) * m.cfg.vocab;
|
| 571 |
+
let best = 0, bv = -1e30;
|
| 572 |
+
for (let j = 0; j < m.cfg.vocab; j++) if (logits[row + j] > bv) { bv = logits[row + j]; best = j; }
|
| 573 |
+
ids.push(best);
|
| 574 |
+
}
|
| 575 |
+
return decode(ids);
|
| 576 |
+
}
|
| 577 |
+
|
| 578 |
+
const api = { init, trainStep, applyUpdate, getFlatParams, setFlatParams, generate,
|
| 579 |
+
streamFineWebEdu, streamDataset, datasetName, loadTokenizer, loadTokenizerData,
|
| 580 |
+
vocabSize, tokenizerName, encode, decode };
|
| 581 |
+
if (typeof module !== "undefined" && module.exports) { TC = require("./traincore.js"); V = require("./verified_core.js"); module.exports = api; }
|
| 582 |
+
else { TC = root.TrainCore; V = root.Verified; root.Transformer = api; }
|
| 583 |
+
})(typeof self !== "undefined" ? self : this);
|
web/public/verified_core.js
CHANGED
|
@@ -1,519 +1,519 @@
|
|
| 1 |
-
// Verified INT8 compute — the emulated GPU logic, in the browser.
|
| 2 |
-
// A layer's forward runs THROUGH the units: quantize -> LUT multiply -> requant
|
| 3 |
-
// -> optional ReLU -> dequant. Backward is a straight-through estimator (the
|
| 4 |
-
// integer path has no gradient), so ordinary float weights still learn.
|
| 5 |
-
// Same units as the Python/Docker DaisyChain; here they're lookup tables.
|
| 6 |
-
(function (root) {
|
| 7 |
-
"use strict";
|
| 8 |
-
|
| 9 |
-
let TC; // TrainCore (matmul/transpose) — resolved per environment at the end
|
| 10 |
-
|
| 11 |
-
// A float -> int8 conversion has no answer for NaN/Inf, and the quantizers below give a
|
| 12 |
-
// silent wrong one: an Inf makes the |max| scale Infinity and EVERY value quantizes to 0
|
| 13 |
-
// ([0.5, Inf, 3] -> [0, 0, 0]); a NaN is skipped by `a > mx` and then stored as 0 by the
|
| 14 |
-
// Int8Array. The Python trainer refuses the same input (qat.py); so does this. One pass,
|
| 15 |
-
// no effect on finite inputs, so builds with and without it still co-train.
|
| 16 |
-
function assertFinite(X, where) {
|
| 17 |
-
for (let i = 0; i < X.length; i++) {
|
| 18 |
-
if (!Number.isFinite(X[i])) {
|
| 19 |
-
let bad = 0;
|
| 20 |
-
for (let j = i; j < X.length; j++) if (!Number.isFinite(X[j])) bad++;
|
| 21 |
-
const e = new Error("non-finite input to " + where + " (" + bad + " of " + X.length +
|
| 22 |
-
" values): the int8 quantize would turn it into 0");
|
| 23 |
-
e.name = "NonFiniteError";
|
| 24 |
-
throw e;
|
| 25 |
-
}
|
| 26 |
-
}
|
| 27 |
-
}
|
| 28 |
-
|
| 29 |
-
function quantize(X) {
|
| 30 |
-
assertFinite(X, "quantize");
|
| 31 |
-
let mx = 0; for (let i = 0; i < X.length; i++) { const a = Math.abs(X[i]); if (a > mx) mx = a; }
|
| 32 |
-
const scale = Math.max(mx / 127, 1e-8);
|
| 33 |
-
const q = new Int8Array(X.length);
|
| 34 |
-
for (let i = 0; i < X.length; i++) { let v = Math.round(X[i] / scale); q[i] = v < -128 ? -128 : v > 127 ? 127 : v; }
|
| 35 |
-
return { q, scale };
|
| 36 |
-
}
|
| 37 |
-
|
| 38 |
-
// int8 matmul via the verified multiply LUT: acc(m×n) = sum_k mulLUT[Xq,Wq]
|
| 39 |
-
function lutMatmulJS(Xq, Wq, m, k, n, L) {
|
| 40 |
-
const C = new Int32Array(m * n), mul = L.mul;
|
| 41 |
-
for (let i = 0; i < m; i++) {
|
| 42 |
-
for (let p = 0; p < k; p++) {
|
| 43 |
-
const au = (Xq[i * k + p] & 0xFF) * 256, wo = p * n, co = i * n;
|
| 44 |
-
for (let j = 0; j < n; j++) C[co + j] += mul[au + (Wq[wo + j] & 0xFF)];
|
| 45 |
-
}
|
| 46 |
-
}
|
| 47 |
-
return C;
|
| 48 |
-
}
|
| 49 |
-
|
| 50 |
-
// ---- 3xINT8 fast-accurate GEMM --------------------------------------------
|
| 51 |
-
// The CUTLASS example-27 "3xTF32" scheme, ported to the verified units:
|
| 52 |
-
// split each float into a coarse int8 part plus an int8-quantized residual,
|
| 53 |
-
// run three EXACT LUT GEMMs (hi·hi, hi·lo, lo·hi), drop the negligible
|
| 54 |
-
// lo·lo, and recombine. Same big/small decomposition NVIDIA uses to recover
|
| 55 |
-
// near-fp32 accuracy from TF32 tensor cores — here it recovers ~14-bit
|
| 56 |
-
// accuracy from the 8-bit units, at 3× the unit ops. Every product still
|
| 57 |
-
// goes through the verified mul8 LUT.
|
| 58 |
-
function quantize2(X) {
|
| 59 |
-
const hi = quantize(X);
|
| 60 |
-
const r = new Float32Array(X.length);
|
| 61 |
-
for (let i = 0; i < X.length; i++) r[i] = X[i] - hi.q[i] * hi.scale;
|
| 62 |
-
const lo = quantize(r);
|
| 63 |
-
return { hi, lo };
|
| 64 |
-
}
|
| 65 |
-
function combine3(hh, hl, lh, x, w, len) {
|
| 66 |
-
const out = new Float32Array(len);
|
| 67 |
-
const shh = x.hi.scale * w.hi.scale, shl = x.hi.scale * w.lo.scale, slh = x.lo.scale * w.hi.scale;
|
| 68 |
-
for (let i = 0; i < len; i++) out[i] = hh[i] * shh + hl[i] * shl + lh[i] * slh;
|
| 69 |
-
return out;
|
| 70 |
-
}
|
| 71 |
-
function lutMatmul3JS(Xf, Wf, m, k, n, L) { // sync, CPU LUT path
|
| 72 |
-
const x = quantize2(Xf), w = quantize2(Wf);
|
| 73 |
-
return combine3(lutMatmulJS(x.hi.q, w.hi.q, m, k, n, L),
|
| 74 |
-
lutMatmulJS(x.hi.q, w.lo.q, m, k, n, L),
|
| 75 |
-
lutMatmulJS(x.lo.q, w.hi.q, m, k, n, L), x, w, m * n);
|
| 76 |
-
}
|
| 77 |
-
async function lutMatmul3(Xf, Wf, m, k, n, L, matmulInt8) { // any backend
|
| 78 |
-
const x = quantize2(Xf), w = quantize2(Wf);
|
| 79 |
-
const mm = matmulInt8 || lutMatmulJS;
|
| 80 |
-
const [hh, hl, lh] = await Promise.all([
|
| 81 |
-
mm(x.hi.q, w.hi.q, m, k, n, L),
|
| 82 |
-
mm(x.hi.q, w.lo.q, m, k, n, L),
|
| 83 |
-
mm(x.lo.q, w.hi.q, m, k, n, L),
|
| 84 |
-
]);
|
| 85 |
-
return combine3(hh, hl, lh, x, w, m * n);
|
| 86 |
-
}
|
| 87 |
-
|
| 88 |
-
// ---- block-scaled verified GEMM (CUTLASS ex. 67/81 blockwise scaling) ------
|
| 89 |
-
// Per-ROW scales for the activations and per-COLUMN scales for the weights:
|
| 90 |
-
// the integer math through the mul8 LUT is completely unchanged — only the
|
| 91 |
-
// dequant uses rs[row]·cs[col] instead of one tensor-wide product, so a single
|
| 92 |
-
// outlier no longer crushes the quantization resolution of every other
|
| 93 |
-
// row/column. One LUT pass at this granularity beats the per-tensor 3-pass.
|
| 94 |
-
function quantizeRows(X, rows, cols) {
|
| 95 |
-
assertFinite(X, "quantizeRows");
|
| 96 |
-
const q = new Int8Array(rows * cols), s = new Float32Array(rows);
|
| 97 |
-
for (let r = 0; r < rows; r++) {
|
| 98 |
-
let mx = 0;
|
| 99 |
-
for (let c = 0; c < cols; c++) { const a = Math.abs(X[r * cols + c]); if (a > mx) mx = a; }
|
| 100 |
-
const sc = Math.max(mx / 127, 1e-8); s[r] = sc;
|
| 101 |
-
for (let c = 0; c < cols; c++) {
|
| 102 |
-
const v = Math.round(X[r * cols + c] / sc);
|
| 103 |
-
q[r * cols + c] = v < -128 ? -128 : v > 127 ? 127 : v;
|
| 104 |
-
}
|
| 105 |
-
}
|
| 106 |
-
return { q, s };
|
| 107 |
-
}
|
| 108 |
-
function quantizeCols(W, rows, cols) {
|
| 109 |
-
assertFinite(W, "quantizeCols");
|
| 110 |
-
const q = new Int8Array(rows * cols), s = new Float32Array(cols);
|
| 111 |
-
for (let c = 0; c < cols; c++) {
|
| 112 |
-
let mx = 0;
|
| 113 |
-
for (let r = 0; r < rows; r++) { const a = Math.abs(W[r * cols + c]); if (a > mx) mx = a; }
|
| 114 |
-
s[c] = Math.max(mx / 127, 1e-8);
|
| 115 |
-
}
|
| 116 |
-
for (let r = 0; r < rows; r++)
|
| 117 |
-
for (let c = 0; c < cols; c++) {
|
| 118 |
-
const v = Math.round(W[r * cols + c] / s[c]);
|
| 119 |
-
q[r * cols + c] = v < -128 ? -128 : v > 127 ? 127 : v;
|
| 120 |
-
}
|
| 121 |
-
return { q, s };
|
| 122 |
-
}
|
| 123 |
-
|
| 124 |
-
// ---- B2B MLP chain (CUTLASS ex. 13 two-GEMM fusion + ex. 23 epilogue
|
| 125 |
-
// reduction), respecced for cross-device exactness ---------------------------
|
| 126 |
-
// The MLP is the one back-to-back GEMM pair with no layernorm/softmax between
|
| 127 |
-
// (ReLU is already fused in the epilogue), so the intermediate h1 can be
|
| 128 |
-
// quantized ON the GPU and fed straight to the second GEMM. Two rules make
|
| 129 |
-
// that fleet-safe:
|
| 130 |
-
// 1. The per-row |max| (ex. 23) uses only comparisons — exact on any
|
| 131 |
-
// hardware, order-independent — and comes back to JS as ~1KB.
|
| 132 |
-
// 2. Scale DERIVATION (two divisions) stays in JS f64, which IEEE requires
|
| 133 |
-
// to be exactly rounded and is therefore identical on every device.
|
| 134 |
-
// WGSL division is only 2.5 ULP — a fork waiting to happen — but WGSL
|
| 135 |
-
// multiply/add are correctly rounded and floor/clamp are exact. So the
|
| 136 |
-
// quantize step is respecced from round(x / scale) to
|
| 137 |
-
// floor(f32(x * invScale) + 0.5) — floor(x+0.5) IS Math.round's tie
|
| 138 |
-
// rule — and the fround-stepped mirror below is bit-identical to the
|
| 139 |
-
// GPU kernel, which is exact-gated against it at init.
|
| 140 |
-
// NOTE this changes which int8 a value on a rounding boundary lands on
|
| 141 |
-
// (≤1 step) vs quantizeRows, so old and new builds cannot co-train — the
|
| 142 |
-
// divergence guard stops such mixed groups by design.
|
| 143 |
-
function rowAbsMax(X, rows, cols) {
|
| 144 |
-
assertFinite(X, "rowAbsMax");
|
| 145 |
-
const mx = new Float32Array(rows);
|
| 146 |
-
for (let r = 0; r < rows; r++) {
|
| 147 |
-
let m = 0;
|
| 148 |
-
for (let c = 0; c < cols; c++) { const a = Math.abs(X[r * cols + c]); if (a > m) m = a; }
|
| 149 |
-
mx[r] = m;
|
| 150 |
-
}
|
| 151 |
-
return mx;
|
| 152 |
-
}
|
| 153 |
-
function scalesFromAbsMax(mx) { // f64 divisions: exactly rounded, device-identical
|
| 154 |
-
const scale = new Float32Array(mx.length), inv = new Float32Array(mx.length);
|
| 155 |
-
for (let i = 0; i < mx.length; i++) {
|
| 156 |
-
scale[i] = Math.max(mx[i] / 127, 1e-8);
|
| 157 |
-
inv[i] = 1 / scale[i]; // recip of the STORED f32 scale
|
| 158 |
-
}
|
| 159 |
-
return { scale, inv };
|
| 160 |
-
}
|
| 161 |
-
function quantizeRowsInv(X, rows, cols, inv) { // bit-exact mirror of the GPU quantize kernel
|
| 162 |
-
const q = new Int8Array(rows * cols);
|
| 163 |
-
for (let r = 0; r < rows; r++) {
|
| 164 |
-
const iv = inv[r];
|
| 165 |
-
for (let c = 0; c < cols; c++) {
|
| 166 |
-
const n = Math.floor(f32(f32(X[r * cols + c] * iv) + 0.5));
|
| 167 |
-
q[r * cols + c] = n < -128 ? -128 : n > 127 ? 127 : n;
|
| 168 |
-
}
|
| 169 |
-
}
|
| 170 |
-
return q;
|
| 171 |
-
}
|
| 172 |
-
// the chained MLP: X @ W1 -> ReLU (fused) -> absmax -> quantize -> @ W2.
|
| 173 |
-
// d = { m, k, h, n }; gpuMlp (from webgpu.js) runs both GEMMs + the
|
| 174 |
-
// on-GPU quantize with one tiny absmax readback between; without it the CPU
|
| 175 |
-
// mirror chain runs — SAME math, so mixed GPU/CPU fleets stay bit-identical.
|
| 176 |
-
async function vmlpBlock(Xf, W1f, W2f, d, L, gpuMlp, audit) {
|
| 177 |
-
const x = quantizeRows(Xf, d.m, d.k);
|
| 178 |
-
const w1 = quantizeCols(W1f, d.k, d.h);
|
| 179 |
-
const w2 = quantizeCols(W2f, d.h, d.n);
|
| 180 |
-
if (gpuMlp) {
|
| 181 |
-
const r = await gpuMlp(x.q, w1.q, w2.q, x.s, w1.s, w2.s, d);
|
| 182 |
-
if (audit && audit.due()) {
|
| 183 |
-
// audit BOTH live GEMMs: gemm1 against the units directly; gemm2 by
|
| 184 |
-
// reconstructing its exact operand through the proven quantize mirror
|
| 185 |
-
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);
|
| 186 |
-
if (bad1) { audit.fail("mlp gemm1: " + bad1); return r; }
|
| 187 |
-
const sc = scalesFromAbsMax(rowAbsMax(r.h1, d.m, d.h));
|
| 188 |
-
const hq = quantizeRowsInv(r.h1, d.m, d.h, sc.inv);
|
| 189 |
-
const bad2 = auditTile(hq, w2.q, sc.scale, w2.s, { m: d.m, k: d.h, n: d.n }, r.out, L, audit.cells);
|
| 190 |
-
if (bad2) audit.fail("mlp gemm2: " + bad2);
|
| 191 |
-
}
|
| 192 |
-
return r;
|
| 193 |
-
}
|
| 194 |
-
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);
|
| 195 |
-
const sc = scalesFromAbsMax(rowAbsMax(h1, d.m, d.h));
|
| 196 |
-
const hq = quantizeRowsInv(h1, d.m, d.h, sc.inv);
|
| 197 |
-
const out = bgemmJS(hq, w2.q, sc.scale, w2.s, { m: d.m, k: d.h, n: d.n, batch: 1 }, L);
|
| 198 |
-
return { h1, out };
|
| 199 |
-
}
|
| 200 |
-
|
| 201 |
-
// ---- epilogue mirror -------------------------------------------------------
|
| 202 |
-
// BIT-EXACT mirror of the WGSL epilogue `f32(s) * a * b`. WGSL rounds to f32
|
| 203 |
-
// after the int->float conversion and after EACH multiply; plain JS would do
|
| 204 |
-
// the whole chain in f64 and round once, which differs in the last ulp. That
|
| 205 |
-
// last-ulp gap is what used to force a tolerance into the kernel gates —
|
| 206 |
-
// mirroring the rounding exactly is what lets the gates compare with `!==`.
|
| 207 |
-
const f32 = Math.fround;
|
| 208 |
-
function epi(s, a, b) { return f32(f32(f32(s) * a) * b); }
|
| 209 |
-
|
| 210 |
-
// f32 equality at the BIT level: `!==` says -0 === 0, but replicas are
|
| 211 |
-
// compared by hashing raw bytes, so an audit that can't see the sign of zero
|
| 212 |
-
// could pass a device that later forks the fleet. (Real ISAs have non-IEEE
|
| 213 |
-
// modes that flush -0 to +0 — e.g. RDNA2 output modifiers / legacy muls.)
|
| 214 |
-
const _fb = new Float32Array(1), _ub = new Uint32Array(_fb.buffer);
|
| 215 |
-
function bitDiff(a, b) { _fb[0] = a; const u = _ub[0]; _fb[0] = b; return u !== _ub[0]; }
|
| 216 |
-
|
| 217 |
-
// CPU mirror of the fused GPU kernel: batched int8 GEMM through the LUT with
|
| 218 |
-
// the epilogue (block dequant + optional ReLU) applied before returning —
|
| 219 |
-
// exactly what the WGSL kernel does on-device. d.acc=true returns the raw
|
| 220 |
-
// int32 accumulator instead (the exact oracle the fused kernel normally hides).
|
| 221 |
-
function bgemmJS(Xq, Wq, rs, cs, d, L) {
|
| 222 |
-
const { m, k, n } = d, batch = d.batch || 1, relu = !!d.relu, mul = L.mul;
|
| 223 |
-
const raw = !!d.acc;
|
| 224 |
-
const out = raw ? new Int32Array(batch * m * n) : new Float32Array(batch * m * n);
|
| 225 |
-
const acc = new Int32Array(n);
|
| 226 |
-
for (let bz = 0; bz < batch; bz++) {
|
| 227 |
-
const xo = bz * m * k, wo = bz * k * n, oo = bz * m * n, co = bz * n;
|
| 228 |
-
for (let i = 0; i < m; i++) {
|
| 229 |
-
acc.fill(0);
|
| 230 |
-
const xrow = xo + i * k;
|
| 231 |
-
for (let p = 0; p < k; p++) {
|
| 232 |
-
const au = (Xq[xrow + p] & 0xFF) * 256, wrow = wo + p * n;
|
| 233 |
-
for (let j = 0; j < n; j++) acc[j] += mul[au + (Wq[wrow + j] & 0xFF)];
|
| 234 |
-
}
|
| 235 |
-
const orow = oo + i * n;
|
| 236 |
-
if (raw) { for (let j = 0; j < n; j++) out[orow + j] = acc[j]; continue; }
|
| 237 |
-
const rscale = rs[bz * m + i];
|
| 238 |
-
for (let j = 0; j < n; j++) {
|
| 239 |
-
const v = epi(acc[j], rscale, cs[co + j]);
|
| 240 |
-
out[orow + j] = relu && v < 0 ? 0 : v;
|
| 241 |
-
}
|
| 242 |
-
}
|
| 243 |
-
}
|
| 244 |
-
return out;
|
| 245 |
-
}
|
| 246 |
-
|
| 247 |
-
// Recompute a handful of RANDOM output cells of a live GEMM through the LUT
|
| 248 |
-
// mirror and compare against what the kernel produced. Sampling cells instead
|
| 249 |
-
// of whole matrices makes this cheap enough to run continuously, at the real
|
| 250 |
-
// shapes training uses — not once at boot on toy inputs.
|
| 251 |
-
// STRATIFIED sampling. Uniformly random cells are the wrong instrument for
|
| 252 |
-
// the bugs that actually occur here: a bounds-guard off-by-one or a pack-tail
|
| 253 |
-
// padding bug lives on the LAST row/column, and uniform sampling finds that
|
| 254 |
-
// with probability ~1/n per cell — at the 16512-wide logits GEMM, never. So
|
| 255 |
-
// the first cells are the structurally dangerous ones (corners, last row,
|
| 256 |
-
// last column, last batch) chosen deterministically, and the remainder are
|
| 257 |
-
// random interior cells that catch diffuse bugs. Same principle as poisoning
|
| 258 |
-
// the buffer pool: construct the dangerous case, don't wait to land on it.
|
| 259 |
-
function auditTile(Xq, Wq, rs, cs, d, got, L, nCells) {
|
| 260 |
-
const { m, k, n } = d, batch = d.batch || 1, relu = !!d.relu, mul = L.mul;
|
| 261 |
-
const N = nCells || 8;
|
| 262 |
-
const edges = [[0, m - 1, n - 1], [0, 0, n - 1], [0, m - 1, 0], [0, 0, 0],
|
| 263 |
-
[batch - 1, m - 1, n - 1], [batch - 1, 0, 0]];
|
| 264 |
-
for (let t = 0; t < N; t++) {
|
| 265 |
-
let bz, i, j;
|
| 266 |
-
if (t < edges.length) { bz = edges[t][0]; i = edges[t][1]; j = edges[t][2]; }
|
| 267 |
-
else { bz = (Math.random() * batch) | 0; i = (Math.random() * m) | 0; j = (Math.random() * n) | 0; }
|
| 268 |
-
let acc = 0;
|
| 269 |
-
const xrow = bz * m * k + i * k, wo = bz * k * n;
|
| 270 |
-
for (let p = 0; p < k; p++) acc += mul[(Xq[xrow + p] & 0xFF) * 256 + (Wq[wo + p * n + j] & 0xFF)];
|
| 271 |
-
let v = epi(acc, rs[bz * m + i], cs[bz * n + j]);
|
| 272 |
-
if (relu && v < 0) v = 0;
|
| 273 |
-
const idx = (bz * m + i) * n + j;
|
| 274 |
-
if (bitDiff(got[idx], v))
|
| 275 |
-
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}`;
|
| 276 |
-
}
|
| 277 |
-
return null;
|
| 278 |
-
}
|
| 279 |
-
|
| 280 |
-
// ---- exact mirror of the split-K f32 GEMM ----------------------------------
|
| 281 |
-
// The f32 backward GEMM was the last kernel gated by a TOLERANCE (allclose at
|
| 282 |
-
// 1e-3) — and this project's own gate mutation test shows allclose waving
|
| 283 |
-
// through real bugs. The reason was real though: split-K accumulates in a
|
| 284 |
-
// different ORDER than a naive reference, so bit-equality against the naive
|
| 285 |
-
// one is impossible. The fix is the same as the epilogue mirror: reproduce
|
| 286 |
-
// the kernel's order exactly, then compare with `!==`.
|
| 287 |
-
// partials: for z in 0..S-1, sum p in [z*ks, min(k,(z+1)*ks)) in order
|
| 288 |
-
// reduce: sum the S partials in ascending z
|
| 289 |
-
// `fma` selects the rounding schedule for `s + a*b`: WGSL PERMITS a compiler
|
| 290 |
-
// to contract that into a fused multiply-add (one rounding) instead of two.
|
| 291 |
-
// Which one the device does is a fact about the device, so the gate tries
|
| 292 |
-
// both and reports which matches rather than assuming.
|
| 293 |
-
function fgemmMirror(A, Bm, d, fma) {
|
| 294 |
-
const { m, k, n } = d, transA = !!d.transA;
|
| 295 |
-
const S = k > 4096 ? Math.min(16, Math.ceil(k / 2048)) : 1;
|
| 296 |
-
const ks = Math.ceil(k / S);
|
| 297 |
-
const out = new Float32Array(m * n);
|
| 298 |
-
for (let row = 0; row < m; row++)
|
| 299 |
-
for (let col = 0; col < n; col++) {
|
| 300 |
-
let acc = 0; // reduce pass, ascending z
|
| 301 |
-
for (let z = 0; z < S; z++) {
|
| 302 |
-
const p0 = z * ks, p1 = Math.min(k, p0 + ks);
|
| 303 |
-
let s = 0; // one partial, in order
|
| 304 |
-
for (let p = p0; p < p1; p++) {
|
| 305 |
-
const a = transA ? A[p * m + row] : A[row * k + p];
|
| 306 |
-
s = fma ? f32(s + a * Bm[p * n + col]) // single rounding
|
| 307 |
-
: f32(s + f32(a * Bm[p * n + col]));
|
| 308 |
-
}
|
| 309 |
-
acc = f32(acc + s);
|
| 310 |
-
}
|
| 311 |
-
out[row * n + col] = acc;
|
| 312 |
-
}
|
| 313 |
-
return out;
|
| 314 |
-
}
|
| 315 |
-
|
| 316 |
-
// ---- live audits for the fused attention kernels ---------------------------
|
| 317 |
-
// The attention kernels had exact INIT gates but nothing at live shapes —
|
| 318 |
-
// the exact gap the GEMM audit exists to close, left open on the kernels with
|
| 319 |
-
// the trickiest indexing (head-strided gather, scatter write-back). These
|
| 320 |
-
// recompute individual output cells from the units, stratified like
|
| 321 |
-
// auditTile: last/first token pair, last head, last channel first, then
|
| 322 |
-
// random. Cost is hd (or T) multiply-adds per cell.
|
| 323 |
-
function auditAttScores(qq, kq, qs, ks, d, got, L, nCells) {
|
| 324 |
-
const { B, T, heads, hd } = d, C = heads * hd, mul = L.mul, raw = !!d.acc;
|
| 325 |
-
const N = nCells || 8;
|
| 326 |
-
const edges = [[B - 1, heads - 1, T - 1, T - 1], [0, 0, 0, 0],
|
| 327 |
-
[0, heads - 1, T - 1, 0], [B - 1, 0, 0, T - 1]];
|
| 328 |
-
for (let t = 0; t < N; t++) {
|
| 329 |
-
let bi, h, ti, tj;
|
| 330 |
-
if (t < edges.length) { bi = edges[t][0]; h = edges[t][1]; ti = edges[t][2]; tj = edges[t][3]; }
|
| 331 |
-
else { bi = (Math.random() * B) | 0; h = (Math.random() * heads) | 0;
|
| 332 |
-
ti = (Math.random() * T) | 0; tj = (Math.random() * T) | 0; }
|
| 333 |
-
const bz = bi * heads + h;
|
| 334 |
-
const qo = (bi * T + ti) * C + h * hd, ko = (bi * T + tj) * C + h * hd;
|
| 335 |
-
let acc = 0;
|
| 336 |
-
for (let p = 0; p < hd; p++) acc += mul[(qq[qo + p] & 0xFF) * 256 + (kq[ko + p] & 0xFF)];
|
| 337 |
-
const v = raw ? acc : epi(acc, qs[(bi * T + ti) * heads + h], ks[(bi * T + tj) * heads + h]);
|
| 338 |
-
const idx = (bz * T + ti) * T + tj;
|
| 339 |
-
if (raw ? got[idx] !== v : bitDiff(got[idx], v))
|
| 340 |
-
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}`;
|
| 341 |
-
}
|
| 342 |
-
return null;
|
| 343 |
-
}
|
| 344 |
-
function auditAttCtx(aq, vq, as, vs, d, got, L, nCells) {
|
| 345 |
-
const { B, T, heads, hd } = d, C = heads * hd, mul = L.mul, raw = !!d.acc;
|
| 346 |
-
const N = nCells || 8;
|
| 347 |
-
const edges = [[B - 1, heads - 1, T - 1, hd - 1], [0, 0, 0, 0],
|
| 348 |
-
[0, heads - 1, T - 1, 0], [B - 1, 0, 0, hd - 1]];
|
| 349 |
-
for (let t = 0; t < N; t++) {
|
| 350 |
-
let bi, h, ti, j;
|
| 351 |
-
if (t < edges.length) { bi = edges[t][0]; h = edges[t][1]; ti = edges[t][2]; j = edges[t][3]; }
|
| 352 |
-
else { bi = (Math.random() * B) | 0; h = (Math.random() * heads) | 0;
|
| 353 |
-
ti = (Math.random() * T) | 0; j = (Math.random() * hd) | 0; }
|
| 354 |
-
const bz = bi * heads + h, ao = (bz * T + ti) * T;
|
| 355 |
-
let acc = 0;
|
| 356 |
-
for (let tj = 0; tj < T; tj++)
|
| 357 |
-
acc += mul[(aq[ao + tj] & 0xFF) * 256 + (vq[(bi * T + tj) * C + h * hd + j] & 0xFF)];
|
| 358 |
-
const v = raw ? acc : epi(acc, as[bz * T + ti], vs[(bi * heads + h) * hd + j]);
|
| 359 |
-
const idx = (bi * T + ti) * C + h * hd + j;
|
| 360 |
-
if (raw ? got[idx] !== v : bitDiff(got[idx], v))
|
| 361 |
-
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}`;
|
| 362 |
-
}
|
| 363 |
-
return null;
|
| 364 |
-
}
|
| 365 |
-
|
| 366 |
-
// block-scaled verified GEMM, float in → float out.
|
| 367 |
-
// d = { m, k, n, batch=1, relu=false }; X is (batch·m)×k, W is batch×(k×n)
|
| 368 |
-
// gpuBgemm (from webgpu.js) runs the batched kernel with the fused epilogue;
|
| 369 |
-
// without it the CPU LUT mirror runs. Every product goes through the units.
|
| 370 |
-
async function vgemmBlock(Xf, Wf, d, L, gpuBgemm, audit) {
|
| 371 |
-
const { m, k, n } = d, batch = d.batch || 1;
|
| 372 |
-
const x = quantizeRows(Xf, batch * m, k);
|
| 373 |
-
let wq, ws;
|
| 374 |
-
if (batch === 1) {
|
| 375 |
-
const w = quantizeCols(Wf, k, n); wq = w.q; ws = w.s;
|
| 376 |
-
} else {
|
| 377 |
-
wq = new Int8Array(batch * k * n); ws = new Float32Array(batch * n);
|
| 378 |
-
for (let bz = 0; bz < batch; bz++) {
|
| 379 |
-
const w = quantizeCols(Wf.subarray(bz * k * n, (bz + 1) * k * n), k, n);
|
| 380 |
-
wq.set(w.q, bz * k * n); ws.set(w.s, bz * n);
|
| 381 |
-
}
|
| 382 |
-
}
|
| 383 |
-
if (gpuBgemm) {
|
| 384 |
-
const out = await gpuBgemm(x.q, wq, x.s, ws, d);
|
| 385 |
-
// continuous re-verification at LIVE shapes: the boot gate only ever saw
|
| 386 |
-
// toy inputs, so sample a few real cells against the units as we go
|
| 387 |
-
if (audit && audit.due()) {
|
| 388 |
-
const bad = auditTile(x.q, wq, x.s, ws, d, out, L, audit.cells);
|
| 389 |
-
if (bad) audit.fail(bad);
|
| 390 |
-
}
|
| 391 |
-
return out;
|
| 392 |
-
}
|
| 393 |
-
return bgemmJS(x.q, wq, x.s, ws, d, L);
|
| 394 |
-
}
|
| 395 |
-
|
| 396 |
-
// ---- gather-fused attention through the units (CUTLASS ex. 36/52) ----------
|
| 397 |
-
// The kernels read q/k/v/ctx directly in their natural BT×C layout with
|
| 398 |
-
// head-strided indexing — no JS gather copies, no kᵀ transpose, and the
|
| 399 |
-
// context write scatters straight back into BT×C. Quantization stays
|
| 400 |
-
// block-scaled: q/k/a per (token,head) row, v per (head,channel) column.
|
| 401 |
-
// The (BT·heads)×hd row view of q/k IS the contiguous buffer, so
|
| 402 |
-
// quantizeRows(q, BT·heads, hd) gives per-(token,head) scales for free.
|
| 403 |
-
function quantizeHeadCols(v, B, T, heads, hd) { // per (batch,head,channel) column
|
| 404 |
-
assertFinite(v, "quantizeHeadCols");
|
| 405 |
-
const C = heads * hd;
|
| 406 |
-
const q = new Int8Array(B * T * C), s = new Float32Array(B * heads * hd);
|
| 407 |
-
for (let bi = 0; bi < B; bi++)
|
| 408 |
-
for (let h = 0; h < heads; h++)
|
| 409 |
-
for (let j = 0; j < hd; j++) {
|
| 410 |
-
let mx = 0;
|
| 411 |
-
for (let ti = 0; ti < T; ti++) {
|
| 412 |
-
const a = Math.abs(v[(bi * T + ti) * C + h * hd + j]);
|
| 413 |
-
if (a > mx) mx = a;
|
| 414 |
-
}
|
| 415 |
-
const sc = Math.max(mx / 127, 1e-8);
|
| 416 |
-
s[(bi * heads + h) * hd + j] = sc;
|
| 417 |
-
for (let ti = 0; ti < T; ti++) {
|
| 418 |
-
const idx = (bi * T + ti) * C + h * hd + j;
|
| 419 |
-
const w = Math.round(v[idx] / sc);
|
| 420 |
-
q[idx] = w < -128 ? -128 : w > 127 ? 127 : w;
|
| 421 |
-
}
|
| 422 |
-
}
|
| 423 |
-
return { q, s };
|
| 424 |
-
}
|
| 425 |
-
// scores S[bz,ti,tj] = q_row(bi,ti,h) · k_row(bi,tj,h), every product via the LUT
|
| 426 |
-
// d.acc=true returns the raw int32 accumulator (exact oracle for the kernel gate)
|
| 427 |
-
function attScoresJS(qq, kq, qs, ks, d, L) {
|
| 428 |
-
const { B, T, heads, hd } = d, C = heads * hd, mul = L.mul, raw = !!d.acc;
|
| 429 |
-
const out = raw ? new Int32Array(B * heads * T * T) : new Float32Array(B * heads * T * T);
|
| 430 |
-
for (let bi = 0; bi < B; bi++) for (let h = 0; h < heads; h++) {
|
| 431 |
-
const bz = bi * heads + h;
|
| 432 |
-
for (let ti = 0; ti < T; ti++) {
|
| 433 |
-
const qo = (bi * T + ti) * C + h * hd, rscale = qs[(bi * T + ti) * heads + h];
|
| 434 |
-
for (let tj = 0; tj < T; tj++) {
|
| 435 |
-
const ko = (bi * T + tj) * C + h * hd;
|
| 436 |
-
let acc = 0;
|
| 437 |
-
for (let p = 0; p < hd; p++) acc += mul[(qq[qo + p] & 0xFF) * 256 + (kq[ko + p] & 0xFF)];
|
| 438 |
-
out[(bz * T + ti) * T + tj] = raw ? acc : epi(acc, rscale, ks[(bi * T + tj) * heads + h]);
|
| 439 |
-
}
|
| 440 |
-
}
|
| 441 |
-
}
|
| 442 |
-
return out;
|
| 443 |
-
}
|
| 444 |
-
// ctx[(bi,ti),(h,j)] = Σ_tj a[bz,ti,tj]·v[(bi,tj),(h,j)] — scatter fused into BT×C
|
| 445 |
-
function attCtxJS(aq, vq, as, vs, d, L) {
|
| 446 |
-
const { B, T, heads, hd } = d, C = heads * hd, mul = L.mul, raw = !!d.acc;
|
| 447 |
-
const out = raw ? new Int32Array(B * T * C) : new Float32Array(B * T * C);
|
| 448 |
-
for (let bi = 0; bi < B; bi++) for (let h = 0; h < heads; h++) {
|
| 449 |
-
const bz = bi * heads + h;
|
| 450 |
-
for (let ti = 0; ti < T; ti++) {
|
| 451 |
-
const ao = (bz * T + ti) * T, rscale = as[bz * T + ti];
|
| 452 |
-
for (let j = 0; j < hd; j++) {
|
| 453 |
-
let acc = 0;
|
| 454 |
-
for (let tj = 0; tj < T; tj++)
|
| 455 |
-
acc += mul[(aq[ao + tj] & 0xFF) * 256 + (vq[(bi * T + tj) * C + h * hd + j] & 0xFF)];
|
| 456 |
-
out[(bi * T + ti) * C + h * hd + j] = raw ? acc : epi(acc, rscale, vs[(bi * heads + h) * hd + j]);
|
| 457 |
-
}
|
| 458 |
-
}
|
| 459 |
-
}
|
| 460 |
-
return out;
|
| 461 |
-
}
|
| 462 |
-
|
| 463 |
-
// one verified layer forward; returns float out (+ cache for STE backward).
|
| 464 |
-
// Every product goes through the verified INT8 multiply (mul8 LUT) with exact
|
| 465 |
-
// int32 accumulation — i.e. an emulated INT8 tensor-core GEMM — then dequant.
|
| 466 |
-
async function linearFwd(X, W, m, k, n, L, useRelu, matmulInt8) {
|
| 467 |
-
const xq = quantize(X), wq = quantize(W);
|
| 468 |
-
const acc = await (matmulInt8 || lutMatmulJS)(xq.q, wq.q, m, k, n, L); // verified multiply
|
| 469 |
-
const dq = xq.scale * wq.scale;
|
| 470 |
-
const out = new Float32Array(m * n);
|
| 471 |
-
const mask = useRelu ? new Uint8Array(m * n) : null;
|
| 472 |
-
for (let i = 0; i < m * n; i++) {
|
| 473 |
-
let v = acc[i] * dq;
|
| 474 |
-
if (useRelu) { if (v > 0) mask[i] = 1; else v = 0; }
|
| 475 |
-
out[i] = v;
|
| 476 |
-
}
|
| 477 |
-
return { out, mask };
|
| 478 |
-
}
|
| 479 |
-
|
| 480 |
-
// 2-layer MLP: X→H (relu) →dout. Forward through verified units, MSE loss.
|
| 481 |
-
async function forward(X, y, W1, W2, D, L, matmulInt8) {
|
| 482 |
-
const { n, din, h, dout } = D;
|
| 483 |
-
const l1 = await linearFwd(X, W1, n, din, h, L, true, matmulInt8);
|
| 484 |
-
const l2 = await linearFwd(l1.out, W2, n, h, dout, L, false, matmulInt8);
|
| 485 |
-
const resid = new Float32Array(n * dout); let loss = 0;
|
| 486 |
-
for (let i = 0; i < resid.length; i++) { const r = l2.out[i] - y[i]; resid[i] = r; loss += r * r; }
|
| 487 |
-
loss /= resid.length;
|
| 488 |
-
return { loss, resid, z1: l1.out, mask1: l1.mask };
|
| 489 |
-
}
|
| 490 |
-
|
| 491 |
-
// STE backward (verified matmul treated as float X@W). Returns flat [gW1, gW2].
|
| 492 |
-
function backward(X, W1, W2, fwd, D) {
|
| 493 |
-
const { n, din, h, dout } = D;
|
| 494 |
-
const { resid, z1, mask1 } = fwd;
|
| 495 |
-
const s = 2 / n;
|
| 496 |
-
const dout_ = new Float32Array(resid.length);
|
| 497 |
-
for (let i = 0; i < resid.length; i++) dout_[i] = resid[i] * s;
|
| 498 |
-
const mm = TC.matmul, tr = TC.transpose;
|
| 499 |
-
const gW2 = mm(tr(z1, n, h), dout_, h, n, dout); // z1ᵀ @ dout
|
| 500 |
-
const dz1 = mm(dout_, tr(W2, h, dout), n, dout, h); // dout @ W2ᵀ
|
| 501 |
-
for (let i = 0; i < dz1.length; i++) if (!mask1[i]) dz1[i] = 0; // relu grad
|
| 502 |
-
const gW1 = mm(tr(X, n, din), dz1, din, n, h); // Xᵀ @ dz1
|
| 503 |
-
const g = new Float32Array(gW1.length + gW2.length);
|
| 504 |
-
g.set(gW1, 0); g.set(gW2, gW1.length);
|
| 505 |
-
return g;
|
| 506 |
-
}
|
| 507 |
-
|
| 508 |
-
function splitApply(W1, W2, gAvg, lr) {
|
| 509 |
-
for (let i = 0; i < W1.length; i++) W1[i] -= lr * gAvg[i];
|
| 510 |
-
for (let j = 0; j < W2.length; j++) W2[j] -= lr * gAvg[W1.length + j];
|
| 511 |
-
}
|
| 512 |
-
|
| 513 |
-
const api = { assertFinite, quantize, quantize2, quantizeRows, quantizeCols, quantizeHeadCols, lutMatmulJS, lutMatmul3JS, lutMatmul3,
|
| 514 |
-
bgemmJS, vgemmBlock, auditTile, epi, attScoresJS, attCtxJS, linearFwd, forward, backward, splitApply,
|
| 515 |
-
rowAbsMax, scalesFromAbsMax, quantizeRowsInv, vmlpBlock, bitDiff,
|
| 516 |
-
auditAttScores, auditAttCtx, fgemmMirror };
|
| 517 |
-
if (typeof module !== "undefined" && module.exports) { TC = require("./traincore.js"); module.exports = api; }
|
| 518 |
-
else { TC = root.TrainCore; root.Verified = api; }
|
| 519 |
-
})(typeof self !== "undefined" ? self : this);
|
|
|
|
| 1 |
+
// Verified INT8 compute — the emulated GPU logic, in the browser.
|
| 2 |
+
// A layer's forward runs THROUGH the units: quantize -> LUT multiply -> requant
|
| 3 |
+
// -> optional ReLU -> dequant. Backward is a straight-through estimator (the
|
| 4 |
+
// integer path has no gradient), so ordinary float weights still learn.
|
| 5 |
+
// Same units as the Python/Docker DaisyChain; here they're lookup tables.
|
| 6 |
+
(function (root) {
|
| 7 |
+
"use strict";
|
| 8 |
+
|
| 9 |
+
let TC; // TrainCore (matmul/transpose) — resolved per environment at the end
|
| 10 |
+
|
| 11 |
+
// A float -> int8 conversion has no answer for NaN/Inf, and the quantizers below give a
|
| 12 |
+
// silent wrong one: an Inf makes the |max| scale Infinity and EVERY value quantizes to 0
|
| 13 |
+
// ([0.5, Inf, 3] -> [0, 0, 0]); a NaN is skipped by `a > mx` and then stored as 0 by the
|
| 14 |
+
// Int8Array. The Python trainer refuses the same input (qat.py); so does this. One pass,
|
| 15 |
+
// no effect on finite inputs, so builds with and without it still co-train.
|
| 16 |
+
function assertFinite(X, where) {
|
| 17 |
+
for (let i = 0; i < X.length; i++) {
|
| 18 |
+
if (!Number.isFinite(X[i])) {
|
| 19 |
+
let bad = 0;
|
| 20 |
+
for (let j = i; j < X.length; j++) if (!Number.isFinite(X[j])) bad++;
|
| 21 |
+
const e = new Error("non-finite input to " + where + " (" + bad + " of " + X.length +
|
| 22 |
+
" values): the int8 quantize would turn it into 0");
|
| 23 |
+
e.name = "NonFiniteError";
|
| 24 |
+
throw e;
|
| 25 |
+
}
|
| 26 |
+
}
|
| 27 |
+
}
|
| 28 |
+
|
| 29 |
+
function quantize(X) {
|
| 30 |
+
assertFinite(X, "quantize");
|
| 31 |
+
let mx = 0; for (let i = 0; i < X.length; i++) { const a = Math.abs(X[i]); if (a > mx) mx = a; }
|
| 32 |
+
const scale = Math.max(mx / 127, 1e-8);
|
| 33 |
+
const q = new Int8Array(X.length);
|
| 34 |
+
for (let i = 0; i < X.length; i++) { let v = Math.round(X[i] / scale); q[i] = v < -128 ? -128 : v > 127 ? 127 : v; }
|
| 35 |
+
return { q, scale };
|
| 36 |
+
}
|
| 37 |
+
|
| 38 |
+
// int8 matmul via the verified multiply LUT: acc(m×n) = sum_k mulLUT[Xq,Wq]
|
| 39 |
+
function lutMatmulJS(Xq, Wq, m, k, n, L) {
|
| 40 |
+
const C = new Int32Array(m * n), mul = L.mul;
|
| 41 |
+
for (let i = 0; i < m; i++) {
|
| 42 |
+
for (let p = 0; p < k; p++) {
|
| 43 |
+
const au = (Xq[i * k + p] & 0xFF) * 256, wo = p * n, co = i * n;
|
| 44 |
+
for (let j = 0; j < n; j++) C[co + j] += mul[au + (Wq[wo + j] & 0xFF)];
|
| 45 |
+
}
|
| 46 |
+
}
|
| 47 |
+
return C;
|
| 48 |
+
}
|
| 49 |
+
|
| 50 |
+
// ---- 3xINT8 fast-accurate GEMM --------------------------------------------
|
| 51 |
+
// The CUTLASS example-27 "3xTF32" scheme, ported to the verified units:
|
| 52 |
+
// split each float into a coarse int8 part plus an int8-quantized residual,
|
| 53 |
+
// run three EXACT LUT GEMMs (hi·hi, hi·lo, lo·hi), drop the negligible
|
| 54 |
+
// lo·lo, and recombine. Same big/small decomposition NVIDIA uses to recover
|
| 55 |
+
// near-fp32 accuracy from TF32 tensor cores — here it recovers ~14-bit
|
| 56 |
+
// accuracy from the 8-bit units, at 3× the unit ops. Every product still
|
| 57 |
+
// goes through the verified mul8 LUT.
|
| 58 |
+
function quantize2(X) {
|
| 59 |
+
const hi = quantize(X);
|
| 60 |
+
const r = new Float32Array(X.length);
|
| 61 |
+
for (let i = 0; i < X.length; i++) r[i] = X[i] - hi.q[i] * hi.scale;
|
| 62 |
+
const lo = quantize(r);
|
| 63 |
+
return { hi, lo };
|
| 64 |
+
}
|
| 65 |
+
function combine3(hh, hl, lh, x, w, len) {
|
| 66 |
+
const out = new Float32Array(len);
|
| 67 |
+
const shh = x.hi.scale * w.hi.scale, shl = x.hi.scale * w.lo.scale, slh = x.lo.scale * w.hi.scale;
|
| 68 |
+
for (let i = 0; i < len; i++) out[i] = hh[i] * shh + hl[i] * shl + lh[i] * slh;
|
| 69 |
+
return out;
|
| 70 |
+
}
|
| 71 |
+
function lutMatmul3JS(Xf, Wf, m, k, n, L) { // sync, CPU LUT path
|
| 72 |
+
const x = quantize2(Xf), w = quantize2(Wf);
|
| 73 |
+
return combine3(lutMatmulJS(x.hi.q, w.hi.q, m, k, n, L),
|
| 74 |
+
lutMatmulJS(x.hi.q, w.lo.q, m, k, n, L),
|
| 75 |
+
lutMatmulJS(x.lo.q, w.hi.q, m, k, n, L), x, w, m * n);
|
| 76 |
+
}
|
| 77 |
+
async function lutMatmul3(Xf, Wf, m, k, n, L, matmulInt8) { // any backend
|
| 78 |
+
const x = quantize2(Xf), w = quantize2(Wf);
|
| 79 |
+
const mm = matmulInt8 || lutMatmulJS;
|
| 80 |
+
const [hh, hl, lh] = await Promise.all([
|
| 81 |
+
mm(x.hi.q, w.hi.q, m, k, n, L),
|
| 82 |
+
mm(x.hi.q, w.lo.q, m, k, n, L),
|
| 83 |
+
mm(x.lo.q, w.hi.q, m, k, n, L),
|
| 84 |
+
]);
|
| 85 |
+
return combine3(hh, hl, lh, x, w, m * n);
|
| 86 |
+
}
|
| 87 |
+
|
| 88 |
+
// ---- block-scaled verified GEMM (CUTLASS ex. 67/81 blockwise scaling) ------
|
| 89 |
+
// Per-ROW scales for the activations and per-COLUMN scales for the weights:
|
| 90 |
+
// the integer math through the mul8 LUT is completely unchanged — only the
|
| 91 |
+
// dequant uses rs[row]·cs[col] instead of one tensor-wide product, so a single
|
| 92 |
+
// outlier no longer crushes the quantization resolution of every other
|
| 93 |
+
// row/column. One LUT pass at this granularity beats the per-tensor 3-pass.
|
| 94 |
+
function quantizeRows(X, rows, cols) {
|
| 95 |
+
assertFinite(X, "quantizeRows");
|
| 96 |
+
const q = new Int8Array(rows * cols), s = new Float32Array(rows);
|
| 97 |
+
for (let r = 0; r < rows; r++) {
|
| 98 |
+
let mx = 0;
|
| 99 |
+
for (let c = 0; c < cols; c++) { const a = Math.abs(X[r * cols + c]); if (a > mx) mx = a; }
|
| 100 |
+
const sc = Math.max(mx / 127, 1e-8); s[r] = sc;
|
| 101 |
+
for (let c = 0; c < cols; c++) {
|
| 102 |
+
const v = Math.round(X[r * cols + c] / sc);
|
| 103 |
+
q[r * cols + c] = v < -128 ? -128 : v > 127 ? 127 : v;
|
| 104 |
+
}
|
| 105 |
+
}
|
| 106 |
+
return { q, s };
|
| 107 |
+
}
|
| 108 |
+
function quantizeCols(W, rows, cols) {
|
| 109 |
+
assertFinite(W, "quantizeCols");
|
| 110 |
+
const q = new Int8Array(rows * cols), s = new Float32Array(cols);
|
| 111 |
+
for (let c = 0; c < cols; c++) {
|
| 112 |
+
let mx = 0;
|
| 113 |
+
for (let r = 0; r < rows; r++) { const a = Math.abs(W[r * cols + c]); if (a > mx) mx = a; }
|
| 114 |
+
s[c] = Math.max(mx / 127, 1e-8);
|
| 115 |
+
}
|
| 116 |
+
for (let r = 0; r < rows; r++)
|
| 117 |
+
for (let c = 0; c < cols; c++) {
|
| 118 |
+
const v = Math.round(W[r * cols + c] / s[c]);
|
| 119 |
+
q[r * cols + c] = v < -128 ? -128 : v > 127 ? 127 : v;
|
| 120 |
+
}
|
| 121 |
+
return { q, s };
|
| 122 |
+
}
|
| 123 |
+
|
| 124 |
+
// ---- B2B MLP chain (CUTLASS ex. 13 two-GEMM fusion + ex. 23 epilogue
|
| 125 |
+
// reduction), respecced for cross-device exactness ---------------------------
|
| 126 |
+
// The MLP is the one back-to-back GEMM pair with no layernorm/softmax between
|
| 127 |
+
// (ReLU is already fused in the epilogue), so the intermediate h1 can be
|
| 128 |
+
// quantized ON the GPU and fed straight to the second GEMM. Two rules make
|
| 129 |
+
// that fleet-safe:
|
| 130 |
+
// 1. The per-row |max| (ex. 23) uses only comparisons — exact on any
|
| 131 |
+
// hardware, order-independent — and comes back to JS as ~1KB.
|
| 132 |
+
// 2. Scale DERIVATION (two divisions) stays in JS f64, which IEEE requires
|
| 133 |
+
// to be exactly rounded and is therefore identical on every device.
|
| 134 |
+
// WGSL division is only 2.5 ULP — a fork waiting to happen — but WGSL
|
| 135 |
+
// multiply/add are correctly rounded and floor/clamp are exact. So the
|
| 136 |
+
// quantize step is respecced from round(x / scale) to
|
| 137 |
+
// floor(f32(x * invScale) + 0.5) — floor(x+0.5) IS Math.round's tie
|
| 138 |
+
// rule — and the fround-stepped mirror below is bit-identical to the
|
| 139 |
+
// GPU kernel, which is exact-gated against it at init.
|
| 140 |
+
// NOTE this changes which int8 a value on a rounding boundary lands on
|
| 141 |
+
// (≤1 step) vs quantizeRows, so old and new builds cannot co-train — the
|
| 142 |
+
// divergence guard stops such mixed groups by design.
|
| 143 |
+
function rowAbsMax(X, rows, cols) {
|
| 144 |
+
assertFinite(X, "rowAbsMax");
|
| 145 |
+
const mx = new Float32Array(rows);
|
| 146 |
+
for (let r = 0; r < rows; r++) {
|
| 147 |
+
let m = 0;
|
| 148 |
+
for (let c = 0; c < cols; c++) { const a = Math.abs(X[r * cols + c]); if (a > m) m = a; }
|
| 149 |
+
mx[r] = m;
|
| 150 |
+
}
|
| 151 |
+
return mx;
|
| 152 |
+
}
|
| 153 |
+
function scalesFromAbsMax(mx) { // f64 divisions: exactly rounded, device-identical
|
| 154 |
+
const scale = new Float32Array(mx.length), inv = new Float32Array(mx.length);
|
| 155 |
+
for (let i = 0; i < mx.length; i++) {
|
| 156 |
+
scale[i] = Math.max(mx[i] / 127, 1e-8);
|
| 157 |
+
inv[i] = 1 / scale[i]; // recip of the STORED f32 scale
|
| 158 |
+
}
|
| 159 |
+
return { scale, inv };
|
| 160 |
+
}
|
| 161 |
+
function quantizeRowsInv(X, rows, cols, inv) { // bit-exact mirror of the GPU quantize kernel
|
| 162 |
+
const q = new Int8Array(rows * cols);
|
| 163 |
+
for (let r = 0; r < rows; r++) {
|
| 164 |
+
const iv = inv[r];
|
| 165 |
+
for (let c = 0; c < cols; c++) {
|
| 166 |
+
const n = Math.floor(f32(f32(X[r * cols + c] * iv) + 0.5));
|
| 167 |
+
q[r * cols + c] = n < -128 ? -128 : n > 127 ? 127 : n;
|
| 168 |
+
}
|
| 169 |
+
}
|
| 170 |
+
return q;
|
| 171 |
+
}
|
| 172 |
+
// the chained MLP: X @ W1 -> ReLU (fused) -> absmax -> quantize -> @ W2.
|
| 173 |
+
// d = { m, k, h, n }; gpuMlp (from webgpu.js) runs both GEMMs + the
|
| 174 |
+
// on-GPU quantize with one tiny absmax readback between; without it the CPU
|
| 175 |
+
// mirror chain runs — SAME math, so mixed GPU/CPU fleets stay bit-identical.
|
| 176 |
+
async function vmlpBlock(Xf, W1f, W2f, d, L, gpuMlp, audit) {
|
| 177 |
+
const x = quantizeRows(Xf, d.m, d.k);
|
| 178 |
+
const w1 = quantizeCols(W1f, d.k, d.h);
|
| 179 |
+
const w2 = quantizeCols(W2f, d.h, d.n);
|
| 180 |
+
if (gpuMlp) {
|
| 181 |
+
const r = await gpuMlp(x.q, w1.q, w2.q, x.s, w1.s, w2.s, d);
|
| 182 |
+
if (audit && audit.due()) {
|
| 183 |
+
// audit BOTH live GEMMs: gemm1 against the units directly; gemm2 by
|
| 184 |
+
// reconstructing its exact operand through the proven quantize mirror
|
| 185 |
+
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);
|
| 186 |
+
if (bad1) { audit.fail("mlp gemm1: " + bad1); return r; }
|
| 187 |
+
const sc = scalesFromAbsMax(rowAbsMax(r.h1, d.m, d.h));
|
| 188 |
+
const hq = quantizeRowsInv(r.h1, d.m, d.h, sc.inv);
|
| 189 |
+
const bad2 = auditTile(hq, w2.q, sc.scale, w2.s, { m: d.m, k: d.h, n: d.n }, r.out, L, audit.cells);
|
| 190 |
+
if (bad2) audit.fail("mlp gemm2: " + bad2);
|
| 191 |
+
}
|
| 192 |
+
return r;
|
| 193 |
+
}
|
| 194 |
+
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);
|
| 195 |
+
const sc = scalesFromAbsMax(rowAbsMax(h1, d.m, d.h));
|
| 196 |
+
const hq = quantizeRowsInv(h1, d.m, d.h, sc.inv);
|
| 197 |
+
const out = bgemmJS(hq, w2.q, sc.scale, w2.s, { m: d.m, k: d.h, n: d.n, batch: 1 }, L);
|
| 198 |
+
return { h1, out };
|
| 199 |
+
}
|
| 200 |
+
|
| 201 |
+
// ---- epilogue mirror -------------------------------------------------------
|
| 202 |
+
// BIT-EXACT mirror of the WGSL epilogue `f32(s) * a * b`. WGSL rounds to f32
|
| 203 |
+
// after the int->float conversion and after EACH multiply; plain JS would do
|
| 204 |
+
// the whole chain in f64 and round once, which differs in the last ulp. That
|
| 205 |
+
// last-ulp gap is what used to force a tolerance into the kernel gates —
|
| 206 |
+
// mirroring the rounding exactly is what lets the gates compare with `!==`.
|
| 207 |
+
const f32 = Math.fround;
|
| 208 |
+
function epi(s, a, b) { return f32(f32(f32(s) * a) * b); }
|
| 209 |
+
|
| 210 |
+
// f32 equality at the BIT level: `!==` says -0 === 0, but replicas are
|
| 211 |
+
// compared by hashing raw bytes, so an audit that can't see the sign of zero
|
| 212 |
+
// could pass a device that later forks the fleet. (Real ISAs have non-IEEE
|
| 213 |
+
// modes that flush -0 to +0 — e.g. RDNA2 output modifiers / legacy muls.)
|
| 214 |
+
const _fb = new Float32Array(1), _ub = new Uint32Array(_fb.buffer);
|
| 215 |
+
function bitDiff(a, b) { _fb[0] = a; const u = _ub[0]; _fb[0] = b; return u !== _ub[0]; }
|
| 216 |
+
|
| 217 |
+
// CPU mirror of the fused GPU kernel: batched int8 GEMM through the LUT with
|
| 218 |
+
// the epilogue (block dequant + optional ReLU) applied before returning —
|
| 219 |
+
// exactly what the WGSL kernel does on-device. d.acc=true returns the raw
|
| 220 |
+
// int32 accumulator instead (the exact oracle the fused kernel normally hides).
|
| 221 |
+
function bgemmJS(Xq, Wq, rs, cs, d, L) {
|
| 222 |
+
const { m, k, n } = d, batch = d.batch || 1, relu = !!d.relu, mul = L.mul;
|
| 223 |
+
const raw = !!d.acc;
|
| 224 |
+
const out = raw ? new Int32Array(batch * m * n) : new Float32Array(batch * m * n);
|
| 225 |
+
const acc = new Int32Array(n);
|
| 226 |
+
for (let bz = 0; bz < batch; bz++) {
|
| 227 |
+
const xo = bz * m * k, wo = bz * k * n, oo = bz * m * n, co = bz * n;
|
| 228 |
+
for (let i = 0; i < m; i++) {
|
| 229 |
+
acc.fill(0);
|
| 230 |
+
const xrow = xo + i * k;
|
| 231 |
+
for (let p = 0; p < k; p++) {
|
| 232 |
+
const au = (Xq[xrow + p] & 0xFF) * 256, wrow = wo + p * n;
|
| 233 |
+
for (let j = 0; j < n; j++) acc[j] += mul[au + (Wq[wrow + j] & 0xFF)];
|
| 234 |
+
}
|
| 235 |
+
const orow = oo + i * n;
|
| 236 |
+
if (raw) { for (let j = 0; j < n; j++) out[orow + j] = acc[j]; continue; }
|
| 237 |
+
const rscale = rs[bz * m + i];
|
| 238 |
+
for (let j = 0; j < n; j++) {
|
| 239 |
+
const v = epi(acc[j], rscale, cs[co + j]);
|
| 240 |
+
out[orow + j] = relu && v < 0 ? 0 : v;
|
| 241 |
+
}
|
| 242 |
+
}
|
| 243 |
+
}
|
| 244 |
+
return out;
|
| 245 |
+
}
|
| 246 |
+
|
| 247 |
+
// Recompute a handful of RANDOM output cells of a live GEMM through the LUT
|
| 248 |
+
// mirror and compare against what the kernel produced. Sampling cells instead
|
| 249 |
+
// of whole matrices makes this cheap enough to run continuously, at the real
|
| 250 |
+
// shapes training uses — not once at boot on toy inputs.
|
| 251 |
+
// STRATIFIED sampling. Uniformly random cells are the wrong instrument for
|
| 252 |
+
// the bugs that actually occur here: a bounds-guard off-by-one or a pack-tail
|
| 253 |
+
// padding bug lives on the LAST row/column, and uniform sampling finds that
|
| 254 |
+
// with probability ~1/n per cell — at the 16512-wide logits GEMM, never. So
|
| 255 |
+
// the first cells are the structurally dangerous ones (corners, last row,
|
| 256 |
+
// last column, last batch) chosen deterministically, and the remainder are
|
| 257 |
+
// random interior cells that catch diffuse bugs. Same principle as poisoning
|
| 258 |
+
// the buffer pool: construct the dangerous case, don't wait to land on it.
|
| 259 |
+
function auditTile(Xq, Wq, rs, cs, d, got, L, nCells) {
|
| 260 |
+
const { m, k, n } = d, batch = d.batch || 1, relu = !!d.relu, mul = L.mul;
|
| 261 |
+
const N = nCells || 8;
|
| 262 |
+
const edges = [[0, m - 1, n - 1], [0, 0, n - 1], [0, m - 1, 0], [0, 0, 0],
|
| 263 |
+
[batch - 1, m - 1, n - 1], [batch - 1, 0, 0]];
|
| 264 |
+
for (let t = 0; t < N; t++) {
|
| 265 |
+
let bz, i, j;
|
| 266 |
+
if (t < edges.length) { bz = edges[t][0]; i = edges[t][1]; j = edges[t][2]; }
|
| 267 |
+
else { bz = (Math.random() * batch) | 0; i = (Math.random() * m) | 0; j = (Math.random() * n) | 0; }
|
| 268 |
+
let acc = 0;
|
| 269 |
+
const xrow = bz * m * k + i * k, wo = bz * k * n;
|
| 270 |
+
for (let p = 0; p < k; p++) acc += mul[(Xq[xrow + p] & 0xFF) * 256 + (Wq[wo + p * n + j] & 0xFF)];
|
| 271 |
+
let v = epi(acc, rs[bz * m + i], cs[bz * n + j]);
|
| 272 |
+
if (relu && v < 0) v = 0;
|
| 273 |
+
const idx = (bz * m + i) * n + j;
|
| 274 |
+
if (bitDiff(got[idx], v))
|
| 275 |
+
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}`;
|
| 276 |
+
}
|
| 277 |
+
return null;
|
| 278 |
+
}
|
| 279 |
+
|
| 280 |
+
// ---- exact mirror of the split-K f32 GEMM ----------------------------------
|
| 281 |
+
// The f32 backward GEMM was the last kernel gated by a TOLERANCE (allclose at
|
| 282 |
+
// 1e-3) — and this project's own gate mutation test shows allclose waving
|
| 283 |
+
// through real bugs. The reason was real though: split-K accumulates in a
|
| 284 |
+
// different ORDER than a naive reference, so bit-equality against the naive
|
| 285 |
+
// one is impossible. The fix is the same as the epilogue mirror: reproduce
|
| 286 |
+
// the kernel's order exactly, then compare with `!==`.
|
| 287 |
+
// partials: for z in 0..S-1, sum p in [z*ks, min(k,(z+1)*ks)) in order
|
| 288 |
+
// reduce: sum the S partials in ascending z
|
| 289 |
+
// `fma` selects the rounding schedule for `s + a*b`: WGSL PERMITS a compiler
|
| 290 |
+
// to contract that into a fused multiply-add (one rounding) instead of two.
|
| 291 |
+
// Which one the device does is a fact about the device, so the gate tries
|
| 292 |
+
// both and reports which matches rather than assuming.
|
| 293 |
+
function fgemmMirror(A, Bm, d, fma) {
|
| 294 |
+
const { m, k, n } = d, transA = !!d.transA;
|
| 295 |
+
const S = k > 4096 ? Math.min(16, Math.ceil(k / 2048)) : 1;
|
| 296 |
+
const ks = Math.ceil(k / S);
|
| 297 |
+
const out = new Float32Array(m * n);
|
| 298 |
+
for (let row = 0; row < m; row++)
|
| 299 |
+
for (let col = 0; col < n; col++) {
|
| 300 |
+
let acc = 0; // reduce pass, ascending z
|
| 301 |
+
for (let z = 0; z < S; z++) {
|
| 302 |
+
const p0 = z * ks, p1 = Math.min(k, p0 + ks);
|
| 303 |
+
let s = 0; // one partial, in order
|
| 304 |
+
for (let p = p0; p < p1; p++) {
|
| 305 |
+
const a = transA ? A[p * m + row] : A[row * k + p];
|
| 306 |
+
s = fma ? f32(s + a * Bm[p * n + col]) // single rounding
|
| 307 |
+
: f32(s + f32(a * Bm[p * n + col]));
|
| 308 |
+
}
|
| 309 |
+
acc = f32(acc + s);
|
| 310 |
+
}
|
| 311 |
+
out[row * n + col] = acc;
|
| 312 |
+
}
|
| 313 |
+
return out;
|
| 314 |
+
}
|
| 315 |
+
|
| 316 |
+
// ---- live audits for the fused attention kernels ---------------------------
|
| 317 |
+
// The attention kernels had exact INIT gates but nothing at live shapes —
|
| 318 |
+
// the exact gap the GEMM audit exists to close, left open on the kernels with
|
| 319 |
+
// the trickiest indexing (head-strided gather, scatter write-back). These
|
| 320 |
+
// recompute individual output cells from the units, stratified like
|
| 321 |
+
// auditTile: last/first token pair, last head, last channel first, then
|
| 322 |
+
// random. Cost is hd (or T) multiply-adds per cell.
|
| 323 |
+
function auditAttScores(qq, kq, qs, ks, d, got, L, nCells) {
|
| 324 |
+
const { B, T, heads, hd } = d, C = heads * hd, mul = L.mul, raw = !!d.acc;
|
| 325 |
+
const N = nCells || 8;
|
| 326 |
+
const edges = [[B - 1, heads - 1, T - 1, T - 1], [0, 0, 0, 0],
|
| 327 |
+
[0, heads - 1, T - 1, 0], [B - 1, 0, 0, T - 1]];
|
| 328 |
+
for (let t = 0; t < N; t++) {
|
| 329 |
+
let bi, h, ti, tj;
|
| 330 |
+
if (t < edges.length) { bi = edges[t][0]; h = edges[t][1]; ti = edges[t][2]; tj = edges[t][3]; }
|
| 331 |
+
else { bi = (Math.random() * B) | 0; h = (Math.random() * heads) | 0;
|
| 332 |
+
ti = (Math.random() * T) | 0; tj = (Math.random() * T) | 0; }
|
| 333 |
+
const bz = bi * heads + h;
|
| 334 |
+
const qo = (bi * T + ti) * C + h * hd, ko = (bi * T + tj) * C + h * hd;
|
| 335 |
+
let acc = 0;
|
| 336 |
+
for (let p = 0; p < hd; p++) acc += mul[(qq[qo + p] & 0xFF) * 256 + (kq[ko + p] & 0xFF)];
|
| 337 |
+
const v = raw ? acc : epi(acc, qs[(bi * T + ti) * heads + h], ks[(bi * T + tj) * heads + h]);
|
| 338 |
+
const idx = (bz * T + ti) * T + tj;
|
| 339 |
+
if (raw ? got[idx] !== v : bitDiff(got[idx], v))
|
| 340 |
+
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}`;
|
| 341 |
+
}
|
| 342 |
+
return null;
|
| 343 |
+
}
|
| 344 |
+
function auditAttCtx(aq, vq, as, vs, d, got, L, nCells) {
|
| 345 |
+
const { B, T, heads, hd } = d, C = heads * hd, mul = L.mul, raw = !!d.acc;
|
| 346 |
+
const N = nCells || 8;
|
| 347 |
+
const edges = [[B - 1, heads - 1, T - 1, hd - 1], [0, 0, 0, 0],
|
| 348 |
+
[0, heads - 1, T - 1, 0], [B - 1, 0, 0, hd - 1]];
|
| 349 |
+
for (let t = 0; t < N; t++) {
|
| 350 |
+
let bi, h, ti, j;
|
| 351 |
+
if (t < edges.length) { bi = edges[t][0]; h = edges[t][1]; ti = edges[t][2]; j = edges[t][3]; }
|
| 352 |
+
else { bi = (Math.random() * B) | 0; h = (Math.random() * heads) | 0;
|
| 353 |
+
ti = (Math.random() * T) | 0; j = (Math.random() * hd) | 0; }
|
| 354 |
+
const bz = bi * heads + h, ao = (bz * T + ti) * T;
|
| 355 |
+
let acc = 0;
|
| 356 |
+
for (let tj = 0; tj < T; tj++)
|
| 357 |
+
acc += mul[(aq[ao + tj] & 0xFF) * 256 + (vq[(bi * T + tj) * C + h * hd + j] & 0xFF)];
|
| 358 |
+
const v = raw ? acc : epi(acc, as[bz * T + ti], vs[(bi * heads + h) * hd + j]);
|
| 359 |
+
const idx = (bi * T + ti) * C + h * hd + j;
|
| 360 |
+
if (raw ? got[idx] !== v : bitDiff(got[idx], v))
|
| 361 |
+
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}`;
|
| 362 |
+
}
|
| 363 |
+
return null;
|
| 364 |
+
}
|
| 365 |
+
|
| 366 |
+
// block-scaled verified GEMM, float in → float out.
|
| 367 |
+
// d = { m, k, n, batch=1, relu=false }; X is (batch·m)×k, W is batch×(k×n)
|
| 368 |
+
// gpuBgemm (from webgpu.js) runs the batched kernel with the fused epilogue;
|
| 369 |
+
// without it the CPU LUT mirror runs. Every product goes through the units.
|
| 370 |
+
async function vgemmBlock(Xf, Wf, d, L, gpuBgemm, audit) {
|
| 371 |
+
const { m, k, n } = d, batch = d.batch || 1;
|
| 372 |
+
const x = quantizeRows(Xf, batch * m, k);
|
| 373 |
+
let wq, ws;
|
| 374 |
+
if (batch === 1) {
|
| 375 |
+
const w = quantizeCols(Wf, k, n); wq = w.q; ws = w.s;
|
| 376 |
+
} else {
|
| 377 |
+
wq = new Int8Array(batch * k * n); ws = new Float32Array(batch * n);
|
| 378 |
+
for (let bz = 0; bz < batch; bz++) {
|
| 379 |
+
const w = quantizeCols(Wf.subarray(bz * k * n, (bz + 1) * k * n), k, n);
|
| 380 |
+
wq.set(w.q, bz * k * n); ws.set(w.s, bz * n);
|
| 381 |
+
}
|
| 382 |
+
}
|
| 383 |
+
if (gpuBgemm) {
|
| 384 |
+
const out = await gpuBgemm(x.q, wq, x.s, ws, d);
|
| 385 |
+
// continuous re-verification at LIVE shapes: the boot gate only ever saw
|
| 386 |
+
// toy inputs, so sample a few real cells against the units as we go
|
| 387 |
+
if (audit && audit.due()) {
|
| 388 |
+
const bad = auditTile(x.q, wq, x.s, ws, d, out, L, audit.cells);
|
| 389 |
+
if (bad) audit.fail(bad);
|
| 390 |
+
}
|
| 391 |
+
return out;
|
| 392 |
+
}
|
| 393 |
+
return bgemmJS(x.q, wq, x.s, ws, d, L);
|
| 394 |
+
}
|
| 395 |
+
|
| 396 |
+
// ---- gather-fused attention through the units (CUTLASS ex. 36/52) ----------
|
| 397 |
+
// The kernels read q/k/v/ctx directly in their natural BT×C layout with
|
| 398 |
+
// head-strided indexing — no JS gather copies, no kᵀ transpose, and the
|
| 399 |
+
// context write scatters straight back into BT×C. Quantization stays
|
| 400 |
+
// block-scaled: q/k/a per (token,head) row, v per (head,channel) column.
|
| 401 |
+
// The (BT·heads)×hd row view of q/k IS the contiguous buffer, so
|
| 402 |
+
// quantizeRows(q, BT·heads, hd) gives per-(token,head) scales for free.
|
| 403 |
+
function quantizeHeadCols(v, B, T, heads, hd) { // per (batch,head,channel) column
|
| 404 |
+
assertFinite(v, "quantizeHeadCols");
|
| 405 |
+
const C = heads * hd;
|
| 406 |
+
const q = new Int8Array(B * T * C), s = new Float32Array(B * heads * hd);
|
| 407 |
+
for (let bi = 0; bi < B; bi++)
|
| 408 |
+
for (let h = 0; h < heads; h++)
|
| 409 |
+
for (let j = 0; j < hd; j++) {
|
| 410 |
+
let mx = 0;
|
| 411 |
+
for (let ti = 0; ti < T; ti++) {
|
| 412 |
+
const a = Math.abs(v[(bi * T + ti) * C + h * hd + j]);
|
| 413 |
+
if (a > mx) mx = a;
|
| 414 |
+
}
|
| 415 |
+
const sc = Math.max(mx / 127, 1e-8);
|
| 416 |
+
s[(bi * heads + h) * hd + j] = sc;
|
| 417 |
+
for (let ti = 0; ti < T; ti++) {
|
| 418 |
+
const idx = (bi * T + ti) * C + h * hd + j;
|
| 419 |
+
const w = Math.round(v[idx] / sc);
|
| 420 |
+
q[idx] = w < -128 ? -128 : w > 127 ? 127 : w;
|
| 421 |
+
}
|
| 422 |
+
}
|
| 423 |
+
return { q, s };
|
| 424 |
+
}
|
| 425 |
+
// scores S[bz,ti,tj] = q_row(bi,ti,h) · k_row(bi,tj,h), every product via the LUT
|
| 426 |
+
// d.acc=true returns the raw int32 accumulator (exact oracle for the kernel gate)
|
| 427 |
+
function attScoresJS(qq, kq, qs, ks, d, L) {
|
| 428 |
+
const { B, T, heads, hd } = d, C = heads * hd, mul = L.mul, raw = !!d.acc;
|
| 429 |
+
const out = raw ? new Int32Array(B * heads * T * T) : new Float32Array(B * heads * T * T);
|
| 430 |
+
for (let bi = 0; bi < B; bi++) for (let h = 0; h < heads; h++) {
|
| 431 |
+
const bz = bi * heads + h;
|
| 432 |
+
for (let ti = 0; ti < T; ti++) {
|
| 433 |
+
const qo = (bi * T + ti) * C + h * hd, rscale = qs[(bi * T + ti) * heads + h];
|
| 434 |
+
for (let tj = 0; tj < T; tj++) {
|
| 435 |
+
const ko = (bi * T + tj) * C + h * hd;
|
| 436 |
+
let acc = 0;
|
| 437 |
+
for (let p = 0; p < hd; p++) acc += mul[(qq[qo + p] & 0xFF) * 256 + (kq[ko + p] & 0xFF)];
|
| 438 |
+
out[(bz * T + ti) * T + tj] = raw ? acc : epi(acc, rscale, ks[(bi * T + tj) * heads + h]);
|
| 439 |
+
}
|
| 440 |
+
}
|
| 441 |
+
}
|
| 442 |
+
return out;
|
| 443 |
+
}
|
| 444 |
+
// ctx[(bi,ti),(h,j)] = Σ_tj a[bz,ti,tj]·v[(bi,tj),(h,j)] — scatter fused into BT×C
|
| 445 |
+
function attCtxJS(aq, vq, as, vs, d, L) {
|
| 446 |
+
const { B, T, heads, hd } = d, C = heads * hd, mul = L.mul, raw = !!d.acc;
|
| 447 |
+
const out = raw ? new Int32Array(B * T * C) : new Float32Array(B * T * C);
|
| 448 |
+
for (let bi = 0; bi < B; bi++) for (let h = 0; h < heads; h++) {
|
| 449 |
+
const bz = bi * heads + h;
|
| 450 |
+
for (let ti = 0; ti < T; ti++) {
|
| 451 |
+
const ao = (bz * T + ti) * T, rscale = as[bz * T + ti];
|
| 452 |
+
for (let j = 0; j < hd; j++) {
|
| 453 |
+
let acc = 0;
|
| 454 |
+
for (let tj = 0; tj < T; tj++)
|
| 455 |
+
acc += mul[(aq[ao + tj] & 0xFF) * 256 + (vq[(bi * T + tj) * C + h * hd + j] & 0xFF)];
|
| 456 |
+
out[(bi * T + ti) * C + h * hd + j] = raw ? acc : epi(acc, rscale, vs[(bi * heads + h) * hd + j]);
|
| 457 |
+
}
|
| 458 |
+
}
|
| 459 |
+
}
|
| 460 |
+
return out;
|
| 461 |
+
}
|
| 462 |
+
|
| 463 |
+
// one verified layer forward; returns float out (+ cache for STE backward).
|
| 464 |
+
// Every product goes through the verified INT8 multiply (mul8 LUT) with exact
|
| 465 |
+
// int32 accumulation — i.e. an emulated INT8 tensor-core GEMM — then dequant.
|
| 466 |
+
async function linearFwd(X, W, m, k, n, L, useRelu, matmulInt8) {
|
| 467 |
+
const xq = quantize(X), wq = quantize(W);
|
| 468 |
+
const acc = await (matmulInt8 || lutMatmulJS)(xq.q, wq.q, m, k, n, L); // verified multiply
|
| 469 |
+
const dq = xq.scale * wq.scale;
|
| 470 |
+
const out = new Float32Array(m * n);
|
| 471 |
+
const mask = useRelu ? new Uint8Array(m * n) : null;
|
| 472 |
+
for (let i = 0; i < m * n; i++) {
|
| 473 |
+
let v = acc[i] * dq;
|
| 474 |
+
if (useRelu) { if (v > 0) mask[i] = 1; else v = 0; }
|
| 475 |
+
out[i] = v;
|
| 476 |
+
}
|
| 477 |
+
return { out, mask };
|
| 478 |
+
}
|
| 479 |
+
|
| 480 |
+
// 2-layer MLP: X→H (relu) →dout. Forward through verified units, MSE loss.
|
| 481 |
+
async function forward(X, y, W1, W2, D, L, matmulInt8) {
|
| 482 |
+
const { n, din, h, dout } = D;
|
| 483 |
+
const l1 = await linearFwd(X, W1, n, din, h, L, true, matmulInt8);
|
| 484 |
+
const l2 = await linearFwd(l1.out, W2, n, h, dout, L, false, matmulInt8);
|
| 485 |
+
const resid = new Float32Array(n * dout); let loss = 0;
|
| 486 |
+
for (let i = 0; i < resid.length; i++) { const r = l2.out[i] - y[i]; resid[i] = r; loss += r * r; }
|
| 487 |
+
loss /= resid.length;
|
| 488 |
+
return { loss, resid, z1: l1.out, mask1: l1.mask };
|
| 489 |
+
}
|
| 490 |
+
|
| 491 |
+
// STE backward (verified matmul treated as float X@W). Returns flat [gW1, gW2].
|
| 492 |
+
function backward(X, W1, W2, fwd, D) {
|
| 493 |
+
const { n, din, h, dout } = D;
|
| 494 |
+
const { resid, z1, mask1 } = fwd;
|
| 495 |
+
const s = 2 / n;
|
| 496 |
+
const dout_ = new Float32Array(resid.length);
|
| 497 |
+
for (let i = 0; i < resid.length; i++) dout_[i] = resid[i] * s;
|
| 498 |
+
const mm = TC.matmul, tr = TC.transpose;
|
| 499 |
+
const gW2 = mm(tr(z1, n, h), dout_, h, n, dout); // z1ᵀ @ dout
|
| 500 |
+
const dz1 = mm(dout_, tr(W2, h, dout), n, dout, h); // dout @ W2ᵀ
|
| 501 |
+
for (let i = 0; i < dz1.length; i++) if (!mask1[i]) dz1[i] = 0; // relu grad
|
| 502 |
+
const gW1 = mm(tr(X, n, din), dz1, din, n, h); // Xᵀ @ dz1
|
| 503 |
+
const g = new Float32Array(gW1.length + gW2.length);
|
| 504 |
+
g.set(gW1, 0); g.set(gW2, gW1.length);
|
| 505 |
+
return g;
|
| 506 |
+
}
|
| 507 |
+
|
| 508 |
+
function splitApply(W1, W2, gAvg, lr) {
|
| 509 |
+
for (let i = 0; i < W1.length; i++) W1[i] -= lr * gAvg[i];
|
| 510 |
+
for (let j = 0; j < W2.length; j++) W2[j] -= lr * gAvg[W1.length + j];
|
| 511 |
+
}
|
| 512 |
+
|
| 513 |
+
const api = { assertFinite, quantize, quantize2, quantizeRows, quantizeCols, quantizeHeadCols, lutMatmulJS, lutMatmul3JS, lutMatmul3,
|
| 514 |
+
bgemmJS, vgemmBlock, auditTile, epi, attScoresJS, attCtxJS, linearFwd, forward, backward, splitApply,
|
| 515 |
+
rowAbsMax, scalesFromAbsMax, quantizeRowsInv, vmlpBlock, bitDiff,
|
| 516 |
+
auditAttScores, auditAttCtx, fgemmMirror };
|
| 517 |
+
if (typeof module !== "undefined" && module.exports) { TC = require("./traincore.js"); module.exports = api; }
|
| 518 |
+
else { TC = root.TrainCore; root.Verified = api; }
|
| 519 |
+
})(typeof self !== "undefined" ? self : this);
|