@genai-fi/nanogpt 1.4.1 → 1.5.0

package/dist/models/Llama.js

@@ -0,0 +1,89 @@
+import { di as e, oi as t } from "../dist-Da20xy8E.js";
+import { packingSupported as n } from "../utilities/packed.js";
+import { r, t as i } from "../pack16-BhuXNUS7.js";
+import a from "../layers/RMSNorm.js";
+import o from "../layers/TransformerBlock.js";
+import s from "../layers/TiedEmbedding.js";
+import c from "../layers/RoPECache.js";
+import l from "./model.js";
+import u from "../layers/HeadLayer.js";
+//#region lib/models/Llama.ts
+var d = {
+	modelType: "llama",
+	vocabSize: 2e3,
+	blockSize: 128,
+	nLayer: 6,
+	nHead: 4,
+	nEmbed: 256,
+	mlpFactor: 4
+}, f = class extends l {
+	wte;
+	head;
+	blocks;
+	lnF;
+	ropeCache;
+	constructor(e = {}) {
+		super({
+			...d,
+			...e
+		});
+		let t = {
+			activation: "swiglu",
+			hiddenSize: Math.max(256, Math.ceil(this.config.nEmbed * (8 / 3) / 256) * 256),
+			useGamma: !0,
+			useQKNorm: !1,
+			ropeInterleaved: !1
+		};
+		this.wte = new s(this.config, "token_embedding", this), this.ropeCache = new c(this.config), this.blocks = [];
+		for (let e = 0; e < this.config.nLayer; e++) this.blocks.push(new o(e, this.config, t, this));
+		this.lnF = new a(this.config, t, "final_rms_norm", this), this.head = new u(this.config, "lm_head", this);
+	}
+	getClassName() {
+		return "llama";
+	}
+	inputPhase(t) {
+		return e(() => this.wte.embed(t));
+	}
+	forward(a, o) {
+		return this.validateInput(o), a.ropeCache = this.ropeCache, a.outputEmbeddings && (a.embeddings = []), e(() => {
+			this.startMemory();
+			let e = this.inputPhase(o);
+			if (a.cache && a.cache.length !== this.blocks.length) throw console.error("Cache", a.cache), Error(`Cache length ${a.cache.length} does not match number of blocks ${this.blocks.length}`);
+			let s = a.mixedPrecision === !0 && n(), c = s ? i(e) : e;
+			s && e !== c && e.dispose();
+			for (let e = 0; e < this.blocks.length; e++) {
+				if (a.layerDrop && Math.random() < a.layerDrop * (e / this.blocks.length)) continue;
+				let n = this.blocks[e], r = Math.random() * 1e9, i = {
+					...a,
+					seed: r,
+					pastKV: a.cache ? a.cache[e] : void 0,
+					mixedPrecision: s,
+					epsilon: 1e-6
+				}, o = a.checkpointing && a.training ? n.callCheckpoint(i, c) : n.call(i, c);
+				a.outputEmbeddings ? (t(c), a.embeddings.push({
+					name: `block_output_${e}`,
+					tensor: c
+				})) : c.dispose(), c = o;
+			}
+			if (e = this.lnF.call({
+				...a,
+				mixedPrecision: s
+			}, c), c.dispose(), a.skipLogits) return this.endMemory("Forward"), e;
+			let l = this.head.project(e);
+			a.outputEmbeddings ? (t(e), a.embeddings.push({
+				name: "final_norm_output",
+				tensor: e
+			})) : e.dispose();
+			let u = s ? r(l) : l;
+			return s && l !== u && l.dispose(), u;
+		});
+	}
+	project(t) {
+		return e(() => this.head.project(t));
+	}
+	dispose() {
+		this.weightStore.dispose(), this.wte.dispose(), this.head.dispose(), this.blocks.forEach((e) => e.dispose()), this.lnF.dispose();
+	}
+};
+//#endregion
+export { f as default };

package/dist/models/Qwen3.d.ts

@@ -0,0 +1,16 @@
+import { GPTConfigQwen3 } from './config';
+import { Tensor } from '@tensorflow/tfjs-core';
+import { default as Model, ModelForwardAttributes } from './model';
+export default class Qwen3 extends Model<ModelForwardAttributes, GPTConfigQwen3> {
+    private wte;
+    private head;
+    private blocks;
+    private lnF;
+    private ropeCache?;
+    constructor(config?: Partial<GPTConfigQwen3>);
+    getClassName(): string;
+    private inputPhase;
+    forward(attrs: ModelForwardAttributes, idx: Tensor): Tensor;
+    project(embeddings: Tensor): Tensor;
+    dispose(): void;
+}

package/dist/models/Qwen3.js

@@ -0,0 +1,89 @@
+import { di as e, oi as t } from "../dist-Da20xy8E.js";
+import { packingSupported as n } from "../utilities/packed.js";
+import { r, t as i } from "../pack16-BhuXNUS7.js";
+import a from "../layers/RMSNorm.js";
+import o from "../layers/TransformerBlock.js";
+import s from "../layers/TiedEmbedding.js";
+import c from "../layers/RoPECache.js";
+import l from "./model.js";
+import u from "../layers/HeadLayer.js";
+//#region lib/models/Qwen3.ts
+var d = {
+	modelType: "qwen3",
+	vocabSize: 2e3,
+	blockSize: 128,
+	nLayer: 6,
+	nHead: 4,
+	nEmbed: 256,
+	mlpFactor: 4
+}, f = class extends l {
+	wte;
+	head;
+	blocks;
+	lnF;
+	ropeCache;
+	constructor(e = {}) {
+		super({
+			...d,
+			...e
+		});
+		let t = {
+			activation: "swiglu",
+			hiddenSize: this.config.mlpSize ?? Math.max(256, Math.ceil(this.config.nEmbed * (8 / 3) / 256) * 256),
+			useGamma: !0,
+			useQKNorm: !0,
+			ropeInterleaved: !1,
+			epsilon: this.config.rmsNormEpsilon ?? 1e-6
+		};
+		this.wte = new s(this.config, "token_embedding", this), this.ropeCache = new c(this.config), this.blocks = [];
+		for (let e = 0; e < this.config.nLayer; e++) this.blocks.push(new o(e, this.config, t, this));
+		this.lnF = new a(this.config, t, "final_rms_norm", this), this.config.tieEmbeddings ?? !0 ? this.head = new u(this.config, "lm_head", this) : this.head = null;
+	}
+	getClassName() {
+		return "qwen3";
+	}
+	inputPhase(t) {
+		return e(() => this.wte.embed(t));
+	}
+	forward(a, o) {
+		return this.validateInput(o), a.ropeCache = this.ropeCache, a.outputEmbeddings && (a.embeddings = []), e(() => {
+			this.startMemory();
+			let e = this.inputPhase(o);
+			if (a.cache && a.cache.length !== this.blocks.length) throw console.error("Cache", a.cache), Error(`Cache length ${a.cache.length} does not match number of blocks ${this.blocks.length}`);
+			let s = a.mixedPrecision === !0 && n(), c = s ? i(e) : e;
+			s && e !== c && e.dispose();
+			for (let e = 0; e < this.blocks.length; e++) {
+				if (a.layerDrop && Math.random() < a.layerDrop * (e / this.blocks.length)) continue;
+				let n = this.blocks[e], r = Math.random() * 1e9, i = {
+					...a,
+					seed: r,
+					pastKV: a.cache ? a.cache[e] : void 0,
+					mixedPrecision: s
+				}, o = a.checkpointing && a.training ? n.callCheckpoint(i, c) : n.call(i, c);
+				a.outputEmbeddings ? (t(c), a.embeddings.push({
+					name: `block_output_${e}`,
+					tensor: c
+				})) : c.dispose(), c = o;
+			}
+			if (e = this.lnF.call({
+				...a,
+				mixedPrecision: s
+			}, c), c.dispose(), a.skipLogits) return this.endMemory("Forward"), e;
+			let l = this.head ? this.head.project(e) : this.wte.project(e);
+			a.outputEmbeddings ? (t(e), a.embeddings.push({
+				name: "final_norm_output",
+				tensor: e
+			})) : e.dispose();
+			let u = s ? r(l) : l;
+			return s && l !== u && l.dispose(), u;
+		});
+	}
+	project(t) {
+		return e(() => this.wte.project(t));
+	}
+	dispose() {
+		this.weightStore.dispose(), this.wte.dispose(), this.head && this.head.dispose(), this.blocks.forEach((e) => e.dispose()), this.lnF.dispose();
+	}
+};
+//#endregion
+export { f as default };

package/dist/models/config.d.ts

@@ -13,6 +13,7 @@ export interface GPTConfigBase {
     mlpFactor: number;
     loraConfig?: Map<string, LoRAConfig>;
     loraName?: string;
+    ropeTheta?: number;
 }
 export interface GPTConfigV1 extends GPTConfigBase {
     modelType: 'GenAI_NanoGPT_v1';
@@ -22,6 +23,22 @@ export interface GPTConfigV2 extends GPTConfigBase {
     modelType: 'GenAI_NanoGPT_v2';
     windowSize?: string;
 }
-export type GPTConfig = GPTConfigV1 | GPTConfigV2;
+export interface GPTConfigQwen3 extends GPTConfigBase {
+    modelType: 'qwen3';
+    tieEmbeddings?: boolean;
+    rmsNormEpsilon?: number;
+    kvHeads?: number;
+    headDim?: number;
+    mlpSize?: number;
+}
+export interface GPTConfigLlama extends GPTConfigBase {
+    modelType: 'llama';
+    tieEmbeddings?: boolean;
+    rmsNormEpsilon?: number;
+    kvHeads?: number;
+    headDim?: number;
+    mlpSize?: number;
+}
+export type GPTConfig = GPTConfigV1 | GPTConfigV2 | GPTConfigQwen3 | GPTConfigLlama;
 export declare const defaultConfig: GPTConfig;
 export declare function validateConfig(config: unknown): asserts config is GPTConfig;

package/dist/models/config.js

@@ -31,7 +31,7 @@ function n(e) {
 		if (e.windowSize !== void 0 && typeof e.windowSize != "string") throw Error("Invalid config: \"windowSize\" must be a string for GenAI_NanoGPT_v2.");
 		return;
 	}
-	throw Error("invalid_model_type");
+	if (e.modelType !== "qwen3" && e.modelType !== "llama") throw Error("invalid_model_type");
 }
 //#endregion
 export { e as defaultConfig, n as validateConfig };

package/dist/models/factory.js

@@ -1,13 +1,17 @@
 import e from "./NanoGPTV1.js";
 import t from "./NanoGPTV2.js";
+import n from "./Llama.js";
+import r from "./Qwen3.js";
 //#region lib/models/factory.ts
-function n(n) {
-	let r = n.modelType || "GenAI_NanoGPT_v1";
-	switch (r) {
-		case "GenAI_NanoGPT_v1": return new e(n);
-		case "GenAI_NanoGPT_v2": return new t(n);
-		default: throw Error(`Unsupported model type: ${r}`);
+function i(i) {
+	let a = i.modelType || "GenAI_NanoGPT_v1";
+	switch (a) {
+		case "GenAI_NanoGPT_v1": return new e(i);
+		case "GenAI_NanoGPT_v2": return new t(i);
+		case "qwen3": return new r(i);
+		case "llama": return new n(i);
+		default: throw Error(`Unsupported model type: ${a}`);
 	}
 }
 //#endregion
-export { n as default };
+export { i as default };

package/dist/models/model.d.ts

@@ -1,6 +1,6 @@
 import { Tensor } from '@tensorflow/tfjs-core';
 import { ForwardAttributes, default as BaseLayer } from '../layers/BaseLayer';
-import { AttentionScores, KVCache } from '../layers/CausalSelfAttention';
+import { AttentionScores, KVCache } from '../layers/MHA';
 import { TransformersMetadata } from '../../loader/types';
 import { GPTConfig, LoRAConfig } from './config';
 import { default as LoRA } from '../../layers/LoRA';

package/dist/ops/dot16.js

@@ -2,7 +2,7 @@ import { o as e } from "../tfjs_backend-h5weiy1O.js";
 import { isPackedTensor as t } from "../utilities/packed.js";
 import { transpose16 as n } from "./transpose16.js";
 import { reshape16 as r } from "./reshape16.js";
-import { t as i } from "../matMul16-DGtlziMn.js";
+import { t as i } from "../matMul16-CO1QhsRt.js";
 //#region lib/ops/dot16.ts
 function a(a, o, s = !1, c = !1) {
 	if (!t(a) && !t(o)) return e(a, o);

package/dist/ops/globalNorm.js

@@ -1,19 +1,19 @@
-import { di as e, ii as t, kt as n } from "../dist-Da20xy8E.js";
+import { _n as e, di as t, ii as n, kt as r } from "../dist-Da20xy8E.js";
 //#region lib/ops/globalNorm.ts
-function r(r, i, a = 1) {
-	return e(() => {
-		let e = n([r.length], "int32");
-		return r.map((n, r) => t().runKernel("Norm2", {
-			x: n,
+function i(i, a, o = 1) {
+	return n().backendName === "webgpu" ? t(() => {
+		let e = r([i.length], "int32");
+		return i.map((t, r) => n().runKernel("Norm2", {
+			x: t,
 			output: e
 		}, {
-			invLossScaling: i,
+			invLossScaling: a,
 			index: r
-		})), t().runKernel("ClipScale", { x: e }, {
-			invLossScaling: i,
-			clipNorm: a
+		})), n().runKernel("ClipScale", { x: e }, {
+			invLossScaling: a,
+			clipNorm: o
 		});
-	});
+	}) : e(a);
 }
 //#endregion
-export { r as clipScale };
+export { i as clipScale };

package/dist/ops/grads/attentionMask.js

@@ -1,6 +1,6 @@
 import { Fi as e } from "../../dist-Da20xy8E.js";
 import { transpose16 as t } from "../transpose16.js";
-import { a as n } from "../../matMul16-DGtlziMn.js";
+import { a as n } from "../../matMul16-CO1QhsRt.js";
 //#region lib/ops/grads/attentionMask.ts
 e({
 	kernelName: "AttentionMask",

package/dist/ops/grads/matMul16.js

@@ -1,2 +1,2 @@
-import { o as e } from "../../matMul16-DGtlziMn.js";
+import { o as e } from "../../matMul16-CO1QhsRt.js";
 export { e as matMul16GradConfig };

package/dist/ops/grads/qkv.js

@@ -1,6 +1,6 @@
 import { Fi as e, H as t } from "../../dist-Da20xy8E.js";
 import { sum16 as n } from "../sum16.js";
-import { o as r } from "../../matMul16-DGtlziMn.js";
+import { o as r } from "../../matMul16-CO1QhsRt.js";
 import { concat16 as i } from "../concat16.js";
 //#region lib/ops/grads/qkv.ts
 var a = {

package/dist/ops/grads/rope.js

@@ -1,2 +1,2 @@
-import { n as e } from "../../rope-C8EAiSGY.js";
+import { n as e } from "../../rope-BPu_uooz.js";
 export { e as ropeGradConfig };

package/dist/ops/grads/silu.d.ts

@@ -0,0 +1,2 @@
+import { GradConfig } from '@tensorflow/tfjs-core';
+export declare const siluGradConfig: GradConfig;

package/dist/ops/grads/silu.js

@@ -0,0 +1,2 @@
+import { r as e } from "../../silu-B6bvCgJ7.js";
+export { e as siluGradConfig };

package/dist/ops/matMul16.d.ts

@@ -3,7 +3,7 @@ export declare function matMul16(A: Tensor, B: Tensor, transposeA?: boolean, tra
     scale?: number;
     scaleA?: number;
     scaleB?: number;
-    activation?: 'gelu' | 'relu2' | 'relu';
+    activation?: 'gelu' | 'relu2' | 'relu' | 'silu';
     forceOutputShape?: number[];
     perm?: number[];
     causalMask?: boolean;

package/dist/ops/matMul16.js

@@ -1,2 +1,2 @@
-import { a as e, i as t, n, r, t as i } from "../matMul16-DGtlziMn.js";
+import { a as e, i as t, n, r, t as i } from "../matMul16-CO1QhsRt.js";
 export { i as matMul16, n as matMul16Gelu, r as matMul16ScaleA, t as matMul16ScaleB, e as matMul16Scaled };

package/dist/ops/normRMS.d.ts

@@ -1,2 +1,2 @@
 import { Tensor } from '@tensorflow/tfjs-core';
-export declare function normRMS(x: Tensor, gamma?: Tensor): Tensor;
+export declare function normRMS(x: Tensor, gamma?: Tensor, epsilon?: number): Tensor;

package/dist/ops/normRMS.js

@@ -3,20 +3,20 @@ import "./cpu/normRMS.js";
 import "./webgl/normRMS.js";
 import "./grads/normRMS.js";
 //#region lib/ops/normRMS.ts
-function t(t, n) {
+function t(t, n, r = 1e-8) {
 	if (t.rank === 4) {
-		let [r, i, ...a] = t.shape, o = [r * i, ...a], s = t.reshape(o), c = e().runKernel(n ? "RMSNorm" : "RMSNormNoGamma", n ? {
-			x: s,
+		let [i, a, ...o] = t.shape, s = [i * a, ...o], c = t.reshape(s), l = e().runKernel(n ? "RMSNorm" : "RMSNormNoGamma", n ? {
+			x: c,
 			gamma: n
-		} : { x: s });
-		s.dispose();
-		let l = c.reshape(t.shape);
-		return c.dispose(), l;
+		} : { x: c }, { epsilon: r });
+		c.dispose();
+		let u = l.reshape(t.shape);
+		return l.dispose(), u;
 	}
 	return e().runKernel(n ? "RMSNorm" : "RMSNormNoGamma", n ? {
 		x: t,
 		gamma: n
-	} : { x: t });
+	} : { x: t }, { epsilon: r });
 }
 //#endregion
 export { t as normRMS };

package/dist/ops/rope.d.ts

@@ -1,3 +1,3 @@
 import { default as RoPECache } from '../../layers/RoPECache';
 import { Tensor } from '@tensorflow/tfjs-core';
-export declare function rope(x: Tensor, cache: RoPECache, pastLength: number, negSin?: boolean): Tensor;
+export declare function rope(x: Tensor, cache: RoPECache, pastLength: number, negSin?: boolean, interleaved?: boolean): Tensor;

package/dist/ops/rope.js

@@ -1,2 +1,2 @@
-import { t as e } from "../rope-C8EAiSGY.js";
+import { t as e } from "../rope-BPu_uooz.js";
 export { e as rope };

package/dist/ops/silu.d.ts

@@ -0,0 +1,3 @@
+import { Tensor } from '@tensorflow/tfjs-core';
+export declare function silu(x: Tensor): Tensor;
+export declare function dSilu(dy: Tensor, x: Tensor): Tensor;

package/dist/ops/silu.js

@@ -0,0 +1,2 @@
+import { n as e, t } from "../silu-B6bvCgJ7.js";
+export { t as dSilu, e as silu };

package/dist/ops/webgl/rope.js

@@ -57,11 +57,13 @@ var t = class {
 	}
 };
 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, [
+	let { x: n } = e.inputs, { pastLen: r, ropeCache: i, negSin: a, interleaved: o } = e.attrs, s = a ? i.getNegSin() : i.getSin(), c = i.getCos(), l = e.backend;
+	if (!o) throw Error("Non-interleaved mode is not supported in this WebGL implementation.");
+	let u = n.shape[0], d = n.shape[1], f = n.shape[2], p = n.shape[3], m = new t(u, d, f, p);
+	return l.runWebGLProgram(m, [
 		n,
-		o,
-		s
+		s,
+		c
 	], "float32", [[r]]);
 }
 e({

package/dist/ops/webgpu/attentionMask.js

@@ -1,6 +1,6 @@
 import { Fs as e, Ii as t } from "../../dist-Da20xy8E.js";
 import { isPackedTensor as n } from "../../utilities/packed.js";
-import { t as r } from "../../matMul16-DGtlziMn.js";
+import { t as r } from "../../matMul16-CO1QhsRt.js";
 import i from "./attentionMask32_program.js";
 //#region lib/ops/webgpu/attentionMask.ts
 function a(t) {

package/dist/ops/webgpu/index.js

@@ -5,6 +5,7 @@ import "./sub16.js";
 import "./softmax16grad.js";
 import "./softmax16.js";
 import "./slice16.js";
+import "./silu.js";
 import "./scatterSub.js";
 import "./rope.js";
 import "./qkv.js";

package/dist/ops/webgpu/matMul16.js

@@ -1,74 +1,80 @@
 import { Ii as e, In as t, Wr as n, _n as r, br as i, gr as a, hn as o, lt as s, oc as c, w as l } from "../../dist-Da20xy8E.js";
 import { isPackedTensor as u } from "../../utilities/packed.js";
 import { reshape16 as d } from "../reshape16.js";
-import { matMulMul as f } from "../matMulMul.js";
-import { matMulGelu as p } from "../matMulGelu.js";
-import m from "./matMul16_program.js";
+import { n as f } from "../../silu-B6bvCgJ7.js";
+import { matMulMul as p } from "../matMulMul.js";
+import { matMulGelu as m } from "../matMulGelu.js";
+import h from "./matMul16_program.js";
 //#region lib/ops/webgpu/matMul16.ts
-function h(e) {
-	let { A: h, B: g } = e.inputs, { transposeA: _, transposeB: v, scale: y, activation: b, scaleA: x, scaleB: S, forceOutputShape: C, perm: w, causalMask: T, pastLen: E } = e.attrs, D = e.backend, O = !u(h), k = !u(g);
-	if (O && k) {
-		let e = x === void 0 ? h : n(h, r(x)), t = S === void 0 ? g : n(g, r(S));
-		if (T) throw Error("Causal mask is not supported for unpacked MatMul16.");
+function g(e) {
+	let { A: g, B: _ } = e.inputs, { transposeA: v, transposeB: y, scale: b, activation: x, scaleA: S, scaleB: C, forceOutputShape: w, perm: T, causalMask: E, pastLen: D } = e.attrs, O = e.backend, k = !u(g), A = !u(_);
+	if (k && A) {
+		let e = S === void 0 ? g : n(g, r(S)), t = C === void 0 ? _ : n(_, r(C));
+		if (E) throw Error("Causal mask is not supported for unpacked MatMul16.");
 		let c;
-		if (b === "gelu") c = p(e, t);
-		else if (b === "relu2") {
-			let n = a(e, t, _, v), r = s(n);
+		if (x === "gelu") c = m(e, t);
+		else if (x === "relu2") {
+			let n = a(e, t, v, y), r = s(n);
 			n.dispose(), c = o(r), r.dispose();
-		} else if (b === "relu") {
-			if (c = s(a(e, t, _, v)), y !== void 0) {
-				let e = n(c, r(y));
+		} else if (x === "relu") {
+			if (c = s(a(e, t, v, y)), b !== void 0) {
+				let e = n(c, r(b));
 				c.dispose(), c = e;
 			}
-		} else c = y === void 0 ? a(e, t, _, v) : f(e, t, r(y), _, v);
-		if (w) if (C) {
-			let e = i(c, C);
+		} else if (x === "silu") {
+			if (c = f(a(e, t, v, y)), b !== void 0) {
+				let e = n(c, r(b));
+				c.dispose(), c = e;
+			}
+		} else c = b === void 0 ? a(e, t, v, y) : p(e, t, r(b), v, y);
+		if (T) if (w) {
+			let e = i(c, w);
 			c.dispose();
-			let t = l(e, w);
+			let t = l(e, T);
 			return e.dispose(), t;
 		} else {
-			let e = l(c, w);
+			let e = l(c, T);
 			return c.dispose(), e;
 		}
-		else if (C) {
-			let e = i(c, C);
+		else if (w) {
+			let e = i(c, w);
 			return c.dispose(), e;
 		} else return c;
 	}
-	if (O && !k) throw Error("When using mixed precision, A must be packed if B is packed.");
-	if (!O && k) throw Error("When using mixed precision, B must be packed if A is packed.");
-	let A = h.shape.length, j = g.shape.length, M = h.shape.slice(0, -2), N = g.shape.slice(0, -2), P = c(M), F = c(N), I = t(h.shape.slice(0, -2), g.shape.slice(0, -2)), L = Math.max(P, F), R = h.shape[A - 2], z = g.shape[j - 2], B = h.shape[A - 1] * 2, V = g.shape[j - 1] * 2, H = d(h, [
-		P,
-		h.shape[A - 2],
-		h.shape[A - 1]
-	]), U = d(g, [
+	if (k && !A) throw Error("When using mixed precision, A must be packed if B is packed.");
+	if (!k && A) throw Error("When using mixed precision, B must be packed if A is packed.");
+	let j = g.shape.length, M = _.shape.length, N = g.shape.slice(0, -2), P = _.shape.slice(0, -2), F = c(N), I = c(P), L = t(g.shape.slice(0, -2), _.shape.slice(0, -2)), R = Math.max(F, I), z = g.shape[j - 2], B = _.shape[M - 2], V = g.shape[j - 1] * 2, H = _.shape[M - 1] * 2, U = d(g, [
 		F,
 		g.shape[j - 2],
 		g.shape[j - 1]
-	]), W = new m(L, R, z, B, V, _, v), G = [];
-	y !== void 0 && (W.useScale(), G.push({
+	]), W = d(_, [
+		I,
+		_.shape[M - 2],
+		_.shape[M - 1]
+	]), G = new h(R, z, B, V, H, v, y), K = [];
+	b !== void 0 && (G.useScale(), K.push({
 		type: "float32",
-		data: [y]
-	})), x !== void 0 && (W.useScaleA(), G.push({
-		type: "float32",
-		data: [x]
-	})), S !== void 0 && (W.useScaleB(), G.push({
+		data: [b]
+	})), S !== void 0 && (G.useScaleA(), K.push({
 		type: "float32",
 		data: [S]
-	})), b !== void 0 && W.useActivation(b), T && (W.useCausalMask(), G.push({
+	})), C !== void 0 && (G.useScaleB(), K.push({
+		type: "float32",
+		data: [C]
+	})), x !== void 0 && G.useActivation(x), E && (G.useCausalMask(), K.push({
 		type: "int32",
-		data: [E || 0]
+		data: [D || 0]
 	}));
-	let K = W.outputShape.length;
-	C && (e.attrs.originalShape = W.outputShape);
-	let q = C ?? I.concat([W.outputShape[K - 2], W.outputShape[K - 1]]);
-	W.setOutputShape(q, w);
-	let J = D.runWebGPUProgram(W, [H, U], "packedF16", G.length > 0 ? G : void 0);
-	return H.dispose(), U.dispose(), J;
+	let q = G.outputShape.length;
+	w && (e.attrs.originalShape = G.outputShape);
+	let J = w ?? L.concat([G.outputShape[q - 2], G.outputShape[q - 1]]);
+	G.setOutputShape(J, T);
+	let Y = O.runWebGPUProgram(G, [U, W], "packedF16", K.length > 0 ? K : void 0);
+	return U.dispose(), W.dispose(), Y;
 }
 e({
 	kernelName: "MatMul16",
 	backendName: "webgpu",
-	kernelFunc: h
+	kernelFunc: g
 });
 //#endregion

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

@@ -19,7 +19,7 @@ export default class MatMul16ProgramGeneric implements WebGPUProgram {
     scale: boolean;
     scaleA: boolean;
     scaleB: boolean;
-    activation?: 'gelu' | 'relu2' | 'relu';
+    activation?: 'gelu' | 'relu2' | 'relu' | 'silu';
     causalMask: boolean;
     outputComponent?: number | undefined;
     variableComponents?: number[];
@@ -31,7 +31,7 @@ export default class MatMul16ProgramGeneric implements WebGPUProgram {
     useScale(): void;
     useScaleA(): void;
     useScaleB(): void;
-    useActivation(activation: 'gelu' | 'relu2' | 'relu'): void;
+    useActivation(activation: 'gelu' | 'relu2' | 'relu' | 'silu'): void;
     useCausalMask(): void;
     private activationSnippet;
     private readASnippet;

package/dist/ops/webgpu/matMul16_program.js

@@ -110,7 +110,7 @@ var n = class {
 		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                " : "";
+		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                " : this.activation === "silu" ? "\n                fn activation(x : vec4<f32>) -> vec4<f32> {\n                    return x * (1.0 / (1.0 + exp(-x)));\n                }\n                " : "";
 	}
 	readASnippet() {
 		let e = `