@imaduddin/fastbrain 0.0.0-stage → 1.0.1

package/src/shaders/PF.wgsl

@@ -0,0 +1,19 @@
+@group(0) @binding(0)
+var <storage, read> posVec: array<f32>;
+
+@group(0) @binding(1)
+
+var <storage, read_write> vecs : array<f32>;
+
+
+
+
+@compute @workgroup_size(64)
+
+fn main(@builtin(global_invocation_id) id : vec3<u32>){
+    let x = id.x;
+    if(x >= arrayLength(&vecs)){
+        return;
+    }
+    vecs[x] = vecs[x] + posVec[x];
+}

package/src/shaders/RB.wgsl

@@ -0,0 +1,37 @@
+struct Param {
+    embeddingSize: u32,
+    tokenCount: u32
+}
+
+@group(0) @binding(0)
+
+var<storage, read> input: array<f32>;
+
+@group(0) @binding(1)
+
+var<uniform> params: Param;
+
+@group(0) @binding(2)
+
+var<storage, read_write> output: vec2<f32>;
+
+@group(0) @binding(3)
+
+var<storage, read_write> inputGradient: array<f32>;
+@compute @workgroup_size(256)
+
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    let index = id.x;
+    let total = params.embeddingSize * params.tokenCount;
+    if (index >= total) {
+        return;
+    }
+    let x = input[index];
+    let grad = inputGradient[index];
+    if (x > 0) {
+        output[index] = grad;
+    }
+    else {
+        output[index] = grad * 0.02;
+    }
+}

package/src/shaders/RF.wgsl

@@ -0,0 +1,33 @@
+struct Param {
+    embeddingSize: u32,
+    tokenCount: u32
+}
+
+@group(0) @binding(0)
+
+var<storage, read> input: array<f32>;
+
+@group(0) @binding(1)
+
+var<uniform> params: Param;
+
+@group(0) @binding(2)
+
+var<storage, read_write> output: array<f32>;
+
+@compute @workgroup_size(256)
+
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    let index = id.x;
+    let total = params.embeddingSize * params.tokenCount;
+    if (index >= total) {
+        return;
+    }
+    let x = input[index];
+    if (x > 0) {
+        output[index] = x;
+    }
+    else {
+        output[index] = x * 0.02;
+    }
+}

package/src/shaders/ReductionFunction.wgsl

@@ -0,0 +1,28 @@
+struct Parmas {
+    vocabSize: u32,
+    count: u32
+}
+
+@group(0) @binding(0)
+
+var<storage, read> data: array<f32>;
+
+@group(0) @binding(1)
+
+var<storage, read_write> result: array<f32>;
+
+@group(0) @binding(2)
+
+var<uniform> params: Parmas;
+
+@compute @workgroup_size(1)
+
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    let index = id.x;
+    var sum = 0.0;
+    for (var i = 0u; i < params.count; i++) {
+        sum += data[i];
+    }
+    sum /= f32(params.count);
+    result[index] = sum;
+}

package/src/shaders/SB.wgsl

@@ -0,0 +1,36 @@
+struct Params {
+    vocabSize: u32,
+    tokenCount: u32
+}
+
+@group(0) @binding(0)
+
+var<storage, read_write> probs: array<f32>;
+
+
+@group(0) @binding(2)
+
+var<storage, read_write> targetToken: array<f32>;
+
+@group(0) @binding(3)
+
+var<uniform> params: Params;
+
+@compute @workgroup_size(256)
+
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    let index = id.x;
+    let token = index / params.vocabSize;
+    let vocabIndex = index % params.vocabSize;
+    let start = token * params.vocabSize;
+    let end = start + params.vocabSize;
+    let total = params.tokenCount * params.vocabSize;
+    if (index >= total) {
+      return;
+    }
+    var targetIndex = targetToken[token];
+    probs[index] = probs[index];
+    if(vocabIndex == u32(targetIndex)){
+        probs[index] -= 1;
+    }
+}

package/src/shaders/SF.wgsl

@@ -0,0 +1,75 @@
+struct Params {
+    vocabSize: u32,
+    tokenCount: u32,
+}
+
+@group(0) @binding(0)
+var<storage, read> logits: array<f32>;
+
+@group(0) @binding(1)
+var<storage, read_write> output: array<f32>;
+
+@group(0) @binding(2)
+var<uniform> params: Params;
+
+var<workgroup> sharedMax: array<f32, 256>;
+var<workgroup> sharedSum: array<f32, 256>;
+
+@compute @workgroup_size(256)
+fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+    let token = wg.x;
+    let tid = lid.x;
+    if (token < params.tokenCount) {
+        let start = token * params.vocabSize;
+        var localMax = -3.402823e38;
+        var i = tid;
+
+        while (i < params.vocabSize) {
+            localMax = max(localMax, logits[start + i]);
+            i += 256u;
+        }
+
+        sharedMax[tid] = localMax;
+        workgroupBarrier();
+
+        var stride = 128u;
+        while (stride > 0u) {
+            if (tid < stride) {
+                sharedMax[tid] = max(
+                    sharedMax[tid],
+                    sharedMax[tid + stride]
+                );
+            }
+            workgroupBarrier();
+            stride /= 2u;
+        }
+        let maxLogit = sharedMax[0];
+
+        var localSum = 0.0;
+        i = tid;
+
+        while (i < params.vocabSize) {
+            let e = exp(logits[start + i] - maxLogit);
+            output[start + i] = e;
+            localSum += e;
+            i += 256u;
+        }
+
+        sharedSum[tid] = localSum;
+        workgroupBarrier();
+        stride = 128u;
+        while (stride > 0u) {
+            if (tid < stride) {
+                sharedSum[tid] += sharedSum[tid + stride];
+            }
+            workgroupBarrier();
+            stride /= 2u;
+        }
+        let total = max(sharedSum[0], 1e-20);
+        i = tid;
+        while (i < params.vocabSize) {
+            output[start + i] /= total;
+            i += 256u;
+        }
+    }
+}

package/src/shaders/SoftmaxBackward.wgsl

@@ -0,0 +1,31 @@
+struct Params {
+    tokens: u32,
+    embeddingSize: u32,
+}
+@group(0) @binding(0)
+var<storage, read> P: array<f32>;
+@group(0) @binding(1)
+var<storage, read> dP: array<f32>;
+@group(0) @binding(2)
+var<storage, read_write> output: array<f32>;
+@group(0) @binding(3)
+var<uniform> params: Params;
+
+@compute @workgroup_size(256)
+
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    let index = id.x;
+    let total = params.tokens * params.tokens;
+    if (index >= total) {
+        return;
+    }
+    let row = index / params.tokens;
+    var dot = 0.0;
+    for (var i = 0u; i < params.tokens; i++) {
+        let pos = row * params.tokens + i;
+        dot += P[pos] * dP[pos];
+    }
+    let p = P[index];
+    let dp = dP[index];
+    output[index] = p * (dp - dot);
+}

package/src/shaders/add.wgsl

@@ -0,0 +1,30 @@
+struct Params{
+  size: u32
+}
+
+@group(0) @binding(0)
+
+var <storage, read> a : array<f32>;
+
+@group(0) @binding(1)
+
+var <storage, read> b : array<f32>;
+
+@group(0) @binding(2)
+
+var <storage, read_write> output : array<f32>;
+@group(0) @binding(3)
+
+var <uniform> params : Params;
+
+@compute @workgroup_size(64)
+
+fn main(@builtin(global_invocation_id) id : vec3<u32>){
+    let index = id.x;
+    if(index >= params.size){
+        return;
+    }
+    let mainOut =  a[index] + b[index];
+    output[index] = mainOut; 
+}
+

package/src/shaders/attentionShader/CompueFinalForwardOutput.wgsl

@@ -0,0 +1,29 @@
+struct Params {
+    tokens: u32,
+    embeddingSize: u32,
+}
+@group(0) @binding(0)
+var<storage, read> weights: array<f32>;
+@group(0) @binding(1)
+var<storage, read> value: array<f32>;
+@group(0) @binding(2)
+var<storage, read_write> output: array<f32>;
+@group(0) @binding(3)
+var<uniform> params: Params;
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    let index = id.x;
+    let total =params.tokens * params.embeddingSize;
+    if (index >= total) {
+        return;
+    }
+    let row = index / params.embeddingSize;
+    let dim = index % params.embeddingSize;
+    var sum = 0.0;
+    for (var j = 0u; j < params.tokens; j++) {
+        let weight = weights[row * params.tokens + j];
+        let v =value[j * params.embeddingSize + dim];
+        sum += weight * v;
+    }
+    output[index] = sum;
+}

package/src/shaders/attentionShader/ComputeDA.wgsl

@@ -0,0 +1,30 @@
+struct Params {
+    contextSize: u32,
+    embeddingSize: u32,
+}
+@group(0) @binding(0)
+var<storage, read> dOutput: array<f32>;
+@group(0) @binding(1)
+var<storage, read> V: array<f32>;
+@group(0) @binding(2)
+var<storage, read_write> dA: array<f32>;
+@group(0) @binding(3)
+var<uniform> params: Params;
+@compute @workgroup_size(16, 16)
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    let row = id.x;
+    let col = id.y;
+    let T = params.contextSize;
+    let D = params.embeddingSize;
+    if (row >= T || col >= T) {
+        return;
+    }
+    var sum = 0.0;
+    for (var k = 0u; k < D; k++) {
+        let dOutputIndex = row * D + k;
+        let vIndex = col * D + k;
+        sum += dOutput[dOutputIndex] * V[vIndex];
+    }
+    let outputIndex = row * T + col;
+    dA[outputIndex] = sum;
+}

package/src/shaders/attentionShader/ComputeDV.wgsl

@@ -0,0 +1,37 @@
+struct Params {
+    contextSize: u32,
+    embeddingSize: u32,
+}
+
+@group(0) @binding(0)
+var<storage, read> A: array<f32>;
+
+@group(0) @binding(1)
+var<storage, read> dOutput: array<f32>;
+
+@group(0) @binding(2)
+var<storage, read_write> dV: array<f32>;
+
+@group(0) @binding(3)
+var<uniform> params: Params;
+
+
+@compute @workgroup_size(16, 16)
+fn main(@builtin(global_invocation_id) id: vec3<u32>)
+ {
+    let row = id.x;
+    let col = id.y;
+    let T = params.contextSize;
+    let D = params.embeddingSize;
+    if (row >= T || col >= D) {
+        return;
+    }
+    var sum = 0.0;
+     for (var i = 0u; i < T; i++) {
+        let aIndex = i * T + row;
+        let gradIndex = i * D + col;
+        sum += A[aIndex] * dOutput[gradIndex];
+    }
+    let outputIndex = row * D + col;
+    dV[outputIndex] = sum;
+}

package/src/shaders/attentionShader/ComputeK.wgsl

@@ -0,0 +1,30 @@
+struct Params {
+    contextSize: u32,
+    embeddingSize: u32,
+}
+@group(0) @binding(0)
+var<storage, read> dS: array<f32>;
+@group(0) @binding(1)
+var<storage, read> Q: array<f32>;
+@group(0) @binding(2)
+var<storage, read_write> dK: array<f32>;
+@group(0) @binding(3)
+var<uniform> params: Params;
+@compute @workgroup_size(16, 16)
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    let row = id.x;
+    let col = id.y;
+    let T = params.contextSize;
+    let D = params.embeddingSize;
+    if (row >= T || col >= D) {
+        return;
+    }
+    var sum = 0.0;
+    for (var i = 0u; i < T; i++) {
+        let dsIndex = i * T + row;
+        let qIndex = i * D + col;
+        sum += dS[dsIndex] * Q[qIndex];
+    }
+    let scale = sqrt(f32(D));
+    dK[row * D + col] = sum / scale;
+}

package/src/shaders/attentionShader/ComputedQ.wgsl

@@ -0,0 +1,31 @@
+struct Params {
+    contextSize: u32,
+    embeddingSize: u32,
+}
+@group(0) @binding(0)
+var<storage, read> dS: array<f32>;
+@group(0) @binding(1)
+var<storage, read> K: array<f32>;
+@group(0) @binding(2)
+var<storage, read_write> dQ: array<f32>;
+@group(0) @binding(3)
+var<uniform> params: Params;
+@compute @workgroup_size(16, 16)
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    let row = id.x;
+    let col = id.y;
+    let T = params.contextSize;
+    let D = params.embeddingSize;
+    if (row >= T || col >= D) {
+        return;
+    }
+    var sum = 0.0;
+    for (var i = 0u; i < T; i++) {
+        let dsIndex = row * T + i;
+        let kIndex = i * D + col;
+
+        sum += dS[dsIndex] * K[kIndex];
+    }
+    let scale = sqrt(f32(D));
+    dQ[row * D + col] = sum / scale;
+}

package/src/shaders/attentionShader/attentionSocres.wgsl

@@ -0,0 +1,40 @@
+struct Params {
+    contextSize: u32,
+    embeddingSize: u32,
+}
+
+@group(0) @binding(0)
+var<storage, read> Q: array<f32>;
+
+@group(0) @binding(1)
+var<storage, read> K: array<f32>;
+
+@group(0) @binding(2)
+var<storage, read_write> scores: array<f32>;
+
+@group(0) @binding(3)
+var<uniform> params: Params;
+
+@compute @workgroup_size(16, 16)
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    let row = id.x;
+    let col = id.y;
+
+    if (row >= params.contextSize || col >= params.contextSize) {
+        return;
+    }
+    var sum = 0.0;
+    for (var k = 0u; k < params.embeddingSize; k++) {
+        let qIndex = row * params.embeddingSize + k;
+        let kIndex = col * params.embeddingSize + k;
+        sum += Q[qIndex] * K[kIndex];
+    }
+
+    let outIndex = row * params.contextSize + col;
+    if (col > row) {
+        scores[outIndex] = - 1e9;
+        return;
+    }
+    let scale = sqrt(f32(params.embeddingSize));
+    scores[outIndex] = sum / scale;
+}

package/src/shaders/attentionShader/sumGradient.wgsl

@@ -0,0 +1,24 @@
+struct Params {
+    contextSize: u32,
+    embeddingSize: u32,
+}
+
+@group(0) @binding(0)
+var<storage, read> dXq: array<f32>;
+@group(0) @binding(1)
+var<storage, read> dXk: array<f32>;
+@group(0) @binding(2)
+var<storage, read> dXv: array<f32>;
+@group(0) @binding(3)
+var<storage, read_write> dX: array<f32>;
+@group(0) @binding(4)
+var<uniform> params: Params;
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    let index = id.x;
+    let total = params.contextSize * params.embeddingSize;
+    if (index >= total) {
+        return;
+    }
+    dX[index] =dXq[index] +dXk[index] +dXv[index];
+}

package/src/shaders/calculateTotal.wgsl

@@ -0,0 +1,26 @@
+@group(0) @binding(0)
+var <storage, read> losses : array<f32>;
+
+@group(0) @binding(1)
+
+var <storage, read> acces : array<f32>;
+
+@group(0) @binding(2)
+
+var <storage, read_write> totalLoss : array<f32>;
+
+@group(0) @binding(3)
+
+var <storage, read_write> totalAcc : array<f32>;
+
+@compute @workgroup_size(1)
+
+fn main(){
+    var lossSum = 0.0;
+    for(var a = 0u; a < arrayLength(&losses); a++){
+        lossSum += losses[a];
+    }
+    lossSum /= f32(arrayLength(&losses));
+    totalAcc[0u] += acces[0];
+    totalLoss[0u] += lossSum;
+}

package/README.md

@@ -1,3 +0,0 @@
-# Temporary Holding Version
-
-This version is a temporary placeholder for this package. An operational version to replace this has been submitted for review and is awaiting a staged release.