mambacode.js 1.0.0 → 1.0.1

package/dist/utils/quantization.js

@@ -0,0 +1,116 @@
+/**
+ * quantization.ts – FP16 and Int8 quantization utilities.
+ */
+export function floatToFp16(val) {
+    const buf = new ArrayBuffer(4);
+    const f32 = new Float32Array(buf);
+    const u32 = new Uint32Array(buf);
+    f32[0] = val;
+    const bits = u32[0];
+    const sign = (bits >>> 31) & 0x1;
+    const exponent = (bits >>> 23) & 0xFF;
+    const mantissa = bits & 0x7FFFFF;
+    if (exponent === 255) {
+        return (sign << 15) | 0x7C00 | (mantissa ? 0x200 : 0);
+    }
+    const expAdj = exponent - 127 + 15;
+    if (expAdj >= 31) {
+        return (sign << 15) | 0x7C00;
+    }
+    if (expAdj <= 0) {
+        if (expAdj < -10) {
+            return sign << 15;
+        }
+        const shift = 14 - expAdj;
+        return (sign << 15) | ((mantissa | 0x800000) >> shift);
+    }
+    return (sign << 15) | (expAdj << 10) | (mantissa >> 13);
+}
+export function fp16ToFloat(val) {
+    const sign = (val >>> 15) & 0x1;
+    const exponent = (val >>> 10) & 0x1F;
+    const mantissa = val & 0x3FF;
+    if (exponent === 0) {
+        const f = mantissa / 1024.0;
+        return sign ? -f : f;
+    }
+    if (exponent === 31) {
+        return sign ? -Infinity : (mantissa ? NaN : Infinity);
+    }
+    const expUnbiased = exponent - 15;
+    const f = (1 + mantissa / 1024.0) * Math.pow(2, expUnbiased);
+    return sign ? -f : f;
+}
+export function quantizeFp16(f32) {
+    const out = new Uint16Array(f32.length);
+    for (let i = 0; i < f32.length; i++) {
+        out[i] = floatToFp16(f32[i]);
+    }
+    return out;
+}
+export function dequantizeFp16(fp16) {
+    const out = new Float32Array(fp16.length);
+    for (let i = 0; i < fp16.length; i++) {
+        out[i] = fp16ToFloat(fp16[i]);
+    }
+    return out;
+}
+export function quantizeInt8(f32) {
+    let maxAbs = 0;
+    for (let i = 0; i < f32.length; i++) {
+        const a = Math.abs(f32[i]);
+        if (a > maxAbs)
+            maxAbs = a;
+    }
+    const scale = maxAbs / 127.0 || 1.0;
+    const data = new Int8Array(f32.length);
+    for (let i = 0; i < f32.length; i++) {
+        data[i] = Math.max(-128, Math.min(127, Math.round(f32[i] / scale)));
+    }
+    return { data, scale };
+}
+export function dequantizeInt8(int8, scale) {
+    const out = new Float32Array(int8.length);
+    for (let i = 0; i < int8.length; i++) {
+        out[i] = int8[i] * scale;
+    }
+    return out;
+}
+export function quantizeInt8PerChannel(f32, numChannels) {
+    const channelSize = f32.length / numChannels;
+    const scales = new Float32Array(numChannels);
+    const data = new Int8Array(f32.length);
+    for (let c = 0; c < numChannels; c++) {
+        let maxAbs = 0;
+        const base = c * channelSize;
+        for (let j = 0; j < channelSize; j++) {
+            const a = Math.abs(f32[base + j]);
+            if (a > maxAbs)
+                maxAbs = a;
+        }
+        scales[c] = maxAbs / 127.0 || 1.0;
+        for (let j = 0; j < channelSize; j++) {
+            data[base + j] = Math.max(-128, Math.min(127, Math.round(f32[base + j] / scales[c])));
+        }
+    }
+    return { data, scales };
+}
+export function dequantizeInt8PerChannel(int8, scales, numChannels) {
+    const channelSize = int8.length / numChannels;
+    const out = new Float32Array(int8.length);
+    for (let c = 0; c < numChannels; c++) {
+        const base = c * channelSize;
+        for (let j = 0; j < channelSize; j++) {
+            out[base + j] = int8[base + j] * scales[c];
+        }
+    }
+    return out;
+}
+export function estimateMemory(numElements) {
+    return {
+        fp32: numElements * 4,
+        fp16: numElements * 2,
+        int8: numElements * 1,
+    };
+}
+//# sourceMappingURL=quantization.js.map

package/dist/utils/quantization.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"quantization.js","sourceRoot":"","sources":["../../src/utils/quantization.ts"],"names":[],"mappings":"AAAA;;GAEG;AAkBH,MAAM,UAAU,WAAW,CAAC,GAAW;IACnC,MAAM,GAAG,GAAG,IAAI,WAAW,CAAC,CAAC,CAAC,CAAC;IAC/B,MAAM,GAAG,GAAG,IAAI,YAAY,CAAC,GAAG,CAAC,CAAC;IAClC,MAAM,GAAG,GAAG,IAAI,WAAW,CAAC,GAAG,CAAC,CAAC;IACjC,GAAG,CAAC,CAAC,CAAC,GAAG,GAAG,CAAC;IACb,MAAM,IAAI,GAAG,GAAG,CAAC,CAAC,CAAE,CAAC;IAErB,MAAM,IAAI,GAAO,CAAC,IAAI,KAAK,EAAE,CAAC,GAAG,GAAG,CAAC;IACrC,MAAM,QAAQ,GAAG,CAAC,IAAI,KAAK,EAAE,CAAC,GAAG,IAAI,CAAC;IACtC,MAAM,QAAQ,GAAI,IAAI,GAAW,QAAQ,CAAC;IAE1C,IAAI,QAAQ,KAAK,GAAG,EAAE,CAAC;QACnB,OAAO,CAAC,IAAI,IAAI,EAAE,CAAC,GAAG,MAAM,GAAG,CAAC,QAAQ,CAAC,CAAC,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;IAC1D,CAAC;IAED,MAAM,MAAM,GAAG,QAAQ,GAAG,GAAG,GAAG,EAAE,CAAC;IAEnC,IAAI,MAAM,IAAI,EAAE,EAAE,CAAC;QACf,OAAO,CAAC,IAAI,IAAI,EAAE,CAAC,GAAG,MAAM,CAAC;IACjC,CAAC;IAED,IAAI,MAAM,IAAI,CAAC,EAAE,CAAC;QACd,IAAI,MAAM,GAAG,CAAC,EAAE,EAAE,CAAC;YAAC,OAAO,IAAI,IAAI,EAAE,CAAC;QAAC,CAAC;QACxC,MAAM,KAAK,GAAG,EAAE,GAAG,MAAM,CAAC;QAC1B,OAAO,CAAC,IAAI,IAAI,EAAE,CAAC,GAAG,CAAC,CAAC,QAAQ,GAAG,QAAQ,CAAC,IAAI,KAAK,CAAC,CAAC;IAC3D,CAAC;IAED,OAAO,CAAC,IAAI,IAAI,EAAE,CAAC,GAAG,CAAC,MAAM,IAAI,EAAE,CAAC,GAAG,CAAC,QAAQ,IAAI,EAAE,CAAC,CAAC;AAC5D,CAAC;AAED,MAAM,UAAU,WAAW,CAAC,GAAW;IACnC,MAAM,IAAI,GAAO,CAAC,GAAG,KAAK,EAAE,CAAC,GAAG,GAAG,CAAC;IACpC,MAAM,QAAQ,GAAG,CAAC,GAAG,KAAK,EAAE,CAAC,GAAG,IAAI,CAAC;IACrC,MAAM,QAAQ,GAAI,GAAG,GAAW,KAAK,CAAC;IAEtC,IAAI,QAAQ,KAAK,CAAC,EAAE,CAAC;QACjB,MAAM,CAAC,GAAG,QAAQ,GAAG,MAAM,CAAC;QAC5B,OAAO,IAAI,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;IACzB,CAAC;IAED,IAAI,QAAQ,KAAK,EAAE,EAAE,CAAC;QAClB,OAAO,IAAI,CAAC,CAAC,CAAC,CAAC,QAAQ,CAAC,CAAC,CAAC,CAAC,QAAQ,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,QAAQ,CAAC,CAAC;IAC1D,CAAC;IAED,MAAM,WAAW,GAAG,QAAQ,GAAG,EAAE,CAAC;IAClC,MAAM,CAAC,GAAG,CAAC,CAAC,GAAG,QAAQ,GAAG,MAAM,CAAC,GAAG,IAAI,CAAC,GAAG,CAAC,CAAC,EAAE,WAAW,CAAC,CAAC;IAC7D,OAAO,IAAI,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;AACzB,CAAC;AAED,MAAM,UAAU,YAAY,CAAC,GAAiB;IAC1C,MAAM,GAAG,GAAG,IAAI,WAAW,CAAC,GAAG,CAAC,MAAM,CAAC,CAAC;IACxC,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,GAAG,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;QAClC,GAAG,CAAC,CAAC,CAAC,GAAG,WAAW,CAAC,GAAG,CAAC,CAAC,CAAE,CAAC,CAAC;IAClC,CAAC;IACD,OAAO,GAAG,CAAC;AACf,CAAC;AAED,MAAM,UAAU,cAAc,CAAC,IAAiB;IAC5C,MAAM,GAAG,GAAG,IAAI,YAAY,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;IAC1C,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,IAAI,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;QACnC,GAAG,CAAC,CAAC,CAAC,GAAG,WAAW,CAAC,IAAI,CAAC,CAAC,CAAE,CAAC,CAAC;IACnC,CAAC;IACD,OAAO,GAAG,CAAC;AACf,CAAC;AAED,MAAM,UAAU,YAAY,CAAC,GAAiB;IAC1C,IAAI,MAAM,GAAG,CAAC,CAAC;IACf,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,GAAG,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;QAClC,MAAM,CAAC,GAAG,IAAI,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,CAAE,CAAC,CAAC;QAC5B,IAAI,CAAC,GAAG,MAAM;YAAE,MAAM,GAAG,CAAC,CAAC;IAC/B,CAAC;IAED,MAAM,KAAK,GAAG,MAAM,GAAG,KAAK,IAAI,GAAG,CAAC;IACpC,MAAM,IAAI,GAAI,IAAI,SAAS,CAAC,GAAG,CAAC,MAAM,CAAC,CAAC;IAExC,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,GAAG,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;QAClC,IAAI,CAAC,CAAC,CAAC,GAAG,IAAI,CAAC,GAAG,CAAC,CAAC,GAAG,EAAE,IAAI,CAAC,GAAG,CAAC,GAAG,EAAE,IAAI,CAAC,KAAK,CAAC,GAAG,CAAC,CAAC,CAAE,GAAG,KAAK,CAAC,CAAC,CAAC,CAAC;IACzE,CAAC;IAED,OAAO,EAAE,IAAI,EAAE,KAAK,EAAE,CAAC;AAC3B,CAAC;AAED,MAAM,UAAU,cAAc,CAAC,IAAe,EAAE,KAAa;IACzD,MAAM,GAAG,GAAG,IAAI,YAAY,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;IAC1C,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,IAAI,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;QACnC,GAAG,CAAC,CAAC,CAAC,GAAG,IAAI,CAAC,CAAC,CAAE,GAAG,KAAK,CAAC;IAC9B,CAAC;IACD,OAAO,GAAG,CAAC;AACf,CAAC;AAED,MAAM,UAAU,sBAAsB,CAAC,GAAiB,EAAE,WAAmB;IACzE,MAAM,WAAW,GAAG,GAAG,CAAC,MAAM,GAAG,WAAW,CAAC;IAC7C,MAAM,MAAM,GAAG,IAAI,YAAY,CAAC,WAAW,CAAC,CAAC;IAC7C,MAAM,IAAI,GAAK,IAAI,SAAS,CAAC,GAAG,CAAC,MAAM,CAAC,CAAC;IAEzC,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,WAAW,EAAE,CAAC,EAAE,EAAE,CAAC;QACnC,IAAI,MAAM,GAAG,CAAC,CAAC;QACf,MAAM,IAAI,GAAG,CAAC,GAAG,WAAW,CAAC;QAC7B,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,WAAW,EAAE,CAAC,EAAE,EAAE,CAAC;YACnC,MAAM,CAAC,GAAG,IAAI,CAAC,GAAG,CAAC,GAAG,CAAC,IAAI,GAAG,CAAC,CAAE,CAAC,CAAC;YACnC,IAAI,CAAC,GAAG,MAAM;gBAAE,MAAM,GAAG,CAAC,CAAC;QAC/B,CAAC;QACD,MAAM,CAAC,CAAC,CAAC,GAAG,MAAM,GAAG,KAAK,IAAI,GAAG,CAAC;QAClC,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,WAAW,EAAE,CAAC,EAAE,EAAE,CAAC;YACnC,IAAI,CAAC,IAAI,GAAG,CAAC,CAAC,GAAG,IAAI,CAAC,GAAG,CAAC,CAAC,GAAG,EAAE,IAAI,CAAC,GAAG,CAAC,GAAG,EACxC,IAAI,CAAC,KAAK,CAAC,GAAG,CAAC,IAAI,GAAG,CAAC,CAAE,GAAG,MAAM,CAAC,CAAC,CAAE,CAAC,CAC1C,CAAC,CAAC;QACP,CAAC;IACL,CAAC;IAED,OAAO,EAAE,IAAI,EAAE,MAAM,EAAE,CAAC;AAC5B,CAAC;AAED,MAAM,UAAU,wBAAwB,CAAC,IAAe,EAAE,MAAoB,EAAE,WAAmB;IAC/F,MAAM,WAAW,GAAG,IAAI,CAAC,MAAM,GAAG,WAAW,CAAC;IAC9C,MAAM,GAAG,GAAG,IAAI,YAAY,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;IAE1C,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,WAAW,EAAE,CAAC,EAAE,EAAE,CAAC;QACnC,MAAM,IAAI,GAAG,CAAC,GAAG,WAAW,CAAC;QAC7B,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,WAAW,EAAE,CAAC,EAAE,EAAE,CAAC;YACnC,GAAG,CAAC,IAAI,GAAG,CAAC,CAAC,GAAG,IAAI,CAAC,IAAI,GAAG,CAAC,CAAE,GAAG,MAAM,CAAC,CAAC,CAAE,CAAC;QACjD,CAAC;IACL,CAAC;IAED,OAAO,GAAG,CAAC;AACf,CAAC;AAED,MAAM,UAAU,cAAc,CAAC,WAAmB;IAC9C,OAAO;QACH,IAAI,EAAE,WAAW,GAAG,CAAC;QACrB,IAAI,EAAE,WAAW,GAAG,CAAC;QACrB,IAAI,EAAE,WAAW,GAAG,CAAC;KACxB,CAAC;AACN,CAAC"}

package/package.json

@@ -1,29 +1,30 @@
 {
   "name": "mambacode.js",
-  "version": "1.0.0",
-  "description": "High-performance JavaScript/WGSL Mamba SSM library for browser-based code model training and inference",
-  "main": "src/index.js",
+  "version": "1.0.1",
+  "description": "High-performance TypeScript/WGSL Mamba SSM library for browser-based code model training and inference",
+  "main": "dist/index.js",
+  "module": "dist/index.js",
+  "types": "dist/index.d.ts",
+  "exports": {
+    ".": {
+      "types": "./dist/index.d.ts",
+      "import": "./dist/index.js"
+    }
+  },
   "type": "module",
   "files": [
+    "dist",
     "src",
     "README.md",
     "LICENSE"
   ],
   "scripts": {
+    "build": "tsc",
     "test": "node --experimental-vm-modules node_modules/.bin/jest",
-    "lint": "eslint src/ tests/"
-  },
-  "keywords": [
-    "mamba",
-    "ssm",
-    "state-space-model",
-    "webgpu",
-    "wgsl",
-    "machine-learning",
-    "code-model",
-    "bpe",
-    "transformer-alternative"
-  ],
+    "lint": "eslint src/ tests/",
+    "prepublishOnly": "npm run build"
+  },
+  "keywords": ["mamba", "ssm", "state-space-model", "webgpu", "wgsl", "machine-learning", "code-model", "bpe", "transformer-alternative", "typescript"],
   "author": {
     "name": "Sean Hogg",
     "email": "seanhogg@gmail.com",
@@ -45,10 +46,34 @@
     "access": "public"
   },
   "devDependencies": {
+    "@types/jest": "^29.5.0",
+    "@webgpu/types": "^0.1.49",
+    "eslint": "^8.57.0",
     "jest": "^29.7.0",
-    "eslint": "^8.57.0"
+    "ts-jest": "^29.2.0",
+    "typescript": "^5.0.0"
   },
   "jest": {
-    "transform": {}
+    "preset": "ts-jest/presets/default-esm",
+    "extensionsToTreatAsEsm": [".ts"],
+    "transform": {
+      "^.+\\.tsx?$": [
+        "ts-jest",
+        {
+          "useESM": true,
+          "tsconfig": {
+            "module": "ES2022",
+            "moduleResolution": "bundler",
+            "strict": true,
+            "esModuleInterop": true,
+            "skipLibCheck": true,
+            "types": ["@webgpu/types", "jest"]
+          }
+        }
+      ]
+    },
+    "moduleNameMapper": {
+      "^(\\.{1,2}/.*)\\.js$": "$1"
+    }
   }
 }

package/src/index.ts

@@ -0,0 +1,59 @@
+/**
+ * MambaCode.js – Entry Point
+ */
+
+export { MambaModel }   from './model/mamba_model';
+export { MambaBlock }   from './model/mamba_block';
+
+export { MambaTrainer } from './training/trainer';
+export {
+    Tensor,
+    backward,
+    enableGrad,
+    noGrad,
+    clearTape,
+    recordOperation,
+    crossEntropyLoss,
+    crossEntropyGrad,
+} from './training/autograd';
+
+export { BPETokenizer } from './tokenizer/bpe';
+
+export {
+    initWebGPU,
+    createStorageBuffer,
+    createEmptyStorageBuffer,
+    createUniformBuffer,
+    createComputePipeline,
+    createBindGroup,
+    dispatchKernel,
+    readBuffer,
+    uploadBuffer,
+    cdiv,
+} from './utils/gpu_utils';
+
+export {
+    quantizeFp16,
+    dequantizeFp16,
+    floatToFp16,
+    fp16ToFloat,
+    quantizeInt8,
+    dequantizeInt8,
+    quantizeInt8PerChannel,
+    dequantizeInt8PerChannel,
+    estimateMemory,
+} from './utils/quantization';
+
+export { SELECTIVE_SCAN_FORWARD_WGSL, SELECTIVE_SCAN_BACKWARD_WGSL }
+    from './kernels/selective_scan';
+export { CONV1D_FORWARD_WGSL, CONV1D_BACKWARD_WGSL }
+    from './kernels/conv1d';
+export { LINEAR_FORWARD_WGSL, LINEAR_BACKWARD_WGSL }
+    from './kernels/linear_projection';
+export { WEIGHT_UPDATE_WGSL, GRAD_CLIP_WGSL }
+    from './kernels/weight_update';
+export { ACTIVATIONS_WGSL, ACTIVATIONS_BACKWARD_WGSL }
+    from './kernels/activations';
+
+export const VERSION = '1.0.1';
+export const DESCRIPTION = 'MambaCode.js: WebGPU-accelerated Mamba SSM for browser code models';

package/src/kernels/{activations.js → activations.ts}

@@ -1,7 +1,7 @@
 // Activation function WGSL kernels: SiLU (Swish) and its backward pass.
 // Used in the gating mechanism of the Mamba Mixer Block.
 
-export const ACTIVATIONS_WGSL = /* wgsl */`
+export const ACTIVATIONS_WGSL: string = /* wgsl */`
 
 struct ActParams {
     num_elements : u32,
@@ -61,7 +61,7 @@ fn rmsnorm_forward(
 `;
 
 // ---- Backward for SiLU ----
-export const ACTIVATIONS_BACKWARD_WGSL = /* wgsl */`
+export const ACTIVATIONS_BACKWARD_WGSL: string = /* wgsl */`
 
 struct ActParams {
     num_elements : u32,

package/src/kernels/{linear_projection.js → linear_projection.ts}

@@ -8,7 +8,7 @@
 //   b : (out_features,)                  – bias
 //   Y : (batch * seq_len, out_features)  – output
 
-export const LINEAR_FORWARD_WGSL = /* wgsl */`
+export const LINEAR_FORWARD_WGSL: string = /* wgsl */`
 
 struct LinearParams {
     M : u32,   // number of rows    (batch * seq_len)
@@ -78,7 +78,7 @@ fn linear_forward(
 `;
 
 // ---- Backward pass for linear projection ----
-export const LINEAR_BACKWARD_WGSL = /* wgsl */`
+export const LINEAR_BACKWARD_WGSL: string = /* wgsl */`
 
 struct LinearParams {
     M : u32,

package/src/kernels/{selective_scan.js → selective_scan.ts}

@@ -8,7 +8,7 @@
 //
 // where A_t, B_t, C_t are input-dependent (selective) gate matrices.
 
-export const SELECTIVE_SCAN_FORWARD_WGSL = /* wgsl */`
+export const SELECTIVE_SCAN_FORWARD_WGSL: string = /* wgsl */`
 
 // ---- Binding layout ----
 // group 0: sequence data
@@ -227,7 +227,7 @@ fn forward_reduce(
 // ---- Backward scan kernel (for autograd) ----
 // Computes gradients w.r.t. Δ, A, B, C using the cached hidden states.
 
-export const SELECTIVE_SCAN_BACKWARD_WGSL = /* wgsl */`
+export const SELECTIVE_SCAN_BACKWARD_WGSL: string = /* wgsl */`
 
 struct ScanParams {
     seq_len  : u32,

package/src/kernels/{weight_update.js → weight_update.ts}

@@ -8,7 +8,7 @@
 //   v_hat = v_t / (1 - beta2^t)
 //   theta_t = theta_{t-1} * (1 - lr * weight_decay) - lr * m_hat / (sqrt(v_hat) + eps)
 
-export const WEIGHT_UPDATE_WGSL = /* wgsl */`
+export const WEIGHT_UPDATE_WGSL: string = /* wgsl */`
 
 struct AdamParams {
     num_elements   : u32,
@@ -60,7 +60,7 @@ fn adamw_update(
 
 // Gradient clipping kernel – clips global gradient norm to max_norm.
 // Run before weight updates.  Two-pass: first compute squared norm, then scale.
-export const GRAD_CLIP_WGSL = /* wgsl */`
+export const GRAD_CLIP_WGSL: string = /* wgsl */`
 
 struct ClipParams {
     num_elements : u32,