mlx-bun 0.3.0 → 0.5.0

package/src/backends/mlx/speculative.ts

@@ -10,12 +10,16 @@ import type { SpeculativeTransaction } from "../../inference/rollback";
 import { bindCacheRollback } from "./rollback";
 import type { GraphDescriptor } from "../../inference/graph";
 import { bindLegacyDraftTarget } from "./draft-target";
+import type { MlxForwardWork } from "./mixed-iteration";
+import type { PrefillPolicy } from "../../inference/prefill";
+import { resolveMlxPrefillPolicy } from "./prefill-policy";
 
 /** Bound target operations and draft construction for one speculative run.
  * A replacement graph supplies this entire port, including any hidden taps.
  * The verifier never needs a RuntimeModel or an artifact-family check. */
 export interface MlxSpeculativeBinding {
   readonly runtime?: RuntimeConfig;
+  readonly prefillPolicy?: PrefillPolicy;
   readonly memory?: MlxModelMemory;
   readonly descriptor: GraphDescriptor;
   readonly eosTokenIds: readonly number[];
@@ -23,27 +27,29 @@ export interface MlxSpeculativeBinding {
   makeCache(): Cache[];
   openDraft(sampler: Sampler, caches: Cache[]): DraftSource;
   bindRollback(caches: Cache[]): SpeculativeTransaction;
-  forward(ids: MlxArray, caches: Cache[], tapLayers?: number[]):
+  forward(ids: MlxArray, caches: Cache[], tapLayers?: number[], work?: MlxForwardWork):
     Promise<{ hidden: MlxArray; ctxML: MlxArray | null }>;
   projectLogits(hidden: MlxArray): MlxArray;
   /** Establish the implementation's verify kernel context; restore on close. */
   pinVerify?(): { close(): void };
 }
 
+export type MlxSpeculativeTargetBinding = Omit<MlxSpeculativeBinding, "openDraft">;
+
 /** Legacy mutable tap/kernel fields require the gateway's exclusive lease. */
-export function bindLegacySpeculativeModel(model: RuntimeModel, provider: DraftProvider): MlxSpeculativeBinding {
+export function bindSpeculativeTargetModel(model: RuntimeModel): MlxSpeculativeTargetBinding {
   const runtime = runtimeConfig();
   return {
     runtime,
+    prefillPolicy: resolveMlxPrefillPolicy(model.config, runtime),
     memory: model,
     descriptor: Object.freeze({ id: `legacy-spec:${model.config.modelType}`, backend: "mlx",
       graphAbi: "mlx-hidden-bsh-v1", stateAbi: "legacy-cache-array-v1", artifact: "legacy-resident-model" }),
     eosTokenIds: model.config.eosTokenIds,
     prefillTailSplit: runtime.flag("MLX_BUN_PREFILL_TAIL_SPLIT", true),
     makeCache: model.makeCache.bind(model),
-    openDraft: (sampler, caches) => provider.open({ sampler, target: bindLegacyDraftTarget(model, caches) }),
     bindRollback: bindCacheRollback,
-    forward: (ids, caches, tapLayers) => legacyForwardWithTaps(model, ids, caches, tapLayers),
+    forward: (ids, caches, tapLayers, work) => legacyForwardWithTaps(model, ids, caches, tapLayers, work),
     projectLogits: model.logitsFromHidden.bind(model),
     ...("setSpecKernelPinned" in model ? {
       pinVerify() {
@@ -54,6 +60,12 @@ export function bindLegacySpeculativeModel(model: RuntimeModel, provider: DraftP
   };
 }
 
+/** Serial compatibility binds draft construction separately from target work. */
+export function bindLegacySpeculativeModel(model: RuntimeModel, provider: DraftProvider): MlxSpeculativeBinding {
+  return { ...bindSpeculativeTargetModel(model),
+    openDraft: (sampler, caches) => provider.open({ sampler, target: bindLegacyDraftTarget(model, caches) }) };
+}
+
 export function assertMlxSpeculativeBinding(binding: MlxSpeculativeBinding): void {
   const descriptor = binding.descriptor;
   if (descriptor.backend !== "mlx" || descriptor.graphAbi !== "mlx-hidden-bsh-v1" ||
@@ -71,8 +83,10 @@ async function legacyForwardWithTaps(
   ids: MlxArray,
   caches: Cache[],
   tapLayers: number[] | undefined,
+  work?: MlxForwardWork,
 ): Promise<{ hidden: MlxArray; ctxML: MlxArray | null }> {
   if (!tapLayers) {
+    if (work) return { hidden: await work(ids, caches), ctxML: null };
     const asyncModel = model as RuntimeModel & {
       forwardHiddenAsync?: (
         ids: MlxArray,
@@ -88,10 +102,13 @@ async function legacyForwardWithTaps(
   const m = model;
   const previousTap = m.hiddenTap;
   const cap = new Map<number, MlxArray>();
-  m.hiddenTap = { layers: new Set(tapLayers), captured: cap };
+  const layers = new Set(tapLayers);
+  if (!work) m.hiddenTap = { layers, captured: cap };
   let hidden: MlxArray | null = null;
   try {
-    hidden = model.forwardHidden(ids, caches);
+    hidden = work ? await work(ids, caches, { captureLayer(layer, h) {
+      if (layers.has(layer)) cap.set(layer, ops.contiguous(h));
+    } }) : model.forwardHidden(ids, caches);
     const perLayer = tapLayers.map((li) => {
       const a = cap.get(li);
       if (!a) throw new Error(`spec tap: layer ${li} not captured`);
@@ -109,6 +126,6 @@ async function legacyForwardWithTaps(
     // are already gone (cap cleared, hidden nulled).
     hidden?.dispose();
     for (const [, a] of cap) a.dispose();
-    m.hiddenTap = previousTap;
+    if (!work) m.hiddenTap = previousTap;
   }
 }

package/src/backends/mlx/state-rows.ts

@@ -0,0 +1,104 @@
+import type { BatchableCache, Cache } from "../../model/gemma4-base";
+import { applyStateChanges, cleanupFailure, disposeResources } from "../../engine/resources";
+import type { PreparedStateChange } from "../../contracts/resources";
+
+/** Preserve the graph's layout capabilities when an empty sibling is made. */
+export interface MlxRowLayout<Layout extends BatchableCache> extends BatchableCache {
+  makeEmptyBatch(): Layout;
+}
+
+/** Layer layouts own tensor geometry. This collection only applies the same
+ * row membership operation to every layer. It has no scheduling, sampling,
+ * checkpoint-key or persistence policy. */
+export class MlxStateRows<Layout extends MlxRowLayout<Layout> = BatchableCache> {
+  /** Borrowed by the bound target graph; collection ownership stays here. */
+  readonly caches: Layout[];
+  #rows = 0;
+
+  /** Takes ownership of empty, compatible layer layouts. */
+  constructor(layouts: readonly Layout[]) { this.caches = [...layouts]; }
+  get rowCount(): number { return this.#rows; }
+  /** Retain an immutable batch boundary without extracting or materializing
+   * individual rows. Layer layouts decide how to share their backing arrays. */
+  clone(): MlxStateRows<Layout> {
+    const layouts: Layout[] = [];
+    try {
+      for (const source of this.caches) {
+        const layout = source.makeEmptyBatch(); layouts.push(layout);
+        layout.mergeRows([source]);
+      }
+      const result = new MlxStateRows(layouts); result.#rows = this.#rows;
+      return result;
+    } catch (error) { return cleanupFailure(error, () => disposeResources(layouts)); }
+  }
+  /** Prepare an initial cohort with one merge per layer. */
+  mergeRows(rows: readonly (readonly Cache[])[]): void {
+    const next: Layout[] = [];
+    try {
+      for (let layer = 0; layer < this.caches.length; layer++) {
+        const layout = this.caches[layer]!.makeEmptyBatch(); next.push(layout);
+        layout.mergeRows(rows.map(row => row[layer]!));
+      }
+    } catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
+    const previous = this.caches.splice(0, this.caches.length, ...next);
+    this.#rows = rows.length; disposeResources(previous);
+  }
+
+  /** Borrow a prepared request's state. Build every replacement before
+   * releasing the active group, so a failed admission leaves it intact. */
+  append(row: readonly Cache[]): void {
+    applyStateChanges([() => this.prepareAppend(row)]);
+  }
+
+  /** Stage membership so a method can compose it with companion-state admission. */
+  prepareAppend(row: readonly Cache[]): PreparedStateChange {
+    return this.prepareAppendMany([row]);
+  }
+
+  /** Stage several admissions with one membership operation per layer. */
+  prepareAppendMany(rows: readonly (readonly Cache[])[]): PreparedStateChange {
+    const next: Layout[] = [];
+    try {
+      for (let layer = 0; layer < this.caches.length; layer++) {
+        const layout = this.caches[layer]!;
+        const merged = layout.makeEmptyBatch();
+        next.push(merged);
+        merged.mergeRows([...(this.#rows ? [layout] : []), ...rows.map(row => row[layer]!)]);
+      }
+    } catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
+    return {
+      commit: () => {
+        const previous = this.caches.splice(0, this.caches.length, ...next.splice(0));
+        this.#rows += rows.length;
+        next.push(...previous);
+      },
+      dispose: () => disposeResources(next.splice(0)),
+    };
+  }
+
+  /** The method calls this only between committed steps. A layer failure
+   * invalidates the group; its caller discards the state. */
+  filterRows(keep: readonly number[]): void {
+    if (!keep.length) { this.clear(); return; }
+    for (const layout of this.caches) layout.filterRows(keep);
+    this.#rows = keep.length;
+  }
+
+  /** Returns independent request state for a cache publisher or a new group. */
+  extractRow(row: number): Cache[] {
+    const state: Cache[] = [];
+    try { for (const layout of this.caches) state.push(layout.extractRow(row)); return state; }
+    catch (error) { return cleanupFailure(error, () => disposeResources(state)); }
+  }
+
+  clear(): void {
+    const next: Layout[] = [];
+    try { for (const layout of this.caches) next.push(layout.makeEmptyBatch()); }
+    catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
+    const previous = this.caches.splice(0, this.caches.length, ...next);
+    this.#rows = 0;
+    disposeResources(previous);
+  }
+
+  dispose(): void { this.#rows = 0; disposeResources(this.caches); }
+}

package/src/backends/mlx/state-views.ts

@@ -1,7 +1,8 @@
 import type { StateView } from "../../contracts/resources";
 import { disposeResources, ownResource } from "../../engine/resources";
 import type { MlxArray } from "../../mlx/array";
-import type { Cache } from "../../model/gemma4-base";
+import { isBatchableCache, type Cache } from "../../model/gemma4-base";
+import { cloneKvCaches, type CacheCodecProvider } from "../../kv-store";
 
 /** Compatibility boundary for legacy state(). Numerical code consumes a lease,
  * never the legacy ownership marker. No copy, evaluation or fence is added. */
@@ -24,3 +25,22 @@ export function leaseCacheStates(caches: readonly Cache[]): StateView<MlxArray>
     throw error;
   }
 }
+
+/** Precision/storage transitions constrain reuse independently of scheduling. */
+export function minimumReusableOffset(caches: readonly Pick<Cache, "minimumReusableOffset">[]): number {
+  return Math.max(0, ...caches.map(cache => cache.minimumReusableOffset ?? 0));
+}
+
+/** Snapshot an adopted single-row state into persistence-compatible caches.
+ * Row layouts own compact extraction; ordinary serial caches use their codec. */
+export function cloneSingleRowState(caches: readonly Cache[], codecs?: CacheCodecProvider): Cache[] {
+  const owned: Cache[] = [];
+  try {
+    for (const cache of caches) owned.push(...(isBatchableCache(cache)
+      ? [cache.extractRow(0)] : cloneKvCaches([cache], codecs)));
+    return owned;
+  } catch (error) {
+    try { disposeResources(owned); } catch (cleanupError) { throw new AggregateError([error, cleanupError], "single-row snapshot failed"); }
+    throw error;
+  }
+}

package/src/backends/mlx/target-layout-capability.ts

@@ -0,0 +1,35 @@
+import { Glm52Cache } from "../../model/glm52-cache";
+import { SpeculativeRotatingKVCache } from "../../model/speculative-rotating-kv";
+import { RotatingAffineLayout, SpeculativeRotatingAffineLayout } from "../../model/rotating-kv-layout";
+import type { Cache } from "../../model/gemma4-base";
+import { KVCache, QuantizedKVCache, TurboQuantKVCache, RotatingKVCache, RotatingQuantizedKVCache } from "../../model/gemma4-base";
+import { DelayedRotatingQuantizedKVCache } from "../../model/delayed-rotating-quantized-kv";
+import { DelayedQuantizedKVCache } from "../../model/delayed-quantized-kv";
+import { DelayedTurboQuantKVCache } from "../../model/delayed-turboquant-kv";
+import { SSMCache } from "../../model/qwen3-delta";
+import { BatchedSSMCache } from "../../model/batched-ssm";
+import { BatchedKVCache } from "../../model/batched-kv";
+import { BatchedQuantizedKVCache } from "../../model/batched-quantized-kv";
+import { BatchedTurboQuantKVCache } from "../../model/batched-turboquant-kv";
+
+export type TargetRowLayout = Glm52Cache | DelayedRotatingQuantizedKVCache | RotatingAffineLayout | SpeculativeRotatingKVCache | BatchedKVCache | BatchedQuantizedKVCache | BatchedTurboQuantKVCache |
+  BatchedSSMCache | DelayedTurboQuantKVCache | DelayedQuantizedKVCache;
+
+/** The storage binding advertises a row transaction factory. The gateway can
+ * compose supported model graphs without identifying their model family. */
+export function targetRowLayoutFactory(cache: Cache): (() => TargetRowLayout) | undefined {
+  if (cache instanceof Glm52Cache) return () => cache.makeEmptyBatch();
+  if (cache instanceof DelayedRotatingQuantizedKVCache)
+    return () => new DelayedRotatingQuantizedKVCache(cache.maxSize, cache.groupSize, cache.bits, cache.start, cache.maintain, undefined, true);
+  if (cache instanceof RotatingAffineLayout || cache instanceof RotatingQuantizedKVCache)
+    return () => new SpeculativeRotatingAffineLayout(cache.maxSize, cache.groupSize, cache.bits);
+  if (cache instanceof RotatingKVCache || cache instanceof SpeculativeRotatingKVCache) return () => new SpeculativeRotatingKVCache(cache.maxSize);
+  if (cache instanceof DelayedTurboQuantKVCache || cache instanceof DelayedQuantizedKVCache ||
+      cache instanceof BatchedTurboQuantKVCache || cache instanceof BatchedKVCache ||
+      cache instanceof BatchedQuantizedKVCache || cache instanceof BatchedSSMCache) return () => cache.makeEmptyBatch();
+  if (cache instanceof TurboQuantKVCache) return () => new BatchedTurboQuantKVCache(cache.kBits, cache.vBits, cache.fusedDecode);
+  if (cache instanceof SSMCache) return () => new BatchedSSMCache();
+  if (cache instanceof QuantizedKVCache) return () => new BatchedQuantizedKVCache(cache.groupSize, cache.bits);
+  if (cache instanceof KVCache) return () => new BatchedKVCache();
+  return undefined;
+}

package/src/backends/mlx/transitioning-kv-rows.ts

@@ -0,0 +1,173 @@
+import type { MlxArray } from "../../mlx/array";
+import * as ops from "../../mlx/ops";
+import type { BatchableCache, Cache, Mask } from "../../model/gemma4-base";
+import { KvTensorRows } from "./kv-tensor-rows";
+import { leaseCacheStates, minimumReusableOffset } from "./state-views";
+import { withResource, disposeResources } from "../../engine/resources";
+import { cloneKvCaches } from "../../kv-store";
+
+/** Encoded layouts retain physical row alignment across precision changes. */
+export interface TransitionedKvLayout extends BatchableCache {
+  alignRows?(leftPad: readonly number[]): void;
+}
+export interface SpeculativeTransitionedKvLayout extends TransitionedKvLayout {
+  specRoundBegin(): void;
+  specRoundCommit(): void;
+  specRoundRollback(keep: number | readonly number[]): void;
+}
+export interface KvRowTransition<Layout extends TransitionedKvLayout> {
+  readonly signature: string;
+  readonly conversionOffset?: number;
+  maintain(rows: Cache[]): void;
+  converted(row: Cache): boolean;
+  makeLayout(): Layout;
+  packRows?(layout: Layout, rows: readonly Cache[]): void;
+  prepareRows?(rows: Cache[]): Cache[];
+  extractRow?(row: Cache): Cache;
+  rollbackRow?(row: Cache, before: number, keep: number, preserveColumns: boolean): void;
+}
+
+
+export interface TransitioningKvPositions {
+  readonly rowOffsets: readonly number[];
+  readonly leftPad: readonly number[];
+  readonly offset: number;
+  readonly batchSize: number | null;
+  readonly ropeOffsetArr: MlxArray | undefined;
+  sync(rows: readonly Cache[], preserve: boolean): void;
+  filterRows(keep: readonly number[]): void;
+  makeMask(tokens: number, window: number | null): Mask;
+  dispose(): void;
+}
+export class FullKvPositions extends KvTensorRows implements TransitioningKvPositions {
+  sync(rows: readonly Cache[], preserve: boolean): void {
+    this.mergeRows(preserve
+      ? [{ planes: [], rowOffsets: rows.map(row => row.offset), leftPad: [...this.leftPad] }]
+      : rows.map(row => ({ planes: [], rowOffsets: [row.offset], leftPad: [0] })));
+  }
+}
+
+/** Owns membership, committed positions and precision transitions. Codec and
+ * attention adapters own tensor access; scheduling and persistence use the
+ * existing batch and checkpoint contracts. Row conversion never examines an
+ * uncommitted speculative suffix. */
+export abstract class TransitioningKvRows<Layout extends TransitionedKvLayout> implements BatchableCache {
+  protected readonly positions: TransitioningKvPositions;
+  protected rows: Cache[] = [];
+  protected packed?: Layout;
+  #prefilling = false;
+
+  readonly stateNeedsDispose = true;
+  constructor(readonly transition: KvRowTransition<Layout>, row?: Cache, positions: TransitioningKvPositions = new FullKvPositions()) {
+    this.positions = positions;
+    if (row) { this.rows = transition.prepareRows?.([row]) ?? [row]; this.syncPositions(); }
+  }
+  signature(): string { return this.transition.signature; }
+  get prefillMaintenance() { return this; }
+  beginPrefill(): void { this.#prefilling = true; }
+  commitPrefill(indices: readonly number[]): void {
+    if (this.packed) return;
+    for (const index of indices) {
+      const row = [this.rows[index]!];
+      try { this.transition.maintain(row); }
+      finally { this.rows[index] = row[0]!; }
+    }
+  }
+  endPrefill(): void { this.#prefilling = false; }
+  get minimumReusableOffset(): number { return this.packed?.minimumReusableOffset ?? minimumReusableOffset(this.rows); }
+  get rowOffsets(): readonly number[] { return this.packed?.rowOffsets ?? this.positions.rowOffsets; }
+  get leftPad(): readonly number[] { return this.packed?.leftPad ?? this.positions.leftPad; }
+  get offset(): number { return this.packed?.offset ?? this.positions.offset; }
+  maxAppendTokens(): number {
+    if (this.packed) return this.packed.maxAppendTokens?.() ?? Infinity;
+    const end = this.transition.conversionOffset;
+    if (end === undefined) return Infinity;
+    return Math.min(Infinity, ...this.rowOffsets.map((offset, row) =>
+      offset < end && !this.transition.converted(this.rows[row]!) ? end - offset : Infinity));
+  }
+  get batchSize(): number | null { return this.packed?.batchSize ?? this.positions.batchSize; }
+  get ropeOffsetArr(): MlxArray | undefined { return this.packed ? this.packed.ropeOffsetArr : this.positions.ropeOffsetArr; }
+  bytesPerToken(): number { return this.packed?.bytesPerToken?.() ?? Math.max(0, ...this.rows.map(row => row.bytesPerToken?.() ?? 0)); }
+  projectedBytes(tokens: number): number { return this.bytesPerToken() * tokens; }
+  makeMask(tokens: number, window: number | null): Mask { return (this.packed ?? this.positions).makeMask(tokens, window); }
+  isTrimmable(): boolean { return this.packed?.isTrimmable() ?? this.rows.every(row => row.isTrimmable()); }
+  trim(count: number): void {
+    if (this.packed) this.packed.trim(count);
+    else { for (const row of this.rows) row.trim(count); this.syncPositions(true); }
+  }
+  protected syncPositions(preserve = false): void { this.positions.sync(this.rows, preserve); }
+  protected advance(): void {
+    if (this.packed) return;
+    if (!this.#prefilling) this.transition.maintain(this.rows);
+    if (this.rows.length && this.rows.every(row => this.transition.converted(row))) {
+      const packed = this.transition.makeLayout();
+      try {
+        if (this.transition.packRows) this.transition.packRows(packed, this.rows);
+        else packed.mergeRows(this.rows);
+        packed.alignRows?.(this.positions.leftPad);
+      }
+      catch (error) { packed.dispose(); throw error; }
+      const previous = this.rows; this.rows = []; this.packed = packed;
+      this.positions.dispose(); disposeResources(previous);
+    }
+  }
+  abstract updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray];
+  abstract makeEmptyBatch(): TransitioningKvRows<Layout>;
+  mergeRows(rows: readonly Cache[]): void {
+    const owned: Cache[] = [];
+    try {
+      for (const row of rows) {
+        if (row instanceof TransitioningKvRows) {
+          for (let index = 0; index < row.batchSize!; index++) owned.push(row.extractRow(index));
+        } else owned.push(...cloneKvCaches([row]));
+      }
+    } catch (error) { disposeResources(owned); throw error; }
+    let prepared: Cache[];
+    try { prepared = this.transition.prepareRows?.(owned) ?? owned; }
+    catch (error) { disposeResources(owned); throw error; }
+    this.dispose(); this.rows = prepared; this.syncPositions();
+  }
+  filterRows(keep: readonly number[]): void {
+    if (this.packed) { this.packed.filterRows(keep); return; }
+    const previous = this.rows; this.rows = keep.map(index => previous[index]!);
+    disposeResources(previous.filter((_, index) => !keep.includes(index))); this.positions.filterRows(keep);
+  }
+  extractRow(row: number): Cache { return this.packed?.extractRow(row) ?? (this.transition.extractRow?.(this.rows[row]!) ?? cloneKvCaches([this.rows[row]!])[0]!); }
+  state(): MlxArray[] {
+    return withResource(leaseCacheStates(this.packed ? [this.packed] : this.rows), arrays => {
+      const views: MlxArray[] = [];
+      try { for (const array of arrays) views.push(ops.contiguous(array)); return views; }
+      catch (error) { disposeResources(views); throw error; }
+    });
+  }
+  dispose(): void {
+    this.#prefilling = false;
+    const previous = this.rows; this.rows = [];
+    this.packed?.dispose(); this.packed = undefined; this.positions.dispose(); disposeResources(previous);
+  }
+}
+
+/** Rollback is a method capability, separate from ordinary cache membership. */
+export abstract class SpeculativeTransitioningKvRows<Layout extends SpeculativeTransitionedKvLayout> extends TransitioningKvRows<Layout> {
+  #beforeRound?: number[];
+  specRoundBegin(): void {
+    this.advance(); this.#beforeRound = [...this.rowOffsets]; this.packed?.specRoundBegin();
+  }
+  specRoundCommit(): void {
+    this.packed?.specRoundCommit(); this.#beforeRound = undefined; this.advance();
+  }
+  specRoundRollback(keep: number | readonly number[]): void {
+    if (this.packed) this.packed.specRoundRollback(keep);
+    else {
+      for (const [index, row] of this.rows.entries()) {
+        const count = typeof keep === "number" ? keep : keep[index]!;
+        const before = this.#beforeRound![index]!;
+        if (this.transition.rollbackRow) this.transition.rollbackRow(row, before, count, this.rows.length > 1);
+        else row.trim(row.offset - before - count);
+      }
+      this.syncPositions(true);
+    }
+    this.#beforeRound = undefined; this.advance();
+  }
+  override dispose(): void { this.#beforeRound = undefined; super.dispose(); }
+}