@seanhogg/builderforce-memory-engine 2026.6.18

package/src/kernels/activations.ts

@@ -0,0 +1,174 @@
+// Activation function WGSL kernels: SiLU (Swish) and its backward pass.
+// Used in the gating mechanism of the Mamba Mixer Block.
+
+export const ACTIVATIONS_WGSL: string = /* 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];
+    }
+}
+`;
+
+// ---- Softmax (row-wise with optional causal mask) ----
+// Standalone softmax used by AttentionBlock for the score matrix.
+// Dispatch: (L, H, B) — one workgroup per (row, head, batch).
+// This version is a simple sequential-within-workgroup implementation;
+// for large L prefer the cooperative version in attention.ts.
+export const SOFTMAX_FORWARD_WGSL: string = /* wgsl */`
+struct SoftmaxParams {
+    rows    : u32,   // L
+    cols    : u32,   // L
+    causal  : u32,   // 1 = apply causal mask, 0 = full softmax
+};
+
+@group(0) @binding(0) var<uniform>             sp   : SoftmaxParams;
+@group(0) @binding(1) var<storage, read_write> data : array<f32>;
+
+@compute @workgroup_size(1, 1, 1)
+fn softmax_forward_simple(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let row  = gid.x;
+    let head = gid.y;
+    let bat  = gid.z;
+
+    if (row >= sp.rows) { return; }
+
+    let L    = sp.cols;
+    let base = bat * sp.rows * L + head * L * L + row * L;
+    let lim  = select(L, row + 1u, sp.causal == 1u);
+
+    var max_val = -1e38;
+    for (var c = 0u; c < lim; c = c + 1u) {
+        if (data[base + c] > max_val) { max_val = data[base + c]; }
+    }
+
+    var sum_exp = 0.0;
+    for (var c = 0u; c < lim; c = c + 1u) {
+        let e = exp(data[base + c] - max_val);
+        data[base + c] = e;
+        sum_exp = sum_exp + e;
+    }
+
+    let inv = 1.0 / (sum_exp + 1e-12);
+    for (var c = 0u; c < lim; c = c + 1u) {
+        data[base + c] = data[base + c] * inv;
+    }
+    // Zero out masked positions
+    for (var c = lim; c < L; c = c + 1u) {
+        data[base + c] = 0.0;
+    }
+}
+`;
+
+export const SOFTMAX_BACKWARD_WGSL: string = /* wgsl */`
+struct SoftmaxParams {
+    rows    : u32,
+    cols    : u32,
+    causal  : u32,
+};
+
+@group(0) @binding(0) var<uniform>            sp  : SoftmaxParams;
+@group(0) @binding(1) var<storage, read>      p   : array<f32>;   // post-softmax probs
+@group(0) @binding(2) var<storage, read>      dp  : array<f32>;   // upstream gradient
+@group(0) @binding(3) var<storage, read_write> dx : array<f32>;   // output gradient
+
+@compute @workgroup_size(1, 1, 1)
+fn softmax_backward(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let row  = gid.x;
+    let head = gid.y;
+    let bat  = gid.z;
+
+    if (row >= sp.rows) { return; }
+
+    let L    = sp.cols;
+    let base = bat * sp.rows * L + head * L * L + row * L;
+    let lim  = select(L, row + 1u, sp.causal == 1u);
+
+    // dot = sum_i p[i] * dp[i]
+    var dot = 0.0;
+    for (var i = 0u; i < lim; i = i + 1u) {
+        dot = dot + p[base + i] * dp[base + i];
+    }
+
+    for (var i = 0u; i < lim; i = i + 1u) {
+        dx[base + i] = p[base + i] * (dp[base + i] - dot);
+    }
+}
+`;
+
+// ---- Backward for SiLU ----
+export const ACTIVATIONS_BACKWARD_WGSL: string = /* 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));
+}
+`;

package/src/kernels/attention.ts

@@ -0,0 +1,268 @@
+/**
+ * attention.ts – Causal multi-head self-attention kernels.
+ *
+ * Implements tiled 16×16 causal attention suitable for WebGPU.
+ * No Flash-Attention dependency — straightforward O(L²) with causal mask.
+ *
+ * Buffer layout:
+ *   qkv_in  : [B, L, 3*D_model]   fused Q,K,V after wQKV projection
+ *   out_buf : [B, L, D_model]
+ *   scores  : [B, H, L, L]        intermediate (written then read by kernel)
+ *
+ * Dispatch attention_forward:  (ceil(L/16), H, B)
+ * Dispatch softmax_forward:    (L, H, B)            — one workgroup per row
+ * Dispatch attention_backward: (ceil(L/16), H, B)
+ */
+
+// ── Softmax ───────────────────────────────────────────────────────────────────
+
+export const SOFTMAX_WGSL: string = /* wgsl */`
+struct SoftmaxParams {
+    rows : u32,   // L
+    cols : u32,   // L (score matrix is L×L per head)
+};
+
+@group(0) @binding(0) var<uniform>             params : SoftmaxParams;
+@group(0) @binding(1) var<storage, read_write> data   : array<f32>;
+
+// One workgroup per row; each invocation handles one element within the row.
+// Workgroup size 64 – cooperative reduction for max and sum.
+var<workgroup> wg_max : array<f32, 64>;
+var<workgroup> wg_sum : array<f32, 64>;
+
+@compute @workgroup_size(64, 1, 1)
+fn softmax_forward(@builtin(global_invocation_id) gid: vec3<u32>,
+                   @builtin(local_invocation_id)  lid: vec3<u32>,
+                   @builtin(workgroup_id)          wid: vec3<u32>) {
+    let row  = wid.x;   // L row index
+    let head = wid.y;
+    let bat  = wid.z;
+    let cols = params.cols;
+
+    if (row >= params.rows) { return; }
+
+    let base = (bat * params.rows * cols * /* nHeads from outer dispatch */ 1u)
+             + row * cols;
+
+    // Step 1: find row max (with causal mask: positions > row are -inf)
+    var local_max = -1e38;
+    for (var c = lid.x; c < cols; c = c + 64u) {
+        var v = -1e38;
+        if (c <= row) { v = data[base + c]; }
+        if (v > local_max) { local_max = v; }
+    }
+    wg_max[lid.x] = local_max;
+    workgroupBarrier();
+    for (var s = 32u; s >= 1u; s = s >> 1u) {
+        if (lid.x < s) {
+            if (wg_max[lid.x + s] > wg_max[lid.x]) {
+                wg_max[lid.x] = wg_max[lid.x + s];
+            }
+        }
+        workgroupBarrier();
+    }
+    let row_max = wg_max[0u];
+
+    // Step 2: exp and sum
+    var local_sum = 0.0;
+    for (var c = lid.x; c < cols; c = c + 64u) {
+        if (c <= row) {
+            let e = exp(data[base + c] - row_max);
+            data[base + c] = e;
+            local_sum = local_sum + e;
+        } else {
+            data[base + c] = 0.0;
+        }
+    }
+    wg_sum[lid.x] = local_sum;
+    workgroupBarrier();
+    for (var s = 32u; s >= 1u; s = s >> 1u) {
+        if (lid.x < s) { wg_sum[lid.x] = wg_sum[lid.x] + wg_sum[lid.x + s]; }
+        workgroupBarrier();
+    }
+    let inv_sum = 1.0 / (wg_sum[0u] + 1e-12);
+
+    // Step 3: normalise
+    for (var c = lid.x; c <= row; c = c + 64u) {
+        data[base + c] = data[base + c] * inv_sum;
+    }
+}
+`;
+
+// ── Attention forward ─────────────────────────────────────────────────────────
+
+export const ATTENTION_FORWARD_WGSL: string = /* wgsl */`
+struct AttnParams {
+    batch    : u32,
+    seq_len  : u32,
+    d_model  : u32,
+    n_heads  : u32,
+    d_head   : u32,
+};
+
+@group(0) @binding(0) var<uniform>             params  : AttnParams;
+// Q, K, V packed: [B, L, 3, H, d_head]  (after projection split)
+@group(0) @binding(1) var<storage, read>       Q       : array<f32>; // [B,L,H,dh]
+@group(0) @binding(2) var<storage, read>       K       : array<f32>; // [B,L,H,dh]
+@group(0) @binding(3) var<storage, read>       V       : array<f32>; // [B,L,H,dh]
+@group(0) @binding(4) var<storage, read_write> scores  : array<f32>; // [B,H,L,L]
+@group(0) @binding(5) var<storage, read_write> out_buf : array<f32>; // [B,L,H,dh]
+
+// Tiled 16×16 shared memory for Q row and K col
+var<workgroup> tile_q : array<f32, 256>;  // 16 tokens × 16 d_head
+var<workgroup> tile_k : array<f32, 256>;
+
+@compute @workgroup_size(16, 16, 1)
+fn attention_forward(@builtin(global_invocation_id) gid: vec3<u32>,
+                     @builtin(local_invocation_id)  lid: vec3<u32>,
+                     @builtin(workgroup_id)          wid: vec3<u32>) {
+    let q_tile = wid.x;     // tile index along query (row) dimension
+    let head   = wid.y;
+    let batch  = wid.z;
+
+    let B  = params.batch;
+    let L  = params.seq_len;
+    let H  = params.n_heads;
+    let dh = params.d_head;
+    let inv_sqrt = 1.0 / sqrt(f32(dh));
+
+    let row = q_tile * 16u + lid.x;   // query token index
+    let col = lid.y;                   // key token index offset within tile
+
+    if (row >= L) { return; }
+
+    // ── Phase 1: Compute raw attention scores for all K positions ──────────
+    // scores[batch, head, row, k] = Q[row] · K[k] / sqrt(dh)
+    // We iterate over K tiles
+    let q_base = batch * L * H * dh + row * H * dh + head * dh;
+
+    for (var k_start: u32 = 0u; k_start <= row; k_start = k_start + 16u) {
+        let k_tok = k_start + lid.y;
+
+        // Load Q row tile into shared memory (lid.y = 0..15 element index)
+        if (lid.y < dh && lid.y < 16u) {
+            tile_q[lid.x * 16u + lid.y] = Q[q_base + lid.y];
+        }
+        // Load K col tile
+        if (k_tok < L && lid.x < dh && lid.x < 16u) {
+            let k_base = batch * L * H * dh + k_tok * H * dh + head * dh;
+            tile_k[lid.y * 16u + lid.x] = K[k_base + lid.x];
+        } else if (lid.x < 16u) {
+            tile_k[lid.y * 16u + lid.x] = 0.0;
+        }
+        workgroupBarrier();
+
+        // Dot product: accumulate over dh
+        if (k_tok <= row) {
+            var acc = 0.0;
+            for (var d = 0u; d < min(dh, 16u); d = d + 1u) {
+                acc = acc + tile_q[lid.x * 16u + d] * tile_k[lid.y * 16u + d];
+            }
+            let score_idx = batch * H * L * L + head * L * L + row * L + k_tok;
+            scores[score_idx] = acc * inv_sqrt;
+        }
+        workgroupBarrier();
+    }
+}
+
+// Phase 2: softmax is dispatched separately via softmax_forward kernel.
+
+// Phase 3: weighted sum of V
+@compute @workgroup_size(16, 16, 1)
+fn attention_value(@builtin(global_invocation_id) gid: vec3<u32>,
+                   @builtin(local_invocation_id)  lid: vec3<u32>,
+                   @builtin(workgroup_id)          wid: vec3<u32>) {
+    let q_tile = wid.x;
+    let head   = wid.y;
+    let batch  = wid.z;
+
+    let L  = params.seq_len;
+    let H  = params.n_heads;
+    let dh = params.d_head;
+
+    let row = q_tile * 16u + lid.x;
+    let d   = lid.y;   // d_head dimension
+
+    if (row >= L || d >= dh) { return; }
+
+    var acc = 0.0;
+    for (var k: u32 = 0u; k <= row; k = k + 1u) {
+        let score_idx = batch * H * L * L + head * L * L + row * L + k;
+        let v_idx     = batch * L * H * dh + k * H * dh + head * dh + d;
+        acc = acc + scores[score_idx] * V[v_idx];
+    }
+
+    let out_idx = batch * L * H * dh + row * H * dh + head * dh + d;
+    out_buf[out_idx] = acc;
+}
+`;
+
+// ── Attention backward ────────────────────────────────────────────────────────
+
+export const ATTENTION_BACKWARD_WGSL: string = /* wgsl */`
+struct AttnParams {
+    batch    : u32,
+    seq_len  : u32,
+    d_model  : u32,
+    n_heads  : u32,
+    d_head   : u32,
+};
+
+@group(0) @binding(0) var<uniform>             params    : AttnParams;
+@group(0) @binding(1) var<storage, read>       Q         : array<f32>;
+@group(0) @binding(2) var<storage, read>       K         : array<f32>;
+@group(0) @binding(3) var<storage, read>       V         : array<f32>;
+@group(0) @binding(4) var<storage, read>       scores    : array<f32>; // post-softmax
+@group(0) @binding(5) var<storage, read>       dy        : array<f32>; // [B,L,H,dh]
+@group(0) @binding(6) var<storage, read_write> dQ        : array<f32>;
+@group(0) @binding(7) var<storage, read_write> dK        : array<f32>;
+@group(0) @binding(8) var<storage, read_write> dV        : array<f32>;
+@group(0) @binding(9) var<storage, read_write> dscores   : array<f32>;
+
+@compute @workgroup_size(16, 16, 1)
+fn attention_backward(@builtin(global_invocation_id) gid: vec3<u32>,
+                      @builtin(local_invocation_id)  lid: vec3<u32>,
+                      @builtin(workgroup_id)          wid: vec3<u32>) {
+    let q_tile = wid.x;
+    let head   = wid.y;
+    let batch  = wid.z;
+
+    let L  = params.seq_len;
+    let H  = params.n_heads;
+    let dh = params.d_head;
+    let inv_sqrt = 1.0 / sqrt(f32(dh));
+
+    let row = q_tile * 16u + lid.x;
+    let d   = lid.y;
+
+    if (row >= L || d >= dh) { return; }
+
+    // dV[k, d] += score[row, k] * dy[row, d]
+    // dscores[row, k] += dy[row, d] * V[k, d]  (before softmax backward)
+    for (var k: u32 = 0u; k <= row; k = k + 1u) {
+        let s_idx = batch * H * L * L + head * L * L + row * L + k;
+        let v_idx = batch * L * H * dh + k * H * dh + head * dh + d;
+        let dy_idx = batch * L * H * dh + row * H * dh + head * dh + d;
+
+        dV[v_idx] = dV[v_idx] + scores[s_idx] * dy[dy_idx];
+        dscores[s_idx] = dscores[s_idx] + dy[dy_idx] * V[v_idx];
+    }
+
+    // dQ[row, d] += sum_k dscores_post_softmax[row, k] * K[k, d] * inv_sqrt
+    var dq_acc = 0.0;
+    for (var k: u32 = 0u; k <= row; k = k + 1u) {
+        let ds_idx = batch * H * L * L + head * L * L + row * L + k;
+        let k_idx  = batch * L * H * dh + k * H * dh + head * dh + d;
+        dq_acc = dq_acc + dscores[ds_idx] * K[k_idx];
+    }
+    let q_idx = batch * L * H * dh + row * H * dh + head * dh + d;
+    dQ[q_idx] = dQ[q_idx] + dq_acc * inv_sqrt;
+
+    // dK[k, d] += dscores[row, k] * Q[row, d] * inv_sqrt  (for all rows >= k)
+    for (var k: u32 = 0u; k <= row; k = k + 1u) {
+        let ds_idx = batch * H * L * L + head * L * L + row * L + k;
+        let k_idx  = batch * L * H * dh + k * H * dh + head * dh + d;
+        dK[k_idx] = dK[k_idx] + dscores[ds_idx] * Q[q_idx] * inv_sqrt;
+    }
+}
+`;