@tryhamster/gerbil 1.0.0-rc.9 → 1.0.0

package/docs/gpu-engine/ir.md

@@ -0,0 +1,372 @@
+# IR Deep Dive
+
+The Intermediate Representation (IR) defined in `src/gpu/ir.ts` is the central contract of the GPU engine. Every component -- graph generators, the executor, the kernel registry, the model loader -- speaks this IR.
+
+This document covers how the IR is designed, how to read a `ModelGraph`, how the executor interprets it, and how to add new operation types.
+
+---
+
+## Type Overview
+
+```
+ModelGraph
+  ├── architecture: string          "Qwen2ForCausalLM"
+  ├── config: ModelArchConfig        Resolved model dimensions
+  ├── capabilities: ModelCapabilities  { text, vision, moe }
+  ├── tensors: Record<string, TensorDesc>
+  │     ├── "embed_tokens.weight"    { shape: [vocab, hidden], dtype: "f32", storage: "constant" }
+  │     ├── "embed_out"              { shape: ["T", hidden], dtype: "f32", storage: "activation" }
+  │     ├── "layer0_k_cache"         { shape: ["L_max", kv_dim], dtype: "f32", storage: "kv_cache" }
+  │     └── ...
+  ├── nodes: OpNode[]
+  │     ├── { id: "embed", opType: "Embedding", inputs: [...], outputs: [...] }
+  │     ├── { id: "layer0_norm1", opType: "RMSNorm", inputs: [...], outputs: [...] }
+  │     └── ...
+  ├── executionOrder: string[]       ["embed", "layer0_norm1", "layer0_q_proj", ...]
+  ├── inputs: string[]               ["input_ids"]
+  └── outputs: string[]              ["logits"]
+```
+
+---
+
+## TensorDesc In Detail
+
+```typescript
+interface TensorDesc {
+  name: string;
+  shape: (number | string)[];
+  dtype: DType;
+  storage: TensorStorage;
+  safetensorsKey?: string;
+}
+```
+
+### Shape Dimensions
+
+Shapes can contain concrete numbers and symbolic strings:
+
+| Symbol | Meaning | Resolved When |
+|--------|---------|---------------|
+| `"T"` | Current sequence length | Each forward pass (prefill: prompt length; decode: 1) |
+| `"L_max"` | Total cached sequence length | Each forward pass (seqPos + T) |
+
+Concrete examples:
+- `[151936, 896]` -- embedding weight for Qwen3.5-0.8B (vocab_size x hidden_size)
+- `["T", 896]` -- activation after embedding (seq_len x hidden_size)
+- `["L_max", 256]` -- KV cache for one layer (total_seq x kv_dim)
+
+### Storage Types
+
+| Storage | Lifetime | Allocation | Source |
+|---------|----------|------------|--------|
+| `"constant"` | Entire model session | Exact size from safetensors | Downloaded weights |
+| `"activation"` | Overwritten each forward pass | Max size (T = maxSeqLen) | Computed during forward pass |
+| `"kv_cache"` | Grows during generation | Max size (L_max = maxSeqLen) | Written by attention kernel |
+
+### DType
+
+| DType | Bytes per Element | TypedArray | Notes |
+|-------|------------------|------------|-------|
+| `"f32"` | 4 | `Float32Array` | Default for all computation |
+| `"f16"` | 2 | `Uint16Array` (bitwise) | No native JS typed array |
+| `"i32"` | 4 | `Int32Array` | |
+| `"u32"` | 4 | `Uint32Array` | Used for input_ids |
+| `"i4"` | 0.5 | Packed in `Uint32Array` | 8 values per u32 for INT4 weights |
+
+---
+
+## OpNode In Detail
+
+```typescript
+interface OpNode {
+  id: string;
+  opType: OpType;
+  inputs: string[];
+  outputs: string[];
+  attributes: Record<string, unknown>;
+}
+```
+
+### ID Convention
+
+Node IDs follow a systematic pattern:
+- `"embed"` -- global embedding node
+- `"layer{i}_norm1"` -- input layernorm for layer i
+- `"layer{i}_q_proj"` -- Q projection matmul for layer i
+- `"layer{i}_rope"` -- RoPE for layer i
+- `"layer{i}_attn"` -- attention for layer i
+- `"layer{i}_resid1"` -- first residual add for layer i
+- `"layer{i}_silu"` -- SiLU activation for layer i
+- `"layer{i}_swiglu"` -- gate * up multiply for layer i
+- `"final_norm"` -- final layer norm
+- `"lm_head"` -- language model head matmul
+
+### Input/Output Ordering
+
+The `inputs` and `outputs` arrays define which tensors the kernel reads and writes. **Order matters**: it matches the kernel's binding order.
+
+For example, an RMSNorm node:
+```typescript
+{
+  id: "layer0_norm1",
+  opType: "RMSNorm",
+  inputs: ["embed_out", "layers.0.input_layernorm.weight"],
+  outputs: ["layer0_norm1_out"],
+  attributes: { hidden_size: 896, eps: 1e-6, seq_len_tensor: "embed_out" }
+}
+```
+
+This maps to the RMSNorm WGSL bindings:
+- `@binding(0)` input -> `embed_out` (from `inputs[0]`)
+- `@binding(1)` weight -> `layers.0.input_layernorm.weight` (from `inputs[1]`)
+- `@binding(2)` output -> `layer0_norm1_out` (from `outputs[0]`)
+- `@binding(3)` params -> built from `attributes`
+
+### Attributes
+
+The `attributes` dictionary carries op-specific parameters. Common patterns:
+
+| Attribute | Used By | Purpose |
+|-----------|---------|---------|
+| `hidden_size` | RMSNorm, Add | Dimension of the feature vector |
+| `eps` | RMSNorm, LayerNorm | Numerical stability epsilon |
+| `M_tensor` | MatMul | Name of tensor whose shape gives the M dimension |
+| `K` | MatMul | Inner dimension |
+| `N` | MatMul | Output dimension |
+| `head_dim` | RoPE, Attention | Per-head dimension |
+| `num_heads` | RoPE, Attention | Number of query heads |
+| `num_kv_heads` | RoPE, Attention | Number of key/value heads (GQA) |
+| `rope_base` | RoPE | Base frequency for position encoding |
+| `causal` | Attention | Whether to apply causal mask |
+| `layer_index` | Attention | Which layer (for KV cache buffer lookup) |
+| `count_tensor` | SiLU, Mul, Add | Name of tensor whose element count determines dispatch size |
+
+---
+
+## Concrete Qwen2 Layer Example
+
+For a Qwen3.5-0.8B model (hidden_size=896, num_heads=14, num_kv_heads=2, head_dim=64, intermediate_size=4864), here is layer 0 in full:
+
+### Tensors for Layer 0
+
+```
+Weight tensors (constant):
+  layers.0.input_layernorm.weight         [896]           f32
+  layers.0.self_attn.q_proj.weight        [896, 896]      f32
+  layers.0.self_attn.k_proj.weight        [896, 128]      f32   (2 heads * 64 dim)
+  layers.0.self_attn.v_proj.weight        [896, 128]      f32
+  layers.0.self_attn.o_proj.weight        [896, 896]      f32
+  layers.0.post_attention_layernorm.weight [896]           f32
+  layers.0.mlp.gate_proj.weight           [896, 4864]     f32
+  layers.0.mlp.up_proj.weight             [896, 4864]     f32
+  layers.0.mlp.down_proj.weight           [4864, 896]     f32
+
+Activation tensors:
+  layer0_norm1_out                         [T, 896]        f32
+  layer0_q                                 [T, 896]        f32
+  layer0_k                                 [T, 128]        f32
+  layer0_v                                 [T, 128]        f32
+  layer0_q_rope                            [T, 896]        f32
+  layer0_k_rope                            [T, 128]        f32
+  layer0_attn_out                          [T, 896]        f32
+  layer0_o_proj_out                        [T, 896]        f32
+  layer0_resid1                            [T, 896]        f32
+  layer0_norm2_out                         [T, 896]        f32
+  layer0_gate_out                          [T, 4864]       f32
+  layer0_up_out                            [T, 4864]       f32
+  layer0_silu_out                          [T, 4864]       f32
+  layer0_swiglu_out                        [T, 4864]       f32
+  layer0_mlp_out                           [T, 896]        f32
+  layer0_resid2                            [T, 896]        f32
+
+KV cache tensors:
+  layer0_k_cache                           [L_max, 128]    f32
+  layer0_v_cache                           [L_max, 128]    f32
+```
+
+### Nodes for Layer 0
+
+```
+1.  layer0_norm1     RMSNorm     [embed_out, layers.0.input_layernorm.weight]
+                                 -> [layer0_norm1_out]
+
+2.  layer0_q_proj    MatMul      [layer0_norm1_out, layers.0.self_attn.q_proj.weight]
+                                 -> [layer0_q]       {K:896, N:896}
+
+3.  layer0_k_proj    MatMul      [layer0_norm1_out, layers.0.self_attn.k_proj.weight]
+                                 -> [layer0_k]       {K:896, N:128}
+
+4.  layer0_v_proj    MatMul      [layer0_norm1_out, layers.0.self_attn.v_proj.weight]
+                                 -> [layer0_v]       {K:896, N:128}
+
+5.  layer0_rope      RoPE        [layer0_q, layer0_k]
+                                 -> [layer0_q_rope, layer0_k_rope]
+                                 {head_dim:64, num_heads:14, num_kv_heads:2}
+
+6.  layer0_attn      Attention   [layer0_q_rope, layer0_k_rope, layer0_v,
+                                  layer0_k_cache, layer0_v_cache]
+                                 -> [layer0_attn_out]
+                                 {num_heads:14, num_kv_heads:2, head_dim:64, causal:true}
+
+7.  layer0_o_proj    MatMul      [layer0_attn_out, layers.0.self_attn.o_proj.weight]
+                                 -> [layer0_o_proj_out]  {K:896, N:896}
+
+8.  layer0_resid1    Add         [embed_out, layer0_o_proj_out]
+                                 -> [layer0_resid1]
+
+9.  layer0_norm2     RMSNorm     [layer0_resid1, layers.0.post_attention_layernorm.weight]
+                                 -> [layer0_norm2_out]
+
+10. layer0_gate      MatMul      [layer0_norm2_out, layers.0.mlp.gate_proj.weight]
+                                 -> [layer0_gate_out]    {K:896, N:4864}
+
+11. layer0_up        MatMul      [layer0_norm2_out, layers.0.mlp.up_proj.weight]
+                                 -> [layer0_up_out]      {K:896, N:4864}
+
+12. layer0_silu      SiLU        [layer0_gate_out]
+                                 -> [layer0_silu_out]
+
+13. layer0_swiglu    Mul         [layer0_silu_out, layer0_up_out]
+                                 -> [layer0_swiglu_out]
+
+14. layer0_down      MatMul      [layer0_swiglu_out, layers.0.mlp.down_proj.weight]
+                                 -> [layer0_mlp_out]     {K:4864, N:896}
+
+15. layer0_resid2    Add         [layer0_resid1, layer0_mlp_out]
+                                 -> [layer0_resid2]
+```
+
+Layer 1 takes `layer0_resid2` as its input (the `prevOutput` variable in the generator) and produces `layer1_resid2`, and so on through all 24 layers.
+
+---
+
+## How the Executor Interprets the IR
+
+The executor's `forward()` method iterates through `graph.executionOrder` -- the topologically sorted list of node IDs:
+
+```typescript
+for (const nodeId of this.graph.executionOrder) {
+  const node = this.graph.nodes.find(n => n.id === nodeId)!;
+  this.dispatchOp(pass, node, inputIdsBuffer, resolvedShapes);
+}
+```
+
+For each node, `dispatchOp()`:
+
+1. **Looks up the kernel**: `KERNEL_REGISTRY[node.opType]` returns a `KernelSpec` with the WGSL source code, binding layout, parameter builder, and dispatch size calculator
+2. **Gets or creates the pipeline**: Shader compilation is cached by code+entryPoint
+3. **Builds uniforms**: `spec.buildParams(node, resolvedShapes)` converts the node's attributes and resolved tensor shapes into the packed `ArrayBuffer` that the kernel expects as `@group(0) @binding(N) var<uniform> params`
+4. **Gathers buffers**: For each binding in the kernel spec, finds the corresponding GPU buffer from weight buffers, activation buffers, KV cache buffers, or the input IDs buffer
+5. **Creates bind group**: Binds buffers to their numbered bindings
+6. **Dispatches**: Calls `pass.dispatchWorkgroups(wgX, wgY, wgZ)` with workgroup counts computed by `spec.getDispatchSize(node, resolvedShapes)`
+
+---
+
+## Adding a New Op Type
+
+To add a new operation (e.g., `Conv2d` for vision models):
+
+### Step 1: Add the OpType
+
+In `ir.ts`, add the new type to the `OpType` union:
+
+```typescript
+export type OpType =
+  // ... existing ops ...
+  | "Conv2d";
+```
+
+### Step 2: Write the WGSL Kernel
+
+Create `kernels/wgsl/conv2d.wgsl` with the compute shader. Follow the conventions:
+- Last binding is always the uniform params struct
+- Bindings follow the order: inputs first, then outputs, then uniforms
+- Use `@compute @workgroup_size(...)` to declare workgroup size
+
+### Step 3: Register the Kernel
+
+In the kernel registry (when implemented), add a `KernelSpec` entry:
+
+```typescript
+KERNEL_REGISTRY["Conv2d"] = {
+  shaderCode: conv2dWGSL,
+  entryPoint: "main",
+  bindings: [
+    { type: "storage-read" },   // input
+    { type: "storage-read" },   // weight
+    { type: "storage-rw" },     // output
+    { type: "uniform" },        // params
+  ],
+  buildParams: (node, shapes) => { /* pack uniform struct */ },
+  getDispatchSize: (node, shapes) => { /* compute workgroup counts */ },
+};
+```
+
+### Step 4: Use in a Graph Generator
+
+In a graph generator (e.g., a vision encoder generator), create nodes that reference the new op:
+
+```typescript
+addNode({
+  id: "vision_conv0",
+  opType: "Conv2d",
+  inputs: ["patch_embed_input", "vision_conv0_weight"],
+  outputs: ["vision_conv0_out"],
+  attributes: { in_channels: 3, out_channels: 64, kernel_size: 7, stride: 2, padding: 3 },
+});
+```
+
+The executor will automatically dispatch it using the registered kernel spec.
+
+---
+
+## ModelArchConfig Reference
+
+The `ModelArchConfig` struct captures all the architectural dimensions needed to generate a graph:
+
+```typescript
+interface ModelArchConfig {
+  // Core dimensions
+  hidden_size: number;        // e.g. 896
+  num_layers: number;         // e.g. 24
+  num_heads: number;          // e.g. 14 (query heads)
+  num_kv_heads: number;       // e.g. 2 (for GQA)
+  head_dim: number;           // e.g. 64
+  intermediate_size: number;  // e.g. 4864 (MLP hidden)
+  vocab_size: number;         // e.g. 151936
+  context_length: number;     // e.g. 32768
+
+  // Normalization
+  rms_norm_eps: number;       // e.g. 1e-6
+  norm_type: "rmsnorm" | "layernorm";
+
+  // Positional encoding
+  rope_base: number;          // e.g. 1000000.0
+  rope_dim: number;           // e.g. 64 (usually == head_dim)
+
+  // KV cache
+  kv_layout: "LHSd";         // Only LHSd is currently supported
+
+  // MoE (future)
+  is_moe: boolean;
+  num_experts?: number;
+  top_k_experts?: number;
+
+  // Vision (future)
+  has_vision_tower: boolean;
+  vision_architecture?: string;
+  vision_patch_size?: number;
+  vision_embed_dim?: number;
+}
+```
+
+These values are extracted from HuggingFace `config.json` by the graph generator. For example, the Qwen2 generator maps:
+- `config.hidden_size` -> `hidden_size`
+- `config.num_hidden_layers` -> `num_layers`
+- `config.num_attention_heads` -> `num_heads`
+- `config.num_key_value_heads` -> `num_kv_heads`
+- `config.intermediate_size` -> `intermediate_size`
+- `config.vocab_size` -> `vocab_size`
+- `config.max_position_embeddings` -> `context_length`
+- `config.rms_norm_eps` -> `rms_norm_eps`
+- `config.rope_theta` -> `rope_base`