@seanhogg/builderforce-memory-engine 2026.6.18

package/dist/kernels/activations.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"activations.js","sourceRoot":"","sources":["../../src/kernels/activations.ts"],"names":[],"mappings":"AAAA,wEAAwE;AACxE,yDAAyD;AAEzD,MAAM,CAAC,MAAM,gBAAgB,GAAW,UAAU,CAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CAyDjD,CAAC;AAEF,yDAAyD;AACzD,kEAAkE;AAClE,8DAA8D;AAC9D,uEAAuE;AACvE,8DAA8D;AAC9D,MAAM,CAAC,MAAM,oBAAoB,GAAW,UAAU,CAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CA2CrD,CAAC;AAEF,MAAM,CAAC,MAAM,qBAAqB,GAAW,UAAU,CAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CAkCtD,CAAC;AAEF,8BAA8B;AAC9B,MAAM,CAAC,MAAM,yBAAyB,GAAW,UAAU,CAAA;;;;;;;;;;;;;;;;;;;;;;;;CAwB1D,CAAC"}

package/dist/kernels/attention.d.ts

@@ -0,0 +1,19 @@
+/**
+ * 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)
+ */
+export declare const SOFTMAX_WGSL: string;
+export declare const ATTENTION_FORWARD_WGSL: string;
+export declare const ATTENTION_BACKWARD_WGSL: string;
+//# sourceMappingURL=attention.d.ts.map

package/dist/kernels/attention.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"attention.d.ts","sourceRoot":"","sources":["../../src/kernels/attention.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;GAcG;AAIH,eAAO,MAAM,YAAY,EAAE,MAuE1B,CAAC;AAIF,eAAO,MAAM,sBAAsB,EAAE,MAwGpC,CAAC;AAIF,eAAO,MAAM,uBAAuB,EAAE,MAkErC,CAAC"}

package/dist/kernels/attention.js

@@ -0,0 +1,263 @@
+/**
+ * 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 = /* 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 = /* 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 = /* 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;
+    }
+}
+`;
+//# sourceMappingURL=attention.js.map

package/dist/kernels/attention.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"attention.js","sourceRoot":"","sources":["../../src/kernels/attention.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;GAcG;AAEH,iFAAiF;AAEjF,MAAM,CAAC,MAAM,YAAY,GAAW,UAAU,CAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CAuE7C,CAAC;AAEF,iFAAiF;AAEjF,MAAM,CAAC,MAAM,sBAAsB,GAAW,UAAU,CAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CAwGvD,CAAC;AAEF,iFAAiF;AAEjF,MAAM,CAAC,MAAM,uBAAuB,GAAW,UAAU,CAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CAkExD,CAAC"}

package/dist/kernels/complex_ssd.d.ts

@@ -0,0 +1,33 @@
+/**
+ * complex_ssd.ts – Complex-valued SSD kernels for Mamba-3.
+ *
+ * Three targeted improvements over Mamba-2 SSD:
+ *
+ * 1. Complex-valued states
+ *    h ∈ ℂ^(N/2)  stored as interleaved (real, imag) f32 pairs.
+ *    A ∈ ℂ encoded as A_log[H, 2] = [log|A|, arg(A)].
+ *
+ * 2. Exponential-Trapezoidal (ET) discretisation
+ *    A_bar = exp(Δ · A)                        (complex multiply)
+ *    B_bar = (A_bar − 1) · A⁻¹ · B            (exact, complex division)
+ *
+ * 3. MIMO recurrence (G groups of G inputs/outputs per head)
+ *    Implemented here with G=1 (SISO) as the default; G>1 is a future
+ *    extension that enlarges the B/C projections.
+ *
+ * Buffer layout:
+ *   x          : [B, L, D_inner]       real-valued
+ *   B_proj     : [B, L, n_groups, N*2] interleaved complex (re,im)
+ *   C_proj     : [B, L, n_groups, N*2]
+ *   dt         : [B, L, H]             real-valued
+ *   A_log      : [H, 2]                [log|A|, arg(A)] per head
+ *   dt_bias    : [H]
+ *   D_vec      : [H]
+ *   out        : [B, L, D_inner]       real-valued (Re(C·h))
+ *   state_carry: [n_chunks+1, B, H, N*2, d_head]  complex states
+ *
+ * Dispatch: (n_chunks, H, B)
+ */
+export declare const COMPLEX_SSD_FORWARD_WGSL: string;
+export declare const COMPLEX_SSD_BACKWARD_WGSL: string;
+//# sourceMappingURL=complex_ssd.d.ts.map

package/dist/kernels/complex_ssd.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"complex_ssd.d.ts","sourceRoot":"","sources":["../../src/kernels/complex_ssd.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GA6BG;AAEH,eAAO,MAAM,wBAAwB,EAAE,MAiJtC,CAAC;AAIF,eAAO,MAAM,yBAAyB,EAAE,MA8HvC,CAAC"}

package/dist/kernels/complex_ssd.js

@@ -0,0 +1,305 @@
+/**
+ * complex_ssd.ts – Complex-valued SSD kernels for Mamba-3.
+ *
+ * Three targeted improvements over Mamba-2 SSD:
+ *
+ * 1. Complex-valued states
+ *    h ∈ ℂ^(N/2)  stored as interleaved (real, imag) f32 pairs.
+ *    A ∈ ℂ encoded as A_log[H, 2] = [log|A|, arg(A)].
+ *
+ * 2. Exponential-Trapezoidal (ET) discretisation
+ *    A_bar = exp(Δ · A)                        (complex multiply)
+ *    B_bar = (A_bar − 1) · A⁻¹ · B            (exact, complex division)
+ *
+ * 3. MIMO recurrence (G groups of G inputs/outputs per head)
+ *    Implemented here with G=1 (SISO) as the default; G>1 is a future
+ *    extension that enlarges the B/C projections.
+ *
+ * Buffer layout:
+ *   x          : [B, L, D_inner]       real-valued
+ *   B_proj     : [B, L, n_groups, N*2] interleaved complex (re,im)
+ *   C_proj     : [B, L, n_groups, N*2]
+ *   dt         : [B, L, H]             real-valued
+ *   A_log      : [H, 2]                [log|A|, arg(A)] per head
+ *   dt_bias    : [H]
+ *   D_vec      : [H]
+ *   out        : [B, L, D_inner]       real-valued (Re(C·h))
+ *   state_carry: [n_chunks+1, B, H, N*2, d_head]  complex states
+ *
+ * Dispatch: (n_chunks, H, B)
+ */
+export const COMPLEX_SSD_FORWARD_WGSL = /* wgsl */ `
+struct CssdParams {
+    seq_len    : u32,
+    d_inner    : u32,
+    n_heads    : u32,
+    d_head     : u32,
+    n_groups   : u32,
+    n_complex  : u32,   // N/2 – number of complex state components
+    chunk_len  : u32,
+    n_chunks   : u32,
+    batch      : u32,
+};
+
+@group(0) @binding(0) var<uniform>             params      : CssdParams;
+@group(0) @binding(1) var<storage, read>       x_in        : array<f32>;
+@group(0) @binding(2) var<storage, read>       B_proj      : array<f32>; // complex: N_c*2 per token
+@group(0) @binding(3) var<storage, read>       C_proj      : array<f32>;
+@group(0) @binding(4) var<storage, read>       dt_in       : array<f32>;
+@group(0) @binding(5) var<storage, read>       A_log       : array<f32>; // [H, 2]
+@group(0) @binding(6) var<storage, read>       dt_bias     : array<f32>;
+@group(0) @binding(7) var<storage, read>       D_vec       : array<f32>;
+@group(0) @binding(8) var<storage, read_write> out_buf     : array<f32>;
+@group(0) @binding(9) var<storage, read_write> state_carry : array<f32>; // complex states
+
+fn softplus(v: f32) -> f32 { return log(1.0 + exp(v)); }
+
+// Complex multiply: (ar + i·ai) * (br + i·bi)
+fn cmul_re(ar: f32, ai: f32, br: f32, bi: f32) -> f32 { return ar*br - ai*bi; }
+fn cmul_im(ar: f32, ai: f32, br: f32, bi: f32) -> f32 { return ar*bi + ai*br; }
+
+// Complex exp: exp(x + i·y) = exp(x)*(cos(y) + i*sin(y))
+fn cexp_re(x: f32, y: f32) -> f32 { return exp(x) * cos(y); }
+fn cexp_im(x: f32, y: f32) -> f32 { return exp(x) * sin(y); }
+
+// ET discretisation B_bar = (A_bar - 1) * A^-1 * B
+// A^-1 = 1/A = conj(A)/|A|^2.  Here A = exp(log_mag)*exp(i*phase).
+// |A| = exp(log_mag),  A^-1 = exp(-log_mag)*exp(-i*phase)
+// (A_bar - 1) * A^-1 = scalar complex product computed below.
+fn et_bbar_re(a_bar_re: f32, a_bar_im: f32, log_mag: f32, phase: f32) -> f32 {
+    // (A_bar - 1)
+    let num_re = a_bar_re - 1.0;
+    let num_im = a_bar_im;
+    // A^-1 = exp(-log_mag - i*phase)
+    let inv_re = cexp_re(-log_mag, -phase);
+    let inv_im = cexp_im(-log_mag, -phase);
+    return cmul_re(num_re, num_im, inv_re, inv_im);
+}
+fn et_bbar_im(a_bar_re: f32, a_bar_im: f32, log_mag: f32, phase: f32) -> f32 {
+    let num_re = a_bar_re - 1.0;
+    let num_im = a_bar_im;
+    let inv_re = cexp_re(-log_mag, -phase);
+    let inv_im = cexp_im(-log_mag, -phase);
+    return cmul_im(num_re, num_im, inv_re, inv_im);
+}
+
+@compute @workgroup_size(1, 1, 1)
+fn complex_ssd_forward(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let chunk_id = gid.x;
+    let head_id  = gid.y;
+    let batch_id = gid.z;
+
+    let L  = params.seq_len;
+    let D  = params.d_inner;
+    let H  = params.n_heads;
+    let dh = params.d_head;
+    let G  = params.n_groups;
+    let Nc = params.n_complex;   // complex state count
+    let N2 = Nc * 2u;            // float pairs
+    let CL = params.chunk_len;
+    let B  = params.batch;
+
+    let t_start  = chunk_id * CL;
+    let t_end    = min(t_start + CL, L);
+    let group_id = head_id * G / H;
+
+    // Load A for this head: A = exp(log_mag) * exp(i*phase)
+    let log_mag = A_log[head_id * 2u + 0u];
+    let phase   = A_log[head_id * 2u + 1u];
+    let db      = dt_bias[head_id];
+    let d_skip  = D_vec[head_id];
+
+    // State buffer strides (complex: N2*dh floats per head)
+    let state_stride = B * H * N2 * dh;
+    let state_base_in  = chunk_id * state_stride
+                       + batch_id * H * N2 * dh
+                       + head_id  * N2 * dh;
+    let state_base_out = (chunk_id + 1u) * state_stride
+                       + batch_id * H * N2 * dh
+                       + head_id  * N2 * dh;
+
+    // Copy carry-in to working slot
+    for (var s: u32 = 0u; s < N2 * dh; s = s + 1u) {
+        state_carry[state_base_out + s] = state_carry[state_base_in + s];
+    }
+
+    for (var t: u32 = t_start; t < t_end; t = t + 1u) {
+        let dt_idx = batch_id * L * H + t * H + head_id;
+        let dt_val = softplus(dt_in[dt_idx] + db);
+
+        // A_bar = exp(dt * A) = exp(dt*log_mag + i*dt*phase)
+        let a_bar_re = cexp_re(dt_val * log_mag, dt_val * phase);
+        let a_bar_im = cexp_im(dt_val * log_mag, dt_val * phase);
+
+        // ET B_bar scalar factor (applied per B_proj element)
+        let bbar_factor_re = et_bbar_re(a_bar_re, a_bar_im, log_mag, phase);
+        let bbar_factor_im = et_bbar_im(a_bar_re, a_bar_im, log_mag, phase);
+
+        let x_base = batch_id * L * D + t * D + head_id * dh;
+        // B_proj / C_proj: [B, L, G, N*2] — interleaved re/im
+        let bc_base = batch_id * L * G * N2 + t * G * N2 + group_id * N2;
+
+        for (var i: u32 = 0u; i < dh; i = i + 1u) {
+            let x_val   = x_in[x_base + i];
+            var y_re    = 0.0;
+
+            for (var nc: u32 = 0u; nc < Nc; nc = nc + 1u) {
+                let b_re = B_proj[bc_base + nc * 2u + 0u];
+                let b_im = B_proj[bc_base + nc * 2u + 1u];
+                let c_re = C_proj[bc_base + nc * 2u + 0u];
+                let c_im = C_proj[bc_base + nc * 2u + 1u];
+
+                // B_bar · x  (complex * real = complex scale)
+                let inp_re = cmul_re(bbar_factor_re, bbar_factor_im, b_re, b_im) * x_val;
+                let inp_im = cmul_im(bbar_factor_re, bbar_factor_im, b_re, b_im) * x_val;
+
+                let s_re_idx = state_base_out + nc * 2u * dh + 0u * dh + i;
+                let s_im_idx = state_base_out + nc * 2u * dh + 1u * dh + i;
+
+                // h_t = A_bar * h_{t-1} + B_bar * x
+                let h_prev_re = state_carry[s_re_idx];
+                let h_prev_im = state_carry[s_im_idx];
+                let h_new_re  = cmul_re(a_bar_re, a_bar_im, h_prev_re, h_prev_im) + inp_re;
+                let h_new_im  = cmul_im(a_bar_re, a_bar_im, h_prev_re, h_prev_im) + inp_im;
+                state_carry[s_re_idx] = h_new_re;
+                state_carry[s_im_idx] = h_new_im;
+
+                // y += Re(C · h)
+                y_re = y_re + cmul_re(c_re, -c_im, h_new_re, h_new_im); // C·h real part
+            }
+
+            let out_idx = batch_id * L * D + t * D + head_id * dh + i;
+            out_buf[out_idx] = y_re + d_skip * x_val;
+        }
+    }
+}
+`;
+// ── Backward ──────────────────────────────────────────────────────────────────
+export const COMPLEX_SSD_BACKWARD_WGSL = /* wgsl */ `
+struct CssdParams {
+    seq_len    : u32,
+    d_inner    : u32,
+    n_heads    : u32,
+    d_head     : u32,
+    n_groups   : u32,
+    n_complex  : u32,
+    chunk_len  : u32,
+    n_chunks   : u32,
+    batch      : u32,
+};
+
+@group(0) @binding(0) var<uniform>             params      : CssdParams;
+@group(0) @binding(1) var<storage, read>       x_in        : array<f32>;
+@group(0) @binding(2) var<storage, read>       B_proj      : array<f32>;
+@group(0) @binding(3) var<storage, read>       C_proj      : array<f32>;
+@group(0) @binding(4) var<storage, read>       dt_in       : array<f32>;
+@group(0) @binding(5) var<storage, read>       A_log       : array<f32>;
+@group(0) @binding(6) var<storage, read>       dt_bias     : array<f32>;
+@group(0) @binding(7) var<storage, read>       state_carry : array<f32>;
+@group(0) @binding(8) var<storage, read>       dy          : array<f32>;
+@group(0) @binding(9)  var<storage, read_write> dx         : array<f32>;
+@group(0) @binding(10) var<storage, read_write> dB         : array<f32>;
+@group(0) @binding(11) var<storage, read_write> dC         : array<f32>;
+@group(0) @binding(12) var<storage, read_write> ddt        : array<f32>;
+@group(0) @binding(13) var<storage, read_write> dA_log     : array<f32>;
+@group(0) @binding(14) var<storage, read_write> dD_vec     : array<f32>;
+
+fn softplus(v: f32) -> f32 { return log(1.0 + exp(v)); }
+fn d_softplus(v: f32) -> f32 { return 1.0 / (1.0 + exp(-v)); }
+fn cmul_re(ar: f32, ai: f32, br: f32, bi: f32) -> f32 { return ar*br - ai*bi; }
+fn cmul_im(ar: f32, ai: f32, br: f32, bi: f32) -> f32 { return ar*bi + ai*br; }
+fn cexp_re(x: f32, y: f32) -> f32 { return exp(x) * cos(y); }
+fn cexp_im(x: f32, y: f32) -> f32 { return exp(x) * sin(y); }
+
+@compute @workgroup_size(1, 1, 1)
+fn complex_ssd_backward(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let chunk_id = gid.x;
+    let head_id  = gid.y;
+    let batch_id = gid.z;
+
+    let L  = params.seq_len;
+    let D  = params.d_inner;
+    let H  = params.n_heads;
+    let dh = params.d_head;
+    let G  = params.n_groups;
+    let Nc = params.n_complex;
+    let N2 = Nc * 2u;
+    let CL = params.chunk_len;
+    let B  = params.batch;
+
+    let t_start  = chunk_id * CL;
+    let t_end    = min(t_start + CL, L);
+    let group_id = head_id * G / H;
+
+    let log_mag = A_log[head_id * 2u + 0u];
+    let phase   = A_log[head_id * 2u + 1u];
+    let db      = dt_bias[head_id];
+
+    let state_stride = B * H * N2 * dh;
+
+    for (var t_rev: u32 = 0u; t_rev < t_end - t_start; t_rev = t_rev + 1u) {
+        let t = t_end - 1u - t_rev;
+
+        let dt_idx  = batch_id * L * H + t * H + head_id;
+        let dt_raw  = dt_in[dt_idx] + db;
+        let dt_val  = softplus(dt_raw);
+        let a_bar_re = cexp_re(dt_val * log_mag, dt_val * phase);
+        let a_bar_im = cexp_im(dt_val * log_mag, dt_val * phase);
+
+        let x_base  = batch_id * L * D + t * D + head_id * dh;
+        let bc_base = batch_id * L * G * N2 + t * G * N2 + group_id * N2;
+        let state_base = (chunk_id + 1u) * state_stride
+                        + batch_id * H * N2 * dh
+                        + head_id * N2 * dh;
+        let state_prev = chunk_id * state_stride
+                        + batch_id * H * N2 * dh
+                        + head_id * N2 * dh;
+
+        for (var i: u32 = 0u; i < dh; i = i + 1u) {
+            let dy_val = dy[batch_id * L * D + t * D + head_id * dh + i];
+            let x_val  = x_in[x_base + i];
+
+            dD_vec[head_id] = dD_vec[head_id] + dy_val * x_val;
+            dx[x_base + i]  = dx[x_base + i]  + dy_val;
+
+            for (var nc: u32 = 0u; nc < Nc; nc = nc + 1u) {
+                let c_re = C_proj[bc_base + nc * 2u + 0u];
+                let c_im = C_proj[bc_base + nc * 2u + 1u];
+                let b_re = B_proj[bc_base + nc * 2u + 0u];
+                let b_im = B_proj[bc_base + nc * 2u + 1u];
+
+                let h_re = state_carry[state_base + nc * 2u * dh + 0u * dh + i];
+                let h_im = state_carry[state_base + nc * 2u * dh + 1u * dh + i];
+
+                // dC from Re(C · h) output — gradient of Re(C·h) w.r.t. C is Re(h)
+                dC[bc_base + nc * 2u + 0u] = dC[bc_base + nc * 2u + 0u] + dy_val * h_re;
+                dC[bc_base + nc * 2u + 1u] = dC[bc_base + nc * 2u + 1u] - dy_val * h_im;
+
+                // dh from upstream: dh_re = c_re * dy, dh_im = -c_im * dy (Re(C·h) gradient)
+                let dh_re = c_re * dy_val;
+                let dh_im = -c_im * dy_val;
+
+                // dB: B_bar · x contributed h_new; gradient flows through B_bar
+                // simplified: dB += dh * x  (ignoring complex B_bar Jacobian)
+                dB[bc_base + nc * 2u + 0u] = dB[bc_base + nc * 2u + 0u] + dh_re * x_val;
+                dB[bc_base + nc * 2u + 1u] = dB[bc_base + nc * 2u + 1u] + dh_im * x_val;
+
+                // dx += Re(B_bar* · dh) (simplified)
+                dx[x_base + i] = dx[x_base + i] + cmul_re(b_re, -b_im, dh_re, dh_im);
+
+                // ddt: from A_bar and B_bar dependence on dt
+                let h_prev_re = state_carry[state_prev + nc * 2u * dh + 0u * dh + i];
+                let h_prev_im = state_carry[state_prev + nc * 2u * dh + 1u * dh + i];
+                // dA_bar/ddt = A * A_bar
+                let da_bar_re = cmul_re(cexp_re(log_mag, phase), cexp_im(log_mag, phase), a_bar_re, a_bar_im);
+                let da_bar_im = cmul_im(cexp_re(log_mag, phase), cexp_im(log_mag, phase), a_bar_re, a_bar_im);
+                ddt[dt_idx] = ddt[dt_idx]
+                    + (cmul_re(da_bar_re, da_bar_im, h_prev_re, h_prev_im) * dh_re
+                    -  cmul_im(da_bar_re, da_bar_im, h_prev_re, h_prev_im) * dh_im)
+                    * d_softplus(dt_raw);
+            }
+        }
+    }
+}
+`;
+//# sourceMappingURL=complex_ssd.js.map

package/dist/kernels/complex_ssd.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"complex_ssd.js","sourceRoot":"","sources":["../../src/kernels/complex_ssd.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GA6BG;AAEH,MAAM,CAAC,MAAM,wBAAwB,GAAW,UAAU,CAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CAiJzD,CAAC;AAEF,iFAAiF;AAEjF,MAAM,CAAC,MAAM,yBAAyB,GAAW,UAAU,CAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CA8H1D,CAAC"}

package/dist/kernels/conv1d.d.ts

@@ -0,0 +1,3 @@
+export declare const CONV1D_FORWARD_WGSL = "\n\nstruct ConvParams {\n    seq_len     : u32,   // L\n    d_channels  : u32,   // D (number of depthwise channels in this call)\n    kernel_size : u32,   // K (typically 4)\n    batch       : u32,   // B\n    groups      : u32,   // number of channel groups (1 = standard depthwise)\n};\n\n@group(0) @binding(0) var<uniform>             params   : ConvParams;\n// x      (B, L, D) \u2013 input\n@group(0) @binding(1) var<storage, read>       x        : array<f32>;\n// weight (D, K)    \u2013 depthwise conv weights\n@group(0) @binding(2) var<storage, read>       weight   : array<f32>;\n// bias   (D,)      \u2013 optional bias (zeros if unused)\n@group(0) @binding(3) var<storage, read>       bias     : array<f32>;\n// y      (B, L, D) \u2013 output\n@group(0) @binding(4) var<storage, read_write> y        : array<f32>;\n\n// Dispatch: (ceil(L/16), ceil(D/16), B)\n@compute @workgroup_size(16, 16, 1)\nfn conv1d_forward(\n    @builtin(global_invocation_id) gid : vec3<u32>,\n) {\n    let L  = params.seq_len;\n    let D  = params.d_channels;\n    let K  = params.kernel_size;\n    let B  = params.batch;\n\n    let t  = gid.x;   // time position\n    let d  = gid.y;   // channel\n    let b  = gid.z;   // batch\n\n    if (t >= L || d >= D || b >= B) { return; }\n\n    var acc: f32 = 0.0;\n\n    // Causal: convolve over k = 0..K-1, reading position (t - k)\n    for (var k: u32 = 0u; k < K; k = k + 1u) {\n        let w_idx = d * K + k;\n        let w_val = weight[w_idx];\n\n        // t - k: use causal zero-padding for t < k\n        if (t >= k) {\n            let src = b * L * D + (t - k) * D + d;\n            acc = acc + w_val * x[src];\n        }\n        // else: zero-padding contributes 0\n    }\n\n    acc = acc + bias[d];\n\n    let out = b * L * D + t * D + d;\n    y[out] = acc;\n}\n";
+export declare const CONV1D_BACKWARD_WGSL = "\n\nstruct ConvParams {\n    seq_len     : u32,\n    d_channels  : u32,\n    kernel_size : u32,\n    batch       : u32,\n};\n\n@group(0) @binding(0) var<uniform>              params   : ConvParams;\n@group(0) @binding(1) var<storage, read>        x        : array<f32>;\n@group(0) @binding(2) var<storage, read>        weight   : array<f32>;\n@group(0) @binding(3) var<storage, read>        dy       : array<f32>;\n@group(0) @binding(4) var<storage, read_write>  dx       : array<f32>;\n@group(0) @binding(5) var<storage, read_write>  dweight  : array<f32>;\n@group(0) @binding(6) var<storage, read_write>  dbias    : array<f32>;\n\n// Dispatch: (ceil(L/16), ceil(D/16), B) \u2013 computes dx\n@compute @workgroup_size(16, 16, 1)\nfn conv1d_backward_dx(\n    @builtin(global_invocation_id) gid : vec3<u32>,\n) {\n    let L  = params.seq_len;\n    let D  = params.d_channels;\n    let K  = params.kernel_size;\n    let B  = params.batch;\n\n    let t  = gid.x;\n    let d  = gid.y;\n    let b  = gid.z;\n\n    if (t >= L || d >= D || b >= B) { return; }\n\n    var grad: f32 = 0.0;\n\n    // dx[b, t, d] = sum_{k=0}^{K-1} dy[b, t+k, d] * weight[d, k]\n    for (var k: u32 = 0u; k < K; k = k + 1u) {\n        let tp = t + k;\n        if (tp < L) {\n            let dy_idx = b * L * D + tp * D + d;\n            let w_idx  = d * K + k;\n            grad = grad + dy[dy_idx] * weight[w_idx];\n        }\n    }\n\n    let dx_idx = b * L * D + t * D + d;\n    dx[dx_idx] = grad;\n}\n\n// Dispatch: (K, D, 1) \u2013 accumulates dweight over (B, L)\n@compute @workgroup_size(1, 1, 1)\nfn conv1d_backward_dw(\n    @builtin(global_invocation_id) gid : vec3<u32>,\n) {\n    let L  = params.seq_len;\n    let D  = params.d_channels;\n    let K  = params.kernel_size;\n    let B  = params.batch;\n\n    let k  = gid.x;\n    let d  = gid.y;\n\n    if (k >= K || d >= D) { return; }\n\n    var grad_w: f32 = 0.0;\n    var grad_b: f32 = 0.0;\n\n    for (var b: u32 = 0u; b < B; b = b + 1u) {\n        for (var t: u32 = 0u; t < L; t = t + 1u) {\n            let dy_idx = b * L * D + t * D + d;\n            let dy_val = dy[dy_idx];\n            if (t >= k) {\n                let x_idx = b * L * D + (t - k) * D + d;\n                grad_w = grad_w + dy_val * x[x_idx];\n            }\n            if (k == 0u) {\n                grad_b = grad_b + dy_val;\n            }\n        }\n    }\n\n    dweight[d * K + k] = grad_w;\n    if (k == 0u) {\n        dbias[d] = grad_b;\n    }\n}\n";
+//# sourceMappingURL=conv1d.d.ts.map

package/dist/kernels/conv1d.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"conv1d.d.ts","sourceRoot":"","sources":["../../src/kernels/conv1d.ts"],"names":[],"mappings":"AAaA,eAAO,MAAM,mBAAmB,mxDAwD/B,CAAC;AAGF,eAAO,MAAM,oBAAoB,y6EAsFhC,CAAC"}