@seanhogg/builderforce-memory-engine 2026.7.12 → 2026.9.27

package/src/kernels/slice.ts

@@ -0,0 +1,64 @@
+/**
+ * slice.ts — strided COLUMN slice of a row-major (rows, stride) tensor.
+ *
+ * Every block in this engine projects into one fused tensor and then splits it:
+ * Mamba-1 into (x, z) and (Δ, B, C), Mamba-2/3 into (x, Δ) and (x, B, C),
+ * attention into (Q, K, V). `LINEAR_FORWARD_WGSL` writes its output `(M, N)`
+ * ROW-major, so each of those parts is a contiguous range of COLUMNS **within
+ * every row** — not a contiguous range of the buffer.
+ *
+ * Splitting them with `copyBufferToBuffer` therefore took whole ROWS instead: it
+ * is only the right answer when `M == 1` (batch × seq_len), and for any real
+ * sequence it silently handed the block the wrong tensor. This kernel is the one
+ * correct implementation; every block uses it.
+ */
+
+import { createBindGroup, createUniformBuffer, dispatchKernel, cdiv } from '../utils/gpu_utils.js';
+
+export const COL_SLICE_WGSL: string = /* wgsl */`
+struct SliceParams {
+    rows   : u32,
+    stride : u32,   // source row width
+    offset : u32,   // first column to take
+    width  : u32,   // number of columns to take
+};
+@group(0) @binding(0) var<uniform>             p   : SliceParams;
+@group(0) @binding(1) var<storage, read>       src : array<f32>;
+@group(0) @binding(2) var<storage, read_write> dst : array<f32>;
+
+// Dispatch: (ceil(rows * width / 256), 1, 1)
+@compute @workgroup_size(256)
+fn col_slice(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i >= p.rows * p.width) { return; }
+    let r = i / p.width;
+    let c = i % p.width;
+    dst[i] = src[r * p.stride + p.offset + c];
+}
+`;
+
+/** Entry point name of {@link COL_SLICE_WGSL}. */
+export const COL_SLICE_ENTRY = 'col_slice';
+
+/**
+ * `dst[r, 0..width) = src[r, offset..offset+width)` for every row.
+ *
+ * @param pipeline a pipeline built from {@link COL_SLICE_WGSL} / {@link COL_SLICE_ENTRY}
+ * @param rows     batch * seqLen
+ * @param stride   source row width (the fused projection's output features)
+ */
+export function dispatchColumnSlice(
+    device: GPUDevice,
+    pipeline: GPUComputePipeline,
+    src: GPUBuffer,
+    dst: GPUBuffer,
+    rows: number,
+    stride: number,
+    offset: number,
+    width: number,
+): void {
+    const pBuf = createUniformBuffer(device, new Uint32Array([rows, stride, offset, width]).buffer);
+    const bg = createBindGroup(device, pipeline, [pBuf, src, dst]);
+    dispatchKernel(device, pipeline, bg, [cdiv(rows * width, 256), 1, 1]);
+    pBuf.destroy();
+}

package/src/kernels/weight_update.ts

@@ -79,6 +79,20 @@ fn adamw_update(
 
 // Gradient clipping kernel – clips global gradient norm to max_norm.
 // Run before weight updates.  Two-pass: first compute squared norm, then scale.
+//
+// The accumulator is a REAL atomic. It used to be a plain `array<f32>` written as
+// `norm_sq[0] = norm_sq[0] + local_sq[0]` by the lane-0 invocation of every
+// workgroup — a read-modify-write with no synchronisation, so with more than one
+// workgroup (i.e. any tensor above 256 elements) concurrent workgroups lost each
+// other's contributions and the norm came out too SMALL, under-clipping exactly
+// when clipping matters most. WGSL has no `atomic<f32>`, so the accumulator is an
+// `atomic<u32>` holding the f32 bit pattern and updated with a
+// compare-exchange loop — the standard portable float-atomic-add.
+//
+// Summation ORDER across workgroups is not deterministic (floating-point addition
+// is not associative), so the norm can differ in the last ulp or two between runs.
+// That is fine for a clip threshold and is the accepted trade for correctness;
+// nothing downstream depends on a bit-exact norm.
 export const GRAD_CLIP_WGSL: string = /* wgsl */`
 
 struct ClipParams {
@@ -88,10 +102,23 @@ struct ClipParams {
 
 @group(0) @binding(0) var<uniform>             clip_p  : ClipParams;
 @group(0) @binding(1) var<storage, read_write> grad    : array<f32>;
-@group(0) @binding(2) var<storage, read_write> norm_sq : array<f32>;  // size 1, atomic accumulator
+// size 1 – the f32 bit pattern of the accumulated sum of squares, updated
+// atomically. Zero-initialise it by writing 0.0f (bit pattern 0x00000000).
+@group(0) @binding(2) var<storage, read_write> norm_sq : array<atomic<u32>>;
 
 var<workgroup> local_sq : array<f32, 256>;
 
+// Portable atomic add on an f32 stored as its u32 bit pattern.
+fn atomic_add_f32(value: f32) {
+    var old_bits : u32 = atomicLoad(&norm_sq[0]);
+    loop {
+        let new_bits = bitcast<u32>(bitcast<f32>(old_bits) + value);
+        let res = atomicCompareExchangeWeak(&norm_sq[0], old_bits, new_bits);
+        if (res.exchanged) { break; }
+        old_bits = res.old_value;
+    }
+}
+
 // Pass 1: reduce sum of squares into norm_sq[0]
 @compute @workgroup_size(256, 1, 1)
 fn grad_norm_reduce(
@@ -117,8 +144,7 @@ fn grad_norm_reduce(
     }
 
     if (lid == 0u) {
-        // Non-atomic accumulation (single workgroup assumption for small models)
-        norm_sq[0] = norm_sq[0] + local_sq[0];
+        atomic_add_f32(local_sq[0]);
     }
 }
 
@@ -130,7 +156,7 @@ fn grad_clip_scale(
     let i = gid.x;
     if (i >= clip_p.num_elements) { return; }
 
-    let ns = norm_sq[0];
+    let ns = bitcast<f32>(atomicLoad(&norm_sq[0]));
     if (ns > clip_p.max_norm_sq) {
         let scale = sqrt(clip_p.max_norm_sq / ns);
         grad[i] = grad[i] * scale;