@imaduddin/fastbrain 1.1.1 → 2.0.0

package/shaders/linear/ComputeDinput.wgsl

@@ -0,0 +1,72 @@
+struct Params {
+    inputSize: f32,
+    outputSize: f32,
+    tokenCount: f32,
+    lr : f32
+}
+
+@group(0) @binding(0)
+
+var <storage, read> weights: array<f32>;
+
+
+@group(0) @binding(1)
+
+var <storage, read> outputGradient : array<f32>;
+
+@group(0) @binding(2)
+
+var <storage, read_write> dinput : array<f32>;
+
+@group(0) @binding(3)
+
+var <uniform> params : Params;
+
+var <workgroup> sdata : array<f32, 64>;
+
+@compute @workgroup_size(64)
+
+fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>){
+    var x = wid.x;
+    var y = wid.y;
+    var tid = lid.x;
+    var isCorrect = x >= u32(params.tokenCount) || y >= u32(params.inputSize);
+    if(!isCorrect){
+        var local_sum = 0.0;
+        for(var out = tid; out < u32(params.outputSize); out += 64){
+            let dOut = outputGradient[x * u32(params.outputSize) + out];
+            let w = weights[out * u32(params.inputSize) + y];
+            local_sum += dOut * w;
+        }
+        sdata[tid] = local_sum;
+        workgroupBarrier();
+        
+          if(tid < 32u){
+            sdata[tid] += sdata[tid + 32u];
+        }
+            workgroupBarrier();
+        if(tid < 16u){
+            sdata[tid] += sdata[tid + 16u];
+        }
+            workgroupBarrier();
+        if(tid < 8u){
+            sdata[tid] += sdata[tid + 8u];
+        }
+            workgroupBarrier();
+        if(tid < 4u){
+            sdata[tid] += sdata[tid + 4u];
+        }
+            workgroupBarrier();
+        if(tid < 2u){
+            sdata[tid] += sdata[tid + 2u];
+        }
+            workgroupBarrier();
+        if(tid < 1u){
+            sdata[tid] += sdata[tid + 1u];
+        }
+            workgroupBarrier();
+        if(tid == 0u){
+            dinput[x * u32(params.inputSize) + y] = sdata[0u];
+        }    
+    }
+}

package/shaders/linear/forward.wgsl

@@ -0,0 +1,63 @@
+
+struct Params {
+    inputSize: u32,
+    outputSize: u32,
+    tokenCount: u32,
+};
+@group(0) @binding(0)
+var<storage, read> input: array<f32>;
+@group(0) @binding(1)
+var<storage, read> weights: array<f32>;
+@group(0) @binding(2)
+var<storage, read> bias: array<f32>;
+@group(0) @binding(3)
+var<storage, read_write> output: array<f32>;
+@group(0) @binding(4)
+var<uniform> params: Params;
+
+var <workgroup> sdata: array<f32, 64>;
+@compute @workgroup_size(64)
+fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+    var vocabIndex = wid.x;
+    var tokenIndex = wid.y;
+    var lane = lid.x;
+    let isCorrect = vocabIndex >= params.outputSize || tokenIndex >= params.tokenCount;
+    if (!isCorrect) {
+         var sum = 0.0;
+        for (var i = lane; i < params.inputSize; i+= 64u) {
+            let inputIndex = tokenIndex * params.inputSize + i;
+            let weightIndex = vocabIndex * params.inputSize + i;
+            sum += input[inputIndex] * weights[weightIndex];
+        }
+        sdata[lane] = sum;
+        workgroupBarrier();
+            if(lane < 32u){
+            sdata[lane] += sdata[lane + 32u];
+        }
+        workgroupBarrier();
+         if(lane < 16u){
+            sdata[lane] += sdata[lane + 16u];
+        }
+        workgroupBarrier();
+         if(lane < 8u){
+            sdata[lane] += sdata[lane + 8u];
+        }
+        workgroupBarrier();
+         if(lane < 4u){
+            sdata[lane] += sdata[lane + 4u];
+        }
+        workgroupBarrier();
+         if(lane < 2u){
+            sdata[lane] += sdata[lane + 2u];
+        }
+        workgroupBarrier();
+         if(lane < 1u){
+            sdata[lane] += sdata[lane + 1u];
+        }
+        workgroupBarrier();
+        if(lane == 0u){
+            output[tokenIndex * params.outputSize + vocabIndex] = sdata[0u] + bias[vocabIndex];
+        };
+    }
+    
+}

package/shaders/linear/updateBias.wgsl

@@ -0,0 +1,33 @@
+struct Params {
+    inputSize: f32,
+    outputSize: f32,
+    tokenCount: f32,
+    lr : f32
+}
+@group(0) @binding(0)
+
+var <storage, read_write> bias : array<f32>;
+
+@group(0) @binding(1)
+
+var <storage, read> outputGradient: array<f32>;
+
+@group(0) @binding(2)
+
+var <uniform> params: Params;
+
+
+@compute @workgroup_size(32)
+
+fn main(@builtin(global_invocation_id) gid : vec3<u32>){
+    let vocab = gid.x;
+    if(vocab >= u32(params.outputSize)){
+        return;
+    }
+    var grad = 0.0;
+    for(var a = 0u; a < u32(params.tokenCount); a++){
+        let gradIndex = a * u32(params.outputSize) + vocab;
+        grad += outputGradient[gradIndex];
+    } 
+    bias[vocab] -= params.lr * grad;
+}

package/shaders/linear/updateWeights.wgsl

@@ -0,0 +1,38 @@
+struct Params {
+    inputSize: f32,
+    outputSize: f32,
+    tokenCount: f32,
+    lr : f32
+}
+@group(0) @binding(0)
+
+var <storage, read_write> weights : array<f32>;
+
+@group(0) @binding(1)
+
+var <storage, read> outputGradient: array<f32>;
+
+@group(0) @binding(2)
+
+var <uniform> params: Params;
+
+@group(0) @binding(3)
+
+var <storage, read> input: array<f32>;
+
+@compute @workgroup_size(16, 32)
+
+fn main(@builtin(global_invocation_id) gid : vec3<u32>){
+    let dim = gid.x;
+    let vocab = gid.y;
+    if(dim >= u32(params.inputSize) || vocab >= u32(params.outputSize)){
+        return;
+    }
+    let index = u32(vocab) * u32(params.inputSize) + u32(dim);
+    var grad = 0.0;
+    for(var a = 0u; a < u32(params.tokenCount); a++){
+        let gradIndex = a * u32( params.outputSize) + vocab;
+        grad += outputGradient[gradIndex] * input[a * u32( params.inputSize) + dim];
+    } 
+    weights[index] -= params.lr * grad;
+}

package/shaders/relu/backward.wgsl

@@ -0,0 +1,43 @@
+struct Params{
+    dimSize: u32,
+    tokenCount : u32
+}
+
+@group(0) @binding(0)
+
+var <storage, read> input : array<f32>;
+
+@group(0) @binding(1)
+
+var <storage, read_write> output : array<f32>;
+
+
+@group(0) @binding(2)
+
+var <storage, read_write> grad : array<f32>;
+
+@group(0) @binding(3)
+var <uniform> params : Params;
+
+
+@compute @workgroup_size(16,16)
+fn main(@builtin(global_invocation_id) id : vec3<u32>){
+   var dim = id.x;
+   var token = id.y;
+   if(dim >= params.dimSize || token >= params.tokenCount){
+    return;
+   }
+   var index = token * params.dimSize + dim;
+   let x = input[index];
+   let go = grad[index];
+   let c0 = sqrt(2 / 3.141592653589793);
+   let c1 = 0.044715;
+   let x3 = x * x * x;
+   let inner = c0 * (x + c1 * x3);
+   let tanh_inner = tanh(inner);
+   let left = 0.5 * (1.0 + tanh_inner);
+   let sech2 = 1.0 - tanh_inner * tanh_inner;
+   let right = 0.5 * x * sech2 * c0 * (1.0 + 3.0 * c1 * x * x);
+   let dgelu = left + right;
+   output[index] = go * dgelu;
+}

package/shaders/relu/forward.wgsl

@@ -0,0 +1,28 @@
+struct Params{
+    dimSize: u32,
+    tokenCount : u32
+}
+
+@group(0) @binding(0)
+
+var <storage, read> input : array<f32>;
+
+@group(0) @binding(1)
+
+var <storage, read_write> output : array<f32>;
+
+@group(0) @binding(2)
+
+var <uniform> params : Params;
+
+@compute @workgroup_size(16,16)
+fn main(@builtin(global_invocation_id) id : vec3<u32>){
+   var dim = id.x;
+   var token = id.y;
+   if(dim >= params.dimSize || token >= params.tokenCount){
+    return;
+   }
+   var index = token * params.dimSize + dim;
+   let x = input[index];
+   output[index] = 0.5 * x * (1 + tanh(sqrt(2 / 3.141592653589793) * (x + 0.044715 * x * x * x)));
+}

package/shaders/rmsnorm/backward.wgsl

@@ -0,0 +1,54 @@
+struct Params {
+    dim: u32,
+    tokenCount: u32,
+}
+
+@group(0) @binding(0)
+var<storage, read> input: array<f32>;
+
+@group(0) @binding(1)
+var<storage, read> rstd: array<f32>;
+
+@group(0) @binding(2)
+var<storage, read> dout: array<f32>;
+
+@group(0) @binding(3)
+var<storage, read_write> din: array<f32>;
+
+@group(0) @binding(4)
+var<uniform> params: Params;
+
+var<workgroup> shared1: array<f32, 256>;
+
+@compute @workgroup_size(256)
+fn main(@builtin(local_invocation_id) lid: vec3<u32>, @builtin(workgroup_id) wid: vec3<u32>) {
+    let token = wid.x;
+    let lane  = lid.x;
+
+    if (token >= params.tokenCount) {
+        return;
+    }
+
+    let base = token * params.dim;
+    let inv_rms = rstd[token];
+    var local_sum = 0.0;
+    for (var i = lane; i < params.dim; i += 256u) {
+        local_sum += dout[base + i] * input[base + i];
+    }
+
+    shared1[lane] = local_sum;
+    workgroupBarrier();
+    for (var stride = 128u; stride > 0u; stride >>= 1u) {
+        if (lane < stride) {
+            shared1[lane] += shared1[lane + stride];
+        }
+        workgroupBarrier();
+    }
+    let mean_term = (shared1[0] * inv_rms * inv_rms) / f32(params.dim);
+    for (var i = lane; i < params.dim; i += 256u) {
+        let x  = input[base + i];
+        let dy = dout[base + i];
+
+        din[base + i] = inv_rms * (dy - x * mean_term);
+    }
+}

package/shaders/rmsnorm/forward.wgsl

@@ -0,0 +1,51 @@
+struct Params {
+    dim: u32,
+    tokenCount: u32,
+}
+
+@group(0) @binding(0)
+var<storage, read> input: array<f32>;
+
+@group(0) @binding(1)
+var<storage, read_write> output: array<f32>;
+
+@group(0) @binding(2)
+var<storage, read_write> rstd: array<f32>;
+
+@group(0) @binding(3)
+var<uniform> params: Params;
+
+
+
+var<workgroup> shared1: array<f32, 256>;
+
+@compute @workgroup_size(256)
+fn main(@builtin(local_invocation_id) lid: vec3<u32>, @builtin(workgroup_id) wid: vec3<u32>) {
+    let token = wid.x;
+    let lane  = lid.x;
+    if (token >= params.tokenCount) {
+        return;
+    }
+    let base = token * params.dim;
+    var local_sum = 0.0;
+    for (var i = lane; i < params.dim; i += 256u) {
+        let x = input[base + i];
+        local_sum += x * x;
+    }
+    shared1[lane] = local_sum;
+    workgroupBarrier();
+    for (var stride = 128u; stride > 0u; stride >>= 1u) {
+        if (lane < stride) {
+              shared1[lane] += shared1[lane + stride];
+        }
+        workgroupBarrier();
+    }
+    let inv_rms = 1.0 / sqrt(shared1[0] / f32(params.dim) + 1e-8);
+    if (lane == 0u) {
+        rstd[token] = inv_rms;
+    }
+    workgroupBarrier();
+    for (var i = lane; i < params.dim; i += 256u) {
+        output[base + i] = input[base + i] * inv_rms;
+    }
+}

package/shaders/selfAttention/SB.wgsl

@@ -0,0 +1,30 @@
+
+struct Params {
+    size: u32,        
+    D: u32 
+}
+
+@group(0) @binding(0) 
+var<storage, read> probs: array<f32>;     
+@group(0) @binding(1) 
+var<storage, read> dProbs: array<f32>;    
+@group(0) @binding(2) 
+var<storage, read_write> dScores: 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.y;
+    let col = id.x;
+    if (row >= params.size || col >= params.size) { return; }
+    let idx = row * params.size + col;
+    let p = probs[idx];
+    var sum = 0.0;
+    for (var j = 0u; j < params.size; j++) {
+        sum += probs[row * params.size + j] * dProbs[row * params.size + j];
+    }
+
+    dScores[idx] = p * (dProbs[idx] - sum);
+    dScores[idx] /= sqrt(f32(params.D));
+}

package/shaders/selfAttention/add3.wgsl

@@ -0,0 +1,27 @@
+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> c: array<f32>;
+
+@group(0) @binding(3)
+var<storage, read_write> d: array<f32>;
+
+@group(0) @binding(4)
+var<uniform> params: Params;
+
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    let i = id.x;
+    if (i >= params.size) {
+        return;
+    }
+    d[i] = a[i] + b[i] + c[i];
+}

package/shaders/selfAttention/matMul.wgsl

@@ -0,0 +1,63 @@
+struct Params{
+    M: u32,
+    N: u32,
+    K: 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> c : array<f32>;
+
+@group(0) @binding(3)
+var <uniform> params : Params;
+
+var <workgroup> sdata: array<f32, 64>;
+
+
+@compute @workgroup_size(64)
+
+fn main(@builtin(workgroup_id) id : vec3<u32>, @builtin(local_invocation_id) lid : vec3<u32>){
+    let row = id.y;
+    let col = id.x;
+    let tane = lid.x;
+    let isCorrect = row >= params.M || col >= params.N;
+    if(!isCorrect){
+    var sum = 0.0;
+    for(var k = tane; k < params.K; k+= 64u){
+        sum += a[row * params.K + k] * b[k * params.N + col];
+    }
+    sdata[tane] = sum;
+    workgroupBarrier();
+    if(tane < 32u){
+         sdata[tane] += sdata[tane + 32u];
+    }
+    workgroupBarrier() ; 
+    if(tane < 16u){
+         sdata[tane] += sdata[tane + 16u];
+    }
+    workgroupBarrier();
+    if(tane < 8u){
+         sdata[tane] += sdata[tane + 8u];
+    }
+    workgroupBarrier();  
+    if(tane < 4u){
+         sdata[tane] += sdata[tane + 4u];
+    }
+    workgroupBarrier() ; 
+    if(tane < 2u){
+         sdata[tane] += sdata[tane + 2u];
+    }
+    workgroupBarrier() ; 
+    if(tane < 1u){
+         sdata[tane] += sdata[tane + 1u];
+    }
+    workgroupBarrier() ; 
+    if(tane == 0u){
+        c[row * params.N + col] = sdata[0u];
+    }
+    }
+}

package/shaders/selfAttention/qK.wgsl

@@ -0,0 +1,66 @@
+struct Params{
+    M: u32,
+    N: u32,
+    K: 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> c : array<f32>;
+
+@group(0) @binding(3)
+var <uniform> params : Params;
+
+
+var <workgroup> sdata : array<f32, 64>;
+@compute @workgroup_size(64)
+
+fn main(@builtin(workgroup_id) id : vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>){
+    let row = id.y;
+    let col = id.x;
+    let tane = lid.x;
+    let isCorrect  = row >= params.M || col >= params.N;
+    if(!isCorrect){
+    var sum = 0.0;
+    for(var k = tane; k < params.K; k+= 64u){
+        sum += a[row * params.K + k] * b[k * params.N + col];
+    }
+    sdata[tane] = sum;
+     workgroupBarrier();
+    if(tane < 32u){
+         sdata[tane] += sdata[tane + 32u];
+    }
+    workgroupBarrier() ; 
+    if(tane < 16u){
+         sdata[tane] += sdata[tane + 16u];
+    }
+    workgroupBarrier();
+    if(tane < 8u){
+         sdata[tane] += sdata[tane + 8u];
+    }
+    workgroupBarrier();  
+    if(tane < 4u){
+         sdata[tane] += sdata[tane + 4u];
+    }
+    workgroupBarrier() ; 
+    if(tane < 2u){
+         sdata[tane] += sdata[tane + 2u];
+    }
+    workgroupBarrier() ; 
+    if(tane < 1u){
+         sdata[tane] += sdata[tane + 1u];
+    }
+    workgroupBarrier() ; 
+    if(tane == 0u){
+        sum = sdata[0u] / sqrt(f32(params.K));
+    if(col > row){
+        sum = -1e9;
+    }
+      c[row * params.N + col] = sum ;
+    }
+    }
+}

package/shaders/selfAttention/transpose.wgsl

@@ -0,0 +1,24 @@
+struct Params{
+    M: u32,
+    N: u32,
+    K: u32
+}
+@group(0) @binding(0)
+var <storage, read> a : array<f32>;
+
+@group(0) @binding(1)
+var <storage, read_write> b : array<f32>;
+
+@group(0) @binding(2)
+var <uniform> params : Params;
+
+@compute @workgroup_size(16, 16)
+
+fn main(@builtin(global_invocation_id) id : vec3<u32>){
+    let x = id.x;
+    let y = id.y;
+    if(y >= params.M || x >= params.N){
+        return ;
+    }
+    b[x * params.M + y] = a[y * params.N + x];
+}

package/shaders/softmax/backward.wgsl

@@ -0,0 +1,36 @@
+struct Params{
+    vocabSize: u32,
+    tokenCount: u32
+}
+
+@group(0) @binding(0)
+
+var <storage, read> probs: array<f32>;
+
+@group(0) @binding(1)
+
+var <storage, read_write> dLogits: array<f32>;
+
+@group(0) @binding(2)
+
+var <storage, read> targetBuffer: array<u32>;
+
+@group(0) @binding(3)
+
+var <uniform> params: Params;
+
+@compute @workgroup_size(64, 1)
+
+
+fn main(@builtin(global_invocation_id) id : vec3<u32>){
+    let x = id.x;
+    let y = id.y;
+    if(x >= params.vocabSize || y >= params.tokenCount){
+        return;
+    }
+    let index = y * params.vocabSize + x;
+    dLogits[index] = probs[index];
+    if(x == targetBuffer[y]){
+        dLogits[index] -= 1;
+    }
+}

package/shaders/softmax/crossEntropy.wgsl

@@ -0,0 +1,36 @@
+struct Params{
+    vocabSize: u32,
+    tokenCount: u32
+}
+
+
+@group(0) @binding(0)
+
+var <storage, read> probs: array<f32>;
+
+@group(0) @binding(1)
+
+var <storage, read> targetBuffer: array<u32>;
+
+
+@group(0) @binding(2)
+
+var <storage, read_write> loss: array<f32>;
+
+@group(0) @binding(3)
+
+var <uniform> params: Params;
+
+
+@compute @workgroup_size(16)
+
+
+fn main(@builtin(global_invocation_id) id : vec3<u32>){
+    let token = id.x;
+    if(token >= params.tokenCount){
+        return;
+    }
+    let targetToken = targetBuffer[token];
+    let index = token * params.vocabSize + targetToken;
+    loss[token] = - log(probs[index] + 1e-9);
+}