@seanhogg/builderforce-memory-engine 2026.7.14 → 2026.9.28

package/src/model/mamba1_block.ts

@@ -20,6 +20,7 @@ import { gaussianArray } from '../utils/rng.js';
 import { CONV1D_FORWARD_WGSL }          from '../kernels/conv1d.js';
 import { LINEAR_FORWARD_WGSL }          from '../kernels/linear_projection.js';
 import { ACTIVATIONS_WGSL }             from '../kernels/activations.js';
+import { COL_SLICE_WGSL, COL_SLICE_ENTRY, dispatchColumnSlice } from '../kernels/slice.js';
 
 import type { SequenceLayer, LayerForwardResult, LayerParam } from './sequence_layer.js';
 
@@ -210,6 +211,7 @@ export class Mamba1Block implements SequenceLayer {
             rmsnorm     : createComputePipeline(d, ACTIVATIONS_WGSL,             'rmsnorm_forward'),
             scan_fwd    : createComputePipeline(d, SELECTIVE_SCAN_FORWARD_WGSL,  'forward_scan'),
             scan_reduce : createComputePipeline(d, SELECTIVE_SCAN_FORWARD_WGSL,  'forward_reduce'),
+            colSlice    : createComputePipeline(d, COL_SLICE_WGSL, COL_SLICE_ENTRY),
             elMul       : createComputePipeline(d, MUL_SHADER, 'main'),
             elAdd       : createComputePipeline(d, ADD_SHADER, 'main'),
         };
@@ -253,14 +255,18 @@ export class Mamba1Block implements SequenceLayer {
             dispatchKernel(d, this.pipelines['linear']!, bg, [cdiv(M, 16), cdiv(2 * D, 16), 1]);
         }
 
-        // 3. Split into x and z
+        // 3. Split into x and z.
+        //    in_proj writes (M, 2D) ROW-major, so x and z are interleaved COLUMN
+        //    slices — one per row — not two contiguous halves of the buffer. A
+        //    `copyBufferToBuffer` of the first M*D floats is only the right answer
+        //    when M == 1; for any real sequence it hands the block the first half of
+        //    the ROWS twice. Split with a strided gather instead.
         const xConvIn = createEmptyStorageBuffer(d, M * D * 4, true);
         const zBuf    = createEmptyStorageBuffer(d, M * D * 4, true);
         {
-            const enc = d.createCommandEncoder();
-            enc.copyBufferToBuffer(inProjOut, 0,         xConvIn, 0, M * D * 4);
-            enc.copyBufferToBuffer(inProjOut, M * D * 4, zBuf,    0, M * D * 4);
-            d.queue.submit([enc.finish()]);
+            const slice = this.pipelines['colSlice']!;
+            dispatchColumnSlice(d, slice, inProjOut, xConvIn, M, 2 * D, 0, D);
+            dispatchColumnSlice(d, slice, inProjOut, zBuf, M, 2 * D, D, D);
         }
         inProjOut.destroy();
         cache.zBuf    = zBuf;
@@ -270,7 +276,11 @@ export class Mamba1Block implements SequenceLayer {
         const convOut = createEmptyStorageBuffer(d, M * D * 4, true);
         cache.convOut = convOut;
         {
-            const params = new Uint32Array([L, D, dConv, B]).buffer;
+            // ConvParams is FIVE u32s (seq_len, d_channels, kernel_size, batch,
+            // groups). A four-element uniform is smaller than the struct's minimum
+            // binding size and WebGPU rejects the bind group outright. groups = 1
+            // is standard depthwise, matching Mamba-2/3's call.
+            const params = new Uint32Array([L, D, dConv, B, 1]).buffer;
             const pBuf   = createUniformBuffer(d, params);
             const bg = createBindGroup(d, this.pipelines['conv1d']!,
                 [pBuf, xConvIn, this.gpuWeights['wConv']!, this.gpuWeights['bConv']!, convOut]);
@@ -298,15 +308,17 @@ export class Mamba1Block implements SequenceLayer {
             dispatchKernel(d, this.pipelines['linear']!, bg, [cdiv(M, 16), cdiv(R + 2 * N, 16), 1]);
         }
 
+        //    Same story as the (x, z) split: Δ, B and C are COLUMN slices of the
+        //    (M, R+2N) projection, so they need a strided gather, not a prefix copy.
         const dtRaw = createEmptyStorageBuffer(d, M * R * 4, true);
         const B_raw = createEmptyStorageBuffer(d, B * L * N * 4, true);
         const C_raw = createEmptyStorageBuffer(d, B * L * N * 4, true);
         {
-            const enc = d.createCommandEncoder();
-            enc.copyBufferToBuffer(xProjOut, 0,               dtRaw, 0, M * R * 4);
-            enc.copyBufferToBuffer(xProjOut, M * R * 4,       B_raw, 0, B * L * N * 4);
-            enc.copyBufferToBuffer(xProjOut, M * (R + N) * 4, C_raw, 0, B * L * N * 4);
-            d.queue.submit([enc.finish()]);
+            const P = R + 2 * N;
+            const slice = this.pipelines['colSlice']!;
+            dispatchColumnSlice(d, slice, xProjOut, dtRaw, M, P, 0, R);
+            dispatchColumnSlice(d, slice, xProjOut, B_raw, M, P, R, N);
+            dispatchColumnSlice(d, slice, xProjOut, C_raw, M, P, R + N, N);
         }
         xProjOut.destroy();
         cache.B_raw = B_raw;
@@ -335,7 +347,11 @@ export class Mamba1Block implements SequenceLayer {
             const bg1 = createBindGroup(d, this.pipelines['scan_fwd']!,
                 [pBuf, siluOut, deltaFull, this.gpuWeights['A_log']!, B_raw, C_raw,
                  this.gpuWeights['D_vec']!, scanY, hCache]);
-            dispatchKernel(d, this.pipelines['scan_fwd']!, bg1, [cdiv(D, 8), cdiv(N, 8), B]);
+            // One workgroup per (d, n, batch): the kernel reads `wgid.x` as d and
+            // `wgid.y` as n, so the grid must be D x N x B. Dispatching
+            // ceil(D/8) x ceil(N/8) left all but a 1/64th corner of the state
+            // matrix unscanned (and unwritten in h_cache).
+            dispatchKernel(d, this.pipelines['scan_fwd']!, bg1, [D, N, B]);
 
             const bg2 = createBindGroup(d, this.pipelines['scan_reduce']!,
                 [pBuf, siluOut, deltaFull, this.gpuWeights['A_log']!, B_raw, C_raw,

package/src/model/mamba2_block.ts

@@ -24,6 +24,7 @@ import {
 import { SSD_FORWARD_WGSL }    from '../kernels/ssd.js';
 import { gaussianArray } from '../utils/rng.js';
 import { CONV1D_FORWARD_WGSL } from '../kernels/conv1d.js';
+import { COL_SLICE_WGSL, COL_SLICE_ENTRY, dispatchColumnSlice } from '../kernels/slice.js';
 import { LINEAR_FORWARD_WGSL } from '../kernels/linear_projection.js';
 import { ACTIVATIONS_WGSL }    from '../kernels/activations.js';
 
@@ -127,6 +128,7 @@ export class Mamba2Block implements SequenceLayer {
         this.pipelines = {
             linear   : createComputePipeline(d, LINEAR_FORWARD_WGSL,  'linear_forward'),
             conv1d   : createComputePipeline(d, CONV1D_FORWARD_WGSL,  'conv1d_forward'),
+            colSlice   : createComputePipeline(d, COL_SLICE_WGSL, COL_SLICE_ENTRY),
             rmsnorm  : createComputePipeline(d, ACTIVATIONS_WGSL,     'rmsnorm_forward'),
             ssd_fwd  : createComputePipeline(d, SSD_FORWARD_WGSL,     'ssd_chunk_forward'),
             elAdd    : createComputePipeline(d, ADD_SHADER,           'main'),
@@ -181,10 +183,13 @@ export class Mamba2Block implements SequenceLayer {
         const xConvBuf = createEmptyStorageBuffer(d, M * convD * 4, true);
         const dtBuf    = createEmptyStorageBuffer(d, M * H * 4, true);
         {
-            const enc = d.createCommandEncoder();
-            enc.copyBufferToBuffer(inProjOut, 0,              xConvBuf, 0, M * convD * 4);
-            enc.copyBufferToBuffer(inProjOut, M * convD * 4,  dtBuf,    0, M * H * 4);
-            d.queue.submit([enc.finish()]);
+            // COLUMN split. `linear_forward` writes (M, convD + H) ROW-major, so each
+            // part is a range of columns WITHIN every row — a strided gather, not a
+            // contiguous copy (which only coincides for M == 1).
+            const slice = this.pipelines['colSlice']!;
+            const stride = convD + H;
+            dispatchColumnSlice(d, slice, inProjOut, xConvBuf, M, stride, 0, convD);
+            dispatchColumnSlice(d, slice, inProjOut, dtBuf, M, stride, convD, H);
         }
         inProjOut.destroy();
 
@@ -204,11 +209,11 @@ export class Mamba2Block implements SequenceLayer {
         const bProjBuf = createEmptyStorageBuffer(d, M * G * N * 4, true);
         const cProjBuf = createEmptyStorageBuffer(d, M * G * N * 4, true);
         {
-            const enc = d.createCommandEncoder();
-            enc.copyBufferToBuffer(convOut, 0,            xSsdBuf,  0, M * D * 4);
-            enc.copyBufferToBuffer(convOut, M * D * 4,    bProjBuf, 0, M * G * N * 4);
-            enc.copyBufferToBuffer(convOut, M * (D + G * N) * 4, cProjBuf, 0, M * G * N * 4);
-            d.queue.submit([enc.finish()]);
+            // COLUMN split of the (M, convD) conv output — same reasoning as above.
+            const slice = this.pipelines['colSlice']!;
+            dispatchColumnSlice(d, slice, convOut, xSsdBuf, M, convD, 0, D);
+            dispatchColumnSlice(d, slice, convOut, bProjBuf, M, convD, D, G * N);
+            dispatchColumnSlice(d, slice, convOut, cProjBuf, M, convD, D + G * N, G * N);
         }
         convOut.destroy();
 

package/src/model/mamba3_block.ts

@@ -33,6 +33,7 @@ import {
 import { COMPLEX_SSD_FORWARD_WGSL } from '../kernels/complex_ssd.js';
 import { gaussianArray } from '../utils/rng.js';
 import { CONV1D_FORWARD_WGSL }      from '../kernels/conv1d.js';
+import { COL_SLICE_WGSL, COL_SLICE_ENTRY, dispatchColumnSlice } from '../kernels/slice.js';
 import { LINEAR_FORWARD_WGSL }      from '../kernels/linear_projection.js';
 import { ACTIVATIONS_WGSL }         from '../kernels/activations.js';
 
@@ -145,6 +146,7 @@ export class Mamba3Block implements SequenceLayer {
         this.pipelines = {
             linear     : createComputePipeline(d, LINEAR_FORWARD_WGSL,       'linear_forward'),
             conv1d     : createComputePipeline(d, CONV1D_FORWARD_WGSL,       'conv1d_forward'),
+            colSlice   : createComputePipeline(d, COL_SLICE_WGSL, COL_SLICE_ENTRY),
             rmsnorm    : createComputePipeline(d, ACTIVATIONS_WGSL,          'rmsnorm_forward'),
             cssd_fwd   : createComputePipeline(d, COMPLEX_SSD_FORWARD_WGSL,  'complex_ssd_forward'),
             elAdd      : createComputePipeline(d, ADD_SHADER,                'main'),
@@ -198,10 +200,13 @@ export class Mamba3Block implements SequenceLayer {
         const xConvBuf = createEmptyStorageBuffer(d, M * convD * 4, true);
         const dtBuf    = createEmptyStorageBuffer(d, M * H * 4, true);
         {
-            const enc = d.createCommandEncoder();
-            enc.copyBufferToBuffer(inProjOut, 0,             xConvBuf, 0, M * convD * 4);
-            enc.copyBufferToBuffer(inProjOut, M * convD * 4, dtBuf,    0, M * H * 4);
-            d.queue.submit([enc.finish()]);
+            // COLUMN split. `linear_forward` writes (M, convD + H) ROW-major, so each
+            // part is a range of columns WITHIN every row — a strided gather, not a
+            // contiguous copy (which only coincides for M == 1).
+            const slice = this.pipelines['colSlice']!;
+            const stride = convD + H;
+            dispatchColumnSlice(d, slice, inProjOut, xConvBuf, M, stride, 0, convD);
+            dispatchColumnSlice(d, slice, inProjOut, dtBuf, M, stride, convD, H);
         }
         inProjOut.destroy();
 
@@ -221,11 +226,11 @@ export class Mamba3Block implements SequenceLayer {
         const bProjBuf = createEmptyStorageBuffer(d, M * G * Nc * 2 * 4, true);
         const cProjBuf = createEmptyStorageBuffer(d, M * G * Nc * 2 * 4, true);
         {
-            const enc = d.createCommandEncoder();
-            enc.copyBufferToBuffer(convOut, 0,                       xSsdBuf,  0, M * D * 4);
-            enc.copyBufferToBuffer(convOut, M * D * 4,               bProjBuf, 0, M * G * Nc * 2 * 4);
-            enc.copyBufferToBuffer(convOut, M * (D + G * Nc * 2) * 4, cProjBuf, 0, M * G * Nc * 2 * 4);
-            d.queue.submit([enc.finish()]);
+            // COLUMN split of the (M, convD) conv output — same reasoning as above.
+            const slice = this.pipelines['colSlice']!;
+            dispatchColumnSlice(d, slice, convOut, xSsdBuf, M, convD, 0, D);
+            dispatchColumnSlice(d, slice, convOut, bProjBuf, M, convD, D, G * Nc * 2);
+            dispatchColumnSlice(d, slice, convOut, cProjBuf, M, convD, D + G * Nc * 2, G * Nc * 2);
         }
         convOut.destroy();
 

package/src/moe/index.ts

@@ -13,12 +13,3 @@ export type { MoEConfig, MoEParam, RouteResult } from "./moe_model.js";
 export { MoETrainer } from "./moe_trainer.js";
 export type { MoESample, MoETrainOptions, MoEEpochResult } from "./moe_trainer.js";
 
-export { EvermindModelPackage } from "./moe_package.js";
-export type {
-  EvermindModelManifest,
-  EvermindModelCard,
-  EvermindModelType,
-  EvermindModality,
-  PackageMeta,
-  ValidationResult,
-} from "./moe_package.js";

package/src/optim/adamw.ts

@@ -118,3 +118,63 @@ export class AdamW {
     return n;
   }
 }
+
+// ── The CPU mirror of WEIGHT_UPDATE_WGSL ─────────────────────────────────────
+
+/** Hyperparameters of one {@link adamwUpdateInPlace} step. */
+export interface AdamWKernelStep {
+  lr: number;
+  beta1: number;
+  beta2: number;
+  eps: number;
+  weightDecay: number;
+  /** beta1^t — the precomputed bias-correction term (matches the kernel uniform). */
+  beta1_t: number;
+  /** beta2^t. */
+  beta2_t: number;
+  /** Trust region: max |Δθ| per step. 0 disables the bound. */
+  maxDelta: number;
+  /** Global gradient-clip factor, applied to `grad` on the fly (1 = no clipping). */
+  gradScale?: number;
+}
+
+/**
+ * One AdamW step over a single parameter tensor — the exact CPU mirror of the
+ * `adamw_update` entry point in `WEIGHT_UPDATE_WGSL`, including its two safety
+ * properties:
+ *
+ *   • the **NaN/Inf guard** — a non-finite step is dropped rather than written
+ *     into a weight, so one bad gradient cannot permanently poison the model;
+ *   • the **trust region** — `maxDelta` bounds how far a single step can move a
+ *     weight, which is what keeps repeated write-through adapts reversible.
+ *
+ * Decoupled weight decay (`θ * (1 - lr·wd)`), matching the kernel — deliberately
+ * NOT the coupled form {@link AdamW} uses for the MoE/EvermindLM trainers.
+ *
+ * `param`, `m` and `v` are updated in place.
+ */
+export function adamwUpdateInPlace(
+  param: Float32Array,
+  grad: Float32Array,
+  m: Float32Array,
+  v: Float32Array,
+  hp: AdamWKernelStep,
+): void {
+  const { lr, beta1, beta2, eps, weightDecay, beta1_t, beta2_t, maxDelta } = hp;
+  const gradScale = hp.gradScale ?? 1;
+  const bc1 = 1 - beta1_t;
+  const bc2 = 1 - beta2_t;
+  for (let i = 0; i < param.length; i++) {
+    const g = grad[i]! * gradScale;
+    const mNew = beta1 * m[i]! + (1 - beta1) * g;
+    const vNew = beta2 * v[i]! + (1 - beta2) * g * g;
+    m[i] = mNew;
+    v[i] = vNew;
+
+    let step = (lr * (mNew / bc1)) / (Math.sqrt(vNew / bc2) + eps);
+    if (!Number.isFinite(step)) step = 0;
+    if (maxDelta > 0) step = Math.min(maxDelta, Math.max(-maxDelta, step));
+
+    param[i] = param[i]! * (1 - lr * weightDecay) - step;
+  }
+}

package/src/tokenizer/bpe.ts

@@ -52,6 +52,14 @@ const BYTE_DECODER = new Map([...BYTE_ENCODER].map(([k, v]) => [v, k]));
 const PRE_TOKENIZE_RE =
     /(?:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+/gu;
 
+/** A round-trippable snapshot of a {@link BPETokenizer} — see `toObject`/`deserialize`. */
+export interface BPETokenizerSpec {
+    vocab: Record<string, number>;
+    /** Merge rules in rank order ("a b" strings). */
+    merges: string[];
+    specials?: { bos?: string; eos?: string; pad?: string; unk?: string };
+}
+
 export class BPETokenizer {
     vocab: Map<string, number>;
     idToToken: Map<number, string>;
@@ -225,6 +233,53 @@ export class BPETokenizer {
         this.loadFromObjects(vocabObj, mergeArr);
     }
 
+    /**
+     * The tokenizer as a plain, JSON-serialisable object: vocab, merges and the
+     * four special-token strings.
+     *
+     * This is what lets an `.evermind` package EMBED its tokenizer instead of
+     * expecting a consumer to source the matching vocab separately — token ids are
+     * meaningless without the exact vocabulary that produced them, so shipping a
+     * checkpoint without one is shipping half an artifact.
+     */
+    toObject(): BPETokenizerSpec {
+        return {
+            vocab: Object.fromEntries(this.vocab),
+            merges: [...this.merges.entries()].sort((a, b) => a[1] - b[1]).map(([m]) => m),
+            specials: { bos: this.bosToken, eos: this.eosToken, pad: this.padToken, unk: this.unkToken },
+        };
+    }
+
+    /** Restore a tokenizer from {@link toObject}. */
+    loadFromSpec(spec: BPETokenizerSpec): void {
+        const sp = spec.specials;
+        if (sp?.bos) this.bosToken = sp.bos;
+        if (sp?.eos) this.eosToken = sp.eos;
+        if (sp?.pad) this.padToken = sp.pad;
+        if (sp?.unk) this.unkToken = sp.unk;
+        this.loadFromObjects(spec.vocab, spec.merges);
+    }
+
+    /** UTF-8 JSON bytes of {@link toObject} — the `.evermind` tokenizer section. */
+    serialize(): ArrayBuffer {
+        const bytes = new TextEncoder().encode(JSON.stringify(this.toObject()));
+        // Copy into a standalone ArrayBuffer (the view may be a slice of a pool).
+        const out = new ArrayBuffer(bytes.byteLength);
+        new Uint8Array(out).set(bytes);
+        return out;
+    }
+
+    /** Inverse of {@link serialize}. */
+    static deserialize(buffer: ArrayBuffer): BPETokenizer {
+        const spec = JSON.parse(new TextDecoder().decode(new Uint8Array(buffer))) as BPETokenizerSpec;
+        if (!spec || typeof spec.vocab !== 'object' || !Array.isArray(spec.merges)) {
+            throw new Error('BPETokenizer.deserialize: not a serialised BPE tokenizer (expected { vocab, merges })');
+        }
+        const tok = new BPETokenizer();
+        tok.loadFromSpec(spec);
+        return tok;
+    }
+
     /**
      * Fetch and import a Hugging Face `tokenizer.json` by URL — the IDE's
      * "import merges from a repo" path. Resolves once the tokenizer is loaded.

package/src/training/cpu_ops.ts

@@ -0,0 +1,127 @@
+/**
+ * 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));
+}
+
+/** 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]!;
+        }
+    }
+}