mlx-bun 0.5.0 → 1.0.0

package/src/backends/mlx/speculative.ts

@@ -1,131 +0,0 @@
-import type { MlxArray } from "../../mlx/array";
-import * as ops from "../../mlx/ops";
-import type { Cache } from "../../model/gemma4";
-import type { RuntimeModel } from "../../model/factory";
-import type { Sampler } from "../../sampler";
-import type { DraftProvider, DraftSource } from "../../spec/source";
-import { runtimeConfig, type RuntimeConfig } from "../../runtime-config";
-import type { MlxModelMemory } from "./autoregressive";
-import type { SpeculativeTransaction } from "../../inference/rollback";
-import { bindCacheRollback } from "./rollback";
-import type { GraphDescriptor } from "../../inference/graph";
-import { bindLegacyDraftTarget } from "./draft-target";
-import type { MlxForwardWork } from "./mixed-iteration";
-import type { PrefillPolicy } from "../../inference/prefill";
-import { resolveMlxPrefillPolicy } from "./prefill-policy";
-
-/** Bound target operations and draft construction for one speculative run.
- * A replacement graph supplies this entire port, including any hidden taps.
- * The verifier never needs a RuntimeModel or an artifact-family check. */
-export interface MlxSpeculativeBinding {
-  readonly runtime?: RuntimeConfig;
-  readonly prefillPolicy?: PrefillPolicy;
-  readonly memory?: MlxModelMemory;
-  readonly descriptor: GraphDescriptor;
-  readonly eosTokenIds: readonly number[];
-  readonly prefillTailSplit: boolean;
-  makeCache(): Cache[];
-  openDraft(sampler: Sampler, caches: Cache[]): DraftSource;
-  bindRollback(caches: Cache[]): SpeculativeTransaction;
-  forward(ids: MlxArray, caches: Cache[], tapLayers?: number[], work?: MlxForwardWork):
-    Promise<{ hidden: MlxArray; ctxML: MlxArray | null }>;
-  projectLogits(hidden: MlxArray): MlxArray;
-  /** Establish the implementation's verify kernel context; restore on close. */
-  pinVerify?(): { close(): void };
-}
-
-export type MlxSpeculativeTargetBinding = Omit<MlxSpeculativeBinding, "openDraft">;
-
-/** Legacy mutable tap/kernel fields require the gateway's exclusive lease. */
-export function bindSpeculativeTargetModel(model: RuntimeModel): MlxSpeculativeTargetBinding {
-  const runtime = runtimeConfig();
-  return {
-    runtime,
-    prefillPolicy: resolveMlxPrefillPolicy(model.config, runtime),
-    memory: model,
-    descriptor: Object.freeze({ id: `legacy-spec:${model.config.modelType}`, backend: "mlx",
-      graphAbi: "mlx-hidden-bsh-v1", stateAbi: "legacy-cache-array-v1", artifact: "legacy-resident-model" }),
-    eosTokenIds: model.config.eosTokenIds,
-    prefillTailSplit: runtime.flag("MLX_BUN_PREFILL_TAIL_SPLIT", true),
-    makeCache: model.makeCache.bind(model),
-    bindRollback: bindCacheRollback,
-    forward: (ids, caches, tapLayers, work) => legacyForwardWithTaps(model, ids, caches, tapLayers, work),
-    projectLogits: model.logitsFromHidden.bind(model),
-    ...("setSpecKernelPinned" in model ? {
-      pinVerify() {
-        model.setSpecKernelPinned(true);
-        return { close() { model.setSpecKernelPinned(false); } };
-      },
-    } : {}),
-  };
-}
-
-/** Serial compatibility binds draft construction separately from target work. */
-export function bindLegacySpeculativeModel(model: RuntimeModel, provider: DraftProvider): MlxSpeculativeBinding {
-  return { ...bindSpeculativeTargetModel(model),
-    openDraft: (sampler, caches) => provider.open({ sampler, target: bindLegacyDraftTarget(model, caches) }) };
-}
-
-export function assertMlxSpeculativeBinding(binding: MlxSpeculativeBinding): void {
-  const descriptor = binding.descriptor;
-  if (descriptor.backend !== "mlx" || descriptor.graphAbi !== "mlx-hidden-bsh-v1" ||
-      descriptor.stateAbi !== "legacy-cache-array-v1")
-    throw new Error(`speculative binding ${descriptor.id} has an incompatible backend, graph, or state ABI`);
-}
-
-/** Run one target forward, and when `tapLayers` is set (⟹ a Gemma4 target — a
- *  DSpark-style source), also return the captured multi-layer context
- *  [1,L,m*H] (tapLayers concatenated on the feature axis; index nLayers is the
- *  post-finalNorm sentinel). Non-tapping sources get ctxML=null and never
- *  touch model.hiddenTap. Mirrors generate-dflash.ts forwardTapped. */
-async function legacyForwardWithTaps(
-  model: RuntimeModel,
-  ids: MlxArray,
-  caches: Cache[],
-  tapLayers: number[] | undefined,
-  work?: MlxForwardWork,
-): Promise<{ hidden: MlxArray; ctxML: MlxArray | null }> {
-  if (!tapLayers) {
-    if (work) return { hidden: await work(ids, caches), ctxML: null };
-    const asyncModel = model as RuntimeModel & {
-      forwardHiddenAsync?: (
-        ids: MlxArray,
-        caches: Cache[],
-      ) => Promise<MlxArray>;
-    };
-    const hidden = typeof asyncModel.forwardHiddenAsync === "function"
-      ? await asyncModel.forwardHiddenAsync(ids, caches)
-      : model.forwardHidden(ids, caches);
-    return { hidden, ctxML: null };
-  }
-  if (!("hiddenTap" in model)) throw new Error("target graph does not support hidden taps");
-  const m = model;
-  const previousTap = m.hiddenTap;
-  const cap = new Map<number, MlxArray>();
-  const layers = new Set(tapLayers);
-  if (!work) m.hiddenTap = { layers, captured: cap };
-  let hidden: MlxArray | null = null;
-  try {
-    hidden = work ? await work(ids, caches, { captureLayer(layer, h) {
-      if (layers.has(layer)) cap.set(layer, ops.contiguous(h));
-    } }) : model.forwardHidden(ids, caches);
-    const perLayer = tapLayers.map((li) => {
-      const a = cap.get(li);
-      if (!a) throw new Error(`spec tap: layer ${li} not captured`);
-      return a;
-    });
-    const ctxML = ops.concatAxis(perLayer, 2); // [1,L,m*H]
-    for (const [, a] of cap) a.dispose();
-    cap.clear(); // consumed — the finally must not double-dispose
-    const out = { hidden, ctxML };
-    hidden = null; // ownership returned to the caller
-    return out;
-  } finally {
-    // On any throw (forward mid-capture, missing layer, concat), free the
-    // partially-captured tap tensors and the orphaned hidden. On success both
-    // are already gone (cap cleared, hidden nulled).
-    hidden?.dispose();
-    for (const [, a] of cap) a.dispose();
-    if (!work) m.hiddenTap = previousTap;
-  }
-}

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

@@ -1,104 +0,0 @@
-import type { BatchableCache, Cache } from "../../model/gemma4-base";
-import { applyStateChanges, cleanupFailure, disposeResources } from "../../engine/resources";
-import type { PreparedStateChange } from "../../contracts/resources";
-
-/** Preserve the graph's layout capabilities when an empty sibling is made. */
-export interface MlxRowLayout<Layout extends BatchableCache> extends BatchableCache {
-  makeEmptyBatch(): Layout;
-}
-
-/** Layer layouts own tensor geometry. This collection only applies the same
- * row membership operation to every layer. It has no scheduling, sampling,
- * checkpoint-key or persistence policy. */
-export class MlxStateRows<Layout extends MlxRowLayout<Layout> = BatchableCache> {
-  /** Borrowed by the bound target graph; collection ownership stays here. */
-  readonly caches: Layout[];
-  #rows = 0;
-
-  /** Takes ownership of empty, compatible layer layouts. */
-  constructor(layouts: readonly Layout[]) { this.caches = [...layouts]; }
-  get rowCount(): number { return this.#rows; }
-  /** Retain an immutable batch boundary without extracting or materializing
-   * individual rows. Layer layouts decide how to share their backing arrays. */
-  clone(): MlxStateRows<Layout> {
-    const layouts: Layout[] = [];
-    try {
-      for (const source of this.caches) {
-        const layout = source.makeEmptyBatch(); layouts.push(layout);
-        layout.mergeRows([source]);
-      }
-      const result = new MlxStateRows(layouts); result.#rows = this.#rows;
-      return result;
-    } catch (error) { return cleanupFailure(error, () => disposeResources(layouts)); }
-  }
-  /** Prepare an initial cohort with one merge per layer. */
-  mergeRows(rows: readonly (readonly Cache[])[]): void {
-    const next: Layout[] = [];
-    try {
-      for (let layer = 0; layer < this.caches.length; layer++) {
-        const layout = this.caches[layer]!.makeEmptyBatch(); next.push(layout);
-        layout.mergeRows(rows.map(row => row[layer]!));
-      }
-    } catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
-    const previous = this.caches.splice(0, this.caches.length, ...next);
-    this.#rows = rows.length; disposeResources(previous);
-  }
-
-  /** Borrow a prepared request's state. Build every replacement before
-   * releasing the active group, so a failed admission leaves it intact. */
-  append(row: readonly Cache[]): void {
-    applyStateChanges([() => this.prepareAppend(row)]);
-  }
-
-  /** Stage membership so a method can compose it with companion-state admission. */
-  prepareAppend(row: readonly Cache[]): PreparedStateChange {
-    return this.prepareAppendMany([row]);
-  }
-
-  /** Stage several admissions with one membership operation per layer. */
-  prepareAppendMany(rows: readonly (readonly Cache[])[]): PreparedStateChange {
-    const next: Layout[] = [];
-    try {
-      for (let layer = 0; layer < this.caches.length; layer++) {
-        const layout = this.caches[layer]!;
-        const merged = layout.makeEmptyBatch();
-        next.push(merged);
-        merged.mergeRows([...(this.#rows ? [layout] : []), ...rows.map(row => row[layer]!)]);
-      }
-    } catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
-    return {
-      commit: () => {
-        const previous = this.caches.splice(0, this.caches.length, ...next.splice(0));
-        this.#rows += rows.length;
-        next.push(...previous);
-      },
-      dispose: () => disposeResources(next.splice(0)),
-    };
-  }
-
-  /** The method calls this only between committed steps. A layer failure
-   * invalidates the group; its caller discards the state. */
-  filterRows(keep: readonly number[]): void {
-    if (!keep.length) { this.clear(); return; }
-    for (const layout of this.caches) layout.filterRows(keep);
-    this.#rows = keep.length;
-  }
-
-  /** Returns independent request state for a cache publisher or a new group. */
-  extractRow(row: number): Cache[] {
-    const state: Cache[] = [];
-    try { for (const layout of this.caches) state.push(layout.extractRow(row)); return state; }
-    catch (error) { return cleanupFailure(error, () => disposeResources(state)); }
-  }
-
-  clear(): void {
-    const next: Layout[] = [];
-    try { for (const layout of this.caches) next.push(layout.makeEmptyBatch()); }
-    catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
-    const previous = this.caches.splice(0, this.caches.length, ...next);
-    this.#rows = 0;
-    disposeResources(previous);
-  }
-
-  dispose(): void { this.#rows = 0; disposeResources(this.caches); }
-}

package/src/backends/mlx/state-views.ts

@@ -1,46 +0,0 @@
-import type { StateView } from "../../contracts/resources";
-import { disposeResources, ownResource } from "../../engine/resources";
-import type { MlxArray } from "../../mlx/array";
-import { isBatchableCache, type Cache } from "../../model/gemma4-base";
-import { cloneKvCaches, type CacheCodecProvider } from "../../kv-store";
-
-/** Compatibility boundary for legacy state(). Numerical code consumes a lease,
- * never the legacy ownership marker. No copy, evaluation or fence is added. */
-export function leaseCacheState(cache: Cache): StateView<MlxArray> {
-  return ownResource<readonly MlxArray[]>(cache.state(), cache.stateNeedsDispose ? disposeResources : () => {});
-}
-
-/** Preserve the existing prefill evaluation order: temporary views first,
- * borrowed state second. Partial acquisition releases every acquired view. */
-export function leaseCacheStates(caches: readonly Cache[]): StateView<MlxArray> {
-  const leases: StateView<MlxArray>[] = [];
-  const close = () => disposeResources(leases.map((lease) => ({ dispose: () => lease.close() })));
-  try {
-    for (const owned of [true, false]) {
-      for (const cache of caches) if (!!cache.stateNeedsDispose === owned) leases.push(leaseCacheState(cache));
-    }
-    return ownResource(leases.flatMap((lease) => [...lease.borrow()]), close);
-  } catch (error) {
-    try { close(); } catch (cleanupError) { throw new AggregateError([error, cleanupError], "state view acquisition failed"); }
-    throw error;
-  }
-}
-
-/** Precision/storage transitions constrain reuse independently of scheduling. */
-export function minimumReusableOffset(caches: readonly Pick<Cache, "minimumReusableOffset">[]): number {
-  return Math.max(0, ...caches.map(cache => cache.minimumReusableOffset ?? 0));
-}
-
-/** Snapshot an adopted single-row state into persistence-compatible caches.
- * Row layouts own compact extraction; ordinary serial caches use their codec. */
-export function cloneSingleRowState(caches: readonly Cache[], codecs?: CacheCodecProvider): Cache[] {
-  const owned: Cache[] = [];
-  try {
-    for (const cache of caches) owned.push(...(isBatchableCache(cache)
-      ? [cache.extractRow(0)] : cloneKvCaches([cache], codecs)));
-    return owned;
-  } catch (error) {
-    try { disposeResources(owned); } catch (cleanupError) { throw new AggregateError([error, cleanupError], "single-row snapshot failed"); }
-    throw error;
-  }
-}

package/src/backends/mlx/target-layout-capability.ts

@@ -1,35 +0,0 @@
-import { Glm52Cache } from "../../model/glm52-cache";
-import { SpeculativeRotatingKVCache } from "../../model/speculative-rotating-kv";
-import { RotatingAffineLayout, SpeculativeRotatingAffineLayout } from "../../model/rotating-kv-layout";
-import type { Cache } from "../../model/gemma4-base";
-import { KVCache, QuantizedKVCache, TurboQuantKVCache, RotatingKVCache, RotatingQuantizedKVCache } from "../../model/gemma4-base";
-import { DelayedRotatingQuantizedKVCache } from "../../model/delayed-rotating-quantized-kv";
-import { DelayedQuantizedKVCache } from "../../model/delayed-quantized-kv";
-import { DelayedTurboQuantKVCache } from "../../model/delayed-turboquant-kv";
-import { SSMCache } from "../../model/qwen3-delta";
-import { BatchedSSMCache } from "../../model/batched-ssm";
-import { BatchedKVCache } from "../../model/batched-kv";
-import { BatchedQuantizedKVCache } from "../../model/batched-quantized-kv";
-import { BatchedTurboQuantKVCache } from "../../model/batched-turboquant-kv";
-
-export type TargetRowLayout = Glm52Cache | DelayedRotatingQuantizedKVCache | RotatingAffineLayout | SpeculativeRotatingKVCache | BatchedKVCache | BatchedQuantizedKVCache | BatchedTurboQuantKVCache |
-  BatchedSSMCache | DelayedTurboQuantKVCache | DelayedQuantizedKVCache;
-
-/** The storage binding advertises a row transaction factory. The gateway can
- * compose supported model graphs without identifying their model family. */
-export function targetRowLayoutFactory(cache: Cache): (() => TargetRowLayout) | undefined {
-  if (cache instanceof Glm52Cache) return () => cache.makeEmptyBatch();
-  if (cache instanceof DelayedRotatingQuantizedKVCache)
-    return () => new DelayedRotatingQuantizedKVCache(cache.maxSize, cache.groupSize, cache.bits, cache.start, cache.maintain, undefined, true);
-  if (cache instanceof RotatingAffineLayout || cache instanceof RotatingQuantizedKVCache)
-    return () => new SpeculativeRotatingAffineLayout(cache.maxSize, cache.groupSize, cache.bits);
-  if (cache instanceof RotatingKVCache || cache instanceof SpeculativeRotatingKVCache) return () => new SpeculativeRotatingKVCache(cache.maxSize);
-  if (cache instanceof DelayedTurboQuantKVCache || cache instanceof DelayedQuantizedKVCache ||
-      cache instanceof BatchedTurboQuantKVCache || cache instanceof BatchedKVCache ||
-      cache instanceof BatchedQuantizedKVCache || cache instanceof BatchedSSMCache) return () => cache.makeEmptyBatch();
-  if (cache instanceof TurboQuantKVCache) return () => new BatchedTurboQuantKVCache(cache.kBits, cache.vBits, cache.fusedDecode);
-  if (cache instanceof SSMCache) return () => new BatchedSSMCache();
-  if (cache instanceof QuantizedKVCache) return () => new BatchedQuantizedKVCache(cache.groupSize, cache.bits);
-  if (cache instanceof KVCache) return () => new BatchedKVCache();
-  return undefined;
-}

package/src/backends/mlx/text-engine.ts

@@ -1,73 +0,0 @@
-import { createCompletionMethod } from "../../serve/session-completion-engine";
-import type { ExecutionPlanner, InferenceEngine, Timer } from "../../contracts/generation";
-import { createInferenceEngine, type EngineOptions } from "../../engine/engine";
-import { throwIfCancelled } from "../../engine/cancellation";
-import type { GenerateOptions, GenerateStats } from "../../generate";
-import type { GenerationGateway, RequestShape } from "../../serve/generation-gateway";
-
-/** Data-only API for text requests whose policy/shape was already
- * resolved by the caller. Native request resources still belong to the prepared
- * completion API; they are deliberately absent from this token API. */
-export type TextGenerationOptions = Omit<GenerateOptions,
-  "signal" | "cache" | "grammar" | "fill" | "adapters" | "promptEmbeddings" |
-  "imageMask" | "multimodalMask" | "visionPixels" | "onPrefillDone" |
-  "onDecodeCheckpoint" | "initialPendingToken" | "initialGeneratedTokens" |
-  "originalPromptTokens" | "snapshotAt" | "checkpointEveryTokens"
->;
-
-export interface TextGenerationRequest {
-  readonly promptIds: readonly number[];
-  readonly options: TextGenerationOptions;
-  readonly shape: RequestShape;
-}
-
-export const hostTimer: Timer = {
-  after(milliseconds, callback) {
-    const timer = setTimeout(callback, milliseconds);
-    return () => { clearTimeout(timer); };
-  },
-};
-
-/** Reuses the gateway's placement, exclusive lease, prompt cache, batching,
- * and quant-specific execution. The portable session never imports them.
- * Closing this engine cancels its sessions; it does not dispose a shared gateway. */
-export function createTextInferenceEngine(
-  gateway: Pick<GenerationGateway, "place" | "run">,
-  options: Omit<EngineOptions, "timer"> & { timer?: Timer } = {},
-): InferenceEngine<TextGenerationRequest, GenerateStats> {
-  const planner: ExecutionPlanner<TextGenerationRequest, GenerateStats> = {
-    async plan(request, cancellation) {
-      throwIfCancelled(cancellation);
-      if (request.shape.hasVision || request.shape.hasAdapters || request.shape.hasGrammar)
-        throw new Error("native request resources require the prepared completion API");
-      const placement = gateway.place(request.shape, request.options);
-      const maxTokens = request.options.maxTokens ?? 512;
-      return {
-        id: `text:${placement.mechanism}`, outputTokenLimit: maxTokens,
-        method: createCompletionMethod(gateway, {
-          prompt: [...request.promptIds], options: request.options,
-          shape: request.shape, placement,
-        }),
-      };
-    },
-  };
-  const engine = createInferenceEngine(planner, { ...options, timer: options.timer ?? hostTimer });
-  return {
-    open(request, control) {
-      // Snapshot before demand-start: callers may reuse/mutate request arrays.
-      // TextGenerationOptions contains data only; no native handles or callbacks.
-      const snapshot = structuredClone(request);
-      if (!snapshot.promptIds.length || snapshot.promptIds.some((id) => !Number.isInteger(id) || id < 0 || id > 0xffffffff))
-        return Promise.reject(new Error("prompt must contain valid token ids"));
-      const maxTokens = snapshot.options.maxTokens ?? 512;
-      if (!Number.isSafeInteger(maxTokens) || maxTokens <= 0)
-        return Promise.reject(new Error("text generation requires a positive token limit"));
-      if ((snapshot.options.topLogprobs ?? 0) > (options.maxTopLogprobs ?? 20))
-        return Promise.reject(new Error("top logprobs exceed session delivery capacity"));
-      if (control.output === "collect" && (snapshot.options.logprobs || snapshot.options.topLogprobs))
-        return Promise.reject(new Error("token logprobs require stream consumption"));
-      return engine.open(snapshot, control);
-    },
-    close: () => engine.close(),
-  };
-}

package/src/backends/mlx/transitioning-kv-rows.ts

@@ -1,173 +0,0 @@
-import type { MlxArray } from "../../mlx/array";
-import * as ops from "../../mlx/ops";
-import type { BatchableCache, Cache, Mask } from "../../model/gemma4-base";
-import { KvTensorRows } from "./kv-tensor-rows";
-import { leaseCacheStates, minimumReusableOffset } from "./state-views";
-import { withResource, disposeResources } from "../../engine/resources";
-import { cloneKvCaches } from "../../kv-store";
-
-/** Encoded layouts retain physical row alignment across precision changes. */
-export interface TransitionedKvLayout extends BatchableCache {
-  alignRows?(leftPad: readonly number[]): void;
-}
-export interface SpeculativeTransitionedKvLayout extends TransitionedKvLayout {
-  specRoundBegin(): void;
-  specRoundCommit(): void;
-  specRoundRollback(keep: number | readonly number[]): void;
-}
-export interface KvRowTransition<Layout extends TransitionedKvLayout> {
-  readonly signature: string;
-  readonly conversionOffset?: number;
-  maintain(rows: Cache[]): void;
-  converted(row: Cache): boolean;
-  makeLayout(): Layout;
-  packRows?(layout: Layout, rows: readonly Cache[]): void;
-  prepareRows?(rows: Cache[]): Cache[];
-  extractRow?(row: Cache): Cache;
-  rollbackRow?(row: Cache, before: number, keep: number, preserveColumns: boolean): void;
-}
-
-
-export interface TransitioningKvPositions {
-  readonly rowOffsets: readonly number[];
-  readonly leftPad: readonly number[];
-  readonly offset: number;
-  readonly batchSize: number | null;
-  readonly ropeOffsetArr: MlxArray | undefined;
-  sync(rows: readonly Cache[], preserve: boolean): void;
-  filterRows(keep: readonly number[]): void;
-  makeMask(tokens: number, window: number | null): Mask;
-  dispose(): void;
-}
-export class FullKvPositions extends KvTensorRows implements TransitioningKvPositions {
-  sync(rows: readonly Cache[], preserve: boolean): void {
-    this.mergeRows(preserve
-      ? [{ planes: [], rowOffsets: rows.map(row => row.offset), leftPad: [...this.leftPad] }]
-      : rows.map(row => ({ planes: [], rowOffsets: [row.offset], leftPad: [0] })));
-  }
-}
-
-/** Owns membership, committed positions and precision transitions. Codec and
- * attention adapters own tensor access; scheduling and persistence use the
- * existing batch and checkpoint contracts. Row conversion never examines an
- * uncommitted speculative suffix. */
-export abstract class TransitioningKvRows<Layout extends TransitionedKvLayout> implements BatchableCache {
-  protected readonly positions: TransitioningKvPositions;
-  protected rows: Cache[] = [];
-  protected packed?: Layout;
-  #prefilling = false;
-
-  readonly stateNeedsDispose = true;
-  constructor(readonly transition: KvRowTransition<Layout>, row?: Cache, positions: TransitioningKvPositions = new FullKvPositions()) {
-    this.positions = positions;
-    if (row) { this.rows = transition.prepareRows?.([row]) ?? [row]; this.syncPositions(); }
-  }
-  signature(): string { return this.transition.signature; }
-  get prefillMaintenance() { return this; }
-  beginPrefill(): void { this.#prefilling = true; }
-  commitPrefill(indices: readonly number[]): void {
-    if (this.packed) return;
-    for (const index of indices) {
-      const row = [this.rows[index]!];
-      try { this.transition.maintain(row); }
-      finally { this.rows[index] = row[0]!; }
-    }
-  }
-  endPrefill(): void { this.#prefilling = false; }
-  get minimumReusableOffset(): number { return this.packed?.minimumReusableOffset ?? minimumReusableOffset(this.rows); }
-  get rowOffsets(): readonly number[] { return this.packed?.rowOffsets ?? this.positions.rowOffsets; }
-  get leftPad(): readonly number[] { return this.packed?.leftPad ?? this.positions.leftPad; }
-  get offset(): number { return this.packed?.offset ?? this.positions.offset; }
-  maxAppendTokens(): number {
-    if (this.packed) return this.packed.maxAppendTokens?.() ?? Infinity;
-    const end = this.transition.conversionOffset;
-    if (end === undefined) return Infinity;
-    return Math.min(Infinity, ...this.rowOffsets.map((offset, row) =>
-      offset < end && !this.transition.converted(this.rows[row]!) ? end - offset : Infinity));
-  }
-  get batchSize(): number | null { return this.packed?.batchSize ?? this.positions.batchSize; }
-  get ropeOffsetArr(): MlxArray | undefined { return this.packed ? this.packed.ropeOffsetArr : this.positions.ropeOffsetArr; }
-  bytesPerToken(): number { return this.packed?.bytesPerToken?.() ?? Math.max(0, ...this.rows.map(row => row.bytesPerToken?.() ?? 0)); }
-  projectedBytes(tokens: number): number { return this.bytesPerToken() * tokens; }
-  makeMask(tokens: number, window: number | null): Mask { return (this.packed ?? this.positions).makeMask(tokens, window); }
-  isTrimmable(): boolean { return this.packed?.isTrimmable() ?? this.rows.every(row => row.isTrimmable()); }
-  trim(count: number): void {
-    if (this.packed) this.packed.trim(count);
-    else { for (const row of this.rows) row.trim(count); this.syncPositions(true); }
-  }
-  protected syncPositions(preserve = false): void { this.positions.sync(this.rows, preserve); }
-  protected advance(): void {
-    if (this.packed) return;
-    if (!this.#prefilling) this.transition.maintain(this.rows);
-    if (this.rows.length && this.rows.every(row => this.transition.converted(row))) {
-      const packed = this.transition.makeLayout();
-      try {
-        if (this.transition.packRows) this.transition.packRows(packed, this.rows);
-        else packed.mergeRows(this.rows);
-        packed.alignRows?.(this.positions.leftPad);
-      }
-      catch (error) { packed.dispose(); throw error; }
-      const previous = this.rows; this.rows = []; this.packed = packed;
-      this.positions.dispose(); disposeResources(previous);
-    }
-  }
-  abstract updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray];
-  abstract makeEmptyBatch(): TransitioningKvRows<Layout>;
-  mergeRows(rows: readonly Cache[]): void {
-    const owned: Cache[] = [];
-    try {
-      for (const row of rows) {
-        if (row instanceof TransitioningKvRows) {
-          for (let index = 0; index < row.batchSize!; index++) owned.push(row.extractRow(index));
-        } else owned.push(...cloneKvCaches([row]));
-      }
-    } catch (error) { disposeResources(owned); throw error; }
-    let prepared: Cache[];
-    try { prepared = this.transition.prepareRows?.(owned) ?? owned; }
-    catch (error) { disposeResources(owned); throw error; }
-    this.dispose(); this.rows = prepared; this.syncPositions();
-  }
-  filterRows(keep: readonly number[]): void {
-    if (this.packed) { this.packed.filterRows(keep); return; }
-    const previous = this.rows; this.rows = keep.map(index => previous[index]!);
-    disposeResources(previous.filter((_, index) => !keep.includes(index))); this.positions.filterRows(keep);
-  }
-  extractRow(row: number): Cache { return this.packed?.extractRow(row) ?? (this.transition.extractRow?.(this.rows[row]!) ?? cloneKvCaches([this.rows[row]!])[0]!); }
-  state(): MlxArray[] {
-    return withResource(leaseCacheStates(this.packed ? [this.packed] : this.rows), arrays => {
-      const views: MlxArray[] = [];
-      try { for (const array of arrays) views.push(ops.contiguous(array)); return views; }
-      catch (error) { disposeResources(views); throw error; }
-    });
-  }
-  dispose(): void {
-    this.#prefilling = false;
-    const previous = this.rows; this.rows = [];
-    this.packed?.dispose(); this.packed = undefined; this.positions.dispose(); disposeResources(previous);
-  }
-}
-
-/** Rollback is a method capability, separate from ordinary cache membership. */
-export abstract class SpeculativeTransitioningKvRows<Layout extends SpeculativeTransitionedKvLayout> extends TransitioningKvRows<Layout> {
-  #beforeRound?: number[];
-  specRoundBegin(): void {
-    this.advance(); this.#beforeRound = [...this.rowOffsets]; this.packed?.specRoundBegin();
-  }
-  specRoundCommit(): void {
-    this.packed?.specRoundCommit(); this.#beforeRound = undefined; this.advance();
-  }
-  specRoundRollback(keep: number | readonly number[]): void {
-    if (this.packed) this.packed.specRoundRollback(keep);
-    else {
-      for (const [index, row] of this.rows.entries()) {
-        const count = typeof keep === "number" ? keep : keep[index]!;
-        const before = this.#beforeRound![index]!;
-        if (this.transition.rollbackRow) this.transition.rollbackRow(row, before, count, this.rows.length > 1);
-        else row.trim(row.offset - before - count);
-      }
-      this.syncPositions(true);
-    }
-    this.#beforeRound = undefined; this.advance();
-  }
-  override dispose(): void { this.#beforeRound = undefined; super.dispose(); }
-}