mlx-bun 0.2.0 → 0.4.0

package/src/{model → lab/paged-kv}/paged-kv.ts

@@ -1,7 +1,7 @@
 // Paged KV cache — vLLM-style block-pool storage behind the standard Cache
-// interface (docs/design/paged-kv-cache.md). OPTIONAL and default-off
-// (`--paged-kv` / GenerateOptions.pagedKv); v1 scope is serial batch=1
-// decode on Gemma4-family plain full-attention layers, bf16 only.
+// interface (docs/design/kv-cache.md). OPTIONAL and default-off
+// (`--paged-kv` / GenerateOptions.pagedKv); Gemma4-family plain full-attention
+// layers, bf16 only. PagedKvRows composes independent pools at B1/B>1.
 //
 // Storage: K/V live in fixed-size per-layer pool tensors
 // [numBlocks, H_kv, blockSize, headDim]; a host-side block table maps the
@@ -11,20 +11,19 @@
 // block axis + transpose + reshape) and hand that to the unchanged
 // ops.sdpa call site — identical bytes to what a plain KVCache would have
 // fetched, so the paged path is gated bit-exact, not KL-tolerated
-// (tests/paged-kv-parity.test.ts). The gather IS the cost: a full K/V copy
+// (tests/parity/paged-kv-parity.test.ts). The gather IS the cost: a full K/V copy
 // per step, pure bandwidth tax at batch=1 — v1 ships the abstraction for
 // the batched/CoW follow-ups, not a speed win (see the design doc's
 // Motivation, which says this honestly).
 //
 // Deliberately NOT a KVCache subclass (the TurboQuantKVCache reasoning):
-// every instanceof gate in the tree — CompiledDecode.supports, the batch
-// gateway's #modelCachesBatchable, generated forwards' #matches() — must
-// EXCLUDE paged caches and fall back to the monolith/uncompiled path.
+// compiled graph adapters exclude the changing block-list geometry. The
+// cache-layout binding supplies dynamic rows independently of scheduling.
 
-import { MlxArray } from "../mlx/array";
-import { Dtype } from "../mlx/ffi";
-import * as ops from "../mlx/ops";
-import { createCausalMask, type Cache, type Mask } from "./gemma4-base";
+import { MlxArray } from "../../mlx/array";
+import { Dtype } from "../../mlx/ffi";
+import * as ops from "../../mlx/ops";
+import { createCausalMask, type Cache, type Mask } from "../../model/gemma4-base";
 
 /** Typed pool-exhaustion error: a generation outgrew its pool. Sizing from
  *  prompt+maxTokens at construction makes this unreachable in practice;
@@ -58,16 +57,24 @@ export class BlockPool {
      *  reach v1's gemma4 scope, but the pool stays shape-honest). */
     vHeadDim?: number;
     dtype: Dtype;
-  }) {
+  }, source?: BlockPool) {
     this.numBlocks = opts.numBlocks;
     this.blockSize = opts.blockSize;
-    this.keys = ops.zeros(
+    this.keys = source ? ops.contiguous(source.keys) : ops.zeros(
       [opts.numBlocks, opts.numKvHeads, opts.blockSize, opts.headDim], opts.dtype);
-    this.values = ops.zeros(
+    this.values = source ? ops.contiguous(source.values) : ops.zeros(
       [opts.numBlocks, opts.numKvHeads, opts.blockSize, opts.vHeadDim ?? opts.headDim],
       opts.dtype);
     // LIFO free list, low indices first — deterministic layout for tests.
-    this.#free = Array.from({ length: opts.numBlocks }, (_, i) => opts.numBlocks - 1 - i);
+    this.#free = source ? [...source.#free] : Array.from({ length: opts.numBlocks }, (_, i) => opts.numBlocks - 1 - i);
+  }
+
+  /** Independent ownership and allocation metadata over immutable MLX buffers.
+   * Functional writes detach storage while another snapshot retains it. */
+  clone(): BlockPool {
+    return new BlockPool({ numBlocks: this.numBlocks, blockSize: this.blockSize,
+      numKvHeads: this.keys.shape[1]!, headDim: this.keys.shape[3]!,
+      vHeadDim: this.values.shape[3]!, dtype: this.keys.dtype }, this);
   }
 
   get freeBlocks(): number {
@@ -134,6 +141,8 @@ export class BlockPool {
  *  shape): head count / head dims / dtype come from the first k/v pair,
  *  so the wiring (maybePageKv) needs no per-model shape plumbing. */
 export class PagedKVCache implements Cache {
+  /** Distinct storage kind; the row adapter preserves its block pools. */
+  signature(): string { return "kv:paged"; }
   /** Matches KVCache.STEP: v1's growth granularity is a permutation of
    *  today's 256-token step into fixed reusable slots, not a new tuning
    *  axis (--paged-kv-block-size overrides for experiments). */
@@ -152,6 +161,14 @@ export class PagedKVCache implements Cache {
     readonly blockSize: number,
   ) {}
 
+  clone(): PagedKVCache {
+    const copy = new PagedKVCache(this.capacityTokens, this.blockSize);
+    copy.offset = this.offset;
+    copy.blockTable = [...this.blockTable];
+    copy.pool = this.pool?.clone() ?? null;
+    return copy;
+  }
+
   get #blockSize(): number {
     return this.blockSize;
   }

package/src/library.ts

@@ -0,0 +1,40 @@
+/** CPU-safe high-level entry. Native modules load only after bootstrap. */
+export { createInferenceEngine } from "./engine/engine";
+export { CancellationSource } from "./engine/cancellation";
+export { createCompletionClient, createDirectHost } from "./client";
+export type * from "./contracts/generation";
+export type * from "./contracts/completion";
+export type * from "./contracts/host";
+
+/** Install/locate the native pack before importing the compatibility API. */
+export async function initializeMlx(): Promise<typeof import("./index")> {
+  const { ensureNativeRuntime } = await import("./native-pack");
+  await ensureNativeRuntime();
+  return import("./index");
+}
+
+/** Own one local worker. Model resolution and native initialization remain in
+ * the existing CLI; importing this module does not start it or download weights. */
+export async function openIsolatedHost(model: string, options: {
+  arguments?: readonly string[];
+  command?: readonly string[];
+  readyTimeoutMs?: number;
+} = {}): Promise<import("./contracts/host").EngineHost<Request, Response>> {
+  const { EngineChild } = await import("./serve/isolate");
+  const { tmpdir } = await import("node:os");
+  const { join } = await import("node:path");
+  const { fileURLToPath } = await import("node:url");
+  const source = fileURLToPath(new URL("./cli.ts", import.meta.url));
+  if (!options.command && source.includes("$bunfs"))
+    throw new Error("a compiled library consumer must supply the mlx-bun executable as command");
+  const command = options.command ?? [process.execPath, source];
+  if (!command.length) throw new Error("engine command must not be empty");
+  const socketPath = join(tmpdir(), `mlx-bun-library-${crypto.randomUUID()}.sock`);
+  const host = new EngineChild({
+    argv: [...command,
+      "serve", "--model", model, ...options.arguments ?? [], "--unix", socketPath, "--no-open"],
+    socketPath, readyTimeoutMs: options.readyTimeoutMs,
+  });
+  try { await host.ready; return host; }
+  catch (error) { await host.close(); throw error; }
+}

package/src/lora.ts

@@ -12,6 +12,7 @@
 
 import { existsSync, readdirSync, statSync } from "node:fs";
 import { resolve } from "node:path";
+import { createHash } from "node:crypto";
 import { ptr, read } from "bun:ffi";
 import { MlxArray, cpuStream } from "./mlx/array";
 import { C } from "./mlx/ffi";
@@ -198,6 +199,7 @@ export class AdapterManager {
   readonly #mounted = new Map<string, AdapterInfo>();
   /** Owned adapter arrays per id (disposed on unmount). */
   readonly #arrays = new Map<string, MlxArray[]>();
+  readonly #revisions = new Map<string, string>();
 
   constructor(model: RuntimeModel) {
     this.#model = model;
@@ -216,6 +218,11 @@ export class AdapterManager {
 
     const weightsFile = adapterWeightsFile(path);
     const { scale, rank: configRank, rsLora } = await readAdapterScale(path);
+    // Hash once at mount, before native allocation. Reusing an adapter name
+    // for different weights or scale must never reuse its prefix/checkpoint.
+    const digest = createHash("sha256").update(JSON.stringify({ scale, rank: configRank, rsLora }));
+    for await (const chunk of Bun.file(weightsFile).stream()) digest.update(chunk);
+    const revision = digest.digest("hex");
     const tensors = loadAdapterTensors(weightsFile);
     const targets = this.#model.loraTargets();
 
@@ -347,6 +354,7 @@ export class AdapterManager {
       ramBytes,
     };
     this.#mounted.set(id, info);
+    this.#revisions.set(id, revision);
     this.#arrays.set(id, validated.flatMap(({ lw }) => [lw.a, lw.b]));
     return info;
   }
@@ -355,6 +363,7 @@ export class AdapterManager {
    *  stay in place for other adapters, like the reference. */
   unmount(id: string): number {
     if (!this.#mounted.delete(id)) return 0;
+    this.#revisions.delete(id);
     let removed = 0;
     for (const linear of this.#model.loraTargets().values()) {
       if (linear.adapters?.delete(id)) removed++;
@@ -366,6 +375,15 @@ export class AdapterManager {
     return removed;
   }
 
+  /** Ordered, content-addressed namespace for the actual mounted composition. */
+  cacheNamespace(ids: readonly string[]): string {
+    return ids.map((id) => {
+      const revision = this.#revisions.get(id);
+      if (!revision) throw new Error(`adapter ${id} is no longer mounted`);
+      return `${id}@${revision}`;
+    }).join("+");
+  }
+
   list(): AdapterInfo[] {
     return [...this.#mounted.values()];
   }

package/src/media-fetch.ts

@@ -17,7 +17,7 @@
 // preserving the logical Host header and HTTPS SNI, so the transport cannot
 // re-resolve to a private address between validation and connect.
 //
-// This module must stay free of mlx imports so tests/media-fetch.test.ts
+// This module must stay free of mlx imports so tests/unit/media-fetch.test.ts
 // runs model-free.
 
 import { lookup } from "node:dns/promises";

package/src/memory/article.ts

@@ -3,7 +3,7 @@
 // The Dreaming write pipeline emits Wikipedia-shaped articles: each entity gets
 // one Markdown file with a fixed top-to-bottom skeleton —
 //   H1 → series banner → infobox → lead → sections → See also → Notes → References
-// (see docs/design/the-dreaming-master-plan.md and write-pipeline-entity.md).
+// (see docs/design/dreaming-nightly-pipeline.md and docs/design/dreaming-nightly-pipeline.md).
 // This module reads that grammar deterministically: the structured outputs the
 // read side traverses (infobox key-value facts, the [[wikilink]] graph, footnote
 // bijection) are parsed here, never by the model.

package/src/memory/chunk.ts

@@ -16,7 +16,7 @@
 //
 // Re-chunk only conversations whose updated_at > chunked_at (NULL chunked_at = a
 // never-segmented conversation). Setting chunked_at after a successful run makes
-// re-running a no-op. See docs/design/the-dreaming-master-plan.md → P5-T2.
+// re-running a no-op. See docs/design/dreaming-nightly-pipeline.md → P5-T2.
 
 import { MemoryStore, chunkId } from "./db";
 import { callLocal, type LocalInput } from "./model";

package/src/memory/model.ts

@@ -24,12 +24,8 @@
 // server and to the trainer's prompt-region render (dataset.ts:98-101).
 
 import { existsSync, readdirSync } from "node:fs";
-import { loadTaskModel, greedyDecodeBitExact, type TaskModel } from "../eval/runner";
 import type { ChatMessage, ChatTemplate } from "../chat-template";
 import type { LoadedTokenizer } from "../tokenizer";
-import { BatchScheduler, type RowSampler } from "../serve/batch-scheduler";
-import * as ops from "../mlx/ops";
-import type { MlxArray } from "../mlx/array";
 import { runtimeValue } from "../runtime-config";
 
 const HF_HUB = `${process.env.HOME}/.cache/huggingface/hub`;
@@ -38,7 +34,7 @@ const E4B_REPO = "models--mlx-community--gemma-4-e4b-it-OptiQ-4bit";
 /** Resolve the e4b snapshot dir by globbing snapshots/ for the one carrying
  *  config.json — so a freshly-downloaded model needs no hardcoded commit hash.
  *  Returns an _unresolved sentinel (precheck → STOP) until the download lands.
- *  Mirrors tests/paths.ts hfSnapshot / SNAPSHOT_E4B, inlined to keep src self-
+ *  Mirrors tests/support/paths.ts hfSnapshot / SNAPSHOT_E4B, inlined to keep src self-
  *  contained (no src→tests dependency). */
 function resolveE4bSnapshot(): string {
   const base = `${HF_HUB}/${E4B_REPO}/snapshots`;
@@ -140,30 +136,18 @@ export function memoryPromptIds(
 /** Async mutex: acquire() resolves to a release fn; releases run FIFO. Serializes
  *  memory generation ops so the global `loraState.active` is owned exclusively for
  *  one op's lifetime (set for the chunk batch, reset after). */
-class AsyncMutex {
-  #tail: Promise<void> = Promise.resolve();
-  acquire(): Promise<() => void> {
-    let release!: () => void;
-    const gate = new Promise<void>((r) => { release = r; });
-    const wait = this.#tail;
-    this.#tail = this.#tail.then(() => gate);
-    return wait.then(() => release);
-  }
-}
-
-interface MemoryRuntime {
-  tm: TaskModel;
-  template: ChatTemplate;
-  scheduler: BatchScheduler;
-  /** True once the `memory-chunk` adapter is mounted on the shared model. */
-  hasChunkAdapter: boolean;
+export interface MemoryCompletionRequest {
+  readonly stage: string;
+  readonly input: LocalInput;
+  readonly maxTokens: number;
 }
+export type MemoryCompletionClient = import("../contracts/completion").BatchCompletionClient<MemoryCompletionRequest, string>;
 
 /** Max rows in a memory batch (continuous-batching width). 1 disables batching
  *  (callLocalBatch then loops on the bit-exact greedy fallback).
  *  Default 1 (serial): batching measured 1.7-1.9x SLOWER for the real
  *  extract/chunk workload (heterogeneous prefills pad) and can diverge on
- *  near-ties — see docs/design/memory-inference-path.md "Verification results".
+ *  near-ties — see docs/design/dreaming-nightly-pipeline.md "Verification results".
  *  Opt back in with MLX_BUN_MEMORY_BATCH=8 if a length-bucketed scheduler
  *  lands. (Decision: Josh, 2026-07-01.) */
 export function memoryBatchSize(): number {
@@ -171,152 +155,29 @@ export function memoryBatchSize(): number {
   return Number.isFinite(n) && n >= 1 ? Math.floor(n) : 1;
 }
 
-function requireModel(): void {
-  if (!existsSync(`${MODEL_ID}/config.json`)) {
-    throw new Error(
-      `memory: Gemma-4-e4b is not downloaded (looked under ${HF_HUB}/${E4B_REPO}/snapshots). ` +
-        `Fetch it first:\n  HF_HUB_DISABLE_XET=1 hf download mlx-community/gemma-4-e4b-it-OptiQ-4bit`,
-    );
-  }
-}
-
-let runtimePromise: Promise<MemoryRuntime> | null = null;
-const opLock = new AsyncMutex(); // serializes whole memory ops (loraState ownership)
-const schedLock = new AsyncMutex(); // the scheduler's exclusive-GPU lock
-
-/** Lazily load the ONE shared model + mount memory-chunk + build the scheduler.
- *  Never auto-downloads; throws (with the fetch hint) until the snapshot lands. */
-async function getRuntime(): Promise<MemoryRuntime> {
-  if (!runtimePromise) {
-    runtimePromise = (async () => {
-      requireModel();
-      const tm = await loadTaskModel(MODEL_ID); // base e4b, no adapter
-      if (!tm.template) throw new Error(`memory: ${MODEL_ID} has no chat template`);
-      let hasChunkAdapter = false;
-      const chunkDir = adapterDirFor("chunk");
-      if (chunkDir) {
-        // Mount ONCE; activation is per-call via loraState.active (see decodeGreedy
-        // / batchDecode). Mounting alone runs base — it must be activated.
-        const { AdapterManager } = await import("../lora");
-        await new AdapterManager(tm.model).mount("memory-chunk", chunkDir);
-        hasChunkAdapter = true;
-      }
-      const scheduler = new BatchScheduler(tm.model, {
-        maxBatch: memoryBatchSize(),
-        lock: { acquire: () => schedLock.acquire() },
-      });
-      return { tm, template: tm.template, scheduler, hasChunkAdapter };
-    })();
-  }
-  return runtimePromise;
-}
-
-/** The adapters to activate for a stage. Only `chunk` is adapter-bound (and only
- *  when `memory-chunk` is mounted); every base stage runs with `active = []`. */
-function activeFor(stage: string, rt: MemoryRuntime): string[] {
-  return stage === "chunk" && rt.hasChunkAdapter ? ["memory-chunk"] : [];
-}
-
-/** One bit-exact greedy decode for a stage (raw-forward argmax, runner parity).
- *  Sets the stage's adapter for the call and resets `active` to [] after. */
-function decodeGreedy(rt: MemoryRuntime, stage: string, input: LocalInput, maxTokens: number): string {
-  const active = activeFor(stage, rt);
-  rt.tm.model.loraState.active = active;
-  const ids = memoryPromptIds(stage, input, rt.tm.tokenizer, rt.template);
-  try {
-    return greedyDecodeBitExact(
-      { ...rt.tm, activeAdapters: active.length ? active : undefined },
-      ids,
-      maxTokens,
-    );
-  } finally {
-    rt.tm.model.loraState.active = [];
-  }
-}
-
-/** Batched decode for N independent inputs of ONE stage. Sets the stage adapter
- *  globally for the whole batch (the scheduler reads `loraState.active` per
- *  forward, so all rows get it uniformly — a batch is necessarily single-adapter),
- *  submits every row to the shared scheduler with a greedy per-row sampler, and
- *  resets `active` after all rows finish. Order-preserving: out[i] ↔ inputs[i].
- *
- *  Decode-path caveat: this samples via forwardHidden/logitsFromHidden (the
- *  scheduler), which can diverge from the runner's raw-forward greedy on near-ties
- *  past ~32 tokens. For the JSON/binary memory stages (validated by
- *  well-formedness) this is acceptable; the single-call fallback stays bit-exact. */
-async function batchDecode(
-  rt: MemoryRuntime,
-  stage: string,
-  inputs: LocalInput[],
-  maxTokens: number,
-): Promise<string[]> {
-  const active = activeFor(stage, rt);
-  rt.tm.model.loraState.active = active;
-  const eos = rt.tm.config.eosTokenIds;
-  const greedySample: RowSampler = (logits1V: MlxArray) => ops.argmaxAxis(logits1V, -1);
-  try {
-    return await Promise.all(
-      inputs.map((inp) => {
-        const ids = memoryPromptIds(stage, inp, rt.tm.tokenizer, rt.template);
-        const toks: number[] = [];
-        return rt.scheduler
-          .submit({
-            promptIds: ids,
-            maxTokens,
-            eosTokenIds: eos,
-            sample: greedySample,
-            onToken: (t) => { toks.push(t); },
-          })
-          .then(() => rt.tm.tokenizer.decode(toks, true));
-      }),
-    );
-  } finally {
-    rt.tm.model.loraState.active = [];
-  }
-}
-
-// ---------------------------------------------------------------------------
-// Public seams
-// ---------------------------------------------------------------------------
 
-/** Call the local model for a pipeline `stage` with a {system?, user} input.
- *  Renders through the shared chat-template path, runs the MODEL precheck, and
- *  decodes ONE completion on the bit-exact greedy path (the safe fallback used by
- *  every single-call site). NEVER auto-downloads. */
-export async function callLocal(
-  stage: string,
-  input: LocalInput,
-  opts?: { maxTokens?: number },
-): Promise<string> {
-  const rt = await getRuntime();
-  const release = await opLock.acquire();
-  try {
-    return decodeGreedy(rt, stage, input, opts?.maxTokens ?? 256);
-  } finally {
-    release();
-  }
+/** One configured client per pipeline. Importing this module does not load MLX. */
+export function createMemoryCalls(client: MemoryCompletionClient) {
+  return {
+    callLocal(stage: string, input: LocalInput, opts?: { maxTokens?: number }): Promise<string> {
+      return client.complete({ stage, input, maxTokens: opts?.maxTokens ?? 256 });
+    },
+    callLocalBatch(stage: string, inputs: LocalInput[], opts?: { maxTokens?: number }): Promise<string[]> {
+      if (!inputs.length) return Promise.resolve([]);
+      return client.completeBatch(inputs.map((input) => ({ stage, input, maxTokens: opts?.maxTokens ?? 256 })));
+    },
+  };
 }
 
-/** N independent calls for ONE stage (homogeneous adapter by construction).
- *  Routed through the shared in-process BatchScheduler when batch>1 AND N>1, else
- *  looped on the bit-exact greedy fallback. Order-preserving: out[i] ↔ inputs[i]. */
-export async function callLocalBatch(
-  stage: string,
-  inputs: LocalInput[],
-  opts?: { maxTokens?: number },
-): Promise<string[]> {
-  if (inputs.length === 0) return [];
-  const rt = await getRuntime();
-  const maxTokens = opts?.maxTokens ?? 256;
-  const release = await opLock.acquire();
-  try {
-    if (inputs.length === 1 || memoryBatchSize() <= 1) {
-      const out: string[] = [];
-      for (const inp of inputs) out.push(decodeGreedy(rt, stage, inp, maxTokens));
-      return out;
-    }
-    return await batchDecode(rt, stage, inputs, maxTokens);
-  } finally {
-    release();
-  }
-}
+const defaultCalls = createMemoryCalls({
+  async complete(request) {
+    const { memoryCompletionClient } = await import("../backends/mlx/memory-client");
+    return memoryCompletionClient.complete(request);
+  },
+  async completeBatch(requests) {
+    const { memoryCompletionClient } = await import("../backends/mlx/memory-client");
+    return memoryCompletionClient.completeBatch(requests);
+  },
+});
+export const callLocal = defaultCalls.callLocal;
+export const callLocalBatch = defaultCalls.callLocalBatch;

package/src/memory/prompts.ts

@@ -7,7 +7,7 @@
 // Topics_to_Ignore.md) and are read in at synthesis time so the system governs
 // itself in its own substrate.
 //
-// See docs/design/memory-system.md → "Synthesis runs on the local model".
+// See docs/design/dreaming-nightly-pipeline.md → "Synthesis runs on the local model".
 
 import { readFileSync } from "node:fs";
 import { join } from "node:path";

package/src/memory/rest.ts

@@ -11,7 +11,7 @@
 // rest of server.ts's /api/* conventions. Route dispatch (matching
 // url.pathname + method) lives in server.ts; this file is the pure handler
 // bodies so they're unit-testable without booting a model context (see
-// tests/memory-rest.test.ts).
+// tests/using/memory-rest.test.ts).
 
 import { realpath } from "node:fs/promises";
 import { tmpdir } from "node:os";

package/src/memory/schedule.ts

@@ -11,7 +11,7 @@
 // it. That job runs the full synthesis DAG and produces/updates articles on the
 // schedule with no further configuration.
 //
-// See docs/design/memory-system.md → "Scheduling (launchd)".
+// See docs/design/dreaming-nightly-pipeline.md → "Scheduling (launchd)".
 
 import { spawn } from "node:child_process";
 import { mkdir, rm, writeFile } from "node:fs/promises";

package/src/memory/stages.ts

@@ -356,7 +356,10 @@ export async function runExtractStage(
   // SERIAL fallback (injected `call`, single chunk, or batching disabled).
   let extracted = 0;
   for (const id of targetIds) {
-    await extractEntities(store, id, { resolver, policy: policy || undefined, call: opts.call });
+    // Pass the (possibly empty) policy through: `|| undefined` made extractEntities
+    // re-load Meta/Entities.md and THROW on a vault without it, while the batched
+    // path above proceeds with "" — the serial path must match.
+    await extractEntities(store, id, { resolver, policy, call: opts.call });
     extracted++;
     if (extracted % 10 === 0) {
       opts.onEvent?.({ type: "log", message: `  extracted ${extracted}/${target} chunks` });

package/src/memory/vault.ts

@@ -9,7 +9,7 @@
 // an already-working system at real scale — they're pure filesystem functions
 // with no MCP coupling, so they drop straight in. We keep lucien's good
 // descriptive concepts (articles, [[wikilinks]], Meta governance pages,
-// conv: citations) and drop the mythology. See docs/design/memory-system.md.
+// conv: citations) and drop the mythology. See docs/design/dreaming-nightly-pipeline.md.
 //
 // M0 is the read path: setup the vault, then list / read / search / index it.
 // Synthesis (writing articles from conversations) is M1 and lives elsewhere.
@@ -987,10 +987,8 @@ const REFERENCE_LINKS: { name: string; source: string }[] = [
   { name: "mlx-bun_Server_API", source: join(REPO_ROOT, "docs", "reference", "server-api.md") },
   { name: "mlx-bun_Server_Config", source: join(REPO_ROOT, "docs", "reference", "server-config.md") },
   { name: "mlx-bun_Library_API", source: join(REPO_ROOT, "docs", "reference", "library-api.md") },
-  { name: "mlx-bun_Embedding", source: join(REPO_ROOT, "docs", "reference", "embedding.md") },
   { name: "mlx-bun_Distribution", source: join(REPO_ROOT, "docs", "reference", "distribution.md") },
   { name: "mlx-bun_Training", source: join(REPO_ROOT, "docs", "reference", "training.md") },
-  { name: "mlx-bun_ORPO_Quickstart", source: join(REPO_ROOT, "docs", "reference", "orpo-quickstart.md") },
   { name: "mlx-bun_Product_Roadmap", source: join(REPO_ROOT, "docs", "planning", "PRODUCT_ROADMAP.md") },
 ];
 

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;
@@ -233,7 +232,7 @@ export class MlxArray {
     );
   }
 
-  /** Raw bytes of the evaluated array (copy). bf16/f16/f32 only. */
+  /** Raw bytes of a supported floating-point or integer array (copy). */
   rawBytes(): Uint8Array {
     return this.rawBytesView().slice();
   }
@@ -254,6 +253,7 @@ 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.int32 ? C.mlx_array_data_int32(this.handle)
       : dt === Dtype.uint8 ? C.mlx_array_data_uint8(this.handle)
       : dt === Dtype.int8 ? C.mlx_array_data_int8(this.handle)
       : null;

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);