mlx-bun 0.0.7 → 0.0.9

package/src/mlx/ffi.ts

@@ -119,6 +119,9 @@ export const C = dlopen(LIBMLXC_PATH, {
   mlx_divide: { args: [P, u64, u64, u64], returns: i32 },
   mlx_tanh: { args: [P, u64, u64], returns: i32 },
   mlx_sigmoid: { args: [P, u64, u64], returns: i32 },
+  mlx_erf: { args: [P, u64, u64], returns: i32 },
+  // (res, c, a, b, alpha, beta, stream) — alpha·(a@b) + beta·c (nn.Linear bias path)
+  mlx_addmm: { args: [P, u64, u64, u64, f32, f32, u64], returns: i32 },
   mlx_power: { args: [P, u64, u64, u64], returns: i32 },
   mlx_take_axis: { args: [P, u64, u64, i32, u64], returns: i32 },
   mlx_reshape: { args: [P, u64, P, u64, u64], returns: i32 },
@@ -142,11 +145,17 @@ export const C = dlopen(LIBMLXC_PATH, {
   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 },
+  mlx_cummax: { args: [P, u64, i32, FFIType.bool, FFIType.bool, u64], returns: i32 },
   mlx_where: { args: [P, u64, u64, u64, u64], returns: i32 },
   mlx_logsumexp_axis: { args: [P, u64, i32, FFIType.bool, u64], returns: i32 },
   mlx_greater_equal: { args: [P, u64, u64, u64], returns: i32 },
   mlx_less: { args: [P, u64, u64, u64], returns: i32 },
+  mlx_less_equal: { args: [P, u64, u64, u64], returns: i32 },
   mlx_logical_and: { args: [P, u64, u64, u64], returns: i32 },
+  mlx_logical_not: { args: [P, u64, u64], returns: i32 },
+  mlx_equal: { args: [P, u64, u64, u64], returns: i32 },
+  mlx_all: { args: [P, u64, FFIType.bool, u64], returns: i32 },
+  mlx_any_axis: { args: [P, u64, i32, FFIType.bool, u64], returns: i32 },
   mlx_random_key: { args: [P, u64], returns: i32 },
   mlx_random_categorical: { args: [P, u64, i32, u64, u64], returns: i32 },
   mlx_astype: { args: [P, u64, i32, u64], returns: i32 },
@@ -181,6 +190,8 @@ export const C = dlopen(LIBMLXC_PATH, {
   // (res, low, high, shape*, n, dtype, key may-be-null, stream) — low/high are
   // scalar mlx_arrays. Used for mlx-lm-faithful LoRA A init (uniform).
   mlx_random_uniform: { args: [P, u64, u64, P, u64, i32, u64, u64], returns: i32 },
+  mlx_random_randint: { args: [P, u64, u64, P, u64, i32, u64, u64], returns: i32 },
+  mlx_random_seed: { args: [u64], returns: i32 },
   mlx_random_split_num: { args: [P, u64, i32, u64], returns: i32 },
   // --- training: loss/optimizer elementwise + reductions
   mlx_stop_gradient: { args: [P, u64, u64], returns: i32 },

package/src/mlx/ops.ts

@@ -111,6 +111,23 @@ export function rope(
   );
 }
 
+/** fast::rope with every knob exposed (traditional + scale), for the
+ *  universal rope factory (rope_utils.initialize_rope port). The plain
+ *  `rope` above stays byte-identical for existing call sites. */
+export function ropeScaled(
+  x: MlxArray, dims: number, traditional: boolean, base: number | null,
+  scale: number, offset: number, freqs: MlxArray | null, s: S = gpuStream,
+): MlxArray {
+  return new MlxArray(
+    outArray("fast_rope", (o) =>
+      C.mlx_fast_rope(
+        o, x.handle, dims, traditional, optFloat(base), scale, offset,
+        freqs?.handle ?? 0n, s,
+      ),
+    ),
+  );
+}
+
 /** fast::rope with the position offset as an ARRAY (int32 scalar) instead
  *  of a baked int — required inside compiled decode graphs, where the
  *  offset changes every step but the graph must not. Same kernel as
@@ -267,6 +284,35 @@ export function geluApprox(x: MlxArray, s: S = gpuStream): MlxArray {
   return out;
 }
 
+export function erf(a: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("erf", (o) => C.mlx_erf(o, a.handle, s)));
+}
+
+/** nn.gelu (PRECISE, erf-based): x · (1 + erf(x/√2)) / 2 — composed exactly
+ *  as mlx's python source (weak scalars promote to x's dtype). Used by the
+ *  gemma-1 gated MLP and starcoder2's plain MLP; distinct from geluApprox
+ *  (gelu_pytorch_tanh). */
+export function geluPrecise(x: MlxArray, s: S = gpuStream): MlxArray {
+  const sqrt2 = scalarLike(Math.sqrt(2), x);
+  const xs = div(x, sqrt2, s); // x / math.sqrt(2) — division, not mul-by-reciprocal
+  const e = erf(xs, s);
+  const one = scalarLike(1, x);
+  const e1 = add(one, e, s);
+  const xe = mul(x, e1, s);
+  const two = scalarLike(2, x);
+  const out = div(xe, two, s); // (...) / 2
+  for (const a of [sqrt2, xs, e, one, e1, xe, two]) a.dispose();
+  return out;
+}
+
+/** mx.addmm(c, a, b): a@b + c fused (α=β=1) — nn.Linear's bias path
+ *  (`mx.addmm(bias, x, weight.T)`). */
+export function addmm(c: MlxArray, a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(
+    outArray("addmm", (o) => C.mlx_addmm(o, c.handle, a.handle, b.handle, 1.0, 1.0, s)),
+  );
+}
+
 export function exp(a: MlxArray, s: S = gpuStream): MlxArray {
   return new MlxArray(outArray("exp", (o) => C.mlx_exp(o, a.handle, s)));
 }
@@ -315,6 +361,13 @@ export function cumsum(a: MlxArray, axis: number, s: S = gpuStream): MlxArray {
   );
 }
 
+/** Cumulative maximum along `axis` (inclusive, forward) — mlx.core.cummax. */
+export function cummax(a: MlxArray, axis: number, s: S = gpuStream): MlxArray {
+  return new MlxArray(
+    outArray("cummax", (o) => C.mlx_cummax(o, a.handle, axis, false, true, s)),
+  );
+}
+
 export function where(cond: MlxArray, x: MlxArray, y: MlxArray, s: S = gpuStream): MlxArray {
   return new MlxArray(
     outArray("where", (o) => C.mlx_where(o, cond.handle, x.handle, y.handle, s)),
@@ -370,6 +423,38 @@ export function less(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
   return new MlxArray(outArray("less", (o) => C.mlx_less(o, a.handle, b.handle, s)));
 }
 
+export function lessEqual(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("less_equal", (o) => C.mlx_less_equal(o, a.handle, b.handle, s)));
+}
+
+export function logicalNot(a: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("logical_not", (o) => C.mlx_logical_not(o, a.handle, s)));
+}
+
+export function equal(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("equal", (o) => C.mlx_equal(o, a.handle, b.handle, s)));
+}
+
+/** Reduce-all to a scalar (mlx.core.all). */
+export function allReduce(a: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("all", (o) => C.mlx_all(o, a.handle, false, s)));
+}
+
+/** any along one axis (mlx.core.any). */
+export function anyAxis(a: MlxArray, axis: number, keepdims: boolean, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("any_axis", (o) => C.mlx_any_axis(o, a.handle, axis, keepdims, s)));
+}
+
+/** Read a 0-d/1-elem boolean array to host (forces eval). */
+export function itemBool(a: MlxArray): boolean {
+  const u = a.astype(Dtype.uint32);
+  try {
+    return itemUint32(u) !== 0;
+  } finally {
+    u.dispose();
+  }
+}
+
 export function logicalAnd(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
   return new MlxArray(
     outArray("logical_and", (o) => C.mlx_logical_and(o, a.handle, b.handle, s)),
@@ -443,6 +528,34 @@ export function randomKey(seed: bigint): MlxArray {
   return new MlxArray(outArray("random_key", (o) => C.mlx_random_key(o, seed)));
 }
 
+/** Set the GLOBAL mlx PRNG seed (mlx.core.random.seed). Calls to randint/etc
+ *  with a null key thread the global key, so seeding then calling in the same
+ *  order as the reference reproduces its draws bit-for-bit. */
+export function randomSeed(seed: bigint): void {
+  if (C.mlx_random_seed(seed) !== 0) throw new Error(`mlx_random_seed failed: ${takeMlxError() ?? ""}`);
+}
+
+/** Uniform integers in [low, high) — mlx.core.random.randint. With key=null it
+ *  uses (and advances) the global key, matching mx.random.randint exactly. */
+export function randint(
+  low: number, high: number, shape: number[], dtype: Dtype,
+  key: MlxArray | null = null, s: S = gpuStream,
+): MlxArray {
+  const buf = new Int32Array(shape);
+  const lo = MlxArray.fromInt32(new Int32Array([low]), []);
+  const hi = MlxArray.fromInt32(new Int32Array([high]), []);
+  try {
+    return new MlxArray(
+      outArray("random_randint", (o) =>
+        C.mlx_random_randint(o, lo.handle, hi.handle, ptr(buf), BigInt(shape.length), dtype, key?.handle ?? 0n, s),
+      ),
+    );
+  } finally {
+    lo.dispose();
+    hi.dispose();
+  }
+}
+
 export function randomCategorical(logits: MlxArray, key: MlxArray | null, s: S = gpuStream): MlxArray {
   return new MlxArray(
     outArray("categorical", (o) =>

package/src/model/compiled-decode.ts

@@ -590,6 +590,29 @@ export class CompiledDecode {
     const evalWith: MlxArray[] = [];
     let carried = cur; // seg 0 consumes the token array; later, hidden state
     let carriedOwned = false; // cur belongs to the caller
+    // Mid-step failure safety (a later segment's closure.apply can throw at
+    // first trace AFTER earlier layers already advanced): the step must be
+    // TRANSACTIONAL, or the caller's uncompiled re-forward of the same token
+    // double-writes the committed layers (+1 offset skew forever, silently
+    // corrupted attention). Two mechanisms, chosen by what each phase allows:
+    //  - ring segments: adoptDecodeStep is a pure host-side handle swap with
+    //    no reader between here and step end (each cache is read only by its
+    //    own layer, inside the closure that produced the update), so the
+    //    adopts are STAGED and committed only once every segment and js layer
+    //    has succeeded. On failure the staged buffers are disposed and the
+    //    ring caches are exactly at their pre-step state (prepareDecodeStep's
+    //    growth/trim/rotate bookkeeping is benign for the uncompiled path —
+    //    the same invariant #stepWhole and generate.ts's fallback rely on).
+    //  - js (concat) layers: layer.forward commits via the cache's real
+    //    updateAndFetch mid-loop, so on failure the commit is ROLLED BACK by
+    //    reverting offset (and ring index) — trim(1) semantics. The written
+    //    KV row beyond the reverted offset is dead data the re-forward
+    //    overwrites at the same position (updateAndFetch writes at offset/
+    //    ringIdx), so post-recovery state equals a single clean write. This
+    //    holds for the post-wrap quantized-under-perf-kernel case too: the
+    //    in-place S=1 write slot is re-derived from the reverted ringIdx.
+    const stagedAdopts: { c: AnyCache; bufs: MlxArray[] }[] = [];
+    const committedJs: { c: AnyCache; offsetBefore: number }[] = [];
     try {
       for (let s = 0; s < segs.length; s++) {
         const { from, to } = segs[s]!;
@@ -636,17 +659,14 @@ export class CompiledDecode {
           carried = outs[0]!;
           carriedOwned = true;
           // ring updates are ancestors of the carried hidden state (the
-          // in-graph fetch reads the updated buffer): no eval roots
+          // in-graph fetch reads the updated buffer): no eval roots. The
+          // host-side adoption is deferred to the end of the step (see above).
           let oi = 1;
           for (let i = from; i < to; i++) {
             const c = anyCaches[i]!;
-            if (c instanceof RotatingQuantizedKVCache) {
-              c.adoptDecodeStep(outs.slice(oi, oi + 6));
-              oi += 6;
-            } else {
-              (c as RotatingKVCache).adoptDecodeStep(outs[oi]!, outs[oi + 1]!);
-              oi += 2;
-            }
+            const width = c instanceof RotatingQuantizedKVCache ? 6 : 2;
+            stagedAdopts.push({ c, bufs: outs.slice(oi, oi + width) });
+            oi += width;
           }
         }
         if (s < jsLayers.length) {
@@ -655,6 +675,7 @@ export class CompiledDecode {
           const c = anyCaches[li]!;
           const window = layer.layerType === "sliding_attention" ? this.model.windowSize : null;
           const mask = c.makeMask(1, window); // N=1 → mode ""
+          committedJs.push({ c, offsetBefore: c.offset });
           const { h: next, shared } = layer.forward(carried, mask, c, null, null);
           if (carriedOwned) carried.dispose();
           carried = next;
@@ -663,6 +684,20 @@ export class CompiledDecode {
           mask.arr?.dispose();
         }
       }
+      // every segment and js layer succeeded — commit the staged ring adopts
+      for (const { c, bufs } of stagedAdopts) {
+        if (c instanceof RotatingQuantizedKVCache) c.adoptDecodeStep(bufs);
+        else (c as RotatingKVCache).adoptDecodeStep(bufs[0]!, bufs[1]!);
+      }
+    } catch (e) {
+      // unwind: free staged (never-adopted) ring buffers and the owned hidden
+      // state, and revert the js layers' committed writes, so the caller's
+      // uncompiled re-forward of this token starts from the pre-step state.
+      for (const { bufs } of stagedAdopts) for (const b of bufs) b.dispose();
+      for (const { c, offsetBefore } of committedJs)
+        if (c.offset > offsetBefore) c.trim(c.offset - offsetBefore);
+      if (carriedOwned) carried.dispose();
+      throw e;
     } finally {
       for (const t of stepTemps) t.dispose();
     }