@imaduddin/fastbrain 1.1.1-patch → 2.0.0

package/src/shaders/Norm/rmsForward.wgsl

@@ -1,30 +0,0 @@
-@group(0) @binding(0)
-var<storage, read> reduce: array<f32>;
-
-@group(0) @binding(1)
-var<storage, read> vector: array<f32>;
-
-@group(0) @binding(2)
-var<storage, read_write> output: array<f32>;
-
-@group(0) @binding(3)
-var<storage, read> gamma: array<f32>;
-
-@compute @workgroup_size(64)
-fn main(@builtin(global_invocation_id) id: vec3<u32>) {
-    let x = id.x;
-
-    if (x >= arrayLength(&vector)) {
-        return;
-    }
-
-    let dim = %embeddingSize%u;
-    let token = x / dim;
-    let d = x % dim;
-
-    let rms = sqrt(
-        reduce[token] / f32(dim) + 1e-6
-    );
-
-    output[x] = (vector[x] / rms) * gamma[d];
-}

package/src/shaders/Norm/sumSquare.wgsl

@@ -1,38 +0,0 @@
-
-@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, 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 base = wid.x * D;
-    var sum = 0.0;
-    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();
-    var stride = 32u;
-    loop {
-        if (stride == 0u) {
-            break;
-        }
-        if (lid < stride) {
-            partial[lid] += partial[lid + stride];
-        }
-        workgroupBarrier();
-        stride /= 2u;
-    }
-    if (lid == 0u) {
-        out[wid.x] = partial[0];
-    }
-}

package/src/shaders/Norm/updateGamma.wgsl

@@ -1,54 +0,0 @@
-struct Params{
-    lr: f32
-}
-@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> sumSquares: array<f32>;
-
-@group(0) @binding(3)
-var<storage, read_write> gamma: array<f32>;
-
-@group(0) @binding(4)
-var<uniform> params: Params;
-
-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 dim = wid.x;
-    let T = arrayLength(&sumSquares);
-    var sum = 0.0;
-    if (dim < D) {
-        var token = lid;
-        while (token < T) {
-            let idx = token * D + dim;
-            let rms = sqrt(sumSquares[token] / f32(D) + EPSILON);
-            sum += gradOutput[idx] * input[idx] / rms;
-            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 && dim < D) {
-        gamma[dim] -= params.lr * partial[0] / f32(T);
-    }
-}

package/src/shaders/PB.wgsl

@@ -1,25 +0,0 @@
-struct Params{
-    lr : f32
-}
-
-@group(0) @binding(0)
-var <storage, read_write> posVec: array<f32>;
-
-@group(0) @binding(1)
-
-var <storage, read> outputGradient : array<f32>;
-
-
-@group(0) @binding(2)
-
-var <uniform> params: Params;
-
-@compute @workgroup_size(64)
-
-fn main(@builtin(global_invocation_id) id : vec3<u32>){
-    let x = id.x;
-    if(x >= arrayLength(&outputGradient)){
-        return;
-    }
-    posVec[x] -= params.lr * outputGradient[x];
-}

package/src/shaders/PF.wgsl

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

package/src/shaders/RB.wgsl

@@ -1,37 +0,0 @@
-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

@@ -1,33 +0,0 @@
-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

@@ -1,28 +0,0 @@
-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

@@ -1,41 +0,0 @@
-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;
-
-
-@group(0) @binding(4)
-
-var<storage, read_write> dLogits: array<f32>;
-
-@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];
-    dLogits[index] = probs[index];
-    if(vocabIndex == u32(targetIndex)){
-        dLogits[index] -= 1;
-    }
-}

package/src/shaders/SF.wgsl

@@ -1,75 +0,0 @@
-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

@@ -1,31 +0,0 @@
-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

@@ -1,30 +0,0 @@
-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

@@ -1,29 +0,0 @@
-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

@@ -1,30 +0,0 @@
-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

@@ -1,37 +0,0 @@
-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

@@ -1,30 +0,0 @@
-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

@@ -1,31 +0,0 @@
-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

@@ -1,40 +0,0 @@
-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

@@ -1,24 +0,0 @@
-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];
-}