@genai-fi/nanogpt 0.6.3 → 0.7.0

package/dist/ops/webgpu/rope.js

@@ -0,0 +1,68 @@
+import { f as h, c as l, g as m } from "../../webgpu_util-g13LvDIv.js";
+import { f as g } from "../../index-C0dhsYom.js";
+class f {
+  variableNames = ["x", "sin", "cos"];
+  outputShape;
+  shaderKey = "Rope";
+  dispatchLayout;
+  dispatch;
+  workgroupSize = [64, 1, 1];
+  size = !0;
+  uniforms = "pastLen: i32";
+  constructor(t, o, s, n) {
+    this.outputShape = [t, o, s, n], this.dispatchLayout = h(this.outputShape), this.dispatch = l(this.dispatchLayout, this.outputShape, this.workgroupSize);
+  }
+  getUserCode() {
+    const t = this.outputShape[3];
+    return `
+        ${m("index")} {
+            if (index < uniforms.size) {
+                let coords = getCoordsFromIndex(index); // [b, h, t, d]
+                let b = coords[0];
+                let h = coords[1];
+                let t = coords[2];
+                let d = coords[3];
+
+                let rotaryDim = ${t};
+
+                var outVal = 0.0;
+
+                if (d < rotaryDim) {
+                    let pairIdx = d / 2;
+                    let cos = getCos(t + uniforms.pastLen, pairIdx, 0);
+                    let sin = getSin(t + uniforms.pastLen, pairIdx, 0);
+
+                    let ownX = getX(b, h, t, d) * cos;
+
+                    if (d % 2 == 0) {
+                        // even index
+                        let even = ownX;
+                        let odd = getX(b, h, t, d + 1);
+                        outVal = even - odd * sin;
+                    } else {
+                        // odd index
+                        let even = getX(b, h, t, d - 1);
+                        let odd = ownX;
+                        outVal = even * sin + odd;
+                    }
+                } else {
+                    // pass through for non-rotary dims
+                    outVal = getX(b, h, t, d);
+                }
+
+                setOutputAtIndex(index, outVal);
+            }
+        }
+        `;
+  }
+}
+function b(e) {
+  const { x: t, sin: o, cos: s } = e.inputs, { pastLen: n } = e.attrs, a = e.backend, i = t.shape[0], r = t.shape[1], d = t.shape[2], p = t.shape[3], u = new f(i, r, d, p), c = [{ type: "int32", data: [n] }];
+  return a.runWebGPUProgram(u, [t, o, s], "float32", c);
+}
+const x = {
+  kernelName: "Rope",
+  backendName: "webgpu",
+  kernelFunc: b
+};
+g(x);

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

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

package/dist/ops/webgpu/scatterSub.js

@@ -0,0 +1,37 @@
+import { f as c, c as u, g as d } from "../../webgpu_util-g13LvDIv.js";
+import { f as p } from "../../index-C0dhsYom.js";
+class h {
+  variableNames = ["labels", "softmaxProbs", "dy"];
+  outputShape;
+  shaderKey = "ScatterSub";
+  dispatchLayout;
+  dispatch;
+  workgroupSize = [64, 1, 1];
+  size = !0;
+  constructor(t, e) {
+    this.outputShape = [t, e], this.dispatchLayout = c(this.outputShape), this.dispatch = u(this.dispatchLayout, this.outputShape, this.workgroupSize);
+  }
+  getUserCode() {
+    return `
+            ${d("index")} {
+                if (index < uniforms.size) {
+                    let coords = getCoordsFromIndex(index); // [batch, depth]
+                    let idx = i32(getLabels(coords[0]));
+                    let prob = getSoftmaxProbsByOutputIndex(index);
+                    let dy = getDy(coords[0]);
+                    setOutputAtIndex(index, select(prob, prob - 1.0, idx == coords[1]) * dy);
+                }
+            }
+            `;
+  }
+}
+function b(o) {
+  const { logits: t, labels: e, dy: s } = o.inputs, a = o.backend, r = e.shape[0], i = t.shape[1], n = new h(r, i);
+  return a.runWebGPUProgram(n, [e, t, s], "float32");
+}
+const l = {
+  kernelName: "EfficientScatterSub",
+  backendName: "webgpu",
+  kernelFunc: b
+};
+p(l);