@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/state/quantized-kv.ts

@@ -0,0 +1,369 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { AttentionRead,BlockAttentionCache,CommittedAttentionCache,CompiledDecodeCache,DecodeSlot,DecodeStepInputs,DecodeStepPlan,DecodeTrace,KvDonorAttention,Mask } from "../contracts/mlx/cache";
+import type { QuantizedAttentionState } from "./kv-storage";
+import { createCausalMask } from "../kernels/attention/masks";
+import { dequantizedFetch, disposeTriple, mapTriple } from "./quantized-tensor";
+import { AttentionMasks, unmaskedLease, withLease } from "./attention-read";
+import { affineCausalLease, affineRead, committedAffineRead, unfusedAffineKernels, viewAffineKernel, type AffineKernels, type AffineKernelSet } from "./affine-attention";
+
+
+/** Quantized KV cache — port of mlx-lm QuantizedKVCache: keys/values
+ *  stored as (packed u32, scales, biases) triples, quantized along
+ *  head_dim. Only full-attention layers convert (mlx-lm's rotating-cache
+ *  quantization is NYI upstream; sliding layers are window-capped
+ *  anyway).
+ *
+ *  Built empty, it quantizes from token zero: the first append is read
+ *  quantized, so the first chunk's numbers differ from today's served path
+ *  and mlx-lm's, which read it bf16 (`Bf16FirstQuantizedKVCache`).
+ *
+ *  Reads: each named read appends through `updateAndFetchQuantized` and
+ *  attends the fetched prefix with the kernel `kernels` fixed for it at
+ *  construction: decode with no mask, a window under the bottom-right causal
+ *  matrix. The composition names the set (`composedAffineKernels`: tiled-causal
+ *  for `--kv-quant config` and `--fused-sdpa on`, unfused otherwise). The
+ *  committed read is `quantizedAppendAttention`. This cache has no sliding
+ *  window (graphs compose `RotatingQuantizedKVCache` for sliding layers). */
+export class QuantizedKVCache implements CompiledDecodeCache, CommittedAttentionCache, BlockAttentionCache {
+  minimumReusableOffset?: number;
+  static readonly STEP = 256;
+  /** Set only by compiled-decode trace adapters (see Cache). */
+  readonly ropeOffsetArr?: MlxArray;
+  keys: ops.QuantizedTensor | null = null;
+  values: ops.QuantizedTensor | null = null;
+  offset = 0;
+
+  constructor(
+    readonly groupSize: number,
+    readonly bits: number,
+    /** The kernel of each read. Row layouts made from this cache keep them. */
+    readonly kernels: AffineKernels,
+    /** The per-forward masks of the model this cache belongs to. Row layouts
+     * made from this cache (`state/layout`) share them. */
+    readonly masks = new AttentionMasks(),
+  ) {}
+
+  /** The kernel set affine storage restored at this cache's layer reads with
+   * (`state/persistence`): this cache's own kernels. */
+  get kernelSet(): AffineKernelSet { return () => this.kernels; }
+
+  get quantizedAttention(): QuantizedAttentionState { return this; }
+
+  signature(): string { return `kv:quant:${this.bits}:${this.groupSize}`; }
+
+  appendDecode(k: MlxArray, v: MlxArray): AttentionRead {
+    const [keys, values] = this.updateAndFetchQuantized(k, v);
+    return affineRead(keys, values, unmaskedLease, this.kernels.decode);
+  }
+
+  appendWindow(k: MlxArray, v: MlxArray): AttentionRead {
+    const L = k.shape[2]!;
+    return withLease(affineCausalLease(this.masks, L, this.offset + L), mask => {
+      const [keys, values] = this.updateAndFetchQuantized(k, v);
+      return affineRead(keys, values, mask, this.kernels.window);
+    });
+  }
+
+  appendCommitted(k: MlxArray, v: MlxArray): AttentionRead {
+    const [keys, values] = this.updateAndFetchQuantized(k, v);
+    return committedAffineRead(keys, values, this.groupSize, this.bits);
+  }
+
+  /** A block read (`BlockAttentionCache.readBlock`): the stored positions,
+   * then the block's keys and values quantized as this cache stores them and
+   * not appended, read with no mask by the masked-window kernel, which takes
+   * any query count. */
+  readBlock(k: MlxArray, v: MlxArray): AttentionRead {
+    if (!this.keys || !this.values) throw new Error("cache is empty");
+    const live = (a: MlxArray) => a.slice([0, 0, 0, 0], [a.shape[0]!, a.shape[1]!, this.offset, a.shape[3]!]);
+    const blockK = ops.quantize(k, this.groupSize, this.bits), blockV = ops.quantize(v, this.groupSize, this.bits);
+    const storedK = mapTriple(this.keys, live), storedV = mapTriple(this.values, live);
+    let keys: ops.QuantizedTensor;
+    try {
+      keys = catAxis2(storedK, blockK);
+      try { return affineRead(keys, catAxis2(storedV, blockV), unmaskedLease, this.kernels.maskedWindow); }
+      catch (error) { disposeTriple(keys); throw error; }
+    } finally { for (const t of [blockK, blockV, storedK, storedV]) disposeTriple(t); }
+  }
+
+  bytesPerToken(): number {
+    const steps = this.keys?.packed.shape[2] ?? 0;
+    if (steps === 0 || !this.keys || !this.values) return 0;
+    const tensors = [
+      this.keys.packed, this.keys.scales, this.keys.biases,
+      this.values.packed, this.values.scales, this.values.biases,
+    ];
+    return tensors.reduce((bytes, array) => bytes + array.nbytes, 0) / steps;
+  }
+
+  /** The dense fetch: appends, then the stored window dequantized. */
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
+    return dequantizedFetch(this.updateAndFetchQuantized(k, v), this.groupSize, this.bits);
+  }
+
+  #grow(triple: ops.QuantizedTensor | null, lastDims: number[], B: number, H: number, steps: number, dtype: Dtype): ops.QuantizedTensor {
+    const mk = (dim: number, dt: Dtype) => ops.zeros([B, H, steps, dim], dt);
+    if (!triple)
+      return { packed: mk(lastDims[0]!, Dtype.uint32), scales: mk(lastDims[1]!, dtype), biases: mk(lastDims[2]!, dtype) };
+    const ext = (a: MlxArray, dim: number, dt: Dtype): MlxArray => {
+      const z = mk(dim, dt);
+      const out = ops.concatAxis([a, z], 2);
+      a.dispose();
+      z.dispose();
+      return out;
+    };
+    return {
+      packed: ext(triple.packed, lastDims[0]!, Dtype.uint32),
+      scales: ext(triple.scales, lastDims[1]!, dtype),
+      biases: ext(triple.biases, lastDims[2]!, dtype),
+    };
+  }
+
+  /** Quantize incoming k/v and append; returns quantized views to offset. */
+  updateAndFetchQuantized(k: MlxArray, v: MlxArray): [ops.QuantizedTensor, ops.QuantizedTensor] {
+    const [B, H, L, kD] = k.shape as [number, number, number, number];
+    const vD = v.shape[3]!;
+    const prev = this.offset;
+    const elPerInt = 32 / this.bits;
+
+    if (!this.keys || prev + L > this.keys.packed.shape[2]!) {
+      const newSteps = Math.floor((QuantizedKVCache.STEP + L - 1) / QuantizedKVCache.STEP) * QuantizedKVCache.STEP;
+      if (this.keys && prev % QuantizedKVCache.STEP !== 0) {
+        const trimTo = (t: ops.QuantizedTensor): ops.QuantizedTensor => {
+          const cut = (a: MlxArray): MlxArray => {
+            const [b, h, , d] = a.shape as [number, number, number, number];
+            const s = a.slice([0, 0, 0, 0], [b, h, prev, d]);
+            a.dispose();
+            return s;
+          };
+          return { packed: cut(t.packed), scales: cut(t.scales), biases: cut(t.biases) };
+        };
+        this.keys = trimTo(this.keys);
+        this.values = trimTo(this.values!);
+      }
+      const dtype = k.dtype;
+      this.keys = this.#grow(this.keys, [kD / elPerInt, kD / this.groupSize, kD / this.groupSize], B, H, newSteps, dtype);
+      this.values = this.#grow(this.values, [vD / elPerInt, vD / this.groupSize, vD / this.groupSize], B, H, newSteps, dtype);
+    }
+
+    this.offset += L;
+    const kq = ops.quantize(k, this.groupSize, this.bits);
+    const vq = ops.quantize(v, this.groupSize, this.bits);
+
+    const writeAll = (dst: ops.QuantizedTensor, src: ops.QuantizedTensor): ops.QuantizedTensor => {
+      const w = (d: MlxArray, s: MlxArray): MlxArray => {
+        const [b, h, , dd] = d.shape as [number, number, number, number];
+        const u = ops.sliceUpdate(d, s, [0, 0, prev, 0], [b, h, this.offset, dd]);
+        d.dispose();
+        s.dispose();
+        return u;
+      };
+      return {
+        packed: w(dst.packed, src.packed),
+        scales: w(dst.scales, src.scales),
+        biases: w(dst.biases, src.biases),
+      };
+    };
+    this.keys = writeAll(this.keys!, kq);
+    this.values = writeAll(this.values!, vq);
+
+    const fetch = (t: ops.QuantizedTensor): ops.QuantizedTensor => {
+      const f = (a: MlxArray): MlxArray => {
+        const [b, h, , d] = a.shape as [number, number, number, number];
+        return a.slice([0, 0, 0, 0], [b, h, this.offset, d]);
+      };
+      return { packed: f(t.packed), scales: f(t.scales), biases: f(t.biases) };
+    };
+    return [fetch(this.keys), fetch(this.values)];
+  }
+
+  makeMask(N: number, windowSize: number | null): Mask {
+    if (N === 1) return { mode: "", arr: null };
+    // Same rule as KVCache.makeMask (mlx-lm cache.py:114-125): windowless
+    // multi-token chunks are the string "causal" at any offset. The old
+    // causalEquivalent escape hatch existed precisely because this case
+    // used to ship an array; only windowed continuations materialize now.
+    if (windowSize === null || (this.offset === 0 && N <= windowSize))
+      return { mode: "causal", arr: null };
+    return { mode: "array", arr: createCausalMask(N, this.offset, windowSize) };
+  }
+
+  /** Compiled decode: growth half of updateAndFetchQuantized (L=1),
+   *  without quantize/write (those live in the compiled graph / after). */
+  prepareDecodeStep(): DecodeStepPlan {
+    const prev = this.offset;
+    if (!this.keys || !this.values) throw new Error("compiled decode on an empty cache");
+    if (prev + 1 > this.keys.packed.shape[2]!) {
+      if (prev % QuantizedKVCache.STEP !== 0) {
+        const trimTo = (t: ops.QuantizedTensor): ops.QuantizedTensor => {
+          const cut = (a: MlxArray): MlxArray => {
+            const [b, h, , d] = a.shape as [number, number, number, number];
+            const s = a.slice([0, 0, 0, 0], [b, h, prev, d]);
+            a.dispose();
+            return s;
+          };
+          return { packed: cut(t.packed), scales: cut(t.scales), biases: cut(t.biases) };
+        };
+        this.keys = trimTo(this.keys);
+        this.values = trimTo(this.values);
+      }
+      const [B, H] = this.keys.packed.shape as [number, number, number, number];
+      const elPerInt = 32 / this.bits;
+      const kD = this.keys.packed.shape[3]! * elPerInt;
+      const vD = this.values.packed.shape[3]! * elPerInt;
+      const dtype = this.keys.scales.dtype;
+      this.keys = this.#grow(this.keys, [kD / elPerInt, kD / this.groupSize, kD / this.groupSize], B, H, QuantizedKVCache.STEP, dtype);
+      this.values = this.#grow(this.values, [vD / elPerInt, vD / this.groupSize, vD / this.groupSize], B, H, QuantizedKVCache.STEP, dtype);
+    }
+    return { fetch: "concat", writePos: prev, activeLen: prev };
+  }
+
+  decodePhase(): DecodeStepPlan["fetch"] { return "concat"; }
+
+  decodeSlot(_plan: DecodeStepPlan): DecodeSlot { return quantConcatDecodeSlot(this.offset, this.groupSize, this.bits); }
+
+  decodeInputs(plan: DecodeStepPlan, step: DecodeStepInputs): MlxArray[] {
+    return activeQuantViews(this.keys!, this.values!, plan.activeLen, step);
+  }
+
+  commitDecodeStep(_slot: DecodeSlot, outputs: MlxArray[]): MlxArray[] {
+    return this.writeDecodeStep(outputs);
+  }
+
+  /** Compiled decode: the write half — six sliceUpdates of the already-
+   *  quantized step row (quantize ran in-graph). Takes ownership of the
+   *  rows; returns the updated buffers to async-eval with the step. */
+  writeDecodeStep(rows: MlxArray[]): MlxArray[] {
+    const prev = this.offset;
+    const w = (d: MlxArray, srcRow: MlxArray): MlxArray => {
+      const [b, h, , dd] = d.shape as [number, number, number, number];
+      const u = ops.sliceUpdate(d, srcRow, [0, 0, prev, 0], [b, h, prev + 1, dd]);
+      d.dispose();
+      srcRow.dispose();
+      return u;
+    };
+    this.keys = {
+      packed: w(this.keys!.packed, rows[0]!),
+      scales: w(this.keys!.scales, rows[1]!),
+      biases: w(this.keys!.biases, rows[2]!),
+    };
+    this.values = {
+      packed: w(this.values!.packed, rows[3]!),
+      scales: w(this.values!.scales, rows[4]!),
+      biases: w(this.values!.biases, rows[5]!),
+    };
+    this.offset = prev + 1;
+    return this.state();
+  }
+
+  state(): MlxArray[] {
+    if (!this.keys || !this.values) return [];
+    return [
+      this.keys.packed, this.keys.scales, this.keys.biases,
+      this.values.packed, this.values.scales, this.values.biases,
+    ];
+  }
+
+  isTrimmable(): boolean {
+    return true;
+  }
+
+  trim(n: number): void {
+    this.offset = Math.max(0, this.offset - n);
+  }
+
+  /** Adopt persisted state (takes ownership of the triples' arrays) —
+   *  the quantized twin of KVCache.restoreState (packages/inference/src/state/persistence.ts). */
+  restoreState(keys: ops.QuantizedTensor, values: ops.QuantizedTensor, offset: number): void {
+    this.dispose();
+    this.keys = keys;
+    this.values = values;
+    this.offset = offset;
+  }
+
+  /** Chronological (K, V) triples sliced to offset — the quantized twin of
+   *  KVCache.temporalView (batched merge/extend/filter reads). Caller owns
+   *  the returned views. */
+  captureDonorAttention(): KvDonorAttention {
+    const [keys, values] = this.temporalView();
+    const B = keys.packed.shape[0]!, width = keys.packed.shape[2]!;
+    return { width, dtype: keys.scales.dtype, offsets: Array(B).fill(this.offset),
+      starts: Array(B).fill(0), ends: Array(B).fill(width),
+      attend: (q, scale, mask) => viewAffineKernel(this.kernels, q.shape[2]!, mask)(q, keys, values, scale, mask),
+      dispose() { disposeTriple(keys); disposeTriple(values); },
+    };
+  }
+
+  temporalView(): [ops.QuantizedTensor, ops.QuantizedTensor] {
+    if (!this.keys || !this.values) throw new Error("cache is empty");
+    const cut = (t: ops.QuantizedTensor): ops.QuantizedTensor => {
+      const f = (a: MlxArray): MlxArray => {
+        const [b, h, , d] = a.shape as [number, number, number, number];
+        return a.slice([0, 0, 0, 0], [b, h, this.offset, d]);
+      };
+      return { packed: f(t.packed), scales: f(t.scales), biases: f(t.biases) };
+    };
+    return [cut(this.keys), cut(this.values)];
+  }
+
+  dispose(): void {
+    for (const a of this.state()) a.dispose();
+    this.keys = this.values = null;
+    this.offset = 0;
+    this.masks.clear();
+  }
+}
+
+/** The six active-prefix planes (keys then values) a concat fetch feeds a compiled decode step. */
+export function activeQuantViews(
+  keys: ops.QuantizedTensor, values: ops.QuantizedTensor, length: number, step: DecodeStepInputs,
+): MlxArray[] {
+  return [keys, values].flatMap((t) => [
+    step.activeView(t.packed, length), step.activeView(t.scales, length), step.activeView(t.biases, length),
+  ]);
+}
+
+const catAxis2 = (a: ops.QuantizedTensor, b: ops.QuantizedTensor): ops.QuantizedTensor => ({
+  packed: ops.concatAxis([a.packed, b.packed], 2),
+  scales: ops.concatAxis([a.scales, b.scales], 2),
+  biases: ops.concatAxis([a.biases, b.biases], 2),
+});
+
+/** The growing quantized cache inside a compiled decode trace: quantize in
+ *  graph, fetch = per-component concat, the six quantized row components are
+ *  closure outputs (see TracedConcatKVCache). */
+class TracedConcatQuantizedKVCache extends QuantizedKVCache implements DecodeTrace {
+  override readonly ropeOffsetArr: MlxArray;
+  outs: MlxArray[] = [];
+  constructor(
+    offset: number, groupSize: number, bits: number,
+    readonly activeKq: ops.QuantizedTensor,
+    readonly activeVq: ops.QuantizedTensor,
+    ropeOffsetArr: MlxArray,
+  ) {
+    // A compiled step is one-token decode, unfused in every composition.
+    super(groupSize, bits, unfusedAffineKernels(bits, groupSize, Dtype.bfloat16));
+    this.offset = offset;
+    this.ropeOffsetArr = ropeOffsetArr;
+  }
+
+  override updateAndFetchQuantized(k: MlxArray, v: MlxArray): [ops.QuantizedTensor, ops.QuantizedTensor] {
+    const kq = ops.quantize(k, this.groupSize, this.bits);
+    const vq = ops.quantize(v, this.groupSize, this.bits);
+    this.outs = [kq.packed, kq.scales, kq.biases, vq.packed, vq.scales, vq.biases];
+    return [catAxis2(this.activeKq, kq), catAxis2(this.activeVq, vq)];
+  }
+}
+
+/** The concat-fetch slot of a quantized cache at `offset` (QuantizedKVCache, and RotatingQuantizedKVCache before the window fills). */
+export function quantConcatDecodeSlot(offset: number, groupSize: number, bits: number): DecodeSlot {
+  const triple = (inputs: readonly MlxArray[], at: number): ops.QuantizedTensor =>
+    ({ packed: inputs[at]!, scales: inputs[at + 1]!, biases: inputs[at + 2]! });
+  return {
+    key: `q-cat:${groupSize}:${bits}`, fetch: "concat", inputs: 6, outputs: 6,
+    trace: (inputs, ropeOffset) =>
+      new TracedConcatQuantizedKVCache(offset, groupSize, bits, triple(inputs, 0), triple(inputs, 3), ropeOffset),
+  };
+}

package/src/state/quantized-tensor.ts

@@ -0,0 +1,26 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+
+
+export const mapTriple = (
+  t: ops.QuantizedTensor, f: (a: MlxArray) => MlxArray,
+): ops.QuantizedTensor => ({ packed: f(t.packed), scales: f(t.scales), biases: f(t.biases) });
+
+export const disposeTriple = (t: ops.QuantizedTensor): void => {
+  t.packed.dispose();
+  t.scales.dispose();
+  t.biases.dispose();
+};
+
+/** The dense fetch of affine storage: the keys and values an affine append
+ * fetched (`updateAndFetchQuantized`), each dequantized (to the activation
+ * dtype its scales carry), as a read that attends keys and values as arrays
+ * takes them. Takes ownership of both triples; the caller owns the result. */
+export function dequantizedFetch([keys, values]: [ops.QuantizedTensor, ops.QuantizedTensor],
+  groupSize: number, bits: number): [MlxArray, MlxArray] {
+  try {
+    const dense = (t: ops.QuantizedTensor) => ops.dequantize(t.packed, t.scales, t.biases, { groupSize, bits, mode: "affine" });
+    const k = dense(keys);
+    try { return [k, dense(values)]; } catch (error) { k.dispose(); throw error; }
+  } finally { disposeTriple(keys); disposeTriple(values); }
+}

package/src/state/request-policy.ts

@@ -0,0 +1,48 @@
+import type { KvStorageScheme, PagedKvComposition } from "../contracts/portable/composition";
+import { type Cache } from "../contracts/mlx/cache";
+import { PagedKVCache, type PagedQuantization, type PagedTurboQuant } from "./paged/cache";
+import type { KVCache } from "./kv";
+import { turboQuantFusedDecode } from "./turboquant-codec";
+import { composedStorageOf, isPlainKvCache } from "./capabilities";
+
+/** The affine pages the composition's KV scheme `kv` names for the layer at
+ * cache index `layer`: every layer of a uniform scheme, the listed layers of a
+ * per-layer one; none (bf16 or TurboQuant pages) otherwise. */
+function pageQuantization(kv: KvStorageScheme, layer: number): PagedQuantization | undefined {
+  if (kv.kind === "affine") return { bits: kv.bits, groupSize: kv.groupSize };
+  if (kv.kind !== "affine-layers") return undefined;
+  const entry = kv.layers.find(item => item.layer === layer);
+  return entry ? { bits: entry.bits, groupSize: entry.groupSize } : undefined;
+}
+
+/** Replace fresh full-attention storage before prefill with pages storing what
+ * the composition's scheme `kv` names for each layer (bf16, affine at the
+ * layer's width, TurboQuant read with the layer's composed read), sized for
+ * `capacityTokens` (the pool grows past it). Pages store their format from the
+ * first position. Sliding layers retain their layout; pre-warmed caches retain
+ * their existing storage. Nothing changes without paging. */
+export function maybePageKv(cache: Cache[], kv: KvStorageScheme, paged: PagedKvComposition | null, capacityTokens: number): void {
+  if (!paged) return;
+  if (cache.some((c) => c.offset > 0)) return;
+  for (let i = 0; i < cache.length; i++) {
+    const storage = composedStorageOf(cache[i]!);
+    if (!isPlainKvCache(storage)) continue;
+    const turbo: PagedTurboQuant | undefined = kv.kind === "turbo" ? Object.freeze({ kBits: kv.kBits, vBits: kv.vBits,
+      fusedDecode: turboQuantFusedDecode(), kernels: (storage as KVCache).turboQuantKernels }) : undefined;
+    cache[i]!.dispose(); // fresh (offset 0) — nothing stored yet
+    cache[i] = new PagedKVCache(capacityTokens, paged.blockSize, paged.direct, pageQuantization(kv, i), undefined, turbo);
+  }
+}
+
+/** The prefix namespace of paged state over a method's namespace `base`: the
+ * block size, what the pages store and their read (bf16 and uniform affine
+ * pages keep the namespace they had before per-layer and TurboQuant pages
+ * existed). */
+export function pagedPrefixNamespace(kv: KvStorageScheme, paged: PagedKvComposition, base: string): string {
+  const direct = paged.direct;
+  if (kv.kind === "affine-layers")
+    return JSON.stringify(["paged-v1", paged.blockSize, "layers", kv.layers.map(entry => [entry.layer, entry.bits, entry.groupSize]), direct, base]);
+  if (kv.kind === "turbo") return JSON.stringify(["paged-v1", paged.blockSize, "turbo", kv.kBits, kv.vBits, direct, base]);
+  const pages = pageQuantization(kv, 0);
+  return JSON.stringify(["paged-v1", paged.blockSize, pages?.bits ?? 0, pages?.groupSize ?? 64, direct, base]);
+}

package/src/state/retention-policy.ts

@@ -0,0 +1,37 @@
+/** RAM policy sees reuse metadata, never tensors, storage or scheduling. */
+export interface RetentionCandidate {
+  bytes: number;
+  cost: number;
+  lastUsed: number;
+  uses: number;
+  priority: number;
+}
+export interface RetentionPolicy {
+  readonly name: string;
+  accessed(entry: RetentionCandidate): void;
+  victim<T extends RetentionCandidate>(entries: readonly T[]): T;
+  evicted(entry: RetentionCandidate): void;
+}
+export class LruRetention implements RetentionPolicy {
+  readonly name = "lru";
+  accessed(_entry: RetentionCandidate): void {}
+  victim<T extends RetentionCandidate>(entries: readonly T[]): T {
+    return entries.reduce((a, b) => a.lastUsed < b.lastUsed ? a : b);
+  }
+  evicted(_entry: RetentionCandidate): void {}
+}
+
+/** GreedyDual size/frequency: saved prefill tokens are the cost proxy.
+ * Advancing age on eviction lets old popularity decay without wall timers. */
+export class CostSizeRetention implements RetentionPolicy {
+  readonly name = "cost-size";
+  #age = 0;
+  accessed(entry: RetentionCandidate): void {
+    entry.priority = this.#age + entry.cost * entry.uses / Math.max(1, entry.bytes);
+  }
+  victim<T extends RetentionCandidate>(entries: readonly T[]): T {
+    return entries.reduce((a, b) => a.priority < b.priority ||
+      (a.priority === b.priority && a.lastUsed < b.lastUsed) ? a : b);
+  }
+  evicted(entry: RetentionCandidate): void { this.#age = Math.max(this.#age, entry.priority); }
+}

package/src/state/rollback.ts

@@ -0,0 +1,131 @@
+import type { SpeculativeTransaction } from "../contracts/portable/rollback";
+import type { Cache } from "../contracts/mlx/cache";
+import type { KvStorageScheme } from "../contracts/portable/composition";
+import { isRotatingPlainCache } from "./capabilities";
+
+/** Whether a serial run's caches, composed with the KV scheme `kv`, can drop
+ * the rejected tail of a verify block (`rollbackVerifyBlock`): bf16 storage
+ * that trims, a plain ring, or a recurrence that restores and replays. */
+export function rollsBackVerifyBlock(kv: KvStorageScheme, caches: readonly Cache[]): boolean {
+  return kv.kind === "bf16" && caches.every(cache =>
+    cache.isTrimmable() || isRotatingPlainCache(cache) || typeof cache.specRoundRollback === "function");
+}
+
+/** Drop the rejected tail of a verify block of `block` positions that one
+ * concat write appended, keeping `accepted`: a recurrence restores its
+ * pre-round snapshot and bit-exactly replays the accepted prefix; other storage
+ * trims the rest. A plain ring stops being trimmable once it wraps, but the
+ * block was one concat write that keeps the newest window-1 positions plus the
+ * block in temporal order, so its rollback trim drops the tail physically
+ * (at least the window remains: `accepted` >= 1). */
+export function rollbackVerifyBlock(caches: readonly Cache[], accepted: number, block: number): void {
+  for (const cache of caches) {
+    if (cache.specRoundRollback) cache.specRoundRollback(accepted);
+    else cache.trim(block - accepted, isRotatingPlainCache(cache));
+  }
+}
+
+interface TransactionLayer<Acceptance> {
+  fits(drafts: number): boolean;
+  begin(drafts: number): 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,
+    // A zero-candidate round consumes only the committed pending token.
+    // It cannot reject any input, so retaining rollback state is unnecessary.
+    begin: drafts => { if (drafts > 0) cache.specRoundBegin(); },
+    resolve(accepted: readonly number[], drafts: number) {
+      if (drafts === 0) return;
+      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));
+}
+
+/** A one-position window whose input may itself be rejected: a step built from a sample not yet read. `resolve`
+ * keeps each row's input (1) or drops it (0), restoring that row's state from before the window. */
+export function bindRowStepRollback(
+  caches: readonly RowSpeculativeState[], rows: number,
+): SpeculativeTransaction<readonly number[]> {
+  return bindTransaction(caches.map(cache => ({
+    fits: () => true,
+    begin: () => cache.specRoundBegin(),
+    resolve(kept: readonly number[]) {
+      if (kept.every(count => count === 1)) cache.specRoundCommit();
+      else cache.specRoundRollback(kept.every(count => count === kept[0]) ? kept[0]! : kept);
+    },
+  })), (kept, inputs) => inputs === 1 && kept.length === rows && kept.every(count => count === 0 || count === 1));
+}
+
+/** 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(drafts); }
+      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/state/rotating-kv-donor.ts

@@ -0,0 +1,35 @@
+import type { KvDonorAttention } from "../contracts/mlx/cache";
+import { BatchedRotatingState } from "./batched-rotating-state";
+import { BatchedRotatingQuantCache } from "./batched-rotating-quant";
+import { plainRowStorage, quantizedRowStorage, temporalStorageView } from "./batched-row-storage";
+import { quantizedDonorAttention } from "./kv-attention-view";
+import { decodedKvDonorAttention } from "./decoded-kv-donor";
+import type { RotatingRing } from "./rotating-row-transaction";
+
+/** Capture the same chronological donor rectangle for plain and encoded rings. */
+export function captureRotatingDonorAttention(ring: RotatingRing): KvDonorAttention {
+  const position = new BatchedRotatingState(ring.maxSize, []);
+  position.restore(ring.positionSnapshot);
+  const offsets = position.offsets.map((_, row) => position.validOffset(row));
+  const ends = offsets.map((offset, row) => position.activeLength - (position.offsets[row]! - offset));
+  // Pending prefill keeps a shared physical rectangle, including when mixed
+  // rows are captured separately. Cropping from each row's valid end would
+  // give siblings different widths and can discard an earlier valid window.
+  const from = position.hasPendingPadding ? 0 : Math.max(0, position.activeLength - ring.maxSize);
+  const range = { from, to: position.activeLength };
+  const validity = { offsets,
+    starts: position.leftPad.map((pad, row) => Math.max(0, pad - from, ends[row]! - from - ring.maxSize)),
+    ends: ends.map(end => Math.max(0, end - from)) };
+  if (ring instanceof BatchedRotatingQuantCache) {
+    const keys = temporalStorageView(quantizedRowStorage, ring.keys!, position, range);
+    let values: import("@mlx-bun/mlx/ops").QuantizedTensor;
+    try { values = temporalStorageView(quantizedRowStorage, ring.values!, position, range); }
+    catch (error) { quantizedRowStorage.dispose(keys); throw error; }
+    return quantizedDonorAttention(keys, values, validity, ring.kernels);
+  }
+  const keys = temporalStorageView(plainRowStorage, ring.keys!, position, range);
+  let values: import("@mlx-bun/mlx/array").MlxArray;
+  try { values = temporalStorageView(plainRowStorage, ring.values!, position, range); }
+  catch (error) { keys.dispose(); throw error; }
+  return decodedKvDonorAttention({ keys, values, ...validity });
+}