mlx-bun 0.3.0 → 0.5.0

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

@@ -0,0 +1,46 @@
+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

@@ -0,0 +1,120 @@
+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

@@ -5,6 +5,9 @@ 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 {
@@ -33,16 +36,33 @@ export interface MlxDecodeStep {
 /** 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 {
-  maxChunkSize(state: readonly Cache[]): number;
+  /** 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;
@@ -63,6 +83,7 @@ export function bindLegacyAutoregressiveModel(model: RuntimeModel): MlxAutoregre
   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