@genai-fi/nanogpt 0.4.1 → 0.4.3

package/dist/ops/webgl/appendCache.js

@@ -1,12 +1,12 @@
-import { r as h } from "../../index-XjBAhiFO.js";
+import { r as p } from "../../index-C4JCoBvj.js";
 class m {
   variableNames = ["cache", "item"];
   outputShape;
   userCode;
   customUniforms = [{ name: "cacheT", type: "int" }];
-  constructor(t, a, o, s, n) {
-    const c = Math.min(o + 1, n);
-    this.outputShape = [t, a, c, s], this.userCode = `
+  constructor(t, a, n, o, c) {
+    const s = Math.min(n + 1, c);
+    this.outputShape = [t, a, s, o], this.userCode = `
         void main() {
             ivec4 coords = getOutputCoords(); // [b, h, t, d]
             int b = coords.x;
@@ -15,7 +15,7 @@ class m {
             int d = coords.w;
 
             int itemT = 1;
-            int maxSize = ${n};
+            int maxSize = ${c};
             int totalT = cacheT + itemT;
             int start = totalT >= maxSize ? 1 : 0;
 
@@ -23,21 +23,22 @@ class m {
             float val = 0.0;
             if (srcT < cacheT) {
                 val = getCache(b, h, srcT, d);
-            } else {
+            } else if (srcT == cacheT) {
                 val = getItem(b, h, 0, d);
-            }
+            } else {
+                val = 0.0;}
             setOutput(val);
         }
         `;
   }
 }
-function p(e) {
-  const { cache: t, item: a } = e.inputs, { maxSize: o } = e.attrs, s = e.backend, n = t.shape[0], c = t.shape[2], r = t.shape[1], i = new m(n, r, c, a.shape[3], o);
-  return s.runWebGLProgram(i, [t, a], "float32", [[c]]);
+function d(e) {
+  const { cache: t, item: a } = e.inputs, { maxSize: n, pastLen: o } = e.attrs, c = e.backend, s = t.shape[0], r = t.shape[2], i = t.shape[1], h = new m(s, i, r, a.shape[3], n);
+  return c.runWebGLProgram(h, [t, a], "float32", [[o]]);
 }
-const d = {
+const l = {
   kernelName: "AppendCache",
   backendName: "webgl",
-  kernelFunc: p
+  kernelFunc: d
 };
-h(d);
+p(l);

package/dist/ops/webgl/attentionMask.js

@@ -1,11 +1,14 @@
-import { r as m } from "../../index-XjBAhiFO.js";
-class k {
-  variableNames = ["q", "k", "mask"];
+import { r as h } from "../../index-C4JCoBvj.js";
+class l {
+  variableNames = ["q", "k"];
   outputShape;
   userCode;
-  customUniforms = [{ name: "divisor", type: "float" }];
-  constructor(t, a, o, e) {
-    this.outputShape = [t, a, o, o], this.userCode = `
+  customUniforms = [
+    { name: "divisor", type: "float" },
+    { name: "pastLen", type: "int" }
+  ];
+  constructor(t, s, e, n, a) {
+    this.outputShape = [t, s, e, n], this.userCode = `
         void main() {
             ivec4 coords = getOutputCoords(); // [batch, nh, t1, t2]
             int b = coords.x;
@@ -14,30 +17,28 @@ class k {
             int t2 = coords.w;
 
             float sum = 0.0;
-            for (int i = 0; i < ${e}; ++i) {
+            for (int i = 0; i < ${a}; ++i) {
                 float qv = getQ(b, h, t1, i);
-                float kv = getK(b, h, t2, i); // k is transposed on last two dims
+                float kv = getK(b, h, t2, i);
                 sum += qv * kv;
             }
 
             // Scale by divisor
             float scaled = sum * divisor;
 
-            // Add mask
-            float maskVal = getMask(t1, t2); // mask is [T,T]
-
-            setOutput(scaled + maskVal);
+            // Mask out future positions
+            setOutput((t2 > t1 + pastLen) ? -1.0/0.0 : scaled);
         }
         `;
   }
 }
-function l(s) {
-  const { q: t, k: a, mask: o } = s.inputs, { divisor: e } = s.attrs, n = s.backend, i = t.shape[0], r = t.shape[2], c = t.shape[1], d = new k(i, c, r, t.shape[3]);
-  return n.runWebGLProgram(d, [t, a, o], "float32", [[e]]);
+function m(o) {
+  const { q: t, k: s } = o.inputs, { divisor: e, pastLen: n } = o.attrs, a = o.backend, i = t.shape[0], r = t.shape[2], c = s.shape[2], u = t.shape[1], p = t.shape[3], d = new l(i, u, r, c, p);
+  return a.runWebGLProgram(d, [t, s], "float32", [[e], [n]]);
 }
-const u = {
+const k = {
   kernelName: "AttentionMask",
   backendName: "webgl",
-  kernelFunc: l
+  kernelFunc: m
 };
-m(u);
+h(k);