@mlx-bun/inference 0.0.0-stage → 1.0.0

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

@@ -0,0 +1,41 @@
+// Draft-selection policy for the speculative engine.
+//
+// Verification is exact token match: at each window position the TARGET
+// samples its own token with the step key (seed, step), independent of the
+// draft, and the draft is accepted only when equal. Draft policy therefore
+// cannot change the output distribution, only acceptance and draft cost.
+// It is NOT output-neutral trajectory by trajectory: a token's bf16 logits
+// depend on which row of the verify window computed it, so a different
+// acceptance pattern re-rolls some seeded trajectories. Compare draft-side
+// arms by ms per round and tokens per round; compare tok/s only on items whose
+// outputs are byte-identical.
+//
+// The default draft uses the request sampler with the SAME step key as the
+// target position it predicts. MLX's categorical over a [1, V] input draws by
+// INVERSE CDF in token-id order from ONE uniform of that key (see
+// draft-vocab.ts), so draft and target read the same point of their own CDFs:
+// a coupled draw whose match rate is far above independent sampling. Every
+// knob here keeps that key; each is default off and needs a paired microbench
+// win before any default (flag-and-try-both).
+import { runtimeFlag, runtimeValue } from "../../runtime/config";
+import type { SamplerOptions } from "../../sampling/index";
+
+/** Sampler options for the DRAFT of a request. `greedy` is the composition's
+ *  `greedyDraft`: argmax drafts, which drop the key coupling.
+ *  `MLX_BUN_SPEC_DRAFT_TEMPERATURE=<t>` samples the draft at its own
+ *  temperature: a head flatter than its target matches more often when
+ *  sharpened (t < request temperature); t=0 is argmax.
+ *  `MLX_BUN_SPEC_DRAFT_FILTER=0` drops top-p/top-k/min-p/XTC from the draft:
+ *  the coupled inverse-CDF draw over the full vocabulary, a quarter of the
+ *  per-step sampling cost, which at low temperature rarely leaves the
+ *  filtered set. */
+export function draftSamplerOptions<T extends SamplerOptions>(options: T, greedy: boolean): T {
+  if (greedy) return { ...options, temperature: 0 };
+  let next = options;
+  const temperature = runtimeValue("MLX_BUN_SPEC_DRAFT_TEMPERATURE");
+  if (temperature !== undefined && Number.isFinite(Number(temperature)) && Number(temperature) >= 0)
+    next = { ...next, temperature: Number(temperature) };
+  if (!runtimeFlag("MLX_BUN_SPEC_DRAFT_FILTER", true))
+    next = { ...next, topP: 0, topK: 0, minP: 0, xtcProbability: 0 };
+  return next;
+}

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

@@ -0,0 +1,118 @@
+// The draft providers a server can load, each owning how its artifact is
+// recognized and how it is loaded. The registry is model-free: it never loads
+// MLX itself, so callers detect kinds cheaply and pay for a provider only when
+// they load one. Applications use the library's default registry
+// (`defaultDraftProviders`), extend it, or hand a ready DraftProvider in
+// instead; none names a provider class.
+import { declaredGraph } from "../../models/capabilities";
+import type { DraftProvider } from "./source";
+
+/** A draft artifact directory, read lazily and at most once. */
+export interface DraftArtifact {
+  readonly dir: string;
+  has(file: string): Promise<boolean>;
+  /** config.json, or null when absent or unreadable. */
+  config(): Promise<Record<string, unknown> | null>;
+}
+
+export function draftArtifact(dir: string): DraftArtifact {
+  let config: Promise<Record<string, unknown> | null> | undefined;
+  return {
+    dir,
+    has: file => Bun.file(`${dir}/${file}`).exists(),
+    config: () => config ??= Bun.file(`${dir}/config.json`).json().then(
+      value => value as Record<string, unknown>, () => null),
+  };
+}
+
+/** What the target offers a draft that must agree with it. */
+export interface DraftTarget {
+  readonly vocabSize: number;
+  readonly tokenizer: { encode(text: string): readonly number[] };
+}
+
+export interface DraftLoadRequest {
+  /** The artifact directory; undefined for a kind with no artifact. */
+  readonly dir: string | undefined;
+  readonly target: DraftTarget;
+  /** Requested draft tokens per round; each kind applies its own default and bound. */
+  readonly numDraftTokens?: number;
+  /** Prompt-lookup window, for the kind that has one. */
+  readonly ngram?: { readonly max?: number; readonly min?: number };
+}
+
+export interface LoadedDraft {
+  readonly provider: DraftProvider;
+  readonly numDraftTokens: number;
+  /** The registered kind that loaded it (a graph may declare an adaptive width for it). */
+  readonly kind?: string;
+}
+
+export interface DraftProviderKind {
+  readonly kind: string;
+  /** False for a model-free kind: it takes no artifact and is chosen only by name. */
+  readonly artifact: boolean;
+  /** Subdirectory of a target artifact that bundles this kind's companion. */
+  readonly bundledCompanion?: string;
+  /** Whether the artifact is this kind's. Absent: the kind is chosen only by
+   * name. Kinds are asked in registration order. */
+  detect?(artifact: DraftArtifact): Promise<boolean>;
+  /** A catch-all: asked only after every other kind declined, wherever it was
+   * registered, so a kind added later is still detected ahead of it. */
+  readonly fallback?: boolean;
+  /** Load the provider and settle its round width; the provider is owned by the
+   * caller on return and disposed here on failure. */
+  load(request: DraftLoadRequest): Promise<LoadedDraft>;
+}
+
+/** A provider that binds to graphs declaring `GraphCapabilities.nativeDraft`
+ * with this kind: the draft head the graph's own checkpoint carries. */
+export interface NativeDraftKind {
+  readonly kind: string;
+  create(graph: object): Promise<DraftProvider>;
+}
+
+export class DraftProviderRegistry {
+  readonly #kinds = new Map<string, DraftProviderKind>();
+  readonly #native = new Map<string, NativeDraftKind>();
+
+  register(kind: DraftProviderKind): this {
+    if (this.#kinds.has(kind.kind)) throw new Error(`draft provider kind ${kind.kind} is already registered`);
+    this.#kinds.set(kind.kind, kind);
+    return this;
+  }
+
+  registerNative(kind: NativeDraftKind): this {
+    if (this.#native.has(kind.kind)) throw new Error(`native draft kind ${kind.kind} is already registered`);
+    this.#native.set(kind.kind, kind);
+    return this;
+  }
+
+  kinds(): string[] { return [...this.#kinds.keys()]; }
+
+  get(kind: string): DraftProviderKind | undefined { return this.#kinds.get(kind); }
+
+  /** The first registered kind that recognizes the artifact. Reads files only. */
+  async detect(dir: string): Promise<string> {
+    const artifact = draftArtifact(dir);
+    const asked = [...this.#kinds.values()].filter(kind => kind.detect);
+    for (const kind of [...asked.filter(kind => !kind.fallback), ...asked.filter(kind => kind.fallback)])
+      if (await kind.detect!(artifact)) return kind.kind;
+    throw new Error(`${dir}: no registered draft provider recognizes this artifact`);
+  }
+
+  async load(kind: string, request: DraftLoadRequest): Promise<LoadedDraft> {
+    const entry = this.#kinds.get(kind);
+    if (!entry) throw new Error(`unknown draft provider kind ${kind} (registered: ${this.kinds().join(", ")})`);
+    return { ...await entry.load(request), kind };
+  }
+
+  /** The draft the graph declares it carries, or null when it declares none. */
+  async native(graph: object): Promise<LoadedDraft | null> {
+    const declared = declaredGraph(graph).graphCapabilities.nativeDraft;
+    if (!declared) return null;
+    const kind = this.#native.get(declared.kind);
+    if (!kind) throw new Error(`the graph declares native draft ${declared.kind}, which no provider is registered for`);
+    return { provider: await kind.create(graph), numDraftTokens: declared.numDraftTokens, kind: declared.kind };
+  }
+}

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

@@ -0,0 +1,110 @@
+// Frequency-ranked draft vocabulary (FR-Spec, Zhao et al. 2025).
+//
+// A native MTP head spends most of each draft step in the target's vocabulary
+// projection (Qwen3.8: 248,320 rows). Drafts only need to be LIKELY tokens, so
+// the draft projects onto the N most frequent rows of the same quantized head.
+// The target still verifies and samples over the full vocabulary: the output
+// distribution is unchanged (a different acceptance pattern can re-roll a seeded
+// trajectory), a token outside the list simply cannot be drafted (one miss).
+//
+// Coupling is preserved. For a [1, V] input MLX's categorical draws by INVERSE
+// CDF in token-id order with ONE uniform from the step key (random.cpp
+// categorical_inverse_cdf), and the default draft uses the same key as the target
+// position it predicts, so both read the same point of their own CDFs. The list
+// is kept in ascending id order, so a categorical draw over the listed rows IS a
+// full-vocabulary draw with zero mass outside the list: same uniform, same order,
+// same coupling. No noise reconstruction is involved.
+//
+// A list shipped beside the draft companion (`draft_vocab.json`) loads
+// automatically; MLX_BUN_SPEC_DRAFT_VOCAB=<path to json {"ids":[...]}> selects
+// another list and =0/=off disables it (see configuredDraftVocabulary).
+import { existsSync } from "node:fs";
+import { join } from "node:path";
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { DRAFT_VOCAB_FILE, readDraftVocabulary } from "../../artifacts/draft-vocabulary";
+import { runtimeValue } from "../../runtime/config";
+import { applyMinP, applyTopK, applyTopP, stepKey, toLogprobs, type SamplerOptions } from "../../sampling/index";
+
+export { DRAFT_VOCAB_FILE };
+
+export interface QuantizedHeadState {
+  readonly w: MlxArray;
+  readonly scales: MlxArray;
+  readonly biases: MlxArray | null;
+  readonly spec: ops.QuantSpec;
+}
+
+/** The draft vocabulary for this target, validated against its vocabulary size,
+ *  or null. Precedence: `MLX_BUN_SPEC_DRAFT_VOCAB=<path>` selects a list,
+ *  `=0`/`=off` disables, and otherwise the list shipped beside the draft
+ *  companion (`<companionDir>/draft_vocab.json`) is used when present. */
+export function configuredDraftVocabulary(vocabSize: number, companionDir?: string): readonly number[] | null {
+  const setting = runtimeValue("MLX_BUN_SPEC_DRAFT_VOCAB");
+  if (setting === "0" || setting === "off") return null;
+  let path = setting;
+  if (!path && companionDir) {
+    const shipped = join(companionDir, DRAFT_VOCAB_FILE);
+    if (existsSync(shipped)) path = shipped;
+  }
+  return path ? readDraftVocabulary(path, vocabSize) : null;
+}
+
+/** Rows of a quantized head (affine packs along the input axis, so a row
+ *  gather is exact). Owns its gathered arrays; borrows nothing afterwards. */
+export class DraftVocabularyHead {
+  readonly ids: MlxArray; // [N] int32, ascending token ids
+  readonly size: number;
+  readonly #w: MlxArray;
+  readonly #scales: MlxArray;
+  readonly #biases: MlxArray | null;
+  readonly #spec: ops.QuantSpec;
+
+  constructor(head: QuantizedHeadState, ids: readonly number[]) {
+    this.size = ids.length;
+    this.ids = ops.fromInt32([...ids], [ids.length]);
+    const rows = (a: MlxArray) => { using taken = ops.takeAxis(a, this.ids, 0); return ops.contiguous(taken); };
+    this.#w = rows(head.w);
+    this.#scales = rows(head.scales);
+    this.#biases = head.biases ? rows(head.biases) : null;
+    this.#spec = head.spec;
+    ops.evalAll([this.ids, this.#w, this.#scales, ...(this.#biases ? [this.#biases] : [])]);
+  }
+
+  /** hidden [..., H] → logits [..., N] over the listed ids. */
+  project(hidden: MlxArray): MlxArray {
+    return ops.quantizedMatmul(hidden, this.#w, this.#scales, this.#biases, this.#spec, true);
+  }
+
+  dispose(): void {
+    this.ids.dispose(); this.#w.dispose(); this.#scales.dispose(); this.#biases?.dispose();
+  }
+}
+
+export type SubsetDraftSampler = (logits: MlxArray, head: DraftVocabularyHead, step: number) => MlxArray;
+
+/** One request's coupled draw over a draft vocabulary: subset logits [1, N] →
+ *  owned token id [1] (a FULL-vocabulary id). Null when the options need a
+ *  sampler this path does not reproduce (curve, HLG, XTC). */
+export function makeSubsetDraftSampler(options: SamplerOptions): SubsetDraftSampler | null {
+  const { temperature = 0, topP = 0, topK = 0, minP = 0, minTokensToKeep = 1, xtcProbability = 0, seed = 0 } = options;
+  if (options.curve || options.hlg?.enabled === true || xtcProbability > 0) return null;
+  const pick = (index: MlxArray, head: DraftVocabularyHead): MlxArray => {
+    using flat = ops.reshape(index, [1]);
+    return ops.takeAxis(head.ids, flat, 0);
+  };
+  if (temperature === 0) return (logits, head) => { using index = ops.argmaxAxis(logits, -1); return pick(index, head); };
+  return (logits, head, step) => {
+    const owned: MlxArray[] = [];
+    try {
+      let current = toLogprobs(logits); owned.push(current);
+      if (topP > 0 && topP < 1) { current = applyTopP(current, topP); owned.push(current); }
+      if (minP !== 0) { current = applyMinP(current, minP, minTokensToKeep); owned.push(current); }
+      if (topK > 0 && topK < head.size) { current = applyTopK(current, topK); owned.push(current); }
+      using scaled = ops.mulScalar(current, 1 / temperature);
+      using key = stepKey(seed, step);
+      using index = ops.randomCategorical(scaled, key);
+      return pick(index, head);
+    } finally { for (const array of owned) array.dispose(); }
+  };
+}

package/src/generation/speculative/index.ts

@@ -0,0 +1,17 @@
+export * from "./source";
+export * from "./acceptance";
+export * from "./output";
+export * from "./draft-policy";
+export * from "./draft-registry";
+export * from "./draft-kind";
+export * from "./draft-vocab";
+export * from "./adaptive-draft-gate";
+export * from "./sources/assistant-source";
+export * from "./sources/two-model";
+export * from "./sources/qwen-mtp-source";
+export * from "./sources/glm52-mtp-source";
+export * from "./sources/deepspec-source";
+export * from "./sources/dflash-source";
+export * from "./sources/dflash2-source";
+export * from "./sources/ngram-source";
+export { projectedDraftGroups, type ProjectedDraftMethod } from "./bindings/projected-draft-rows";

package/src/generation/speculative/output.ts

@@ -0,0 +1,49 @@
+import type { Cancellation, CancelReason, GenerationOutput, TokenLogprobs } from "../../contracts/portable/generation";
+import type { DraftAcceptance } from "./acceptance";
+
+export interface DraftOutputRow {
+  readonly acceptance: DraftAcceptance;
+  readonly logprobs?: readonly (TokenLogprobs | undefined)[];
+  readonly output: Pick<GenerationOutput, "commit">;
+  readonly cancellation?: Cancellation;
+}
+
+/** Counts apply to this round. Accepted counts describe processed proposed
+ * inputs, excluding an unprocessed correction/bonus and any unpublished EOS.
+ * Failed/cancelled requests must be discarded, never checkpointed. */
+export type DraftOutputResult = { readonly generated: number; readonly accepted: number } & (
+  | { readonly kind: "continue"; readonly pending: number }
+  | { readonly kind: "stop" | "length" }
+  | { readonly kind: "cancelled"; readonly reason: CancelReason }
+  | { readonly kind: "failed"; readonly error: unknown }
+);
+
+/** Deliver verified content through the output interface before resolving
+ * retained state. A stopped or failed consumer affects only its own request.
+ * Token-sized publications preserve callbacks that stop inside a draft burst. */
+export async function deliverDraftOutputs(rows: readonly DraftOutputRow[]): Promise<DraftOutputResult[]> {
+  const results: DraftOutputResult[] = [];
+  for (const row of rows) results.push(await deliver(row));
+  return results;
+}
+
+async function deliver(row: DraftOutputRow): Promise<DraftOutputResult> {
+  const { acceptance } = row;
+  let generated = 0;
+  const retained = () => ({ generated, accepted: Math.min(generated, acceptance.accepted) });
+  try {
+    for (const token of acceptance.emitted) {
+      const reason = row.cancellation?.reason;
+      if (reason) return { kind: "cancelled", reason, generated, accepted: 0 };
+      const metadata = row.logprobs?.[generated];
+      const more = await row.output.commit([token], metadata ? [metadata] : undefined);
+      generated++;
+      if (more === false) return { kind: "stop", ...retained() };
+    }
+    const reason = row.cancellation?.reason;
+    if (reason) return { kind: "cancelled", reason, generated, accepted: 0 };
+    if (acceptance.sawEos || acceptance.grammarDone) return { kind: "stop", ...retained() };
+    if (generated >= acceptance.remaining) return { kind: "length", ...retained() };
+    return { kind: "continue", pending: acceptance.correction!, ...retained() };
+  } catch (error) { return { kind: "failed", error, generated, accepted: 0 }; }
+}

package/src/generation/speculative/round.ts

@@ -0,0 +1,307 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { MlxVerifyWindow } from "../phase-dispatch";
+import type { Cancellation, GenerationOutput, TokenLogprobs } from "../../contracts/portable/generation";
+import type { DisposableResource } from "../../contracts/portable/resources";
+import { cleanupFailure, disposeResources } from "../../runtime/resources";
+import type { SpeculativeTransaction } from "../../contracts/portable/rollback";
+import { DraftAcceptance } from "./acceptance";
+import { deliverDraftOutputs, type DraftOutputResult } from "./output";
+import { readStepExtras, type NumberStepSampler } from "../../sampling/index";
+import { runtimeFlag } from "../../runtime/config";
+import type { SpecPhaseMs } from "../index";
+
+/** Diagnostic phase timing: evaluates the verify logits before sampling so
+ * draft / verify / sample / commit wall time can be attributed separately.
+ * Off by default because that evaluation removes the draft-verify overlap. */
+export function specPhaseTiming(): boolean {
+  return runtimeFlag("MLX_BUN_SPEC_PHASE_TIMING", false);
+}
+
+/** 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[][]>;
+  /** Optional: the same proposals left on the device as [B, depth] ids, resolved
+   *  from one later host readback. Lets the round build its verify graph while
+   *  the GPU is still drafting. Null when unavailable for this call. */
+  draftDevice?(pending: readonly number[], depth: number, steps: readonly number[]):
+    { tokens: MlxArray; resolve(read: readonly number[]): number[][] } | null;
+  commit(accepted: readonly number[], context: Context): void | Promise<void>;
+  /** Diagnostic barrier for retained lazy state; never called by normal rounds. */
+  materialize?(): 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[]>;
+  /** The verify forward over `ids` [sequences, positions], the window the round built. */
+  forward(ids: MlxArray, window: MlxVerifyWindow): Promise<{ logits: MlxArray; context: Context }>;
+  /** Diagnostic barrier for retained target state after verification/rollback. */
+  materialize?(): void;
+}
+
+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[]; phaseMs?: SpecPhaseMs }> {
+  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, ...(round.phaseMs ? { phaseMs: round.phaseMs } : {}) };
+}
+
+/** 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>;
+  /** Present under `MLX_BUN_SPEC_PHASE_TIMING=1`; `commit` fills in after commit(). */
+  readonly phaseMs?: SpecPhaseMs;
+}
+
+/** 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 timing = specPhaseTiming();
+  const phaseMs: SpecPhaseMs | undefined = timing ? { draft: 0, verify: 0, sample: 0, commit: 0, rounds: 1 } : undefined;
+  let mark = timing ? performance.now() : 0;
+  const lap = (): number => { const now = performance.now(); const elapsed = now - mark; mark = now; return elapsed; };
+  // Device-first round: leave the draft tokens on the GPU, assemble the verify
+  // ids there, start the GPU on the draft chain, and build the ~2.5k-node verify
+  // graph while it runs. Drafts and window samples then come back in ONE read.
+  // Same ops and values as the host-first order, so outputs are unchanged.
+  // Diagnostics (phase timing) keep the host-first order so phases stay separable.
+  const lazy = !phaseMs && depth > 0 && runtimeFlag("MLX_BUN_SPEC_DEVICE_ROUND", true) && draft.draftDevice &&
+    rows.every(row => row.sampling.independent || row.sampling.positional)
+    ? draft.draftDevice(rows.map(row => row.pending), depth, rows.map(row => row.step)) : null;
+  if (lazy) return prepareDeviceRound(rows, depth, lazy, draft, target);
+  const inventory = phaseMs && runtimeFlag("MLX_BUN_SPEC_OP_INVENTORY", false);
+  const counting = globalThis as Record<string, unknown>;
+  if (inventory) counting.__opCount = {};
+  let proposals: number[][];
+  try { proposals = await draft.draft(rows.map(row => row.pending), depth, rows.map(row => row.step)); }
+  finally { if (inventory && phaseMs) { phaseMs.draftOps = counting.__opCount as Record<string, number>; counting.__opCount = undefined; } }
+  if (phaseMs) phaseMs.draft = lap();
+  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);
+  // Component attribution INSIDE the served verify forward. The model's
+  // profiler hooks add an evaluation barrier per op, so this distorts the
+  // verify total; it is only ever on with phase timing, for diagnosis.
+  const layers: Record<string, number> | undefined =
+    phaseMs && runtimeFlag("MLX_BUN_SPEC_LAYER_PROFILE", false) ? {} : undefined;
+  const profiled = globalThis as Record<string, unknown>;
+  if (layers) profiled.__deltaProf = layers;
+  if (inventory) counting.__opCount = {};
+  let verification: Awaited<ReturnType<typeof target.forward>>;
+  try { verification = await target.forward(ids, { sequences: B, positions: verifyDepth + 1 }); }
+  finally {
+    if (layers) profiled.__deltaProf = undefined;
+    if (inventory && phaseMs) { phaseMs.verifyOps = counting.__opCount as Record<string, number>; counting.__opCount = undefined; }
+  }
+  if (layers && phaseMs) phaseMs.layers = layers;
+  const dispose = () => disposeResources([verification.context, verification.logits]);
+  try {
+    if (phaseMs) {
+      // Force the verify forward to completion so its time is not attributed
+      // to the sampling readback. Production leaves it lazy on purpose.
+      ops.evalAll([verification.logits]);
+      target.materialize?.();
+      phaseMs.verify = lap();
+    }
+    const results = await sampleSpeculativeRows(rows, proposals, verification.logits, ids);
+    if (phaseMs) phaseMs.sample = lap();
+    return {
+      results,
+      async commit(accepted = results.map(result => result.acceptance.accepted)) {
+        if (phaseMs) lap();
+        target.transaction.resolve(accepted);
+        await draft.commit(accepted, verification.context);
+        if (phaseMs) {
+          target.materialize?.();
+          draft.materialize?.();
+          phaseMs.commit = lap();
+        }
+      },
+      dispose,
+      ...(phaseMs ? { phaseMs } : {}),
+    };
+  } catch (error) { return cleanupFailure(error, dispose); }
+}
+
+async function prepareDeviceRound<Context extends DisposableResource>(
+  rows: readonly MlxSpeculativeRow[], depth: number,
+  lazy: { tokens: MlxArray; resolve(read: readonly number[]): number[][] },
+  draft: MlxDraftRows<NoInfer<Context>>, target: MlxRowVerification<Context>,
+): Promise<MlxPreparedSpeculativeRound> {
+  const B = rows.length;
+  let verification: { logits: MlxArray; context: Context } | null = null;
+  const dispose = () => { if (verification) disposeResources([verification.context, verification.logits]); };
+  try {
+    ops.asyncEvalAll([lazy.tokens]); // the GPU drafts while the verify graph is built
+    using pendingIds = ops.fromInt32(rows.map(row => row.pending), [B, 1]);
+    using pending = pendingIds.astype(lazy.tokens.dtype);
+    using ids = ops.concatAxis([pending, lazy.tokens], 1);
+    target.transaction.begin(depth);
+    verification = await target.forward(ids, { sequences: B, positions: depth + 1 });
+    const sampled = sampleWindowOnDevice(rows, verification.logits, ids);
+    if (!sampled) throw new Error("device round requires stateless window sampling");
+    let read: number[];
+    try {
+      using drafted = ops.reshape(lazy.tokens, [B * depth]);
+      using draftedSameType = drafted.astype(sampled.dtype);
+      using both = ops.concatAxis([draftedSameType, sampled], 0);
+      read = both.toIntTokens();
+    } finally { sampled.dispose(); }
+    const proposals = lazy.resolve(read.slice(0, B * depth));
+    const results = await walkSpeculativeRows(rows, proposals, verification.logits, read.slice(B * depth));
+    const held = verification;
+    return {
+      results,
+      async commit(accepted = results.map(result => result.acceptance.accepted)) {
+        target.transaction.resolve(accepted);
+        await draft.commit(accepted, held.context);
+      },
+      dispose,
+    };
+  } catch (error) { return cleanupFailure(error, dispose); }
+  finally { lazy.tokens.dispose(); }
+}
+
+/** Sample a borrowed verify window over its borrowed inputs `ids` [B, width]. 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, ids: MlxArray,
+): Promise<MlxSpeculativeRowResult[]> {
+  let tokens: number[] | undefined;
+  const sampled = sampleWindowOnDevice(rows, logits, ids);
+  if (sampled) { try { tokens = sampled.toIntTokens(); } finally { sampled.dispose(); } }
+  return walkSpeculativeRows(rows, proposals, logits, tokens);
+}
+
+/** Every window position of every row as ONE owned device array [B * width]
+ *  (row-major), or null when some row's sampling depends on a grammar or
+ *  probability capture (those walk position by position).
+ *  Plain greedy rows share one stateless argmax; other rows (temperature/top-p/top-k,
+ *  processor penalties) draw every position with its own (seed, step) key over
+ *  the history the window's inputs `ids` [B, width] (borrowed) give it, so the
+ *  walk still stops at each row's first rejection with the draws sequential
+ *  sampling would make, and unused suffix draws consume nothing. */
+export function sampleWindowOnDevice(rows: readonly MlxSpeculativeRow[], logits: MlxArray, ids: MlxArray): MlxArray | null {
+  const V = logits.shape[2]!, width = logits.shape[1]!;
+  const independent = rows[0]?.sampling.independent;
+  if (independent && rows.every(row => row.sampling.independent === independent)) {
+    using flat = ops.reshape(logits, [rows.length * width, V]);
+    return independent.sample(flat);
+  }
+  if (rows.length === 0 || !rows.every(row => row.sampling.positional)) return null;
+  const perRow: MlxArray[] = [];
+  try {
+    for (const [row, request] of rows.entries()) {
+      using slice = logits.slice([row, 0, 0], [row + 1, width, V]);
+      using scores = ops.reshape(slice, [width, V]);
+      using inputs = request.sampling.needsHistory ? ids.slice([row, 0], [row + 1, width]) : null;
+      using window = inputs ? ops.reshape(inputs, [width]) : null;
+      perRow.push(request.sampling.positional!.sample(
+        scores, Array.from({ length: width }, (_, position) => request.step + position), window ?? undefined));
+    }
+    return perRow.length === 1 ? perRow.pop()! : ops.concatAxis(perRow, 0);
+  } finally { for (const sampled of perRow) sampled.dispose(); }
+}
+
+/** The host accept walk over already-read window samples (`tokens`, row-major
+ *  [B * width]), whose used draws it commits to each row's history, or, without
+ *  them, position-by-position sampling. */
+export async function walkSpeculativeRows(
+  rows: readonly MlxSpeculativeRow[], proposals: readonly (readonly number[])[], logits: MlxArray,
+  tokens: readonly number[] | undefined,
+): Promise<MlxSpeculativeRowResult[]> {
+  const V = logits.shape[2]!, width = logits.shape[1]!;
+  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);
+    }
+    if (tokens) request.sampling.positional?.commit(tokens.slice(row * width, row * width + acceptance.position));
+    results.push({ drafts: proposals[row]!, acceptance, logprobs });
+  }
+  return results;
+}