@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/state/training-cache.ts

@@ -0,0 +1,300 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+import { createCausalMask } from "../kernels/attention/masks";
+import { flashAttention, flashSupported, getTrainingAttn } from "../kernels/attention/flash";
+import { gatedDeltaUpdate } from "../kernels/delta/gated";
+import type { AttentionRead, Cache, GatedDeltaCache, GatedDeltaParameters, Mask, TrainingAttentionCache } from "../contracts/mlx/cache";
+import { ropeBlocks } from "../layers/prefix-layout";
+import { AttentionMasks, causalLease, unmaskedLease, withLease, type MaskLease } from "./attention-read";
+import type { PrefixLayout } from "../contracts/mlx/trainable";
+import { isRecurrentCache } from "./capabilities";
+
+/** The kernel a training cache's reads attend with, fixed when the cache is
+ *  built. `sdpa` is the stock fused SDPA, the kernel mlx-lm's tuner
+ *  differentiates. `flash` is the opt-in O(L)-memory kernel
+ *  (`MLX_BUN_TRAIN_ATTN=flash`): it serves a causal or masked read of a query
+ *  it supports (`flashSupported`: head size, at least 32 positions) as causal
+ *  attention within the cache's sliding window, and the SDPA serves the rest,
+ *  the rule the training forwards applied per call before the caches owned
+ *  their reads. */
+export type TrainingAttention = "sdpa" | "flash";
+
+/** The training attention the trainer selected for its run (`setTrainingAttn`,
+ *  set around a run from `MLX_BUN_TRAIN_ATTN` once the graph's training
+ *  declaration allows it). Read where a forward's training caches are built,
+ *  never inside a read. */
+export function selectedTrainingAttention(): TrainingAttention {
+  return getTrainingAttn() === "flash" ? "flash" : "sdpa";
+}
+
+/** A training cache's read: pass-through views of `k` and `v` (borrowed, as
+ *  every append's are) under `mask`, which the read owns. */
+function trainingRead(k: MlxArray, v: MlxArray, mask: MaskLease, attention: TrainingAttention, window: number | null): AttentionRead {
+  return withLease(mask, held => {
+    const keys = k.slice([0, 0, 0, 0], k.shape);
+    let values: MlxArray;
+    try { values = v.slice([0, 0, 0, 0], v.shape); } catch (error) { keys.dispose(); throw error; }
+    const { mode, arr } = held.mask;
+    const flash = attention === "flash" && (mode === "causal" || mode === "array");
+    return {
+      attend: (q, scale) => flash && flashSupported(q)
+        ? flashAttention(q, keys, values, scale, true, window ?? 0)
+        : ops.sdpa(q, keys, values, scale, mode, arr),
+      dispose() { keys.dispose(); values.dispose(); held.release(); },
+    };
+  });
+}
+
+/** The padded-batch training mask over `L` positions: [B, 1, L, L] bool, true
+ *  where query i of row b may attend key j:
+ *    (j <= i) AND (j < lengths[b]) [AND (i < j + window) under a sliding window]
+ *  createCausalMask(L, 0, window) broadcast over the rows and masked by each
+ *  row's key validity. Caller owns the result. */
+export function paddedTrainingMask(B: number, L: number, lengths: readonly number[], window: number | null): MlxArray {
+  const causal2d = createCausalMask(L, 0, window);
+  const causal4d = ops.reshape(causal2d, [1, 1, L, L]);
+  causal2d.dispose();
+  const keyValidData = new Int32Array(B * L);
+  for (let b = 0; b < B; b++) {
+    const length = lengths[b]!;
+    for (let j = 0; j < L; j++) keyValidData[b * L + j] = j < length ? 1 : 0;
+  }
+  const keyValidI = MlxArray.fromInt32(keyValidData, [B, 1, 1, L]);
+  const zero = MlxArray.fromInt32(new Int32Array([0]), []);
+  const keyValid = ops.less(zero, keyValidI);
+  zero.dispose();
+  keyValidI.dispose();
+  const allow = ops.logicalAnd(causal4d, keyValid);
+  causal4d.dispose();
+  keyValid.dispose();
+  return allow;
+}
+
+/** Stateless cache for the training and scoring forward. Training is always a
+ *  single offset-0 full-sequence pass, so nothing is stored and every append
+ *  reads exactly the keys and values it is handed; the offset stays 0. The
+ *  statelessness is REQUIRED: gradient checkpointing re-runs a layer's closure
+ *  during the backward recompute, and a layer that shares another layer's keys
+ *  and values reads them through that layer's cache again, so an append must
+ *  leave the cache as it found it.
+ *
+ *  Reads: pass-through views of `k` and `v`, attended with `attention`.
+ *  Without `rowLengths`: no mask for decode, the fused causal mask for a
+ *  window, and the offset-0 sliding-window mask once a window is longer than
+ *  `window`. With `rowLengths` (a right-padded batch whose rows hold that many
+ *  real positions), both reads take the padded-batch mask
+ *  (`paddedTrainingMask`) for the cache's window, so no real query sees a
+ *  padded key. A layer with a sliding window gets a cache composed with it; the
+ *  caches of one forward share `masks`, so its layers build each mask once. */
+export class TrainingCache implements TrainingAttentionCache {
+  constructor(
+    /** The layer's sliding window; null for full attention. */
+    readonly window: number | null = null,
+    /** The per-forward masks of the model this cache belongs to. */
+    readonly masks = new AttentionMasks(),
+    /** The kernel this cache's reads attend with. */
+    readonly attention: TrainingAttention = "sdpa",
+    /** Real positions per row of a right-padded batch; null for unpadded rows. */
+    readonly rowLengths: readonly number[] | null = null,
+  ) {}
+  /** Training-only adapter; never admitted, merged, or persisted. */
+  signature(): string { return "train:training"; }
+  offset = 0;
+  /** Uniform positions from the pinned offset. */
+  ropePositions(x: MlxArray, rope: (block: MlxArray, offset: number) => MlxArray): MlxArray { return rope(x, this.offset); }
+  appendDecode(k: MlxArray, v: MlxArray): AttentionRead {
+    return trainingRead(k, v, this.rowLengths ? this.#padded(k.shape[2]!) : unmaskedLease, this.attention, this.window);
+  }
+  appendWindow(k: MlxArray, v: MlxArray): AttentionRead {
+    const L = k.shape[2]!, window = this.window;
+    const mask = this.rowLengths ? this.#padded(L)
+      : window === null || L <= window ? causalLease
+      : this.masks.lease(`training-window:${window}`, `${L}`,
+        () => ({ mode: "array", arr: createCausalMask(L, 0, window) }));
+    return trainingRead(k, v, mask, this.attention, window);
+  }
+  #padded(L: number): MaskLease {
+    const lengths = this.rowLengths!;
+    return this.masks.lease(`training-padded:${this.window}`, `${L}|${lengths.join(",")}`,
+      () => ({ mode: "array", arr: paddedTrainingMask(lengths.length, L, lengths, this.window) }));
+  }
+  /** The dense fetch and mask the softcapped read (`state/softcap-attention.ts`)
+   *  attends through: pass-through views and the mask the window read takes. */
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
+    return [k.slice([0, 0, 0, 0], k.shape), v.slice([0, 0, 0, 0], v.shape)];
+  }
+  makeMask(N: number, windowSize: number | null): Mask {
+    if (this.rowLengths) return { mode: "array", arr: paddedTrainingMask(this.rowLengths.length, N, this.rowLengths, this.window) };
+    if (N === 1) return { mode: "", arr: null };
+    if (windowSize === null || N <= windowSize) return { mode: "causal", arr: null };
+    return { mode: "array", arr: createCausalMask(N, 0, windowSize) };
+  }
+  state(): MlxArray[] {
+    return [];
+  }
+  isTrimmable(): boolean {
+    return true;
+  }
+  trim(_n: number): void {
+    /* offset is pinned at 0 */
+  }
+  dispose(): void {
+    // Owns no arrays; drops the memo's holds on its masks.
+    this.masks.clear();
+  }
+}
+
+/** Stateless DeltaNet stand-in (same re-runnability contract as
+ *  TrainingCache): state WRITES are discarded immediately, so a gradient-
+ *  checkpoint recompute re-enters with the identical null-state conditions;
+ *  reads always yield null (zeros inside the kernel) — a full-sequence
+ *  forward from t=0, exactly the training/perplexity semantics. Before this
+ *  existed, qwen3_5 routed its DeltaNet layers into plain TrainingCache and
+ *  threw on the missing conv/recurrent/advance surface (the recorded
+ *  "mlx-bun perplexity cannot score qwen3_5" gap).
+ *
+ *  Its `recurDecode` and `recurWindow` run the recurrence with a backward
+ *  attached (`gatedDeltaUpdate`'s `differentiable`), since the inference
+ *  kernel has no gradient; the DeltaNet block hands them its projections as it
+ *  does to any recurrent cache. */
+export class TrainingSSMCache implements GatedDeltaCache {
+  /** Training-only adapter; never admitted, merged, or persisted. */
+  signature(): string { return "train:training-ssm"; }
+  offset = 0;
+  specRound: null = null;
+  recurDecode(qkv: MlxArray, a: MlxArray, b: MlxArray, layer: GatedDeltaParameters): MlxArray {
+    return this.#recur(qkv, a, b, layer);
+  }
+  recurWindow(qkv: MlxArray, a: MlxArray, b: MlxArray, layer: GatedDeltaParameters): MlxArray {
+    return this.#recur(qkv, a, b, layer);
+  }
+  /** The training recurrence: a zero conv state, the
+   *  causal convolution, the block's glue, and the differentiable recurrence
+   *  from a zero recurrent state. Nothing is stored, so the new recurrent
+   *  state is disposed at once and no conv tail is kept. */
+  #recur(qkv: MlxArray, a: MlxArray, b: MlxArray, layer: GatedDeltaParameters): MlxArray {
+    const [B, , D] = qkv.shape as [number, number, number];
+    const convState = ops.zeros([B, layer.convWeight.shape[1]! - 1, D], qkv.dtype);
+    const input = ops.concatAxis([convState, qkv], 1);
+    convState.dispose();
+    qkv.dispose();
+    const conv = ops.conv1d(input, layer.convWeight, 1, 0, 1, D);
+    input.dispose();
+    const { q, k, v } = layer.heads(conv);
+    conv.dispose();
+    const [out, state] = gatedDeltaUpdate(q, k, v, a, b, layer.aLog, layer.dtBias, null, null, true);
+    for (const array of [q, k, v, a, b, state]) array.dispose();
+    return out;
+  }
+  advance(_n: number): void { /* offset pinned at 0 */ }
+  rowOffset(_i: number): number { return 0; }
+  state(): MlxArray[] { return []; }
+  isTrimmable(): boolean { return true; }
+  trim(_n: number): void { /* stateless */ }
+  dispose(): void { /* owns no arrays */ }
+}
+
+/** Block-sparse attention mask [T,T] (bool, true = attend) for the prefix-shared
+ *  concat [prompt(P); chosen(Rc); rejected(Rr)]: physical causal AND NOT
+ *  (rejected row -> chosen col). Chosen rows already can't see rejected cols
+ *  (rejected come later), so the only extra cut is rejected->chosen. With a
+ *  sliding `window`, one more AND on LOGICAL positions (`logpos[row] -
+ *  logpos[col] < window`, `logpos[i] = i` for prompt+chosen and `i - Rc` for the
+ *  rejected block, reset to P): a rejected token at physical P+Rc+k has logical
+ *  position P+k, so its window to the prompt tail differs from its physical
+ *  distance by Rc. Caller owns the result. */
+export function blockSparsePrefixMask(layout: PrefixLayout, window: number | null): MlxArray {
+  const { P, Rc, Rr } = layout;
+  const T = P + Rc + Rr;
+  const causal = createCausalMask(T, 0, null); // [T,T] physical causal (a window is applied separately, on logical positions)
+  const idxFlat = ops.arange(0, T, 1, Dtype.int32);
+  const row = ops.reshape(idxFlat, [T, 1]);
+  const col = ops.reshape(idxFlat, [1, T]);
+  const pp = ops.fromInt32([P], []);
+  const pRc = ops.fromInt32([P + Rc], []);
+  // notForbid = NOT(rejRow AND chosenCol) = (i < P+Rc) OR (j < P) OR (j >= P+Rc)
+  const notRejRow = ops.less(row, pRc);
+  const colLtP = ops.less(col, pp);
+  const colGePRc = ops.greaterEqual(col, pRc);
+  const notChosenCol = ops.logicalOr(colLtP, colGePRc);
+  const notForbid = ops.logicalOr(notRejRow, notChosenCol);
+  let allow = ops.logicalAnd(causal, notForbid); // [T,T] bool
+
+  if (window !== null) {
+    const logposArr = new Int32Array(T);
+    for (let i = 0; i < T; i++) logposArr[i] = i < P + Rc ? i : i - Rc;
+    const logpos = MlxArray.fromInt32(logposArr, [T]);
+    const lrow = ops.reshape(logpos, [T, 1]);
+    const lcol = ops.reshape(logpos, [1, T]);
+    const dist = ops.sub(lrow, lcol); // logpos[row] - logpos[col] (>=0 wherever causal allows)
+    const w = ops.fromInt32([window], []);
+    const slidingOK = ops.less(dist, w); // attend only within `window` logical positions
+    const next = ops.logicalAnd(allow, slidingOK);
+    allow.dispose();
+    allow = next;
+    for (const a of [logpos, lrow, lcol, dist, w, slidingOK]) a.dispose();
+  }
+
+  for (const a of [causal, idxFlat, row, col, pp, pRc, notRejRow, colLtP, colGePRc, notChosenCol, notForbid]) a.dispose();
+  return allow;
+}
+
+/** Stateless pass-through cache (offset 0, like TrainingCache) for the
+ *  prefix-shared forward over the concat [prompt(P); chosen(Rc); rejected(Rr)]:
+ *  its window read takes exactly the concat's T positions under the
+ *  block-sparse mask (`blockSparsePrefixMask`) for the cache's sliding window,
+ *  attended with `attention`. The caches of one pass share `masks`, so each
+ *  window's mask is built once per pass, across its layers, segments and
+ *  backward recomputes. A decode read is not defined over a concat; it reads as
+ *  the window does and refuses anything but the whole concat. */
+export class PrefixSharedCache implements TrainingAttentionCache {
+  /** Training-only adapter; never admitted, merged, or persisted. */
+  signature(): string { return "train:prefix-shared"; }
+  offset = 0;
+  constructor(
+    private readonly layout: PrefixLayout,
+    /** The layer's sliding window, on logical positions; null for full attention. */
+    readonly window: number | null = null,
+    /** The per-pass masks of the model this cache belongs to. */
+    readonly masks = new AttentionMasks(),
+    /** The kernel this cache's reads attend with. */
+    readonly attention: TrainingAttention = "sdpa",
+  ) {}
+  /** The concat's block positions: the prompt from 0, each response restarted at P. */
+  ropePositions(x: MlxArray, rope: (block: MlxArray, offset: number) => MlxArray): MlxArray {
+    return ropeBlocks(x, this.layout, rope);
+  }
+  appendDecode(k: MlxArray, v: MlxArray): AttentionRead { return this.appendWindow(k, v); }
+  appendWindow(k: MlxArray, v: MlxArray): AttentionRead {
+    const { P, Rc, Rr } = this.layout, N = k.shape[2]!;
+    if (N !== P + Rc + Rr) throw new Error(`PrefixSharedCache: N=${N} != P+Rc+Rr=${P + Rc + Rr}`);
+    const mask = this.masks.lease(`prefix-shared:${this.window}`, `${P},${Rc},${Rr}`,
+      () => ({ mode: "array", arr: blockSparsePrefixMask(this.layout, this.window) }));
+    return trainingRead(k, v, mask, this.attention, this.window);
+  }
+  state(): MlxArray[] { return []; }
+  isTrimmable(): boolean { return true; }
+  trim(_n: number): void { /* offset pinned at 0 */ }
+  dispose(): void { this.masks.clear(); }
+}
+
+/** Prefix-shared caches for one pass, one per entry of `windows` (each layer's
+ *  sliding window, or null), sharing one set of masks and attending with
+ *  `attention`. */
+export function prefixSharedCaches(layout: PrefixLayout, windows: readonly (number | null)[],
+  attention: TrainingAttention = "sdpa"): Cache[] {
+  const masks = new AttentionMasks();
+  return windows.map(window => new PrefixSharedCache(layout, window, masks, attention));
+}
+
+/** The stateless stand-in a training or scoring forward runs in place of one
+ * of the graph's fresh caches: recurrent state takes `TrainingSSMCache` (fixed
+ * size and causal, so right padding never reaches a real position and needs no
+ * mask); an attention cache takes `attention(window)`, the caller's stand-in
+ * over that cache's sliding window (its ring size; null for full attention). */
+export function trainingStandIn(cache: Cache, attention: (window: number | null) => Cache): Cache {
+  if (isRecurrentCache(cache)) return new TrainingSSMCache();
+  const maxSize = (cache as { maxSize?: number }).maxSize;
+  return attention(typeof maxSize === "number" ? maxSize : null);
+}

package/src/state/transitioning-kv-rows.ts

@@ -0,0 +1,267 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { BatchableCache, Cache, CompiledDecodeCache, DecodeSlot, DecodeStepInputs, DecodeStepPlan, Mask } from "../contracts/mlx/cache";
+import { KvTensorRows } from "./kv-tensor-rows";
+import type { KvStorage } from "./kv-storage";
+import { leaseCacheStates } from "./leases";
+import { minimumReusableOffset } from "./views";
+import { withResource, disposeResources } from "../runtime/resources";
+import { cloneKvCaches } from "./persistence";
+
+/** Encoded layouts retain physical row alignment across precision changes. */
+export interface TransitionedKvLayout extends BatchableCache, KvStorage {
+  alignRows?(leftPad: readonly number[]): void;
+}
+export interface SpeculativeTransitionedKvLayout extends TransitionedKvLayout {
+  specRoundBegin(): void;
+  specRoundCommit(): void;
+  specRoundRollback(keep: number | readonly number[]): void;
+}
+/** One precision transition over a row layout, built with the layout's scheme:
+ * the same conversion the single-sequence legos make (`state/bf16-first-kv`). */
+export interface KvRowTransition<Layout extends TransitionedKvLayout> {
+  readonly signature: string;
+  readonly conversionOffset?: number;
+  /** Converts, in place, every row whose stored length has reached the
+   * conversion offset. Converted rows, and rows below the offset, stay. */
+  settle(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;
+  /** How the lego these rows came from composes storage for serial use
+   * (`ComposedStorage.over`): a row this layout hands back to serial use is
+   * that lego again. Absent for rows that came as storage. */
+  compose?(storage: Cache): Cache;
+  /** The lego of this transition over one row's storage (taking ownership):
+   * what a compiled decode step drives this layout's lone row through. */
+  lone(storage: Cache): Cache & CompiledDecodeCache;
+  /** `row` handed to serial use as it is (ownership moves), or null when the
+   * row is a wrapper whose storage must be extracted. Absent: every row is
+   * extracted. */
+  adopt?(row: Cache): Cache | null;
+}
+
+
+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.
+ *
+ * Rows convert before an append outside prefill, and during prefill only at
+ * each row's committed chunk boundary (`prefillMaintenance.commitPrefill`), so
+ * a sibling that splits a row's planned chunk across forwards does not move
+ * that row's precision boundary. This is why the transition runs here rather
+ * than at the end of each append, as the single-sequence legos run it.
+ *
+ * Compiled decode: a layout holding one row (the lone row a batch adopted)
+ * drives it through a compiled decode step as the lego of its transition
+ * (`KvRowTransition.lone`): the first step that prepares this layout moves the
+ * row's storage into that lego, which converts at the same position the rows
+ * convert (the append that reaches the start is read before the conversion), and
+ * every compiled step after reaches the row there. Its offsets, state and
+ * signature read through the lego; any other operation (an ordinary append, a
+ * merge, an extraction, a speculative round) first moves the storage back into
+ * the rows. A one-row layout reads each decode position with the same bits as
+ * its lego (`tests/state/compiled-decode-protocol.test.ts`). */
+export abstract class TransitioningKvRows<Layout extends TransitionedKvLayout> implements BatchableCache, CompiledDecodeCache {
+  readonly #positions: TransitioningKvPositions;
+  #rows: Cache[] = [];
+  #packed?: Layout;
+  /** The lone row as its lego, while a compiled decode step drives it. */
+  #lone: (Cache & CompiledDecodeCache) | null = null;
+  #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(); }
+  }
+  // Every access to the rows, the packed layout or the positions finds the row
+  // home first (see the class).
+  protected get rows(): Cache[] { this.#home(); return this.#rows; }
+  protected set rows(rows: Cache[]) { this.#home(); this.#rows = rows; }
+  protected get packed(): Layout | undefined { this.#home(); return this.#packed; }
+  protected set packed(packed: Layout | undefined) { this.#home(); this.#packed = packed; }
+  protected get positions(): TransitioningKvPositions { this.#home(); return this.#positions; }
+  /** The lone row's storage, moved back from its lego into the rows. */
+  #home(): void {
+    const lone = this.#lone;
+    if (!lone) return;
+    this.#lone = null;
+    const storage = lone.composed ? lone.composed.release() : lone;
+    if (storage !== lone) lone.dispose();
+    this.#rows = this.transition.prepareRows?.([storage]) ?? [storage];
+    this.#positions.sync(this.#rows, false);
+  }
+  /** The lone row as its lego, moving its storage there on the first call. */
+  #loneLego(): Cache & CompiledDecodeCache {
+    if (this.#lone) return this.#lone;
+    if (this.batchSize !== 1) throw new Error(`compiled decode drives a layout of one row; this one holds ${this.batchSize}`);
+    const packed = this.#packed, row = this.#rows[0]!;
+    const adopted = packed ? null : this.transition.adopt?.(row) ?? null;
+    const storage = adopted ?? this.extractStorageRow(0);
+    let lego: Cache & CompiledDecodeCache;
+    try { lego = this.transition.lone(storage); }
+    catch (error) { if (!adopted) storage.dispose(); throw error; }
+    this.#rows = []; this.#packed = undefined;
+    if (packed) packed.dispose();
+    else if (!adopted) row.dispose();
+    this.#lone = lego;
+    return lego;
+  }
+  /** The fetch of the coming compiled step: the lego's while it holds the row,
+   * else the rows' (`decodePhaseOfRows`), without moving the row. */
+  decodePhase(): DecodeStepPlan["fetch"] { return this.#lone ? this.#lone.decodePhase() : this.decodePhaseOfRows(); }
+  /** The fetch a lone row's lego would choose, from this layout's positions. */
+  protected decodePhaseOfRows(): DecodeStepPlan["fetch"] { return "concat"; }
+  prepareDecodeStep(): DecodeStepPlan { return this.#loneLego().prepareDecodeStep(); }
+  decodeSlot(plan: DecodeStepPlan): DecodeSlot { return this.#loneLego().decodeSlot(plan); }
+  decodeInputs(plan: DecodeStepPlan, step: DecodeStepInputs): MlxArray[] { return this.#loneLego().decodeInputs(plan, step); }
+  commitDecodeStep(slot: DecodeSlot, outputs: MlxArray[]): MlxArray[] { return this.#loneLego().commitDecodeStep(slot, outputs); }
+  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.settle(row); }
+      finally { this.rows[index] = row[0]!; }
+    }
+  }
+  endPrefill(): void { this.#prefilling = false; }
+  get minimumReusableOffset(): number {
+    if (this.#lone) return this.#lone.minimumReusableOffset ?? 0;
+    return this.packed?.minimumReusableOffset ?? minimumReusableOffset(this.rows);
+  }
+  get rowOffsets(): readonly number[] { return this.#lone ? [this.#lone.offset] : this.packed?.rowOffsets ?? this.positions.rowOffsets; }
+  get leftPad(): readonly number[] { return this.#lone ? [0] : this.packed?.leftPad ?? this.positions.leftPad; }
+  get offset(): number { return this.#lone ? this.#lone.offset : 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.#lone ? 1 : this.packed?.batchSize ?? this.positions.batchSize; }
+  get ropeOffsetArr(): MlxArray | undefined {
+    if (this.#lone) return this.#lone.ropeOffsetArr;
+    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.settle(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.extractStorageRow(index));
+        } else owned.push(...cloneKvCaches([row.composed ? row.composed.storage() : 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);
+  }
+  /** A row's storage, owned by the caller: what another row layout merges. */
+  extractStorageRow(row: number): Cache {
+    return this.packed?.extractRow(row) ?? (this.transition.extractRow?.(this.rows[row]!) ?? cloneKvCaches([this.rows[row]!])[0]!);
+  }
+  /** A row handed back to serial use: its storage, composed again by the lego
+   * the row came from (`KvRowTransition.compose`). */
+  extractRow(row: number): Cache {
+    const storage = this.extractStorageRow(row);
+    return this.transition.compose ? this.transition.compose(storage) : storage;
+  }
+  state(): MlxArray[] {
+    return withResource(leaseCacheStates(this.#lone ? [this.#lone] : 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 lone = this.#lone; this.#lone = null; lone?.dispose();
+    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(); }
+}

package/src/state/turboquant-attention.ts

@@ -0,0 +1,97 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { AttentionRead } from "../contracts/mlx/cache";
+import * as tq from "../kernels/turboquant/ops";
+import { sdpaRead, withLease, type MaskLease } from "./attention-read";
+
+/** One append's window decoded with its values still in the codec's rotated
+ * domain, and the transform that returns an attention output over those
+ * values to the values' own domain. The caller owns `keys` and `values`. */
+export interface RotatedWindow {
+  readonly keys: MlxArray;
+  readonly values: MlxArray;
+  /** Borrows an attention output [B, Hq, L, Dv] and returns an owned one. */
+  readonly restore: (output: MlxArray) => MlxArray;
+}
+
+/** The two ways TurboQuant storage can append one call's keys and values and
+ * decode its window. A read calls exactly one of them, once. */
+export interface TurboQuantFetches {
+  /** Append, then decode the window with keys and values in their own domain
+   * (`updateAndFetch`). The caller owns both. */
+  decoded(): [MlxArray, MlxArray];
+  /** Append, then decode the window with values left rotated
+   * (`updateAndFetchDeferredV`). */
+  rotated(): RotatedWindow;
+}
+
+/** How a TurboQuant cache attends, fixed when the cache is built, like the
+ * affine kernel sets. The cache leases the mask for the phase and offers its
+ * fetches; the set decides which fetch runs and what attends over it. Both sets
+ * below are a graph's served numbers today, so the graph that composes the
+ * cache names its set: `denseTurboQuantKernels` for Qwen 3.5's,
+ * `deferredTurboQuantKernels` for Gemma 4's and MiniCPM5's. */
+export interface TurboQuantKernels {
+  /** The fetch `read` calls: the window decoded to the values' own domain, or
+   * with values left rotated. A compiled decode step decodes the stored prefix
+   * in this domain outside its recorded graph (`TurboQuantKVCache`). */
+  readonly fetch: "decoded" | "rotated";
+  /** Appends through `fetches` and returns the read under `mask`. Takes
+   * ownership of `mask` (released if the fetch throws) and of what the fetch
+   * returns. */
+  read(fetches: TurboQuantFetches, mask: MaskLease): AttentionRead;
+}
+
+/** The read over fetched keys and values still in the codec's rotated domain:
+ * the stock fused SDPA under the leased mask, then `restore` on its output
+ * (the inverse value rotation, f32 through the transform, unless the storage
+ * says otherwise). Attention is linear in V, so the caller receives the output
+ * in the values' own domain. Takes ownership of `keys`, `values` and `mask`. */
+function rotatedValueRead(keys: MlxArray, values: MlxArray, mask: MaskLease,
+  restore: (output: MlxArray) => MlxArray = tq.unrotateValues): AttentionRead {
+  return {
+    attend(q, scale) {
+      const rotated = ops.sdpa(q, keys, values, scale, mask.mask.mode, mask.mask.arr);
+      try { return restore(rotated); } finally { rotated.dispose(); }
+    },
+    dispose() { keys.dispose(); values.dispose(); mask.release(); },
+  };
+}
+
+/** Decode keys and values back to their own domain, then the stock fused SDPA
+ * under the mask. The Qwen graphs (Qwen3, Qwen3-MoE, Qwen 3.5 and the written
+ * Trellis graphs) and the universal graph compose this set; it computes their
+ * served numbers. */
+export const denseTurboQuantKernels: TurboQuantKernels = Object.freeze({
+  fetch: "decoded" as const,
+  read: (fetches: TurboQuantFetches, mask: MaskLease) => withLease(mask, held => {
+    const [keys, values] = fetches.decoded();
+    return sdpaRead(keys, values, held);
+  }),
+});
+
+/** Leave values rotated, run the stock fused SDPA under the mask on them, then
+ * undo the rotation on the output: attention is linear in V, so one inverse
+ * rotation per query row replaces one per cached token. The Gemma 4 graphs
+ * (generic and written), DiffusionGemma and MiniCPM5 compose this set; it computes their served
+ * numbers. */
+export const deferredTurboQuantKernels: TurboQuantKernels = Object.freeze({
+  fetch: "rotated" as const,
+  read: (fetches: TurboQuantFetches, mask: MaskLease) => withLease(mask, held => {
+    const { keys, values, restore } = fetches.rotated();
+    return rotatedValueRead(keys, values, held, restore);
+  }),
+});
+
+/** The set of plain storage whose graph composes no TurboQuant read: every
+ * graph but Qwen 3.5's, Gemma 4's and MiniCPM5's reads converted TurboQuant
+ * storage through the deprecated fetches, which take no set. A TurboQuant cache
+ * converted from such storage serves those fetches; its named reads throw. */
+export const uncomposedTurboQuantKernels: TurboQuantKernels = Object.freeze({
+  fetch: "decoded" as const,
+  read(_fetches: TurboQuantFetches, mask: MaskLease): AttentionRead {
+    mask.release();
+    throw new Error("this TurboQuant cache was converted from storage whose graph composed no TurboQuant read: " +
+      "the graph's makeCache passes denseTurboQuantKernels or deferredTurboQuantKernels (state/turboquant-attention.ts) to its KVCache");
+  },
+});