@tryhamster/gerbil 1.0.0-rc.9 → 1.0.0

package/docs/research/nemotron-mamba2-inference.md

@@ -0,0 +1,910 @@
+# Nemotron-3 & Mamba-2: Inference Engine Research
+
+**Research for Gerbil WebGPU inference engine optimization.**
+**Date: March 2026**
+
+---
+
+## Table of Contents
+
+1. [Executive Summary](#1-executive-summary)
+2. [Nemotron-3 Model Family](#2-nemotron-3-model-family)
+3. [Hybrid Mamba-2 + Transformer Architecture](#3-hybrid-mamba-2--transformer-architecture)
+4. [Mamba-2 / SSD Deep Dive](#4-mamba-2--ssd-deep-dive)
+5. [Mixture of Experts (MoE)](#5-mixture-of-experts-moe)
+6. [LatentMoE](#6-latentmoe)
+7. [Multi-Token Prediction (MTP)](#7-multi-token-prediction-mtp)
+8. [Quantization: NVFP4 vs INT4](#8-quantization-nvfp4-vs-int4)
+9. [Nemotron-3 Nano 30B-A3B Config Reference](#9-nemotron-3-nano-30b-a3b-config-reference)
+10. [Gerbil Implications & Implementation Roadmap](#10-gerbil-implications--implementation-roadmap)
+11. [Small Model Viability: Nano vs Qwen 3.5 0.8B](#11-small-model-viability-nano-vs-qwen-35-08b)
+12. [Sources](#12-sources)
+
+---
+
+## 1. Executive Summary
+
+NVIDIA's Nemotron-3 family represents a paradigm shift for local inference: **hybrid Mamba-2 + Transformer + MoE** models that are dramatically faster than pure Transformers while matching or exceeding their accuracy. The key innovations relevant to Gerbil:
+
+| Innovation | What It Does | Gerbil Impact |
+|-----------|-------------|---------------|
+| **Mamba-2 layers** | Replace 92% of attention with constant-memory SSM | Eliminates KV cache growth, unlocks long context on mobile |
+| **LatentMoE** | 4x bandwidth reduction in expert routing | Direct throughput multiplier on bandwidth-bound M4 Max |
+| **MTP** | Native speculative decoding (97% accept rate) | ~2x decode throughput at batch-1, no draft model needed |
+| **NVFP4** | 4-bit floating point with micro-block scaling | Better accuracy/size than our INT4 per-group quantization |
+| **Open weights** | Apache 2.0 / nvidia-open-model-license | Can ship Nemotron models directly in Gerbil |
+
+The smallest model, **Nemotron-3 Nano** (30B total / 3B active), is the primary target for Gerbil. No sub-1B Nemotron model exists yet, so Qwen 3.5 0.8B remains our small-model option. But Nano's 3B active params with MoE sparsity could be feasible on M4 Max (128GB unified memory) and represents a massive capability upgrade.
+
+---
+
+## 2. Nemotron-3 Model Family
+
+| Model | Total Params | Active Params | Architecture | License |
+|-------|-------------|---------------|-------------|---------|
+| **Nano** | 31.6B | 3.6B (w/ embeddings) | Hybrid Mamba-2 + MoE | nvidia-open-model-license |
+| **Super** | ~120B | ~12B | Hybrid Mamba-2 + LatentMoE | nvidia-open-model-license |
+| **Ultra** | ~500B | ~50B | Hybrid Mamba-2 + LatentMoE | nvidia-open-model-license |
+
+### Nano Performance Claims
+- **3.3x faster** than Qwen3-30B-A3B on H200 (8k input / 16k output)
+- **2.2x faster** than GPT-OSS-20B
+- Trained on **25 trillion tokens**
+- Supports **1M token context**
+- Best-in-class on reasoning, coding, math, agentic tasks in its size class
+
+### Quantization Variants Available
+- **BF16** (full precision, official)
+- **FP8** (official)
+- **NVFP4** (official)
+- **GGUF** variants (community: IQ4_XS, MXFP4_MOE via Unsloth)
+
+---
+
+## 3. Hybrid Mamba-2 + Transformer Architecture
+
+### The Core Insight
+
+Pure Transformers have **linearly growing** KV cache and per-token attention compute during decode. Nemotron replaces ~92% of attention layers with Mamba-2, which has **constant** memory and compute per generated token.
+
+### Layer Pattern (Nemotron-3 Nano)
+
+From the `config.json`:
+```
+hybrid_override_pattern: "MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"
+```
+
+Where:
+- **M** = Mamba-2 layer
+- **E** = MoE (expert) layer (replaces standard FFN)
+- **\*** = Attention layer (grouped-query attention with 2 KV heads)
+
+**52 total layers:**
+- **23 Mamba-2 layers** (constant-memory decode, no KV cache)
+- **23 MoE layers** (128 experts, top-6 routing + 1 shared expert)
+- **6 Attention layers** (GQA with 32 heads, 2 KV groups — these need KV cache)
+
+### Why This Matters for Gerbil
+
+Only 6 out of 52 layers need KV cache. For Gerbil on mobile:
+- KV cache memory drops by **~88%** vs pure Transformer
+- Long context becomes feasible within iOS/mobile memory budgets
+- Mamba-2 decode is a simple matrix-vector multiply (no softmax, no causal masking, no KV append)
+
+### Nemotron-H Architecture Search Results
+
+From the Nemotron-H paper, the optimal ratio is **~8% attention layers** across all model sizes:
+- Nemotron-H-8B: 4 attention out of 52 layers
+- Nemotron-H-56B: 10 attention out of 118 layers
+- Nemotron-3 Nano 30B: 6 attention out of 52 layers (slightly higher, ~11.5%)
+
+First layer is always Mamba-2, last layer is always FFN/MoE. Attention layers are evenly dispersed throughout.
+
+### Inference Throughput (Nemotron-H benchmarks, H100)
+
+| Model | vs Transformer Baseline | Context |
+|-------|------------------------|---------|
+| Nemotron-H-8B | **1.8-3x faster** than Qwen-2.5-7B/Llama-3.1-8B | 65k input |
+| Nemotron-H-56B | **2.4x faster** than Qwen-2.5-72B/Llama-3.1-70B | 65k input |
+| Nemotron-H-56B | **19.6x faster** than Llama-3.1-405B | 65k input |
+
+---
+
+## 4. Mamba-2 / SSD Deep Dive
+
+### 4.1 SSM Equations
+
+The Selective State Space Model at the core of Mamba-2:
+
+```
+State update:  h_t = a_t · h_{t-1} + B_t · x_t
+Output:        y_t = C_t^T · h_t
+```
+
+Where per timestep t:
+- `x_t ∈ ℝ^P` — input (P = head_dim, e.g. 64)
+- `h_t ∈ ℝ^(P × N)` — hidden state (N = state_size, e.g. 128)
+- `y_t ∈ ℝ^P` — output
+- `a_t ∈ ℝ` — **scalar** decay (Mamba-2 key simplification)
+- `B_t ∈ ℝ^N` — input projection (input-dependent)
+- `C_t ∈ ℝ^N` — output projection (input-dependent)
+
+### 4.2 Mamba-2 vs Mamba-1
+
+| Aspect | Mamba-1 (S6) | Mamba-2 (SSD) |
+|--------|-------------|---------------|
+| A matrix | Diagonal `(N,)` per head | **Scalar** per head (scalar × identity) |
+| State size N | 16 | **64-256** (much larger) |
+| Training algorithm | Custom CUDA scan | **Matmul-based** (tensor core friendly) |
+| Training speed | Baseline | **2-8x faster** |
+| B, C generation | Sequential (after conv) | **Parallel** (with X) |
+
+The critical insight: by restricting A to a scalar (instead of diagonal), the recurrence dynamics are shared across all N state dimensions. This means the state update can be expressed as matrix multiplications, enabling use of tensor cores during training.
+
+### 4.3 The SSD (Structured State Space Duality)
+
+The SSM recurrence can be equivalently written as a matrix multiplication:
+
+```
+Y = M · X
+
+where M = L ⊙ (C · B^T)    [semiseparable matrix]
+
+L[i,j] = a_i · a_{i-1} · ... · a_{j+1}    for i > j
+L[i,i] = 1
+L[i,j] = 0                                  for i < j
+```
+
+This is structurally identical to **causal linear attention** when all `a_t = 1` (L becomes the causal mask). The scalar `a_t` values act as input-dependent relative positional encodings / decay factors.
+
+### 4.4 Decode Mode (Single Token — Gerbil Hot Path)
+
+For autoregressive generation, Mamba-2 is a simple recurrence:
+
+```python
+# Per token, per head:
+h = a_t * h + outer(B_t, x_t)    # state update: (P, N) = scalar * (P, N) + (N,) ⊗ (P,)
+y = h @ C_t                       # output: (P,) = (P, N) @ (N,)
+```
+
+**Per-token cost per head:**
+- 1 scalar multiply of state: `P × N` multiplies
+- 1 outer product + add: `P × N` multiply-adds
+- 1 matrix-vector product: `P × N` multiply-adds
+- **Total: ~3PN FLOPs per head**
+
+**For Nemotron-3 Nano** (64 heads, P=64, N=128):
+- Per Mamba layer: `64 × 3 × 64 × 128 = 1.57M FLOPs`
+- Compare to attention: KV cache read + softmax + output = much more for long sequences
+
+**State memory per Mamba layer:**
+- `num_heads × head_dim × state_size = 64 × 64 × 128 = 524,288` values
+- At f32: **2MB per layer**, at f16: **1MB per layer**
+- For 23 Mamba layers: **~23MB total** (constant, regardless of sequence length!)
+
+Compare to attention KV cache for 1024 tokens: `6 layers × 2(K+V) × 32 heads × 128 dim × 1024 tokens × 2 bytes = ~100MB` and growing linearly.
+
+### 4.5 Prefill Mode (Parallel — Processing Prompt)
+
+The chunked SSD algorithm processes the prompt in parallel using 4 steps:
+
+```
+Given: X (T, P), A (T,), B (T, N), C (T, N)
+Chunk into blocks of size Q (default 64-128):
+  X_chunks: (T/Q, Q, P)
+  A_chunks: (T/Q, Q), etc.
+
+Step 1 — Intra-chunk (parallel, uses matmuls):
+  Y_diag = einsum("bclhn, bcshn, bhcls, bcshp -> bclhp", C, B, L, X)
+  # L is the Q×Q semiseparable mask from cumsum of A within each chunk
+
+Step 2 — Chunk states (parallel):
+  states = einsum("bclhn, bhcl, bclhp -> bchpn", B, decay, X)
+  # Final state of each chunk assuming zero initial state
+
+Step 3 — Inter-chunk recurrence (sequential, but over T/Q chunks only):
+  new_states[c] = decay_chunk[c] * new_states[c-1] + states[c]
+  # Simple scan over reduced sequence length T/Q
+
+Step 4 — Output correction (parallel):
+  Y_off = einsum("bclhn, bchpn, bhcl -> bclhp", C, states, state_decay)
+  # Add contribution from initial states to each chunk's output
+
+Final: Y = Y_diag + Y_off
+```
+
+**Key property:** Steps 1, 2, 4 are all matmuls (tensor core / GPU friendly). Step 3 is a short sequential scan over T/Q elements (e.g., for 2048 tokens with Q=128, only 16 sequential steps).
+
+### 4.6 Mamba-2 Block Architecture
+
+```
+Input x (d_model)
+    │
+    ├──→ Linear projection → (expand * d_model) ──→ split into:
+    │        │                                         │
+    │        ├── z (gate, expand * d_model)            │
+    │        └── x' (SSM input, expand * d_model)      │
+    │                    │                              │
+    │              Conv1d(kernel=4, causal)             │
+    │                    │                              │
+    │                 SiLU(x')                          │
+    │                    │                              │
+    │           ┌── B_t = Linear(x')  (N per head)     │
+    │           ├── C_t = Linear(x')  (N per head)     │
+    │           ├── dt  = Linear(x')  → softplus → a_t │
+    │           │                                       │
+    │           └──── SSM(x', a_t, B_t, C_t) ──→ y    │
+    │                                              │    │
+    │                                     y * SiLU(z)  │ (gating)
+    │                                              │    │
+    │                                    Linear(out) ──→ output (d_model)
+```
+
+**Nemotron-3 Nano Mamba-2 config:**
+- `expand = 2` → inner dim = 2 × 2688 = 5376
+- `mamba_head_dim = 64` (P)
+- `mamba_num_heads = 64` → 64 × 64 = 4096 (but expand × hidden = 5376?)
+- `ssm_state_size = 128` (N)
+- `conv_kernel = 4`
+- `n_groups = 8` (Mamba groups, likely for B/C sharing)
+- `chunk_size = 128` (Q for chunked prefill)
+
+### 4.7 Weight Tensors for Mamba-2 Layer
+
+Expected HuggingFace weight names per layer:
+```
+model.layers.{i}.mamba.in_proj.weight      # (expand*d + expand*d, d_model) or split
+model.layers.{i}.mamba.conv1d.weight       # (inner_dim, 1, conv_kernel)
+model.layers.{i}.mamba.conv1d.bias         # (inner_dim,)
+model.layers.{i}.mamba.out_proj.weight     # (d_model, inner_dim)
+model.layers.{i}.mamba.dt_bias             # (num_heads,) — added to dt before softplus
+model.layers.{i}.mamba.A_log              # (num_heads,) — log-space A parameter
+model.layers.{i}.mamba.D                  # (num_heads,) — skip connection scalar
+model.layers.{i}.mamba.norm.weight        # (inner_dim,) — RMSNorm before output
+```
+
+Note: `A_log` is stored in log-space for numerical stability. At runtime: `A = -exp(A_log)` (negative to ensure decay).
+
+### 4.8 Decode Path Implementation Pseudocode
+
+```python
+def mamba2_decode_step(x, layer, state):
+    """Single token decode for one Mamba-2 layer.
+
+    x: (d_model,) — input token embedding
+    state.h: (num_heads, head_dim, state_size) — SSM hidden state
+    state.conv: (inner_dim, conv_kernel-1) — conv1d rolling buffer
+    """
+    # 1. Input projection
+    xz = layer.in_proj @ x                    # (2 * inner_dim,)
+    x_inner, z = split(xz, inner_dim)         # each (inner_dim,)
+
+    # 2. Causal conv1d (shift buffer, apply kernel)
+    state.conv = roll_left(state.conv)
+    state.conv[:, -1] = x_inner
+    x_conv = sum(state.conv * layer.conv1d_weight, dim=-1) + layer.conv1d_bias
+
+    # 3. SiLU activation
+    x_act = x_conv * sigmoid(x_conv)          # SiLU = x * sigmoid(x)
+
+    # 4. Generate B, C, dt from activated input
+    # (These come from the in_proj output, exact split depends on implementation)
+    B = ...                                     # (num_heads, state_size)
+    C = ...                                     # (num_heads, state_size)
+    dt = softplus(layer.dt_proj(x_act) + layer.dt_bias)  # (num_heads,)
+
+    # 5. Discretize A
+    A = -exp(layer.A_log)                       # (num_heads,) — negative real
+    a = exp(dt * A)                             # (num_heads,) — decay per head
+
+    # 6. SSM state update (THE HOT LOOP)
+    # state.h shape: (num_heads, head_dim, state_size)
+    # a shape: (num_heads,) — broadcast over head_dim and state_size
+    x_heads = reshape(x_act, (num_heads, head_dim))  # (num_heads, head_dim)
+
+    for each head h:
+        state.h[h] = a[h] * state.h[h] + outer(x_heads[h], B[h])
+        # (head_dim, state_size) = scalar * (head_dim, state_size) + (head_dim,1) @ (1, state_size)
+
+    # 7. SSM output
+    y_heads = einsum("hpn, hn -> hp", state.h, C)  # (num_heads, head_dim)
+    y = reshape(y_heads, (inner_dim,))
+
+    # 8. Add D skip connection
+    y = y + reshape(layer.D, ...) * x_act
+
+    # 9. RMSNorm
+    y = layer.norm(y)
+
+    # 10. Gate and output
+    y = y * (z * sigmoid(z))                    # gate with SiLU(z)
+    output = layer.out_proj @ y                  # (d_model,)
+
+    return output, state
+```
+
+### 4.9 Memory Bandwidth Analysis (Decode)
+
+For one Mamba-2 layer in Nemotron Nano, per token:
+
+**Weights to read:**
+- `in_proj`: 2688 × 5376 × 2 bytes = ~28.9 MB (at f16)
+- `conv1d`: 5376 × 4 × 2 = ~43 KB
+- `out_proj`: 5376 × 2688 × 2 = ~28.9 MB
+- Other (dt_bias, A_log, D, norm): negligible
+- **Total weights: ~58 MB per Mamba layer**
+
+**State to read/write:**
+- SSM state: 64 × 64 × 128 × 4 = 2 MB (at f32)
+- Conv buffer: 5376 × 3 × 2 = ~32 KB
+- **Total state: ~2 MB per layer**
+
+**Comparison with attention layer:**
+- Attention weights: QKV proj + output proj ≈ similar to Mamba projections
+- **Plus KV cache read**: grows with sequence length (at 1024 tokens: ~1 MB per layer per read)
+- At long sequences, KV cache read dominates — Mamba avoids this entirely
+
+---
+
+## 5. Mixture of Experts (MoE)
+
+### Nemotron-3 Nano MoE Configuration
+
+```json
+{
+  "n_routed_experts": 128,
+  "num_experts_per_tok": 6,
+  "n_shared_experts": 1,
+  "moe_intermediate_size": 1856,
+  "moe_shared_expert_intermediate_size": 3712,
+  "routed_scaling_factor": 2.5,
+  "norm_topk_prob": true,
+  "topk_group": 1
+}
+```
+
+### How MoE Works at Inference
+
+For each token:
+1. Router (small MLP) computes scores over all 128 experts
+2. Top-6 experts selected
+3. Each selected expert processes the token independently
+4. Outputs are weighted-summed by router scores
+5. Shared expert (always active) output is added
+
+**Active parameters per token:** 6 × (2688 × 1856 × 2) + 1 × (2688 × 3712 × 2) ≈ **80M params** (per MoE layer)
+
+### Memory Implications for Gerbil
+
+All 128 expert weights must be in GPU memory, but only 6 are read per token:
+- Per MoE layer total weights: 128 × 2 × 2688 × 1856 × 2 bytes ≈ **2.5 GB** (at f16)
+- Per MoE layer per-token read: 6 × 2 × 2688 × 1856 × 2 ≈ **120 MB**
+- Plus shared expert: 2 × 2688 × 3712 × 2 ≈ **40 MB**
+- **Total per MoE layer per token: ~160 MB bandwidth**
+
+With 23 MoE layers: **~3.7 GB bandwidth per token** — this is the bottleneck.
+
+At M4 Max 546 GB/s: theoretical minimum **~6.8ms per token** just for MoE weight reads (assuming perfect bandwidth utilization).
+
+---
+
+## 6. LatentMoE
+
+### Architecture
+
+LatentMoE reduces the memory bandwidth of MoE by projecting tokens to a smaller latent space before expert routing:
+
+```
+Standard MoE:
+  x (d) → Router → Expert_i(x) (d→m→d) → weighted sum
+
+LatentMoE:
+  x (d) → Router (in d-space)
+         → W_down @ x (d→ℓ)
+         → Expert_i(x_latent) (ℓ→m→ℓ)
+         → W_up @ y_latent (ℓ→d)
+         → weighted sum
+```
+
+### Key Properties
+
+- **Compression ratio α = d/ℓ = 4** (typical, validated up to 4x without quality loss)
+- **Router stays in d-space** — routing decisions use full-dimensional information
+- **Experts operate in ℓ-space** — 4x smaller input/output dimensions
+- **Expert count scales up**: N′ = α × N (e.g., 128 → 512 experts)
+- **Active experts scale up**: K′ = α × K (e.g., 6 → 24 active)
+- **Net effect**: same compute, same params, but 4x less bandwidth per expert
+
+### Bandwidth Savings
+
+| Metric | Standard MoE | LatentMoE (α=4) |
+|--------|-------------|-----------------|
+| Expert weight read per token | `K × 2 × d × m` | `K' × 2 × ℓ × m` (same total) |
+| All-to-all communication | `K × d` per token | `K × ℓ` per token (**4x less**) |
+| Routing payload | `d` per token per expert | `ℓ` per token per expert (**4x less**) |
+
+### Performance
+
+- At iso-accuracy: **up to 3.5x throughput improvement** over standard MoE
+- At iso-FLOP: significant accuracy gains (MMLU-Pro: +5.65 points)
+- Overhead of down/up projection: **<9%** of total compute
+
+### Gerbil Relevance
+
+Nemotron-3 Nano uses standard MoE (not LatentMoE). LatentMoE is used in Super and Ultra models. However:
+- If NVIDIA releases a Nano-class model with LatentMoE, the bandwidth savings would be transformative for WebGPU
+- The down/up projections are simple matmuls — trivial to implement in WGSL
+- Could reduce our per-token MoE bandwidth from ~160MB to ~40MB per layer
+
+---
+
+## 7. Multi-Token Prediction (MTP)
+
+### How It Works
+
+Instead of predicting only the next token, the model has N additional prediction heads that predict tokens 2, 3, ..., N+1 steps ahead.
+
+```
+Backbone output (hidden states)
+    │
+    ├── Head 0 (standard): predict token t+1
+    ├── Head 1 (MTP): predict token t+2
+    ├── Head 2 (MTP): predict token t+3
+    └── Head 3 (MTP): predict token t+4
+```
+
+### Training Benefits
+- Richer training signal (predicting further ahead encourages planning)
+- ~2.4% average improvement across benchmarks
+- Minimal additional FLOPs during training
+
+### Inference: Self-Speculative Decoding
+
+The MTP heads enable speculative decoding **without a separate draft model**:
+
+```
+1. Run one forward pass → get predictions from all heads
+2. Head 0 predicts token t+1 (high confidence)
+3. Head 1 predicts token t+2 (draft)
+4. Head 2 predicts token t+3 (draft)
+5. Verify drafts by running a single forward pass on [t+1, t+2, t+3]
+6. Accept all tokens up to first mismatch
+```
+
+### Key Numbers
+- **97% acceptance rate** on first two predicted tokens (Nemotron-3)
+- **Up to 3x faster** inference with 4-token prediction heads (Meta paper)
+- Particularly effective at **batch-size-1** (exactly Gerbil's use case)
+
+### Gerbil Implementation
+
+For decode, MTP would:
+1. Add lightweight prediction heads (shared backbone, separate output projections)
+2. Generate 2-4 draft tokens per forward pass
+3. Verify all drafts in a single prefill-style forward pass
+4. Accept valid tokens, fall back to first mismatch
+
+**Expected speedup: ~1.5-2x** decode throughput at batch-1, with near-zero memory overhead (heads are just additional small linear projections over vocab).
+
+**Note:** Nemotron-3 Nano's config doesn't explicitly show MTP configuration, so this may be a Super/Ultra feature or baked into the architecture at the weights level. Needs verification.
+
+---
+
+## 8. Quantization: NVFP4 vs INT4
+
+### Our Current INT4
+
+```
+Format: 4-bit unsigned integer, packed 8 nibbles per u32
+Dequant: output = (nibble - zero) * scale
+Group size: 128 (configurable 32/64)
+Scale/zero: f32 per group
+```
+
+### NVFP4
+
+```
+Format: E2M1 (2-bit exponent, 1-bit mantissa)
+Values: {0, 0.5, 1, 1.5, 2, 3, 4, 6} × sign
+Block scaling: E4M3 format, per 16-element micro-block
+Global scaling: FP32, second-level scale
+2D block scaling for weights
+```
+
+### Comparison
+
+| Aspect | Our INT4 | NVFP4 |
+|--------|---------|-------|
+| Element format | 4-bit unsigned int | E2M1 (4-bit float) |
+| Group/block size | 128 elements | **16 elements** (8x finer) |
+| Scale format | f32 | E4M3 + FP32 (two-level) |
+| Zero point | Per-group | Not needed (symmetric) |
+| Accuracy (vs BF16) | ~decent (no published numbers) | **<1% relative loss difference** |
+| Selective precision | No | **Last 15% of layers in BF16** |
+
+### Implications for Gerbil
+
+- NVFP4's 16-element micro-blocks are **8x finer granularity** than our 128-element groups
+- This means 8x more scale factors to read, but much better accuracy
+- E2M1 dequantization is simpler than INT4 (no zero point subtraction)
+- Could implement as: read 2 packed u32s (16 nibbles) + 1 E4M3 scale + share FP32 global scale
+- NVFP4 weights from HuggingFace could be loaded directly without requantization
+
+---
+
+## 9. Nemotron-3 Nano 30B-A3B Config Reference
+
+Complete architecture config from HuggingFace:
+
+```json
+{
+  "architectures": ["NemotronHForCausalLM"],
+  "model_type": "nemotron_h",
+
+  "hidden_size": 2688,
+  "num_hidden_layers": 52,
+  "intermediate_size": 1856,
+  "vocab_size": 131072,
+  "max_position_embeddings": 262144,
+
+  "hybrid_override_pattern": "MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME",
+
+  "num_attention_heads": 32,
+  "num_key_value_heads": 2,
+  "head_dim": 128,
+  "rope_theta": 10000,
+
+  "mamba_num_heads": 64,
+  "mamba_head_dim": 64,
+  "ssm_state_size": 128,
+  "expand": 2,
+  "conv_kernel": 4,
+  "n_groups": 8,
+  "chunk_size": 128,
+  "mamba_hidden_act": "silu",
+
+  "n_routed_experts": 128,
+  "num_experts_per_tok": 6,
+  "n_shared_experts": 1,
+  "moe_intermediate_size": 1856,
+  "moe_shared_expert_intermediate_size": 3712,
+  "routed_scaling_factor": 2.5,
+
+  "mlp_hidden_act": "relu2",
+  "layer_norm_epsilon": 1e-05,
+  "torch_dtype": "bfloat16"
+}
+```
+
+### Layer Pattern Decoded
+
+```
+Position:  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19
+Layer:     M  E  M  E  M  *  E  M  E  M  E  M  *  E  M  E  M  E  M  *
+Type:      m  x  m  x  m  a  x  m  x  m  x  m  a  x  m  x  m  x  m  a
+
+Position: 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39
+Layer:     E  M  E  M  E  M  *  E  M  E  M  E  M  E  M  *  E  M  E  M
+Type:      x  m  x  m  x  m  a  x  m  x  m  x  m  x  m  a  x  m  x  m
+
+Position: 40 41 42 43 44 45 46 47 48 49 50 51
+Layer:     E  M  E  M  E  M  E  M  E  M  E  M  E
+Type:      x  m  x  m  x  m  x  m  a  m  x  m  x  (approximate — pattern varies at end)
+
+m = Mamba-2, x = MoE, a = Attention (GQA)
+```
+
+### Derived Dimensions
+
+| Component | Dimension | Notes |
+|-----------|-----------|-------|
+| Mamba inner dim | 2 × 2688 = 5376 | expand=2 |
+| Mamba state per layer | 64 heads × 64 head_dim × 128 state = 524K values | ~2MB at f32 |
+| Mamba conv buffer | 5376 × 3 values | conv_kernel=4, store last 3 |
+| Attention KV per token | 2 KV heads × 128 dim × 2(K+V) = 512 values | ~1KB at f16 |
+| MoE per-token read | 6 experts × 2 × 2688 × 1856 = ~60M values | ~120MB at f16 |
+| Total model weights (BF16) | ~63 GB | All 128 experts × 23 layers |
+| Total model weights (INT4) | ~16 GB | With quantization |
+| Active weights per token | ~1.8 GB | Only active experts + Mamba/attention |
+
+### Weight Size Estimates (INT4/FP4 quantized)
+
+| Component | Per Layer | Total |
+|-----------|-----------|-------|
+| Mamba in_proj | 2688 × 5376 / 2 = 7.2 MB | 23 × 7.2 = 166 MB |
+| Mamba out_proj | 5376 × 2688 / 2 = 7.2 MB | 23 × 7.2 = 166 MB |
+| Mamba conv + small | ~11 KB | ~0.25 MB |
+| MoE routed experts | 128 × 2 × 2688 × 1856 / 2 = 1.22 GB | 23 × 1.22 = **28 GB** |
+| MoE shared expert | 2 × 2688 × 3712 / 2 = 10 MB | 23 × 10 = 230 MB |
+| Attention (QKV+O) | 4 × 2688 × 2688 / 2 = 14.5 MB | 6 × 14.5 = 87 MB |
+| Embeddings + head | 131072 × 2688 × 2 = 670 MB | 670 MB |
+| **Total (INT4)** | | **~29.3 GB** |
+
+This fits in M4 Max 128GB unified memory, but is tight for 64GB configs and impossible for mobile.
+
+---
+
+## 10. Gerbil Implications & Implementation Roadmap
+
+### Phase 1: Mamba-2 Kernel (WGSL)
+
+**New kernels needed:**
+
+1. **Mamba2Decode** (hot path — single token)
+   - SSM state update: `h = a*h + B⊗x` per head
+   - SSM output: `y = h @ C` per head
+   - Fuse with SiLU gate and output projection
+   - Workgroup: one workgroup per head (64 threads for head_dim=64)
+   - State in storage buffer (persistent across tokens)
+
+2. **Mamba2Conv1d** (per token)
+   - 1D causal convolution, kernel=4
+   - Maintain rolling buffer of last 3 activations
+   - Simple: read 4 values, multiply-accumulate
+
+3. **Mamba2Prefill** (prompt processing)
+   - Chunked SSD algorithm (4 steps)
+   - Steps 1,2,4 are matmuls — reuse existing MatMul kernels
+   - Step 3 is short scan — new small kernel
+   - Chunk size 128 (from config)
+
+4. **Mamba2InputProj / OutputProj**
+   - Standard MatVec (decode) or MatMul (prefill)
+   - Reuse existing kernels, just need correct dispatch
+
+**State management changes:**
+- New persistent buffer type: SSM state (per layer, constant size)
+- Conv buffer: rolling window (per layer, width=3)
+- These replace KV cache for Mamba layers (massive memory savings)
+
+### Phase 2: MoE Routing
+
+**New kernels:**
+1. **TopKRouter** — compute router scores, find top-6 experts
+2. **MoEGather** — select expert weights for active experts
+3. **MoEScatter** — weighted sum of expert outputs
+
+**Architecture changes:**
+- Expert weights stored as large buffers (128 experts × weight matrix)
+- Router dispatches only active expert kernels per token
+- Shared expert always dispatched in parallel
+
+### Phase 3: Graph Generator for NemotronH
+
+New architecture handler:
+```typescript
+function generateNemotronHGraph(config: NemotronHConfig): ModelGraph {
+  const pattern = config.hybrid_override_pattern;
+  for (let i = 0; i < config.num_hidden_layers; i++) {
+    if (pattern[i] === 'M') {
+      // Mamba-2 layer: in_proj → conv1d → SSM → gate → out_proj
+      addMamba2Layer(graph, i);
+    } else if (pattern[i] === '*') {
+      // Attention layer: QKV → attention → out_proj (existing kernels)
+      addAttentionLayer(graph, i);
+    } else if (pattern[i] === 'E') {
+      // MoE layer: router → top-k gather → expert FFN → scatter
+      addMoELayer(graph, i);
+    }
+  }
+}
+```
+
+### Phase 4: MTP Speculative Decoding (if applicable)
+
+- Add lightweight prediction heads
+- Implement verify-and-accept loop
+- Expected ~1.5-2x decode speedup
+
+### Priority Order
+
+1. **NemotronH graph generator** — ships the architecture adapter, validates IR design
+2. **Mamba-2 decode kernel** — gates everything, simplest to validate
+3. **MoE routing + dispatch** — needed for any Nemotron model
+4. **Mamba-2 prefill (chunked SSD)** — needed for prompt processing
+5. **NVFP4 dequantization** — use their quantized weights directly
+6. **MTP speculative decode** — optimization on top
+
+---
+
+## 10.5. NemotronH Graph Generator — Detailed Implementation Plan
+
+### Verified Model Structure (from `modeling_nemotron_h.py`)
+
+The actual HuggingFace implementation reveals key differences from assumptions:
+
+**Weight prefix is `backbone.` not `model.`:**
+```
+backbone.embeddings.weight
+backbone.layers.{i}.norm.weight
+backbone.layers.{i}.mixer.{component}
+backbone.norm_f.weight
+lm_head.weight
+```
+
+**All layer types share the `.mixer.` namespace** — Mamba, Attention, and MLP/MoE all live under `layers.{i}.mixer.*`.
+
+**Pattern characters map to:**
+- **M** → `NemotronHMamba2Mixer` (standard Mamba-2 SSM)
+- **E** → `NemotronHMLP` (8B dense model) / MoE (30B Nano model)
+- **\*** → `NemotronHAttention` (GQA)
+
+### Verified Safetensors Key Names
+
+**Mamba layers (`M`):**
+```
+backbone.layers.{i}.norm.weight                    # pre-norm
+backbone.layers.{i}.mixer.in_proj.weight           # fused [d_mlp, d_mlp, gate, x_B_C, dt]
+backbone.layers.{i}.mixer.conv1d.weight            # depthwise conv1d
+backbone.layers.{i}.mixer.conv1d.bias
+backbone.layers.{i}.mixer.dt_bias                  # (num_heads,)
+backbone.layers.{i}.mixer.A_log                    # (num_heads,) log-space decay
+backbone.layers.{i}.mixer.D                        # (num_heads,) skip connection
+backbone.layers.{i}.mixer.norm.weight              # gated RMSNorm before output
+backbone.layers.{i}.mixer.out_proj.weight          # back to hidden_size
+```
+
+**Attention layers (`*`):**
+```
+backbone.layers.{i}.norm.weight
+backbone.layers.{i}.mixer.q_proj.weight
+backbone.layers.{i}.mixer.k_proj.weight
+backbone.layers.{i}.mixer.v_proj.weight
+backbone.layers.{i}.mixer.o_proj.weight
+```
+
+**MLP layers (`E` in 8B dense):**
+```
+backbone.layers.{i}.norm.weight
+backbone.layers.{i}.mixer.up_proj.weight
+backbone.layers.{i}.mixer.down_proj.weight
+```
+
+**MoE layers (`E` in 30B Nano) — needs verification:**
+```
+backbone.layers.{i}.norm.weight
+backbone.layers.{i}.mixer.gate.weight              # router
+backbone.layers.{i}.mixer.experts.{j}.up_proj.weight
+backbone.layers.{i}.mixer.experts.{j}.down_proj.weight
+backbone.layers.{i}.mixer.shared_experts.up_proj.weight    # shared expert
+backbone.layers.{i}.mixer.shared_experts.down_proj.weight
+```
+
+### Mamba-2 in_proj Split (from source code)
+
+NemotronH's Mamba-2 fuses everything into a single `in_proj`:
+
+```python
+# in_proj output is split as:
+d_mlp, d_mlp, gate, hidden_states_B_C, dt = projected_states.split(
+    [d_mlp, d_mlp, intermediate_size, conv_dim, num_heads], dim=-1
+)
+
+# where:
+#   intermediate_size = mamba_num_heads * mamba_head_dim = 64 * 64 = 4096
+#   conv_dim = intermediate_size + 2 * n_groups * ssm_state_size
+#            = 4096 + 2 * 8 * 128 = 6144
+#   num_heads = 64 (for dt)
+#   d_mlp = any remaining dimensions / 2 (optional MLP component)
+
+# Then hidden_states_B_C is further split:
+hidden_states, B, C = split(hidden_states_B_C,
+    [intermediate_size, n_groups*ssm_state_size, n_groups*ssm_state_size])
+```
+
+This differs from Qwen3.5's separate `in_proj_qkv`, `in_proj_a`, `in_proj_b`, `in_proj_z` projections.
+
+### Files to Change
+
+| File | Change | Size |
+|------|--------|------|
+| `src/gpu/architectures/nemotron_h.ts` | **New file** — graph generator | ~400 lines |
+| `src/gpu/architectures/index.ts` | Register `NemotronHForCausalLM` | 3 lines |
+| `src/gpu/ir.ts` | Add NemotronH canonical keys + HF key mapper | ~40 lines |
+
+### Existing Op Reuse
+
+| Component | Exists in Gerbil? | Reusable? |
+|-----------|-------------------|-----------|
+| MatMul / MatMulInt4 | Yes | Direct reuse for all projections |
+| Attention + KV cache | Yes | Direct reuse for \* layers |
+| RMSNorm / ResidualRMSNorm | Yes | Direct reuse |
+| SiLU / SwiGLU | Yes | Direct reuse for gating |
+| CausalConv1d | Yes (Qwen3.5) | Direct reuse for Mamba conv |
+| ConvStateUpdate | Yes (Qwen3.5) | Direct reuse |
+| MambaSSM | **Partial** | Qwen3.5 implements Gated DeltaNet. NemotronH needs standard Mamba-2 SSM (different recurrence: `h = a*h + outer(B,x)` vs delta rule). Needs new kernel variant. |
+| ReLU² activation | **No** | Trivial — `relu(x)²` or `x * relu(x)`, ~20 lines WGSL |
+| MoERouter (top-k) | **No** (stubbed in IR) | New kernel: compute router logits, select top-6 |
+| MoE Expert dispatch | **No** (stubbed in IR) | New kernel: route tokens to experts, weighted combine |
+
+### What the Graph Generator Produces (But Can't Run Yet)
+
+The graph generator will emit correct IR nodes referencing ops whose kernels don't exist yet. This is the same pattern used for Qwen3.5 — the graph is correct documentation of the compute graph, and when kernels are implemented, inference works automatically.
+
+**Ops that block end-to-end execution:**
+1. **MoE kernels** (23 E-layers × 128 experts) — the big one
+2. **Standard Mamba-2 SSM kernel** — differs from Qwen3.5's DeltaNet kernel
+3. **ReLU² activation** — trivial to add
+
+### Approach
+
+Ship the graph generator now (Option A). It serves as:
+- Verified architecture documentation (weight names, dimensions, layer pattern)
+- Correct IR that will work when kernels land
+- Test target for kernel development (load config → generate graph → inspect nodes)
+
+---
+
+## 11. Small Model Viability: Nano vs Qwen 3.5 0.8B
+
+### Does Nano Compete at the Small End?
+
+**No.** Nemotron-3 Nano is not a small model replacement for Qwen 3.5 0.8B:
+
+| Aspect | Qwen 3.5 0.8B | Nemotron-3 Nano 30B-A3B |
+|--------|--------------|------------------------|
+| Total params | 0.8B | 31.6B |
+| Active params | 0.8B (dense) | 3.6B (MoE) |
+| Weights (INT4) | ~0.4 GB | ~29 GB |
+| Target hardware | Mobile, browser, edge | M4 Max, RTX, DGX Spark |
+| Context length | 32K | 1M |
+| Capability | Basic chat/generation | Full reasoning, coding, agentic |
+
+**The use cases are different:**
+- **Qwen 3.5 0.8B**: ultra-lightweight, runs on any device, good enough for simple tasks
+- **Nemotron Nano**: dramatically more capable, but needs 29GB for weights alone
+
+### No Sub-1B Nemotron Models Exist
+
+NVIDIA has not released any Nemotron model with fewer than 3B active parameters. The family starts at Nano (3.6B active / 31.6B total). There are no announced plans for smaller variants.
+
+### Gerbil Strategy
+
+**Keep both:**
+- **Qwen 3.5 0.8B** for mobile/browser (our current optimized path)
+- **Nemotron-3 Nano** for desktop with M4 Max or similar (29GB fits in 128GB unified memory)
+- Watch for smaller Nemotron or community-distilled variants
+
+### Could Nemotron Nano Run on M4 Max via Gerbil?
+
+**Theoretically yes**, with caveats:
+- INT4 weights: ~29 GB (fits in 128GB, tight for 64GB)
+- Active weight reads per token: ~1.8 GB of bandwidth
+- At 546 GB/s (M4 Max): theoretical ~300 tok/s **if perfectly bandwidth-bound**
+- Realistic with overhead: **50-150 tok/s** estimated
+- Would need MoE expert weight management (not all 128 experts in active cache)
+
+Compared to our current 160 tok/s on Qwen 0.8B, Nano would be slower but **dramatically more capable** — it's a reasoning model that can do multi-step coding, math, and tool use.
+
+**Important:** MoE means **all 30B weights must reside in memory** even though only 3.5B are active per token. The router selects different experts per token unpredictably. The saving is in compute FLOPs, not memory bandwidth — this is crucial for bandwidth-bound WebGPU inference.
+
+### Community Small Mamba-2 Models
+
+If smaller Mamba-2 models are of future interest:
+
+| Model | Params | Type | Notes |
+|-------|--------|------|-------|
+| `state-spaces/mamba2-2.7b` | 2.7B | Pure Mamba-2, dense | Base model only (no instruct) |
+| `state-spaces/mamba2-1.3b` | 1.3B | Pure Mamba-2, dense | Base model only |
+| `state-spaces/mamba2-780m` | 780M | Pure Mamba-2, dense | Base model only |
+| `state-spaces/mamba2attn-2.7b` | 2.7B | Mamba-2 + Attention hybrid | Minimal documentation |
+| `Zyphra/Zamba2-1.2B-Instruct-v2` | 1.2B | Mamba-2 + shared attention | Apache 2.0, instruct-tuned |
+
+None of these are Nemotron-class models, but they prove the Mamba-2 kernel would have broader utility beyond just Nemotron.
+
+---
+
+## 12. Sources
+
+- [NVIDIA Nemotron 3: Efficient and Open Intelligence (arxiv 2512.20856)](https://arxiv.org/abs/2512.20856)
+- [Nemotron-H: A Family of Accurate and Efficient Hybrid Mamba-Transformer Models (arxiv 2504.03624)](https://arxiv.org/abs/2504.03624)
+- [Nemotron 3 Nano: Open, Efficient MoE Hybrid Mamba-Transformer (arxiv 2512.20848)](https://arxiv.org/abs/2512.20848)
+- [Mamba-2 / SSD: Transformers are SSMs (arxiv 2405.21060)](https://arxiv.org/abs/2405.21060)
+- [Tri Dao blog: Mamba-2 Part I - The Model](https://tridao.me/blog/2024/mamba2-part1-model/)
+- [Tri Dao blog: Mamba-2 Part III - The Algorithm](https://tridao.me/blog/2024/mamba2-part3-algorithm/)
+- [LatentMoE: Toward Optimal Accuracy per FLOP (arxiv 2601.18089)](https://arxiv.org/abs/2601.18089)
+- [LatentMoE NVIDIA Research Page](https://research.nvidia.com/labs/nemotron/LatentMoE/)
+- [Meta: Better & Faster LLMs via Multi-token Prediction (arxiv 2404.19737)](https://arxiv.org/abs/2404.19737)
+- [Nemotron-3 Nano HuggingFace Blog](https://huggingface.co/blog/nvidia/nemotron-3-nano-efficient-open-intelligent-models)
+- [Nemotron-3 Nano FP8 Config.json](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8/blob/main/config.json)
+- [NVIDIA Nemotron-3 Super Blog](https://developer.nvidia.com/blog/introducing-nemotron-3-super-an-open-hybrid-mamba-transformer-moe-for-agentic-reasoning/)