mlx-bun 0.2.0 → 0.4.0

package/src/engine/execution-plan.ts

@@ -0,0 +1,53 @@
+import type { ExecutionCapabilities, ExecutionFeatures, ExecutionRequirements, ResolvedExecution } from "../contracts/execution";
+
+/** Select the request method, then place it using executor capabilities. */
+export function resolveExecution(
+  request: ExecutionRequirements, capabilities: ExecutionCapabilities,
+  features: ExecutionFeatures = { pagedKv: false, fill: false },
+): ResolvedExecution {
+  const reasons: string[] = [];
+  const sharedRequestExclusions = [
+    [!capabilities.continuous, "continuous-unavailable"],
+    [request.hasVision, "media-requires-serial"],
+    [request.hasAdapters && !capabilities.adapterBatch, "adapters-require-serial"],
+    [request.kvQuant && !capabilities.quantizedBatch, "kv-scheme-requires-serial"],
+    [request.turboQuant && !capabilities.turboQuantBatch, "turbo-kv-requires-serial"],
+    [request.hasGrammar && !capabilities.grammarBatch, "grammar-batching-disabled"],
+    [features.pagedKv && !capabilities.pagedBatch, "paged-kv-requires-serial"],
+  ] as const;
+  const sharedAdapterMethod = capabilities.sharedSpeculativeAdapters === true &&
+    !sharedRequestExclusions.some(([excluded]) => excluded) &&
+    capabilities.groupedMethods?.includes("speculative") === true;
+  const speculative = capabilities.method === "autoregressive" && request.hasDraft &&
+    !request.hasVision && (!request.hasAdapters || sharedAdapterMethod) && (!request.wantsLogprobs || capabilities.speculativeLogprobs === true) &&
+    (!request.kvQuant || capabilities.speculativeKvQuant === true) &&
+    (!request.turboQuant || capabilities.speculativeTurboQuant === true) && !features.pagedKv;
+  const method = speculative ? "speculative" : capabilities.method;
+  const continuousExclusions = [sharedRequestExclusions[0],
+    [!(capabilities.groupedMethods ?? ["autoregressive"]).includes(method), "method-requires-serial"],
+    ...sharedRequestExclusions.slice(1),
+  ] as const;
+  for (const [excluded, reason] of continuousExclusions) if (excluded) reasons.push(reason);
+  const mechanism = reasons.length ? "serial" : "continuous";
+  if (request.hasDraft && !speculative) reasons.push("draft-incompatible-with-request");
+  const pagedKv = features.pagedKv && !request.hasVision && !request.hasAdapters;
+  if (features.pagedKv && !pagedKv) reasons.push("paged-kv-bypassed-for-media-or-adapters");
+  const promptCache = method !== "speculative" && !request.hasVision && !pagedKv;
+  const fill = features.fill && method === "autoregressive" && mechanism === "serial" &&
+    !request.hasDraft && !request.hasVision && !request.userSeed && !request.hasGrammar &&
+    !request.wantsLogprobs && !request.kvQuant && !request.turboQuant;
+  if (features.fill && !fill) reasons.push("fill-incompatible-with-request");
+  const checkpoint = capabilities.checkpoints && method === "autoregressive" &&
+    (mechanism === "serial" || capabilities.sharedCheckpoints === true) && promptCache && !request.hasGrammar &&
+    !features.fill && !request.wantsLogprobs;
+  const compiledDecode = features.compiledDecode === true && capabilities.compiledDecode === true &&
+    method === "autoregressive" && !request.hasAdapters && !pagedKv;
+  if (features.compiledDecode && !compiledDecode) reasons.push("compiled-decode-unavailable-for-request");
+  const grammarJump = features.grammarJump === true && request.hasGrammar &&
+    method === "autoregressive" && mechanism === "serial" && !request.wantsLogprobs;
+  if (features.grammarJump && request.hasGrammar && !grammarJump)
+    reasons.push("grammar-jump-incompatible-with-request");
+  return Object.freeze({ method, mechanism, pagedKv, promptCache, checkpoint, fill,
+    compiledDecode, grammarJump,
+    reasons: Object.freeze(reasons) });
+}

package/src/engine/resources.ts

@@ -0,0 +1,53 @@
+import type { DisposableResource, PreparedStateChange, ResourceOwner } from "../contracts/resources";
+
+export function ownResource<T>(value: T, release: (value: T) => void): ResourceOwner<T> {
+  let held: { value: T } | undefined = { value };
+  const borrow = () => {
+    if (!held) throw new Error("resource ownership has ended");
+    return held.value;
+  };
+  return {
+    borrow,
+    transfer() { const result = borrow(); held = undefined; return result; },
+    close() {
+      if (!held) return;
+      const result = held.value;
+      held = undefined;
+      release(result);
+    },
+  };
+}
+
+/** Attempt every release once. One failed destructor must not leak siblings. */
+export function disposeResources(resources: Iterable<DisposableResource>): void {
+  const errors: unknown[] = [];
+  for (const resource of new Set(resources)) {
+    try { resource.dispose(); } catch (error) { errors.push(error); }
+  }
+  if (errors.length === 1) throw errors[0];
+  if (errors.length) throw new AggregateError(errors, "resource cleanup failed");
+}
+
+/** Used at failure boundaries to retain the execution error as well as cleanup. */
+export function cleanupFailure(error: unknown, cleanup: () => void): never {
+  try { cleanup(); }
+  catch (cleanupError) { throw new AggregateError([error, cleanupError], "execution and cleanup failed"); }
+  throw error;
+}
+
+/** Complete every preparation before publishing any owner's replacement.
+ * Old state is released only after all owners have committed. */
+export function applyStateChanges(prepare: readonly (() => PreparedStateChange)[]): void {
+  const changes: PreparedStateChange[] = [];
+  try {
+    for (const build of prepare) changes.push(build());
+    for (const change of changes) change.commit();
+  } catch (error) { return cleanupFailure(error, () => disposeResources(changes.splice(0))); }
+  disposeResources(changes);
+}
+
+export function withResource<T, R>(owner: Pick<ResourceOwner<T>, "borrow" | "close">, use: (value: T) => R): R {
+  try { return use(owner.borrow()); }
+  catch (error) { cleanupFailure(error, () => owner.close()); }
+  finally { owner.close(); }
+}

package/src/engine/scheduler.ts

@@ -0,0 +1,72 @@
+import type { ExecutionGroup, SchedulingClock } from "../contracts/scheduling";
+import { disposeResources } from "./resources";
+
+/** Continuous scheduling policy. Each advance is a backend-defined safe unit;
+ * this driver neither handles native state nor forces pending work to evaluate. */
+export async function driveExecutionGroup(
+  group: ExecutionGroup, clock: SchedulingClock, yieldAfterPreparation = false,
+): Promise<void> {
+  let execution: (() => void) | undefined;
+  let residency: (() => void) | undefined;
+  let lastYield = 0;
+  let terminalError: unknown = new Error("scheduler closed");
+  const release = () => {
+    const releases = [execution, residency];
+    execution = residency = undefined;
+    disposeResources(releases.filter((fn): fn is () => void => !!fn).map((dispose) => ({ dispose })));
+  };
+  try {
+    while (!group.closed) {
+      group.pruneCancelled();
+      const held = group.admissionHeld;
+      if (!group.active && !group.preparing && (held || !group.queued)) {
+        release();
+        await group.waitForWork();
+        continue;
+      }
+      residency ??= group.reserveResidency();
+      if (!execution && group.acquireExecution) execution = await group.acquireExecution();
+      if (group.closed) break;
+
+      if (!group.preparing && !held && group.queued && group.active < group.maxActive)
+        group.admitNext();
+      while (group.preparing && group.canPrepareMore && !held && group.queued &&
+          group.active + (group.preparingRows ?? 1) < group.maxActive && group.canBurst() &&
+          (group.preparingTokens ?? 0) + (group.nextPreparationTokens ?? 0) <= (group.maxPreparationTokens ?? Infinity)) {
+        if (!group.admitNext()) break;
+      }
+      if (group.preparing) {
+        const activeBefore = group.active;
+        await group.advancePreparation();
+        // A bounded preparation may publish output before its state work ends.
+        // Flush it before the next work unit; queued short admissions still
+        // form a group across yields. Re-enter policy to observe new work,
+        // cancellation, admission holds and shutdown.
+        if (group.preparationPublishedOutput ||
+            (yieldAfterPreparation && !activeBefore && !group.preparing && group.active === 1 && !group.queued)) {
+          lastYield = clock.now();
+          await clock.yield();
+          continue;
+        }
+      }
+
+      // Finish short admissions before decode; long preparation interleaves
+      // with active work because the group still reports preparing=true.
+      if (!group.preparing && !held && group.queued && group.active < group.maxActive && group.canBurst())
+        continue;
+      if (group.active) {
+        try { await group.advance(); }
+        catch (error) { group.failActive(error); }
+      }
+      // Preserve the serial responsiveness budget for the single active run.
+      if (group.active === 1 && !group.queued && !group.preparing &&
+          !group.admissionHeld && clock.now() - lastYield < 25) continue;
+      lastYield = clock.now();
+      await clock.yield();
+    }
+  } catch (error) {
+    terminalError = error;
+  } finally {
+    disposeResources([{ dispose: () => group.failAll(terminalError) }, { dispose: release }]);
+  }
+}

package/src/engine/session.ts

@@ -0,0 +1,247 @@
+import type {
+  Cancellation, CancelReason, ExecutionPlan, GenerationEvent, GenerationOutcome,
+  GenerationOutput, GenerationSession, MethodResult, MethodRun, RunControl,
+  SessionState, Timer, TokenLogprobs,
+} from "../contracts/generation";
+import { CancellationSource, throwIfCancelled } from "./cancellation";
+
+export interface SessionResources {
+  readonly timer: Timer;
+  readonly idleTimeoutMs: number;
+  readonly maxQueuedEvents: number;
+  readonly maxQueuedTokens: number;
+  readonly maxTopLogprobs: number;
+  reserveOutput(tokens: number): { buffer: Uint32Array; release(): void };
+}
+
+/** Lifecycle and bounded delivery only. No graph, cache, sampler, or method branches. */
+export class EngineSession<Metrics> implements GenerationSession<Metrics> {
+  readonly result: Promise<GenerationOutcome<Metrics>>;
+  readonly events: AsyncIterable<GenerationEvent>;
+  #state: SessionState = "created";
+  #cancel = new CancellationSource();
+  #settle!: (outcome: GenerationOutcome<Metrics>) => void;
+  #started = false;
+  #readerClaimed = false;
+  #reading = false;
+  #publishing = false;
+  #queue: GenerationEvent[] = [];
+  #queuedTokens = 0;
+  #committed = 0;
+  #sequence = 0;
+  #wakeReader: (() => void) | undefined;
+  #wakeProducer: (() => void) | undefined;
+  #stopTimer: (() => void) | undefined;
+  #unsubscribe: (() => void) | undefined;
+  #buffer: Uint32Array | undefined;
+  #limit = 0;
+  #outputStopped = false;
+
+  constructor(
+    readonly id: string,
+    private readonly prepare: (cancellation: Cancellation) => Promise<ExecutionPlan<Metrics>>,
+    private readonly control: RunControl,
+    private readonly resources: SessionResources,
+    private readonly onTerminal: () => void,
+  ) {
+    this.result = new Promise((resolve) => { this.#settle = resolve; });
+    this.events = { [Symbol.asyncIterator]: () => this.#iterator() };
+    const unsubscribe = control.cancellation?.subscribe((reason) => { void this.cancel(reason); });
+    if (this.#state === "terminal") unsubscribe?.();
+    else this.#unsubscribe = unsubscribe;
+    if (this.#state !== "terminal") {
+      if (control.output !== "stream") this.#start();
+      else this.#armIdle();
+    }
+  }
+
+  get state(): SessionState { return this.#state; }
+
+  async cancel(reason: CancelReason = "requested"): Promise<void> {
+    if (this.#state !== "terminal") {
+      this.#cancel.cancel(reason);
+      this.#queue = [];
+      this.#queuedTokens = 0;
+      this.#wakeReader?.();
+      this.#wakeProducer?.();
+      this.#stopTimer?.();
+      if (!this.#started) this.#finish({ status: "cancelled", reason, committedTokens: 0 });
+    }
+    await this.result;
+  }
+
+  close(): Promise<void> { return this.cancel("consumer_closed"); }
+
+  #armIdle(): void {
+    this.#stopTimer?.();
+    this.#stopTimer = this.resources.timer.after(this.resources.idleTimeoutMs,
+      () => { void this.cancel("consumer_idle"); });
+  }
+
+  #iterator(): AsyncIterator<GenerationEvent> {
+    if (this.control.output !== "stream") throw new Error("collect sessions return output through result");
+    if (this.#readerClaimed) throw new Error("generation events have one consumer");
+    this.#readerClaimed = true;
+    return {
+      next: async () => {
+        if (this.#reading) throw new Error("concurrent next() is not supported");
+        this.#reading = true;
+        this.#stopTimer?.();
+        try {
+          this.#start();
+          while (!this.#queue.length && this.#state !== "terminal") {
+            await new Promise<void>((resolve) => { this.#wakeReader = resolve; });
+            this.#wakeReader = undefined;
+          }
+          const event = this.#queue.shift();
+          if (!event) return { done: true, value: undefined };
+          if (event.type === "committed") this.#queuedTokens -= event.tokenIds.length;
+          this.#wakeProducer?.();
+          if (this.#state !== "terminal") this.#armIdle();
+          return { done: false, value: event };
+        } finally { this.#reading = false; }
+      },
+      return: async () => { await this.close(); return { done: true, value: undefined }; },
+      throw: async (error) => { await this.close(); throw error; },
+    };
+  }
+
+  #start(): void {
+    if (this.#started || this.#state === "terminal") return;
+    this.#started = true;
+    void this.#execute();
+  }
+
+  async #execute(): Promise<void> {
+    let run: MethodRun<Metrics> | undefined;
+    let reservation: ReturnType<SessionResources["reserveOutput"]> | undefined;
+    let result: MethodResult<Metrics> | undefined;
+    let failure: unknown;
+    let failed = false;
+    try {
+      this.#state = "preparing";
+      throwIfCancelled(this.#cancel);
+      const plan = await this.prepare(this.#cancel);
+      throwIfCancelled(this.#cancel);
+      if (!Number.isSafeInteger(plan.outputTokenLimit) || plan.outputTokenLimit < 0)
+        throw new Error("invalid planned output token limit");
+      this.#limit = plan.outputTokenLimit;
+      if (this.control.output === "collect") {
+        reservation = this.resources.reserveOutput(this.#limit);
+        this.#buffer = reservation.buffer;
+      }
+      run = await plan.method.createRun(this.#cancel);
+      throwIfCancelled(this.#cancel);
+      this.#state = "running";
+      const output: GenerationOutput = {
+        commit: (ids, logprobs) => this.#publish(ids, logprobs),
+        progress: (completed, total) => this.#progress(completed, total),
+      };
+      result = await run.execute(output, this.#cancel);
+    } catch (error) { failure = error; failed = true; }
+    this.#state = "settling";
+    // Terminal means resources are safe to release, not merely that JS stopped reading.
+    let cleanupFailed = false;
+    let cleanupError: string | undefined;
+    try { await run?.close(); }
+    catch (error) {
+      cleanupError = error instanceof Error ? error.message : String(error);
+      if (!failed) failure = error;
+      failed = true;
+      cleanupFailed = true;
+    }
+    reservation?.release(); // Completed output ownership transfers to the caller.
+    const common = { committedTokens: this.#committed,
+      ...(this.#buffer ? { output: this.#buffer.subarray(0, this.#committed) } : {}) };
+    if (this.#cancel.reason && !cleanupFailed) {
+      this.#finish({ status: "cancelled", reason: this.#cancel.reason, ...common });
+    } else if (failed) {
+      this.#finish({ status: "failed", error: { code: cleanupFailed ? "cleanup_failed" : "execution_failed",
+        message: failure instanceof Error ? failure.message : String(failure),
+        ...(cleanupError === undefined ? {} : { cleanupError }) }, ...common });
+    } else {
+      this.#finish({ status: "completed", result: result!, ...common });
+    }
+  }
+
+  async #publish(ids: readonly number[], logprobs?: readonly (TokenLogprobs | undefined)[]): Promise<void | false> {
+    if (this.#publishing) throw new Error("method must await each publication");
+    this.#publishing = true;
+    try {
+      throwIfCancelled(this.#cancel);
+      if (this.#state !== "running") throw new Error("publication outside method execution");
+      if (this.#outputStopped) throw new Error("method published after the consumer requested a stop");
+      if (ids.length > this.#limit - this.#committed) throw new Error("method exceeded planned output limit");
+      if (logprobs && logprobs.length !== ids.length) throw new Error("logprobs must align with committed tokens");
+      if (logprobs?.some((item) => (item?.top?.length ?? 0) > this.resources.maxTopLogprobs))
+        throw new Error("top logprobs exceed session delivery capacity");
+      for (const id of ids)
+        if (!Number.isInteger(id) || id < 0 || id > 0xffffffff) throw new Error("invalid committed token id");
+      if (this.control.output === "callback") {
+        this.#committed += ids.length;
+        const result = await this.control.onTokens(ids, logprobs);
+        if (result === false) this.#outputStopped = true;
+        return result;
+      }
+      if (this.#buffer) {
+        if (logprobs) throw new Error("token logprobs require stream consumption");
+        this.#buffer.set(ids, this.#committed);
+        this.#committed += ids.length;
+        return;
+      }
+      for (let offset = 0; offset < ids.length;) {
+        const length = Math.min(ids.length - offset, this.resources.maxQueuedTokens);
+        await this.#space(length);
+        // Snapshot producer-owned arrays; later scratch-buffer reuse cannot alter events.
+        const event: GenerationEvent = { type: "committed", sequence: this.#sequence++,
+          tokenIds: ids.slice(offset, offset + length),
+          ...(logprobs ? { logprobs: logprobs.slice(offset, offset + length).map((item) => item && ({
+            ...(item.logprob === undefined ? {} : { logprob: item.logprob }),
+            ...(item.top ? { top: item.top.map((entry) => ({ id: entry.id, logprob: entry.logprob })) } : {}),
+          })) } : {}) };
+        this.#queue.push(event);
+        this.#queuedTokens += length;
+        this.#committed += length;
+        offset += length;
+        this.#wakeReader?.();
+      }
+    } finally { this.#publishing = false; }
+  }
+
+  async #progress(completed: number, total?: number): Promise<void> {
+    if (this.#publishing) throw new Error("method must await each publication");
+    this.#publishing = true;
+    try {
+      throwIfCancelled(this.#cancel);
+      if (this.#state !== "running") throw new Error("publication outside method execution");
+      if (!Number.isFinite(completed) || completed < 0 ||
+          (total !== undefined && (!Number.isFinite(total) || total < completed)))
+        throw new Error("invalid progress");
+      if (this.control.output !== "stream") return;
+      await this.#space(0);
+      this.#queue.push({ type: "progress", completed, ...(total === undefined ? {} : { total }) });
+      this.#wakeReader?.();
+    } finally { this.#publishing = false; }
+  }
+
+  async #space(tokens: number): Promise<void> {
+    while (this.#queue.length >= this.resources.maxQueuedEvents ||
+           this.#queuedTokens + tokens > this.resources.maxQueuedTokens) {
+      throwIfCancelled(this.#cancel);
+      this.#armIdle();
+      await new Promise<void>((resolve) => { this.#wakeProducer = resolve; });
+      this.#wakeProducer = undefined;
+    }
+    throwIfCancelled(this.#cancel);
+  }
+
+  #finish(outcome: GenerationOutcome<Metrics>): void {
+    if (this.#state === "terminal") return;
+    this.#state = "terminal";
+    this.#stopTimer?.();
+    this.#unsubscribe?.();
+    this.#settle(outcome);
+    this.#wakeReader?.();
+    this.onTerminal();
+  }
+}

package/src/eval/README.md

@@ -29,6 +29,45 @@ repo rule: free-running greedy "measures chaos").
   `KL(compat ‖ fused)` is literally "measure the output difference."
 - **two-model (optiq-style absolute quality):** a separate reference model
   (bf16 if it fits, else uniform-4bit), both resident.
+- **saved-teacher (`--reference-logits <dir>`):** the reference is a
+  **dump on disk**, not a resident model — the only way to score against a
+  bf16 teacher on 32 GB. A Python-side script forwards the teacher once and
+  writes, per teacher-forced position, the top-2048 logits plus the
+  **full-vocab logsumexp**; every candidate is then scored through OUR
+  engine. Format (`format_version` 1, read by `reference-logits.ts`):
+
+  ```
+  runs/kl-teacher/<tag>/
+    manifest.json  {format_version, model, corpus_sha256, ctx_len, n_seqs,
+                    top_k, positions_per_seq, created, notes}
+    tokens.bin     int32 LE [n_seqs][ctx_len]
+    seq-<i>.bin     record r (r = 0…ctx_len-2, the distribution over token
+                    r+1, i.e. candidate POSITION r):
+                      int32[top_k] indices, descending by logit
+                      float16[top_k] reference logits
+                      float32 full-vocab logsumexp
+  ```
+
+  `KL = Σ_{j∈topK} p_j·(log p_j − log q_j)`, `p_j = exp(ref_j − ref_lse)`,
+  `log q_j = cand_j − logsumexp(cand, full vocab)`, float32 throughout.
+  The sum is truncated to top-k, so the run also reports **captured mass**
+  `Σ_j p_j` (how much of `p` the number covers — a low-mass run is not a
+  comparable KL) and **top-1 agreement** (argmax over the full vocab).
+  The forward is the SERVING path (`makeCache` + chunked `forwardHidden` +
+  `logitsFromHidden` per chunk, `--chunk` positions at a time so a 4k
+  context's full-vocab logits never exist at once), **not** `trainForward` —
+  which is why this works on qwen3_5, where `mlx-bun perplexity` cannot go
+  (docs/design/turboquant.md, "Known engine gaps"). KV stays bf16: this is a
+  weights-only instrument.
+
+  Noise floor (measured, M1 Max, Qwen3.5-0.8B-OptiQ-4bit, 2×128 ctx,
+  top-64): scoring a dump against the model that produced it gives KL
+  **exactly 0** at the same `--chunk`, and mean 9.3e-6 / p95 8.3e-7 across
+  chunk 64 → 40. float16 storage is lossless for a bf16 teacher (bf16's
+  7 mantissa bits fit float16's 10, and logits sit far inside float16's
+  65504 range), so the residual is only the engine's own L-dependent bf16
+  rounding — the same effect that makes prefill chunk size ulp-visible
+  (`generate.ts`, prefill-tail-split). Compare candidates at one `--chunk`.
 
 ## Run (on a cleared machine — loads models, not from an agent session)
 
@@ -42,6 +81,10 @@ bun scripts/eval.ts kl --candidate e4b --self MLX_BUN_NO_FUSED_SDPA --ref-value
 # optiq-style absolute drift vs a reference model:
 bun scripts/eval.ts kl --candidate e4b --reference <bf16-or-4bit-id>
 
+# vs a saved bf16 teacher-logits dump (--n caps sequences, --chunk positions
+# per forward; the dump brings its own corpus, so --prompts-file/--seq don't apply):
+bun scripts/eval.ts kl --candidate <model> --reference-logits runs/kl-teacher/<tag>
+
 # higher-fidelity prompts (optiq's calibration mix):
 bun scripts/eval.ts kl --candidate e4b --prompts-file \
   /Users/joshrossi/Code/mlx-lm/.venv/.../optiq/calibration/data/optiq.jsonl
@@ -71,7 +114,7 @@ the throughput `runs` table). Each row records the active perf levers
   perf levers (NO_FUSED_SDPA) actually bite.
 - **M1 (done):** all six capability tasks ported + runtime-validated, plus
   `capability` aggregation and the `smoketest` (KL + GSM8K-50) gate. Datasets
-  exported by `scripts/eval/export-datasets.py` → `~/.cache/mlx-bun/eval-data/`.
+  exported by `scripts/oracle/export-datasets.py` → `~/.cache/mlx-bun/eval-data/`.
   Tasks run through the real quantized-KV path (`runner.ts`). HumanEval executes
   in a sandbox-exec + venv-python subprocess; HashHop reconstructs Magic AI's
   generator (`MultiHopEval` isn't on PyPI). **Open quality checks** (need a real