@genai-fi/nanogpt 0.19.1 → 0.20.1

package/dist/ops/webgpu/matMul16.js

@@ -1,65 +1,70 @@
-import { c as V, m as U, a as g, n as X, U as y, _ as Y } from "../../index-CUXkjxiT.js";
-import { isPackedTensor as x } from "../../utilities/packed.js";
-import { reshape16 as C } from "../reshape16.js";
-import { matMulMul as Z } from "../matMulMul.js";
-import { matMulGelu as $ } from "../matMulGelu.js";
-import ee from "./matMul16_program.js";
-import { m as v } from "../../mat_mul-DhG0Newp.js";
-import { r as E } from "../../relu-J_X6MUzx.js";
-import { r as F } from "../../reshape-BDOuCSNW.js";
-import { t as W } from "../../transpose-9kRxIXWR.js";
-function se(m) {
-  const { A: s, B: t } = m.inputs, { transposeA: p, transposeB: i, scale: h, activation: c, scaleA: d, scaleB: f, forceOutputShape: r, perm: k, causalMask: w, pastLen: z } = m.attrs, I = m.backend, B = !x(s), b = !x(t);
-  if (B && b) {
-    const u = d !== void 0 ? U(s, g(d)) : s, l = f !== void 0 ? U(t, g(f)) : t;
-    if (w)
-      throw new Error("Causal mask is not supported for unpacked MatMul16.");
-    let e;
-    if (h !== void 0)
-      e = Z(u, l, g(h), p, i);
-    else if (c === "gelu")
-      e = $(u, l);
-    else if (c === "relu2") {
-      const o = v(u, l, p, i), A = E(o);
-      o.dispose(), e = X(A), A.dispose();
-    } else c === "relu" ? e = E(v(u, l, p, i)) : e = v(u, l, p, i);
-    if (k)
-      if (r) {
-        const o = F(e, r);
-        e.dispose();
-        const A = W(o, k);
-        return o.dispose(), A;
-      } else {
-        const o = W(e, k);
-        return e.dispose(), o;
-      }
-    else if (r) {
-      const o = F(e, r);
-      return e.dispose(), o;
-    } else
-      return e;
-  }
-  if (B && !b)
-    throw new Error("When using mixed precision, A must be packed if B is packed.");
-  if (!B && b)
-    throw new Error("When using mixed precision, B must be packed if A is packed.");
-  const M = s.shape.length, S = t.shape.length, N = s.shape.slice(0, -2), q = t.shape.slice(0, -2), R = y(N), D = y(q), K = Y(s.shape.slice(0, -2), t.shape.slice(0, -2)), L = Math.max(R, D), T = s.shape[M - 2], _ = t.shape[S - 2], j = s.shape[M - 1] * 2, H = t.shape[S - 1] * 2, G = C(s, [R, s.shape[M - 2], s.shape[M - 1]]), O = C(t, [D, t.shape[S - 2], t.shape[S - 1]]), a = new ee(L, T, _, j, H, p, i), n = [];
-  h !== void 0 && (a.useScale(), n.push({ type: "float32", data: [h] })), d !== void 0 && (a.useScaleA(), n.push({ type: "float32", data: [d] })), f !== void 0 && (a.useScaleB(), n.push({ type: "float32", data: [f] })), c !== void 0 && a.useActivation(c), w && (a.useCausalMask(), n.push({ type: "int32", data: [z || 0] }));
-  const P = a.outputShape.length;
-  r && (m.attrs.originalShape = a.outputShape);
-  const J = r ?? K.concat([a.outputShape[P - 2], a.outputShape[P - 1]]);
-  a.setOutputShape(J, k);
-  const Q = I.runWebGPUProgram(
-    a,
-    [G, O],
-    "packedF16",
-    n.length > 0 ? n : void 0
-  );
-  return G.dispose(), O.dispose(), Q;
+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-BewPQWjc.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";
+//#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.");
+		let c;
+		if (y !== void 0) c = f(e, t, r(y), _, v);
+		else if (b === "gelu") c = p(e, t);
+		else if (b === "relu2") {
+			let n = a(e, t, _, v), r = s(n);
+			n.dispose(), c = o(r), r.dispose();
+		} else c = b === "relu" ? s(a(e, t, _, v)) : a(e, t, _, v);
+		if (w) if (C) {
+			let e = i(c, C);
+			c.dispose();
+			let t = l(e, w);
+			return e.dispose(), t;
+		} else {
+			let e = l(c, w);
+			return c.dispose(), e;
+		}
+		else if (C) {
+			let e = i(c, C);
+			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, [
+		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({
+		type: "float32",
+		data: [y]
+	})), x !== void 0 && (W.useScaleA(), G.push({
+		type: "float32",
+		data: [x]
+	})), S !== void 0 && (W.useScaleB(), G.push({
+		type: "float32",
+		data: [S]
+	})), b !== void 0 && W.useActivation(b), T && (W.useCausalMask(), G.push({
+		type: "int32",
+		data: [E || 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;
 }
-const te = {
-  kernelName: "MatMul16",
-  backendName: "webgpu",
-  kernelFunc: se
-};
-V(te);
+e({
+	kernelName: "MatMul16",
+	backendName: "webgpu",
+	kernelFunc: h
+});
+//#endregion

package/dist/ops/webgpu/matMul16_program.js

@@ -1,149 +1,119 @@
-import { U as f } from "../../index-CUXkjxiT.js";
-import { e as p } from "../../webgpu_program-B4HmApL1.js";
-class B {
-  variableNames = ["A", "B"];
-  outputShape;
-  shaderKey = "MatMul16TB";
-  dispatchLayout;
-  dispatch;
-  workgroupSize = [8, 8, 1];
-  // 8x8 threads for 32x32 tile
-  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, o, a, i, r = !1, s = !1) {
-    if (this.transposeA = r, this.transposeB = s, this.variableComponents = [2, 2], this.outputComponent = 2, this.shaderKey = `MatMul16TB_${t}_${o}_${a}_${i}_${r ? "TA" : ""}${s ? "TB" : ""}`, r) {
-      if (this.outputShape = [e, a, i / 2], this.dimInner = t, t !== o)
-        throw new Error("Inner dimensions of A and B must match for MatMul16 transposeA");
-    } else if (s) {
-      if (this.outputShape = [e, t, o / 2], this.dimInner = i, i !== a)
-        throw new Error("Inner dimensions of A and B must match for MatMul16 transposeB");
-    } else if (this.outputShape = [e, t, i / 2], this.dimInner = a, a !== o)
-      throw new Error("Inner dimensions of A and B must match for MatMul16");
-    if (this.dimInner % this.tileInner !== 0)
-      throw new 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)),
-      // 4 unpacked cols per thread = 2 packed cols
-      Math.ceil(this.outputShape[1] / (this.workgroupSize[1] * 4)),
-      // 4 rows per thread
-      this.outputShape[0]
-    ], i % 32 !== 0)
-      throw new Error("Head size must be even for MatMul16 transposeB");
-    if (a % 32 !== 0)
-      throw new Error("Head size must be even for MatMul16 transposeB");
-    if (t % 32 !== 0)
-      throw new Error("Sequence length must be multiple of 32 for MatMul16 transposeB");
-    if (o % 32 !== 0)
-      throw new 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;
-  }
-  /* Note: this is done after constructor because it shouldn't affect dispatch */
-  setOutputShape(e, t) {
-    const o = f(e), a = f(this.outputShape);
-    if (o !== a)
-      throw new Error(`New shape size ${o} must match current size ${a}`);
-    function i(c, u) {
-      return [`${c} / ${u}`, `${c} % ${u}`];
-    }
-    const r = this.outputShape;
-    let s = [];
-    if (e.length === r.length + 1)
-      if (e[0] * e[1] === r[0])
-        s = [
-          ...i("batch", e[1]),
-          // batch / B2, batch % B2
-          "gRow",
-          "gColPacked"
-        ], this.shaderKey += `_batchSplit_${e[1]}`;
-      else if (e[e.length - 2] * e[e.length - 1] === r[r.length - 1])
-        s = [
-          "batch",
-          "gRow",
-          ...i("gColPacked", e[e.length - 1])
-          // gColPacked / N2, gColPacked % N2
-        ], this.shaderKey += `_colSplit_${e[e.length - 1]}`;
-      else
-        throw new Error("Unsupported output shape split");
-    else if (e.length === r.length)
-      s = ["batch", "gRow", "gColPacked"];
-    else if (e.length === 2 && r[0] === 1)
-      s = ["gRow", "gColPacked"], this.shaderKey += "_batchRemoved";
-    else
-      throw new Error(`Unsupported output shape rank change: ${r.length} -> ${e.length}}`);
-    let n = [];
-    if (t) {
-      if (t.length !== e.length)
-        throw new Error("Permutation length must match output rank");
-      n = t.map((c) => s[c]), this.shaderKey += `_perm_${t.join("")}`;
-    } else
-      n = s;
-    const l = n.findIndex((c) => c === "gRow"), h = `vec${e.length}<i32>(${n.join(", ")})`;
-    this.outputIndexSnippet = `var idx0: i32 = getOutputIndexFromCoords(${h});`, this.outputStrideSnippet = `idx0 = idx0 + uniforms.outShapeStrides${l === 0 ? "" : `[${l}]`}; `, t ? this.outputShape = t.map((c) => e[c]) : this.outputShape = e;
-  }
-  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" ? `
-                // TODO: revisit after https://github.com/gpuweb/gpuweb/issues/4458 is resolved
-                fn tanhComplete(x: vec4<f32>) -> vec4<f32> {
-                    return vec4<f32>(
-                        select(tanh(x.x), sign(x.x), abs(x.x) > 15.0f),
-                        select(tanh(x.y), sign(x.y), abs(x.y) > 15.0f),
-                        select(tanh(x.z), sign(x.z), abs(x.z) > 15.0f),
-                        select(tanh(x.w), sign(x.w), abs(x.w) > 15.0f),
-                    );
-                }
-                fn activation(x : vec4<f32>) -> vec4<f32> {
-                    let x3 = x * x * x;
-                    var inner = fma(vec4<f32>(${0.044715}f), x3, x);
-                    inner = ${0.7978845608028654}f * inner;
-                    inner = tanhComplete(inner);
-                    inner = 0.5f * (1.0f + inner);
-                    return x * inner;
-                }
-                ` : this.activation === "relu2" ? `
-                fn activation(x : vec4<f32>) -> vec4<f32> {
-                    let y = max(x, vec4<f32>(0.0));
-                    return y * y;
-                }
-                ` : this.activation === "relu" ? `
-                fn activation(x : vec4<f32>) -> vec4<f32> {
-                    return max(x, vec4<f32>(0.0));
-                }
-                ` : "";
-  }
-  /* Transpose when writing to shared memory */
-  readASnippet() {
-    const e = `
+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];
@@ -172,7 +142,7 @@ class B {
             ${this.scaleA ? "Arow3 = Arow3 * uniforms.scaleA;" : ""}
             ${this.scaleA ? "Arow4 = Arow4 * uniforms.scaleA;" : ""}
         `;
-    return this.transposeA ? `{
+		return this.transposeA ? `{
                 ${e}
                 mm_Asub[row][col] = Arow1;
                 mm_Asub[row + 1][col] = Arow2;
@@ -189,10 +159,9 @@ class B {
                 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);
         }`;
-  }
-  /* Transpose when writing to shared memory */
-  readBSnippet() {
-    const e = `
+	}
+	readBSnippet() {
+		let e = `
             var col = i32(localId.x);
             var row = i32(localId.y) * 4;
             var packedB: vec2<i32> = B[offsetB + row * strideB + col];
@@ -221,7 +190,7 @@ class B {
             ${this.scaleB ? "Brow3 = Brow3 * uniforms.scaleB;" : ""}
             ${this.scaleB ? "Brow4 = Brow4 * uniforms.scaleB;" : ""}
         `;
-    return this.transposeB ? `{
+		return this.transposeB ? `{
                 ${e}
                 
                 col = i32(localId.x) * 4;
@@ -238,47 +207,46 @@ class B {
                 mm_Bsub[row + 2][col] = Brow3;
                 mm_Bsub[row + 3][col] = Brow4;
             }`;
-  }
-  baseIndexSnippets() {
-    const e = `
+	}
+	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;
-        `;
-    let t = "";
-    this.transposeB ? t = "let baseB = getIndexFromCoords3D(vec3<i32>(batchB, globalColStart, 0), vec3<i32>(uniforms.bShape.x, uniforms.bShape.y, strideB));" : t = "let baseB = getIndexFromCoords3D(vec3<i32>(batchB, 0, globalColStart / 4), vec3<i32>(uniforms.bShape.x, uniforms.bShape.y, strideB));";
-    let o = "";
-    return this.transposeA ? o = "let baseA = getIndexFromCoords3D(vec3<i32>(batchA, 0, globalRowStart / 4), vec3<i32>(uniforms.aShape.x, uniforms.aShape.y, strideA));" : o = "let baseA = getIndexFromCoords3D(vec3<i32>(batchA, globalRowStart, 0), vec3<i32>(uniforms.aShape.x, uniforms.aShape.y, strideA));", `
-            ${e}
-            ${o}
+        
             ${t}
+            ${e}
         `;
-  }
-  offsetSnippets() {
-    let e = "";
-    this.transposeA ? e = "let offsetA = baseA + kStart * strideA;" : e = "let offsetA = baseA + kStart / 4;";
-    let t = "";
-    return this.transposeB ? t = "let offsetB = baseB + kStart / 4;" : t = "let offsetB = baseB + kStart * strideB;", `
+	}
+	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() {
-    const e = this.transposeA, t = this.tileInner, o = this.workgroupSize[1] * 4, a = this.workgroupSize[0] * 4, i = e ? o : t, r = e ? t : o, s = this.dimInner, n = Math.ceil(s / t);
-    return `
-            var<workgroup> mm_Asub : array<array<vec4<f32>, ${i / 4 + (this.transposeA ? 0 : 1)}>, ${r}>;
-            var<workgroup> mm_Bsub : array<array<vec4<f32>, ${a / 4 + (this.transposeB ? 1 : 0)}>, ${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() : ""}
 
-            ${p()} {
+            ${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) * ${o};
-                let globalColStart = i32(workgroupId.x) * ${a};
+                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)
@@ -288,16 +256,16 @@ class B {
 
                 ${this.baseIndexSnippets()}
 
-                for (var t = 0; t < ${n}; t++) {
+                for (var t = 0; t < ${c}; t++) {
                     ${this.offsetSnippets()}
 
                     ${this.readASnippet()}
                     ${this.readBSnippet()}
 
-                    kStart = kStart + ${t};
+                    kStart = kStart + ${n};
                     workgroupBarrier();
 
-                    for (var k = 0; k < ${t}; k++) {
+                    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];
@@ -316,14 +284,7 @@ class B {
                 
                 ${this.outputIndexSnippet}
                 for (var i = 0; i < 4; i = i + 1) {
-                    ${this.causalMask ? `
-                    // Causal Masking: mask if col > row + pastLen
-                    let r = gRow + i;
-                    let cBase = gColPacked * 2;
-                    let cVec = vec4<i32>(cBase, cBase + 1, cBase + 2, cBase + 3);
-                    let mask = cVec > vec4<i32>(r + uniforms.pastLen);
-                    acc[i] = select(acc[i], vec4<f32>(-uniforms.INFINITY), mask);
-                    ` : ""}
+                    ${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;" : ""}
@@ -336,8 +297,7 @@ class B {
                 }
             }
         `;
-  }
-}
-export {
-  B as default
+	}
 };
+//#endregion
+export { n as default };

package/dist/ops/webgpu/mul16.js

@@ -1,13 +1,14 @@
-import { c as t } from "../../index-CUXkjxiT.js";
-import { BinaryOpScalarProgram as s, BinaryOpProgram as c } from "./utils/binary_op.js";
-import { B as a } from "../../binary_op_util-pKXltfxI.js";
-function m(n) {
-  const { a: e, b: r } = n.inputs, o = n.backend, p = r.shape.length === 0 ? new s(a.MUL, e.shape) : new c(a.MUL, e.shape, r.shape);
-  return o.runWebGPUProgram(p, [e, r], "packedF16");
+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");
 }
-const i = {
-  kernelName: "Mul16",
-  backendName: "webgpu",
-  kernelFunc: m
-};
-t(i);
+e({
+	kernelName: "Mul16",
+	backendName: "webgpu",
+	kernelFunc: i
+});
+//#endregion