Chat Template API, 1024 Context Window, & Critical GPU Fixes

This release transforms bitnet.js from a raw inference engine into a developer-friendly library, introducing a proper Chat API, expanding the context window, and fixing a critical GPU math bug that caused generation collapse.

New Features

• Chat Template API: Added engine.chat() method with native support for Llama 3 chat formatting (<|start_header_id|>, <|eot_id|>, etc.).
• Expanded Vocabulary: Updated Python extraction scripts to include 10 Llama 3 special tokens (vocab size increased from 16,385 to 16,395).
• Expanded Context Window: Increased MAX_SEQ_LEN from 128 to 1024 tokens, allowing for much longer multi-turn conversations.
• Sampling Parameters: Wired temperature, topK, topP, and repetitionPenalty through the .chat() options (previously hardcoded).
• Examples Directory: Added a new examples/chat/ directory featuring a fully interactive, streaming multi-turn chat UI.

Bug Fixes

• CRITICAL: Strided Softmax Bug: Fixed a severe WGSL shader bug in SHADER_GQA_ATTENTION where tokens beyond index 127 were never normalized due to a flat workgroup size mismatch. The softmax phases (max, exp+sum, normalize) now correctly use strided loops.
• Streaming UI Fix: Fixed a DOM firstChild null reference error in the chat example that occurred during token streaming.

Structural Changes

• Landing Page: Completely redesigned the root index.html into a proper library landing page with Quick Start, API docs, and example cards (using Feather Icons).
• Test Suite Relocation: Moved the automated GPU validation tests from the root to a dedicated tests/ directory.

Rebrand to bitnet.js

What's Changed

• Bump to v0.11.1 and rebrand to bitnet.js by @qwatts-dev in #26

Full Changelog: v0.11.0...v0.11.1

The Engine Refactor

This release transforms the procedural WebGPU proof-of-concept into a clean, object-oriented JavaScript library — now branded as bitnet.js.

Highlights

• ~200 ms/token on MacBook Pro M2 Max (~5 tok/s, 300 GPU dispatches per token)
• Coherent, grammatically correct English output through all 30 transformer layers

Architectural Changes

• BitNetEngine class — encapsulates all WebGPU state, tokenizer, KV cache, and generation loop into a single class with a clean public API: init(), generate(), reset()
• Renamed to bitnet.js — file, package name, and all branding updated to reflect library-grade positioning
• npm run setup — one-command model extraction via scripts/setup_model.py master orchestrator
• npm start — launches the local dev server

Repository Hygiene

• Moved all 9 Python extraction scripts into scripts/
• Centralized all model assets into weights/ (embeddings, LM head, vocab map, ternary weights, norms, scales)
• Deleted stale root-level binaries (sparse_embeddings.bin, sparse_lm_head.bin, vocab_map.json)
• Added __pycache__/ to .gitignore

Documentation

• README fully rewritten: updated Quick Start, project structure, benchmarks section, npm scripts
• docs/TESTING.md — test suite documentation extracted into dedicated file
• package.json — license corrected to MIT, keywords updated, description refreshed

The Awakening

This release marks the successful completion of the core Proof of Concept. The custom WebGPU runtime is now running the full 30-layer microsoft/bitnet-b1.58-2B-4T model natively in the browser, outputting coherent, context-aware English at ~5 tokens per second.

Key Architectural Updates:

• Full Neural Stack: Upgraded the engine from a single layer to the complete 30-layer architecture.
• Llama 3 RoPE: Fixed the Rotary Positional Embedding (RoPE) implementation to correctly use Llama-style rotate_half dimension pairing (paired within heads) and updated the base frequency (ROPE_THETA) to 500000.0.
• BitLinear Weight Scales: Extracted and implemented the floating-point weight_scale for all 210 ternary matrices, preventing activation explosion during integer math.
• Sub-Layer Normalization (SubLN): Extracted and wired the 60 learned SubLN vectors required to stabilize the 1.58-bit architecture.
• Bit-Packing Alignment: Fixed a critical bug in the Python extraction script to correctly interpret Hugging Face's uint8 offset ternary encoding (0b00 -> -1, 0b10 -> +1), restoring the AI's semantic reasoning.
• Squared ReLU: Swapped the standard SiLU activation for the BitNet-specific ReLU² function.

Performance Benchmarks (Warmed-up Inference):

• MacBook Pro M2 Max: ~190 ms/token (~5.2 tokens/sec) for the full 2B parameter pipeline.

The Sampling Engine

This release replaces rigid greedy decoding (argmax) with a robust LLM sampling engine, enabling dynamic and varied text generation while successfully breaking infinite hallucination loops.

Key Architectural Updates:

• Advanced Token Sampling: Implemented sampleToken() featuring Temperature scaling, Top-K filtering, and Top-P (Nucleus) sampling.
• Frequency Penalty: Replaced the flat presence penalty with an exponentially scaling frequency penalty (Map<sparseIdx, count>). This aggressively taxes repeated tokens (penalty^count), successfully forcing the single-layer (Layer 0) engine out of its mathematical feedback loop ("mass mass mass").
• W1.58A8 Parked: The W1.58A8 Activation Quantization shader was successfully written but parked. WebGPU does not yet expose hardware-accelerated integer dot products (DP4a) on Apple Silicon, meaning the FP32 pipeline remains mathematically faster for now. The 2-bit packed weights (W1.58) are already memory-efficient enough (~442 MB for 26 layers) to fit within the iOS 1.5 GB memory limit.

Performance Benchmarks (Layer 0):

• MacBook Pro M2 Max: Sustained ~43.4 ms/token (~23 tokens per second) utilizing the restored FP32 WGSL pipeline.

Auto-Regressive Generation Loop

The engine is no longer a one-shot calculator — it now generates text token-by-token,
streaming results live to the browser.

What's New

Auto-Regressive Pipeline

• forward(tokenId) — Single-token forward pass through the full Layer 0 transformer
  (Embedding → RMSNorm → Self-Attention → Residual → RMSNorm → SwiGLU MLP → Residual → RMSNorm → LM Head).
  Advances the KV cache automatically on each call.
• generateText(prompt, maxTokens, onToken) — Two-phase generation:
  • Prefill: tokenizes the prompt and feeds every token through forward() to warm the KV cache
  • Decode: greedy argmax loop with EOS detection (Llama 3 tokens 128001, 128009) and
    KV cache bounds checking (MAX_SEQ_LEN = 128)
• getEmbeddingByTokenId(tokenId) — Direct embedding lookup by raw token ID, bypassing
  the tokenizer (used internally by the decode loop)

Streaming UI

• Button renamed Compute → Generate
• Prompt echoed in gray, generated tokens stream in green as they arrive
• Final stats show total time and average ms/token (including prefill)

Performance (Layer 0 only, greedy argmax)

Device	ms/token (avg)	20 tokens total
MacBook M2 Max	44.4 ms	~1.0 s

Known Limitations

• Single layer (1/26) — repetitive output is expected; the model lacks the depth
  to produce coherent multi-token sequences. This is not a loop bug.
• Greedy decoding only — no temperature, top-k, or top-p sampling yet.
• Next milestone: multi-layer stacking to unlock real generation quality.

End-to-End Token Prediction on WebGPU

This release completes the first full inference path: Embed → SwiGLU MLP → RMSNorm → LM Head → Argmax → Decode, enabling real word prediction directly in browser WebGPU.

Highlights

• Added dense FP32 LM Head compute path and vocab-logit decoding flow.
• Added LM Head extraction pipeline (tied embeddings) and generated sparse_lm_head.bin (FP16 vocab slice).
• Fixed WebGPU storage-buffer limit handling by requesting adapter hardware limits (unblocks 160 MB LM Head buffer).
• Added RMSNorm stabilization before LM Head to prevent FP32 overflow and NaN/zero logits.
• Verified cross-device deterministic output (M2 Max, M3 iPad, A16 iPhone): identical predicted token (" volume", ID 8286, logit 207.9285).
• Updated docs with mobile/ngrok testing flow and benchmark results for v0.6.0.

Full SwiGLU MLP Block + Unified GPU Orchestration

Complete Layer 0 MLP Block

This release implements the complete SwiGLU Multi-Layer Perceptron block for Layer 0 of microsoft/bitnet-b1.58-2B-4T, running natively in the browser via WebGPU. The full pipeline:

embedding(2560) → gate_proj(6912) ─→ SiLU·mul(6912) → down_proj(2560)
                → up_proj(6912)   ─┘

All three ternary weight matrices are extracted, bit-packed (16 weights per u32), and served as static .bin files:

• gate_proj — 6912×2560 (4,320 KB)
• up_proj — 6912×2560 (4,320 KB)
• down_proj — 2560×6912 (4,320 KB)

Unified GPU Orchestration (Zero-Copy Pipeline)

The critical architectural optimization in this release: all four compute passes execute in a single WebGPU command encoder submission. Intermediate tensors (gate_out, up_out, silu_out) never leave VRAM — eliminating the GPU↔CPU "ping-pong" that was adding ~35ms of PCIe bus latency per step.

Before (v0.4.0 approach): 55.1ms — 4 separate submissions with CPU readback between each
After (unified): 7.2ms — single submission, single readback at the end

Benchmark Results (Warmed Cache)

Compute time for the full SwiGLU MLP block (52.7M ternary parameters, 4 compute passes):

Device	GPU	Setup	Compute	Total
MacBook Pro 14" (2023)	M2 Max 30-core	0.8ms	6.4ms	7.2ms
iPad 13"	M3 10-core	2.0ms	16.0ms	18.0ms
iPhone 14 Pro Max	A16 5-core	9.0ms	29.0ms	38.0ms

Bit-exact determinism — output values are identical across all three devices down to the last decimal place.

Real Semantic Embeddings (from v0.4.0)

Carried forward from the real-embeddings sprint:

• 16,384-word "Vocab-Slice" dictionary using sparse FP16 embeddings
• Successfully bypasses mobile WebKit out-of-memory (OOM) limits
• Browser Cache API (304 Not Modified) for instant reload

New Files

• extract_full_mlp.py — Python script to extract & pack all 3 MLP weight matrices from HuggingFace safetensors
• bitnet_layer_0_gate_proj.bin — packed gate_proj weights
• bitnet_layer_0_up_proj.bin — packed up_proj weights
• SHADER_SILU_MUL — new WGSL compute shader for SwiGLU activation fusion

What's Next

• LM Head — map the MLP output back to vocabulary logits so the model can "speak"
• RMSNorm — proper layer normalization for correct magnitudes
• Attention block — the other half of a transformer layer

Real Semantic Embeddings & Mobile Optimization

This release marks a major milestone: the WebGPU pipeline is no longer running on mock data. It now successfully extracts, loads, and processes the actual high-precision semantic embeddings from Microsoft's bitnet-b1.58-2B-4T model, running end-to-end on desktop, tablet, and mobile devices.

Key Features

• Real Semantic Embeddings: Replaced the PRNG mock embeddings with the true 2,560-dimensional vectors used by the BitNet model. When you type "Hello", the GPU processes the exact mathematical concept of that word.
• Mobile Memory Wall Conquered: Implemented a "Vocab-Slice" strategy to bypass strict iOS WebKit memory limits. By extracting the top 16,384 most common English tokens (plus domain-specific tokens like "WebGPU") and converting them to Float16, the embedding dictionary size was reduced from a crushing 1.3 GB down to a mobile-friendly 80 MB.
• Cross-Device Determinism: The pipeline produces bit-exact identical output across an iPhone 14 Pro Max (A16), iPad Air (M3), and MacBook Pro (M2 Max).
• Interactive Latency on Edge Devices: An iPhone 14 Pro Max successfully processes the 17.7 million ternary parameters of the down_proj layer in just 27 ms.

Technical Changes

• Added extract_sparse_embeddings.py to slice the Hugging Face model.embed_tokens.weight tensor into a sparse FP16 binary and generate a vocab_map.json.
• Updated bitnet-kernel.js to fetch the sparse dictionary, perform OOV (Out-of-Vocabulary) fallback, convert FP16 to Float32 on the fly, and zero-pad the vector to match the GPU kernel's expected 6,912 input dimension.
• Configured Git LFS to track the .bin weight and embedding files.

Benchmarks (Interactive Mode)

• MacBook M2 Max: 6.3 ms compute
• iPad 13" M3: 10.0 ms compute
• iPhone 14 Pro Max: 27.0 ms compute

Micro-Bundle Tokenization & Cache API Integration

This release drastically reduces the project's dependency footprint while maintaining full end-to-end interactive inference capabilities, thanks to a direct architecture suggestion from the creator of transformers.js, @xenova.

Key Updates:

• Micro-Bundle Tokenization: Replaced the full @huggingface/transformers library with the newly released, standalone @huggingface/tokenizers package. This shrinks the required dependency bundle from ~1.2MB down to a microscopic ~8.3kB.
• Browser Cache API: Implemented a custom fetchWithCache wrapper utilizing the native browser Cache API (caches.open). This ensures the Llama 3 tokenizer.json and config files are only downloaded once from the Hub and are instantly loaded from local storage on subsequent runs.
• Syntax Updates: Refactored the text-to-tensor mock bridge to utilize the updated tokenizer.encode() syntax required by the standalone module.
• Verified Determinism: Confirmed that the new, ultra-lightweight tokenization pipeline still feeds perfect, deterministic seeds into the WebGPU engine, maintaining bit-exact output across A16, M3, and M2 Max architectures.