web: gate-the-gate boundary discrimination in the B2B init gate
Browse files- web/public/webgpu.js +29 -0
- web/test_b2b.js +3 -0
- 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);
|