Quazim0t0 commited on
Commit
b23b595
·
verified ·
1 Parent(s): b004b80

web: gate-the-gate boundary discrimination in the B2B init gate

Browse files
Files changed (3) hide show
  1. web/public/webgpu.js +29 -0
  2. web/test_b2b.js +3 -0
  3. web/webgpu.js +751 -0
web/public/webgpu.js CHANGED
@@ -416,6 +416,35 @@
416
  for (let i = 0; i < ref.out.length; i++)
417
  if (hw.out[i] !== ref.out[i]) return `out mismatch @${i} (${shape}): ${hw.out[i]} vs ${ref.out[i]}`;
418
  }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
419
  return null;
420
  }
421
  const lutMlpEnv = { dp4: false, gemm: bgLutPipe, rowmax: rowmaxPipe, quant: quantI32Pipe, lutBuf };
 
416
  for (let i = 0; i < ref.out.length; i++)
417
  if (hw.out[i] !== ref.out[i]) return `out mismatch @${i} (${shape}): ${hw.out[i]} vs ${ref.out[i]}`;
418
  }
419
+ // DISCRIMINATING case: the sweep above passes vacuously if no value
420
+ // lands on a rounding boundary — the old round(x/scale) spec would
421
+ // pass it too. So hunt (in fast JS) for an input where the two specs
422
+ // actually disagree, then check the GPU sides with the RESPEC. This
423
+ // gates the gate: a pass must be something the old spec would fail.
424
+ const V = root.Verified;
425
+ const d1 = { m: 16, k: 32, h: 64, n: 16 };
426
+ const rnd1 = (len) => Float32Array.from({ length: len }, () => Math.random() * 4 - 2);
427
+ for (let t = 0; t < 800; t++) {
428
+ const Xf = rnd1(d1.m * d1.k), W1 = rnd1(d1.k * d1.h), W2 = rnd1(d1.h * d1.n);
429
+ const x = V.quantizeRows(Xf, d1.m, d1.k), w1 = V.quantizeCols(W1, d1.k, d1.h);
430
+ const h1 = V.bgemmJS(x.q, w1.q, x.s, w1.s, { m: d1.m, k: d1.k, n: d1.h, batch: 1, relu: true }, L);
431
+ const sc = V.scalesFromAbsMax(V.rowAbsMax(h1, d1.m, d1.h));
432
+ const qNew = V.quantizeRowsInv(h1, d1.m, d1.h, sc.inv);
433
+ const qOld = V.quantizeRows(h1, d1.m, d1.h).q;
434
+ let boundary = false;
435
+ for (let i = 0; i < qNew.length; i++) if (qNew[i] !== qOld[i]) { boundary = true; break; }
436
+ if (!boundary) continue;
437
+ const w2 = V.quantizeCols(W2, d1.h, d1.n);
438
+ const gpu = await mlpFn(x.q, w1.q, w2.q, x.s, w1.s, w2.s, d1);
439
+ const refNew = V.bgemmJS(qNew, w2.q, sc.scale, w2.s, { m: d1.m, k: d1.h, n: d1.n, batch: 1 }, L);
440
+ const refOld = V.bgemmJS(qOld, w2.q, sc.scale, w2.s, { m: d1.m, k: d1.h, n: d1.n, batch: 1 }, L);
441
+ let eqNew = true, eqOld = true;
442
+ for (let i = 0; i < refNew.length; i++) { if (gpu.out[i] !== refNew[i]) eqNew = false; if (gpu.out[i] !== refOld[i]) eqOld = false; }
443
+ if (!eqNew) return "discriminating boundary case: GPU chain does not match the respec mirror";
444
+ if (eqOld) return "discriminating boundary case: GPU chain matches the OLD quantize spec";
445
+ return null; // proven: respec, not merely gate-compatible
446
+ }
447
+ console.warn("B2B gate: no rounding-boundary input found in 800 trials — respec discrimination unproven this boot (sweep still exact)");
448
  return null;
449
  }
450
  const lutMlpEnv = { dp4: false, gemm: bgLutPipe, rowmax: rowmaxPipe, quant: quantI32Pipe, lutBuf };
web/test_b2b.js CHANGED
@@ -61,6 +61,9 @@ const ok = (c, msg) => { console.log(`${c ? " ok " : " FAIL"} ${msg}`); if (
61
  }
62
  }
63
  ok(maxd <= 1, `respecced quantize differs by at most 1 step (max ${maxd}; ${diff}/${n} = ${(100 * diff / n).toFixed(3)}% of values moved)`);
 
 
 
64
  }
65
  // 3) chain == its parts, and gemm1 is byte-identical to the un-chained GEMM
66
  {
 
61
  }
62
  }
63
  ok(maxd <= 1, `respecced quantize differs by at most 1 step (max ${maxd}; ${diff}/${n} = ${(100 * diff / n).toFixed(3)}% of values moved)`);
64
+ // and it must differ SOMEWHERE: if the respec silently became the old
65
+ // spec, every "bit-identical to the mirror" claim downstream is untested
66
+ ok(diff > 0, `the respec is a real change (${diff} boundary values moved) — the mirror equivalence is not vacuous`);
67
  }
68
  // 3) chain == its parts, and gemm1 is byte-identical to the un-chained GEMM
69
  {
web/webgpu.js ADDED
@@ -0,0 +1,751 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ // WebGPU INT8 matmul via the verified multiply LUT — the emulated GPU logic
2
+ // running on the browser's GPU. Automatic CPU fallback (same LUT) for machines
3
+ // without WebGPU (e.g. old PCs via Supermium). initCompute() returns
4
+ // { backend, label, matmulInt8(Xq, Wq, m, k, n, L) -> Int32Array }
5
+ // matching Verified.lutMatmulJS, so the trainer is device-blind.
6
+ //
7
+ // High-throughput path: when the browser exposes WGSL's
8
+ // packed_4x8_integer_dot_product feature, we use dot4I8Packed — it compiles to
9
+ // the GPU's DP4A/INT8 dot-product hardware (the same units tensor-core INT8
10
+ // paths are built on): 4 exact int8 MACs per instruction, int32 accumulation,
11
+ // and 4× less memory traffic from packing. Because int8×int8→int32 is exact,
12
+ // it is bit-identical to the verified mul8 LUT — and we PROVE that at init by
13
+ // cross-checking random matmuls against the LUT before trusting it. If the
14
+ // hardware ever disagrees with the units, we fall back to the LUT shader.
15
+ (function (root) {
16
+ "use strict";
17
+
18
+ const WGSL_LUT = `
19
+ @group(0) @binding(0) var<storage, read> Xq : array<i32>; // int8 byte per elem
20
+ @group(0) @binding(1) var<storage, read> Wq : array<i32>;
21
+ @group(0) @binding(2) var<storage, read> lut : array<i32>; // 65536 signed products
22
+ @group(0) @binding(3) var<storage, read_write> C : array<i32>;
23
+ @group(0) @binding(4) var<uniform> dims : vec3<u32>; // m, k, n
24
+ @compute @workgroup_size(8, 8)
25
+ fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
26
+ let m = dims.x; let k = dims.y; let n = dims.z;
27
+ let row = gid.x; let col = gid.y;
28
+ if (row >= m || col >= n) { return; }
29
+ var s : i32 = 0;
30
+ for (var p = 0u; p < k; p = p + 1u) {
31
+ let au = u32(Xq[row * k + p] & 255);
32
+ let bu = u32(Wq[p * n + col] & 255);
33
+ s = s + lut[au * 256u + bu];
34
+ }
35
+ C[row * n + col] = s;
36
+ }`;
37
+
38
+ // ---- B2B MLP chain kernels (CUTLASS ex. 13 + 23) ---------------------------
39
+ // ROWMAX: per-row |max| of a GEMM's f32 output, fused into the same command
40
+ // encoder (ex. 23 epilogue reduction). Non-negative f32 bit patterns order
41
+ // like u32, so atomicMax on bitcast(abs(v)) computes an EXACT max, in any
42
+ // execution order, on any hardware — nothing here can round.
43
+ const WGSL_ROWMAX = `
44
+ @group(0) @binding(0) var<storage, read> O : array<f32>;
45
+ @group(0) @binding(1) var<storage, read_write> MX : array<atomic<u32>>;
46
+ @group(0) @binding(2) var<uniform> dims : vec4<u32>; // m, n, _, _
47
+ @compute @workgroup_size(8, 8, 1)
48
+ fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
49
+ let m = dims.x; let n = dims.y;
50
+ let row = gid.x; let col = gid.y;
51
+ if (row >= m || col >= n) { return; }
52
+ atomicMax(&MX[row], bitcast<u32>(abs(O[row * n + col])));
53
+ }`;
54
+ // QUANT: h1 (f32, still on the GPU) -> int8 by MULTIPLY with a JS-computed
55
+ // inverse scale. floor(f32(x*inv)+0.5) uses only ops WGSL guarantees exact
56
+ // or correctly rounded (mul, add, floor, clamp) — division is 2.5 ULP and
57
+ // never runs on the GPU. Bit-identical to Verified.quantizeRowsInv, and
58
+ // exact-gated against it at init. pack=true emits 4 bytes per u32 for the
59
+ // DP4A kernel; pack=false emits one i32 per element for the LUT kernel.
60
+ const WGSL_QUANT = (pack) => `
61
+ @group(0) @binding(0) var<storage, read> H : array<f32>;
62
+ @group(0) @binding(1) var<storage, read> inv : array<f32>; // per row
63
+ @group(0) @binding(2) var<storage, read_write> Q : array<${pack ? "u32" : "i32"}>;
64
+ @group(0) @binding(3) var<uniform> dims : vec4<u32>; // m, k, kw, _
65
+ @compute @workgroup_size(64)
66
+ fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
67
+ let m = dims.x; let k = dims.y; let kw = dims.z;
68
+ let idx = gid.x;
69
+ ${pack ? `
70
+ if (idx >= m * kw) { return; }
71
+ let row = idx / kw;
72
+ var acc : u32 = 0u;
73
+ for (var b = 0u; b < 4u; b = b + 1u) {
74
+ let c = (idx % kw) * 4u + b;
75
+ var q : i32 = 0;
76
+ if (c < k) {
77
+ let v = clamp(floor(H[row * k + c] * inv[row] + 0.5), -128.0, 127.0);
78
+ q = i32(v);
79
+ }
80
+ acc = acc | ((u32(q) & 255u) << (8u * b));
81
+ }
82
+ Q[idx] = acc;` : `
83
+ if (idx >= m * k) { return; }
84
+ let row = idx / k;
85
+ let v = clamp(floor(H[idx] * inv[row] + 0.5), -128.0, 127.0);
86
+ Q[idx] = i32(v);`}
87
+ }`;
88
+
89
+ // NOTE: the un-batched DP4A matmul that used to live here was removed. It was
90
+ // the only kernel with an exact gate, but the transformer stopped calling it
91
+ // when the block-scaled path landed — so it sat here passing its own gate
92
+ // while verifying nothing that ran. A gate on a kernel nobody calls is worse
93
+ // than no gate: it reads like coverage. The batched kernels below are the
94
+ // ones training uses, and they now carry that exact gate instead.
95
+
96
+ // Batched block-scaled GEMM with a FUSED EPILOGUE (CUTLASS ex. 05/24 + 12):
97
+ // grid z = batch index, so ALL attention heads run in ONE dispatch, and the
98
+ // epilogue (block dequant rs·cs + optional ReLU) happens before the data
99
+ // leaves the GPU — f32 out, no int32 readback, no second pass in JS.
100
+ // Each kernel is emitted in two variants from ONE source string: the live one
101
+ // (fused epilogue, f32 out) and a `verify` one that writes the raw int32
102
+ // accumulator instead. The indexing — the part that actually goes wrong — is
103
+ // textually identical, so gating the verify variant genuinely gates the live
104
+ // kernel, and the comparison is EXACT (int8xint8->int32 has no rounding to
105
+ // hide in) instead of an allclose that whole bug classes walk through.
106
+ const OUT_DECL = (v, b) => `@group(0) @binding(${b}) var<storage, read_write> O : array<${v ? "i32" : "f32"}>;`;
107
+
108
+ const WGSL_BG_LUT = (verify) => `
109
+ @group(0) @binding(0) var<storage, read> Xq : array<i32>; // int8 byte per elem
110
+ @group(0) @binding(1) var<storage, read> Wq : array<i32>;
111
+ @group(0) @binding(2) var<storage, read> lut : array<i32>;
112
+ @group(0) @binding(3) var<storage, read> rs : array<f32>; // per (batch,row)
113
+ @group(0) @binding(4) var<storage, read> cs : array<f32>; // per (batch,col)
114
+ ${OUT_DECL(verify, 5)}
115
+ @group(0) @binding(6) var<uniform> dims : vec4<u32>; // m, k, n, flags(1=relu)
116
+ @compute @workgroup_size(8, 8, 1)
117
+ fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
118
+ let m = dims.x; let k = dims.y; let n = dims.z;
119
+ let row = gid.x; let col = gid.y; let bz = gid.z;
120
+ if (row >= m || col >= n) { return; }
121
+ var s : i32 = 0;
122
+ let xo = (bz * m + row) * k;
123
+ let wo = bz * k * n + col;
124
+ for (var p = 0u; p < k; p = p + 1u) {
125
+ let au = u32(Xq[xo + p] & 255);
126
+ let bu = u32(Wq[wo + p * n] & 255);
127
+ s = s + lut[au * 256u + bu];
128
+ }
129
+ ${verify ? `O[(bz * m + row) * n + col] = s;` : `
130
+ var v = f32(s) * rs[bz * m + row] * cs[bz * n + col];
131
+ if ((dims.w & 1u) == 1u && v < 0.0) { v = 0.0; }
132
+ O[(bz * m + row) * n + col] = v;`}
133
+ }`;
134
+
135
+ // same fused/batched kernel on the DP4A hardware path (Wᵀ packed per batch)
136
+ const WGSL_BG_DP4 = (verify) => `
137
+ @group(0) @binding(0) var<storage, read> Xp : array<u32>;
138
+ @group(0) @binding(1) var<storage, read> Wp : array<u32>; // per-batch Wᵀ, packed
139
+ @group(0) @binding(2) var<storage, read> rs : array<f32>;
140
+ @group(0) @binding(3) var<storage, read> cs : array<f32>;
141
+ ${OUT_DECL(verify, 4)}
142
+ @group(0) @binding(5) var<uniform> dims : vec4<u32>; // m, kw, n, flags
143
+ @compute @workgroup_size(8, 8, 1)
144
+ fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
145
+ let m = dims.x; let kw = dims.y; let n = dims.z;
146
+ let row = gid.x; let col = gid.y; let bz = gid.z;
147
+ if (row >= m || col >= n) { return; }
148
+ var s : i32 = 0;
149
+ let xo = (bz * m + row) * kw;
150
+ let wo = (bz * n + col) * kw;
151
+ for (var p = 0u; p < kw; p = p + 1u) {
152
+ s = s + dot4I8Packed(Xp[xo + p], Wp[wo + p]);
153
+ }
154
+ ${verify ? `O[(bz * m + row) * n + col] = s;` : `
155
+ var v = f32(s) * rs[bz * m + row] * cs[bz * n + col];
156
+ if ((dims.w & 1u) == 1u && v < 0.0) { v = 0.0; }
157
+ O[(bz * m + row) * n + col] = v;`}
158
+ }`;
159
+
160
+ // Gather-fused attention (CUTLASS ex. 36/52): kernels index q/k/v directly in
161
+ // their BT×C layout (head-strided) — no gather copies, no kᵀ transpose — and
162
+ // the ctx kernel scatters straight back into BT×C. int8×int8→i32 is exact, so
163
+ // these are bit-identical to the LUT mirrors (proved at init before use).
164
+ const WGSL_ATT_SCORES = (verify) => `
165
+ @group(0) @binding(0) var<storage, read> Q : array<i32>; // int8 per elem, BT×C
166
+ @group(0) @binding(1) var<storage, read> K : array<i32>;
167
+ @group(0) @binding(2) var<storage, read> qs : array<f32>; // per (token,head)
168
+ @group(0) @binding(3) var<storage, read> ks : array<f32>;
169
+ ${OUT_DECL(verify, 4)}
170
+ @group(0) @binding(5) var<uniform> dims : vec4<u32>; // T, heads, hd, _
171
+ @compute @workgroup_size(8, 8, 1)
172
+ fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
173
+ let T = dims.x; let heads = dims.y; let hd = dims.z;
174
+ let ti = gid.x; let tj = gid.y; let bz = gid.z;
175
+ if (ti >= T || tj >= T) { return; }
176
+ let bi = bz / heads; let h = bz % heads;
177
+ let C = heads * hd;
178
+ let qo = (bi * T + ti) * C + h * hd;
179
+ let ko = (bi * T + tj) * C + h * hd;
180
+ var s : i32 = 0;
181
+ for (var p = 0u; p < hd; p = p + 1u) { s = s + Q[qo + p] * K[ko + p]; }
182
+ ${verify ? `O[(bz * T + ti) * T + tj] = s;`
183
+ : `O[(bz * T + ti) * T + tj] = f32(s) * qs[(bi * T + ti) * heads + h] * ks[(bi * T + tj) * heads + h];`}
184
+ }`;
185
+ const WGSL_ATT_CTX = (verify) => `
186
+ @group(0) @binding(0) var<storage, read> A : array<i32>; // int8, BH×T×T
187
+ @group(0) @binding(1) var<storage, read> V : array<i32>; // int8, BT×C
188
+ @group(0) @binding(2) var<storage, read> as_ : array<f32>; // per (bz,row)
189
+ @group(0) @binding(3) var<storage, read> vs : array<f32>; // per (batch,head,chan)
190
+ ${OUT_DECL(verify, 4)} // BT×C (scatter fused)
191
+ @group(0) @binding(5) var<uniform> dims : vec4<u32>; // T, heads, hd, _
192
+ @compute @workgroup_size(8, 8, 1)
193
+ fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
194
+ let T = dims.x; let heads = dims.y; let hd = dims.z;
195
+ let ti = gid.x; let j = gid.y; let bz = gid.z;
196
+ if (ti >= T || j >= hd) { return; }
197
+ let bi = bz / heads; let h = bz % heads;
198
+ let C = heads * hd;
199
+ let ao = (bz * T + ti) * T;
200
+ var s : i32 = 0;
201
+ for (var tj = 0u; tj < T; tj = tj + 1u) { s = s + A[ao + tj] * V[(bi * T + tj) * C + h * hd + j]; }
202
+ ${verify ? `O[(bi * T + ti) * C + h * hd + j] = s;`
203
+ : `O[(bi * T + ti) * C + h * hd + j] = f32(s) * as_[bz * T + ti] * vs[(bi * heads + h) * hd + j];`}
204
+ }`;
205
+
206
+ // Split-K f32 GEMM (CUTLASS ex. 06) for the STE BACKWARD only — the backward
207
+ // was always float (the integer path has no gradient); this just moves that
208
+ // exact float math off the JS thread. Split-K matters for dlnf: M=256, N=32,
209
+ // K=16512 — 8k outputs with a huge inner loop would idle the GPU, so slices
210
+ // of K run on separate workgroups and a second tiny pass reduces partials.
211
+ const WGSL_FGEMM = `
212
+ @group(0) @binding(0) var<storage, read> A : array<f32>;
213
+ @group(0) @binding(1) var<storage, read> Bm : array<f32>;
214
+ @group(0) @binding(2) var<storage, read_write> P : array<f32>; // S partials
215
+ @group(0) @binding(3) var<uniform> dims : vec4<u32>; // m, k, n, flags(bit0=transA, rest=S)
216
+ @compute @workgroup_size(8, 8, 1)
217
+ fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
218
+ let m = dims.x; let k = dims.y; let n = dims.z;
219
+ let transA = (dims.w & 1u) == 1u;
220
+ let S = dims.w >> 1u;
221
+ let row = gid.x; let col = gid.y; let z = gid.z;
222
+ if (row >= m || col >= n) { return; }
223
+ let ks = (k + S - 1u) / S;
224
+ let p0 = z * ks;
225
+ let p1 = min(k, p0 + ks);
226
+ var s : f32 = 0.0;
227
+ for (var p = p0; p < p1; p = p + 1u) {
228
+ let a = select(A[row * k + p], A[p * m + row], transA);
229
+ s = s + a * Bm[p * n + col];
230
+ }
231
+ P[(z * m + row) * n + col] = s;
232
+ }`;
233
+ const WGSL_FREDUCE = `
234
+ @group(0) @binding(0) var<storage, read> P : array<f32>;
235
+ @group(0) @binding(1) var<storage, read_write> O : array<f32>;
236
+ @group(0) @binding(2) var<uniform> dims : vec4<u32>; // mn, S, _, _
237
+ @compute @workgroup_size(64)
238
+ fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
239
+ let mn = dims.x; let S = dims.y;
240
+ let i = gid.x;
241
+ if (i >= mn) { return; }
242
+ var s : f32 = 0.0;
243
+ for (var z = 0u; z < S; z = z + 1u) { s = s + P[z * mn + i]; }
244
+ O[i] = s;
245
+ }`;
246
+
247
+ async function loadLUTs(base) {
248
+ base = base || "";
249
+ const [mulB, reqB, reluB, meta] = await Promise.all([
250
+ fetch(base + "mul_lut.bin").then(r => r.arrayBuffer()),
251
+ fetch(base + "requant_lut.bin").then(r => r.arrayBuffer()),
252
+ fetch(base + "relu_lut.bin").then(r => r.arrayBuffer()),
253
+ fetch(base + "luts_meta.json").then(r => r.json()),
254
+ ]);
255
+ return { mul: new Int16Array(mulB), requant: new Int8Array(reqB),
256
+ relu: new Int8Array(reluB), shift: meta.shift };
257
+ }
258
+
259
+ // pack a row-major int8 matrix (rows×cols) into u32 words of 4 bytes along
260
+ // cols, zero-padded to kw words per row (zeros contribute 0 to the dot)
261
+ function packRows(Q, rows, cols, kw) {
262
+ const out = new Uint32Array(rows * kw);
263
+ const bytes = new Uint8Array(out.buffer);
264
+ for (let r = 0; r < rows; r++)
265
+ for (let c = 0; c < cols; c++) bytes[(r * kw * 4) + c] = Q[r * cols + c] & 0xFF;
266
+ return out;
267
+ }
268
+ function transposeI8(Q, rows, cols) {
269
+ const out = new Int8Array(rows * cols);
270
+ for (let r = 0; r < rows; r++) for (let c = 0; c < cols; c++) out[c * rows + r] = Q[r * cols + c];
271
+ return out;
272
+ }
273
+
274
+ async function initCompute(L) {
275
+ const cpu = { backend: "cpu", label: "CPU (JS)",
276
+ matmulInt8: (Xq, Wq, m, k, n, LL) => root.Verified.lutMatmulJS(Xq, Wq, m, k, n, LL) };
277
+ if (!(root.navigator && navigator.gpu)) return cpu;
278
+ try {
279
+ const adapter = await navigator.gpu.requestAdapter();
280
+ if (!adapter) return cpu;
281
+ const device = await adapter.requestDevice();
282
+ const info = adapter.info || {};
283
+ const gpuName = info.description || info.vendor || "WebGPU";
284
+
285
+ // LUT pipeline (always built — the fallback and the verification oracle)
286
+ const lutModule = device.createShaderModule({ code: WGSL_LUT });
287
+ const lutPipe = device.createComputePipeline({ layout: "auto", compute: { module: lutModule, entryPoint: "main" } });
288
+ const lut32 = new Int32Array(L.mul); // widen int16 -> int32
289
+ const lutBuf = device.createBuffer({ size: lut32.byteLength, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST });
290
+ device.queue.writeBuffer(lutBuf, 0, lut32);
291
+ const mkPipe = (code) => device.createComputePipeline({ layout: "auto",
292
+ compute: { module: device.createShaderModule({ code }), entryPoint: "main" } });
293
+ // The verify variant doesn't reference the scale buffers, so `layout:auto`
294
+ // would strip those bindings and the bind group would silently mismatch.
295
+ // An EXPLICIT layout keeps both variants binding-compatible — which is the
296
+ // point: they must differ only in the final write, nothing else.
297
+ const mkLayout = (spec) => device.createBindGroupLayout({
298
+ entries: spec.map((t, i) => ({ binding: i, visibility: GPUShaderStage.COMPUTE,
299
+ buffer: { type: t === "u" ? "uniform" : t === "rw" ? "storage" : "read-only-storage" } })) });
300
+ const mkPipeL = (code, layout) => device.createComputePipeline({
301
+ layout: device.createPipelineLayout({ bindGroupLayouts: [layout] }),
302
+ compute: { module: device.createShaderModule({ code }), entryPoint: "main" } });
303
+ const bgLutLayout = mkLayout(["r", "r", "r", "r", "r", "rw", "u"]);
304
+ const bgDp4Layout = mkLayout(["r", "r", "r", "r", "rw", "u"]);
305
+ const attLayout = mkLayout(["r", "r", "r", "r", "rw", "u"]);
306
+ const rowmaxLayout = mkLayout(["r", "rw", "u"]);
307
+ const quantLayout = mkLayout(["r", "r", "rw", "u"]);
308
+ const rowmaxPipe = mkPipeL(WGSL_ROWMAX, rowmaxLayout);
309
+ const quantI32Pipe = mkPipeL(WGSL_QUANT(false), quantLayout);
310
+ const quantPackPipe = mkPipeL(WGSL_QUANT(true), quantLayout);
311
+ // live + verify variants, compiled from the same source (see WGSL_* above)
312
+ const bgLutPipe = mkPipeL(WGSL_BG_LUT(false), bgLutLayout), bgLutVPipe = mkPipeL(WGSL_BG_LUT(true), bgLutLayout);
313
+ const scoresPipe = mkPipeL(WGSL_ATT_SCORES(false), attLayout), scoresVPipe = mkPipeL(WGSL_ATT_SCORES(true), attLayout);
314
+ const ctxPipe = mkPipeL(WGSL_ATT_CTX(false), attLayout), ctxVPipe = mkPipeL(WGSL_ATT_CTX(true), attLayout);
315
+
316
+ // gather-fused attention kernels. The gate runs the VERIFY variant of the
317
+ // same source and compares the int32 accumulator with `!==` — exact, no
318
+ // tolerance — then checks the fused epilogue against a bit-exact JS mirror
319
+ // of the WGSL rounding. Swept over several shapes, incl. odd/ragged ones,
320
+ // because head-strided addressing is where these kernels can go wrong.
321
+ let att = { scores: (qq, kq, qs, ks, d) => gpuAttScores(device, d.acc ? scoresVPipe : scoresPipe, qq, kq, qs, ks, d),
322
+ ctx: (aq, vq, as, vs, d) => gpuAttCtx(device, d.acc ? ctxVPipe : ctxPipe, aq, vq, as, vs, d) };
323
+ try {
324
+ for (const d0 of [{ B: 2, T: 8, heads: 2, hd: 8 }, { B: 1, T: 32, heads: 2, hd: 16 },
325
+ { B: 3, T: 7, heads: 3, hd: 5 }, { B: 2, T: 33, heads: 4, hd: 8 }]) {
326
+ const nQ = d0.B * d0.T * d0.heads * d0.hd;
327
+ const qq = new Int8Array(nQ), kq = new Int8Array(nQ), vq = new Int8Array(nQ);
328
+ for (let i = 0; i < nQ; i++) { qq[i] = (Math.random() * 256 - 128) | 0; kq[i] = (Math.random() * 256 - 128) | 0; vq[i] = (Math.random() * 256 - 128) | 0; }
329
+ const qs = Float32Array.from({ length: d0.B * d0.T * d0.heads }, () => Math.random() + 0.5);
330
+ const ks = Float32Array.from({ length: d0.B * d0.T * d0.heads }, () => Math.random() + 0.5);
331
+ const aq = new Int8Array(d0.B * d0.heads * d0.T * d0.T);
332
+ for (let i = 0; i < aq.length; i++) aq[i] = (Math.random() * 127) | 0;
333
+ const as = Float32Array.from({ length: d0.B * d0.heads * d0.T }, () => Math.random() + 0.5);
334
+ const vs = Float32Array.from({ length: d0.B * d0.heads * d0.hd }, () => Math.random() + 0.5);
335
+ const dv = { ...d0, acc: true };
336
+ const [accS, accC, hwS, hwC] = await Promise.all([
337
+ att.scores(qq, kq, qs, ks, dv), att.ctx(aq, vq, as, vs, dv),
338
+ att.scores(qq, kq, qs, ks, d0), att.ctx(aq, vq, as, vs, d0)]);
339
+ const refAccS = root.Verified.attScoresJS(qq, kq, qs, ks, dv, L);
340
+ const refAccC = root.Verified.attCtxJS(aq, vq, as, vs, dv, L);
341
+ for (let i = 0; i < refAccS.length; i++) if (accS[i] !== refAccS[i]) throw new Error(`scores accumulator mismatch @${i} shape ${JSON.stringify(d0)}`);
342
+ for (let i = 0; i < refAccC.length; i++) if (accC[i] !== refAccC[i]) throw new Error(`ctx accumulator mismatch @${i} shape ${JSON.stringify(d0)}`);
343
+ const refS = root.Verified.attScoresJS(qq, kq, qs, ks, d0, L);
344
+ const refC = root.Verified.attCtxJS(aq, vq, as, vs, d0, L);
345
+ for (let i = 0; i < refS.length; i++) if (hwS[i] !== refS[i]) throw new Error(`scores epilogue mismatch @${i}`);
346
+ for (let i = 0; i < refC.length; i++) if (hwC[i] !== refC[i]) throw new Error(`ctx epilogue mismatch @${i}`);
347
+ }
348
+ } catch (e) {
349
+ console.warn("fused attention kernels failed verification — using CPU LUT mirrors:", e.message);
350
+ att = null;
351
+ }
352
+
353
+ // split-K f32 GEMM for the STE backward (self-tested vs JS float matmul)
354
+ const fPipes = { gemm: mkPipe(WGSL_FGEMM), reduce: mkPipe(WGSL_FREDUCE) };
355
+ let fgemm = (A, Bm, d) => gpuFgemm(device, fPipes, A, Bm, d);
356
+ try {
357
+ const m0 = 7, k0 = 4500, n0 = 5; // k big enough to exercise split-K
358
+ const A = Float32Array.from({ length: m0 * k0 }, () => Math.random() - 0.5);
359
+ const Bm = Float32Array.from({ length: k0 * n0 }, () => Math.random() - 0.5);
360
+ const hw = await fgemm(A, Bm, { m: m0, k: k0, n: n0 });
361
+ const ref = root.TrainCore.matmul(A, Bm, m0, k0, n0);
362
+ for (let i = 0; i < ref.length; i++)
363
+ if (Math.abs(hw[i] - ref[i]) > Math.abs(ref[i]) * 1e-3 + 1e-3) throw new Error("fgemm mismatch");
364
+ } catch (e) {
365
+ console.warn("split-K f32 GEMM failed verification — backward stays in JS:", e.message);
366
+ fgemm = null;
367
+ }
368
+
369
+ // Shared exact gate for a bgemm implementation. Sweeps shapes (including
370
+ // ragged ones and a k long enough to matter), compares the int32
371
+ // accumulator from the verify variant with `!==`, then compares the fused
372
+ // f32 epilogue against the bit-exact JS mirror. Returns null if clean.
373
+ async function gateBgemm(bgFn) {
374
+ for (const d0 of [{ m: 5, k: 9, n: 6, batch: 3, relu: true },
375
+ { m: 32, k: 64, n: 32, batch: 1, relu: false },
376
+ { m: 7, k: 253, n: 5, batch: 2, relu: true },
377
+ { m: 1, k: 4, n: 1, batch: 1, relu: false },
378
+ { m: 17, k: 33, n: 9, batch: 1, relu: true }]) {
379
+ const Xq = new Int8Array(d0.batch * d0.m * d0.k), Wq = new Int8Array(d0.batch * d0.k * d0.n);
380
+ for (let i = 0; i < Xq.length; i++) Xq[i] = (Math.random() * 256 - 128) | 0;
381
+ for (let i = 0; i < Wq.length; i++) Wq[i] = (Math.random() * 256 - 128) | 0;
382
+ const rs = Float32Array.from({ length: d0.batch * d0.m }, () => Math.random() + 0.5);
383
+ const cs = Float32Array.from({ length: d0.batch * d0.n }, () => Math.random() + 0.5);
384
+ const shape = `${d0.m}x${d0.k}x${d0.n}b${d0.batch}`;
385
+ const accHw = await bgFn(Xq, Wq, rs, cs, { ...d0, acc: true });
386
+ const accRef = root.Verified.bgemmJS(Xq, Wq, rs, cs, { ...d0, acc: true }, L);
387
+ for (let i = 0; i < accRef.length; i++)
388
+ if (accHw[i] !== accRef[i]) return `accumulator mismatch @${i} (${shape}): ${accHw[i]} vs ${accRef[i]}`;
389
+ const hw = await bgFn(Xq, Wq, rs, cs, d0);
390
+ const ref = root.Verified.bgemmJS(Xq, Wq, rs, cs, d0, L);
391
+ for (let i = 0; i < ref.length; i++)
392
+ if (hw[i] !== ref[i]) return `epilogue mismatch @${i} (${shape}): ${hw[i]} vs ${ref[i]}`;
393
+ }
394
+ return null;
395
+ }
396
+
397
+ const bgLut = (Xq, Wq, rs, cs, d) => gpuBgemmLUT(device, d.acc ? bgLutVPipe : bgLutPipe, lutBuf, Xq, Wq, rs, cs, d);
398
+ // the LUT bgemm is the fallback AND the oracle's shader twin — gate it too
399
+ const lutBad = await gateBgemm(bgLut);
400
+ if (lutBad) { console.warn("LUT bgemm shader failed verification — CPU mirrors only:", lutBad); return cpu; }
401
+
402
+ // B2B MLP chain gate (CUTLASS ex. 13+23): run the WHOLE chain — gemm1 +
403
+ // ReLU + on-GPU rowmax + on-GPU quantize + gemm2 — against the pure-JS
404
+ // mirror chain, exact `!==` on both h1 and the final output. Sweeps
405
+ // ragged shapes; h not a multiple of 4 exercises the pack-tail padding.
406
+ async function gateMlp(mlpFn) {
407
+ for (const d0 of [{ m: 6, k: 8, h: 12, n: 5 }, { m: 5, k: 16, h: 6, n: 3 },
408
+ { m: 17, k: 33, h: 10, n: 9 }, { m: 32, k: 64, h: 128, n: 32 }]) {
409
+ const rnd = (len) => Float32Array.from({ length: len }, () => Math.random() * 2 - 1);
410
+ const Xf = rnd(d0.m * d0.k), W1 = rnd(d0.k * d0.h), W2 = rnd(d0.h * d0.n);
411
+ const hw = await root.Verified.vmlpBlock(Xf, W1, W2, d0, L, mlpFn, null);
412
+ const ref = await root.Verified.vmlpBlock(Xf, W1, W2, d0, L, null, null);
413
+ const shape = `${d0.m}x${d0.k}x${d0.h}x${d0.n}`;
414
+ for (let i = 0; i < ref.h1.length; i++)
415
+ if (hw.h1[i] !== ref.h1[i]) return `h1 mismatch @${i} (${shape}): ${hw.h1[i]} vs ${ref.h1[i]}`;
416
+ for (let i = 0; i < ref.out.length; i++)
417
+ if (hw.out[i] !== ref.out[i]) return `out mismatch @${i} (${shape}): ${hw.out[i]} vs ${ref.out[i]}`;
418
+ }
419
+ // DISCRIMINATING case: the sweep above passes vacuously if no value
420
+ // lands on a rounding boundary — the old round(x/scale) spec would
421
+ // pass it too. So hunt (in fast JS) for an input where the two specs
422
+ // actually disagree, then check the GPU sides with the RESPEC. This
423
+ // gates the gate: a pass must be something the old spec would fail.
424
+ const V = root.Verified;
425
+ const d1 = { m: 16, k: 32, h: 64, n: 16 };
426
+ const rnd1 = (len) => Float32Array.from({ length: len }, () => Math.random() * 4 - 2);
427
+ for (let t = 0; t < 800; t++) {
428
+ const Xf = rnd1(d1.m * d1.k), W1 = rnd1(d1.k * d1.h), W2 = rnd1(d1.h * d1.n);
429
+ const x = V.quantizeRows(Xf, d1.m, d1.k), w1 = V.quantizeCols(W1, d1.k, d1.h);
430
+ const h1 = V.bgemmJS(x.q, w1.q, x.s, w1.s, { m: d1.m, k: d1.k, n: d1.h, batch: 1, relu: true }, L);
431
+ const sc = V.scalesFromAbsMax(V.rowAbsMax(h1, d1.m, d1.h));
432
+ const qNew = V.quantizeRowsInv(h1, d1.m, d1.h, sc.inv);
433
+ const qOld = V.quantizeRows(h1, d1.m, d1.h).q;
434
+ let boundary = false;
435
+ for (let i = 0; i < qNew.length; i++) if (qNew[i] !== qOld[i]) { boundary = true; break; }
436
+ if (!boundary) continue;
437
+ const w2 = V.quantizeCols(W2, d1.h, d1.n);
438
+ const gpu = await mlpFn(x.q, w1.q, w2.q, x.s, w1.s, w2.s, d1);
439
+ const refNew = V.bgemmJS(qNew, w2.q, sc.scale, w2.s, { m: d1.m, k: d1.h, n: d1.n, batch: 1 }, L);
440
+ const refOld = V.bgemmJS(qOld, w2.q, sc.scale, w2.s, { m: d1.m, k: d1.h, n: d1.n, batch: 1 }, L);
441
+ let eqNew = true, eqOld = true;
442
+ for (let i = 0; i < refNew.length; i++) { if (gpu.out[i] !== refNew[i]) eqNew = false; if (gpu.out[i] !== refOld[i]) eqOld = false; }
443
+ if (!eqNew) return "discriminating boundary case: GPU chain does not match the respec mirror";
444
+ if (eqOld) return "discriminating boundary case: GPU chain matches the OLD quantize spec";
445
+ return null; // proven: respec, not merely gate-compatible
446
+ }
447
+ console.warn("B2B gate: no rounding-boundary input found in 800 trials — respec discrimination unproven this boot (sweep still exact)");
448
+ return null;
449
+ }
450
+ const lutMlpEnv = { dp4: false, gemm: bgLutPipe, rowmax: rowmaxPipe, quant: quantI32Pipe, lutBuf };
451
+ let mlpLut = (xq, w1q, w2q, xs, w1s, w2s, d) => gpuMlpChain(device, lutMlpEnv, xq, w1q, w2q, xs, w1s, w2s, d);
452
+ const mlpLutBad = await gateMlp(mlpLut);
453
+ if (mlpLutBad) { console.warn("B2B MLP chain (LUT) failed verification — MLP stays on the CPU mirror chain:", mlpLutBad); mlpLut = null; }
454
+
455
+ const viaLUT = { backend: "webgpu", label: `${gpuName} (LUT shader · exact-gated)`,
456
+ matmulInt8: (Xq, Wq, m, k, n) => gpuMatmulLUT(device, lutPipe, lutBuf, Xq, Wq, m, k, n),
457
+ bgemm: bgLut, att, fgemm, mlp: mlpLut };
458
+
459
+ // DP4A pipeline — only if the WGSL feature exists AND its batched kernel
460
+ // reproduces the verified units exactly across the shape sweep
461
+ if (!(navigator.gpu.wgslLanguageFeatures && navigator.gpu.wgslLanguageFeatures.has("packed_4x8_integer_dot_product")))
462
+ return viaLUT;
463
+ const bgDp4Pipe = mkPipeL(WGSL_BG_DP4(false), bgDp4Layout), bgDp4VPipe = mkPipeL(WGSL_BG_DP4(true), bgDp4Layout);
464
+ const bg = (Xq, Wq, rs, cs, d) => gpuBgemmDP4(device, d.acc ? bgDp4VPipe : bgDp4Pipe, Xq, Wq, rs, cs, d);
465
+ const dp4Bad = await gateBgemm(bg);
466
+ if (dp4Bad) {
467
+ console.warn("batched DP4A disagreed with the verified units — using LUT bgemm:", dp4Bad);
468
+ return viaLUT;
469
+ }
470
+ const dp4MlpEnv = { dp4: true, gemm: bgDp4Pipe, rowmax: rowmaxPipe, quant: quantPackPipe };
471
+ let mlpDp4 = (xq, w1q, w2q, xs, w1s, w2s, d) => gpuMlpChain(device, dp4MlpEnv, xq, w1q, w2q, xs, w1s, w2s, d);
472
+ const mlpDp4Bad = await gateMlp(mlpDp4);
473
+ if (mlpDp4Bad) { console.warn("B2B MLP chain (DP4A) failed verification — using the LUT chain:", mlpDp4Bad); mlpDp4 = mlpLut; }
474
+ return { backend: "webgpu", label: `${gpuName} (DP4A int8 dot · exact-gated vs units)`,
475
+ bgemm: bg, att, fgemm, mlp: mlpDp4 };
476
+ } catch (e) { console.warn("WebGPU init failed, CPU fallback:", e); return cpu; }
477
+ }
478
+
479
+ // shared dispatch/readback plumbing
480
+ async function runPass(device, pipeline, entries, m, n) {
481
+ const bytesC = m * n * 4;
482
+ const bufC = mk(device, bytesC, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
483
+ const bind = device.createBindGroup({ layout: pipeline.getBindGroupLayout(0),
484
+ entries: entries(bufC) });
485
+ const enc = device.createCommandEncoder();
486
+ const pass = enc.beginComputePass();
487
+ pass.setPipeline(pipeline); pass.setBindGroup(0, bind);
488
+ pass.dispatchWorkgroups(Math.ceil(m / 8), Math.ceil(n / 8)); pass.end();
489
+ const read = mk(device, bytesC, GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ);
490
+ enc.copyBufferToBuffer(bufC, 0, read, 0, bytesC);
491
+ device.queue.submit([enc.finish()]);
492
+ await read.mapAsync(GPUMapMode.READ);
493
+ const out = new Int32Array(read.getMappedRange().slice(0));
494
+ read.unmap();
495
+ return { out, bufC, read };
496
+ }
497
+
498
+ async function gpuMatmulLUT(device, pipeline, lutBuf, Xq, Wq, m, k, n) {
499
+ const X32 = Int32Array.from(Xq), W32 = Int32Array.from(Wq); // byte -> i32
500
+ const bufX = up(device, X32, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
501
+ const bufW = up(device, W32, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
502
+ const bufD = up(device, new Uint32Array([m, k, n, 0]), GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST);
503
+ const r = await runPass(device, pipeline, (bufC) => [
504
+ { binding: 0, resource: { buffer: bufX } }, { binding: 1, resource: { buffer: bufW } },
505
+ { binding: 2, resource: { buffer: lutBuf } }, { binding: 3, resource: { buffer: bufC } },
506
+ { binding: 4, resource: { buffer: bufD } } ], m, n);
507
+ [bufX, bufW, bufD, r.bufC, r.read].forEach(b => b.destroy());
508
+ return r.out;
509
+ }
510
+
511
+ // attention kernels: int8 (widened i32) in, f32 out, strided head indexing
512
+ async function gpuAttScores(device, pipeline, qq, kq, qs, ks, d) {
513
+ const { B, T, heads } = d;
514
+ const bufQ = up(device, Int32Array.from(qq), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
515
+ const bufK = up(device, Int32Array.from(kq), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
516
+ const bufQs = up(device, qs, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
517
+ const bufKs = up(device, ks, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
518
+ const bufD = up(device, new Uint32Array([d.T, d.heads, d.hd, 0]), GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST);
519
+ const r = await runBgPass(device, pipeline, (bufO) => [
520
+ { binding: 0, resource: { buffer: bufQ } }, { binding: 1, resource: { buffer: bufK } },
521
+ { binding: 2, resource: { buffer: bufQs } }, { binding: 3, resource: { buffer: bufKs } },
522
+ { binding: 4, resource: { buffer: bufO } }, { binding: 5, resource: { buffer: bufD } } ], T, T, B * heads, d.acc);
523
+ [bufQ, bufK, bufQs, bufKs, bufD, r.bufO, r.read].forEach(b => b.destroy());
524
+ return r.out;
525
+ }
526
+ async function gpuAttCtx(device, pipeline, aq, vq, as, vs, d) {
527
+ const { B, T, heads, hd } = d;
528
+ const bufA = up(device, Int32Array.from(aq), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
529
+ const bufV = up(device, Int32Array.from(vq), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
530
+ const bufAs = up(device, as, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
531
+ const bufVs = up(device, vs, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
532
+ const bufD = up(device, new Uint32Array([T, heads, hd, 0]), GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST);
533
+ const r = await runBgPass(device, pipeline, (bufO) => [
534
+ { binding: 0, resource: { buffer: bufA } }, { binding: 1, resource: { buffer: bufV } },
535
+ { binding: 2, resource: { buffer: bufAs } }, { binding: 3, resource: { buffer: bufVs } },
536
+ { binding: 4, resource: { buffer: bufO } }, { binding: 5, resource: { buffer: bufD } } ], T, hd, B * heads, d.acc);
537
+ [bufA, bufV, bufAs, bufVs, bufD, r.bufO, r.read].forEach(b => b.destroy());
538
+ return r.out;
539
+ }
540
+
541
+ // split-K f32 GEMM (backward): partial pass + reduce pass
542
+ async function gpuFgemm(device, pipes, A, Bm, d) {
543
+ const { m, k, n } = d, transA = d.transA ? 1 : 0;
544
+ const S = k > 4096 ? Math.min(16, Math.ceil(k / 2048)) : 1;
545
+ const bufA = up(device, A, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
546
+ const bufB = up(device, Bm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
547
+ const bufP = mk(device, S * m * n * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
548
+ const bufD1 = up(device, new Uint32Array([m, k, n, transA | (S << 1)]), GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST);
549
+ const bufO = mk(device, m * n * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
550
+ const bufD2 = up(device, new Uint32Array([m * n, S, 0, 0]), GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST);
551
+ const enc = device.createCommandEncoder();
552
+ let pass = enc.beginComputePass();
553
+ pass.setPipeline(pipes.gemm);
554
+ pass.setBindGroup(0, device.createBindGroup({ layout: pipes.gemm.getBindGroupLayout(0), entries: [
555
+ { binding: 0, resource: { buffer: bufA } }, { binding: 1, resource: { buffer: bufB } },
556
+ { binding: 2, resource: { buffer: bufP } }, { binding: 3, resource: { buffer: bufD1 } } ] }));
557
+ pass.dispatchWorkgroups(Math.ceil(m / 8), Math.ceil(n / 8), S); pass.end();
558
+ pass = enc.beginComputePass();
559
+ pass.setPipeline(pipes.reduce);
560
+ pass.setBindGroup(0, device.createBindGroup({ layout: pipes.reduce.getBindGroupLayout(0), entries: [
561
+ { binding: 0, resource: { buffer: bufP } }, { binding: 1, resource: { buffer: bufO } },
562
+ { binding: 2, resource: { buffer: bufD2 } } ] }));
563
+ pass.dispatchWorkgroups(Math.ceil(m * n / 64)); pass.end();
564
+ const read = mk(device, m * n * 4, GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ);
565
+ enc.copyBufferToBuffer(bufO, 0, read, 0, m * n * 4);
566
+ device.queue.submit([enc.finish()]);
567
+ await read.mapAsync(GPUMapMode.READ);
568
+ const out = new Float32Array(read.getMappedRange().slice(0));
569
+ read.unmap();
570
+ [bufA, bufB, bufP, bufD1, bufO, bufD2, read].forEach(b => b.destroy());
571
+ return out;
572
+ }
573
+
574
+ // fused batched dispatch: f32 out, epilogue done on-device (raw=true reads the
575
+ // verify variant's int32 accumulator instead)
576
+ async function runBgPass(device, pipeline, entries, m, n, batch, raw) {
577
+ const bytesO = batch * m * n * 4;
578
+ const bufO = mk(device, bytesO, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
579
+ const bind = device.createBindGroup({ layout: pipeline.getBindGroupLayout(0), entries: entries(bufO) });
580
+ const enc = device.createCommandEncoder();
581
+ const pass = enc.beginComputePass();
582
+ pass.setPipeline(pipeline); pass.setBindGroup(0, bind);
583
+ pass.dispatchWorkgroups(Math.ceil(m / 8), Math.ceil(n / 8), batch); pass.end();
584
+ const read = mk(device, bytesO, GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ);
585
+ enc.copyBufferToBuffer(bufO, 0, read, 0, bytesO);
586
+ device.queue.submit([enc.finish()]);
587
+ await read.mapAsync(GPUMapMode.READ);
588
+ const buf = read.getMappedRange().slice(0);
589
+ const out = raw ? new Int32Array(buf) : new Float32Array(buf);
590
+ read.unmap();
591
+ return { out, bufO, read };
592
+ }
593
+
594
+ async function gpuBgemmLUT(device, pipeline, lutBuf, Xq, Wq, rs, cs, d) {
595
+ const { m, k, n } = d, batch = d.batch || 1, flags = d.relu ? 1 : 0;
596
+ const bufX = up(device, Int32Array.from(Xq), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
597
+ const bufW = up(device, Int32Array.from(Wq), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
598
+ const bufR = up(device, rs, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
599
+ const bufS = up(device, cs, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
600
+ const bufD = up(device, new Uint32Array([m, k, n, flags]), GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST);
601
+ const r = await runBgPass(device, pipeline, (bufO) => [
602
+ { binding: 0, resource: { buffer: bufX } }, { binding: 1, resource: { buffer: bufW } },
603
+ { binding: 2, resource: { buffer: lutBuf } }, { binding: 3, resource: { buffer: bufR } },
604
+ { binding: 4, resource: { buffer: bufS } }, { binding: 5, resource: { buffer: bufO } },
605
+ { binding: 6, resource: { buffer: bufD } } ], m, n, batch, d.acc);
606
+ [bufX, bufW, bufR, bufS, bufD, r.bufO, r.read].forEach(b => b.destroy());
607
+ return r.out;
608
+ }
609
+
610
+ async function gpuBgemmDP4(device, pipeline, Xq, Wq, rs, cs, d) {
611
+ const { m, k, n } = d, batch = d.batch || 1, flags = d.relu ? 1 : 0;
612
+ const kw = Math.ceil(k / 4);
613
+ const Xp = packRows(Xq, batch * m, k, kw); // rows are (batch·m)
614
+ const Wp = new Uint32Array(batch * n * kw); // per-batch Wᵀ, packed
615
+ for (let bz = 0; bz < batch; bz++) {
616
+ const wt = transposeI8(Wq.subarray(bz * k * n, (bz + 1) * k * n), k, n);
617
+ Wp.set(packRows(wt, n, k, kw), bz * n * kw);
618
+ }
619
+ const bufX = up(device, Xp, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
620
+ const bufW = up(device, Wp, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
621
+ const bufR = up(device, rs, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
622
+ const bufS = up(device, cs, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
623
+ const bufD = up(device, new Uint32Array([m, kw, n, flags]), GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST);
624
+ const r = await runBgPass(device, pipeline, (bufO) => [
625
+ { binding: 0, resource: { buffer: bufX } }, { binding: 1, resource: { buffer: bufW } },
626
+ { binding: 2, resource: { buffer: bufR } }, { binding: 3, resource: { buffer: bufS } },
627
+ { binding: 4, resource: { buffer: bufO } }, { binding: 5, resource: { buffer: bufD } } ], m, n, batch, d.acc);
628
+ [bufX, bufW, bufR, bufS, bufD, r.bufO, r.read].forEach(b => b.destroy());
629
+ return r.out;
630
+ }
631
+
632
+ // B2B MLP chain: gemm1 (ReLU fused) + rowmax in one encoder, a 4·m-byte
633
+ // absmax readback, then quantize + gemm2 in a second encoder. h1 comes back
634
+ // because the STE backward needs it, but it never goes UP again — gemm2's
635
+ // left operand is produced and consumed entirely on the GPU.
636
+ async function gpuMlpChain(device, env, xq, w1q, w2q, xs, w1s, w2s, d) {
637
+ const { m, k, h, n } = d, dp4 = !!env.dp4;
638
+ const SU = GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST;
639
+ let bufX, bufW1, bufW2, kw1, hw;
640
+ if (dp4) {
641
+ kw1 = Math.ceil(k / 4); hw = Math.ceil(h / 4);
642
+ bufX = up(device, packRows(xq, m, k, kw1), SU);
643
+ bufW1 = up(device, packRows(transposeI8(w1q, k, h), h, k, kw1), SU);
644
+ bufW2 = up(device, packRows(transposeI8(w2q, h, n), n, h, hw), SU);
645
+ } else {
646
+ kw1 = k; hw = h;
647
+ bufX = up(device, Int32Array.from(xq), SU);
648
+ bufW1 = up(device, Int32Array.from(w1q), SU);
649
+ bufW2 = up(device, Int32Array.from(w2q), SU);
650
+ }
651
+ const bufRs = up(device, xs, SU), bufCs1 = up(device, w1s, SU), bufCs2 = up(device, w2s, SU);
652
+ const bufH = mk(device, m * h * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
653
+ const bufMX = mk(device, m * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC); // zero-initialized
654
+ const UU = GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST;
655
+ const bufD1 = up(device, new Uint32Array([m, kw1, h, 1]), UU); // flags=1: fused ReLU
656
+ const bufDM = up(device, new Uint32Array([m, h, 0, 0]), UU);
657
+ const gemmBind = (bufA, bufB, bufR, bufC, bufO, bufD) => device.createBindGroup({
658
+ layout: env.gemm.getBindGroupLayout(0),
659
+ entries: (dp4 ? [bufA, bufB, bufR, bufC, bufO, bufD]
660
+ : [bufA, bufB, env.lutBuf, bufR, bufC, bufO, bufD])
661
+ .map((b, i) => ({ binding: i, resource: { buffer: b } })) });
662
+ const enc1 = device.createCommandEncoder();
663
+ let pass = enc1.beginComputePass();
664
+ pass.setPipeline(env.gemm);
665
+ pass.setBindGroup(0, gemmBind(bufX, bufW1, bufRs, bufCs1, bufH, bufD1));
666
+ pass.dispatchWorkgroups(Math.ceil(m / 8), Math.ceil(h / 8), 1); pass.end();
667
+ pass = enc1.beginComputePass();
668
+ pass.setPipeline(env.rowmax);
669
+ pass.setBindGroup(0, device.createBindGroup({ layout: env.rowmax.getBindGroupLayout(0), entries: [
670
+ { binding: 0, resource: { buffer: bufH } }, { binding: 1, resource: { buffer: bufMX } },
671
+ { binding: 2, resource: { buffer: bufDM } } ] }));
672
+ pass.dispatchWorkgroups(Math.ceil(m / 8), Math.ceil(h / 8), 1); pass.end();
673
+ const readH = mk(device, m * h * 4, GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ);
674
+ const readM = mk(device, m * 4, GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ);
675
+ enc1.copyBufferToBuffer(bufH, 0, readH, 0, m * h * 4);
676
+ enc1.copyBufferToBuffer(bufMX, 0, readM, 0, m * 4);
677
+ device.queue.submit([enc1.finish()]);
678
+ await Promise.all([readH.mapAsync(GPUMapMode.READ), readM.mapAsync(GPUMapMode.READ)]);
679
+ const h1 = new Float32Array(readH.getMappedRange().slice(0)); readH.unmap();
680
+ // the atomicMax'ed u32 bit patterns ARE the f32 |max| values
681
+ const mx = new Float32Array(readM.getMappedRange().slice(0)); readM.unmap();
682
+ // scale derivation in JS f64 — exactly rounded, identical on every device
683
+ // (WGSL division is 2.5 ULP, which is why it never runs on the GPU)
684
+ const sc = root.Verified.scalesFromAbsMax(mx);
685
+ const bufInv = up(device, sc.inv, SU), bufHs = up(device, sc.scale, SU);
686
+ const bufQ = mk(device, m * hw * 4, GPUBufferUsage.STORAGE);
687
+ const bufDQ = up(device, new Uint32Array([m, h, hw, 0]), UU);
688
+ const bufD2 = up(device, new Uint32Array([m, hw, n, 0]), UU);
689
+ const bufO = mk(device, m * n * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
690
+ const enc2 = device.createCommandEncoder();
691
+ pass = enc2.beginComputePass();
692
+ pass.setPipeline(env.quant);
693
+ pass.setBindGroup(0, device.createBindGroup({ layout: env.quant.getBindGroupLayout(0), entries: [
694
+ { binding: 0, resource: { buffer: bufH } }, { binding: 1, resource: { buffer: bufInv } },
695
+ { binding: 2, resource: { buffer: bufQ } }, { binding: 3, resource: { buffer: bufDQ } } ] }));
696
+ pass.dispatchWorkgroups(Math.ceil((m * (dp4 ? hw : h)) / 64)); pass.end();
697
+ pass = enc2.beginComputePass();
698
+ pass.setPipeline(env.gemm);
699
+ pass.setBindGroup(0, gemmBind(bufQ, bufW2, bufHs, bufCs2, bufO, bufD2));
700
+ pass.dispatchWorkgroups(Math.ceil(m / 8), Math.ceil(n / 8), 1); pass.end();
701
+ const readO = mk(device, m * n * 4, GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ);
702
+ enc2.copyBufferToBuffer(bufO, 0, readO, 0, m * n * 4);
703
+ device.queue.submit([enc2.finish()]);
704
+ await readO.mapAsync(GPUMapMode.READ);
705
+ const out = new Float32Array(readO.getMappedRange().slice(0)); readO.unmap();
706
+ [bufX, bufW1, bufW2, bufRs, bufCs1, bufCs2, bufH, bufMX, bufD1, bufDM,
707
+ bufInv, bufHs, bufQ, bufDQ, bufD2, bufO, readH, readM, readO].forEach(b => b.destroy());
708
+ return { h1, out };
709
+ }
710
+
711
+ function mk(device, size, usage) { return device.createBuffer({ size, usage }); }
712
+ function up(device, arr, usage) {
713
+ const b = mk(device, Math.max(16, arr.byteLength), usage);
714
+ device.queue.writeBuffer(b, 0, arr);
715
+ return b;
716
+ }
717
+
718
+ // ---- canonical kernel probe -------------------------------------------------
719
+ // The weight hash CANNOT catch a device whose kernel is wrong: weights only
720
+ // depend on the gradient bytes everyone receives, so a fleet averaging one
721
+ // device's bad gradient stays bit-identical and perfectly happy. This is the
722
+ // check that can. Every device runs the SAME seeded int8 GEMM through its own
723
+ // live kernel (verify variant -> raw int32 accumulator, exact on every
724
+ // backend — GPU DP4A, GPU LUT, CPU mirror alike) and hashes the result. Same
725
+ // input + correct kernels => same hash, regardless of hardware. A device that
726
+ // disagrees is computing different arithmetic than the rest of the fleet.
727
+ const PROBE = { m: 24, k: 96, n: 24, batch: 2 };
728
+ function probeInputs() {
729
+ let seed = 0x5EED; // fixed: identical on every device
730
+ const rnd = () => { seed = (Math.imul(seed, 1103515245) + 12345) & 0x7fffffff; return seed / 0x7fffffff; };
731
+ const { m, k, n, batch } = PROBE;
732
+ const Xq = new Int8Array(batch * m * k), Wq = new Int8Array(batch * k * n);
733
+ for (let i = 0; i < Xq.length; i++) Xq[i] = Math.round(rnd() * 254 - 127);
734
+ for (let i = 0; i < Wq.length; i++) Wq[i] = Math.round(rnd() * 254 - 127);
735
+ const rs = Float32Array.from({ length: batch * m }, () => 1);
736
+ const cs = Float32Array.from({ length: batch * n }, () => 1);
737
+ return { Xq, Wq, rs, cs, d: { ...PROBE, acc: true } };
738
+ }
739
+ async function kernelProbe(compute, L) {
740
+ const { Xq, Wq, rs, cs, d } = probeInputs();
741
+ const out = compute && compute.bgemm
742
+ ? await compute.bgemm(Xq, Wq, rs, cs, d)
743
+ : root.Verified.bgemmJS(Xq, Wq, rs, cs, d, L);
744
+ let h = 0x811c9dc5; // FNV-1a over the exact int32 results
745
+ const b = new Uint8Array(out.buffer, out.byteOffset, out.byteLength);
746
+ for (let i = 0; i < b.length; i++) { h ^= b[i]; h = Math.imul(h, 0x01000193); }
747
+ return h >>> 0;
748
+ }
749
+
750
+ root.Compute = { initCompute, loadLUTs, kernelProbe };
751
+ })(self);