@genai-fi/nanogpt 1.4.0 → 1.5.0

package/dist/layers/{CausalSelfAttention.js → GQA.js}

@@ -2,18 +2,18 @@ import { T as e, di as t, oi as n, pt as r } from "../dist-Da20xy8E.js";
 import { isPackedTensor as i } from "../utilities/packed.js";
 import { transpose16 as a } from "../ops/transpose16.js";
 import { reshape16 as o } from "../ops/reshape16.js";
-import { t as s } from "../matMul16-DGtlziMn.js";
+import { t as s } from "../matMul16-CO1QhsRt.js";
 import { t as c } from "../pack16-BhuXNUS7.js";
 import { attentionMask as l } from "../ops/attentionMask.js";
 import u from "./BaseLayer.js";
-import { t as d } from "../rope-C8EAiSGY.js";
+import { t as d } from "../rope-BPu_uooz.js";
 import { appendCache as f } from "../ops/appendCache.js";
 import { softmax16 as p } from "../ops/softmax16.js";
 import { dot16 as m } from "../ops/dot16.js";
 import { qkv as h } from "../ops/qkv.js";
-import { normRMS as g } from "../ops/normRMS.js";
-import { dropout16 as _ } from "../ops/dropout16.js";
-//#region lib/layers/CausalSelfAttention.ts
+import { dropout16 as g } from "../ops/dropout16.js";
+import _ from "./RMSNorm.js";
+//#region lib/layers/GQA.ts
 var v = class extends u {
 	attentionConfig;
 	divisor;
@@ -22,8 +22,10 @@ var v = class extends u {
 	projUnits;
 	ATTN;
 	PROJ;
+	qNorm = null;
+	kNorm = null;
 	constructor(e, t, n, r) {
-		super(t, r), this.attentionConfig = n, this.index = e, this.units = t.nEmbed * 3, this.projUnits = t.nEmbed, this.ATTN = `block_${this.index}_cAttn`, this.PROJ = `block_${this.index}_cProj`, this.addVariable(this.ATTN), this.addVariable(this.PROJ), this.divisor = 1 / Math.sqrt(t.nEmbed / t.nHead);
+		super(t, r), this.attentionConfig = n, this.index = e, this.units = t.nEmbed * 3, this.projUnits = t.nEmbed, this.ATTN = `block_${this.index}_cAttn`, this.PROJ = `block_${this.index}_cProj`, this.addVariable(this.ATTN), this.addVariable(this.PROJ), this.divisor = 1 / Math.sqrt(t.nEmbed / t.nHead), (this.attentionConfig.useQKNorm ?? !1) && (this.qNorm = new _(t, n, `block_${this.index}_qNorm`, this), this.kNorm = new _(t, n, `block_${this.index}_kNorm`, this));
 	}
 	build() {
 		this.hasVariable(this.ATTN) === !1 && this.setVariable(this.ATTN, e(r([this.config.nEmbed, this.units], 0, .02), !0, this.ATTN)), this.hasVariable(this.PROJ) === !1 && this.setVariable(this.PROJ, e(r([this.projUnits, this.config.nEmbed], 0, .02), !0, this.PROJ));
@@ -62,20 +64,20 @@ var v = class extends u {
 	forward(e, r) {
 		return t(() => {
 			this.startMemory();
-			let [t, i, a] = this.getQKV(r), o = e.pastKV ? e.pastKV.cumulativeLength : e.ropePositionOffset || 0, c = e.ropeCache, l = c ? d(t, c, o) : t, u = c ? d(i, c, o) : i, f = this.attentionConfig.useQKNorm ?? !1, p = f ? g(l) : l;
-			f && l.dispose();
-			let m = f ? g(u) : u;
-			f && u.dispose(), c && (t.dispose(), i.dispose());
+			let [t, i, a] = this.getQKV(r), o = e.pastKV ? e.pastKV.cumulativeLength : e.ropePositionOffset || 0, c = e.ropeCache, l = this.attentionConfig.ropeInterleaved ?? !0, u = c ? d(t, c, o, !1, l) : t, f = c ? d(i, c, o, !1, l) : i, p = this.qNorm ? this.qNorm.call(e, u) : u;
+			this.qNorm && u.dispose();
+			let m = this.kNorm ? this.kNorm.call(e, f) : f;
+			this.kNorm && f.dispose(), c && (t.dispose(), i.dispose());
 			let h = e.pastKV ? e.pastKV.length : 0;
 			e.pastKV && !e.training && this.updateCache(m, a, e.pastKV);
-			let v = e.pastKV?.k ? e.pastKV.k : m, y = e.pastKV?.v ? e.pastKV.v : a, b;
-			b = h > 0 ? this.getAttentionScores(p, v, h) : this.getAttentionScores(p, v), p.dispose(), e.pastKV || v.dispose();
-			let x = s(b, y), S = e.attentionScores !== void 0 && e.attentionScores.attentionOut !== void 0;
-			S || b.dispose(), e.pastKV || y.dispose();
-			let C = this.getOutputProjection(x);
-			if (x.dispose(), S && e.attentionScores && e.attentionScores.attentionOut !== void 0) {
-				let t = b.shape[1], r = b.shape[2];
-				e.attentionScores.attentionOut?.push(n(b.slice([
+			let _ = e.pastKV?.k ? e.pastKV.k : m, v = e.pastKV?.v ? e.pastKV.v : a, y;
+			y = h > 0 ? this.getAttentionScores(p, _, h) : this.getAttentionScores(p, _), p.dispose(), e.pastKV || _.dispose();
+			let b = s(y, v), x = e.attentionScores !== void 0 && e.attentionScores.attentionOut !== void 0;
+			x || y.dispose(), e.pastKV || v.dispose();
+			let S = this.getOutputProjection(b);
+			if (b.dispose(), x && e.attentionScores && e.attentionScores.attentionOut !== void 0) {
+				let t = y.shape[1], r = y.shape[2];
+				e.attentionScores.attentionOut?.push(n(y.slice([
 					0,
 					0,
 					0,
@@ -91,7 +93,7 @@ var v = class extends u {
 					-1
 				])));
 			}
-			return this.endMemory("CausalSelfAttention"), e.dropout && e.dropout > 0 ? _(C, e.dropout) : C;
+			return this.endMemory("CausalSelfAttention"), e.dropout && e.dropout > 0 ? g(S, e.dropout) : S;
 		});
 	}
 };

package/dist/layers/HeadLayer.d.ts

@@ -0,0 +1,12 @@
+import { Tensor } from '@tensorflow/tfjs-core';
+import { default as BaseLayer, ForwardAttributes } from './BaseLayer';
+import { GPTConfig } from '../../models/config';
+export default class HeadLayer extends BaseLayer {
+    private vocabSize;
+    private embedDim;
+    private initializer;
+    private WEIGHTS;
+    constructor(config: GPTConfig, name: string, parent?: BaseLayer);
+    project(inputs: Tensor): Tensor;
+    forward(_: ForwardAttributes, x: Tensor): Tensor;
+}

package/dist/layers/HeadLayer.js

@@ -0,0 +1,28 @@
+import { T as e } from "../dist-Da20xy8E.js";
+import { n as t } from "../dist-CwK5S7Ls.js";
+import { isPackedTensor as n } from "../utilities/packed.js";
+import { t as r } from "../pack16-BhuXNUS7.js";
+import i from "./BaseLayer.js";
+import { dot16 as a } from "../ops/dot16.js";
+//#region lib/layers/HeadLayer.ts
+var o = class extends i {
+	vocabSize;
+	embedDim;
+	initializer;
+	WEIGHTS;
+	constructor(n, r, i) {
+		super(n, i), this.WEIGHTS = r, this.vocabSize = n.vocabSize, this.embedDim = n.nEmbed, this.initializer = t({
+			mean: 0,
+			stddev: .02
+		}), this.addVariable(this.WEIGHTS, e(this.initializer.apply([this.embedDim, this.vocabSize]), !0, r));
+	}
+	project(e) {
+		let t = n(e) ? r(this.getVariable(this.WEIGHTS), void 0, 32) : this.getVariable(this.WEIGHTS), i = a(e, t);
+		return n(e) && t.dispose(), i;
+	}
+	forward(e, t) {
+		return this.project(t);
+	}
+};
+//#endregion
+export { o as default };

package/dist/layers/{CausalSelfAttention.d.ts → MHA.d.ts}

@@ -1,6 +1,7 @@
 import { default as BaseLayer, ForwardAttributes } from './BaseLayer';
 import { Tensor } from '@tensorflow/tfjs-core';
 import { GPTConfig } from '../../models/config';
+import { RMSNormConfig } from './RMSNorm';
 export interface KVCache {
     k?: Tensor;
     v?: Tensor;
@@ -17,10 +18,11 @@ interface AttentionForwardAttributes extends ForwardAttributes {
     seed?: number;
     ropePositionOffset?: number;
 }
-export interface CausalSelfAttentionConfig {
+export interface MHAConfig extends RMSNormConfig {
     useQKNorm?: boolean;
+    ropeInterleaved?: boolean;
 }
-export default class CausalSelfAttention extends BaseLayer<AttentionForwardAttributes> {
+export default class MHA extends BaseLayer<AttentionForwardAttributes> {
     private readonly attentionConfig;
     private divisor;
     private index;
@@ -28,7 +30,9 @@ export default class CausalSelfAttention extends BaseLayer<AttentionForwardAttri
     private projUnits;
     private ATTN;
     private PROJ;
-    constructor(index: number, config: GPTConfig, attentionConfig: CausalSelfAttentionConfig, parent?: BaseLayer);
+    private qNorm;
+    private kNorm;
+    constructor(index: number, config: GPTConfig, attentionConfig: MHAConfig, parent?: BaseLayer);
     protected build(): void;
     private getAttentionScores;
     private getQKV;

package/dist/layers/MHA.js

@@ -0,0 +1,101 @@
+import { T as e, di as t, oi as n, pt as r } from "../dist-Da20xy8E.js";
+import { isPackedTensor as i } from "../utilities/packed.js";
+import { transpose16 as a } from "../ops/transpose16.js";
+import { reshape16 as o } from "../ops/reshape16.js";
+import { t as s } from "../matMul16-CO1QhsRt.js";
+import { t as c } from "../pack16-BhuXNUS7.js";
+import { attentionMask as l } from "../ops/attentionMask.js";
+import u from "./BaseLayer.js";
+import { t as d } from "../rope-BPu_uooz.js";
+import { appendCache as f } from "../ops/appendCache.js";
+import { softmax16 as p } from "../ops/softmax16.js";
+import { dot16 as m } from "../ops/dot16.js";
+import { qkv as h } from "../ops/qkv.js";
+import { dropout16 as g } from "../ops/dropout16.js";
+import _ from "./RMSNorm.js";
+//#region lib/layers/MHA.ts
+var v = class extends u {
+	attentionConfig;
+	divisor;
+	index;
+	units;
+	projUnits;
+	ATTN;
+	PROJ;
+	qNorm = null;
+	kNorm = null;
+	constructor(e, t, n, r) {
+		super(t, r), this.attentionConfig = n, this.index = e, this.units = t.nEmbed * 3, this.projUnits = t.nEmbed, this.ATTN = `block_${this.index}_cAttn`, this.PROJ = `block_${this.index}_cProj`, this.addVariable(this.ATTN), this.addVariable(this.PROJ), this.divisor = 1 / Math.sqrt(t.nEmbed / t.nHead), (this.attentionConfig.useQKNorm ?? !1) && (this.qNorm = new _(t, n, `block_${this.index}_qNorm`, this), this.kNorm = new _(t, n, `block_${this.index}_kNorm`, this));
+	}
+	build() {
+		this.hasVariable(this.ATTN) === !1 && this.setVariable(this.ATTN, e(r([this.config.nEmbed, this.units], 0, .02), !0, this.ATTN)), this.hasVariable(this.PROJ) === !1 && this.setVariable(this.PROJ, e(r([this.projUnits, this.config.nEmbed], 0, .02), !0, this.PROJ));
+	}
+	getAttentionScores(e, t, n) {
+		let r = l(e, t, this.divisor, n), i = p(r);
+		return r.dispose(), i;
+	}
+	getQKV(e) {
+		let t = i(e) ? c(this.getVariable(this.ATTN)) : this.getVariable(this.ATTN), n = h(e, t, this.config.nHead);
+		return i(e) && t.dispose(), n;
+	}
+	getOutputProjection(e) {
+		let t = e.shape[0], n = e.shape[2], r = this.config.nEmbed, s = i(e), l = a(e, [
+			0,
+			2,
+			1,
+			3
+		]), u = o(l, [
+			t,
+			n,
+			s ? r / 2 : r
+		]);
+		l.dispose();
+		let d = s ? c(this.getVariable(this.PROJ)) : this.getVariable(this.PROJ), f = m(u, d);
+		return s && d.dispose(), u.dispose(), f;
+	}
+	updateCache(e, t, r) {
+		let i = this.config.blockSize, a = e.shape[2], o = r.length || 0, s = f(e, i, o, r.k);
+		e.dispose(), r.k && r.k.dispose();
+		let c = f(t, i, o, r.v);
+		t.dispose(), r.v && r.v.dispose();
+		let l = Math.min(o + a, i), u = r.cumulativeLength + a;
+		r.length = l, r.cumulativeLength = u, r.k = n(s), r.v = n(c);
+	}
+	forward(e, r) {
+		return t(() => {
+			this.startMemory();
+			let [t, i, a] = this.getQKV(r), o = e.pastKV ? e.pastKV.cumulativeLength : e.ropePositionOffset || 0, c = e.ropeCache, l = this.attentionConfig.ropeInterleaved ?? !0, u = c ? d(t, c, o, !1, l) : t, f = c ? d(i, c, o, !1, l) : i, p = this.qNorm ? this.qNorm.call(e, u) : u;
+			this.qNorm && u.dispose();
+			let m = this.kNorm ? this.kNorm.call(e, f) : f;
+			this.kNorm && f.dispose(), c && (t.dispose(), i.dispose());
+			let h = e.pastKV ? e.pastKV.length : 0;
+			e.pastKV && !e.training && this.updateCache(m, a, e.pastKV);
+			let _ = e.pastKV?.k ? e.pastKV.k : m, v = e.pastKV?.v ? e.pastKV.v : a, y;
+			y = h > 0 ? this.getAttentionScores(p, _, h) : this.getAttentionScores(p, _), p.dispose(), e.pastKV || _.dispose();
+			let b = s(y, v), x = e.attentionScores !== void 0 && e.attentionScores.attentionOut !== void 0;
+			x || y.dispose(), e.pastKV || v.dispose();
+			let S = this.getOutputProjection(b);
+			if (b.dispose(), x && e.attentionScores && e.attentionScores.attentionOut !== void 0) {
+				let t = y.shape[1], r = y.shape[2];
+				e.attentionScores.attentionOut?.push(n(y.slice([
+					0,
+					0,
+					0,
+					0
+				], [
+					1,
+					-1,
+					-1,
+					-1
+				]).reshape([
+					t,
+					r,
+					-1
+				])));
+			}
+			return this.endMemory("CausalSelfAttention"), e.dropout && e.dropout > 0 ? g(S, e.dropout) : S;
+		});
+	}
+};
+//#endregion
+export { v as default };

package/dist/layers/MLP.d.ts

@@ -2,14 +2,18 @@ import { Tensor } from '@tensorflow/tfjs-core';
 import { default as BaseLayer, ForwardAttributes } from './BaseLayer';
 import { GPTConfig } from '../../models/config';
 export interface MLPConfig {
-    activation?: 'gelu' | 'relu2';
+    activation?: 'gelu' | 'relu2' | 'swiglu';
     hiddenFactor?: number;
+    hiddenSize?: number;
 }
+export declare function swigluHidden(C: number, factor?: number, multipleOf?: number): number;
 export default class MLP extends BaseLayer {
     private index;
     private hiddenUnits;
     private MLPHIDDEN;
     private MLPOUT;
+    private MLPGATE;
+    private MLPUP;
     private mlpConfig;
     constructor(index: number, config: GPTConfig, mlpConfig: MLPConfig, parent?: BaseLayer);
     protected build(): void;

package/dist/layers/MLP.js

@@ -1,34 +1,46 @@
 import { T as e, di as t, pt as n } from "../dist-Da20xy8E.js";
 import { reshape16 as r } from "../ops/reshape16.js";
-import { t as i } from "../matMul16-DGtlziMn.js";
-import a from "./BaseLayer.js";
-import { dropout16 as o } from "../ops/dropout16.js";
+import { mul16 as i } from "../ops/mul16.js";
+import { t as a } from "../matMul16-CO1QhsRt.js";
+import o from "./BaseLayer.js";
+import { dropout16 as s } from "../ops/dropout16.js";
 //#region lib/layers/MLP.ts
-var s = class extends a {
+function c(e, t = 8 / 3, n = 32) {
+	let r = Math.round(t * e);
+	return Math.max(n, Math.round(r / n) * n);
+}
+var l = class extends o {
 	index;
 	hiddenUnits;
 	MLPHIDDEN;
 	MLPOUT;
+	MLPGATE;
+	MLPUP;
 	mlpConfig;
 	constructor(e, t, n, r) {
-		super(t, r), this.index = e, this.mlpConfig = n, this.hiddenUnits = (n.hiddenFactor ?? t.mlpFactor) * t.nEmbed, this.MLPHIDDEN = `block_${this.index}_mlpHidden`, this.MLPOUT = `block_${this.index}_mlpOut`, this.addVariable(this.MLPHIDDEN), this.addVariable(this.MLPOUT);
+		super(t, r), this.index = e, this.mlpConfig = n, this.hiddenUnits = n.hiddenSize ?? c(t.nEmbed, n.hiddenFactor ?? t.mlpFactor), this.MLPHIDDEN = `block_${this.index}_mlpHidden`, this.MLPOUT = `block_${this.index}_mlpOut`, this.addVariable(this.MLPOUT), this.MLPGATE = `block_${this.index}_mlpGate`, this.MLPUP = `block_${this.index}_mlpUp`, this.mlpConfig.activation === "swiglu" ? (this.addVariable(this.MLPGATE), this.addVariable(this.MLPUP)) : this.addVariable(this.MLPHIDDEN);
 	}
 	build() {
-		this.hasVariable(this.MLPHIDDEN) === !1 && this.setVariable(this.MLPHIDDEN, e(n([this.config.nEmbed, this.hiddenUnits], 0, .02), !0, this.MLPHIDDEN)), this.hasVariable(this.MLPOUT) === !1 && this.setVariable(this.MLPOUT, e(n([this.hiddenUnits, this.config.nEmbed], 0, .02 / Math.sqrt(2 * this.config.nLayer)), !0, this.MLPOUT));
+		this.mlpConfig.activation === "swiglu" ? (this.hasVariable(this.MLPGATE) === !1 && this.setVariable(this.MLPGATE, e(n([this.config.nEmbed, this.hiddenUnits], 0, .02), !0, this.MLPGATE)), this.hasVariable(this.MLPUP) === !1 && this.setVariable(this.MLPUP, e(n([this.config.nEmbed, this.hiddenUnits], 0, .02), !0, this.MLPUP))) : this.hasVariable(this.MLPHIDDEN) === !1 && this.setVariable(this.MLPHIDDEN, e(n([this.config.nEmbed, this.hiddenUnits], 0, .02), !0, this.MLPHIDDEN)), this.hasVariable(this.MLPOUT) === !1 && this.setVariable(this.MLPOUT, e(n([this.hiddenUnits, this.config.nEmbed], 0, .02 / Math.sqrt(2 * this.config.nLayer)), !0, this.MLPOUT));
 	}
 	forward(e, n) {
 		return t(() => {
 			this.startMemory();
-			let [t, a, s] = n.shape, c = i(r(n, [t * a, s]), this.getVariable(this.MLPHIDDEN), !1, !1, { activation: this.mlpConfig.activation ?? "gelu" }), l = i(c, this.getVariable(this.MLPOUT));
-			c.dispose();
-			let u = r(l, [
+			let [t, o, c] = n.shape, l = r(n, [t * o, c]), u;
+			if (this.mlpConfig.activation === "swiglu") {
+				let e = a(l, this.getVariable(this.MLPGATE), !1, !1, { activation: "silu" }), t = a(l, this.getVariable(this.MLPUP));
+				u = i(e, t), e.dispose(), t.dispose();
+			} else u = a(l, this.getVariable(this.MLPHIDDEN), !1, !1, { activation: this.mlpConfig.activation ?? "gelu" });
+			let d = a(u, this.getVariable(this.MLPOUT));
+			u.dispose();
+			let f = r(d, [
 				t,
-				a,
-				s
+				o,
+				c
 			]);
-			return this.endMemory("MLP"), e.dropout && e.dropout > 0 ? o(u, e.dropout) : u;
+			return this.endMemory("MLP"), e.dropout && e.dropout > 0 ? s(f, e.dropout) : f;
 		});
 	}
 };
 //#endregion
-export { s as default };
+export { l as default, c as swigluHidden };

package/dist/layers/RMSNorm.d.ts

@@ -3,10 +3,12 @@ import { default as BaseLayer, ForwardAttributes } from './BaseLayer';
 import { GPTConfig } from '../../models/config';
 export interface RMSNormConfig {
     useGamma?: boolean;
+    epsilon?: number;
 }
 export default class RMSNorm extends BaseLayer {
     private GAMMA;
     private rmsConfig;
     constructor(config: GPTConfig, rmsConfig: RMSNormConfig, name?: string, parent?: BaseLayer);
+    protected buildGamma(size: number): void;
     forward(_: ForwardAttributes, x: Tensor): Tensor;
 }

package/dist/layers/RMSNorm.js

@@ -1,20 +1,24 @@
 import { Ot as e, T as t, di as n } from "../dist-Da20xy8E.js";
-import r from "./BaseLayer.js";
-import { normRMS as i } from "../ops/normRMS.js";
+import { isPackedTensor as r } from "../utilities/packed.js";
+import i from "./BaseLayer.js";
+import { normRMS as a } from "../ops/normRMS.js";
 //#region lib/layers/RMSNorm.ts
-var a = class extends r {
+var o = class extends i {
 	GAMMA;
 	rmsConfig;
-	constructor(n, r, i = "", a) {
-		super(n, a), this.GAMMA = i, this.rmsConfig = r, (this.rmsConfig.useGamma ?? !0) && this.addVariable(this.GAMMA, t(e([n.nEmbed]), !0, this.GAMMA, "float32"));
+	constructor(e, t, n = "", r) {
+		super(e, r), this.GAMMA = n, this.rmsConfig = t, (this.rmsConfig.useGamma ?? !0) && this.addVariable(this.GAMMA);
+	}
+	buildGamma(n) {
+		(this.rmsConfig.useGamma ?? !0) && this.hasVariable(this.GAMMA) === !1 && this.setVariable(this.GAMMA, t(e([n]), !0, this.GAMMA, "float32"));
 	}
 	forward(e, t) {
 		return n(() => {
-			this.startMemory();
-			let e = i(t, this.rmsConfig.useGamma ?? !0 ? this.getVariable(this.GAMMA) : void 0);
+			this.startMemory(), this.buildGamma(r(t) ? t.shape[t.shape.length - 1] * 2 : t.shape[t.shape.length - 1]);
+			let e = a(t, this.rmsConfig.useGamma ?? !0 ? this.getVariable(this.GAMMA) : void 0, this.rmsConfig.epsilon);
 			return this.endMemory("RMSNorm"), e;
 		});
 	}
 };
 //#endregion
-export { a as default };
+export { o as default };

package/dist/layers/RoPECache.js

@@ -306,9 +306,10 @@ var wt = class {
 		this.config = e;
 		let t = this.config.nEmbed / this.config.nHead;
 		if (this.rotaryDim = t, this.rotaryDim % 2 != 0) throw Error("rotaryDim must be even");
-		this.ropeBase = 1e4;
-		let n = J(0, this.rotaryDim, 2, "float32"), r = u(n, W(this.rotaryDim, "float32")), i = Q(W(this.ropeBase, "float32"), r);
-		this.ropeInvFreq = Z(i), r.dispose(), i.dispose(), n.dispose(), (this.config.modelType === "GenAI_NanoGPT_v1" ? this.config.useRope ?? !0 : !0) === !1 ? (this.ropeCos = null, this.ropeSin = null, this.ropeNegSin = null, this.ropeCacheLen = 0) : K(() => {
+		let n = this.config.ropeTheta ?? 1e4;
+		this.ropeBase = n;
+		let r = J(0, this.rotaryDim, 2, "float32"), i = u(r, W(this.rotaryDim, "float32")), a = Q(W(this.ropeBase, "float32"), i);
+		this.ropeInvFreq = Z(a), i.dispose(), a.dispose(), r.dispose(), (this.config.modelType === "GenAI_NanoGPT_v1" ? this.config.useRope ?? !0 : !0) === !1 ? (this.ropeCos = null, this.ropeSin = null, this.ropeNegSin = null, this.ropeCacheLen = 0) : K(() => {
 			this.ensureRopeCache(this.config.blockSize * 4);
 		});
 	}

package/dist/layers/TransformerBlock.d.ts

@@ -1,4 +1,4 @@
-import { AttentionScores, CausalSelfAttentionConfig, KVCache } from './CausalSelfAttention';
+import { AttentionScores, MHAConfig, KVCache } from './MHA';
 import { MLPConfig } from './MLP';
 import { RMSNormConfig } from './RMSNorm';
 import { default as BaseLayer, ForwardAttributes } from './BaseLayer';
@@ -10,7 +10,7 @@ interface BlockAttributes extends ForwardAttributes {
     attentionScores?: AttentionScores;
     ropePositionOffset?: number;
 }
-export type TransformerBlockConfig = MLPConfig & RMSNormConfig & CausalSelfAttentionConfig;
+export type TransformerBlockConfig = MLPConfig & RMSNormConfig & MHAConfig;
 export default class Block extends BaseLayer<BlockAttributes> {
     private ln1;
     private attn;

package/dist/layers/TransformerBlock.js

@@ -1,8 +1,8 @@
 import { di as e, oi as t } from "../dist-Da20xy8E.js";
 import n from "./BaseLayer.js";
-import r from "./CausalSelfAttention.js";
-import i from "./MLP.js";
-import a from "./RMSNorm.js";
+import r from "./RMSNorm.js";
+import i from "./MHA.js";
+import a from "./MLP.js";
 import { add16 as o } from "../ops/add16.js";
 //#region lib/layers/TransformerBlock.ts
 var s = class extends n {
@@ -14,7 +14,7 @@ var s = class extends n {
 	skipped = !1;
 	blockConfig;
 	constructor(e, t, n, o) {
-		super(t, o), this.index = e, this.blockConfig = n, this.ln1 = new a(t, this.blockConfig, `block_${this.index}_rms1`, this), this.attn = new r(this.index, t, this.blockConfig, this), this.ln2 = new a(t, this.blockConfig, `block_${this.index}_rms2`, this), this.mlp = new i(this.index, t, this.blockConfig, this);
+		super(t, o), this.index = e, this.blockConfig = n, this.ln1 = new r(t, this.blockConfig, `block_${this.index}_rms1`, this), this.attn = new i(this.index, t, this.blockConfig, this), this.ln2 = new r(t, this.blockConfig, `block_${this.index}_rms2`, this), this.mlp = new a(this.index, t, this.blockConfig, this);
 	}
 	getMLPOutput(e, n) {
 		let r = this.ln2.call({ training: n.training }, e), i = this.mlp.call(n, r);

package/dist/layers.d.ts

@@ -1,4 +1,4 @@
-export { default as CausalSelfAttention } from './layers/CausalSelfAttention';
+export { default as CausalSelfAttention } from './layers/MHA';
 export { default as MLP } from './layers/MLP';
 export { default as TransformerBlock } from './layers/TransformerBlock';
 export { default as RoPECache } from './layers/RoPECache';

package/dist/layers.js

@@ -1,5 +1,5 @@
 import { n as e } from "./chunk-CWhphoD1.js";
-import t from "./layers/CausalSelfAttention.js";
+import t from "./layers/MHA.js";
 import n from "./layers/MLP.js";
 import r from "./layers/TransformerBlock.js";
 import i from "./layers/RoPECache.js";

package/dist/loader/loadTransformers.js

@@ -16,7 +16,7 @@ function o(e) {
 		nEmbed: e.hidden_size,
 		mlpFactor: e.mlpFactor,
 		useRope: e.useRope
-	} : {
+	} : e.model_type === "GenAI_NanoGPT_v2" ? {
 		modelType: "GenAI_NanoGPT_v2",
 		vocabSize: e.vocab_size,
 		blockSize: e.block_size,
@@ -27,6 +27,22 @@ function o(e) {
 		loraConfig: e.loraConfig ? new Map(Object.entries(e.loraConfig)) : void 0,
 		loraName: e.loraName,
 		windowSize: e.windowSize
+	} : e.model_type === "qwen3" ? {
+		modelType: "qwen3",
+		vocabSize: e.vocab_size,
+		blockSize: e.block_size,
+		nLayer: e.num_hidden_layers,
+		nHead: e.num_attention_heads,
+		nEmbed: e.hidden_size,
+		mlpFactor: e.mlpFactor
+	} : {
+		modelType: "llama",
+		vocabSize: e.vocab_size,
+		blockSize: e.block_size,
+		nLayer: e.num_hidden_layers,
+		nHead: e.num_attention_heads,
+		nEmbed: e.hidden_size,
+		mlpFactor: e.mlpFactor
 	}, t;
 }
 async function s(s, c, l, u) {

package/dist/loader/save.js

@@ -37,7 +37,7 @@ async function o(e, n, o, s) {
 		block_size: e.config.blockSize,
 		mlpFactor: e.config.mlpFactor,
 		useRope: e.config.useRope
-	} : {
+	} : f === "GenAI_NanoGPT_v2" ? {
 		model_type: "GenAI_NanoGPT_v2",
 		vocab_size: n.getVocab().length,
 		hidden_size: e.config.nEmbed,
@@ -46,8 +46,23 @@ async function o(e, n, o, s) {
 		block_size: e.config.blockSize,
 		mlpFactor: e.config.mlpFactor,
 		loraConfig: e.config.loraConfig ? Object.fromEntries(e.config.loraConfig) : void 0,
-		loraName: e.config.loraName,
-		windowSize: e.config.windowSize
+		loraName: e.config.loraName
+	} : f === "qwen3" ? {
+		model_type: "qwen3",
+		vocab_size: n.getVocab().length,
+		hidden_size: e.config.nEmbed,
+		num_hidden_layers: e.config.nLayer,
+		num_attention_heads: e.config.nHead,
+		block_size: e.config.blockSize,
+		mlpFactor: e.config.mlpFactor
+	} : {
+		model_type: "llama",
+		vocab_size: n.getVocab().length,
+		hidden_size: e.config.nEmbed,
+		num_hidden_layers: e.config.nLayer,
+		num_attention_heads: e.config.nHead,
+		block_size: e.config.blockSize,
+		mlpFactor: e.config.mlpFactor
 	}, l.file("config.json", JSON.stringify(p, void 0, 4), { binary: !1 });
 	let m = {
 		version: 2,

package/dist/loader/types.d.ts

@@ -5,7 +5,7 @@ import { ITokeniser } from '../../tokeniser/type';
 import { AdamWOptimizer } from '../../training/AdamW';
 import { TrainingLogEntry, TrainingOptions } from '../../training/types';
 export interface TransformersConfigBase {
-    model_type: 'GenAI_NanoGPT_v1' | 'GenAI_NanoGPT_v2';
+    model_type: 'GenAI_NanoGPT_v1' | 'GenAI_NanoGPT_v2' | 'qwen3' | 'llama';
     vocab_size: number;
     hidden_size: number;
     num_hidden_layers: number;
@@ -23,7 +23,13 @@ export interface TransformersConfigV2 extends TransformersConfigBase {
     model_type: 'GenAI_NanoGPT_v2';
     windowSize?: string;
 }
-export type TransformersConfig = TransformersConfigV1 | TransformersConfigV2;
+export interface TransformersConfigQwen3 extends TransformersConfigBase {
+    model_type: 'qwen3';
+}
+export interface TransformersConfigLlama extends TransformersConfigBase {
+    model_type: 'llama';
+}
+export type TransformersConfig = TransformersConfigV1 | TransformersConfigV2 | TransformersConfigQwen3 | TransformersConfigLlama;
 export interface TransformersTokeniser {
     type: 'char' | 'bpe';
     vocab: string[];

package/dist/main.js

@@ -11,13 +11,13 @@ import { $ as ii, A as ai, B as oi, C as si, D as ci, E as li, F as ui, G as di,
 import { S as ia, _ as aa, a as oa, c as sa, d as ca, f as la, g as ua, h as da, i as G, l as fa, m as pa, n as K, o as ma, p as ha, r as ga, s as _a, t as va, u as ya, v as q, x as ba, y as xa } from "./kernel_funcs_utils-ql8Y8qPn.js";
 import "./gelu-B6637MKa.js";
 import "./ops/transpose16.js";
-import "./matMul16-DGtlziMn.js";
+import "./matMul16-CO1QhsRt.js";
 import { n as Sa, r as Ca, t as J } from "./Reshape-BlylqwWy.js";
 import { i as wa } from "./matMulGelu-CBoqTZM7.js";
 import "./ops/matMulGelu.js";
 import { r as Ta, t as Ea } from "./pack16-BhuXNUS7.js";
 import "./ops/attentionMask.js";
-import "./rope-C8EAiSGY.js";
+import "./rope-BPu_uooz.js";
 import "./ops/appendCache.js";
 import "./ops/softmax16.js";
 import "./ops/qkv.js";

package/dist/matMul16-CO1QhsRt.js

@@ -0,0 +1,85 @@
+import { Fi as e, ii as t } from "./dist-Da20xy8E.js";
+import { t as n } from "./gelu-B6637MKa.js";
+import { isPackedTensor as r } from "./utilities/packed.js";
+import { transpose16 as i } from "./ops/transpose16.js";
+import { reshape16 as a } from "./ops/reshape16.js";
+import { mul16 as o } from "./ops/mul16.js";
+import { t as s } from "./silu-B6bvCgJ7.js";
+import "./ops/webgl/matMul16.js";
+import "./ops/cpu/matMul16.js";
+import { t as c } from "./pack16-BhuXNUS7.js";
+//#region lib/ops/grads/matMul16.ts
+var l = {
+	kernelName: "MatMul16",
+	inputsToSave: ["A", "B"],
+	outputsToSave: [],
+	gradFunc: (e, t, r) => {
+		let [c, l] = t;
+		if (Array.isArray(e)) throw Error("Expected dy to be a single Tensor");
+		let d = e, { transposeA: m, transposeB: h, scale: g, activation: _, originalShape: v, perm: y } = r;
+		if (y && v) {
+			let e = Array(y.length);
+			for (let t = 0; t < y.length; ++t) e[y[t]] = t;
+			let t = d;
+			d = i(d, e), t.dispose();
+		}
+		if (v) {
+			let e = d;
+			d = a(d, v), e.dispose();
+		}
+		if (_ === "gelu") {
+			let e = d, t = u(c, l, m, h);
+			d = n(e, t), e.dispose(), t.dispose();
+		} else if (_ === "relu2") {
+			let e = d, t = u(c, l, m, h, {
+				activation: "relu",
+				scale: 2
+			});
+			d = o(e, t), e.dispose(), t.dispose();
+		} else if (_ === "silu") {
+			let e = d, t = u(c, l, m, h);
+			d = s(e, t), e.dispose(), t.dispose();
+		}
+		if (!m && !h) return {
+			A: () => g === void 0 ? u(d, l, !1, !0) : f(d, l, g, !1, !0),
+			B: () => g === void 0 ? u(c, d, !0, !1) : p(c, d, g, !0, !1)
+		};
+		if (!m && h) return {
+			A: () => g === void 0 ? u(d, l, !1, !1) : f(d, l, g, !1, !1),
+			B: () => g === void 0 ? u(c, d, !0, !1) : p(c, d, g, !0, !1)
+		};
+		if (m && !h) return {
+			A: () => g === void 0 ? u(l, d, !1, !0) : p(l, d, g, !1, !0),
+			B: () => g === void 0 ? u(c, d, !1, !1) : p(c, d, g, !1, !1)
+		};
+		throw Error("Gradient for transposeA=true and transposeB=true is not supported yet.");
+	}
+};
+e(l);
+//#endregion
+//#region lib/ops/matMul16.ts
+function u(e, n, i = !1, a = !1, o = {}) {
+	let s = r(e), l = r(n), u = s || l, d = !u || s ? e : c(e), f = !u || l ? n : c(n), p = t().runKernel("MatMul16", {
+		A: d,
+		B: f
+	}, {
+		transposeA: i,
+		transposeB: a,
+		...o
+	});
+	return u && !s && d.dispose(), u && !l && f.dispose(), p;
+}
+function d(e, t, n, r = !1, i = !1) {
+	return u(e, t, r, i, { scale: n });
+}
+function f(e, t, n, r = !1, i = !1) {
+	return u(e, t, r, i, { scaleA: n });
+}
+function p(e, t, n, r = !1, i = !1) {
+	return u(e, t, r, i, { scaleB: n });
+}
+function m(e, t, n = !1, r = !1) {
+	return u(e, t, n, r, { activation: "gelu" });
+}
+//#endregion
+export { d as a, p as i, m as n, l as o, f as r, u as t };