Web demo: block-scaled INT8 + batched/fused GEMM kernels

web/public/app.js

@@ -806,7 +806,7 @@ function cfgSliders() { return [ui.cfgC, ui.cfgT, ui.cfgB, ui.cfgSteps, ui.cfgLr
 // (re)build the mini transformer — deterministic shared init (same seeds on
 // every peer); each device samples its own batch windows from the shared corpus
 function buildModel(cfg) {
-  model = Transformer.init(cfg, L, compute.matmulInt8);
+  model = Transformer.init(cfg, L, compute);       // compute.bgemm = fused batched GPU path
   trainedSteps = 0;
 }
 
@@ -856,7 +856,7 @@ function buildModel(cfg) {
     ui.start.disabled = true;
     return;                                              // no signaling, no training
   }
-  ui.backend.textContent = `${compute.backend.toUpperCase()} — ${compute.label} · 3×INT8 fast-accurate through verified units`;
+  ui.backend.textContent = `${compute.backend.toUpperCase()} — ${compute.label} · block-scaled INT8, batched + fused epilogue, through verified units`;
   // tokenizer must load BEFORE the model is built (it sets the vocab size)
   try {
     const name = await Transformer.loadTokenizer();

web/public/transformer.js

@@ -111,17 +111,17 @@
   const streamFineWebEdu = () => streamDataset(DEFAULT_DS);
   function datasetName() { return DATASET; }
 
-  // ---- verified matmul: float in -> quantize -> LUT multiply -> dequant -----
-  // weights/projections may go through WebGPU; attention (per-head, many small
-  // matmuls) uses the CPU LUT path — same verified units, no dispatch overhead.
-  // 3xINT8 fast-accurate GEMM (CUTLASS 3xTF32 scheme through the verified
-  // units): three exact LUT GEMMs recover ~14-bit precision from the 8-bit
-  // units — see Verified.lutMatmul3.
-  async function vmm(Xf, Wf, m, k, n, ctx) {
-    return V.lutMatmul3(Xf, Wf, m, k, n, ctx.L, ctx.matmulInt8);
+  // ---- verified matmul: block-scaled INT8 through the units ------------------
+  // CUTLASS ex. 67/81 blockwise scaling: per-row activation scales × per-column
+  // weight scales, one exact LUT/DP4A GEMM, dequant (+ optional fused ReLU) in
+  // the kernel epilogue (ex. 12). Replaces the per-tensor 3-pass: same outlier
+  // robustness at one third of the unit ops.
+  async function vmm(Xf, Wf, m, k, n, ctx, relu) {
+    return V.vgemmBlock(Xf, Wf, { m, k, n, batch: 1, relu: !!relu }, ctx.L, ctx.bgemm);
   }
-  function vmmCPU(Xf, Wf, m, k, n, ctx) {
-    return V.lutMatmul3JS(Xf, Wf, m, k, n, ctx.L);
+  // batched: one dispatch for all B×H attention-head problems (ex. 05/24)
+  async function vmmBatched(Xf, Wf, m, k, n, batch, ctx) {
+    return V.vgemmBlock(Xf, Wf, { m, k, n, batch }, ctx.L, ctx.bgemm);
   }
 
   // ---- layernorm (no affine) -------------------------------------------------
@@ -148,7 +148,9 @@
 
   // ---- model -----------------------------------------------------------------
   // cfg: { c: width, t: seq len, b: batch/device, layers, heads, steps, lr }
-  function init(cfg, L, matmulInt8) {
+  // engine: the Compute backend object ({bgemm} for the fused WebGPU path) or a
+  // legacy matmulInt8 function (Node tests, inference kit) -> CPU LUT mirror
+  function init(cfg, L, engine) {
     const c = cfg.c, layers = cfg.layers || 2, heads = cfg.heads || 2, hidden = 2 * c;
     let seed = 100;
     const mk = (nEl, scale) => { const w = randn(nEl, mulberry32(seed++)); for (let i = 0; i < nEl; i++) w[i] *= scale; return w; };
@@ -156,7 +158,7 @@
     const add = (name, w) => { params.push(w); names.push(name); return w; };
     const m = {
       cfg: { ...cfg, layers, heads, hidden, vocab: vocabSize() },
-      ctx: { L, matmulInt8 },
+      ctx: { L, bgemm: (engine && engine.bgemm) || null },
       emb: add("emb", mk(vocabSize() * c, 0.08)),
       pos: add("pos", mk(cfg.t * c, 0.02)),
       blocks: [], params, names,
@@ -217,22 +219,43 @@
       const ctxOut = new Float32Array(BT * C);
       cb.att = [];                                     // per (b,h): softmax probs
       const scale = 1 / Math.sqrt(hd);
+      // gather every head into batched contiguous layouts, then run ALL B×H
+      // score GEMMs in ONE batched dispatch (and all attn·V in another) —
+      // instead of 2·B·H tiny sequential matmuls
+      const BH = B * heads;
+      const qb = new Float32Array(BH * T * hd), kb = new Float32Array(BH * T * hd);
+      const vb = new Float32Array(BH * T * hd), ktb = new Float32Array(BH * hd * T);
       for (let bi = 0; bi < B; bi++) for (let h = 0; h < heads; h++) {
-        const qh = headSlice(q, bi, h, T, C, hd), kh = headSlice(k, bi, h, T, C, hd), vh = headSlice(v, bi, h, T, C, hd);
-        const kT = TC.transpose(kh, T, hd);            // hd×T
-        const s = vmmCPU(qh, kT, T, hd, T, ctx);       // scores through the units
-        for (let i = 0; i < T * T; i++) s[i] *= scale;
-        const a = new Float32Array(T * T);             // causal softmax
+        const bz = bi * heads + h;
+        for (let ti = 0; ti < T; ti++) for (let j = 0; j < hd; j++) {
+          const src = (bi * T + ti) * C + h * hd + j;
+          qb[(bz * T + ti) * hd + j] = q[src];
+          kb[(bz * T + ti) * hd + j] = k[src];
+          vb[(bz * T + ti) * hd + j] = v[src];
+          ktb[(bz * hd + j) * T + ti] = k[src];        // kᵀ per head, for scores
+        }
+      }
+      const sAll = await vmmBatched(qb, ktb, T, hd, T, BH, ctx);   // scores through the units
+      const aAll = new Float32Array(BH * T * T);                   // causal softmax
+      for (let bz = 0; bz < BH; bz++) {
+        const so = bz * T * T;
         for (let ti = 0; ti < T; ti++) {
           let mx = -1e30;
-          for (let tj = 0; tj <= ti; tj++) mx = Math.max(mx, s[ti * T + tj]);
+          for (let tj = 0; tj <= ti; tj++) mx = Math.max(mx, sAll[so + ti * T + tj] * scale);
           let z = 0;
-          for (let tj = 0; tj <= ti; tj++) { const e = Math.exp(s[ti * T + tj] - mx); a[ti * T + tj] = e; z += e; }
-          for (let tj = 0; tj <= ti; tj++) a[ti * T + tj] /= z;
+          for (let tj = 0; tj <= ti; tj++) { const e = Math.exp(sAll[so + ti * T + tj] * scale - mx); aAll[so + ti * T + tj] = e; z += e; }
+          for (let tj = 0; tj <= ti; tj++) aAll[so + ti * T + tj] /= z;
         }
-        const ch = vmmCPU(a, vh, T, T, hd, ctx);       // attn·V through the units
-        headUnslice(ctxOut, ch, bi, h, T, C, hd, false);
-        cb.att.push({ a, qh, kh, vh });
+      }
+      const cAll = await vmmBatched(aAll, vb, T, T, hd, BH, ctx);  // attn·V through the units
+      for (let bi = 0; bi < B; bi++) for (let h = 0; h < heads; h++) {
+        const bz = bi * heads + h;
+        for (let ti = 0; ti < T; ti++) for (let j = 0; j < hd; j++)
+          ctxOut[(bi * T + ti) * C + h * hd + j] = cAll[(bz * T + ti) * hd + j];
+        cb.att.push({ a: aAll.subarray(bz * T * T, (bz + 1) * T * T),
+                      qh: qb.subarray(bz * T * hd, (bz + 1) * T * hd),
+                      kh: kb.subarray(bz * T * hd, (bz + 1) * T * hd),
+                      vh: vb.subarray(bz * T * hd, (bz + 1) * T * hd) });
       }
       cb.ctxOut = ctxOut;
       const attnOut = await vmm(ctxOut, bl.Wo, BT, C, C, ctx);
@@ -240,9 +263,9 @@
       for (let i = 0; i < x2.length; i++) x2[i] = x[i] + attnOut[i];
       cb.x2 = x2;
       const l2 = lnFwd(x2, BT, C); cb.ln2 = l2;
-      const h1 = await vmm(l2.y, bl.W1, BT, C, hidden, ctx);
+      const h1 = await vmm(l2.y, bl.W1, BT, C, hidden, ctx, true);  // ReLU fused in the epilogue
       const mask = new Uint8Array(h1.length);
-      for (let i = 0; i < h1.length; i++) { if (h1[i] > 0) mask[i] = 1; else h1[i] = 0; }
+      for (let i = 0; i < h1.length; i++) if (h1[i] > 0) mask[i] = 1;
       cb.h1 = h1; cb.mask = mask;
       const mlpOut = await vmm(h1, bl.W2, BT, hidden, C, ctx);
       x = new Float32Array(BT * C);

web/public/verified_core.js

@@ -66,6 +66,87 @@
     return combine3(hh, hl, lh, x, w, m * n);
   }
 
+  // ---- block-scaled verified GEMM (CUTLASS ex. 67/81 blockwise scaling) ------
+  // Per-ROW scales for the activations and per-COLUMN scales for the weights:
+  // the integer math through the mul8 LUT is completely unchanged — only the
+  // dequant uses rs[row]·cs[col] instead of one tensor-wide product, so a single
+  // outlier no longer crushes the quantization resolution of every other
+  // row/column. One LUT pass at this granularity beats the per-tensor 3-pass.
+  function quantizeRows(X, rows, cols) {
+    const q = new Int8Array(rows * cols), s = new Float32Array(rows);
+    for (let r = 0; r < rows; r++) {
+      let mx = 0;
+      for (let c = 0; c < cols; c++) { const a = Math.abs(X[r * cols + c]); if (a > mx) mx = a; }
+      const sc = Math.max(mx / 127, 1e-8); s[r] = sc;
+      for (let c = 0; c < cols; c++) {
+        const v = Math.round(X[r * cols + c] / sc);
+        q[r * cols + c] = v < -128 ? -128 : v > 127 ? 127 : v;
+      }
+    }
+    return { q, s };
+  }
+  function quantizeCols(W, rows, cols) {
+    const q = new Int8Array(rows * cols), s = new Float32Array(cols);
+    for (let c = 0; c < cols; c++) {
+      let mx = 0;
+      for (let r = 0; r < rows; r++) { const a = Math.abs(W[r * cols + c]); if (a > mx) mx = a; }
+      s[c] = Math.max(mx / 127, 1e-8);
+    }
+    for (let r = 0; r < rows; r++)
+      for (let c = 0; c < cols; c++) {
+        const v = Math.round(W[r * cols + c] / s[c]);
+        q[r * cols + c] = v < -128 ? -128 : v > 127 ? 127 : v;
+      }
+    return { q, s };
+  }
+
+  // CPU mirror of the fused GPU kernel: batched int8 GEMM through the LUT with
+  // the epilogue (block dequant + optional ReLU) applied before returning —
+  // exactly what the WGSL kernel does on-device.
+  function bgemmJS(Xq, Wq, rs, cs, d, L) {
+    const { m, k, n } = d, batch = d.batch || 1, relu = !!d.relu, mul = L.mul;
+    const out = new Float32Array(batch * m * n);
+    const acc = new Int32Array(n);
+    for (let bz = 0; bz < batch; bz++) {
+      const xo = bz * m * k, wo = bz * k * n, oo = bz * m * n, co = bz * n;
+      for (let i = 0; i < m; i++) {
+        acc.fill(0);
+        const xrow = xo + i * k;
+        for (let p = 0; p < k; p++) {
+          const au = (Xq[xrow + p] & 0xFF) * 256, wrow = wo + p * n;
+          for (let j = 0; j < n; j++) acc[j] += mul[au + (Wq[wrow + j] & 0xFF)];
+        }
+        const rscale = rs[bz * m + i], orow = oo + i * n;
+        for (let j = 0; j < n; j++) {
+          const v = acc[j] * rscale * cs[co + j];
+          out[orow + j] = relu && v < 0 ? 0 : v;
+        }
+      }
+    }
+    return out;
+  }
+
+  // block-scaled verified GEMM, float in → float out.
+  //   d = { m, k, n, batch=1, relu=false }; X is (batch·m)×k, W is batch×(k×n)
+  // gpuBgemm (from webgpu.js) runs the batched kernel with the fused epilogue;
+  // without it the CPU LUT mirror runs. Every product goes through the units.
+  async function vgemmBlock(Xf, Wf, d, L, gpuBgemm) {
+    const { m, k, n } = d, batch = d.batch || 1;
+    const x = quantizeRows(Xf, batch * m, k);
+    let wq, ws;
+    if (batch === 1) {
+      const w = quantizeCols(Wf, k, n); wq = w.q; ws = w.s;
+    } else {
+      wq = new Int8Array(batch * k * n); ws = new Float32Array(batch * n);
+      for (let bz = 0; bz < batch; bz++) {
+        const w = quantizeCols(Wf.subarray(bz * k * n, (bz + 1) * k * n), k, n);
+        wq.set(w.q, bz * k * n); ws.set(w.s, bz * n);
+      }
+    }
+    if (gpuBgemm) return gpuBgemm(x.q, wq, x.s, ws, d);
+    return bgemmJS(x.q, wq, x.s, ws, d, L);
+  }
+
   // one verified layer forward; returns float out (+ cache for STE backward).
   // Every product goes through the verified INT8 multiply (mul8 LUT) with exact
   // int32 accumulation — i.e. an emulated INT8 tensor-core GEMM — then dequant.
@@ -116,7 +197,8 @@
     for (let j = 0; j < W2.length; j++) W2[j] -= lr * gAvg[W1.length + j];
   }
 
-  const api = { quantize, quantize2, lutMatmulJS, lutMatmul3JS, lutMatmul3, linearFwd, forward, backward, splitApply };
+  const api = { quantize, quantize2, quantizeRows, quantizeCols, lutMatmulJS, lutMatmul3JS, lutMatmul3,
+                bgemmJS, vgemmBlock, linearFwd, forward, backward, splitApply };
   if (typeof module !== "undefined" && module.exports) { TC = require("./traincore.js"); module.exports = api; }
   else { TC = root.TrainCore; root.Verified = api; }
 })(typeof self !== "undefined" ? self : this);

web/public/webgpu.js

@@ -54,6 +54,60 @@
       C[row * n + col] = s;
     }`;
 
+  // Batched block-scaled GEMM with a FUSED EPILOGUE (CUTLASS ex. 05/24 + 12):
+  // grid z = batch index, so ALL attention heads run in ONE dispatch, and the
+  // epilogue (block dequant rs·cs + optional ReLU) happens before the data
+  // leaves the GPU — f32 out, no int32 readback, no second pass in JS.
+  const WGSL_BG_LUT = `
+    @group(0) @binding(0) var<storage, read> Xq  : array<i32>;   // int8 byte per elem
+    @group(0) @binding(1) var<storage, read> Wq  : array<i32>;
+    @group(0) @binding(2) var<storage, read> lut : array<i32>;
+    @group(0) @binding(3) var<storage, read> rs  : array<f32>;   // per (batch,row)
+    @group(0) @binding(4) var<storage, read> cs  : array<f32>;   // per (batch,col)
+    @group(0) @binding(5) var<storage, read_write> O : array<f32>;
+    @group(0) @binding(6) var<uniform> dims : vec4<u32>;         // m, k, n, flags(1=relu)
+    @compute @workgroup_size(8, 8, 1)
+    fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+      let m = dims.x; let k = dims.y; let n = dims.z;
+      let row = gid.x; let col = gid.y; let bz = gid.z;
+      if (row >= m || col >= n) { return; }
+      var s : i32 = 0;
+      let xo = (bz * m + row) * k;
+      let wo = bz * k * n + col;
+      for (var p = 0u; p < k; p = p + 1u) {
+        let au = u32(Xq[xo + p] & 255);
+        let bu = u32(Wq[wo + p * n] & 255);
+        s = s + lut[au * 256u + bu];
+      }
+      var v = f32(s) * rs[bz * m + row] * cs[bz * n + col];
+      if ((dims.w & 1u) == 1u && v < 0.0) { v = 0.0; }
+      O[(bz * m + row) * n + col] = v;
+    }`;
+
+  // same fused/batched kernel on the DP4A hardware path (Wᵀ packed per batch)
+  const WGSL_BG_DP4 = `
+    @group(0) @binding(0) var<storage, read> Xp : array<u32>;
+    @group(0) @binding(1) var<storage, read> Wp : array<u32>;    // per-batch Wᵀ, packed
+    @group(0) @binding(2) var<storage, read> rs : array<f32>;
+    @group(0) @binding(3) var<storage, read> cs : array<f32>;
+    @group(0) @binding(4) var<storage, read_write> O : array<f32>;
+    @group(0) @binding(5) var<uniform> dims : vec4<u32>;         // m, kw, n, flags
+    @compute @workgroup_size(8, 8, 1)
+    fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+      let m = dims.x; let kw = dims.y; let n = dims.z;
+      let row = gid.x; let col = gid.y; let bz = gid.z;
+      if (row >= m || col >= n) { return; }
+      var s : i32 = 0;
+      let xo = (bz * m + row) * kw;
+      let wo = (bz * n + col) * kw;
+      for (var p = 0u; p < kw; p = p + 1u) {
+        s = s + dot4I8Packed(Xp[xo + p], Wp[wo + p]);
+      }
+      var v = f32(s) * rs[bz * m + row] * cs[bz * n + col];
+      if ((dims.w & 1u) == 1u && v < 0.0) { v = 0.0; }
+      O[(bz * m + row) * n + col] = v;
+    }`;
+
   async function loadLUTs(base) {
     base = base || "";
     const [mulB, reqB, reluB, meta] = await Promise.all([
@@ -98,8 +152,11 @@
       const lut32 = new Int32Array(L.mul);   // widen int16 -> int32
       const lutBuf = device.createBuffer({ size: lut32.byteLength, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST });
       device.queue.writeBuffer(lutBuf, 0, lut32);
+      const bgLutModule = device.createShaderModule({ code: WGSL_BG_LUT });
+      const bgLutPipe = device.createComputePipeline({ layout: "auto", compute: { module: bgLutModule, entryPoint: "main" } });
       const viaLUT = { backend: "webgpu", label: `${gpuName} (LUT shader)`,
-                       matmulInt8: (Xq, Wq, m, k, n) => gpuMatmulLUT(device, lutPipe, lutBuf, Xq, Wq, m, k, n) };
+                       matmulInt8: (Xq, Wq, m, k, n) => gpuMatmulLUT(device, lutPipe, lutBuf, Xq, Wq, m, k, n),
+                       bgemm: (Xq, Wq, rs, cs, d) => gpuBgemmLUT(device, bgLutPipe, lutBuf, Xq, Wq, rs, cs, d) };
 
       // DP4A pipeline — only if the WGSL feature exists AND it reproduces the
       // verified units exactly on random self-tests
@@ -118,7 +175,29 @@
         for (let i = 0; i < ref.length; i++)
           if (hw[i] !== ref[i]) { console.warn("DP4A disagreed with the verified units — using LUT shader"); return viaLUT; }
       }
-      return { backend: "webgpu", label: `${gpuName} (DP4A int8 dot — HW verified vs units)`, matmulInt8: dp4mm };
+      // fused/batched DP4A kernel, gated by the same verification: it must
+      // reproduce the CPU LUT mirror bit-for-bit on the integer accumulate
+      const bgDp4Module = device.createShaderModule({ code: WGSL_BG_DP4 });
+      const bgDp4Pipe = device.createComputePipeline({ layout: "auto", compute: { module: bgDp4Module, entryPoint: "main" } });
+      const bg = (Xq, Wq, rs, cs, d) => gpuBgemmDP4(device, bgDp4Pipe, Xq, Wq, rs, cs, d);
+      {
+        const d0 = { m: 5, k: 9, n: 6, batch: 3, relu: true };
+        const Xq = new Int8Array(d0.batch * d0.m * d0.k), Wq = new Int8Array(d0.batch * d0.k * d0.n);
+        for (let i = 0; i < Xq.length; i++) Xq[i] = (Math.random() * 256 - 128) | 0;
+        for (let i = 0; i < Wq.length; i++) Wq[i] = (Math.random() * 256 - 128) | 0;
+        const rs = new Float32Array(d0.batch * d0.m).fill(0).map(() => Math.random() + 0.5);
+        const cs = new Float32Array(d0.batch * d0.n).fill(0).map(() => Math.random() + 0.5);
+        const hw = await bg(Xq, Wq, rs, cs, d0);
+        const ref = root.Verified.bgemmJS(Xq, Wq, rs, cs, d0, L);
+        for (let i = 0; i < ref.length; i++)
+          if (Math.abs(hw[i] - ref[i]) > Math.abs(ref[i]) * 1e-6 + 1e-6) {
+            console.warn("batched DP4A disagreed with the verified units — using LUT bgemm");
+            return { backend: "webgpu", label: `${gpuName} (DP4A int8 dot — HW verified vs units)`,
+                     matmulInt8: dp4mm, bgemm: viaLUT.bgemm };
+          }
+      }
+      return { backend: "webgpu", label: `${gpuName} (DP4A int8 dot — HW verified vs units)`,
+               matmulInt8: dp4mm, bgemm: bg };
     } catch (e) { console.warn("WebGPU init failed, CPU fallback:", e); return cpu; }
   }
 
@@ -168,6 +247,62 @@
     return r.out;
   }
 
+  // fused batched dispatch: f32 out, epilogue done on-device
+  async function runBgPass(device, pipeline, entries, m, n, batch) {
+    const bytesO = batch * m * n * 4;
+    const bufO = mk(device, bytesO, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
+    const bind = device.createBindGroup({ layout: pipeline.getBindGroupLayout(0), entries: entries(bufO) });
+    const enc = device.createCommandEncoder();
+    const pass = enc.beginComputePass();
+    pass.setPipeline(pipeline); pass.setBindGroup(0, bind);
+    pass.dispatchWorkgroups(Math.ceil(m / 8), Math.ceil(n / 8), batch); pass.end();
+    const read = mk(device, bytesO, GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ);
+    enc.copyBufferToBuffer(bufO, 0, read, 0, bytesO);
+    device.queue.submit([enc.finish()]);
+    await read.mapAsync(GPUMapMode.READ);
+    const out = new Float32Array(read.getMappedRange().slice(0));
+    read.unmap();
+    return { out, bufO, read };
+  }
+
+  async function gpuBgemmLUT(device, pipeline, lutBuf, Xq, Wq, rs, cs, d) {
+    const { m, k, n } = d, batch = d.batch || 1, flags = d.relu ? 1 : 0;
+    const bufX = up(device, Int32Array.from(Xq), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
+    const bufW = up(device, Int32Array.from(Wq), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
+    const bufR = up(device, rs, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
+    const bufS = up(device, cs, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
+    const bufD = up(device, new Uint32Array([m, k, n, flags]), GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST);
+    const r = await runBgPass(device, pipeline, (bufO) => [
+      { binding: 0, resource: { buffer: bufX } }, { binding: 1, resource: { buffer: bufW } },
+      { binding: 2, resource: { buffer: lutBuf } }, { binding: 3, resource: { buffer: bufR } },
+      { binding: 4, resource: { buffer: bufS } }, { binding: 5, resource: { buffer: bufO } },
+      { binding: 6, resource: { buffer: bufD } } ], m, n, batch);
+    [bufX, bufW, bufR, bufS, bufD, r.bufO, r.read].forEach(b => b.destroy());
+    return r.out;
+  }
+
+  async function gpuBgemmDP4(device, pipeline, Xq, Wq, rs, cs, d) {
+    const { m, k, n } = d, batch = d.batch || 1, flags = d.relu ? 1 : 0;
+    const kw = Math.ceil(k / 4);
+    const Xp = packRows(Xq, batch * m, k, kw);            // rows are (batch·m)
+    const Wp = new Uint32Array(batch * n * kw);           // per-batch Wᵀ, packed
+    for (let bz = 0; bz < batch; bz++) {
+      const wt = transposeI8(Wq.subarray(bz * k * n, (bz + 1) * k * n), k, n);
+      Wp.set(packRows(wt, n, k, kw), bz * n * kw);
+    }
+    const bufX = up(device, Xp, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
+    const bufW = up(device, Wp, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
+    const bufR = up(device, rs, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
+    const bufS = up(device, cs, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
+    const bufD = up(device, new Uint32Array([m, kw, n, flags]), GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST);
+    const r = await runBgPass(device, pipeline, (bufO) => [
+      { binding: 0, resource: { buffer: bufX } }, { binding: 1, resource: { buffer: bufW } },
+      { binding: 2, resource: { buffer: bufR } }, { binding: 3, resource: { buffer: bufS } },
+      { binding: 4, resource: { buffer: bufO } }, { binding: 5, resource: { buffer: bufD } } ], m, n, batch);
+    [bufX, bufW, bufR, bufS, bufD, r.bufO, r.read].forEach(b => b.destroy());
+    return r.out;
+  }
+
   function mk(device, size, usage) { return device.createBuffer({ size, usage }); }
   function up(device, arr, usage) {
     const b = mk(device, Math.max(16, arr.byteLength), usage);