@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/generation/grammar-step.ts

@@ -0,0 +1,39 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { Cache } from "../contracts/mlx/cache";
+import { cleanupFailure, disposeResources } from "../runtime/resources";
+import type { MlxPhaseDispatch } from "./phase-dispatch";
+
+/** Append an already accepted grammar span in one graph call: known tokens
+ * (`appendKnown`), a single position through `decode`. This operation
+ * intentionally does not use the fill append policy: splitting the span changes
+ * the historical grammar path's arithmetic. Inputs/state are borrowed. Record
+ * committed IDs before projection so failures cannot hide written state.
+ * The sampler must retain its result before returning: tensor cleanup can fail
+ * after sampling. All temporary tensors belong to this operation. */
+export async function appendGrammarSpan(
+  phases: MlxPhaseDispatch, state: Cache[], ids: readonly number[],
+  committed: (ids: readonly number[]) => void, sampleLast?: (logits: MlxArray) => void,
+): Promise<void> {
+  let input: MlxArray | null = null, hidden: MlxArray | null = null;
+  let last: MlxArray | null = null, logits: MlxArray | null = null;
+  const close = () => {
+    const owned = [input, hidden, last, logits];
+    input = hidden = last = logits = null;
+    disposeResources(owned.filter((array): array is MlxArray => array !== null));
+  };
+  try {
+    input = ops.fromInt32([...ids], [1, ids.length]);
+    hidden = await (ids.length === 1 ? phases.decode(1)(input, state, null) : phases.known(ids.length)(input, state));
+    const consumed = input; input = null; consumed.dispose();
+    committed(ids);
+    if (!sampleLast) return;
+    const [, length, width] = hidden.shape as [number, number, number];
+    last = hidden.slice([0, length - 1, 0], [1, length, width]);
+    const full = hidden; hidden = null; full.dispose();
+    logits = phases.logits(1)(last);
+    const selected = last; last = null; selected.dispose();
+    sampleLast(logits);
+  } catch (error) { return cleanupFailure(error, close); }
+  finally { close(); }
+}

package/src/generation/index.ts

@@ -0,0 +1,14 @@
+export * from "./types";
+export * from "./result";
+export * from "./generate";
+export * from "./autoregressive";
+export * from "./denoising";
+export * from "./diffusion";
+export { maybePageKv } from "../state/request-policy";
+export { evalCacheState } from "./prefill";
+export type { TokenLogprobs } from "../contracts/portable/generation";
+export { acquireModelWiredLimit, wiredWorkingSetBytes, modelNeedsWiredLimit, withModelWiredLimit, withModelUsageFlush } from "./scopes";
+
+export { generationCheckpointKey } from "./checkpoint-identity";
+
+export { MlxPhaseDispatch, knownOperationFor } from "./phase-dispatch";

package/src/generation/phase-dispatch.ts

@@ -0,0 +1,115 @@
+// The scheduler's side of the graph contract: it names the phase it is in and
+// the rows it is loading, and this resolves the operation the graph built for
+// them. Bound once per graph; nothing here inspects a tensor.
+import { type MlxCapturingForward, type MlxHeadForward, type MlxLayerCapture, type MlxPromptRange, type MlxPhaseForward, type MlxTokenPhases, type MlxKnownOperations } from "../contracts/mlx/graph";
+import type { PrefillPhase } from "../contracts/portable/prefill";
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import type { MlxForwardStep } from "../contracts/mlx/forward-work";
+import { rowOperationFor } from "../runtime/row-operations";
+
+/** Index the declared whole-span operation without changing its geometry. */
+export function knownOperationFor(operations: MlxKnownOperations, rows: number): MlxPhaseForward {
+  try { return rowOperationFor(operations, rows); }
+  catch { throw new Error(`graph has no known-span operation for ${rows} rows`); }
+}
+
+/** One prompt step the scheduler runs: the phase its planner named and its
+ *  rows, `sequences` sequences of `tokens` positions each. */
+export interface MlxPromptStep {
+  readonly phase: PrefillPhase;
+  readonly sequences: number;
+  readonly tokens: number;
+}
+
+/** A verify window: `sequences` sequences of `positions` (depth + 1) each. */
+export interface MlxVerifyWindow {
+  readonly sequences: number;
+  readonly positions: number;
+}
+
+/** A prompt step as the mixed iteration's hook takes it: one position per
+ *  sequence is a decode group, any other piece a window. */
+export function promptWorkStep(step: MlxPromptStep, capture: MlxLayerCapture | null,
+  own: () => Promise<MlxArray>): MlxForwardStep {
+  const rows = step.sequences * step.tokens;
+  return { phase: step.phase === "decode" ? "decode" : "window", rows, capture, own };
+}
+
+/** A verify window as the mixed iteration's hook takes it. */
+export function verifyWorkStep(window: MlxVerifyWindow, capture: MlxLayerCapture | null,
+  own: () => Promise<MlxArray>): MlxForwardStep {
+  const rows = window.sequences * window.positions;
+  return { phase: "verify", rows, capture, own };
+}
+
+/** The phase operations of one graph, indexed by what the scheduler decides.
+ *
+ *  The Neural Engine is a unit the scheduler owns: read once, here, which
+ *  Neural Engine operations the composition built and the row range each takes.
+ *  A prompt piece of one sequence, a whole chunk or a tail, whose positions one
+ *  range holds runs on that operation; every other piece runs on the GPU. One
+ *  scheduler drives one dispatch at a time, so the unit serves one piece at a
+ *  time. */
+export class MlxPhaseDispatch {
+  readonly #phases: Partial<MlxTokenPhases> | null;
+
+  /** `graph`: the phases `tokenPhasesOf` resolved, or null for a graph whose
+   *  method runs none (denoising), which then refuses every token phase. */
+  constructor(graph: Partial<MlxTokenPhases> | null) {
+    this.#phases = graph;
+  }
+
+  get #graph(): Partial<MlxTokenPhases> {
+    if (!this.#phases) throw new Error("this graph's method runs no token phases");
+    return this.#phases;
+  }
+
+  /** Whether the composition built a Neural Engine operation. */
+  get neuralEngine(): boolean {
+    const graph = this.#phases;
+    return [graph?.prompt, graph?.promptChunk, graph?.multiSequencePrompt, graph?.multiSequencePromptChunk]
+      .some(ranges => ranges?.some(range => range.unit === "ane"));
+  }
+
+  /** One prompt step of `sequences` sequences of `tokens` positions each, in
+   *  the `phase` its planner named. */
+  prompt(phase: PrefillPhase, sequences: number, tokens: number): MlxCapturingForward {
+    if (phase === "decode") return this.decode(sequences);
+    const positions = sequences > 1 ? sequences * tokens : tokens;
+    const ranges = sequences > 1
+      ? phase === "chunk" ? this.#graph.multiSequencePromptChunk : this.#graph.multiSequencePrompt
+      : phase === "chunk" ? this.#graph.promptChunk : this.#graph.prompt;
+    if (!ranges) throw new Error(`graph has no ${sequences > 1 ? "multi-sequence " : ""}${phase} operation table`);
+    const range = ranges.find(range => positions >= range.from && positions <= range.to);
+    if (!range) throw new Error(`graph has no prompt operation for ${tokens} positions`);
+    return range.operation;
+  }
+
+  /** One new token for each of `sequences` sequences. */
+  decode(sequences: number): MlxCapturingForward {
+    const operations = this.#graph.decode;
+    if (!operations) throw new Error("graph has no decode operation table");
+    return rowOperationFor(operations, sequences);
+  }
+
+  /** A verify window of `positions` (depth + 1) per sequence. */
+  verify(sequences: number, positions: number): MlxCapturingForward {
+    const operations = this.#graph.verify;
+    if (!operations) throw new Error("graph has no verify operation table");
+    return rowOperationFor(operations, sequences * positions);
+  }
+
+  /** `tokens` known tokens appended to one sequence. */
+  known(tokens: number): MlxPhaseForward {
+    const operations = this.#graph.appendKnown;
+    if (!operations) throw new Error("graph has no known-span operation table");
+    return knownOperationFor(operations, tokens);
+  }
+
+  /** The head over a hidden of `rows` rows. */
+  logits(rows: number): MlxHeadForward {
+    const operations = this.#graph.logits;
+    if (!operations) throw new Error("graph has no logits operation table");
+    return rowOperationFor(operations, rows);
+  }
+}

package/src/generation/prefill.ts

@@ -0,0 +1,33 @@
+import type { PrefillStep } from "../contracts/portable/prefill";
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import type { Cache } from "../contracts/mlx/cache";
+import * as ops from "@mlx-bun/mlx/ops";
+import { clearCache } from "@mlx-bun/mlx/ffi";
+import { leaseCacheStates } from "../state/leases";
+import { withResource } from "../runtime/resources";
+import type { MlxPhaseDispatch } from "./phase-dispatch";
+
+/** State views with explicit ownership are released even on deferred errors.
+ * Owned views precede borrowed state, preserving the serial eval ordering. */
+export function evalCacheState(cache: Cache[]): void {
+  withResource(leaseCacheStates(cache), (state) => { if (state.length) ops.evalAll([...state]); });
+}
+
+/** Forward one prompt step of one sequence through the operation its planned
+ * phase names. Drains evaluate state (the composed caches converted inside
+ * their own appends), then clear the allocator cache. Finals return owned
+ * hidden state without forcing eval. */
+export async function executeMlxPrefillStep(
+  phases: MlxPhaseDispatch, cache: Cache[], prompt: readonly number[], step: PrefillStep,
+): Promise<MlxArray | null> {
+  const chunk = prompt.slice(step.start, step.end);
+  const ids = ops.fromInt32(chunk, [1, chunk.length]);
+  let hidden: MlxArray;
+  try { hidden = await phases.prompt(step.phase, 1, chunk.length)(ids, cache, null); }
+  finally { ids.dispose(); }
+  if (step.kind === "final") return hidden;
+  hidden.dispose();
+  evalCacheState(cache);
+  clearCache();
+  return null;
+}

package/src/generation/result.ts

@@ -0,0 +1,49 @@
+// Streaming generation — port of mlx-lm's generate_step:
+// - prefill in chunks; cache state evaluated per chunk (bounded transient
+//   memory), logits never computed for non-final prefill positions
+// - decode pipelining via mx.async_eval: step n+1's graph is built and
+//   dispatched before step n's token is read back, so the GPU never idles
+//   on the JS round-trip
+// - opt-in early token-zero yield reduces latency before that pipeline starts
+// - sampling stays on-device; only the chosen token id crosses to JS
+
+import { runtimeConfig,withRuntimeConfig } from "../runtime/config";
+import { GenerateStats,GeneratedToken } from "./types";
+
+export class Generation implements AsyncIterable<GeneratedToken> {
+  stats: GenerateStats | null = null;
+  readonly #next: () => Promise<IteratorResult<GeneratedToken, GenerateStats>>;
+  readonly #return: () => Promise<IteratorResult<GeneratedToken, GenerateStats>>;
+
+  constructor(iter: AsyncGenerator<GeneratedToken, GenerateStats>, runtime = runtimeConfig()) {
+    // Async generators resume in the caller's context on every next/return.
+    // Bind both operations so consumer awaits and early close cannot change
+    // the runtime snapshot seen by kernels or their cleanup.
+    const next = iter.next.bind(iter);
+    const close = () => iter.return(undefined as unknown as GenerateStats);
+    this.#next = () => withRuntimeConfig(runtime, next);
+    this.#return = () => withRuntimeConfig(runtime, close);
+  }
+
+  async *[Symbol.asyncIterator](): AsyncGenerator<GeneratedToken> {
+    try {
+      while (true) {
+        const r = await this.#next();
+        if (r.done) {
+          this.stats = r.value;
+          return;
+        }
+        yield r.value;
+      }
+    } finally {
+      // Consumer broke early (e.g. a decoded-text stop sequence fired):
+      // drive the inner generator's shutdown so its finallys run (array
+      // disposal, wired/adapter scopes) and capture the stats its
+      // early-return path still reports.
+      if (this.stats === null) {
+        const r = await this.#return();
+        if (r.done && r.value) this.stats = r.value;
+      }
+    }
+  }
+}

package/src/generation/scopes.ts

@@ -0,0 +1,157 @@
+// Streaming generation — port of mlx-lm's generate_step:
+// - prefill in chunks; cache state evaluated per chunk (bounded transient
+//   memory), logits never computed for non-final prefill positions
+// - decode pipelining via mx.async_eval: step n+1's graph is built and
+//   dispatched before step n's token is read back, so the GPU never idles
+//   on the JS round-trip
+// - opt-in early token-zero yield reduces latency before that pipeline starts
+// - sampling stays on-device; only the chosen token id crosses to JS
+
+import { gpuStream } from "@mlx-bun/mlx/array";
+import {
+maxRecommendedWorkingSetSize,
+setWiredLimit,
+synchronize
+} from "@mlx-bun/mlx/ffi";
+import { GenerateStats,GeneratedToken } from "./types";
+import {
+type MlxModelMemory
+} from "./bindings/autoregressive";
+import { runtimeValue } from "../runtime/config";
+
+// Scoped wired limit, raised only for near-ceiling models. mlx-lm's
+// wired_limit context wires unconditionally per generation; we deviate
+// with a measured justification:
+// - 26B-A4B (16.4 GB = 92% of the 17.8 GiB working set) NEEDS wiring —
+//   8.6 tok/s without, 32.3 with (Metal evicts weight buffers per token).
+// - 12B/e4b (≤47%) hit reference parity WITHOUT wiring, and wiring in a
+//   multi-model process (the test suite) pins memory the OTHER resident
+//   models need — async GPU exec OOM, which is uncatchable (the mlx
+//   completion-handler throw terminates the process).
+// Scope semantics match the reference: set → generate → synchronize →
+// restore; nothing stays pinned between generations. Re-entrant: only
+// the outermost wiring scope touches the limit.
+// macOS 26.6 reports a 26.8e9 B (24.96 GiB, 0.78 × RAM) recommended set on
+// a 32 GB M1 Max (2026-09-08; the earlier "24 GB machine" attribution was
+// wrong — 24.96 GiB exceeds a 24 GB box's RAM). At 0.75 × that the old
+// fraction stopped wiring the 13-16 GiB Qwen/GLM models even though they
+// page heavily without an explicit wired limit.
+// Keep smaller 8-9 GiB models unwired while covering the large-model class.
+// A graph whose composition measured a gain below the line declares
+// `wiredWeights` and is wired whatever its share.
+const WIRE_THRESHOLD = 0.5;
+let wiredScopeDepth = 0;
+let wiredOldLimit = 0;
+
+type WiredModelMemory = MlxModelMemory;
+
+/** Bytes whose MLX graph must stay resident while a model executes. */
+export function wiredWorkingSetBytes(model: WiredModelMemory): number {
+  const planned = model.expertRuntime?.plan.plannedBytes;
+  return typeof planned === "number" && Number.isFinite(planned) && planned > 0
+    ? Math.max(model.weightsBytes, planned)
+    : model.weightsBytes;
+}
+
+export function modelNeedsWiredLimit(
+  model: WiredModelMemory,
+  recommendedBytes = maxRecommendedWorkingSetSize(),
+  force = runtimeValue("MLX_BUN_FORCE_WIRE") === "1",
+): boolean {
+  return force || model.wiredWeights === true ||
+    wiredWorkingSetBytes(model) > WIRE_THRESHOLD * recommendedBytes;
+}
+
+function enterWiredScope(): void {
+  if (wiredScopeDepth++ === 0)
+    wiredOldLimit = setWiredLimit(maxRecommendedWorkingSetSize());
+}
+function exitWiredScope(): void {
+  if (--wiredScopeDepth === 0) {
+    synchronize(gpuStream);
+    setWiredLimit(wiredOldLimit);
+  }
+}
+
+/** Hold the process wired limit while one model owns GPU execution. */
+export function acquireModelWiredLimit(model: WiredModelMemory): () => void {
+  if (!modelNeedsWiredLimit(model)) return () => {};
+  enterWiredScope();
+  let released = false;
+  return () => {
+    if (released) return;
+    released = true;
+    exitWiredScope();
+  };
+}
+
+/** Hold the model's active-adapter list for exactly this generation. */
+export async function* adapterScoped(
+  model: { loraState: { active: string[] } },
+  adapters: string[],
+  inner: AsyncGenerator<GeneratedToken, GenerateStats>,
+): AsyncGenerator<GeneratedToken, GenerateStats> {
+  const previous = model.loraState.active;
+  model.loraState.active = adapters;
+  try {
+    return yield* inner;
+  } finally {
+    model.loraState.active = previous;
+  }
+}
+
+/** Wrap the generator so the wired limit is held exactly while it runs
+ *  (incl. early break/return/throw — finally fires on .return()). */
+export async function* wiredScoped(
+  inner: AsyncGenerator<GeneratedToken, GenerateStats>,
+): AsyncGenerator<GeneratedToken, GenerateStats> {
+  enterWiredScope();
+  try {
+    return yield* inner;
+  } finally {
+    exitWiredScope();
+  }
+}
+
+/** Publish the streamed model's route ledger on every generator exit path. */
+export async function* usageScoped(
+  model: WiredModelMemory,
+  inner: AsyncGenerator<GeneratedToken, GenerateStats>,
+): AsyncGenerator<GeneratedToken, GenerateStats> {
+  try {
+    return yield* inner;
+  } finally {
+    if (model.expertRuntime?.finishUsage)
+      await model.expertRuntime.finishUsage();
+    else
+      model.expertRuntime?.flushUsage?.();
+  }
+}
+
+/** Apply generate()'s scoped wiring policy to non-generator execution paths. */
+export async function withModelWiredLimit<T>(
+  model: WiredModelMemory,
+  run: () => Promise<T>,
+): Promise<T> {
+  const release = acquireModelWiredLimit(model);
+  try {
+    return await run();
+  } finally {
+    release();
+  }
+}
+
+/** Apply generate()'s usage safe-point to direct/non-generator paths. */
+export async function withModelUsageFlush<T>(
+  model: WiredModelMemory,
+  run: () => Promise<T>,
+): Promise<T> {
+  try {
+    return await run();
+  } finally {
+    if (model.expertRuntime?.finishUsage)
+      await model.expertRuntime.finishUsage();
+    else
+      model.expertRuntime?.flushUsage?.();
+  }
+}

package/src/generation/speculative/acceptance.ts

@@ -0,0 +1,41 @@
+/** One request's exact-token speculative acceptance walk. The backend owns
+ * logits and sampling; this state owns only the accepted prefix and stop facts.
+ * A group advances each live walk independently, regardless of batch size. */
+export class DraftAcceptance {
+  position = 0;
+  accepted = 0;
+  correction: number | null = null;
+  readonly emitted: number[] = [];
+  sawEos = false;
+  grammarDone = false;
+  done = false;
+
+  constructor(
+    readonly drafts: readonly number[],
+    readonly remaining: number,
+    readonly eosTokenIds: readonly number[],
+  ) {}
+
+  /** Supply one target sample after its sampler has updated request history
+   * and grammar. Once done, this request needs no further target samples. */
+  accept(token: number, grammarDone = false): void {
+    const matches = this.position < this.drafts.length && token === this.drafts[this.position];
+    this.position++;
+    if (matches) {
+      this.accepted++;
+      if (this.eosTokenIds.includes(token)) {
+        this.sawEos = this.done = true;
+        return;
+      }
+      this.emitted.push(token);
+      this.grammarDone = grammarDone;
+      this.done = grammarDone || this.emitted.length >= this.remaining;
+      return;
+    }
+    this.correction = token;
+    this.sawEos = this.eosTokenIds.includes(token);
+    if (!this.sawEos) this.emitted.push(token);
+    this.grammarDone = grammarDone;
+    this.done = true;
+  }
+}

package/src/generation/speculative/adaptive-draft-gate.ts

@@ -0,0 +1,166 @@
+/** Cost-aware speculative draft scheduler.
+ *
+ * Adapted from mlx-optiq 0.5.8's AcceptanceGate (MIT). The scheduler models
+ * the expected emitted tokens for a proposal depth and divides that by the
+ * measured cost of its round, by the rows its forward carries: a group of B
+ * sequences drafting k each verifies B(k + 1) rows, and its round without
+ * drafts is a decode step of B rows. It uses acceptance rather than
+ * observed wall-clock throughput as the online signal, so unrelated machine
+ * load does not make it oscillate.
+ */
+
+import type { AdaptiveDraftGateSettings } from "../../contracts/mlx/graph";
+
+export type VerifyCostModel = ReadonlyMap<number, number>;
+
+/** The cost of a forward of `positions` rows: the measured one, interpolated between measured rows, and past the
+ *  widest measured forward the widest one's cost per row (a forward's cost per row does not rise with its rows, so
+ *  an unmeasured wider forward is charged at most that). */
+function interpolatedCost(costs: VerifyCostModel, positions: number): number {
+  const keys = [...costs.keys()].sort((a, b) => a - b);
+  if (keys.length < 2 || !costs.has(1))
+    throw new Error("adaptive draft costs need position 1 and at least one larger position");
+  const exact = costs.get(positions);
+  if (exact !== undefined) return exact;
+  if (positions > keys.at(-1)!) return costs.get(keys.at(-1)!)! * positions / keys.at(-1)!;
+  const lo = keys.filter((key) => key < positions).at(-1)!;
+  const hi = keys.find((key) => key > positions)!;
+  return costs.get(lo)! +
+    ((costs.get(hi)! - costs.get(lo)!) / (hi - lo)) * (positions - lo);
+}
+
+/** The defaults of each setting a graph may declare for its machine (`AdaptiveDraftGateSettings`):
+ *  - `priorAccepted`, `priorRejected`: every position's counts before any round is observed, 4.5 and 0.5 (an
+ *    acceptance of 0.9, so a new gate drafts deeply and backs off as rejections arrive);
+ *  - `decay` 0.85, and `idleDecay` 1: a round without drafts leaves the counts as they are;
+ *  - `margin` 0.03;
+ *  - `probeInterval`: rounds without drafts before `next` drafts one token to observe acceptance again (8); each probe
+ *    after which the gate still drafts nothing doubles the interval, up to `maxProbeInterval` (32);
+ *  - `probeDepth`: the depth a drafter that proposes without committing drafts a probe at (the round's budget). */
+export interface AdaptiveDraftGateOptions extends AdaptiveDraftGateSettings {
+  readonly maxDraftTokens: number;
+}
+
+/** Acceptance is modeled per draft position: the probability that the draft at position i is accepted given that
+ *  every draft before it was. Each position's decayed counts are shrunk toward the position before it (weight 1;
+ *  position 1 toward the prior), so a position seldom reached follows the ones that are, and a drafter whose later
+ *  positions fall off (a recurrent MTP head trained for two steps) and one that stays flat (a block drafter) are
+ *  both modeled. A round decays the positions it drafted; with `idleDecay` a round without drafts decays them all,
+ *  so what the gate observed ages with the text it decodes and not only with the rounds it drafts. */
+export class AdaptiveDraftGate {
+  /** Decayed accepted and rejected drafts at each position (index 0 is position 1). */
+  readonly #accepted: number[];
+  readonly #rejected: number[];
+  readonly #prior: number;
+  readonly #priorWeight: number;
+  readonly #decay: number;
+  readonly #idleDecay: number;
+  readonly #margin: number;
+  readonly #maxDraftTokens: number;
+  readonly #probeInterval: number;
+  readonly #maxProbeInterval: number;
+  /** The depth of a probe drafted without a verify round. */
+  readonly probeDepth: number;
+  /** Rounds in a row `next` has run without drafts, and the rounds before its next probe. */
+  #idle = 0;
+  #interval: number;
+  #probing = false;
+
+  constructor(readonly costs: VerifyCostModel, options: AdaptiveDraftGateOptions) {
+    this.#maxDraftTokens = options.maxDraftTokens;
+    const accepted = options.priorAccepted ?? 4.5, rejected = options.priorRejected ?? 0.5;
+    this.#prior = accepted / (accepted + rejected);
+    this.#priorWeight = accepted + rejected;
+    this.#accepted = Array<number>(options.maxDraftTokens).fill(accepted);
+    this.#rejected = Array<number>(options.maxDraftTokens).fill(rejected);
+    this.#decay = options.decay ?? 0.85;
+    this.#idleDecay = options.idleDecay ?? 1;
+    this.#margin = options.margin ?? 0.03;
+    this.#probeInterval = this.#interval = options.probeInterval ?? 8;
+    this.#maxProbeInterval = options.maxProbeInterval ?? 32;
+    this.probeDepth = options.probeDepth ?? Infinity;
+  }
+
+  /** Forget what the gate observed: back to the prior, probing at the first interval. For new text (a request
+   *  joining), whose acceptance the rounds before say little about. */
+  reset(): void {
+    this.#accepted.fill(this.#prior * this.#priorWeight);
+    this.#rejected.fill((1 - this.#prior) * this.#priorWeight);
+    this.#idle = 0;
+    this.#interval = this.#probeInterval;
+  }
+
+  /** The probability that the draft at `position` (1-based) is accepted, given that the drafts before it were. */
+  acceptance(position: number): number {
+    let shrunk = this.#prior;
+    for (let i = 0; i < position; i++)
+      shrunk = (this.#accepted[i]! + shrunk) / (this.#accepted[i]! + this.#rejected[i]! + 1);
+    return shrunk;
+  }
+
+  /** One round's outcome: `drafted` proposals, the first `accepted` of them accepted. Positions 1..accepted count an
+   *  acceptance and the first rejected position a rejection, after the counts decay: every position's when every
+   *  draft was accepted (the text may have turned predictable, so the positions past the round follow the ones it
+   *  accepted), only the drafted positions' otherwise (a rejected probe says nothing about deeper drafts). A round
+   *  without drafts observes nothing. */
+  observe(drafted: number, accepted: number): void {
+    const decayed = drafted > 0 && accepted === drafted ? this.#maxDraftTokens : Math.min(drafted, this.#maxDraftTokens);
+    for (let i = 0; i < decayed; i++) {
+      this.#accepted[i] = this.#accepted[i]! * this.#decay + Number(i < Math.min(accepted, drafted));
+      this.#rejected[i] = this.#rejected[i]! * this.#decay + Number(i === accepted && accepted < drafted);
+    }
+  }
+
+  /** Expected emitted tokens per sequence of a round drafting `k` for each of `rows` sequences, per unit of cost,
+   *  where a round without drafts over those sequences costs 1 (a plain decode step emits 1 per 1). */
+  expectedRate(k: number, rows = 1): number {
+    // survive: the probability that every draft up to a position is accepted; the round emits one token more
+    // than it accepts, so tokens = 1 + the sum of survive over positions 1..k.
+    let survive = 1, tokens = 1;
+    for (let position = 1; position <= k; position++) {
+      survive *= this.acceptance(position);
+      tokens += survive;
+    }
+    return tokens * interpolatedCost(this.costs, rows) / interpolatedCost(this.costs, rows * (k + 1));
+  }
+
+  /** The next round's depth: `choose`'s, except that after a run of rounds without drafts it drafts one token,
+   *  the cheapest round that observes acceptance, so the estimate keeps following the text (a round without drafts
+   *  observes nothing). The run starts at `probeInterval` rounds and doubles after each probe that leaves the gate
+   *  drafting nothing, up to `maxProbeInterval`: a probe costs more than the plain step it replaces, so text that
+   *  stays hard to draft pays for fewer of them. `rows`: the sequences the round drafts for. */
+  next(budget: number, rows = 1): number {
+    this.#probing = false;
+    const depth = this.choose(budget, rows);
+    if (depth > 0 || budget < 1) { this.#idle = 0; this.#interval = this.#probeInterval; return depth; }
+    if (this.#idleDecay < 1) for (let i = 0; i < this.#maxDraftTokens; i++) {
+      this.#accepted[i] = this.#accepted[i]! * this.#idleDecay;
+      this.#rejected[i] = this.#rejected[i]! * this.#idleDecay;
+    }
+    if (++this.#idle <= this.#interval) return 0;
+    this.#idle = 0;
+    this.#interval = Math.min(2 * this.#interval, this.#maxProbeInterval);
+    this.#probing = true;
+    return 1;
+  }
+
+  /** Whether the depth `next` last returned is a probe: drafts the estimate does not pay for, drafted only to observe
+   *  acceptance. A drafter that can propose without committing may observe them without a verify round. */
+  get probing(): boolean { return this.#probing; }
+
+  /** The depth whose expected rate for a group of `rows` sequences beats a round without drafts by the margin, the
+   *  best such, up to `budget`; none when no depth does. */
+  choose(budget: number, rows = 1): number {
+    const limit = Math.min(this.#maxDraftTokens, budget);
+    let bestDepth = 0;
+    let bestRate = 1 + this.#margin;
+    for (let depth = 1; depth <= limit; depth++) {
+      const rate = this.expectedRate(depth, rows);
+      if (rate > bestRate) {
+        bestDepth = depth;
+        bestRate = rate;
+      }
+    }
+    return bestDepth;
+  }
+}