mlx-bun 0.0.4

package/src/mlx/ffi.ts

@@ -0,0 +1,433 @@
+// Raw bun:ffi bindings to libmlxc.
+//
+// Calling convention (verified in Phase 0, see PLAN.md):
+// - Every mlx-c handle is a one-pointer struct `{ void* ctx }` — passed and
+//   returned by value in a single register on arm64, so FFIType.u64 works.
+// - Out-params (`mlx_array* res`) are a pointer to that slot: allocate a
+//   BigUint64Array(1), pass ptr(), reread after the call via read.u64 —
+//   NOT slot[0] (DFG stale-read bug, see rule above outArray).
+// - Ops return int status; 0 = success.
+// - 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 { existsSync } from "node:fs";
+import { dirname, join } from "node:path";
+import { nativePackDir } from "../native-pack";
+
+const { ptr: P, i32, u64, f32, 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
+ *  (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;
+  if (env) return env;
+  const beside = join(dirname(process.execPath), "libmlxc.dylib");
+  if (existsSync(beside)) return beside;
+  const cached = join(nativePackDir(), "libmlxc.dylib");
+  if (existsSync(cached)) return cached;
+  for (const p of ["/opt/homebrew/lib/libmlxc.dylib", "/usr/local/lib/libmlxc.dylib"])
+    if (existsSync(p)) return p;
+  // let dlopen produce the canonical error message for the default path
+  return "/opt/homebrew/lib/libmlxc.dylib";
+}
+
+export const LIBMLXC_PATH = resolveLibmlxc();
+
+export const C = dlopen(LIBMLXC_PATH, {
+  // array lifecycle
+  mlx_array_new: { args: [], returns: u64 },
+  mlx_array_new_data: { args: [P, P, i32, i32], returns: u64 },
+  mlx_array_new_data_managed_payload: { args: [P, P, i32, i32, P, P], returns: u64 },
+  mlx_array_free: { args: [u64], returns: i32 },
+  mlx_array_eval: { args: [u64], returns: i32 },
+  // array metadata
+  mlx_array_ndim: { args: [u64], returns: u64 },
+  mlx_array_shape: { args: [u64], returns: P },
+  mlx_array_dim: { args: [u64, i32], returns: i32 },
+  mlx_array_dtype: { args: [u64], returns: i32 },
+  mlx_array_size: { args: [u64], returns: u64 },
+  mlx_array_nbytes: { args: [u64], returns: u64 },
+  // data access (forces eval/copy to host-visible memory)
+  mlx_array_data_float32: { args: [u64], returns: P },
+  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_item_float32: { args: [P, u64], returns: i32 },
+  // streams
+  mlx_default_gpu_stream_new: { args: [], returns: u64 },
+  mlx_default_cpu_stream_new: { args: [], returns: u64 },
+  mlx_stream_free: { args: [u64], returns: i32 },
+  // eval control (both take mlx_vector_array)
+  mlx_eval: { args: [u64], returns: i32 },
+  mlx_async_eval: { args: [u64], returns: i32 },
+  // memory introspection / limits
+  mlx_get_active_memory: { args: [P], returns: i32 },
+  mlx_get_peak_memory: { args: [P], returns: i32 },
+  mlx_reset_peak_memory: { args: [], returns: i32 },
+  // mlx_device_info is a one-pointer struct like every other handle
+  mlx_device_info_new: { args: [], returns: u64 },
+  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_get_default_device: { args: [P], returns: i32 },
+  mlx_device_free: { args: [u64], returns: i32 },
+  mlx_synchronize: { args: [u64], returns: i32 },
+  mlx_get_cache_memory: { args: [P], returns: i32 },
+  mlx_clear_cache: { args: [], returns: i32 },
+  mlx_set_memory_limit: { args: [P, u64], returns: i32 },
+  mlx_set_wired_limit: { args: [P, u64], returns: i32 },
+  // io: native lazy safetensors loader (Load primitive is CPU-only —
+  // always pass the CPU stream; GPU ops consume the unified buffers)
+  mlx_load_safetensors: { args: [P, P, cstring, u64], returns: i32 },
+  mlx_map_string_to_array_new: { args: [], returns: u64 },
+  mlx_map_string_to_array_free: { args: [u64], returns: i32 },
+  mlx_map_string_to_array_get: { args: [P, u64, cstring], returns: i32 },
+  mlx_map_string_to_string_new: { args: [], returns: u64 },
+  mlx_map_string_to_string_free: { args: [u64], returns: i32 },
+  // fast fused ops
+  // (res, x, weight may-be-null, eps, stream)
+  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_matmul: { args: [P, u64, u64, u64], returns: i32 },
+  mlx_contiguous: { args: [P, u64, FFIType.bool, 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 },
+  mlx_vector_array_new: { args: [], returns: u64 },
+  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 },
+  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 },
+  mlx_multiply: { args: [P, u64, u64, u64], returns: i32 },
+  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_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 },
+  mlx_transpose_axes: { args: [P, u64, P, u64, u64], returns: i32 },
+  mlx_argmax_axis: { args: [P, u64, i32, FFIType.bool, u64], returns: i32 },
+  mlx_concatenate_axis: { args: [P, u64, i32, u64], returns: i32 },
+  mlx_vector_array_new_data: { args: [P, u64], returns: u64 },
+  mlx_vector_array_free: { args: [u64], returns: i32 },
+  mlx_array_item_uint32: { args: [P, u64], returns: i32 },
+  mlx_array_item_int32: { args: [P, u64], returns: i32 },
+  mlx_slice: { args: [P, u64, P, u64, P, u64, P, u64, u64], returns: i32 },
+  // (res, src, update, start*, n, stop*, n, strides*, n, stream)
+  mlx_slice_update: { args: [P, u64, u64, P, u64, P, u64, P, u64, u64], returns: i32 },
+  mlx_zeros: { args: [P, P, u64, i32, u64], returns: i32 },
+  mlx_exp: { args: [P, u64, u64], returns: i32 },
+  mlx_cos: { args: [P, u64, u64], returns: i32 },
+  mlx_sin: { args: [P, u64, u64], returns: i32 },
+  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_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_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_logical_and: { args: [P, u64, u64, 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 },
+  // (res, w, scales, biases, opt group_size, opt bits, mode, global_scale, opt dtype, stream)
+  mlx_dequantize: { args: [P, u64, u64, u64, u64, u64, cstring, u64, u64, u64], returns: i32 },
+  mlx_quantized_matmul: { args: [P, u64, u64, u64, u64, FFIType.bool, u64, u64, cstring, u64], returns: i32 },
+  // (res, x, w, scales, biases?, lhs_indices?, rhs_indices?, transpose,
+  //  opt group_size, opt bits, mode, sorted_indices, stream)
+  mlx_gather_qmm: { args: [P, u64, u64, u64, u64, u64, u64, FFIType.bool, u64, u64, cstring, FFIType.bool, u64], returns: i32 },
+  mlx_floor_divide: { args: [P, u64, u64, u64], returns: i32 },
+  // (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.
+  // 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).
+  mlx_value_and_grad: { args: [P, u64, P, u64], returns: i32 },
+  mlx_closure_value_and_grad_new: { args: [], returns: u64 },
+  // (value: vector_array*, grads: vector_array*, cls, input: vector_array)
+  mlx_closure_value_and_grad_apply: { args: [P, P, u64, u64], returns: i32 },
+  mlx_closure_value_and_grad_free: { args: [u64], returns: i32 },
+  // mlx_vjp(outputs: vector_array*, vjps: vector_array*, fun: mlx_closure,
+  //   primals: vector_array, cotangents: vector_array) — reverse-mode AD given an
+  //   explicit cotangent per output. The natural segmented-backward primitive:
+  //   fun = segment forward, cotangent = dh, vjps = [dh_in, ...dLoRA]. No
+  //   surrogate scalar, no value_and_grad closure (see src/train/segmented.ts).
+  mlx_vjp: { args: [P, P, u64, u64, u64], returns: i32 },
+  // --- training: random init (LoRA A ~ normal; B stays zeros)
+  // (res, shape*, n, dtype, loc, scale, key may-be-null, stream)
+  mlx_random_normal: { args: [P, P, u64, i32, f32, f32, u64, u64], returns: i32 },
+  // (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_split_num: { args: [P, u64, i32, u64], returns: i32 },
+  // --- training: loss/optimizer elementwise + reductions
+  mlx_stop_gradient: { args: [P, u64, u64], returns: i32 },
+  mlx_log: { args: [P, u64, u64], returns: i32 },
+  mlx_logaddexp: { args: [P, u64, u64, u64], returns: i32 },
+  mlx_mean: { args: [P, u64, FFIType.bool, u64], returns: i32 },
+  mlx_mean_axis: { args: [P, u64, i32, FFIType.bool, u64], returns: i32 },
+  mlx_sqrt: { args: [P, u64, u64], returns: i32 },
+  mlx_rsqrt: { args: [P, u64, u64], returns: i32 },
+  mlx_square: { args: [P, u64, u64], returns: i32 },
+  mlx_abs: { args: [P, u64, u64], returns: i32 },
+  // (res, a, a_min may-be-null, a_max may-be-null, stream)
+  mlx_clip: { args: [P, u64, u64, u64, u64], returns: i32 },
+  // --- quantization: native safetensors writer (io.h) + map inserts (map.h)
+  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).
+  // 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 },
+  mlx_closure_new_func_payload: { args: [P, P, P], returns: u64 },
+  mlx_closure_free: { args: [u64], returns: i32 },
+  mlx_closure_apply: { args: [P, u64, u64], returns: i32 },
+  // gradient checkpointing (transforms.h): wrap a closure so its forward
+  // activations are dropped and recomputed during the enclosing backward.
+  // mlx_checkpoint(res: mlx_closure*, fun: mlx_closure) → checkpointed closure.
+  mlx_checkpoint: { args: [P, u64], returns: i32 },
+  // custom_vjp (transforms.h): attach a hand-written backward to a forward whose
+  // primitives have no usable vjp (e.g. a fused metal_kernel). Needed for the L2
+  // flash-attention training op. mlx_custom_vjp(res: mlx_closure*, fun: mlx_closure,
+  //   fun_vjp: mlx_closure_custom). The custom-vjp closure's callback is
+  //   (res*, primals, cotangents, outputs, payload) → status.
+  mlx_custom_vjp: { args: [P, u64, u64], returns: i32 },
+  mlx_closure_custom_new_func_payload: { args: [P, P, P], returns: u64 },
+  mlx_closure_custom_free: { args: [u64], returns: i32 },
+  mlx_compile: { args: [P, u64, FFIType.bool], returns: i32 },
+  mlx_set_compile_mode: { args: [i32], returns: i32 },
+  mlx_vector_array_size: { args: [u64], returns: u64 },
+  mlx_vector_array_set_data: { args: [P, P, u64], returns: i32 },
+  // dynamic-start variants: per-decode-step offsets enter the graph as
+  // array VALUES instead of baked ints, so one compiled graph serves
+  // every step (no per-step retrace).
+  // (res, x, dims, traditional, opt base, scale, offset ARRAY, freqs may-be-null, stream)
+  mlx_fast_rope_dynamic: { args: [P, u64, i32, FFIType.bool, u64, f32, u64, u64, u64], returns: i32 },
+  // (res, src, update, start ARRAY, axes*, axes_num, stream)
+  mlx_slice_update_dynamic: { args: [P, u64, u64, u64, P, u64, u64], returns: i32 },
+  // (res, a, start ARRAY, axes*, axes_num, slice_size*, n, stream)
+  mlx_slice_dynamic: { args: [P, u64, u64, P, u64, P, u64, u64], returns: i32 },
+  // custom Metal kernels (Phase E — mx.fast.metal_kernel from Bun)
+  mlx_vector_string_new_data: { args: [P, u64], returns: u64 },
+  mlx_vector_string_free: { args: [u64], returns: i32 },
+  // (name, input_names, output_names, source, header, ensure_row_contiguous, atomic_outputs)
+  mlx_fast_metal_kernel_new: { args: [cstring, u64, u64, cstring, cstring, FFIType.bool, FFIType.bool], returns: u64 },
+  mlx_fast_metal_kernel_free: { args: [u64], returns: FFIType.void },
+  // (outputs: mlx_vector_array*, kernel, inputs, config, stream)
+  mlx_fast_metal_kernel_apply: { args: [P, u64, u64, u64, u64], returns: i32 },
+  mlx_fast_metal_kernel_config_new: { args: [], returns: u64 },
+  mlx_fast_metal_kernel_config_free: { args: [u64], returns: FFIType.void },
+  mlx_fast_metal_kernel_config_add_output_arg: { args: [u64, P, u64, i32], returns: i32 },
+  mlx_fast_metal_kernel_config_set_grid: { args: [u64, i32, i32, i32], returns: i32 },
+  mlx_fast_metal_kernel_config_set_thread_group: { args: [u64, i32, i32, i32], returns: i32 },
+  mlx_fast_metal_kernel_config_set_init_value: { args: [u64, f32], returns: i32 },
+  mlx_fast_metal_kernel_config_add_template_arg_int: { args: [u64, cstring, i32], returns: i32 },
+  mlx_fast_metal_kernel_config_add_template_arg_dtype: { args: [u64, cstring, i32], returns: i32 },
+  // 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 },
+}).symbols;
+
+export type MlxHandle = bigint;
+
+// --- error handling -------------------------------------------------------
+// mlx-c's default error handler prints and aborts the process. Install a
+// recording handler instead: ops then return non-zero status and outArray
+// throws a JS Error carrying the mlx message + JS stack. mlx-c invokes the
+// handler on the calling thread (exceptions are caught in the C wrappers),
+// so a non-threadsafe JSCallback is safe here.
+import { CString, JSCallback } from "bun:ffi";
+
+let lastMlxError: string | null = null;
+
+const errorHandler = new JSCallback(
+  (msgPtr: number) => {
+    lastMlxError = msgPtr ? new CString(msgPtr as never).toString() : "unknown mlx error";
+  },
+  { args: ["ptr", "ptr"], returns: "void" },
+);
+
+const errLib = dlopen(LIBMLXC_PATH, {
+  mlx_set_error_handler: { args: [P, P, P], returns: FFIType.void },
+}).symbols;
+errLib.mlx_set_error_handler(errorHandler.ptr, null, null);
+
+/** Consume the last recorded mlx error message (if any). */
+export function takeMlxError(): string | null {
+  const e = lastMlxError;
+  lastMlxError = null;
+  return e;
+}
+
+/** mlx_dtype enum values (array.h). */
+export const enum Dtype {
+  bool = 0, uint8 = 1, uint16 = 2, uint32 = 3, uint64 = 4,
+  int8 = 5, int16 = 6, int32 = 7, int64 = 8,
+  float16 = 9, float32 = 10, float64 = 11, bfloat16 = 12, complex64 = 13,
+}
+
+export const DTYPE_NAMES: Record<number, string> = {
+  0: "bool", 1: "uint8", 2: "uint16", 3: "uint32", 4: "uint64",
+  5: "int8", 6: "int16", 7: "int32", 8: "int64",
+  9: "float16", 10: "float32", 11: "float64", 12: "bfloat16", 13: "complex64",
+};
+
+/** Pack mlx_optional_int for by-value passing. */
+export function optInt(value: number | null): bigint {
+  if (value === null) return 0n;
+  return (BigInt(value >>> 0) & 0xffffffffn) | (1n << 32n);
+}
+
+/** Pack mlx_optional_dtype (same layout: enum + has_value flag). */
+export const optDtype = optInt;
+
+/** Pack mlx_optional_float `{float value; bool has_value}` for by-value passing. */
+const f32Bits = new DataView(new ArrayBuffer(4));
+export function optFloat(value: number | null): bigint {
+  if (value === null) return 0n;
+  f32Bits.setFloat32(0, value, true);
+  return BigInt(f32Bits.getUint32(0, true)) | (1n << 32n);
+}
+
+/** Null handle for optional mlx_array parameters. */
+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
+// 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.
+
+/** Run an op that writes its result through an mlx_array* out-param. */
+export function outArray(op: string, call: (slotPtr: number) => number): MlxHandle {
+  const slot = new BigUint64Array([C.mlx_array_new()]);
+  const slotPtr = ptr(slot);
+  const status = call(slotPtr);
+  const handle = read.u64(slotPtr, 0);
+  if (status !== 0) {
+    C.mlx_array_free(handle);
+    throw new Error(`${op} failed: ${takeMlxError() ?? "(no mlx message)"}`);
+  }
+  return handle;
+}
+
+export function activeMemory(): number {
+  const out = new BigUint64Array(1);
+  const p = ptr(out);
+  C.mlx_get_active_memory(p);
+  return Number(read.u64(p, 0));
+}
+
+export function cacheMemory(): number {
+  const out = new BigUint64Array(1);
+  const p = ptr(out);
+  C.mlx_get_cache_memory(p);
+  return Number(read.u64(p, 0));
+}
+
+export function peakMemory(): number {
+  const out = new BigUint64Array(1);
+  const p = ptr(out);
+  C.mlx_get_peak_memory(p);
+  return Number(read.u64(p, 0));
+}
+
+/** mx.reset_peak_memory — zero the peak counter (e.g. after model load,
+ *  so a subsequent peakMemory() reads generation-only peak). */
+export function resetPeakMemory(): void {
+  C.mlx_reset_peak_memory();
+}
+
+/** mx.clear_cache — release the allocator's cached buffers back to the
+ *  OS. mlx-lm calls this after every prefill chunk (and every 256 decode
+ *  steps); without it the first decode step after a long prefill pays a
+ *  one-shot allocator-reclaim stall (~800 ms after an 8k prefill —
+ *  measured, scripts/decode-split.ts; the root cause of the
+ *  context-scaling decode gap). */
+export function clearCache(): void {
+  C.mlx_clear_cache();
+}
+
+/** Metal's recommended max working-set size for the default device. */
+export function maxRecommendedWorkingSetSize(): number {
+  const devSlot = new BigUint64Array(1);
+  const 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()]);
+  const infoPtr = ptr(infoSlot);
+  try {
+    if (C.mlx_device_info_get(infoPtr, dev) !== 0)
+      throw new Error(`mlx_device_info_get failed: ${takeMlxError() ?? ""}`);
+    const info = read.u64(infoPtr, 0);
+    const out = new BigUint64Array(1);
+    const outPtr = ptr(out);
+    const key = Buffer.from("max_recommended_working_set_size\0", "utf8");
+    if (C.mlx_device_info_get_size(outPtr, info, ptr(key)) !== 0)
+      throw new Error("max_recommended_working_set_size not in device info");
+    return Number(read.u64(outPtr, 0));
+  } finally {
+    C.mlx_device_info_free(read.u64(infoPtr, 0));
+    C.mlx_device_free(dev);
+  }
+}
+
+/** mx.set_wired_limit — returns the previous limit. Models close to the
+ *  working-set ceiling decode ~4x slower without this (Metal evicts and
+ *  re-faults weight buffers every token; mlx-lm's wired_limit context
+ *  is the reference behavior). MUST be scoped (set → generate →
+ *  synchronize → restore), never process-permanent: a permanent limit
+ *  pins buffers across idle periods and OOM-kills multi-model processes
+ *  (the test suite) via uncatchable async GPU errors. */
+export function setWiredLimit(bytes: number): number {
+  const out = new BigUint64Array(1);
+  const p = ptr(out);
+  if (C.mlx_set_wired_limit(p, BigInt(Math.floor(bytes))) !== 0)
+    throw new Error(`mlx_set_wired_limit failed: ${takeMlxError() ?? ""}`);
+  return Number(read.u64(p, 0));
+}
+
+/** mx.set_memory_limit — returns the previous limit. At the limit the
+ *  allocator reclaims/waits instead of ballooning past it. Defense in
+ *  depth under admission control ONLY: it does NOT make Metal
+ *  command-buffer OOM catchable (Phase 6 finding — that throw comes
+ *  from a completion handler and is std::terminate). */
+export function setMemoryLimit(bytes: number): number {
+  const out = new BigUint64Array(1);
+  const p = ptr(out);
+  if (C.mlx_set_memory_limit(p, BigInt(Math.floor(bytes))) !== 0)
+    throw new Error(`mlx_set_memory_limit failed: ${takeMlxError() ?? ""}`);
+  return Number(read.u64(p, 0));
+}
+
+/** mx.synchronize(stream) — wired limit must not change mid-async-eval. */
+export function synchronize(stream: MlxHandle): void {
+  if (C.mlx_synchronize(stream) !== 0)
+    throw new Error(`mlx_synchronize failed: ${takeMlxError() ?? ""}`);
+}

package/src/mlx/metal-kernel.ts

@@ -0,0 +1,130 @@
+// mx.fast.metal_kernel from Bun (docs/design/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
+// `*_shape`/`*_strides`/`*_ndim` helpers when referenced). Dispatched
+// like any lazy op: apply() builds graph nodes, eval runs the kernel.
+
+import { ptr, read } from "bun:ffi";
+import { C, Dtype, type MlxHandle, takeMlxError } from "./ffi";
+import { MlxArray, gpuStream } from "./array";
+
+function vectorString(strings: string[]): { vec: MlxHandle; keepAlive: Buffer[] } {
+  const bufs = strings.map((s) => Buffer.from(s + "\0", "utf8"));
+  const ptrs = new BigUint64Array(bufs.map((b) => BigInt(ptr(b))));
+  const vec = C.mlx_vector_string_new_data(ptr(ptrs), BigInt(bufs.length));
+  return { vec, keepAlive: bufs };
+}
+
+export interface MetalKernelSpec {
+  name: string;
+  inputNames: string[];
+  outputNames: string[];
+  /** Kernel BODY (not a full function) in MSL. */
+  source: string;
+  /** Optional includes/defines prepended outside the kernel. */
+  header?: string;
+  ensureRowContiguous?: boolean;
+  atomicOutputs?: boolean;
+}
+
+export interface MetalKernelCall {
+  outputs: { shape: number[]; dtype: Dtype }[];
+  grid: [number, number, number];
+  threadGroup: [number, number, number];
+  templateInts?: Record<string, number>;
+  templateDtypes?: Record<string, Dtype>;
+  initValue?: number;
+  stream?: MlxHandle;
+}
+
+export class MetalKernel {
+  #kernel: MlxHandle;
+  #disposed = false;
+
+  constructor(spec: MetalKernelSpec) {
+    const inNames = vectorString(spec.inputNames);
+    const outNames = vectorString(spec.outputNames);
+    const name = Buffer.from(spec.name + "\0", "utf8");
+    const source = Buffer.from(spec.source + "\0", "utf8");
+    const header = Buffer.from((spec.header ?? "") + "\0", "utf8");
+    this.#kernel = C.mlx_fast_metal_kernel_new(
+      ptr(name), inNames.vec, outNames.vec, ptr(source), ptr(header),
+      spec.ensureRowContiguous ?? true, spec.atomicOutputs ?? false,
+    );
+    C.mlx_vector_string_free(inNames.vec);
+    C.mlx_vector_string_free(outNames.vec);
+  }
+
+  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();
+    try {
+      for (const o of call.outputs) {
+        const shape = new Int32Array(o.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_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");
+        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");
+        C.mlx_fast_metal_kernel_config_add_template_arg_dtype(cfg, ptr(kb), v);
+      }
+
+      const handles = new BigUint64Array(inputs.map((a) => a.handle));
+      const inVec = C.mlx_vector_array_new_data(ptr(handles), BigInt(inputs.length));
+      const slot = new BigUint64Array([C.mlx_vector_array_new()]);
+      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);
+      C.mlx_vector_array_free(inVec);
+      if (status !== 0) {
+        C.mlx_vector_array_free(outVec);
+        throw new Error(`metal_kernel apply failed: ${takeMlxError() ?? "(no message)"}`);
+      }
+      const n = Number(C.mlx_vector_array_size(outVec));
+      const outs: MlxArray[] = [];
+      const aSlot = new BigUint64Array(1);
+      const aPtr = ptr(aSlot);
+      for (let i = 0; i < n; i++) {
+        aSlot[0] = C.mlx_array_new();
+        if (C.mlx_vector_array_get(aPtr, outVec, BigInt(i)) !== 0) {
+          C.mlx_vector_array_free(outVec);
+          throw new Error(`metal_kernel output get(${i}) failed: ${takeMlxError() ?? ""}`);
+        }
+        outs.push(new MlxArray(read.u64(aPtr, 0)));
+      }
+      C.mlx_vector_array_free(outVec);
+      return outs;
+    } finally {
+      C.mlx_fast_metal_kernel_config_free(cfg);
+    }
+  }
+
+  dispose(): void {
+    if (this.#disposed) return;
+    this.#disposed = true;
+    C.mlx_fast_metal_kernel_free(this.#kernel);
+  }
+}
+
+/** mlx_metal_start/stop_capture — wrap a single decode step to inspect
+ *  its command buffer in Xcode (Phase E step 2: size the prize). Needs
+ *  MTL_CAPTURE_ENABLED=1 in the environment. */
+export function metalCapture(path: string, fn: () => void): void {
+  const p = Buffer.from(path + "\0", "utf8");
+  if (C.mlx_metal_start_capture(ptr(p)) !== 0)
+    throw new Error(`metal_start_capture failed: ${takeMlxError() ?? ""}`);
+  try {
+    fn();
+  } finally {
+    C.mlx_metal_stop_capture();
+  }
+}