@genai-fi/nanogpt 0.20.0 → 0.20.2

package/dist/ops/webgpu/matMul16_program.js

@@ -0,0 +1,303 @@
+import { oc as e } from "../../dist-BewPQWjc.js";
+import { s as t } from "../../webgpu_program-WOyIVMlZ.js";
+//#region lib/ops/webgpu/matMul16_program.ts
+var n = class {
+	variableNames = ["A", "B"];
+	outputShape;
+	shaderKey = "MatMul16TB";
+	dispatchLayout;
+	dispatch;
+	workgroupSize = [
+		8,
+		8,
+		1
+	];
+	dimInner;
+	transposeA = !1;
+	transposeB = !0;
+	broadcastBatch = !0;
+	tileInner = 32;
+	uniforms;
+	scale = !1;
+	scaleA = !1;
+	scaleB = !1;
+	activation;
+	causalMask = !1;
+	outputComponent;
+	variableComponents;
+	outputIndexSnippet;
+	outputStrideSnippet;
+	constructor(e, t, n, r, i, a = !1, o = !1) {
+		if (this.transposeA = a, this.transposeB = o, this.variableComponents = [2, 2], this.outputComponent = 2, this.shaderKey = `MatMul16TB_${t}_${n}_${r}_${i}_${a ? "TA" : ""}${o ? "TB" : ""}`, a) {
+			if (this.outputShape = [
+				e,
+				r,
+				i / 2
+			], this.dimInner = t, t !== n) throw Error("Inner dimensions of A and B must match for MatMul16 transposeA");
+		} else if (o) {
+			if (this.outputShape = [
+				e,
+				t,
+				n / 2
+			], this.dimInner = i, i !== r) throw Error("Inner dimensions of A and B must match for MatMul16 transposeB");
+		} else if (this.outputShape = [
+			e,
+			t,
+			i / 2
+		], this.dimInner = r, r !== n) throw Error("Inner dimensions of A and B must match for MatMul16");
+		if (this.dimInner % this.tileInner !== 0) throw Error(`Inner dimension ${this.dimInner} must be multiple of ${this.tileInner}`);
+		if (this.dispatchLayout = {
+			x: [2],
+			y: [1],
+			z: [0]
+		}, this.dispatch = [
+			Math.ceil(this.outputShape[2] / (this.workgroupSize[0] * 2)),
+			Math.ceil(this.outputShape[1] / (this.workgroupSize[1] * 4)),
+			this.outputShape[0]
+		], i % 32 != 0 || r % 32 != 0) throw Error("Head size must be even for MatMul16 transposeB");
+		if (t % 32 != 0 || n % 32 != 0) throw Error("Sequence length must be multiple of 32 for MatMul16 transposeB");
+		this.outputIndexSnippet = "var idx0 = getOutputIndexFromCoords(vec3<i32>(batch, gRow, gColPacked));", this.outputStrideSnippet = "idx0 = idx0 + uniforms.outShapeStrides[1];  // Next row";
+	}
+	addUniform(e) {
+		this.uniforms ? this.uniforms += `, ${e}` : this.uniforms = e;
+	}
+	setOutputShape(t, n) {
+		let r = e(t), i = e(this.outputShape);
+		if (r !== i) throw Error(`New shape size ${r} must match current size ${i}`);
+		function a(e, t) {
+			return [`${e} / ${t}`, `${e} % ${t}`];
+		}
+		let o = this.outputShape, s = [];
+		if (t.length === o.length + 1) if (t[0] * t[1] === o[0]) s = [
+			...a("batch", t[1]),
+			"gRow",
+			"gColPacked"
+		], this.shaderKey += `_batchSplit_${t[1]}`;
+		else if (t[t.length - 2] * t[t.length - 1] === o[o.length - 1]) s = [
+			"batch",
+			"gRow",
+			...a("gColPacked", t[t.length - 1])
+		], this.shaderKey += `_colSplit_${t[t.length - 1]}`;
+		else throw Error("Unsupported output shape split");
+		else if (t.length === o.length) s = [
+			"batch",
+			"gRow",
+			"gColPacked"
+		];
+		else if (t.length === 2 && o[0] === 1) s = ["gRow", "gColPacked"], this.shaderKey += "_batchRemoved";
+		else throw Error(`Unsupported output shape rank change: ${o.length} -> ${t.length}}`);
+		let c = [];
+		if (n) {
+			if (n.length !== t.length) throw Error("Permutation length must match output rank");
+			c = n.map((e) => s[e]), this.shaderKey += `_perm_${n.join("")}`;
+		} else c = s;
+		let l = c.findIndex((e) => e === "gRow"), u = `vec${t.length}<i32>(${c.join(", ")})`;
+		this.outputIndexSnippet = `var idx0: i32 = getOutputIndexFromCoords(${u});`, this.outputStrideSnippet = `idx0 = idx0 + uniforms.outShapeStrides${l === 0 ? "" : `[${l}]`}; `, n ? this.outputShape = n.map((e) => t[e]) : this.outputShape = t;
+	}
+	useScale() {
+		this.addUniform("scale: f32"), this.scale = !0, this.shaderKey += "_scaled";
+	}
+	useScaleA() {
+		this.addUniform("scaleA: f32"), this.scaleA = !0, this.shaderKey += "_scaledA";
+	}
+	useScaleB() {
+		this.addUniform("scaleB: f32"), this.scaleB = !0, this.shaderKey += "_scaledB";
+	}
+	useActivation(e) {
+		this.activation = e, this.shaderKey += `_${e}`;
+	}
+	useCausalMask() {
+		this.causalMask = !0, this.addUniform("pastLen: i32"), this.shaderKey += "_causalMask";
+	}
+	activationSnippet() {
+		return this.activation === "gelu" ? "\n                // TODO: revisit after https://github.com/gpuweb/gpuweb/issues/4458 is resolved\n                fn tanhComplete(x: vec4<f32>) -> vec4<f32> {\n                    return vec4<f32>(\n                        select(tanh(x.x), sign(x.x), abs(x.x) > 15.0f),\n                        select(tanh(x.y), sign(x.y), abs(x.y) > 15.0f),\n                        select(tanh(x.z), sign(x.z), abs(x.z) > 15.0f),\n                        select(tanh(x.w), sign(x.w), abs(x.w) > 15.0f),\n                    );\n                }\n                fn activation(x : vec4<f32>) -> vec4<f32> {\n                    let x3 = x * x * x;\n                    var inner = fma(vec4<f32>(0.044715f), x3, x);\n                    inner = 0.7978845608028654f * inner;\n                    inner = tanhComplete(inner);\n                    inner = 0.5f * (1.0f + inner);\n                    return x * inner;\n                }\n                " : this.activation === "relu2" ? "\n                fn activation(x : vec4<f32>) -> vec4<f32> {\n                    let y = max(x, vec4<f32>(0.0));\n                    return y * y;\n                }\n                " : this.activation === "relu" ? "\n                fn activation(x : vec4<f32>) -> vec4<f32> {\n                    return max(x, vec4<f32>(0.0));\n                }\n                " : "";
+	}
+	readASnippet() {
+		let e = `
+            var col = i32(localId.x);
+            var row = i32(localId.y) * 4;
+            var packedA: vec2<i32> = A[offsetA + row * strideA + col];
+            var Arow1 = vec4<f32>(
+                unpack2x16float(u32(packedA.x)),
+                unpack2x16float(u32(packedA.y))
+            );
+            packedA = A[offsetA + (row + 1) * strideA + col];
+            var Arow2 = vec4<f32>(
+                unpack2x16float(u32(packedA.x)),
+                unpack2x16float(u32(packedA.y))
+            );
+            packedA = A[offsetA + (row + 2) * strideA + col];
+            var Arow3 = vec4<f32>(
+                unpack2x16float(u32(packedA.x)),
+                unpack2x16float(u32(packedA.y))
+            );
+            packedA = A[offsetA + (row + 3) * strideA + col];
+            var Arow4 = vec4<f32>(
+                unpack2x16float(u32(packedA.x)),
+                unpack2x16float(u32(packedA.y))
+            );
+            
+            ${this.scaleA ? "Arow1 = Arow1 * uniforms.scaleA;" : ""}
+            ${this.scaleA ? "Arow2 = Arow2 * uniforms.scaleA;" : ""}
+            ${this.scaleA ? "Arow3 = Arow3 * uniforms.scaleA;" : ""}
+            ${this.scaleA ? "Arow4 = Arow4 * uniforms.scaleA;" : ""}
+        `;
+		return this.transposeA ? `{
+                ${e}
+                mm_Asub[row][col] = Arow1;
+                mm_Asub[row + 1][col] = Arow2;
+                mm_Asub[row + 2][col] = Arow3;
+                mm_Asub[row + 3][col] = Arow4;
+        }` : `{
+                ${e}
+                
+                col = i32(localId.x) * 4;
+                row = i32(localId.y);
+                
+                mm_Asub[col][row] = vec4<f32>(Arow1.x, Arow2.x, Arow3.x, Arow4.x);
+                mm_Asub[col + 1][row] = vec4<f32>(Arow1.y, Arow2.y, Arow3.y, Arow4.y);
+                mm_Asub[col + 2][row] = vec4<f32>(Arow1.z, Arow2.z, Arow3.z, Arow4.z);
+                mm_Asub[col + 3][row] = vec4<f32>(Arow1.w, Arow2.w, Arow3.w, Arow4.w);
+        }`;
+	}
+	readBSnippet() {
+		let e = `
+            var col = i32(localId.x);
+            var row = i32(localId.y) * 4;
+            var packedB: vec2<i32> = B[offsetB + row * strideB + col];
+            var Brow1 = vec4<f32>(
+                unpack2x16float(u32(packedB.x)),
+                unpack2x16float(u32(packedB.y))
+            );
+            packedB = B[offsetB + (row + 1) * strideB + col];
+            var Brow2 = vec4<f32>(
+                unpack2x16float(u32(packedB.x)),
+                unpack2x16float(u32(packedB.y))
+            );
+            packedB = B[offsetB + (row + 2) * strideB + col];
+            var Brow3 = vec4<f32>(
+                unpack2x16float(u32(packedB.x)),
+                unpack2x16float(u32(packedB.y))
+            );
+            packedB = B[offsetB + (row + 3) * strideB + col];
+            var Brow4 = vec4<f32>(
+                unpack2x16float(u32(packedB.x)),
+                unpack2x16float(u32(packedB.y))
+            );
+            
+            ${this.scaleB ? "Brow1 = Brow1 * uniforms.scaleB;" : ""}
+            ${this.scaleB ? "Brow2 = Brow2 * uniforms.scaleB;" : ""}
+            ${this.scaleB ? "Brow3 = Brow3 * uniforms.scaleB;" : ""}
+            ${this.scaleB ? "Brow4 = Brow4 * uniforms.scaleB;" : ""}
+        `;
+		return this.transposeB ? `{
+                ${e}
+                
+                col = i32(localId.x) * 4;
+                row = i32(localId.y);
+                
+                mm_Bsub[col][row] = vec4<f32>(Brow1.x, Brow2.x, Brow3.x, Brow4.x);
+                mm_Bsub[col + 1][row] = vec4<f32>(Brow1.y, Brow2.y, Brow3.y, Brow4.y);
+                mm_Bsub[col + 2][row] = vec4<f32>(Brow1.z, Brow2.z, Brow3.z, Brow4.z);
+                mm_Bsub[col + 3][row] = vec4<f32>(Brow1.w, Brow2.w, Brow3.w, Brow4.w);
+            }` : `{
+                ${e}
+                mm_Bsub[row][col] = Brow1;
+                mm_Bsub[row + 1][col] = Brow2;
+                mm_Bsub[row + 2][col] = Brow3;
+                mm_Bsub[row + 3][col] = Brow4;
+            }`;
+	}
+	baseIndexSnippets() {
+		let e = "";
+		e = this.transposeB ? "let baseB = getIndexFromCoords3D(vec3<i32>(batchB, globalColStart, 0), vec3<i32>(uniforms.bShape.x, uniforms.bShape.y, strideB));" : "let baseB = getIndexFromCoords3D(vec3<i32>(batchB, 0, globalColStart / 4), vec3<i32>(uniforms.bShape.x, uniforms.bShape.y, strideB));";
+		let t = "";
+		return t = this.transposeA ? "let baseA = getIndexFromCoords3D(vec3<i32>(batchA, 0, globalRowStart / 4), vec3<i32>(uniforms.aShape.x, uniforms.aShape.y, strideA));" : "let baseA = getIndexFromCoords3D(vec3<i32>(batchA, globalRowStart, 0), vec3<i32>(uniforms.aShape.x, uniforms.aShape.y, strideA));", `
+            
+            let strideA = uniforms.aShape.z / 2;
+            let strideB = uniforms.bShape.z / 2;
+        
+            ${t}
+            ${e}
+        `;
+	}
+	offsetSnippets() {
+		let e = "";
+		e = this.transposeA ? "let offsetA = baseA + kStart * strideA;" : "let offsetA = baseA + kStart / 4;";
+		let t = "";
+		return t = this.transposeB ? "let offsetB = baseB + kStart / 4;" : "let offsetB = baseB + kStart * strideB;", `
+            ${e}
+            ${t}
+        `;
+	}
+	getUserCode() {
+		let e = this.transposeA, n = this.tileInner, r = this.workgroupSize[1] * 4, i = this.workgroupSize[0] * 4, a = e ? r : n, o = e ? n : r, s = this.dimInner, c = Math.ceil(s / n);
+		return `
+            var<workgroup> mm_Asub : array<array<vec4<f32>, ${a / 4 + +!this.transposeA}>, ${o}>;
+            var<workgroup> mm_Bsub : array<array<vec4<f32>, ${i / 4 + +!!this.transposeB}>, ${n}>;
+
+            ${this.activation ? this.activationSnippet() : ""}
+
+            ${t()} {
+                let batch = i32(globalId.z);
+                let batchA = ${this.broadcastBatch ? "batch % uniforms.aShape[0]" : "batch"};
+                let batchB = ${this.broadcastBatch ? "batch % uniforms.bShape[0]" : "batch"};
+                var kStart = 0;
+                let localRow = i32(localId.y);
+                let localCol = i32(localId.x);
+                let globalRowStart = i32(workgroupId.y) * ${r};
+                let globalColStart = i32(workgroupId.x) * ${i};
+
+                // 4 rows x 4 cols accumulator
+                // acc[i] holds row i (4 cols)
+                var acc = array<vec4<f32>, 4>(
+                    vec4<f32>(0.0), vec4<f32>(0.0), vec4<f32>(0.0), vec4<f32>(0.0)
+                );
+
+                ${this.baseIndexSnippets()}
+
+                for (var t = 0; t < ${c}; t++) {
+                    ${this.offsetSnippets()}
+
+                    ${this.readASnippet()}
+                    ${this.readBSnippet()}
+
+                    kStart = kStart + ${n};
+                    workgroupBarrier();
+
+                    for (var k = 0; k < ${n}; k++) {
+                        // Load 4 columns of B as a vec4
+                        let bVec = mm_Bsub[k][localCol];
+                        let aVec = mm_Asub[k][localRow];
+
+                        // Compute 4 rows
+                        for (var r = 0; r < 4; r = r + 1) {
+                            acc[r] = fma(vec4<f32>(aVec[r]), bVec, acc[r]);
+                        }
+                    }
+                    workgroupBarrier();
+                }
+
+                // Write out 4 rows x 2 packed cols (4 unpacked cols)
+                let gRow = globalRowStart + localRow * 4;
+                let gColPacked = i32(workgroupId.x) * ${this.workgroupSize[0] * 2} + localCol * 2;
+                
+                ${this.outputIndexSnippet}
+                for (var i = 0; i < 4; i = i + 1) {
+                    ${this.causalMask ? "\n                    // Causal Masking: mask if col > row + pastLen\n                    let r = gRow + i;\n                    let cBase = gColPacked * 2;\n                    let cVec = vec4<i32>(cBase, cBase + 1, cBase + 2, cBase + 3);\n                    let mask = cVec > vec4<i32>(r + uniforms.pastLen);\n                    acc[i] = select(acc[i], vec4<f32>(-uniforms.INFINITY), mask);\n                    " : ""}
+
+                    ${this.activation ? "acc[i] = activation(acc[i]);" : ""}
+                    ${this.scale ? "acc[i] = acc[i] * uniforms.scale;" : ""}
+                    result[idx0 / 2] = vec2<i32>(
+                        i32(pack2x16float(acc[i].xy)),
+                        i32(pack2x16float(acc[i].zw))
+                    );
+
+                    ${this.outputStrideSnippet}
+                }
+            }
+        `;
+	}
+};
+//#endregion
+export { n as default };

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

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

package/dist/ops/webgpu/mul16.js

@@ -0,0 +1,14 @@
+import { Ii as e } from "../../dist-BewPQWjc.js";
+import { t } from "../../binary_op_util-CrYk9LXL.js";
+import { BinaryOpProgram as n, BinaryOpScalarProgram as r } from "./utils/binary_op.js";
+//#region lib/ops/webgpu/mul16.ts
+function i(e) {
+	let { a: i, b: a } = e.inputs, o = e.backend, s = a.shape.length === 0 ? new r(t.MUL, i.shape) : new n(t.MUL, i.shape, a.shape);
+	return o.runWebGPUProgram(s, [i, a], "packedF16");
+}
+e({
+	kernelName: "Mul16",
+	backendName: "webgpu",
+	kernelFunc: i
+});
+//#endregion

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

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

package/dist/ops/webgpu/norm2.js

@@ -0,0 +1,46 @@
+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        ";
+	}
+};
+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.d.ts

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

package/dist/ops/webgpu/normRMS.js

@@ -0,0 +1,26 @@
+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;
+}
+t({
+	kernelName: "RMSNorm",
+	backendName: "webgpu",
+	kernelFunc: l
+}), t({
+	kernelName: "RMSNormNoGamma",
+	backendName: "webgpu",
+	kernelFunc: l
+});
+//#endregion

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

@@ -0,0 +1,10 @@
+import { backend_util } from '@tensorflow/tfjs-core';
+import { ReduceProgram } from './utils/reductions';
+import { DeviceInformation } from './utils/deviceInfo';
+export default class RMSProgram16 extends ReduceProgram {
+    private hasGamma;
+    constructor(deviceInfo: DeviceInformation, reduceInfo: backend_util.ReduceInfo, hasGamma?: boolean);
+    getPreprocessSnippet(): string;
+    getPostprocessSnippet(): string;
+    getWriteSnippet(): string;
+}

package/dist/ops/webgpu/normRMS16_program.js

@@ -0,0 +1,28 @@
+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));
+        `;
+	}
+};
+//#endregion
+export { t as default };

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

@@ -0,0 +1,10 @@
+import { backend_util } from '@tensorflow/tfjs-core';
+import { ReduceProgram } from './utils/reductions';
+import { DeviceInformation } from './utils/deviceInfo';
+export default class RMSProgram32 extends ReduceProgram {
+    private hasGamma;
+    constructor(deviceInfo: DeviceInformation, reduceInfo: backend_util.ReduceInfo, hasGamma?: boolean);
+    protected getPreprocessSnippet(): string;
+    protected getPostprocessSnippet(): string;
+    protected getWriteSnippet(): string;
+}

package/dist/ops/webgpu/normRMS32_program.js

@@ -0,0 +1,28 @@
+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);
+        `;
+	}
+};
+//#endregion
+export { t as default };

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

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