@imaduddin/fastbrain 1.1.1 → 2.0.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);
+}

package/shaders/softmax/forward.wgsl

@@ -0,0 +1,68 @@
+struct Params {
+    vocabSize: u32,
+    tokenCount: u32,
+};
+
+@group(0) @binding(0)
+var<storage, read> logits: array<f32>;
+
+@group(0) @binding(1)
+var<storage, read_write> probs: 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, 1, 1)
+fn main(@builtin(local_invocation_id) lid: vec3<u32>, @builtin(workgroup_id) wid: vec3<u32>) {
+    let lane = lid.x;
+    let token = wid.x;
+
+    if (token >= params.tokenCount) {
+        return;
+    }
+
+    let vocabSize = params.vocabSize;
+    let base = token * vocabSize;
+    var localMax = -3.402823466e+38;
+
+    for (var i = lane; i < vocabSize; i += 256u) {
+        localMax = max(localMax, logits[base + i]);
+    }
+
+    sharedMax[lane] = localMax;
+    workgroupBarrier();
+
+    for (var stride = 128u; stride > 0u; stride >>= 1u) {
+        if (lane < stride) {
+            sharedMax[lane] = max(
+                sharedMax[lane],
+                sharedMax[lane + stride]
+            );
+        }
+        workgroupBarrier();
+    }
+
+    let maxValue = sharedMax[0];
+    var localSum = 0.0;
+
+    for (var i = lane; i < vocabSize; i += 256u) {
+        localSum += exp(logits[base + i] - maxValue);
+    }
+
+    sharedSum[lane] = localSum;
+    workgroupBarrier();
+    for (var stride = 128u; stride > 0u; stride >>= 1u) {
+        if (lane < stride) {
+            sharedSum[lane] += sharedSum[lane + stride];
+        }
+        workgroupBarrier();
+    }
+    let denominator = sharedSum[0];
+
+    for (var i = lane; i < vocabSize; i += 256u) {
+        probs[base + i] = exp(logits[base + i] - maxValue) / denominator;
+    }
+}

package/src/Embedding/index.ts

@@ -0,0 +1,81 @@
+import { GetVectorsForFile, SaveFileVector, initRandomVector } from '../fileManager.js'
+import { Buffer, GPUManager } from '@imaduddin/webgpu-manager';
+
+
+export class Embedding {
+    vectors: Float32Array;
+    vectorsBuffer: Buffer | undefined;
+    gpu: GPUManager | undefined;
+    param: Buffer | undefined;
+    paramBackward: Buffer | undefined;
+    tokenBuffer: Buffer | undefined;
+    output: Buffer | undefined;
+    forwardShader: GPUComputePipeline | undefined;
+    backwardPipline: GPUComputePipeline | undefined;
+    constructor(public device: GPUDevice, public embeddingSize: number, public vocabSize: number, public configs: {
+        filename: string
+    }) {
+        if (!this.configs) {
+            this.vectors = initRandomVector(embeddingSize * vocabSize);
+        } else {
+            if (!this.configs.filename) {
+                this.vectors = initRandomVector(embeddingSize * vocabSize)
+            } else {
+                const fileVec = GetVectorsForFile(configs.filename, embeddingSize * vocabSize, false, embeddingSize, vocabSize);
+                if (!fileVec) {
+                    this.vectors = initRandomVector(embeddingSize * vocabSize)
+                    SaveFileVector(this.configs.filename, this.vectors)
+                } else {
+                    this.vectors = fileVec;
+                }
+            }
+        }
+        this.gpu = new GPUManager(device, {
+            shaderDir: "./shaders"
+        })
+    }
+    async forward(input: number[], encoderShared: GPUCommandEncoder | undefined) {
+        let token = new Uint32Array(input)
+        let encoder = encoderShared ? encoderShared : await this.gpu?.CreateEncoder()
+        if (!this.vectorsBuffer) {
+            this.vectorsBuffer = await this.gpu?.createBuffer(this.vectors, "storage")
+            this.forwardShader = await this.gpu?.getPiplineShader("embedding/forward.wgsl")
+        }
+        if (!this.param) this.param = await this.gpu?.createBuffer(new Uint32Array([this.embeddingSize, this.vocabSize]), "uniform");
+        if (!this.tokenBuffer) this.tokenBuffer = await this.gpu?.createBuffer(token, "storage")
+        if (this.tokenBuffer?.buffer.size !== token.byteLength) {
+            this.tokenBuffer?.destroy()
+            this.tokenBuffer = await this.gpu?.createBuffer(token, "storage")
+        }
+        this.tokenBuffer?.changeData(token)
+        if (!this.output) this.output = await this.gpu?.createBuffer(new Float32Array(this.embeddingSize * input.length), "storage");
+        if (this.output?.buffer.size !== input.length * this.embeddingSize * 4 || this.output.isDestroy()) {
+            this.output?.destroy();
+            this.output = await this.gpu?.createBuffer(new Float32Array(this.embeddingSize * input.length), "storage");
+        }
+        const bindGroup: any = await this.gpu?.createBindGroup(this.forwardShader!, [this.vectorsBuffer!, this.tokenBuffer!, this.output!, this.param!], 0)
+        await this.gpu?.RunPipline(encoder!, this.forwardShader!, bindGroup, Math.ceil(this.embeddingSize / 64), Math.ceil(input.length / 4))
+        if (!encoderShared) await this.gpu?.SumbitEncoder(encoder!)
+        return this.output
+
+    }
+    async backward(dInput: Buffer, lr: number, encoderShared: GPUCommandEncoder | undefined) {
+        const encoder = encoderShared ? encoderShared : await this.gpu?.CreateEncoder()
+        const tokenCount  = dInput.buffer.size / (this.embeddingSize * 4)
+        if (!this.backwardPipline) {
+            this.backwardPipline = await this.gpu?.getPiplineShader("embedding/backward.wgsl")
+        }
+        if (!this.paramBackward) {
+            this.paramBackward = await this.gpu?.createBuffer(new Float32Array([this.embeddingSize, this.vocabSize, lr]), "uniform")
+        } else {
+            this.paramBackward.changeData(new Float32Array([this.embeddingSize, this.vocabSize, lr]))
+        }
+        const bindGroup = await this.gpu?.createBindGroup(this.backwardPipline!, [this.vectorsBuffer!, this.tokenBuffer!, dInput, this.paramBackward!], 0)
+        await this.gpu?.RunPipline(encoder!, this.backwardPipline!,bindGroup!, Math.ceil(this.embeddingSize / 64), Math.ceil(tokenCount / 4))
+        if(!encoderShared) await this.gpu?.SumbitEncoder(encoder!)
+    }
+    async Save() {
+        let vec: any = await this.vectorsBuffer?.read()
+        await SaveFileVector(this.configs.filename, vec)
+    }
+}

package/src/FFN/index.ts

@@ -0,0 +1,47 @@
+import { Linear } from "../Linear/index.js";
+import { Gelu as Relu } from '../Gelu/index.js'
+import { join } from "node:path";
+import { Buffer, GPUManager } from '@imaduddin/webgpu-manager'
+export class FFN {
+    linear1: Linear | undefined
+    linear2: Linear | undefined
+    activation: Relu | undefined
+    gpu: GPUManager | undefined
+    constructor(public device: GPUDevice, public dimSize: number, public configs: {
+        dirname: string
+    }) {
+        this.linear1 = new Linear(device, dimSize, dimSize * 4, {
+            filename: [join(configs.dirname, "LINEAR_EXPAND_FFN_WEIGHTS_1.bin"), join(configs.dirname, "LINEAR_EXPAND_FFN_BIAS_1.bin")]
+        })
+        this.linear2 = new Linear(device, dimSize * 4, dimSize, {
+            filename: [join(configs.dirname, "LINEAR_EXPAND_FFN_WEIGHTS_2.bin"), join(configs.dirname, "LINEAR_EXPAND_FFN_BIAS_2.bin")]
+        })
+        this.activation = new Relu(device, dimSize * 4)
+        this.gpu = new GPUManager(device, {
+            shaderDir: "./shaders"
+        })
+    }
+    async forward(input: Buffer, encoderShared: GPUCommandEncoder | undefined) {
+        const encoder = encoderShared ? encoderShared : await this.gpu?.CreateEncoder()
+        const hidden1 = await this.linear1?.forward(input, encoder!)
+        const hiddenActivation = await this.activation?.forward(hidden1!, encoder!)
+        const hidden2 = await this.linear2?.forward(hiddenActivation!, encoder!)
+        if (!encoderShared) {
+            await this.gpu?.SumbitEncoder(encoder!)
+        }
+        return hidden2
+    }
+    async backward(outputGradient : Buffer, lr: number, encoderShared : GPUCommandEncoder | undefined){
+        const encoder = encoderShared ? encoderShared : await this.gpu?.CreateEncoder()
+        const dActivation = await this.linear2?.backward(outputGradient, lr, encoder)
+        const dHidden2 = await this.activation?.backward(dActivation!, encoder)
+        const dHidden1 = await this.linear1?.backward(dHidden2!, lr, encoder)
+        await this.linear2?.updateParameter(outputGradient, lr, encoder)
+        await this.linear1?.updateParameter(dHidden2!, lr, encoder)
+        return dHidden1
+    }
+    async Save(){
+        await this.linear1?.Save()
+        await this.linear2?.Save()
+    }
+}

package/src/Gelu/index.ts

@@ -0,0 +1,53 @@
+import { Buffer, GPUManager } from "@imaduddin/webgpu-manager";
+
+
+export class Gelu {
+    input: Buffer | undefined;
+    output: Buffer | undefined;
+    outputGrad: Buffer | undefined;
+    params: Buffer | undefined;
+    forwardShader: GPUComputePipeline | undefined;
+    backwardShader: GPUComputePipeline | undefined;
+    gpu: GPUManager | undefined;
+    constructor(public device: GPUDevice, public hiddenSize: number) {
+        this.gpu = new GPUManager(device, {
+            shaderDir: "./shaders/relu"
+        })
+    }
+    async forward(input: Buffer, encoderShared: GPUCommandEncoder | undefined) {
+        let tokenCount = input.buffer.size / (this.hiddenSize * 4)
+        this.input = input;
+        if (!this.output || this.output.buffer.size !== input.buffer.size) {
+            if (this.output) this.output.destroy()
+            this.output = await this.gpu?.createBuffer(new Float32Array(this.hiddenSize * tokenCount), "storage")
+        }
+        if (!this.forwardShader) {
+            this.forwardShader = await this.gpu?.getPiplineShader("forward.wgsl")
+            this.backwardShader = await this.gpu?.getPiplineShader("backward.wgsl")
+
+        }
+        if (!this.params) {
+            this.params = await this.gpu?.createBuffer(new Uint32Array([this.hiddenSize, tokenCount]), "uniform")
+        } else {
+            await this.params.changeData(new Uint32Array([this.hiddenSize, tokenCount]))
+        }
+        let encoder = encoderShared ? encoderShared : await this.gpu?.CreateEncoder()
+        const bindGroup = await this.gpu?.createBindGroup(this.forwardShader!, [this.input!, this.output!, this.params!], 0)
+        await this.gpu?.RunPipline(encoder!, this.forwardShader!, bindGroup!, Math.ceil(this.hiddenSize / 16), Math.ceil(tokenCount / 16))
+        if (!encoderShared) await this.gpu?.SumbitEncoder(encoder!)
+        return this.output
+    }
+    async backward(gradient: Buffer, encoderShared: GPUCommandEncoder | undefined) {
+        let tokenCount = gradient.buffer.size / (this.hiddenSize * 4)
+        let encoder = encoderShared ? encoderShared : await this.gpu?.CreateEncoder()
+        if(!this.outputGrad || this.outputGrad.buffer.size !== gradient.buffer.size){
+            if(this.outputGrad) this.outputGrad.destroy()
+            this.outputGrad = await this.gpu?.createBuffer(new Float32Array(gradient.buffer.size / 4), "storage")
+        }
+        const bindGroup = await this.gpu?.createBindGroup(this.backwardShader!, [this.input!, this.outputGrad!, gradient!, this.params!], 0)
+        await this.gpu?.RunPipline(encoder!, this.backwardShader!, bindGroup!, Math.ceil(this.hiddenSize / 16), Math.ceil(tokenCount  / 16))
+        if(!encoderShared) await this.gpu?.SumbitEncoder(encoder!)
+        return this.outputGrad
+    }
+
+}

package/src/Linear/index.ts

@@ -0,0 +1,122 @@
+import { GPUManager, Buffer } from "@imaduddin/webgpu-manager";
+import { GetVectorsForFile, SaveFileVector, initRandomVector, getParameter } from '../fileManager.js'
+
+
+export class Linear {
+    bias: Float32Array | undefined;
+    weights: Float32Array | undefined;
+    wBuffer: Buffer | undefined;
+    bBuffer: Buffer | undefined;
+    gpu: GPUManager | undefined;
+    paramsForward: Buffer | undefined;
+    paramsBackward: Buffer | undefined;
+    output: Buffer | undefined;
+    forwardPipline: GPUComputePipeline | undefined;
+    backwardPipline: GPUComputePipeline | undefined;
+    weightUpdate: GPUComputePipeline | undefined;
+    biasUpdate: GPUComputePipeline | undefined;
+    gradOut: Buffer | undefined;
+    input: Buffer | undefined
+    constructor(public device: GPUDevice, public embeddingSize: number, public vocabSize: number, public configs: {
+        filename: string[]
+    }) {
+        if (!this.configs) {
+            this.weights = initRandomVector(embeddingSize * vocabSize);
+            this.bias = initRandomVector(vocabSize);
+        } else {
+            if (!this.configs.filename) {
+                this.weights = initRandomVector(embeddingSize * vocabSize);
+                this.bias = initRandomVector(vocabSize);
+            } else {
+                const fileVec = GetVectorsForFile(configs.filename[0], embeddingSize * vocabSize, false, embeddingSize, vocabSize);
+                const fileBias = GetVectorsForFile(configs.filename[1], vocabSize, true, embeddingSize, vocabSize);
+                if (!fileVec) {
+                    this.weights = initRandomVector(embeddingSize * vocabSize);
+                    this.bias = initRandomVector(vocabSize);
+                    SaveFileVector(this.configs.filename[0], this.weights)
+                    SaveFileVector(this.configs.filename[1], this.bias)
+                } else {
+                    this.weights = fileVec;
+                    this.bias = fileBias!;
+                }
+            }
+        }
+        this.gpu = new GPUManager(device, {
+            shaderDir: "./shaders"
+        })
+    }
+    getParameter(){
+        getParameter()
+    }
+    async forward(input: Buffer, en: GPUCommandEncoder | undefined) {
+        const encoder = en ? en : await this.gpu?.CreateEncoder()
+        this.input = input;
+        if (!this.wBuffer || !this.bBuffer) {
+            this.wBuffer = await this.gpu?.createBuffer(this.weights!, "storage")
+            this.bBuffer = await this.gpu?.createBuffer(this.bias!, "storage")
+            this.forwardPipline = await this.gpu?.getPiplineShader("linear/forward.wgsl")
+        }
+        let tokenCount = input.buffer.size / (this.embeddingSize * 4)
+        if (!this.paramsForward) {
+            this.paramsForward = await this.gpu?.createBuffer(new Uint32Array([this.embeddingSize, this.vocabSize, tokenCount]), "uniform")
+        } else {
+            await this.paramsForward.changeData(new Uint32Array([this.embeddingSize, this.vocabSize, tokenCount]))
+        }
+        if (!this.output || this.output.buffer.size !== tokenCount * this.vocabSize * 4) {
+            if (this.output) this.output?.destroy();
+            this.output = await this.gpu?.createBuffer(new Float32Array(this.vocabSize * tokenCount), "storage")
+        }
+        const bindGroup = await this.gpu?.createBindGroup(this.forwardPipline!, [input, this.wBuffer!, this.bBuffer!, this.output!, this.paramsForward!], 0)
+        await this.gpu?.RunPipline(encoder, this.forwardPipline!, bindGroup!, this.vocabSize , tokenCount)
+        if (!en) await this.gpu?.SumbitEncoder(encoder!)
+        return this.output
+    }
+    async backward(outputGradient: Buffer, lr: number, encoderShared: GPUCommandEncoder | undefined) {
+        let tokenCount = outputGradient.buffer.size / (4 * this.vocabSize);
+        let encoder = encoderShared ? encoderShared : await this.gpu?.CreateEncoder()
+        if (!this.weightUpdate) {
+            this.weightUpdate = await this.gpu?.getPiplineShader("linear/updateWeights.wgsl")
+            this.biasUpdate = await this.gpu?.getPiplineShader("linear/updateBias.wgsl")
+            this.backwardPipline = await this.gpu?.getPiplineShader("linear/ComputeDinput.wgsl")
+        }
+        if (!this.paramsBackward) {
+            this.paramsBackward = await this.gpu?.createBuffer(new Float32Array([this.embeddingSize, this.vocabSize, tokenCount, lr]), "uniform")
+        } else {
+            await this.paramsBackward.changeData(new Float32Array([this.embeddingSize, this.vocabSize, tokenCount, lr]))
+        }
+        if (!this.gradOut || this.gradOut.buffer.size !== tokenCount * this.embeddingSize * 4) {
+            if (this.gradOut) this.gradOut.destroy()
+            this.gradOut = await this.gpu?.createBuffer(new Float32Array(this.embeddingSize * tokenCount), "storage")
+        }
+        const bindGroup2 = await this.gpu?.createBindGroup(this.backwardPipline!, [this.wBuffer!, outputGradient, this.gradOut!, this.paramsBackward!], 0)
+        await this.gpu?.RunPipline(encoder!, this.backwardPipline!, bindGroup2!, tokenCount, this.embeddingSize)
+        if (!encoderShared) await this.gpu?.SumbitEncoder(encoder!)
+        return this.gradOut
+    }
+    async updateParameter(outputGradient: Buffer, lr: number, encoderShared: GPUCommandEncoder | undefined) {
+        let tokenCount = outputGradient.buffer.size / (4 * this.vocabSize);
+        let encoder = encoderShared ? encoderShared : await this.gpu?.CreateEncoder()
+        if (!this.weightUpdate) {
+            this.weightUpdate = await this.gpu?.getPiplineShader("linear/updateWeights.wgsl")
+            this.biasUpdate = await this.gpu?.getPiplineShader("linear/updateBias.wgsl")
+        }
+        if (!this.paramsBackward) {
+            this.paramsBackward = await this.gpu?.createBuffer(new Float32Array([this.embeddingSize, this.vocabSize, tokenCount, lr]), "uniform")
+        } else {
+            await this.paramsBackward.changeData(new Float32Array([this.embeddingSize, this.vocabSize, tokenCount, lr]))
+        }
+
+        const bindGroup = await this.gpu?.createBindGroup(this.weightUpdate!, [this.wBuffer!, outputGradient, this.paramsBackward!, this.input!], 0)
+        const bindGroup1 = await this.gpu?.createBindGroup(this.biasUpdate!, [this.bBuffer!, outputGradient, this.paramsBackward!], 0)
+        await this.gpu?.RunPipline(encoder!, this.weightUpdate!, bindGroup!, Math.ceil(this.embeddingSize / 16), Math.ceil(this.vocabSize / 32))
+        await this.gpu?.RunPipline(encoder!, this.biasUpdate!, bindGroup1!, Math.ceil(this.vocabSize / 32))
+        if (!encoderShared) await this.gpu?.SumbitEncoder(encoder!)
+        return this.gradOut
+    }
+    async Save() {
+        const bias: any = await this.bBuffer?.read()
+        const weights: any = await this.wBuffer?.read()
+        SaveFileVector(this.configs.filename[0], weights!)
+        SaveFileVector(this.configs.filename[1], bias!)
+    }
+}

package/src/RMSNorm/index.ts

@@ -0,0 +1,56 @@
+import { GPUManager, Buffer } from "@imaduddin/webgpu-manager";
+
+
+export class RMSNorm {
+    rms: Buffer | undefined
+    input: Buffer | undefined
+    dInput: Buffer | undefined
+    params: Buffer | undefined
+    output: Buffer | undefined
+    forwardShader: GPUComputePipeline | undefined
+    backwardShader: GPUComputePipeline | undefined
+    gpu: GPUManager | undefined
+    constructor(public device: GPUDevice, public embeddingSize: number) {
+        this.gpu = new GPUManager(device, {
+            shaderDir: "./shaders/rmsnorm"
+        })
+    }
+    async forward(input: Buffer, encoderShared: GPUCommandEncoder | undefined) {
+        this.input = input;
+        let tokenCount = input.buffer.size / (this.embeddingSize * 4)
+        if (!this.rms || this.rms.buffer.size !== tokenCount * 4) {
+            if (this.rms) this.rms.destroy();
+            this.rms = await this.gpu?.createBuffer(new Float32Array(tokenCount), "storage")
+        }
+        if(!this.output || this.output.buffer.size !== tokenCount * this.embeddingSize * 4){
+            if(this.output) this.output.destroy();
+            if(this.dInput) this.dInput.destroy();
+            this.output = await this.gpu?.createBuffer(new Float32Array(tokenCount * this.embeddingSize), "storage")
+            this.dInput = await this.gpu?.createBuffer(new Float32Array(tokenCount * this.embeddingSize), "storage")
+        }
+        if(!this.params){
+            this.params = await this.gpu?.createBuffer(new Uint32Array([this.embeddingSize, tokenCount,0 ]), "uniform")
+        }else{
+            await this.params.changeData(new Uint32Array([this.embeddingSize, tokenCount, 0]))
+        }
+        if(!this.forwardShader){
+            this.forwardShader = await this.gpu?.getPiplineShader("forward.wgsl")
+            this.backwardShader = await this.gpu?.getPiplineShader("backward.wgsl")
+        }
+        const encoder = encoderShared ? encoderShared : await this.gpu?.CreateEncoder()
+        const bindGroup = await this.gpu?.createBindGroup(this.forwardShader!, [this.input!, this.output!, this.rms!, this.params!], 0)
+        await this.gpu?.RunPipline(encoder!, this.forwardShader!, bindGroup!, tokenCount);
+        if(!encoderShared){
+            this.gpu?.SumbitEncoder(encoder!)
+        }
+        return this.output
+    }
+    async backward(outputGradient : Buffer, encoderShared : GPUCommandEncoder | undefined){
+        const tokenCount = outputGradient.buffer.size / (this.embeddingSize * 4);
+        const encoder = encoderShared ? encoderShared : await this.gpu?.CreateEncoder()
+        const bindGroup = await this.gpu?.createBindGroup(this.backwardShader!, [this.input!, this.rms!, outputGradient, this.dInput!, this.params!] , 0)
+        await this.gpu?.RunPipline(encoder!, this.backwardShader!, bindGroup!, tokenCount)
+        if(!encoderShared) await this.gpu?.SumbitEncoder(encoder!)
+        return this.dInput
+    }
+}

package/src/SelfAttention/index.ts

@@ -0,0 +1,177 @@
+import { join } from "node:path";
+import { Linear } from "../Linear/index.js";
+import { GPUManager, Buffer } from "@imaduddin/webgpu-manager";
+
+export class SelfAttention {
+    query: Linear | undefined
+    value: Linear | undefined
+    key: Linear | undefined
+    v: Buffer | undefined
+    q: Buffer | undefined
+    weights: Buffer | undefined
+    k: Buffer | undefined
+    gpu: GPUManager | undefined
+    params: Buffer | undefined
+    Softparams: Buffer | undefined
+    Softbackparams: Buffer | undefined
+    Transparams: Buffer | undefined
+    matrixMat: GPUComputePipeline | undefined
+    transposeShader: GPUComputePipeline | undefined
+    addShader: GPUComputePipeline | undefined
+    softmaxShader: GPUComputePipeline | undefined
+    softmaxBackShader: GPUComputePipeline | undefined
+    qkT: Buffer | undefined
+    transPosedKey: Buffer | undefined
+    output: Buffer | undefined
+    qKTShader: GPUComputePipeline | undefined
+    dQ: Buffer | undefined
+    dK: Buffer | undefined
+    dV: Buffer | undefined
+    dWeights: Buffer | undefined
+    dScores: Buffer | undefined
+    dInput: Buffer | undefined
+    addParams: Buffer | undefined
+    constructor(public device: GPUDevice, public embeddingSize: number, public configs: {
+        dirname: string
+    }) {
+        this.query = new Linear(device, embeddingSize, embeddingSize, {
+            filename: [join(configs.dirname, "LINEAR_QUERY_WEIGHTS.bin"), join(configs.dirname, "LINEAR_QUERY_BIAS.bin")]
+        })
+        this.value = new Linear(device, embeddingSize, embeddingSize, {
+            filename: [join(configs.dirname, "LINEAR_VALUE_WEIGHTS.bin"), join(configs.dirname, "LINEAR_VALUE_BIAS.bin")]
+        })
+        this.key = new Linear(device, embeddingSize, embeddingSize, {
+            filename: [join(configs.dirname, "LINEAR_KEY_WEIGHTS.bin"), join(configs.dirname, "LINEAR_KEY_BIAS.bin")]
+        })
+        this.gpu = new GPUManager(device, {
+            shaderDir: "./shaders"
+        })
+    }
+    async softmaxBackward(weights: Buffer, dWeigths: Buffer, output: Buffer, tokenCount: number, enocder: GPUCommandEncoder) {
+        if (!this.Softbackparams) {
+            this.Softbackparams = await this.gpu?.createBuffer(new Uint32Array([tokenCount, this.embeddingSize, 0]), "uniform")
+        } else {
+            await this.Softbackparams.changeData(new Uint32Array([tokenCount, this.embeddingSize, 0]))
+        }
+        if (!this.softmaxBackShader) {
+            this.softmaxBackShader = await this.gpu?.getPiplineShader("selfAttention/SB.wgsl")
+        }
+        const bind = await this.gpu?.createBindGroup(this.softmaxBackShader!, [weights!, dWeigths!, output!, this.Softbackparams!], 0)
+        await this.gpu?.RunPipline(enocder, this.softmaxBackShader!, bind!, Math.ceil(tokenCount / 16), Math.ceil(tokenCount / 16))
+
+    }
+    async MatMut(a: Buffer, b: Buffer, output: Buffer, M: number, N: number, K: number, encoderS: GPUCommandEncoder, isqK: boolean) {
+        if (!this.params) {
+            this.params = await this.gpu?.createBuffer(new Uint32Array([M, N, K]), "uniform")
+            this.qKTShader = await this.gpu?.getPiplineShader("selfAttention/qK.wgsl")
+            this.matrixMat = await this.gpu?.getPiplineShader("selfAttention/matMul.wgsl")
+        } else {
+            this.params.changeData(new Uint32Array([M, N, K]))
+        }
+        const encoder = encoderS ? encoderS : await this.gpu?.CreateEncoder()
+        const bindGroup = await this.gpu?.createBindGroup(isqK ? this.qKTShader! : this.matrixMat!, [a!, b!, output!, this.params!], 0)
+        await this.gpu?.RunPipline(encoder, isqK ? this.qKTShader! : this.matrixMat!, bindGroup!, N, M)
+    }
+    async transpose(a: Buffer, output: Buffer, M: number, N: number, encoder: GPUCommandEncoder) {
+        if (!this.Transparams) {
+            this.Transparams = await this.gpu?.createBuffer(new Uint32Array([M, N, 0]), "uniform")
+            this.matrixMat = await this.gpu?.getPiplineShader("selfAttention/matMul.wgsl")
+            this.transposeShader = await this.gpu?.getPiplineShader("selfAttention/transpose.wgsl")
+        } else {
+            this.Transparams.changeData(new Uint32Array([M, N, 0]))
+        }
+        const bindGroup = await this.gpu?.createBindGroup(this.transposeShader!, [a!, output!, this.Transparams!], 0)
+        await this.gpu?.RunPipline(encoder!, this.transposeShader!, bindGroup!, N, M)
+
+    }
+    async forward(input: Buffer) {
+        let tokenCount = input.buffer.size / (this.embeddingSize * 4)
+        this.q = await this.query?.forward(input, undefined)
+        this.k = await this.key?.forward(input, undefined)
+        this.v = await this.value?.forward(input, undefined)
+        if (!this.softmaxShader) {
+            this.softmaxShader = await this.gpu?.getPiplineShader("softmax/forward.wgsl")
+        }
+        if (!this.qkT || this.qkT.buffer.size !== tokenCount * tokenCount * 4) {
+            if (this.qkT) this.qkT.destroy();
+            this.qkT = await this.gpu?.createBuffer(new Float32Array(tokenCount * tokenCount), "storage")
+        }
+        if (!this.transPosedKey || this.transPosedKey.buffer.size !== tokenCount * this.embeddingSize * 4) {
+            if (this.transPosedKey) this.transPosedKey.destroy();
+            this.transPosedKey = await this.gpu?.createBuffer(new Float32Array(tokenCount * this.embeddingSize), "storage")
+        }
+        if (!this.weights || this.weights.buffer.size !== tokenCount * tokenCount * 4) {
+            if (this.weights) this.weights.destroy()
+            this.weights = await this.gpu?.createBuffer(new Float32Array(tokenCount * tokenCount), "storage")
+        }
+        if (!this.output || this.output.buffer.size !== tokenCount * this.embeddingSize * 4) {
+            if (this.output) this.output.destroy()
+            this.output = await this.gpu?.createBuffer(new Float32Array(tokenCount * this.embeddingSize), "storage")
+        }
+        if (!this.Softparams) {
+            this.Softparams = await this.gpu?.createBuffer(new Uint32Array([tokenCount, tokenCount, 0]), "uniform")
+        } else {
+            this.Softparams.changeData(new Uint32Array([tokenCount, tokenCount, 0]))
+        }
+        const encoder = await this.gpu?.CreateEncoder()
+        await this.transpose(this.k!, this.transPosedKey!, this.embeddingSize, tokenCount, encoder!)
+        await this.MatMut(this.q!, this.transPosedKey!, this.qkT!, tokenCount, tokenCount, this.embeddingSize, encoder!, true)
+        const bindGroup = await this.gpu?.createBindGroup(this.softmaxShader!, [this.qkT!, this.weights!, this.Softparams!], 0)
+        await this.gpu?.RunPipline(encoder!, this.softmaxShader!, bindGroup!, tokenCount)
+        await this.MatMut(this.weights!, this.v!, this.output!, tokenCount, this.embeddingSize, tokenCount, encoder!, false)
+        await this.gpu?.SumbitEncoder(encoder!)
+        return this.output
+    }
+    async Save() {
+        await this.query?.Save()
+        await this.key?.Save()
+        await this.value?.Save()
+    }
+    async add3(a: Buffer, b: Buffer, c: Buffer, out: Buffer, size: number, encoder: GPUCommandEncoder) {
+        if (!this.addShader) {
+            this.addShader = await this.gpu!.getPiplineShader("selfAttention/add3.wgsl");
+            this.addParams = await this.gpu!.createBuffer(new Uint32Array([size]), "uniform");
+        } else {
+            await this.addParams!.changeData(new Uint32Array([size]));
+        }
+        const bindGroup = await this.gpu!.createBindGroup(this.addShader!, [a, b, c, out, this.addParams!], 0);
+        await this.gpu!.RunPipline(encoder, this.addShader!, bindGroup, Math.ceil(size / 256));
+    }
+    async backward(outputGradient: Buffer, lr: number) {
+        const D = this.embeddingSize;
+        const T = outputGradient.buffer.size / (D * 4);   
+        if (!this.dV || this.dV.buffer.size !== T * D * 4) {
+            this.dV?.destroy();
+            this.dQ?.destroy();
+            this.dK?.destroy();
+            this.dWeights?.destroy();
+            this.dScores?.destroy();
+            this.dInput?.destroy();
+
+            this.dV = await this.gpu!.createBuffer(new Float32Array(T * D), "storage");
+            this.dQ = await this.gpu!.createBuffer(new Float32Array(T * D), "storage");
+            this.dK = await this.gpu!.createBuffer(new Float32Array(T * D), "storage");
+            this.dWeights = await this.gpu!.createBuffer(new Float32Array(T * T), "storage");
+            this.dScores = await this.gpu!.createBuffer(new Float32Array(T * T), "storage");
+            this.dInput = await this.gpu!.createBuffer(new Float32Array(T * D), "storage");
+        }
+        const encoder = await this.gpu!.CreateEncoder();
+        await this.transpose(this.weights!, this.qkT!, T, T, encoder);
+        await this.MatMut(this.qkT!, outputGradient, this.dV!, T, D, T, encoder, false);
+        await this.transpose(this.v!, this.transPosedKey!, D, T, encoder);
+        await this.MatMut(outputGradient, this.transPosedKey!, this.dWeights!, T, T, D, encoder, false);
+        await this.softmaxBackward(this.weights!, this.dWeights!, this.dScores!, T, encoder);
+        await this.MatMut(this.dScores!, this.k!, this.dQ!, T, D, T, encoder, false);
+        await this.transpose(this.dScores!, this.qkT!, T, T, encoder);       
+        await this.MatMut(this.qkT!, this.q!, this.dK!, T, D, T, encoder, false);
+        const dInputQ = await this.query!.backward(this.dQ!, lr, encoder);
+        const dInputK = await this.key!.backward(this.dK!, lr, encoder);
+        const dInputV = await this.value!.backward(this.dV!, lr, encoder);
+        await this.add3(dInputQ!, dInputK!, dInputV!, this.dInput!, T * D, encoder);
+        await this.query!.updateParameter(this.dQ!, lr, encoder);
+        await this.key!.updateParameter(this.dK!, lr, encoder);
+        await this.value!.updateParameter(this.dV!, lr, encoder);
+        await this.gpu!.SumbitEncoder(encoder!);
+        return this.dInput;
+    }
+}