mambacode.js 1.0.0 → 1.0.2

package/README.md

@@ -1,14 +1,22 @@
-# Mamba
-MambaCode.js — WebGPU-accelerated Mamba SSM library for browser-based code model training and inference.
+# MambaCode.js
 
-## Overview
+> WebGPU-accelerated Mamba State Space Model library — written in **TypeScript**, compiled for use in any JavaScript application.
 
-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.
+[![npm](https://img.shields.io/npm/v/mambacode.js)](https://www.npmjs.com/package/mambacode.js)
+[![license](https://img.shields.io/badge/license-MIT-blue)](./LICENSE)
 
-### Key features
+MambaCode.js is a **TypeScript-first** library that brings the [Mamba SSM](https://arxiv.org/abs/2312.00752) architecture to the browser via WebGPU. It targets the Qwen3.5-Coder-0.8B model shape and supports full **on-device training** (backpropagation), allowing models to adapt to a user's private codebase locally — without any data leaving the browser.
+
+> 📖 **New to MambaCode.js?** Start with the [Getting Started Guide](./docs/getting-started.md).
+
+---
+
+## Key Features
 
 | Feature | Detail |
 |---|---|
+| **TypeScript-first** | Full type declarations shipped with the package |
+| **Plain JS compatible** | Import the compiled `dist/` in any JavaScript project — no TypeScript toolchain required |
 | **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) |
@@ -20,91 +28,138 @@ MambaCode.js is a pure JavaScript/WGSL implementation of the **Mamba State Space
 
 ---
 
-## Architecture
+## Installation
 
+```bash
+npm install mambacode.js
 ```
-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
+
+Build the library from source:
+
+```bash
+npm run build   # compiles TypeScript → dist/
 ```
 
 ---
 
+## Documentation
+
+| Guide | Description |
+|---|---|
+| **[Getting Started](docs/getting-started.md)** | Beginner-friendly introduction — what LLMs are, how Qwen fits in, the full model lifecycle, and what to do next |
+| **[Integration & Architecture](docs/integration-architecture.md)** | Production architecture guide — embedding Mamba as a unified brain + memory system, integration patterns, advanced use cases, and design tradeoffs |
+| **[Weight Lifecycle](docs/weight-lifecycle.md)** | Complete guide to obtaining Qwen vocabulary files, loading pre-trained checkpoints, fine-tuning, exporting weights, and sharing with your team |
+| **[API Reference](docs/api-reference.md)** | Full technical reference — every exported class, interface, and function with TypeScript and JavaScript examples |
+| **[MambaKit PRD](docs/mamba-kit-prd.md)** | Product requirements document for MambaKit — an opinionated, zero-boilerplate facade over MambaCode.js |
+
+---
+
 ## Quick Start
 
-```js
-import { MambaModel, MambaTrainer, BPETokenizer, initWebGPU } from './src/index.js';
+### TypeScript
+
+```ts
+import {
+  MambaModel,
+  MambaTrainer,
+  BPETokenizer,
+  initWebGPU,
+  type MambaModelConfig,
+  type TrainOptions,
+} from 'mambacode.js';
 
 // 1. Initialise WebGPU
 const { device } = await initWebGPU();
 
-// 2. Load tokenizer
+// 2. Load tokenizer (vocab.json + merges.txt from Qwen3.5-Coder)
 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
+const config: MambaModelConfig = {
+  vocabSize : tokenizer.vocabSize,   // 151936 for Qwen3.5-Coder
   dModel    : 512,
   numLayers : 8,
   dState    : 16,
   dConv     : 4,
   expand    : 2,
+};
+const model = new MambaModel(device, config);
+
+// 4. Load a pre-trained checkpoint
+const response = await fetch('/models/mamba-coder-checkpoint.bin');
+await model.loadWeights(await response.arrayBuffer());
+
+// 5. Fine-tune on local code
+const trainer = new MambaTrainer(model, tokenizer);
+const opts: TrainOptions = {
+  learningRate : 1e-4,
+  epochs       : 5,
+  onEpochEnd   : (epoch, loss) => console.log(`Epoch ${epoch}: loss=${loss.toFixed(4)}`),
+};
+const losses = await trainer.train(myCodeString, opts);
+
+// 6. Generate code
+const promptIds = tokenizer.encode('function fibonacci(');
+const outputIds = await model.generate(promptIds, 200, { temperature: 0.8 });
+console.log(tokenizer.decode(outputIds));
+
+// 7. Save fine-tuned weights for next session
+const checkpoint = await model.exportWeights();
+```
+
+### JavaScript (ESM)
+
+The compiled output in `dist/` is plain JavaScript with no TypeScript runtime dependency:
+
+```js
+import {
+  MambaModel,
+  MambaTrainer,
+  BPETokenizer,
+  initWebGPU,
+} from 'mambacode.js';
+
+// 1. Initialise WebGPU
+const { device } = await initWebGPU();
+
+// 2. Load tokenizer (vocab.json + merges.txt from Qwen3.5-Coder)
+const tokenizer = new BPETokenizer();
+await tokenizer.load('/vocab.json', '/merges.txt');
+
+// 3. Create model
+const model = new MambaModel(device, {
+  vocabSize : tokenizer.vocabSize,
+  dModel    : 512,
+  numLayers : 8,
 });
 
-// 4. Train on local code
+// 4. Load a pre-trained checkpoint
+const response = await fetch('/models/mamba-coder-checkpoint.bin');
+await model.loadWeights(await response.arrayBuffer());
+
+// 5. Fine-tune on local code
 const trainer = new MambaTrainer(model, tokenizer);
-const losses  = await trainer.train(myCodeString, {
+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
+// 6. Generate code
 const promptIds = tokenizer.encode('function fibonacci(');
 const outputIds = await model.generate(promptIds, 200, { temperature: 0.8 });
 console.log(tokenizer.decode(outputIds));
+
+// 7. Save fine-tuned weights for next session
+const checkpoint = await model.exportWeights();
 ```
 
 ### WSLA (Weight-Selective Local Adaptation)
 
 Fine-tune only the B and C matrices for rapid private-codebase adaptation:
 
-```js
+```ts
 await trainer.train(apiUsageExamples, {
   learningRate : 1e-4,
   epochs       : 3,
@@ -114,40 +169,95 @@ await trainer.train(apiUsageExamples, {
 
 ---
 
+## 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
+```
+
+---
+
 ## File Structure
 
 ```
-src/
-├── index.js                      ← public API entry point
+src/                                ← TypeScript source (edit here)
+├── index.ts                        ← 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
+│   ├── selective_scan.ts           ← WGSL: S6 forward + backward (Kogge-Stone)
+│   ├── conv1d.ts                   ← WGSL: 1D causal convolution
+│   ├── linear_projection.ts        ← WGSL: tiled matrix multiplication
+│   ├── weight_update.ts            ← WGSL: AdamW optimizer + gradient clipping
+│   └── activations.ts              ← WGSL: SiLU, RMSNorm
 ├── model/
-│   ├── mamba_block.js            ← Mamba Mixer Block (forward pass)
-│   └── mamba_model.js            ← Full stacked model + generation
+│   ├── mamba_block.ts              ← Mamba Mixer Block (forward pass)
+│   └── mamba_model.ts              ← Full stacked model + generation
 ├── training/
-│   ├── autograd.js               ← Tape-based AD engine + loss helpers
-│   └── trainer.js                ← MambaTrainer class
+│   ├── autograd.ts                 ← Tape-based AD engine + loss helpers
+│   └── trainer.ts                  ← MambaTrainer class
 ├── tokenizer/
-│   └── bpe.js                    ← Browser-side BPE tokenizer
+│   └── bpe.ts                      ← Browser-side BPE tokenizer
 └── utils/
-    ├── gpu_utils.js              ← WebGPU device/buffer management
-    └── quantization.js           ← FP16 / Int8 quantization utilities
+    ├── gpu_utils.ts                ← WebGPU device/buffer management
+    └── quantization.ts             ← FP16 / Int8 quantization utilities
+
+dist/                               ← Compiled output (JS + .d.ts, gitignored)
+├── index.js                        ← ESM entry point for JS consumers
+├── index.d.ts                      ← TypeScript declarations for TS consumers
+└── ...                             ← mirrored sub-folders
+
 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
+├── kernels.test.ts                 ← WGSL kernel source smoke tests
+├── autograd.test.ts                ← Autograd engine unit tests
+├── bpe.test.ts                     ← BPE tokenizer unit tests
+└── quantization.test.ts            ← Quantization round-trip tests
+
+docs/
+├── getting-started.md              ← Step-by-step guide (TS & JS)
+├── integration-architecture.md     ← Brain + Memory architecture guide
+├── weight-lifecycle.md             ← Weight loading, fine-tuning, export
+└── api-reference.md                ← Full API reference
 ```
 
 ---
 
 ## WGSL Kernels
 
-### Parallel Selective Scan (`selective_scan.js`)
+### Parallel Selective Scan (`selective_scan.ts`)
 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:
@@ -158,13 +268,13 @@ The associative operator for the recurrence `h_t = Ā·h_{t-1} + B̄·x_t` is:
 
 Tiles are chained via a carry-in state, covering arbitrarily long sequences.
 
-### 1D Causal Convolution (`conv1d.js`)
+### 1D Causal Convolution (`conv1d.ts`)
 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`)
+### Linear Projection (`linear_projection.ts`)
 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`)
+### AdamW Optimizer (`weight_update.ts`)
 Fused single-kernel AdamW update with decoupled weight decay. Includes a two-pass gradient norm clipping kernel (reduce → scale).
 
 ---
@@ -172,10 +282,12 @@ Fused single-kernel AdamW update with decoupled weight decay. Includes a two-pas
 ## Testing
 
 ```bash
-npm test
+npm test        # run 58 unit tests (no GPU required)
+npm run build   # compile TypeScript → dist/
+npm run lint    # ESLint on src/ and tests/
 ```
 
-Runs 58 unit tests covering quantization, BPE tokenization, autograd, and WGSL kernel source validation. GPU execution tests require a real browser with WebGPU support.
+Unit tests cover quantization, BPE tokenization, autograd, and WGSL kernel source validation. GPU execution tests require a real browser with WebGPU support.
 
 ---
 
@@ -191,6 +303,16 @@ Runs 58 unit tests covering quantization, BPE tokenization, autograd, and WGSL k
 
 ---
 
+## Acknowledgements
+
+This library builds on the Mamba selective state space model research. Special credit to:
+
+- **Mamba 3** — Tri Dao's blog post [*Mamba 3, Part 1*](https://tridao.me/blog/2026/mamba3-part1/) (2026), which describes the latest architectural refinements.
+- **Mamba 3 paper** — [*Mamba: The Hard Way* (arXiv 2603.15569)](https://arxiv.org/abs/2603.15569), the accompanying technical paper.
+- Original **Mamba SSM** paper — [*Mamba: Linear-Time Sequence Modeling with Selective State Spaces* (arXiv 2312.00752)](https://arxiv.org/abs/2312.00752) by Gu & Dao (2023).
+
+---
+
 ## License
 
 MIT

package/dist/index.d.ts

@@ -0,0 +1,19 @@
+/**
+ * MambaCode.js – Entry Point
+ */
+export { MambaModel } from './model/mamba_model.js';
+export { MambaBlock } from './model/mamba_block.js';
+export { MambaTrainer } from './training/trainer.js';
+export { Tensor, backward, enableGrad, noGrad, clearTape, recordOperation, crossEntropyLoss, crossEntropyGrad, } from './training/autograd.js';
+export { BPETokenizer } from './tokenizer/bpe.js';
+export type { MambaModelConfig, SamplingOptions } from './model/mamba_model.js';
+export { initWebGPU, createStorageBuffer, createEmptyStorageBuffer, createUniformBuffer, createComputePipeline, createBindGroup, dispatchKernel, readBuffer, uploadBuffer, cdiv, } from './utils/gpu_utils.js';
+export { quantizeFp16, dequantizeFp16, floatToFp16, fp16ToFloat, quantizeInt8, dequantizeInt8, quantizeInt8PerChannel, dequantizeInt8PerChannel, estimateMemory, } from './utils/quantization.js';
+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';
+export declare const VERSION = "1.0.2";
+export declare const DESCRIPTION = "MambaCode.js: WebGPU-accelerated Mamba SSM for browser code models";
+//# sourceMappingURL=index.d.ts.map

package/dist/index.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA;;GAEG;AAEH,OAAO,EAAE,UAAU,EAAE,MAAQ,wBAAwB,CAAC;AACtD,OAAO,EAAE,UAAU,EAAE,MAAQ,wBAAwB,CAAC;AAEtD,OAAO,EAAE,YAAY,EAAE,MAAM,uBAAuB,CAAC;AACrD,OAAO,EACH,MAAM,EACN,QAAQ,EACR,UAAU,EACV,MAAM,EACN,SAAS,EACT,eAAe,EACf,gBAAgB,EAChB,gBAAgB,GACnB,MAAM,wBAAwB,CAAC;AAEhC,OAAO,EAAE,YAAY,EAAE,MAAM,oBAAoB,CAAC;AAElD,YAAY,EAAE,gBAAgB,EAAE,eAAe,EAAE,MAAM,wBAAwB,CAAC;AAEhF,OAAO,EACH,UAAU,EACV,mBAAmB,EACnB,wBAAwB,EACxB,mBAAmB,EACnB,qBAAqB,EACrB,eAAe,EACf,cAAc,EACd,UAAU,EACV,YAAY,EACZ,IAAI,GACP,MAAM,sBAAsB,CAAC;AAE9B,OAAO,EACH,YAAY,EACZ,cAAc,EACd,WAAW,EACX,WAAW,EACX,YAAY,EACZ,cAAc,EACd,sBAAsB,EACtB,wBAAwB,EACxB,cAAc,GACjB,MAAM,yBAAyB,CAAC;AAEjC,OAAO,EAAE,2BAA2B,EAAE,4BAA4B,EAAE,MAC3D,6BAA6B,CAAC;AACvC,OAAO,EAAE,mBAAmB,EAAE,oBAAoB,EAAE,MAC3C,qBAAqB,CAAC;AAC/B,OAAO,EAAE,mBAAmB,EAAE,oBAAoB,EAAE,MAC3C,gCAAgC,CAAC;AAC1C,OAAO,EAAE,kBAAkB,EAAE,cAAc,EAAE,MACpC,4BAA4B,CAAC;AACtC,OAAO,EAAE,gBAAgB,EAAE,yBAAyB,EAAE,MAC7C,0BAA0B,CAAC;AAEpC,eAAO,MAAM,OAAO,UAAU,CAAC;AAC/B,eAAO,MAAM,WAAW,uEAAuE,CAAC"}

package/dist/index.js

@@ -0,0 +1,18 @@
+/**
+ * MambaCode.js – Entry Point
+ */
+export { MambaModel } from './model/mamba_model.js';
+export { MambaBlock } from './model/mamba_block.js';
+export { MambaTrainer } from './training/trainer.js';
+export { Tensor, backward, enableGrad, noGrad, clearTape, recordOperation, crossEntropyLoss, crossEntropyGrad, } from './training/autograd.js';
+export { BPETokenizer } from './tokenizer/bpe.js';
+export { initWebGPU, createStorageBuffer, createEmptyStorageBuffer, createUniformBuffer, createComputePipeline, createBindGroup, dispatchKernel, readBuffer, uploadBuffer, cdiv, } from './utils/gpu_utils.js';
+export { quantizeFp16, dequantizeFp16, floatToFp16, fp16ToFloat, quantizeInt8, dequantizeInt8, quantizeInt8PerChannel, dequantizeInt8PerChannel, estimateMemory, } from './utils/quantization.js';
+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';
+export const VERSION = '1.0.2';
+export const DESCRIPTION = 'MambaCode.js: WebGPU-accelerated Mamba SSM for browser code models';
+//# sourceMappingURL=index.js.map

package/dist/index.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"index.js","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA;;GAEG;AAEH,OAAO,EAAE,UAAU,EAAE,MAAQ,wBAAwB,CAAC;AACtD,OAAO,EAAE,UAAU,EAAE,MAAQ,wBAAwB,CAAC;AAEtD,OAAO,EAAE,YAAY,EAAE,MAAM,uBAAuB,CAAC;AACrD,OAAO,EACH,MAAM,EACN,QAAQ,EACR,UAAU,EACV,MAAM,EACN,SAAS,EACT,eAAe,EACf,gBAAgB,EAChB,gBAAgB,GACnB,MAAM,wBAAwB,CAAC;AAEhC,OAAO,EAAE,YAAY,EAAE,MAAM,oBAAoB,CAAC;AAIlD,OAAO,EACH,UAAU,EACV,mBAAmB,EACnB,wBAAwB,EACxB,mBAAmB,EACnB,qBAAqB,EACrB,eAAe,EACf,cAAc,EACd,UAAU,EACV,YAAY,EACZ,IAAI,GACP,MAAM,sBAAsB,CAAC;AAE9B,OAAO,EACH,YAAY,EACZ,cAAc,EACd,WAAW,EACX,WAAW,EACX,YAAY,EACZ,cAAc,EACd,sBAAsB,EACtB,wBAAwB,EACxB,cAAc,GACjB,MAAM,yBAAyB,CAAC;AAEjC,OAAO,EAAE,2BAA2B,EAAE,4BAA4B,EAAE,MAC3D,6BAA6B,CAAC;AACvC,OAAO,EAAE,mBAAmB,EAAE,oBAAoB,EAAE,MAC3C,qBAAqB,CAAC;AAC/B,OAAO,EAAE,mBAAmB,EAAE,oBAAoB,EAAE,MAC3C,gCAAgC,CAAC;AAC1C,OAAO,EAAE,kBAAkB,EAAE,cAAc,EAAE,MACpC,4BAA4B,CAAC;AACtC,OAAO,EAAE,gBAAgB,EAAE,yBAAyB,EAAE,MAC7C,0BAA0B,CAAC;AAEpC,MAAM,CAAC,MAAM,OAAO,GAAG,OAAO,CAAC;AAC/B,MAAM,CAAC,MAAM,WAAW,GAAG,oEAAoE,CAAC"}

package/dist/kernels/activations.d.ts

@@ -0,0 +1,3 @@
+export declare const ACTIVATIONS_WGSL: string;
+export declare const ACTIVATIONS_BACKWARD_WGSL: string;
+//# sourceMappingURL=activations.d.ts.map

package/dist/kernels/activations.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"activations.d.ts","sourceRoot":"","sources":["../../src/kernels/activations.ts"],"names":[],"mappings":"AAGA,eAAO,MAAM,gBAAgB,EAAE,MAyD9B,CAAC;AAGF,eAAO,MAAM,yBAAyB,EAAE,MAwBvC,CAAC"}

package/dist/kernels/activations.js

@@ -0,0 +1,87 @@
+// 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));
+}
+`;
+//# sourceMappingURL=activations.js.map

package/dist/kernels/activations.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"activations.js","sourceRoot":"","sources":["../../src/kernels/activations.ts"],"names":[],"mappings":"AAAA,wEAAwE;AACxE,yDAAyD;AAEzD,MAAM,CAAC,MAAM,gBAAgB,GAAW,UAAU,CAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CAyDjD,CAAC;AAEF,8BAA8B;AAC9B,MAAM,CAAC,MAAM,yBAAyB,GAAW,UAAU,CAAA;;;;;;;;;;;;;;;;;;;;;;;;CAwB1D,CAAC"}

package/dist/kernels/conv1d.d.ts

@@ -0,0 +1,3 @@
+export declare const CONV1D_FORWARD_WGSL = "\n\nstruct ConvParams {\n    seq_len     : u32,   // L\n    d_channels  : u32,   // D (number of depthwise channels)\n    kernel_size : u32,   // K (typically 4)\n    batch       : u32,   // B\n};\n\n@group(0) @binding(0) var<uniform>             params   : ConvParams;\n// x      (B, L, D) \u2013 input\n@group(0) @binding(1) var<storage, read>       x        : array<f32>;\n// weight (D, K)    \u2013 depthwise conv weights\n@group(0) @binding(2) var<storage, read>       weight   : array<f32>;\n// bias   (D,)      \u2013 optional bias (zeros if unused)\n@group(0) @binding(3) var<storage, read>       bias     : array<f32>;\n// y      (B, L, D) \u2013 output\n@group(0) @binding(4) var<storage, read_write> y        : array<f32>;\n\n// Dispatch: (ceil(L/16), ceil(D/16), B)\n@compute @workgroup_size(16, 16, 1)\nfn conv1d_forward(\n    @builtin(global_invocation_id) gid : vec3<u32>,\n) {\n    let L  = params.seq_len;\n    let D  = params.d_channels;\n    let K  = params.kernel_size;\n    let B  = params.batch;\n\n    let t  = gid.x;   // time position\n    let d  = gid.y;   // channel\n    let b  = gid.z;   // batch\n\n    if (t >= L || d >= D || b >= B) { return; }\n\n    var acc: f32 = 0.0;\n\n    // Causal: convolve over k = 0..K-1, reading position (t - k)\n    for (var k: u32 = 0u; k < K; k = k + 1u) {\n        let w_idx = d * K + k;\n        let w_val = weight[w_idx];\n\n        // t - k: use causal zero-padding for t < k\n        if (t >= k) {\n            let src = b * L * D + (t - k) * D + d;\n            acc = acc + w_val * x[src];\n        }\n        // else: zero-padding contributes 0\n    }\n\n    acc = acc + bias[d];\n\n    let out = b * L * D + t * D + d;\n    y[out] = acc;\n}\n";
+export declare const CONV1D_BACKWARD_WGSL = "\n\nstruct ConvParams {\n    seq_len     : u32,\n    d_channels  : u32,\n    kernel_size : u32,\n    batch       : u32,\n};\n\n@group(0) @binding(0) var<uniform>              params   : ConvParams;\n@group(0) @binding(1) var<storage, read>        x        : array<f32>;\n@group(0) @binding(2) var<storage, read>        weight   : array<f32>;\n@group(0) @binding(3) var<storage, read>        dy       : array<f32>;\n@group(0) @binding(4) var<storage, read_write>  dx       : array<f32>;\n@group(0) @binding(5) var<storage, read_write>  dweight  : array<f32>;\n@group(0) @binding(6) var<storage, read_write>  dbias    : array<f32>;\n\n// Dispatch: (ceil(L/16), ceil(D/16), B) \u2013 computes dx\n@compute @workgroup_size(16, 16, 1)\nfn conv1d_backward_dx(\n    @builtin(global_invocation_id) gid : vec3<u32>,\n) {\n    let L  = params.seq_len;\n    let D  = params.d_channels;\n    let K  = params.kernel_size;\n    let B  = params.batch;\n\n    let t  = gid.x;\n    let d  = gid.y;\n    let b  = gid.z;\n\n    if (t >= L || d >= D || b >= B) { return; }\n\n    var grad: f32 = 0.0;\n\n    // dx[b, t, d] = sum_{k=0}^{K-1} dy[b, t+k, d] * weight[d, k]\n    for (var k: u32 = 0u; k < K; k = k + 1u) {\n        let tp = t + k;\n        if (tp < L) {\n            let dy_idx = b * L * D + tp * D + d;\n            let w_idx  = d * K + k;\n            grad = grad + dy[dy_idx] * weight[w_idx];\n        }\n    }\n\n    let dx_idx = b * L * D + t * D + d;\n    dx[dx_idx] = grad;\n}\n\n// Dispatch: (K, D, 1) \u2013 accumulates dweight over (B, L)\n@compute @workgroup_size(1, 1, 1)\nfn conv1d_backward_dw(\n    @builtin(global_invocation_id) gid : vec3<u32>,\n) {\n    let L  = params.seq_len;\n    let D  = params.d_channels;\n    let K  = params.kernel_size;\n    let B  = params.batch;\n\n    let k  = gid.x;\n    let d  = gid.y;\n\n    if (k >= K || d >= D) { return; }\n\n    var grad_w: f32 = 0.0;\n    var grad_b: f32 = 0.0;\n\n    for (var b: u32 = 0u; b < B; b = b + 1u) {\n        for (var t: u32 = 0u; t < L; t = t + 1u) {\n            let dy_idx = b * L * D + t * D + d;\n            let dy_val = dy[dy_idx];\n            if (t >= k) {\n                let x_idx = b * L * D + (t - k) * D + d;\n                grad_w = grad_w + dy_val * x[x_idx];\n            }\n            if (k == 0u) {\n                grad_b = grad_b + dy_val;\n            }\n        }\n    }\n\n    dweight[d * K + k] = grad_w;\n    if (k == 0u) {\n        dbias[d] = grad_b;\n    }\n}\n";
+//# sourceMappingURL=conv1d.d.ts.map

package/dist/kernels/conv1d.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"conv1d.d.ts","sourceRoot":"","sources":["../../src/kernels/conv1d.ts"],"names":[],"mappings":"AAQA,eAAO,MAAM,mBAAmB,urDAuD/B,CAAC;AAGF,eAAO,MAAM,oBAAoB,y6EAsFhC,CAAC"}