| |
| |
| |
| |
| |
| (function (root) { |
| "use strict"; |
|
|
| let TC; |
|
|
| function quantize(X) { |
| let mx = 0; for (let i = 0; i < X.length; i++) { const a = Math.abs(X[i]); if (a > mx) mx = a; } |
| const scale = Math.max(mx / 127, 1e-8); |
| const q = new Int8Array(X.length); |
| for (let i = 0; i < X.length; i++) { let v = Math.round(X[i] / scale); q[i] = v < -128 ? -128 : v > 127 ? 127 : v; } |
| return { q, scale }; |
| } |
|
|
| |
| function lutMatmulJS(Xq, Wq, m, k, n, L) { |
| const C = new Int32Array(m * n), mul = L.mul; |
| for (let i = 0; i < m; i++) { |
| for (let p = 0; p < k; p++) { |
| const au = (Xq[i * k + p] & 0xFF) * 256, wo = p * n, co = i * n; |
| for (let j = 0; j < n; j++) C[co + j] += mul[au + (Wq[wo + j] & 0xFF)]; |
| } |
| } |
| return C; |
| } |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| function quantize2(X) { |
| const hi = quantize(X); |
| const r = new Float32Array(X.length); |
| for (let i = 0; i < X.length; i++) r[i] = X[i] - hi.q[i] * hi.scale; |
| const lo = quantize(r); |
| return { hi, lo }; |
| } |
| function combine3(hh, hl, lh, x, w, len) { |
| const out = new Float32Array(len); |
| const shh = x.hi.scale * w.hi.scale, shl = x.hi.scale * w.lo.scale, slh = x.lo.scale * w.hi.scale; |
| for (let i = 0; i < len; i++) out[i] = hh[i] * shh + hl[i] * shl + lh[i] * slh; |
| return out; |
| } |
| function lutMatmul3JS(Xf, Wf, m, k, n, L) { |
| const x = quantize2(Xf), w = quantize2(Wf); |
| return combine3(lutMatmulJS(x.hi.q, w.hi.q, m, k, n, L), |
| lutMatmulJS(x.hi.q, w.lo.q, m, k, n, L), |
| lutMatmulJS(x.lo.q, w.hi.q, m, k, n, L), x, w, m * n); |
| } |
| async function lutMatmul3(Xf, Wf, m, k, n, L, matmulInt8) { |
| const x = quantize2(Xf), w = quantize2(Wf); |
| const mm = matmulInt8 || lutMatmulJS; |
| const [hh, hl, lh] = await Promise.all([ |
| mm(x.hi.q, w.hi.q, m, k, n, L), |
| mm(x.hi.q, w.lo.q, m, k, n, L), |
| mm(x.lo.q, w.hi.q, m, k, n, L), |
| ]); |
| return combine3(hh, hl, lh, x, w, m * n); |
| } |
|
|
| |
| |
| |
| |
| |
| |
| function quantizeRows(X, rows, cols) { |
| const q = new Int8Array(rows * cols), s = new Float32Array(rows); |
| for (let r = 0; r < rows; r++) { |
| let mx = 0; |
| for (let c = 0; c < cols; c++) { const a = Math.abs(X[r * cols + c]); if (a > mx) mx = a; } |
| const sc = Math.max(mx / 127, 1e-8); s[r] = sc; |
| for (let c = 0; c < cols; c++) { |
| const v = Math.round(X[r * cols + c] / sc); |
| q[r * cols + c] = v < -128 ? -128 : v > 127 ? 127 : v; |
| } |
| } |
| return { q, s }; |
| } |
| function quantizeCols(W, rows, cols) { |
| const q = new Int8Array(rows * cols), s = new Float32Array(cols); |
| for (let c = 0; c < cols; c++) { |
| let mx = 0; |
| for (let r = 0; r < rows; r++) { const a = Math.abs(W[r * cols + c]); if (a > mx) mx = a; } |
| s[c] = Math.max(mx / 127, 1e-8); |
| } |
| for (let r = 0; r < rows; r++) |
| for (let c = 0; c < cols; c++) { |
| const v = Math.round(W[r * cols + c] / s[c]); |
| q[r * cols + c] = v < -128 ? -128 : v > 127 ? 127 : v; |
| } |
| return { q, s }; |
| } |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| function rowAbsMax(X, rows, cols) { |
| const mx = new Float32Array(rows); |
| for (let r = 0; r < rows; r++) { |
| let m = 0; |
| for (let c = 0; c < cols; c++) { const a = Math.abs(X[r * cols + c]); if (a > m) m = a; } |
| mx[r] = m; |
| } |
| return mx; |
| } |
| function scalesFromAbsMax(mx) { |
| const scale = new Float32Array(mx.length), inv = new Float32Array(mx.length); |
| for (let i = 0; i < mx.length; i++) { |
| scale[i] = Math.max(mx[i] / 127, 1e-8); |
| inv[i] = 1 / scale[i]; |
| } |
| return { scale, inv }; |
| } |
| function quantizeRowsInv(X, rows, cols, inv) { |
| const q = new Int8Array(rows * cols); |
| for (let r = 0; r < rows; r++) { |
| const iv = inv[r]; |
| for (let c = 0; c < cols; c++) { |
| const n = Math.floor(f32(f32(X[r * cols + c] * iv) + 0.5)); |
| q[r * cols + c] = n < -128 ? -128 : n > 127 ? 127 : n; |
| } |
| } |
| return q; |
| } |
| |
| |
| |
| |
| async function vmlpBlock(Xf, W1f, W2f, d, L, gpuMlp, audit) { |
| const x = quantizeRows(Xf, d.m, d.k); |
| const w1 = quantizeCols(W1f, d.k, d.h); |
| const w2 = quantizeCols(W2f, d.h, d.n); |
| if (gpuMlp) { |
| const r = await gpuMlp(x.q, w1.q, w2.q, x.s, w1.s, w2.s, d); |
| if (audit && audit.due()) { |
| |
| |
| const bad1 = auditTile(x.q, w1.q, x.s, w1.s, { m: d.m, k: d.k, n: d.h, relu: true }, r.h1, L, audit.cells); |
| if (bad1) { audit.fail("mlp gemm1: " + bad1); return r; } |
| const sc = scalesFromAbsMax(rowAbsMax(r.h1, d.m, d.h)); |
| const hq = quantizeRowsInv(r.h1, d.m, d.h, sc.inv); |
| const bad2 = auditTile(hq, w2.q, sc.scale, w2.s, { m: d.m, k: d.h, n: d.n }, r.out, L, audit.cells); |
| if (bad2) audit.fail("mlp gemm2: " + bad2); |
| } |
| return r; |
| } |
| const h1 = bgemmJS(x.q, w1.q, x.s, w1.s, { m: d.m, k: d.k, n: d.h, batch: 1, relu: true }, L); |
| const sc = scalesFromAbsMax(rowAbsMax(h1, d.m, d.h)); |
| const hq = quantizeRowsInv(h1, d.m, d.h, sc.inv); |
| const out = bgemmJS(hq, w2.q, sc.scale, w2.s, { m: d.m, k: d.h, n: d.n, batch: 1 }, L); |
| return { h1, out }; |
| } |
|
|
| |
| |
| |
| |
| |
| |
| const f32 = Math.fround; |
| function epi(s, a, b) { return f32(f32(f32(s) * a) * b); } |
|
|
| |
| |
| |
| |
| const _fb = new Float32Array(1), _ub = new Uint32Array(_fb.buffer); |
| function bitDiff(a, b) { _fb[0] = a; const u = _ub[0]; _fb[0] = b; return u !== _ub[0]; } |
|
|
| |
| |
| |
| |
| function bgemmJS(Xq, Wq, rs, cs, d, L) { |
| const { m, k, n } = d, batch = d.batch || 1, relu = !!d.relu, mul = L.mul; |
| const raw = !!d.acc; |
| const out = raw ? new Int32Array(batch * m * n) : new Float32Array(batch * m * n); |
| const acc = new Int32Array(n); |
| for (let bz = 0; bz < batch; bz++) { |
| const xo = bz * m * k, wo = bz * k * n, oo = bz * m * n, co = bz * n; |
| for (let i = 0; i < m; i++) { |
| acc.fill(0); |
| const xrow = xo + i * k; |
| for (let p = 0; p < k; p++) { |
| const au = (Xq[xrow + p] & 0xFF) * 256, wrow = wo + p * n; |
| for (let j = 0; j < n; j++) acc[j] += mul[au + (Wq[wrow + j] & 0xFF)]; |
| } |
| const orow = oo + i * n; |
| if (raw) { for (let j = 0; j < n; j++) out[orow + j] = acc[j]; continue; } |
| const rscale = rs[bz * m + i]; |
| for (let j = 0; j < n; j++) { |
| const v = epi(acc[j], rscale, cs[co + j]); |
| out[orow + j] = relu && v < 0 ? 0 : v; |
| } |
| } |
| } |
| return out; |
| } |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| function auditTile(Xq, Wq, rs, cs, d, got, L, nCells) { |
| const { m, k, n } = d, batch = d.batch || 1, relu = !!d.relu, mul = L.mul; |
| const N = nCells || 8; |
| const edges = [[0, m - 1, n - 1], [0, 0, n - 1], [0, m - 1, 0], [0, 0, 0], |
| [batch - 1, m - 1, n - 1], [batch - 1, 0, 0]]; |
| for (let t = 0; t < N; t++) { |
| let bz, i, j; |
| if (t < edges.length) { bz = edges[t][0]; i = edges[t][1]; j = edges[t][2]; } |
| else { bz = (Math.random() * batch) | 0; i = (Math.random() * m) | 0; j = (Math.random() * n) | 0; } |
| let acc = 0; |
| const xrow = bz * m * k + i * k, wo = bz * k * n; |
| for (let p = 0; p < k; p++) acc += mul[(Xq[xrow + p] & 0xFF) * 256 + (Wq[wo + p * n + j] & 0xFF)]; |
| let v = epi(acc, rs[bz * m + i], cs[bz * n + j]); |
| if (relu && v < 0) v = 0; |
| const idx = (bz * m + i) * n + j; |
| if (bitDiff(got[idx], v)) |
| return `GEMM audit failed at [b${bz},${i},${j}] shape ${m}x${k}x${n}: kernel ${Object.is(got[idx], -0) ? "-0" : got[idx]} vs units ${Object.is(v, -0) ? "-0" : v}`; |
| } |
| return null; |
| } |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| function fgemmMirror(A, Bm, d, fma) { |
| const { m, k, n } = d, transA = !!d.transA; |
| const S = k > 4096 ? Math.min(16, Math.ceil(k / 2048)) : 1; |
| const ks = Math.ceil(k / S); |
| const out = new Float32Array(m * n); |
| for (let row = 0; row < m; row++) |
| for (let col = 0; col < n; col++) { |
| let acc = 0; |
| for (let z = 0; z < S; z++) { |
| const p0 = z * ks, p1 = Math.min(k, p0 + ks); |
| let s = 0; |
| for (let p = p0; p < p1; p++) { |
| const a = transA ? A[p * m + row] : A[row * k + p]; |
| s = fma ? f32(s + a * Bm[p * n + col]) |
| : f32(s + f32(a * Bm[p * n + col])); |
| } |
| acc = f32(acc + s); |
| } |
| out[row * n + col] = acc; |
| } |
| return out; |
| } |
|
|
| |
| |
| |
| |
| |
| |
| |
| function auditAttScores(qq, kq, qs, ks, d, got, L, nCells) { |
| const { B, T, heads, hd } = d, C = heads * hd, mul = L.mul, raw = !!d.acc; |
| const N = nCells || 8; |
| const edges = [[B - 1, heads - 1, T - 1, T - 1], [0, 0, 0, 0], |
| [0, heads - 1, T - 1, 0], [B - 1, 0, 0, T - 1]]; |
| for (let t = 0; t < N; t++) { |
| let bi, h, ti, tj; |
| if (t < edges.length) { bi = edges[t][0]; h = edges[t][1]; ti = edges[t][2]; tj = edges[t][3]; } |
| else { bi = (Math.random() * B) | 0; h = (Math.random() * heads) | 0; |
| ti = (Math.random() * T) | 0; tj = (Math.random() * T) | 0; } |
| const bz = bi * heads + h; |
| const qo = (bi * T + ti) * C + h * hd, ko = (bi * T + tj) * C + h * hd; |
| let acc = 0; |
| for (let p = 0; p < hd; p++) acc += mul[(qq[qo + p] & 0xFF) * 256 + (kq[ko + p] & 0xFF)]; |
| const v = raw ? acc : epi(acc, qs[(bi * T + ti) * heads + h], ks[(bi * T + tj) * heads + h]); |
| const idx = (bz * T + ti) * T + tj; |
| if (raw ? got[idx] !== v : bitDiff(got[idx], v)) |
| return `att.scores audit failed at [b${bi},h${h},${ti},${tj}] B${B}T${T}H${heads}d${hd}: kernel ${got[idx]} vs units ${v}`; |
| } |
| return null; |
| } |
| function auditAttCtx(aq, vq, as, vs, d, got, L, nCells) { |
| const { B, T, heads, hd } = d, C = heads * hd, mul = L.mul, raw = !!d.acc; |
| const N = nCells || 8; |
| const edges = [[B - 1, heads - 1, T - 1, hd - 1], [0, 0, 0, 0], |
| [0, heads - 1, T - 1, 0], [B - 1, 0, 0, hd - 1]]; |
| for (let t = 0; t < N; t++) { |
| let bi, h, ti, j; |
| if (t < edges.length) { bi = edges[t][0]; h = edges[t][1]; ti = edges[t][2]; j = edges[t][3]; } |
| else { bi = (Math.random() * B) | 0; h = (Math.random() * heads) | 0; |
| ti = (Math.random() * T) | 0; j = (Math.random() * hd) | 0; } |
| const bz = bi * heads + h, ao = (bz * T + ti) * T; |
| let acc = 0; |
| for (let tj = 0; tj < T; tj++) |
| acc += mul[(aq[ao + tj] & 0xFF) * 256 + (vq[(bi * T + tj) * C + h * hd + j] & 0xFF)]; |
| const v = raw ? acc : epi(acc, as[bz * T + ti], vs[(bi * heads + h) * hd + j]); |
| const idx = (bi * T + ti) * C + h * hd + j; |
| if (raw ? got[idx] !== v : bitDiff(got[idx], v)) |
| return `att.ctx audit failed at [b${bi},h${h},${ti},${j}] B${B}T${T}H${heads}d${hd}: kernel ${got[idx]} vs units ${v}`; |
| } |
| return null; |
| } |
|
|
| |
| |
| |
| |
| async function vgemmBlock(Xf, Wf, d, L, gpuBgemm, audit) { |
| const { m, k, n } = d, batch = d.batch || 1; |
| const x = quantizeRows(Xf, batch * m, k); |
| let wq, ws; |
| if (batch === 1) { |
| const w = quantizeCols(Wf, k, n); wq = w.q; ws = w.s; |
| } else { |
| wq = new Int8Array(batch * k * n); ws = new Float32Array(batch * n); |
| for (let bz = 0; bz < batch; bz++) { |
| const w = quantizeCols(Wf.subarray(bz * k * n, (bz + 1) * k * n), k, n); |
| wq.set(w.q, bz * k * n); ws.set(w.s, bz * n); |
| } |
| } |
| if (gpuBgemm) { |
| const out = await gpuBgemm(x.q, wq, x.s, ws, d); |
| |
| |
| if (audit && audit.due()) { |
| const bad = auditTile(x.q, wq, x.s, ws, d, out, L, audit.cells); |
| if (bad) audit.fail(bad); |
| } |
| return out; |
| } |
| return bgemmJS(x.q, wq, x.s, ws, d, L); |
| } |
|
|
| |
| |
| |
| |
| |
| |
| |
| function quantizeHeadCols(v, B, T, heads, hd) { |
| const C = heads * hd; |
| const q = new Int8Array(B * T * C), s = new Float32Array(B * heads * hd); |
| for (let bi = 0; bi < B; bi++) |
| for (let h = 0; h < heads; h++) |
| for (let j = 0; j < hd; j++) { |
| let mx = 0; |
| for (let ti = 0; ti < T; ti++) { |
| const a = Math.abs(v[(bi * T + ti) * C + h * hd + j]); |
| if (a > mx) mx = a; |
| } |
| const sc = Math.max(mx / 127, 1e-8); |
| s[(bi * heads + h) * hd + j] = sc; |
| for (let ti = 0; ti < T; ti++) { |
| const idx = (bi * T + ti) * C + h * hd + j; |
| const w = Math.round(v[idx] / sc); |
| q[idx] = w < -128 ? -128 : w > 127 ? 127 : w; |
| } |
| } |
| return { q, s }; |
| } |
| |
| |
| function attScoresJS(qq, kq, qs, ks, d, L) { |
| const { B, T, heads, hd } = d, C = heads * hd, mul = L.mul, raw = !!d.acc; |
| const out = raw ? new Int32Array(B * heads * T * T) : new Float32Array(B * heads * T * T); |
| for (let bi = 0; bi < B; bi++) for (let h = 0; h < heads; h++) { |
| const bz = bi * heads + h; |
| for (let ti = 0; ti < T; ti++) { |
| const qo = (bi * T + ti) * C + h * hd, rscale = qs[(bi * T + ti) * heads + h]; |
| for (let tj = 0; tj < T; tj++) { |
| const ko = (bi * T + tj) * C + h * hd; |
| let acc = 0; |
| for (let p = 0; p < hd; p++) acc += mul[(qq[qo + p] & 0xFF) * 256 + (kq[ko + p] & 0xFF)]; |
| out[(bz * T + ti) * T + tj] = raw ? acc : epi(acc, rscale, ks[(bi * T + tj) * heads + h]); |
| } |
| } |
| } |
| return out; |
| } |
| |
| function attCtxJS(aq, vq, as, vs, d, L) { |
| const { B, T, heads, hd } = d, C = heads * hd, mul = L.mul, raw = !!d.acc; |
| const out = raw ? new Int32Array(B * T * C) : new Float32Array(B * T * C); |
| for (let bi = 0; bi < B; bi++) for (let h = 0; h < heads; h++) { |
| const bz = bi * heads + h; |
| for (let ti = 0; ti < T; ti++) { |
| const ao = (bz * T + ti) * T, rscale = as[bz * T + ti]; |
| for (let j = 0; j < hd; j++) { |
| let acc = 0; |
| for (let tj = 0; tj < T; tj++) |
| acc += mul[(aq[ao + tj] & 0xFF) * 256 + (vq[(bi * T + tj) * C + h * hd + j] & 0xFF)]; |
| out[(bi * T + ti) * C + h * hd + j] = raw ? acc : epi(acc, rscale, vs[(bi * heads + h) * hd + j]); |
| } |
| } |
| } |
| return out; |
| } |
|
|
| |
| |
| |
| async function linearFwd(X, W, m, k, n, L, useRelu, matmulInt8) { |
| const xq = quantize(X), wq = quantize(W); |
| const acc = await (matmulInt8 || lutMatmulJS)(xq.q, wq.q, m, k, n, L); |
| const dq = xq.scale * wq.scale; |
| const out = new Float32Array(m * n); |
| const mask = useRelu ? new Uint8Array(m * n) : null; |
| for (let i = 0; i < m * n; i++) { |
| let v = acc[i] * dq; |
| if (useRelu) { if (v > 0) mask[i] = 1; else v = 0; } |
| out[i] = v; |
| } |
| return { out, mask }; |
| } |
|
|
| |
| async function forward(X, y, W1, W2, D, L, matmulInt8) { |
| const { n, din, h, dout } = D; |
| const l1 = await linearFwd(X, W1, n, din, h, L, true, matmulInt8); |
| const l2 = await linearFwd(l1.out, W2, n, h, dout, L, false, matmulInt8); |
| const resid = new Float32Array(n * dout); let loss = 0; |
| for (let i = 0; i < resid.length; i++) { const r = l2.out[i] - y[i]; resid[i] = r; loss += r * r; } |
| loss /= resid.length; |
| return { loss, resid, z1: l1.out, mask1: l1.mask }; |
| } |
|
|
| |
| function backward(X, W1, W2, fwd, D) { |
| const { n, din, h, dout } = D; |
| const { resid, z1, mask1 } = fwd; |
| const s = 2 / n; |
| const dout_ = new Float32Array(resid.length); |
| for (let i = 0; i < resid.length; i++) dout_[i] = resid[i] * s; |
| const mm = TC.matmul, tr = TC.transpose; |
| const gW2 = mm(tr(z1, n, h), dout_, h, n, dout); |
| const dz1 = mm(dout_, tr(W2, h, dout), n, dout, h); |
| for (let i = 0; i < dz1.length; i++) if (!mask1[i]) dz1[i] = 0; |
| const gW1 = mm(tr(X, n, din), dz1, din, n, h); |
| const g = new Float32Array(gW1.length + gW2.length); |
| g.set(gW1, 0); g.set(gW2, gW1.length); |
| return g; |
| } |
|
|
| function splitApply(W1, W2, gAvg, lr) { |
| for (let i = 0; i < W1.length; i++) W1[i] -= lr * gAvg[i]; |
| for (let j = 0; j < W2.length; j++) W2[j] -= lr * gAvg[W1.length + j]; |
| } |
|
|
| const api = { quantize, quantize2, quantizeRows, quantizeCols, quantizeHeadCols, lutMatmulJS, lutMatmul3JS, lutMatmul3, |
| bgemmJS, vgemmBlock, auditTile, epi, attScoresJS, attCtxJS, linearFwd, forward, backward, splitApply, |
| rowAbsMax, scalesFromAbsMax, quantizeRowsInv, vmlpBlock, bitDiff, |
| auditAttScores, auditAttCtx, fgemmMirror }; |
| if (typeof module !== "undefined" && module.exports) { TC = require("./traincore.js"); module.exports = api; } |
| else { TC = root.TrainCore; root.Verified = api; } |
| })(typeof self !== "undefined" ? self : this); |
|
|