File size: 25,242 Bytes
4fd620e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4cf5cdc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
32f59c2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4a8b60f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
93d9b2e
 
 
 
 
 
 
 
 
75f9477
 
 
 
 
 
 
32f59c2
 
93d9b2e
 
32f59c2
 
93d9b2e
 
32f59c2
 
 
 
 
 
 
 
 
 
93d9b2e
 
 
32f59c2
93d9b2e
32f59c2
 
 
 
 
 
 
93d9b2e
 
 
 
d5a32d6
 
 
 
 
 
 
 
93d9b2e
 
d5a32d6
 
 
 
 
 
 
93d9b2e
 
 
 
 
 
75f9477
 
93d9b2e
 
 
 
7fcb6bd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
32f59c2
 
 
 
93d9b2e
32f59c2
 
 
 
 
 
 
 
 
 
 
 
93d9b2e
 
 
 
 
 
 
 
 
 
32f59c2
 
 
f92cfa9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
93d9b2e
f92cfa9
93d9b2e
 
f92cfa9
 
 
 
 
 
 
 
93d9b2e
f92cfa9
 
 
 
 
 
 
93d9b2e
 
f92cfa9
 
 
 
 
 
 
 
93d9b2e
f92cfa9
 
 
 
 
 
4fd620e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f92cfa9
4a8b60f
7fcb6bd
 
4fd620e
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
// Verified INT8 compute — the emulated GPU logic, in the browser.
// A layer's forward runs THROUGH the units: quantize -> LUT multiply -> requant
// -> optional ReLU -> dequant. Backward is a straight-through estimator (the
// integer path has no gradient), so ordinary float weights still learn.
// Same units as the Python/Docker DaisyChain; here they're lookup tables.
(function (root) {
  "use strict";

  let TC;   // TrainCore (matmul/transpose) — resolved per environment at the end

  function quantize(X) {
    let mx = 0; for (let i = 0; i < X.length; i++) { const a = Math.abs(X[i]); if (a > mx) mx = a; }
    const scale = Math.max(mx / 127, 1e-8);
    const q = new Int8Array(X.length);
    for (let i = 0; i < X.length; i++) { let v = Math.round(X[i] / scale); q[i] = v < -128 ? -128 : v > 127 ? 127 : v; }
    return { q, scale };
  }

  // int8 matmul via the verified multiply LUT: acc(m×n) = sum_k mulLUT[Xq,Wq]
  function lutMatmulJS(Xq, Wq, m, k, n, L) {
    const C = new Int32Array(m * n), mul = L.mul;
    for (let i = 0; i < m; i++) {
      for (let p = 0; p < k; p++) {
        const au = (Xq[i * k + p] & 0xFF) * 256, wo = p * n, co = i * n;
        for (let j = 0; j < n; j++) C[co + j] += mul[au + (Wq[wo + j] & 0xFF)];
      }
    }
    return C;
  }

  // ---- 3xINT8 fast-accurate GEMM --------------------------------------------
  // The CUTLASS example-27 "3xTF32" scheme, ported to the verified units:
  // split each float into a coarse int8 part plus an int8-quantized residual,
  // run three EXACT LUT GEMMs (hi·hi, hi·lo, lo·hi), drop the negligible
  // lo·lo, and recombine. Same big/small decomposition NVIDIA uses to recover
  // near-fp32 accuracy from TF32 tensor cores — here it recovers ~14-bit
  // accuracy from the 8-bit units, at 3× the unit ops. Every product still
  // goes through the verified mul8 LUT.
  function quantize2(X) {
    const hi = quantize(X);
    const r = new Float32Array(X.length);
    for (let i = 0; i < X.length; i++) r[i] = X[i] - hi.q[i] * hi.scale;
    const lo = quantize(r);
    return { hi, lo };
  }
  function combine3(hh, hl, lh, x, w, len) {
    const out = new Float32Array(len);
    const shh = x.hi.scale * w.hi.scale, shl = x.hi.scale * w.lo.scale, slh = x.lo.scale * w.hi.scale;
    for (let i = 0; i < len; i++) out[i] = hh[i] * shh + hl[i] * shl + lh[i] * slh;
    return out;
  }
  function lutMatmul3JS(Xf, Wf, m, k, n, L) {          // sync, CPU LUT path
    const x = quantize2(Xf), w = quantize2(Wf);
    return combine3(lutMatmulJS(x.hi.q, w.hi.q, m, k, n, L),
                    lutMatmulJS(x.hi.q, w.lo.q, m, k, n, L),
                    lutMatmulJS(x.lo.q, w.hi.q, m, k, n, L), x, w, m * n);
  }
  async function lutMatmul3(Xf, Wf, m, k, n, L, matmulInt8) {  // any backend
    const x = quantize2(Xf), w = quantize2(Wf);
    const mm = matmulInt8 || lutMatmulJS;
    const [hh, hl, lh] = await Promise.all([
      mm(x.hi.q, w.hi.q, m, k, n, L),
      mm(x.hi.q, w.lo.q, m, k, n, L),
      mm(x.lo.q, w.hi.q, m, k, n, L),
    ]);
    return combine3(hh, hl, lh, x, w, m * n);
  }

  // ---- block-scaled verified GEMM (CUTLASS ex. 67/81 blockwise scaling) ------
  // Per-ROW scales for the activations and per-COLUMN scales for the weights:
  // the integer math through the mul8 LUT is completely unchanged — only the
  // dequant uses rs[row]·cs[col] instead of one tensor-wide product, so a single
  // outlier no longer crushes the quantization resolution of every other
  // row/column. One LUT pass at this granularity beats the per-tensor 3-pass.
  function quantizeRows(X, rows, cols) {
    const q = new Int8Array(rows * cols), s = new Float32Array(rows);
    for (let r = 0; r < rows; r++) {
      let mx = 0;
      for (let c = 0; c < cols; c++) { const a = Math.abs(X[r * cols + c]); if (a > mx) mx = a; }
      const sc = Math.max(mx / 127, 1e-8); s[r] = sc;
      for (let c = 0; c < cols; c++) {
        const v = Math.round(X[r * cols + c] / sc);
        q[r * cols + c] = v < -128 ? -128 : v > 127 ? 127 : v;
      }
    }
    return { q, s };
  }
  function quantizeCols(W, rows, cols) {
    const q = new Int8Array(rows * cols), s = new Float32Array(cols);
    for (let c = 0; c < cols; c++) {
      let mx = 0;
      for (let r = 0; r < rows; r++) { const a = Math.abs(W[r * cols + c]); if (a > mx) mx = a; }
      s[c] = Math.max(mx / 127, 1e-8);
    }
    for (let r = 0; r < rows; r++)
      for (let c = 0; c < cols; c++) {
        const v = Math.round(W[r * cols + c] / s[c]);
        q[r * cols + c] = v < -128 ? -128 : v > 127 ? 127 : v;
      }
    return { q, s };
  }

  // ---- B2B MLP chain (CUTLASS ex. 13 two-GEMM fusion + ex. 23 epilogue
  // reduction), respecced for cross-device exactness ---------------------------
  // The MLP is the one back-to-back GEMM pair with no layernorm/softmax between
  // (ReLU is already fused in the epilogue), so the intermediate h1 can be
  // quantized ON the GPU and fed straight to the second GEMM. Two rules make
  // that fleet-safe:
  //   1. The per-row |max| (ex. 23) uses only comparisons — exact on any
  //      hardware, order-independent — and comes back to JS as ~1KB.
  //   2. Scale DERIVATION (two divisions) stays in JS f64, which IEEE requires
  //      to be exactly rounded and is therefore identical on every device.
  //      WGSL division is only 2.5 ULP — a fork waiting to happen — but WGSL
  //      multiply/add are correctly rounded and floor/clamp are exact. So the
  //      quantize step is respecced from  round(x / scale)  to
  //      floor(f32(x * invScale) + 0.5)  — floor(x+0.5) IS Math.round's tie
  //      rule — and the fround-stepped mirror below is bit-identical to the
  //      GPU kernel, which is exact-gated against it at init.
  // NOTE this changes which int8 a value on a rounding boundary lands on
  // (≤1 step) vs quantizeRows, so old and new builds cannot co-train — the
  // divergence guard stops such mixed groups by design.
  function rowAbsMax(X, rows, cols) {
    const mx = new Float32Array(rows);
    for (let r = 0; r < rows; r++) {
      let m = 0;
      for (let c = 0; c < cols; c++) { const a = Math.abs(X[r * cols + c]); if (a > m) m = a; }
      mx[r] = m;
    }
    return mx;
  }
  function scalesFromAbsMax(mx) {                // f64 divisions: exactly rounded, device-identical
    const scale = new Float32Array(mx.length), inv = new Float32Array(mx.length);
    for (let i = 0; i < mx.length; i++) {
      scale[i] = Math.max(mx[i] / 127, 1e-8);
      inv[i] = 1 / scale[i];                     // recip of the STORED f32 scale
    }
    return { scale, inv };
  }
  function quantizeRowsInv(X, rows, cols, inv) { // bit-exact mirror of the GPU quantize kernel
    const q = new Int8Array(rows * cols);
    for (let r = 0; r < rows; r++) {
      const iv = inv[r];
      for (let c = 0; c < cols; c++) {
        const n = Math.floor(f32(f32(X[r * cols + c] * iv) + 0.5));
        q[r * cols + c] = n < -128 ? -128 : n > 127 ? 127 : n;
      }
    }
    return q;
  }
  // the chained MLP: X @ W1 -> ReLU (fused) -> absmax -> quantize -> @ W2.
  //   d = { m, k, h, n };  gpuMlp (from webgpu.js) runs both GEMMs + the
  // on-GPU quantize with one tiny absmax readback between; without it the CPU
  // mirror chain runs — SAME math, so mixed GPU/CPU fleets stay bit-identical.
  async function vmlpBlock(Xf, W1f, W2f, d, L, gpuMlp, audit) {
    const x = quantizeRows(Xf, d.m, d.k);
    const w1 = quantizeCols(W1f, d.k, d.h);
    const w2 = quantizeCols(W2f, d.h, d.n);
    if (gpuMlp) {
      const r = await gpuMlp(x.q, w1.q, w2.q, x.s, w1.s, w2.s, d);
      if (audit && audit.due()) {
        // audit BOTH live GEMMs: gemm1 against the units directly; gemm2 by
        // reconstructing its exact operand through the proven quantize mirror
        const bad1 = auditTile(x.q, w1.q, x.s, w1.s, { m: d.m, k: d.k, n: d.h, relu: true }, r.h1, L, audit.cells);
        if (bad1) { audit.fail("mlp gemm1: " + bad1); return r; }
        const sc = scalesFromAbsMax(rowAbsMax(r.h1, d.m, d.h));
        const hq = quantizeRowsInv(r.h1, d.m, d.h, sc.inv);
        const bad2 = auditTile(hq, w2.q, sc.scale, w2.s, { m: d.m, k: d.h, n: d.n }, r.out, L, audit.cells);
        if (bad2) audit.fail("mlp gemm2: " + bad2);
      }
      return r;
    }
    const h1 = bgemmJS(x.q, w1.q, x.s, w1.s, { m: d.m, k: d.k, n: d.h, batch: 1, relu: true }, L);
    const sc = scalesFromAbsMax(rowAbsMax(h1, d.m, d.h));
    const hq = quantizeRowsInv(h1, d.m, d.h, sc.inv);
    const out = bgemmJS(hq, w2.q, sc.scale, w2.s, { m: d.m, k: d.h, n: d.n, batch: 1 }, L);
    return { h1, out };
  }

  // ---- epilogue mirror -------------------------------------------------------
  // BIT-EXACT mirror of the WGSL epilogue `f32(s) * a * b`. WGSL rounds to f32
  // after the int->float conversion and after EACH multiply; plain JS would do
  // the whole chain in f64 and round once, which differs in the last ulp. That
  // last-ulp gap is what used to force a tolerance into the kernel gates —
  // mirroring the rounding exactly is what lets the gates compare with `!==`.
  const f32 = Math.fround;
  function epi(s, a, b) { return f32(f32(f32(s) * a) * b); }

  // f32 equality at the BIT level: `!==` says -0 === 0, but replicas are
  // compared by hashing raw bytes, so an audit that can't see the sign of zero
  // could pass a device that later forks the fleet. (Real ISAs have non-IEEE
  // modes that flush -0 to +0 — e.g. RDNA2 output modifiers / legacy muls.)
  const _fb = new Float32Array(1), _ub = new Uint32Array(_fb.buffer);
  function bitDiff(a, b) { _fb[0] = a; const u = _ub[0]; _fb[0] = b; return u !== _ub[0]; }

  // CPU mirror of the fused GPU kernel: batched int8 GEMM through the LUT with
  // the epilogue (block dequant + optional ReLU) applied before returning —
  // exactly what the WGSL kernel does on-device. d.acc=true returns the raw
  // int32 accumulator instead (the exact oracle the fused kernel normally hides).
  function bgemmJS(Xq, Wq, rs, cs, d, L) {
    const { m, k, n } = d, batch = d.batch || 1, relu = !!d.relu, mul = L.mul;
    const raw = !!d.acc;
    const out = raw ? new Int32Array(batch * m * n) : new Float32Array(batch * m * n);
    const acc = new Int32Array(n);
    for (let bz = 0; bz < batch; bz++) {
      const xo = bz * m * k, wo = bz * k * n, oo = bz * m * n, co = bz * n;
      for (let i = 0; i < m; i++) {
        acc.fill(0);
        const xrow = xo + i * k;
        for (let p = 0; p < k; p++) {
          const au = (Xq[xrow + p] & 0xFF) * 256, wrow = wo + p * n;
          for (let j = 0; j < n; j++) acc[j] += mul[au + (Wq[wrow + j] & 0xFF)];
        }
        const orow = oo + i * n;
        if (raw) { for (let j = 0; j < n; j++) out[orow + j] = acc[j]; continue; }
        const rscale = rs[bz * m + i];
        for (let j = 0; j < n; j++) {
          const v = epi(acc[j], rscale, cs[co + j]);
          out[orow + j] = relu && v < 0 ? 0 : v;
        }
      }
    }
    return out;
  }

  // Recompute a handful of RANDOM output cells of a live GEMM through the LUT
  // mirror and compare against what the kernel produced. Sampling cells instead
  // of whole matrices makes this cheap enough to run continuously, at the real
  // shapes training uses — not once at boot on toy inputs.
  // STRATIFIED sampling. Uniformly random cells are the wrong instrument for
  // the bugs that actually occur here: a bounds-guard off-by-one or a pack-tail
  // padding bug lives on the LAST row/column, and uniform sampling finds that
  // with probability ~1/n per cell — at the 16512-wide logits GEMM, never. So
  // the first cells are the structurally dangerous ones (corners, last row,
  // last column, last batch) chosen deterministically, and the remainder are
  // random interior cells that catch diffuse bugs. Same principle as poisoning
  // the buffer pool: construct the dangerous case, don't wait to land on it.
  function auditTile(Xq, Wq, rs, cs, d, got, L, nCells) {
    const { m, k, n } = d, batch = d.batch || 1, relu = !!d.relu, mul = L.mul;
    const N = nCells || 8;
    const edges = [[0, m - 1, n - 1], [0, 0, n - 1], [0, m - 1, 0], [0, 0, 0],
                   [batch - 1, m - 1, n - 1], [batch - 1, 0, 0]];
    for (let t = 0; t < N; t++) {
      let bz, i, j;
      if (t < edges.length) { bz = edges[t][0]; i = edges[t][1]; j = edges[t][2]; }
      else { bz = (Math.random() * batch) | 0; i = (Math.random() * m) | 0; j = (Math.random() * n) | 0; }
      let acc = 0;
      const xrow = bz * m * k + i * k, wo = bz * k * n;
      for (let p = 0; p < k; p++) acc += mul[(Xq[xrow + p] & 0xFF) * 256 + (Wq[wo + p * n + j] & 0xFF)];
      let v = epi(acc, rs[bz * m + i], cs[bz * n + j]);
      if (relu && v < 0) v = 0;
      const idx = (bz * m + i) * n + j;
      if (bitDiff(got[idx], v))
        return `GEMM audit failed at [b${bz},${i},${j}] shape ${m}x${k}x${n}: kernel ${Object.is(got[idx], -0) ? "-0" : got[idx]} vs units ${Object.is(v, -0) ? "-0" : v}`;
    }
    return null;
  }

  // ---- exact mirror of the split-K f32 GEMM ----------------------------------
  // The f32 backward GEMM was the last kernel gated by a TOLERANCE (allclose at
  // 1e-3) — and this project's own gate mutation test shows allclose waving
  // through real bugs. The reason was real though: split-K accumulates in a
  // different ORDER than a naive reference, so bit-equality against the naive
  // one is impossible. The fix is the same as the epilogue mirror: reproduce
  // the kernel's order exactly, then compare with `!==`.
  //   partials: for z in 0..S-1, sum p in [z*ks, min(k,(z+1)*ks)) in order
  //   reduce:   sum the S partials in ascending z
  // `fma` selects the rounding schedule for `s + a*b`: WGSL PERMITS a compiler
  // to contract that into a fused multiply-add (one rounding) instead of two.
  // Which one the device does is a fact about the device, so the gate tries
  // both and reports which matches rather than assuming.
  function fgemmMirror(A, Bm, d, fma) {
    const { m, k, n } = d, transA = !!d.transA;
    const S = k > 4096 ? Math.min(16, Math.ceil(k / 2048)) : 1;
    const ks = Math.ceil(k / S);
    const out = new Float32Array(m * n);
    for (let row = 0; row < m; row++)
      for (let col = 0; col < n; col++) {
        let acc = 0;                                    // reduce pass, ascending z
        for (let z = 0; z < S; z++) {
          const p0 = z * ks, p1 = Math.min(k, p0 + ks);
          let s = 0;                                    // one partial, in order
          for (let p = p0; p < p1; p++) {
            const a = transA ? A[p * m + row] : A[row * k + p];
            s = fma ? f32(s + a * Bm[p * n + col])      // single rounding
                    : f32(s + f32(a * Bm[p * n + col]));
          }
          acc = f32(acc + s);
        }
        out[row * n + col] = acc;
      }
    return out;
  }

  // ---- live audits for the fused attention kernels ---------------------------
  // The attention kernels had exact INIT gates but nothing at live shapes —
  // the exact gap the GEMM audit exists to close, left open on the kernels with
  // the trickiest indexing (head-strided gather, scatter write-back). These
  // recompute individual output cells from the units, stratified like
  // auditTile: last/first token pair, last head, last channel first, then
  // random. Cost is hd (or T) multiply-adds per cell.
  function auditAttScores(qq, kq, qs, ks, d, got, L, nCells) {
    const { B, T, heads, hd } = d, C = heads * hd, mul = L.mul, raw = !!d.acc;
    const N = nCells || 8;
    const edges = [[B - 1, heads - 1, T - 1, T - 1], [0, 0, 0, 0],
                   [0, heads - 1, T - 1, 0], [B - 1, 0, 0, T - 1]];
    for (let t = 0; t < N; t++) {
      let bi, h, ti, tj;
      if (t < edges.length) { bi = edges[t][0]; h = edges[t][1]; ti = edges[t][2]; tj = edges[t][3]; }
      else { bi = (Math.random() * B) | 0; h = (Math.random() * heads) | 0;
             ti = (Math.random() * T) | 0; tj = (Math.random() * T) | 0; }
      const bz = bi * heads + h;
      const qo = (bi * T + ti) * C + h * hd, ko = (bi * T + tj) * C + h * hd;
      let acc = 0;
      for (let p = 0; p < hd; p++) acc += mul[(qq[qo + p] & 0xFF) * 256 + (kq[ko + p] & 0xFF)];
      const v = raw ? acc : epi(acc, qs[(bi * T + ti) * heads + h], ks[(bi * T + tj) * heads + h]);
      const idx = (bz * T + ti) * T + tj;
      if (raw ? got[idx] !== v : bitDiff(got[idx], v))
        return `att.scores audit failed at [b${bi},h${h},${ti},${tj}] B${B}T${T}H${heads}d${hd}: kernel ${got[idx]} vs units ${v}`;
    }
    return null;
  }
  function auditAttCtx(aq, vq, as, vs, d, got, L, nCells) {
    const { B, T, heads, hd } = d, C = heads * hd, mul = L.mul, raw = !!d.acc;
    const N = nCells || 8;
    const edges = [[B - 1, heads - 1, T - 1, hd - 1], [0, 0, 0, 0],
                   [0, heads - 1, T - 1, 0], [B - 1, 0, 0, hd - 1]];
    for (let t = 0; t < N; t++) {
      let bi, h, ti, j;
      if (t < edges.length) { bi = edges[t][0]; h = edges[t][1]; ti = edges[t][2]; j = edges[t][3]; }
      else { bi = (Math.random() * B) | 0; h = (Math.random() * heads) | 0;
             ti = (Math.random() * T) | 0; j = (Math.random() * hd) | 0; }
      const bz = bi * heads + h, ao = (bz * T + ti) * T;
      let acc = 0;
      for (let tj = 0; tj < T; tj++)
        acc += mul[(aq[ao + tj] & 0xFF) * 256 + (vq[(bi * T + tj) * C + h * hd + j] & 0xFF)];
      const v = raw ? acc : epi(acc, as[bz * T + ti], vs[(bi * heads + h) * hd + j]);
      const idx = (bi * T + ti) * C + h * hd + j;
      if (raw ? got[idx] !== v : bitDiff(got[idx], v))
        return `att.ctx audit failed at [b${bi},h${h},${ti},${j}] B${B}T${T}H${heads}d${hd}: kernel ${got[idx]} vs units ${v}`;
    }
    return null;
  }

  // block-scaled verified GEMM, float in → float out.
  //   d = { m, k, n, batch=1, relu=false }; X is (batch·m)×k, W is batch×(k×n)
  // gpuBgemm (from webgpu.js) runs the batched kernel with the fused epilogue;
  // without it the CPU LUT mirror runs. Every product goes through the units.
  async function vgemmBlock(Xf, Wf, d, L, gpuBgemm, audit) {
    const { m, k, n } = d, batch = d.batch || 1;
    const x = quantizeRows(Xf, batch * m, k);
    let wq, ws;
    if (batch === 1) {
      const w = quantizeCols(Wf, k, n); wq = w.q; ws = w.s;
    } else {
      wq = new Int8Array(batch * k * n); ws = new Float32Array(batch * n);
      for (let bz = 0; bz < batch; bz++) {
        const w = quantizeCols(Wf.subarray(bz * k * n, (bz + 1) * k * n), k, n);
        wq.set(w.q, bz * k * n); ws.set(w.s, bz * n);
      }
    }
    if (gpuBgemm) {
      const out = await gpuBgemm(x.q, wq, x.s, ws, d);
      // continuous re-verification at LIVE shapes: the boot gate only ever saw
      // toy inputs, so sample a few real cells against the units as we go
      if (audit && audit.due()) {
        const bad = auditTile(x.q, wq, x.s, ws, d, out, L, audit.cells);
        if (bad) audit.fail(bad);
      }
      return out;
    }
    return bgemmJS(x.q, wq, x.s, ws, d, L);
  }

  // ---- gather-fused attention through the units (CUTLASS ex. 36/52) ----------
  // The kernels read q/k/v/ctx directly in their natural BT×C layout with
  // head-strided indexing — no JS gather copies, no kᵀ transpose, and the
  // context write scatters straight back into BT×C. Quantization stays
  // block-scaled: q/k/a per (token,head) row, v per (head,channel) column.
  // The (BT·heads)×hd row view of q/k IS the contiguous buffer, so
  // quantizeRows(q, BT·heads, hd) gives per-(token,head) scales for free.
  function quantizeHeadCols(v, B, T, heads, hd) {   // per (batch,head,channel) column
    const C = heads * hd;
    const q = new Int8Array(B * T * C), s = new Float32Array(B * heads * hd);
    for (let bi = 0; bi < B; bi++)
      for (let h = 0; h < heads; h++)
        for (let j = 0; j < hd; j++) {
          let mx = 0;
          for (let ti = 0; ti < T; ti++) {
            const a = Math.abs(v[(bi * T + ti) * C + h * hd + j]);
            if (a > mx) mx = a;
          }
          const sc = Math.max(mx / 127, 1e-8);
          s[(bi * heads + h) * hd + j] = sc;
          for (let ti = 0; ti < T; ti++) {
            const idx = (bi * T + ti) * C + h * hd + j;
            const w = Math.round(v[idx] / sc);
            q[idx] = w < -128 ? -128 : w > 127 ? 127 : w;
          }
        }
    return { q, s };
  }
  // scores S[bz,ti,tj] = q_row(bi,ti,h) · k_row(bi,tj,h), every product via the LUT
  // d.acc=true returns the raw int32 accumulator (exact oracle for the kernel gate)
  function attScoresJS(qq, kq, qs, ks, d, L) {
    const { B, T, heads, hd } = d, C = heads * hd, mul = L.mul, raw = !!d.acc;
    const out = raw ? new Int32Array(B * heads * T * T) : new Float32Array(B * heads * T * T);
    for (let bi = 0; bi < B; bi++) for (let h = 0; h < heads; h++) {
      const bz = bi * heads + h;
      for (let ti = 0; ti < T; ti++) {
        const qo = (bi * T + ti) * C + h * hd, rscale = qs[(bi * T + ti) * heads + h];
        for (let tj = 0; tj < T; tj++) {
          const ko = (bi * T + tj) * C + h * hd;
          let acc = 0;
          for (let p = 0; p < hd; p++) acc += mul[(qq[qo + p] & 0xFF) * 256 + (kq[ko + p] & 0xFF)];
          out[(bz * T + ti) * T + tj] = raw ? acc : epi(acc, rscale, ks[(bi * T + tj) * heads + h]);
        }
      }
    }
    return out;
  }
  // ctx[(bi,ti),(h,j)] = Σ_tj a[bz,ti,tj]·v[(bi,tj),(h,j)] — scatter fused into BT×C
  function attCtxJS(aq, vq, as, vs, d, L) {
    const { B, T, heads, hd } = d, C = heads * hd, mul = L.mul, raw = !!d.acc;
    const out = raw ? new Int32Array(B * T * C) : new Float32Array(B * T * C);
    for (let bi = 0; bi < B; bi++) for (let h = 0; h < heads; h++) {
      const bz = bi * heads + h;
      for (let ti = 0; ti < T; ti++) {
        const ao = (bz * T + ti) * T, rscale = as[bz * T + ti];
        for (let j = 0; j < hd; j++) {
          let acc = 0;
          for (let tj = 0; tj < T; tj++)
            acc += mul[(aq[ao + tj] & 0xFF) * 256 + (vq[(bi * T + tj) * C + h * hd + j] & 0xFF)];
          out[(bi * T + ti) * C + h * hd + j] = raw ? acc : epi(acc, rscale, vs[(bi * heads + h) * hd + j]);
        }
      }
    }
    return out;
  }

  // one verified layer forward; returns float out (+ cache for STE backward).
  // Every product goes through the verified INT8 multiply (mul8 LUT) with exact
  // int32 accumulation — i.e. an emulated INT8 tensor-core GEMM — then dequant.
  async function linearFwd(X, W, m, k, n, L, useRelu, matmulInt8) {
    const xq = quantize(X), wq = quantize(W);
    const acc = await (matmulInt8 || lutMatmulJS)(xq.q, wq.q, m, k, n, L);  // verified multiply
    const dq = xq.scale * wq.scale;
    const out = new Float32Array(m * n);
    const mask = useRelu ? new Uint8Array(m * n) : null;
    for (let i = 0; i < m * n; i++) {
      let v = acc[i] * dq;
      if (useRelu) { if (v > 0) mask[i] = 1; else v = 0; }
      out[i] = v;
    }
    return { out, mask };
  }

  // 2-layer MLP: X→H (relu) →dout. Forward through verified units, MSE loss.
  async function forward(X, y, W1, W2, D, L, matmulInt8) {
    const { n, din, h, dout } = D;
    const l1 = await linearFwd(X, W1, n, din, h, L, true, matmulInt8);
    const l2 = await linearFwd(l1.out, W2, n, h, dout, L, false, matmulInt8);
    const resid = new Float32Array(n * dout); let loss = 0;
    for (let i = 0; i < resid.length; i++) { const r = l2.out[i] - y[i]; resid[i] = r; loss += r * r; }
    loss /= resid.length;
    return { loss, resid, z1: l1.out, mask1: l1.mask };
  }

  // STE backward (verified matmul treated as float X@W). Returns flat [gW1, gW2].
  function backward(X, W1, W2, fwd, D) {
    const { n, din, h, dout } = D;
    const { resid, z1, mask1 } = fwd;
    const s = 2 / n;
    const dout_ = new Float32Array(resid.length);
    for (let i = 0; i < resid.length; i++) dout_[i] = resid[i] * s;
    const mm = TC.matmul, tr = TC.transpose;
    const gW2 = mm(tr(z1, n, h), dout_, h, n, dout);          // z1ᵀ @ dout
    const dz1 = mm(dout_, tr(W2, h, dout), n, dout, h);        // dout @ W2ᵀ
    for (let i = 0; i < dz1.length; i++) if (!mask1[i]) dz1[i] = 0;  // relu grad
    const gW1 = mm(tr(X, n, din), dz1, din, n, h);            // Xᵀ @ dz1
    const g = new Float32Array(gW1.length + gW2.length);
    g.set(gW1, 0); g.set(gW2, gW1.length);
    return g;
  }

  function splitApply(W1, W2, gAvg, lr) {
    for (let i = 0; i < W1.length; i++) W1[i] -= lr * gAvg[i];
    for (let j = 0; j < W2.length; j++) W2[j] -= lr * gAvg[W1.length + j];
  }

  const api = { quantize, quantize2, quantizeRows, quantizeCols, quantizeHeadCols, lutMatmulJS, lutMatmul3JS, lutMatmul3,
                bgemmJS, vgemmBlock, auditTile, epi, attScoresJS, attCtxJS, linearFwd, forward, backward, splitApply,
                rowAbsMax, scalesFromAbsMax, quantizeRowsInv, vmlpBlock, bitDiff,
                auditAttScores, auditAttCtx, fgemmMirror };
  if (typeof module !== "undefined" && module.exports) { TC = require("./traincore.js"); module.exports = api; }
  else { TC = root.TrainCore; root.Verified = api; }
})(typeof self !== "undefined" ? self : this);