mlx-bun 0.5.0 → 1.0.0

package/src/backends/mlx/assistant-target.ts

@@ -1,46 +0,0 @@
-import type { MlxArray } from "../../mlx/array";
-import { Dtype } from "../../mlx/ffi";
-import * as ops from "../../mlx/ops";
-import type { Cache, KvDonorAttention } from "../../model/gemma4-base";
-import { captureKvDonorAttention } from "../../model/kv-attention-view";
-import type { AssistantAttention } from "../../spec/drafter";
-import { disposeResources } from "../../engine/resources";
-
-/** Bind captured cache validity to Q-only attention. Encoded layouts can
- * supply this same numerical port without changing the draft graph. */
-export function readAssistantDonors(sliding: Cache, full: Cache) {
-  const views: KvDonorAttention[] = [], masks: MlxArray[] = [];
-  const dispose = () => disposeResources([...masks, ...views]);
-  try {
-    for (const cache of [sliding, full]) views.push(captureKvDonorAttention(cache));
-    const bind = (view: KvDonorAttention): AssistantAttention => {
-      // All draft steps borrow the same target snapshot. Build each validity
-      // mask once, then reuse it across matching layers and the entire chain.
-      const byWindow = new Map<number | null, MlxArray | null>();
-      return { attend(query, scale, window) {
-        if (!byWindow.has(window)) {
-          const width = view.width, B = view.offsets.length;
-          const starts = view.starts.map((start, row) => Math.max(start,
-            window === null ? start : view.ends[row]! - window));
-          let mask: MlxArray | null = null;
-          if (starts.some(start => start !== 0) || view.ends.some(end => end !== width)) {
-            using columns = ops.arange(0, width, 1, Dtype.int32);
-            using lower = ops.fromInt32(starts, [B, 1, 1, 1]);
-            using upper = ops.fromInt32([...view.ends], [B, 1, 1, 1]);
-            using before = ops.less(columns, lower);
-            using after = ops.greaterEqual(columns, upper);
-            using invalid = ops.logicalOr(before, after);
-            using floats = invalid.astype(view.dtype);
-            mask = ops.mulScalar(floats, -1e9);
-            masks.push(mask);
-          }
-          byWindow.set(window, mask);
-        }
-        const mask = byWindow.get(window)!;
-        return view.attend(query, scale, { mode: mask ? "array" : "", arr: mask });
-      } };
-    };
-    return { positions: views[0]!.offsets.map(offset => offset - 1),
-      sliding: bind(views[0]!), full: bind(views[1]!), dispose };
-  } catch (error) { dispose(); throw error; }
-}

package/src/backends/mlx/autoregressive-draft-rows.ts

@@ -1,120 +0,0 @@
-import type { MlxArray } from "../../mlx/array";
-import * as ops from "../../mlx/ops";
-import type { AutoregressiveGraph } from "../../inference/graph";
-import type { Cache, BatchableCache, PaddedPrefillCache } from "../../model/gemma4-base";
-import type { DraftRowSampling } from "../../spec/source";
-import { toLogprobs } from "../../sampler";
-import { disposeResources } from "../../engine/resources";
-import { MlxStateRows } from "./state-rows";
-import type { PreparedStateChange } from "../../contracts/resources";
-import { evalCacheState } from "./prefill";
-
-export interface MlxDraftLayout extends BatchableCache, PaddedPrefillCache {
-  makeEmptyBatch(): MlxDraftLayout;
-}
-
-/** The draft graph and sampler operate at B. State snapshots own rewind and
- * membership independently of the graph, target verifier and scheduler. */
-export class MlxAutoregressiveDraftRows {
-  #history: MlxStateRows[] = [];
-  #drafts: number[][] = [];
-  #depth = 0;
-  #pendingLast: Array<number | null>;
-  constructor(readonly graph: AutoregressiveGraph<MlxArray, Cache[], MlxArray>,
-    readonly state: MlxStateRows<MlxDraftLayout>, readonly sampling: DraftRowSampling | null) {
-    this.#pendingLast = Array.from({ length: state.rowCount }, () => null);
-  }
-  get pendingLast(): readonly (number | null)[] { return this.#pendingLast; }
-  prepareAppend(rows: readonly (readonly Cache[])[], pending: readonly (number | null)[]): PreparedStateChange {
-    const state = this.state.prepareAppendMany(rows), next = [...this.#pendingLast, ...pending];
-    return { commit: () => { state.commit(); this.#pendingLast = next; }, dispose: () => state.dispose() };
-  }
-  filterRows(keep: readonly number[]): void {
-    this.state.filterRows(keep); this.#pendingLast = keep.map(row => this.#pendingLast[row]!);
-  }
-
-  async #forward(ids: MlxArray, lengths?: readonly number[]): Promise<MlxArray> {
-    const width = ids.shape[1]!, padded = lengths?.some(length => length !== width);
-    if (padded) for (const cache of this.state.caches) cache.preparePrefill({
-      lengths: [...lengths!], rightPadding: lengths!.map(length => width - length),
-    });
-    const hidden = await this.graph.forwardHidden(ids, this.state.caches);
-    try {
-      if (padded) for (const cache of this.state.caches) cache.finalizePrefill();
-      return hidden;
-    } catch (error) { hidden.dispose(); throw error; }
-  }
-
-  /** A restored full-accept boundary may have one unprocessed draft token.
-   * Combine it with the incoming block on device, retaining per-row lengths. */
-  async prefill(tokens: MlxArray): Promise<void> {
-    const [B, N] = tokens.shape as [number, number];
-    if (this.#pendingLast.some(token => token !== null)) {
-      using previous = ops.fromInt32(this.#pendingLast.map(token => token ?? 0), [B, 1]);
-      using joined = ops.concatAxis([previous, tokens], 1);
-      const lengths = this.#pendingLast.map(token => N + Number(token !== null));
-      using indices = ops.fromInt32(lengths.flatMap(length => Array.from({ length: N + 1 },
-        (_, position) => length === N + 1 ? position : Math.min(position + 1, N))), [B, N + 1]);
-      using inputs = ops.takeAlongAxis(joined, indices, 1);
-      (await this.#forward(inputs, lengths)).dispose();
-    } else (await this.#forward(tokens)).dispose();
-    this.#pendingLast.fill(null); this.materialize();
-  }
-  materialize(): void { evalCacheState(this.state.caches); }
-
-  async draft(pending: readonly number[], depth: number, steps: readonly number[]): Promise<number[][]> {
-    this.#releaseHistory();
-    this.#depth = depth;
-    const tokens: MlxArray[] = [];
-    const feeds = pending.map((token, row) => this.#pendingLast[row] == null ? [token] : [this.#pendingLast[row]!, token]);
-    const lengths = feeds.map(feed => feed.length), width = Math.max(...lengths);
-    let ids = ops.fromInt32(feeds.flatMap(feed => [...feed, ...Array(width - feed.length).fill(0)]), [pending.length, width]);
-    try {
-      // With zero proposals, consume the pending input to keep companion
-      // coverage aligned with the target's one-token verification.
-      for (let position = 0; position < Math.max(1, depth); position++) {
-        const padded = position === 0 && lengths.some(length => length !== width);
-        using hidden = await this.#forward(ids, position === 0 ? lengths : undefined);
-        this.#history.push(this.state.clone());
-        if (depth === 0) break;
-        using selection = padded ? ops.fromInt32(lengths.map(length => length - 1), [pending.length, 1, 1]) : null;
-        using selected = selection ? ops.takeAlongAxis(hidden, selection, 1) : null;
-        using logits = this.graph.projectLogits(selected ?? hidden, { type: "last" });
-        using flat = ops.reshape(logits, [pending.length, logits.shape.at(-1)!]);
-        using logprobs = toLogprobs(flat);
-        const token = this.sampling!.sample(logprobs, steps.map(step => step + position));
-        tokens.push(token);
-        ids.dispose(); ids = ops.reshape(token, [pending.length, 1]);
-      }
-      if (depth === 0) this.#drafts = pending.map(() => []);
-      else {
-        using packed = ops.concatAxis(tokens, 0);
-        const values = packed.toIntTokens();
-        this.#drafts = pending.map((_, row) => Array.from({ length: depth },
-          (_, position) => values[position * pending.length + row]!));
-      }
-      return this.#drafts;
-    } catch (error) { this.#releaseHistory(); throw error; }
-    finally { ids.dispose(); disposeResources(tokens); }
-  }
-
-  /** Preserve mlx-lm's full-accept re-feed shape: the final draft remains
-   * pending and joins the next target token in one forward. Recurrent state
-   * uses saved boundaries, never trimming. The lag is companion state. */
-  async commit(accepted: readonly number[]): Promise<void> {
-    const selected: Cache[][] = [];
-    try {
-      if (!accepted.every(count => count >= this.#depth - 1)) {
-        for (const [row, count] of accepted.entries())
-          selected.push(this.#history[Math.min(count, Math.max(0, this.#depth - 1))]!.extractRow(row));
-        this.state.mergeRows(selected);
-      }
-    } finally { disposeResources(selected.flat()); this.#releaseHistory(); }
-
-    this.#pendingLast = accepted.map((count, row) => this.#depth > 0 && count === this.#depth ? this.#drafts[row]!.at(-1)! : null);
-    this.#drafts = []; this.#depth = 0;
-  }
-
-  #releaseHistory(): void { disposeResources(this.#history.splice(0)); }
-  dispose(): void { this.#releaseHistory(); this.#drafts = []; this.state.dispose(); }
-}

package/src/backends/mlx/autoregressive.ts

@@ -1,129 +0,0 @@
-import type { AutoregressiveGraph } from "../../inference/graph";
-import type { MlxArray } from "../../mlx/array";
-import type { Cache, Gemma4Model } from "../../model/gemma4";
-import type { RuntimeModel } from "../../model/factory";
-import { CompiledDecode } from "../../model/compiled-decode";
-import { runtimeConfig, type RuntimeConfig } from "../../runtime-config";
-import { bindMlxGraph } from "./graph";
-import type { PrefillPolicy } from "../../inference/prefill";
-import { resolveMlxPrefillPolicy } from "./prefill-policy";
-import type { KvSchemeOptions } from "../../kv-scheme";
-
-/** One declaration shared by planning and both native execution lanes. */
-export function legacyCompiledDecodeAvailable(model: RuntimeModel): boolean {
-  return model.config?.modelType?.startsWith("gemma4") === true && !model.config.text.enableMoeBlock;
-}
-
-export interface MlxModelMemory {
-  readonly weightsBytes: number;
-  readonly expertRuntime?: {
-    readonly plan: { readonly plannedBytes: number };
-    flushUsage?: () => void;
-    finishUsage?: () => Promise<void>;
-  } | null;
-}
-
-export interface MlxDecodeStep {
-  /** Consume one pending token and advance state once. Returned arrays belong
-   * to the caller. null permits ordinary graph execution and MUST leave state
-   * unchanged. Errors propagate unless this implementation can roll back. */
-  tryStep(token: MlxArray, state: Cache[]): { logits: MlxArray; evalWith: MlxArray[] } | null;
-  /** Release per-run workspaces after pending outputs are dropped and before
-   * cache disposal. Wait for native work if releasing requires completion. */
-  close(): void | Promise<void>;
-}
-
-/** Model-owned execution for committed tokens. Recheck the chunk limit after
- * each forward: native arithmetic can change at a cache-length boundary. */
-export interface MlxTokenAppend {
-  /** Affine formats whose committed append retains one-token arithmetic. */
-  readonly affineKvBits?: readonly number[];
-  readonly turboQuantFormats?: readonly { readonly kBits: number; readonly vBits: number }[];
-  /** Maximum positions per row at this cohort size; omitted rows means B=1. */
-  maxChunkSize(state: readonly Cache[], rows?: number): number;
-  forwardHidden(ids: MlxArray, state: Cache[]): MlxArray | Promise<MlxArray>;
-}
-
-/** The model declares which codecs retain its committed-span arithmetic. */
-export function supportsCommittedAppendCache(
-  append: Pick<MlxTokenAppend, "affineKvBits" | "turboQuantFormats"> | null | undefined,
-  options: Pick<KvSchemeOptions, "kvBits" | "kvConfig" | "turboQuant">,
-): boolean {
-  if (options.turboQuant) return !!append?.turboQuantFormats?.some(format =>
-    format.kBits === options.turboQuant!.kBits && format.vBits === options.turboQuant!.vBits);
-  return options.kvConfig?.length
-    ? options.kvConfig.every(layer => append?.affineKvBits?.includes(layer.bits))
-    : !options.kvBits || !!append?.affineKvBits?.includes(options.kvBits);
-}
-
-/** One replaceable binding owns all graph-specific operations used by the AR
- * loop. The current method requires MLX tensors and the legacy Cache[] ABI.
- * Weights are borrowed for the binding's lifetime; generated caches are owned
- * by the run. Caller-provided caches and media remain borrowed. */
-export interface MlxAutoregressiveBinding {
-  readonly runtime?: RuntimeConfig;
-  readonly prefillPolicy?: PrefillPolicy;
-  readonly graph: AutoregressiveGraph<MlxArray, Cache[], MlxArray>;
-  readonly eosTokenIds: readonly number[];
-  readonly memory: MlxModelMemory;
-  readonly adapters?: { active: string[] };
-  makeCache(): Cache[];
-  forwardEmbeddings?(
-    embeddings: MlxArray, state: Cache[], imageMask: MlxArray | null,
-    ids: MlxArray, multimodalMask: MlxArray | null,
-  ): MlxArray;
-  /** Called once after prefill. A compiled/fused decoder belongs to THIS graph,
-   * never to a model inferred from a name or inherited by a replacement. */
-  createDecode?(policy: { hasAdapters: boolean; pagedKv: boolean }): MlxDecodeStep | null;
-  createAppend?(policy: { hasAdapters: boolean; pagedKv: boolean }): MlxTokenAppend | null;
-}
-
-/** Keep concrete model and compiled-decode decisions at the legacy boundary. */
-export function bindLegacyAutoregressiveModel(model: RuntimeModel): MlxAutoregressiveBinding {
-  const runtime = runtimeConfig();
-  return {
-    runtime,
-    prefillPolicy: resolveMlxPrefillPolicy(model.config, runtime),
-    graph: bindMlxGraph<Cache[]>(model, {
-      id: `legacy:${model.config.modelType}`, artifact: "legacy-resident-model",
-      stateAbi: "legacy-cache-array-v1", // not a persistence identity
-    }),
-    eosTokenIds: model.config.eosTokenIds,
-    memory: model,
-    adapters: model.loraState,
-    makeCache: model.makeCache.bind(model),
-    forwardEmbeddings: model.forwardEmbeddings?.bind(model),
-    createAppend: "createAppend" in model ? model.createAppend.bind(model) : undefined,
-    createDecode(policy) {
-      // Preserve the existing Gemma path and exclusions. MoE shapeless replay
-      // retraces growing windows; adapters would bake residuals into the tape.
-      if (!runtime.flag("MLX_BUN_COMPILED_DECODE", true) || policy.hasAdapters || policy.pagedKv ||
-          !legacyCompiledDecodeAvailable(model))
-        return null;
-      let compiled: CompiledDecode | null = CompiledDecode.for(model as Gemma4Model);
-      return {
-        tryStep(token, state) {
-          if (!compiled || !CompiledDecode.supports(state)) return null;
-          try { return compiled.step(token, state); }
-          catch (error) {
-            // CompiledDecode restores committed writes on failure, so retrying
-            // this token through the ordinary graph is safe. Other decoders
-            // must establish that guarantee themselves before returning null.
-            compiled = null;
-            console.warn(`compiled decode disabled for this generation: ${error}`);
-            return null;
-          }
-        },
-        // Compiled closures belong to the model's existing unload lifecycle.
-        close() { compiled = null; },
-      };
-    },
-  };
-}
-
-export function assertMlxAutoregressiveBinding(binding: MlxAutoregressiveBinding): void {
-  const descriptor = binding.graph.descriptor;
-  if (descriptor.backend !== "mlx" || descriptor.graphAbi !== "mlx-hidden-bsh-v1" ||
-      descriptor.stateAbi !== "legacy-cache-array-v1")
-    throw new Error(`AR binding ${descriptor.id} has an incompatible backend, graph, or state ABI`);
-}