@genai-fi/nanogpt 0.7.0 → 0.7.1

package/dist/ops/webgpu/adamAdjust.js

@@ -0,0 +1,52 @@
+import { g as p } from "../../webgpu_program-WaoMq-WD.js";
+import { f as d, c as l } from "../../webgpu_util-DhSeP4b6.js";
+import { f } from "../../index-UdZhlibC.js";
+class c {
+  variableNames = ["moments", "value"];
+  outputShape;
+  shaderKey = "AdamAdjust";
+  dispatchLayout;
+  dispatch;
+  workgroupSize = [64, 1, 1];
+  size = !0;
+  uniforms = "invbeta1: f32, invbeta2: f32, learningRate: f32, epsilon: f32";
+  outputComponent = 1;
+  variableComponents = [2, 1];
+  constructor(e) {
+    this.outputShape = e, this.dispatchLayout = d(this.outputShape), this.dispatch = l(this.dispatchLayout, this.outputShape, this.workgroupSize);
+  }
+  getUserCode() {
+    return `
+        ${p("index")} {
+            if (index < uniforms.size) {
+                let moments: vec2<f32> = moments[index];
+                let value: f32 = value[index];
+
+                let m1Hat = moments.x * uniforms.invbeta1;
+                let m2Hat = moments.y * uniforms.invbeta2;
+
+                let invSqrt = inverseSqrt(max(m2Hat, 1e-30));
+                let invDenom = invSqrt / (1.0 + uniforms.epsilon * invSqrt);
+                let adjustedValue = fma(-uniforms.learningRate * m1Hat, invDenom, value);
+
+                setOutputAtIndex(index, adjustedValue);
+            }
+        }
+    `;
+  }
+}
+function h(t) {
+  const { moments: e, value: a } = t.inputs, { beta1: n, beta2: o, learningRate: i, epsilon: s } = t.attrs, r = t.backend, u = new c(a.shape), m = [
+    { type: "float32", data: [1 / n] },
+    { type: "float32", data: [1 / o] },
+    { type: "float32", data: [i] },
+    { type: "float32", data: [s] }
+  ];
+  return r.runWebGPUProgram(u, [e, a], "float32", m);
+}
+const v = {
+  kernelName: "AdamAdjust",
+  backendName: "webgpu",
+  kernelFunc: h
+};
+f(v);

package/dist/ops/webgpu/adamMoments.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/ops/webgpu/adamMoments.js

@@ -0,0 +1,51 @@
+import { g as m } from "../../webgpu_program-WaoMq-WD.js";
+import { f as u, c as p } from "../../webgpu_util-DhSeP4b6.js";
+import { f as c } from "../../index-UdZhlibC.js";
+class d {
+  variableNames = ["moments", "gradient"];
+  outputShape;
+  shaderKey = "AdamMoments";
+  dispatchLayout;
+  dispatch;
+  workgroupSize = [64, 1, 1];
+  size = !0;
+  uniforms = "beta1: f32, beta2: f32";
+  outputComponent = 2;
+  variableComponents = [2, 1];
+  constructor(e) {
+    this.outputShape = e, this.dispatchLayout = u(this.outputShape.slice(0, -1)), this.dispatch = p(
+      this.dispatchLayout,
+      this.outputShape.slice(0, -1),
+      this.workgroupSize,
+      [1, 1, 1]
+    );
+  }
+  getUserCode() {
+    return `
+        ${m("index")} {
+            if (index < uniforms.size) {
+                let m: vec2<f32> = moments[index];
+                let g: f32 = gradient[index];
+
+                let newM1 = fma(m.x, uniforms.beta1, g * (1.0 - uniforms.beta1));
+                let newM2 = fma(m.y, uniforms.beta2, g * g * (1.0 - uniforms.beta2));
+
+                setOutputAtIndex(index, vec2<f32>(newM1, newM2));
+            }
+        }
+    `;
+  }
+}
+function f(t) {
+  const { moments: e, gradient: a } = t.inputs, { beta1: n, beta2: o } = t.attrs, s = t.backend, i = new d(e.shape), r = [
+    { type: "float32", data: [n] },
+    { type: "float32", data: [o] }
+  ];
+  return s.runWebGPUProgram(i, [e, a], "float32", r);
+}
+const h = {
+  kernelName: "AdamMoments",
+  backendName: "webgpu",
+  kernelFunc: f
+};
+c(h);

package/dist/ops/webgpu/appendCache.js

@@ -1,5 +1,6 @@
-import { f as u, c as d, g as l } from "../../webgpu_util-g13LvDIv.js";
-import { f as m } from "../../index-C0dhsYom.js";
+import { g as u } from "../../webgpu_program-WaoMq-WD.js";
+import { f as d, c as m } from "../../webgpu_util-DhSeP4b6.js";
+import { f as l } from "../../index-UdZhlibC.js";
 class f {
   variableNames = ["cache", "item"];
   outputShape;
@@ -9,14 +10,14 @@ class f {
   workgroupSize = [64, 1, 1];
   size = !0;
   uniforms = "cacheT: i32";
-  constructor(t, a, s, o, c) {
-    const i = Math.min(s + 1, c);
-    this.outputShape = [t, a, i, o], this.dispatchLayout = u(this.outputShape), this.dispatch = d(this.dispatchLayout, this.outputShape, this.workgroupSize);
+  constructor(t, a, o, c, i) {
+    const s = Math.min(o + 1, i);
+    this.shaderKey = `AppendCache_${s}`, this.outputShape = [t, a, s, c], this.dispatchLayout = d(this.outputShape), this.dispatch = m(this.dispatchLayout, this.outputShape, this.workgroupSize);
   }
   getUserCode() {
     const t = this.outputShape[2];
     return `
-        ${l("index")} {
+        ${u("index")} {
             if (index < uniforms.size) {
                 let coords = getCoordsFromIndex(index); // [b, h, t, d]
                 let b = coords[0];
@@ -33,11 +34,11 @@ class f {
                 var val = 0.0;
                 if (srcT < uniforms.cacheT) {
                     val = getCache(b, h, srcT, d);
-                } else if (srcT == uniforms.cacheT) {
+                }
+                if (srcT == uniforms.cacheT) {
                     val = getItem(b, h, 0, d);
-                } else {
-                    val = 0.0;
                 }
+
                 setOutputAtIndex(index, val);
             }
         }
@@ -45,12 +46,12 @@ class f {
   }
 }
 function T(e) {
-  const { cache: t, item: a } = e.inputs, { maxSize: s, pastLen: o } = e.attrs, c = e.backend, i = t.shape[0], r = t.shape[2], n = t.shape[1], h = new f(i, n, r, a.shape[3], s), p = [{ type: "int32", data: [o] }];
-  return c.runWebGPUProgram(h, [t, a], "float32", p);
+  const { cache: t, item: a } = e.inputs, { maxSize: o, pastLen: c } = e.attrs, i = e.backend, s = t.shape[0], r = t.shape[2], n = t.shape[1], h = new f(s, n, r, a.shape[3], o), p = [{ type: "int32", data: [c] }];
+  return i.runWebGPUProgram(h, [t, a], "float32", p);
 }
 const g = {
   kernelName: "AppendCache",
   backendName: "webgpu",
   kernelFunc: T
 };
-m(g);
+l(g);

package/dist/ops/webgpu/attentionMask.js

@@ -1,5 +1,6 @@
-import { f } from "../../index-C0dhsYom.js";
-import { f as m, c as k, g as l } from "../../webgpu_util-g13LvDIv.js";
+import { f } from "../../index-UdZhlibC.js";
+import { g as m } from "../../webgpu_program-WaoMq-WD.js";
+import { f as k, c as l } from "../../webgpu_util-DhSeP4b6.js";
 class g {
   variableNames = ["q", "k"];
   outputShape;
@@ -13,13 +14,13 @@ class g {
   nh;
   T1;
   T2;
-  constructor(t, e, o, i, a) {
-    if (this.outputShape = [t, e, o, i], this.hs = a, this.nh = e, this.T1 = o, this.T2 = i, this.dispatchLayout = m(this.outputShape), this.dispatch = k(this.dispatchLayout, this.outputShape, this.workgroupSize), a % 4 !== 0)
+  constructor(t, e, i, r, s) {
+    if (this.shaderKey = `AttentionMask_${s}`, this.outputShape = [t, e, i, r], this.hs = s, this.nh = e, this.T1 = i, this.T2 = r, this.dispatchLayout = k(this.outputShape), this.dispatch = l(this.dispatchLayout, this.outputShape, this.workgroupSize), s % 4 !== 0)
       throw new Error("Head size must be a multiple of 4 for AttentionMaskProgram");
   }
   getUserCode() {
     return `
-            ${l("index")} {
+            ${m("index")} {
                 
                 let coords = getCoordsFromIndex(index);
                 let b = coords[0];
@@ -48,13 +49,13 @@ class g {
         `;
   }
 }
-function q(s) {
-  const { q: t, k: e } = s.inputs, { divisor: o, pastLen: i } = s.attrs, a = s.backend, n = t.shape[0], r = t.shape[2], u = e.shape[2], c = t.shape[1], d = t.shape[3], h = new g(n, c, r, u, d), p = [
-    { type: "float32", data: [o] },
-    { type: "int32", data: [i] },
+function q(o) {
+  const { q: t, k: e } = o.inputs, { divisor: i, pastLen: r } = o.attrs, s = o.backend, a = t.shape[0], n = t.shape[2], u = e.shape[2], d = t.shape[1], h = t.shape[3], c = new g(a, d, n, u, h), p = [
+    { type: "float32", data: [i] },
+    { type: "int32", data: [r] },
     { type: "float32", data: [Number.NEGATIVE_INFINITY] }
   ];
-  return a.runWebGPUProgram(h, [t, e], "float32", p);
+  return s.runWebGPUProgram(c, [t, e], "float32", p);
 }
 const v = {
   kernelName: "AttentionMask",

package/dist/ops/webgpu/gatherSub.js

@@ -1,6 +1,7 @@
-import { f as u, c as n, g as c } from "../../webgpu_util-g13LvDIv.js";
-import { f as p } from "../../index-C0dhsYom.js";
-class d {
+import { g as r } from "../../webgpu_program-WaoMq-WD.js";
+import { f as u, c as d } from "../../webgpu_util-DhSeP4b6.js";
+import { f as p } from "../../index-UdZhlibC.js";
+class l {
   variableNames = ["labels", "logits", "values"];
   outputShape;
   shaderKey = "GatherSub";
@@ -9,29 +10,43 @@ class d {
   workgroupSize = [64, 1, 1];
   size = !0;
   constructor(e) {
-    this.outputShape = [e], this.dispatchLayout = u(this.outputShape), this.dispatch = n(this.dispatchLayout, this.outputShape, this.workgroupSize);
+    this.outputShape = [e], this.dispatchLayout = u(this.outputShape), this.dispatch = d(this.dispatchLayout, this.outputShape, this.workgroupSize);
   }
   getUserCode() {
     return `
-        ${c("index")} {
+        ${r("index")} {
             if (index < uniforms.size) {
                 let coords = getCoordsFromIndex(index);
                 let idx = i32(getLabelsByOutputIndex(index));
                 let val = getValuesByOutputIndex(index);
-                let logit = getLogits(coords, idx);
-                setOutputAtIndex(index, val - logit);
+
+                if (isnan(val)) {
+                    setOutputAtIndex(index, 0.0);
+                    return;
+                }
+
+                if (idx < uniforms.logitsShape[1] && idx >= 0) {
+                    let logit = getLogits(coords, idx);
+                    if (isnan(logit)) {
+                        setOutputAtIndex(index, 0.0);
+                        return;
+                    }
+                    setOutputAtIndex(index, val - logit);
+                } else {
+                    setOutputAtIndex(index, 0.0);
+                }
             }
         }
     `;
   }
 }
-function l(t) {
-  const { logits: e, labels: a, values: s } = t.inputs, i = t.backend, o = a.shape[0], r = new d(o);
-  return i.runWebGPUProgram(r, [a, e, s], "float32");
+function c(t) {
+  const { logits: e, labels: i, values: s } = t.inputs, a = t.backend, n = i.shape[0], o = new l(n);
+  return a.runWebGPUProgram(o, [i, e, s], "float32");
 }
 const h = {
   kernelName: "EfficientGatherSub",
   backendName: "webgpu",
-  kernelFunc: l
+  kernelFunc: c
 };
 p(h);

package/dist/ops/webgpu/gelu.js

@@ -1,5 +1,6 @@
-import { f as i } from "../../index-C0dhsYom.js";
-import { f as o, c as s, g as p } from "../../webgpu_util-g13LvDIv.js";
+import { f as i } from "../../index-UdZhlibC.js";
+import { g as o } from "../../webgpu_program-WaoMq-WD.js";
+import { f as s, c as p } from "../../webgpu_util-DhSeP4b6.js";
 const u = 0.7978845608028654, a = 0.044715;
 class c {
   outputShape;
@@ -10,7 +11,7 @@ class c {
   workgroupSize;
   size = !0;
   constructor(e) {
-    this.workgroupSize = [128, 1, 1], this.outputShape = e, this.dispatchLayout = o(this.outputShape), this.dispatch = s(this.dispatchLayout, this.outputShape, this.workgroupSize), this.shaderKey = "unary_gelu";
+    this.workgroupSize = [128, 1, 1], this.outputShape = e, this.dispatchLayout = s(this.outputShape), this.dispatch = p(this.dispatchLayout, this.outputShape, this.workgroupSize), this.shaderKey = "unary_gelu";
   }
   getUserCode() {
     return `
@@ -22,7 +23,7 @@ class c {
         inner = 0.5 * (1.0 + inner);
         return x * inner;
       }
-      ${p("index")} {
+      ${o("index")} {
         if (index < uniforms.size) {
           let a = getAByOutputIndex(index);
           setOutputAtIndex(index, unaryOperation(a));
@@ -51,11 +52,11 @@ class x {
   workgroupSize = [128, 1, 1];
   size = !0;
   constructor(e) {
-    this.outputShape = e, this.dispatchLayout = o(this.outputShape), this.dispatch = s(this.dispatchLayout, this.outputShape, this.workgroupSize);
+    this.outputShape = e, this.dispatchLayout = s(this.outputShape), this.dispatch = p(this.dispatchLayout, this.outputShape, this.workgroupSize);
   }
   getUserCode() {
     return `
-            ${p("index")} {
+            ${o("index")} {
                 if (index < uniforms.size) {
                     let X  = getXByOutputIndex(index);
                     let x2 = X * X;

package/dist/ops/webgpu/index.js

@@ -1,8 +1,11 @@
 import "./attentionMask.js";
 import "./normRMS.js";
+import "./normRMSGrad.js";
 import "./rope.js";
 import "./appendCache.js";
 import "./scatterSub.js";
 import "./gatherSub.js";
 import "./qkv.js";
 import "./gelu.js";
+import "./adamMoments.js";
+import "./adamAdjust.js";

package/dist/ops/webgpu/normRMS.js

@@ -1,115 +1,41 @@
-import { f as m, c as p, g as c } from "../../webgpu_util-g13LvDIv.js";
-import { f as l, e as k } from "../../index-C0dhsYom.js";
-import { s as M } from "../../sum-B8wEpKsg.js";
-class N {
-  variableNames = ["x", "meanSquare", "gamma"];
+import { f as p } from "../../webgpu_util-DhSeP4b6.js";
+import { f as s } from "../../index-UdZhlibC.js";
+import { createReduceInfo as u, reduce as c, createReductionShader as m } from "./utils/reductions.js";
+class d {
   outputShape;
   shaderKey = "RMSNorm";
   dispatchLayout;
   dispatch;
   workgroupSize = [64, 1, 1];
+  variableNames = ["x", "gamma"];
+  uniforms = "reduceSize : i32,";
+  inputShape;
   size = !0;
-  constructor(t, e, a) {
-    this.outputShape = [t, e, a], this.dispatchLayout = m(this.outputShape), this.dispatch = p(this.dispatchLayout, this.outputShape, this.workgroupSize);
+  constructor(e) {
+    this.inputShape = [e.batchSize, e.inSize], this.outputShape = this.inputShape, this.dispatchLayout = p(this.outputShape), this.dispatch = [e.batchSize, 1, 1];
   }
   getUserCode() {
-    return `
-        ${c("index")} {
-            if (index < uniforms.size) {
-                let coords = getCoordsFromIndex(index);
-                let x = getXByOutputIndex(index);
-                let meanSquare = getMeanSquare(coords[0], coords[1], 0);
-                let gamma = getGammaByOutputIndex(index);
-                let invRms = inverseSqrt(meanSquare + 1e-8);
-                let normalized = x * invRms;
-                let outVal = normalized * gamma;
-                setOutputAtIndex(index, outVal);
-            }
-        }
-        `;
+    const e = this.workgroupSize[0];
+    return m(e, "mean", `
+            candidate = candidate * candidate;
+        `, `
+            bestValue = inverseSqrt(bestValue + 1e-8);
+        `, `
+            let X = f32(x[offset + k]);
+            let gamma = gamma[k];
+            let normalized = X * bestValue;
+            let outVal = normalized * gamma;
+            result[offset + k] = f32(outVal);
+        `);
   }
 }
-function b(s) {
-  const { x: t, gamma: e } = s.inputs, a = s.backend, o = t.shape[0], n = t.shape[1], i = t.shape[2], u = t.square().mean(-1, !0), r = new N(o, n, i);
-  return a.runWebGPUProgram(r, [t, u, e], "float32");
+function S(t) {
+  const { x: e, gamma: i } = t.inputs, n = t.backend, a = [e, i], o = u(a, -1), r = new d(o);
+  return c(r, a, n);
 }
-const z = {
+const h = {
   kernelName: "RMSNorm",
   backendName: "webgpu",
-  kernelFunc: b
+  kernelFunc: S
 };
-l(z);
-class R {
-  variableNames = ["x", "meanSquare", "dyGamma", "dyXMean"];
-  outputShape;
-  shaderKey = "RMSNormGradX";
-  dispatchLayout;
-  dispatch;
-  workgroupSize = [64, 1, 1];
-  size = !0;
-  C;
-  constructor(t, e, a) {
-    this.outputShape = [t, e, a], this.dispatchLayout = m(this.outputShape), this.dispatch = p(this.dispatchLayout, this.outputShape, this.workgroupSize), this.C = a;
-  }
-  getUserCode() {
-    return `
-        ${c("index")} {
-            if (index < uniforms.size) {
-                let coords = getCoordsFromIndex(index);
-                let x = getXByOutputIndex(index);
-                let meanSquare = getMeanSquare(coords[0], coords[1], 0) + 1e-8;
-                let dyGamma = getDyGammaByOutputIndex(index);
-                let dyXMean = getDyXMean(coords[0], coords[1], 0) / ${this.C}.0;
-                let invRms = inverseSqrt(meanSquare);
-                let dx = dyGamma * invRms - x * dyXMean * invRms / meanSquare;
-                setOutputAtIndex(index, dx);
-            }
-        }
-        `;
-  }
-}
-class v {
-  variableNames = ["x", "meanSquare", "dy"];
-  outputShape;
-  shaderKey = "RMSNormGradGamma";
-  dispatchLayout;
-  dispatch;
-  workgroupSize = [64, 1, 1];
-  size = !0;
-  constructor(t, e, a) {
-    this.outputShape = [t, e, a], this.dispatchLayout = m(this.outputShape), this.dispatch = p(this.dispatchLayout, this.outputShape, this.workgroupSize);
-  }
-  getUserCode() {
-    return `
-        ${c("index")} {
-            if (index < uniforms.size) {
-                let coords = getCoordsFromIndex(index);
-                let x = getXByOutputIndex(index);
-                let meanSquare = getMeanSquare(coords[0], coords[1], 0) + 1e-8;
-                let dy = getDyByOutputIndex(index);
-                let invRms = inverseSqrt(meanSquare);
-                let dGamma = dy * (x * invRms);
-                setOutputAtIndex(index,dGamma);
-            }
-        }
-        `;
-  }
-}
-function I(s) {
-  const { dy: t, x: e, gamma: a } = s.inputs, o = s.backend, n = e.shape[0], i = e.shape[1], u = e.shape[2], r = t.mul(a), h = r.mul(e), g = h.sum(-1, !0);
-  h.dispose();
-  const S = e.square(), d = S.mean(-1, !0);
-  S.dispose();
-  const y = new R(n, i, u), G = o.runWebGPUProgram(y, [e, d, r, g], "float32");
-  r.dispose(), g.dispose();
-  const q = new v(n, i, u), x = o.runWebGPUProgram(q, [e, d, t], "float32");
-  d.dispose();
-  const f = M(k().makeTensorFromTensorInfo(x), [0, 1]);
-  return o.disposeData(x), [G, f];
-}
-const P = {
-  kernelName: "RMSNormGrad",
-  backendName: "webgpu",
-  kernelFunc: I
-};
-l(P);
+s(h);

package/dist/ops/webgpu/normRMSGrad.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/ops/webgpu/normRMSGrad.js

@@ -0,0 +1,128 @@
+import { f as h, e as S } from "../../index-UdZhlibC.js";
+import { createReduceInfo as g } from "./utils/reductions.js";
+import { f as k } from "../../webgpu_util-DhSeP4b6.js";
+import { g as x } from "../../webgpu_program-WaoMq-WD.js";
+import { s as c } from "../../slice-BHbDHObE.js";
+import { s as z } from "../../sum-NLYbiDag.js";
+class w {
+  outputShape;
+  shaderKey = "RMSNormGrad";
+  dispatchLayout;
+  dispatch;
+  workgroupSize = [64, 1, 1];
+  variableNames = ["x", "gamma", "dy"];
+  uniforms = "reduceSize : i32, batchSize: i32";
+  inputShape;
+  size = !1;
+  rowsPerWorkgroup;
+  constructor(r, e = 4) {
+    if (this.shaderKey = `RMSNormGrad_${e}`, this.rowsPerWorkgroup = e, this.inputShape = [r.batchSize, r.inSize], this.outputShape = [r.batchSize + r.batchSize / this.rowsPerWorkgroup, r.inSize], this.dispatchLayout = k(this.outputShape), this.dispatch = [r.batchSize / this.rowsPerWorkgroup, 1, 1], r.batchSize % this.rowsPerWorkgroup !== 0)
+      throw new Error(
+        `RMSNormGradProgram: batch size ${r.batchSize} must be divisible by rowsPerWorkgroup ${this.rowsPerWorkgroup}`
+      );
+    if (r.inSize > 1024)
+      throw new Error(`RMSNormGradProgram: inSize ${r.inSize} exceeds max of 1024`);
+  }
+  getUserCode() {
+    const r = this.workgroupSize[0], e = this.rowsPerWorkgroup;
+    return `
+          fn DIV_CEIL(a : u32, b : u32) -> u32 {
+            return ((a - 1u) / b + 1u);
+          }
+
+          ${`
+          var<workgroup> partials : array<vec2<f32>, ${r}>;
+          var<workgroup> accumulation: array<f32, 1024>;
+        `}
+
+          ${x("index")} {
+            // One workgroup per row (batch).
+            let Length = uniforms.reduceSize;
+            let BatchSize = uniforms.batchSize;
+            for (var k = i32(localId.x); k < Length; k = k + ${r}) {
+                accumulation[k] = 0.0;
+            }
+
+            for (var rowOff = 0; rowOff < ${e}; rowOff = rowOff + 1) {
+                let row = i32(workgroupId.x) * ${e} + rowOff;
+                let offset = row * Length;
+                
+                var sum_x2 = 0.0;
+                var sum_dygx = 0.0;
+
+                for (var k = i32(localId.x); k < Length; k = k + ${r}) {
+                    let X = f32(x[offset + k]);
+                    let DY = f32(dy[offset + k]);
+                    let G  = f32(gamma[k]);
+                    sum_x2 = fma(X, X, sum_x2);
+                    sum_dygx = fma(DY * G, X, sum_dygx);
+                }
+
+                partials[localId.x] = vec2<f32>(sum_x2, sum_dygx);
+                workgroupBarrier();
+
+                var reduceSize = min(u32(Length), ${r}u);
+                for (var currentSize = reduceSize / 2u; reduceSize > 1u; currentSize = reduceSize / 2u) {
+                    let interval = DIV_CEIL(reduceSize, 2u);
+                    if (localId.x < currentSize) {
+                        partials[localId.x] = partials[localId.x] + partials[localId.x + interval];
+                    }
+                    reduceSize = interval;
+                    workgroupBarrier();
+                }
+
+                let invN = 1.0 / f32(Length);
+                let mean_x2   = fma(partials[0].x, invN, 1e-8);
+                let mean_dygx = partials[0].y * invN;
+
+                let invRMS = inverseSqrt(mean_x2);
+                let scale = (mean_dygx / (mean_x2)) * invRMS;
+
+                // write dx and dGamma.
+                for (var k = i32(localId.x); k < Length; k = k + ${r}) {
+                    let X  = f32(x[offset + k]);
+                    let DY = f32(dy[offset + k]);
+                    let G  = f32(gamma[k]);
+
+                    let dyGamma = DY * G;
+                    let dx = fma(dyGamma, invRMS, -X * scale);
+
+                    result[offset + k] = dx;
+
+                    // dGamma
+                    accumulation[k] = fma(DY, X * invRMS, accumulation[k]);
+                }
+
+                workgroupBarrier();
+            }
+
+            // Write out the partially accumulated dGamma
+            let outDgBase = BatchSize * Length + i32(workgroupId.x) * Length;
+            for (var k = i32(localId.x); k < Length; k = k + ${r}) {
+                result[outDgBase + k] = accumulation[k];
+            }
+          }
+        `;
+  }
+}
+function v(t) {
+  const { dy: r, x: e, gamma: o } = t.inputs, i = 4, l = t.backend, a = g([e, o, r], -1), s = new w(a, i), d = [
+    { type: "int32", data: [s.inputShape[1]] },
+    // Reduce size
+    { type: "int32", data: [s.inputShape[0]] }
+    // Batch size
+  ], n = S().makeTensorFromTensorInfo(
+    l.runWebGPUProgram(s, [e, o, r], "float32", d)
+  ), u = c(n, [0, 0], [a.batchSize, a.inSize]), m = c(n, [a.batchSize, 0], [a.batchSize / i, a.inSize]);
+  n.dispose();
+  const p = u.reshape(e.shape);
+  u.dispose();
+  const f = z(m, [0]);
+  return m.dispose(), [p, f];
+}
+const y = {
+  kernelName: "RMSNormGrad",
+  backendName: "webgpu",
+  kernelFunc: v
+};
+h(y);

package/dist/ops/webgpu/qkv.js

@@ -1,6 +1,7 @@
-import { f as c, c as d, g as h } from "../../webgpu_util-g13LvDIv.js";
-import { f as p } from "../../index-C0dhsYom.js";
-class l {
+import { g as c } from "../../webgpu_program-WaoMq-WD.js";
+import { f as d, c as h } from "../../webgpu_util-DhSeP4b6.js";
+import { f as p } from "../../index-UdZhlibC.js";
+class m {
   variableNames = ["x", "kernel"];
   outputShape;
   shaderKey = "QKV";
@@ -11,12 +12,12 @@ class l {
   size = !0;
   constructor(t, e, o, s) {
     const r = s / e;
-    this.outputShape = [t, e, o, r], this.dispatchLayout = c(this.outputShape), this.dispatch = d(this.dispatchLayout, this.outputShape, this.workgroupSize);
+    this.shaderKey = `QKV_${e}_${r}`, this.outputShape = [t, e, o, r], this.dispatchLayout = d(this.outputShape), this.dispatch = h(this.dispatchLayout, this.outputShape, this.workgroupSize);
   }
   getUserCode() {
     const t = this.outputShape[1], e = this.outputShape[3], o = t * e;
     return `
-        ${h("index")} {
+        ${c("index")} {
             if (index < uniforms.size) {
                 let coords = getCoordsFromIndex(index); // [b, h, t, d]
                 let b = coords[0];
@@ -40,8 +41,8 @@ class l {
         `;
   }
 }
-function m(a) {
-  const { x: t, kernel: e } = a.inputs, { heads: o } = a.attrs, s = a.backend, r = t.shape[0], i = t.shape[1], u = t.shape[2], n = new l(r, o, i, u);
+function l(a) {
+  const { x: t, kernel: e } = a.inputs, { heads: o } = a.attrs, s = a.backend, r = t.shape[0], i = t.shape[1], u = t.shape[2], n = new m(r, o, i, u);
   return [
     s.runWebGPUProgram(n, [t, e], "float32", [{ type: "int32", data: [0] }]),
     s.runWebGPUProgram(n, [t, e], "float32", [{ type: "int32", data: [1] }]),
@@ -51,6 +52,6 @@ function m(a) {
 const f = {
   kernelName: "QKV",
   backendName: "webgpu",
-  kernelFunc: m
+  kernelFunc: l
 };
 p(f);