@seanhogg/builderforce-memory-engine 2026.10.1 → 2026.10.3

package/src/training/cpu_ops.ts

@@ -1,141 +1,141 @@
-/**
- * cpu_ops.ts — the small set of exact CPU primitives the reference forward and
- * backward passes share (linear, RMSNorm, SiLU, softplus).
- *
- * Each one mirrors a WGSL kernel exactly, so a gradient computed here is the
- * gradient of the function the GPU actually evaluates. They are deliberately
- * plain loops over `Float32Array`: correctness first, and small enough that the
- * finite-difference check in `tests/backprop.test.ts` is meaningful.
- */
-
-/** Epsilon used by `rmsnorm_forward` in `ACTIVATIONS_WGSL`. */
-export const RMSNORM_EPS = 1e-6;
-
-export function sigmoid(x: number): number {
-    return 1 / (1 + Math.exp(-x));
-}
-
-/** Numerically-stable softmax over a flat array. */
-export function softmax(logits: ArrayLike<number>): Float32Array {
-    let max = -Infinity;
-    for (let i = 0; i < logits.length; i++) if (logits[i]! > max) max = logits[i]!;
-    const out = new Float32Array(logits.length);
-    let sum = 0;
-    for (let i = 0; i < logits.length; i++) {
-        out[i] = Math.exp(logits[i]! - max);
-        sum += out[i]!;
-    }
-    for (let i = 0; i < logits.length; i++) out[i] = out[i]! / sum;
-    return out;
-}
-
-/** SiLU(x) = x * sigmoid(x) — mirrors `silu_forward`. */
-export function silu(x: number): number {
-    return x / (1 + Math.exp(-x));
-}
-
-/** d/dx SiLU(x) = sigmoid(x) * (1 + x * (1 - sigmoid(x))). */
-export function siluGrad(x: number): number {
-    const s = sigmoid(x);
-    return s * (1 + x * (1 - s));
-}
-
-/** Numerically stable softplus — mirrors the kernel's `max(x,0) + log1p(exp(-|x|))`. */
-export function softplus(x: number): number {
-    return Math.max(x, 0) + Math.log1p(Math.exp(-Math.abs(x)));
-}
-
-/**
- * Y = X @ W^T + b, mirroring `linear_forward`.
- *
- * @param X   (M, K)
- * @param W   (N, K) row-major
- * @param b   (N,)
- * @param out (M, N) — written, not accumulated
- */
-export function linearForward(
-    X: Float32Array, W: Float32Array, b: Float32Array,
-    M: number, K: number, N: number, out: Float32Array,
-): void {
-    for (let r = 0; r < M; r++) {
-        const xo = r * K;
-        const yo = r * N;
-        for (let n = 0; n < N; n++) {
-            let acc = b[n]!;
-            const wo = n * K;
-            for (let k = 0; k < K; k++) acc += X[xo + k]! * W[wo + k]!;
-            out[yo + n] = acc;
-        }
-    }
-}
-
-/**
- * Backward of {@link linearForward}. `dX` is written; `dW` and `db` are
- * ACCUMULATED into (so one parameter can collect gradient from several call
- * sites without a temporary).
- */
-export function linearBackward(
-    dY: Float32Array, X: Float32Array, W: Float32Array,
-    M: number, K: number, N: number,
-    dX: Float32Array, dW: Float32Array, db: Float32Array,
-): void {
-    dX.fill(0);
-    for (let r = 0; r < M; r++) {
-        const xo = r * K;
-        const yo = r * N;
-        for (let n = 0; n < N; n++) {
-            const g = dY[yo + n]!;
-            if (g === 0) continue;
-            db[n] = db[n]! + g;
-            const wo = n * K;
-            for (let k = 0; k < K; k++) {
-                dX[xo + k] = dX[xo + k]! + g * W[wo + k]!;
-                dW[wo + k] = dW[wo + k]! + g * X[xo + k]!;
-            }
-        }
-    }
-}
-
-/**
- * RMSNorm forward — mirrors `rmsnorm_forward`:
- *   inv[r] = 1 / sqrt(mean(x[r]²) + eps);  y[r,i] = x[r,i] * inv[r] * w[i]
- */
-export function rmsNormForward(
-    x: Float32Array, w: Float32Array, M: number, D: number,
-    out: Float32Array, inv: Float32Array, eps: number = RMSNORM_EPS,
-): void {
-    for (let r = 0; r < M; r++) {
-        const base = r * D;
-        let sq = 0;
-        for (let i = 0; i < D; i++) { const v = x[base + i]!; sq += v * v; }
-        const iv = 1 / Math.sqrt(sq / D + eps);
-        inv[r] = iv;
-        for (let i = 0; i < D; i++) out[base + i] = x[base + i]! * iv * w[i]!;
-    }
-}
-
-/**
- * RMSNorm backward. `dx` and `dw` are ACCUMULATED into — `dx` in particular is
- * usually the residual path, which already carries a gradient.
- */
-export function rmsNormBackward(
-    dY: Float32Array, x: Float32Array, w: Float32Array, inv: Float32Array,
-    M: number, D: number,
-    dx: Float32Array, dw: Float32Array, _eps: number = RMSNORM_EPS,
-): void {
-    for (let r = 0; r < M; r++) {
-        const base = r * D;
-        const iv = inv[r]!;
-        let dot = 0;   // Σ_i (dY[r,i] * w[i]) * x[r,i]
-        for (let i = 0; i < D; i++) {
-            const g = dY[base + i]! * w[i]!;
-            dot += g * x[base + i]!;
-            dw[i] = dw[i]! + dY[base + i]! * x[base + i]! * iv;
-        }
-        const scale = (iv * iv * iv) / D;
-        for (let i = 0; i < D; i++) {
-            const g = dY[base + i]! * w[i]!;
-            dx[base + i] = dx[base + i]! + g * iv - scale * dot * x[base + i]!;
-        }
-    }
-}
+/**
+ * cpu_ops.ts — the small set of exact CPU primitives the reference forward and
+ * backward passes share (linear, RMSNorm, SiLU, softplus).
+ *
+ * Each one mirrors a WGSL kernel exactly, so a gradient computed here is the
+ * gradient of the function the GPU actually evaluates. They are deliberately
+ * plain loops over `Float32Array`: correctness first, and small enough that the
+ * finite-difference check in `tests/backprop.test.ts` is meaningful.
+ */
+
+/** Epsilon used by `rmsnorm_forward` in `ACTIVATIONS_WGSL`. */
+export const RMSNORM_EPS = 1e-6;
+
+export function sigmoid(x: number): number {
+    return 1 / (1 + Math.exp(-x));
+}
+
+/** Numerically-stable softmax over a flat array. */
+export function softmax(logits: ArrayLike<number>): Float32Array {
+    let max = -Infinity;
+    for (let i = 0; i < logits.length; i++) if (logits[i]! > max) max = logits[i]!;
+    const out = new Float32Array(logits.length);
+    let sum = 0;
+    for (let i = 0; i < logits.length; i++) {
+        out[i] = Math.exp(logits[i]! - max);
+        sum += out[i]!;
+    }
+    for (let i = 0; i < logits.length; i++) out[i] = out[i]! / sum;
+    return out;
+}
+
+/** SiLU(x) = x * sigmoid(x) — mirrors `silu_forward`. */
+export function silu(x: number): number {
+    return x / (1 + Math.exp(-x));
+}
+
+/** d/dx SiLU(x) = sigmoid(x) * (1 + x * (1 - sigmoid(x))). */
+export function siluGrad(x: number): number {
+    const s = sigmoid(x);
+    return s * (1 + x * (1 - s));
+}
+
+/** Numerically stable softplus — mirrors the kernel's `max(x,0) + log1p(exp(-|x|))`. */
+export function softplus(x: number): number {
+    return Math.max(x, 0) + Math.log1p(Math.exp(-Math.abs(x)));
+}
+
+/**
+ * Y = X @ W^T + b, mirroring `linear_forward`.
+ *
+ * @param X   (M, K)
+ * @param W   (N, K) row-major
+ * @param b   (N,)
+ * @param out (M, N) — written, not accumulated
+ */
+export function linearForward(
+    X: Float32Array, W: Float32Array, b: Float32Array,
+    M: number, K: number, N: number, out: Float32Array,
+): void {
+    for (let r = 0; r < M; r++) {
+        const xo = r * K;
+        const yo = r * N;
+        for (let n = 0; n < N; n++) {
+            let acc = b[n]!;
+            const wo = n * K;
+            for (let k = 0; k < K; k++) acc += X[xo + k]! * W[wo + k]!;
+            out[yo + n] = acc;
+        }
+    }
+}
+
+/**
+ * Backward of {@link linearForward}. `dX` is written; `dW` and `db` are
+ * ACCUMULATED into (so one parameter can collect gradient from several call
+ * sites without a temporary).
+ */
+export function linearBackward(
+    dY: Float32Array, X: Float32Array, W: Float32Array,
+    M: number, K: number, N: number,
+    dX: Float32Array, dW: Float32Array, db: Float32Array,
+): void {
+    dX.fill(0);
+    for (let r = 0; r < M; r++) {
+        const xo = r * K;
+        const yo = r * N;
+        for (let n = 0; n < N; n++) {
+            const g = dY[yo + n]!;
+            if (g === 0) continue;
+            db[n] = db[n]! + g;
+            const wo = n * K;
+            for (let k = 0; k < K; k++) {
+                dX[xo + k] = dX[xo + k]! + g * W[wo + k]!;
+                dW[wo + k] = dW[wo + k]! + g * X[xo + k]!;
+            }
+        }
+    }
+}
+
+/**
+ * RMSNorm forward — mirrors `rmsnorm_forward`:
+ *   inv[r] = 1 / sqrt(mean(x[r]²) + eps);  y[r,i] = x[r,i] * inv[r] * w[i]
+ */
+export function rmsNormForward(
+    x: Float32Array, w: Float32Array, M: number, D: number,
+    out: Float32Array, inv: Float32Array, eps: number = RMSNORM_EPS,
+): void {
+    for (let r = 0; r < M; r++) {
+        const base = r * D;
+        let sq = 0;
+        for (let i = 0; i < D; i++) { const v = x[base + i]!; sq += v * v; }
+        const iv = 1 / Math.sqrt(sq / D + eps);
+        inv[r] = iv;
+        for (let i = 0; i < D; i++) out[base + i] = x[base + i]! * iv * w[i]!;
+    }
+}
+
+/**
+ * RMSNorm backward. `dx` and `dw` are ACCUMULATED into — `dx` in particular is
+ * usually the residual path, which already carries a gradient.
+ */
+export function rmsNormBackward(
+    dY: Float32Array, x: Float32Array, w: Float32Array, inv: Float32Array,
+    M: number, D: number,
+    dx: Float32Array, dw: Float32Array, _eps: number = RMSNORM_EPS,
+): void {
+    for (let r = 0; r < M; r++) {
+        const base = r * D;
+        const iv = inv[r]!;
+        let dot = 0;   // Σ_i (dY[r,i] * w[i]) * x[r,i]
+        for (let i = 0; i < D; i++) {
+            const g = dY[base + i]! * w[i]!;
+            dot += g * x[base + i]!;
+            dw[i] = dw[i]! + dY[base + i]! * x[base + i]! * iv;
+        }
+        const scale = (iv * iv * iv) / D;
+        for (let i = 0; i < D; i++) {
+            const g = dY[base + i]! * w[i]!;
+            dx[base + i] = dx[base + i]! + g * iv - scale * dot * x[base + i]!;
+        }
+    }
+}