@genai-fi/nanogpt 0.9.0 → 0.10.0

package/dist/ops/webgpu/matMul16_program.js

@@ -0,0 +1,336 @@
+import "../../index-ZyQhjEPo.js";
+import { e as h } from "../../webgpu_program-Cigz-7RF.js";
+import { s as f } from "../../tensor-DdQUJZlz.js";
+class A {
+  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"), p = `vec${e.length}<i32>(${n.join(", ")})`;
+    this.outputIndexSnippet = `var idx0: i32 = getOutputIndexFromCoords(${p});`, 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;
+                }
+                ` : "";
+  }
+  /* Transpose when writing to shared memory */
+  readASnippet() {
+    const e = `
+            var col = i32(localId.x);
+            var row = i32(localId.y) * 4;
+            var packedA: vec2<i32> = A[offsetA + row * strideA + col];
+            var Arow1 = vec4<f32>(
+                unpack2x16float(u32(packedA.x)),
+                unpack2x16float(u32(packedA.y))
+            );
+            packedA = A[offsetA + (row + 1) * strideA + col];
+            var Arow2 = vec4<f32>(
+                unpack2x16float(u32(packedA.x)),
+                unpack2x16float(u32(packedA.y))
+            );
+            packedA = A[offsetA + (row + 2) * strideA + col];
+            var Arow3 = vec4<f32>(
+                unpack2x16float(u32(packedA.x)),
+                unpack2x16float(u32(packedA.y))
+            );
+            packedA = A[offsetA + (row + 3) * strideA + col];
+            var Arow4 = vec4<f32>(
+                unpack2x16float(u32(packedA.x)),
+                unpack2x16float(u32(packedA.y))
+            );
+            
+            ${this.scaleA ? "Arow1 = Arow1 * uniforms.scaleA;" : ""}
+            ${this.scaleA ? "Arow2 = Arow2 * uniforms.scaleA;" : ""}
+            ${this.scaleA ? "Arow3 = Arow3 * uniforms.scaleA;" : ""}
+            ${this.scaleA ? "Arow4 = Arow4 * uniforms.scaleA;" : ""}
+        `;
+    return this.transposeA ? `{
+                ${e}
+                mm_Asub[row][col] = Arow1;
+                mm_Asub[row + 1][col] = Arow2;
+                mm_Asub[row + 2][col] = Arow3;
+                mm_Asub[row + 3][col] = Arow4;
+        }` : `{
+                ${e}
+                
+                col = i32(localId.x) * 4;
+                row = i32(localId.y);
+                
+                mm_Asub[col][row] = vec4<f32>(Arow1.x, Arow2.x, Arow3.x, Arow4.x);
+                mm_Asub[col + 1][row] = vec4<f32>(Arow1.y, Arow2.y, Arow3.y, Arow4.y);
+                mm_Asub[col + 2][row] = vec4<f32>(Arow1.z, Arow2.z, Arow3.z, Arow4.z);
+                mm_Asub[col + 3][row] = vec4<f32>(Arow1.w, Arow2.w, Arow3.w, Arow4.w);
+        }`;
+  }
+  /* Transpose when writing to shared memory */
+  readBSnippet() {
+    const e = `
+            var col = i32(localId.x);
+            var row = i32(localId.y) * 4;
+            var packedB: vec2<i32> = B[offsetB + row * strideB + col];
+            var Brow1 = vec4<f32>(
+                unpack2x16float(u32(packedB.x)),
+                unpack2x16float(u32(packedB.y))
+            );
+            packedB = B[offsetB + (row + 1) * strideB + col];
+            var Brow2 = vec4<f32>(
+                unpack2x16float(u32(packedB.x)),
+                unpack2x16float(u32(packedB.y))
+            );
+            packedB = B[offsetB + (row + 2) * strideB + col];
+            var Brow3 = vec4<f32>(
+                unpack2x16float(u32(packedB.x)),
+                unpack2x16float(u32(packedB.y))
+            );
+            packedB = B[offsetB + (row + 3) * strideB + col];
+            var Brow4 = vec4<f32>(
+                unpack2x16float(u32(packedB.x)),
+                unpack2x16float(u32(packedB.y))
+            );
+            
+            ${this.scaleB ? "Brow1 = Brow1 * uniforms.scaleB;" : ""}
+            ${this.scaleB ? "Brow2 = Brow2 * uniforms.scaleB;" : ""}
+            ${this.scaleB ? "Brow3 = Brow3 * uniforms.scaleB;" : ""}
+            ${this.scaleB ? "Brow4 = Brow4 * uniforms.scaleB;" : ""}
+        `;
+    return this.transposeB ? `{
+                ${e}
+                
+                col = i32(localId.x) * 4;
+                row = i32(localId.y);
+                
+                mm_Bsub[col][row] = vec4<f32>(Brow1.x, Brow2.x, Brow3.x, Brow4.x);
+                mm_Bsub[col + 1][row] = vec4<f32>(Brow1.y, Brow2.y, Brow3.y, Brow4.y);
+                mm_Bsub[col + 2][row] = vec4<f32>(Brow1.z, Brow2.z, Brow3.z, Brow4.z);
+                mm_Bsub[col + 3][row] = vec4<f32>(Brow1.w, Brow2.w, Brow3.w, Brow4.w);
+            }` : `{
+                ${e}
+                mm_Bsub[row][col] = Brow1;
+                mm_Bsub[row + 1][col] = Brow2;
+                mm_Bsub[row + 2][col] = Brow3;
+                mm_Bsub[row + 3][col] = Brow4;
+            }`;
+  }
+  baseIndexSnippets() {
+    const e = `
+            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}
+        `;
+  }
+  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;", `
+            ${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}>;
+
+            ${this.activation ? this.activationSnippet() : ""}
+
+            ${h()} {
+                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};
+
+                // 4 rows x 4 cols accumulator
+                // acc[i] holds row i (4 cols)
+                var acc = array<vec4<f32>, 4>(
+                    vec4<f32>(0.0), vec4<f32>(0.0), vec4<f32>(0.0), vec4<f32>(0.0)
+                );
+
+                ${this.baseIndexSnippets()}
+
+                for (var t = 0; t < ${n}; t++) {
+                    ${this.offsetSnippets()}
+
+                    ${this.readASnippet()}
+                    ${this.readBSnippet()}
+
+                    kStart = kStart + ${t};
+                    workgroupBarrier();
+
+                    for (var k = 0; k < ${t}; k++) {
+                        // Load 4 columns of B as a vec4
+                        let bVec = mm_Bsub[k][localCol];
+                        let aVec = mm_Asub[k][localRow];
+
+                        // Compute 4 rows
+                        for (var r = 0; r < 4; r = r + 1) {
+                            acc[r] = fma(vec4<f32>(aVec[r]), bVec, acc[r]);
+                        }
+                    }
+                    workgroupBarrier();
+                }
+
+                // Write out 4 rows x 2 packed cols (4 unpacked cols)
+                let gRow = globalRowStart + localRow * 4;
+                let gColPacked = i32(workgroupId.x) * ${this.workgroupSize[0] * 2} + localCol * 2;
+                
+                ${this.outputIndexSnippet}
+                for (var i = 0; i < 4; i = i + 1) {
+                    ${this.scale ? "acc[i] = acc[i] * uniforms.scale;" : ""}
+
+                    ${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.activation ? "acc[i] = activation(acc[i]);" : ""}
+                    result[idx0 / 2] = vec2<i32>(
+                        i32(pack2x16float(acc[i].xy)),
+                        i32(pack2x16float(acc[i].zw))
+                    );
+
+                    ${this.outputStrideSnippet}
+                }
+            }
+        `;
+  }
+}
+export {
+  A as default
+};

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

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

package/dist/ops/webgpu/mul16.js

@@ -0,0 +1,14 @@
+import "../../index-ZyQhjEPo.js";
+import { BinaryOpProgram as m } from "./utils/binary_op.js";
+import { B as p } from "../../binary_op_util-pKXltfxI.js";
+import { r as c } from "../../tensor_util-DV-FP5Q3.js";
+function i(r) {
+  const { a: e, b: n } = r.inputs, t = r.backend, a = new m(p.MUL, e.shape, n.shape), o = t.runWebGPUProgram(a, [e, n], "int32");
+  return o.packed = !0, o;
+}
+const s = {
+  kernelName: "Mul16",
+  backendName: "webgpu",
+  kernelFunc: i
+};
+c(s);

package/dist/ops/webgpu/normRMS.js

@@ -1,49 +1,30 @@
-import { f as n } from "../../webgpu_util-pLEV9tks.js";
-import { f as p, a4 as h } from "../../index-BzFyqcy-.js";
-import { createReduceInfo as u, reduce as c, createReductionShader as m } from "./utils/reductions.js";
-class d {
-  outputShape;
-  shaderKey = "RMSNorm";
-  dispatchLayout;
-  dispatch;
-  workgroupSize = [64, 1, 1];
-  variableNames = ["x", "gamma"];
-  uniforms = "reduceSize : i32,";
-  inputShape;
-  size = !0;
-  constructor(e) {
-    this.inputShape = [e.batchSize, e.inSize], this.outputShape = this.inputShape, this.dispatchLayout = n(this.outputShape), this.dispatch = [e.batchSize, 1, 1];
-  }
-  getUserCode() {
-    const e = this.workgroupSize[0];
-    return m(e, "mean", `
-            candidate = candidate * candidate;
-        `, `
-            bestValue = inverseSqrt(bestValue + 1e-8);
-        `, `
-            let X = f32(x[offset + k]);
-            let gamma = gamma[k];
-            let normalized = X * bestValue;
-            let outVal = normalized * gamma;
-            result[offset + k] = f32(outVal);
-        `);
-  }
-}
-function S(a) {
-  const { x: e, gamma: r } = a.inputs, o = a.backend, s = [e, r], t = u(s, -1), i = new d(t);
-  if (h(r.shape, [e.shape[e.shape.length - 1]], "Error in RMSNorm: "), e.shape.length !== 3)
+import "../../index-ZyQhjEPo.js";
+import { createReduceInfo as g, reduce as l } from "./utils/reductions.js";
+import { j as w } from "../../tensor-DdQUJZlz.js";
+import { isPackedTensor as f } from "../../utilities/packed.js";
+import { p as k } from "../../pack16-CFUqumar.js";
+import S from "./normRMS16_program.js";
+import z from "./normRMS32_program.js";
+import N from "./utils/deviceInfo.js";
+import { r as b } from "../../tensor_util-DV-FP5Q3.js";
+function P(m) {
+  const { x: e, gamma: n } = m.inputs, c = m.backend, i = N(c), s = f(e), a = f(n), o = s || a, r = !o || s ? e : k(e), p = !o || a ? n : k(n), h = [r, p], t = g(h, -1), u = o ? new S(i, t) : new z(i, t);
+  if (w(p.shape, [r.shape[r.shape.length - 1]], "Error in RMSNorm: "), e.shape.length !== 3)
     throw new Error(`rmsNormGPU: input rank ${e.shape.length} not supported, only rank 3 is supported`);
-  if (t.inSize !== e.shape[2])
-    throw new Error(`rmsNormGPU: reduction size ${t.inSize} does not match expected size ${e.shape[2]}`);
+  if (t.inSize !== r.shape[r.shape.length - 1])
+    throw new Error(
+      `rmsNormGPU: reduction size ${t.inSize} does not match expected size ${r.shape[r.shape.length - 1]}`
+    );
   if (t.batchSize !== e.shape[0] * e.shape[1])
     throw new Error(
       `rmsNormGPU: batch size ${t.batchSize} does not match expected size ${e.shape[0] * e.shape[1]}`
     );
-  return c(i, s, o);
+  const d = l(u, h, c);
+  return d.packed = o, o && !s && r.dispose(), o && !a && p.dispose(), d;
 }
-const l = {
+const G = {
   kernelName: "RMSNorm",
   backendName: "webgpu",
-  kernelFunc: S
+  kernelFunc: P
 };
-p(l);
+b(G);

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

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

package/dist/ops/webgpu/normRMS16_program.js

@@ -0,0 +1,24 @@
+import { ReduceProgram as a } from "./utils/reductions.js";
+class o extends a {
+  constructor(e, t) {
+    super(e, t, { reductionOp: "mean", elementwise: !0 }, !0), this.shaderKey = "RMSNorm16", this.variableNames.push("gamma"), this.variableComponents = [1, 1];
+  }
+  getPreprocessSnippet() {
+    return "candidate = candidate * candidate;";
+  }
+  getPostprocessSnippet() {
+    return "bestValue = inverseSqrt(bestValue + 1e-8);";
+  }
+  getWriteSnippet() {
+    return `
+            let X = unpack2x16float(u32(x[offset + k]));
+            let gamma = unpack2x16float(u32(gamma[k]));
+            let normalized = X * bestValue;
+            let outVal = normalized * gamma;
+            result[offset + k] = i32(pack2x16float(outVal));
+        `;
+  }
+}
+export {
+  o as default
+};

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

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

package/dist/ops/webgpu/normRMS32_program.js

@@ -0,0 +1,24 @@
+import { ReduceProgram as a } from "./utils/reductions.js";
+class o extends a {
+  constructor(e, t) {
+    super(e, t, { reductionOp: "mean", elementwise: !0 }, !1), this.shaderKey = "RMSNorm32", this.variableNames.push("gamma"), this.variableComponents = [1, 1];
+  }
+  getPreprocessSnippet() {
+    return "candidate = candidate * candidate;";
+  }
+  getPostprocessSnippet() {
+    return "bestValue = inverseSqrt(bestValue + 1e-8);";
+  }
+  getWriteSnippet() {
+    return `
+            let X = f32(x[offset + k]);
+            let gamma = gamma[k];
+            let normalized = X * bestValue;
+            let outVal = normalized * gamma;
+            result[offset + k] = f32(outVal);
+        `;
+  }
+}
+export {
+  o as default
+};