@genai-fi/nanogpt 0.19.1 → 0.20.1

package/dist/ops/webgpu/norm2.js

@@ -1,76 +1,46 @@
-import { reduce as g, ReduceProgram as h } from "./utils/reductions.js";
-import { c as w, U as S } from "../../index-CUXkjxiT.js";
-import k from "./utils/deviceInfo.js";
-class v extends h {
-  shaderKey = "norm2";
-  atomic = !0;
-  utilityFunctions = `
-        fn atomicAddF32(sum: ptr<storage, atomic<i32>, read_write>, value: f32) -> f32 {
-            var old = atomicLoad(sum);
-            loop {
-                let new_value = value + bitcast<f32>(old);
-                let exchange_result = atomicCompareExchangeWeak(sum, old, bitcast<i32>(new_value));
-                if (exchange_result.exchanged) {
-                    return new_value;
-                }
-                old = exchange_result.old_value;
-            }
-        }
-    `;
-  constructor(o, r, i) {
-    super(
-      o,
-      r,
-      {
-        reductionOp: "sum",
-        elementwise: !1,
-        forceWorkgroupSize: i
-      },
-      !1
-    ), this.uniforms += "lossScaling: f32, index: i32";
-  }
-  getPreprocessSnippet() {
-    return `
-            candidate = candidate * uniforms.lossScaling;
-            candidate = candidate * candidate;
-        `;
-  }
-  getWriteSnippet() {
-    return `
-            if (tid == 0) {
-                atomicAddF32(&result[uniforms.index], bestValue);
-            }
-        `;
-  }
-}
-function x(t) {
-  const { x: o, output: r } = t.inputs, { invLossScaling: i, index: c } = t.attrs, n = t.backend, u = [], d = k(n);
-  let e = Math.min(512, n.device.limits.maxComputeWorkgroupSizeX);
-  const a = 4, s = S(o.shape);
-  for (; s % (e * a) !== 0 && e > 1; )
-    e /= 2;
-  if (e === 1)
-    throw new Error(`Cannot find suitable workgroup size for Norm2Program with reduce size ${s}`);
-  const l = {
-    inSize: e * a,
-    outSize: 1,
-    batchSize: s / (e * a),
-    windowSize: e
-  }, m = new v(d, l, e), p = g(
-    m,
-    [o],
-    n,
-    [
-      { type: "float32", data: [i] },
-      { type: "int32", data: [c] }
-    ],
-    r
-  );
-  return u.forEach((f) => f.dispose()), p;
-}
-const z = {
-  kernelName: "Norm2",
-  backendName: "webgpu",
-  kernelFunc: x
+import { Ii as e, oc as t } from "../../dist-BewPQWjc.js";
+import { ReduceProgram as n, reduce as r } from "./utils/reductions.js";
+import i from "./utils/deviceInfo.js";
+//#region lib/ops/webgpu/norm2.ts
+var a = class extends n {
+	shaderKey = "norm2";
+	atomic = !0;
+	utilityFunctions = "\n        fn atomicAddF32(sum: ptr<storage, atomic<i32>, read_write>, value: f32) -> f32 {\n            var old = atomicLoad(sum);\n            loop {\n                let new_value = value + bitcast<f32>(old);\n                let exchange_result = atomicCompareExchangeWeak(sum, old, bitcast<i32>(new_value));\n                if (exchange_result.exchanged) {\n                    return new_value;\n                }\n                old = exchange_result.old_value;\n            }\n        }\n    ";
+	constructor(e, t, n) {
+		super(e, t, {
+			reductionOp: "sum",
+			elementwise: !1,
+			forceWorkgroupSize: n
+		}, !1), this.uniforms += "lossScaling: f32, index: i32";
+	}
+	getPreprocessSnippet() {
+		return "\n            candidate = candidate * uniforms.lossScaling;\n            candidate = candidate * candidate;\n        ";
+	}
+	getWriteSnippet() {
+		return "\n            if (tid == 0) {\n                atomicAddF32(&result[uniforms.index], bestValue);\n            }\n        ";
+	}
 };
-w(z);
+function o(e) {
+	let { x: n, output: o } = e.inputs, { invLossScaling: s, index: c } = e.attrs, l = e.backend, u = [], d = i(l), f = Math.min(512, l.device.limits.maxComputeWorkgroupSizeX), p = t(n.shape);
+	for (; p % (f * 4) != 0 && f > 1;) f /= 2;
+	if (f === 1) throw Error(`Cannot find suitable workgroup size for Norm2Program with reduce size ${p}`);
+	let m = r(new a(d, {
+		inSize: f * 4,
+		outSize: 1,
+		batchSize: p / (f * 4),
+		windowSize: f
+	}, f), [n], l, [{
+		type: "float32",
+		data: [s]
+	}, {
+		type: "int32",
+		data: [c]
+	}], o);
+	return u.forEach((e) => e.dispose()), m;
+}
+e({
+	kernelName: "Norm2",
+	backendName: "webgpu",
+	kernelFunc: o
+});
+//#endregion

package/dist/ops/webgpu/normRMS.js

@@ -1,34 +1,26 @@
-import { c as k, a6 as N } from "../../index-CUXkjxiT.js";
-import { createReduceInfo as w, reduce as S } from "./utils/reductions.js";
-import { isPackedTensor as d } from "../../utilities/packed.js";
-import { p as f } from "../../pack16-bqltoUlR.js";
-import b from "./normRMS16_program.js";
-import z from "./normRMS32_program.js";
-import G from "./utils/deviceInfo.js";
-function u(t) {
-  const { x: e, gamma: o } = t.inputs, m = t.backend, c = G(m), p = d(e), i = o ? d(o) : !1, n = p || i, r = !n || n ? e : f(e), a = !n || !o ? o : f(o), h = a ? [r, a] : [r], s = w(h, -1), l = n ? new b(c, s, !!o) : new z(c, s, !!o);
-  if (a && N(a.shape, [r.shape[r.shape.length - 1]], "Error in RMSNorm: "), e.shape.length !== 3)
-    throw new Error(`rmsNormGPU: input rank ${e.shape.length} not supported, only rank 3 is supported`);
-  if (s.inSize !== r.shape[r.shape.length - 1])
-    throw new Error(
-      `rmsNormGPU: reduction size ${s.inSize} does not match expected size ${r.shape[r.shape.length - 1]}`
-    );
-  if (s.batchSize !== e.shape[0] * e.shape[1])
-    throw new Error(
-      `rmsNormGPU: batch size ${s.batchSize} does not match expected size ${e.shape[0] * e.shape[1]}`
-    );
-  const g = S(l, h, m);
-  return n && !p && r.dispose(), n && !i && a && a.dispose(), g;
+import { Fs as e, Ii as t } from "../../dist-BewPQWjc.js";
+import { isPackedTensor as n } from "../../utilities/packed.js";
+import { t as r } from "../../pack16-Ck-spx_F.js";
+import { createReduceInfo as i, reduce as a } from "./utils/reductions.js";
+import o from "./utils/deviceInfo.js";
+import s from "./normRMS32_program.js";
+import c from "./normRMS16_program.js";
+//#region lib/ops/webgpu/normRMS.ts
+function l(t) {
+	let { x: l, gamma: u } = t.inputs, d = t.backend, f = o(d), p = n(l), m = u ? n(u) : !1, h = p || m, g = !h || h ? l : r(l), _ = !h || !u ? u : r(u), v = _ ? [g, _] : [g], y = i(v, -1), b = h ? new c(f, y, !!u) : new s(f, y, !!u);
+	if (_ && e(_.shape, [g.shape[g.shape.length - 1]], "Error in RMSNorm: "), l.shape.length !== 3) throw Error(`rmsNormGPU: input rank ${l.shape.length} not supported, only rank 3 is supported`);
+	if (y.inSize !== g.shape[g.shape.length - 1]) throw Error(`rmsNormGPU: reduction size ${y.inSize} does not match expected size ${g.shape[g.shape.length - 1]}`);
+	if (y.batchSize !== l.shape[0] * l.shape[1]) throw Error(`rmsNormGPU: batch size ${y.batchSize} does not match expected size ${l.shape[0] * l.shape[1]}`);
+	let x = a(b, v, d);
+	return h && !p && g.dispose(), h && !m && _ && _.dispose(), x;
 }
-const P = {
-  kernelName: "RMSNorm",
-  backendName: "webgpu",
-  kernelFunc: u
-};
-k(P);
-const M = {
-  kernelName: "RMSNormNoGamma",
-  backendName: "webgpu",
-  kernelFunc: u
-};
-k(M);
+t({
+	kernelName: "RMSNorm",
+	backendName: "webgpu",
+	kernelFunc: l
+}), t({
+	kernelName: "RMSNormNoGamma",
+	backendName: "webgpu",
+	kernelFunc: l
+});
+//#endregion

package/dist/ops/webgpu/normRMS16_program.js

@@ -1,25 +1,28 @@
-import { ReduceProgram as s } from "./utils/reductions.js";
-class i extends s {
-  hasGamma;
-  constructor(a, t, e = !1) {
-    super(a, t, { reductionOp: "mean", elementwise: !0 }, !0), this.shaderKey = "RMSNorm16", e && this.variableNames.push("gamma"), this.variableComponents = [1, 1], this.hasGamma = e;
-  }
-  getPreprocessSnippet() {
-    return "candidate = candidate * candidate;";
-  }
-  getPostprocessSnippet() {
-    return "bestValue = inverseSqrt(bestValue + 1e-8);";
-  }
-  getWriteSnippet() {
-    return `
+import { ReduceProgram as e } from "./utils/reductions.js";
+//#region lib/ops/webgpu/normRMS16_program.ts
+var t = class extends e {
+	hasGamma;
+	constructor(e, t, n = !1) {
+		super(e, t, {
+			reductionOp: "mean",
+			elementwise: !0
+		}, !0), this.shaderKey = "RMSNorm16", n && this.variableNames.push("gamma"), this.variableComponents = [1, 1], this.hasGamma = n;
+	}
+	getPreprocessSnippet() {
+		return "candidate = candidate * candidate;";
+	}
+	getPostprocessSnippet() {
+		return "bestValue = inverseSqrt(bestValue + 1e-8);";
+	}
+	getWriteSnippet() {
+		return `
             let X = unpack2x16float(u32(x[offset + k]));
             ${this.hasGamma ? "let gamma = unpack2x16float(u32(gamma[k]));" : ""}
             let normalized = X * bestValue;
             let outVal = normalized ${this.hasGamma ? " * gamma" : ""};
             result[offset + k] = i32(pack2x16float(outVal));
         `;
-  }
-}
-export {
-  i as default
+	}
 };
+//#endregion
+export { t as default };

package/dist/ops/webgpu/normRMS32_program.js

@@ -1,25 +1,28 @@
-import { ReduceProgram as s } from "./utils/reductions.js";
-class i extends s {
-  hasGamma;
-  constructor(a, t, e = !1) {
-    super(a, t, { reductionOp: "mean", elementwise: !0 }, !1), this.shaderKey = "RMSNorm32", e && this.variableNames.push("gamma"), this.variableComponents = [1, 1], this.hasGamma = e;
-  }
-  getPreprocessSnippet() {
-    return "candidate = candidate * candidate;";
-  }
-  getPostprocessSnippet() {
-    return "bestValue = inverseSqrt(bestValue + 1e-8);";
-  }
-  getWriteSnippet() {
-    return `
+import { ReduceProgram as e } from "./utils/reductions.js";
+//#region lib/ops/webgpu/normRMS32_program.ts
+var t = class extends e {
+	hasGamma;
+	constructor(e, t, n = !1) {
+		super(e, t, {
+			reductionOp: "mean",
+			elementwise: !0
+		}, !1), this.shaderKey = "RMSNorm32", n && this.variableNames.push("gamma"), this.variableComponents = [1, 1], this.hasGamma = n;
+	}
+	getPreprocessSnippet() {
+		return "candidate = candidate * candidate;";
+	}
+	getPostprocessSnippet() {
+		return "bestValue = inverseSqrt(bestValue + 1e-8);";
+	}
+	getWriteSnippet() {
+		return `
             let X = f32(x[offset + k]);
             ${this.hasGamma ? "let gamma = gamma[k];" : ""}
             let normalized = X * bestValue;
             let outVal = normalized ${this.hasGamma ? " * gamma" : ""};
             result[offset + k] = f32(outVal);
         `;
-  }
-}
-export {
-  i as default
+	}
 };
+//#endregion
+export { t as default };

package/dist/ops/webgpu/normRMSGrad.js

@@ -1,91 +1,55 @@
-import { c as $, a6 as y, h as L } from "../../index-CUXkjxiT.js";
-import { createReduceInfo as M } from "./utils/reductions.js";
-import { f as w } from "../../webgpu_util-DYlGSwOJ.js";
-import { e as _ } from "../../webgpu_program-B4HmApL1.js";
-import { p as x, u as R } from "../../pack16-bqltoUlR.js";
-import { isPackedTensor as h } from "../../utilities/packed.js";
-import { reshape16 as z } from "../reshape16.js";
-import { sum16 as N } from "../sum16.js";
-import { slice16 as v } from "../slice16.js";
-class Y {
-  outputShape;
-  shaderKey = "RMSNormGrad";
-  dispatchLayout;
-  dispatch;
-  workgroupSize = [64, 1, 1];
-  variableNames = ["x", "gamma", "dy"];
-  uniforms = "reduceSize : i32, batchSize: i32";
-  inputShape;
-  size = !1;
-  rowsPerWorkgroup;
-  packed = !1;
-  outputComponent;
-  constructor(e, a = 4, t = !1) {
-    if (this.packed = t, this.shaderKey = `RMSNormGrad_${a}`, this.rowsPerWorkgroup = a, this.inputShape = [e.batchSize, e.inSize], this.outputShape = [e.batchSize + e.batchSize / this.rowsPerWorkgroup, e.inSize], this.dispatchLayout = w(this.outputShape), this.dispatch = [e.batchSize / this.rowsPerWorkgroup, 1, 1], e.batchSize % this.rowsPerWorkgroup !== 0)
-      throw new Error(
-        `RMSNormGradProgram: batch size ${e.batchSize} must be divisible by rowsPerWorkgroup ${this.rowsPerWorkgroup}`
-      );
-    if (e.inSize > 1024)
-      throw new Error(`RMSNormGradProgram: inSize ${e.inSize} exceeds max of 1024`);
-  }
-  getUserCode() {
-    const e = this.workgroupSize[0], a = this.rowsPerWorkgroup, t = `
+import { Fs as e, Ii as t, ii as n } from "../../dist-BewPQWjc.js";
+import { isPackedTensor as r } from "../../utilities/packed.js";
+import { reshape16 as i } from "../reshape16.js";
+import { sum16 as a } from "../sum16.js";
+import { r as o, t as s } from "../../pack16-Ck-spx_F.js";
+import { s as c } from "../../webgpu_program-WOyIVMlZ.js";
+import { c as l } from "../../webgpu_util-B_F3SShA.js";
+import { slice16 as u } from "../slice16.js";
+import { createReduceInfo as d } from "./utils/reductions.js";
+//#region lib/ops/webgpu/normRMSGrad.ts
+var f = class {
+	outputShape;
+	shaderKey = "RMSNormGrad";
+	dispatchLayout;
+	dispatch;
+	workgroupSize = [
+		64,
+		1,
+		1
+	];
+	variableNames = [
+		"x",
+		"gamma",
+		"dy"
+	];
+	uniforms = "reduceSize : i32, batchSize: i32";
+	inputShape;
+	size = !1;
+	rowsPerWorkgroup;
+	packed = !1;
+	outputComponent;
+	constructor(e, t = 4, n = !1) {
+		if (this.packed = n, this.shaderKey = `RMSNormGrad_${t}`, this.rowsPerWorkgroup = t, this.inputShape = [e.batchSize, e.inSize], this.outputShape = [e.batchSize + e.batchSize / this.rowsPerWorkgroup, e.inSize], this.dispatchLayout = l(this.outputShape), this.dispatch = [
+			e.batchSize / this.rowsPerWorkgroup,
+			1,
+			1
+		], e.batchSize % this.rowsPerWorkgroup !== 0) throw Error(`RMSNormGradProgram: batch size ${e.batchSize} must be divisible by rowsPerWorkgroup ${this.rowsPerWorkgroup}`);
+		if (e.inSize > 1024) throw Error(`RMSNormGradProgram: inSize ${e.inSize} exceeds max of 1024`);
+	}
+	getUserCode() {
+		let e = this.workgroupSize[0], t = this.rowsPerWorkgroup, n = `
           var<workgroup> partials : array<vec2<f32>, ${e}>;
           var<workgroup> accumulation: array<${this.packed ? "vec2<f32>" : "f32"}, 1024>;
-        `, u = this.packed ? `
-                let X = unpack2x16float(u32(x[offset + k]));
-                let DY = unpack2x16float(u32(dy[offset + k]));
-                let G  = unpack2x16float(u32(gamma[k]));
-                sum_x2 = fma(X.x, X.x, sum_x2);
-                sum_x2 = fma(X.y, X.y, sum_x2);
-                sum_dygx = fma(DY.x * G.x, X.x, sum_dygx);
-                sum_dygx = fma(DY.y * G.y, X.y, sum_dygx);
-        ` : `
-                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);
-            `, m = this.packed ? `
-                let X  = unpack2x16float(u32(x[offset + k]));
-                let DY = unpack2x16float(u32(dy[offset + k]));
-                let G  = unpack2x16float(u32(gamma[k]));
-
-                let dyGamma = DY * G;
-                let dx = vec2<f32>(
-                    fma(dyGamma.x, invRMS, -X.x * scale),
-                    fma(dyGamma.y, invRMS, -X.y * scale)
-                );
-
-                result[offset + k] = i32(pack2x16float(dx));
-
-                // dGamma
-                accumulation[k] = fma(DY, X * invRMS, accumulation[k]);
-        ` : `
-                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]);
-        `, n = this.packed ? `
-                result[outDgBase + k] = i32(pack2x16float(accumulation[k]));
-        ` : `
-                result[outDgBase + k] = accumulation[k];
-        `;
-    return `
+        `, r = this.packed ? "\n                let X = unpack2x16float(u32(x[offset + k]));\n                let DY = unpack2x16float(u32(dy[offset + k]));\n                let G  = unpack2x16float(u32(gamma[k]));\n                sum_x2 = fma(X.x, X.x, sum_x2);\n                sum_x2 = fma(X.y, X.y, sum_x2);\n                sum_dygx = fma(DY.x * G.x, X.x, sum_dygx);\n                sum_dygx = fma(DY.y * G.y, X.y, sum_dygx);\n        " : "\n                let X = f32(x[offset + k]);\n                let DY = f32(dy[offset + k]);\n                let G  = f32(gamma[k]);\n                sum_x2 = fma(X, X, sum_x2);\n                sum_dygx = fma(DY * G, X, sum_dygx);\n            ", i = this.packed ? "\n                let X  = unpack2x16float(u32(x[offset + k]));\n                let DY = unpack2x16float(u32(dy[offset + k]));\n                let G  = unpack2x16float(u32(gamma[k]));\n\n                let dyGamma = DY * G;\n                let dx = vec2<f32>(\n                    fma(dyGamma.x, invRMS, -X.x * scale),\n                    fma(dyGamma.y, invRMS, -X.y * scale)\n                );\n\n                result[offset + k] = i32(pack2x16float(dx));\n\n                // dGamma\n                accumulation[k] = fma(DY, X * invRMS, accumulation[k]);\n        " : "\n                let X  = f32(x[offset + k]);\n                let DY = f32(dy[offset + k]);\n                let G  = f32(gamma[k]);\n\n                let dyGamma = DY * G;\n                let dx = fma(dyGamma, invRMS, -X * scale);\n\n                result[offset + k] = dx;\n\n                // dGamma\n                accumulation[k] = fma(DY, X * invRMS, accumulation[k]);\n        ", a = this.packed ? "\n                result[outDgBase + k] = i32(pack2x16float(accumulation[k]));\n        " : "\n                result[outDgBase + k] = accumulation[k];\n        ";
+		return `
           fn DIV_CEIL(a : u32, b : u32) -> u32 {
             return ((a - 1u) / b + 1u);
           }
 
-          ${t}
+          ${n}
 
-          ${_("index")} {
+          ${c("index")} {
             // One workgroup per row (batch).
             let Length = uniforms.reduceSize;
             let BatchSize = uniforms.batchSize;
@@ -93,15 +57,15 @@ class Y {
                 accumulation[k] = ${this.packed ? "vec2<f32>(0.0f)" : "0.0f"};
             }
 
-            for (var rowOff = 0; rowOff < ${a}; rowOff = rowOff + 1) {
-                let row = i32(workgroupId.x) * ${a} + rowOff;
+            for (var rowOff = 0; rowOff < ${t}; rowOff = rowOff + 1) {
+                let row = i32(workgroupId.x) * ${t} + rowOff;
                 let offset = row * Length;
                 
                 var sum_x2 = 0.0f;
                 var sum_dygx = 0.0f;
 
                 for (var k = i32(localId.x); k < Length; k = k + ${e}) {
-                    ${u}
+                    ${r}
                 }
 
                 partials[localId.x] = vec2<f32>(sum_x2, sum_dygx);
@@ -126,7 +90,7 @@ class Y {
 
                 // write dx and dGamma.
                 for (var k = i32(localId.x); k < Length; k = k + ${e}) {
-                    ${m}
+                    ${i}
                 }
 
                 workgroupBarrier();
@@ -135,68 +99,46 @@ class Y {
             // Write out the partially accumulated dGamma
             let outDgBase = BatchSize * Length + i32(workgroupId.x) * Length;
             for (var k = i32(localId.x); k < Length; k = k + ${e}) {
-                ${n}
+                ${a}
             }
           }
         `;
-  }
-}
-class I {
-  outputShape;
-  shaderKey = "RMSNormGrad";
-  dispatchLayout;
-  dispatch;
-  workgroupSize = [64, 1, 1];
-  variableNames = ["x", "dy"];
-  uniforms = "reduceSize : i32, batchSize: i32";
-  inputShape;
-  size = !1;
-  packed = !1;
-  outputComponent;
-  constructor(e, a = !1) {
-    this.packed = a, this.shaderKey = "RMSNormGrad_NoGamma", this.inputShape = [e.batchSize, e.inSize], this.outputShape = [e.batchSize, e.inSize], this.dispatchLayout = w(this.outputShape), this.dispatch = [e.batchSize, 1, 1];
-  }
-  getUserCode() {
-    const e = this.workgroupSize[0], a = `
+	}
+}, p = class {
+	outputShape;
+	shaderKey = "RMSNormGrad";
+	dispatchLayout;
+	dispatch;
+	workgroupSize = [
+		64,
+		1,
+		1
+	];
+	variableNames = ["x", "dy"];
+	uniforms = "reduceSize : i32, batchSize: i32";
+	inputShape;
+	size = !1;
+	packed = !1;
+	outputComponent;
+	constructor(e, t = !1) {
+		this.packed = t, this.shaderKey = "RMSNormGrad_NoGamma", this.inputShape = [e.batchSize, e.inSize], this.outputShape = [e.batchSize, e.inSize], this.dispatchLayout = l(this.outputShape), this.dispatch = [
+			e.batchSize,
+			1,
+			1
+		];
+	}
+	getUserCode() {
+		let e = this.workgroupSize[0], t = `
           var<workgroup> partials : array<vec2<f32>, ${e}>;
-        `, t = this.packed ? `
-                let X = unpack2x16float(u32(x[offset + k]));
-                let DY = unpack2x16float(u32(dy[offset + k]));
-                sum_x2 = fma(X.x, X.x, sum_x2);
-                sum_x2 = fma(X.y, X.y, sum_x2);
-                sum_dygx = fma(DY.x, X.x, sum_dygx);
-                sum_dygx = fma(DY.y, X.y, sum_dygx);
-        ` : `
-                let X = f32(x[offset + k]);
-                let DY = f32(dy[offset + k]);
-                sum_x2 = fma(X, X, sum_x2);
-                sum_dygx = fma(DY, X, sum_dygx);
-            `, u = this.packed ? `
-                let X  = unpack2x16float(u32(x[offset + k]));
-                let DY = unpack2x16float(u32(dy[offset + k]));
-
-                let dx = vec2<f32>(
-                    fma(DY.x, invRMS, -X.x * scale),
-                    fma(DY.y, invRMS, -X.y * scale)
-                );
-
-                result[offset + k] = i32(pack2x16float(dx));
-        ` : `
-                let X  = f32(x[offset + k]);
-                let DY = f32(dy[offset + k]);
-
-                let dx = fma(DY, invRMS, -X * scale);
-
-                result[offset + k] = dx;
-        `;
-    return `
+        `, n = this.packed ? "\n                let X = unpack2x16float(u32(x[offset + k]));\n                let DY = unpack2x16float(u32(dy[offset + k]));\n                sum_x2 = fma(X.x, X.x, sum_x2);\n                sum_x2 = fma(X.y, X.y, sum_x2);\n                sum_dygx = fma(DY.x, X.x, sum_dygx);\n                sum_dygx = fma(DY.y, X.y, sum_dygx);\n        " : "\n                let X = f32(x[offset + k]);\n                let DY = f32(dy[offset + k]);\n                sum_x2 = fma(X, X, sum_x2);\n                sum_dygx = fma(DY, X, sum_dygx);\n            ", r = this.packed ? "\n                let X  = unpack2x16float(u32(x[offset + k]));\n                let DY = unpack2x16float(u32(dy[offset + k]));\n\n                let dx = vec2<f32>(\n                    fma(DY.x, invRMS, -X.x * scale),\n                    fma(DY.y, invRMS, -X.y * scale)\n                );\n\n                result[offset + k] = i32(pack2x16float(dx));\n        " : "\n                let X  = f32(x[offset + k]);\n                let DY = f32(dy[offset + k]);\n\n                let dx = fma(DY, invRMS, -X * scale);\n\n                result[offset + k] = dx;\n        ";
+		return `
           fn DIV_CEIL(a : u32, b : u32) -> u32 {
             return ((a - 1u) / b + 1u);
           }
 
-          ${a}
+          ${t}
 
-          ${_("index")} {
+          ${c("index")} {
             // One workgroup per row (batch).
             let Length = uniforms.reduceSize;
             let BatchSize = uniforms.batchSize;
@@ -208,7 +150,7 @@ class I {
                 var sum_dygx = 0.0f;
 
                 for (var k = i32(localId.x); k < Length; k = k + ${e}) {
-                    ${t}
+                    ${n}
                 }
 
                 partials[localId.x] = vec2<f32>(sum_x2, sum_dygx);
@@ -233,52 +175,51 @@ class I {
 
                 // write dx and dGamma.
                 for (var k = i32(localId.x); k < Length; k = k + ${e}) {
-                    ${u}
+                    ${r}
                 }
           }
         `;
-  }
-}
-function P(d) {
-  const { dy: e, x: a, gamma: t } = d.inputs, u = 4;
-  y(a.shape, e.shape, "Error in RMSNormGrad dy: ");
-  const m = h(a), n = t ? h(t) : !1, f = h(e), r = m || n || f, i = !r || m ? a : x(a), s = !r || n || !t ? t : x(t), c = !r || f ? e : x(e);
-  s && y(s.shape, [i.shape[i.shape.length - 1]], "Error in RMSNormGrad gamma: ");
-  const G = d.backend, o = M(s ? [i, s, c] : [i, c], -1), k = t ? new Y(o, u, r) : new I(o, r), X = [
-    { type: "int32", data: [k.inputShape[1]] },
-    // Reduce size
-    { type: "int32", data: [k.inputShape[0]] }
-    // Batch size
-  ];
-  if (o.inSize > 1024)
-    throw new Error(`rmsNormGradGPU: inSize ${o.inSize} exceeds max of 1024`);
-  const D = G.runWebGPUProgram(
-    k,
-    s ? [i, s, c] : [i, c],
-    r ? "packedF16" : "float32",
-    X
-  );
-  r && !m && i.dispose(), r && !n && s && s.dispose(), r && !f && c.dispose();
-  const l = L().makeTensorFromTensorInfo(D);
-  if (t) {
-    const p = v(l, [0, 0], [o.batchSize, o.inSize]), S = v(
-      l,
-      [o.batchSize, 0],
-      [o.batchSize / u, o.inSize]
-    );
-    l.dispose();
-    const b = z(p, a.shape);
-    p.dispose();
-    const g = N(S, [0]);
-    return S.dispose(), [b, !r || n ? g : R(g)];
-  } else {
-    const p = z(l, a.shape);
-    return l.dispose(), [p];
-  }
-}
-const C = {
-  kernelName: "RMSNormGrad",
-  backendName: "webgpu",
-  kernelFunc: P
+	}
 };
-$(C);
+function m(t) {
+	let { dy: c, x: l, gamma: m } = t.inputs;
+	e(l.shape, c.shape, "Error in RMSNormGrad dy: ");
+	let h = r(l), g = m ? r(m) : !1, _ = r(c), v = h || g || _, y = !v || h ? l : s(l), b = !v || g || !m ? m : s(m), x = !v || _ ? c : s(c);
+	b && e(b.shape, [y.shape[y.shape.length - 1]], "Error in RMSNormGrad gamma: ");
+	let S = t.backend, C = d(b ? [
+		y,
+		b,
+		x
+	] : [y, x], -1), w = m ? new f(C, 4, v) : new p(C, v), T = [{
+		type: "int32",
+		data: [w.inputShape[1]]
+	}, {
+		type: "int32",
+		data: [w.inputShape[0]]
+	}];
+	if (C.inSize > 1024) throw Error(`rmsNormGradGPU: inSize ${C.inSize} exceeds max of 1024`);
+	let E = S.runWebGPUProgram(w, b ? [
+		y,
+		b,
+		x
+	] : [y, x], v ? "packedF16" : "float32", T);
+	v && !h && y.dispose(), v && !g && b && b.dispose(), v && !_ && x.dispose();
+	let D = n().makeTensorFromTensorInfo(E);
+	if (m) {
+		let e = u(D, [0, 0], [C.batchSize, C.inSize]), t = u(D, [C.batchSize, 0], [C.batchSize / 4, C.inSize]);
+		D.dispose();
+		let n = i(e, l.shape);
+		e.dispose();
+		let r = a(t, [0]);
+		return t.dispose(), [n, !v || g ? r : o(r)];
+	} else {
+		let e = i(D, l.shape);
+		return D.dispose(), [e];
+	}
+}
+t({
+	kernelName: "RMSNormGrad",
+	backendName: "webgpu",
+	kernelFunc: m
+});
+//#endregion