mlx-bun 0.4.0 → 1.0.0

package/src/spec/dspark/sample.ts

@@ -1,112 +0,0 @@
-// Sampling support for DSpark speculative decoding at temperature > 0.
-//
-// Lossless speculative SAMPLING (Leviathan/Chen): a drafted token x sampled
-// from draft dist q is accepted w.p. min(1, p(x)/q(x)) where p is the target
-// dist; on rejection x is resampled from the residual norm(relu(p − q)); if the
-// whole block is accepted a bonus token is sampled from p. The emitted stream is
-// then distributed EXACTLY as if the target sampled token-by-token from p —
-// independent of q (q only moves the acceptance rate τ). This is the
-// temperature analog of the greedy longest-prefix gate.
-//
-// `p` and `q` must be the SAME processing applied to target/draft logits — so we
-// reuse the server sampler's exact top-p/top-k/temperature path (sampler.ts).
-// HLG/curve samplers are out of scope here; the caller rejects them.
-
-import { MlxArray } from "../../mlx/array";
-import { Dtype } from "../../mlx/ffi";
-import * as ops from "../../mlx/ops";
-import { applyTopP, applyTopK, toLogprobs } from "../../sampler";
-
-export interface DSparkSampleConfig {
-  /** temperature > 0 enables sampling; 0 ⇒ greedy (handled by the caller). */
-  temperature: number;
-  topP?: number;
-  topK?: number;
-  seed?: number;
-}
-
-const GOLDEN = 0x9e3779b97f4a7c15n;
-const U64 = 0xffffffffffffffffn;
-
-/** Seeded stream of mlx random keys — reproducible draws without sharing global
- *  RNG state (same derivation as sampler.ts stepKey). */
-export class KeyStream {
-  #seed: bigint;
-  #ctr = 0;
-  constructor(seed = 0) { this.#seed = BigInt(seed >>> 0); }
-  next(): MlxArray {
-    const mixed = (this.#seed ^ ((BigInt(this.#ctr++) + 1n) * GOLDEN)) & U64;
-    return ops.randomKey(mixed);
-  }
-}
-
-/**
- * Processed categorical logits for one position: logits [1,V] → masked +
- * temperature-scaled logits [1,V] (the input to randomCategorical;
- * softmax(·) is the sampling distribution). Mirrors makeSampler's top-p/top-k/
- * temperature path exactly. temperature must be > 0.
- */
-export function processLogits(logits: MlxArray, cfg: DSparkSampleConfig): MlxArray {
-  if (cfg.temperature <= 0) throw new Error("processLogits requires temperature > 0");
-  const lp = toLogprobs(logits); // [1,V]
-  let cur = lp;
-  const owned: MlxArray[] = [lp];
-  if (cfg.topP && cfg.topP > 0 && cfg.topP < 1) { cur = applyTopP(cur, cfg.topP); owned.push(cur); }
-  if (cfg.topK && cfg.topK > 0) { cur = applyTopK(cur, cfg.topK); owned.push(cur); }
-  const scaled = ops.mulScalar(cur, 1 / cfg.temperature);
-  for (const a of owned) a.dispose();
-  return scaled;
-}
-
-/** softmax of processed logits → probability vector [1,V]. */
-export function probsOf(scaled: MlxArray): MlxArray {
-  return ops.softmaxAxis(scaled, 1, true);
-}
-
-/** Sample one token id from processed logits [1,V] (∝ softmax). */
-export function sampleToken(scaled: MlxArray, key: MlxArray): number {
-  const t = ops.randomCategorical(scaled, key); // [1] uint32
-  const id = ops.itemUint32(t);
-  t.dispose();
-  return id;
-}
-
-/** Probability mass a [1,V] distribution assigns to token id (host scalar). */
-export function probAtToken(probs: MlxArray, id: number): number {
-  const sl = probs.slice([0, id], [1, id + 1]); // [1,1]
-  const v = sl.toFloat32()[0]!;
-  sl.dispose();
-  return v;
-}
-
-/** Sample one token from the residual norm(relu(p − q)) over [1,V]. Falls back
- *  to sampling from p if the residual is degenerate (mass ≈ 0). */
-export function sampleResidual(p: MlxArray, q: MlxArray, key: MlxArray): number {
-  const diff = ops.sub(p, q);
-  const zero = MlxArray.fromFloat32(new Float32Array([0]), []);
-  const relu = ops.maximum(diff, zero);
-  diff.dispose(); zero.dispose();
-  const mass = ops.sumAxis(relu, 1, true); // [1,1]
-  const massV = mass.toFloat32()[0]!;
-  if (!(massV > 1e-12)) {
-    // degenerate (p ⪯ q everywhere) — sample from p directly.
-    relu.dispose(); mass.dispose();
-    const lp = ops.log(p);
-    const id = sampleToken(lp, key);
-    lp.dispose();
-    return id;
-  }
-  const resid = ops.div(relu, mass); // normalized [1,V]
-  relu.dispose(); mass.dispose();
-  // randomCategorical samples ∝ exp(logits) → feed log(resid).
-  const eps = MlxArray.fromFloat32(new Float32Array([1e-30]), []);
-  const safe = ops.add(resid, eps);
-  resid.dispose(); eps.dispose();
-  const logits = ops.log(safe);
-  safe.dispose();
-  const id = sampleToken(logits, key);
-  logits.dispose();
-  return id;
-}
-
-export { Dtype };

package/src/spec/generate.ts

@@ -1,276 +0,0 @@
-// Speculative decoding loop — port of optiq runtime/spec/runtime.py
-// (γ-draft / verify-in-one-pass / accept-longest-prefix / rollback).
-//
-// Correctness oracle: OPTIQ's spec_generate, NOT stock decode. Assistant-
-// drafter speculation is an optiq feature (mlx-lm has only generic
-// two-model spec, and can't drive a KV-borrowing drafter). Both optiq AND
-// mlx-lm BATCH the verify lm-head (one matmul over the whole γ+1 window,
-// then argmax per position) — so neither is bit-exact to stock token-at-
-// a-time decode; they diverge from it at bf16 knife-edges. We do the same
-// (picksBatched): greedy spec is bit-exact to optiq's spec, and agrees
-// with stock only on tie-free prompts (where batched == per-position).
-// An earlier version verified per-position to stay bit-exact to STOCK —
-// that matched no real oracle and cost an extra γ× read of the lm-head
-// weight per step; superseded. See docs/design/speculative-decoding.md.
-//
-// Limitation (same as the reference): partial-accept rollback requires
-// trimmable caches; RotatingKVCache loses trimability once its ring
-// wraps (offset ≥ sliding_window). We throw, as the reference does.
-
-import { MlxArray } from "../mlx/array";
-import * as ops from "../mlx/ops";
-import { Gemma4Model, type Cache } from "../model/gemma4";
-import type { GemmaAssistantDrafter } from "./drafter";
-
-export interface SpecOptions {
-  gamma?: number;
-  maxTokens?: number;
-  eosTokenIds?: number[];
-}
-
-export interface SpecStats {
-  emitted: number;
-  drafted: number;
-  accepted: number;
-  targetCalls: number;
-  decodeMs: number;
-  prefillMs: number;
-  acceptanceRate: number;
-  decodeTps: number;
-}
-
-export interface SpecResult {
-  tokens: number[];
-  stats: SpecStats;
-}
-
-/** Single greedy pick from one hidden position [1,1,H] (lm-head + argmax).
- *  Used for the NON-speculative prefill token (matches stock there). */
-function pickFromHidden(model: Gemma4Model, hidden: MlxArray, pos: number): number {
-  const H = hidden.shape[2]!;
-  const hSl = hidden.slice([0, pos, 0], [1, pos + 1, H]);
-  const logits = model.logitsFromHidden(hSl);
-  hSl.dispose();
-  const V = logits.shape[2]!;
-  const flat = ops.reshape(logits, [1, V]);
-  logits.dispose();
-  const am = ops.argmaxAxis(flat, -1);
-  flat.dispose();
-  const t = ops.itemUint32(am);
-  am.dispose();
-  return t;
-}
-
-/** optiq-faithful batched verify: ONE lm-head over the whole γ+1 window
- *  (reads the lm-head weight ONCE, as optiq and mlx-lm both do), then
- *  argmax per position. This is what makes greedy spec bit-exact to
- *  optiq's spec_generate. argmax is invariant to the logsumexp shift and
- *  to the monotone final-logit softcap, so raw-logit argmax here picks the
- *  same token optiq does. */
-function picksBatched(model: Gemma4Model, hidden: MlxArray, count: number): number[] {
-  const logits = model.logitsFromHidden(hidden); // [1, L, V] — one matmul
-  const V = logits.shape[2]!;
-  const out: number[] = [];
-  for (let k = 0; k < count; k++) {
-    const sl = logits.slice([0, k, 0], [1, k + 1, V]);
-    const flat = ops.reshape(sl, [1, V]);
-    sl.dispose();
-    const am = ops.argmaxAxis(flat, -1);
-    flat.dispose();
-    out.push(ops.itemUint32(am));
-    am.dispose();
-  }
-  logits.dispose();
-  return out;
-}
-
-/** Donor caches: the target's LAST sliding and LAST full cache owners
- *  (port of optiq kv_view.find_donor_layers). */
-function donorIndices(model: Gemma4Model): { sliding: number; full: number } {
-  let sliding = -1;
-  let full = -1;
-  for (let i = 0; i < model.numDonors; i++) {
-    if (model.layers[i]!.layerType === "sliding_attention") sliding = i;
-    else full = i;
-  }
-  if (sliding < 0 || full < 0) throw new Error("missing donor layer type");
-  return { sliding, full };
-}
-
-export function specGenerate(
-  model: Gemma4Model,
-  drafter: GemmaAssistantDrafter,
-  promptTokens: number[],
-  options: SpecOptions = {},
-): SpecResult {
-  const { gamma = 2, maxTokens = 256, eosTokenIds = model.config.eosTokenIds } = options;
-  if (gamma < 1) throw new Error("gamma must be >= 1");
-
-  const caches = model.makeCache();
-  const donors = donorIndices(model);
-  const stats: SpecStats = {
-    emitted: 0, drafted: 0, accepted: 0, targetCalls: 0,
-    decodeMs: 0, prefillMs: 0, acceptanceRate: 0, decodeTps: 0,
-  };
-  const out: number[] = [];
-
-  const embedScaled = (token: number): MlxArray => {
-    const ids = ops.fromInt32([token], [1, 1]);
-    const e = model.embed.encode(ids);
-    ids.dispose();
-    const s = ops.mulScalar(e, model.embedScale);
-    e.dispose();
-    return s;
-  };
-
-  const readDonors = () => ({
-    sliding: (caches[donors.sliding] as never as { temporalView(): [MlxArray, MlxArray] }).temporalView(),
-    full: (caches[donors.full] as never as { temporalView(): [MlxArray, MlxArray] }).temporalView(),
-  });
-
-  // Round/prefill scratch tensors, hoisted so the finally frees them on a
-  // mid-round throw (the 2f not-trimmable throw, or any mlx op) — try-body
-  // locals are invisible to the finally. Each is nulled at its normal
-  // disposal, so the finally never double-frees. Same discipline as
-  // src/spec/serve-loop.ts / generate-dflash.ts (2026-07-06 review fix);
-  // see memory [[mlx-inline-slice-leak-pattern]].
-  let prefillH: MlxArray | null = null;
-  let lastHidden: MlxArray | null = null;
-  let vHidden: MlxArray | null = null;
-
-  try {
-    // 1. prefill — FULL-prompt single forward, deliberately. Per-oracle
-    // convention (2026-07-07 re-anchor): mlx-lm's speculative path drains to
-    // len-1 with no step-0 (the SERVE loop, whose oracle it is, now mirrors
-    // that — src/spec/serve-loop.ts); THIS loop's oracle is optiq
-    // spec_generate, which prefills the whole prompt in one forward and
-    // samples token0 from it (runtime/spec/runtime.py:162-170, read from the
-    // installed oracle venv) — the 2026-06-14 bit-exact gate was established
-    // against exactly that shape. Do not tail-split here.
-    const t0 = performance.now();
-    const ids = ops.fromInt32(promptTokens, [1, promptTokens.length]);
-    prefillH = model.forwardHidden(ids, caches);
-    ids.dispose();
-    stats.targetCalls++;
-    const H = prefillH.shape[2]!;
-    const Lp = prefillH.shape[1]!;
-    let next = pickFromHidden(model, prefillH, Lp - 1);
-    lastHidden = prefillH.slice([0, Lp - 1, 0], [1, Lp, H]);
-    prefillH.dispose();
-    prefillH = null;
-    stats.prefillMs = performance.now() - t0;
-
-    // EOS convention: an EOS id stops generation and counts toward
-    // stats.emitted (reference-faithful: optiq runtime.py counts it) but
-    // is NOT part of the returned content — matching our generate(),
-    // which never yields EOS. optiq's spec yields it only as a stream
-    // EVENT (clients see the stop); its token array role here is content.
-    const nextIsEos = eosTokenIds.includes(next);
-    if (!nextIsEos) out.push(next);
-    stats.emitted++;
-    if (nextIsEos) return { tokens: out, stats };
-
-    // 2. outer loop
-    const tDecode = performance.now();
-    outer: while (stats.emitted < maxTokens) {
-      const position = caches[0]!.offset - 1;
-      const shared = readDonors();
-
-      // 2a. draft γ tokens against the frozen donor views. try/finally so a
-      // mid-draft throw still frees the chained hiddens + donor views + the
-      // in-flight emb (assistant-source.ts draft() discipline).
-      const drafts: number[] = [];
-      let dTok = next;
-      let dHid = lastHidden!; // borrowed for k=0
-      const ownedHiddens: MlxArray[] = [];
-      let emb: MlxArray | null = null;
-      try {
-        for (let k = 0; k < gamma; k++) {
-          emb = embedScaled(dTok);
-          const step = drafter.forward(emb, dHid, shared, position + k);
-          emb.dispose();
-          emb = null;
-          drafts.push(step.token);
-          stats.drafted++;
-          ownedHiddens.push(step.nextHidden);
-          dTok = step.token;
-          dHid = step.nextHidden;
-        }
-      } finally {
-        emb?.dispose();
-        for (const a of ownedHiddens) a.dispose();
-        for (const [k, v] of [shared.sliding, shared.full]) {
-          k.dispose();
-          v.dispose();
-        }
-      }
-
-      // 2b. verify all γ drafts (+ the pending token) in one forward
-      const vIds = ops.fromInt32([next, ...drafts], [1, gamma + 1]);
-      vHidden = model.forwardHidden(vIds, caches);
-      vIds.dispose();
-      stats.targetCalls++;
-
-      // 2c. accept the longest matching prefix. Batched verify lm-head
-      // (optiq-faithful — see picksBatched).
-      const gt = picksBatched(model, vHidden, gamma + 1);
-      let kAccept = 0;
-      while (kAccept < gamma && drafts[kAccept] === gt[kAccept]) kAccept++;
-      stats.accepted += kAccept;
-
-      // 2d. emit accepted drafts
-      for (let k = 0; k < kAccept; k++) {
-        const isEos = eosTokenIds.includes(drafts[k]!);
-        if (!isEos) out.push(drafts[k]!);
-        stats.emitted++;
-        if (isEos || stats.emitted >= maxTokens) {
-          vHidden.dispose();
-          vHidden = null;
-          break outer;
-        }
-      }
-
-      // 2e. emit the correction (or bonus when all accepted)
-      const emit = gt[kAccept]!;
-      const emitIsEos = eosTokenIds.includes(emit);
-      if (!emitIsEos) out.push(emit);
-      stats.emitted++;
-
-      // 2f. roll back rejected entries
-      if (kAccept < gamma) {
-        const n = gamma - kAccept;
-        for (const c of caches) {
-          if (!c.isTrimmable())
-            throw new Error(`${c.constructor.name} not trimmable — context exceeded the sliding window during spec decode`);
-          c.trim(n);
-        }
-      }
-
-      if (emitIsEos || stats.emitted >= maxTokens) {
-        vHidden.dispose();
-        vHidden = null;
-        break;
-      }
-
-      // 2g. chain state
-      next = emit;
-      lastHidden!.dispose();
-      lastHidden = vHidden.slice([0, kAccept, 0], [1, kAccept + 1, H]);
-      vHidden.dispose();
-      vHidden = null;
-    }
-    stats.decodeMs = performance.now() - tDecode;
-    lastHidden.dispose();
-    lastHidden = null;
-  } finally {
-    // Survivors of a mid-round throw (each nulled at its normal disposal).
-    prefillH?.dispose();
-    lastHidden?.dispose();
-    vHidden?.dispose();
-    for (const c of caches) c.dispose();
-  }
-
-  stats.acceptanceRate = stats.drafted ? stats.accepted / stats.drafted : 0;
-  stats.decodeTps = stats.decodeMs ? ((stats.emitted - 1) / stats.decodeMs) * 1000 : 0;
-  return { tokens: out, stats };
-}

package/src/spec/glm52-mtp-graph.ts

@@ -1,29 +0,0 @@
-import type { MlxArray } from "../mlx/array";
-import * as ops from "../mlx/ops";
-import type { Glm52DecoderLayer, Glm52Model } from "../model/glm52";
-import type { MLACache } from "../model/glm52-cache";
-import { rmsNormF32Mlx } from "../model/glm52-mla";
-
-/** Native MTP numerical graph. Tokens, anchor hidden and compressed state may
- * have one or several rows; sampling and state retention belong to callers. */
-export class Glm52MtpGraph {
-  constructor(readonly model: Glm52Model, readonly layer: Glm52DecoderLayer) {}
-
-  async forward(ids: MlxArray, hidden: MlxArray, cache: MLACache): Promise<MlxArray> {
-    const { glmConfig: config, weights } = this.model;
-    const prefix = `model.layers.${config.numHiddenLayers}`;
-    using embedded = weights.embedding(ids, "model.embed_tokens.weight", config.vocabSize, config.hiddenSize);
-    using embeddedNorm = rmsNormF32Mlx(embedded, weights.tensor(`${prefix}.enorm.weight`), config.rmsNormEps);
-    using hiddenNorm = rmsNormF32Mlx(hidden, weights.tensor(`${prefix}.hnorm.weight`), config.rmsNormEps);
-    using joined = ops.concatAxis([embeddedNorm, hiddenNorm], 2);
-    using projected = weights.linear(joined, `${prefix}.eh_proj.weight`, config.hiddenSize, 2 * config.hiddenSize);
-    return await this.layer.forwardAsync(projected, cache, null);
-  }
-
-  project(hidden: MlxArray): MlxArray {
-    const { glmConfig: config, weights } = this.model;
-    using normalized = rmsNormF32Mlx(hidden,
-      weights.tensor(`model.layers.${config.numHiddenLayers}.shared_head.norm.weight`), config.rmsNormEps);
-    return this.model.logitsFromHidden(normalized);
-  }
-}

package/src/spec/glm52-mtp-rows.ts

@@ -1,132 +0,0 @@
-import type { MlxArray } from "../mlx/array";
-import * as ops from "../mlx/ops";
-import { Dtype } from "../mlx/ffi";
-import { materializeCopy } from "../mlx/materialize";
-import { MLACache } from "../model/glm52-cache";
-import { toLogprobs } from "../sampler";
-import { applyStateChanges, disposeResources } from "../engine/resources";
-import type { PreparedStateChange } from "../contracts/resources";
-import type { DraftRowSampling } from "./source";
-import type { Glm52MtpGraph } from "./glm52-mtp-graph";
-
-export interface Glm52MtpRowState {
-  readonly cache: MLACache;
-  readonly hidden: MlxArray;
-  readonly processedTokens: number;
-}
-
-/** Native Colibri draft history starts empty, independently of target prompt
- * length. Accepted draft rows are rebuilt from true target hidden states. */
-export class Glm52MtpRows {
-  #cache: MLACache;
-  #hidden: MlxArray | null = null;
-  #processed: number[] = [];
-  #drafts: number[][] = [];
-  #depth = 0;
-
-  constructor(readonly graph: Glm52MtpGraph, readonly sampling: DraftRowSampling | null) {
-    this.#cache = this.makeCache();
-  }
-  makeCache(): MLACache {
-    const c = this.graph.model.glmConfig;
-    return new MLACache({ kvLoraRank: c.kvLoraRank, ropeHeadDim: c.qkRopeHeadDim,
-      maxTokens: c.maxPositionEmbeddings, role: "mtp" });
-  }
-  get rowCount(): number { return this.#processed.length; }
-  append(states: readonly (Glm52MtpRowState | null)[]): void {
-    applyStateChanges([() => this.prepareAppend(states)]);
-  }
-  prepareAppend(states: readonly (Glm52MtpRowState | null)[]): PreparedStateChange {
-    if (!states.length) return { commit() {}, dispose() {} };
-    let cache: MLACache | null = this.makeCache(), hidden: MlxArray | null = null;
-    const held: { dispose(): void }[] = [];
-    const processed = [...this.#processed, ...states.map(state => state?.processedTokens ?? 0)];
-    try {
-      cache.mergeRows([...(this.rowCount ? [this.#cache] : []), ...states.map(state => {
-        if (state) return state.cache;
-        const empty = this.makeCache(); held.push(empty); return empty;
-      })]);
-      hidden = ops.concatAxis([...(this.#hidden ? [this.#hidden] : []), ...states.map(state => {
-        if (state) return state.hidden;
-        const zero = ops.zeros([1, 1, this.graph.model.glmConfig.hiddenSize], Dtype.float32);
-        held.push(zero); return zero;
-      })], 0);
-    } catch (error) { cache.dispose(); hidden?.dispose(); throw error; }
-    finally { disposeResources(held); }
-    return {
-      commit: () => {
-        const priorCache = this.#cache, priorHidden = this.#hidden;
-        this.#cache = cache!; this.#hidden = hidden; this.#processed = processed;
-        cache = priorCache; hidden = priorHidden;
-      },
-      dispose: () => { cache?.dispose(); hidden?.dispose(); cache = null; hidden = null; },
-    };
-  }
-  prefill(tokens: MlxArray, context: MlxArray): void {
-    const [batch, length] = tokens.shape as [number, number];
-    using tail = context.slice([0, length - 1, 0], [batch, length, context.shape[2]!]);
-    const hidden = materializeCopy(tail);
-    this.#hidden?.dispose(); this.#hidden = hidden;
-    this.#processed = this.#processed.map(count => count + length);
-  }
-  materialize(): void {
-    ops.evalAll([...this.#cache.state(), ...(this.#hidden ? [this.#hidden] : [])]);
-  }
-  filterRows(keep: readonly number[]): void {
-    if (!keep.length) { this.dispose(); return; }
-    using indices = ops.fromInt32([...keep], [keep.length]);
-    const hidden = ops.takeAxis(this.#hidden!, indices, 0);
-    try { this.#cache.filterRows(keep); } catch (error) { hidden.dispose(); throw error; }
-    this.#hidden!.dispose(); this.#hidden = hidden;
-    this.#processed = keep.map(row => this.#processed[row]!);
-  }
-  async draft(pending: readonly number[], depth: number, steps: readonly number[]): Promise<number[][]> {
-    const B = pending.length, tokens: MlxArray[] = [];
-    let ids: MlxArray | null = ops.fromInt32([...pending], [B, 1]), chained: MlxArray | null = null;
-    this.#cache.specRoundBegin();
-    try {
-      for (let position = 0; position < depth; position++) {
-        const hidden = await this.graph.forward(ids!, chained ?? this.#hidden!, this.#cache);
-        chained?.dispose(); chained = hidden;
-        using logits = this.graph.project(hidden);
-        using flat = ops.reshape(logits, [B, logits.shape.at(-1)!]);
-        using logprobs = toLogprobs(flat);
-        const next = this.sampling!.sample(logprobs, steps.map(step => step + position));
-        tokens.push(next); ids!.dispose(); ids = ops.reshape(next, [B, 1]);
-      }
-      using packed = depth ? ops.concatAxis(tokens, 0) : null;
-      const read = packed?.toIntTokens() ?? [];
-      this.#drafts = Array.from({ length: B }, (_, row) =>
-        Array.from({ length: depth }, (_, position) => read[position * B + row]!));
-      this.#depth = depth; return this.#drafts;
-    } finally { ids?.dispose(); chained?.dispose(); disposeResources(tokens); }
-  }
-  async commit(accepted: readonly number[], context: MlxArray): Promise<void> {
-    const B = accepted.length, width = Math.max(...accepted);
-    if (this.#depth) this.#cache.specRoundRollback(1);
-    else this.#cache.specRoundCommit();
-    if (width) {
-      using ids = ops.fromInt32(this.#drafts.flatMap(tokens => Array.from({ length: width }, (_, i) => tokens[i] ?? 0)), [B, width]);
-      using hidden = context.slice([0, 0, 0], [B, width, context.shape[2]!]);
-      this.#cache.specRoundBegin();
-      using output = await this.graph.forward(ids, hidden, this.#cache);
-      output.eval(); this.#cache.specRoundRollback(accepted);
-    }
-    using positions = ops.fromInt32([...accepted], [B, 1, 1]);
-    const hidden = ops.takeAlongAxis(context, positions, 1);
-    this.#hidden!.dispose(); this.#hidden = hidden;
-    this.#processed = this.#processed.map((count, row) => count + 1 + accepted[row]!);
-    this.#drafts = []; this.#depth = 0;
-  }
-  extractRow(row: number): Glm52MtpRowState {
-    const cache = this.#cache.extractRow(row);
-    try {
-      using hidden = this.#hidden!.slice([row, 0, 0], [row + 1, 1, this.#hidden!.shape[2]!]);
-      return { cache, hidden: materializeCopy(hidden), processedTokens: this.#processed[row]! };
-    } catch (error) { cache.dispose(); throw error; }
-  }
-  dispose(): void {
-    this.#cache.dispose(); this.#hidden?.dispose(); this.#hidden = null;
-    this.#processed = []; this.#drafts = []; this.#depth = 0;
-  }
-}