Web demo: split-K backward + gather-fused attention

web/public/transformer.js

@@ -119,10 +119,6 @@
   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);
   }
-  // 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) -------------------------------------------------
   function lnFwd(x, rows, C) {
@@ -158,7 +154,8 @@
     const add = (name, w) => { params.push(w); names.push(name); return w; };
     const m = {
       cfg: { ...cfg, layers, heads, hidden, vocab: vocabSize() },
-      ctx: { L, bgemm: (engine && engine.bgemm) || null },
+      ctx: { L, bgemm: (engine && engine.bgemm) || null,
+             att: (engine && engine.att) || null, fgemm: (engine && engine.fgemm) || null },
       emb: add("emb", mk(vocabSize() * c, 0.08)),
       pos: add("pos", mk(cfg.t * c, 0.02)),
       blocks: [], params, names,
@@ -218,26 +215,16 @@
       const k = await vmm(l1.y, bl.Wk, BT, C, C, ctx);
       const v = await vmm(l1.y, bl.Wv, BT, C, C, ctx);
       cb.q = q; cb.k = k; cb.v = v;
-      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 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
+      // gather-FUSED attention (CUTLASS ex. 36/52): the kernels read q/k/v in
+      // their natural BT×C layout with head-strided indexing and scatter ctx
+      // straight back — no JS gather copies, no kᵀ transpose. All B×H heads in
+      // one dispatch per stage, every product through the verified units.
+      const BH = B * heads, dAtt = { B, T, heads, hd };
+      // per-(token,head) row quantization: the (BT·heads)×hd view IS the buffer
+      const qq = V.quantizeRows(q, BT * heads, hd), kq = V.quantizeRows(k, BT * heads, hd);
+      const sAll = ctx.att ? await ctx.att.scores(qq.q, kq.q, qq.s, kq.s, dAtt)
+                           : V.attScoresJS(qq.q, kq.q, qq.s, kq.s, dAtt, ctx.L);
       const aAll = new Float32Array(BH * T * T);                   // causal softmax
       for (let bz = 0; bz < BH; bz++) {
         const so = bz * T * T;
@@ -249,16 +236,11 @@
           for (let tj = 0; tj <= ti; tj++) aAll[so + ti * T + tj] /= z;
         }
       }
-      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) });
-      }
+      const aq = V.quantizeRows(aAll, BH * T, T);
+      const vq = V.quantizeHeadCols(v, B, T, heads, hd);
+      const ctxOut = ctx.att ? await ctx.att.ctx(aq.q, vq.q, aq.s, vq.s, dAtt)
+                             : V.attCtxJS(aq.q, vq.q, aq.s, vq.s, dAtt, ctx.L);
+      cb.aAll = aAll;                                  // backward slices heads from q/k/v/aAll
       cb.ctxOut = ctxOut;
       const attnOut = await vmm(ctxOut, bl.Wo, BT, C, C, ctx);
       const x2 = new Float32Array(BT * C);
@@ -295,15 +277,26 @@
   }
 
   // ---- STE backward (float), mirrors forward exactly --------------------------
-  function backward(m, cache) {
+  // The two vocab-sized matmuls run on the split-K f32 GPU kernel when
+  // available (CUTLASS ex. 06) — same float math, off the JS thread.
+  async function backward(m, cache) {
     const { c: C, t: T, b: B, layers, heads, hidden, vocab } = m.cfg;
     const BT = B * T, hd = C / heads, mm = TC.matmul, tr = TC.transpose;
     const g = m.params.map(p => new Float32Array(p.length));
     const gi = Object.fromEntries(m.names.map((n, i) => [n, i]));
     // tied unembed: logits = lnf @ embᵀ, so the unembedding gradient flows
     // straight into emb — dlogitsᵀ @ lnf is V×C, emb's own shape
-    g[gi.emb] = mm(tr(cache.dlogits, BT, vocab), cache.lnf.y, vocab, BT, C);
-    let dx = lnBwd(mm(cache.dlogits, m.emb, BT, vocab, C), cache.lnf.y, cache.lnf.sig, BT, C);
+    let dlnfIn;
+    if (m.ctx.fgemm) {
+      [g[gi.emb], dlnfIn] = await Promise.all([
+        m.ctx.fgemm(cache.dlogits, cache.lnf.y, { m: vocab, k: BT, n: C, transA: true }),
+        m.ctx.fgemm(cache.dlogits, m.emb, { m: BT, k: vocab, n: C }),      // split-K shape
+      ]);
+    } else {
+      g[gi.emb] = mm(tr(cache.dlogits, BT, vocab), cache.lnf.y, vocab, BT, C);
+      dlnfIn = mm(cache.dlogits, m.emb, BT, vocab, C);
+    }
+    let dx = lnBwd(dlnfIn, cache.lnf.y, cache.lnf.sig, BT, C);
     const scale = 1 / Math.sqrt(hd);
     for (let l = layers - 1; l >= 0; l--) {
       const bl = m.blocks[l], cb = cache.blocks[l];
@@ -320,9 +313,12 @@
       g[gi[`b${l}.Wo`]] = mm(tr(cb.ctxOut, BT, C), dx2, C, BT, C);
       const dctx = mm(dx2, tr(bl.Wo, C, C), BT, C, C);
       const dq = new Float32Array(BT * C), dk = new Float32Array(BT * C), dv = new Float32Array(BT * C);
-      let ai = 0;
       for (let bi = 0; bi < B; bi++) for (let h = 0; h < heads; h++) {
-        const { a, qh, kh, vh } = cb.att[ai++];
+        const bz = bi * heads + h;
+        const a = cb.aAll.subarray(bz * T * T, (bz + 1) * T * T);
+        const qh = headSlice(cb.q, bi, h, T, C, hd);
+        const kh = headSlice(cb.k, bi, h, T, C, hd);
+        const vh = headSlice(cb.v, bi, h, T, C, hd);
         const dch = headSlice(dctx, bi, h, T, C, hd);        // T×hd
         const dvh = mm(tr(a, T, T), dch, T, T, hd);          // aᵀ @ dctx
         const da = mm(dch, tr(vh, T, hd), T, hd, T);         // dctx @ vᵀ
@@ -364,7 +360,7 @@
   async function trainStep(m) {
     const { X, Y } = sampleBatch(m.cfg);
     const { loss, cache } = await forward(m, X, Y);
-    const grad = backward(m, cache);
+    const grad = await backward(m, cache);
     return { loss, grad };
   }
 

web/public/verified_core.js

@@ -147,6 +147,71 @@
     return bgemmJS(x.q, wq, x.s, ws, d, L);
   }
 
+  // ---- gather-fused attention through the units (CUTLASS ex. 36/52) ----------
+  // The kernels read q/k/v/ctx directly in their natural BT×C layout with
+  // head-strided indexing — no JS gather copies, no kᵀ transpose, and the
+  // context write scatters straight back into BT×C. Quantization stays
+  // block-scaled: q/k/a per (token,head) row, v per (head,channel) column.
+  // The (BT·heads)×hd row view of q/k IS the contiguous buffer, so
+  // quantizeRows(q, BT·heads, hd) gives per-(token,head) scales for free.
+  function quantizeHeadCols(v, B, T, heads, hd) {   // per (batch,head,channel) column
+    const C = heads * hd;
+    const q = new Int8Array(B * T * C), s = new Float32Array(B * heads * hd);
+    for (let bi = 0; bi < B; bi++)
+      for (let h = 0; h < heads; h++)
+        for (let j = 0; j < hd; j++) {
+          let mx = 0;
+          for (let ti = 0; ti < T; ti++) {
+            const a = Math.abs(v[(bi * T + ti) * C + h * hd + j]);
+            if (a > mx) mx = a;
+          }
+          const sc = Math.max(mx / 127, 1e-8);
+          s[(bi * heads + h) * hd + j] = sc;
+          for (let ti = 0; ti < T; ti++) {
+            const idx = (bi * T + ti) * C + h * hd + j;
+            const w = Math.round(v[idx] / sc);
+            q[idx] = w < -128 ? -128 : w > 127 ? 127 : w;
+          }
+        }
+    return { q, s };
+  }
+  // scores S[bz,ti,tj] = q_row(bi,ti,h) · k_row(bi,tj,h), every product via the LUT
+  function attScoresJS(qq, kq, qs, ks, d, L) {
+    const { B, T, heads, hd } = d, C = heads * hd, mul = L.mul;
+    const out = new Float32Array(B * heads * T * T);
+    for (let bi = 0; bi < B; bi++) for (let h = 0; h < heads; h++) {
+      const bz = bi * heads + h;
+      for (let ti = 0; ti < T; ti++) {
+        const qo = (bi * T + ti) * C + h * hd, rscale = qs[(bi * T + ti) * heads + h];
+        for (let tj = 0; tj < T; tj++) {
+          const ko = (bi * T + tj) * C + h * hd;
+          let acc = 0;
+          for (let p = 0; p < hd; p++) acc += mul[(qq[qo + p] & 0xFF) * 256 + (kq[ko + p] & 0xFF)];
+          out[(bz * T + ti) * T + tj] = acc * rscale * ks[(bi * T + tj) * heads + h];
+        }
+      }
+    }
+    return out;
+  }
+  // ctx[(bi,ti),(h,j)] = Σ_tj a[bz,ti,tj]·v[(bi,tj),(h,j)] — scatter fused into BT×C
+  function attCtxJS(aq, vq, as, vs, d, L) {
+    const { B, T, heads, hd } = d, C = heads * hd, mul = L.mul;
+    const out = new Float32Array(B * T * C);
+    for (let bi = 0; bi < B; bi++) for (let h = 0; h < heads; h++) {
+      const bz = bi * heads + h;
+      for (let ti = 0; ti < T; ti++) {
+        const ao = (bz * T + ti) * T, rscale = as[bz * T + ti];
+        for (let j = 0; j < hd; j++) {
+          let acc = 0;
+          for (let tj = 0; tj < T; tj++)
+            acc += mul[(aq[ao + tj] & 0xFF) * 256 + (vq[(bi * T + tj) * C + h * hd + j] & 0xFF)];
+          out[(bi * T + ti) * C + h * hd + j] = acc * rscale * vs[(bi * heads + h) * hd + j];
+        }
+      }
+    }
+    return out;
+  }
+
   // one verified layer forward; returns float out (+ cache for STE backward).
   // Every product goes through the verified INT8 multiply (mul8 LUT) with exact
   // int32 accumulation — i.e. an emulated INT8 tensor-core GEMM — then dequant.
@@ -197,8 +262,8 @@
     for (let j = 0; j < W2.length; j++) W2[j] -= lr * gAvg[W1.length + j];
   }
 
-  const api = { quantize, quantize2, quantizeRows, quantizeCols, lutMatmulJS, lutMatmul3JS, lutMatmul3,
-                bgemmJS, vgemmBlock, linearFwd, forward, backward, splitApply };
+  const api = { quantize, quantize2, quantizeRows, quantizeCols, quantizeHeadCols, lutMatmulJS, lutMatmul3JS, lutMatmul3,
+                bgemmJS, vgemmBlock, attScoresJS, attCtxJS, 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

@@ -108,6 +108,91 @@
       O[(bz * m + row) * n + col] = v;
     }`;
 
+  // Gather-fused attention (CUTLASS ex. 36/52): kernels index q/k/v directly in
+  // their BT×C layout (head-strided) — no gather copies, no kᵀ transpose — and
+  // the ctx kernel scatters straight back into BT×C. int8×int8→i32 is exact, so
+  // these are bit-identical to the LUT mirrors (proved at init before use).
+  const WGSL_ATT_SCORES = `
+    @group(0) @binding(0) var<storage, read> Q : array<i32>;     // int8 per elem, BT×C
+    @group(0) @binding(1) var<storage, read> K : array<i32>;
+    @group(0) @binding(2) var<storage, read> qs : array<f32>;    // per (token,head)
+    @group(0) @binding(3) var<storage, read> ks : array<f32>;
+    @group(0) @binding(4) var<storage, read_write> O : array<f32>;
+    @group(0) @binding(5) var<uniform> dims : vec4<u32>;         // T, heads, hd, _
+    @compute @workgroup_size(8, 8, 1)
+    fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+      let T = dims.x; let heads = dims.y; let hd = dims.z;
+      let ti = gid.x; let tj = gid.y; let bz = gid.z;
+      if (ti >= T || tj >= T) { return; }
+      let bi = bz / heads; let h = bz % heads;
+      let C = heads * hd;
+      let qo = (bi * T + ti) * C + h * hd;
+      let ko = (bi * T + tj) * C + h * hd;
+      var s : i32 = 0;
+      for (var p = 0u; p < hd; p = p + 1u) { s = s + Q[qo + p] * K[ko + p]; }
+      O[(bz * T + ti) * T + tj] = f32(s) * qs[(bi * T + ti) * heads + h] * ks[(bi * T + tj) * heads + h];
+    }`;
+  const WGSL_ATT_CTX = `
+    @group(0) @binding(0) var<storage, read> A : array<i32>;     // int8, BH×T×T
+    @group(0) @binding(1) var<storage, read> V : array<i32>;     // int8, BT×C
+    @group(0) @binding(2) var<storage, read> as_ : array<f32>;   // per (bz,row)
+    @group(0) @binding(3) var<storage, read> vs : array<f32>;    // per (batch,head,chan)
+    @group(0) @binding(4) var<storage, read_write> O : array<f32>;  // BT×C (scatter fused)
+    @group(0) @binding(5) var<uniform> dims : vec4<u32>;         // T, heads, hd, _
+    @compute @workgroup_size(8, 8, 1)
+    fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+      let T = dims.x; let heads = dims.y; let hd = dims.z;
+      let ti = gid.x; let j = gid.y; let bz = gid.z;
+      if (ti >= T || j >= hd) { return; }
+      let bi = bz / heads; let h = bz % heads;
+      let C = heads * hd;
+      let ao = (bz * T + ti) * T;
+      var s : i32 = 0;
+      for (var tj = 0u; tj < T; tj = tj + 1u) { s = s + A[ao + tj] * V[(bi * T + tj) * C + h * hd + j]; }
+      O[(bi * T + ti) * C + h * hd + j] = f32(s) * as_[bz * T + ti] * vs[(bi * heads + h) * hd + j];
+    }`;
+
+  // Split-K f32 GEMM (CUTLASS ex. 06) for the STE BACKWARD only — the backward
+  // was always float (the integer path has no gradient); this just moves that
+  // exact float math off the JS thread. Split-K matters for dlnf: M=256, N=32,
+  // K=16512 — 8k outputs with a huge inner loop would idle the GPU, so slices
+  // of K run on separate workgroups and a second tiny pass reduces partials.
+  const WGSL_FGEMM = `
+    @group(0) @binding(0) var<storage, read> A : array<f32>;
+    @group(0) @binding(1) var<storage, read> Bm : array<f32>;
+    @group(0) @binding(2) var<storage, read_write> P : array<f32>;   // S partials
+    @group(0) @binding(3) var<uniform> dims : vec4<u32>;             // m, k, n, flags(bit0=transA, rest=S)
+    @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 transA = (dims.w & 1u) == 1u;
+      let S = dims.w >> 1u;
+      let row = gid.x; let col = gid.y; let z = gid.z;
+      if (row >= m || col >= n) { return; }
+      let ks = (k + S - 1u) / S;
+      let p0 = z * ks;
+      let p1 = min(k, p0 + ks);
+      var s : f32 = 0.0;
+      for (var p = p0; p < p1; p = p + 1u) {
+        let a = select(A[row * k + p], A[p * m + row], transA);
+        s = s + a * Bm[p * n + col];
+      }
+      P[(z * m + row) * n + col] = s;
+    }`;
+  const WGSL_FREDUCE = `
+    @group(0) @binding(0) var<storage, read> P : array<f32>;
+    @group(0) @binding(1) var<storage, read_write> O : array<f32>;
+    @group(0) @binding(2) var<uniform> dims : vec4<u32>;             // mn, S, _, _
+    @compute @workgroup_size(64)
+    fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+      let mn = dims.x; let S = dims.y;
+      let i = gid.x;
+      if (i >= mn) { return; }
+      var s : f32 = 0.0;
+      for (var z = 0u; z < S; z = z + 1u) { s = s + P[z * mn + i]; }
+      O[i] = s;
+    }`;
+
   async function loadLUTs(base) {
     base = base || "";
     const [mulB, reqB, reluB, meta] = await Promise.all([
@@ -154,9 +239,55 @@
       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" } });
+
+      // gather-fused attention kernels — must reproduce the LUT mirrors exactly
+      const mkPipe = (code) => device.createComputePipeline({ layout: "auto",
+        compute: { module: device.createShaderModule({ code }), entryPoint: "main" } });
+      const scoresPipe = mkPipe(WGSL_ATT_SCORES), ctxPipe = mkPipe(WGSL_ATT_CTX);
+      let att = { scores: (qq, kq, qs, ks, d) => gpuAttScores(device, scoresPipe, qq, kq, qs, ks, d),
+                  ctx: (aq, vq, as, vs, d) => gpuAttCtx(device, ctxPipe, aq, vq, as, vs, d) };
+      try {
+        const d0 = { B: 2, T: 8, heads: 2, hd: 8 };
+        const nQ = d0.B * d0.T * d0.heads * d0.hd;
+        const qq = new Int8Array(nQ), kq = new Int8Array(nQ), vq = new Int8Array(nQ);
+        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; }
+        const qs = Float32Array.from({ length: d0.B * d0.T * d0.heads }, () => Math.random() + 0.5);
+        const ks = Float32Array.from(qs, () => Math.random() + 0.5);
+        const aq = new Int8Array(d0.B * d0.heads * d0.T * d0.T);
+        for (let i = 0; i < aq.length; i++) aq[i] = (Math.random() * 127) | 0;
+        const as = Float32Array.from({ length: d0.B * d0.heads * d0.T }, () => Math.random() + 0.5);
+        const vs = Float32Array.from({ length: d0.B * d0.heads * d0.hd }, () => Math.random() + 0.5);
+        const [hwS, hwC] = await Promise.all([att.scores(qq, kq, qs, ks, d0), att.ctx(aq, vq, as, vs, d0)]);
+        const refS = root.Verified.attScoresJS(qq, kq, qs, ks, d0, L);
+        const refC = root.Verified.attCtxJS(aq, vq, as, vs, d0, L);
+        const close = (a, b) => Math.abs(a - b) <= Math.abs(b) * 1e-6 + 1e-6;
+        for (let i = 0; i < refS.length; i++) if (!close(hwS[i], refS[i])) throw new Error("scores mismatch");
+        for (let i = 0; i < refC.length; i++) if (!close(hwC[i], refC[i])) throw new Error("ctx mismatch");
+      } catch (e) {
+        console.warn("fused attention kernels failed verification — using CPU LUT mirrors:", e.message);
+        att = null;
+      }
+
+      // split-K f32 GEMM for the STE backward (self-tested vs JS float matmul)
+      const fPipes = { gemm: mkPipe(WGSL_FGEMM), reduce: mkPipe(WGSL_FREDUCE) };
+      let fgemm = (A, Bm, d) => gpuFgemm(device, fPipes, A, Bm, d);
+      try {
+        const m0 = 7, k0 = 4500, n0 = 5;                      // k big enough to exercise split-K
+        const A = Float32Array.from({ length: m0 * k0 }, () => Math.random() - 0.5);
+        const Bm = Float32Array.from({ length: k0 * n0 }, () => Math.random() - 0.5);
+        const hw = await fgemm(A, Bm, { m: m0, k: k0, n: n0 });
+        const ref = root.TrainCore.matmul(A, Bm, m0, k0, n0);
+        for (let i = 0; i < ref.length; i++)
+          if (Math.abs(hw[i] - ref[i]) > Math.abs(ref[i]) * 1e-3 + 1e-3) throw new Error("fgemm mismatch");
+      } catch (e) {
+        console.warn("split-K f32 GEMM failed verification — backward stays in JS:", e.message);
+        fgemm = null;
+      }
+
       const viaLUT = { backend: "webgpu", label: `${gpuName} (LUT shader)`,
                        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) };
+                       bgemm: (Xq, Wq, rs, cs, d) => gpuBgemmLUT(device, bgLutPipe, lutBuf, Xq, Wq, rs, cs, d),
+                       att, fgemm };
 
       // DP4A pipeline — only if the WGSL feature exists AND it reproduces the
       // verified units exactly on random self-tests
@@ -193,11 +324,11 @@
           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 };
+                     matmulInt8: dp4mm, bgemm: viaLUT.bgemm, att, fgemm };
           }
       }
       return { backend: "webgpu", label: `${gpuName} (DP4A int8 dot — HW verified vs units)`,
-               matmulInt8: dp4mm, bgemm: bg };
+               matmulInt8: dp4mm, bgemm: bg, att, fgemm };
     } catch (e) { console.warn("WebGPU init failed, CPU fallback:", e); return cpu; }
   }
 
@@ -247,6 +378,69 @@
     return r.out;
   }
 
+  // attention kernels: int8 (widened i32) in, f32 out, strided head indexing
+  async function gpuAttScores(device, pipeline, qq, kq, qs, ks, d) {
+    const { B, T, heads } = d;
+    const bufQ = up(device, Int32Array.from(qq), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
+    const bufK = up(device, Int32Array.from(kq), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
+    const bufQs = up(device, qs, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
+    const bufKs = up(device, ks, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
+    const bufD = up(device, new Uint32Array([d.T, d.heads, d.hd, 0]), GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST);
+    const r = await runBgPass(device, pipeline, (bufO) => [
+      { binding: 0, resource: { buffer: bufQ } }, { binding: 1, resource: { buffer: bufK } },
+      { binding: 2, resource: { buffer: bufQs } }, { binding: 3, resource: { buffer: bufKs } },
+      { binding: 4, resource: { buffer: bufO } }, { binding: 5, resource: { buffer: bufD } } ], T, T, B * heads);
+    [bufQ, bufK, bufQs, bufKs, bufD, r.bufO, r.read].forEach(b => b.destroy());
+    return r.out;
+  }
+  async function gpuAttCtx(device, pipeline, aq, vq, as, vs, d) {
+    const { B, T, heads, hd } = d;
+    const bufA = up(device, Int32Array.from(aq), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
+    const bufV = up(device, Int32Array.from(vq), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
+    const bufAs = up(device, as, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
+    const bufVs = up(device, vs, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
+    const bufD = up(device, new Uint32Array([T, heads, hd, 0]), GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST);
+    const r = await runBgPass(device, pipeline, (bufO) => [
+      { binding: 0, resource: { buffer: bufA } }, { binding: 1, resource: { buffer: bufV } },
+      { binding: 2, resource: { buffer: bufAs } }, { binding: 3, resource: { buffer: bufVs } },
+      { binding: 4, resource: { buffer: bufO } }, { binding: 5, resource: { buffer: bufD } } ], T, hd, B * heads);
+    [bufA, bufV, bufAs, bufVs, bufD, r.bufO, r.read].forEach(b => b.destroy());
+    return r.out;
+  }
+
+  // split-K f32 GEMM (backward): partial pass + reduce pass
+  async function gpuFgemm(device, pipes, A, Bm, d) {
+    const { m, k, n } = d, transA = d.transA ? 1 : 0;
+    const S = k > 4096 ? Math.min(16, Math.ceil(k / 2048)) : 1;
+    const bufA = up(device, A, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
+    const bufB = up(device, Bm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
+    const bufP = mk(device, S * m * n * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
+    const bufD1 = up(device, new Uint32Array([m, k, n, transA | (S << 1)]), GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST);
+    const bufO = mk(device, m * n * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
+    const bufD2 = up(device, new Uint32Array([m * n, S, 0, 0]), GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST);
+    const enc = device.createCommandEncoder();
+    let pass = enc.beginComputePass();
+    pass.setPipeline(pipes.gemm);
+    pass.setBindGroup(0, device.createBindGroup({ layout: pipes.gemm.getBindGroupLayout(0), entries: [
+      { binding: 0, resource: { buffer: bufA } }, { binding: 1, resource: { buffer: bufB } },
+      { binding: 2, resource: { buffer: bufP } }, { binding: 3, resource: { buffer: bufD1 } } ] }));
+    pass.dispatchWorkgroups(Math.ceil(m / 8), Math.ceil(n / 8), S); pass.end();
+    pass = enc.beginComputePass();
+    pass.setPipeline(pipes.reduce);
+    pass.setBindGroup(0, device.createBindGroup({ layout: pipes.reduce.getBindGroupLayout(0), entries: [
+      { binding: 0, resource: { buffer: bufP } }, { binding: 1, resource: { buffer: bufO } },
+      { binding: 2, resource: { buffer: bufD2 } } ] }));
+    pass.dispatchWorkgroups(Math.ceil(m * n / 64)); pass.end();
+    const read = mk(device, m * n * 4, GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ);
+    enc.copyBufferToBuffer(bufO, 0, read, 0, m * n * 4);
+    device.queue.submit([enc.finish()]);
+    await read.mapAsync(GPUMapMode.READ);
+    const out = new Float32Array(read.getMappedRange().slice(0));
+    read.unmap();
+    [bufA, bufB, bufP, bufD1, bufO, bufD2, read].forEach(b => b.destroy());
+    return 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;