mlx-bun 0.0.9 → 0.0.11

package/src/memory/vault.ts

@@ -218,6 +218,17 @@ export async function readArticle(root: string, article: string): Promise<{ path
   return { path, content: await readFile(path, "utf8") };
 }
 
+/** Vault-relative path (`articles/<stem>.md` or `Reference/<stem>.md`) for an
+ *  existing article/reference doc — the safe argv piece for git plumbing
+ *  (articleHistory/articleDiff), resolved via the same name normalization/
+ *  existence check as readArticle so a caller can't smuggle a path-traversal
+ *  or shell-metacharacter name through. Throws if the doc doesn't exist. */
+export async function resolveArticleRelPath(root: string, article: string): Promise<string> {
+  const doc = normalizeMemoryDocId(article);
+  await resolveArticlePath(root, article); // existence check, reuses the same normalization
+  return doc.kind === "reference" ? join("Reference", `${doc.stem}.md`) : join("articles", `${doc.stem}.md`);
+}
+
 /** Filename stems for every `*.md` under articles/, sorted alphabetically. */
 export async function listArticles(root: string): Promise<string[]> {
   let names: string[];
@@ -1013,6 +1024,69 @@ function runGit(args: string[], cwd: string): Promise<boolean> {
   });
 }
 
+// ---- read-only git plumbing (per-article history/diff) ----------------
+//
+// Used by the memory REST surface (GET /api/memory/history|diff) to render
+// "watch it self-heal" — the article's commit history and a specific
+// commit's diff. Read-only plumbing only (log/show), args always an
+// explicit argv array via Bun.spawn (never a shell string), and callers
+// MUST validate `rev`/`name` before calling these (see server.ts) — this
+// layer trusts its `relPath`/`rev` inputs are already safe.
+
+export interface ArticleHistoryEntry {
+  hash: string;
+  date: string;
+  subject: string;
+}
+
+/** `git log` for one file, newest first, argv-only (never a shell string).
+ *  Returns [] when the vault isn't a git repo, the file has no history, or
+ *  git is unavailable. `relPath` must already be a repo-relative path (no
+ *  `..`, not absolute) — callers validate before calling. */
+export async function articleHistory(root: string, relPath: string, limit = 50): Promise<ArticleHistoryEntry[]> {
+  if (!(await pathExists(join(root, ".git")))) return [];
+  try {
+    const proc = Bun.spawn(
+      ["git", "log", `--max-count=${limit}`, "--date=short", "--pretty=format:%H%x1f%ad%x1f%s%x1e", "--", relPath],
+      { cwd: root, stdout: "pipe", stderr: "ignore" },
+    );
+    const out = await new Response(proc.stdout).text();
+    if ((await proc.exited) !== 0) return [];
+    return out
+      .split("\x1e")
+      .map((rec) => rec.trim())
+      .filter((rec) => rec.length > 0)
+      .map((rec) => {
+        const [hash = "", date = "", subject = ""] = rec.split("\x1f");
+        return { hash, date, subject };
+      })
+      .filter((e) => e.hash.length > 0);
+  } catch {
+    return [];
+  }
+}
+
+/** `git show <rev> -- <relPath>` — the diff for one commit touching one file.
+ *  `rev` must already be validated as a hex commit hash by the caller (never
+ *  a ref expression like `HEAD~1` or `main^`, which could be abused as an
+ *  argument-injection vector); `relPath` must already be repo-relative.
+ *  Returns null when the vault isn't a git repo or the commit/diff can't be
+ *  produced. */
+export async function articleDiff(root: string, rev: string, relPath: string): Promise<string | null> {
+  if (!(await pathExists(join(root, ".git")))) return null;
+  try {
+    const proc = Bun.spawn(
+      ["git", "show", rev, "--pretty=format:", "--", relPath],
+      { cwd: root, stdout: "pipe", stderr: "ignore" },
+    );
+    const out = await new Response(proc.stdout).text();
+    if ((await proc.exited) !== 0) return null;
+    return out;
+  } catch {
+    return null;
+  }
+}
+
 export interface SetupResult {
   root: string;
   created: string[];

package/src/mlx/array.ts

@@ -1,8 +1,8 @@
 // MlxArray: ownership wrapper around an mlx_array handle.
 // Explicit .dispose() is the contract; a FinalizationRegistry backstop
-// frees leaked handles on GC (verified in spikes/phase0-memory.ts).
+// frees leaked handles on GC (verified in lab/spikes/phase0-memory.ts).
 
-import { JSCallback, ptr, toArrayBuffer } from "bun:ffi";
+import { dlopen, ptr, toArrayBuffer } from "bun:ffi";
 import { C, Dtype, DTYPE_NAMES, type MlxHandle, optInt, outArray, takeMlxError } from "./ffi";
 import type { SafetensorsDtype } from "../safetensors";
 
@@ -24,17 +24,39 @@ export const SAFETENSORS_TO_MLX: Record<SafetensorsDtype, Dtype> = {
 
 // --- zero-copy pinning ---------------------------------------------------
 // Buffers handed to mlx via new_data_managed_payload must outlive the mlx
-// array. We pin the JS view in a map keyed by a payload id; the shared dtor
-// callback unpins when mlx drops its last reference. threadsafe: mlx may
-// release buffers from its eval thread.
+// array. fromView pins the JS view in this map (a GC root, nothing more);
+// release is EXPLICIT and JS-thread-only via unpinHostBuffer() — never an
+// mlx dtor. See the hazard note below.
 
 const pinned = new Map<number, Uint8Array>();
 let nextPinId = 1;
 
-const unpinCallback = new JSCallback(
-  (payloadId: number) => { pinned.delete(payloadId); },
-  { args: ["u64"], returns: "void", threadsafe: true },
-);
+// HAZARD (root cause of the 2026-07-06 restart-restore hang): mlx releases
+// buffer Data wherever the LAST shared_ptr drops — often the Metal
+// COMPLETION thread (gpu::eval retains buffers until the command buffer
+// finishes, i.e. past dispose()). A bun:ffi JSCallback dtor invoked from
+// that thread deadlocks when the JS thread is inside a blocking FFI eval
+// (completion waits on JS, JS waits on completion) and SIGTRAPs mid-GC.
+// So NO array here ever gets a JS dtor: fromPointer AND fromView both hand
+// mlx the native no-op dtor below (payload 0), and host-buffer lifetime is
+// the CALLER's contract, released only from the JS thread (unpinHostBuffer)
+// or never (process-lifetime mmaps: expert-offload).
+
+// Native no-op dtor: payload is always 0, and libc free(NULL) is defined to
+// do nothing — a dtor mlx can safely call from ANY thread. The pointed-to
+// memory's lifetime is the CALLER's contract (weight mmaps live for the
+// process; restored KV is COPIED at load since 2026-07-07 — see kv-store.ts
+// loadKvCache — so no restore mapping needs pinning anymore). The address
+// comes via dlsym because
+// bun:ffi's symbol .ptr returns the pointer bit-cast to float64
+// (lab/repro/bun-ffi-f64) — passing that back truncates to NULL.
+const libcDlsym = dlopen("/usr/lib/libSystem.B.dylib", {
+  dlsym: { args: ["u64", "ptr"], returns: "u64" },
+});
+const FREE_NAME = Buffer.from("free\0");
+const RTLD_DEFAULT = 0xfffffffffffffffen;
+const freeFnAddr = Number(libcDlsym.symbols.dlsym(RTLD_DEFAULT, ptr(FREE_NAME)));
+if (!freeFnAddr) throw new Error("dlsym(free) failed — zero-copy arrays need a native no-op dtor");
 
 export function pinnedBufferCount(): number {
   return pinned.size;
@@ -50,31 +72,61 @@ export class MlxArray {
   #handle: MlxHandle;
   #disposed = false;
   #token = {};
+  #pinId = 0; // fromView only: key into the pinned-view map (0 = no pin)
 
   constructor(handle: MlxHandle) {
     this.#handle = handle;
     registry.register(this, handle, this.#token);
   }
 
-  /** Zero-copy: wrap an existing buffer (e.g. an mmap view). The view is
-   *  pinned until mlx releases it. */
+  /** Zero-copy: wrap an existing JS buffer. The view is pinned (GC-rooted)
+   *  in a process-side map; mlx gets the NATIVE no-op dtor (free(NULL)) —
+   *  never a JS callback, so the last release may happen on the Metal
+   *  completion thread without deadlock (see hazard note above). Buffer
+   *  lifetime is therefore the CALLER's contract:
+   *   - mmap-backed views (toArrayBuffer aliases): the mapping owns the
+   *     memory; prefer fromPointer, which skips the pin entirely.
+   *   - JS-heap buffers: the pin keeps the buffer alive; call
+   *     unpinHostBuffer() only after a JS-side sync point proves mlx is
+   *     done with it (everything depending on the array evaluated AND the
+   *     array disposed), or never (process-lifetime, e.g. short scripts). */
   static fromView(view: Uint8Array, shape: number[], dtype: Dtype): MlxArray {
     const sb = shapeBuf(shape);
     const id = nextPinId++;
     pinned.set(id, view);
     const handle = C.mlx_array_new_data_managed_payload(
-      ptr(view), ptr(sb), shape.length, dtype, id, unpinCallback.ptr,
+      ptr(view), ptr(sb), shape.length, dtype, 0, freeFnAddr,
     );
-    return new MlxArray(handle);
+    const arr = new MlxArray(handle);
+    arr.#pinId = id;
+    return arr;
+  }
+
+  /** Release the GC root fromView placed on this array's host buffer.
+   *  JS-thread-only by construction (it's a plain method) — this is the
+   *  replacement for the deadlocking JSCallback unpin dtor. Idempotent;
+   *  no-op for non-fromView arrays. Caller contract: only call once mlx can
+   *  no longer touch the buffer (see fromView doc), or when the memory is
+   *  owned elsewhere (an mmap kept mapped past every use). */
+  unpinHostBuffer(): void {
+    if (this.#pinId !== 0) {
+      pinned.delete(this.#pinId);
+      this.#pinId = 0;
+    }
   }
 
   /** Zero-copy: wrap a raw pointer (e.g. into an mmap'd weight shard).
-   *  The caller guarantees the memory outlives the array — weight mmaps
-   *  live for the process, so the dtor is a no-op unpin of id 0. */
+   *  The caller guarantees the memory outlives EVERY mlx
+   *  reference to the array — including GPU command buffers that retain the
+   *  buffer past dispose() — i.e. mmaps backing these arrays stay mapped for
+   *  the process. The dtor is free(NULL) (payload 0): native + no-op, so
+   *  mlx may run it from the Metal completion thread without touching the
+   *  JS runtime (the old JSCallback dtor deadlocked serving / SIGTRAPed
+   *  under GC when that thread held the last reference — 2026-07-06). */
   static fromPointer(dataPtr: number, shape: number[], dtype: Dtype): MlxArray {
     const sb = shapeBuf(shape);
     const handle = C.mlx_array_new_data_managed_payload(
-      dataPtr, ptr(sb), shape.length, dtype, 0, unpinCallback.ptr,
+      dataPtr, ptr(sb), shape.length, dtype, 0, freeFnAddr,
     );
     return new MlxArray(handle);
   }
@@ -108,6 +160,27 @@ export class MlxArray {
     );
   }
 
+  /** Graph-free LEAF with this array's exact bytes — ONE copy, straight from
+   *  the (evaluated) mlx buffer into a fresh mlx-owned array, no JS-heap
+   *  intermediate (`fromBytesCopy(rawBytes())` is TWO copies: device→JS heap
+   *  →new leaf). mlx_array_new_data copies synchronously, so the source only
+   *  has to outlive this call. The buffer is read LINEARLY — the caller must
+   *  pass a row-major contiguous array (ops.contiguous first for views), same
+   *  contract as rawBytes. Segmented-backward boundary detach hot path. */
+  detachCopy(): MlxArray {
+    this.eval();
+    const dt = this.dtype;
+    const p =
+      dt === Dtype.float32 ? C.mlx_array_data_float32(this.handle)
+      : dt === Dtype.float16 ? C.mlx_array_data_float16(this.handle)
+      : dt === Dtype.bfloat16 ? C.mlx_array_data_bfloat16(this.handle)
+      : dt === Dtype.uint32 ? C.mlx_array_data_uint32(this.handle)
+      : null;
+    if (p === null) throw new Error(`detachCopy: unsupported dtype ${this.dtypeName}`);
+    const sb = shapeBuf(this.shape);
+    return new MlxArray(C.mlx_array_new_data(p as never, ptr(sb), this.ndim, dt));
+  }
+
   get handle(): MlxHandle {
     if (this.#disposed) throw new Error("MlxArray used after dispose");
     return this.#handle;
@@ -117,7 +190,7 @@ export class MlxArray {
     return Number(C.mlx_array_ndim(this.handle));
   }
 
-  // NOTE on the DFG stale-read bug (repro/bun-ffi-f64/ISSUE.md): the
+  // NOTE on the DFG stale-read bug (lab/repro/bun-ffi-f64/ISSUE.md): the
   // toArrayBuffer readbacks below (shape/rawBytes/toFloat32) are safe.
   // The hazard is reading a *pre-existing* typed array after an FFI call
   // wrote through its pointer — the JIT forwards a stale value from before
@@ -162,6 +235,18 @@ export class MlxArray {
 
   /** Raw bytes of the evaluated array (copy). bf16/f16/f32 only. */
   rawBytes(): Uint8Array {
+    return this.rawBytesView().slice();
+  }
+
+  /** ZERO-COPY view of the evaluated array's bytes — aliases the mlx
+   *  buffer directly, no JS-heap copy. Valid only while THIS array is
+   *  alive and unmutated: do not retain past dispose(). Built for
+   *  bulk-write paths (kv-store's streamed writer) where per-tensor
+   *  rawBytes() copies pile up in the JS heap faster than GC collects
+   *  them — the residual A7 whole-entry RSS spike (2026-07-07). The
+   *  array must be contiguous (ops.contiguous first for views), same
+   *  contract as rawBytes. */
+  rawBytesView(): Uint8Array {
     this.eval();
     const dt = this.dtype;
     const p =
@@ -169,9 +254,32 @@ export class MlxArray {
       : dt === Dtype.float16 ? C.mlx_array_data_float16(this.handle)
       : dt === Dtype.bfloat16 ? C.mlx_array_data_bfloat16(this.handle)
       : dt === Dtype.uint32 ? C.mlx_array_data_uint32(this.handle)
+      : dt === Dtype.uint8 ? C.mlx_array_data_uint8(this.handle)
+      : dt === Dtype.int8 ? C.mlx_array_data_int8(this.handle)
       : null;
-    if (p === null) throw new Error(`rawBytes: unsupported dtype ${this.dtypeName}`);
-    return new Uint8Array(toArrayBuffer(p!, 0, this.nbytes)).slice();
+    if (p === null) throw new Error(`rawBytesView: unsupported dtype ${this.dtypeName}`);
+    return new Uint8Array(toArrayBuffer(p!, 0, this.nbytes));
+  }
+
+  /** Read back an INTEGER array without enqueueing a cast kernel: eval +
+   *  direct buffer read for uint32/int32. A cast (`astype`) is a NEW op that
+   *  queues BEHIND everything already dispatched on the GPU stream — on the
+   *  batched decode hot path, reading the pipeline register via toFloat32
+   *  stalled the "overlapped" token read for a FULL step of GPU work every
+   *  token (the core of the Phase-0 B=1 gap; unified-engine plan Phase 2).
+   *  Non-integer dtypes fall back to the rounded float readback. */
+  toIntTokens(): number[] {
+    if (this.dtype === Dtype.uint32) {
+      this.eval();
+      const p = C.mlx_array_data_uint32(this.handle);
+      return [...new Uint32Array(toArrayBuffer(p!, 0, this.size * 4))];
+    }
+    if (this.dtype === Dtype.int32) {
+      this.eval();
+      const p = C.mlx_array_data_int32(this.handle);
+      return [...new Int32Array(toArrayBuffer(p!, 0, this.size * 4))];
+    }
+    return [...this.toFloat32()].map((x) => Math.round(x));
   }
 
   /** Read back as float32 (casts on GPU if needed, then copies out). */

package/src/mlx/autograd.ts

@@ -1,6 +1,6 @@
 // Native autograd via mlx-c value_and_grad over a JS loss closure.
 //
-// Encapsulates the exact FFI lifecycle proven in spikes/phase-train-vag.ts:
+// Encapsulates the exact FFI lifecycle proven in lab/spikes/phase-train-vag.ts:
 //   JS loss closure (JSCallback) → mlx_closure_new_func_payload →
 //   mlx_value_and_grad → mlx_closure_value_and_grad_apply → read grads.
 //
@@ -45,7 +45,7 @@ export class ValueAndGrad {
           // Read the primals out of the input vector_array. mlx_vector_array_get
           // returns a fresh (ref-counted) handle, so disposing the wrapper
           // afterward is correct — it drops our reference, not mlx's graph
-          // reference (proven in spikes/phase-train-vag.ts).
+          // reference (proven in lab/spikes/phase-train-vag.ts).
           const n = Number(C.mlx_vector_array_size(inVec));
           const primals: MlxArray[] = [];
           for (let i = 0; i < n; i++) {

package/src/mlx/ffi.ts

@@ -15,7 +15,7 @@ import { existsSync } from "node:fs";
 import { dirname, join } from "node:path";
 import { nativePackDir } from "../native-pack";
 
-const { ptr: P, i32, u64, f32, cstring } = FFIType;
+const { ptr: P, i32, u64, f32, f64, cstring } = FFIType;
 
 /** libmlxc resolution, in order: explicit env override → next to the
  *  executable (the `bun build --compile` sidecar layout, where
@@ -57,6 +57,9 @@ export const C = dlopen(LIBMLXC_PATH, {
   mlx_array_data_float16: { args: [u64], returns: P },
   mlx_array_data_bfloat16: { args: [u64], returns: P },
   mlx_array_data_uint32: { args: [u64], returns: P },
+  mlx_array_data_int32: { args: [u64], returns: P },
+  mlx_array_data_uint8: { args: [u64], returns: P },
+  mlx_array_data_int8: { args: [u64], returns: P },
   mlx_array_item_float32: { args: [P, u64], returns: i32 },
   // streams
   mlx_default_gpu_stream_new: { args: [], returns: u64 },
@@ -142,6 +145,7 @@ export const C = dlopen(LIBMLXC_PATH, {
   mlx_negative: { args: [P, u64, u64], returns: i32 },
   mlx_argsort_axis: { args: [P, u64, i32, u64], returns: i32 },
   mlx_argpartition_axis: { args: [P, u64, i32, i32, u64], returns: i32 },
+  mlx_arange: { args: [P, f64, f64, f64, i32, u64], returns: i32 },
   mlx_take_along_axis: { args: [P, u64, u64, i32, u64], returns: i32 },
   mlx_put_along_axis: { args: [P, u64, u64, u64, i32, u64], returns: i32 },
   mlx_cumsum: { args: [P, u64, i32, FFIType.bool, FFIType.bool, u64], returns: i32 },
@@ -169,7 +173,17 @@ export const C = dlopen(LIBMLXC_PATH, {
   // (res, input, weight, stride, padding, dilation, groups, stream) — depthwise
   // causal conv1d for Qwen3.5 gated-DeltaNet (groups == channels, padding 0).
   mlx_conv1d: { args: [P, u64, u64, i32, i32, i32, i32, u64], returns: i32 },
-  // --- training: autograd (value_and_grad) — proven in spikes/phase-train-vag.ts.
+  // C signature: (res, input, weight, stride_0, stride_1, padding_0,
+  // padding_1, dilation_0, dilation_1, groups, stream) — NHWC conv2d for the
+  // gemma-4 audio SSCP subsampler.
+  // HAZARD (lab/repro/bun-ffi-stack-args): Apple arm64 packs sub-8-byte STACK
+  // args at natural size (dilation_1@sp+0, groups@sp+4, stream@sp+8) but
+  // bun:ffi 1.3.14 writes one 8-byte slot per arg — the callee reads shifted
+  // garbage and mlx segfaults on a bogus stream. Workaround: declare the two
+  // stack ints as ONE u64 (dilation_1 | groups<<32) so Bun's slot layout
+  // byte-matches the Apple ABI. ops.conv2d does the packing.
+  mlx_conv2d: { args: [P, u64, u64, i32, i32, i32, i32, i32, u64, u64], returns: i32 },
+  // --- training: autograd (value_and_grad) — proven in lab/spikes/phase-train-vag.ts.
   // mlx_value_and_grad(res: mlx_closure_value_and_grad*, fun: mlx_closure,
   //   argnums: int*, n) → builds a value+grad closure differentiating the
   //   selected inputs of `fun` (a mlx_closure_new_func_payload over a JS loss).
@@ -260,6 +274,16 @@ export const C = dlopen(LIBMLXC_PATH, {
   // GPU trace capture (Phase E step 2: size the prize before the kernel)
   mlx_metal_start_capture: { args: [cstring], returns: i32 },
   mlx_metal_stop_capture: { args: [], returns: i32 },
+  // TurboQuant KV codec (Phase 13): FWHT rotation + per-group affine/Lloyd-Max
+  // (res, a, opt scale, stream) — 1/sqrt(dim) normalization when scale is null.
+  mlx_hadamard_transform: { args: [P, u64, u64, u64], returns: i32 },
+  mlx_min_axis: { args: [P, u64, i32, FFIType.bool, u64], returns: i32 },
+  mlx_round: { args: [P, u64, i32, u64], returns: i32 },
+  mlx_greater: { args: [P, u64, u64, u64], returns: i32 },
+  mlx_bitwise_and: { args: [P, u64, u64, u64], returns: i32 },
+  mlx_bitwise_or: { args: [P, u64, u64, u64], returns: i32 },
+  mlx_left_shift: { args: [P, u64, u64, u64], returns: i32 },
+  mlx_right_shift: { args: [P, u64, u64, u64], returns: i32 },
 }).symbols;
 
 export type MlxHandle = bigint;
@@ -329,7 +353,7 @@ export const NULL_HANDLE = 0n;
 // RULE: never read a typed array that a bun:ffi call wrote through a
 // pointer — once the calling function is DFG-compiled, the JIT eliminates
 // the load across the native call and `buf[i]` returns stale values
-// (repro/bun-ffi-f64/ISSUE.md; PLAN.md Phase 4 findings). Out-param slots
+// (lab/repro/bun-ffi-f64/ISSUE.md; PLAN.md Phase 4 findings). Out-param slots
 // must be read back with bun:ffi `read.*` (verified safe). Initializing the
 // slot via the typed-array *constructor* is fine: that store happens in
 // host code the JIT can't elide.

package/src/mlx/metal-kernel.ts

@@ -89,15 +89,18 @@ export class MetalKernel {
       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 ?? {})) {
+      // Dtypes-first, then ints, to MATCH mlx's own specialization-name order
+      // (mlx-lm's gated_delta_step passes [("InT",…),("StT",…),("Dk",…),…]).
+      // Order affects only the generated name/cache-key, not binding (named args).
+      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_int(cfg, ptr(kb), v);
+        C.mlx_fast_metal_kernel_config_add_template_arg_dtype(cfg, ptr(kb), v);
       }
-      for (const [k, v] of Object.entries(call.templateDtypes ?? {})) {
+      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_dtype(cfg, ptr(kb), v);
+        C.mlx_fast_metal_kernel_config_add_template_arg_int(cfg, ptr(kb), v);
       }
 
       const handles = new BigUint64Array(inputs.map((a) => a.handle));