mlx-bun 0.2.0 → 0.4.0

package/src/backends/mlx/prefill-rows.ts

@@ -0,0 +1,200 @@
+import type { BatchableCache, Cache } from "../../model/gemma4-base";
+import type { MlxArray } from "../../mlx/array";
+import * as ops from "../../mlx/ops";
+import { clearCache } from "../../mlx/ffi";
+import type { PrefillStep } from "../../inference/prefill";
+import { cloneKvCaches, type CacheCodecProvider } from "../../kv-store";
+import { disposeResources, cleanupFailure } from "../../engine/resources";
+import type { MlxGroupPreparation, Row } from "./batch-group";
+import { evalCacheState } from "./prefill";
+import { prefillCacheLayout } from "./cache-layout";
+import { MlxStateRows } from "./state-rows";
+import type { KvMaintenance } from "./kv-maintenance";
+
+export interface MlxPrefillStep extends PrefillStep {
+  /** Some methods maintain their final target chunk before sampling. */
+  readonly maintain?: boolean;
+}
+
+export interface MlxPrefillState {
+  row: Row;
+  solo: Cache[];
+  pos: number;
+  retain?: () => void;
+  planned?: MlxPrefillStep;
+}
+
+/** The method owns initialization, planning, companion state, snapshots and
+ * completion. This driver owns common target batching and work boundaries.
+ * Checkpoint factories return owned cache state; close/dispose are idempotent. */
+export interface MlxPrefillRowsHost<State extends MlxPrefillState> {
+  stateCodecs?: CacheCodecProvider;
+  maintain?: KvMaintenance;
+  open(row: Row): State;
+  plan(state: State): MlxPrefillStep;
+  /** Already prepared state can complete without any target forward. */
+  ready?(state: State): boolean;
+  forward(ids: MlxArray, caches: Cache[], states: readonly State[]): Promise<MlxArray>;
+  project(hidden: MlxArray, caches: Cache[], completed: readonly State[]): MlxArray | null;
+  /** Consume method context after target maintenance, before snapshots. */
+  afterForward?(ids: MlxArray, caches: Cache[], states: readonly State[], hidden: MlxArray): void | Promise<void>;
+  checkpoint(state: State, capture: () => Cache[], row: number): void;
+  complete(state: State, logits: MlxArray | null, row: number): Promise<void>;
+  reject(row: Row, error: unknown): void;
+  close(state: State): void;
+  filterRows?(keep: readonly number[]): void;
+  dispose?(): void;
+}
+
+/** One target preparation lifecycle at B1/B>1. Late joins retain each row's
+ * planned chunk endpoint, precision transition and checkpoint boundary. */
+export class MlxPrefillRows<State extends MlxPrefillState> implements MlxGroupPreparation {
+  #states: State[] = [];
+  #rows: Row[] = [];
+  #stateRows?: MlxStateRows;
+  constructor(readonly operations: MlxPrefillRowsHost<State>) {}
+  get rows(): readonly Row[] { return this.#rows; }
+  get canAdmit(): boolean { return this.#rows.length > 0; }
+  /** Keep a request's prefill weight until retirement so a nearly completed
+   * request does not repeatedly attract another long prefill into its tail. */
+  get tokenWeight(): number { return this.#states.reduce((sum, state) => sum + state.row.promptTokens - state.row.cachedTokens, 0); }
+
+  admit(row: Row): void {
+    this.#states.push(this.operations.open(row));
+    this.#rows.push(row);
+  }
+
+  #release(state: State): void {
+    const caches = state.solo; state.solo = [];
+    const retain = state.retain; state.retain = undefined;
+    disposeResources([...caches, { dispose: () => retain?.() }, { dispose: () => this.operations.close(state) }]);
+  }
+  #filter(keep: number[]): void {
+    if (keep.length) this.#stateRows?.filterRows(keep);
+    else { const rows = this.#stateRows; this.#stateRows = undefined; rows?.dispose(); }
+    this.operations.filterRows?.(keep);
+    this.#states = keep.map(row => this.#states[row]!);
+    this.#rows = keep.map(row => this.#rows[row]!);
+  }
+  /** Admission changes membership only between forwards. Existing planned
+   * endpoints survive a join, including partially consumed drain chunks. */
+  #synchronizeRows(): void {
+    if (!this.#stateRows && this.#states.length === 1) return;
+    if (!this.#stateRows) {
+      for (const state of this.#states) this.operations.maintain?.preparePrefill?.(state.solo);
+      const layouts: BatchableCache[] = [];
+      try { for (const cache of this.#states[0]!.solo) layouts.push(prefillCacheLayout(cache)); }
+      catch (error) { return cleanupFailure(error, () => disposeResources(layouts)); }
+      this.#stateRows = new MlxStateRows(layouts);
+      this.#stateRows.mergeRows(this.#states.map(state => state.solo));
+    } else {
+      for (const state of this.#states.slice(this.#stateRows.rowCount)) {
+        this.operations.maintain?.preparePrefill?.(state.solo);
+        this.#stateRows.append(state.solo);
+      }
+    }
+    // Resolve input-backed copies before releasing SSD restore leases.
+    if (this.#states.some(state => state.solo.length)) {
+      evalCacheState(this.#stateRows.caches);
+      for (const state of this.#states) {
+        const caches = state.solo; state.solo = [];
+        const retain = state.retain; state.retain = undefined;
+        disposeResources([...caches, { dispose: () => retain?.() }]);
+      }
+    }
+    for (const cache of this.#stateRows.caches) cache.prefillMaintenance?.beginPrefill();
+  }
+
+  async advance(): Promise<boolean> {
+    // Complete restored inputs before merging: their cache coverage must not
+    // introduce padding or a zero-width forward into unrelated cold rows.
+    if (this.operations.ready && this.#states.some(state => this.operations.ready!(state))) {
+      const remaining: State[] = [];
+      for (const state of this.#states) {
+        if (!this.operations.ready(state)) { remaining.push(state); continue; }
+        try {
+          state.row.req.signal?.throwIfAborted();
+          await this.operations.complete(state, null, -1);
+          state.solo = []; state.retain = undefined;
+        } catch (error) { this.operations.reject(state.row, error); this.#release(state); }
+        finally { this.operations.close(state); }
+      }
+      this.#states = remaining; this.#rows = remaining.map(state => state.row);
+      if (!remaining.length) { this.dispose(); return true; }
+    }
+    this.#synchronizeRows();
+    const keep: number[] = [];
+    for (const [index, state] of this.#states.entries()) {
+      if (!state.row.req.signal?.aborted) { keep.push(index); continue; }
+      this.#release(state); this.operations.reject(state.row, state.row.req.signal.reason);
+    }
+    if (keep.length !== this.#states.length) this.#filter(keep);
+    if (!this.#states.length) { this.dispose(); return true; }
+    while (this.#states.length) {
+      for (const state of this.#states) state.planned ??= this.operations.plan(state);
+      const count = Math.min(...this.#states.map(state => state.planned!.end - state.pos));
+      const closes = this.#states.map(state => state.row.req.trace?.begin("prefill.chunk", {
+        mechanism: "continuous", startToken: state.pos, batchSize: this.#states.length,
+      }));
+      const caches = this.#stateRows?.caches ?? this.#states[0]!.solo;
+      try {
+        using ids = count ? ops.fromInt32(this.#states.flatMap(state => state.row.req.promptIds.slice(state.pos, state.pos + count)),
+          [this.#states.length, count]) : null;
+        using hidden = ids ? await this.operations.forward(ids, caches, this.#states) : null;
+        const drains: number[] = [], finals: number[] = [];
+        let batchYield = false;
+        for (const [row, state] of this.#states.entries()) {
+          state.pos += count;
+          if (state.pos !== state.planned!.end) continue;
+          if (state.planned!.kind === "final") finals.push(row);
+          else drains.push(row);
+          batchYield ||= state.planned!.batchYield;
+        }
+        const maintained = [...drains, ...finals.filter(row => this.#states[row]!.planned!.maintain)];
+        if (count && (maintained.length || !finals.length)) {
+          evalCacheState(caches);
+          if (this.#stateRows) {
+            for (const cache of caches) cache.prefillMaintenance?.commitPrefill(maintained);
+          } else if (maintained.length) this.operations.maintain?.(this.#states[0]!.solo);
+          clearCache();
+        }
+        if (ids) {
+          const companion = this.operations.afterForward?.(ids, caches, this.#states, hidden!);
+          if (companion) await companion;
+        }
+        for (const row of [...drains, ...finals]) {
+          const state = this.#states[row]!;
+          if (state.planned!.snapshot) this.operations.checkpoint(state, () =>
+            this.#stateRows ? this.#stateRows.extractRow(row) : cloneKvCaches(state.solo, this.operations.stateCodecs), row);
+          if (!finals.includes(row)) state.planned = undefined;
+        }
+        using logits = hidden && finals.length ? this.operations.project(hidden, caches, finals.map(row => this.#states[row]!)) : null;
+        for (const row of finals) {
+          const state = this.#states[row]!;
+          if (this.#stateRows) state.solo = this.#stateRows.extractRow(row);
+          using part = logits ? logits.slice([row, logits.shape[1]! - 1, 0], [row + 1, logits.shape[1]!, logits.shape[2]!]) : null;
+          try {
+            state.row.req.signal?.throwIfAborted();
+            await this.operations.complete(state, part, row);
+            state.solo = []; state.retain = undefined;
+          } catch (error) {
+            this.operations.reject(state.row, error);
+            this.#release(state);
+          } finally { this.operations.close(state); }
+        }
+        for (const cache of caches) (cache as { releaseRopeArr?: () => void }).releaseRopeArr?.();
+        if (finals.length) this.#filter(this.#states.flatMap((_, row) => finals.includes(row) ? [] : [row]));
+        if (!this.#states.length) { this.dispose(); return true; }
+        if (batchYield || finals.length) return false;
+      } finally { for (const close of closes) close?.(); }
+    }
+    return true;
+  }
+
+  dispose(): void {
+    const states = this.#states; this.#states = []; this.#rows = [];
+    const rows = this.#stateRows; this.#stateRows = undefined;
+    disposeResources([...(rows ? [rows] : []), ...states.map(state => ({ dispose: () => this.#release(state) })),
+      { dispose: () => this.operations.dispose?.() }]);
+  }
+}

package/src/backends/mlx/prefill.ts

@@ -0,0 +1,33 @@
+import type { PrefillStep } from "../../inference/prefill";
+import type { MlxArray } from "../../mlx/array";
+import type { Cache } from "../../model/gemma4";
+import * as ops from "../../mlx/ops";
+import { clearCache } from "../../mlx/ffi";
+import { leaseCacheStates } from "./state-views";
+import { withResource } from "../../engine/resources";
+
+/** 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 once. Drains evaluate state, apply KV maintenance, then clear the
+ * allocator cache. Finals return owned hidden state without forcing eval. */
+export async function executeMlxPrefillStep(
+  forward: (ids: MlxArray, cache: Cache[]) => MlxArray | Promise<MlxArray>,
+  cache: Cache[], prompt: readonly number[], step: PrefillStep,
+  maintain: () => void,
+): 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 forward(ids, cache); }
+  finally { ids.dispose(); }
+  if (step.kind === "final") return hidden;
+  hidden.dispose();
+  evalCacheState(cache);
+  maintain();
+  clearCache();
+  return null;
+}

package/src/backends/mlx/projected-context-rows.ts

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

package/src/backends/mlx/request-policy.ts

@@ -0,0 +1,20 @@
+import type { GenerateOptions } from "../../generate";
+
+/** Copy policy data without cloning native caches, controllers or callbacks. */
+export function snapshotGenerationPolicy<T extends GenerateOptions & { stopSequences?: readonly string[] }>(options: T): T {
+  const snapshot = { ...options };
+  const freeze = (value: unknown): void => {
+    if (!value || typeof value !== "object") return;
+    for (const item of Object.values(value)) freeze(item);
+    Object.freeze(value);
+  };
+  for (const key of ["stopSequences", "eosTokenIds", "adapters", "xtcSpecialTokens", "logitBias",
+    "kvConfig", "turboQuant", "pagedKv", "hlg", "curve", "decodePolicy"] as const) {
+    if (snapshot[key] === undefined) continue;
+    const value = structuredClone(snapshot[key]);
+    freeze(value);
+    Object.assign(snapshot, { [key]: value });
+  }
+  return Object.freeze(snapshot);
+}
+

package/src/backends/mlx/request-state-policy.ts

@@ -0,0 +1,41 @@
+import type { GenerateOptions } from "../../generate";
+import type { RuntimeModel } from "../../model/factory";
+import { KVCache, type Cache } from "../../model/gemma4-base";
+import { PagedKVCache } from "../../lab/paged-kv/paged-kv";
+import type { RowPromptCache } from "./batch-group";
+import { disposeResources } from "../../engine/resources";
+
+/** Request state construction and reusable-prefix policy. Equal keys declare
+ * compatible state layouts; the executor need not inspect their settings. */
+export interface MlxRequestStatePolicy {
+  readonly key: string;
+  readonly promptCache?: RowPromptCache;
+  create(): Cache[];
+}
+
+/** Replace fresh full-attention storage before prefill. Sliding layers retain
+ * their layout; pre-warmed library caches retain their existing storage. */
+export function maybePageKv(
+  cache: Cache[], options: GenerateOptions, capacityTokens: number,
+): void {
+  if (!options.pagedKv) return;
+  if (cache.some((c) => c.offset > 0)) return;
+  const blockSize = options.pagedKv.blockSize ?? PagedKVCache.DEFAULT_BLOCK_SIZE;
+  for (let i = 0; i < cache.length; i++) {
+    if (cache[i] instanceof KVCache) {
+      cache[i]!.dispose(); // fresh (offset 0) — nothing stored yet
+      cache[i] = new PagedKVCache(capacityTokens, blockSize);
+    }
+  }
+}
+
+export function bindPagedRequestState(model: RuntimeModel, options: GenerateOptions,
+  capacityTokens: number): MlxRequestStatePolicy | undefined {
+  if (!options.pagedKv) return undefined;
+  const blockSize = options.pagedKv.blockSize ?? PagedKVCache.DEFAULT_BLOCK_SIZE;
+  return { key: `paged:${blockSize}`, create() {
+    const caches = model.makeCache();
+    try { maybePageKv(caches, { pagedKv: { blockSize } }, capacityTokens); return caches; }
+    catch (error) { disposeResources(caches); throw error; }
+  } };
+}

package/src/backends/mlx/rollback.ts

@@ -0,0 +1,89 @@
+import type { SpeculativeTransaction } from "../../inference/rollback";
+import type { Cache } from "../../model/gemma4-base";
+
+interface TransactionLayer<Acceptance> {
+  fits(drafts: number): boolean;
+  begin(): void;
+  resolve(accepted: Acceptance, drafts: number): void;
+}
+
+/** State storage that can retain a different verified prefix for each row.
+ * Both attention and recurrent layouts implement this same transaction. */
+export interface RowSpeculativeState {
+  specRoundBegin(): void;
+  specRoundCommit(): void;
+  specRoundRollback(keep: number | readonly number[]): void;
+}
+
+export function bindRowCacheRollback(
+  caches: readonly RowSpeculativeState[], rows: number,
+): SpeculativeTransaction<readonly number[]> {
+  return bindTransaction(caches.map(cache => ({
+    fits: () => true,
+    begin: () => cache.specRoundBegin(),
+    resolve(accepted: readonly number[], drafts: number) {
+      if (accepted.every(count => count === drafts)) cache.specRoundCommit();
+      else cache.specRoundRollback(accepted.every(count => count === accepted[0])
+        ? accepted[0]! + 1 : accepted.map(count => count + 1));
+    },
+  })), (accepted, drafts) => accepted.length === rows &&
+    accepted.every(count => Number.isSafeInteger(count) && count >= 0 && count <= drafts));
+}
+
+/** Bind the recurrent transaction or the trim implementation once.
+ * Ring-capacity checks still happen at each safe boundary before any writes. */
+export function bindCacheRollback(caches: readonly Cache[]): SpeculativeTransaction {
+  const layers = caches.map((cache) => {
+    const { specRoundBegin: begin, specRoundCommit: commit, specRoundRollback: rollback } = cache;
+    const hasRound = [begin, commit, rollback].some((method) => method !== undefined);
+    if (hasRound && !(typeof begin === "function" && typeof commit === "function" && typeof rollback === "function"))
+      throw new Error("cache exposes an incomplete speculative transaction");
+    return {
+      fits(drafts: number) {
+        if ("maxSize" in cache && typeof cache.maxSize === "number")
+          return cache.offset + drafts + 1 < cache.maxSize;
+        return hasRound || cache.isTrimmable();
+      },
+      begin: () => begin?.call(cache),
+      resolve(accepted: number, drafts: number) {
+        if (accepted < drafts) {
+          if (rollback) rollback.call(cache, accepted + 1);
+          else cache.trim(drafts - accepted);
+        } else commit?.call(cache);
+      },
+    };
+  });
+  return bindTransaction(layers, (accepted, drafts) =>
+    Number.isSafeInteger(accepted) && accepted >= 0 && accepted <= drafts);
+}
+
+/** One transaction lifecycle for scalar and row-wise state bindings. */
+function bindTransaction<Acceptance>(
+  layers: readonly TransactionLayer<Acceptance>[],
+  accepts: (accepted: Acceptance, drafts: number) => boolean,
+): SpeculativeTransaction<Acceptance> {
+  let active: number | undefined;
+  let invalid = false;
+  const validate = (drafts: number) => {
+    if (invalid) throw new Error("speculative state is invalid and must be discarded");
+    if (!Number.isSafeInteger(drafts) || drafts < 0) throw new Error("invalid speculative draft count");
+  };
+  return {
+    canBegin(drafts) { validate(drafts); return active === undefined && layers.every((layer) => layer.fits(drafts)); },
+    begin(drafts) {
+      validate(drafts);
+      if (active !== undefined || !layers.every((layer) => layer.fits(drafts)))
+        throw new Error("speculative state cannot begin this round");
+      active = drafts;
+      try { for (const layer of layers) layer.begin(); }
+      catch (error) { invalid = true; throw error; }
+    },
+    resolve(accepted) {
+      if (invalid) throw new Error("speculative state is invalid and must be discarded");
+      if (active === undefined || !accepts(accepted, active)) throw new Error("invalid speculative acceptance count");
+      try { for (const layer of layers) layer.resolve(accepted, active); }
+      catch (error) { invalid = true; throw error; }
+      active = undefined;
+    },
+  };
+}

package/src/backends/mlx/row-sampling.ts

@@ -0,0 +1,49 @@
+import type { MlxArray } from "../../mlx/array";
+import * as ops from "../../mlx/ops";
+import type { TokenLogprobs } from "../../contracts/generation";
+import { readStepExtras, stepExtrasArrays, disposeStepExtras,
+  type DeviceStepSampler, type StepExtras } from "../../sampler";
+import { disposeResources, cleanupFailure } from "../../engine/resources";
+
+type TokenSink = (token: number, logprobs?: TokenLogprobs) => void | boolean | Promise<void | boolean>;
+
+/** Request-owned sampling and output metadata. Execution groups receive token
+ * tensors and an output callback; they need no logprob or retention policy. */
+export function createRowSampling(
+  sampler: DeviceStepSampler, onToken: TokenSink,
+) {
+  if (!sampler.capturesLogprobs) return {
+    sample: (logits: MlxArray, step: number) => sampler.sample(logits, step).token,
+    plainGreedy: sampler.isPlainGreedy, onToken, dispose: () => sampler.dispose(),
+  };
+
+  // A pipelined group can sample the next token before emitting the previous
+  // one. Keep captures by request-local step, independently of row position.
+  const pending = new Map<number, StepExtras>();
+  let emitted = 0;
+  return {
+    sample(logits: MlxArray, step: number): MlxArray {
+      const result = sampler.sample(logits, step);
+      try {
+        ops.asyncEvalAll([result.token, ...stepExtrasArrays(result.extras)]);
+        if (result.extras) pending.set(step, result.extras);
+        return result.token;
+      } catch (error) {
+        return cleanupFailure(error, () => disposeResources([result.token,
+          { dispose: () => disposeStepExtras(result.extras) }]));
+      }
+    },
+    // The vectorized token-only shortcut bypasses the per-request sampler.
+    plainGreedy: false,
+    onToken(token: number) {
+      const extras = pending.get(emitted) ?? null;
+      pending.delete(emitted++);
+      return onToken(token, readStepExtras(extras));
+    },
+    dispose() {
+      const extras = [...pending.values()];
+      pending.clear();
+      disposeResources([sampler, ...extras.map(value => ({ dispose: () => disposeStepExtras(value) }))]);
+    },
+  };
+}