@mlx-bun/inference 0.0.0-stage → 1.0.0

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

@@ -0,0 +1,95 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { KVCache } from "../../../state/kv";
+import { type Cache } from "../../../contracts/mlx/cache";
+import { MlxProjectedContextRows } from "./projected-context-rows";
+import { captureCacheAttachment, restoreCacheAttachment } from "../../../state/attachment";
+import { applyStateChanges, disposeResources } from "../../../runtime/resources";
+import type { PreparedStateChange } from "../../../contracts/portable/resources";
+import type { DraftPrefillGroup, DraftRowGroup, GroupedDraftProvider, TargetView } from "../source";
+import type { DraftRowCheckpoint } from "../../../contracts/mlx/draft-checkpoint";
+
+export interface ProjectedDraftMethod {
+  namespace: string;
+  schema: string;
+  layers: number;
+  project(hidden: MlxArray, positions: number | MlxArray): { k: MlxArray; v: MlxArray }[];
+  draft(context: ReturnType<MlxProjectedContextRows["readAttention"]>, pending: readonly number[], positions: number | MlxArray, depth: number): number[][];
+  /** Optional device-resident proposals [B, depth], row-major: the round reads
+   *  them back together with the verify samples. */
+  draftDevice?(context: ReturnType<MlxProjectedContextRows["readAttention"]>, pending: readonly number[], positions: number | MlxArray, depth: number): MlxArray;
+}
+
+/** Provider numerics compose with shared context membership and persistence. */
+class ProjectedDraftRows implements DraftPrefillGroup, DraftRowGroup {
+  readonly prefillMode = "tail-split" as const;
+  readonly tapLayers: readonly number[];
+  readonly context: MlxProjectedContextRows;
+  constructor(readonly method: ProjectedDraftMethod, tapLayers: readonly number[],
+    checkpoints: readonly (DraftRowCheckpoint | null)[]) {
+    this.tapLayers = tapLayers;
+    this.context = new MlxProjectedContextRows(method.layers, { project: method.project });
+    try { this.append(checkpoints); } catch (error) { this.dispose(); throw error; }
+  }
+  get namespace(): string { return this.method.namespace; }
+  get rowCount(): number { return this.context.state.rowCount; }
+  prepareAppend(checkpoints: readonly (DraftRowCheckpoint | null)[]): PreparedStateChange {
+    const restored: Cache[][] = [];
+    const empty = () => Array.from({ length: this.method.layers }, () => new KVCache());
+    try {
+      for (const checkpoint of checkpoints) restored.push(checkpoint ? restoreCacheAttachment(checkpoint.attachment, empty) : empty());
+      return this.context.prepareAppend(restored);
+    } finally { disposeResources(restored.flat()); }
+  }
+  append(checkpoints: readonly (DraftRowCheckpoint | null)[]): void {
+    if (checkpoints.length) applyStateChanges([() => this.prepareAppend(checkpoints)]);
+  }
+  prefill(tokens: MlxArray, hidden?: MlxArray): void {
+    this.context.append(hidden!, Array(this.rowCount).fill(tokens.shape[1]!)); this.materialize();
+  }
+  materialize(): void { this.context.materialize(); }
+  filterRows(keep: readonly number[]): void { this.context.state.filterRows(keep); }
+  capture(row: number): DraftRowCheckpoint {
+    const caches = this.context.state.extractRow(row);
+    try {
+      const processedTokens = caches[0]!.offset;
+      return { processedTokens, attachment: captureCacheAttachment(this.method.schema, caches, { processedTokens }) };
+    } finally { disposeResources(caches); }
+  }
+  draft(pending: readonly number[], depth: number): number[][] {
+    if (!depth) return pending.map(() => []);
+    const positions = this.context.positions, donors = this.context.readAttention();
+    using offsets = positions.every(position => position === positions[0]) ? null : ops.fromInt32([...positions], [pending.length]);
+    try { return this.method.draft(donors, pending, offsets ?? positions[0]!, depth); }
+    finally { donors.dispose(); }
+  }
+  draftDevice(pending: readonly number[], depth: number):
+    { tokens: MlxArray; resolve(read: readonly number[]): number[][] } | null {
+    if (!depth || !this.method.draftDevice) return null;
+    const positions = this.context.positions, donors = this.context.readAttention();
+    using offsets = positions.every(position => position === positions[0]) ? null : ops.fromInt32([...positions], [pending.length]);
+    try {
+      const tokens = this.method.draftDevice(donors, pending, offsets ?? positions[0]!, depth);
+      const width = tokens.shape[1]!;
+      return { tokens, resolve: read => pending.map((_, row) => read.slice(row * width, (row + 1) * width)) };
+    } finally { donors.dispose(); }
+  }
+  preview(pending: readonly number[], depth: number): MlxArray | null {
+    return this.draftDevice(pending, depth)?.tokens ?? null;
+  }
+  commit(accepted: readonly number[], hidden: MlxArray): void {
+    this.context.append(hidden, accepted.map(count => count + 1));
+  }
+  dispose(): void { this.context.dispose(); }
+}
+
+/** Projected rows tap the target: the provider declares the one immutable
+ * layer list its rows consume, so bindings check it before any row opens. */
+export function projectedDraftGroups(namespace: string, tapLayers: readonly number[],
+  bind: (target: TargetView) => ProjectedDraftMethod): GroupedDraftProvider {
+  const taps = Object.freeze([...tapLayers]);
+  const open = ({ target, checkpoints }: {
+    target: TargetView; checkpoints: readonly (DraftRowCheckpoint | null)[];
+  }) => new ProjectedDraftRows(bind(target), taps, checkpoints);
+  return { checkpointNamespace: () => namespace, targetTapLayers: () => taps, open, openPrefill: open };
+}

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

@@ -0,0 +1,336 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { materializeCopy } from "@mlx-bun/mlx/materialize";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { PreparedStateChange } from "../../../contracts/portable/resources";
+import { normalizedArgmax } from "../../../kernels/sampling/normalized-argmax";
+import type { RecurrentMtpModule } from "../../../contracts/mlx/drafter";
+import { type MlxRowTable } from "../../../contracts/mlx/graph";
+import { mapOperations, rowOperationFor } from "../../../runtime/row-operations";
+import { applyStateChanges,cleanupFailure,disposeResources } from "../../../runtime/resources";
+import { toLogprobs } from "../../../sampling/index";
+import { BatchedKVCache } from "../../../state/batched-kv";
+import { KVCache } from "../../../state/kv";
+import { DraftVocabularyHead } from "../draft-vocab";
+import type { DraftRowSampling,RecurrentMtpTarget } from "../source";
+
+/** Sampling owns each request's RNG policy. Scores are [B,V] log-probabilities;
+ * the returned owned [B] token array stays on device through the draft chain. */
+export type MtpRowSampling = DraftRowSampling;
+
+/** One draft step's token selection: the module output [B, 1, H] (borrowed) and
+ *  each row's step in, owned full-vocabulary ids [B] out. */
+type MtpSelect = (output: MlxArray, steps: readonly number[]) => MlxArray;
+
+/** The selection the draft group bound, per exact row count of the rows it drafts
+ *  for: the listed vocabulary's projection and each row's coupled draw over it,
+ *  the full head and its fused normalized argmax, or the full head and each
+ *  row's sampler over its log-probabilities. Resolved once when the rows open;
+ *  `owned` receives the listed projection, which the rows dispose. */
+function bindMtpSelection(target: RecurrentMtpTarget, sampling: MtpRowSampling,
+  owned: { vocabulary: DraftVocabularyHead | null }): MlxRowTable<MtpSelect> {
+  const selection = sampling.selection;
+  if (selection.kind === "listed") {
+    // Frequency-ranked draft vocabulary: project onto the listed rows only.
+    // The target still verifies over the full vocabulary.
+    if (!target.vocabularyHead) throw new Error("a listed draft vocabulary needs the target's quantized vocabulary head");
+    const head = owned.vocabulary = new DraftVocabularyHead(target.vocabularyHead, selection.ids);
+    return Object.freeze({
+      default: (output: MlxArray, steps: readonly number[]) => {
+      using subsetLogits = head.project(output);
+      using subsetFlat = ops.reshape(subsetLogits, [steps.length, head.size]);
+      return selection.sample(subsetFlat, head, steps);
+    },
+    });
+  }
+  // Greedy drafting: one fused normalized-argmax over [B,V]; no logsumexp
+  // pass, no per-row top-p/top-k sort, no categorical draw.
+  if (selection.kind === "greedy") return mapOperations(target.logits, head => (output: MlxArray, steps: readonly number[]) => {
+    using logits = head(output);
+    using flat = ops.reshape(logits, [steps.length, logits.shape.at(-1)!]);
+    return normalizedArgmax(flat);
+  }) as MlxRowTable<MtpSelect>;
+  return mapOperations(target.logits, head => (output: MlxArray, steps: readonly number[]) => {
+    using logits = head(output);
+    using flat = ops.reshape(logits, [steps.length, logits.shape.at(-1)!]);
+    using logprobs = toLogprobs(flat);
+    return sampling.sample(logprobs, steps);
+  }) as MlxRowTable<MtpSelect>;
+}
+
+/** Prefill membership drafts nothing, so its rows bind no selection. */
+const prefillOnly: MlxRowTable<MtpSelect> = Object.freeze({
+      default: () => {
+  throw new Error("Qwen MTP prefill rows do not draft");
+},
+    });
+
+/** The same draft graph and state lifecycle for one or several requests.
+ * Admission, target verification, output publication and persistence belong to
+ * the caller. Rows remain fixed within a draft/commit round. */
+export class QwenMtpRows {
+  #cache = new BatchedKVCache();
+  #hidden: MlxArray | null = null;
+  #prefilled: boolean[] = [];
+  #drafts: number[][] = [];
+  #depth = 0;
+  /** The listed vocabulary's row subset of the target's projection, when the
+   *  bound selection drafts over it. */
+  readonly #owned: { vocabulary: DraftVocabularyHead | null } = { vocabulary: null };
+  /** The draft token selection, bound once from the group's sampling. */
+  readonly #select: MlxRowTable<MtpSelect>;
+
+  constructor(
+    readonly target: RecurrentMtpTarget,
+    readonly module: RecurrentMtpModule,
+    sampling: MtpRowSampling | null,
+    states: readonly (MtpRowState | null)[],
+  ) {
+    this.#select = sampling ? bindMtpSelection(target, sampling, this.#owned) : prefillOnly;
+    try { this.append(states); }
+    catch (error) { cleanupFailure(error, () => this.dispose()); }
+  }
+
+  get rowCount(): number { return this.#cache.rowOffsets.length; }
+
+  /** Add committed state at a method boundary. Inputs remain caller-owned. */
+  append(states: readonly (MtpRowState | null)[]): void {
+    if (!states.length) return;
+    applyStateChanges([() => this.prepareAppend(states)]);
+  }
+
+  prepareAppend(states: readonly (MtpRowState | null)[]): PreparedStateChange {
+    if (!states.length) return { commit() {}, dispose() {} };
+    const prefilled = [...this.#prefilled, ...states.map(state => state !== null)];
+    const example = this.#hidden ?? states.find(state => state !== null)?.hidden;
+    let hidden: MlxArray | null = null;
+    let cache: BatchedKVCache | null = new BatchedKVCache();
+    const held: { dispose(): void }[] = [];
+    try {
+      if (example) {
+        const zero = (rows: number) => {
+          const array = ops.zeros([rows, 1, this.target.hiddenSize], example.dtype);
+          held.push(array); return array;
+        };
+        hidden = ops.concatAxis([
+          ...(this.rowCount ? [this.#hidden ?? zero(this.rowCount)] : []),
+          ...states.map(state => state?.hidden ?? zero(1)),
+        ], 0);
+      }
+      cache.mergeRows([...(this.rowCount ? [this.#cache] : []), ...states.map(state => {
+        if (state) return state.cache;
+        const empty = new KVCache(); held.push(empty); return empty;
+      })]);
+    } catch (error) { return cleanupFailure(error, () => disposeResources([cache!, ...(hidden ? [hidden] : [])])); }
+    finally { disposeResources(held); }
+    return {
+      commit: () => {
+        const previousCache = this.#cache, previousHidden = this.#hidden;
+        this.#cache = cache!; this.#hidden = hidden;
+        this.#prefilled = prefilled;
+        cache = previousCache; hidden = previousHidden;
+      },
+      dispose: () => {
+        const held = [cache, hidden].filter(value => value !== null);
+        cache = null; hidden = null;
+        disposeResources(held);
+      },
+    };
+  }
+
+  /** Retire/reorder committed requests without involving their sampler policy. */
+  filterRows(keep: readonly number[]): void {
+    // A drained group keeps its bound selection for the requests that rejoin it.
+    if (!keep.length) { this.#releaseRows(); return; }
+    using indices = ops.fromInt32([...keep], [keep.length]);
+    const hidden = this.#hidden ? ops.takeAxis(this.#hidden, indices, 0) : null;
+    try { this.#cache.filterRows(keep); }
+    catch (error) { hidden?.dispose(); throw error; }
+    this.#hidden?.dispose(); this.#hidden = hidden;
+    this.#prefilled = keep.map(row => this.#prefilled[row]!);
+  }
+
+  /** Build companion KV from borrowed new target tokens/taps at the same B.
+   * The first token bridges to each restored row's preceding true hidden.
+   * Cold rows discard that bridge. All later rows pair token[p+1] with h[p].
+   * Only KV and the final true hidden are evaluated; unused attention/MLP
+   * outputs remain outside the evaluated prefill graph. */
+  prefill(tokens: MlxArray, context: MlxArray): void {
+    const [B, length] = tokens.shape as [number, number];
+    const H = context.shape[2]!;
+    if (this.#prefilled.some(Boolean)) {
+      using first = tokens.slice([0, 0], [B, 1]);
+      using embeds = this.target.embed(first);
+      const mixed = this.#prefilled.some(value => !value);
+      if (mixed) this.#cache.specRoundBegin();
+      rowOperationFor(this.module.decode, B)(embeds, this.#hidden!, this.#cache).dispose();
+      if (mixed) this.#cache.specRoundRollback(this.#prefilled.map(Number));
+    }
+    for (let pos = 0; pos + 1 < length; pos += 2048) {
+      const end = Math.min(pos + 2048, length - 1);
+      using ids = tokens.slice([0, pos + 1], [B, end + 1]);
+      using embeds = this.target.embed(ids);
+      using hidden = context.slice([0, pos, 0], [B, end, H]);
+      // A single position per row appends through the one-position read.
+      (end - pos === 1 ? rowOperationFor(this.module.decode, B) : this.module.window)(embeds, hidden, this.#cache).dispose();
+      ops.evalAll(this.#cache.state());
+    }
+    using tail = context.slice([0, length - 1, 0], [B, length, H]);
+    const hidden = materializeCopy(tail);
+    this.#hidden?.dispose(); this.#hidden = hidden;
+    ops.evalAll([hidden, ...this.#cache.state()]);
+    this.#prefilled.fill(true);
+  }
+
+  materialize(): void {
+    const state = [...this.#cache.state(), ...(this.#hidden ? [this.#hidden] : [])];
+    if (state.length) ops.evalAll(state);
+  }
+
+  /** Align the companion to verified external continuations through the same
+   * true-hidden prefill graph. Different rows retain different prefix lengths. */
+  consume(tokens: MlxArray, context: MlxArray, lengths: readonly number[]): void {
+    this.#cache.specRoundBegin();
+    this.prefill(tokens, context);
+    this.#cache.specRoundRollback([...lengths]);
+    using positions = ops.fromInt32(lengths.map(length => length - 1), [lengths.length, 1, 1]);
+    const hidden = ops.takeAlongAxis(context, positions, 1);
+    this.#hidden!.dispose(); this.#hidden = hidden;
+  }
+
+  /** Nonnegative depth, with every request's first pending token already known.
+   * Build the complete device dependency chain before reading proposal IDs. */
+  draft(pending: readonly number[], depth: number, steps: readonly number[]): number[][] {
+    if (depth === 0) return this.#chain(pending, depth, steps) as number[][];
+    using packed = this.#chain(pending, depth, steps) as MlxArray;
+    return this.#resolve(packed.toIntTokens(), pending.length, depth, true);
+  }
+
+  /** The same chain with the proposal IDs left on the device ([B, depth]), so the
+   *  caller can build its verify graph before anything is read back. `resolve`
+   *  takes the row-major host copy of `tokens`. Null at depth 0. */
+  draftDevice(pending: readonly number[], depth: number, steps: readonly number[]):
+    { tokens: MlxArray; resolve(read: readonly number[]): number[][] } | null {
+    if (depth === 0) return null;
+    const B = pending.length;
+    using packed = this.#chain(pending, depth, steps) as MlxArray; // [depth * B], position-major
+    using grid = ops.reshape(packed, [depth, B]);
+    using rows = ops.transposeAxes(grid, [1, 0]);
+    return { tokens: ops.contiguous(rows), resolve: read => this.#resolve(read, B, depth, false) };
+  }
+
+  #resolve(read: readonly number[], B: number, depth: number, positionMajor: boolean): number[][] {
+    this.#drafts = Array.from({ length: B }, (_, row) =>
+      Array.from({ length: depth }, (_, position) => read[positionMajor ? position * B + row : row * depth + position]!));
+    return this.#drafts;
+  }
+
+  /** The first proposal `draftDevice` would make from `pending` ([B, 1], owned, left on the device), drafted under a
+   *  round of the head's own cache that rolls back, so the rows' state is unchanged (`DraftRowGroup.preview`). One
+   *  position whatever `depth` asks: each further one costs another head step, and one is what a verify round's
+   *  probe drafts. Null at depth 0. */
+  preview(pending: readonly number[], depth: number, steps: readonly number[]): MlxArray | null {
+    if (depth === 0) return null;
+    const B = pending.length;
+    this.#cache.specRoundBegin();
+    try {
+      using packed = this.#proposals(pending, 1, steps); // [B]
+      return ops.reshape(packed, [B, 1]);
+    } finally { this.#cache.specRoundRollback(0); }
+  }
+
+  /** Build the draft chain. Returns the packed device tokens [depth * B]
+   *  (position-major), or the empty proposals at depth 0. */
+  #chain(pending: readonly number[], depth: number, steps: readonly number[]): MlxArray | number[][] {
+    const B = pending.length;
+    this.#cache.specRoundBegin();
+    // A zero-proposal round still processes the pending target token. Keep
+    // its companion KV aligned so a later round can resume speculation.
+    if (depth === 0) {
+      using ids = ops.fromInt32([...pending], [B, 1]);
+      using embeds = this.target.embed(ids);
+      rowOperationFor(this.module.decode, B)(embeds, this.#hidden!, this.#cache).dispose();
+      this.#drafts = Array.from({ length: B }, () => []);
+      this.#depth = 0;
+      this.#hidden!.dispose(); this.#hidden = null;
+      return this.#drafts;
+    }
+    const packed = this.#proposals(pending, depth, steps);
+    this.#depth = depth;
+    this.#hidden!.dispose(); this.#hidden = null;
+    return packed;
+  }
+
+  /** The draft steps from each row's pending token and the true hidden before it, appended to the head's cache:
+   *  the packed device tokens [depth * B] (position-major), owned. Changes no other state. */
+  #proposals(pending: readonly number[], depth: number, steps: readonly number[]): MlxArray {
+    const B = pending.length;
+    const tokens: MlxArray[] = [];
+    let ids: MlxArray | null = ops.fromInt32([...pending], [B, 1]);
+    let chained: MlxArray | null = null;
+    try {
+      const select = rowOperationFor(this.#select, B);
+      for (let position = 0; position < depth; position++) {
+        using embeds = this.target.embed(ids!);
+        const output = rowOperationFor(this.module.decode, B)(embeds, chained ?? this.#hidden!, this.#cache);
+        chained?.dispose(); chained = output;
+        const token = select(output, steps.map(step => step + position));
+        tokens.push(token);
+        ids!.dispose(); ids = ops.reshape(token, [B, 1]);
+      }
+      return ops.concatAxis(tokens, 0);
+    } finally {
+      ids?.dispose(); chained?.dispose();
+      for (const token of tokens) token.dispose();
+    }
+  }
+
+  /** Resolve draft storage using each request's accepted prefix. The verify
+   * context is borrowed [B,depth+1,H], before the target's final norm. */
+  commit(accepted: readonly number[], context: MlxArray): void {
+    const B = accepted.length, depth = this.#depth;
+    const appended = Math.max(1, depth);
+    const kept = accepted.map(count => Math.min(count + 1, appended));
+    if (kept.every(count => count === appended)) this.#cache.specRoundCommit();
+    else this.#cache.specRoundRollback(kept);
+
+    // Full acceptance needs the last accepted token's missing draft row,
+    // using the TRUE target hidden at its preceding position. Other requests
+    // keep their current coverage; their temporary append is never retained.
+    const extra = accepted.map(count => Number(depth > 0 && count === depth));
+    if (extra.some(Boolean)) {
+      using ids = ops.fromInt32(this.#drafts.map(row => row[depth - 1]!), [B, 1]);
+      using embeds = this.target.embed(ids);
+      using hidden = context.slice([0, depth - 1, 0], [B, depth, context.shape[2]!]);
+      this.#cache.specRoundBegin();
+      rowOperationFor(this.module.decode, B)(embeds, hidden, this.#cache).dispose();
+      this.#cache.specRoundRollback(extra);
+    }
+    using positions = ops.fromInt32([...accepted], [B, 1, 1]);
+    this.#hidden = ops.takeAlongAxis(context, positions, 1);
+    this.#drafts = []; this.#depth = 0;
+  }
+
+  /** Only between committed rounds. Returned state is independent of the
+   * group's storage and uses the existing single-request cache representation. */
+  extractRow(row: number): MtpRowState {
+    const cache = this.#cache.extractRow(row);
+    try {
+      using view = this.#hidden!.slice([row, 0, 0], [row + 1, 1, this.#hidden!.shape[2]!]);
+      return { cache, hidden: materializeCopy(view) };
+    } catch (error) { cache.dispose(); throw error; }
+  }
+
+  /** Release every row's state; the bound selection stays. */
+  #releaseRows(): void {
+    this.#cache.dispose(); this.#hidden?.dispose(); this.#hidden = null;
+    this.#drafts = []; this.#depth = 0; this.#prefilled = [];
+  }
+
+  dispose(): void {
+    this.#owned.vocabulary?.dispose(); this.#owned.vocabulary = null;
+    this.#releaseRows();
+  }
+}
+
+import { MtpRowState } from "../../../state/speculative/qwen-mtp-state";
+export { type MtpRowState } from "../../../state/speculative/qwen-mtp-state";

package/src/generation/speculative/bindings/two-model-rows.ts

@@ -0,0 +1,63 @@
+import type { MlxAutoregressiveBinding } from "../../bindings/autoregressive";
+import { MlxAutoregressiveDraftRows, type MlxDraftLayout } from "./autoregressive-rows";
+import { prefillCacheLayout } from "../../../state/layout";
+import { MlxStateRows } from "../../../state/rows";
+import { captureCacheAttachment, restoreCacheAttachment } from "../../../state/attachment";
+import { applyStateChanges, cleanupFailure, disposeResources } from "../../../runtime/resources";
+import type { Cache } from "../../../contracts/mlx/cache";
+import type { DraftPrefillGroup, DraftRowGroup, DraftRowSampling, GroupedDraftProvider } from "../source";
+import type { DraftRowCheckpoint } from "../../../contracts/mlx/draft-checkpoint";
+import type { PreparedStateChange } from "../../../contracts/portable/resources";
+import type { MlxArray } from "@mlx-bun/mlx/array";
+
+export interface StandaloneDraftRows extends DraftPrefillGroup {
+  draft: DraftRowGroup["draft"];
+  commit(accepted: readonly number[], context?: MlxArray): Promise<void>;
+}
+
+/** The standalone model supplies graph/state capabilities once. Target
+ * verification, scheduling and cache tiers use the existing provider ports. */
+export function openStandaloneDraftRows(binding: MlxAutoregressiveBinding, namespace: string,
+  sampling: DraftRowSampling | null, checkpoints: readonly (DraftRowCheckpoint | null)[]): StandaloneDraftRows {
+    const empty = binding.makeCache(), layouts: MlxDraftLayout[] = [];
+    try { for (const cache of empty) layouts.push(prefillCacheLayout(cache) as MlxDraftLayout); }
+    catch (error) { disposeResources(layouts); throw error; }
+    finally { disposeResources(empty); }
+    const state = new MlxStateRows(layouts), rows = new MlxAutoregressiveDraftRows(binding.phases, state, sampling);
+    const prepareAppend = (checkpoints: readonly (DraftRowCheckpoint | null)[]): PreparedStateChange => {
+      const restored: Cache[][] = [], pending: Array<number | null> = [];
+      try {
+        for (const checkpoint of checkpoints) {
+          restored.push(checkpoint ? restoreCacheAttachment(checkpoint.attachment, binding.makeCache) : binding.makeCache());
+          pending.push(checkpoint && checkpoint.attachment.metadata.pendingLast !== false ? Number(checkpoint.attachment.metadata.pendingLast) : null);
+        }
+        return rows.prepareAppend(restored, pending);
+      } finally { disposeResources(restored.flat()); }
+    };
+    const append = (checkpoints: readonly (DraftRowCheckpoint | null)[]) => applyStateChanges([() => prepareAppend(checkpoints)]);
+    try { append(checkpoints); }
+    catch (error) { return cleanupFailure(error, () => rows.dispose()); }
+    return {
+      namespace, prefillMode: "tail-split", tapLayers: [],
+      get rowCount() { return state.rowCount; },
+      append, prepareAppend, prefill: tokens => rows.prefill(tokens), materialize: () => rows.materialize(),
+      filterRows: keep => rows.filterRows(keep), draft: rows.draft.bind(rows), commit: accepted => rows.commit(accepted),
+      capture(row): DraftRowCheckpoint {
+        const caches = state.extractRow(row), pendingLast = rows.pendingLast[row] ?? null;
+        try {
+          const processedTokens = caches[0]!.offset + Number(pendingLast !== null);
+          return { processedTokens, attachment: captureCacheAttachment("standalone-draft-v1", caches,
+            { processedTokens, pendingLast: pendingLast ?? false }) };
+        } finally { disposeResources(caches); }
+      },
+      dispose: () => rows.dispose(),
+    };
+}
+
+export function standaloneDraftGroups(binding: MlxAutoregressiveBinding, namespace: string): GroupedDraftProvider {
+  return {
+    checkpointNamespace: () => namespace,
+    open: options => openStandaloneDraftRows(binding, namespace, options.sampling, options.checkpoints),
+    openPrefill: options => openStandaloneDraftRows(binding, namespace, null, options.checkpoints),
+  };
+}

package/src/generation/speculative/cache-identity.ts

@@ -0,0 +1,16 @@
+import type { GenerateOptions } from "../index";
+
+/** Snapshot the request's configuration when it is bound. The KV scheme is
+ * the composition's, shared by every request, so nothing here copies it. */
+export function captureSpeculativeOptions(options: GenerateOptions): GenerateOptions {
+  return { ...options };
+}
+
+/** The prompt-cache namespace of a speculative method's entries: its method
+ * namespace and adapters. Execution and persistence consume the same captured
+ * settings; the store holding the entries is keyed by the composition's KV
+ * scheme. */
+export function speculativePrefixNamespace(method: string | undefined, adapters: string | string[],
+  _options: GenerateOptions): string {
+  return JSON.stringify({ method, adapters });
+}

package/src/generation/speculative/draft-kind.ts

@@ -0,0 +1,149 @@
+import { isAssistantModelType, isDeepspecDrafterConfig, isDflash2DrafterConfig, isMtpModelType, NATIVE_MTP_DRAFT } from "../../models/drafters";
+import type { DraftProvider } from "./source";
+import { DraftProviderRegistry, type DraftLoadRequest, type DraftProviderKind, type LoadedDraft } from "./draft-registry";
+
+export const DRAFT_KINDS = ["dspark", "deepspec", "dflash2", "assistant", "mtp", "two-model", "ngram"] as const;
+export type DraftKind = typeof DRAFT_KINDS[number];
+
+/** The library's draft providers. Each entry is the provider's own convention
+ *  for recognizing its artifact (files only, never MLX) and for loading it; the
+ *  provider modules load on demand. All providers share ONE serve loop
+ *  (run.ts). "ngram" has no artifact (model-free prompt lookup,
+ *  sources/ngram-source.ts) and is chosen only by name.
+ *  The two-model provider, a full second model, is the catch-all: it accepts
+ *  whatever no other provider recognizes. */
+export function defaultDraftProviders(): DraftProviderRegistry {
+  return new DraftProviderRegistry()
+    .register(dspark).register(deepspec).register(dflash2).register(assistant).register(mtp).register(twoModel).register(ngram)
+    .registerNative(nativeMtp);
+}
+
+/** Detect the draft artifact's kind among the library's providers. */
+export async function detectDraftKind(dir: string): Promise<DraftKind> {
+  return await defaultDraftProviders().detect(dir) as DraftKind;
+}
+
+const artifactDir = (request: DraftLoadRequest): string => {
+  if (request.dir === undefined) throw new Error("this draft provider loads an artifact directory");
+  return request.dir;
+};
+const fixedWidth = (request: DraftLoadRequest, fallback: number) => Math.max(1, request.numDraftTokens ?? fallback);
+/** A block-trained drafter never verifies more positions than it was trained for. */
+const pinned = (request: DraftLoadRequest, gamma: number) => Math.max(1, Math.min(request.numDraftTokens ?? gamma, gamma));
+
+/** Our trained module: dspark.json beside the weights. */
+const dspark: DraftProviderKind = {
+  kind: "dspark", artifact: true,
+  detect: artifact => artifact.has("dspark.json"),
+  async load(request) {
+    const { DflashProvider } = await import("./sources/dflash-source");
+    const provider = await DflashProvider.load(artifactDir(request));
+    return { provider, numDraftTokens: pinned(request, provider.gamma) };
+  },
+};
+
+/** DeepSeek's released DSpark drafters (DeepSpec reference): no dspark.json,
+ *  a plain HF config stamped Gemma4DSparkModel. The width pins to their
+ *  config's block_size (e.g. 7 for dspark_gemma4_12b_block7). */
+const deepspec: DraftProviderKind = {
+  kind: "deepspec", artifact: true,
+  detect: async artifact => isDeepspecDrafterConfig(await artifact.config()),
+  async load(request) {
+    const { DeepspecProvider } = await import("./sources/deepspec-source");
+    const provider = await DeepspecProvider.load(artifactDir(request));
+    return { provider, numDraftTokens: pinned(request, provider.gamma) };
+  },
+};
+
+/** DFlash 2 block drafters (z-lab/incoai): a plain HF config stamped
+ *  DFlash2DraftModel. The width pins to block_size − 1 (7). Projections are
+ *  quantized to 4 bits at load (1.0 GB resident instead of the 3.85 GB BF16
+ *  checkpoint). */
+const dflash2: DraftProviderKind = {
+  kind: "dflash2", artifact: true,
+  detect: async artifact => isDflash2DrafterConfig(await artifact.config()),
+  async load(request) {
+    const { Dflash2Provider } = await import("./sources/dflash2-source");
+    const provider = await Dflash2Provider.load(artifactDir(request), { bits: 4 });
+    return { provider, numDraftTokens: pinned(request, provider.gamma) };
+  },
+};
+
+/** The optiq KV-borrowing drafter (a *_assistant config). */
+const assistant: DraftProviderKind = {
+  kind: "assistant", artifact: true,
+  detect: async artifact => isAssistantModelType(String((await artifact.config())?.model_type ?? "")),
+  async load(request) {
+    const { AssistantProvider } = await import("./sources/assistant-source");
+    return { provider: await AssistantProvider.load(artifactDir(request)), numDraftTokens: fixedWidth(request, 3) };
+  },
+};
+
+/** Native MTP heads split from a qwen3_5-family release
+ *  (mlx-community/Qwen3.8-27B-MTP-*). The target's recurrent DeltaNet caches
+ *  roll back via the serve loop's spec-round snapshot/replay contract
+ *  (SSMCache.specRound*). A target artifact may bundle the head at `mtp/`. */
+const mtp: DraftProviderKind = {
+  kind: "mtp", artifact: true, bundledCompanion: "mtp",
+  detect: async artifact => isMtpModelType(String((await artifact.config())?.model_type ?? "")),
+  async load(request) {
+    const dir = artifactDir(request);
+    const { QwenMtpProvider } = await import("./sources/qwen-mtp-source");
+    const provider = await QwenMtpProvider.load(dir);
+    // Default the round width to the head's trained block (block_size 3 → 2
+    // recursive drafts + the pending row per round; the head was trained
+    // multi-step, so an explicit larger width is allowed but acceptance
+    // decides whether it pays). A graph may declare the adaptive width it
+    // measured for this kind (`adaptiveDraftWidths`).
+    const block = (await Bun.file(`${dir}/config.json`).json() as { block_size?: number }).block_size;
+    const numDraftTokens = request.numDraftTokens === undefined && typeof block === "number"
+      ? Math.max(1, block - 1) : fixedWidth(request, 3);
+    return { provider, numDraftTokens };
+  },
+};
+
+const PROBE = "The 3 quick brown foxes jumped över the lazy dog?! 🦊";
+
+/** A full second model (mlx-lm parity). It ships its own tokenizer, and exact
+ *  token-match acceptance is meaningful only when both models tokenize
+ *  identically: a probe that ENCODES differently means different families —
+ *  refuse instead of silently accepting ~0% of drafts. */
+const twoModel: DraftProviderKind = {
+  kind: "two-model", artifact: true, fallback: true,
+  detect: async () => true,
+  async load(request) {
+    const dir = artifactDir(request);
+    const { TwoModelProvider } = await import("./sources/two-model");
+    const { loadTokenizer } = await import("../../input");
+    const provider = await TwoModelProvider.load(dir, request.target.vocabSize);
+    try {
+      const draftTokenizer = await loadTokenizer(dir);
+      if (JSON.stringify(request.target.tokenizer.encode(PROBE)) !== JSON.stringify(draftTokenizer.encode(PROBE)))
+        throw new Error(
+          `the draft model's tokenizer differs from the target's (probe string encodes ` +
+            `differently) — speculation needs the same tokenizer family`);
+    } catch (error) { provider.dispose(); throw error; }
+    return { provider, numDraftTokens: fixedWidth(request, 3) };
+  },
+};
+
+/** Model-free prompt lookup: no artifact, weightsBytes 0, open() never throws.
+ *  Default γ=10 per the reference implementation — drafting is free, so wide
+ *  blocks cost only verify-window width when wrong. */
+const ngram: DraftProviderKind = {
+  kind: "ngram", artifact: false,
+  async load(request): Promise<LoadedDraft> {
+    const { NgramProvider } = await import("./sources/ngram-source");
+    return { provider: new NgramProvider({ max: request.ngram?.max, min: request.ngram?.min }), numDraftTokens: fixedWidth(request, 10) };
+  },
+};
+
+/** The checkpoint-native MTP row of a graph that declares it: the graph
+ *  declares the head and its width; this binds the provider that drives it. */
+const nativeMtp = {
+  kind: NATIVE_MTP_DRAFT,
+  async create(graph: object): Promise<DraftProvider> {
+    const { Glm52NativeMtpProvider } = await import("./sources/glm52-mtp-source");
+    return new Glm52NativeMtpProvider(graph);
+  },
+};