@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/generation/speculative/bindings/assistant-rows.ts

@@ -0,0 +1,102 @@
+import { rowOperationFor } from "../../../runtime/row-operations";
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { materializeCopy } from "@mlx-bun/mlx/materialize";
+import { disposeResources, applyStateChanges } from "../../../runtime/resources";
+import type { PreparedStateChange } from "../../../contracts/portable/resources";
+import { targetLacks, type AssistantRowsTarget, type DraftPrefillGroup, type DraftRowGroup, type GroupedDraftProvider, type TargetView } from "../source";
+import { DRAFT_CHECKPOINT_SCHEMA, type DraftRowCheckpoint } from "../../../contracts/mlx/draft-checkpoint";
+import type { AssistantDrafterModel } from "../../../contracts/mlx/drafter";
+
+/** The assistant owns only its last true target hidden and coverage. Donor
+ * attention stays with target storage; checkpoints contain no duplicate KV. */
+class AssistantRows implements DraftPrefillGroup, DraftRowGroup {
+  readonly prefillMode = "tail-split" as const;
+  readonly tapLayers: readonly number[] = [];
+  #hidden: MlxArray | null = null;
+  #processed: number[] = [];
+  constructor(readonly target: AssistantRowsTarget, readonly drafter: AssistantDrafterModel,
+    readonly namespace: string, checkpoints: readonly (DraftRowCheckpoint | null)[]) { this.append(checkpoints); }
+  get rowCount(): number { return this.#processed.length; }
+  append(checkpoints: readonly (DraftRowCheckpoint | null)[]): void {
+    if (checkpoints.length) applyStateChanges([() => this.prepareAppend(checkpoints)]);
+  }
+  prepareAppend(checkpoints: readonly (DraftRowCheckpoint | null)[]): PreparedStateChange {
+    const processed = [...this.#processed, ...checkpoints.map(checkpoint => checkpoint?.processedTokens ?? 0)];
+    const example = this.#hidden ?? checkpoints.find(checkpoint => checkpoint !== null)?.attachment.tensors[0];
+    const held: MlxArray[] = []; let hidden: MlxArray | null = null;
+    try {
+      for (const checkpoint of checkpoints) if (checkpoint && (checkpoint.attachment.schema !== DRAFT_CHECKPOINT_SCHEMA.assistant ||
+        checkpoint.attachment.metadata.processedTokens !== checkpoint.processedTokens))
+        throw new Error("invalid assistant checkpoint alignment");
+      if (example) {
+        const zero = (rows: number) => { const value = ops.zeros([rows,1,this.target.hiddenSize],example.dtype); held.push(value); return value; };
+        hidden = ops.concatAxis([...(this.rowCount ? [this.#hidden ?? zero(this.rowCount)] : []),
+          ...checkpoints.map(checkpoint => checkpoint?.attachment.tensors[0] ?? zero(1))],0);
+      }
+    } finally { disposeResources(held); }
+    return { commit: () => { const old = this.#hidden; this.#hidden=hidden; hidden=old; this.#processed=processed; },
+      dispose: () => { hidden?.dispose(); hidden=null; } };
+  }
+  filterRows(keep: readonly number[]): void {
+    if (!keep.length) { this.dispose(); return; }
+    using indices = ops.fromInt32([...keep],[keep.length]);
+    const next = this.#hidden && keep.length ? ops.takeAxis(this.#hidden,indices,0) : null;
+    this.#hidden?.dispose(); this.#hidden=next; this.#processed=keep.map(row=>this.#processed[row]!);
+  }
+  prefill(tokens: MlxArray, context?: MlxArray): void {
+    const [B,N] = tokens.shape as [number,number];
+    using tail = context!.slice([0,N-1,0],[B,N,context!.shape[2]!]);
+    const hidden = materializeCopy(tail);
+    this.#hidden?.dispose(); this.#hidden=hidden; this.#processed=this.#processed.map(count=>count+N);
+  }
+  materialize(): void { this.#hidden?.eval(); }
+  capture(row: number): DraftRowCheckpoint {
+    const processedTokens=this.#processed[row]!;
+    const hidden=this.#hidden!.slice([row,0,0],[row+1,1,this.target.hiddenSize]);
+    return { processedTokens, attachment: { schema:DRAFT_CHECKPOINT_SCHEMA.assistant,metadata:{processedTokens},tensors:[hidden] } };
+  }
+  draft(pending: readonly number[], depth: number, _steps: readonly number[]): number[][] {
+    if (!depth) return pending.map(()=>[]);
+    const B=pending.length, owned: MlxArray[]=[];
+    const donors=this.target.readDonors();
+    let ids=ops.fromInt32([...pending],[B,1]), hidden=this.#hidden!;
+    const tokens: MlxArray[]=[];
+    const forward=rowOperationFor(this.drafter.forwardRows, B);
+    try {
+      for(let step=0;step<depth;step++) {
+        using emb=this.target.embed(ids);
+        const positions=donors.positions.map(position=>position+step);
+        const common=positions.every(position=>position===positions[0]);
+        using offsets=common ? null : ops.fromInt32(positions,[B]);
+        const next=forward(emb,hidden,donors,offsets ?? positions[0]!);
+        owned.push(next.nextHidden); tokens.push(next.tokens); hidden=next.nextHidden;
+        ids.dispose(); ids=ops.reshape(next.tokens,[B,1]);
+      }
+      using packed=ops.concatAxis(tokens,0);
+      const values=packed.toIntTokens();
+      return pending.map((_,row)=>Array.from({length:depth},(_,step)=>values[step*B+row]!));
+    } finally { ids.dispose(); donors.dispose(); disposeResources([...owned,...tokens]); }
+  }
+  commit(accepted: readonly number[], context: MlxArray): void {
+    using indices=ops.fromInt32([...accepted],[accepted.length,1,1]);
+    const hidden=ops.takeAlongAxis(context,indices,1);
+    this.#hidden?.dispose(); this.#hidden=hidden;
+    this.#processed=this.#processed.map((count,row)=>count+accepted[row]!+1);
+  }
+  dispose(): void { this.#hidden?.dispose(); this.#hidden=null; this.#processed=[]; }
+}
+
+/** A target without the donor port refuses the pairing before any row state exists. */
+function assistantTarget(target: TargetView): AssistantRowsTarget {
+  if (!target.assistantRows) throw targetLacks("assistantRows");
+  return target.assistantRows;
+}
+
+export function assistantGroups(drafter: AssistantDrafterModel, namespace: string): GroupedDraftProvider {
+  return {
+    checkpointNamespace: () => namespace,
+    openPrefill: ({target,checkpoints}) => new AssistantRows(assistantTarget(target),drafter,namespace,checkpoints),
+    open: ({target,checkpoints}) => new AssistantRows(assistantTarget(target),drafter,namespace,checkpoints),
+  };
+}

package/src/generation/speculative/bindings/autoregressive-rows.ts

@@ -0,0 +1,124 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { MlxPhaseDispatch } from "../../phase-dispatch";
+import type { Cache, BatchableCache, PaddedPrefillCache } from "../../../contracts/mlx/cache";
+import type { DraftRowSampling } from "../source";
+import { toLogprobs } from "../../../sampling/index";
+import { disposeResources } from "../../../runtime/resources";
+import { MlxStateRows } from "../../../state/rows";
+import type { PreparedStateChange } from "../../../contracts/portable/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 phases: MlxPhaseDispatch,
+    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]!);
+  }
+
+  /** One draft-model forward of `sequences` rows of `width` positions: a block of
+   *  known tokens (`prefillTail`), or one position per row (`decode`). */
+  async #forward(ids: MlxArray, sequences: number, width: number, lengths?: readonly number[]): Promise<MlxArray> {
+    const 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.phases.prompt(width === 1 ? "decode" : "tail", sequences, width)(ids, this.state.caches, null);
+    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, B, N + 1, lengths)).dispose();
+    } else (await this.#forward(tokens, B, N)).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);
+        const positions = position === 0 ? width : 1;
+        using hidden = await this.#forward(ids, pending.length, positions, 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)
+          : hidden.slice([0, positions - 1, 0], [pending.length, positions, hidden.shape[2]!]);
+        using logits = this.phases.logits(pending.length)(selected);
+        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/generation/speculative/bindings/binding.ts

@@ -0,0 +1,106 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { Cache } from "../../../contracts/mlx/cache";
+import { type MlxLayerCapture } from "../../../contracts/mlx/graph";
+import type { MlxTokenGraph } from "../../../models/graph";
+import { declaredGraph, tokenPhasesOf } from "../../../models/capabilities";
+import type { MlxForwardWork } from "../../../contracts/mlx/forward-work";
+import { COMPOSITION_DEFAULTS } from "../../../contracts/portable/composition";
+import { MlxPhaseDispatch, promptWorkStep, verifyWorkStep, type MlxPromptStep, type MlxVerifyWindow } from "../../phase-dispatch";
+
+/** A target forward's post-final-norm hidden and, when the drafter taps layers,
+ *  the tapped layers concatenated on the feature axis [B, L, m*H] (null
+ *  otherwise). Both belong to the caller. */
+export interface MlxTargetForward {
+  readonly hidden: MlxArray;
+  readonly ctxML: MlxArray | null;
+}
+
+/** The target operations the speculative group and the fill group call, bound
+ * once from the graph's phases. A replacement graph supplies this entire port,
+ * including any hidden taps. The verifier never needs a MlxTokenGraph or an
+ * artifact-family check. Each forward names its phase; `work`, when given, is
+ * the mixed iteration's hook and runs the same rows inside a shared forward. */
+export interface MlxSpeculativeTargetBinding {
+  /** Whether prompts leave their last token pending for the first verify window
+   *  (the composition's `prefillTailSplit`). */
+  readonly prefillTailSplit: boolean;
+  /** The graph's composed caches for one row. */
+  makeCache(): Cache[];
+  /** One prompt step in the phase its planner named, capturing `tapLayers`. */
+  prefill(ids: MlxArray, caches: Cache[], step: MlxPromptStep, tapLayers?: readonly number[],
+    work?: MlxForwardWork): Promise<MlxTargetForward>;
+  /** One new position per sequence through the graph's decode phase, capturing `tapLayers`: a round without
+   *  drafts, the step plain decode runs. */
+  decode(ids: MlxArray, caches: Cache[], tapLayers?: readonly number[], work?: MlxForwardWork): Promise<MlxTargetForward>;
+  /** A verify window, capturing `tapLayers`. */
+  verify(ids: MlxArray, caches: Cache[], window: MlxVerifyWindow, tapLayers?: readonly number[],
+    work?: MlxForwardWork): Promise<MlxTargetForward>;
+  /** Whether the forwards can capture these hidden layers (a provider's declared taps). */
+  supportsTapLayers(layers: readonly number[]): boolean;
+  /** The head over a hidden of `rows` rows (sequences times positions). */
+  projectLogits(hidden: MlxArray, rows: number): MlxArray;
+}
+
+/** The target graph's phase operations for speculation, bound once. */
+export function bindSpeculativeTargetModel(model: MlxTokenGraph): MlxSpeculativeTargetBinding {
+  const dispatch = new MlxPhaseDispatch(tokenPhasesOf(model));
+  /** The decoder's layer count, the post-final-norm capture's index; read where a tap needs it. */
+  const layerCount = () => model.config.text.numHiddenLayers;
+  return {
+    prefillTailSplit: (model.composition ?? COMPOSITION_DEFAULTS).prefillTailSplit,
+    makeCache: model.makeCache.bind(model),
+    prefill: (ids, caches, step, tapLayers, work) => withCapture(tapLayers, layerCount, capture => {
+      const own = async () => dispatch.prompt(step.phase, step.sequences, step.tokens)(ids, caches, capture);
+      return work ? work(ids, caches, promptWorkStep(step, capture, own)) : own();
+    }),
+    decode: (ids, caches, tapLayers, work) => withCapture(tapLayers, layerCount, capture => {
+      const sequences = ids.shape[0]!;
+      const own = async () => dispatch.decode(sequences)(ids, caches, capture);
+      return work ? work(ids, caches, { phase: "decode", rows: sequences, capture, own }) : own();
+    }),
+    verify: (ids, caches, window, tapLayers, work) => withCapture(tapLayers, layerCount, capture => {
+      const own = async () => dispatch.verify(window.sequences, window.positions)(ids, caches, capture);
+      return work ? work(ids, caches, verifyWorkStep(window, capture, own)) : own();
+    }),
+    // The capture sees layers 0..nLayers-1 and then the post-final-norm output as index nLayers.
+    supportsTapLayers: layers => !layers.length || (declaredGraph(model).graphCapabilities.hiddenLayerTaps && layers.every(layer =>
+      Number.isInteger(layer) && layer >= 0 && layer <= layerCount())),
+    projectLogits: (hidden, rows) => dispatch.logits(rows)(hidden),
+  };
+}
+
+/** Run `forward` with a capture of `tapLayers` (none without taps) and return
+ *  its hidden and the tapped layers [B, L, m*H], concatenated in `tapLayers`
+ *  order on the feature axis. The capture copies each tapped layer as the
+ *  graph hands it over (index `layerCount` is the post-final-norm output), so
+ *  every graph and the mixed iteration capture the same way. */
+async function withCapture(tapLayers: readonly number[] | undefined, layerCount: () => number,
+  forward: (capture: MlxLayerCapture | null) => MlxArray | Promise<MlxArray>): Promise<MlxTargetForward> {
+  if (!tapLayers?.length) return { hidden: await forward(null), ctxML: null };
+  const cap = new Map<number, MlxArray>();
+  const layers = new Set(tapLayers);
+  let hidden: MlxArray | null = null;
+  try {
+    hidden = await forward((layer, h) => { if (layers.has(layer)) cap.set(layer, ops.contiguous(h)); });
+    const final = layerCount();
+    if (layers.has(final) && !cap.has(final)) cap.set(final, ops.contiguous(hidden));
+    const perLayer = tapLayers.map(layer => {
+      const captured = cap.get(layer);
+      if (!captured) throw new Error(`spec tap: layer ${layer} not captured`);
+      return captured;
+    });
+    const ctxML = ops.concatAxis(perLayer, 2);
+    for (const [, a] of cap) a.dispose();
+    cap.clear();
+    const out = { hidden, ctxML };
+    hidden = null; // ownership returned to the caller
+    return out;
+  } finally {
+    // On any throw, free the partial captures and the orphaned hidden.
+    hidden?.dispose();
+    for (const [, a] of cap) a.dispose();
+  }
+}
+
+export { bindLegacyDraftTarget } from "./draft-target";

package/src/generation/speculative/bindings/deepspec-rows.ts

@@ -0,0 +1,33 @@
+import { projectedDraftGroups } from "./projected-draft-rows";
+import type { GroupedDraftProvider, TargetView } from "../source";
+import { targetLacks } from "../source";
+import type { DeepspecDrafterModel } from "../../../contracts/mlx/drafter";
+import { rowOperationFor } from "../../../runtime/row-operations";
+
+/** Existing pairing requirements are checked once when the provider opens. */
+export function bindDeepspecTarget(target: TargetView, drafter: Pick<DeepspecDrafterModel, "cfg">): void {
+  if (!target.hiddenLayerTaps) throw targetLacks("hiddenLayerTaps");
+  const layers = drafter.cfg.num_target_layers;
+  if (target.hiddenLayerTaps.layerCount !== layers) throw new Error(
+    `DeepSpec drafter was trained for a ${layers}-layer target; ` +
+    `this model has ${target.hiddenLayerTaps.layerCount} layers — wrong (target, drafter) pairing`);
+}
+
+export function deepspecGroups(drafter: DeepspecDrafterModel, namespace: string): GroupedDraftProvider {
+  return projectedDraftGroups(namespace, drafter.tapLayers, target => {
+    bindDeepspecTarget(target, drafter);
+    return {
+      namespace, schema: "deepspec-context-v1",
+      layers: drafter.cfg.num_hidden_layers,
+      // The B × L context rows a prefill or a commit appends index the drafter's context tables.
+      project(hidden, positions) {
+        const rows = hidden.shape[0]! * hidden.shape[1]!;
+        using projected = rowOperationFor(drafter.projectContext, rows)(hidden);
+        return rowOperationFor(drafter.projectContextKVRows, rows)(projected, positions);
+      },
+      draft(context, pending, positions, depth) {
+        return drafter.draftRows(context, pending, positions).tokens.map(row => row.slice(0, depth));
+      },
+    };
+  });
+}

package/src/generation/speculative/bindings/draft-target.ts

@@ -0,0 +1,10 @@
+import { type Cache } from "../../../contracts/mlx/cache";
+import type { MlxTokenGraph } from "../../../models/graph";
+import { declaredGraph } from "../../../models/capabilities";
+import type { TargetView } from "../source";
+
+/** The ports a graph declares over its live caches for draft sources; a graph
+ * that declares none is a plain target identified by itself. */
+export function bindLegacyDraftTarget(model: MlxTokenGraph, caches: Cache[]): TargetView {
+  return declaredGraph(model).draftTarget?.(caches) ?? Object.freeze({ identity: model });
+}

package/src/generation/speculative/bindings/glm52-mtp-rows.ts

@@ -0,0 +1,128 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import { materializeCopy } from "@mlx-bun/mlx/materialize";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { PreparedStateChange } from "../../../contracts/portable/resources";
+import type { NativeMtpHead } from "../../../contracts/mlx/drafter";
+import { applyStateChanges,disposeResources } from "../../../runtime/resources";
+import { toLogprobs } from "../../../sampling/index";
+import { MLACache } from "../../../state/glm52-cache";
+import type { DraftRowSampling } from "../source";
+
+/** Native Colibri draft history starts empty, independently of target prompt
+ * length. Accepted draft rows are rebuilt from true target hidden states. */
+export class Glm52MtpRows {
+  #cache: MLACache;
+  #hidden: MlxArray | null = null;
+  #processed: number[] = [];
+  #drafts: number[][] = [];
+  #depth = 0;
+
+  constructor(readonly graph: NativeMtpHead, readonly sampling: DraftRowSampling | null) {
+    this.#cache = this.makeCache();
+  }
+  makeCache(): MLACache {
+    return new MLACache({ ...this.graph.cache, role: "mtp" });
+  }
+  get rowCount(): number { return this.#processed.length; }
+  append(states: readonly (Glm52MtpRowState | null)[]): void {
+    applyStateChanges([() => this.prepareAppend(states)]);
+  }
+  prepareAppend(states: readonly (Glm52MtpRowState | null)[]): PreparedStateChange {
+    if (!states.length) return { commit() {}, dispose() {} };
+    let cache: MLACache | null = this.makeCache(), hidden: MlxArray | null = null;
+    const held: { dispose(): void }[] = [];
+    const processed = [...this.#processed, ...states.map(state => state?.processedTokens ?? 0)];
+    try {
+      cache.mergeRows([...(this.rowCount ? [this.#cache] : []), ...states.map(state => {
+        if (state) return state.cache;
+        const empty = this.makeCache(); held.push(empty); return empty;
+      })]);
+      hidden = ops.concatAxis([...(this.#hidden ? [this.#hidden] : []), ...states.map(state => {
+        if (state) return state.hidden;
+        const zero = ops.zeros([1, 1, this.graph.hiddenSize], Dtype.float32);
+        held.push(zero); return zero;
+      })], 0);
+    } catch (error) { cache.dispose(); hidden?.dispose(); throw error; }
+    finally { disposeResources(held); }
+    return {
+      commit: () => {
+        const priorCache = this.#cache, priorHidden = this.#hidden;
+        this.#cache = cache!; this.#hidden = hidden; this.#processed = processed;
+        cache = priorCache; hidden = priorHidden;
+      },
+      dispose: () => { cache?.dispose(); hidden?.dispose(); cache = null; hidden = null; },
+    };
+  }
+  prefill(tokens: MlxArray, context: MlxArray): void {
+    const [batch, length] = tokens.shape as [number, number];
+    using tail = context.slice([0, length - 1, 0], [batch, length, context.shape[2]!]);
+    const hidden = materializeCopy(tail);
+    this.#hidden?.dispose(); this.#hidden = hidden;
+    this.#processed = this.#processed.map(count => count + length);
+  }
+  materialize(): void {
+    ops.evalAll([...this.#cache.state(), ...(this.#hidden ? [this.#hidden] : [])]);
+  }
+  filterRows(keep: readonly number[]): void {
+    if (!keep.length) { this.dispose(); return; }
+    using indices = ops.fromInt32([...keep], [keep.length]);
+    const hidden = ops.takeAxis(this.#hidden!, indices, 0);
+    try { this.#cache.filterRows(keep); } catch (error) { hidden.dispose(); throw error; }
+    this.#hidden!.dispose(); this.#hidden = hidden;
+    this.#processed = keep.map(row => this.#processed[row]!);
+  }
+  async draft(pending: readonly number[], depth: number, steps: readonly number[]): Promise<number[][]> {
+    const B = pending.length, tokens: MlxArray[] = [];
+    let ids: MlxArray | null = ops.fromInt32([...pending], [B, 1]), chained: MlxArray | null = null;
+    this.#cache.specRoundBegin();
+    try {
+      for (let position = 0; position < depth; position++) {
+        const hidden = await this.graph.decode(ids!, chained ?? this.#hidden!, this.#cache);
+        chained?.dispose(); chained = hidden;
+        using logits = this.graph.project(hidden);
+        using flat = ops.reshape(logits, [B, logits.shape.at(-1)!]);
+        using logprobs = toLogprobs(flat);
+        const next = this.sampling!.sample(logprobs, steps.map(step => step + position));
+        tokens.push(next); ids!.dispose(); ids = ops.reshape(next, [B, 1]);
+      }
+      using packed = depth ? ops.concatAxis(tokens, 0) : null;
+      const read = packed?.toIntTokens() ?? [];
+      this.#drafts = Array.from({ length: B }, (_, row) =>
+        Array.from({ length: depth }, (_, position) => read[position * B + row]!));
+      this.#depth = depth; return this.#drafts;
+    } finally { ids?.dispose(); chained?.dispose(); disposeResources(tokens); }
+  }
+  async commit(accepted: readonly number[], context: MlxArray): Promise<void> {
+    const B = accepted.length, width = Math.max(...accepted);
+    if (this.#depth) this.#cache.specRoundRollback(1);
+    else this.#cache.specRoundCommit();
+    if (width) {
+      using ids = ops.fromInt32(this.#drafts.flatMap(tokens => Array.from({ length: width }, (_, i) => tokens[i] ?? 0)), [B, width]);
+      using hidden = context.slice([0, 0, 0], [B, width, context.shape[2]!]);
+      this.#cache.specRoundBegin();
+      // The window it replays names the phase: one accepted token per row is a decode step.
+      using output = await (width === 1 ? this.graph.decode(ids, hidden, this.#cache) : this.graph.window(ids, hidden, this.#cache));
+      output.eval(); this.#cache.specRoundRollback(accepted);
+    }
+    using positions = ops.fromInt32([...accepted], [B, 1, 1]);
+    const hidden = ops.takeAlongAxis(context, positions, 1);
+    this.#hidden!.dispose(); this.#hidden = hidden;
+    this.#processed = this.#processed.map((count, row) => count + 1 + accepted[row]!);
+    this.#drafts = []; this.#depth = 0;
+  }
+  extractRow(row: number): Glm52MtpRowState {
+    const cache = this.#cache.extractRow(row);
+    try {
+      using hidden = this.#hidden!.slice([row, 0, 0], [row + 1, 1, this.#hidden!.shape[2]!]);
+      return { cache, hidden: materializeCopy(hidden), processedTokens: this.#processed[row]! };
+    } catch (error) { cache.dispose(); throw error; }
+  }
+  dispose(): void {
+    this.#cache.dispose(); this.#hidden?.dispose(); this.#hidden = null;
+    this.#processed = []; this.#drafts = []; this.#depth = 0;
+  }
+}
+
+import { Glm52MtpRowState } from "../../../state/speculative/glm52-mtp-state";
+export { type Glm52MtpRowState } from "../../../state/speculative/glm52-mtp-state";

package/src/generation/speculative/bindings/projected-context-rows.ts

@@ -0,0 +1,78 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+import { BatchedKVCache } from "../../../state/batched-kv";
+import { AttentionMasks } from "../../../state/attention-read";
+import type { Cache } from "../../../contracts/mlx/cache";
+import { MlxStateRows } from "../../../state/rows";
+import { disposeResources } from "../../../runtime/resources";
+
+export interface ProjectedContext { k: MlxArray; v: MlxArray; }
+export interface ContextProjection {
+  project(hidden: MlxArray, positions: number | MlxArray): ProjectedContext[];
+}
+
+/** Incremental projected context reuses ordinary full-attention row storage.
+ * The projection graph owns numerics; this object owns positions and validity.
+ * Context enters through the rows' window read, disposed unattended (the
+ * contract's append without attending); the drafter reads the stored rows as
+ * donor rows (`captureDonorRows`). The layers share one mask memo, so an append
+ * builds its per-row mask once. */
+export class MlxProjectedContextRows {
+  readonly state: MlxStateRows<BatchedKVCache>;
+  constructor(layers: number, readonly graph: ContextProjection) {
+    const masks = new AttentionMasks();
+    this.state = new MlxStateRows(Array.from({ length: layers }, () => new BatchedKVCache(masks)));
+  }
+  get positions(): readonly number[] { return this.state.caches[0]!.rowOffsets; }
+  prepareAppend(rows: readonly (readonly Cache[])[]) { return this.state.prepareAppendMany(rows); }
+  append(hidden: MlxArray, lengths: readonly number[]): void {
+    const width = Math.max(...lengths), B = lengths.length, positions = this.positions;
+    using context = hidden.slice([0, 0, 0], [B, width, hidden.shape[2]!]);
+    using offsets = positions.every(position => position === positions[0]) ? null : ops.fromInt32([...positions], [B]);
+    const projected = this.graph.project(context, offsets ?? positions[0]!);
+    try {
+      for (const [layer, cache] of this.state.caches.entries()) {
+        cache.specRoundBegin();
+        cache.appendWindow(projected[layer]!.k, projected[layer]!.v).dispose();
+        if (lengths.every(length => length === width)) cache.specRoundCommit();
+        else cache.specRoundRollback(lengths);
+      }
+    } finally { disposeResources(projected.flatMap(pair => [pair.k, pair.v])); }
+  }
+  materialize(): void { ops.evalAll(this.state.caches.flatMap(cache => cache.state())); }
+  readAttention() {
+    const views: import("../../../contracts/mlx/cache").KvDonorRows[] = [];
+    try { for (const cache of this.state.caches) views.push(cache.captureDonorRows()); }
+    catch (error) { disposeResources(views.flatMap(view => [view.keys, view.values])); throw error; }
+    const masks = new Map<string, MlxArray | null>();
+    return {
+      attend(layer: number, query: MlxArray, blockKeys: MlxArray, blockValues: MlxArray, scale = 1): MlxArray {
+        const view = views[layer]!, width = view.keys.shape[2]!, block = blockKeys.shape[2]!, B = view.offsets.length;
+        // Membership and coverage advance together in every context layer.
+        const key = `${width}:${block}:${view.keys.dtype}`;
+        if (!masks.has(key)) {
+          let mask: MlxArray | null = null;
+          if (view.starts.some(start => start !== 0) || view.ends.some(end => end !== width)) {
+            using columns = ops.arange(0, width, 1, Dtype.int32);
+            using lower = ops.fromInt32([...view.starts], [B, 1, 1, 1]);
+            using upper = ops.fromInt32([...view.ends], [B, 1, 1, 1]);
+            using before = ops.less(columns, lower), after = ops.greaterEqual(columns, upper);
+            using invalid = ops.logicalOr(before, after);
+            using floats = invalid.astype(view.keys.dtype);
+            using prefix = ops.mulScalar(floats, -1e9);
+            using suffix = ops.zeros([B, 1, 1, block], view.keys.dtype);
+            mask = ops.concatAxis([prefix, suffix], 3);
+          }
+          masks.set(key, mask);
+        }
+        using keys = ops.concatAxis([view.keys, blockKeys], 2);
+        using values = ops.concatAxis([view.values, blockValues], 2);
+        const mask = masks.get(key)!;
+        return ops.sdpa(query, keys, values, scale, mask ? "array" : "", mask);
+      },
+      dispose() { disposeResources([...views.flatMap(view => [view.keys, view.values]), ...[...masks.values()].filter(value => value !== null)]); },
+    };
+  }
+  dispose(): void { this.state.dispose(); }
+}