@imaduddin/fastbrain 0.0.0-stage → 1.0.0

package/src/PositionEmbedding/index.js

@@ -0,0 +1,65 @@
+import { OperationManager } from "../OperationManager.js"
+import { getVector, SaveVector } from "../getVec.js"
+function GenerateVectors(embeddingSize, contextSize) {
+    const vec = new Float32Array(embeddingSize * contextSize).map(() => (Math.random() * 2 - 1) * 0.02)
+    return vec
+}
+
+export class PositionEmbedding {
+    vectors;
+    vecBuffer;
+    gpu;
+    embeddingSize;
+    contextSize;
+    params;
+    forwardShader;
+    backwardShader;
+    configs;
+    constructor(device, embeddingSize, contextSize, configs) {
+        this.embeddingSize = embeddingSize;
+        const total = embeddingSize * contextSize;
+        this.configs = configs;
+        this.contextSize = contextSize;
+        if (configs && configs.save) {
+            const vec = getVector(configs.save.filename, total);
+            if (vec) {
+                this.vectors = vec;
+
+            } else {
+                this.vectors = GenerateVectors(embeddingSize, contextSize)
+                SaveVector(this.configs.save.filename, this.vectors)
+            }
+
+        } else {
+            this.vectors = GenerateVectors(embeddingSize, contextSize)
+
+        }
+        this.gpu = new OperationManager(device)
+        this.forwardShader = this.gpu.getShaderPipline("./src/shaders/PF.wgsl")
+        this.backwardShader = this.gpu.getShaderPipline("./src/shaders/PB.wgsl")
+    }
+    async forward(input, encoderMain) {
+        const encoder = encoderMain ? encoderMain : this.gpu.CreateEncoder()
+        if (!this.params) this.params = await this.gpu.createUniFormBuffer(4 * 4)
+        if (!this.vecBuffer) {
+            this.vecBuffer = await this.gpu.createBuffer(Float32Array.BYTES_PER_ELEMENT * this.contextSize * this.embeddingSize);
+            await this.gpu.WriteBuffer(this.vecBuffer, this.vectors)
+        }
+        const size = Math.ceil((input.size / 4) / 64);
+        const bindGroup = await this.gpu.createBindGroup(this.forwardShader, 0, [this.vecBuffer, input.mainBuffer])
+        await this.gpu.RunPipline(encoder, this.forwardShader, bindGroup, size);
+        if(!encoderMain) await this.gpu.submitQueue(encoder);
+        return input
+    }
+    async backward(outputGradient, lr, encoderMain){
+        const encoder = encoderMain ? encoderMain : await this.gpu.CreateEncoder()
+        await this.gpu.WriteBuffer(this.params, new Float32Array([lr,0]))
+        const bindGroup = await this.gpu.createBindGroup(this.backwardShader, 0, [this.vecBuffer, outputGradient.mainBuffer, this.params])
+        await this.gpu.RunPipline(encoder, this.backwardShader, bindGroup, Math.ceil((outputGradient.size / 4) / 64))
+        if(!encoderMain) await this.gpu.submitQueue(encoder);
+    }
+    async Save() {
+        const vec = await this.gpu.readBuffer(Float32Array.BYTES_PER_ELEMENT * this.contextSize * this.embeddingSize, this.vecBuffer)
+        SaveVector(this.configs.save.filename, vec)
+    }
+}

package/src/RMSNorm/index.d.ts

@@ -0,0 +1,30 @@
+export class RMSNorm {
+    constructor(device: GPUAdapter, embeddingSize: number, {
+        save: {
+            filename: string
+        }
+    });
+    async forward(input: {
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }): Promise<{
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }>
+    async backward(outputGradient: {
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }): Promise<{
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }>
+    async Save(): Promise<void>
+}

package/src/RMSNorm/index.js

@@ -0,0 +1,96 @@
+import { OperationManager } from '../OperationManager.js'
+import fs from 'fs'
+import { SaveVector, getVector } from '../getVec.js'
+
+export class RMSNorm {
+  embeddingSize;
+  reductionShader;
+  forwardShader;
+  backwardShader;
+  op;
+  gamma;
+  gammabuffer;
+  configs;
+  constructor(device, embeddingSize, configs) {
+    this.configs = configs
+    this.op = new OperationManager(device);
+    this.embeddingSize = embeddingSize;
+    this.reductionShader = this.op.getShaderPipline("./src/shaders/Norm/sumSquare.wgsl", {
+      embeddingSize: embeddingSize,
+    })
+    this.forwardShader = this.op.getShaderPipline("./src/shaders/Norm/rmsForward.wgsl", {
+      embeddingSize: embeddingSize
+    })
+    this.backwardShader = this.op.getShaderPipline("./src/shaders/Norm/rmsBackward.wgsl", {
+      embeddingSize: embeddingSize
+    })
+    this.updateGamma = this.op.getShaderPipline("./src/shaders/Norm/updateGamma.wgsl", {
+      embeddingSize: embeddingSize
+    })
+    if (configs) {
+      const ga = getVector(configs.save.filename, embeddingSize);
+      if (!ga) {
+        this.gamma = new Float32Array(embeddingSize).fill(1);
+        SaveVector(configs.save.filename, this.gamma)
+      } else {
+        this.gamma = ga
+      }
+    } else {
+      this.gamma = new Float32Array(embeddingSize).fill(1)
+    }
+  }
+  async forward(input) {
+    this.input = input;
+    if (!this.gammabuffer) {
+      this.gammabuffer = await this.op.createBuffer(4 * this.embeddingSize);
+      await this.op.WriteBuffer(this.gammabuffer, this.gamma);
+    }
+    let context = input.size / (this.embeddingSize * 4)
+    if (!this.rms) this.rms = await this.op.createBuffer(4 * context);
+    if (this.rms.size !== 4 * context) {
+      this.rms.destroy()
+      this.rms = await this.op.createBuffer(4 * context);
+    }
+    if (!this.output) this.output = await this.op.ManageBuffer(input.size);
+    if (this.output.size !== 4 * this.embeddingSize * context || this.output.isDestroyed() || this.output.size == 0) {
+      this.output.destroy()
+      this.output = await this.op.ManageBuffer(input.size)
+    }
+    await this.output.count()
+    const bindGroup = await this.op.createBindGroup(this.reductionShader, 0, [input.mainBuffer, this.rms])
+    const encoder = await this.op.CreateEncoder()
+    await this.op.RunPipline(encoder, this.reductionShader, bindGroup, context)
+    const bindGroup1 = await this.op.createBindGroup(this.forwardShader, 0, [this.rms, input.mainBuffer, this.output.mainBuffer, this.gammabuffer])
+    await this.op.RunPipline(encoder, this.forwardShader, bindGroup1, Math.ceil((input.size / 4) / 64));
+    await this.op.submitQueue(encoder);
+    return this.output
+  }
+  async backward(outputGradient, lr) {
+    if (!this.input || !this.input.mainBuffer || this.input.isDestroyed?.()) {
+      throw new Error("RMSNorm.backward: input buffer is gone");
+    }
+    if (!this.rms || this.rms.isDestroyed?.()) {
+      throw new Error("RMSNorm.backward: rms buffer is gone");
+    }
+    if (!this.gradOut) this.gradOut = await this.op.ManageBuffer(outputGradient.size);
+    if (this.gradOut.size !== outputGradient.size || this.gradOut.isDestroyed() || this.gradOut.size == 0) {
+      this.gradOut.destroy()
+      this.gradOut = await this.op.ManageBuffer(outputGradient.size)
+    }
+    await this.gradOut.count()
+    if (!this.params) this.params = await this.op.createBuffer(4 * 2)
+    await this.op.WriteBuffer(this.params, new Float32Array([lr, 0]))
+    const bindGroup = await this.op.createBindGroup(this.backwardShader, 0, [this.input.mainBuffer, outputGradient.mainBuffer, this.rms, this.gammabuffer, this.gradOut.mainBuffer]);
+    const bindGroup2 = await this.op.createBindGroup(this.updateGamma, 0, [this.input.mainBuffer, outputGradient.mainBuffer, this.rms, this.gammabuffer, this.params]);
+    const encoder = await this.op.CreateEncoder()
+    await this.op.RunPipline(encoder, this.backwardShader, bindGroup, Math.ceil((this.input.size) / (this.embeddingSize * 4)))
+    await this.op.RunPipline(encoder, this.updateGamma, bindGroup2, Math.ceil(this.embeddingSize / 64))
+    await this.op.submitQueue(encoder)
+    return this.gradOut
+  }
+  async Save() {
+    const read = await this.op.readBuffer(4 * this.embeddingSize, this.gammabuffer);
+    SaveVector(this.configs.save.filename, read)
+  }
+
+}

package/src/Relu/index.d.ts

@@ -0,0 +1,25 @@
+export class Relu {
+    constructor(device: GPUAdapter, embeddingSize: number);
+    async forward(input: {
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }): Promise<{
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }>
+    async backward(outGradinet: {
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }): Promise<{
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }>
+}

package/src/Relu/index.js

@@ -0,0 +1,110 @@
+import { OperationManager } from "../OperationManager.js";
+
+export class Relu {
+    embeddingSize;
+    paramBuffer;
+    oldParam;
+    forwardPipline;
+    gpu;
+    backwardPipline;
+    input;
+    constructor(device, embeddingSize) {
+        this.embeddingSize = embeddingSize
+        this.gpu = new OperationManager(device);
+        this.device = device;
+        this.forwardPipline = this.gpu.getShaderPipline("./src/shaders/RF.wgsl")
+        this.backwardPipline = this.gpu.getShaderPipline("./src/shaders/RB.wgsl")
+    }
+    async forward(input, encoderMain) {
+        const encoder = encoderMain ? encoderMain : this.device.createCommandEncoder()
+        const tokenCount = input.mainBuffer.size / (this.embeddingSize * 4)
+        if (!this.oldParam) {
+            this.oldParam = new Uint32Array([0, 0, 0]);
+            this.paramBuffer = await this.gpu.createUniFormBuffer(4 * 3);
+        }
+        if (!this.output) this.output = await this.gpu.ManageBuffer(input.mainBuffer.size);
+        if (this.output.mainBuffer.size !== input.mainBuffer.size || this.output.isDestroyed() || this.output.size == 0) {
+            this.output.destroy()
+            this.output = await this.gpu.ManageBuffer(input.mainBuffer.size);
+        }
+        await this.output.count()
+        if (this.oldParam[1] !== tokenCount) {
+            this.oldParam = new Uint32Array([this.embeddingSize, tokenCount, 0])
+            await this.gpu.WriteBuffer(this.paramBuffer, this.oldParam)
+        }
+        const bind = await this.device.createBindGroup({
+            layout: await this.forwardPipline.getBindGroupLayout(0),
+            entries: [
+                {
+                    binding: 0,
+                    resource: {
+                        buffer: input.mainBuffer
+                    }
+                },
+                {
+                    binding: 1,
+                    resource: {
+                        buffer: this.paramBuffer
+                    }
+                }, {
+                    binding: 2,
+                    resource: {
+                        buffer: this.output.mainBuffer
+                    }
+                }
+            ]
+        })
+        await this.gpu.RunPipline(encoder, this.forwardPipline, bind, Math.ceil((tokenCount * this.embeddingSize) / 256))
+        if(!encoderMain) this.device.queue.submit([encoder.finish()])
+        this.input = input
+        return this.output
+    }
+    async backward(inputGradient, encoderMain) {
+        const encoder = encoderMain ? encoderMain : this.device.createCommandEncoder()
+        const tokenCount = inputGradient.size / (this.embeddingSize * 4)
+        if (!this.oldParam) {
+            this.oldParam = new Uint32Array([0, 0, 0]);
+            this.paramBuffer = await this.gpu.createUniFormBuffer(4 * 3);
+        }
+        if(!this.gradOut) this.gradOut = await this.gpu.ManageBuffer(inputGradient.size);
+        if(this.gradOut.size !== inputGradient.size || this.gradOut.isDestroyed() || this.gradOut.size == 0){
+            this.gradOut.destroy();
+            this.gradOut = await this.gpu.ManageBuffer(inputGradient.size)
+        }
+        await this.gradOut.count()
+        if (this.oldParam[1] !== tokenCount) {
+            this.oldParam = new Uint32Array([this.embeddingSize, tokenCount, 0])
+            await this.gpu.WriteBuffer(this.paramBuffer, this.oldParam)
+        }
+        const bind = await this.device.createBindGroup({
+            layout: await this.backwardPipline.getBindGroupLayout(0),
+            entries: [
+                {
+                    binding: 0,
+                    resource: {
+                        buffer: this.input.mainBuffer
+                    }
+                },
+                {
+                    binding: 1,
+                    resource: {
+                        buffer: this.paramBuffer
+                    }
+                }, {
+                    binding: 2,
+                    resource: {
+                        buffer: this.gradOut.mainBuffer
+                    }
+                }, {
+                    binding: 3,
+                    resource: {
+                        buffer: inputGradient.mainBuffer
+                    }
+                }
+            ]
+        })
+        await this.gpu.RunPipline(encoder, this.backwardPipline, bind, Math.ceil((tokenCount * this.embeddingSize) / 256))
+        if(!encoderMain) await this.device.queue.submit([encoder.finish()])
+        return this.gradOut
+    }
+}

package/src/SelfAttention/index.d.ts

@@ -0,0 +1,28 @@
+export class SelfAttention {
+    constructor(device: GPUAdapter, embeddingSize: number, configs: {
+        dirname: string
+    })
+    async forward(vec: {
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }): Promise<{
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }>
+    async backward(dO: {
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }, lr: number): Promise<{
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }>
+    async Save(): Promise<void>
+} 

package/src/SelfAttention/index.js

@@ -0,0 +1,138 @@
+import { Linear } from "../../index.js";
+import { join } from 'path'
+import { OperationManager } from "../OperationManager.js";
+
+export class SelfAttention {
+    query;
+    value;
+    key;
+    q;
+    k;
+    v;
+    embeddingSize;
+    params;
+    attentionscoresShader;
+    softmaxForward;
+    gpu;
+    input;
+    constructor(device, embeddingSize, configs) {
+        this.embeddingSize = embeddingSize;
+        const dd = [1, 2, 3].map((e) => {
+            return {
+                save: {
+                    filename: [join(configs.dirname, `SELF-W-${e}.bin`), join(configs.dirname, `SELF-B-${e}.bin`)]
+                }
+            }
+        })
+        this.q = new Linear(device, embeddingSize, embeddingSize, dd[0])
+        this.k = new Linear(device, embeddingSize, embeddingSize, dd[1])
+        this.v = new Linear(device, embeddingSize, embeddingSize, dd[2])
+        this.gpu = new OperationManager(device)
+        this.attentionscoresShader = this.gpu.getShaderPipline("./src/shaders/attentionShader/attentionSocres.wgsl")
+        this.softmaxForward = this.gpu.getShaderPipline("./src/shaders/SF.wgsl")
+        this.softmaxBackward = this.gpu.getShaderPipline("./src/shaders/SoftmaxBackward.wgsl")
+        this.outputForward = this.gpu.getShaderPipline("./src/shaders/attentionShader/CompueFinalForwardOutput.wgsl")
+        this.dVShader = this.gpu.getShaderPipline("./src/shaders/attentionShader/ComputeDV.wgsl")
+        this.dAshader = this.gpu.getShaderPipline("./src/shaders/attentionShader/ComputeDA.wgsl")
+        this.dQshader = this.gpu.getShaderPipline("./src/shaders/attentionShader/ComputedQ.wgsl")
+        this.dKShader = this.gpu.getShaderPipline("./src/shaders/attentionShader/ComputeK.wgsl")
+        this.sumShader = this.gpu.getShaderPipline("./src/shaders/attentionShader/sumGradient.wgsl")
+    }
+    async forward(vec, encoderMain) {
+        this.input = vec.mainBuffer;
+        const tokenCount = vec.mainBuffer.size / (this.embeddingSize * 4)
+        if (!this.asParams) this.asParams = await this.gpu.createUniFormBuffer(4 * 3);
+        if (!this.softmaxParams) this.softmaxParams = await this.gpu.createUniFormBuffer(4 * 3);
+        const encoder = encoderMain ? encoderMain : await this.gpu.CreateEncoder()
+        if (!this.scroes) {
+            this.scroes = await this.gpu.createBuffer(tokenCount * tokenCount * 4)
+        }
+        if (this.scroes.size !== tokenCount * tokenCount * 4) {
+            this.scroes.destroy()
+            this.scroes = await this.gpu.createBuffer(tokenCount * tokenCount * 4)
+        }
+        await this.gpu.WriteBuffer(this.asParams, new Uint32Array([tokenCount, this.embeddingSize, 0]))
+        await this.gpu.WriteBuffer(this.softmaxParams, new Uint32Array([tokenCount, tokenCount, 0]))
+        this.query = await this.q.forward(vec, encoderMain)
+        this.key = await this.k.forward(vec, encoderMain)
+        this.value = await this.v.forward(vec, encoderMain)
+        const bindAttention = await this.gpu.createBindGroup(this.attentionscoresShader, 0, [this.query.mainBuffer, this.key.mainBuffer, this.scroes, this.asParams])
+        await this.gpu.RunPipline(encoder, this.attentionscoresShader, bindAttention, Math.ceil(tokenCount / 16), Math.ceil(tokenCount / 16))
+        if (this.weights) this.weights.destroy()
+        this.weights = await this.gpu.createBuffer(tokenCount * tokenCount * 4)
+        const bindSF = await this.gpu.createBindGroup(this.softmaxForward, 0, [this.scroes, this.weights, this.softmaxParams]);
+        await this.gpu.RunPipline(encoder, this.softmaxForward, bindSF, tokenCount)
+        if (!this.finalOutput) {
+            this.finalOutput = await this.gpu.ManageBuffer(tokenCount * this.embeddingSize * 4)
+        }
+        if (this.finalOutput.mainBuffer.size !== tokenCount * this.embeddingSize * 4 || this.finalOutput.size == 0 || this.finalOutput.isDestroyed()) {
+            this.finalOutput = await this.gpu.ManageBuffer(tokenCount * this.embeddingSize * 4)
+        }
+        await this.finalOutput.count()
+        const outputBind = await this.gpu.createBindGroup(this.outputForward, 0, [this.weights, this.value.mainBuffer, this.finalOutput.mainBuffer, this.asParams]);
+        await this.gpu.RunPipline(encoder, this.outputForward, outputBind, Math.ceil((tokenCount * this.embeddingSize) / 256))
+        if(!encoderMain) await this.gpu.submitQueue(encoder)
+        return this.finalOutput
+    }
+    async backward(dO, lr, encoderMain) {
+        const tokenCount = dO.size / (this.embeddingSize * 4);
+        const encoder = encoderMain ? encoderMain : await this.gpu.CreateEncoder();
+        if (!this.dV) this.dV = await this.gpu.createBuffer(dO.size);
+        if (!this.dA) this.dA = await this.gpu.createBuffer(tokenCount * tokenCount * 4);
+        if (!this.dS) this.dS = await this.gpu.createBuffer(tokenCount * tokenCount * 4);
+        if (!this.dQ) this.dQ = await this.gpu.createBuffer(tokenCount * this.embeddingSize * 4);
+        if (!this.dK) this.dK = await this.gpu.createBuffer(this.embeddingSize * tokenCount * 4)
+        if (this.dV.size !== dO.size || this.dV.size == 0) {
+            this.dV.destroy()
+            this.dV = await this.gpu.createBuffer(dO.size);
+        }
+        if (this.dA.size !== tokenCount * tokenCount * 4 || this.dA.size == 0) {
+            this.dA.destroy()
+            this.dA = await this.gpu.createBuffer(tokenCount * tokenCount * 4);
+        }
+        if (this.dS.size !== tokenCount * tokenCount * 4 || this.dS.size == 0) {
+            this.dS.destroy()
+            this.dS = await this.gpu.createBuffer(tokenCount * tokenCount * 4);
+        }
+        if (this.dQ.size !== tokenCount * this.embeddingSize * 4 || this.dQ.size == 0) {
+            this.dQ.destroy()
+            this.dQ = await this.gpu.createBuffer(tokenCount * this.embeddingSize * 4);
+        }
+        if (this.dK.size !== tokenCount * this.embeddingSize * 4 || this.dK.size == 0) {
+            this.dK.destroy()
+            this.dK = await this.gpu.createBuffer(tokenCount * this.embeddingSize * 4);
+        }
+        const DvBind = await this.gpu.createBindGroup(this.dVShader, 0, [this.weights, dO.mainBuffer, this.dV, this.asParams]);
+        await this.gpu.RunPipline(encoder, this.dVShader, DvBind, Math.ceil(tokenCount / 16), Math.ceil(this.embeddingSize / 16));
+        const DaBind = await this.gpu.createBindGroup(this.dAshader, 0, [dO.mainBuffer, this.value.mainBuffer, this.dA, this.asParams]);
+        await this.gpu.RunPipline(encoder, this.dAshader, DaBind, Math.ceil(tokenCount / 16), Math.ceil(tokenCount / 16))
+        const dSBind = await this.gpu.createBindGroup(this.softmaxBackward, 0, [this.weights, this.dA, this.dS, this.asParams])
+        await this.gpu.RunPipline(encoder, this.softmaxBackward, dSBind, Math.ceil((tokenCount * tokenCount) / 256))
+        const dQbind = await this.gpu.createBindGroup(this.dQshader, 0, [this.dS, this.key.mainBuffer, this.dQ, this.asParams]);
+        await this.gpu.RunPipline(encoder, this.dQshader, dQbind, Math.ceil(tokenCount / 16), Math.ceil(this.embeddingSize / 16))
+        const dKbind = await this.gpu.createBindGroup(this.dKShader, 0, [this.dS, this.query.mainBuffer, this.dK, this.asParams]);
+        await this.gpu.RunPipline(encoder, this.dKShader, dKbind, Math.ceil(tokenCount / 16), Math.ceil(this.embeddingSize / 16))
+        const dq = await this.q.backward({ mainBuffer: this.dQ, size: this.dQ.size}, lr, encoderMain)
+        const dk = await this.k.backward({ mainBuffer: this.dK, size: this.dK.size}, lr, encoderMain)
+        const dv = await this.v.backward({ mainBuffer: this.dV, size: this.dV.size}, lr, encoderMain)
+        if(!this.gradOut) this.gradOut = await this.gpu.ManageBuffer(tokenCount * this.embeddingSize * 4)
+        if(this.gradOut.size !== tokenCount * 4 * this.embeddingSize || this.gradOut.isDestroyed() || this.gradOut.size == 0){
+            this.gradOut.destroy()
+            this.gradOut = await this.gpu.ManageBuffer(tokenCount * 4 * this.embeddingSize)
+        }
+        await this.gradOut.count()
+        if(!dq || !dk || !dq || !this.gradOut){
+            throw new Error("Something is wrong!")
+        }
+        const outputBind = await this.gpu.createBindGroup(this.sumShader, 0, [dq.mainBuffer, dk.mainBuffer, dv.mainBuffer, this.gradOut.mainBuffer, this.asParams]);
+        await this.gpu.RunPipline(encoder, this.sumShader, outputBind, Math.ceil((tokenCount * this.embeddingSize / 256)));
+        if(!encoderMain) await this.gpu.submitQueue(encoder)
+        return this.gradOut
+
+    }
+    async Save() {
+        await this.q.Save()
+        await this.k.Save()
+        await this.v.Save()
+    }
+}

package/src/Softmax/index.d.ts

@@ -0,0 +1,26 @@
+export class Softmax {
+    constructor(device: GPUAdapter, embeddingSize: number, vocabSize: number);
+    async forward(logit: {
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }, isCpuReadBack): Promise<{
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }>;
+    async backward(probs: {
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }, target: number[]): Promise<{
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }>
+    async crossEntropy(probs: GPUBuffer, target: number[]): Promise<{ loss: Function, acc: Function, clear : Function }>
+}