@mineproness/tensorloop 2.1.0 → 2.1.2

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@mineproness/tensorloop",
-  "version": "2.1.0",
+  "version": "2.1.2",
   "description": "Raw-JAVASCRIPT Based Matirx and AI LLM Package called Tensorloop.",
   "keywords": [
     "ai",

package/src/SelfAttention/index.js

@@ -1,4 +1,5 @@
 import { Softmax } from '../../index.js';
+import { GPUAttention } from '../gpiInit.js';
 import { Linear } from '../Linear/index.js'
 
 export class SelfAttention {
@@ -13,6 +14,7 @@ export class SelfAttention {
     score;
     softmax;
     weights;
+    gpu;
     constructor(embeddingSize, configs) {
         this.configs = configs
         this.softmax = new Softmax()
@@ -22,96 +24,29 @@ export class SelfAttention {
         this.value = new Linear(embeddingSize, embeddingSize, { save: { filename: [`${configs.dirname}/value.bin`, `${configs.dirname}/valueBias.bin`] } })
     }
     forward(vectors) {
+        if(!this.gpu) this.gpu = new GPUAttention(vectors[0].length, vectors.length);
         this.q = this.query.forward(vectors)
         this.v = this.value.forward(vectors)
         this.k = this.key.forward(vectors)
-        const attentionSocres = []
-        for (let i = 0; i < this.q.length; i++) {
-            attentionSocres[i] = new Float32Array(this.k.length);
-
-            for (let j = 0; j < this.k.length; j++) {
-                if (j > i) {
-                    attentionSocres[i][j] = -1e9; // mask future
-                    continue;
-                }
-
-                let dot = 0;
-
-                for (let d = 0; d < this.q[i].length; d++) {
-                    dot += this.q[i][d] * this.k[j][d];
-                }
-
-                attentionSocres[i][j] = dot / this.score;
-            }
-        }
-
+        const attentionSocres = this.gpu.attentionScores(this.q, this.k, this.score);
         const weights = this.softmax.forward(attentionSocres)
         this.weights = weights
-        const output = [];
-
-        for (let i = 0; i < weights.length; i++) {
-            output[i] = new Float32Array(this.v[0].length);
-
-            for (let j = 0; j < this.v.length; j++) {
-                for (let d = 0; d < this.v[j].length; d++) {
-                    output[i][d] += weights[i][j] * this.v[j][d];
-                }
-            }
-        }
-
+        const output = this.gpu.attentionOutput(weights, this.v)
         return output;
     }
-    softmaxBack(probs, dOut) {
-        const dIn = new Float32Array(probs.length);
 
-        for (let i = 0; i < probs.length; i++) {
-            let sum = 0;
-            for (let j = 0; j < probs.length; j++) {
-                const jacobian =
-                    i === j
-                        ? probs[i] * (1 - probs[i])
-                        : -probs[i] * probs[j];
-
-                sum += jacobian * dOut[j];
-            }
-            dIn[i] = sum;
-        }
-
-        return dIn;
-    }
     backward(dOutput, learningRate) {
-        const dV = this.v.map((e) => new Float32Array(e.length));
-        const dWeights = []
-        for (let a = 0; a < this.weights.length; a++) {
-            dWeights[a] = new Float32Array(this.weights[a].length)
-            for (let j = 0; j < this.v.length; j++) {
-                for (let r = 0; r < this.v[j].length; r++) {
-                    dWeights[a][j] += dOutput[a][r] * this.v[j][r]
-                    dV[j][r] += dOutput[a][r] * this.weights[a][j]
-                }
-            }
-        }
-        const dScores = []
-        for (let t = 0; t < dWeights.length; t++) {
-            dScores[t] = this.softmaxBack(this.weights[t], dWeights[t])
-        }
-        const dQ = this.q.map(q => new Float32Array(q.length));
-        const dK = this.k.map(k => new Float32Array(k.length));
-        for (let a = 0; a < this.q.length; a++) {
-            for (let j = 0; j < this.k.length; j++) {
-                if (j > a) continue;
-                for (let r = 0; r < this.q[a].length; r++) {
-                    dQ[a][r] += dScores[a][j] * this.k[j][r] / this.score;
-                    dK[j][r] += dScores[a][j] * this.q[a][r] / this.score;
-                }
-            }
-        }
+        const dWeights = this.gpu.dWeightsKernel(dOutput, this.v)
+        const dV = this.gpu.dVKernel(dOutput, this.weights);
+        const dScores = this.gpu.softmaxBackKernel(this.weights, dWeights)
+        const dQ = this.gpu.dQKernel(dScores, this.k, this.score)
+        const dK = this.gpu.dKKernel(dScores, this.q, this.score)
         const dInput = this.q.map(() => new Float32Array(this.q[0].length))
         const dq = this.query.backward(dQ, learningRate)
         const dk = this.key.backward(dK, learningRate)
         const dv = this.value.backward(dV, learningRate)
         for (let a = 0; a < this.q.length; a++) {
-            for (let b = 0; b < dq.length; b++) {
+            for (let b = 0; b < dq[a].length; b++) {
                 dInput[a][b] = dq[a][b] + dk[a][b] + dv[a][b]
             }
         }

package/src/gpiInit.js

@@ -261,12 +261,14 @@ export class GPUSoftmax {
         const maxLogits = this.cacheMaxCache(logits)
         const sumes = this.cacheSumCalculate(maxLogits, logits)
         const splitedLogits = split2DByColumns(logits)
-        const result = Array.from({ length: this.contextSize }, () => [])
+        const result = Array.from({ length: this.contextSize }, () => new Float32Array(logits[0].length))
+        let offset = 0
         for (let i = 0; i < splitedLogits.length; i++) {
             const out = this.cacheSF[i](splitedLogits[i], sumes, maxLogits)
             for (let a = 0; a < this.contextSize; a++) {
-                result[a].push(...out[a])
+                result[a].set(out[a] , offset)
             }
+            offset += out[0].length
         }
         maxLogits.delete()
         sumes.delete()
@@ -274,12 +276,14 @@ export class GPUSoftmax {
     }
     backward(probs, target) {
         const splitedProbs = split2DByColumns(probs)
-        const result = Array.from({ length: this.contextSize }, () => [])
+        const result = Array.from({ length: this.contextSize }, () => new Float32Array(probs[0].length))
+        let offset = 0;
         for (let i = 0; i < this.cacheSB.length; i++) {
             const out = this.cacheSB[i](splitedProbs[i], target)
             for (let a = 0; a < this.contextSize; a++) {
-                result[a].push(...out[a])
+                result[a].set(out[a], offset)
             }
+            offset += out[0].length
         }
         return result
     }
@@ -601,4 +605,121 @@ export class GELUGPU {
     backward(dOutput, input) {
         return this.cacheGB(input, dOutput)
     }
+}
+
+export class GPUAttention {
+    attentionScores;
+    attentionOutput;
+    dWeightsKernel;
+    dVKernel;
+    softmaxBackKernel;
+    dQKernel;
+    dKKernel;
+    constructor(embeddingSize, contextSize) {
+        this.attentionScores = gpu.createKernel(function (q, k, scale) {
+            const i = this.thread.y;
+            const j = this.thread.x;
+            if (j > i) {
+                return -1e9;
+            }
+            let sum = 0;
+            for (let d = 0; d < this.constants.D; d++) {
+                sum += q[i][d] * k[j][d];
+            }
+            return sum / scale;
+        }, {
+            output: [contextSize, contextSize],
+            constants: {
+                D: embeddingSize
+            }
+        });
+        this.attentionOutput = gpu.createKernel(function (weights, v) {
+            const i = this.thread.y;
+            const d = this.thread.x;
+            let sum = 0;
+            for (let j = 0; j < this.constants.CONTEXT; j++) {
+                sum += weights[i][j] * v[j][d];
+            }
+            return sum;
+        }, {
+            output: [embeddingSize, contextSize],
+            constants: {
+                CONTEXT: contextSize
+            },
+        });
+        this.dWeightsKernel = gpu.createKernel(function (dOutput, v) {
+            const a = this.thread.y;
+            const j = this.thread.x;
+            let sum = 0;
+            for (let r = 0; r < this.constants.D; r++) {
+                sum += dOutput[a][r] * v[j][r];
+            }
+            return sum;
+        }, {
+            output: [contextSize, contextSize],
+            pipeline: true,
+            immutable: true,
+            constants: {
+                D: embeddingSize
+            }
+        });
+        this.dVKernel = gpu.createKernel(function (dOutput, weights) {
+            const j = this.thread.y;
+            const r = this.thread.x;
+            let sum = 0;
+            for (let a = 0; a < this.constants.C; a++) {
+                sum += weights[a][j] * dOutput[a][r];
+            }
+            return sum;
+        }, {
+            output: [embeddingSize, contextSize],
+
+            constants: {
+                C: contextSize
+            }
+        });
+        this.softmaxBackKernel = gpu.createKernel(function (probs, dOut) {
+            const a = this.thread.y;
+            const i = this.thread.x;
+            let dot = 0;
+            for (let j = 0; j < this.constants.C; j++) {
+                dot += probs[a][j] * dOut[a][j];
+            }
+            return probs[a][i] * (dOut[a][i] - dot);
+        }, {
+            output: [contextSize, contextSize],
+            pipeline: true,
+            immutable: true,
+            constants: {
+                C: contextSize
+            }
+        });
+        this.dQKernel = gpu.createKernel(function (dScores, k, scale) {
+            const a = this.thread.y;
+            const r = this.thread.x;
+            let sum = 0;
+            for (let j = 0; j <= a; j++) {
+                sum += dScores[a][j] * k[j][r];
+            }
+            return sum / scale;
+        }, {
+            output: [embeddingSize, contextSize],
+
+        });
+        this.dKKernel = gpu.createKernel(function (dScores, q, scale) {
+            const j = this.thread.y;
+            const r = this.thread.x;
+            let sum = 0;
+            for (let a = j; a < this.constants.C; a++) {
+                sum += dScores[a][j] * q[a][r];
+            }
+            return sum / scale;
+        }, {
+            output: [embeddingSize, contextSize],
+
+            constants: {
+                C: contextSize
+            }
+        });
+    }
 }