mambacode.js 1.0.0 → 1.0.1

package/src/utils/quantization.ts

@@ -0,0 +1,154 @@
+/**
+ * quantization.ts – FP16 and Int8 quantization utilities.
+ */
+
+export interface QuantizeInt8Result {
+  data: Int8Array;
+  scale: number;
+}
+
+export interface QuantizeInt8PerChannelResult {
+  data: Int8Array;
+  scales: Float32Array;
+}
+
+export interface MemoryEstimate {
+  fp32: number;
+  fp16: number;
+  int8: number;
+}
+
+export function floatToFp16(val: number): number {
+    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: number): number {
+    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: Float32Array): Uint16Array {
+    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: Uint16Array): Float32Array {
+    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: Float32Array): QuantizeInt8Result {
+    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: Int8Array, scale: number): Float32Array {
+    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: Float32Array, numChannels: number): QuantizeInt8PerChannelResult {
+    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: Int8Array, scales: Float32Array, numChannels: number): Float32Array {
+    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: number): MemoryEstimate {
+    return {
+        fp32: numElements * 4,
+        fp16: numElements * 2,
+        int8: numElements * 1,
+    };
+}

package/src/index.js

@@ -1,89 +0,0 @@
-/**
- * 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/tokenizer/bpe.js

@@ -1,256 +0,0 @@
-/**
- * bpe.js – Browser-side Byte Pair Encoding (BPE) tokenizer.
- *
- * Compatible with Qwen3.5-Coder's vocabulary and merge rules.
- * Implements the standard Hugging Face / tiktoken BPE algorithm:
- *
- *   1. Pre-tokenize text into Unicode code point sequences.
- *   2. Apply byte-level encoding for unknown characters.
- *   3. Iteratively merge the highest-priority adjacent byte-pair
- *      according to the learnt merge table.
- *
- * Usage
- * -----
- *   const tokenizer = new BPETokenizer();
- *   await tokenizer.load(vocabUrl, mergesUrl);
- *
- *   const ids  = tokenizer.encode("function foo() {}");
- *   const text = tokenizer.decode(ids);
- */
-
-// Byte-level fallback map (matches GPT-2 / Qwen convention).
-// Maps raw bytes 0-255 to printable Unicode characters so that every
-// byte sequence has a valid string representation.
-function buildByteEncoder() {
-    /** @type {Map<number, string>} */
-    const enc = new Map();
-    const ranges = [
-        [0x21, 0x7E],   // ! → ~
-        [0xA1, 0xAC],   // ¡ → ¬
-        [0xAE, 0xFF],   // ® → ÿ
-    ];
-    let n = 0;
-    for (const [lo, hi] of ranges) {
-        for (let b = lo; b <= hi; b++) {
-            enc.set(b, String.fromCodePoint(b));
-        }
-    }
-    for (let b = 0; b < 256; b++) {
-        if (!enc.has(b)) {
-            enc.set(b, String.fromCodePoint(256 + n));
-            n++;
-        }
-    }
-    return enc;
-}
-
-const BYTE_ENCODER = buildByteEncoder();
-const BYTE_DECODER = new Map([...BYTE_ENCODER].map(([k, v]) => [v, k]));
-
-// Pre-tokenisation regex: matches words, numbers, punctuation, whitespace
-// (closely mirrors tiktoken's GPT-2/Qwen pre-tokenizer pattern).
-const PRE_TOKENIZE_RE =
-    /(?:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+/gu;
-
-export class BPETokenizer {
-    constructor() {
-        /** @type {Map<string, number>} token → id */
-        this.vocab  = new Map();
-        /** @type {Map<number, string>} id → token */
-        this.idToToken = new Map();
-        /** @type {Map<string, number>} "a b" → rank (lower = higher priority) */
-        this.merges = new Map();
-
-        // Special tokens (added after loading vocab)
-        this.bosToken   = '<|im_start|>';
-        this.eosToken   = '<|im_end|>';
-        this.padToken   = '<|endoftext|>';
-        this.unkToken   = '<unk>';
-
-        this.bosId      = null;
-        this.eosId      = null;
-        this.padId      = null;
-    }
-
-    /**
-     * Load vocabulary and merge rules from JSON / text URLs.
-     *
-     * vocab.json  – { "token": id, ... }
-     * merges.txt  – one merge per line: "a b"  (sorted by rank)
-     *
-     * @param {string|object} vocab   – URL string or pre-parsed vocab object
-     * @param {string|string[]} merges – URL string or array of merge strings
-     * @returns {Promise<void>}
-     */
-    async load(vocab, merges) {
-        // Load vocab
-        let vocabObj;
-        if (typeof vocab === 'string') {
-            const res = await fetch(vocab);
-            vocabObj = await res.json();
-        } else {
-            vocabObj = vocab;
-        }
-        this.vocab     = new Map(Object.entries(vocabObj).map(([k, v]) => [k, Number(v)]));
-        this.idToToken = new Map([...this.vocab].map(([k, v]) => [v, k]));
-
-        // Load merges
-        let mergeLines;
-        if (typeof merges === 'string') {
-            const res = await fetch(merges);
-            const txt = await res.text();
-            mergeLines = txt.split('\n').filter(l => l && !l.startsWith('#'));
-        } else {
-            mergeLines = merges;
-        }
-        this.merges = new Map();
-        mergeLines.forEach((line, rank) => {
-            this.merges.set(line.trim(), rank);
-        });
-
-        // Resolve special token ids
-        this.bosId = this.vocab.get(this.bosToken) ?? null;
-        this.eosId = this.vocab.get(this.eosToken) ?? null;
-        this.padId = this.vocab.get(this.padToken) ?? null;
-    }
-
-    /**
-     * Load from plain JavaScript objects (no network fetch).
-     * Useful for bundling a small vocabulary directly.
-     *
-     * @param {Object} vocabObj   – { token: id }
-     * @param {string[]} mergeArr – ["a b", "c d", ...]
-     */
-    loadFromObjects(vocabObj, mergeArr) {
-        this.vocab     = new Map(Object.entries(vocabObj).map(([k, v]) => [k, Number(v)]));
-        this.idToToken = new Map([...this.vocab].map(([k, v]) => [v, k]));
-        this.merges    = new Map(mergeArr.map((m, i) => [m, i]));
-        this.bosId = this.vocab.get(this.bosToken) ?? null;
-        this.eosId = this.vocab.get(this.eosToken) ?? null;
-        this.padId = this.vocab.get(this.padToken) ?? null;
-    }
-
-    // ── Encoding ──────────────────────────────────────────────────────────────
-
-    /**
-     * Encode a string to an array of token IDs.
-     *
-     * @param {string} text
-     * @param {{ addBos?: boolean, addEos?: boolean }} [opts]
-     * @returns {number[]}
-     */
-    encode(text, opts = {}) {
-        const words = text.match(PRE_TOKENIZE_RE) ?? [];
-        const ids   = [];
-
-        if (opts.addBos && this.bosId !== null) ids.push(this.bosId);
-
-        for (const word of words) {
-            // Convert to byte-level encoding
-            const bytes    = new TextEncoder().encode(word);
-            const byteStr  = Array.from(bytes).map(b => BYTE_ENCODER.get(b) ?? '?').join('');
-            const bpeTokens = this._bpe(byteStr);
-
-            for (const tok of bpeTokens) {
-                const id = this.vocab.get(tok);
-                if (id !== undefined) {
-                    ids.push(id);
-                } else {
-                    // Fallback: encode each character individually
-                    for (const ch of tok) {
-                        const cid = this.vocab.get(ch);
-                        if (cid !== undefined) ids.push(cid);
-                    }
-                }
-            }
-        }
-
-        if (opts.addEos && this.eosId !== null) ids.push(this.eosId);
-        return ids;
-    }
-
-    /**
-     * Decode an array of token IDs back to a string.
-     *
-     * @param {number[]} ids
-     * @returns {string}
-     */
-    decode(ids) {
-        let byteStr = '';
-        for (const id of ids) {
-            const tok = this.idToToken.get(id);
-            if (tok !== undefined) byteStr += tok;
-        }
-        // Convert byte-level string back to raw bytes then UTF-8 decode
-        const bytes = new Uint8Array(
-            [...byteStr].map(ch => BYTE_DECODER.get(ch) ?? ch.codePointAt(0))
-        );
-        try {
-            return new TextDecoder('utf-8').decode(bytes);
-        } catch {
-            return byteStr;
-        }
-    }
-
-    // ── BPE merge algorithm ───────────────────────────────────────────────────
-
-    /**
-     * Apply BPE merges to a byte-encoded word.
-     *
-     * @param {string} word  – Space-free string of byte-level characters
-     * @returns {string[]}   – Merged token pieces
-     */
-    _bpe(word) {
-        if (this.vocab.has(word)) return [word];
-
-        // Start: each character is its own symbol
-        let symbols = [...word];
-
-        // Iteratively merge the highest-priority pair
-        while (symbols.length > 1) {
-            let bestRank = Infinity;
-            let bestIdx  = -1;
-
-            for (let i = 0; i < symbols.length - 1; i++) {
-                const pair = symbols[i] + ' ' + symbols[i + 1];
-                const rank = this.merges.get(pair);
-                if (rank !== undefined && rank < bestRank) {
-                    bestRank = rank;
-                    bestIdx  = i;
-                }
-            }
-
-            if (bestIdx === -1) break;  // no more merges available
-
-            // Merge pair at bestIdx
-            const merged = symbols[bestIdx] + symbols[bestIdx + 1];
-            symbols = [
-                ...symbols.slice(0, bestIdx),
-                merged,
-                ...symbols.slice(bestIdx + 2),
-            ];
-        }
-
-        return symbols;
-    }
-
-    // ── Padding / truncation helpers ─────────────────────────────────────────
-
-    /**
-     * Pad or truncate a sequence to a fixed length.
-     *
-     * @param {number[]} ids
-     * @param {number}   maxLen
-     * @param {'right'|'left'} [side='right']
-     * @returns {number[]}
-     */
-    padOrTruncate(ids, maxLen, side = 'right') {
-        if (ids.length >= maxLen) return ids.slice(0, maxLen);
-        const padId = this.padId ?? 0;
-        const pad   = new Array(maxLen - ids.length).fill(padId);
-        return side === 'right' ? [...ids, ...pad] : [...pad, ...ids];
-    }
-
-    /** @returns {number} */
-    get vocabSize() { return this.vocab.size; }
-}