mlx-bun 0.5.0 → 1.0.0

package/src/backends/mlx/rollback.ts

@@ -1,92 +0,0 @@
-import type { SpeculativeTransaction } from "../../inference/rollback";
-import type { Cache } from "../../model/gemma4-base";
-
-interface TransactionLayer<Acceptance> {
-  fits(drafts: number): boolean;
-  begin(drafts: number): void;
-  resolve(accepted: Acceptance, drafts: number): void;
-}
-
-/** State storage that can retain a different verified prefix for each row.
- * Both attention and recurrent layouts implement this same transaction. */
-export interface RowSpeculativeState {
-  specRoundBegin(): void;
-  specRoundCommit(): void;
-  specRoundRollback(keep: number | readonly number[]): void;
-}
-
-export function bindRowCacheRollback(
-  caches: readonly RowSpeculativeState[], rows: number,
-): SpeculativeTransaction<readonly number[]> {
-  return bindTransaction(caches.map(cache => ({
-    fits: () => true,
-    // A zero-candidate round consumes only the committed pending token.
-    // It cannot reject any input, so retaining rollback state is unnecessary.
-    begin: drafts => { if (drafts > 0) cache.specRoundBegin(); },
-    resolve(accepted: readonly number[], drafts: number) {
-      if (drafts === 0) return;
-      if (accepted.every(count => count === drafts)) cache.specRoundCommit();
-      else cache.specRoundRollback(accepted.every(count => count === accepted[0])
-        ? accepted[0]! + 1 : accepted.map(count => count + 1));
-    },
-  })), (accepted, drafts) => accepted.length === rows &&
-    accepted.every(count => Number.isSafeInteger(count) && count >= 0 && count <= drafts));
-}
-
-/** Bind the recurrent transaction or the trim implementation once.
- * Ring-capacity checks still happen at each safe boundary before any writes. */
-export function bindCacheRollback(caches: readonly Cache[]): SpeculativeTransaction {
-  const layers = caches.map((cache) => {
-    const { specRoundBegin: begin, specRoundCommit: commit, specRoundRollback: rollback } = cache;
-    const hasRound = [begin, commit, rollback].some((method) => method !== undefined);
-    if (hasRound && !(typeof begin === "function" && typeof commit === "function" && typeof rollback === "function"))
-      throw new Error("cache exposes an incomplete speculative transaction");
-    return {
-      fits(drafts: number) {
-        if ("maxSize" in cache && typeof cache.maxSize === "number")
-          return cache.offset + drafts + 1 < cache.maxSize;
-        return hasRound || cache.isTrimmable();
-      },
-      begin: () => begin?.call(cache),
-      resolve(accepted: number, drafts: number) {
-        if (accepted < drafts) {
-          if (rollback) rollback.call(cache, accepted + 1);
-          else cache.trim(drafts - accepted);
-        } else commit?.call(cache);
-      },
-    };
-  });
-  return bindTransaction(layers, (accepted, drafts) =>
-    Number.isSafeInteger(accepted) && accepted >= 0 && accepted <= drafts);
-}
-
-/** One transaction lifecycle for scalar and row-wise state bindings. */
-function bindTransaction<Acceptance>(
-  layers: readonly TransactionLayer<Acceptance>[],
-  accepts: (accepted: Acceptance, drafts: number) => boolean,
-): SpeculativeTransaction<Acceptance> {
-  let active: number | undefined;
-  let invalid = false;
-  const validate = (drafts: number) => {
-    if (invalid) throw new Error("speculative state is invalid and must be discarded");
-    if (!Number.isSafeInteger(drafts) || drafts < 0) throw new Error("invalid speculative draft count");
-  };
-  return {
-    canBegin(drafts) { validate(drafts); return active === undefined && layers.every((layer) => layer.fits(drafts)); },
-    begin(drafts) {
-      validate(drafts);
-      if (active !== undefined || !layers.every((layer) => layer.fits(drafts)))
-        throw new Error("speculative state cannot begin this round");
-      active = drafts;
-      try { for (const layer of layers) layer.begin(drafts); }
-      catch (error) { invalid = true; throw error; }
-    },
-    resolve(accepted) {
-      if (invalid) throw new Error("speculative state is invalid and must be discarded");
-      if (active === undefined || !accepts(accepted, active)) throw new Error("invalid speculative acceptance count");
-      try { for (const layer of layers) layer.resolve(accepted, active); }
-      catch (error) { invalid = true; throw error; }
-      active = undefined;
-    },
-  };
-}

package/src/backends/mlx/row-sampling.ts

@@ -1,49 +0,0 @@
-import type { MlxArray } from "../../mlx/array";
-import * as ops from "../../mlx/ops";
-import type { TokenLogprobs } from "../../contracts/generation";
-import { readStepExtras, stepExtrasArrays, disposeStepExtras,
-  type DeviceStepSampler, type StepExtras } from "../../sampler";
-import { disposeResources, cleanupFailure } from "../../engine/resources";
-
-type TokenSink = (token: number, logprobs?: TokenLogprobs) => void | boolean | Promise<void | boolean>;
-
-/** Request-owned sampling and output metadata. Execution groups receive token
- * tensors and an output callback; they need no logprob or retention policy. */
-export function createRowSampling(
-  sampler: DeviceStepSampler, onToken: TokenSink,
-) {
-  if (!sampler.capturesLogprobs) return {
-    sample: (logits: MlxArray, step: number) => sampler.sample(logits, step).token,
-    plainGreedy: sampler.isPlainGreedy, onToken, dispose: () => sampler.dispose(),
-  };
-
-  // A pipelined group can sample the next token before emitting the previous
-  // one. Keep captures by request-local step, independently of row position.
-  const pending = new Map<number, StepExtras>();
-  let emitted = 0;
-  return {
-    sample(logits: MlxArray, step: number): MlxArray {
-      const result = sampler.sample(logits, step);
-      try {
-        ops.asyncEvalAll([result.token, ...stepExtrasArrays(result.extras)]);
-        if (result.extras) pending.set(step, result.extras);
-        return result.token;
-      } catch (error) {
-        return cleanupFailure(error, () => disposeResources([result.token,
-          { dispose: () => disposeStepExtras(result.extras) }]));
-      }
-    },
-    // The vectorized token-only shortcut bypasses the per-request sampler.
-    plainGreedy: false,
-    onToken(token: number) {
-      const extras = pending.get(emitted) ?? null;
-      pending.delete(emitted++);
-      return onToken(token, readStepExtras(extras));
-    },
-    dispose() {
-      const extras = [...pending.values()];
-      pending.clear();
-      disposeResources([sampler, ...extras.map(value => ({ dispose: () => disposeStepExtras(value) }))]);
-    },
-  };
-}

package/src/backends/mlx/serial-executor.ts

@@ -1,251 +0,0 @@
-import { pagedPrefixNamespace } from "./request-state-policy";
-import type { SpecRunServices } from "../../spec/serve-loop";
-import { cleanupFailure, disposeResources, ownResource } from "../../engine/resources";
-import { bindGeneration } from "../../generate";
-import type { RuntimeModel } from "../../model/factory";
-import { Qwen35Model } from "../../model/qwen3_5";
-import type { Cache } from "../../model/gemma4";
-import type { MlxPrefixCache } from "./checkpoint-state";
-import type { SsdCacheStore } from "../../ssd-cache";
-import type { DraftProvider } from "../../spec/source";
-import type { SerialRun, Vision } from "../../serve/generation-gateway";
-import { bindContinuationPolicy } from "./continuation";
-import { runtimeConfig, withRuntimeConfig, type RuntimeConfig } from "../../runtime-config";
-import { bindLegacySpeculativeModel } from "./speculative";
-
-/** Native serial execution depends on this bound port, never a model union.
- * Weights remain borrowed. The gateway supplies the exclusive runtime lease. */
-export interface MlxSerialBinding {
-  readonly runtime: RuntimeConfig;
-  readonly generate: ReturnType<typeof bindGeneration>;
-  readonly speculate?: (
-    promptIds: Parameters<SerialRun>[0], options: Parameters<SerialRun>[1],
-    onToken: Parameters<SerialRun>[2], services?: SpecRunServices,
-  ) => ReturnType<SerialRun>;
-  makeCache(): Cache[];
-  enterMedia?(vision?: Vision): () => void;
-}
-
-/** Family-specific context is bound once at the compatibility boundary. */
-export function bindLegacySerialModel(
-  model: RuntimeModel,
-  draft?: { provider: DraftProvider; numDraftTokens: number },
-): MlxSerialBinding {
-  const speculative = draft ? bindLegacySpeculativeModel(model, draft.provider) : undefined;
-  return {
-    runtime: runtimeConfig(),
-    generate: bindGeneration(model),
-    makeCache: model.makeCache.bind(model),
-    ...(speculative && draft ? { speculate: (async (prompt, options, onToken, services) => {
-      const { specRun } = await import("../../spec/serve-loop");
-      return specRun(speculative, draft.numDraftTokens, prompt, options, onToken, services);
-    }) satisfies NonNullable<MlxSerialBinding["speculate"]> } : {}),
-    ...(model instanceof Qwen35Model ? { enterMedia(vision?: Vision) {
-      const previous = model.mrope;
-      if (vision?.mrope) model.mrope = vision.mrope;
-      return () => { model.mrope = previous; };
-    } } : {}),
-  };
-}
-
-export interface MlxSerialServices {
-  readonly promptCache: MlxPrefixCache & { readonly maxBytes?: number };
-  readonly checkpoints: Pick<SsdCacheStore, "findGenerationCheckpoint" | "restore" |
-    "storeGenerationCheckpoint" | "removeGenerationCheckpoints"> | null;
-  readonly checkpointEveryTokens?: number;
-  readonly checkpointPersistence?: import("./continuation-persistence").ContinuationPersistence;
-  /** Artifact, implementation, state ABI and codec identity captured at load. */
-  readonly identity: unknown;
-  adapterNamespace(adapters: string[]): string;
-  cloneState(caches: Cache[]): Cache[];
-}
-
-/** Cache lookup, replay, prompt snapshots and checkpoint persistence belong to
- * native execution. HTTP only composes this service and frames its output.
- * The caller transfers prepared native inputs; caches created/taken here are
- * owned until returned to the prefix store or disposed after execution. */
-export function createMlxSerialExecutor(binding: MlxSerialBinding, services: MlxSerialServices): SerialRun {
-  const { promptCache, checkpoints: ssdStore } = services;
-  return (promptIds, options, onToken, vision, trace, execution) => withRuntimeConfig(binding.runtime, async () => {
-    if (promptCache.reclaim) {
-      const publish = onToken;
-      let steps = 0;
-      onToken = (token, logprobs) => {
-        if (steps++ % 256 === 0) promptCache.reclaim!();
-        return publish(token, logprobs);
-      };
-    }
-    let releaseContinuation: (() => void) | undefined;
-    let caches: Cache[] = [];
-    let retain: (() => void) | undefined;
-    let closeMedia: (() => void) | undefined;
-    const cleanup = ownResource(null, () => disposeResources([
-      { dispose() {
-        disposeResources(caches);
-        retain?.(); // backing release requires successful cache disposal
-      } },
-      { dispose: () => closeMedia?.() },
-      { dispose: () => releaseContinuation?.() },
-      ...[options.grammar, vision?.embeddings, vision?.imageMask,
-        vision?.multimodalMask, options.visionPixels].filter((value) => value != null),
-    ]));
-    try {
-      options.signal?.throwIfAborted();
-      if (!execution) throw new Error("serial execution requires a resolved plan");
-      if (execution.method === "speculative") {
-        if (!binding.speculate) throw new Error("resolved speculation requires a bound verifier");
-        return await binding.speculate(promptIds, options, onToken, {
-          ...(binding.runtime.flag("MLX_BUN_MTP_PROMPT_CACHE", true) && services.promptCache.maxBytes !== 0
-            ? { prefixCache: services.promptCache } : {}),
-          cloneState: services.cloneState,
-          cacheNamespace: options.adapters?.length ? services.adapterNamespace(options.adapters) : "",
-        });
-      }
-      // Cache entries are adapter-specific: KV computed under one adapter
-      // must never seed another's (or the base's) prefill.
-      const baseNs = options.adapters?.length ? services.adapterNamespace(options.adapters) : "";
-      const cacheNs = execution.pagedKv ? pagedPrefixNamespace(options, baseNs) : baseNs;
-      // Paged-KV request scope (docs/design/kv-cache.md): media
-      // prompts (bidir overlay) and LoRA-adapter requests are v1 non-goals —
-      // they run the PLAIN cache path even under --paged-kv (scope the flag
-      // per request, never 400). Effective value computed ONCE so the
-      // prompt-cache bypass below and the generate() options can't disagree.
-      const pagedKv = execution.pagedKv ? options.pagedKv : undefined;
-      // The resolved policy controls reuse; paged layouts have their own
-      // numerical namespace and use the same RAM/SSD cache interface.
-      const skipPromptCache = !execution.promptCache;
-      const checkpointEvery = services.checkpointEveryTokens;
-      const checkpointEligible = execution.checkpoint;
-      const continuation = checkpointEligible ? bindContinuationPolicy({ store: ssdStore!,
-        restore: entry => ssdStore!.restore(entry, binding), prompt: promptIds,
-        options, execution, identity: services.identity, namespace: cacheNs, persistence: services.checkpointPersistence }) : null;
-      releaseContinuation = continuation?.release;
-      const checkpointKey = continuation?.key;
-      // Both tiers in one call (Layer 0): take() prefers a strictly-longer
-      // SSD prefix, restores it zero-copy, and trims — see PromptCache.take.
-      const closeCacheLookup = trace?.begin("cache.lookup_restore", {
-        mechanism: "serial",
-        bypassed: skipPromptCache,
-      });
-      const checkpoint = continuation?.restore();
-      caches = checkpoint?.caches ?? [];
-      const resuming = Boolean(checkpoint);
-      const generationPromptIds = checkpoint?.cacheTokens ?? promptIds;
-      const entry = skipPromptCache || resuming
-        ? null
-        : promptCache.take(promptIds, cacheNs, options.cacheSessionId);
-      if (entry) { caches = entry.caches; retain = entry.retain; }
-      if (!checkpoint && !entry) caches = binding.makeCache();
-      closeCacheLookup?.();
-      // Prompt-boundary snapshot (the multi-turn agent fix, 2026-07-04): the
-      // prompt+gen entry put() below is UNTRIMMABLE at context > sliding
-      // window (wrapped rings) and under quantized KV (mid-group), so any
-      // decode→encode roundtrip drift in the reply the client sends back
-      // turns the next turn into a total miss (measured: 12B turn-2 TTFT
-      // 8.9 s instead of ~0.2 s). A prompt-ONLY entry is always an exact
-      // prefix of the next turn's rendering regardless of reply drift.
-      // Zero-copy (cloneKvCaches = slice views); only for substantial cold
-      // prefills, where the re-prefill it saves is worth an extra entry.
-      // The oracle invariant (mlx-lm insert_segments): a trim-free STRICT
-      // prefix of the prompt exists for EVERY substantial request — cap the
-      // boundary at len-1 so even a stableLen == len prompt (e4b: the
-      // template tail survives the probe render) snapshots prompt[:-1]. An
-      // exact repeat then matches with trimNeeded == 0, bypassing
-      // isTrimmable() entirely — the only reuse path a wrapped ring has.
-      const boundary = Math.min(options.snapshotAt ?? promptIds.length, promptIds.length - 1);
-      // Re-snapshot on EVERY substantial request whose stable boundary extends
-      // past the cached prefix; the clone is zero-copy views, so re-putting
-      // is ~free.
-      const snapshotBoundary =
-        !skipPromptCache && !resuming && boundary >= 256 &&
-        boundary > (entry?.tokens.length ?? 0);
-      closeMedia = binding.enterMedia?.(vision);
-      if (resuming) {
-        const replay = generationPromptIds.slice(promptIds.length);
-        console.log(
-          `[generation-checkpoint] resuming ${checkpointKey} at ` +
-          `${replay.length} emitted tokens`,
-        );
-        for (const token of replay) {
-          options.signal?.throwIfAborted();
-          if ((await onToken(token)) === false)
-            throw new Error("saved generation prefix triggered a terminal stop while replaying");
-        }
-      }
-      const gen = binding.generate(generationPromptIds, {
-        ...options,
-        decodePolicy: execution,
-        ...(resuming ? { seed: checkpoint!.seed } : {}),
-        fill: execution.fill ? options.fill : undefined,
-        pagedKv, // request-scoped (undefined strips the server-wide flag)
-        cache: caches,
-        ...(resuming
-          ? {
-              initialPendingToken: checkpoint!.pendingToken,
-              initialGeneratedTokens: checkpoint!.generatedTokens,
-              originalPromptTokens: promptIds.length,
-            }
-          : {}),
-        ...(checkpointEligible
-          ? {
-              checkpointEveryTokens: checkpointEvery,
-              onDecodeCheckpoint: async (state: {
-                cacheTokens: number[];
-                caches: Cache[];
-                generatedTokens: number;
-                pendingToken: number;
-              }) => {
-                if (services.checkpointPersistence)
-                  continuation!.captureOwned({ ...state, caches: services.cloneState(state.caches) });
-                else await continuation!.capture(state);
-              },
-            }
-          : {}),
-        ...(snapshotBoundary
-          ? {
-              // snapshotAt MUST travel with the hook: generate() splits the
-              // prefill at exactly this many tokens and fires the hook while
-              // the caches hold exactly that prefix — putting boundary
-              // tokens against caches at any other offset is silent KV
-              // corruption.
-              snapshotAt: boundary,
-              onPrefillDone: () => {
-                try {
-                  const snapshot = ownResource(services.cloneState(caches), disposeResources);
-                  try {
-                    promptCache.put(promptIds.slice(0, boundary), snapshot.borrow(), cacheNs, undefined, undefined, options.cacheSessionId);
-                    snapshot.transfer();
-                  } finally { snapshot.close(); }
-                } catch (e) {
-                  console.warn(`prompt-boundary snapshot skipped: ${(e as Error).message}`);
-                }
-              },
-            }
-          : {}),
-        ...(vision
-          ? {
-              promptEmbeddings: vision.embeddings,
-              ...(vision.imageMask ? { imageMask: vision.imageMask } : {}),
-              ...(vision.multimodalMask ? { multimodalMask: vision.multimodalMask } : {}),
-            }
-          : {}),
-      }, { trace, mechanism: "serial" });
-      for await (const t of gen) {
-        if ((await onToken(t.token, t.logprobs)) === false) break;
-      }
-      const s = gen.stats!; // set on completion AND on early break
-      continuation?.complete();
-      if (!skipPromptCache) {
-        // put() fires onPut → the debounced write-behind SSD snapshot
-        // (wired below), covering the batch lane's puts too.
-        promptCache.put(s.cacheTokens, caches, cacheNs, retain, undefined, options.cacheSessionId);
-        caches = []; retain = undefined; // ownership returned to the prefix store
-      }
-      return s;
-    } catch (error) {
-      cleanupFailure(error, () => cleanup.close());
-    } finally {
-      cleanup.close();
-    }
-  });
-}