@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/state/checkpoint.ts

@@ -0,0 +1,46 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { Cache } from "../contracts/mlx/cache";
+import { cleanupFailure,disposeResources,withResource } from "../runtime/resources";
+import { leaseCacheStates } from "./leases";
+
+export function attachmentBytes(attachments: readonly CheckpointAttachment[] = []): number {
+  return attachments.reduce((sum, entry) =>
+    sum + entry.tensors.reduce((bytes, tensor) => bytes + tensor.nbytes, 0), 0);
+}
+
+export function disposeAttachments(attachments: readonly CheckpointAttachment[] = []): void {
+  disposeResources(attachments.flatMap((entry) => entry.tensors));
+}
+
+export function cloneAttachments(attachments: readonly CheckpointAttachment[] = []): CheckpointAttachment[] {
+  const held: MlxArray[] = [];
+  try {
+    return attachments.map((entry) => ({
+      schema: entry.schema, metadata: { ...entry.metadata },
+      tensors: entry.tensors.map((tensor) => {
+        const copy = tensor.slice(tensor.shape.map(() => 0), [...tensor.shape]);
+        held.push(copy);
+        return copy;
+      }),
+    }));
+  } catch (error) { return cleanupFailure(error, () => disposeResources(held)); }
+}
+
+/** Shared cache port for MLX execution methods and their companion state. */
+export type MlxPrefixCache = import("../contracts/portable/prefix-cache").PrefixCache<
+  import("../contracts/mlx/cache").Cache[], CheckpointAttachment[]
+>;
+
+/** Resolve snapshot copies before storage retains them. A deferred compact
+ * row still owns its full parent batch; persistence may wait until idle.
+ * This only settles GPU state and performs no serialization or disk I/O. */
+export function materializeCheckpoint(caches: Cache[], attachments: readonly CheckpointAttachment[] = []): void {
+  withResource(leaseCacheStates(caches), state => {
+    const tensors = [...state, ...attachments.flatMap(attachment => attachment.tensors)];
+    if (tensors.length) ops.evalAll(tensors);
+  });
+}
+
+import { CheckpointAttachment } from "../contracts/mlx/checkpoint";
+export { type CheckpointAttachment } from "../contracts/mlx/checkpoint";

package/src/state/continuation-persistence.ts

@@ -0,0 +1,126 @@
+import { SpillQueue } from "./persistence";
+import { type SpillItem } from "./persistence-types";
+import { cacheBytes } from "./prefix-cache";
+import { disposeResources } from "../runtime/resources";
+import type { Cache } from "../contracts/mlx/cache";
+import type { SsdCacheStore } from "./ssd-cache";
+
+/** The saved-state operations generation checkpoints use. */
+export type ContinuationStore = Pick<SsdCacheStore, "findGenerationCheckpoint" | "restore" |
+  "storeGenerationCheckpoint" | "removeGenerationCheckpoints">;
+
+/** A replayable generation's state: caches covering the prompt and the
+ * generated tokens, the token sampled from them but not yet emitted, and the
+ * seed its sampler continues from. */
+export interface ContinuationState {
+  caches: Cache[];
+  cacheTokens: number[];
+  generatedTokens: number;
+  pendingToken: number;
+  seed: number;
+}
+
+type Metadata = Parameters<ContinuationStore["storeGenerationCheckpoint"]>[2];
+/** One request's checkpoint writes under its key. */
+export interface ContinuationAttempt { key: string; completed: boolean; pending: number; failed: boolean; version: number; retired: boolean }
+type Attempt = ContinuationAttempt;
+interface Job { attempt: Attempt; metadata?: Metadata; cleanup?: number }
+
+/** One coordinator per checkpoint store. Queue owns snapshots through disk
+ * completion; per-key attempts order cleanup before subsequent writes. */
+export class ContinuationPersistence {
+  readonly #attempts = new Map<string, Attempt>();
+  readonly #jobs = new WeakMap<SpillItem["caches"], Job>();
+  readonly #queue: SpillQueue;
+  readonly #cleanupFailures = new Map<string, Set<number>>();
+  readonly #cleanupPending = new Map<string, Set<number>>();
+  readonly #cleanupSucceeded = new Map<string, number>();
+  #nextBarrier = 0;
+  constructor(readonly store: ContinuationStore, options: {
+    maxBytes: number;
+    /** Optional owner-thread preparation hook; disk writes use the CPU worker. */
+    runStep?: <T>(step: () => T) => Promise<T>;
+  }) {
+    const runStep = options.runStep ?? (async <T>(step: () => T): Promise<T> => step());
+    this.#queue = new SpillQueue(options.maxBytes, cacheBytes, async item => {
+      const job = this.#jobs.get(item.caches)!;
+      if (job.cleanup) {
+        await runStep(() => store.removeGenerationCheckpoints(job.attempt.key));
+        return true;
+      }
+      const current = () => this.#attempts.get(job.attempt.key) === job.attempt && !job.attempt.completed;
+      if (!current()) return true;
+      const stored = await store.storeGenerationCheckpoint(item.tokens, item.caches, job.metadata!, options.runStep);
+      // No newer queued write can run until this callback returns. An old
+      // in-flight rename can therefore be removed without deleting new work.
+      if (!current()) await runStep(() => store.removeGenerationCheckpoints(job.attempt.key));
+      return stored;
+    }, disposeResources, item => !this.#jobs.get(item.caches)?.cleanup);
+  }
+  /** A completed response is ineligible immediately, before queued unlink. */
+  canRestore(key: string): boolean { return !this.#cleanupPending.has(key) && !this.#cleanupFailures.has(key); }
+  begin(key: string): Attempt {
+    const attempt = { key, completed: false, pending: 0, failed: false, version: 0, retired: false };
+    this.#attempts.set(key, attempt);
+    this.#queue.cancelWhere(item => !this.#jobs.get(item.caches)?.cleanup && this.#jobs.get(item.caches)?.attempt.key === key);
+    return attempt;
+  }
+  /** Takes ownership even when a completed/superseded attempt declines it. */
+  enqueue(attempt: Attempt, state: Omit<ContinuationState, "seed">, metadata: Metadata): void {
+    if (attempt.completed || this.#attempts.get(attempt.key) !== attempt) {
+      disposeResources(state.caches); return;
+    }
+    this.#jobs.set(state.caches, { attempt, metadata: { ...metadata } });
+    attempt.pending++;
+    const version = ++attempt.version;
+    let completion: Promise<boolean>;
+    try { completion = this.#queue.enqueue({ caches: state.caches, tokens: [...state.cacheTokens], ns: attempt.key }); }
+    catch (error) { attempt.pending--; disposeResources(state.caches); throw error; }
+    void completion.then(stored => {
+      attempt.pending--;
+      if (version === attempt.version && !attempt.completed) attempt.failed = !stored;
+      this.#releaseSettled(attempt);
+    });
+  }
+  complete(attempt: Attempt): void {
+    attempt.completed = true;
+    this.#queue.cancelWhere(item => !this.#jobs.get(item.caches)?.cleanup && this.#jobs.get(item.caches)?.attempt === attempt);
+    if (this.#attempts.get(attempt.key) === attempt) {
+      this.#attempts.delete(attempt.key);
+      const barrier = ++this.#nextBarrier;
+      const pending = this.#cleanupPending.get(attempt.key) ?? new Set<number>();
+      pending.add(barrier); this.#cleanupPending.set(attempt.key, pending);
+      const caches: SpillItem["caches"] = [];
+      this.#jobs.set(caches, { attempt, cleanup: barrier });
+      void this.#queue.enqueue({ caches, tokens: [], ns: attempt.key }).then(stored => {
+        const key = attempt.key;
+        const outstanding = this.#cleanupPending.get(key)!;
+        outstanding.delete(barrier);
+        if (!outstanding.size) this.#cleanupPending.delete(key);
+        const failures = this.#cleanupFailures.get(key) ?? new Set<number>();
+        if (stored) {
+          this.#cleanupSucceeded.set(key, Math.max(barrier, this.#cleanupSucceeded.get(key) ?? 0));
+          // A successful unlink covers older failures, never a newer barrier.
+          for (const failed of failures) if (failed <= barrier) failures.delete(failed);
+        } else if (barrier > (this.#cleanupSucceeded.get(key) ?? 0)) failures.add(barrier);
+        if (failures.size) this.#cleanupFailures.set(key, failures);
+        else this.#cleanupFailures.delete(key);
+        if (!outstanding.size) this.#cleanupSucceeded.delete(key);
+      });
+    }
+  }
+  /** Request cancellation/disposal preserves queued and durable state. */
+  release(attempt: Attempt): void { attempt.retired = true; this.#releaseSettled(attempt); }
+  #releaseSettled(attempt: Attempt): void {
+    if (attempt.retired && attempt.pending === 0 && !attempt.failed && this.#attempts.get(attempt.key) === attempt)
+      this.#attempts.delete(attempt.key);
+  }
+  get stats() { return { pendingCount: this.#queue.pendingCount, pendingBytes: this.#queue.pendingBytes,
+    dropped: this.#queue.droppedCount, failed: this.#queue.failedCount }; }
+  async flush(): Promise<{ durable: boolean; pendingBytes: number; failed: number }> {
+    await this.#queue.drain();
+    const active = [...this.#attempts.values()].filter(attempt => !attempt.completed);
+    return { durable: this.#cleanupFailures.size === 0 && active.every(attempt => attempt.pending === 0 && !attempt.failed),
+      pendingBytes: this.#queue.pendingBytes, failed: this.#cleanupFailures.size + active.filter(attempt => attempt.failed).length };
+  }
+}

package/src/state/decode-step-inputs.ts

@@ -0,0 +1,31 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { DecodeStepInputs } from "../contracts/mlx/cache";
+
+/** The per-step arrays a cache assembles its compiled-decode closure inputs from.
+ *  Everything it creates is pushed onto `temps`, which the caller releases after the step. */
+export function decodeStepInputs(temps: MlxArray[]): DecodeStepInputs {
+  const writePositions = new Map<number, MlxArray>();
+  return {
+    activeView(array, length) {
+      const stop = [...array.shape];
+      stop[2] = length;
+      const view = array.slice(stop.map(() => 0), stop);
+      temps.push(view);
+      return view;
+    },
+    own(array) {
+      temps.push(array);
+      return array;
+    },
+    writePosition(position) {
+      let a = writePositions.get(position);
+      if (!a) {
+        a = ops.fromInt32([position], [1]);
+        writePositions.set(position, a);
+        temps.push(a);
+      }
+      return a;
+    },
+  };
+}

package/src/state/decoded-kv-donor.ts

@@ -0,0 +1,12 @@
+import * as ops from "@mlx-bun/mlx/ops";
+import type { KvDonorAttention, KvDonorRows } from "../contracts/mlx/cache";
+
+/** Takes ownership of decoded donor planes. Assistant attention uses decoded
+ * values, preserving its SDPA arithmetic rather than moving inverse rotation. */
+export function decodedKvDonorAttention(view: KvDonorRows): KvDonorAttention {
+  const { keys, values, offsets, starts, ends } = view;
+  return { width: keys.shape[2]!, dtype: keys.dtype, offsets, starts, ends,
+    attend: (q, scale, mask) => ops.sdpa(q, keys, values, scale, mask.mode, mask.arr),
+    dispose() { keys.dispose(); values.dispose(); },
+  };
+}

package/src/state/delayed-quantized-kv.ts

@@ -0,0 +1,141 @@
+import { captureFullKvDonorAttention } from "./full-kv-row-donor";
+import { appendFullKvRows } from "./full-kv-row-append";
+import { FullPrefillRow, fullRowPadding, fullRowPhysicalLength } from "./full-prefill-row";
+import { captureKvAttention } from "./kv-attention-view";
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+import { FullTransitioningKvRows } from "./full-transitioning-kv-rows";
+import { BatchedQuantizedKVCache } from "./batched-quantized-kv";
+import { KVCache, type AffineKernelSet } from "./kv";
+import { QuantizedKVCache } from "./quantized-kv";
+import { Bf16FirstQuantizedKVCache, convertToAffine, transitionDue } from "./bf16-first-kv";
+import { type AttentionRead, type Cache, type CommittedAttentionCache, type Mask, type KvAttentionView } from "../contracts/mlx/cache";
+import { kvStorage, type KvAttentionState } from "./kv-storage";
+
+/** The bf16 source of a full-attention row: a `KVCache`, or a padded
+ * `FullPrefillRow` over one. Converted rows have none. */
+function affineSource(row: Cache): KVCache | FullPrefillRow | undefined {
+  return row instanceof FullPrefillRow ? row.affineConversion : row instanceof KVCache ? row : undefined;
+}
+
+/** The affine transition of `Bf16FirstQuantizedKVCache` over full-attention
+ * rows: each row converts with its own `toQuantized` once it holds `start`
+ * positions, and attends with `kernelSet` after. */
+function settleAffineRows(rows: Cache[], groupSize: number, bits: number, start: number, kernelSet: AffineKernelSet): void {
+  for (let index = 0; index < rows.length; index++) {
+    const source = affineSource(rows[index]!);
+    if (source && transitionDue(source.offset, start)) rows[index] = convertToAffine<Cache>(source, groupSize, bits, start, kernelSet);
+  }
+}
+
+/** The rows of the affine bf16-first lego (`Bf16FirstQuantizedKVCache`) for
+ * continuous batching. Each row is bf16 until it holds `start` positions and
+ * converts by the same transition; rows cross independently, a converted row
+ * keeps its native quantized-attention arithmetic, and once all rows convert,
+ * their packed planes use the existing batched attention implementation. When
+ * a row converts is described on `TransitioningKvRows`. The reads are the
+ * graph's call today: this layout's mask before the append, then its attention
+ * view. Converted rows, the packed layout and every view attend with
+ * `kernelSet`, taken when the layout is built. A row this layout hands back
+ * to serial use is the lego it came from again (`compose`), or its storage when
+ * it came as storage. */
+export class DelayedQuantizedKVCache extends FullTransitioningKvRows<BatchedQuantizedKVCache> implements KvAttentionState, CommittedAttentionCache {
+  constructor(readonly groupSize: number, readonly bits: number, readonly start: number,
+    /** The affine kernel set converted rows and every read attend with, named
+     * by the composition (`composedAffineKernels`). */
+    readonly kernelSet: AffineKernelSet, row?: Cache,
+    /** How the lego these rows came from composes the rows this layout hands
+     * back to serial use (`ComposedStorage.over`). */
+    readonly compose?: (storage: Cache) => Cache) {
+    super({ signature: `kv:delayed-quant:${bits}:${groupSize}:${start}`, conversionOffset: start,
+      settle: rows => settleAffineRows(rows, groupSize, bits, start, kernelSet),
+      converted: row => row instanceof QuantizedKVCache,
+      makeLayout: () => new BatchedQuantizedKVCache(groupSize, bits, kernelSet(bits, groupSize, Dtype.bfloat16)),
+      lone: storage => new Bf16FirstQuantizedKVCache(groupSize, bits, start, kernelSet, undefined, storage),
+      ...(compose ? { compose } : {}) }, row);
+  }
+  get attentionState(): KvAttentionState { return this; }
+  captureDonorAttention() {
+    return this.packed?.captureDonorAttention() ?? captureFullKvDonorAttention(this.rows, this.leftPad, this.offset);
+  }
+  makeEmptyBatch(): DelayedQuantizedKVCache {
+    return new DelayedQuantizedKVCache(this.groupSize, this.bits, this.start, this.kernelSet, undefined, this.compose);
+  }
+
+  appendDecode(k: MlxArray, v: MlxArray): AttentionRead { return this.#read(k, v, false); }
+  appendWindow(k: MlxArray, v: MlxArray): AttentionRead { return this.#read(k, v, false); }
+  /** The committed read (Qwen 3.5's `appendKnown` under `knownSpans`): the same
+   * mask and attention view, attended with independent positions, so a
+   * converted row keeps the single-position reduction order. */
+  appendCommitted(k: MlxArray, v: MlxArray): AttentionRead { return this.#read(k, v, true); }
+  /** The read: the mask from this layout before the append (no
+   * sliding window), then the attention view the append returns. */
+  #read(k: MlxArray, v: MlxArray, independentPositions: boolean): AttentionRead {
+    const mask = this.makeMask(k.shape[2]!, null);
+    let view: KvAttentionView;
+    try { view = this.appendAndFetch(k, v); } catch (error) { mask.arr?.dispose(); throw error; }
+    return { attend: (q, scale) => view.attend(q, scale, mask, independentPositions),
+      dispose() { view.dispose(); mask.arr?.dispose(); } };
+  }
+  /** The dense fetch at the model's B (the softcapped read's): after the
+   * transition the append runs, every row's keys and values as its storage
+   * fetches them dense (a converted row's dequantized), joined left-padded as
+   * the plain attention view joins them, or the packed layout's dense fetch
+   * once every row is converted; the caller owns them. */
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
+    this.advance();
+    if (this.packed) return kvStorage(this.packed).updateAndFetch(k, v);
+    return this.#appendPlain(k, v);
+  }
+  #appendPlain(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
+    const [keys, values] = appendFullKvRows(this.rows, this.leftPad, this.offset + k.shape[2]!, k, v);
+    try { this.syncPositions(true); } catch (error) { keys.dispose(); values.dispose(); throw error; }
+    return [keys, values];
+  }
+  appendAndFetch(k: MlxArray, v: MlxArray): KvAttentionView {
+    this.advance();
+    const kernels = this.kernelSet(this.bits, this.groupSize, k.dtype);
+    if (this.packed) return captureKvAttention(this.packed, k, v, kernels);
+    if (this.rows.every(row => !kvStorage(row).quantizedAttention)) {
+      const [keys, values] = this.#appendPlain(k, v);
+      return { attend: (q, scale, mask) => ops.sdpa(q, keys, values, scale, mask.mode, mask.arr),
+        dispose() { keys.dispose(); values.dispose(); } };
+    }
+    const views: KvAttentionView[] = [], pads = this.rows.map((row, index) => this.leftPad[index]! - fullRowPadding(row));
+    const lengths = this.rows.map(cache => fullRowPhysicalLength(cache) + k.shape[2]!);
+    try {
+      for (const [row, cache] of this.rows.entries()) {
+        const slice = (a: MlxArray) => a.slice([row, 0, 0, 0], [row + 1, a.shape[1]!, a.shape[2]!, a.shape[3]!]);
+        using kr = slice(k), vr = slice(v);
+        views.push(captureKvAttention(cache, kr, vr, kernels));
+      }
+      this.syncPositions(true);
+    } catch (error) { for (const view of views) view.dispose(); throw error; }
+    return {
+      attend(q, scale, mask, independentPositions) {
+        const outputs: MlxArray[] = [];
+        try {
+          for (const [row, view] of views.entries()) {
+            using qr = q.slice([row, 0, 0, 0], [row + 1, q.shape[1]!, q.shape[2]!, q.shape[3]!]);
+            // Capture alignment with the KV tensors, never consult live rows.
+            const start = mask.arr?.shape.map(() => 0), end = mask.arr ? [...mask.arr.shape] : undefined;
+            if (start && end) {
+              if (end.length === 4 && end[0]! > 1) { start[0] = row; end[0] = row + 1; }
+              start[start.length - 1] = pads[row]!;
+              end[end.length - 1] = pads[row]! + lengths[row]!;
+            }
+            using arr = mask.arr ? mask.arr.slice(start!, end!) : null;
+            outputs.push(view.attend(qr, scale, { mode: mask.mode, arr }, independentPositions));
+          }
+          return outputs.length === 1 ? ops.contiguous(outputs[0]!) : ops.concatAxis(outputs, 0);
+        } finally { for (const output of outputs) output.dispose(); }
+      },
+      dispose() { for (const view of views) view.dispose(); },
+    };
+  }
+  updateAndAttend(q: MlxArray, k: MlxArray, v: MlxArray, scale: number, mask: Mask): MlxArray {
+    const view = this.appendAndFetch(k, v);
+    try { return view.attend(q, scale, mask); } finally { view.dispose(); }
+  }
+}

package/src/state/delayed-rotating-quantized-kv.ts

@@ -0,0 +1,178 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+import { RotatingKVCache } from "./rotating-kv";
+import { type AttentionRead, type Cache, type CommittedAttentionCache, type KvAttentionView, type PaddedPrefillCache, type PrefillPadding } from "../contracts/mlx/cache";
+import { kvStorage, type KvAttentionState } from "./kv-storage";
+import { SpeculativeTransitioningKvRows } from "./transitioning-kv-rows";
+import { AlignedRotatingCache, alignRotatingRows, RotatingAffineLayout, SpeculativeRotatingAffineLayout, RotatingKvPositions } from "./rotating-kv-layout";
+import { BatchedRotatingQuantCache } from "./batched-rotating-quant";
+import { captureKvAttention, combineKvDonorAttention } from "./kv-attention-view";
+import { Bf16FirstRotatingQuantizedKVCache, convertToAffine, transitionDue } from "./bf16-first-kv";
+import type { AffineKernelSet } from "./kv";
+import type { DecodeStepPlan } from "../contracts/mlx/cache";
+
+/** The affine transition of `Bf16FirstRotatingQuantizedKVCache` over aligned
+ * ring rows: a plain ring converts its physical columns with its own
+ * `toQuantized` once it holds `start` positions, and attends with `kernelSet`
+ * after. */
+function settleAffineRings(rows: Cache[], groupSize: number, bits: number, start: number, kernelSet: AffineKernelSet): void {
+  for (let index = 0; index < rows.length; index++) {
+    const source = (rows[index] as AlignedRotatingCache).affineConversion;
+    if (source && transitionDue(source.offset, start)) rows[index] = convertToAffine(source, groupSize, bits, start, kernelSet);
+  }
+}
+
+/** The rows of the rotating affine bf16-first lego
+ * (`Bf16FirstRotatingQuantizedKVCache`) for continuous batching. Precision
+ * changes preserve each row's physical columns. The model owns queries and
+ * scale; the shared lifecycle owns membership and the transition (when a row
+ * converts is described on `TransitioningKvRows`). A `start` of `Infinity`
+ * never converts: the prefill cohort's layout for plain rings. The reads are
+ * the graph's call today: this layout's mask before the append, under the
+ * ring's window, then its attention view. Converted rings and the packed
+ * layout attend with `kernelSet`, taken when the layout is built; the rows a
+ * cohort extracts carry it into the layouts they join. */
+export class DelayedRotatingQuantizedKVCache extends SpeculativeTransitioningKvRows<RotatingAffineLayout>
+  implements KvAttentionState, PaddedPrefillCache, CommittedAttentionCache {
+  constructor(readonly maxSize: number, readonly groupSize: number, readonly bits: number, readonly start: number,
+    /** The affine kernel set converted rings and every read attend with,
+     * named by the composition (`composedAffineKernels`). */
+    readonly kernelSet: AffineKernelSet, row?: Cache, readonly speculative = false,
+    /** How the lego these rings came from composes the rings this layout
+     * hands back to serial use (`ComposedStorage.over`). */
+    readonly compose?: (storage: Cache) => Cache) {
+    super({ signature: `kv:delayed-rotating-quant:${maxSize}:${bits}:${groupSize}:${start}`, conversionOffset: start,
+      ...(compose ? { compose } : {}),
+      settle: rows => settleAffineRings(rows, groupSize, bits, start, kernelSet),
+      converted: row => row instanceof AlignedRotatingCache && row.inner instanceof BatchedRotatingQuantCache,
+      makeLayout: () => { const kernels = kernelSet(bits, groupSize, Dtype.bfloat16);
+        return speculative ? new SpeculativeRotatingAffineLayout(maxSize, groupSize, bits, kernels) : new RotatingAffineLayout(maxSize, groupSize, bits, kernels); },
+      prepareRows: alignRotatingRows,
+      lone: storage => new Bf16FirstRotatingQuantizedKVCache(maxSize, groupSize, bits, start, kernelSet, undefined, storage),
+      packRows: (layout, rows) => layout.adoptAlignedRows(rows as readonly AlignedRotatingCache[]),
+      extractRow: row => (row as AlignedRotatingCache).extract(speculative ? maxSize : undefined),
+      rollbackRow: (row, before, keep, preserve) => (row as AlignedRotatingCache).rollback(before, keep, preserve) }, row, new RotatingKvPositions(maxSize));
+  }
+  /** A ring's lone row fetches by its ring: concat until the next write lands at
+   * the window, the ring after (`RotatingKVCache.decodePhase`). */
+  protected override decodePhaseOfRows(): DecodeStepPlan["fetch"] {
+    return this.rowOffsets[0]! + 1 < this.maxSize ? "concat" : "ring";
+  }
+  get attentionState(): KvAttentionState { return this; }
+  captureDonorAttention(): import("../contracts/mlx/cache").KvDonorAttention {
+    if (this.packed) return this.packed.captureDonorAttention();
+    const views: import("../contracts/mlx/cache").KvDonorAttention[] = [];
+    try {
+      for (const row of this.rows) views.push((row as AlignedRotatingCache).captureDonorAttention());
+      return combineKvDonorAttention(views);
+    } catch (error) { for (const view of views) view.dispose(); throw error; }
+  }
+  override makeMask(tokens: number, window: number | null) {
+    if (this.speculative && this.leftPad.every(pad => pad === 0) &&
+        this.offset + tokens <= Math.min(this.maxSize, window ?? this.maxSize))
+      return { mode: tokens === 1 ? "" as const : "causal" as const, arr: null };
+    return super.makeMask(tokens, window);
+  }
+
+  appendDecode(k: MlxArray, v: MlxArray): AttentionRead { return this.#read(k, v, false); }
+  appendWindow(k: MlxArray, v: MlxArray): AttentionRead { return this.#read(k, v, false); }
+  /** The committed read (Qwen 3.5's `appendKnown` under `knownSpans`): the same
+   * mask and attention view, attended with independent positions, which the
+   * packed layout's view honours and the per-row views ignore. */
+  appendCommitted(k: MlxArray, v: MlxArray): AttentionRead { return this.#read(k, v, true); }
+  /** The read: the mask from this layout before the append, with
+   * the ring's window, then the attention view the append returns. */
+  #read(k: MlxArray, v: MlxArray, independentPositions: boolean): AttentionRead {
+    const mask = this.makeMask(k.shape[2]!, this.maxSize);
+    let view: KvAttentionView;
+    try { view = this.appendAndFetch(k, v); } catch (error) { mask.arr?.dispose(); throw error; }
+    return { attend: (q, scale) => view.attend(q, scale, mask, independentPositions),
+      dispose() { view.dispose(); mask.arr?.dispose(); } };
+  }
+
+  preparePrefill(padding: PrefillPadding): void {
+    if (!this.batchSize) {
+      const rows = padding.lengths.map(() => new RotatingKVCache(this.maxSize));
+      try { this.mergeRows(rows); } finally { for (const row of rows) row.dispose(); }
+    }
+    if (this.packed) { this.packed.preparePrefill(padding); return; }
+    // Every physical row follows the same block geometry, even when only a
+    // sibling has right padding. The token lengths still belong to each row.
+    const rightPaddedBatch = padding.rightPadding?.some(pad => pad > 0);
+    for (const [row, cache] of this.rows.entries()) (cache as AlignedRotatingCache).inner.preparePrefill({
+      lengths: [padding.lengths[row]!],
+      ...(padding.leftPadding ? { leftPadding: [padding.leftPadding[row]!] } : {}),
+      ...(padding.rightPadding ? { rightPadding: [padding.rightPadding[row]!] } : {}),
+    }, rightPaddedBatch);
+    this.syncPositions(true);
+  }
+  finalizePrefill(): void {
+    if (this.packed) { this.packed.finalizePrefill(); return; }
+    for (const cache of this.rows) (cache as AlignedRotatingCache).inner.finalizePrefill();
+    this.syncPositions(true);
+  }
+  get maxTokens(): number { return this.maxSize; }
+  projectedBytes(tokens: number): number { return this.bytesPerToken() * Math.min(tokens, this.maxSize); }
+  makeEmptyBatch(): DelayedRotatingQuantizedKVCache {
+    return new DelayedRotatingQuantizedKVCache(this.maxSize, this.groupSize, this.bits, this.start, this.kernelSet, undefined,
+      this.speculative, this.compose);
+  }
+  /** The dense fetch (the softcapped read's): after the transition the append
+   * runs, each row's own ring answers as its storage fetches it dense (a
+   * converted ring dequantized; one row keeps its physical columns and phase;
+   * several share one aligned geometry) and rows join along the batch, or the
+   * packed layout's dense fetch once every row is converted; the caller owns them. */
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
+    this.advance();
+    if (this.packed) return kvStorage(this.packed).updateAndFetch(k, v);
+    const rows = this.rows.map(kvStorage);
+    if (rows.length === 1) {
+      const [keys, values] = rows[0]!.updateAndFetch(k, v);
+      try { this.syncPositions(true); } catch (error) { keys.dispose(); values.dispose(); throw error; }
+      return [keys, values];
+    }
+    const parts: MlxArray[][] = [];
+    try {
+      for (const [row, cache] of rows.entries()) {
+        const slice = (a: MlxArray) => a.slice([row, 0, 0, 0], [row + 1, a.shape[1]!, a.shape[2]!, a.shape[3]!]);
+        using kr = slice(k), vr = slice(v);
+        parts.push(cache.updateAndFetch(kr, vr));
+      }
+      this.syncPositions(true);
+      const keys = ops.concatAxis(parts.map(part => part[0]!), 0);
+      try { return [keys, ops.concatAxis(parts.map(part => part[1]!), 0)]; } catch (error) { keys.dispose(); throw error; }
+    } finally { for (const array of parts.flat()) array.dispose(); }
+  }
+  releaseRopeArr(): void { this.packed?.releaseRopeArr(); }
+  appendAndFetch(k: MlxArray, v: MlxArray): KvAttentionView {
+    this.advance();
+    const kernels = this.kernelSet(this.bits, this.groupSize, k.dtype);
+    if (this.packed) return captureKvAttention(this.packed, k, v, kernels);
+    const views: KvAttentionView[] = [];
+    try {
+      for (const [row, cache] of this.rows.entries()) {
+        const slice = (a: MlxArray) => a.slice([row, 0, 0, 0], [row + 1, a.shape[1]!, a.shape[2]!, a.shape[3]!]);
+        using kr = slice(k), vr = slice(v);
+        views.push(captureKvAttention(cache, kr, vr, kernels));
+      }
+      this.syncPositions(true);
+    } catch (error) { for (const view of views) view.dispose(); throw error; }
+    return {
+      attend(q, scale, mask) {
+        const outputs: MlxArray[] = [];
+        try {
+          for (const [row, view] of views.entries()) {
+            using qr = q.slice([row, 0, 0, 0], [row + 1, q.shape[1]!, q.shape[2]!, q.shape[3]!]);
+            const start = mask.arr?.shape.map(() => 0), end = mask.arr ? [...mask.arr.shape] : undefined;
+            if (start && end && end.length === 4 && end[0]! > 1) { start[0] = row; end[0] = row + 1; }
+            using arr = mask.arr ? mask.arr.slice(start!, end!) : null;
+            outputs.push(view.attend(qr, scale, { mode: mask.mode, arr }));
+          }
+          return outputs.length === 1 ? ops.contiguous(outputs[0]!) : ops.concatAxis(outputs, 0);
+        } finally { for (const output of outputs) output.dispose(); }
+      },
+      dispose() { for (const view of views) view.dispose(); },
+    };
+  }
+}