mlx-bun 0.0.13 → 0.3.0

package/src/mlx/array.ts

@@ -46,17 +46,16 @@ let nextPinId = 1;
 // 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" },
+// loadKvCache — so no restore mapping needs pinning anymore). Bun >= 1.4
+// returns a numeric function address from .ptr; verified against dlsym on
+// 1.4.0 and 1.4.2. Keep the library open for every native destructor's lifetime.
+const libcFree = dlopen("/usr/lib/libSystem.B.dylib", {
+  free: { args: ["ptr"], returns: "void" },
 });
-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");
+// The pinned bun-types predates the runtime's numeric .ptr property.
+const freeFnAddr = Number(Reflect.get(libcFree.symbols.free, "ptr"));
+if (!Number.isSafeInteger(freeFnAddr) || freeFnAddr <= 0)
+  throw new Error("free.ptr is invalid — zero-copy arrays need a native no-op dtor");
 
 export function pinnedBufferCount(): number {
   return pinned.size;

package/src/mlx/compile.ts

@@ -23,8 +23,14 @@ import { MlxArray } from "./array";
 
 export type TraceFn = (inputs: MlxArray[]) => MlxArray[];
 
-const traceRegistry = new Map<number, { fn: TraceFn; count: number }>();
+const traceRegistry = new Map<number, { fn: TraceFn; count: number; shapeless: boolean }>();
 let nextTraceId = 1;
+let shapelessTracing = false;
+
+/** Fixed-shape custom kernels must fall back while a shapeless graph traces. */
+export function isShapelessTracing(): boolean {
+  return shapelessTracing;
+}
 /** Error thrown by a trace fn, carried across the C boundary (the
  *  trampoline must not throw into mlx; it returns non-zero instead). */
 let traceError: unknown = null;
@@ -40,9 +46,11 @@ export let traceCalls = 0;
 // error handler in ffi.ts).
 const trampoline = new JSCallback(
   (resPtr: number, argsVec: bigint, payload: number): number => {
+    const previousShapeless = shapelessTracing;
     try {
       const entry = traceRegistry.get(payload);
       if (!entry) throw new Error(`compile trampoline: unknown trace id ${payload}`);
+      shapelessTracing ||= entry.shapeless;
       const fn = entry.fn;
       entry.count++;
       traceCalls++;
@@ -66,6 +74,8 @@ const trampoline = new JSCallback(
     } catch (e) {
       traceError = e;
       return 1;
+    } finally {
+      shapelessTracing = previousShapeless;
     }
   },
   { args: ["ptr", "u64", "ptr"], returns: "i32" },
@@ -95,7 +105,7 @@ export class CompiledFunction {
 
   constructor(fn: TraceFn, shapeless = true) {
     this.#traceId = nextTraceId++;
-    traceRegistry.set(this.#traceId, { fn, count: 0 });
+    traceRegistry.set(this.#traceId, { fn, count: 0, shapeless });
     // payload is the trace id smuggled through the void* (never deref'd);
     // dtor is null — nothing to free on the C side.
     const raw = C.mlx_closure_new_func_payload(trampoline.ptr, this.#traceId as never, null);

package/src/mlx/ffi.ts

@@ -10,21 +10,22 @@
 // - mlx_optional_int is `{ int32 value; bool has_value }` = 8 bytes, by
 //   value in one register: value in bits 0..31, has_value at bit 32.
 
-import { dlopen, FFIType, ptr, read } from "bun:ffi";
+import { dlopen, FFIType, ptr, read, type Pointer } from "bun:ffi";
 import { existsSync } from "node:fs";
 import { dirname, join } from "node:path";
-import { nativePackDir } from "../native-pack";
+import { MLX_CORE_VERSION, nativePackDir } from "../native-pack";
+import { runtimeValue } from "../runtime-config";
 
 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
  *  libmlxc.dylib + libmlx.dylib + mlx.metallib ship beside the binary —
- *  see docs/reference/embedding.md) → the downloaded native-pack cache
+ *  see docs/reference/distribution.md) → the downloaded native-pack cache
  *  (src/native-pack.ts; populated on first run by the CLI) → homebrew
  *  (arm64, then Intel prefix). Keep in sync with nativeRuntimeDir(). */
 function resolveLibmlxc(): string {
-  const env = process.env.MLX_BUN_LIBMLXC;
+  const env = runtimeValue("MLX_BUN_LIBMLXC");
   if (env) return env;
   const beside = join(dirname(process.execPath), "libmlxc.dylib");
   if (existsSync(beside)) return beside;
@@ -37,8 +38,12 @@ function resolveLibmlxc(): string {
 }
 
 export const LIBMLXC_PATH = resolveLibmlxc();
-
 export const C = dlopen(LIBMLXC_PATH, {
+  // linked MLX core version, independent of the C API family
+  mlx_version: { args: [P], returns: i32 },
+  mlx_string_new: { args: [], returns: u64 },
+  mlx_string_data: { args: [u64], returns: cstring },
+  mlx_string_free: { args: [u64], returns: i32 },
   // array lifecycle
   mlx_array_new: { args: [], returns: u64 },
   mlx_array_new_data: { args: [P, P, i32, i32], returns: u64 },
@@ -77,6 +82,7 @@ export const C = dlopen(LIBMLXC_PATH, {
   mlx_device_info_get: { args: [P, u64], returns: i32 },
   mlx_device_info_free: { args: [u64], returns: i32 },
   mlx_device_info_get_size: { args: [P, u64, cstring], returns: i32 },
+  mlx_device_info_get_string: { args: [P, u64, cstring], returns: i32 },
   mlx_get_default_device: { args: [P], returns: i32 },
   mlx_device_free: { args: [u64], returns: i32 },
   mlx_synchronize: { args: [u64], returns: i32 },
@@ -97,8 +103,13 @@ export const C = dlopen(LIBMLXC_PATH, {
   mlx_fast_rms_norm: { args: [P, u64, u64, f32, u64], returns: i32 },
   // (res, x, weight may-be-null, bias may-be-null, eps, stream)
   mlx_fast_layer_norm: { args: [P, u64, u64, u64, f32, u64], returns: i32 },
+  // MLX 0.32.2: mask, sinks, force_fused, stream.
+  mlx_fast_scaled_dot_product_attention: {
+    args: [P, u64, u64, u64, f32, cstring, u64, u64, FFIType.bool, u64], returns: i32,
+  },
   mlx_matmul: { args: [P, u64, u64, u64], returns: i32 },
   mlx_contiguous: { args: [P, u64, FFIType.bool, u64], returns: i32 },
+  mlx_copy: { args: [P, u64, u64], returns: i32 },
   // (res: mlx_vector_array*, w, opt group_size, opt bits, mode, global_scale, stream)
   mlx_quantize: { args: [P, u64, u64, u64, cstring, u64, u64], returns: i32 },
   mlx_vector_array_get: { args: [P, u64, u64], returns: i32 },
@@ -106,15 +117,14 @@ export const C = dlopen(LIBMLXC_PATH, {
   mlx_softmax_axis: { args: [P, u64, i32, FFIType.bool, u64], returns: i32 },
   mlx_max_axis: { args: [P, u64, i32, FFIType.bool, u64], returns: i32 },
   mlx_sum_axis: { args: [P, u64, i32, FFIType.bool, u64], returns: i32 },
+  // Optional dtype uses the eight-byte value/has_value layout of optional int.
+  mlx_cumsum_axis: { args: [P, u64, i32, FFIType.bool, FFIType.bool, u64, u64], returns: i32 },
+  mlx_cummax_axis: { args: [P, u64, i32, FFIType.bool, FFIType.bool, u64], returns: i32 },
   mlx_maximum: { args: [P, u64, u64, u64], returns: i32 },
   mlx_expand_dims: { args: [P, u64, i32, u64], returns: i32 },
   mlx_logical_or: { args: [P, u64, u64, u64], returns: i32 },
   // (res, x, dims, traditional, opt base, scale, offset, freqs may-be-null, stream)
   mlx_fast_rope: { args: [P, u64, i32, FFIType.bool, u64, f32, i32, u64, u64], returns: i32 },
-  // (res, q, k, v, scale, mask_mode, mask may-be-null, sinks may-be-null, stream)
-  mlx_fast_scaled_dot_product_attention: {
-    args: [P, u64, u64, u64, f32, cstring, u64, u64, u64], returns: i32,
-  },
   // ops (grown as phases need them)
   mlx_add: { args: [P, u64, u64, u64], returns: i32 },
   mlx_subtract: { args: [P, u64, u64, u64], returns: i32 },
@@ -148,8 +158,6 @@ export const C = dlopen(LIBMLXC_PATH, {
   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 },
-  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 },
@@ -173,16 +181,11 @@ 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 },
-  // 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 },
+  // Natural mlx-c signatures. Bun 1.4 fixed the macOS arm64 stack packing
+  // bug recorded in lab/repro/bun-ffi-stack-args, so these must match the
+  // header exactly. mlx-bun now requires Bun 1.4 or newer.
+  mlx_conv2d: { args: [P, u64, u64, i32, i32, i32, i32, i32, i32, i32, u64], returns: i32 },
+  mlx_conv3d: { args: [P, u64, u64, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, 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
@@ -223,7 +226,7 @@ export const C = dlopen(LIBMLXC_PATH, {
   mlx_save_safetensors: { args: [cstring, u64, u64], returns: i32 },
   mlx_map_string_to_array_insert: { args: [u64, cstring, u64], returns: i32 },
   mlx_map_string_to_string_insert: { args: [u64, cstring, cstring], returns: i32 },
-  // compile via closures (Phase A compiled decode — docs/design/optimization_plan.md).
+  // compile via closures (Phase A compiled decode — docs/archive/investigations/optimization_plan.md).
   // mlx_closure is the usual one-pointer struct; the func-payload variant
   // carries an id we use to route the trace callback back to JS.
   mlx_closure_new: { args: [], returns: u64 },
@@ -317,6 +320,22 @@ export function takeMlxError(): string | null {
   return e;
 }
 
+function readRuntimeVersion(): string {
+  const slot = new BigUint64Array([C.mlx_string_new()]);
+  const p = ptr(slot);
+  try {
+    if (C.mlx_version(p) !== 0)
+      throw new Error(`mlx_version failed: ${takeMlxError() ?? "(no mlx message)"}`);
+    return C.mlx_string_data(read.u64(p, 0)).toString();
+  } finally { C.mlx_string_free(read.u64(p, 0)); }
+}
+
+/** Read once at startup; kernel eligibility must follow the linked core. */
+export const MLX_VERSION = readRuntimeVersion();
+if (MLX_VERSION !== MLX_CORE_VERSION) {
+  throw new Error(`mlx-bun requires its bundled MLX ${MLX_CORE_VERSION}; ${LIBMLXC_PATH} loaded MLX ${MLX_VERSION}`);
+}
+
 /** mlx_dtype enum values (array.h). */
 export const enum Dtype {
   bool = 0, uint8 = 1, uint16 = 2, uint32 = 3, uint64 = 4,
@@ -408,6 +427,31 @@ export function clearCache(): void {
   C.mlx_clear_cache();
 }
 
+let cachedDeviceArchitecture: string | undefined;
+
+/** Native architecture identifier, used to match MLX's hardware dispatch. */
+export function deviceArchitecture(): string {
+  if (cachedDeviceArchitecture !== undefined) return cachedDeviceArchitecture;
+  const devSlot = new BigUint64Array(1), devPtr = ptr(devSlot);
+  if (C.mlx_get_default_device(devPtr) !== 0)
+    throw new Error(`mlx_get_default_device failed: ${takeMlxError() ?? ""}`);
+  const dev = read.u64(devPtr, 0);
+  const infoSlot = new BigUint64Array([C.mlx_device_info_new()]), infoPtr = ptr(infoSlot);
+  try {
+    if (C.mlx_device_info_get(infoPtr, dev) !== 0)
+      throw new Error(`mlx_device_info_get failed: ${takeMlxError() ?? ""}`);
+    const out = new BigUint64Array(1), outPtr = ptr(out);
+    const key = Buffer.from("architecture\0", "utf8");
+    if (C.mlx_device_info_get_string(outPtr, read.u64(infoPtr, 0), ptr(key)) !== 0)
+      throw new Error("architecture not in device info");
+    // Copy before the device-info object releases the native string.
+    return cachedDeviceArchitecture = new CString(Number(read.u64(outPtr, 0)) as Pointer).toString();
+  } finally {
+    C.mlx_device_info_free(read.u64(infoPtr, 0));
+    C.mlx_device_free(dev);
+  }
+}
+
 /** Metal's recommended max working-set size for the default device. */
 export function maxRecommendedWorkingSetSize(): number {
   const devSlot = new BigUint64Array(1);

package/src/mlx/materialize.ts

@@ -0,0 +1,23 @@
+// Explicit compact storage for a value that must release a larger backing
+// buffer. MLX 0.31.2 Copy::eval shares storage; contiguous() may also alias.
+// MLX's dynamic slice allocates the logical output size and has a native VJP.
+import { MlxArray } from "./array";
+import * as ops from "./ops";
+
+let zero: MlxArray | undefined;
+
+/** Borrow an array; return an owned, row-contiguous copy on the GPU.
+ * Values and dtype are unchanged. Evaluation stays lazy and never reads data
+ * back to the host. Use only when independent compact storage is required;
+ * immutable array aliases normally avoid this allocation and dispatch.
+ * Slice sizes are concrete, so use shape-specialized compiled graphs. */
+export function materializeCopy(input: MlxArray): MlxArray {
+  zero ??= MlxArray.fromInt32(new Int32Array([0]), []);
+  if (input.ndim > 0) return ops.sliceDynamic(input, zero, [0], input.shape);
+  const vector = ops.reshape(input, [1]);
+  try {
+    const copy = ops.sliceDynamic(vector, zero, [0], [1]);
+    try { return ops.reshape(copy, []); }
+    finally { copy.dispose(); }
+  } finally { vector.dispose(); }
+}

package/src/mlx/metal-kernel.ts

@@ -1,4 +1,4 @@
-// mx.fast.metal_kernel from Bun (docs/design/optimization_plan.md Phase E): the
+// mx.fast.metal_kernel from Bun (docs/archive/investigations/optimization_plan.md Phase E): the
 // intended MLX extension point for custom Metal kernels — no fork. The
 // kernel body is Metal Shading Language; mlx generates the signature
 // from input/output names (inputs by name, outputs by name, plus
@@ -29,15 +29,15 @@ export interface MetalKernelSpec {
 }
 
 export interface MetalKernelCall {
-  /** 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?"). */
+  /** Concrete output shapes/dtypes. Provide this OR `outputShapeFn`.
+   *  Both are captured when apply() builds the custom primitive. */
   outputs?: { shape: number[]; dtype: Dtype }[];
-  /** Derive output shapes/dtypes from the (current/traced) input shapes. */
+  /** Derive outputs during eager construction or a shape-specific compile trace.
+   *  This JS callback is not native shape inference: MLX's CustomKernel lacks
+   *  output_shapes(), so it cannot participate in shapeless replay. */
   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). */
+  /** Concrete launch grid, OR derive it during eager construction or tracing.
+   *  Native replay reuses the captured grid without invoking this callback. */
   grid: [number, number, number] | ((inputs: MlxArray[]) => [number, number, number]);
   threadGroup: [number, number, number];
   templateInts?: Record<string, number>;
@@ -73,9 +73,8 @@ export class MetalKernel {
     // flash-CCE pin bug). Pure-additive: holds references, changes no logic.
     const keepAlive: unknown[] = [];
     try {
-      // 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.
+      // Resolve callbacks once for this eager call or shape-specific trace.
+      // A compiled graph captures these values; replay does not call JS again.
       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;

package/src/mlx/ops.ts

@@ -113,6 +113,15 @@ export function matmul(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
   return new MlxArray(outArray("matmul", (o) => C.mlx_matmul(o, a.handle, b.handle, s)));
 }
 
+/** Bit-identical copy through a FRESH primitive node (mx.copy). Unlike
+ *  `contiguous`, this always creates a new graph node even for an already-
+ *  contiguous input — the tool for DETACHING an array from a multi-output
+ *  primitive's sibling set (a retained sibling output pins every co-output's
+ *  buffer natively; see gatedDeltaUpdate's state detach, 2026-08-20). */
+export function copyOf(a: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("copy", (o) => C.mlx_copy(o, a.handle, s)));
+}
+
 /** Materialize a row-major contiguous copy (needed before rawBytes on views). */
 export function contiguous(a: MlxArray, s: S = gpuStream): MlxArray {
   return new MlxArray(
@@ -161,7 +170,7 @@ export function ropeScaled(
  *  of a baked int — required inside compiled decode graphs, where the
  *  offset changes every step but the graph must not. Same kernel as
  *  `rope`; bit-exactness vs the static form is asserted in
- *  tests/compile.test.ts. */
+ *  tests/unit/compile.test.ts. */
 export function ropeDynamic(
   x: MlxArray, dims: number, base: number | null, offset: MlxArray,
   freqs: MlxArray | null, s: S = gpuStream,
@@ -180,7 +189,7 @@ export function ropeDynamic(
  *  the dynamic twin of `ropeScaled`, for the universal rope factory's
  *  batched-decode path (per-row [B] offsets under left-padding). Same
  *  kernel as `ropeDynamic`; static/dynamic bit-exactness is asserted in
- *  tests/compile.test.ts. */
+ *  tests/unit/compile.test.ts. */
 export function ropeScaledDynamic(
   x: MlxArray, dims: number, traditional: boolean, base: number | null,
   scale: number, offset: MlxArray, freqs: MlxArray | null, s: S = gpuStream,
@@ -204,7 +213,7 @@ export function sdpa(
     outArray("fast_sdpa", (o) =>
       C.mlx_fast_scaled_dot_product_attention(
         o, q.handle, k.handle, v.handle, scale, ptr(cstr(maskMode)),
-        maskArr?.handle ?? 0n, 0n, s,
+        maskArr?.handle ?? 0n, 0n, false, s,
       ),
     ),
   );
@@ -305,7 +314,8 @@ export function takeAxis(a: MlxArray, indices: MlxArray, axis: number, s: S = gp
 }
 
 export function reshape(a: MlxArray, shape: number[], s: S = gpuStream): MlxArray {
-  const buf = new Int32Array(shape);
+  // bun:ffi cannot take a pointer to an empty view; ndim=0 ignores the dummy.
+  const buf = new Int32Array(shape.length ? shape : [0]);
   return new MlxArray(
     outArray("reshape", (o) => C.mlx_reshape(o, a.handle, ptr(buf), BigInt(shape.length), s)),
   );
@@ -445,14 +455,14 @@ export function putAlongAxis(
 
 export function cumsum(a: MlxArray, axis: number, s: S = gpuStream): MlxArray {
   return new MlxArray(
-    outArray("cumsum", (o) => C.mlx_cumsum(o, a.handle, axis, false, true, s)),
+    outArray("cumsum", (o) => C.mlx_cumsum_axis(o, a.handle, axis, false, true, 0n, s)),
   );
 }
 
 /** 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)),
+    outArray("cummax", (o) => C.mlx_cummax_axis(o, a.handle, axis, false, true, s)),
   );
 }
 
@@ -467,10 +477,13 @@ export function where(cond: MlxArray, x: MlxArray, y: MlxArray, s: S = gpuStream
 // (lab/repro/bun-ffi-f64/) traced the root cause to a Bun bug: after DFG
 // tier-up, typed-array reads following a bun:ffi call return stale values
 // (the JIT eliminates the load across the native call). Not f64 marshaling
-// — args reach C intact. See PLAN.md Phase 4 findings. Large constant
-// ranges (topP's vocab arange) are cached for the process lifetime.
-const arangeCache = new Map<string, MlxArray>();
+// — args reach C intact. See PLAN.md Phase 4 findings. Reuse large ranges
+// across attention layers and sampling, but bound their retained storage:
+// causal masks request a new range on every round once context exceeds 64K.
+const arangeCache = new Map<string, { array: MlxArray; bytes: number }>();
 const ARANGE_CACHE_MIN = 65536;
+const ARANGE_CACHE_BYTES = 8 * 1024 * 1024;
+let arangeCacheBytes = 0;
 
 export function arange(start: number, stop: number, step: number, dtype: Dtype, s: S = gpuStream): MlxArray {
   if (!Number.isInteger(start) || !Number.isInteger(step) || step === 0)
@@ -478,7 +491,11 @@ export function arange(start: number, stop: number, step: number, dtype: Dtype,
   const n = Math.max(0, Math.ceil((stop - start) / step));
   const key = `${start}|${step}|${n}|${dtype}`;
   const cached = arangeCache.get(key);
-  if (cached) return cached.slice([0], [n], s);
+  if (cached) {
+    arangeCache.delete(key);
+    arangeCache.set(key, cached);
+    return cached.array.slice([0], [n], s);
+  }
 
   let arr: MlxArray;
   if (dtype === Dtype.int32) {
@@ -494,8 +511,20 @@ export function arange(start: number, stop: number, step: number, dtype: Dtype,
     arr = new MlxArray(outArray("arange", (o) => C.mlx_arange(o, start, stop, step, dtype, s)));
   }
   if (n >= ARANGE_CACHE_MIN) {
-    arangeCache.set(key, arr);
-    return arr.slice([0], [n], s);
+    const bytes = arr.nbytes;
+    if (bytes <= ARANGE_CACHE_BYTES) {
+      while (arangeCacheBytes + bytes > ARANGE_CACHE_BYTES) {
+        const [oldKey, old] = arangeCache.entries().next().value!;
+        arangeCache.delete(oldKey);
+        arangeCacheBytes -= old.bytes;
+        // Returned slices retain their own native reference, including lazy
+        // views that have not been evaluated when this owner is evicted.
+        old.array.dispose();
+      }
+      arangeCache.set(key, { array: arr, bytes });
+      arangeCacheBytes += bytes;
+      return arr.slice([0], [n], s);
+    }
   }
   return arr;
 }
@@ -602,7 +631,7 @@ export function sliceDynamic(
 /** slice_update with the start index as an ARRAY (one int32 per entry in
  *  `axes`) — the compiled-decode form of the per-step cache write. Same
  *  write as `sliceUpdate`; bit-exactness vs the static form is asserted
- *  in tests/compile.test.ts. */
+ *  in tests/unit/compile.test.ts. */
 export function sliceUpdateDynamic(
   src: MlxArray, update: MlxArray, start: MlxArray, axes: number[], s: S = gpuStream,
 ): MlxArray {
@@ -880,20 +909,36 @@ export function conv1d(
 
 /** nn.Conv2d: input [B, H, W, C_in] (NHWC), weight [C_out, kH, kW, C_in/groups]
  *  → [B, H', W', C_out]. mlx weight layout — the oracle transposes PyTorch's
- *  [C_out, C_in, kH, kW] at load. For the gemma-4 audio SSCP subsampler.
- *  dilation_1 and groups ride in one packed u64 — see the mlx_conv2d hazard
- *  note in ffi.ts (bun:ffi stack-arg ABI bug). */
+ *  [C_out, C_in, kH, kW] at load. For the gemma-4 audio SSCP subsampler. */
 export function conv2d(
   input: MlxArray, weight: MlxArray, stride: [number, number] = [1, 1],
   padding: [number, number] = [0, 0], dilation: [number, number] = [1, 1],
   groups = 1, s: S = gpuStream,
 ): MlxArray {
-  const dilation1AndGroups =
-    (BigInt(dilation[1] >>> 0) & 0xffffffffn) | (BigInt(groups >>> 0) << 32n);
   return new MlxArray(
     outArray("conv2d", (o) =>
       C.mlx_conv2d(o, input.handle, weight.handle, stride[0], stride[1],
-        padding[0], padding[1], dilation[0], dilation1AndGroups, s),
+        padding[0], padding[1], dilation[0], dilation[1], groups, s),
+    ),
+  );
+}
+
+/** nn.Conv3d: input [B, D, H, W, C_in] (NDHWC), weight [C_out, kD, kH, kW,
+ *  C_in/groups] → [B, D', H', W', C_out]. The Qwen3-VL vision patch embed
+ *  (kernel == stride) — the conv KERNEL's accumulation order is the parity-
+ *  relevant arithmetic; a value-equivalent GEMM differs by bf16 ulp. */
+export function conv3d(
+  input: MlxArray, weight: MlxArray,
+  stride: [number, number, number] = [1, 1, 1],
+  padding: [number, number, number] = [0, 0, 0],
+  dilation: [number, number, number] = [1, 1, 1],
+  groups = 1, s: S = gpuStream,
+): MlxArray {
+  return new MlxArray(
+    outArray("conv3d", (o) =>
+      C.mlx_conv3d(o, input.handle, weight.handle, stride[0], stride[1],
+        stride[2], padding[0], padding[1], padding[2], dilation[0],
+        dilation[1], dilation[2], groups, s),
     ),
   );
 }

package/src/mlx/quantized-rows.ts

@@ -0,0 +1,36 @@
+import type { MlxArray } from "./array";
+import * as ops from "./ops";
+
+/** Run independent affine M=1 matrix-vector products in one native batched call.
+ * Weight/scales/bias views share their original buffers. MLX 0.32.2's
+ * flattened M>=2 path uses different arithmetic, so appending known tokens
+ * must explicitly select this layout when preserving decode state.
+ */
+export function quantizedMatmulRows(
+  x: MlxArray, w: MlxArray, scales: MlxArray, biases: MlxArray | null,
+  spec: ops.QuantSpec,
+): MlxArray {
+  const shape = x.shape, width = shape[shape.length - 1]!;
+  const rows = shape.reduce((a, b) => a * b, 1) / width;
+  if (rows === 1) return ops.quantizedMatmul(x, w, scales, biases, spec);
+  if (spec.mode !== "affine")
+    throw new Error("Independent-row batched matmul requires affine quantization");
+  if (w.shape.length !== 2 || scales.shape.length !== 2 || (biases && biases.shape.length !== 2))
+    throw new Error("Independent-row matmul requires shared two-dimensional weights");
+  const views: MlxArray[] = [];
+  const view = (value: MlxArray, newShape: number[]) => {
+    const result = ops.reshape(value, newShape); views.push(result); return result;
+  };
+  let output: MlxArray | undefined;
+  try {
+    const input = view(x, [rows, 1, width]);
+    const weight = view(w, [1, ...w.shape]);
+    const scale = view(scales, [1, ...scales.shape]);
+    const bias = biases ? view(biases, [1, ...biases.shape]) : null;
+    output = ops.quantizedMatmul(input, weight, scale, bias, spec);
+    return ops.reshape(output, [...shape.slice(0, -1), output.shape.at(-1)!]);
+  } finally {
+    output?.dispose();
+    for (const value of views) value.dispose();
+  }
+}

package/src/mlx/token-bitmask.ts

@@ -0,0 +1,52 @@
+// XGrammar's additive 0/-Infinity mask contract, decoded from packed bits.
+// Reference: mlc-ai/xgrammar b5048784, apply_token_bitmask_mlx.py (Apache-2.0).
+import { MlxArray } from "./array";
+import { CompiledFunction } from "./compile";
+import { Dtype } from "./ffi";
+import { MetalKernel } from "./metal-kernel";
+
+let compiled: CompiledFunction | undefined;
+let kernel: MetalKernel | undefined;
+
+function maskGraph(logits: MlxArray, bits: MlxArray): MlxArray {
+  kernel ??= new MetalKernel({
+    name: "mlx_bun_token_bitmask",
+    inputNames: ["logits", "bits"],
+    outputNames: ["out"],
+    source: String.raw`
+      const uint i = thread_position_in_grid.x;
+      if (i >= SIZE) return;
+      const uint id = i % V;
+      const uint word = uint(bits[id >> 5]);
+      const bool valid = ((word >> (id & 31)) & 1u) != 0;
+      out[i] = logits[i] + T(valid ? 0.0f : -INFINITY);
+    `,
+  });
+  return kernel.apply([logits, bits], {
+    outputs: [{ shape: logits.shape, dtype: logits.dtype }],
+    grid: [logits.size, 1, 1],
+    threadGroup: [128, 1, 1],
+    templateDtypes: { T: logits.dtype },
+    templateInts: { V: logits.shape.at(-1)!, SIZE: logits.size },
+  })[0]!;
+}
+
+/** Apply one little-endian token bitmask to every row along the final axis.
+ * Borrow logits; return an owned lazy result. Copy the mutable host bits now,
+ * so a matcher can advance before GPU evaluation. Missing words reject tokens.
+ * Preserve the additive-mask expression, including NaNs and signed zero.
+ * Shape-specialized compilation handles bf16/f16/f32 and strided inputs. */
+export function applyTokenBitmask(logits: MlxArray, bits: Int32Array): MlxArray {
+  const vocab = logits.shape.at(-1);
+  if (!vocab || logits.size === 0 || logits.size > 0x7fffffff)
+    throw new Error("token bitmask requires nonempty logits with at most 2^31-1 elements");
+  if (![Dtype.bfloat16, Dtype.float16, Dtype.float32].includes(logits.dtype))
+    throw new Error("token bitmask requires bf16, f16 or f32 logits");
+  const copy = new Int32Array(Math.ceil(vocab / 32));
+  copy.set(bits.subarray(0, copy.length));
+  const deviceBits = MlxArray.fromInt32(copy, [copy.length]);
+  try {
+    compiled ??= new CompiledFunction(([x, mask]) => [maskGraph(x!, mask!)], false);
+    return compiled.apply([logits, deviceBits])[0]!;
+  } finally { deviceBits.dispose(); }
+}