mambacode.js 1.0.0

package/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Sean Hogg
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

package/README.md

@@ -0,0 +1,196 @@
+# Mamba
+MambaCode.js — WebGPU-accelerated Mamba SSM library for browser-based code model training and inference.
+
+## Overview
+
+MambaCode.js is a pure JavaScript/WGSL implementation of the **Mamba State Space Model (SSM)** architecture, optimised for on-device code model training and inference in the browser. It targets the Qwen3.5-Coder-0.8B logic and supports full **on-device training** (backpropagation) via WebGPU, allowing models to adapt to a user's private codebase locally — without any data leaving the browser.
+
+### Key features
+
+| Feature | Detail |
+|---|---|
+| **Architecture** | Selective State Space Model (S6) — linear O(N) context scaling |
+| **Hardware target** | WebGPU (WGSL) — Chrome 113+, Edge 113+, Firefox Nightly |
+| **Memory ceiling** | ≤ 3 GB VRAM (Chrome/Edge/Firefox stable) |
+| **No heavy frameworks** | Zero TensorFlow.js / Transformers.js dependencies |
+| **On-device training** | Tape-based autograd + AdamW GPU optimizer |
+| **Quantization** | FP16 weights, Int8 activations |
+| **Tokenizer** | Browser-side BPE (Qwen3.5-Coder compatible) |
+| **WSLA mode** | Fine-tune only B & C matrices for rapid local adaptation |
+
+---
+
+## Architecture
+
+```
+Token IDs
+    │
+    ▼
+Embedding Lookup (GPU gather kernel)
+    │
+    ▼  ┌─────────────────────────────────────────┐
+       │           Mamba Block × N               │
+       │                                         │
+       │  Input ──► RMSNorm                      │
+       │               │                         │
+       │      ┌────────┴────────┐                │
+       │      ▼                 ▼                │
+       │  in_proj(x)       in_proj(z)  [gate]    │
+       │      │                                  │
+       │  Conv1D (causal, K=4)                   │
+       │      │                                  │
+       │   SiLU activation                       │
+       │      │                                  │
+       │  x_proj → Δ, B, C  (selective)          │
+       │      │                                  │
+       │  Δ → dt_proj (full D_inner width)        │
+       │      │                                  │
+       │  ┌───▼──────────────────────────────┐   │
+       │  │  Selective Scan S6               │   │
+       │  │  (Kogge-Stone parallel prefix)   │   │
+       │  │  h_t = Ā·h_{t-1} + B̄·x_t        │   │
+       │  │  y_t = C·h_t + D·x_t            │   │
+       │  └──────────────────────────────────┘   │
+       │      │                                  │
+       │  Gate: y * SiLU(z)                      │
+       │      │                                  │
+       │  out_proj → residual add ──► output     │
+       └─────────────────────────────────────────┘
+    │
+    ▼
+Final RMSNorm → LM Head (tied embedding) → Logits
+```
+
+---
+
+## Quick Start
+
+```js
+import { MambaModel, MambaTrainer, BPETokenizer, initWebGPU } from './src/index.js';
+
+// 1. Initialise WebGPU
+const { device } = await initWebGPU();
+
+// 2. Load tokenizer
+const tokenizer = new BPETokenizer();
+await tokenizer.load('/vocab.json', '/merges.txt');
+
+// 3. Create model
+const model = new MambaModel(device, {
+  vocabSize : tokenizer.vocabSize,   // e.g. 151936 for Qwen3.5-Coder
+  dModel    : 512,
+  numLayers : 8,
+  dState    : 16,
+  dConv     : 4,
+  expand    : 2,
+});
+
+// 4. Train on local code
+const trainer = new MambaTrainer(model, tokenizer);
+const losses  = await trainer.train(myCodeString, {
+  learningRate : 1e-4,
+  epochs       : 5,
+  device       : 'webgpu',
+  onEpochEnd   : (epoch, loss) => console.log(`Epoch ${epoch}: loss=${loss.toFixed(4)}`),
+});
+
+// 5. Generate code
+const promptIds = tokenizer.encode('function fibonacci(');
+const outputIds = await model.generate(promptIds, 200, { temperature: 0.8 });
+console.log(tokenizer.decode(outputIds));
+```
+
+### WSLA (Weight-Selective Local Adaptation)
+
+Fine-tune only the B and C matrices for rapid private-codebase adaptation:
+
+```js
+await trainer.train(apiUsageExamples, {
+  learningRate : 1e-4,
+  epochs       : 3,
+  wsla         : true,   // only B and C matrices are updated
+});
+```
+
+---
+
+## File Structure
+
+```
+src/
+├── index.js                      ← public API entry point
+├── kernels/
+│   ├── selective_scan.js         ← WGSL: S6 forward + backward (Kogge-Stone)
+│   ├── conv1d.js                 ← WGSL: 1D causal convolution
+│   ├── linear_projection.js      ← WGSL: tiled matrix multiplication
+│   ├── weight_update.js          ← WGSL: AdamW optimizer + gradient clipping
+│   └── activations.js            ← WGSL: SiLU, RMSNorm
+├── model/
+│   ├── mamba_block.js            ← Mamba Mixer Block (forward pass)
+│   └── mamba_model.js            ← Full stacked model + generation
+├── training/
+│   ├── autograd.js               ← Tape-based AD engine + loss helpers
+│   └── trainer.js                ← MambaTrainer class
+├── tokenizer/
+│   └── bpe.js                    ← Browser-side BPE tokenizer
+└── utils/
+    ├── gpu_utils.js              ← WebGPU device/buffer management
+    └── quantization.js           ← FP16 / Int8 quantization utilities
+tests/
+├── kernels.test.js               ← WGSL kernel source smoke tests
+├── autograd.test.js              ← Autograd engine unit tests
+├── bpe.test.js                   ← BPE tokenizer unit tests
+└── quantization.test.js          ← Quantization round-trip tests
+```
+
+---
+
+## WGSL Kernels
+
+### Parallel Selective Scan (`selective_scan.js`)
+Implements the S6 core using a **Kogge-Stone parallel prefix-sum** inside each workgroup tile. Each tile of 64 time steps is scanned in log₂(64) = 6 GPU barrier rounds, giving O(log N) wall-clock time on the GPU.
+
+The associative operator for the recurrence `h_t = Ā·h_{t-1} + B̄·x_t` is:
+
+```
+(a₁, b₁) ∘ (a₂, b₂) = (a₁·a₂, a₁·b₂ + b₁)
+```
+
+Tiles are chained via a carry-in state, covering arbitrarily long sequences.
+
+### 1D Causal Convolution (`conv1d.js`)
+Depthwise 1D causal conv (kernel size K=4) with zero left-padding. Enforces causality by only reading positions `t-k` for `k ≥ 0`, contributing 0 for `t < k`.
+
+### Linear Projection (`linear_projection.js`)
+Tiled 16×16 GEMM in WGSL using workgroup shared memory. Handles arbitrary (M, K) × (N, K) → (M, N) shapes with boundary guards.
+
+### AdamW Optimizer (`weight_update.js`)
+Fused single-kernel AdamW update with decoupled weight decay. Includes a two-pass gradient norm clipping kernel (reduce → scale).
+
+---
+
+## Testing
+
+```bash
+npm test
+```
+
+Runs 58 unit tests covering quantization, BPE tokenization, autograd, and WGSL kernel source validation. GPU execution tests require a real browser with WebGPU support.
+
+---
+
+## Browser Compatibility
+
+| Browser | Version | Status |
+|---|---|---|
+| Chrome | 113+ | ✅ Supported |
+| Edge | 113+ | ✅ Supported |
+| Firefox | Nightly | ✅ Supported (flag: `dom.webgpu.enabled`) |
+| Safari | 18+ | ⚠️ Partial (WebGPU in preview) |
+| Node.js | — | ❌ Not supported (no `navigator.gpu`) |
+
+---
+
+## License
+
+MIT

package/package.json

@@ -0,0 +1,54 @@
+{
+  "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",
+  "type": "module",
+  "files": [
+    "src",
+    "README.md",
+    "LICENSE"
+  ],
+  "scripts": {
+    "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"
+  ],
+  "author": {
+    "name": "Sean Hogg",
+    "email": "seanhogg@gmail.com",
+    "url": "https://builderforce.ai"
+  },
+  "license": "MIT",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/SeanHogg/Mamba.git"
+  },
+  "bugs": {
+    "url": "https://github.com/SeanHogg/Mamba/issues"
+  },
+  "homepage": "https://github.com/SeanHogg/Mamba#readme",
+  "engines": {
+    "node": ">=18.0.0"
+  },
+  "publishConfig": {
+    "access": "public"
+  },
+  "devDependencies": {
+    "jest": "^29.7.0",
+    "eslint": "^8.57.0"
+  },
+  "jest": {
+    "transform": {}
+  }
+}

package/src/index.js

@@ -0,0 +1,89 @@
+/**
+ * MambaCode.js – Entry Point
+ *
+ * High-performance JavaScript/WGSL Mamba SSM library for browser-based
+ * code model training and inference.
+ *
+ * Quick-start example
+ * -------------------
+ * import { MambaModel, MambaTrainer, BPETokenizer, initWebGPU } from 'mambacode.js';
+ *
+ * const { device }   = await initWebGPU();
+ * const tokenizer    = new BPETokenizer();
+ * await tokenizer.load('/vocab.json', '/merges.txt');
+ *
+ * const model = new MambaModel(device, {
+ *   vocabSize : tokenizer.vocabSize,
+ *   dModel    : 512,
+ *   numLayers : 8,
+ * });
+ *
+ * const trainer = new MambaTrainer(model, tokenizer);
+ * const losses  = await trainer.train(myCodeString, { learningRate: 1e-4, epochs: 5 });
+ *
+ * const generated = await model.generate(tokenizer.encode('function '), 100);
+ * console.log(tokenizer.decode(generated));
+ */
+
+// ── Core model ────────────────────────────────────────────────────────────────
+export { MambaModel }   from './model/mamba_model.js';
+export { MambaBlock }   from './model/mamba_block.js';
+
+// ── Training ──────────────────────────────────────────────────────────────────
+export { MambaTrainer } from './training/trainer.js';
+export {
+    Tensor,
+    backward,
+    enableGrad,
+    noGrad,
+    clearTape,
+    recordOperation,
+    crossEntropyLoss,
+    crossEntropyGrad,
+} from './training/autograd.js';
+
+// ── Tokenizer ─────────────────────────────────────────────────────────────────
+export { BPETokenizer } from './tokenizer/bpe.js';
+
+// ── WebGPU utilities ──────────────────────────────────────────────────────────
+export {
+    initWebGPU,
+    createStorageBuffer,
+    createEmptyStorageBuffer,
+    createUniformBuffer,
+    createComputePipeline,
+    createBindGroup,
+    dispatchKernel,
+    readBuffer,
+    uploadBuffer,
+    cdiv,
+} from './utils/gpu_utils.js';
+
+// ── Quantization utilities ────────────────────────────────────────────────────
+export {
+    quantizeFp16,
+    dequantizeFp16,
+    floatToFp16,
+    fp16ToFloat,
+    quantizeInt8,
+    dequantizeInt8,
+    quantizeInt8PerChannel,
+    dequantizeInt8PerChannel,
+    estimateMemory,
+} from './utils/quantization.js';
+
+// ── Raw WGSL kernel sources (for advanced users / custom pipelines) ───────────
+export { SELECTIVE_SCAN_FORWARD_WGSL, SELECTIVE_SCAN_BACKWARD_WGSL }
+    from './kernels/selective_scan.js';
+export { CONV1D_FORWARD_WGSL, CONV1D_BACKWARD_WGSL }
+    from './kernels/conv1d.js';
+export { LINEAR_FORWARD_WGSL, LINEAR_BACKWARD_WGSL }
+    from './kernels/linear_projection.js';
+export { WEIGHT_UPDATE_WGSL, GRAD_CLIP_WGSL }
+    from './kernels/weight_update.js';
+export { ACTIVATIONS_WGSL, ACTIVATIONS_BACKWARD_WGSL }
+    from './kernels/activations.js';
+
+// ── Library metadata ──────────────────────────────────────────────────────────
+export const VERSION = '0.1.0';
+export const DESCRIPTION = 'MambaCode.js: WebGPU-accelerated Mamba SSM for browser code models';

package/src/kernels/activations.js

@@ -0,0 +1,88 @@
+// 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 */`
+
+struct ActParams {
+    num_elements : u32,
+};
+
+@group(0) @binding(0) var<uniform>             p    : ActParams;
+@group(0) @binding(1) var<storage, read>       x    : array<f32>;
+@group(0) @binding(2) var<storage, read_write> y    : array<f32>;
+
+// SiLU(x) = x * sigmoid(x)
+@compute @workgroup_size(256, 1, 1)
+fn silu_forward(
+    @builtin(global_invocation_id) gid : vec3<u32>,
+) {
+    let i = gid.x;
+    if (i >= p.num_elements) { return; }
+    let v = x[i];
+    y[i] = v / (1.0 + exp(-v));
+}
+
+// RMSNorm forward:  y = x / rms(x) * weight
+// Requires separate uniform for rms norm params.
+struct RMSNormParams {
+    num_rows  : u32,   // number of vectors (batch * seq_len)
+    dim       : u32,   // feature dimension
+    eps       : f32,
+};
+
+@group(0) @binding(0) var<uniform>             rms_p    : RMSNormParams;
+@group(0) @binding(1) var<storage, read>       rms_x    : array<f32>;
+@group(0) @binding(2) var<storage, read>       rms_w    : array<f32>;   // scale (dim,)
+@group(0) @binding(3) var<storage, read_write> rms_y    : array<f32>;
+@group(0) @binding(4) var<storage, read_write> rms_inv  : array<f32>;   // cache 1/rms per row
+
+@compute @workgroup_size(64, 1, 1)
+fn rmsnorm_forward(
+    @builtin(global_invocation_id) gid : vec3<u32>,
+) {
+    let row = gid.x;
+    if (row >= rms_p.num_rows) { return; }
+
+    let D = rms_p.dim;
+    let base = row * D;
+
+    var sq_sum: f32 = 0.0;
+    for (var i: u32 = 0u; i < D; i = i + 1u) {
+        let v = rms_x[base + i];
+        sq_sum = sq_sum + v * v;
+    }
+    let inv_rms = 1.0 / sqrt(sq_sum / f32(D) + rms_p.eps);
+    rms_inv[row] = inv_rms;
+
+    for (var i: u32 = 0u; i < D; i = i + 1u) {
+        rms_y[base + i] = rms_x[base + i] * inv_rms * rms_w[i];
+    }
+}
+`;
+
+// ---- Backward for SiLU ----
+export const ACTIVATIONS_BACKWARD_WGSL = /* wgsl */`
+
+struct ActParams {
+    num_elements : u32,
+};
+
+@group(0) @binding(0) var<uniform>            p   : ActParams;
+@group(0) @binding(1) var<storage, read>      x   : array<f32>;
+@group(0) @binding(2) var<storage, read>      dy  : array<f32>;
+@group(0) @binding(3) var<storage, read_write> dx : array<f32>;
+
+// d/dx [x * sigmoid(x)] = sigmoid(x) + x * sigmoid(x) * (1 - sigmoid(x))
+//                        = silu(x)/x  + sigmoid(x) * (1 - sigmoid(x)) * x
+//                        simplified:  sigmoid(x) * (1 + x*(1 - sigmoid(x)))
+@compute @workgroup_size(256, 1, 1)
+fn silu_backward(
+    @builtin(global_invocation_id) gid : vec3<u32>,
+) {
+    let i = gid.x;
+    if (i >= p.num_elements) { return; }
+    let v   = x[i];
+    let sig = 1.0 / (1.0 + exp(-v));
+    dx[i] = dy[i] * sig * (1.0 + v * (1.0 - sig));
+}
+`;

package/src/kernels/conv1d.js

@@ -0,0 +1,153 @@
+// 1D Causal Convolution WGSL Kernel
+// Implements a depthwise 1D causal convolution over the sequence dimension.
+// "Causal" means the output at position t only depends on positions <= t,
+// which is enforced by left-padding with (kernel_size - 1) zeros.
+//
+// Forward:  y[b, t, d] = sum_{k=0}^{K-1} weight[d, k] * x[b, t-k, d]
+//           where x[b, t', d] = 0 for t' < 0  (causal padding)
+
+export const CONV1D_FORWARD_WGSL = /* wgsl */`
+
+struct ConvParams {
+    seq_len     : u32,   // L
+    d_channels  : u32,   // D (number of depthwise channels)
+    kernel_size : u32,   // K (typically 4)
+    batch       : u32,   // B
+};
+
+@group(0) @binding(0) var<uniform>             params   : ConvParams;
+// x      (B, L, D) – input
+@group(0) @binding(1) var<storage, read>       x        : array<f32>;
+// weight (D, K)    – depthwise conv weights
+@group(0) @binding(2) var<storage, read>       weight   : array<f32>;
+// bias   (D,)      – optional bias (zeros if unused)
+@group(0) @binding(3) var<storage, read>       bias     : array<f32>;
+// y      (B, L, D) – output
+@group(0) @binding(4) var<storage, read_write> y        : array<f32>;
+
+// Dispatch: (ceil(L/16), ceil(D/16), B)
+@compute @workgroup_size(16, 16, 1)
+fn conv1d_forward(
+    @builtin(global_invocation_id) gid : vec3<u32>,
+) {
+    let L  = params.seq_len;
+    let D  = params.d_channels;
+    let K  = params.kernel_size;
+    let B  = params.batch;
+
+    let t  = gid.x;   // time position
+    let d  = gid.y;   // channel
+    let b  = gid.z;   // batch
+
+    if (t >= L || d >= D || b >= B) { return; }
+
+    var acc: f32 = 0.0;
+
+    // Causal: convolve over k = 0..K-1, reading position (t - k)
+    for (var k: u32 = 0u; k < K; k = k + 1u) {
+        let w_idx = d * K + k;
+        let w_val = weight[w_idx];
+
+        // t - k: use causal zero-padding for t < k
+        if (t >= k) {
+            let src = b * L * D + (t - k) * D + d;
+            acc = acc + w_val * x[src];
+        }
+        // else: zero-padding contributes 0
+    }
+
+    acc = acc + bias[d];
+
+    let out = b * L * D + t * D + d;
+    y[out] = acc;
+}
+`;
+
+// ---- Backward kernel for 1D convolution ----
+export const CONV1D_BACKWARD_WGSL = /* wgsl */`
+
+struct ConvParams {
+    seq_len     : u32,
+    d_channels  : u32,
+    kernel_size : u32,
+    batch       : u32,
+};
+
+@group(0) @binding(0) var<uniform>              params   : ConvParams;
+@group(0) @binding(1) var<storage, read>        x        : array<f32>;
+@group(0) @binding(2) var<storage, read>        weight   : array<f32>;
+@group(0) @binding(3) var<storage, read>        dy       : array<f32>;
+@group(0) @binding(4) var<storage, read_write>  dx       : array<f32>;
+@group(0) @binding(5) var<storage, read_write>  dweight  : array<f32>;
+@group(0) @binding(6) var<storage, read_write>  dbias    : array<f32>;
+
+// Dispatch: (ceil(L/16), ceil(D/16), B) – computes dx
+@compute @workgroup_size(16, 16, 1)
+fn conv1d_backward_dx(
+    @builtin(global_invocation_id) gid : vec3<u32>,
+) {
+    let L  = params.seq_len;
+    let D  = params.d_channels;
+    let K  = params.kernel_size;
+    let B  = params.batch;
+
+    let t  = gid.x;
+    let d  = gid.y;
+    let b  = gid.z;
+
+    if (t >= L || d >= D || b >= B) { return; }
+
+    var grad: f32 = 0.0;
+
+    // dx[b, t, d] = sum_{k=0}^{K-1} dy[b, t+k, d] * weight[d, k]
+    for (var k: u32 = 0u; k < K; k = k + 1u) {
+        let tp = t + k;
+        if (tp < L) {
+            let dy_idx = b * L * D + tp * D + d;
+            let w_idx  = d * K + k;
+            grad = grad + dy[dy_idx] * weight[w_idx];
+        }
+    }
+
+    let dx_idx = b * L * D + t * D + d;
+    dx[dx_idx] = grad;
+}
+
+// Dispatch: (K, D, 1) – accumulates dweight over (B, L)
+@compute @workgroup_size(1, 1, 1)
+fn conv1d_backward_dw(
+    @builtin(global_invocation_id) gid : vec3<u32>,
+) {
+    let L  = params.seq_len;
+    let D  = params.d_channels;
+    let K  = params.kernel_size;
+    let B  = params.batch;
+
+    let k  = gid.x;
+    let d  = gid.y;
+
+    if (k >= K || d >= D) { return; }
+
+    var grad_w: f32 = 0.0;
+    var grad_b: f32 = 0.0;
+
+    for (var b: u32 = 0u; b < B; b = b + 1u) {
+        for (var t: u32 = 0u; t < L; t = t + 1u) {
+            let dy_idx = b * L * D + t * D + d;
+            let dy_val = dy[dy_idx];
+            if (t >= k) {
+                let x_idx = b * L * D + (t - k) * D + d;
+                grad_w = grad_w + dy_val * x[x_idx];
+            }
+            if (k == 0u) {
+                grad_b = grad_b + dy_val;
+            }
+        }
+    }
+
+    dweight[d * K + k] = grad_w;
+    if (k == 0u) {
+        dbias[d] = grad_b;
+    }
+}
+`;