@imaduddin/fastbrain 1.1.1-patch → 2.0.0

package/dist/Embedding/index.d.ts

@@ -0,0 +1,24 @@
+import { Buffer, GPUManager } from '@imaduddin/webgpu-manager';
+export declare class Embedding {
+    device: GPUDevice;
+    embeddingSize: number;
+    vocabSize: number;
+    configs: {
+        filename: string;
+    };
+    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(device: GPUDevice, embeddingSize: number, vocabSize: number, configs: {
+        filename: string;
+    });
+    forward(input: number[], encoderShared: GPUCommandEncoder | undefined): Promise<Buffer | undefined>;
+    backward(dInput: Buffer, lr: number, encoderShared: GPUCommandEncoder | undefined): Promise<void>;
+    Save(): Promise<void>;
+}

package/dist/Embedding/index.js

@@ -0,0 +1,93 @@
+import { GetVectorsForFile, SaveFileVector, initRandomVector } from '../fileManager.js';
+import { GPUManager } from '@imaduddin/webgpu-manager';
+export class Embedding {
+    device;
+    embeddingSize;
+    vocabSize;
+    configs;
+    vectors;
+    vectorsBuffer;
+    gpu;
+    param;
+    paramBackward;
+    tokenBuffer;
+    output;
+    forwardShader;
+    backwardPipline;
+    constructor(device, embeddingSize, vocabSize, configs) {
+        this.device = device;
+        this.embeddingSize = embeddingSize;
+        this.vocabSize = vocabSize;
+        this.configs = configs;
+        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, encoderShared) {
+        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 = 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, lr, encoderShared) {
+        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 = await this.vectorsBuffer?.read();
+        await SaveFileVector(this.configs.filename, vec);
+    }
+}

package/dist/FFN/index.d.ts

@@ -0,0 +1,20 @@
+import { Linear } from "../Linear/index.js";
+import { Gelu as Relu } from '../Gelu/index.js';
+import { Buffer, GPUManager } from '@imaduddin/webgpu-manager';
+export declare class FFN {
+    device: GPUDevice;
+    dimSize: number;
+    configs: {
+        dirname: string;
+    };
+    linear1: Linear | undefined;
+    linear2: Linear | undefined;
+    activation: Relu | undefined;
+    gpu: GPUManager | undefined;
+    constructor(device: GPUDevice, dimSize: number, configs: {
+        dirname: string;
+    });
+    forward(input: Buffer, encoderShared: GPUCommandEncoder | undefined): Promise<Buffer | undefined>;
+    backward(outputGradient: Buffer, lr: number, encoderShared: GPUCommandEncoder | undefined): Promise<Buffer | undefined>;
+    Save(): Promise<void>;
+}

package/dist/FFN/index.js

@@ -0,0 +1,51 @@
+import { Linear } from "../Linear/index.js";
+import { Gelu as Relu } from '../Gelu/index.js';
+import { join } from "node:path";
+import { GPUManager } from '@imaduddin/webgpu-manager';
+export class FFN {
+    device;
+    dimSize;
+    configs;
+    linear1;
+    linear2;
+    activation;
+    gpu;
+    constructor(device, dimSize, configs) {
+        this.device = device;
+        this.dimSize = dimSize;
+        this.configs = configs;
+        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, encoderShared) {
+        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, lr, encoderShared) {
+        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/dist/Gelu/index.d.ts

@@ -0,0 +1,15 @@
+import { Buffer, GPUManager } from "@imaduddin/webgpu-manager";
+export declare class Gelu {
+    device: GPUDevice;
+    hiddenSize: number;
+    input: Buffer | undefined;
+    output: Buffer | undefined;
+    outputGrad: Buffer | undefined;
+    params: Buffer | undefined;
+    forwardShader: GPUComputePipeline | undefined;
+    backwardShader: GPUComputePipeline | undefined;
+    gpu: GPUManager | undefined;
+    constructor(device: GPUDevice, hiddenSize: number);
+    forward(input: Buffer, encoderShared: GPUCommandEncoder | undefined): Promise<Buffer | undefined>;
+    backward(gradient: Buffer, encoderShared: GPUCommandEncoder | undefined): Promise<Buffer | undefined>;
+}

package/dist/Gelu/index.js

@@ -0,0 +1,58 @@
+import { GPUManager } from "@imaduddin/webgpu-manager";
+export class Gelu {
+    device;
+    hiddenSize;
+    input;
+    output;
+    outputGrad;
+    params;
+    forwardShader;
+    backwardShader;
+    gpu;
+    constructor(device, hiddenSize) {
+        this.device = device;
+        this.hiddenSize = hiddenSize;
+        this.gpu = new GPUManager(device, {
+            shaderDir: "./shaders/relu"
+        });
+    }
+    async forward(input, encoderShared) {
+        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, encoderShared) {
+        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/dist/Linear/index.d.ts

@@ -0,0 +1,31 @@
+import { GPUManager, Buffer } from "@imaduddin/webgpu-manager";
+export declare class Linear {
+    device: GPUDevice;
+    embeddingSize: number;
+    vocabSize: number;
+    configs: {
+        filename: string[];
+    };
+    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(device: GPUDevice, embeddingSize: number, vocabSize: number, configs: {
+        filename: string[];
+    });
+    getParameter(): void;
+    forward(input: Buffer, en: GPUCommandEncoder | undefined): Promise<Buffer | undefined>;
+    backward(outputGradient: Buffer, lr: number, encoderShared: GPUCommandEncoder | undefined): Promise<Buffer | undefined>;
+    updateParameter(outputGradient: Buffer, lr: number, encoderShared: GPUCommandEncoder | undefined): Promise<Buffer | undefined>;
+    Save(): Promise<void>;
+}

package/dist/Linear/index.js

@@ -0,0 +1,136 @@
+import { GPUManager } from "@imaduddin/webgpu-manager";
+import { GetVectorsForFile, SaveFileVector, initRandomVector, getParameter } from '../fileManager.js';
+export class Linear {
+    device;
+    embeddingSize;
+    vocabSize;
+    configs;
+    bias;
+    weights;
+    wBuffer;
+    bBuffer;
+    gpu;
+    paramsForward;
+    paramsBackward;
+    output;
+    forwardPipline;
+    backwardPipline;
+    weightUpdate;
+    biasUpdate;
+    gradOut;
+    input;
+    constructor(device, embeddingSize, vocabSize, configs) {
+        this.device = device;
+        this.embeddingSize = embeddingSize;
+        this.vocabSize = vocabSize;
+        this.configs = configs;
+        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, en) {
+        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, lr, encoderShared) {
+        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, lr, encoderShared) {
+        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 = await this.bBuffer?.read();
+        const weights = await this.wBuffer?.read();
+        SaveFileVector(this.configs.filename[0], weights);
+        SaveFileVector(this.configs.filename[1], bias);
+    }
+}

package/dist/RMSNorm/index.d.ts

@@ -0,0 +1,16 @@
+import { GPUManager, Buffer } from "@imaduddin/webgpu-manager";
+export declare class RMSNorm {
+    device: GPUDevice;
+    embeddingSize: number;
+    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(device: GPUDevice, embeddingSize: number);
+    forward(input: Buffer, encoderShared: GPUCommandEncoder | undefined): Promise<Buffer | undefined>;
+    backward(outputGradient: Buffer, encoderShared: GPUCommandEncoder | undefined): Promise<Buffer | undefined>;
+}

package/dist/RMSNorm/index.js

@@ -0,0 +1,63 @@
+import { GPUManager } from "@imaduddin/webgpu-manager";
+export class RMSNorm {
+    device;
+    embeddingSize;
+    rms;
+    input;
+    dInput;
+    params;
+    output;
+    forwardShader;
+    backwardShader;
+    gpu;
+    constructor(device, embeddingSize) {
+        this.device = device;
+        this.embeddingSize = embeddingSize;
+        this.gpu = new GPUManager(device, {
+            shaderDir: "./shaders/rmsnorm"
+        });
+    }
+    async forward(input, encoderShared) {
+        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, encoderShared) {
+        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/dist/Relu/index.d.ts

@@ -0,0 +1,15 @@
+import { Buffer, GPUManager } from "@imaduddin/webgpu-manager";
+export declare class Relu {
+    device: GPUDevice;
+    hiddenSize: number;
+    input: Buffer | undefined;
+    output: Buffer | undefined;
+    outputGrad: Buffer | undefined;
+    params: Buffer | undefined;
+    forwardShader: GPUComputePipeline | undefined;
+    backwardShader: GPUComputePipeline | undefined;
+    gpu: GPUManager | undefined;
+    constructor(device: GPUDevice, hiddenSize: number);
+    forward(input: Buffer, encoderShared: GPUCommandEncoder | undefined): Promise<Buffer | undefined>;
+    backward(gradient: Buffer, encoderShared: GPUCommandEncoder | undefined): Promise<Buffer | undefined>;
+}

package/dist/Relu/index.js

@@ -0,0 +1,58 @@
+import { GPUManager } from "@imaduddin/webgpu-manager";
+export class Relu {
+    device;
+    hiddenSize;
+    input;
+    output;
+    outputGrad;
+    params;
+    forwardShader;
+    backwardShader;
+    gpu;
+    constructor(device, hiddenSize) {
+        this.device = device;
+        this.hiddenSize = hiddenSize;
+        this.gpu = new GPUManager(device, {
+            shaderDir: "./shaders/relu"
+        });
+    }
+    async forward(input, encoderShared) {
+        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, encoderShared) {
+        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/dist/SelfAttention/index.d.ts

@@ -0,0 +1,47 @@
+import { Linear } from "../Linear/index.js";
+import { GPUManager, Buffer } from "@imaduddin/webgpu-manager";
+export declare class SelfAttention {
+    device: GPUDevice;
+    embeddingSize: number;
+    configs: {
+        dirname: string;
+    };
+    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(device: GPUDevice, embeddingSize: number, configs: {
+        dirname: string;
+    });
+    softmaxBackward(weights: Buffer, dWeigths: Buffer, output: Buffer, tokenCount: number, enocder: GPUCommandEncoder): Promise<void>;
+    MatMut(a: Buffer, b: Buffer, output: Buffer, M: number, N: number, K: number, encoderS: GPUCommandEncoder, isqK: boolean): Promise<void>;
+    transpose(a: Buffer, output: Buffer, M: number, N: number, encoder: GPUCommandEncoder): Promise<void>;
+    forward(input: Buffer): Promise<Buffer | undefined>;
+    Save(): Promise<void>;
+    add3(a: Buffer, b: Buffer, c: Buffer, out: Buffer, size: number, encoder: GPUCommandEncoder): Promise<void>;
+    backward(outputGradient: Buffer, lr: number): Promise<Buffer | undefined>;
+}