@mineproness/tensorloop 2.1.72 → 2.1.74

package/package.json

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

package/src/RMSNorm/index.js

@@ -49,7 +49,7 @@ export class RMSNorm {
     const output = await this.op.createBuffer(input.size);
     const bindGroup = await this.op.createBindGroup(this.reductionShader, 0, [input, outputReduce])
     const encoder = await this.op.CreateEncoder()
-    await this.op.RunPipline(encoder, this.reductionShader, bindGroup, Math.ceil(input.size / (this.embeddingSize * 4)))
+    await this.op.RunPipline(encoder, this.reductionShader, bindGroup, context)
     const bindGroup1 = await this.op.createBindGroup(this.forwardShader, 0, [outputReduce, input, output, this.gammabuffer])
     await this.op.RunPipline(encoder, this.forwardShader, bindGroup1, Math.ceil((input.size / 4) / 64));
     await this.op.submitQueue(encoder);
@@ -69,7 +69,7 @@ export class RMSNorm {
     const bindGroup2 = await this.op.createBindGroup(this.updateGamma, 0, [this.input, outputGradient, this.rms, this.gammabuffer, params]);
     const encoder = await this.op.CreateEncoder()
     await this.op.RunPipline(encoder, this.backwardShader, bindGroup, Math.ceil((this.input.size) / (this.embeddingSize * 4)))
-    await this.op.RunPipline(encoder, this.updateGamma, bindGroup2, this.embeddingSize)
+    await this.op.RunPipline(encoder, this.updateGamma, bindGroup2, Math.ceil(this.embeddingSize / 64))
     await this.op.submitQueue(encoder)
     if (this.prevOut) this.prevOut.destroy()
     this.prevOut = outputGradient

package/src/Softmax/index.js

@@ -111,7 +111,7 @@ export class Softmax {
                 }
             ]
         })
-        await this.gpu.RunPipline(encoder, this.backwardPipline, bind, Math.ceil((tokenCount * this.vocabSize) / 256))
+        await this.gpu.RunPipline(encoder, this.backwardPipline, bind, tokenCount)
         this.device.queue.submit([encoder.finish()])
         probs.destroy()
         targetBuffer.destroy()

package/src/shaders/Norm/rmsBackward.wgsl

@@ -1,11 +1,11 @@
+
 @group(0) @binding(0)
 var<storage, read> input: array<f32>;
 
 @group(0) @binding(1)
 var<storage, read> gradOutput: array<f32>;
-
 @group(0) @binding(2)
-var<storage, read> rms: array<f32>;
+var<storage, read> sumSquares: array<f32>;
 
 @group(0) @binding(3)
 var<storage, read> gamma: array<f32>;
@@ -15,28 +15,51 @@ var<storage, read_write> gradInput: array<f32>;
 
 var<workgroup> partial: array<f32, 64>;
 
+const D: u32 = %embeddingSize%u;
+const EPSILON: f32 = 1e-6;
 @compute @workgroup_size(64)
 fn main(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wid: vec3<u32>) {
-    let D = %embeddingSize%u;
     let token = wid.x;
-    let idx = token * D + lid;
-    let x = input[idx];
-    let dy = gradOutput[idx];
-    let g = gamma[lid];
-    let r = rms[token];
-    partial[lid] = dy * g * x;
+    if (token >= arrayLength(&sumSquares)) {
+        return;
+    }
+    let base = token * D;
+    let r = sqrt(sumSquares[token] / f32(D) + EPSILON);
+    var localSum = 0.0;
+    var d = lid;
+    while (d < D) {
+        let idx = base + d;
+        localSum += gradOutput[idx] * gamma[d] * input[idx];
+        d += 64u;
+    }
+
+    partial[lid] = localSum;
     workgroupBarrier();
-    var stride = D / 2u;
+    var stride = 32u;
     loop {
         if (stride == 0u) {
             break;
         }
+
         if (lid < stride) {
             partial[lid] += partial[lid + stride];
         }
+
         workgroupBarrier();
         stride /= 2u;
     }
     let dot = partial[0];
-    gradInput[idx] = (dy * g) / r - x * dot / (f32(D) * r * r * r);
+    d = lid;
+    while (d < D) {
+        let idx = base + d;
+        let x = input[idx];
+        let dy = gradOutput[idx];
+        let g = gamma[d];
+
+        gradInput[idx] =
+            (dy * g) / r
+            - (x * dot) / (f32(D) * r * r * r);
+
+        d += 64u;
+    }
 }

package/src/shaders/Norm/sumSquare.wgsl

@@ -1,20 +1,23 @@
+
 @group(0) @binding(0)
 var<storage, read> input: array<f32>;
 
 @group(0) @binding(1)
 var<storage, read_write> out: array<f32>;
 
-
-var<workgroup> partial: array<f32, %embeddingSize%>;
-
-const embeddingSize : u32 = %embeddingSize%u;
+var<workgroup> partial: array<f32, 64>;
+const D: u32 = %embeddingSize%u;
 @compute @workgroup_size(64)
 fn main(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wid: vec3<u32>) {
-    let dim = embeddingSize;
-    let index = wid.x * dim + lid;
+    let base = wid.x * D;
     var sum = 0.0;
-    if (lid < dim && index < arrayLength(&input)) {
-        sum = input[index] * input[index];
+    var d = lid;
+    while (d < D) {
+        let idx = base + d;
+        if (idx < arrayLength(&input)) {
+            sum += input[idx] * input[idx];
+        }
+        d += 64u;
     }
     partial[lid] = sum;
     workgroupBarrier();

package/src/shaders/Norm/updateGamma.wgsl

@@ -1,3 +1,4 @@
+
 @group(0) @binding(0)
 var<storage, read> input: array<f32>;
 
@@ -5,7 +6,7 @@ var<storage, read> input: array<f32>;
 var<storage, read> gradOutput: array<f32>;
 
 @group(0) @binding(2)
-var<storage, read> rms: array<f32>;
+var<storage, read> sumSquares: array<f32>;
 
 @group(0) @binding(3)
 var<storage, read_write> gamma: array<f32>;
@@ -13,49 +14,40 @@ var<storage, read_write> gamma: array<f32>;
 @group(0) @binding(4)
 var<storage, read> params: array<f32>;
 
-var<workgroup> partial: array<f32, %embeddingSize%>;
+var<workgroup> partial: array<f32, 64>;
 
+const D: u32 = %embeddingSize%u;
+const EPSILON: f32 = 1e-6;
 @compute @workgroup_size(64)
 fn main(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wid: vec3<u32>) {
-    let D = %embeddingSize%u;
-    let dim = wid.x;
-
+    let dim = wid.x * 64u + local_invocation_index;
     if (dim >= D) {
         return;
     }
-
-    let T = arrayLength(&rms);
+    let T = arrayLength(&sumSquares);
     let lr = params[0];
-
     var sum = 0.0;
     var token = lid;
-
     while (token < T) {
         let idx = token * D + dim;
-
-        sum += gradOutput[idx] * input[idx] / rms[token];
-
+        let r = sqrt(sumSquares[token] / f32(D) + EPSILON);
+        sum += gradOutput[idx] * input[idx] / r;
         token += 64u;
     }
-
     partial[lid] = sum;
     workgroupBarrier();
-
     var stride = 32u;
     loop {
         if (stride == 0u) {
             break;
         }
-
         if (lid < stride) {
             partial[lid] += partial[lid + stride];
         }
-
         workgroupBarrier();
         stride /= 2u;
     }
     if (lid == 0u) {
-        let gradGamma = partial[0];
-        gamma[dim] -= lr * gradGamma;
+        gamma[dim] -= lr * partial[0];
     }
 }