mlx-bun 0.0.4 → 0.0.6

package/src/memory/wikify.ts

@@ -0,0 +1,23 @@
+// mlx-bun memory — synthesis stage 5: WIKIFY (deterministic).  ⚠️ M1 STUB.
+//
+// Non-model editorial cleanup after synthesis: canonicalize [[wikilinks]]
+// (spaces↔underscores, case, resolve to existing stems), repair footnotes, and
+// run the deterministic editorial gate. Non-fatal — a failure here logs and
+// continues; it never corrupts an article. Mirrors lucien's wikify.ts +
+// normalize-wikilinks.ts + normalize-footnotes.ts.
+//
+// See docs/design/memory-system.md → "The nightly pipeline" (wikify/normalize).
+
+import type { SynthesisEvent } from "./pipeline";
+
+export interface WikifyResult {
+  articlesNormalized: number;
+  linksFixed: number;
+}
+
+/** Deterministic editorial gate + link/footnote normalize. STUB. */
+export async function wikify(_onEvent?: (e: SynthesisEvent) => void): Promise<WikifyResult> {
+  // TODO(M1): resolveWikilinkToStem across touched articles (vault.ts already
+  // has the resolver), footnote renumber/repair, {{stub}} marking.
+  throw new Error("memory.wikify is not implemented yet (M1)");
+}

package/src/mlx/metal-kernel.ts

@@ -29,8 +29,16 @@ export interface MetalKernelSpec {
 }
 
 export interface MetalKernelCall {
-  outputs: { shape: number[]; dtype: Dtype }[];
-  grid: [number, number, number];
+  /** Concrete output shapes/dtypes. Provide this OR `outputShapeFn` (the latter
+   *  makes the kernel derive its outputs from the input shapes, like a real MLX
+   *  primitive — required to compose inside an mx.compile'd / shapeless closure,
+   *  where the wrapper must answer "given these input shapes, what comes out?"). */
+  outputs?: { shape: number[]; dtype: Dtype }[];
+  /** Derive output shapes/dtypes from the (current/traced) input shapes. */
+  outputShapeFn?: (inputs: MlxArray[]) => { shape: number[]; dtype: Dtype }[];
+  /** Concrete launch grid, OR derive it from inputs (so it tracks shape under
+   *  compile replay the same way `outputShapeFn` tracks output shapes). */
+  grid: [number, number, number] | ((inputs: MlxArray[]) => [number, number, number]);
   threadGroup: [number, number, number];
   templateInts?: Record<string, number>;
   templateDtypes?: Record<string, Dtype>;
@@ -59,31 +67,50 @@ export class MetalKernel {
   apply(inputs: MlxArray[], call: MetalKernelCall): MlxArray[] {
     if (this.#disposed) throw new Error("MetalKernel used after dispose");
     const cfg = C.mlx_fast_metal_kernel_config_new();
+    // bun:ffi ptr() does NOT retain its argument, so any JS buffer whose pointer
+    // crosses an FFI call below is anchored here until that call returns —
+    // otherwise GC could free it mid-call (the lifetime class behind the
+    // flash-CCE pin bug). Pure-additive: holds references, changes no logic.
+    const keepAlive: unknown[] = [];
     try {
-      for (const o of call.outputs) {
+      // Derive outputs + grid from the (current) input shapes when a function is
+      // given — so the SAME kernel object answers correctly for any input shape
+      // during a compile/replay trace, instead of baking in one call's shapes.
+      const outputs = call.outputShapeFn ? call.outputShapeFn(inputs) : call.outputs;
+      if (!outputs) throw new Error("metal_kernel: provide outputs or outputShapeFn");
+      const grid = typeof call.grid === "function" ? call.grid(inputs) : call.grid;
+      for (const o of outputs) {
         const shape = new Int32Array(o.shape);
+        keepAlive.push(shape);
         if (C.mlx_fast_metal_kernel_config_add_output_arg(cfg, ptr(shape), BigInt(o.shape.length), o.dtype) !== 0)
           throw new Error(`metal_kernel add_output_arg failed: ${takeMlxError() ?? ""}`);
       }
-      C.mlx_fast_metal_kernel_config_set_grid(cfg, call.grid[0], call.grid[1], call.grid[2]);
+      C.mlx_fast_metal_kernel_config_set_grid(cfg, grid[0], grid[1], grid[2]);
       C.mlx_fast_metal_kernel_config_set_thread_group(cfg, call.threadGroup[0], call.threadGroup[1], call.threadGroup[2]);
       if (call.initValue !== undefined)
         C.mlx_fast_metal_kernel_config_set_init_value(cfg, call.initValue);
       for (const [k, v] of Object.entries(call.templateInts ?? {})) {
         const kb = Buffer.from(k + "\0", "utf8");
+        keepAlive.push(kb);
         C.mlx_fast_metal_kernel_config_add_template_arg_int(cfg, ptr(kb), v);
       }
       for (const [k, v] of Object.entries(call.templateDtypes ?? {})) {
         const kb = Buffer.from(k + "\0", "utf8");
+        keepAlive.push(kb);
         C.mlx_fast_metal_kernel_config_add_template_arg_dtype(cfg, ptr(kb), v);
       }
 
       const handles = new BigUint64Array(inputs.map((a) => a.handle));
+      keepAlive.push(handles);
       const inVec = C.mlx_vector_array_new_data(ptr(handles), BigInt(inputs.length));
       const slot = new BigUint64Array([C.mlx_vector_array_new()]);
+      keepAlive.push(slot);
       const slotPtr = ptr(slot);
       const status = C.mlx_fast_metal_kernel_apply(slotPtr, this.#kernel, inVec, cfg, call.stream ?? gpuStream);
       const outVec = read.u64(slotPtr, 0);
+      // anchor: all ptr()'d JS buffers above are done being read by native code by
+      // this point; touching keepAlive keeps the JIT from freeing them earlier.
+      if (keepAlive.length < 0) throw new Error("unreachable");
       C.mlx_vector_array_free(inVec);
       if (status !== 0) {
         C.mlx_vector_array_free(outVec);

package/src/model/flash-attention.ts

@@ -29,10 +29,12 @@ function forwardTiles(headDim: number): [number, number] {
   throw new Error(`unsupported head_dim=${headDim}`);
 }
 
-/** uint32 [n] array (matches optiq's shape/causal input dtype). */
+/** uint32 [n] array (matches optiq's shape/causal input dtype). COPY into an
+ *  mlx-owned array (not the pinning fromView) — these tiny kernel-arg arrays are
+ *  disposed before the lazy kernel evaluates; pinning leaks the host buffer per
+ *  call (see flash-cce.ts u32 for the full diagnosis). */
 function u32(values: number[]): MlxArray {
-  const u = new Uint32Array(values);
-  return MlxArray.fromView(new Uint8Array(u.buffer.slice(0)), [values.length], Dtype.uint32);
+  return MlxArray.fromBytesCopy(new Uint8Array(new Uint32Array(values).buffer), [values.length], Dtype.uint32);
 }
 
 // Forward shader — verbatim from optiq flash_attention_metal.py _KERNEL_SOURCE.
@@ -406,7 +408,13 @@ for (uint kj = 0u; kj < n_kv_tiles; ++kj) {
     }
     threadgroup_barrier(mem_flags::mem_threadgroup);
 
-    if (is_causal && kv_base_row > q_row_global) {
+    // Skip a KV tile only when it is entirely beyond this Q tile's last row
+    // (fully causally masked). The condition is UNIFORM across the threadgroup
+    // (q_tile_idx·BQ is the same for every thread), so the barrier below is hit
+    // by all threads — using the per-thread q_row_global here would diverge the
+    // barrier (undefined behavior). The per-element mask below still handles the
+    // partial diagonal tile.
+    if (is_causal && kv_base_row >= (q_tile_idx * BQ + BQ)) {
         threadgroup_barrier(mem_flags::mem_threadgroup);
         continue;
     }
@@ -533,16 +541,12 @@ export function flashBackward(
       grid: [nQTiles * bq, Hq, B],
       threadGroup: [bq, 1, 1],
     });
-    // The dKV kernel emits dK with its (Tkv, D) axes transposed in the buffer
-    // (verified: dK[i][j] == correct[j][i] at T=D, a [Tkv,D]↔[D,Tkv] swap in
-    // general — scripts/flash-dk-debug.ts JS-FIX 0.07% vs autograd). dV is
-    // unaffected. Recover the correct layout: transpose → materialize → reshape.
-    const dkFixed = ops.reshape(
-      ops.contiguous(ops.transposeAxes(dkv[0]!, [0, 1, 3, 2])),
-      [B, Hkv, Tkv, D],
-    );
-    dkv[0]!.dispose();
-    return [dq[0]!, dkFixed, dkv[1]!];
+    // dK is written by the kernel in the correct [B,Hkv,Tkv,D] layout
+    // (dK_j[d] → [kv_row][d]); no post-transpose. (An earlier "fix" transposed
+    // it on the false belief the buffer was swapped — that transpose is only the
+    // identity at Tkv==D and corrupted dK for Tkv≠D; removed after
+    // scripts/experiments/flash-dkv-debug.ts showed the raw output is correct.)
+    return [dq[0]!, dkv[0]!, dkv[1]!];
   } finally {
     Dvec.dispose();
     shape.dispose();
@@ -607,11 +611,25 @@ export function flashAttention(
 }
 
 // ---------------------------------------------------------------------------
-// Training-attention routing. ops.sdpa's vjp gives a wrong dK (it's an
-// inference op), so the training forward must use a known-correct attention:
-// the flash kernel for causal/full layers (memory-efficient, validated), and a
-// manual materialized SDPA for sliding-window (array-mask) layers. Set by the
-// trainer; null = inference (use ops.sdpa as before).
+// Training-attention routing (opt-in; default null = ops.sdpa).
+//
+// Parity oracle, made explicit:
+//  - DEFAULT (ops.sdpa) is the L1 path: mlx-lm's tuner differentiates the same
+//    mx.fast.scaled_dot_product_attention, so training through it is bit-faithful
+//    to mlx-lm. Exact dQ/dK/dV (finite-difference cross-checked), fast, but
+//    O(L²) backward memory (materializes the score matrix).
+//  - This flash kernel is the L2 path: a port of mlx-optiq's
+//    flash_attention_metal (O(L) backward memory). Its parity oracle is OPTIQ,
+//    not finite differences. Confirmed: optiq's flash dK == ops.sdpa to f16
+//    (scripts/flash-optiq-check.py, rel 0.0%), and mlx-bun's flash == ops.sdpa
+//    (tests/flash-attention.test.ts) ⟹ mlx-bun flash == optiq flash (L2 parity).
+//
+// Two real port bugs were fixed to reach that parity (the original diverged
+// ~100% from optiq on dK): a spurious dK transpose in flashBackward (corrupted
+// dK for Tkv≠D), and a per-thread (divergent) threadgroup_barrier in the dQ
+// causal tile-skip. flash is opt-in (MLX_BUN_TRAIN_ATTN=flash) for memory-bound
+// long context; it is ~30× slower than ops.sdpa. flashSupported gates
+// head_dim/seq; sliding-window (array-mask) layers fall back to manualSdpa.
 // ---------------------------------------------------------------------------
 
 let trainAttnMode: "flash" | "manual" | null = null;

package/src/model/fused-geglu-kernel.ts

@@ -16,6 +16,8 @@
 import { MlxArray } from "../mlx/array";
 import { Dtype } from "../mlx/ffi";
 import { MetalKernel } from "../mlx/metal-kernel";
+import { CustomVjp } from "../mlx/custom-vjp";
+import * as ops from "../mlx/ops";
 
 /** ON by default: the match-compat kernel is BIT-EXACT with the spelled-out
  *  path (kl --decode = 0 on e4b), so per the bit-parity-floor principle it
@@ -28,6 +30,24 @@ export function fusedGeluEnabled(): boolean {
 /** Dispatch counter (a gate test asserts the kernel actually ran). */
 export let fusedGeluCalls = 0;
 
+// Training-mode flag: when set, model GeGLU sites use the DIFFERENTIABLE wrapper
+// (Metal-kernel forward + hand-derived vjp) instead of the plain `fusedGeglu`
+// (a bare CustomKernel with no vjp — feeding it to autograd throws
+// "[Primitive::vjp] Not implemented for CustomKernel"). The differentiable
+// wrapper also recomputes the gelu in the backward from the primal `a` instead
+// of retaining the ~9 forward intermediates, so it lowers the MLP activation
+// memory the backward holds — the dominant resident at long context under
+// ops.sdpa (which never materializes the O(T²) scores). Set by the trainer
+// around the train loop; cleared in its finally. Off → inference behaviour
+// (plain `fusedGeglu`) is byte-for-byte unchanged.
+let _trainingGeglu = false;
+export function setFusedGeluTraining(on: boolean): void {
+  _trainingGeglu = on;
+}
+export function fusedGeluTraining(): boolean {
+  return _trainingGeglu;
+}
+
 // Pointwise: one thread per element. gelu_approx =
 //   0.5 x (1 + tanh(√(2/π)·(x + 0.044715 x³)))  — composed in f32.
 const SOURCE = String.raw`
@@ -85,3 +105,94 @@ export function fusedGeglu(a: MlxArray, b: MlxArray): MlxArray {
   });
   return out!;
 }
+
+// ---------------------------------------------------------------------------
+// Differentiable wrapper: Metal kernel forward + pure-MLX backward
+// ---------------------------------------------------------------------------
+
+// GELU tanh-approx constants
+const SQRT_2_PI = 0.7978845608028654;
+const C_CUBE = 0.044715;
+
+/** Compute gelu_approx(a) = 0.5 * a * (1 + tanh(sqrt(2/pi) * (a + 0.044715*a^3)))
+ *  and gelu'(a) using the MLX ops layer. All intermediates are disposed
+ *  before returning; caller owns the two returned arrays. */
+function geluAndGrad(a: MlxArray): { gelu: MlxArray; geluGrad: MlxArray } {
+  // a^3 = a * a * a
+  const a2 = ops.mul(a, a);
+  const a3 = ops.mul(a2, a);
+  // inner = a + 0.044715 * a^3
+  const cx3 = ops.mulScalar(a3, C_CUBE);
+  const inner = ops.add(a, cx3);
+  // z = sqrt(2/pi) * inner
+  const z = ops.mulScalar(inner, SQRT_2_PI);
+  // t = tanh(z)
+  const t = ops.tanh(z);
+  // gelu = 0.5 * a * (1 + t)
+  const one = ops.scalarLike(1, a);
+  const t1 = ops.add(one, t);
+  const halfA = ops.mulScalar(a, 0.5);
+  const gelu = ops.mul(halfA, t1);
+
+  // gelu'(a) = 0.5*(1+t) + 0.5*a*(1-t^2)*dz/da
+  // dz/da = sqrt(2/pi) * (1 + 3*0.044715*a^2)
+  // 1 - t^2 = sech^2(z)
+  const halfT1 = ops.mulScalar(t1, 0.5);            // 0.5*(1+t)
+  const t2 = ops.mul(t, t);                         // t^2
+  const sech2 = ops.sub(one, t2);                   // 1 - t^2
+  const halfASech2 = ops.mul(halfA, sech2);         // 0.5*a*(1-t^2)
+  // dz/da
+  const a2scaled = ops.mulScalar(a2, 3 * C_CUBE);   // 3*0.044715*a^2
+  const dzda = ops.add(one, a2scaled);               // 1 + 3*0.044715*a^2
+  const dzdaScaled = ops.mulScalar(dzda, SQRT_2_PI);// sqrt(2/pi)*dz_unnorm
+  // chain: 0.5*a*(1-t^2)*dz/da
+  const chainTerm = ops.mul(halfASech2, dzdaScaled);
+  const geluGrad = ops.add(halfT1, chainTerm);       // 0.5*(1+t) + chain
+
+  // dispose all intermediates
+  for (const arr of [a2, a3, cx3, inner, z, t, one, t1, halfA, halfT1, t2, sech2,
+                      halfASech2, a2scaled, dzda, dzdaScaled, chainTerm]) {
+    arr.dispose();
+  }
+
+  return { gelu, geluGrad };
+}
+
+let _differentiableKernel: CustomVjp | null = null;
+
+function getDifferentiableKernel(): CustomVjp {
+  if (!_differentiableKernel) {
+    _differentiableKernel = new CustomVjp(
+      // forward: use the Metal kernel
+      ([a, b]: MlxArray[]) => [fusedGeglu(a!, b!)],
+      // backward
+      (primals: MlxArray[], cotangents: MlxArray[], _outputs: MlxArray[]) => {
+        const [a, b] = primals as [MlxArray, MlxArray];
+        const [dc] = cotangents as [MlxArray];
+
+        const { gelu, geluGrad } = geluAndGrad(a);
+
+        // grad_b = dc * gelu(a)
+        const grad_b = ops.mul(dc, gelu);
+
+        // grad_a = dc * b * gelu'(a)
+        const dcb = ops.mul(dc, b);
+        const grad_a = ops.mul(dcb, geluGrad);
+
+        gelu.dispose();
+        geluGrad.dispose();
+        dcb.dispose();
+
+        return [grad_a, grad_b];
+      },
+    );
+  }
+  return _differentiableKernel;
+}
+
+/** gelu_approx(a) · b with a hand-derived vjp. Uses the Metal kernel forward,
+ *  MLX-ops backward. For training only (inference uses the non-differentiable
+ *  fusedGeglu). */
+export function fusedGegluDifferentiable(a: MlxArray, b: MlxArray): MlxArray {
+  return getDifferentiableKernel().apply([a, b])[0]!;
+}

package/src/model/gemma4-base.ts

@@ -46,6 +46,31 @@ export interface LoraWeights {
  *  (PLAN Phase 8 decision). Set/restored by generate(). */
 export class LoraState {
   active: string[] = [];
+  /** Training-only LoRA-input dropout. `rate` is the drop probability; `seed` is
+   *  set per micro-step by the trainer and is CONSTANT across that step's forward
+   *  and any recompute (segmented / gradient-checkpoint), so each layer's mask —
+   *  keyed by (seed, the linear's dropoutId) — is reproduced exactly in the
+   *  backward. `seed === null` (the default, and for inference) disables it. */
+  dropoutRate = 0;
+  dropoutSeed: number | null = null;
+}
+
+/** Inverted LoRA-input dropout, keyed by (seed, id) so it is deterministic —
+ *  the same (seed, id, shape) reproduces the mask, which is what makes the
+ *  segmented / gradient-checkpoint recompute correct. kept ⇒ x/(1-p),
+ *  dropped ⇒ 0 (preserves the expectation). Caller owns the result. */
+export function loraInputDropout(x: MlxArray, p: number, seed: number, id: number): MlxArray {
+  const key = ops.randomKey(BigInt(seed) * 100003n + BigInt(id));
+  const u = ops.randomUniform(x.shape, Dtype.float32, 0, 1, key); // [shape] in [0,1)
+  key.dispose();
+  const pArr = ops.scalarLike(p, u);
+  const keep = ops.less(pArr, u); // u > p ⇒ keep (bool [shape])
+  const scaleX = ops.scalarLike(1 / (1 - p), x);
+  const zeroX = ops.scalarLike(0, x);
+  const mask = ops.where(keep, scaleX, zeroX); // x.dtype
+  const xd = ops.mul(x, mask);
+  for (const a of [u, pArr, keep, scaleX, zeroX, mask]) a.dispose();
+  return xd;
 }
 
 export class QuantizedLinear {
@@ -53,6 +78,9 @@ export class QuantizedLinear {
   adapters: Map<string, LoraWeights> | null = null;
   /** Shared per-model active state (wired by AdapterManager.mount). */
   loraState: LoraState | null = null;
+  /** Stable per-target index for keying training-only LoRA dropout (set by
+   *  attachForTraining). Gives each adapted linear an independent dropout mask. */
+  dropoutId = 0;
 
   constructor(
     readonly w: MlxArray,
@@ -92,10 +120,19 @@ export class QuantizedLinear {
     // cast form is what the adapters were trained behind.)
     const st = this.loraState;
     if (st && st.active.length > 0 && this.adapters && this.adapters.size > 0) {
+      // Training-only LoRA-input dropout (PEFT applies dropout to x before A;
+      // the base quantized path is untouched). Keyed by (seed, dropoutId) so the
+      // backward recompute reproduces the exact mask.
+      let xLora = x;
+      let xDrop: MlxArray | null = null;
+      if (st.dropoutSeed !== null && st.dropoutRate > 0) {
+        xDrop = loraInputDropout(x, st.dropoutRate, st.dropoutSeed, this.dropoutId);
+        xLora = xDrop;
+      }
       for (const id of st.active) {
         const lw = this.adapters.get(id);
         if (!lw) continue;
-        const xa = ops.matmul(x, lw.a);
+        const xa = ops.matmul(xLora, lw.a);
         const z = ops.matmul(xa, lw.b);
         xa.dispose();
         const zs = ops.mulScalar(z, lw.scale);
@@ -105,6 +142,7 @@ export class QuantizedLinear {
         out = disposing(out, ops.add(out, zc));
         zc.dispose();
       }
+      xDrop?.dispose();
     }
     return out;
   }