@seanhogg/builderforce-memory-engine 2026.7.14 → 2026.9.27

package/src/training/mamba1_cpu.ts

@@ -0,0 +1,421 @@
+/**
+ * mamba1_cpu.ts — exact CPU forward + backward for a Mamba-1 (S6) block.
+ *
+ * This is the REFERENCE implementation of the same maths the WGSL kernels run
+ * (`rmsnorm_forward`, `linear_forward`, `conv1d_forward`, `silu_forward`,
+ * `forward_scan` + `forward_reduce`). It exists so real per-parameter gradients
+ * can be computed — and, crucially, PROVEN correct by finite differences —
+ * without a WebGPU device. {@link MambaTrainer} consumes it as its gradient
+ * engine: the optimizer (AdamW + trust region + NaN guard) stays on the GPU,
+ * the gradients are exact.
+ *
+ * Every formula below is a line-for-line mirror of the corresponding kernel; if
+ * a kernel changes, this file changes with it and `tests/backprop.test.ts`
+ * catches the drift (the gradient check fails the moment forward and backward
+ * disagree).
+ *
+ * Layout conventions (identical to the GPU buffers):
+ *   x        (M, dModel)        M = batch * seqLen, row r = b * seqLen + t
+ *   wInProj  (2*dInner, dModel) row-major, Y = X @ W^T + b
+ *   wConv    (dInner, dConv)    depthwise, causal
+ *   wXProj   (dtRank + 2*dState, dInner)
+ *   wDtProj  (dInner, dtRank)
+ *   A_log    (dInner, dState)
+ *   wOutProj (dModel, dInner)
+ *   hCache   (B, L, dInner, dState)
+ */
+
+import {
+    linearForward,
+    linearBackward,
+    rmsNormForward,
+    rmsNormBackward,
+    silu,
+    siluGrad,
+    softplus,
+    sigmoid,
+    RMSNORM_EPS,
+} from './cpu_ops.js';
+
+/** The 13 learnable tensors of a Mamba-1 block, keyed by the SAME names
+ *  {@link Mamba1Block.parameters} publishes. */
+export interface Mamba1CpuWeights {
+    wInProj   : Float32Array;
+    bInProj   : Float32Array;
+    wConv     : Float32Array;
+    bConv     : Float32Array;
+    wXProj    : Float32Array;
+    bXProj    : Float32Array;
+    wDtProj   : Float32Array;
+    bDtProj   : Float32Array;
+    A_log     : Float32Array;
+    D_vec     : Float32Array;
+    wOutProj  : Float32Array;
+    bOutProj  : Float32Array;
+    normWeight: Float32Array;
+}
+
+/** Parameter names of a Mamba-1 block, in `parameters()` order. */
+export const MAMBA1_PARAM_NAMES = [
+    'wInProj', 'bInProj', 'wConv', 'bConv', 'wXProj', 'bXProj',
+    'wDtProj', 'bDtProj', 'A_log', 'D_vec', 'wOutProj', 'bOutProj', 'normWeight',
+] as const satisfies ReadonlyArray<keyof Mamba1CpuWeights>;
+
+export interface Mamba1CpuDims {
+    dModel : number;
+    dState : number;
+    dConv  : number;
+    dInner : number;
+    dtRank : number;
+    batch  : number;
+    seqLen : number;
+}
+
+/** Everything the backward pass needs from the forward pass. */
+export interface Mamba1CpuCache {
+    x        : Float32Array;   // (M, dModel) block input (also the residual)
+    normInv  : Float32Array;   // (M,)
+    normOut  : Float32Array;   // (M, dModel)
+    xConvIn  : Float32Array;   // (M, D)
+    z        : Float32Array;   // (M, D)
+    convOut  : Float32Array;   // (M, D)
+    u        : Float32Array;   // (M, D) = SiLU(convOut)
+    dtRaw    : Float32Array;   // (M, R)
+    bRaw     : Float32Array;   // (M, N)
+    cRaw     : Float32Array;   // (M, N)
+    deltaFull: Float32Array;   // (M, D)
+    dv       : Float32Array;   // (M, D) = softplus(deltaFull)
+    aCont    : Float32Array;   // (D, N)
+    aBar     : Float32Array;   // (M, D, N)
+    bBar     : Float32Array;   // (M, D, N)
+    h        : Float32Array;   // (M, D, N)
+    scanY    : Float32Array;   // (M, D)
+    siluZ    : Float32Array;   // (M, D)
+    gated    : Float32Array;   // (M, D)
+}
+
+export interface Mamba1CpuForwardResult {
+    output: Float32Array;      // (M, dModel)
+    cache : Mamba1CpuCache;
+}
+
+/** Gradients for every Mamba-1 parameter, same keys as {@link Mamba1CpuWeights}. */
+export type Mamba1CpuGrads = Record<keyof Mamba1CpuWeights, Float32Array>;
+
+/**
+ * Clamp bounds on `A_log` applied by `discretise_A`/`discretise_B` in
+ * `SELECTIVE_SCAN_FORWARD_WGSL`. Outside this range the kernel's output is
+ * constant in `A_log`, so the gradient is exactly zero — mirrored here so the
+ * analytic gradient matches the function the model actually computes.
+ */
+export const A_LOG_CLAMP_LO = -10;
+export const A_LOG_CLAMP_HI = 5;
+
+/** Allocate a zeroed gradient set matching `dims`. */
+export function zeroMamba1Grads(dims: Mamba1CpuDims): Mamba1CpuGrads {
+    const { dModel, dState: N, dConv: K, dInner: D, dtRank: R } = dims;
+    return {
+        wInProj   : new Float32Array(2 * D * dModel),
+        bInProj   : new Float32Array(2 * D),
+        wConv     : new Float32Array(D * K),
+        bConv     : new Float32Array(D),
+        wXProj    : new Float32Array((R + 2 * N) * D),
+        bXProj    : new Float32Array(R + 2 * N),
+        wDtProj   : new Float32Array(D * R),
+        bDtProj   : new Float32Array(D),
+        A_log     : new Float32Array(D * N),
+        D_vec     : new Float32Array(D),
+        wOutProj  : new Float32Array(dModel * D),
+        bOutProj  : new Float32Array(dModel),
+        normWeight: new Float32Array(dModel),
+    };
+}
+
+/**
+ * Exact CPU forward for one Mamba-1 block.
+ *
+ * Mirrors `Mamba1Block.forward` step for step:
+ *   RMSNorm → in_proj → split(x, z) → causal conv1d → SiLU → x_proj(Δ, B, C)
+ *   → dt_proj → selective scan → gate by SiLU(z) → out_proj → residual add.
+ */
+export function mamba1CpuForward(
+    x: Float32Array,
+    w: Mamba1CpuWeights,
+    dims: Mamba1CpuDims,
+): Mamba1CpuForwardResult {
+    const { dModel, dState: N, dConv: K, dInner: D, dtRank: R, batch: B, seqLen: L } = dims;
+    const M = B * L;
+    const P = R + 2 * N;
+
+    // 1. Pre-block RMSNorm.
+    const normOut = new Float32Array(M * dModel);
+    const normInv = new Float32Array(M);
+    rmsNormForward(x, w.normWeight, M, dModel, normOut, normInv);
+
+    // 2. Input projection → (x, z).
+    const inProj = new Float32Array(M * 2 * D);
+    linearForward(normOut, w.wInProj, w.bInProj, M, dModel, 2 * D, inProj);
+
+    // 3. COLUMN split: x = inProj[:, 0:D], z = inProj[:, D:2D].
+    const xConvIn = new Float32Array(M * D);
+    const z       = new Float32Array(M * D);
+    for (let r = 0; r < M; r++) {
+        for (let d = 0; d < D; d++) {
+            xConvIn[r * D + d] = inProj[r * 2 * D + d]!;
+            z[r * D + d]       = inProj[r * 2 * D + D + d]!;
+        }
+    }
+
+    // 4. Depthwise causal conv1d over the sequence axis.
+    const convOut = new Float32Array(M * D);
+    for (let b = 0; b < B; b++) {
+        for (let t = 0; t < L; t++) {
+            const base = (b * L + t) * D;
+            for (let d = 0; d < D; d++) {
+                let acc = w.bConv[d]!;
+                for (let k = 0; k < K; k++) {
+                    if (t >= k) acc += w.wConv[d * K + k]! * xConvIn[(b * L + t - k) * D + d]!;
+                }
+                convOut[base + d] = acc;
+            }
+        }
+    }
+
+    // 5. SiLU.
+    const u = new Float32Array(M * D);
+    for (let i = 0; i < M * D; i++) u[i] = silu(convOut[i]!);
+
+    // 6. x_proj → Δ_raw, B, C (COLUMN slices of the (M, R+2N) projection).
+    const xProj = new Float32Array(M * P);
+    linearForward(u, w.wXProj, w.bXProj, M, D, P, xProj);
+    const dtRaw = new Float32Array(M * R);
+    const bRaw  = new Float32Array(M * N);
+    const cRaw  = new Float32Array(M * N);
+    for (let r = 0; r < M; r++) {
+        for (let i = 0; i < R; i++) dtRaw[r * R + i] = xProj[r * P + i]!;
+        for (let n = 0; n < N; n++) {
+            bRaw[r * N + n] = xProj[r * P + R + n]!;
+            cRaw[r * N + n] = xProj[r * P + R + N + n]!;
+        }
+    }
+
+    // 7. dt_proj: expand Δ to the inner dimension.
+    const deltaFull = new Float32Array(M * D);
+    linearForward(dtRaw, w.wDtProj, w.bDtProj, M, R, D, deltaFull);
+
+    // 8. Selective scan (S6). ZOH discretisation, exactly as the kernel does it.
+    const dv = new Float32Array(M * D);
+    for (let i = 0; i < M * D; i++) dv[i] = softplus(deltaFull[i]!);
+
+    const aCont = new Float32Array(D * N);
+    for (let i = 0; i < D * N; i++) {
+        aCont[i] = -Math.exp(Math.min(A_LOG_CLAMP_HI, Math.max(A_LOG_CLAMP_LO, w.A_log[i]!)));
+    }
+
+    const aBar  = new Float32Array(M * D * N);
+    const bBar  = new Float32Array(M * D * N);
+    const h     = new Float32Array(M * D * N);
+    const scanY = new Float32Array(M * D);
+
+    for (let b = 0; b < B; b++) {
+        for (let t = 0; t < L; t++) {
+            const r = b * L + t;
+            for (let d = 0; d < D; d++) {
+                const dvv = dv[r * D + d]!;
+                const uv  = u[r * D + d]!;
+                let acc = w.D_vec[d]! * uv;
+                for (let n = 0; n < N; n++) {
+                    const ac = aCont[d * N + n]!;
+                    const ab = Math.exp(dvv * ac);
+                    const bb = ((ab - 1) / ac) * bRaw[r * N + n]!;
+                    const idx = (r * D + d) * N + n;
+                    aBar[idx] = ab;
+                    bBar[idx] = bb;
+                    const hPrev = t === 0 ? 0 : h[((r - 1) * D + d) * N + n]!;
+                    const hv = ab * hPrev + bb * uv;
+                    h[idx] = hv;
+                    acc += cRaw[r * N + n]! * hv;
+                }
+                scanY[r * D + d] = acc;
+            }
+        }
+    }
+
+    // 9. Gate: y ⊗ SiLU(z).
+    const siluZ = new Float32Array(M * D);
+    const gated = new Float32Array(M * D);
+    for (let i = 0; i < M * D; i++) {
+        siluZ[i] = silu(z[i]!);
+        gated[i] = scanY[i]! * siluZ[i]!;
+    }
+
+    // 10 + 11. Output projection, then the residual add.
+    const output = new Float32Array(M * dModel);
+    linearForward(gated, w.wOutProj, w.bOutProj, M, D, dModel, output);
+    for (let i = 0; i < M * dModel; i++) output[i] = output[i]! + x[i]!;
+
+    return {
+        output,
+        cache: { x, normInv, normOut, xConvIn, z, convOut, u, dtRaw, bRaw, cRaw, deltaFull, dv, aCont, aBar, bBar, h, scanY, siluZ, gated },
+    };
+}
+
+/**
+ * Exact CPU backward for one Mamba-1 block.
+ *
+ * @param dOut  dL/d(block output), (M, dModel)
+ * @param grads accumulator — parameter gradients are ADDED into it
+ * @returns dL/d(block input), (M, dModel)
+ */
+export function mamba1CpuBackward(
+    dOut: Float32Array,
+    w: Mamba1CpuWeights,
+    cache: Mamba1CpuCache,
+    dims: Mamba1CpuDims,
+    grads: Mamba1CpuGrads,
+): Float32Array {
+    const { dModel, dState: N, dConv: K, dInner: D, dtRank: R, batch: B, seqLen: L } = dims;
+    const M = B * L;
+    const P = R + 2 * N;
+
+    // 11. Residual: the block input receives dOut directly, plus whatever flows
+    //     back through the mixer branch (accumulated at the very end).
+    const dx = new Float32Array(M * dModel);
+    dx.set(dOut);
+
+    // 10. out_proj.
+    const dGated = new Float32Array(M * D);
+    linearBackward(dOut, cache.gated, w.wOutProj, M, D, dModel, dGated, grads.wOutProj, grads.bOutProj);
+
+    // 9. Gate: gated = scanY * SiLU(z).
+    const dScanY = new Float32Array(M * D);
+    const dz     = new Float32Array(M * D);
+    for (let i = 0; i < M * D; i++) {
+        dScanY[i] = dGated[i]! * cache.siluZ[i]!;
+        dz[i]     = dGated[i]! * cache.scanY[i]! * siluGrad(cache.z[i]!);
+    }
+
+    // 8. Selective scan, walked backwards in time.
+    const du    = new Float32Array(M * D);
+    const dBRaw = new Float32Array(M * N);
+    const dCRaw = new Float32Array(M * N);
+    const dDv   = new Float32Array(M * D);      // dL/d softplus(Δ)
+    const dACont = new Float32Array(D * N);
+
+    // dhNext[d*N+n] carries dL/dh_t propagated from step t+1 for the current batch row.
+    const dhNext = new Float32Array(D * N);
+    for (let b = 0; b < B; b++) {
+        dhNext.fill(0);
+        for (let t = L - 1; t >= 0; t--) {
+            const r = b * L + t;
+            for (let d = 0; d < D; d++) {
+                const dY  = dScanY[r * D + d]!;
+                const uv  = cache.u[r * D + d]!;
+                const dvv = cache.dv[r * D + d]!;
+                // Skip connection: y_t[d] += D_vec[d] * u_t[d].
+                grads.D_vec[d] = grads.D_vec[d]! + dY * uv;
+                du[r * D + d] = du[r * D + d]! + dY * w.D_vec[d]!;
+
+                let ddvAcc = 0;
+                for (let n = 0; n < N; n++) {
+                    const idx = (r * D + d) * N + n;
+                    const hv  = cache.h[idx]!;
+                    const ab  = cache.aBar[idx]!;
+                    const bb  = cache.bBar[idx]!;
+                    const ac  = cache.aCont[d * N + n]!;
+                    const cv  = cache.cRaw[r * N + n]!;
+
+                    // y_t[d] = Σ_n C_t[n] * h_t[d,n]
+                    dCRaw[r * N + n] = dCRaw[r * N + n]! + dY * hv;
+                    const dh = dY * cv + dhNext[d * N + n]!;
+
+                    // h_t = a_bar * h_{t-1} + b_bar * u_t
+                    const hPrev = t === 0 ? 0 : cache.h[((r - 1) * D + d) * N + n]!;
+                    const dABar = dh * hPrev;
+                    const dBBar = dh * uv;
+                    du[r * D + d] = du[r * D + d]! + dh * bb;
+                    dhNext[d * N + n] = dh * ab;   // flows to step t-1
+
+                    // a_bar = exp(dv * a_cont);  b_bar = (a_bar - 1) / a_cont * B_raw
+                    const bv = cache.bRaw[r * N + n]!;
+                    ddvAcc += dABar * ac * ab + dBBar * ab * bv;
+                    dACont[d * N + n] = dACont[d * N + n]!
+                        + dABar * dvv * ab
+                        + dBBar * bv * ((dvv * ab * ac - (ab - 1)) / (ac * ac));
+                    dBRaw[r * N + n] = dBRaw[r * N + n]! + dBBar * ((ab - 1) / ac);
+                }
+                dDv[r * D + d] = dDv[r * D + d]! + ddvAcc;
+            }
+        }
+    }
+
+    // a_cont = -exp(clamp(A_log)) ⇒ d a_cont / d A_log = a_cont inside the clamp,
+    // and exactly 0 outside it (the kernel's output is constant there).
+    for (let i = 0; i < D * N; i++) {
+        const al = w.A_log[i]!;
+        if (al <= A_LOG_CLAMP_LO || al >= A_LOG_CLAMP_HI) continue;
+        grads.A_log[i] = grads.A_log[i]! + dACont[i]! * cache.aCont[i]!;
+    }
+
+    // softplus: dv = softplus(deltaFull)
+    const dDeltaFull = new Float32Array(M * D);
+    for (let i = 0; i < M * D; i++) dDeltaFull[i] = dDv[i]! * sigmoid(cache.deltaFull[i]!);
+
+    // 7. dt_proj.
+    const dDtRaw = new Float32Array(M * R);
+    linearBackward(dDeltaFull, cache.dtRaw, w.wDtProj, M, R, D, dDtRaw, grads.wDtProj, grads.bDtProj);
+
+    // 6. x_proj — re-assemble the (M, P) column slices, then one linear backward.
+    const dXProj = new Float32Array(M * P);
+    for (let r = 0; r < M; r++) {
+        for (let i = 0; i < R; i++) dXProj[r * P + i] = dDtRaw[r * R + i]!;
+        for (let n = 0; n < N; n++) {
+            dXProj[r * P + R + n]     = dBRaw[r * N + n]!;
+            dXProj[r * P + R + N + n] = dCRaw[r * N + n]!;
+        }
+    }
+    const duFromXProj = new Float32Array(M * D);
+    linearBackward(dXProj, cache.u, w.wXProj, M, D, P, duFromXProj, grads.wXProj, grads.bXProj);
+    for (let i = 0; i < M * D; i++) du[i] = du[i]! + duFromXProj[i]!;
+
+    // 5. SiLU.
+    const dConvOut = new Float32Array(M * D);
+    for (let i = 0; i < M * D; i++) dConvOut[i] = du[i]! * siluGrad(cache.convOut[i]!);
+
+    // 4. Depthwise causal conv1d.
+    const dXConvIn = new Float32Array(M * D);
+    for (let b = 0; b < B; b++) {
+        for (let t = 0; t < L; t++) {
+            const r = b * L + t;
+            for (let d = 0; d < D; d++) {
+                const g = dConvOut[r * D + d]!;
+                if (g === 0) continue;
+                grads.bConv[d] = grads.bConv[d]! + g;
+                for (let k = 0; k < K; k++) {
+                    if (t < k) continue;
+                    const src = (b * L + t - k) * D + d;
+                    grads.wConv[d * K + k] = grads.wConv[d * K + k]! + g * cache.xConvIn[src]!;
+                    dXConvIn[src] = dXConvIn[src]! + g * w.wConv[d * K + k]!;
+                }
+            }
+        }
+    }
+
+    // 3. Un-split the (x, z) columns.
+    const dInProj = new Float32Array(M * 2 * D);
+    for (let r = 0; r < M; r++) {
+        for (let d = 0; d < D; d++) {
+            dInProj[r * 2 * D + d]     = dXConvIn[r * D + d]!;
+            dInProj[r * 2 * D + D + d] = dz[r * D + d]!;
+        }
+    }
+
+    // 2. in_proj.
+    const dNormOut = new Float32Array(M * dModel);
+    linearBackward(dInProj, cache.normOut, w.wInProj, M, dModel, 2 * D, dNormOut, grads.wInProj, grads.bInProj);
+
+    // 1. RMSNorm — adds into the residual path already sitting in `dx`.
+    rmsNormBackward(dNormOut, cache.x, w.normWeight, cache.normInv, M, dModel, dx, grads.normWeight, RMSNORM_EPS);
+
+    return dx;
+}