@genai-fi/nanogpt 0.20.0 → 0.20.2

package/dist/ops/webgl/normRMS.js

@@ -0,0 +1,114 @@
+import { Ii as e, ii as t, mn as n } from "../../dist-BewPQWjc.js";
+//#region lib/ops/webgl/normRMS.ts
+var r = class {
+	variableNames = ["x", "meanSquare"];
+	outputShape;
+	userCode;
+	constructor(e, t, n, r = !0) {
+		r && this.variableNames.push("gamma"), this.outputShape = [
+			e,
+			t,
+			n
+		], this.userCode = `
+        void main() {
+            ivec3 coords = getOutputCoords();
+            float x = getXAtOutCoords();
+            float meanSquare = getMeanSquare(coords.x, coords.y, 0);
+            ${r ? "float gamma = getGammaAtOutCoords();" : ""}
+            float invRms = inversesqrt(meanSquare + 1e-8);
+            float normalized = x * invRms;
+            float outVal = normalized ${r ? " * gamma" : ""};
+            setOutput(outVal);
+        }
+        `;
+	}
+};
+function i(e) {
+	let { x: t, gamma: n } = e.inputs, i = e.backend, a = t.shape[0], o = t.shape[1], s = t.shape[2], c = t.square().mean(-1, !0), l = new r(a, o, s, n !== void 0);
+	return i.runWebGLProgram(l, n ? [
+		t,
+		c,
+		n
+	] : [t, c], "float32");
+}
+var a = {
+	kernelName: "RMSNorm",
+	backendName: "webgl",
+	kernelFunc: i
+}, o = {
+	kernelName: "RMSNormNoGamma",
+	backendName: "webgl",
+	kernelFunc: i
+};
+e(a), e(o);
+var s = class {
+	variableNames = [
+		"x",
+		"meanSquare",
+		"dyGamma",
+		"dyXMean"
+	];
+	outputShape;
+	userCode;
+	constructor(e, t, n) {
+		this.outputShape = [
+			e,
+			t,
+			n
+		], this.userCode = `
+        void main() {
+            ivec3 coords = getOutputCoords();
+            float x = getXAtOutCoords();
+            float meanSquare = getMeanSquare(coords.x, coords.y, 0) + 1e-8;
+            float dyGamma = getDyGammaAtOutCoords();
+            float dyXMean = getDyXMean(coords.x, coords.y, 0) / ${n}.0;
+            float invRms = inversesqrt(meanSquare);
+            float dx = dyGamma * invRms - x * dyXMean * invRms / meanSquare;
+            setOutput(dx);
+        }
+        `;
+	}
+}, c = class {
+	variableNames = [
+		"x",
+		"meanSquare",
+		"dy"
+	];
+	outputShape;
+	userCode;
+	constructor(e, t, n) {
+		this.outputShape = [
+			e,
+			t,
+			n
+		], this.userCode = "\n        void main() {\n            ivec3 coords = getOutputCoords();\n            float x = getXAtOutCoords();\n            float meanSquare = getMeanSquare(coords.x, coords.y, 0) + 1e-8;\n            float dy = getDyAtOutCoords();\n            float invRms = inversesqrt(meanSquare);\n            float dGamma = dy * (x * invRms);\n            setOutput(dGamma);\n        }\n        ";
+	}
+};
+function l(e) {
+	let { dy: r, x: i, gamma: a } = e.inputs, o = e.backend, l = i.shape[0], u = i.shape[1], d = i.shape[2], f = a ? r.mul(a) : r, p = f.mul(i), m = p.sum(-1, !0);
+	p.dispose();
+	let h = i.square(), g = h.mean(-1, !0);
+	h.dispose();
+	let _ = new s(l, u, d), v = o.runWebGLProgram(_, [
+		i,
+		g,
+		f,
+		m
+	], "float32");
+	if (a && f.dispose(), m.dispose(), a) {
+		let e = new c(l, u, d), a = o.runWebGLProgram(e, [
+			i,
+			g,
+			r
+		], "float32");
+		g.dispose();
+		let s = n(t().makeTensorFromTensorInfo(a), [0, 1]);
+		return o.disposeIntermediateTensorInfo(a), [v, s];
+	} else return g.dispose(), [v];
+}
+e({
+	kernelName: "RMSNormGrad",
+	backendName: "webgl",
+	kernelFunc: l
+});
+//#endregion

package/dist/ops/webgl/qkv.d.ts

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

package/dist/ops/webgl/qkv.js

@@ -0,0 +1,54 @@
+import { Ii as e } from "../../dist-BewPQWjc.js";
+//#region lib/ops/webgl/qkv.ts
+var t = class {
+	variableNames = ["x", "kernel"];
+	outputShape;
+	userCode;
+	customUniforms = [{
+		name: "mode",
+		type: "int"
+	}];
+	constructor(e, t, n, r) {
+		let i = r / t;
+		this.outputShape = [
+			e,
+			t,
+			n,
+			i
+		], this.userCode = `
+        void main() {
+            ivec4 coords = getOutputCoords(); // [b, h, t, d]
+            int b = coords.x;
+            int h = coords.y;
+            int t = coords.z;
+            int d = coords.w;
+
+            // Compute output channel index in fused kernel
+            int out_offset = mode * ${t} * ${i} + h * ${i} + d;
+
+            float sum = 0.0;
+            for (int c = 0; c < ${r}; ++c) {
+                float xval = getX(b, t, c); // fetch from x
+                float kval = getKernel(c, out_offset); // fetch from kernel
+                sum += xval * kval;
+            }
+
+            setOutput(sum);
+        }
+        `;
+	}
+};
+function n(e) {
+	let { x: n, kernel: r } = e.inputs, { heads: i } = e.attrs, a = e.backend, o = n.shape[0], s = n.shape[1], c = n.shape[2], l = new t(o, i, s, c);
+	return [
+		a.runWebGLProgram(l, [n, r], "float32", [[0]]),
+		a.runWebGLProgram(l, [n, r], "float32", [[1]]),
+		a.runWebGLProgram(l, [n, r], "float32", [[2]])
+	];
+}
+e({
+	kernelName: "QKV",
+	backendName: "webgl",
+	kernelFunc: n
+});
+//#endregion

package/dist/ops/webgl/rope.d.ts

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

package/dist/ops/webgl/rope.js

@@ -0,0 +1,72 @@
+import { Ii as e } from "../../dist-BewPQWjc.js";
+//#region lib/ops/webgl/rope.ts
+var t = class {
+	variableNames = [
+		"x",
+		"sin",
+		"cos"
+	];
+	outputShape;
+	userCode;
+	customUniforms = [{
+		name: "pastLen",
+		type: "int"
+	}];
+	constructor(e, t, n, r) {
+		this.outputShape = [
+			e,
+			t,
+			n,
+			r
+		], this.userCode = `
+        void main() {
+            ivec4 coords = getOutputCoords(); // [b, h, t, d]
+            int b = coords.x;
+            int h = coords.y;
+            int t = coords.z;
+            int d = coords.w;
+
+            int rotaryDim = ${r};
+
+            float outVal = 0.0;
+
+            if (d < rotaryDim) {
+                int pairIdx = d / 2;
+                float cos = getCos(t + pastLen, pairIdx, 0);
+                float sin = getSin(t + pastLen, pairIdx, 0);
+
+                if (d % 2 == 0) {
+                    // even index
+                    float even = getX(b, h, t, d);
+                    float odd = getX(b, h, t, d + 1);
+                    outVal = even * cos - odd * sin;
+                } else {
+                    // odd index
+                    float even = getX(b, h, t, d - 1);
+                    float odd = getX(b, h, t, d);
+                    outVal = even * sin + odd * cos;
+                }
+            } else {
+                // pass through for non-rotary dims
+                outVal = getX(b, h, t, d);
+            }
+
+            setOutput(outVal);
+        }
+        `;
+	}
+};
+function n(e) {
+	let { x: n } = e.inputs, { pastLen: r, ropeCache: i, negSin: a } = e.attrs, o = a ? i.getNegSin() : i.getSin(), s = i.getCos(), c = e.backend, l = n.shape[0], u = n.shape[1], d = n.shape[2], f = n.shape[3], p = new t(l, u, d, f);
+	return c.runWebGLProgram(p, [
+		n,
+		o,
+		s
+	], "float32", [[r]]);
+}
+e({
+	kernelName: "Rope",
+	backendName: "webgl",
+	kernelFunc: n
+});
+//#endregion

package/dist/ops/webgl/scatterSub.d.ts

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

package/dist/ops/webgl/scatterSub.js

@@ -0,0 +1,28 @@
+import { Ii as e } from "../../dist-BewPQWjc.js";
+//#region lib/ops/webgl/scatterSub.ts
+var t = class {
+	variableNames = [
+		"labels",
+		"softmaxProbs",
+		"dy"
+	];
+	outputShape;
+	userCode;
+	constructor(e, t) {
+		this.outputShape = [e, t], this.userCode = "\n            void main() {\n                ivec2 coords = getOutputCoords();\n                int index = int(getLabels(coords.x));\n                float prob = getSoftmaxProbsAtOutCoords();\n                float dy = getDy(coords.x);\n                setOutput((index == coords.y ? prob - 1.0 : prob) * dy);\n            }\n            ";
+	}
+};
+function n(e) {
+	let { logits: n, labels: r, dy: i } = e.inputs, a = e.backend, o = r.shape[0], s = n.shape[1], c = new t(o, s);
+	return a.runWebGLProgram(c, [
+		r,
+		n,
+		i
+	], "float32");
+}
+e({
+	kernelName: "EfficientScatterSub",
+	backendName: "webgl",
+	kernelFunc: n
+});
+//#endregion

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

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

package/dist/ops/webgpu/adamAdjust.js

@@ -0,0 +1,77 @@
+import { Fs as e, Ii as t } from "../../dist-BewPQWjc.js";
+import { s as n } from "../../webgpu_program-WOyIVMlZ.js";
+import { c as r, i } from "../../webgpu_util-B_F3SShA.js";
+//#region lib/ops/webgpu/adamAdjust.ts
+var a = class {
+	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];
+	useWeightDecay;
+	constructor(e, t) {
+		this.outputShape = e, this.dispatchLayout = r(this.outputShape), this.dispatch = i(this.dispatchLayout, this.outputShape, this.workgroupSize), this.useWeightDecay = t, t && (this.uniforms += ", weightDecay: f32");
+	}
+	getUserCode() {
+		return `
+        ${n("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 / fma(uniforms.epsilon, invSqrt, 1.0);
+                var adjustedValue = fma(-uniforms.learningRate * m1Hat, invDenom, value);
+
+                ${this.useWeightDecay ? "adjustedValue = adjustedValue - uniforms.learningRate * uniforms.weightDecay * value;" : ""}
+
+                setOutputAtIndex(index, adjustedValue);
+            }
+        }
+    `;
+	}
+};
+function o(t) {
+	let { moments: n, value: r } = t.inputs, { beta1: i, beta2: o, learningRate: s, epsilon: c, weightDecay: l } = t.attrs, u = t.backend;
+	e(n.shape, [...r.shape, 2], "Error in AdamAdjust: ");
+	let d = new a(r.shape, l > 0), f = [
+		{
+			type: "float32",
+			data: [1 / i]
+		},
+		{
+			type: "float32",
+			data: [1 / o]
+		},
+		{
+			type: "float32",
+			data: [s]
+		},
+		{
+			type: "float32",
+			data: [c]
+		}
+	];
+	return l > 0 && f.push({
+		type: "float32",
+		data: [l]
+	}), u.runWebGPUProgram(d, [n, r], "float32", f);
+}
+t({
+	kernelName: "AdamAdjust",
+	backendName: "webgpu",
+	kernelFunc: o
+});
+//#endregion

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

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

package/dist/ops/webgpu/adamMoments.js

@@ -0,0 +1,76 @@
+import { Fs as e, Ii as t } from "../../dist-BewPQWjc.js";
+import { s as n } from "../../webgpu_program-WOyIVMlZ.js";
+import { c as r, i } from "../../webgpu_util-B_F3SShA.js";
+//#region lib/ops/webgpu/adamMoments.ts
+var a = class {
+	variableNames = [
+		"moments",
+		"gradient",
+		"scaling"
+	];
+	outputShape;
+	shaderKey = "AdamMoments";
+	dispatchLayout;
+	dispatch;
+	workgroupSize = [
+		64,
+		1,
+		1
+	];
+	size = !0;
+	uniforms = "beta1: f32, beta2: f32";
+	outputComponent = 2;
+	variableComponents = [
+		2,
+		1,
+		1
+	];
+	constructor(e) {
+		this.outputShape = e, this.dispatchLayout = r(this.outputShape.slice(0, -1)), this.dispatch = i(this.dispatchLayout, this.outputShape.slice(0, -1), this.workgroupSize, [
+			1,
+			1,
+			1
+		]);
+	}
+	getUserCode() {
+		return `
+        ${n("index")} {
+            if (index < uniforms.size) {
+                let m: vec2<f32> = moments[index];
+
+                // Loss and clip scaling.
+                let g: f32 = gradient[index] * scaling[0];
+
+                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 o(t) {
+	let { moments: n, gradient: r, scaling: i } = t.inputs, { beta1: o, beta2: s } = t.attrs, c = t.backend;
+	if (r.dtype !== "float32") throw Error(`Gradient must be float32, but got ${r.dtype}`);
+	if (e(n.shape, [...r.shape, 2], "Error in AdamMoments: "), o < 0 || o >= 1) throw Error(`Invalid beta1 value: ${o}. Must be in the range [0, 1).`);
+	if (s < 0 || s >= 1) throw Error(`Invalid beta2 value: ${s}. Must be in the range [0, 1).`);
+	let l = new a(n.shape), u = [{
+		type: "float32",
+		data: [o]
+	}, {
+		type: "float32",
+		data: [s]
+	}];
+	return c.runWebGPUProgram(l, [
+		n,
+		r,
+		i
+	], "float32", u);
+}
+t({
+	kernelName: "AdamMoments",
+	backendName: "webgpu",
+	kernelFunc: o
+});
+//#endregion

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

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

package/dist/ops/webgpu/add16.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 } from "./utils/binary_op.js";
+//#region lib/ops/webgpu/add16.ts
+function r(e) {
+	let { a: r, b: i } = e.inputs, a = e.backend, o = new n(t.ADD, r.shape, i.shape);
+	return a.runWebGPUProgram(o, [r, i], "packedF16");
+}
+e({
+	kernelName: "Add16",
+	backendName: "webgpu",
+	kernelFunc: r
+});
+//#endregion

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

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

package/dist/ops/webgpu/appendCache.js

@@ -0,0 +1,130 @@
+import { Fs as e, Ii as t } from "../../dist-BewPQWjc.js";
+import { isPackedTensor as n } from "../../utilities/packed.js";
+import { s as r } from "../../webgpu_program-WOyIVMlZ.js";
+import { c as i, i as a } from "../../webgpu_util-B_F3SShA.js";
+//#region lib/ops/webgpu/appendCache.ts
+var o = class {
+	variableNames = ["cache", "item"];
+	outputShape;
+	shaderKey = "AppendCache";
+	dispatchLayout;
+	dispatch;
+	workgroupSize = [
+		64,
+		1,
+		1
+	];
+	size = !0;
+	uniforms = "cacheT: i32";
+	constructor(e, t, n, r, o) {
+		let s = Math.min(n + 1, o);
+		this.shaderKey = `AppendCache_${s}`, this.outputShape = [
+			e,
+			t,
+			s,
+			r
+		], this.dispatchLayout = i(this.outputShape), this.dispatch = a(this.dispatchLayout, this.outputShape, this.workgroupSize);
+	}
+	getUserCode() {
+		let e = this.outputShape[2];
+		return `
+        ${r("index")} {
+            if (index < uniforms.size) {
+                let coords = getCoordsFromIndex(index); // [b, h, t, d]
+                let b = coords[0];
+                let h = coords[1];
+                let t = coords[2];
+                let d = coords[3];
+
+                let itemT = 1;
+                let maxSize = ${e};
+                let totalT = uniforms.cacheT + itemT;
+                let start = select(0, 1, totalT >= maxSize);
+
+                let srcT = t + start;
+                var val = 0.0;
+                if (srcT < uniforms.cacheT) {
+                    val = getCache(b, h, srcT, d);
+                }
+                if (srcT == uniforms.cacheT) {
+                    val = getItem(b, h, 0, d);
+                }
+
+                setOutputAtIndex(index, val);
+            }
+        }
+        `;
+	}
+}, s = class {
+	variableNames = ["cache", "item"];
+	outputShape;
+	shaderKey = "AppendCache";
+	dispatchLayout;
+	dispatch;
+	workgroupSize = [
+		64,
+		1,
+		1
+	];
+	size = !0;
+	uniforms = "cacheT: i32";
+	constructor(e, t, n, r, o) {
+		let s = Math.min(n + 1, o);
+		this.shaderKey = `AppendCache_${s}`, this.outputShape = [
+			e,
+			t,
+			s,
+			r
+		], this.dispatchLayout = i(this.outputShape), this.dispatch = a(this.dispatchLayout, this.outputShape, this.workgroupSize);
+	}
+	getUserCode() {
+		let e = this.outputShape[2];
+		return `
+        ${r("index")} {
+            if (index < uniforms.size) {
+                let coords = getCoordsFromIndex(index); // [b, h, t, d]
+                let b = coords[0];
+                let h = coords[1];
+                let t = coords[2];
+                let d = coords[3];
+
+                let itemT = 1;
+                let maxSize = ${e};
+                let totalT = uniforms.cacheT + itemT;
+                let start = select(0, 1, totalT >= maxSize);
+
+                let srcT = t + start;
+                var val: i32 = 0i;
+                if (srcT < uniforms.cacheT) {
+                    val = cache[getIndexFromCoords4D(vec4<i32>(b, h, srcT, d), uniforms.cacheShape)];
+                }
+                if (srcT == uniforms.cacheT) {
+                    val = item[getIndexFromCoords4D(vec4<i32>(b, h, 0, d), uniforms.itemShape)];
+                }
+
+                result[index] = val;
+            }
+        }
+        `;
+	}
+};
+function c(t) {
+	let { cache: r, item: i } = t.inputs, { maxSize: a, pastLen: c } = t.attrs, l = t.backend, u = n(r), d = r.shape[0], f = r.shape[2], p = r.shape[1];
+	if (e(i.shape, [
+		d,
+		p,
+		1,
+		i.shape[3]
+	], "Error in AppendCache: "), c < 0 || c > a) throw Error(`Invalid pastLen value: ${c}. Must be in the range [0, ${a}].`);
+	let m = u ? new s(d, p, f, i.shape[3], a) : new o(d, p, f, i.shape[3], a), h = [{
+		type: "int32",
+		data: [c]
+	}], g = u ? "packedF16" : r.dtype;
+	return l.runWebGPUProgram(m, [r, i], g, h);
+}
+t({
+	kernelName: "AppendCache",
+	backendName: "webgpu",
+	kernelFunc: c
+});
+//#endregion

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

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

package/dist/ops/webgpu/attentionMask.js

@@ -0,0 +1,42 @@
+import { Fs as e, Ii as t } from "../../dist-BewPQWjc.js";
+import { isPackedTensor as n } from "../../utilities/packed.js";
+import { t as r } from "../../matMul16-BNfZSnNM.js";
+import i from "./attentionMask32_program.js";
+//#region lib/ops/webgpu/attentionMask.ts
+function a(t) {
+	let { q: a, k: o } = t.inputs, { divisor: s, pastLen: c } = t.attrs, l = t.backend;
+	if (n(a) && n(o)) return r(a, o, !1, !0, {
+		causalMask: !0,
+		pastLen: c,
+		scale: s
+	});
+	let u = a.shape[0], d = a.shape[2], f = o.shape[2], p = a.shape[1], m = a.shape[3];
+	if (e(o.shape, [
+		u,
+		p,
+		f,
+		m
+	], "Error in AttentionMask: "), s === 0) throw Error("Divisor must be non-zero in AttentionMask");
+	if (c < 0) throw Error("pastLen must be non-negative in AttentionMask");
+	let h = new i(u, p, d, f, m), g = [
+		{
+			type: "float32",
+			data: [s]
+		},
+		{
+			type: "int32",
+			data: [c]
+		},
+		{
+			type: "float32",
+			data: [-Infinity]
+		}
+	], _ = a.dtype;
+	return l.runWebGPUProgram(h, [a, o], _, g);
+}
+t({
+	kernelName: "AttentionMask",
+	backendName: "webgpu",
+	kernelFunc: a
+});
+//#endregion

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

@@ -0,0 +1,19 @@
+import { WebGPUProgram } from '@tensorflow/tfjs-backend-webgpu';
+export default class AttentionMaskProgram32 implements WebGPUProgram {
+    variableNames: string[];
+    outputShape: number[];
+    shaderKey: string;
+    dispatchLayout: {
+        x: number[];
+    };
+    dispatch: [number, number, number];
+    uniforms: string;
+    workgroupSize: [number, number, number];
+    size: boolean;
+    hs: number;
+    nh: number;
+    T1: number;
+    T2: number;
+    constructor(batch: number, nh: number, T1: number, T2: number, hs: number);
+    getUserCode(): string;
+}