mlx-bun 0.3.0 → 0.4.0

package/src/backends/mlx/speculative-group.ts

@@ -0,0 +1,359 @@
+import {captureSpeculativeOptions, speculativePrefixNamespace} from "../../spec/cache-identity";
+import type { RuntimeModel } from "../../model/factory";
+import type { GenerateOptions } from "../../generate";
+import type { DraftProvider, DraftRowGroup, DraftPrefillGroup, DraftRowCheckpoint } from "../../spec/source";
+import { makeSampler, makeStepSampler, readStepExtras, type NumberStepSampler, type Sampler } from "../../sampler";
+import { MlxArray } from "../../mlx/array";
+import * as ops from "../../mlx/ops";
+import { clearCache } from "../../mlx/ffi";
+import { applyStateChanges, cleanupFailure, disposeResources } from "../../engine/resources";
+import { cloneKvCaches } from "../../kv-store";
+import type { Cache } from "../../model/gemma4-base";
+import { targetCacheLayout } from "./cache-layout";
+import { MlxStateRows } from "./state-rows";
+import { bindLegacySpeculativeModel, type MlxSpeculativeBinding } from "./speculative";
+import { bindLegacyDraftTarget } from "./draft-target";
+import { bindRowCacheRollback } from "./rollback";
+import { createKvMaintenance } from "./kv-maintenance";
+import { disposeAttachments, type CheckpointAttachment } from "./checkpoint-state";
+import { advanceSpeculativeOutputs } from "./speculative-round";
+import { MlxPrefillRows, type MlxPrefillState } from "./prefill-rows";
+import type { MlxGroupMethodHost, MlxGroupMethodRequest, MlxGroupPreparation, MlxGroupedMethod, Row } from "./batch-group";
+
+interface RequestState {
+  sampling: NumberStepSampler;
+  draftSampling: Sampler;
+  retain?: () => void;
+  namespace: string;
+  prefixLength: number;
+  /** Generated inputs retained by committed target/draft rounds. The pending
+   * correction is excluded until a later round actually processes it. */
+  processed?: number[];
+}
+
+/** Binding owns graph/layout selection. The executor receives only the method
+ * key and lifecycle; sampling, checkpoints and numerical state remain ports. */
+export function bindSpeculativeGroupRequests(model: RuntimeModel, provider: DraftProvider, depth: number) {
+  const binding = bindLegacySpeculativeModel(model, provider);
+  return (input: GenerateOptions): MlxGroupMethodRequest => {
+    const options = captureSpeculativeOptions(input);
+    return {
+      key: JSON.stringify(["speculative", provider.id, depth, options.kvBits ?? null,
+        options.kvGroupSize ?? 64, options.quantizedKvStart ?? null, options.turboQuant ?? null,
+        ...(options.kvConfig?.length ? [options.kvConfig] : [])]),
+      data: options,
+      open: host => new SpeculativeGroup(host, model, provider, binding, depth),
+    };
+  };
+}
+
+class SpeculativeGroup implements MlxGroupedMethod {
+  #target: MlxStateRows<ReturnType<typeof targetCacheLayout>> | null = null;
+  #draft: DraftRowGroup | null = null;
+  readonly #requests = new Map<Row, RequestState>();
+  #steps = 0;
+
+  constructor(readonly host: MlxGroupMethodHost, readonly model: RuntimeModel,
+    readonly provider: DraftProvider, readonly binding: MlxSpeculativeBinding, readonly depth: number) {}
+
+  prepare(first: Row): MlxGroupPreparation {
+    const method = this;
+    const options = first.req.method!.data as GenerateOptions;
+    const maintain = createKvMaintenance(options);
+    interface Preparation extends MlxPrefillState {
+      request: RequestState;
+      end: number;
+      chunkSize: number;
+      pendingPrompt: boolean;
+      boundary: number;
+      transferred: boolean;
+      closed: boolean;
+    }
+    interface Ready {
+      row: Row;
+      request: RequestState;
+      caches: Cache[];
+      draft: DraftRowCheckpoint;
+      retain?: () => void;
+      transferred: boolean;
+    }
+    let prefix: DraftPrefillGroup | undefined;
+    let context: MlxArray | null = null;
+    const ready: Ready[] = [];
+    const closeReady = (state: Ready) => {
+      const caches = state.caches.splice(0), attachments = [state.draft.attachment];
+      const release = state.retain; state.retain = undefined;
+      disposeResources([...caches, { dispose: () => disposeAttachments(attachments) },
+        { dispose: () => release?.() }, ...(!state.transferred ? [state.request.sampling] : [])]);
+    };
+    const capture = (state: Preparation, snapshot: () => Cache[], index: number) => {
+      let caches: Cache[] = [], attachments: CheckpointAttachment[] = [];
+      try {
+        caches = snapshot(); attachments.push(prefix!.capture(index).attachment);
+        method.host.promptCache!.put(state.row.req.promptIds.slice(0, state.pos), caches,
+          state.request.namespace, undefined, attachments);
+        caches = []; attachments = [];
+      } finally { disposeResources([...caches, { dispose: () => disposeAttachments(attachments) }]); }
+    };
+    const target = new MlxPrefillRows<Preparation>({
+      maintain,
+      open(row) {
+        const options = row.req.method!.data as GenerateOptions;
+        row.spec = { drafted: 0, accepted: 0, targetCalls: 0, rejected: 0, rounds: 0,
+          acceptanceLengths: [], tokensPerForward: 0, forwardsSaved: 0, draftedByPos: [], acceptedByPos: [] };
+        const sampler = makeSampler(options);
+        const request: RequestState = { draftSampling: sampler, namespace: "", prefixLength: 0,
+          processed: method.host.promptCache ? [] : undefined,
+          sampling: makeStepSampler(options, { tokenRepresentation: "number", grammarWait: "before-sample",
+            historyUpdate: "after-sample", initialHistory: row.req.promptIds, acceptGrammar: true,
+            eosTokenIds: row.req.eosTokenIds, captureSelectedLogprob: options.logprobs === true,
+            captureTopLogprobs: options.topLogprobs }) };
+        let caches: Cache[] = [], retained: (() => void) | undefined;
+        let appended = false;
+        const previous = prefix?.rowCount ?? 0;
+        try {
+          caches = method.binding.makeCache();
+          prefix ??= method.provider.grouped!.openPrefill({ target: bindLegacyDraftTarget(method.model, caches), checkpoints: [] });
+          const prompt = row.req.promptIds;
+          const pendingPrompt = prefix.prefillMode !== "full" && method.binding.prefillTailSplit && prompt.length > 1;
+          request.prefixLength = prompt.length - Number(pendingPrompt);
+          request.namespace = speculativePrefixNamespace(prefix.namespace, row.cacheNamespace ?? "", options);
+          const hit = method.host.promptCache?.take(prompt, request.namespace);
+          if (hit) {
+            try {
+              disposeResources(caches); caches = hit.caches.splice(0);
+              retained = hit.retain; hit.retain = undefined;
+              prefix.append([{ processedTokens: hit.tokens.length, attachment: hit.attachments![0]! }]);
+              appended = true;
+              // Target membership later releases backing leases after its own
+              // copies resolve. Settle companion copies under the same lease.
+              if (retained) prefix.materialize();
+              row.cachedTokens = hit.tokens.length;
+            } finally { disposeResources([...hit.caches, { dispose: () => disposeAttachments(hit.attachments) },
+              { dispose: () => hit.retain?.() }]); }
+          } else { prefix.append([null]); appended = true; }
+          return { row, request, solo: caches, retain: retained, pos: row.cachedTokens,
+            end: request.prefixLength, pendingPrompt, transferred: false, closed: false,
+            chunkSize: options.prefillChunkSize ?? method.host.prefillChunkSize,
+            boundary: method.host.promptCache ? Math.min(row.req.snapshotAt ?? prompt.length, prompt.length - 1) : -1 };
+        } catch (error) {
+          return cleanupFailure(error, () => disposeResources([
+            { dispose: () => { if (appended) prefix!.filterRows(Array.from({ length: previous }, (_, i) => i)); } },
+            ...caches, request.sampling, { dispose: () => retained?.() },
+          ]));
+        }
+      },
+      plan(state) {
+        const end = Math.min(state.pos + state.chunkSize, state.end, state.boundary > state.pos ? state.boundary : Infinity);
+        return { start: state.pos, end, kind: end === state.end ? "final" : "drain",
+          snapshot: end === state.boundary && end > state.row.cachedTokens, batchYield: true, maintain: true };
+      },
+      async forward(ids, caches) {
+        const taps = prefix!.tapLayers;
+        const output = await method.binding.forward(ids, caches, taps.length ? [...taps] : undefined);
+        context = output.ctxML;
+        return output.hidden;
+      },
+      afterForward(ids, _caches, _states, hidden) {
+        // Tapped providers consume their captured layers; assistant drafting
+        // consumes the true target hidden. Other providers ignore context.
+        const captured = context; context = null;
+        try {
+          const pending = prefix!.prefill(ids, captured ?? hidden);
+          if (pending) return pending.finally(() => captured?.dispose());
+        } catch (error) { captured?.dispose(); throw error; }
+        captured?.dispose();
+      },
+      project(hidden, _caches, completed) {
+        if (completed.every(state => state.pendingPrompt)) return null;
+        using last = hidden.slice([0, hidden.shape[1]! - 1, 0], [hidden.shape[0]!, hidden.shape[1]!, hidden.shape[2]!]);
+        return method.binding.projectLogits(last);
+      },
+      checkpoint: capture,
+      async complete(state, logits, index) {
+        const row = state.row;
+        row.spec!.targetCalls++;
+        if (state.pendingPrompt) { row.current = row.req.promptIds[state.end]!; row.generated = row.sampled = 0; }
+        else {
+          using scores = ops.reshape(logits!, [1, logits!.shape.at(-1)!]);
+          const sampled = await state.request.sampling.sample(scores, 0);
+          row.current = sampled.token; row.generated = row.sampled = 1;
+          if (row.req.eosTokenIds.includes(row.current)) {
+            if (method.host.promptCache) capture(state, () => cloneKvCaches(state.solo), index);
+            disposeResources(state.solo.splice(0)); state.retain?.(); state.retain = undefined;
+            method.host.finish(row, "stop"); return;
+          }
+          const more = await method.host.publish(row, row.current, readStepExtras(sampled.extras));
+          if (more === false || row.req.grammar?.isTerminated || row.generated >= row.req.maxTokens) {
+            if (method.host.promptCache && !row.req.signal?.aborted) capture(state, () => cloneKvCaches(state.solo), index);
+            disposeResources(state.solo.splice(0)); state.retain?.(); state.retain = undefined;
+            method.host.finish(row, more === false || row.req.grammar?.isTerminated ? "stop" : "length"); return;
+          }
+        }
+        ready.push({ row, request: state.request, caches: state.solo, retain: state.retain,
+          draft: prefix!.capture(index), transferred: false });
+        state.transferred = true;
+      },
+      reject(row, error) { row.reject(error); },
+      close(state) {
+        if (state.closed) return;
+        state.closed = true;
+        if (!state.transferred) state.request.sampling.dispose();
+      },
+      filterRows(keep) { prefix!.filterRows(keep); },
+      dispose() {
+        const held = [prefix, context].filter(value => value != null);
+        prefix = undefined; context = null; disposeResources(held);
+      },
+    });
+    const preparation: MlxGroupPreparation = {
+      get rows() { return [...target.rows, ...ready.map(state => state.row)]; },
+      get canAdmit() { return target.canAdmit; },
+      get tokenWeight() { return target.tokenWeight + ready.reduce((sum, state) => sum + state.row.promptTokens - state.row.cachedTokens, 0); },
+      admit: row => target.admit(row),
+      advance: async () => {
+        // Publication returns to scheduling before prepared state joins active
+        // decode. This preserves the first-output flush boundary at every B.
+        if (ready.length) {
+          while (ready.length) {
+            const state = ready[0]!;
+            try {
+              if (state.row.req.signal?.aborted) {
+                state.row.reject(state.row.req.signal.reason);
+                continue;
+              }
+              maintain.prepareBatch?.(state.caches);
+              this.#target ??= new MlxStateRows(state.caches.map(targetCacheLayout));
+              this.#draft ??= this.provider.grouped!.open({ target: bindLegacyDraftTarget(this.model, this.#target.caches),
+                checkpoints: [], sampling: { sample: (lp, steps) => this.#sampleDraftRows(lp, steps) } });
+              applyStateChanges([() => this.#target!.prepareAppend(state.caches),
+                () => this.#draft!.prepareAppend([state.draft]), () => ({ commit: () => {
+                  state.request.retain = state.retain; state.retain = undefined;
+                  this.#requests.set(state.row, state.request); state.transferred = true;
+                  this.host.join(state.row);
+                }, dispose() {} })]);
+            } finally { ready.shift(); closeReady(state); }
+          }
+          return target.rows.length === 0;
+        }
+        const done = await target.advance();
+        return done && ready.length === 0;
+      },
+      dispose() {
+        const held = ready.splice(0);
+        disposeResources([target, ...held.map(state => ({ dispose: () => closeReady(state) }))]);
+      },
+    };
+    try { target.admit(first); return preparation; }
+    catch (error) { return cleanupFailure(error, () => preparation.dispose()); }
+  }
+
+  #sampleDraftRows(logprobs: MlxArray, steps: readonly number[]): MlxArray {
+    const tokens: MlxArray[] = [];
+    try {
+      for (let row = 0; row < this.host.rows.length; row++) {
+        using scores = logprobs.slice([row, 0], [row + 1, logprobs.shape[1]!]);
+        tokens.push(this.#requests.get(this.host.rows[row]!)!.draftSampling(scores, steps[row]!));
+      }
+      return ops.concatAxis(tokens, 0);
+    } finally { disposeResources(tokens); }
+  }
+
+  async advance(): Promise<void> {
+    const active = [...this.host.rows];
+    const live = active.flatMap((row, index) => row.req.signal?.aborted ? [] : [index]);
+    if (live.length !== active.length) {
+      for (const row of active) if (row.req.signal?.aborted) row.reject(row.req.signal.reason);
+      this.host.filterRows(live);
+    }
+    const rows = [...this.host.rows];
+    if (!rows.length) return;
+    const depth = Math.min(this.depth, Math.max(...rows.map(row => row.req.maxTokens - row.generated)));
+    const halted = new Set<Row>();
+    const pin = this.binding.pinVerify?.();
+    try {
+      const completed = await advanceSpeculativeOutputs(rows.map(row => ({ pending: row.current,
+        step: row.sampled, remaining: row.req.maxTokens - row.generated, eosTokenIds: row.req.eosTokenIds,
+        sampling: this.#requests.get(row)!.sampling, grammarDone: () => row.req.grammar?.isTerminated === true,
+        output: { commit: async (ids, metadata) => {
+          row.generated++;
+          const more = await this.host.publish(row, ids[0]!, metadata?.[0]);
+          if (more === false) { halted.add(row); return false; }
+        } },
+      })), depth, this.#draft!, {
+        transaction: bindRowCacheRollback(this.#target!.caches, rows.length),
+        forward: async ids => {
+          const taps = this.#draft!.tapLayers;
+          const result = await this.binding.forward(ids, this.#target!.caches, taps.length ? [...taps] : undefined);
+          const context = result.ctxML ?? result.hidden;
+          try { return { logits: this.binding.projectLogits(result.hidden), context }; }
+          catch (error) { context.dispose(); throw error; }
+          finally { if (result.ctxML) result.hidden.dispose(); }
+        },
+      });
+      const keep: number[] = [];
+      for (let index = 0; index < rows.length; index++) {
+        const row = rows[index]!, output = completed.outputs[index]!, round = completed.rounds[index]!;
+        const request = this.#requests.get(row)!;
+        if (output.kind !== "failed" && output.kind !== "cancelled")
+          request.processed?.push(row.current, ...round.drafts.slice(0, output.accepted));
+        const stats = row.spec!;
+        stats.drafted += round.drafts.length;
+        stats.accepted += round.acceptance.accepted;
+        stats.rejected = stats.drafted - stats.accepted;
+        stats.rounds!++; stats.targetCalls++;
+        stats.acceptanceLengths!.push(round.acceptance.accepted);
+        for (let p = 0; p < round.drafts.length; p++) stats.draftedByPos![p] = (stats.draftedByPos![p] ?? 0) + 1;
+        for (let p = 0; p < round.acceptance.accepted; p++) stats.acceptedByPos![p] = (stats.acceptedByPos![p] ?? 0) + 1;
+        row.sampled += round.acceptance.emitted.length + Number(round.acceptance.sawEos);
+        // A consumer may stop inside a verified burst before its terminal EOS.
+        if (output.kind === "stop" && !halted.has(row) && round.acceptance.sawEos &&
+          output.generated === round.acceptance.emitted.length && row.generated < row.req.maxTokens)
+          row.generated++;
+        stats.tokensPerForward = row.generated / stats.rounds!;
+        stats.forwardsSaved = Math.max(0, row.generated - 1 - stats.rounds!);
+        if (output.kind === "continue") { row.current = output.pending; keep.push(index); }
+        else if (output.kind === "failed") row.reject(output.error);
+        else if (output.kind === "cancelled") row.reject(row.req.signal?.reason ?? new Error(output.reason));
+        else {
+          if (!row.req.signal?.aborted) this.#publishCheckpoint(row, index, request);
+          this.host.finish(row, output.kind);
+        }
+      }
+      if (keep.length !== rows.length) this.host.filterRows(keep);
+      if (++this.#steps % 256 === 0) clearCache();
+    } finally { pin?.close(); }
+  }
+
+  /** The method aligns token coverage and captures its state. The cache owns
+   * RAM retention and eventual SSD persistence; no storage work is awaited. */
+  #publishCheckpoint(row: Row, index: number, request: RequestState): void {
+    if (!this.host.promptCache || !request.processed) return;
+    let target: Cache[] = [], attachments: CheckpointAttachment[] = [];
+    try {
+      target = this.#target!.extractRow(index);
+      attachments.push(this.#draft!.capture(index).attachment);
+      this.host.promptCache.put([...row.req.promptIds.slice(0, request.prefixLength), ...request.processed], target,
+        request.namespace, undefined, attachments);
+      target = []; attachments = [];
+    } finally { disposeResources([...target, { dispose: () => disposeAttachments(attachments) }]); }
+  }
+
+  filterRows(keep: readonly number[], _discard: boolean): void {
+    this.#target?.filterRows(keep); this.#draft?.filterRows(keep);
+    const retained = new Set(keep.map(index => this.host.rows[index]!));
+    const releases: { dispose(): void }[] = [];
+    for (const [row, request] of this.#requests) if (!retained.has(row)) {
+      this.#requests.delete(row);
+      releases.push(request.sampling, { dispose: () => request.retain?.() });
+    }
+    disposeResources(releases);
+  }
+
+  dispose(): void {
+    const resources = [...(this.#target ? [this.#target] : []), ...(this.#draft ? [this.#draft] : []),
+      ...[...this.#requests.values()].flatMap(request => [request.sampling, { dispose: () => request.retain?.() }])];
+    this.#target = null; this.#draft = null; this.#requests.clear();
+    disposeResources(resources);
+  }
+}

package/src/backends/mlx/speculative-round.ts

@@ -0,0 +1,161 @@
+import type { MlxArray } from "../../mlx/array";
+import * as ops from "../../mlx/ops";
+import type { Cancellation, GenerationOutput, TokenLogprobs } from "../../contracts/generation";
+import type { DisposableResource } from "../../contracts/resources";
+import { cleanupFailure, disposeResources } from "../../engine/resources";
+import type { SpeculativeTransaction } from "../../inference/rollback";
+import { DraftAcceptance } from "../../inference/draft-acceptance";
+import { deliverDraftOutputs, type DraftOutputResult } from "../../inference/draft-output";
+import { readStepExtras, type NumberStepSampler } from "../../sampler";
+
+/** Draft state owns proposal generation and its committed hidden/KV state.
+ * It borrows verification context and knows nothing about target storage. */
+export interface MlxDraftRows<Context = MlxArray> {
+  /** Each row returns 0..depth proposals. The method verifies the returned
+   * lengths; an empty row still consumes its pending token. */
+  draft(pending: readonly number[], depth: number, steps: readonly number[]): number[][] | Promise<number[][]>;
+  commit(accepted: readonly number[], context: Context): void | Promise<void>;
+}
+
+/** A bound graph transfers both tensors on success and releases partial
+ * allocations on failure. State retention is a separate injected interface. */
+export interface MlxRowVerification<Context extends DisposableResource = MlxArray> {
+  readonly transaction: SpeculativeTransaction<readonly number[]>;
+  forward(ids: MlxArray): Promise<{ logits: MlxArray; context: Context }>;
+}
+
+export interface MlxSpeculativeRow {
+  readonly pending: number;
+  readonly step: number;
+  readonly remaining: number;
+  readonly eosTokenIds: readonly number[];
+  readonly sampling: NumberStepSampler;
+  readonly grammarDone?: () => boolean;
+}
+
+export interface MlxSpeculativeRowResult {
+  readonly drafts: readonly number[];
+  readonly acceptance: DraftAcceptance;
+  /** One entry per content token; stopping EOS is excluded. */
+  readonly logprobs: readonly (TokenLogprobs | undefined)[];
+}
+
+export interface MlxSpeculativeOutputRow extends MlxSpeculativeRow {
+  readonly output: Pick<GenerationOutput, "commit">;
+  readonly cancellation?: Cancellation;
+}
+
+/** Execute one method step through output delivery and aligned state commit.
+ * The caller retires terminal rows and publishes eligible completed state
+ * before changing membership. Consumer failures are returned per request;
+ * a graph/state failure throws and requires discarding the whole group. */
+export async function advanceSpeculativeOutputs<Context extends DisposableResource = MlxArray>(
+  rows: readonly MlxSpeculativeOutputRow[], depth: number,
+  draft: MlxDraftRows<NoInfer<Context>>, target: MlxRowVerification<Context>,
+): Promise<{ rounds: readonly MlxSpeculativeRowResult[]; outputs: readonly DraftOutputResult[] }> {
+  const round = await prepareSpeculativeRows(rows, depth, draft, target);
+  let outputs: DraftOutputResult[];
+  try {
+    outputs = await deliverDraftOutputs(round.results.map((result, row) => ({
+      acceptance: result.acceptance, logprobs: result.logprobs,
+      output: rows[row]!.output, cancellation: rows[row]!.cancellation,
+    })));
+    await round.commit(outputs.map(output => output.accepted));
+  } catch (error) { return cleanupFailure(error, () => round.dispose()); }
+  round.dispose();
+  return { rounds: round.results, outputs };
+}
+
+/** Retention stays unresolved while the method delivers output. A terminal
+ * row may keep fewer accepted inputs for an output-aligned checkpoint;
+ * continuing rows retain their complete accepted prefix. Commit exactly once,
+ * then dispose; disposing an uncommitted round requires discarding its state. */
+export interface MlxPreparedSpeculativeRound extends DisposableResource {
+  readonly results: readonly MlxSpeculativeRowResult[];
+  commit(keepAccepted?: readonly number[]): Promise<void>;
+}
+
+/** One committed method step at B=1 or B>1. The caller supplies live rows with
+ * a positive token budget and a common nonnegative draft depth. Admission,
+ * cancellation boundaries, output delivery and cache publication stay outside
+ * this step. On failure the caller discards the participating method state. */
+export async function advanceSpeculativeRows<Context extends DisposableResource = MlxArray>(
+  rows: readonly MlxSpeculativeRow[], depth: number,
+  draft: MlxDraftRows<NoInfer<Context>>, target: MlxRowVerification<Context>,
+): Promise<readonly MlxSpeculativeRowResult[]> {
+  const round = await prepareSpeculativeRows(rows, depth, draft, target);
+  try { await round.commit(); }
+  catch (error) { return cleanupFailure(error, () => round.dispose()); }
+  round.dispose();
+  return round.results;
+}
+
+/** Draft, verify and sample without deciding which output-aligned state to
+ * retain. Graph-specific context is opaque to this operation. */
+export async function prepareSpeculativeRows<Context extends DisposableResource = MlxArray>(
+  rows: readonly MlxSpeculativeRow[], depth: number,
+  draft: MlxDraftRows<NoInfer<Context>>, target: MlxRowVerification<Context>,
+): Promise<MlxPreparedSpeculativeRound> {
+  const proposals = await draft.draft(rows.map(row => row.pending), depth, rows.map(row => row.step));
+  const B = rows.length;
+  const verifyDepth = proposals.reduce((max, tokens) => Math.max(max, tokens.length), 0);
+  using ids = ops.fromInt32(proposals.flatMap((tokens, row) => {
+    const window = [rows[row]!.pending, ...tokens];
+    // Causal right padding only fills the rectangular graph. Sampling stops
+    // at this row's own bonus position; commit discards every padded input.
+    while (window.length <= verifyDepth) window.push(rows[row]!.pending);
+    return window;
+  }), [B, verifyDepth + 1]);
+  target.transaction.begin(verifyDepth);
+  const verification = await target.forward(ids);
+  const dispose = () => disposeResources([verification.context, verification.logits]);
+  try {
+    const results = await sampleSpeculativeRows(rows, proposals, verification.logits);
+    return {
+      results,
+      async commit(accepted = results.map(result => result.acceptance.accepted)) {
+        target.transaction.resolve(accepted);
+        await draft.commit(accepted, verification.context);
+      },
+      dispose,
+    };
+  } catch (error) { return cleanupFailure(error, dispose); }
+}
+
+/** Sample a borrowed verify window. Each request owns its processor history,
+ * grammar and RNG stream; completing one walk never samples its unused suffix.
+ * The existing request adapter and grouped method use this same operation. */
+export async function sampleSpeculativeRows(
+  rows: readonly MlxSpeculativeRow[], proposals: readonly (readonly number[])[], logits: MlxArray,
+): Promise<MlxSpeculativeRowResult[]> {
+  const V = logits.shape[2]!, width = logits.shape[1]!;
+  const independent = rows[0]?.sampling.independent;
+  let tokens: number[] | undefined;
+  if (independent && rows.every(row => row.sampling.independent === independent)) {
+    using flat = ops.reshape(logits, [rows.length * width, V]);
+    using sampled = independent.sample(flat);
+    tokens = sampled.toIntTokens();
+  }
+  const results: MlxSpeculativeRowResult[] = [];
+  for (let row = 0; row < rows.length; row++) {
+    const request = rows[row]!;
+    const acceptance = new DraftAcceptance(proposals[row]!, request.remaining, request.eosTokenIds);
+    const logprobs: (TokenLogprobs | undefined)[] = [];
+    while (!acceptance.done) {
+      const position = acceptance.position;
+      let token: number, metadata: TokenLogprobs | undefined;
+      if (tokens) token = tokens[row * width + position]!;
+      else {
+        using slice = logits.slice([row, position, 0], [row + 1, position + 1, V]);
+        using scores = ops.reshape(slice, [1, V]);
+        const sampled = await request.sampling.sample(scores, request.step + acceptance.emitted.length);
+        token = sampled.token;
+        metadata = readStepExtras(sampled.extras);
+      }
+      acceptance.accept(token, request.grammarDone?.() ?? false);
+      if (!acceptance.sawEos) logprobs.push(metadata);
+    }
+    results.push({ drafts: proposals[row]!, acceptance, logprobs });
+  }
+  return results;
+}

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;
+  }
+}