WebGPU-accelerated Mamba-1/2/3 and Hybrid SSM library — written in TypeScript, compiled for use in any JavaScript application.
MambaCode.js is a TypeScript-first library that brings the Mamba family of State Space Models to the browser via WebGPU. Version 2.0.0 adds Mamba-2 (SSD), Mamba-3 (complex-valued MIMO + ET discretisation), and hybrid attention layers, while remaining fully backward-compatible with Mamba-1 checkpoints.
📖 New to MambaCode.js? Start with the Getting Started Guide.
| Feature | Detail |
|---|---|
| Mamba-2 (SSD) | Structured State Space Duality — chunked matmul scan, multi-head, scalar A, inner RMSNorm |
| Mamba-3 | Complex-valued states (ℂ^N), ET discretisation, MIMO recurrence, 2× smaller state size |
| AttentionBlock | Causal multi-head attention for hybrid (Jamba/Zamba) layer schedules |
| HybridMambaModel | Per-layer type schedule — mix mamba1/2/3/attention freely |
| MBJS v2 format | Layer-type metadata in checkpoint header; v1 files still load unchanged |
MambaBlock alias |
Kept as deprecated alias for Mamba1Block until 3.0.0 |
| Feature | Detail |
|---|---|
| TypeScript-first | Full type declarations shipped with the package |
| Plain JS compatible | Import the compiled dist/ in any JavaScript project |
| SSM variants | Mamba-1 (S6), Mamba-2 (SSD), Mamba-3 (complex MIMO+ET) |
| Hybrid models | Jamba/Zamba-style mixed SSM + attention schedules |
| Hardware target | WebGPU (WGSL) — Chrome 113+, Edge 113+, Firefox Nightly |
| 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 (Qwen2.5-Coder compatible) |
| WSLA mode | Fast-adapt: trains only the selective projection rows |
npm install mambacode.jsBuild from source:
npm run build # compiles TypeScript → dist/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,
});
await model.loadWeights(await (await fetch('/checkpoint.bin')).arrayBuffer());
const ids = await model.generate(tokenizer.encode('function add('), 200);
console.log(tokenizer.decode(ids));import { HybridMambaModel } from 'mambacode.js';
const model = new HybridMambaModel(device, {
vocabSize : tokenizer.vocabSize,
dModel : 512,
numLayers : 8,
nHeads : 8,
layers : Array(8).fill({ type: 'mamba2' }),
});const model = new HybridMambaModel(device, {
vocabSize : tokenizer.vocabSize,
dModel : 512,
numLayers : 8,
nHeads : 8,
layers : Array(8).fill({ type: 'mamba3' }),
});const model = new HybridMambaModel(device, {
vocabSize : tokenizer.vocabSize,
dModel : 512,
numLayers : 12,
nHeads : 8,
layers : Array.from({ length: 12 }, (_, i) => ({
type: i % 4 === 3 ? 'attention' : 'mamba2',
})),
});Input (B, L, D)
└─ RMSNorm
└─ in_proj → x, z (gate)
x → conv1d → SiLU → x_proj → Δ, B, C
Δ → dt_proj → softplus
Selective Scan S6
h_t = Ā·h_{t-1} + B̄·x_t
y_t = C·h_t + D·x_t
└─ y * SiLU(z)
└─ out_proj + residual
Input (B, L, D)
└─ RMSNorm
└─ in_proj → [x (D_inner), B (G·N), C (G·N), dt (H)]
conv1d over x, B, C (fused)
SSD scan: A_bar = exp(-softplus(A) · softplus(dt))
h_t = A_bar · h_{t-1} + B · x_t
y_t = C · h_t
└─ inner RMSNorm
└─ out_proj + residual
Input (B, L, D)
└─ Same structure as Mamba-2 but:
• A ∈ ℂ (log|A|, arg(A)) per head
• A_bar = exp(Δ·A) [complex]
• B_bar = (A_bar − 1)·A⁻¹·B [ET, exact]
• h_t ∈ ℂ^(N/2), y_t = Re(C·h_t)
Input (B, L, D)
└─ RMSNorm
└─ wQKV → Q, K, V (B, L, H, d_head)
└─ scores = Q·Kᵀ / √d_head (causal mask)
└─ softmax → weighted V sum
└─ concat heads → wO + residual
[optional FFN sublayer]
src/
├── index.ts ← public API entry point (v2.0.0)
├── kernels/
│ ├── selective_scan.ts ← WGSL: S6 forward/backward (Mamba-1)
│ ├── ssd.ts ← WGSL: chunked SSD forward/backward (Mamba-2)
│ ├── complex_ssd.ts ← WGSL: complex SSD + ET + MIMO (Mamba-3)
│ ├── attention.ts ← WGSL: tiled causal MHA forward/backward
│ ├── conv1d.ts ← WGSL: 1D causal convolution (+ groups param)
│ ├── linear_projection.ts ← WGSL: tiled GEMM
│ ├── weight_update.ts ← WGSL: AdamW + gradient clipping
│ └── activations.ts ← WGSL: SiLU, RMSNorm, Softmax
├── model/
│ ├── sequence_layer.ts ← SequenceLayer interface (LayerType, LayerParam)
│ ├── mamba1_block.ts ← Mamba1Block (renamed from MambaBlock)
│ ├── mamba2_block.ts ← Mamba2Block (SSD)
│ ├── mamba3_block.ts ← Mamba3Block (complex + MIMO + ET)
│ ├── attention_block.ts ← AttentionBlock (causal MHA)
│ └── mamba_model.ts ← HybridMambaModel + MambaModel alias
├── training/
│ ├── autograd.ts ← Tape-based AD + loss helpers
│ └── trainer.ts ← MambaTrainer (AdamW, WSLA)
├── tokenizer/
│ └── bpe.ts ← Browser-side BPE tokenizer
└── utils/
├── gpu_utils.ts ← WebGPU device/buffer management
└── quantization.ts ← FP16 / Int8 quantization
tools/ ← Model building & checkpoint tooling
├── generate-bin.js ← CLI: generate an MBJS v2 checkpoint from scratch
├── pretrain.html ← Browser: pretrain a model on a text corpus
└── convert.html ← Browser: convert HuggingFace Mamba → MBJS format
tests/
├── kernels.test.ts
├── autograd.test.ts
├── bpe.test.ts
└── quantization.test.ts
docs/
├── getting-started.md
├── integration-architecture.md
├── weight-lifecycle.md
├── api-reference.md
└── prd-mambacode-v2-v3-hybrid.md ← PRD: Mamba-2/3/hybrid implementation spec
The tools/ directory contains model-building and checkpoint utilities that operate at the mambacode.js level. These are not part of the MambaKit API — they are for authors who want to build, pretrain, or convert model weights.
Creates a properly-shaped MBJS v2 .bin file with randomly initialised weights. Useful as a starting point before pretraining.
node tools/generate-bin.js # nano → model.bin
node tools/generate-bin.js --size small # small preset
node tools/generate-bin.js --size nano --out my.binThe weights are not pretrained — use pretrain.html to run language-model training.
In-browser training loop over a text corpus. Requires a WebGPU-capable browser.
npm run build
npm run serve
# Open http://localhost:3000/tools/pretrain.html
# Load a corpus (e.g. TinyStories), configure size/epochs, click Start Training
# Download the resulting .bin checkpointConverts state-spaces/mamba safetensors checkpoints to MBJS format.
# Open http://localhost:3000/tools/convert.html
# Drop model.safetensors from huggingface.co/state-spaces/mamba-130m
# Download converted .bin| Kernel file | Entry points | Used by |
|---|---|---|
selective_scan.ts |
forward_scan, forward_reduce, selective_scan_backward |
Mamba-1 |
ssd.ts |
ssd_chunk_forward, ssd_chunk_backward |
Mamba-2 |
complex_ssd.ts |
complex_ssd_forward, complex_ssd_backward |
Mamba-3 |
attention.ts |
attention_forward, attention_value, attention_backward |
Attention |
conv1d.ts |
conv1d_forward, conv1d_backward_dx, conv1d_backward_dw |
All SSM |
linear_projection.ts |
linear_forward, linear_backward_dX, linear_backward_dW |
All layers |
activations.ts |
silu_forward, rmsnorm_forward, softmax_forward_simple |
All layers |
weight_update.ts |
adamw_update, grad_norm_reduce, grad_clip_scale |
Training |
[0..3] magic = 0x4D424A53 ('MBJS')
[4..7] version = 1
[8..11] nParams : uint32
[12 ..] numel[i] : uint32 (×nParams)
[data] float32 values
[0..3] magic = 0x4D424A53
[4..7] version = 2
[8..11] nLayers : uint32
[12 ..] layerType[i] : uint8 (0=mamba1, 1=mamba2, 2=mamba3, 3=attention)
[pad] aligned to 4 bytes
[next4] nParams : uint32
[next..] numel[i] : uint32 (×nParams)
[data] float32 values
Version 1 files are loaded transparently — all layers assumed mamba1.
// v1.x — no change needed (mamba1 default is preserved)
const model = new MambaModel(device, config);
// v2.x — opt into Mamba-2
const model = new HybridMambaModel(device, { ...config, layers: Array(8).fill({ type: 'mamba2' }) });
// v2.x — MambaBlock is a deprecated alias for Mamba1Block; both still work
import { MambaBlock, Mamba1Block } from 'mambacode.js';npm test # unit tests (no GPU required)
npm run build # compile TypeScript → dist/
npm run lint # ESLint| Browser | Version | Status |
|---|---|---|
| Chrome | 113+ | ✅ Supported |
| Edge | 113+ | ✅ Supported |
| Firefox | Nightly | ✅ (flag: dom.webgpu.enabled) |
| Safari | 18+ | |
| Node.js | — | ❌ Not supported |
- Mamba-3 — Lahoti et al., Mamba: The Hard Way (arXiv 2603.15569, ICLR 2026)
- Mamba-2 — Dao & Gu, Transformers are SSMs (arXiv 2405.21060, 2024)
- Mamba-1 — Gu & Dao, Mamba: Linear-Time Sequence Modeling with Selective State Spaces (arXiv 2312.00752, 2023)
Want managed infrastructure for your MambaCode.js models?
Builderforce.ai is the professional enterprise platform built on MambaCode.js. It provides:
- In-browser LoRA training — fine-tune up to 2B-parameter models on instruction datasets using the MambaCode.js WebGPU kernels, entirely client-side
- Hybrid Local Brain — the Mamba State Engine runs a selective scan alongside Transformers.js inference for persistent agent memory, powered by MambaCode.js WGSL kernels
- Dataset generation — LLM-assisted JSONL instruction dataset creation with streaming progress
- Workforce Registry — publish trained models as specialist AI agents; discoverable and hirable by the community
- Agent portability —
AgentPackagebundles the LoRA adapter,MambaStateSnapshot, and agent profile into a single portable JSON artifact - CoderClaw mesh — trained agents deploy as self-hosted coding agents via CoderClaw, orchestrated from Builderforce
Use MambaCode.js to build and experiment locally. Use Builderforce.ai to deploy, manage, and share at scale.
MambaCode.js (WebGPU kernels)
↓
SSM.js (session API + runtime + memory)
↓
Builderforce.ai (enterprise IDE + training + registry)
↓
CoderClaw (self-hosted agent mesh)
MIT