@seanhogg/builderforce-memory-engine 2026.6.18

package/src/kernels/complex_ssd.ts

@@ -0,0 +1,307 @@
+/**
+ * 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: string = /* 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: string = /* 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);
+            }
+        }
+    }
+}
+`;

package/src/kernels/conv1d.ts

@@ -0,0 +1,159 @@
+// 1D Causal Convolution WGSL Kernel
+// Implements a depthwise 1D causal convolution over the sequence dimension.
+// "Causal" means the output at position t only depends on positions <= t,
+// which is enforced by left-padding with (kernel_size - 1) zeros.
+//
+// Forward:  y[b, t, d] = sum_{k=0}^{K-1} weight[d, k] * x[b, t-k, d]
+//           where x[b, t', d] = 0 for t' < 0  (causal padding)
+//
+// The `groups` uniform is included for Mamba-2/3 compatibility where conv runs
+// over a fused (x, B, C) buffer.  For standard Mamba-1 usage set groups = 1.
+// The kernel math is identical; the field is reserved for future grouped
+// weight-sharing variants.
+
+export const CONV1D_FORWARD_WGSL = /* wgsl */`
+
+struct ConvParams {
+    seq_len     : u32,   // L
+    d_channels  : u32,   // D (number of depthwise channels in this call)
+    kernel_size : u32,   // K (typically 4)
+    batch       : u32,   // B
+    groups      : u32,   // number of channel groups (1 = standard depthwise)
+};
+
+@group(0) @binding(0) var<uniform>             params   : ConvParams;
+// x      (B, L, D) – input
+@group(0) @binding(1) var<storage, read>       x        : array<f32>;
+// weight (D, K)    – depthwise conv weights
+@group(0) @binding(2) var<storage, read>       weight   : array<f32>;
+// bias   (D,)      – optional bias (zeros if unused)
+@group(0) @binding(3) var<storage, read>       bias     : array<f32>;
+// y      (B, L, D) – output
+@group(0) @binding(4) var<storage, read_write> y        : array<f32>;
+
+// Dispatch: (ceil(L/16), ceil(D/16), B)
+@compute @workgroup_size(16, 16, 1)
+fn conv1d_forward(
+    @builtin(global_invocation_id) gid : vec3<u32>,
+) {
+    let L  = params.seq_len;
+    let D  = params.d_channels;
+    let K  = params.kernel_size;
+    let B  = params.batch;
+
+    let t  = gid.x;   // time position
+    let d  = gid.y;   // channel
+    let b  = gid.z;   // batch
+
+    if (t >= L || d >= D || b >= B) { return; }
+
+    var acc: f32 = 0.0;
+
+    // Causal: convolve over k = 0..K-1, reading position (t - k)
+    for (var k: u32 = 0u; k < K; k = k + 1u) {
+        let w_idx = d * K + k;
+        let w_val = weight[w_idx];
+
+        // t - k: use causal zero-padding for t < k
+        if (t >= k) {
+            let src = b * L * D + (t - k) * D + d;
+            acc = acc + w_val * x[src];
+        }
+        // else: zero-padding contributes 0
+    }
+
+    acc = acc + bias[d];
+
+    let out = b * L * D + t * D + d;
+    y[out] = acc;
+}
+`;
+
+// ---- Backward kernel for 1D convolution ----
+export const CONV1D_BACKWARD_WGSL = /* wgsl */`
+
+struct ConvParams {
+    seq_len     : u32,
+    d_channels  : u32,
+    kernel_size : u32,
+    batch       : u32,
+};
+
+@group(0) @binding(0) var<uniform>              params   : ConvParams;
+@group(0) @binding(1) var<storage, read>        x        : array<f32>;
+@group(0) @binding(2) var<storage, read>        weight   : array<f32>;
+@group(0) @binding(3) var<storage, read>        dy       : array<f32>;
+@group(0) @binding(4) var<storage, read_write>  dx       : array<f32>;
+@group(0) @binding(5) var<storage, read_write>  dweight  : array<f32>;
+@group(0) @binding(6) var<storage, read_write>  dbias    : array<f32>;
+
+// Dispatch: (ceil(L/16), ceil(D/16), B) – computes dx
+@compute @workgroup_size(16, 16, 1)
+fn conv1d_backward_dx(
+    @builtin(global_invocation_id) gid : vec3<u32>,
+) {
+    let L  = params.seq_len;
+    let D  = params.d_channels;
+    let K  = params.kernel_size;
+    let B  = params.batch;
+
+    let t  = gid.x;
+    let d  = gid.y;
+    let b  = gid.z;
+
+    if (t >= L || d >= D || b >= B) { return; }
+
+    var grad: f32 = 0.0;
+
+    // dx[b, t, d] = sum_{k=0}^{K-1} dy[b, t+k, d] * weight[d, k]
+    for (var k: u32 = 0u; k < K; k = k + 1u) {
+        let tp = t + k;
+        if (tp < L) {
+            let dy_idx = b * L * D + tp * D + d;
+            let w_idx  = d * K + k;
+            grad = grad + dy[dy_idx] * weight[w_idx];
+        }
+    }
+
+    let dx_idx = b * L * D + t * D + d;
+    dx[dx_idx] = grad;
+}
+
+// Dispatch: (K, D, 1) – accumulates dweight over (B, L)
+@compute @workgroup_size(1, 1, 1)
+fn conv1d_backward_dw(
+    @builtin(global_invocation_id) gid : vec3<u32>,
+) {
+    let L  = params.seq_len;
+    let D  = params.d_channels;
+    let K  = params.kernel_size;
+    let B  = params.batch;
+
+    let k  = gid.x;
+    let d  = gid.y;
+
+    if (k >= K || d >= D) { return; }
+
+    var grad_w: f32 = 0.0;
+    var grad_b: f32 = 0.0;
+
+    for (var b: u32 = 0u; b < B; b = b + 1u) {
+        for (var t: u32 = 0u; t < L; t = t + 1u) {
+            let dy_idx = b * L * D + t * D + d;
+            let dy_val = dy[dy_idx];
+            if (t >= k) {
+                let x_idx = b * L * D + (t - k) * D + d;
+                grad_w = grad_w + dy_val * x[x_idx];
+            }
+            if (k == 0u) {
+                grad_b = grad_b + dy_val;
+            }
+        }
+    }
+
+    dweight[d * K + k] = grad_w;
+    if (k == 0u) {
+        dbias[d] = grad_b;
+    }
+}
+`;

package/src/kernels/linear_projection.ts

@@ -0,0 +1,220 @@
+// Linear Projection WGSL Kernel
+// General-purpose matrix multiplication: Y = X @ W^T + b
+// Supports the up-projection and down-projection linear layers in the Mamba block.
+//
+// Shapes:
+//   X : (batch * seq_len, in_features)   – input (rows)
+//   W : (out_features, in_features)      – weight matrix (row-major)
+//   b : (out_features,)                  – bias
+//   Y : (batch * seq_len, out_features)  – output
+
+export const LINEAR_FORWARD_WGSL: string = /* wgsl */`
+
+struct LinearParams {
+    M : u32,   // number of rows    (batch * seq_len)
+    K : u32,   // in_features
+    N : u32,   // out_features
+};
+
+@group(0) @binding(0) var<uniform>             params : LinearParams;
+@group(0) @binding(1) var<storage, read>       X      : array<f32>;   // (M, K)
+@group(0) @binding(2) var<storage, read>       W      : array<f32>;   // (N, K)
+@group(0) @binding(3) var<storage, read>       bias   : array<f32>;   // (N,)
+@group(0) @binding(4) var<storage, read_write> Y      : array<f32>;   // (M, N)
+
+// Tiled matmul using workgroup shared memory (16x16 tiles)
+var<workgroup> tile_X : array<f32, 256>;  // 16 * 16
+var<workgroup> tile_W : array<f32, 256>;
+
+@compute @workgroup_size(16, 16, 1)
+fn linear_forward(
+    @builtin(global_invocation_id)   gid : vec3<u32>,
+    @builtin(local_invocation_id)    lid : vec3<u32>,
+    @builtin(workgroup_id)           wid : vec3<u32>,
+) {
+    let M = params.M;
+    let K = params.K;
+    let N = params.N;
+
+    let row = gid.x;   // output row (M dimension)
+    let col = gid.y;   // output col (N dimension)
+
+    var acc: f32 = 0.0;
+    let TILE: u32 = 16u;
+    let num_tiles = (K + TILE - 1u) / TILE;
+
+    for (var tile_idx: u32 = 0u; tile_idx < num_tiles; tile_idx = tile_idx + 1u) {
+        // Load X tile: shape (TILE_M, TILE_K)
+        let x_col = tile_idx * TILE + lid.y;
+        let x_row = wid.x * TILE + lid.x;
+        if (x_row < M && x_col < K) {
+            tile_X[lid.x * TILE + lid.y] = X[x_row * K + x_col];
+        } else {
+            tile_X[lid.x * TILE + lid.y] = 0.0;
+        }
+
+        // Load W tile: shape (TILE_N, TILE_K)  — W is (N, K)
+        let w_col = tile_idx * TILE + lid.x;  // K dimension
+        let w_row = wid.y * TILE + lid.y;     // N dimension
+        if (w_row < N && w_col < K) {
+            tile_W[lid.y * TILE + lid.x] = W[w_row * K + w_col];
+        } else {
+            tile_W[lid.y * TILE + lid.x] = 0.0;
+        }
+
+        workgroupBarrier();
+
+        // Dot product within tile
+        for (var k: u32 = 0u; k < TILE; k = k + 1u) {
+            acc = acc + tile_X[lid.x * TILE + k] * tile_W[lid.y * TILE + k];
+        }
+        workgroupBarrier();
+    }
+
+    if (row < M && col < N) {
+        Y[row * N + col] = acc + bias[col];
+    }
+}
+`;
+
+// ---- Backward pass for linear projection ----
+export const LINEAR_BACKWARD_WGSL: string = /* wgsl */`
+
+struct LinearParams {
+    M : u32,
+    K : u32,
+    N : u32,
+};
+
+@group(0) @binding(0) var<uniform>             params : LinearParams;
+@group(0) @binding(1) var<storage, read>       X      : array<f32>;   // (M, K)
+@group(0) @binding(2) var<storage, read>       W      : array<f32>;   // (N, K)
+@group(0) @binding(3) var<storage, read>       dY     : array<f32>;   // (M, N)
+@group(0) @binding(4) var<storage, read_write> dX     : array<f32>;   // (M, K)
+@group(0) @binding(5) var<storage, read_write> dW     : array<f32>;   // (N, K)
+@group(0) @binding(6) var<storage, read_write> db     : array<f32>;   // (N,)
+
+// Dispatch: (ceil(M/16), ceil(K/16), 1)  – computes dX = dY @ W
+var<workgroup> tile_dY : array<f32, 256>;
+var<workgroup> tile_W  : array<f32, 256>;
+
+@compute @workgroup_size(16, 16, 1)
+fn linear_backward_dX(
+    @builtin(global_invocation_id) gid : vec3<u32>,
+    @builtin(local_invocation_id)  lid : vec3<u32>,
+    @builtin(workgroup_id)         wid : vec3<u32>,
+) {
+    let M = params.M;
+    let K = params.K;
+    let N = params.N;
+
+    let row = gid.x;   // M
+    let col = gid.y;   // K
+
+    var acc: f32 = 0.0;
+    let TILE: u32 = 16u;
+    let num_tiles = (N + TILE - 1u) / TILE;
+
+    for (var tile_idx: u32 = 0u; tile_idx < num_tiles; tile_idx = tile_idx + 1u) {
+        // tile_dY: (M, TILE_N) slice
+        let dy_col = tile_idx * TILE + lid.y;
+        let dy_row = wid.x * TILE + lid.x;
+        if (dy_row < M && dy_col < N) {
+            tile_dY[lid.x * TILE + lid.y] = dY[dy_row * N + dy_col];
+        } else {
+            tile_dY[lid.x * TILE + lid.y] = 0.0;
+        }
+
+        // tile_W: (TILE_N, K) slice  — W[n, k]
+        let w_row = tile_idx * TILE + lid.x;   // N
+        let w_col = wid.y * TILE + lid.y;      // K
+        if (w_row < N && w_col < K) {
+            tile_W[lid.x * TILE + lid.y] = W[w_row * K + w_col];
+        } else {
+            tile_W[lid.x * TILE + lid.y] = 0.0;
+        }
+
+        workgroupBarrier();
+
+        for (var n: u32 = 0u; n < TILE; n = n + 1u) {
+            acc = acc + tile_dY[lid.x * TILE + n] * tile_W[n * TILE + lid.y];
+        }
+        workgroupBarrier();
+    }
+
+    if (row < M && col < K) {
+        dX[row * K + col] = acc;
+    }
+}
+
+// Dispatch: (ceil(N/16), ceil(K/16), 1)  – computes dW = dY^T @ X
+var<workgroup> tile_dY2 : array<f32, 256>;
+var<workgroup> tile_X2  : array<f32, 256>;
+
+@compute @workgroup_size(16, 16, 1)
+fn linear_backward_dW(
+    @builtin(global_invocation_id) gid : vec3<u32>,
+    @builtin(local_invocation_id)  lid : vec3<u32>,
+    @builtin(workgroup_id)         wid : vec3<u32>,
+) {
+    let M = params.M;
+    let K = params.K;
+    let N = params.N;
+
+    let row = gid.x;   // N
+    let col = gid.y;   // K
+
+    var acc: f32 = 0.0;
+    let TILE: u32 = 16u;
+    let num_tiles = (M + TILE - 1u) / TILE;
+
+    for (var tile_idx: u32 = 0u; tile_idx < num_tiles; tile_idx = tile_idx + 1u) {
+        // dY^T tile: [N, M] accessed as dY[m, n]
+        let m_idx = tile_idx * TILE + lid.y;
+        let n_idx = wid.x * TILE + lid.x;
+        if (n_idx < N && m_idx < M) {
+            tile_dY2[lid.x * TILE + lid.y] = dY[m_idx * N + n_idx];
+        } else {
+            tile_dY2[lid.x * TILE + lid.y] = 0.0;
+        }
+
+        // X tile: [M, K]
+        let xm = tile_idx * TILE + lid.x;
+        let xk = wid.y * TILE + lid.y;
+        if (xm < M && xk < K) {
+            tile_X2[lid.x * TILE + lid.y] = X[xm * K + xk];
+        } else {
+            tile_X2[lid.x * TILE + lid.y] = 0.0;
+        }
+
+        workgroupBarrier();
+
+        for (var m: u32 = 0u; m < TILE; m = m + 1u) {
+            acc = acc + tile_dY2[lid.x * TILE + m] * tile_X2[m * TILE + lid.y];
+        }
+        workgroupBarrier();
+    }
+
+    if (row < N && col < K) {
+        dW[row * K + col] = acc;
+    }
+}
+
+// Dispatch: (N, 1, 1) – accumulates db = sum_M dY
+@compute @workgroup_size(64, 1, 1)
+fn linear_backward_db(
+    @builtin(global_invocation_id) gid : vec3<u32>,
+) {
+    let M = params.M;
+    let N = params.N;
+
+    let n = gid.x;
+    if (n >= N) { return; }
+
+    var acc: f32 = 0.0;
+    for (var m: u32 = 0u; m < M; m = m + 1u) {
+        acc = acc + dY[m * N + n];
+    }
+    db[n] = acc;
+}
+`;