mlx-bun 0.4.0 → 1.0.0

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

@@ -1,200 +0,0 @@
-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

@@ -1,33 +0,0 @@
-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

@@ -1,72 +0,0 @@
-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

@@ -1,20 +0,0 @@
-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

@@ -1,41 +0,0 @@
-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

@@ -1,89 +0,0 @@
-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

@@ -1,49 +0,0 @@
-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) }))]);
-    },
-  };
-}