mlx-bun 0.3.0 → 0.5.0

package/src/engine/execution-plan.ts

@@ -1,48 +1,61 @@
 import type { ExecutionCapabilities, ExecutionFeatures, ExecutionRequirements, ResolvedExecution } from "../contracts/execution";
 
-/** Preserve the shipped support matrix. A configured draft still forces serial
- * placement even when this request must use the ordinary AR method. */
+/** 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 continuousExclusions = [
+  const sharedRequestExclusions = [
     [!capabilities.continuous, "continuous-unavailable"],
-    [capabilities.method !== "autoregressive", "method-requires-serial"],
-    [request.hasVision, "media-requires-serial"],
-    [request.hasAdapters, "adapters-require-serial"],
-    [request.wantsLogprobs, "logprobs-require-serial"],
-    [request.userSeed, "explicit-seed-requires-serial"],
+    [request.hasVision && !capabilities.mediaBatch, "media-requires-serial"],
+    [request.hasAdapters && !capabilities.adapterBatch, "adapters-require-serial"],
     [request.kvQuant && !capabilities.quantizedBatch, "kv-scheme-requires-serial"],
-    [request.turboQuant, "turbo-kv-requires-serial"],
-    [request.hasDraft, "configured-draft-requires-serial"],
+    [request.turboQuant && !capabilities.turboQuantBatch, "turbo-kv-requires-serial"],
     [request.hasGrammar && !capabilities.grammarBatch, "grammar-batching-disabled"],
-    [features.pagedKv, "paged-kv-requires-serial"],
+    [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 grammarProposals = features.grammarJump === true && request.hasGrammar && !request.hasDraft && !request.hasVision &&
+    (!request.wantsLogprobs || capabilities.speculativeLogprobs === true) && !features.pagedKv && capabilities.method === "autoregressive" &&
+    capabilities.sharedGrammarProposals === true && capabilities.groupedMethods?.includes("speculative") === true &&
+    (!request.hasAdapters || sharedAdapterMethod) && (!request.kvQuant || capabilities.speculativeKvQuant === true) &&
+    (!request.turboQuant || capabilities.speculativeTurboQuant === true) &&
+    !sharedRequestExclusions.some(([excluded]) => excluded);
+  const speculative = grammarProposals || (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";
-  const speculative = capabilities.method === "autoregressive" && request.hasDraft &&
-    !request.hasVision && !request.hasAdapters && !request.wantsLogprobs &&
-    (!request.kvQuant || capabilities.speculativeKvQuant === true) &&
-    !request.turboQuant && !features.pagedKv;
   if (request.hasDraft && !speculative) reasons.push("draft-incompatible-with-request");
-  const method = speculative ? "speculative" : capabilities.method;
   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;
+  const promptCache = method !== "speculative" && (!request.hasVision ||
+    (mechanism === "continuous" && request.hasPreparedPrefixIdentity === true && capabilities.mediaPrefixCache === true));
+  const sharedFill = mechanism === "continuous" && capabilities.sharedFill === true && !pagedKv;
+  const speculativeEcho = method === "speculative" && mechanism === "continuous" &&
+    capabilities.sharedSpeculativeEcho === true;
+  const fill = features.fill && !request.hasVision && !request.hasGrammar && !request.wantsLogprobs &&
+    (speculativeEcho || (method === "autoregressive" && (mechanism === "serial" || sharedFill) &&
+      !request.hasDraft && (!request.userSeed || sharedFill)));
+  // Cache-format eligibility belongs to the method's append binding.
   if (features.fill && !fill) reasons.push("fill-incompatible-with-request");
-  const checkpoint = capabilities.checkpoints && method === "autoregressive" &&
-    mechanism === "serial" && promptCache && !request.hasGrammar &&
+  const checkpoint = capabilities.checkpoints && method === "autoregressive" && !request.hasVision &&
+    (mechanism === "serial" || capabilities.sharedCheckpoints === true) && promptCache && !pagedKv && !request.hasGrammar &&
     !features.fill && !request.wantsLogprobs;
   const compiledDecode = features.compiledDecode === true && capabilities.compiledDecode === true &&
-    method === "autoregressive" && !request.hasAdapters && !pagedKv;
+    method === "autoregressive" && !request.hasAdapters && !pagedKv && !(sharedFill && fill);
   if (features.compiledDecode && !compiledDecode) reasons.push("compiled-decode-unavailable-for-request");
-  const grammarJump = features.grammarJump === true && request.hasGrammar &&
-    method === "autoregressive" && mechanism === "serial" && !request.wantsLogprobs;
+  const grammarJump = grammarProposals || (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,

package/src/engine/execution-tasks.ts

@@ -0,0 +1,44 @@
+interface Task {
+  run(): Promise<void>;
+  reject(error: unknown): void;
+}
+
+/** Work whose owner requires an execution boundary, without owning a decode
+ * row. The driver supplies the lease; this queue owns waiting and cancellation.
+ * Once started, work settles before its caller can release native resources. */
+export class ExecutionTasks {
+  #pending: Task[] = [];
+  get pending(): number { return this.#pending.length; }
+
+  enqueue<T>(work: () => Promise<T>, cancellation?: Cancellation): Promise<T> {
+    if (cancellation?.reason !== undefined) return Promise.reject(new GenerationCancelled(cancellation.reason));
+    return new Promise<T>((resolve, reject) => {
+      let unsubscribe = () => {};
+      const cleanup = () => unsubscribe();
+      const task: Task = {
+        async run() {
+          cleanup();
+          try { resolve(await work()); } catch (error) { reject(error); }
+        },
+        reject(error) { cleanup(); reject(error); },
+      };
+      this.#pending.push(task);
+      unsubscribe = cancellation?.subscribe(reason => {
+        const index = this.#pending.indexOf(task);
+        if (index < 0) return;
+        this.#pending.splice(index, 1);
+        task.reject(new GenerationCancelled(reason));
+      }) ?? (() => {});
+      if (!this.#pending.includes(task)) unsubscribe();
+    });
+  }
+
+  async advance(): Promise<void> { await this.#pending.shift()?.run(); }
+
+  rejectAll(error: unknown): void {
+    const tasks = this.#pending; this.#pending = [];
+    for (const task of tasks) task.reject(error);
+  }
+}
+import type { Cancellation } from "../contracts/generation";
+import { GenerationCancelled } from "./cancellation";

package/src/engine/namespaced-cache.ts

@@ -0,0 +1,13 @@
+import type { PrefixCache } from "../contracts/prefix-cache";
+
+/** Compose numerical identity with a cache without exposing its storage tiers. */
+export function namespacedCache<State, Attachment>(cache: PrefixCache<State, Attachment>,
+  namespace: (base: string) => string): PrefixCache<State, Attachment> {
+  return {
+    closeSession: id => cache.closeSession?.(id),
+    reclaim: () => cache.reclaim?.(),
+    ...(cache.prefetch ? { prefetch: (tokens: number[], ns = "", sessionId?: string) => cache.prefetch!(tokens, namespace(ns), sessionId) } : {}),
+    take: (tokens, ns = "", sessionId) => cache.take(tokens, namespace(ns), sessionId),
+    put: (tokens, state, ns = "", retain, attachments, sessionId) => cache.put(tokens, state, namespace(ns), retain, attachments, sessionId),
+  };
+}

package/src/engine/resources.ts

@@ -1,4 +1,4 @@
-import type { DisposableResource, ResourceOwner } from "../contracts/resources";
+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 };
@@ -35,6 +35,17 @@ export function cleanupFailure(error: unknown, cleanup: () => void): never {
   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()); }

package/src/engine/scheduler.ts

@@ -18,8 +18,8 @@ export async function driveExecutionGroup(
   try {
     while (!group.closed) {
       group.pruneCancelled();
-      const held = group.admissionHeld;
-      if (!group.active && !group.preparing && (held || !group.queued)) {
+      let held = group.admissionHeld;
+      if (!group.active && !group.preparing && (held || (!group.queued && !group.pendingTasks))) {
         release();
         await group.waitForWork();
         continue;
@@ -28,17 +28,42 @@ export async function driveExecutionGroup(
       if (!execution && group.acquireExecution) execution = await group.acquireExecution();
       if (group.closed) break;
 
+      // One preparation task per iteration. An exclusive model mutation
+      // keeps its existing drain priority; preparation cannot prolong it.
+      if (!held && group.pendingTasks && group.advanceTask) {
+        await group.advanceTask();
+        if (group.closed) break;
+        group.pruneCancelled();
+        held = group.admissionHeld;
+      }
+
       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;
+      }
+      const mixed = group.preparing && group.active ? group.mixedPreparation : undefined;
+      if (mixed && group.advanceMixed) {
+        // Reserve the method's running-token demand first; prompt work fills
+        // the rest. A budget smaller than one token per row still progresses.
+        const budget = Math.max(group.maxIterationTokens ?? Infinity,
+          mixed.runningTokens + mixed.minimumPreparationTokens);
+        await group.advanceMixed(budget);
+        lastYield = clock.now();
+        await clock.yield();
+        continue;
+      }
       if (group.preparing) {
         const activeBefore = group.active;
         await group.advancePreparation();
-        // A backend may emit its first output while preparing the first row.
-        // Flush a lone admission before decode; queued short admissions still
-        // group together. Re-enter policy after yielding to observe new work,
+        // 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 (yieldAfterPreparation && !activeBefore && !group.preparing &&
-            group.active === 1 && !group.queued) {
+        if (group.preparationPublishedOutput ||
+            (yieldAfterPreparation && !activeBefore && !group.preparing && group.active === 1 && !group.queued)) {
           lastYield = clock.now();
           await clock.yield();
           continue;

package/src/fill/fill-session.ts

@@ -20,7 +20,7 @@
 // sources in priority order, clamps what they return, and counts. It never
 // touches MLX, the tokenizer, or the template. Row COMPILATION is
 // ./schema-rows.ts.
-import { flagOn, runtimeNumber, runtimeValue } from "../runtime-config";
+import { runtimeFlag, runtimeNumber, runtimeValue } from "../runtime-config";
 import { EchoSource, type EchoConfig } from "./echo-index";
 import type {
   Proposal, ProposalOrigin, ProposalPolicy, ProposalSource, TokenView,
@@ -153,7 +153,7 @@ export function fillEchoConfig(): FillEchoConfig {
  *  the engine's fill path. Off by default (the invariant it checks is the
  *  easiest bug in the feature, and the check costs a JS compare). */
 export function fillTraceEnabled(): boolean {
-  return flagOn("MLX_BUN_FILL_TRACE", false) || fillTracePath() !== null;
+  return runtimeFlag("MLX_BUN_FILL_TRACE", false) || fillTracePath() !== null;
 }
 
 /** MLX_BUN_FILL_TRACE=<file.jsonl> — additionally append one record per
@@ -319,12 +319,23 @@ export class FillSession {
   /** Feed one SAMPLED token, then ask the sources what follows. `budget` is
    *  the caller's remaining allowance (max_tokens − generated). */
   push(tokenId: number, budget: number = Number.POSITIVE_INFINITY): Proposal | null {
+    this.observe(tokenId);
+    return this.propose(budget);
+  }
+
+  /** Record sampled output independently of when the method requests proposals. */
+  observe(tokenId: number): void {
     this.stats.decodeSteps++;
     this.#append([tokenId]);
+  }
+
+  /** Select a continuation after a committed output boundary. Methods may
+   * request only verified proposals while retaining ordinary sampled policy. */
+  propose(budget: number = Number.POSITIVE_INFINITY, policy?: ProposalPolicy): Proposal | null {
     this.#budget = budget;
     for (const source of this.#sources) {
       const proposal = source.propose(this);
-      if (!proposal) continue;
+      if (!proposal || (policy && proposal.policy !== policy)) continue;
       if (proposal.branchStop) this.stats.branchStops++;
       const clamped = this.#clamp(proposal.ids, budget, proposal.origin === "echo");
       if (!clamped) continue;

package/src/fit.ts

@@ -13,12 +13,12 @@ import type { ModelConfig } from "./config";
 import type { MemoryPlan } from "./admission-plan";
 import {
   kvBytesAt,
-  kvGeometry,
   kvQuantBytesPerElement,
+  sdpaFallbackBytes,
   type KvSchemeOptions,
 } from "./kv-scheme";
 
-export { kvBytesAt, kvQuantBytesPerElement } from "./kv-scheme";
+export { kvBytesAt, kvQuantBytesPerElement, sdpaFallbackBytes } from "./kv-scheme";
 export type FitKvScheme = KvSchemeOptions;
 
 /** Decode-efficiency vs theoretical bandwidth ceiling, measured on the
@@ -161,23 +161,37 @@ export function fit(
   kvScheme?: FitKvScheme,
 ): FitReport {
   const usable = usableBytes ?? machine.ramBytes * WIRED_FRACTION;
-  const transient = Math.min(chunk, ctx) * TRANSIENT_PER_TOKEN;
+  // Prefill transient = the calibrated per-chunk-token constant + the
+  // SDPA-fallback scores tensor (heads × min(chunk, ctx) × ctx × 2 B) for
+  // attention layers whose head dim MLX cannot fuse (kv-scheme.ts). The
+  // second term is what grows with context on Qwen3.x (head_dim 256) and
+  // Gemma 4 (256/512) and was unbilled before 2026-09-08.
+  const transientAt = (n: number) =>
+    Math.min(chunk, n) * TRANSIENT_PER_TOKEN + sdpaFallbackBytes(config, chunk, n);
+  const transient = transientAt(ctx);
   const kv = kvBytesAt(config, ctx, kvScheme);
   const total = weightsBytes + kv + transient;
 
-  // solve max context: weights + kv(ctx) + transient ≤ usable.
-  // Below the window both KV terms are linear in ctx; above it the
-  // sliding term saturates and only full-attention layers keep growing.
-  const g = kvGeometry(config, kvScheme);
-  const fixed = weightsBytes + chunk * TRANSIENT_PER_TOKEN + g.linearStateBytes;
+  // solve max context: weights + kv(n) + transient(n) ≤ usable. The same
+  // formula `fits` uses, so the solved ceiling always fits. It is monotone
+  // in n (KV linear, sliding saturating at the window, scores quadratic
+  // below the chunk then linear), so bisect over [0, maxPositionEmbeddings].
+  const totalAt = (n: number) => weightsBytes + kvBytesAt(config, n, kvScheme) + transientAt(n);
   let maxCtx = 0;
-  if (usable > fixed) {
-    const budget = usable - fixed;
-    const linear = Math.floor(budget / (g.fullBytesPerToken + g.slidingBytesPerToken));
-    maxCtx = linear <= g.window
-      ? linear
-      : Math.floor((budget - g.slidingBytesPerToken * g.window) / g.fullBytesPerToken);
-    maxCtx = Math.min(maxCtx, config.text.maxPositionEmbeddings);
+  {
+    let lo = 0;
+    let hi = Math.max(0, Math.floor(config.text.maxPositionEmbeddings));
+    if (totalAt(hi) <= usable) lo = hi;
+    else {
+      // invariant: totalAt(lo) ≤ usable (lo = 0 may itself not fit → 0)
+      if (totalAt(0) > usable) hi = 0;
+      while (hi - lo > 1) {
+        const mid = Math.floor((lo + hi) / 2);
+        if (totalAt(mid) <= usable) lo = mid;
+        else hi = mid;
+      }
+    }
+    maxCtx = lo;
   }
 
   // decode reads all weights + the KV cache once per token — except MoE